@rigstate/mcp 0.4.6 → 0.5.1

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@rigstate/mcp",
-    "version": "0.4.6",
+    "version": "0.5.1",
     "description": "Rigstate MCP Server - Model Context Protocol for AI Editors",
     "type": "module",
     "main": "./dist/index.js",
@@ -15,9 +15,12 @@
         "test": "node dist/index.js --help"
     },
     "dependencies": {
+        "@ai-sdk/google": "^3.0.13",
+        "@google/generative-ai": "^0.24.1",
         "@modelcontextprotocol/sdk": "^1.0.0",
-        "@supabase/supabase-js": "^2.39.0",
         "@rigstate/rules-engine": "0.1.0",
+        "@supabase/supabase-js": "^2.39.0",
+        "ai": "^6.0.5",
         "dotenv": "^17.2.3",
         "uuid": "^9.0.0",
         "zod": "^3.22.4"

package/src/lib/schemas.ts

@@ -8,7 +8,7 @@ export const QueryBrainInputSchema = z.object({
     projectId: z.string().uuid('Invalid project ID'),
     query: z.string().min(1, 'Query is required'),
     limit: z.number().min(1).max(20).optional().default(8),
-    threshold: z.number().min(0).max(1).optional().default(0.5)
+    threshold: z.number().min(0).max(1).optional().default(0.1)
 });
 
 export const GetProjectContextInputSchema = z.object({

package/src/tools/query-brain.ts

@@ -34,15 +34,27 @@ architecture rules, decisions, and constraints.`,
     }
 });
 
-// Simple embedding generation using text normalization
-async function generateQueryEmbedding(query: string): Promise<number[]> {
-    // For MCP server, we'll use the match_memories_text RPC if available,
-    // or fall back to keyword-based search.
-    // This is a placeholder - the actual embedding should be done via the Supabase edge function
-    // or passed through to the main Rigstate API.
-
-    // For now, return null to trigger keyword fallback
-    return [];
+// Simple embedding generation using Google Gemini
+async function generateQueryEmbedding(query: string): Promise<number[] | null> {
+    const apiKey = process.env.GOOGLE_GENERATIVE_AI_API_KEY;
+    if (!apiKey) {
+        console.warn('GOOGLE_GENERATIVE_AI_API_KEY not found, skipping vector search.');
+        return null;
+    }
+
+    try {
+        const { google } = await import('@ai-sdk/google');
+        const { embed } = await import('ai');
+
+        const { embedding } = await embed({
+            model: google.embedding('text-embedding-004'),
+            value: query.replace(/\n/g, ' '),
+        });
+        return embedding;
+    } catch (error) {
+        console.error('Failed to generate embedding for search:', error);
+        return null;
+    }
 }
 
 export async function queryBrain(
@@ -67,45 +79,24 @@ export async function queryBrain(
 
     // Try semantic search first using the match_memories RPC
     // This requires the embedding to be generated, so we'll try
-    // a text-based search as fallback
+    // Generate embedding if possible for semantic search
+    const embedding = await generateQueryEmbedding(query);
     let memories: MemoryRecord[] = [];
 
-    // Attempt keyword-based search using ilike on content
-    // This is a simpler approach that works without embeddings
-    const searchTerms = query.toLowerCase().split(/\s+/).filter(t => t.length > 2);
-
-    if (searchTerms.length > 0) {
-        // Build OR condition for fuzzy matching
-        const orConditions = searchTerms.map(term => `content.ilike.%${term}%`).join(',');
-
-        const { data: keywordMatches, error: searchError } = await supabase
-            .from('project_memories')
-            .select('id, content, category, tags, importance, created_at')
-            .eq('project_id', projectId)
-            .eq('is_active', true)
-            .or(orConditions)
-            .order('importance', { ascending: false, nullsFirst: false })
-            .limit(limit);
-
-        if (searchError) {
-            console.error('Search error:', searchError);
-        }
-
-        if (keywordMatches) {
-            memories = keywordMatches.map(m => ({
-                id: m.id,
-                content: m.content,
-                category: m.category || 'general',
-                tags: m.tags || [],
-                netVotes: m.importance || 0,
-                createdAt: m.created_at
-            }));
-        }
-    }
-
-    // If no keyword matches, fetch recent memories
-    if (memories.length === 0) {
-        const { data: recentMemories, error: recentError } = await supabase
+    // Use the hybrid search RPC
+    const { data: searchResults, error: searchError } = await supabase
+        .rpc('hybrid_search_memories', {
+            p_project_id: projectId,
+            p_query: query,
+            p_embedding: embedding,
+            p_limit: limit,
+            p_similarity_threshold: threshold || 0.2
+        });
+
+    if (searchError) {
+        console.error('Hybrid search error:', searchError);
+        // Fallback to basic recent fetch if RPC fails
+        const { data: recentMemories } = await supabase
             .from('project_memories')
             .select('id, content, category, tags, importance, created_at')
             .eq('project_id', projectId)
@@ -113,10 +104,6 @@ export async function queryBrain(
             .order('created_at', { ascending: false })
             .limit(limit);
 
-        if (recentError) {
-            console.error('Recent error:', recentError);
-        }
-
         if (recentMemories) {
             memories = recentMemories.map(m => ({
                 id: m.id,
@@ -127,17 +114,30 @@ export async function queryBrain(
                 createdAt: m.created_at
             }));
         }
+    } else if (searchResults) {
+        memories = searchResults.map((m: any) => ({
+            id: m.id,
+            content: m.content,
+            category: m.category,
+            tags: m.tags,
+            netVotes: m.importance,
+            createdAt: m.created_at
+        }));
     }
 
     // Format memories into a readable context block
     const contextLines = memories.map((m) => {
         const voteIndicator = m.netVotes && m.netVotes < 0 ? ` [⚠️ POORLY RATED: ${m.netVotes}]` : '';
-        const tagStr = m.tags.length > 0 ? ` [${m.tags.join(', ')}]` : '';
-        return `- [${m.category.toUpperCase()}]${tagStr}${voteIndicator}: ${m.content}`;
+        const tagStr = m.tags && m.tags.length > 0 ? ` [${m.tags.join(', ')}]` : '';
+        const category = m.category ? m.category.toUpperCase() : 'GENERAL';
+        return `- [${category}]${tagStr}${voteIndicator}: ${m.content}`;
     });
 
+    const searchType = embedding ? 'TRIPLE-HYBRID (Vector + FTS + Fuzzy)' : 'HYBRID (FTS + Fuzzy)';
+
     const formatted = memories.length > 0
         ? `=== PROJECT BRAIN: RELEVANT MEMORIES ===
+Search Mode: ${searchType}
 Query: "${query}"
 Found ${memories.length} relevant memories: