From 19a905d9d1717eeffd42a8cecc7df7e78c54a5e2 Mon Sep 17 00:00:00 2001 From: Ajax Davis Date: Fri, 23 Jan 2026 10:54:02 +1000 Subject: [PATCH] refactor(omega): use BM25 auto-loading instead of meta-tools --- apps/web/package.json | 2 - .../conversations/[id]/messages/route.ts | 280 ++++++++++++++---- .../src/app/omega/[conversationId]/page.tsx | 266 ++++++++++------- apps/web/src/lib/omega/system-prompt.ts | 49 +-- pnpm-lock.yaml | 6 - 5 files changed, 401 insertions(+), 202 deletions(-) diff --git a/apps/web/package.json b/apps/web/package.json index ca3ce25..d4be2fe 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -33,8 +33,6 @@ "@tpmjs/env": "workspace:*", "@tpmjs/npm-client": "workspace:*", "@tpmjs/package-executor": "workspace:*", - "@tpmjs/registry-execute": "workspace:*", - "@tpmjs/registry-search": "workspace:*", "@tpmjs/types": "workspace:*", "@tpmjs/ui": "workspace:*", "@tpmjs/utils": "workspace:*", diff --git a/apps/web/src/app/api/omega/conversations/[id]/messages/route.ts b/apps/web/src/app/api/omega/conversations/[id]/messages/route.ts index 3ecc071..437e240 100644 --- a/apps/web/src/app/api/omega/conversations/[id]/messages/route.ts +++ b/apps/web/src/app/api/omega/conversations/[id]/messages/route.ts @@ -4,15 +4,13 @@ * POST: Send a message and stream the AI response via SSE * * This endpoint implements the core Omega chat functionality with: - * - Registry search tool for discovering tools - * - Tool executor for running discovered tools + * - Automatic BM25 search to find relevant tools based on user message + * - Dynamic tool loading - top 15 matching tools become available to the AI * - SSE streaming for real-time updates */ import { Prisma, prisma } from '@tpmjs/db'; -import { registryExecuteTool } from '@tpmjs/registry-execute'; -import { registrySearchTool } from '@tpmjs/registry-search'; -import type { ModelMessage } from 'ai'; +import { jsonSchema, type ModelMessage } from 'ai'; import { type NextRequest, NextResponse } from 'next/server'; import { z } from 'zod'; import { authenticateRequest } from '~/lib/api-keys/middleware'; @@ -40,6 +38,126 @@ const SendMessageSchema = z.object({ message: z.string().min(1).max(10000), }); +// Executor service URL +const EXECUTOR_URL = process.env.TPMJS_EXECUTOR_URL || 'https://executor.tpmjs.com'; + +/** + * Search for relevant tools based on user query + */ +async function searchRelevantTools( + query: string, + limit = 15 +): Promise< + Array<{ + toolId: string; + packageName: string; + name: string; + description: string; + version: string; + importUrl: string; + inputSchema?: unknown; + }> +> { + const params = new URLSearchParams({ + q: query, + limit: String(limit), + }); + + // Use internal API (same server) + const baseUrl = process.env.VERCEL_URL + ? `https://${process.env.VERCEL_URL}` + : 'http://localhost:3000'; + + const response = await fetch(`${baseUrl}/api/tools/search?${params}`); + + if (!response.ok) { + console.error(`Tool search failed: ${response.status} ${response.statusText}`); + return []; + } + + // biome-ignore lint/suspicious/noExplicitAny: API response types vary + const data = (await response.json()) as any; + const toolsArray = data.results?.tools || []; + + // biome-ignore lint/suspicious/noExplicitAny: API response types vary + return toolsArray.map((tool: any) => ({ + toolId: `${tool.package.npmPackageName}::${tool.name}`, + packageName: tool.package.npmPackageName, + name: tool.name, + description: tool.description || `Tool: ${tool.name}`, + version: tool.package.npmVersion, + importUrl: `https://esm.sh/${tool.package.npmPackageName}@${tool.package.npmVersion}`, + inputSchema: tool.inputSchema, + })); +} + +/** + * Create a dynamic tool wrapper that executes via the sandbox executor + */ +function createDynamicTool(toolMeta: { + toolId: string; + packageName: string; + name: string; + description: string; + version: string; + importUrl: string; + inputSchema?: unknown; +}) { + // Import tool() dynamically to avoid top-level await + const { tool } = require('ai'); + + return tool({ + description: toolMeta.description, + inputSchema: toolMeta.inputSchema + ? jsonSchema(toolMeta.inputSchema as Parameters[0]) + : jsonSchema({ + type: 'object', + properties: {}, + additionalProperties: true, + }), + // biome-ignore lint/suspicious/noExplicitAny: Dynamic tool params + execute: async (params: any) => { + console.log(`🚀 Executing ${toolMeta.packageName}/${toolMeta.name} with params:`, params); + + const response = await fetch(`${EXECUTOR_URL}/execute-tool`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + packageName: toolMeta.packageName, + name: toolMeta.name, + version: toolMeta.version, + importUrl: toolMeta.importUrl, + params, + env: {}, // TODO: Pass user's env vars if needed + }), + }); + + // biome-ignore lint/suspicious/noExplicitAny: API response types vary + const result = (await response.json()) as any; + + if (!result.success) { + console.error(`❌ Tool execution failed: ${result.error}`); + throw new Error(result.error || 'Tool execution failed'); + } + + console.log(`✅ Tool executed in ${result.executionTimeMs}ms`); + return result.output; + }, + }); +} + +/** + * Sanitize tool name to be a valid JS identifier + */ +function sanitizeToolName(name: string): string { + return name + .replace(/@/g, '') + .replace(/\//g, '_') + .replace(/-/g, '_') + .replace(/::/g, '_') + .replace(/[^a-zA-Z0-9_]/g, ''); +} + /** * POST /api/omega/conversations/[id]/messages * Send a message and stream the AI response via SSE @@ -126,6 +244,27 @@ export async function POST(request: NextRequest, context: RouteContext): Promise }, }); + // 🔍 Search for relevant tools based on user's message + console.log(`🔍 Searching for tools matching: "${parsed.data.message}"`); + const relevantTools = await searchRelevantTools(parsed.data.message, 15); + console.log(`📦 Found ${relevantTools.length} relevant tools`); + + // Create dynamic tool wrappers for each found tool + // biome-ignore lint/suspicious/noExplicitAny: Dynamic tool types + const tools: Record = {}; + + for (const toolMeta of relevantTools) { + const sanitizedName = sanitizeToolName(toolMeta.toolId); + try { + tools[sanitizedName] = createDynamicTool(toolMeta); + console.log(`✅ Loaded tool: ${sanitizedName}`); + } catch (error) { + console.error(`❌ Failed to create tool wrapper for ${toolMeta.toolId}:`, error); + } + } + + console.log(`🔧 ${Object.keys(tools).length} tools available for this request`); + // Fetch recent messages for context const recentMessages = await prisma.omegaMessage.findMany({ where: { conversationId }, @@ -137,11 +276,35 @@ export async function POST(request: NextRequest, context: RouteContext): Promise // Build AI SDK messages const messages: ModelMessage[] = []; - // Add system prompt - const systemPrompt = buildSystemPrompt({ + // Build tool list for system prompt + const toolsList = Object.entries(tools) + .map(([name, t]) => { + const tool = t as { description?: string }; + return `- ${name}: ${tool.description || 'No description'}`; + }) + .join('\n'); + + // Add system prompt with available tools + const baseSystemPrompt = buildSystemPrompt({ customSystemPrompt: userSettings?.customSystemPrompt, pinnedToolIds: userSettings?.pinnedToolIds || [], }); + + const systemPrompt = `${baseSystemPrompt} + +## Available Tools + +The following tools have been automatically loaded based on the user's request. Use them directly to accomplish the task: + +${toolsList || 'No tools matched this query. Try to help the user with general knowledge.'} + +## Instructions + +1. If a tool is available that can help, USE IT immediately +2. Don't describe what tools could do - actually call them +3. After calling a tool, explain the results to the user +4. If no tools match, help the user with general knowledge`; + messages.push({ role: 'system', content: systemPrompt }); // Add conversation history @@ -204,12 +367,6 @@ export async function POST(request: NextRequest, context: RouteContext): Promise // Add new user message messages.push({ role: 'user', content: parsed.data.message }); - // Build Omega tools using published TPMJS packages - const tools = { - registrySearch: registrySearchTool, - registryExecute: registryExecuteTool, - }; - // Get the provider model (using OpenAI by default) const { createOpenAI } = await import('@ai-sdk/openai'); const apiKey = process.env.OPENAI_API_KEY; @@ -253,40 +410,34 @@ export async function POST(request: NextRequest, context: RouteContext): Promise let inputTokens = 0; let outputTokens = 0; - // Stream the response with up to 10 tool call iterations - const result = await streamText({ + const result = streamText({ model, messages, - tools, - stopWhen: stepCountIs(10), - onChunk: async ({ chunk }) => { - if (chunk.type === 'tool-call') { - const input = 'args' in chunk ? chunk.args : chunk.input; - - console.log('[Omega] Tool call:', { - toolName: chunk.toolName, - toolCallId: chunk.toolCallId, - }); - - toolCallsMap.set(chunk.toolCallId, { - toolCallId: chunk.toolCallId, - toolName: chunk.toolName, - args: input, - }); + tools: Object.keys(tools).length > 0 ? tools : undefined, + stopWhen: stepCountIs(5), // Allow up to 5 tool calls + onChunk: async (chunk) => { + // Handle tool call (complete tool call with args) + if (chunk.chunk.type === 'tool-call') { + const input = + 'input' in chunk.chunk + ? chunk.chunk.input + : 'args' in chunk.chunk + ? (chunk.chunk as { args: unknown }).args + : {}; // Create tool run record await prisma.omegaToolRun.create({ data: { conversationId, - toolName: chunk.toolName, + toolName: chunk.chunk.toolName, input: input as Prisma.InputJsonValue, status: 'running', }, }); sendEvent('run.step.tool.started', { - toolCallId: chunk.toolCallId, - toolName: chunk.toolName, + toolCallId: chunk.chunk.toolCallId, + toolName: chunk.chunk.toolName, input, }); } @@ -404,37 +555,48 @@ export async function POST(request: NextRequest, context: RouteContext): Promise }); } - // Update conversation - const updateData: Prisma.OmegaConversationUpdateInput = { - executionState: 'idle', - inputTokensTotal: { increment: inputTokens }, - outputTokensTotal: { increment: outputTokens }, - updatedAt: new Date(), - }; - - // Auto-generate title from first message if not set - if (conversation.title === null) { - updateData.title = parsed.data.message.slice(0, 100); - } - + // Update conversation token totals await prisma.omegaConversation.update({ where: { id: conversationId }, - data: updateData, + data: { + executionState: 'idle', + inputTokensTotal: { increment: inputTokens }, + outputTokensTotal: { increment: outputTokens }, + }, }); - const executionTimeMs = Date.now() - startTime; + // Update title if this is the first message pair + const messageCount = await prisma.omegaMessage.count({ where: { conversationId } }); + if (messageCount <= 3 && !conversation.title) { + // Use first 50 chars of user message as title + const title = + parsed.data.message.slice(0, 50) + (parsed.data.message.length > 50 ? '...' : ''); + await prisma.omegaConversation.update({ + where: { id: conversationId }, + data: { title }, + }); + } sendEvent('run.completed', { messageId: assistantMessage.id, - conversationId, inputTokens, outputTokens, - executionTimeMs, + toolCallCount: allToolCalls.length, + toolsLoaded: Object.keys(tools).length, + loadedToolNames: Object.keys(tools), + bm25Results: relevantTools.map((t) => ({ + toolId: t.toolId, + name: t.name, + packageName: t.packageName, + description: t.description, + })), }); - } catch (error) { - console.error('[Omega] Stream error:', error); - // Update conversation state + controller.close(); + } catch (error) { + console.error('Stream error:', error); + + // Update conversation state on error await prisma.omegaConversation.update({ where: { id: conversationId }, data: { executionState: 'idle' }, @@ -443,13 +605,12 @@ export async function POST(request: NextRequest, context: RouteContext): Promise sendEvent('run.failed', { error: error instanceof Error ? error.message : 'Unknown error', }); - } finally { controller.close(); } }, }); - return new NextResponse(stream, { + return new Response(stream, { headers: { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', @@ -457,7 +618,7 @@ export async function POST(request: NextRequest, context: RouteContext): Promise }, }); } catch (error) { - console.error('[Omega] Failed to process message:', error); + console.error('Message handler error:', error); // Reset conversation state await prisma.omegaConversation.update({ @@ -466,10 +627,7 @@ export async function POST(request: NextRequest, context: RouteContext): Promise }); return NextResponse.json( - { - success: false, - error: error instanceof Error ? error.message : 'Failed to process message', - }, + { success: false, error: error instanceof Error ? error.message : 'Unknown error' }, { status: 500 } ); } diff --git a/apps/web/src/app/omega/[conversationId]/page.tsx b/apps/web/src/app/omega/[conversationId]/page.tsx index a57ed19..3f76f09 100644 --- a/apps/web/src/app/omega/[conversationId]/page.tsx +++ b/apps/web/src/app/omega/[conversationId]/page.tsx @@ -7,7 +7,6 @@ import { Textarea } from '@tpmjs/ui/Textarea/Textarea'; import Link from 'next/link'; import { useParams, useRouter } from 'next/navigation'; import { useCallback, useEffect, useRef, useState } from 'react'; -import { Virtuoso, type VirtuosoHandle } from 'react-virtuoso'; import { Streamdown } from 'streamdown'; import { AppHeader } from '~/components/AppHeader'; @@ -159,6 +158,7 @@ function ToolCallCard({ /** * Omega Chat Page */ +// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: Complex chat page with multiple UI states and SSE handling export default function OmegaChatPage(): React.ReactElement { const params = useParams(); const router = useRouter(); @@ -174,11 +174,17 @@ export default function OmegaChatPage(): React.ReactElement { const [toolCalls, setToolCalls] = useState([]); const [expandedToolCalls, setExpandedToolCalls] = useState>(new Set()); const [viewMode, setViewMode] = useState<'chat' | 'debug'>('chat'); + const [lastRunInfo, setLastRunInfo] = useState<{ + loadedToolNames?: string[]; + bm25Results?: Array<{ + toolId: string; + name: string; + packageName: string; + description: string; + }>; + } | null>(null); - // Track first item index for prepending (Virtuoso pattern) - const [firstItemIndex, setFirstItemIndex] = useState(10000); - - const virtuosoRef = useRef(null); + const messagesContainerRef = useRef(null); const inputRef = useRef(null); const toggleToolCall = (toolCallId: string) => { @@ -214,7 +220,6 @@ export default function OmegaChatPage(): React.ReactElement { const { messages: messageList, ...conversationData } = data.data; setConversation(conversationData); setMessages(messageList || []); - setFirstItemIndex(10000); } catch (err) { console.error('Failed to fetch conversation:', err); setError(err instanceof Error ? err.message : 'Failed to fetch conversation'); @@ -227,6 +232,14 @@ export default function OmegaChatPage(): React.ReactElement { fetchConversation(); }, [fetchConversation]); + // Auto-scroll to bottom when messages change or streaming + // biome-ignore lint/correctness/useExhaustiveDependencies: We intentionally trigger scroll when messages/streamingContent change + useEffect(() => { + if (messagesContainerRef.current) { + messagesContainerRef.current.scrollTop = messagesContainerRef.current.scrollHeight; + } + }, [messages, streamingContent]); + // Check for initial prompt from landing page // biome-ignore lint/correctness/useExhaustiveDependencies: Only run on mount and when messages load useEffect(() => { @@ -332,6 +345,11 @@ export default function OmegaChatPage(): React.ReactElement { ); break; case 'run.completed': + // Capture tool info for debug view + setLastRunInfo({ + loadedToolNames: data.loadedToolNames, + bm25Results: data.bm25Results, + }); // Refresh messages await fetchConversation(); setStreamingContent(''); @@ -477,31 +495,81 @@ export default function OmegaChatPage(): React.ReactElement { {/* Debug JSON View */} {viewMode === 'debug' && (
-
-
+
+ {/* BM25 Search Results */} + {lastRunInfo?.bm25Results && lastRunInfo.bm25Results.length > 0 && (

- Raw Messages Array ({messages.length} messages) + BM25 Search Results ({lastRunInfo.bm25Results.length} tools found)

-

- Messages are ordered by createdAt. Each message includes role, content, and tool - call data. +

+ Tools discovered from the registry based on the user's message.

+
+ {lastRunInfo.bm25Results.map((tool, i) => ( +
+ {i + 1}.{' '} + {tool.packageName} + :: + {tool.name} +

{tool.description}

+
+ ))} +
- + )} + + {/* Loaded Tool Names */} + {lastRunInfo?.loadedToolNames && lastRunInfo.loadedToolNames.length > 0 && ( +
+

+ Loaded Tools ({lastRunInfo.loadedToolNames.length} tools) +

+

+ Sanitized tool names available to the AI for this request. +

+
+
+ {lastRunInfo.loadedToolNames.map((name) => ( + + {name} + + ))} +
+
+
+ )} + + {/* Messages Array */} +
+
+
+

+ Raw Messages Array ({messages.length} messages) +

+

+ Messages are ordered by createdAt. Each message includes role, content, and + tool call data. +

+
+ +
+ 
+                  {JSON.stringify(messages, null, 2)}
+
- 
-                {JSON.stringify(messages, null, 2)}
-
)} @@ -527,103 +595,101 @@ export default function OmegaChatPage(): React.ReactElement {
) : ( - ( -
- {/* Live tool calls during streaming */} - {toolCalls.length > 0 && ( -
- {toolCalls.map((tc) => ( -
-
- toggleToolCall(tc.toolCallId)} - /> -
-
- ))} -
- )} - - {streamingContent && ( -
-
-
- {streamingContent} -
- +
+
+ {/* Messages */} + {messages.map((message) => ( +
+ {/* USER message */} + {message.role === 'USER' && ( +
+
+
{message.content}
)} - {isSending && !streamingContent && toolCalls.length === 0 && ( + {/* ASSISTANT message */} + {message.role === 'ASSISTANT' && message.content && (
-
-
- - Omega is thinking... +
+
+ {message.content}
+ {/* Token usage for debugging */} + {(message.inputTokens || message.outputTokens) && ( +
+ {message.inputTokens && In: {message.inputTokens}} + {message.inputTokens && message.outputTokens && | } + {message.outputTokens && Out: {message.outputTokens}} +
+ )} +
+
+ )} + + {/* TOOL message - show as collapsed tool result */} + {message.role === 'TOOL' && ( +
+
+
+ Tool Results +
+ 
+                                {message.content}
+
)}
- ), - }} - itemContent={(_index, message) => ( -
- {/* USER message */} - {message.role === 'USER' && ( -
-
-
{message.content}
-
-
- )} + ))} - {/* ASSISTANT message */} - {message.role === 'ASSISTANT' && message.content && ( + {/* Live tool calls during streaming */} + {toolCalls.length > 0 && ( +
+ {toolCalls.map((tc) => ( +
+
+ toggleToolCall(tc.toolCallId)} + /> +
+
+ ))} +
+ )} + + {/* Streaming content */} + {streamingContent && ( +
- {message.content} + {streamingContent}
- {/* Token usage for debugging */} - {(message.inputTokens || message.outputTokens) && ( -
- {message.inputTokens && In: {message.inputTokens}} - {message.inputTokens && message.outputTokens && | } - {message.outputTokens && Out: {message.outputTokens}} -
- )} +
- )} +
+ )} - {/* TOOL message - show as collapsed tool result */} - {message.role === 'TOOL' && ( + {/* Thinking indicator */} + {isSending && !streamingContent && toolCalls.length === 0 && ( +
-
-
- Tool Results +
+
+ + Omega is thinking...
- 
-                              {message.content}
-
- )} -
- )} - /> +
+ )} +
+
)}
diff --git a/apps/web/src/lib/omega/system-prompt.ts b/apps/web/src/lib/omega/system-prompt.ts index 89aabab..112006a 100644 --- a/apps/web/src/lib/omega/system-prompt.ts +++ b/apps/web/src/lib/omega/system-prompt.ts @@ -4,49 +4,32 @@ * Defines the default behavior for the Omega AI agent. */ -export const OMEGA_SYSTEM_PROMPT = `You are Omega, an AI assistant with access to tools from the TPMJS registry - a collection of AI-ready tools that can help you complete tasks. +export const OMEGA_SYSTEM_PROMPT = `You are Omega, an AI assistant powered by the TPMJS tool registry - a collection of 1M+ AI-ready tools. -## Your Capabilities +## How It Works -You have access to two special tools: +When you receive a message, relevant tools are automatically loaded based on the user's request. You'll see a list of available tools in each response - USE THEM. -1. **registrySearch** - Search the TPMJS tool registry to find tools that can help complete a task. Returns tool metadata including toolId for use with registryExecute. -2. **registryExecute** - Execute a tool from the registry by its toolId with the required parameters. Tools run in a secure sandbox. +## Your Job -## How to Help Users - -When a user asks you to complete a task: - -1. **Analyze the request** - Understand what the user wants to accomplish -2. **Search for tools** - Use registrySearch to find relevant tools that can help -3. **Select the best tool(s)** - Choose the most appropriate tool(s) based on the search results and their toolIds -4. **Execute tools** - Use registryExecute to run the selected tools with the correct parameters (toolId and params) -5. **Synthesize results** - Combine tool outputs into a helpful, clear response +1. **Look at the available tools** - They've been selected based on what the user asked for +2. **Call the appropriate tool(s)** - Don't just describe what they do, actually use them +3. **Explain the results** - After a tool returns, summarize what happened for the user ## Best Practices -- **Be transparent** - Always explain which tools you're using and why -- **Handle errors gracefully** - If a tool fails, explain what went wrong and try alternatives -- **Validate inputs** - Make sure you have all required parameters before executing a tool -- **Iterate when needed** - Some tasks may require multiple tool executions -- **Ask for clarification** - If the user's request is ambiguous, ask questions before proceeding +- **Take action** - If a tool can help, call it immediately +- **Be transparent** - Tell the user which tool you're using +- **Handle errors** - If a tool fails, explain and try an alternative +- **Ask for clarity** - If you need more info, ask before proceeding -## Tool Search Tips +## Response Style -When searching for tools: -- Use descriptive keywords related to the task (e.g., "web scraping", "image processing", "API call") -- If you don't find the right tool, try different search terms -- Consider the tool's input schema to ensure you can provide the required parameters +- Keep responses concise and helpful +- Present tool outputs in a clear, readable format +- Offer to do more if the user might need it -## Response Format - -When presenting results: -- Summarize what you did and what tools you used -- Present the output in a clear, readable format -- If the output is large, highlight the most relevant parts -- Offer to do more or explain further if needed - -Remember: You're here to help users accomplish tasks efficiently by leveraging the vast TPMJS tool ecosystem. Be helpful, be clear, and be thorough.`; +Remember: Your value is in EXECUTING tools to get real results, not describing what tools could do.`; /** * Generate a custom system prompt with user preferences diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index da75409..0d956b4 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -266,12 +266,6 @@ importers: '@tpmjs/package-executor': specifier: workspace:* version: link:../../packages/package-executor - '@tpmjs/registry-execute': - specifier: workspace:* - version: link:../../packages/tools/registryExecute - '@tpmjs/registry-search': - specifier: workspace:* - version: link:../../packages/tools/registrySearch '@tpmjs/types': specifier: workspace:* version: link:../../packages/types