KnowledgeBase/server/src/chat/chat.service.ts

import { Injectable, Logger, Inject, forwardRef } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { ChatOpenAI } from '@langchain/openai';
import { PromptTemplate } from '@langchain/core/prompts';
import { ElasticsearchService } from '../elasticsearch/elasticsearch.service';
import { EmbeddingService } from '../knowledge-base/embedding.service';
import { ModelConfigService } from '../model-config/model-config.service';
import { KnowledgeGroupService } from '../knowledge-group/knowledge-group.service';
import { SearchHistoryService } from '../search-history/search-history.service';
import { ModelConfig, ModelType } from '../types';
import { RagService } from '../rag/rag.service';

import { DEFAULT_VECTOR_DIMENSIONS, DEFAULT_LANGUAGE } from '../common/constants';
import { I18nService } from '../i18n/i18n.service';
import { TenantService } from '../tenant/tenant.service';
import { UserSettingService } from '../user/user-setting.service';

export interface ChatMessage {
  role: 'user' | 'assistant';
  content: string;
}

@Injectable()
export class ChatService {
  private readonly logger = new Logger(ChatService.name);
  private readonly defaultDimensions: number;

  constructor(
    @Inject(forwardRef(() => ElasticsearchService))
    private elasticsearchService: ElasticsearchService,
    private embeddingService: EmbeddingService,
    private modelConfigService: ModelConfigService,
    @Inject(forwardRef(() => KnowledgeGroupService))
    private knowledgeGroupService: KnowledgeGroupService,
    private searchHistoryService: SearchHistoryService,
    private configService: ConfigService,
    private ragService: RagService,
    private i18nService: I18nService,
    private tenantService: TenantService,
    private userSettingService: UserSettingService,
  ) {
    this.defaultDimensions = parseInt(
      this.configService.get<string>('DEFAULT_VECTOR_DIMENSIONS', String(DEFAULT_VECTOR_DIMENSIONS)),
    );
  }

  async *streamChat(
    message: string,
    history: ChatMessage[],
    userId: string,
    modelConfig: ModelConfig,
    userLanguage: string = DEFAULT_LANGUAGE,
    selectedEmbeddingId?: string,
    selectedGroups?: string[], // New: Selected groups
    selectedFiles?: string[], // New: Selected files
    historyId?: string, // New: Chat history ID
    enableRerank: boolean = false,
    selectedRerankId?: string,
    temperature?: number, // New: temperature parameter
    maxTokens?: number, // New: maxTokens parameter
    topK?: number, // New: topK parameter
    similarityThreshold?: number, // New: similarityThreshold parameter
    rerankSimilarityThreshold?: number, // New: rerankSimilarityThreshold parameter
    enableQueryExpansion?: boolean, // New
    enableHyDE?: boolean, // New
    tenantId?: string // New: tenant isolation
  ): AsyncGenerator<{ type: 'content' | 'sources' | 'historyId'; data: any }> {
    console.log('=== ChatService.streamChat ===');
    console.log('User ID:', userId);
    console.log('User language:', userLanguage);
    console.log('Selected embedding model ID:', selectedEmbeddingId);
    console.log('Selected groups:', selectedGroups);
    console.log('Selected files:', selectedFiles);
    console.log('History ID:', historyId);
    console.log('Temperature:', temperature);
    console.log('Max Tokens:', maxTokens);
    console.log('Top K:', topK);
    console.log('Similarity threshold:', similarityThreshold);
    console.log('Rerank threshold:', rerankSimilarityThreshold);
    console.log('Query expansion:', enableQueryExpansion);
    console.log('HyDE:', enableHyDE);
    console.log('Model configuration:', {
      name: modelConfig.name,
      modelId: modelConfig.modelId,
      baseUrl: modelConfig.baseUrl,
    });
    console.log('API Key prefix:', modelConfig.apiKey?.substring(0, 10) + '...');
    console.log('API Key length:', modelConfig.apiKey?.length);

    // Get current language setting (keeping LANGUAGE_CONFIG for backward compatibility, now uses i18n service)
    // Use actual language based on user settings
    const effectiveUserLanguage = userLanguage || DEFAULT_LANGUAGE;

    let currentHistoryId = historyId;
    let fullResponse = '';

    try {
      // Create new chat history if no historyId
      if (!currentHistoryId) {
        const searchHistory = await this.searchHistoryService.create(
          userId,
          tenantId || 'default', // New
          message,
          selectedGroups,
        );
        currentHistoryId = searchHistory.id;
        console.log(this.i18nService.getMessage('creatingHistory', effectiveUserLanguage) + currentHistoryId);
        yield { type: 'historyId', data: currentHistoryId };
      }

      // Save user message
      await this.searchHistoryService.addMessage(currentHistoryId, 'user', message);
      // 1. Get user's embedding model settings
      let embeddingModel: any;

      if (selectedEmbeddingId) {
        // Find specifically selected model
        embeddingModel = await this.modelConfigService.findOne(selectedEmbeddingId, userId, tenantId || 'default');
      } else {
        // Use organization's default from Index Chat Config (strict)
        embeddingModel = await this.modelConfigService.findDefaultByType(tenantId || 'default', ModelType.EMBEDDING);
      }

      console.log(this.i18nService.getMessage('usingEmbeddingModel', effectiveUserLanguage) + embeddingModel.name + ' ' + embeddingModel.modelId + ' ID:' + embeddingModel.id);

      // 2. Search using user's query directly
      console.log(this.i18nService.getMessage('startingSearch', effectiveUserLanguage));
      yield { type: 'content', data: this.i18nService.getMessage('searching', effectiveUserLanguage) + '\n' };

      let searchResults: any[] = [];
      let context = '';

      try {
        // 3. If knowledge groups are selected, get file IDs from those groups first
        let effectiveFileIds = selectedFiles; // Prioritize explicitly specified files
        if (!effectiveFileIds && selectedGroups && selectedGroups.length > 0) {
          // Get file IDs from knowledge groups
          effectiveFileIds = await this.knowledgeGroupService.getFileIdsByGroups(selectedGroups, userId, tenantId as string);
        }

        // 3. Use RagService for search (supports hybrid search + Rerank)
        const ragResults = await this.ragService.searchKnowledge(
          message,
          userId,
          topK,
          similarityThreshold,
          embeddingModel.id,
          true, // enableFullTextSearch (Chat defaults to hybrid)
          enableRerank,
          selectedRerankId,
          undefined, // selectedGroups
          effectiveFileIds,
          rerankSimilarityThreshold,
          tenantId,
          enableQueryExpansion,
          enableHyDE
        );

        // Convert RagSearchResult to format needed by ChatService (any[])
        // HybridSearch returns ES hit structure, but RagSearchResult is normalized
        // BuildContext expects {fileName, content}. RagSearchResult has these
        searchResults = ragResults;
        console.log(this.i18nService.getMessage('searchResultsCount', effectiveUserLanguage) + searchResults.length);

        // 4. Build context
        context = this.buildContext(searchResults, effectiveUserLanguage);

        if (searchResults.length === 0) {
          if (selectedGroups && selectedGroups.length > 0) {
            // User selected knowledge groups but no matches found
            const noMatchMsg = this.i18nService.getMessage('noMatchInKnowledgeGroup', effectiveUserLanguage);
            yield { type: 'content', data: `⚠️ ${noMatchMsg}\n\n` };
          } else {
            yield { type: 'content', data: this.i18nService.getMessage('noResults', effectiveUserLanguage) + '\n\n' };
          }
          yield { type: 'content', data: `[Debug] ${this.i18nService.getMessage('searchScope', effectiveUserLanguage)}: ${selectedFiles ? selectedFiles.length + ' ' + this.i18nService.getMessage('files', effectiveUserLanguage) : selectedGroups ? selectedGroups.length + ' ' + this.i18nService.getMessage('notebooks', effectiveUserLanguage) : this.i18nService.getMessage('all', effectiveUserLanguage)}\n` };
          yield { type: 'content', data: `[Debug] ${this.i18nService.getMessage('searchResults', effectiveUserLanguage)}: 0 ${this.i18nService.getMessage('items', effectiveUserLanguage)}\n` };
        } else {
          yield {
            type: 'content',
            data: `${searchResults.length} ${this.i18nService.getMessage('relevantInfoFound', effectiveUserLanguage)}。${this.i18nService.getMessage('generatingResponse', effectiveUserLanguage)}...\n\n`,
          };
          // Debug info
          const scores = searchResults.map(r => {
            if (r.originalScore !== undefined && r.originalScore !== r.score) {
              return `${r.originalScore.toFixed(2)} → ${r.score.toFixed(2)}`;
            }
            return r.score.toFixed(2);
          }).join(', ');
          const files = [...new Set(searchResults.map(r => r.fileName))].join(', ');
          yield { type: 'content', data: `> [Debug] ${this.i18nService.getMessage('searchHits', effectiveUserLanguage)}: ${searchResults.length} ${this.i18nService.getMessage('items', effectiveUserLanguage)}\n` };
          yield { type: 'content', data: `> [Debug] ${this.i18nService.getMessage('relevance', effectiveUserLanguage)}: ${scores}\n` };
          yield { type: 'content', data: `> [Debug] ${this.i18nService.getMessage('sourceFiles', effectiveUserLanguage)}: ${files}\n\n---\n\n` };
        }
      } catch (searchError) {
        console.error(this.i18nService.getMessage('searchFailedLog', effectiveUserLanguage) + ':', searchError);
        yield { type: 'content', data: this.i18nService.getMessage('searchFailed', effectiveUserLanguage) + '\n\n' };
      }

      // 5. Stream response generation
      this.logger.log(this.i18nService.formatMessage('modelCall', {
        type: 'LLM',
        model: `${modelConfig.name} (${modelConfig.modelId})`,
        user: userId
      }, effectiveUserLanguage));
      const llm = new ChatOpenAI({
        apiKey: modelConfig.apiKey || 'ollama',
        streaming: true,
        temperature: temperature !== undefined ? temperature : 0.3,
        maxTokens: maxTokens !== undefined ? maxTokens : undefined,
        modelName: modelConfig.modelId,
        configuration: {
          baseURL: modelConfig.baseUrl || 'http://localhost:11434/v1',
        },
      });

      const promptTemplate =
        context.length > 0
          ? this.i18nService.getPrompt(
            effectiveUserLanguage,
            'withContext',
            selectedGroups && selectedGroups.length > 0
          )
          : this.i18nService.getPrompt(effectiveUserLanguage, 'withoutContext');

      const prompt = PromptTemplate.fromTemplate(promptTemplate);

      const chain = prompt.pipe(llm);

      const stream = await chain.stream({
        context,
        history: this.formatHistory(history, userLanguage),
        question: message,
      });

      for await (const chunk of stream) {
        if (chunk.content) {
          fullResponse += chunk.content;
          yield { type: 'content', data: chunk.content };
        }
      }

      // Save AI response
      await this.searchHistoryService.addMessage(
        currentHistoryId,
        'assistant',
        fullResponse,
        searchResults.map((result) => ({
          fileName: result.fileName,
          title: result.metadata?.title || result.metadata?.originalName, // ES metadata contains these
          content: String(result.content).substring(0, 200) + '...',
          score: result.score,
          chunkIndex: result.chunkIndex,
          fileId: result.fileId,
        })),
      );

      // 7. Auto-generate chat title (executed after first exchange)
      const messagesInHistory = await this.searchHistoryService.findOne(currentHistoryId, userId, tenantId);
      if (messagesInHistory.messages.length === 2) {
        this.generateChatTitle(currentHistoryId, userId, tenantId).catch((err) => {
          this.logger.error(`Failed to generate chat title for ${currentHistoryId}`, err);
        });
      }

      // 6. Return sources
      yield {
        type: 'sources',
        data: searchResults.map((result) => ({
          fileName: result.fileName,
          content: String(result.content).substring(0, 200) + '...',
          score: result.score,
          chunkIndex: result.chunkIndex,
          fileId: result.fileId,
        })),
      };
    } catch (error) {
      this.logger.error(this.i18nService.getMessage('chatStreamError', effectiveUserLanguage), error);
      yield { type: 'content', data: `${this.i18nService.getMessage('error', effectiveUserLanguage)}: ${error.message}` };
    }
  }

  async *streamAssist(
    instruction: string,
    context: string,
    modelConfig: ModelConfig,
    userLanguage: string = DEFAULT_LANGUAGE,
  ): AsyncGenerator<{ type: 'content'; data: any }> {
    try {
      this.logger.log(this.i18nService.formatMessage('modelCall', {
        type: 'LLM (Assist)',
        model: `${modelConfig.name} (${modelConfig.modelId})`,
        user: 'N/A'
      }, userLanguage));
      const llm = new ChatOpenAI({
        apiKey: modelConfig.apiKey || 'ollama',
        streaming: true,
        temperature: 0.7,
        modelName: modelConfig.modelId,
        configuration: {
          baseURL: modelConfig.baseUrl || 'http://localhost:11434/v1',
        },
      });

      const systemPrompt = `${this.i18nService.getMessage('intelligentAssistant', userLanguage)}
${this.i18nService.getMessage('assistSystemPrompt', userLanguage)}

${this.i18nService.getMessage('contextLabel', userLanguage)}:
${context}

${this.i18nService.getMessage('userInstructionLabel', userLanguage)}:
${instruction}`;

      const stream = await llm.stream(systemPrompt);

      for await (const chunk of stream) {
        if (chunk.content) {
          yield { type: 'content', data: chunk.content };
        }
      }
    } catch (error) {
      this.logger.error(this.i18nService.getMessage('assistStreamError', userLanguage), error);
      yield { type: 'content', data: `${this.i18nService.getMessage('error', userLanguage)}: ${error.message}` };
    }
  }

  private async hybridSearch(
    keywords: string[],
    userId: string,
    embeddingModelId?: string,
    selectedGroups?: string[], // New parameter
    explicitFileIds?: string[], // New parameter
    tenantId?: string, // Added
    userLanguage: string = DEFAULT_LANGUAGE,
  ): Promise<any[]> {
    try {
      // Join keywords into search string
      const combinedQuery = keywords.join(' ');
      console.log(this.i18nService.getMessage('searchString', userLanguage) + combinedQuery);

      // Check if embedding model ID is provided
      if (!embeddingModelId) {
        console.log(this.i18nService.getMessage('embeddingModelIdNotProvided', userLanguage));
        return [];
      }

      // Use actual embedding vector
      console.log(this.i18nService.getMessage('generatingEmbeddings', userLanguage));
      const queryEmbedding = await this.embeddingService.getEmbeddings(
        [combinedQuery],
        userId,
        embeddingModelId,
      );
      const queryVector = queryEmbedding[0];
      console.log(this.i18nService.getMessage('embeddingsGenerated', userLanguage) + this.i18nService.getMessage('dimensions', userLanguage) + ':', queryVector.length);

      // Hybrid search
      console.log(this.i18nService.getMessage('performingHybridSearch', userLanguage));
      const results = await this.elasticsearchService.hybridSearch(
        queryVector,
        combinedQuery,
        userId,
        10,
        0.6,
        selectedGroups, // Pass selected groups
        explicitFileIds, // Pass explicit file IDs
        tenantId, // Pass tenant ID
      );
      console.log(this.i18nService.getMessage('esSearchCompleted', userLanguage) + this.i18nService.getMessage('resultsCount', userLanguage) + ':', results.length);

      return results.slice(0, 10);
    } catch (error) {
      console.error(this.i18nService.getMessage('hybridSearchFailed', userLanguage) + ':', error);
      return [];
    }
  }

  private buildContext(results: any[], language: string = 'ja'): string {
    return results
      .map(
        (result, index) =>
          `[${index + 1}] ${this.i18nService.getMessage('file', language)}：${result.fileName}\n${this.i18nService.getMessage('content', language)}：${result.content}\n`,
      )
      .join('\n');
  }

  private formatHistory(
    history: ChatMessage[],
    userLanguage: string = 'ja',
  ): string {
    const userLabel = this.i18nService.getMessage('userLabel', userLanguage);
    const assistantLabel = this.i18nService.getMessage('assistantLabel', userLanguage);

    return history
      .slice(-6)
      .map(
        (msg) =>
          `${msg.role === 'user' ? userLabel : assistantLabel}：${msg.content}`,
      )
      .join('\n');
  }
  async getContextForTopic(topic: string, userId: string, tenantId?: string, groupId?: string, fileIds?: string[], userLanguage: string = DEFAULT_LANGUAGE): Promise<string> {
    try {
      // Use organization's default embedding from Index Chat Config (strict)
      const embeddingModel = await this.modelConfigService.findDefaultByType(tenantId || 'default', ModelType.EMBEDDING);

      const results = await this.hybridSearch(
        [topic],
        userId,
        embeddingModel.id,
        groupId ? [groupId] : undefined,
        fileIds,
        tenantId,
        userLanguage,
      );

      return this.buildContext(results);
    } catch (err) {
      this.logger.error(`${this.i18nService.getMessage('getContextForTopicFailed', userLanguage)}: ${err.message}`);
      return '';
    }
  }

  async generateSimpleChat(
    messages: ChatMessage[],
    userId: string,
    tenantId?: string,
    modelConfig?: ModelConfig, // Optional, looks up if not provided
    userLanguage: string = DEFAULT_LANGUAGE,
  ): Promise<string> {
    try {
      let config = modelConfig;
      if (!config) {
        // Use organization's default LLM from Index Chat Config (strict)
        const found = await this.modelConfigService.findDefaultByType(tenantId || 'default', ModelType.LLM);
        config = found as unknown as ModelConfig;
      }

      this.logger.log(this.i18nService.formatMessage('modelCall', {
        type: 'LLM (Simple)',
        model: `${config.name} (${config.modelId})`,
        user: userId
      }, 'ja'));
      const settings = await this.tenantService.getSettings(tenantId || 'default');
      const llm = new ChatOpenAI({
        apiKey: config.apiKey || 'ollama',
        temperature: settings?.temperature ?? 0.7,
        modelName: config.modelId,
        configuration: {
          baseURL: config.baseUrl || 'http://localhost:11434/v1',
        },
      });

      const response = await llm.invoke(
        messages.map(m => [m.role, m.content])
      );

      return String(response.content);
    } catch (error) {
      this.logger.error(this.i18nService.getMessage('simpleChatGenerationError', userLanguage), error);
      throw error;
    }
  }

  /**
   * Automatically generate chat title based on conversation content
   */
  async generateChatTitle(historyId: string, userId: string, tenantId?: string): Promise<string | null> {
    this.logger.log(`Generating automatic title for chat session ${historyId}`);

    try {
      const history = await this.searchHistoryService.findOne(historyId, userId, tenantId || 'default');
      if (!history || history.messages.length < 2) {
        return null;
      }

      const userMessage = history.messages.find(m => m.role === 'user')?.content || '';
      const aiResponse = history.messages.find(m => m.role === 'assistant')?.content || '';

      if (!userMessage || !aiResponse) {
        return null;
      }

      // Get language from user settings
      const userSettings = await this.userSettingService.getByUser(userId);
      const language = userSettings?.language || 'ja';

      // Build prompt
      const prompt = this.i18nService.getChatTitlePrompt(language, userMessage, aiResponse);

      // Call LLM to generate title
      const generatedTitle = await this.generateSimpleChat(
        [{ role: 'user', content: prompt }],
        userId,
        tenantId || 'default'
      );

      if (generatedTitle && generatedTitle.trim().length > 0) {
        // Remove extra quotes
        const cleanedTitle = generatedTitle.trim().replace(/^["']|["']$/g, '').substring(0, 50);
        await this.searchHistoryService.updateTitle(historyId, cleanedTitle);
        this.logger.log(`Successfully generated title for chat ${historyId}: ${cleanedTitle}`);
        return cleanedTitle;
      }
    } catch (error) {
      this.logger.error(`Failed to generate chat title for ${historyId}`, error);
    }

    return null;
  }
}