v3-api/src/ai-v2/ai-v2.service.ts

import { HttpException, HttpStatus, Injectable } from '@nestjs/common';
import { InjectModel } from '@nestjs/mongoose';
import { Request, Response } from 'express';
import { Model, Types } from 'mongoose';
import { BaseResponseDTO } from 'src/common/dto/base-response.dto';
import { TimeHelper } from 'src/common/tools/time-helper';
import { ReactEnum } from 'src/common/types/react.type';
import { Sender } from 'src/common/types/sender.type';
import { ReactsModel } from 'src/database/model/botReact.model';
import { SessionModel } from 'src/database/model/sessions.model';
import {
  AI_UNAVAILABLE_FALLBACK,
  AskCompatiblePayload,
  buildBotMessage,
  buildUserMessage,
  isEscalationOffered,
  resolveAskText,
  resolveSessionId,
  toAskCompatiblePayload,
  wrapAskResponse,
} from './ai-v2-ask.mapper';
import { AiV2Client, AiV2RunResult } from './ai-v2.client';
import { AiV2Exception } from './ai-v2.exception';
import { CreateFeedbackDto, CreateRunDto } from './dto/threads.dto';

@Injectable()
export class AiV2Service {
  constructor(
    private readonly client: AiV2Client,
    @InjectModel(SessionModel.name)
    private readonly session: Model<SessionModel>,
    @InjectModel(ReactsModel.name)
    private readonly reacts: Model<ReactsModel>,
  ) {}

  identityUserId(user: any): string {
    const id = user?._id ?? user?.userData?._id ?? user?.id;
    if (!id) {
      throw new AiV2Exception(HttpStatus.UNAUTHORIZED, 'user_id_required');
    }
    return String(id);
  }

  private userObjectId(user: any): Types.ObjectId {
    return new Types.ObjectId(this.identityUserId(user));
  }

  async run(body: CreateRunDto, user: any, req: Request, res: Response) {
    console.log('run service called with body:', body);
    const question = resolveAskText(body);
    if (!question) {
      throw new AiV2Exception(HttpStatus.BAD_REQUEST, 'question_required');
    }

    const now = Date.now() / 1000;
    const sessionId = resolveSessionId(body.sessionId);
    const isNewSession = !sessionId;
    const session = isNewSession
      ? await this.createSession(question, user, now)
      : await this.appendUserTurn(sessionId, question, user, now);

    const stream = this.client.isStreamToFrontendEnabled();
    if (stream) {
      this.beginSse(res);
      this.writeSse(res, 'meta', {
        sessionId: String(session._id),
        isNewSession,
      });
    }

    let result: AiV2RunResult | null = null;
    let status: AskCompatiblePayload['status'] = 'answered';
    let answer = AI_UNAVAILABLE_FALLBACK;
    let runId: string | null = null;

    try {
      result = await this.client.consumeRun(
        String(session._id),
        { message: question, user_id: this.identityUserId(user) },
        {
          onToken: stream
            ? (text) => this.writeSse(res, 'token', { text })
            : undefined,
          onClientClose: (abort) => {
            res.on('close', () => {
              if (!res.writableEnded) abort();
            });
          },
        },
      );
      answer = result.message;
      status = isEscalationOffered(result)
        ? 'escalation_offered'
        : 'answered';
      runId = result.run_id;
    } catch (err) {
      if (
        err instanceof HttpException &&
        err.getStatus() < 500 &&
        err.getStatus() !== HttpStatus.BAD_GATEWAY
      ) {
        throw err;
      }
      status = 'ai_unavailable';
    }

    const sessionObjectId = new Types.ObjectId(String(session._id));
    const botMessage = buildBotMessage({
      text: answer,
      now,
      runId,
      status,
      escalation: result?.escalation,
    });
    await this.session.updateOne(
      { _id: sessionObjectId },
      { $push: { messages: botMessage } },
    );

    const updated = await this.session.findById(sessionObjectId).exec();
    if (!updated) {
      throw new HttpException('session_not_found', HttpStatus.NOT_FOUND);
    }
    const botCount = updated.messages.filter(
      (message) => message.sender === Sender.Bot,
    ).length;
    const sessionPatch: Record<string, unknown> = {};
    if (botCount >= 20) sessionPatch.chatClosed = true;
    if (status === 'escalation_offered') sessionPatch.escalationOffered = true;
    if (Object.keys(sessionPatch).length) {
      await this.session.updateOne(
        { _id: sessionObjectId },
        { $set: sessionPatch },
      );
      if (sessionPatch.chatClosed) (updated as any).chatClosed = true;
      if (sessionPatch.escalationOffered) {
        (updated as any).escalationOffered = true;
      }
    }

    const payload = toAskCompatiblePayload({
      session: updated,
      now,
      question,
      answer,
      messageId: String(botMessage.messageId),
      runId,
      status,
      escalation: result?.escalation,
      isNewSession,
    });

    if (stream) {
      this.writeSse(res, 'result', payload);
      res.end();
      return;
    }

    res.status(HttpStatus.OK).json(wrapAskResponse(payload));
  }

  async feedback(body: CreateFeedbackDto, user: any) {
    const sessionId = resolveSessionId(body.sessionId);
    if (!sessionId || !Types.ObjectId.isValid(sessionId)) {
      throw new AiV2Exception(HttpStatus.BAD_REQUEST, 'invalid_session_id');
    }
    if (!Types.ObjectId.isValid(body.messageId)) {
      throw new AiV2Exception(HttpStatus.BAD_REQUEST, 'invalid_message_id');
    }

    const sessionObjectId = new Types.ObjectId(sessionId);
    const messageObjectId = new Types.ObjectId(body.messageId);
    const userObjectId = this.userObjectId(user);
    const react = body.helpful ? ReactEnum.like : ReactEnum.dislike;

    const session = await this.session
      .findOne({
        _id: sessionObjectId,
        userId: userObjectId,
      })
      .exec();
    if (!session) {
      throw new HttpException('session_not_found', HttpStatus.NOT_FOUND);
    }

    const botMessageIndex = session.messages.findIndex(
      (message) => String(message.messageId) === String(body.messageId),
    );
    if (botMessageIndex < 0) {
      throw new HttpException('message_not_found', HttpStatus.NOT_FOUND);
    }

    const botMessage = session.messages[botMessageIndex];
    if (botMessage.runId) {
      await this.client.createFeedback(String(session._id), botMessage.runId, {
        helpful: body.helpful,
        user_id: this.identityUserId(user),
        comment: body.comment ?? null,
      });
    }

    await this.session.updateOne(
      {
        _id: sessionObjectId,
        'messages.messageId': messageObjectId,
      },
      {
        $set: { 'messages.$.react': react },
      },
    );

    let questionText: string | null = null;
    for (let i = botMessageIndex - 1; i >= 0; i--) {
      if (session.messages[i].sender === Sender.User) {
        questionText = session.messages[i].text;
        break;
      }
    }

    await this.reacts.findOneAndUpdate(
      {
        sessionId: sessionObjectId,
        messageId: messageObjectId,
        userId: userObjectId,
      },
      {
        $set: {
          question: questionText,
          answer: botMessage.text,
          sender: botMessage.sender,
          react,
          createdAt: TimeHelper.unix2PersianTimeAndDate(Date.now() / 1000),
          createdISO: new Date(),
        },
      },
      { upsert: true, new: true },
    );

    const updated = await this.session.findById(sessionObjectId).exec();
    if (!updated) {
      throw new HttpException('session_not_found', HttpStatus.NOT_FOUND);
    }
    return this.ok({
      sessionId: updated._id,
      messageId: String(body.messageId),
      runId: botMessage.runId || null,
      helpful: body.helpful,
      react,
      messages: updated.messages,
    });
  }

  ok<T>(data: T, status: HttpStatus = HttpStatus.OK, message = 'SUCCESS') {
    return new BaseResponseDTO(status, message, data);
  }

  created<T>(data: T) {
    return new BaseResponseDTO(HttpStatus.CREATED, 'SUCCESS', data);
  }

  private async createSession(question: string, user: any, now: number) {
    return this.session.create({
      _id: new Types.ObjectId(),
      userId: this.userObjectId(user),
      chatTitle: question,
      connectedToExpert: false,
      onlineChatClosed: false,
      chatClosed: false,
      escalationOffered: false,
      expert: '',
      expertRate: null,
      createdAt: TimeHelper.unix2PersianTimeAndDate(now),
      createdISO: Date.now(),
      messages: [buildUserMessage(question, now)],
    });
  }

  private async appendUserTurn(
    sessionId: string,
    question: string,
    user: any,
    now: number,
  ) {
    if (!Types.ObjectId.isValid(sessionId)) {
      throw new AiV2Exception(HttpStatus.BAD_REQUEST, 'invalid_session_id');
    }
    const session = await this.session
      .findOne({
        _id: new Types.ObjectId(sessionId),
        userId: this.userObjectId(user),
      })
      .exec();

    if (!session) {
      throw new HttpException('session_not_found', HttpStatus.NOT_FOUND);
    }
    if (session.chatClosed) {
      throw new HttpException('session_limit_reached', HttpStatus.BAD_REQUEST);
    }

    await this.session.updateOne(
      { _id: session._id },
      { $push: { messages: buildUserMessage(question, now) } },
    );
    return session;
  }

  private beginSse(res: Response) {
    res.status(200);
    res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
    res.setHeader('Cache-Control', 'no-cache, no-transform');
    res.setHeader('Connection', 'keep-alive');
    res.setHeader('X-Accel-Buffering', 'no');
    if (typeof (res as any).flushHeaders === 'function') {
      (res as any).flushHeaders();
    }
    res.socket?.setNoDelay?.(true);
  }

  private writeSse(res: Response, event: string, data: unknown) {
    res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`);
    if (typeof (res as any).flush === 'function') {
      (res as any).flush();
    }
  }
}