import {ChatMessage, CreateSessionResponse, DiagramSession, DiagramToolCall, SessionEntity, SessionUsage, SyncEntitiesResponse, ToolResult} from "../types/chatTypes"; import {clearDiagram, connectEntities, createEntity, getCameraPosition, listEntities, modifyConnection, modifyEntity, removeEntity} from "./entityBridge"; import {v4 as uuidv4} from 'uuid'; import log from 'loglevel'; const logger = log.getLogger('diagramAI'); logger.setLevel('debug'); // Session management let currentSessionId: string | null = null; // Model management export type AIProvider = 'claude' | 'ollama' | 'cloudflare'; export interface ModelInfo { id: string; name: string; description: string; provider: AIProvider; } const AVAILABLE_MODELS: ModelInfo[] = [ // Claude models { id: 'claude-sonnet-4-20250514', name: 'Claude Sonnet 4', description: 'Balanced performance and speed (default)', provider: 'claude' }, { id: 'claude-opus-4-20250514', name: 'Claude Opus 4', description: 'Most capable, best for complex tasks', provider: 'claude' }, { id: 'claude-haiku-3-5-20241022', name: 'Claude Haiku 3.5', description: 'Fastest responses, good for simple tasks', provider: 'claude' }, // Cloudflare Workers AI models - with tool support { id: '@cf/mistralai/mistral-small-3.1-24b-instruct', name: 'Mistral Small 3.1 (CF)', description: 'Best CF model - supports diagram tools', provider: 'cloudflare' }, { id: '@hf/nousresearch/hermes-2-pro-mistral-7b', name: 'Hermes 2 Pro (CF)', description: 'Lightweight - supports diagram tools', provider: 'cloudflare' }, // Cloudflare models WITHOUT tool support - chat only { id: '@cf/meta/llama-3.3-70b-instruct-fp8-fast', name: 'Llama 3.3 70B (CF)', description: 'Powerful but NO tool support', provider: 'cloudflare' }, { id: '@cf/meta/llama-3.1-8b-instruct', name: 'Llama 3.1 8B (CF)', description: 'Fast/cheap but NO tool support', provider: 'cloudflare' }, { id: '@cf/deepseek-ai/deepseek-r1-distill-qwen-32b', name: 'DeepSeek R1 (CF)', description: 'Reasoning but NO tool support', provider: 'cloudflare' }, { id: '@cf/qwen/qwen2.5-coder-32b-instruct', name: 'Qwen 2.5 Coder (CF)', description: 'Code-focused but NO tool support', provider: 'cloudflare' }, // Ollama models (local) { id: 'llama3.1', name: 'Llama 3.1', description: 'Local model with function calling support', provider: 'ollama' }, { id: 'mistral', name: 'Mistral', description: 'Fast local model with good tool support', provider: 'ollama' }, { id: 'qwen2.5', name: 'Qwen 2.5', description: 'Capable local model with function calling', provider: 'ollama' } ]; let currentModelId: string = '@cf/mistralai/mistral-small-3.1-24b-instruct'; /** * Get the API endpoint for the current model's provider */ function getApiEndpoint(): string { const model = getCurrentModel(); if (model.provider === 'ollama') { return '/api/ollama/v1/messages'; } if (model.provider === 'cloudflare') { return '/api/cloudflare/v1/messages'; } return '/api/claude/v1/messages'; } /** * Get available models */ export function getAvailableModels(): ModelInfo[] { return [...AVAILABLE_MODELS]; } /** * Get current model */ export function getCurrentModel(): ModelInfo { return AVAILABLE_MODELS.find(m => m.id === currentModelId) || AVAILABLE_MODELS[0]; } /** * Set current model */ export function setCurrentModel(modelId: string): boolean { const model = AVAILABLE_MODELS.find(m => m.id === modelId); if (model) { currentModelId = modelId; logger.info('Model changed to:', model.name); return true; } logger.warn('Invalid model ID:', modelId); return false; } /** * Create a new session or resume existing one for a diagram */ export async function createOrResumeSession(diagramId: string): Promise<{ session: DiagramSession; isNew: boolean }> { const response = await fetch('/api/session/create', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ diagramId }) }); if (!response.ok) { throw new Error(`Failed to create session: ${response.status}`); } const data: CreateSessionResponse = await response.json(); currentSessionId = data.session.id; return data; } /** * Get current session ID */ export function getCurrentSessionId(): string | null { return currentSessionId; } /** * Sync entities to the current session */ export async function syncEntitiesToSession(entities: SessionEntity[]): Promise { if (!currentSessionId) { console.warn('No active session to sync entities to'); return; } const response = await fetch(`/api/session/${currentSessionId}/sync`, { method: 'PUT', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ entities }) }); if (!response.ok) { console.error('Failed to sync entities:', response.status); } } /** * Get session with conversation history */ export async function getSessionHistory(): Promise { if (!currentSessionId) { return null; } const response = await fetch(`/api/session/${currentSessionId}`); if (!response.ok) { return null; } const data = await response.json(); return data.session; } /** * Clear session history */ export async function clearSessionHistory(): Promise { if (!currentSessionId) { return; } await fetch(`/api/session/${currentSessionId}/history`, { method: 'DELETE' }); } /** * Get token usage for current session */ export async function getSessionUsage(): Promise { if (!currentSessionId) { return null; } try { const response = await fetch(`/api/session/${currentSessionId}/usage`); if (!response.ok) { return null; } return await response.json(); } catch (err) { logger.error('Failed to fetch session usage:', err); return null; } } const SYSTEM_PROMPT = `You are a 3D diagram assistant. You MUST use tools to perform actions - never just describe what you would do. ## CRITICAL RULES 1. When the user asks to create, add, or make something → CALL create_entity tool 2. When the user asks to connect things → CALL connect_entities tool 3. When the user asks to change, modify, move, resize, or rotate → CALL modify_entity tool 4. When the user asks to remove or delete → CALL remove_entity tool 5. When the user asks what exists or to list → CALL list_entities tool 6. When the user uses directions (left, right, forward, in front of me) → CALL get_camera_position FIRST 7. When diagramming unfamiliar topics → CALL search_wikipedia FIRST to research the concept DO NOT just describe actions. DO NOT say "I will create..." without calling a tool. ALWAYS call the appropriate tool. ## Research First When asked to diagram a technical concept, architecture, or system you're not fully familiar with, use search_wikipedia to research it first. This ensures accurate and comprehensive diagrams. ## Available Shapes box, sphere, cylinder, cone, plane, person ## Available Colors red, blue, green, yellow, orange, purple, cyan, pink, white, black, brown, gray (or hex like #ff5500) ## Position Coordinates - x: left/right - y: up/down (1.5 = eye level, 0 = floor) - z: forward/backward ## Layout Guidelines - Spread entities apart (at least 0.5 units) - Use y=1.5 for eye-level entities - For relative directions, call get_camera_position first Be concise. Call tools immediately when the user requests an action.`; const TOOLS = [ { name: "create_entity", description: "Create a 3D shape in the diagram. Use this to add new elements like boxes, spheres, cylinders, etc.", input_schema: { type: "object", properties: { shape: { type: "string", enum: ["box", "sphere", "cylinder", "cone", "plane", "person"], description: "The type of 3D shape to create" }, color: { type: "string", description: "Color name (red, blue, green, etc.) or hex code (#ff0000)" }, label: { type: "string", description: "Text label to display on or near the entity" }, position: { type: "object", properties: { x: {type: "number", description: "Left/right position"}, y: {type: "number", description: "Up/down position (1.5 = eye level)"}, z: {type: "number", description: "Forward/backward position"} }, description: "3D position. If not specified, defaults to (0, 1.5, 2)" } }, required: ["shape"] } }, { name: "connect_entities", description: "Draw a connection line between two entities. Use entity IDs or labels to identify them.", input_schema: { type: "object", properties: { from: { type: "string", description: "ID or label of the source entity" }, to: { type: "string", description: "ID or label of the target entity" }, color: { type: "string", description: "Color of the connection line" } }, required: ["from", "to"] } }, { name: "list_entities", description: "List all entities currently in the diagram. Use this to see what exists before connecting or modifying.", input_schema: { type: "object", properties: {} } }, { name: "remove_entity", description: "Remove an entity from the diagram by its ID or label.", input_schema: { type: "object", properties: { target: { type: "string", description: "ID or label of the entity to remove" } }, required: ["target"] } }, { name: "modify_entity", description: "Modify an existing entity's properties like color, label, position, scale, or rotation.", input_schema: { type: "object", properties: { target: { type: "string", description: "ID or label of the entity to modify" }, color: { type: "string", description: "New color for the entity" }, label: { type: "string", description: "New label text. Use empty string \"\" to remove the label." }, position: { type: "object", properties: { x: {type: "number"}, y: {type: "number"}, z: {type: "number"} } }, scale: { oneOf: [ { type: "number", description: "Uniform scale factor (e.g., 0.2 = double size, 0.05 = half size). Default is 0.1." }, { type: "object", properties: { x: {type: "number"}, y: {type: "number"}, z: {type: "number"} }, description: "Non-uniform scale as {x, y, z}. Default is {x: 0.1, y: 0.1, z: 0.1}." } ], description: "Scale/size of the entity. Use a number for uniform scaling or {x, y, z} for non-uniform." }, rotation: { oneOf: [ { type: "number", description: "Y-axis rotation in degrees (e.g., 90 = turn right 90°, -90 = turn left 90°)" }, { type: "object", properties: { x: {type: "number", description: "Pitch in degrees"}, y: {type: "number", description: "Yaw in degrees"}, z: {type: "number", description: "Roll in degrees"} }, description: "Full 3D rotation in degrees as {x, y, z}" } ], description: "Rotation in degrees. Use a number for Y-axis rotation or {x, y, z} for full 3D rotation." } }, required: ["target"] } }, { name: "modify_connection", description: "Modify a connection's label or color. Connections can be identified by their label (e.g., 'Server to Database') or by specifying the from/to entities.", input_schema: { type: "object", properties: { target: { type: "string", description: "Label of the connection to modify, or use from/to to identify it" }, from: { type: "string", description: "ID or label of the source entity (alternative to target)" }, to: { type: "string", description: "ID or label of the destination entity (alternative to target)" }, label: { type: "string", description: "New label text for the connection. Use empty string \"\" to remove the label." }, color: { type: "string", description: "New color for the connection" } } } }, { name: "clear_diagram", description: "DESTRUCTIVE: Permanently delete ALL entities from the diagram and clear the session. This cannot be undone. IMPORTANT: Before calling this tool, you MUST first ask the user to confirm by saying something like 'Are you sure you want to clear the entire diagram? This will permanently delete all entities and cannot be undone.' Only call this tool with confirmed=true AFTER the user explicitly confirms (e.g., says 'yes', 'confirm', 'do it', etc.).", input_schema: { type: "object", properties: { confirmed: { type: "boolean", description: "Must be true to execute. Only set to true after user has explicitly confirmed the deletion." } }, required: ["confirmed"] } }, { name: "get_camera_position", description: "Get the current camera/viewer position and orientation in the 3D scene. Use this to understand where the user is looking and to position new entities relative to their view. Returns position (x, y, z) and forward direction vector.", input_schema: { type: "object", properties: {} } }, { name: "list_models", description: "List all available AI models that can be used for this conversation.", input_schema: { type: "object", properties: {} } }, { name: "get_current_model", description: "Get information about the currently active AI model.", input_schema: { type: "object", properties: {} } }, { name: "set_model", description: "Change the AI model used for this conversation. Use list_models first to see available options.", input_schema: { type: "object", properties: { model_id: { type: "string", description: "The model ID to switch to. Claude models: 'claude-sonnet-4-20250514', 'claude-opus-4-20250514', 'claude-haiku-3-5-20241022'. Ollama models (local): 'llama3.1', 'mistral', 'qwen2.5'" } }, required: ["model_id"] } }, { name: "clear_conversation", description: "Clear the conversation history to start fresh. Use this when the conversation has become too long, confusing, or when the user wants to start over. This preserves the diagram entities but clears chat history.", input_schema: { type: "object", properties: {} } }, { name: "search_wikipedia", description: "Search Wikipedia for information about a topic. Use this to research concepts, architectures, technologies, or anything else that would help create more accurate and detailed diagrams. Returns summaries from multiple related articles.", input_schema: { type: "object", properties: { topic: { type: "string", description: "The topic or concept to search for (e.g., 'microservices architecture', 'neural network', 'kubernetes')" } }, required: ["topic"] } } ]; interface ClaudeMessage { role: 'user' | 'assistant'; content: string | ClaudeContentBlock[]; } interface ClaudeContentBlock { type: 'text' | 'tool_use' | 'tool_result'; text?: string; id?: string; name?: string; input?: Record; tool_use_id?: string; content?: string; } interface ClaudeResponse { content: ClaudeContentBlock[]; stop_reason: 'end_turn' | 'tool_use' | 'max_tokens'; } async function executeToolCall(toolCall: DiagramToolCall): Promise { logger.debug('[executeToolCall] Executing:', toolCall.name, toolCall.input); let result: ToolResult; switch (toolCall.name) { case 'create_entity': result = createEntity(toolCall.input); break; case 'connect_entities': result = await connectEntities(toolCall.input); break; case 'remove_entity': result = removeEntity(toolCall.input); break; case 'modify_entity': result = modifyEntity(toolCall.input); break; case 'modify_connection': result = await modifyConnection(toolCall.input); break; case 'list_entities': result = await listEntities(); break; case 'clear_diagram': result = await clearDiagram(toolCall.input); // If successful, also clear the session (history and entities) if (result.success && currentSessionId) { await clearSessionHistory(); // Sync empty entity list to clear server-side entity cache await syncEntitiesToSession([]); } break; case 'get_camera_position': result = await getCameraPosition(); break; case 'list_models': { const models = getAvailableModels(); const current = getCurrentModel(); const modelList = models.map(m => `• ${m.name} (${m.id})${m.id === current.id ? ' [CURRENT]' : ''}\n ${m.description}` ).join('\n\n'); result = { toolName: 'list_models', success: true, message: `Available models:\n\n${modelList}` }; break; } case 'get_current_model': { const model = getCurrentModel(); result = { toolName: 'get_current_model', success: true, message: `Current model: ${model.name} (${model.id})\n${model.description}` }; break; } case 'set_model': { const success = setCurrentModel(toolCall.input.model_id); if (success) { const model = getCurrentModel(); result = { toolName: 'set_model', success: true, message: `Model changed to: ${model.name}\n${model.description}\n\nNote: The new model will be used starting from the next message.` }; } else { const models = getAvailableModels(); result = { toolName: 'set_model', success: false, message: `Invalid model ID: "${toolCall.input.model_id}"\n\nAvailable models: ${models.map(m => m.id).join(', ')}` }; } break; } case 'clear_conversation': { const cleared = await clearConversationHistory(); result = { toolName: 'clear_conversation', success: cleared, message: cleared ? 'Conversation history cleared. Starting fresh while keeping the diagram intact.' : 'Failed to clear conversation history (no active session).' }; break; } case 'search_wikipedia': { const searchResult = await searchWikipedia(toolCall.input.topic); result = { toolName: 'search_wikipedia', success: searchResult.success, message: searchResult.message }; break; } default: result = { toolName: 'unknown', success: false, message: 'Unknown tool' }; } logger.debug('[executeToolCall] Result:', result); return result; } /** * Clear conversation history on the server */ export async function clearConversationHistory(): Promise { if (!currentSessionId) { logger.warn('[clearConversationHistory] No active session'); return false; } try { const response = await fetch(`/api/session/${currentSessionId}/history`, { method: 'DELETE' }); if (response.ok) { logger.info('[clearConversationHistory] History cleared for session:', currentSessionId); return true; } else { logger.error('[clearConversationHistory] Failed:', response.status); return false; } } catch (error) { logger.error('[clearConversationHistory] Error:', error); return false; } } /** * Search Wikipedia for a topic using multiple query variations * Returns combined results for context */ async function searchWikipedia(topic: string): Promise<{ success: boolean; message: string }> { logger.info('[searchWikipedia] Searching for:', topic); // Generate 3 search variations const searchQueries = [ topic, `${topic} overview`, `${topic} definition concept` ]; const results: string[] = []; for (const query of searchQueries) { try { // Search Wikipedia API const searchUrl = `https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch=${encodeURIComponent(query)}&srlimit=2&format=json&origin=*`; const searchResponse = await fetch(searchUrl); const searchData = await searchResponse.json(); if (searchData.query?.search?.length > 0) { // Get summaries for top results for (const result of searchData.query.search.slice(0, 2)) { const title = result.title; // Skip if we already have this article if (results.some(r => r.includes(`**${title}**`))) { continue; } // Get article summary const summaryUrl = `https://en.wikipedia.org/api/rest_v1/page/summary/${encodeURIComponent(title)}`; const summaryResponse = await fetch(summaryUrl); if (summaryResponse.ok) { const summaryData = await summaryResponse.json(); const extract = summaryData.extract || 'No summary available.'; // Truncate long summaries const truncated = extract.length > 500 ? extract.substring(0, 500) + '...' : extract; results.push(`**${title}**\n${truncated}`); } } } } catch (error) { logger.warn('[searchWikipedia] Error searching for:', query, error); } } if (results.length === 0) { return { success: false, message: `No Wikipedia articles found for "${topic}". Try a different search term.` }; } // Combine results const combinedResults = results.slice(0, 5).join('\n\n---\n\n'); return { success: true, message: `## Wikipedia Research: ${topic}\n\nFound ${results.length} relevant articles:\n\n${combinedResults}\n\n---\nUse this information to create accurate diagrams. You can now create entities based on this research.` }; } export async function sendMessage( userMessage: string, conversationHistory: ChatMessage[], onToolResult?: (result: ToolResult) => void ): Promise<{ response: string; toolResults: ToolResult[] }> { logger.debug('[sendMessage] Starting with message:', userMessage); logger.debug('[sendMessage] Session ID:', currentSessionId); // When using sessions, we don't need to send full history - server manages it // Just send the new message const messages: ClaudeMessage[] = currentSessionId ? [{ role: 'user', content: userMessage }] : conversationHistory .filter(m => !m.isLoading) .map(m => ({ role: m.role, content: m.content })); if (!currentSessionId) { messages.push({role: 'user', content: userMessage}); } logger.debug('[sendMessage] Messages to send:', messages.length); const allToolResults: ToolResult[] = []; let finalResponse = ''; let continueLoop = true; let loopCount = 0; const MAX_TOOL_ITERATIONS = 10; // Safety limit to prevent infinite loops while (continueLoop) { loopCount++; logger.debug(`[sendMessage] Loop iteration ${loopCount}`); // Safety check: prevent infinite tool calling loops if (loopCount > MAX_TOOL_ITERATIONS) { logger.warn(`[sendMessage] Max tool iterations (${MAX_TOOL_ITERATIONS}) reached, breaking loop`); break; } const requestBody = { model: currentModelId, max_tokens: 4096, system: SYSTEM_PROMPT, tools: TOOLS, messages, ...(currentSessionId && { sessionId: currentSessionId }) }; logger.debug('[sendMessage] Request body:', JSON.stringify(requestBody, null, 2).substring(0, 500) + '...'); const apiEndpoint = getApiEndpoint(); logger.debug('[sendMessage] Using API endpoint:', apiEndpoint); const response = await fetch(apiEndpoint, { method: 'POST', headers: { 'Content-Type': 'application/json', }, body: JSON.stringify(requestBody) }); logger.debug('[sendMessage] Response status:', response.status); if (!response.ok) { const errorText = await response.text(); logger.error('[sendMessage] API error:', response.status, errorText); throw new Error(`API error: ${response.status} - ${errorText}`); } const data: ClaudeResponse = await response.json(); logger.debug('[sendMessage] Response data:', JSON.stringify(data, null, 2).substring(0, 500) + '...'); logger.debug('[sendMessage] Stop reason:', data.stop_reason); const textBlocks = data.content.filter(b => b.type === 'text'); const toolBlocks = data.content.filter(b => b.type === 'tool_use'); logger.debug('[sendMessage] Text blocks:', textBlocks.length, 'Tool blocks:', toolBlocks.length); if (textBlocks.length > 0) { finalResponse = textBlocks.map(b => b.text).join('\n'); logger.debug('[sendMessage] Final response:', finalResponse.substring(0, 200) + '...'); } if (data.stop_reason === 'tool_use' && toolBlocks.length > 0) { logger.debug('[sendMessage] Processing tool calls...'); messages.push({ role: 'assistant', content: data.content }); const toolResults: ClaudeContentBlock[] = []; let modelSwitched = false; for (const toolBlock of toolBlocks) { logger.debug('[sendMessage] Tool call:', toolBlock.name, JSON.stringify(toolBlock.input)); const toolCall: DiagramToolCall = { name: toolBlock.name as DiagramToolCall['name'], input: toolBlock.input as DiagramToolCall['input'] }; // Track if we're switching models if (toolBlock.name === 'set_model') { modelSwitched = true; } const result = await executeToolCall(toolCall); logger.debug('[sendMessage] Tool result:', result); allToolResults.push(result); onToolResult?.(result); toolResults.push({ type: 'tool_result', tool_use_id: toolBlock.id, content: result.message }); } messages.push({ role: 'user', content: toolResults }); // If model was switched, break the loop - the new model won't have proper context // and continuing could cause infinite loops with Cloudflare models if (modelSwitched) { logger.debug('[sendMessage] Model switched, ending loop to prevent context issues'); // Set a response so user knows what happened if (!finalResponse) { const model = getCurrentModel(); finalResponse = `Switched to ${model.name}. The new model is now active for your next message.`; } continueLoop = false; } else { logger.debug('[sendMessage] Added tool results, continuing loop...'); } } else { logger.debug('[sendMessage] No more tool calls, ending loop'); continueLoop = false; } } logger.debug('[sendMessage] Complete. Final response length:', finalResponse.length, 'Tool results:', allToolResults.length); return {response: finalResponse, toolResults: allToolResults}; } export function createUserMessage(content: string): ChatMessage { return { id: uuidv4(), role: 'user', content, timestamp: new Date() }; } export function createAssistantMessage(content: string, toolResults?: ToolResult[]): ChatMessage { return { id: uuidv4(), role: 'assistant', content, timestamp: new Date(), toolResults }; } export function createLoadingMessage(): ChatMessage { return { id: uuidv4(), role: 'assistant', content: '', timestamp: new Date(), isLoading: true }; }