Tools/Notes
1
0
Fork 0
mirror of https://github.com/TriliumNext/Notes.git synced 2025-09-23 21:17:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

break up the rest_chat_service

Browse Source
This commit is contained in:
perf3ct 2025-04-16 19:35:09 +00:00
parent 77e637384d
commit 534396bce5
No known key found for this signature in database
GPG Key ID: 569C4EEC436F5232
10 changed files with 1664 additions and 1845 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

168
src/services/llm/chat/handlers/context-handler.ts Normal file
View File

@ -0,0 +1,168 @@
/**
* Handler for LLM context management
*/
import log from "../../../log.js";
import becca from "../../../../becca/becca.js";
import vectorStore from "../../embeddings/index.js";
import providerManager from "../../providers/providers.js";
import contextService from "../../context/services/context_service.js";
import type { NoteSource } from "../interfaces/session.js";
import { SEARCH_CONSTANTS } from '../../constants/search_constants.js';
/**
* Handles context management for LLM chat
*/
export class ContextHandler {
/**
* Find relevant notes based on search query
* @param content The search content
* @param contextNoteId Optional note ID for context
* @param limit Maximum number of results to return
* @returns Array of relevant note sources
*/
static async findRelevantNotes(content: string, contextNoteId: string | null = null, limit = 5): Promise<NoteSource[]> {
try {
// If content is too short, don't bother
if (content.length < 3) {
return [];
}
// Check if embeddings are available
const enabledProviders = await providerManager.getEnabledEmbeddingProviders();
if (enabledProviders.length === 0) {
log.info("No embedding providers available, can't find relevant notes");
return [];
}
// Get the embedding for the query
const provider = enabledProviders[0];
const embedding = await provider.generateEmbeddings(content);
let results;
if (contextNoteId) {
// For branch context, get notes specifically from that branch
const contextNote = becca.notes[contextNoteId];
if (!contextNote) {
return [];
}
const sql = require("../../../../services/sql.js").default;
const childBranches = await sql.getRows(`
SELECT branches.* FROM branches
WHERE branches.parentNoteId = ?
AND branches.isDeleted = 0
`, [contextNoteId]);
const childNoteIds = childBranches.map((branch: any) => branch.noteId);
// Include the context note itself
childNoteIds.push(contextNoteId);
// Find similar notes in this context
results = [];
for (const noteId of childNoteIds) {
const noteEmbedding = await vectorStore.getEmbeddingForNote(
noteId,
provider.name,
provider.getConfig().model
);
if (noteEmbedding) {
const similarity = vectorStore.cosineSimilarity(
embedding,
noteEmbedding.embedding
);
if (similarity > SEARCH_CONSTANTS.VECTOR_SEARCH.EXACT_MATCH_THRESHOLD) {
results.push({
noteId,
similarity
});
}
}
}
// Sort by similarity
results.sort((a, b) => b.similarity - a.similarity);
results = results.slice(0, limit);
} else {
// General search across all notes
results = await vectorStore.findSimilarNotes(
embedding,
provider.name,
provider.getConfig().model,
limit
);
}
// Format the results
const sources: NoteSource[] = [];
for (const result of results) {
const note = becca.notes[result.noteId];
if (!note) continue;
let noteContent: string | undefined = undefined;
if (note.type === 'text') {
const content = note.getContent();
// Handle both string and Buffer types
noteContent = typeof content === 'string' ? content :
content instanceof Buffer ? content.toString('utf8') : undefined;
}
sources.push({
noteId: result.noteId,
title: note.title,
content: noteContent,
similarity: result.similarity,
branchId: note.getBranches()[0]?.branchId
});
}
return sources;
} catch (error: any) {
log.error(`Error finding relevant notes: ${error.message}`);
return [];
}
}
/**
* Process enhanced context using the context service
* @param query Query to process
* @param contextNoteId Optional note ID for context
* @param showThinking Whether to show thinking process
*/
static async processEnhancedContext(query: string, llmService: any, options: {
contextNoteId?: string,
showThinking?: boolean
}) {
// Use the Trilium-specific approach
const contextNoteId = options.contextNoteId || null;
const showThinking = options.showThinking || false;
// Log that we're calling contextService with the parameters
log.info(`Using enhanced context with: noteId=${contextNoteId}, showThinking=${showThinking}`);
// Call context service for processing
const results = await contextService.processQuery(
query,
llmService,
{
contextNoteId,
showThinking
}
);
// Return the generated context and sources
return {
context: results.context,
sources: results.sources.map(source => ({
noteId: source.noteId,
title: source.title,
content: source.content || undefined, // Convert null to undefined
similarity: source.similarity
}))
};
}
}

368
src/services/llm/chat/handlers/stream-handler.ts Normal file
View File

@ -0,0 +1,368 @@
/**
* Handler for streaming LLM responses
*/
import log from "../../../log.js";
import type { Response } from "express";
import type { StreamChunk } from "../../ai_interface.js";
import type { LLMStreamMessage } from "../interfaces/ws-messages.js";
import type { ChatSession } from "../interfaces/session.js";
/**
* Handles streaming of LLM responses via WebSocket
*/
export class StreamHandler {
/**
* Handle streaming response via WebSocket
*
* This method processes LLM responses and sends them incrementally via WebSocket
* to the client, supporting both text content and tool execution status updates.
*
* @param res Express response object
* @param aiMessages Messages to send to the LLM
* @param chatOptions Options for the chat completion
* @param service LLM service to use
* @param session Chat session for storing the response
*/
static async handleStreamingResponse(
res: Response,
aiMessages: any[],
chatOptions: any,
service: any,
session: ChatSession
): Promise<void> {
// The client receives a success response for their HTTP request,
// but the actual content will be streamed via WebSocket
res.json({ success: true, message: 'Streaming response started' });
// Import the WebSocket service
const wsService = (await import('../../../ws.js')).default;
let messageContent = '';
const sessionId = session.id;
// Immediately send an initial message to confirm WebSocket connection is working
// This helps prevent timeouts on the client side
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
thinking: 'Preparing response...'
} as LLMStreamMessage);
try {
// Import the tool handler
const { ToolHandler } = await import('./tool-handler.js');
// Generate the LLM completion with streaming enabled
const response = await service.generateChatCompletion(aiMessages, {
...chatOptions,
stream: true
});
// If the model doesn't support streaming via .stream() method or returns tool calls,
// we'll handle it specially
if (response.tool_calls && response.tool_calls.length > 0) {
// Send thinking state notification via WebSocket
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
thinking: 'Analyzing tools needed for this request...'
} as LLMStreamMessage);
try {
// Execute the tools
const toolResults = await ToolHandler.executeToolCalls(response, sessionId);
// For each tool execution, send progress update via WebSocket
for (const toolResult of toolResults) {
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
toolExecution: {
action: 'complete',
tool: toolResult.name,
result: toolResult.content.substring(0, 100) + (toolResult.content.length > 100 ? '...' : '')
}
} as LLMStreamMessage);
}
// Make follow-up request with tool results
const toolMessages = [...aiMessages, {
role: 'assistant',
content: response.text || '',
tool_calls: response.tool_calls
}, ...toolResults];
// Preserve streaming for follow-up if it was enabled in the original request
const followUpOptions = {
...chatOptions,
// Only disable streaming if it wasn't explicitly requested
stream: chatOptions.stream === true,
// Allow tools but track iterations to prevent infinite loops
enableTools: true,
maxToolIterations: chatOptions.maxToolIterations || 5,
currentToolIteration: 1 // Start counting tool iterations
};
const followUpResponse = await service.generateChatCompletion(toolMessages, followUpOptions);
await this.processStreamedResponse(
followUpResponse,
wsService,
sessionId,
session,
toolMessages,
followUpOptions,
service
);
} catch (toolError: any) {
log.error(`Error executing tools: ${toolError.message}`);
// Send error via WebSocket with done flag
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
error: `Error executing tools: ${toolError instanceof Error ? toolError.message : 'Unknown error'}`,
done: true
} as LLMStreamMessage);
}
} else if (response.stream) {
// Handle standard streaming through the stream() method
log.info(`Provider ${service.getName ? service.getName() : 'unknown'} supports streaming via stream() method`);
// Store information about the model and provider in session metadata
session.metadata.model = response.model || session.metadata.model;
session.metadata.provider = response.provider || session.metadata.provider;
session.metadata.lastUpdated = new Date().toISOString();
await this.processStreamedResponse(
response,
wsService,
sessionId,
session
);
} else {
log.info(`Provider ${service.getName ? service.getName() : 'unknown'} does not support streaming via stream() method, falling back to single response`);
// If streaming isn't available, send the entire response at once
messageContent = response.text || '';
// Send via WebSocket - include both content and done flag in same message
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
content: messageContent,
done: true
} as LLMStreamMessage);
log.info(`Complete response sent`);
// Store the full response in the session
session.messages.push({
role: 'assistant',
content: messageContent,
timestamp: new Date()
});
}
} catch (streamingError: any) {
log.error(`Streaming error: ${streamingError.message}`);
// Import the WebSocket service directly in case it wasn't imported earlier
const wsService = (await import('../../../ws.js')).default;
// Send error via WebSocket
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
error: `Error generating response: ${streamingError instanceof Error ? streamingError.message : 'Unknown error'}`
} as LLMStreamMessage);
// Signal completion
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
done: true
} as LLMStreamMessage);
}
}
/**
* Process a streamed response from an LLM
*/
private static async processStreamedResponse(
response: any,
wsService: any,
sessionId: string,
session: ChatSession,
toolMessages?: any[],
followUpOptions?: any,
service?: any
): Promise<void> {
// Import tool handler lazily to avoid circular dependencies
const { ToolHandler } = await import('./tool-handler.js');
let messageContent = '';
try {
await response.stream(async (chunk: StreamChunk) => {
if (chunk.text) {
messageContent += chunk.text;
// Enhanced logging for each chunk
log.info(`Received stream chunk with ${chunk.text.length} chars of text, done=${!!chunk.done}`);
// Send each individual chunk via WebSocket as it arrives
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
content: chunk.text,
done: !!chunk.done, // Include done flag with each chunk
// Include any raw data from the provider that might contain thinking/tool info
...(chunk.raw ? { raw: chunk.raw } : {})
} as LLMStreamMessage);
// Log the first chunk (useful for debugging)
if (messageContent.length === chunk.text.length) {
log.info(`First stream chunk received: "${chunk.text.substring(0, 50)}${chunk.text.length > 50 ? '...' : ''}"`);
}
}
// If the provider indicates this is "thinking" state, relay that
if (chunk.raw?.thinking) {
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
thinking: chunk.raw.thinking
} as LLMStreamMessage);
}
// If the provider indicates tool execution, relay that
if (chunk.raw?.toolExecution) {
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
toolExecution: chunk.raw.toolExecution
} as LLMStreamMessage);
}
// Handle direct tool_calls in the response (for OpenAI)
if (chunk.tool_calls && chunk.tool_calls.length > 0) {
log.info(`Detected direct tool_calls in stream chunk: ${chunk.tool_calls.length} tools`);
// Send tool execution notification
wsService.sendMessageToAllClients({
type: 'tool_execution_start',
sessionId
} as LLMStreamMessage);
// Process each tool call
for (const toolCall of chunk.tool_calls) {
// Process arguments
let args = toolCall.function?.arguments;
if (typeof args === 'string') {
try {
args = JSON.parse(args);
} catch (e) {
log.info(`Could not parse tool arguments as JSON: ${e}`);
args = { raw: args };
}
}
// Format into a standardized tool execution message
wsService.sendMessageToAllClients({
type: 'tool_result',
sessionId,
toolExecution: {
action: 'executing',
tool: toolCall.function?.name || 'unknown',
toolCallId: toolCall.id,
args: args
}
} as LLMStreamMessage);
}
}
// Signal completion when done
if (chunk.done) {
log.info(`Stream completed, total content: ${messageContent.length} chars`);
// Check if there are more tool calls to execute (recursive tool calling)
if (service && toolMessages && followUpOptions &&
response.tool_calls && response.tool_calls.length > 0 &&
followUpOptions.currentToolIteration < followUpOptions.maxToolIterations) {
log.info(`Found ${response.tool_calls.length} more tool calls in iteration ${followUpOptions.currentToolIteration}`);
// Execute these tool calls in another iteration
const assistantMessage = {
role: 'assistant' as const,
content: messageContent,
tool_calls: response.tool_calls
};
// Execute the next round of tools
const nextToolResults = await ToolHandler.executeToolCalls(response, sessionId);
// Create a new messages array with the latest tool results
const nextToolMessages = [...toolMessages, assistantMessage, ...nextToolResults];
// Increment the tool iteration counter for the next call
const nextFollowUpOptions = {
...followUpOptions,
currentToolIteration: followUpOptions.currentToolIteration + 1
};
log.info(`Making another follow-up request (iteration ${nextFollowUpOptions.currentToolIteration}/${nextFollowUpOptions.maxToolIterations})`);
// Make another follow-up request
const nextResponse = await service.generateChatCompletion(nextToolMessages, nextFollowUpOptions);
// Process the next response recursively
await this.processStreamedResponse(
nextResponse,
wsService,
sessionId,
session,
nextToolMessages,
nextFollowUpOptions,
service
);
} else {
// Only send final done message if it wasn't already sent with content
// This ensures we don't duplicate the content but still mark completion
if (!chunk.text) {
log.info(`No content in final chunk, sending explicit completion message`);
// Send final message with done flag only (no content)
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
done: true
} as LLMStreamMessage);
}
// Store the full response in the session
session.messages.push({
role: 'assistant',
content: messageContent,
timestamp: new Date()
});
}
}
});
log.info(`Streaming completed successfully`);
} catch (streamError: any) {
log.error(`Error during streaming: ${streamError.message}`);
// Report the error to the client
wsService.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
error: `Error during streaming: ${streamError instanceof Error ? streamError.message : 'Unknown error'}`,
done: true
} as LLMStreamMessage);
throw streamError;
}
}
}

181
src/services/llm/chat/handlers/tool-handler.ts Normal file
View File

@ -0,0 +1,181 @@
/**
* Handler for LLM tool executions
*/
import log from "../../../log.js";
import type { Message } from "../../ai_interface.js";
import SessionsStore from "../sessions-store.js";
/**
* Handles the execution of LLM tools
*/
export class ToolHandler {
/**
* Execute tool calls from the LLM response
* @param response The LLM response containing tool calls
* @param sessionId Optional session ID for tracking
*/
static async executeToolCalls(response: any, sessionId?: string): Promise<Message[]> {
log.info(`========== TOOL EXECUTION FLOW ==========`);
if (!response.tool_calls || response.tool_calls.length === 0) {
log.info(`No tool calls to execute, returning early`);
return [];
}
log.info(`Executing ${response.tool_calls.length} tool calls`);
try {
// Import tool registry directly to avoid circular dependencies
const toolRegistry = (await import('../../tools/tool_registry.js')).default;
// Check if tools are available
const availableTools = toolRegistry.getAllTools();
log.info(`Available tools in registry: ${availableTools.length}`);
if (availableTools.length === 0) {
log.error('No tools available in registry for execution');
// Try to initialize tools
try {
// Ensure tools are initialized
const initResult = await this.ensureToolsInitialized();
if (!initResult) {
throw new Error('Failed to initialize tools');
}
} catch (error: unknown) {
const errorMessage = error instanceof Error ? error.message : String(error);
log.error(`Failed to initialize tools: ${errorMessage}`);
throw new Error('Tool execution failed: No tools available');
}
}
// Execute each tool call and collect results
const toolResults = await Promise.all(response.tool_calls.map(async (toolCall: any) => {
try {
log.info(`Executing tool: ${toolCall.function.name}, ID: ${toolCall.id || 'unknown'}`);
// Get the tool from registry
const tool = toolRegistry.getTool(toolCall.function.name);
if (!tool) {
throw new Error(`Tool not found: ${toolCall.function.name}`);
}
// Parse arguments
let args;
if (typeof toolCall.function.arguments === 'string') {
try {
args = JSON.parse(toolCall.function.arguments);
} catch (e: unknown) {
log.error(`Failed to parse tool arguments: ${e instanceof Error ? e.message : String(e)}`);
// Try cleanup and retry
try {
const cleaned = toolCall.function.arguments
.replace(/^['"]|['"]$/g, '') // Remove surrounding quotes
.replace(/\\"/g, '"') // Replace escaped quotes
.replace(/([{,])\s*'([^']+)'\s*:/g, '$1"$2":') // Replace single quotes around property names
.replace(/([{,])\s*(\w+)\s*:/g, '$1"$2":'); // Add quotes around unquoted property names
args = JSON.parse(cleaned);
} catch (cleanErr) {
// If all parsing fails, use as-is
args = { text: toolCall.function.arguments };
}
}
} else {
args = toolCall.function.arguments;
}
// Log what we're about to execute
log.info(`Executing tool with arguments: ${JSON.stringify(args)}`);
// Execute the tool and get result
const startTime = Date.now();
const result = await tool.execute(args);
const executionTime = Date.now() - startTime;
log.info(`Tool execution completed in ${executionTime}ms`);
// Log the result
const resultPreview = typeof result === 'string'
? result.substring(0, 100) + (result.length > 100 ? '...' : '')
: JSON.stringify(result).substring(0, 100) + '...';
log.info(`Tool result: ${resultPreview}`);
// Record tool execution in session if session ID is provided
if (sessionId) {
SessionsStore.recordToolExecution(sessionId, toolCall, typeof result === 'string' ? result : JSON.stringify(result));
}
// Format result as a proper message
return {
role: 'tool',
content: typeof result === 'string' ? result : JSON.stringify(result),
name: toolCall.function.name,
tool_call_id: toolCall.id || `tool-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`
};
} catch (error: any) {
log.error(`Error executing tool ${toolCall.function.name}: ${error.message}`);
// Record error in session if session ID is provided
if (sessionId) {
SessionsStore.recordToolExecution(sessionId, toolCall, '', error.message);
}
// Return error as tool result
return {
role: 'tool',
content: `Error: ${error.message}`,
name: toolCall.function.name,
tool_call_id: toolCall.id || `tool-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`
};
}
}));
log.info(`Completed execution of ${toolResults.length} tools`);
return toolResults;
} catch (error: any) {
log.error(`Error in tool execution handler: ${error.message}`);
throw error;
}
}
/**
* Ensure LLM tools are initialized
*/
static async ensureToolsInitialized(): Promise<boolean> {
try {
log.info("Checking LLM tool initialization...");
// Import tool registry
const toolRegistry = (await import('../../tools/tool_registry.js')).default;
// Check if tools are already initialized
const registeredTools = toolRegistry.getAllTools();
if (registeredTools.length === 0) {
log.info("No tools found in registry.");
log.info("Note: Tools should be initialized in the AIServiceManager constructor.");
// Create AI service manager instance to trigger tool initialization
const aiServiceManager = (await import('../../ai_service_manager.js')).default;
aiServiceManager.getInstance();
// Check again after AIServiceManager instantiation
const tools = toolRegistry.getAllTools();
log.info(`After AIServiceManager instantiation: ${tools.length} tools available`);
} else {
log.info(`LLM tools already initialized: ${registeredTools.length} tools available`);
}
// Get all available tools for logging
const availableTools = toolRegistry.getAllTools().map(t => t.definition.function.name);
log.info(`Available tools: ${availableTools.join(', ')}`);
log.info("LLM tools initialized successfully: " + availableTools.length + " tools available");
return true;
} catch (error) {
log.error(`Failed to initialize LLM tools: ${error}`);
return false;
}
}
}

29
src/services/llm/chat/index.ts Normal file
View File

@ -0,0 +1,29 @@
/**
* Chat module export
*/
import restChatService from './rest-chat-service.js';
import sessionsStore from './sessions-store.js';
import { ContextHandler } from './handlers/context-handler.js';
import { ToolHandler } from './handlers/tool-handler.js';
import { StreamHandler } from './handlers/stream-handler.js';
import * as messageFormatter from './utils/message-formatter.js';
import type { ChatSession, ChatMessage, NoteSource } from './interfaces/session.js';
import type { LLMStreamMessage } from './interfaces/ws-messages.js';
// Export components
export {
restChatService as default,
sessionsStore,
ContextHandler,
ToolHandler,
StreamHandler,
messageFormatter
};
// Export types
export type {
ChatSession,
ChatMessage,
NoteSource,
LLMStreamMessage
};

37
src/services/llm/chat/interfaces/session.ts Normal file
View File

@ -0,0 +1,37 @@
/**
* Interfaces for chat sessions and related data
*/
import type { Message } from "../../ai_interface.js";
/**
* Represents a source note from which context is drawn
*/
export interface NoteSource {
noteId: string;
title: string;
content?: string;
similarity?: number;
branchId?: string;
}
/**
* Represents a chat session with message history
*/
export interface ChatSession {
id: string;
title: string;
messages: ChatMessage[];
createdAt: Date;
lastActive: Date;
noteContext?: string;
metadata: Record<string, any>;
}
/**
* Represents a single chat message
*/
export interface ChatMessage {
role: 'user' | 'assistant' | 'system';
content: string;
timestamp?: Date;
}

24
src/services/llm/chat/interfaces/ws-messages.ts Normal file
View File

@ -0,0 +1,24 @@
/**
* Interfaces for WebSocket LLM streaming messages
*/
/**
* Interface for WebSocket LLM streaming messages
*/
export interface LLMStreamMessage {
type: 'llm-stream' | 'tool_execution_start' | 'tool_result' | 'tool_execution_error' | 'tool_completion_processing';
sessionId: string;
content?: string;
thinking?: string;
toolExecution?: {
action?: string;
tool?: string;
toolCallId?: string;
result?: string | Record<string, any>;
error?: string;
args?: Record<string, unknown>;
};
done?: boolean;
error?: string;
raw?: unknown;
}

562
src/services/llm/chat/rest-chat-service.ts Normal file
View File

@ -0,0 +1,562 @@
/**
* Service to handle chat API interactions
*/
import log from "../../log.js";
import type { Request, Response } from "express";
import type { Message, ChatCompletionOptions } from "../ai_interface.js";
import { AIServiceManager } from "../ai_service_manager.js";
import { ChatPipeline } from "../pipeline/chat_pipeline.js";
import type { ChatPipelineInput } from "../pipeline/interfaces.js";
import options from "../../options.js";
import { SEARCH_CONSTANTS } from '../constants/search_constants.js';
// Import our refactored modules
import { ContextHandler } from "./handlers/context-handler.js";
import { ToolHandler } from "./handlers/tool-handler.js";
import { StreamHandler } from "./handlers/stream-handler.js";
import SessionsStore from "./sessions-store.js";
import * as MessageFormatter from "./utils/message-formatter.js";
import type { NoteSource } from "./interfaces/session.js";
import type { LLMStreamMessage } from "./interfaces/ws-messages.js";
/**
* Service to handle chat API interactions
*/
class RestChatService {
/**
* Check if the database is initialized
*/
isDatabaseInitialized(): boolean {
try {
options.getOption('initialized');
return true;
} catch (error) {
return false;
}
}
/**
* Check if AI services are available
*/
safelyUseAIManager(): boolean {
// Only use AI manager if database is initialized
if (!this.isDatabaseInitialized()) {
log.info("AI check failed: Database is not initialized");
return false;
}
// Try to access the manager - will create instance only if needed
try {
// Create local instance to avoid circular references
const aiManager = new AIServiceManager();
if (!aiManager) {
log.info("AI check failed: AI manager module is not available");
return false;
}
const isAvailable = aiManager.isAnyServiceAvailable();
log.info(`AI service availability check result: ${isAvailable}`);
if (isAvailable) {
// Additional diagnostics
try {
const providers = aiManager.getAvailableProviders();
log.info(`Available AI providers: ${providers.join(', ')}`);
} catch (err) {
log.info(`Could not get available providers: ${err}`);
}
}
return isAvailable;
} catch (error) {
log.error(`Error accessing AI service manager: ${error}`);
return false;
}
}
/**
* Handle a message sent to an LLM and get a response
*/
async handleSendMessage(req: Request, res: Response) {
log.info("=== Starting handleSendMessage ===");
try {
// Extract parameters differently based on the request method
let content, useAdvancedContext, showThinking, sessionId;
if (req.method === 'POST') {
// For POST requests, get content from the request body
const requestBody = req.body || {};
content = requestBody.content;
useAdvancedContext = requestBody.useAdvancedContext || false;
showThinking = requestBody.showThinking || false;
// Add logging for POST requests
log.info(`LLM POST message: sessionId=${req.params.sessionId}, useAdvancedContext=${useAdvancedContext}, showThinking=${showThinking}, contentLength=${content ? content.length : 0}`);
} else if (req.method === 'GET') {
// For GET (streaming) requests, get parameters from query params and body
// For streaming requests, we need the content from the body
useAdvancedContext = req.query.useAdvancedContext === 'true' || (req.body && req.body.useAdvancedContext === true);
showThinking = req.query.showThinking === 'true' || (req.body && req.body.showThinking === true);
content = req.body && req.body.content ? req.body.content : '';
// Add detailed logging for GET requests
log.info(`LLM GET stream: sessionId=${req.params.sessionId}, useAdvancedContext=${useAdvancedContext}, showThinking=${showThinking}`);
log.info(`Parameters from query: useAdvancedContext=${req.query.useAdvancedContext}, showThinking=${req.query.showThinking}`);
log.info(`Parameters from body: useAdvancedContext=${req.body?.useAdvancedContext}, showThinking=${req.body?.showThinking}, content=${content ? `${content.substring(0, 20)}...` : 'none'}`);
}
// Get sessionId from URL params since it's part of the route
sessionId = req.params.sessionId;
// For GET requests, ensure we have the stream parameter
if (req.method === 'GET' && req.query.stream !== 'true') {
throw new Error('Stream parameter must be set to true for GET/streaming requests');
}
// For POST requests, validate the content
if (req.method === 'POST' && (!content || typeof content !== 'string' || content.trim().length === 0)) {
throw new Error('Content cannot be empty');
}
// Check if session exists, create one if not
let session = SessionsStore.getSession(sessionId);
if (!session) {
if (req.method === 'GET') {
// For GET requests, we must have an existing session
throw new Error('Session not found');
}
// For POST requests, we can create a new session automatically
log.info(`Session ${sessionId} not found, creating a new one automatically`);
session = SessionsStore.createSession({
title: 'Auto-created Session'
});
log.info(`Created new session with ID: ${session.id}`);
}
// Update session last active timestamp
SessionsStore.touchSession(session.id);
// For POST requests, store the user message
if (req.method === 'POST' && content) {
// Add message to session
session.messages.push({
role: 'user',
content,
timestamp: new Date()
});
// Log a preview of the message
log.info(`Processing LLM message: "${content.substring(0, 50)}${content.length > 50 ? '...' : ''}"`);
}
// Check if AI services are enabled before proceeding
const aiEnabled = await options.getOptionBool('aiEnabled');
log.info(`AI enabled setting: ${aiEnabled}`);
if (!aiEnabled) {
log.info("AI services are disabled by configuration");
return {
error: "AI features are disabled. Please enable them in the settings."
};
}
// Check if AI services are available
log.info("Checking if AI services are available...");
if (!this.safelyUseAIManager()) {
log.info("AI services are not available - checking for specific issues");
try {
// Create a direct instance to avoid circular references
const aiManager = new AIServiceManager();
if (!aiManager) {
log.error("AI service manager is not initialized");
return {
error: "AI service is not properly initialized. Please check your configuration."
};
}
const availableProviders = aiManager.getAvailableProviders();
if (availableProviders.length === 0) {
log.error("No AI providers are available");
return {
error: "No AI providers are configured or available. Please check your AI settings."
};
}
} catch (err) {
log.error(`Detailed AI service check failed: ${err}`);
}
return {
error: "AI services are currently unavailable. Please check your configuration."
};
}
// Create direct instance to avoid circular references
const aiManager = new AIServiceManager();
// Get the default service - just use the first available one
const availableProviders = aiManager.getAvailableProviders();
if (availableProviders.length === 0) {
log.error("No AI providers are available after manager check");
return {
error: "No AI providers are configured or available. Please check your AI settings."
};
}
// Use the first available provider
const providerName = availableProviders[0];
log.info(`Using AI provider: ${providerName}`);
// We know the manager has a 'services' property from our code inspection,
// but TypeScript doesn't know that from the interface.
// This is a workaround to access it
const service = (aiManager as any).services[providerName];
if (!service) {
log.error(`AI service for provider ${providerName} not found`);
return {
error: `Selected AI provider (${providerName}) is not available. Please check your configuration.`
};
}
// Initialize tools
log.info("Initializing LLM agent tools...");
// Ensure tools are initialized to prevent tool execution issues
await ToolHandler.ensureToolsInitialized();
// Create and use the chat pipeline instead of direct processing
const pipeline = new ChatPipeline({
enableStreaming: req.method === 'GET',
enableMetrics: true,
maxToolCallIterations: 5
});
log.info("Executing chat pipeline...");
// Create options object for better tracking
const pipelineOptions = {
// Force useAdvancedContext to be a boolean, no matter what
useAdvancedContext: useAdvancedContext === true,
systemPrompt: session.messages.find(m => m.role === 'system')?.content,
temperature: session.metadata.temperature,
maxTokens: session.metadata.maxTokens,
model: session.metadata.model,
// Set stream based on request type, but ensure it's explicitly a boolean value
// GET requests or format=stream parameter indicates streaming should be used
stream: !!(req.method === 'GET' || req.query.format === 'stream' || req.query.stream === 'true'),
// Include sessionId for tracking tool executions
sessionId: sessionId
};
// Log the options to verify what's being sent to the pipeline
log.info(`Pipeline input options: ${JSON.stringify({
useAdvancedContext: pipelineOptions.useAdvancedContext,
stream: pipelineOptions.stream
})}`);
// Import the WebSocket service for direct access
const wsService = await import('../../ws.js');
// Create a stream callback wrapper
// This will ensure we properly handle all streaming messages
let messageContent = '';
// Prepare the pipeline input
const pipelineInput: ChatPipelineInput = {
messages: session.messages.map(msg => ({
role: msg.role as 'user' | 'assistant' | 'system',
content: msg.content
})),
query: content,
noteId: session.noteContext ?? undefined,
showThinking: showThinking,
options: pipelineOptions,
streamCallback: req.method === 'GET' ? (data, done, rawChunk) => {
try {
// Use WebSocket service to send messages
this.handleStreamCallback(
data, done, rawChunk,
wsService.default, sessionId,
messageContent, session, res
);
} catch (error) {
log.error(`Error in stream callback: ${error}`);
// Try to send error message
try {
wsService.default.sendMessageToAllClients({
type: 'llm-stream',
sessionId,
error: `Stream error: ${error instanceof Error ? error.message : 'Unknown error'}`,
done: true
});
// End the response
res.write(`data: ${JSON.stringify({ error: 'Stream error', done: true })}\n\n`);
res.end();
} catch (e) {
log.error(`Failed to send error message: ${e}`);
}
}
} : undefined
};
// Execute the pipeline
const response = await pipeline.execute(pipelineInput);
// Handle the response
if (req.method === 'POST') {
// Add assistant message to session
session.messages.push({
role: 'assistant',
content: response.text || '',
timestamp: new Date()
});
// Extract sources if they're available
const sources = (response as any).sources || [];
// Store sources in the session metadata if they're present
if (sources.length > 0) {
session.metadata.sources = sources;
log.info(`Stored ${sources.length} sources in session metadata`);
}
// Return the response with complete metadata
return {
content: response.text || '',
sources: sources,
metadata: {
model: response.model || session.metadata.model,
provider: response.provider || session.metadata.provider,
temperature: session.metadata.temperature,
maxTokens: session.metadata.maxTokens,
lastUpdated: new Date().toISOString(),
toolExecutions: session.metadata.toolExecutions || []
}
};
} else {
// For streaming requests, we've already sent the response
return null;
}
} catch (processingError: any) {
log.error(`Error processing message: ${processingError}`);
return {
error: `Error processing your request: ${processingError.message}`
};
}
}
/**
* Handle stream callback for WebSocket communication
*/
private handleStreamCallback(
data: string | null,
done: boolean,
rawChunk: any,
wsService: any,
sessionId: string,
messageContent: string,
session: any,
res: Response
) {
// Only accumulate content that's actually text (not tool execution or thinking info)
if (data) {
messageContent += data;
}
// Create a message object with all necessary fields
const message: LLMStreamMessage = {
type: 'llm-stream',
sessionId
};
// Add content if available - either the new chunk or full content on completion
if (data) {
message.content = data;
}
// Add thinking info if available in the raw chunk
if (rawChunk && 'thinking' in rawChunk && rawChunk.thinking) {
message.thinking = rawChunk.thinking as string;
}
// Add tool execution info if available in the raw chunk
if (rawChunk && 'toolExecution' in rawChunk && rawChunk.toolExecution) {
// Transform the toolExecution to match the expected format
const toolExec = rawChunk.toolExecution;
message.toolExecution = {
// Use optional chaining for all properties
tool: typeof toolExec.tool === 'string'
? toolExec.tool
: toolExec.tool?.name,
result: toolExec.result,
// Map arguments to args
args: 'arguments' in toolExec ?
(typeof toolExec.arguments === 'object' ?
toolExec.arguments as Record<string, unknown> : {}) : {},
// Add additional properties if they exist
action: 'action' in toolExec ? toolExec.action as string : undefined,
toolCallId: 'toolCallId' in toolExec ? toolExec.toolCallId as string : undefined,
error: 'error' in toolExec ? toolExec.error as string : undefined
};
}
// Set done flag explicitly
message.done = done;
// On final message, include the complete content too
if (done) {
// Store the response in the session when done
session.messages.push({
role: 'assistant',
content: messageContent,
timestamp: new Date()
});
}
// Send message to all clients
wsService.sendMessageToAllClients(message);
// Log what was sent (first message and completion)
if (message.thinking || done) {
log.info(
`[WS-SERVER] Sending LLM stream message: sessionId=${sessionId}, content=${!!message.content}, contentLength=${message.content?.length || 0}, thinking=${!!message.thinking}, toolExecution=${!!message.toolExecution}, done=${done}`
);
}
// For GET requests, also send as server-sent events
// Prepare response data for JSON event
const responseData: any = {
content: data,
done
};
// Add tool execution if available
if (rawChunk?.toolExecution) {
responseData.toolExecution = rawChunk.toolExecution;
}
// Send the data as a JSON event
res.write(`data: ${JSON.stringify(responseData)}\n\n`);
if (done) {
res.end();
}
}
/**
* Create a new chat session
*/
async createSession(req: Request, res: Response) {
try {
const options: any = req.body || {};
const title = options.title || 'Chat Session';
// Create a new session through our session store
const session = SessionsStore.createSession({
title,
systemPrompt: options.systemPrompt,
contextNoteId: options.contextNoteId,
maxTokens: options.maxTokens,
model: options.model,
provider: options.provider,
temperature: options.temperature
});
return {
id: session.id,
title: session.title,
createdAt: session.createdAt
};
} catch (error: any) {
log.error(`Error creating LLM session: ${error.message || 'Unknown error'}`);
throw new Error(`Failed to create LLM session: ${error.message || 'Unknown error'}`);
}
}
/**
* Get a specific chat session by ID
*/
async getSession(req: Request, res: Response) {
try {
const { sessionId } = req.params;
// Check if session exists
const session = SessionsStore.getSession(sessionId);
if (!session) {
// Instead of throwing an error, return a structured 404 response
// that the frontend can handle gracefully
res.status(404).json({
error: true,
message: `Session with ID ${sessionId} not found`,
code: 'session_not_found',
sessionId
});
return null; // Return null to prevent further processing
}
// Return session with metadata and additional fields
return {
id: session.id,
title: session.title,
createdAt: session.createdAt,
lastActive: session.lastActive,
messages: session.messages,
noteContext: session.noteContext,
// Include additional fields for the frontend
sources: session.metadata.sources || [],
metadata: {
model: session.metadata.model,
provider: session.metadata.provider,
temperature: session.metadata.temperature,
maxTokens: session.metadata.maxTokens,
lastUpdated: session.lastActive.toISOString(),
// Include simplified tool executions if available
toolExecutions: session.metadata.toolExecutions || []
}
};
} catch (error: any) {
log.error(`Error getting LLM session: ${error.message || 'Unknown error'}`);
throw new Error(`Failed to get session: ${error.message || 'Unknown error'}`);
}
}
/**
* Delete a chat session
*/
async deleteSession(req: Request, res: Response) {
try {
const { sessionId } = req.params;
// Delete the session
const success = SessionsStore.deleteSession(sessionId);
if (!success) {
throw new Error(`Session with ID ${sessionId} not found`);
}
return {
success: true,
message: `Session ${sessionId} deleted successfully`
};
} catch (error: any) {
log.error(`Error deleting LLM session: ${error.message || 'Unknown error'}`);
throw new Error(`Failed to delete session: ${error.message || 'Unknown error'}`);
}
}
/**
* Get all sessions
*/
getSessions() {
return SessionsStore.getAllSessions();
}
}
// Create singleton instance
const restChatService = new RestChatService();
export default restChatService;

168
src/services/llm/chat/sessions-store.ts Normal file
View File

@ -0,0 +1,168 @@
/**
* In-memory storage for chat sessions
*/
import log from "../../log.js";
import { LLM_CONSTANTS } from '../constants/provider_constants.js';
import { SEARCH_CONSTANTS } from '../constants/search_constants.js';
import { randomString } from "../../utils.js";
import type { ChatSession, ChatMessage } from './interfaces/session.js';
// In-memory storage for sessions
const sessions = new Map<string, ChatSession>();
// Flag to track if cleanup timer has been initialized
let cleanupInitialized = false;
/**
* Provides methods to manage chat sessions
*/
class SessionsStore {
/**
* Initialize the session cleanup timer to remove old/inactive sessions
*/
initializeCleanupTimer(): void {
if (cleanupInitialized) {
return;
}
// Clean sessions that have expired based on the constants
function cleanupOldSessions() {
const expiryTime = new Date(Date.now() - LLM_CONSTANTS.SESSION.SESSION_EXPIRY_MS);
for (const [sessionId, session] of sessions.entries()) {
if (session.lastActive < expiryTime) {
sessions.delete(sessionId);
}
}
}
// Run cleanup at the configured interval
setInterval(cleanupOldSessions, LLM_CONSTANTS.SESSION.CLEANUP_INTERVAL_MS);
cleanupInitialized = true;
log.info("Session cleanup timer initialized");
}
/**
* Get all sessions
*/
getAllSessions(): Map<string, ChatSession> {
return sessions;
}
/**
* Get a specific session by ID
*/
getSession(sessionId: string): ChatSession | undefined {
return sessions.get(sessionId);
}
/**
* Create a new session
*/
createSession(options: {
title?: string;
systemPrompt?: string;
contextNoteId?: string;
maxTokens?: number;
model?: string;
provider?: string;
temperature?: number;
}): ChatSession {
this.initializeCleanupTimer();
const title = options.title || 'Chat Session';
const sessionId = randomString(16);
const now = new Date();
// Initial system message if provided
const messages: ChatMessage[] = [];
if (options.systemPrompt) {
messages.push({
role: 'system',
content: options.systemPrompt,
timestamp: now
});
}
// Create and store the session
const session: ChatSession = {
id: sessionId,
title,
messages,
createdAt: now,
lastActive: now,
noteContext: options.contextNoteId,
metadata: {
temperature: options.temperature || SEARCH_CONSTANTS.TEMPERATURE.DEFAULT,
maxTokens: options.maxTokens,
model: options.model,
provider: options.provider,
sources: [],
toolExecutions: [],
lastUpdated: now.toISOString()
}
};
sessions.set(sessionId, session);
log.info(`Created new session with ID: ${sessionId}`);
return session;
}
/**
* Update a session's last active timestamp
*/
touchSession(sessionId: string): boolean {
const session = sessions.get(sessionId);
if (!session) {
return false;
}
session.lastActive = new Date();
return true;
}
/**
* Delete a session
*/
deleteSession(sessionId: string): boolean {
return sessions.delete(sessionId);
}
/**
* Record a tool execution in the session metadata
*/
recordToolExecution(sessionId: string, tool: any, result: string, error?: string): void {
if (!sessionId) return;
const session = sessions.get(sessionId);
if (!session) return;
try {
const toolExecutions = session.metadata.toolExecutions || [];
// Format tool execution record
const execution = {
id: tool.id || `tool-${Date.now()}-${Math.random().toString(36).substring(2, 7)}`,
name: tool.function?.name || 'unknown',
arguments: typeof tool.function?.arguments === 'string'
? (() => { try { return JSON.parse(tool.function.arguments); } catch { return tool.function.arguments; } })()
: tool.function?.arguments || {},
result: result,
error: error,
timestamp: new Date().toISOString()
};
// Add to tool executions
toolExecutions.push(execution);
session.metadata.toolExecutions = toolExecutions;
log.info(`Recorded tool execution for ${execution.name} in session ${sessionId}`);
} catch (err) {
log.error(`Failed to record tool execution: ${err}`);
}
}
}
// Create singleton instance
const sessionsStore = new SessionsStore();
export default sessionsStore;

121
src/services/llm/chat/utils/message-formatter.ts Normal file
View File

@ -0,0 +1,121 @@
/**
* Message formatting utilities for different LLM providers
*/
import type { Message } from "../../ai_interface.js";
/**
* Interface for message formatters
*/
interface MessageFormatter {
formatMessages(messages: Message[], systemPrompt?: string, context?: string): Message[];
}
/**
* Factory to get the appropriate message formatter for a given provider
*/
export function getFormatter(providerName: string): MessageFormatter {
// Currently we use a simple implementation that works for most providers
// In the future, this could be expanded to have provider-specific formatters
return {
formatMessages(messages: Message[], systemPrompt?: string, context?: string): Message[] {
// Simple implementation that works for most providers
const formattedMessages: Message[] = [];
// Add system message if context or systemPrompt is provided
if (context || systemPrompt) {
formattedMessages.push({
role: 'system',
content: systemPrompt || (context ? `Use the following context to answer the query: ${context}` : '')
});
}
// Add all other messages
for (const message of messages) {
if (message.role === 'system' && formattedMessages.some(m => m.role === 'system')) {
// Skip duplicate system messages
continue;
}
formattedMessages.push(message);
}
return formattedMessages;
}
};
}
/**
* Build messages with context for a specific LLM provider
*/
export async function buildMessagesWithContext(
messages: Message[],
context: string,
llmService: any
): Promise<Message[]> {
try {
if (!messages || messages.length === 0) {
return [];
}
if (!context || context.trim() === '') {
return messages;
}
// Get the provider name, handling service classes and raw provider names
let providerName: string;
if (typeof llmService === 'string') {
// If llmService is a string, assume it's the provider name
providerName = llmService;
} else if (llmService.constructor && llmService.constructor.name) {
// Extract provider name from service class name (e.g., OllamaService -> ollama)
providerName = llmService.constructor.name.replace('Service', '').toLowerCase();
} else {
// Fallback to default
providerName = 'default';
}
// Get the appropriate formatter for this provider
const formatter = getFormatter(providerName);
// Format messages with context using the provider-specific formatter
const formattedMessages = formatter.formatMessages(
messages,
undefined, // No system prompt override - use what's in the messages
context
);
return formattedMessages;
} catch (error) {
console.error(`Error building messages with context: ${error}`);
// Fallback to original messages in case of error
return messages;
}
}
/**
* Build context from a list of note sources and a query
*/
export function buildContextFromNotes(sources: any[], query: string): string {
if (!sources || sources.length === 0) {
return query || '';
}
const noteContexts = sources
.filter(source => source.content) // Only include sources with content
.map((source) => {
// Format each note with its title as a natural heading and wrap in <note> tags
return `<note>\n### ${source.title}\n${source.content || 'No content available'}\n</note>`;
})
.join('\n\n');
if (!noteContexts) {
return query || '';
}
// Import the CONTEXT_PROMPTS constant
const { CONTEXT_PROMPTS } = require('../../constants/llm_prompt_constants.js');
// Use the template from the constants file, replacing placeholders
return CONTEXT_PROMPTS.CONTEXT_NOTES_WRAPPER
.replace('{noteContexts}', noteContexts)
.replace('{query}', query);
}

1849
src/services/llm/rest_chat_service.ts
View File

File diff suppressed because it is too large Load Diff