code-graph-context 2.9.0 → 2.10.0

package/dist/mcp/service-init.js

@@ -5,6 +5,7 @@
 import fs from 'fs/promises';
 import { join } from 'path';
 import { ensureNeo4jRunning, isDockerInstalled, isDockerRunning } from '../cli/neo4j-docker.js';
+import { isOpenAIEnabled, getEmbeddingDimensions } from '../core/embeddings/embeddings.service.js';
 import { Neo4jService, QUERIES } from '../storage/neo4j/neo4j.service.js';
 import { FILE_PATHS, LOG_CONFIG } from './constants.js';
 import { initializeNaturalLanguageService } from './tools/natural-language-to-cypher.tool.js';
@@ -13,12 +14,32 @@ import { debugLog } from './utils.js';
  * Log startup warnings for missing configuration
  */
 const checkConfiguration = async () => {
-    if (!process.env.OPENAI_API_KEY) {
+    const openai = isOpenAIEnabled();
+    const dims = getEmbeddingDimensions();
+    const provider = openai ? 'openai' : 'local';
+    console.error(JSON.stringify({
+        level: 'info',
+        message: `[code-graph-context] Embedding provider: ${provider} (${dims} dimensions)`,
+    }));
+    await debugLog('Embedding configuration', { provider, dimensions: dims });
+    if (openai && !process.env.OPENAI_API_KEY) {
         console.error(JSON.stringify({
             level: 'warn',
-            message: '[code-graph-context] OPENAI_API_KEY not set. Semantic search and NL queries unavailable.',
+            message: '[code-graph-context] OPENAI_ENABLED=true but OPENAI_API_KEY not set. Embedding calls will fail.',
         }));
-        await debugLog('Configuration warning', { warning: 'OPENAI_API_KEY not set' });
+        await debugLog('Configuration warning', { warning: 'OPENAI_ENABLED=true but OPENAI_API_KEY not set' });
+    }
+    if (!openai) {
+        console.error(JSON.stringify({
+            level: 'info',
+            message: '[code-graph-context] Using local embeddings (Python sidecar). Starts on first embedding request.',
+        }));
+        if (!process.env.OPENAI_API_KEY) {
+            console.error(JSON.stringify({
+                level: 'info',
+                message: '[code-graph-context] natural_language_to_cypher requires OPENAI_API_KEY and is unavailable.',
+            }));
+        }
     }
 };
 /**

package/dist/mcp/tools/search-codebase.tool.js

@@ -39,6 +39,12 @@ export const createSearchCodebaseTool = (server) => {
                 .optional()
                 .describe(`Length of code snippets to include (default: ${DEFAULTS.codeSnippetLength})`)
                 .default(DEFAULTS.codeSnippetLength),
+            topK: z
+                .number()
+                .int()
+                .optional()
+                .describe('Number of top vector matches to return (default: 3, max: 10). The best match is traversed; others shown as alternatives.')
+                .default(3),
             minSimilarity: z
                 .number()
                 .optional()
@@ -50,7 +56,7 @@ export const createSearchCodebaseTool = (server) => {
                 .describe('Use weighted traversal strategy that scores each node for relevance (default: false)')
                 .default(true),
         },
-    }, async ({ projectId, query, maxDepth = DEFAULTS.traversalDepth, maxNodesPerChain = 5, skip = 0, includeCode = true, snippetLength = DEFAULTS.codeSnippetLength, minSimilarity = 0.65, useWeightedTraversal = true, }) => {
+    }, async ({ projectId, query, maxDepth = DEFAULTS.traversalDepth, maxNodesPerChain = 5, skip = 0, includeCode = true, snippetLength = DEFAULTS.codeSnippetLength, topK = 3, minSimilarity = 0.65, useWeightedTraversal = true, }) => {
         const neo4jService = new Neo4jService();
         try {
             // Resolve project ID from name, path, or ID
@@ -63,11 +69,12 @@ export const createSearchCodebaseTool = (server) => {
             const sanitizedMaxNodesPerChain = sanitizeNumericInput(maxNodesPerChain, 5);
             const sanitizedSkip = sanitizeNumericInput(skip, 0);
             const sanitizedSnippetLength = sanitizeNumericInput(snippetLength, DEFAULTS.codeSnippetLength);
+            const sanitizedTopK = sanitizeNumericInput(topK, 3, 10);
             const embeddingsService = new EmbeddingsService();
             const traversalHandler = new TraversalHandler(neo4jService);
             const embedding = await embeddingsService.embedText(query);
             const vectorResults = await neo4jService.run(QUERIES.VECTOR_SEARCH, {
-                limit: 1,
+                limit: sanitizedTopK,
                 embedding,
                 projectId: resolvedProjectId,
                 fetchMultiplier: 10,
@@ -77,28 +84,45 @@ export const createSearchCodebaseTool = (server) => {
                 return createSuccessResponse(`No code found with similarity >= ${minSimilarity}. ` +
                     `Try rephrasing your query or lowering the minSimilarity threshold. Query: "${query}"`);
             }
-            const startNode = vectorResults[0].node;
-            const nodeId = startNode.properties.id;
-            const similarityScore = vectorResults[0].score;
-            // Check if best match meets threshold - prevents traversing low-relevance results
-            if (similarityScore < minSimilarity) {
-                return createSuccessResponse(`No sufficiently relevant code found. Best match score: ${similarityScore.toFixed(3)} ` +
+            // Filter results that meet the similarity threshold
+            const qualifiedResults = vectorResults.filter((r) => r.score >= minSimilarity);
+            if (qualifiedResults.length === 0) {
+                const bestScore = vectorResults[0].score;
+                return createSuccessResponse(`No sufficiently relevant code found. Best match score: ${bestScore.toFixed(3)} ` +
                     `(threshold: ${minSimilarity}). Try rephrasing your query.`);
             }
-            // Include similarity score in the title so users can see relevance
-            const scoreDisplay = typeof similarityScore === 'number' ? similarityScore.toFixed(3) : 'N/A';
-            return await traversalHandler.traverseFromNode(nodeId, embedding, {
+            // Best match — traverse from this node
+            const bestMatch = qualifiedResults[0];
+            const nodeId = bestMatch.node.properties.id;
+            const bestScore = bestMatch.score.toFixed(3);
+            // Build alternative matches summary for the response
+            const alternatives = qualifiedResults.slice(1);
+            const altLines = alternatives.map((r) => {
+                const props = r.node.properties;
+                const name = props.name ?? props.id;
+                const file = props.filePath ? ` (${props.filePath})` : '';
+                return `  - ${name}${file} [score: ${r.score.toFixed(3)}, id: ${props.id}]`;
+            });
+            const altSection = altLines.length > 0
+                ? `\n\nAlternative matches (use traverse_from_node to explore):\n${altLines.join('\n')}`
+                : '';
+            const traversalResult = await traversalHandler.traverseFromNode(nodeId, embedding, {
                 projectId: resolvedProjectId,
                 maxDepth: sanitizedMaxDepth,
-                direction: 'BOTH', // Show both incoming (who calls this) and outgoing (what this calls)
+                direction: 'BOTH',
                 includeCode,
                 maxNodesPerChain: sanitizedMaxNodesPerChain,
                 skip: sanitizedSkip,
                 summaryOnly: false,
                 snippetLength: sanitizedSnippetLength,
-                title: `Search Results (similarity: ${scoreDisplay}) - Starting from: ${nodeId}`,
+                title: `Search Results (${qualifiedResults.length} matches, best: ${bestScore}) - Traversing from: ${nodeId}`,
                 useWeightedTraversal,
             });
+            // Append alternatives to the traversal response
+            if (altSection && traversalResult.content?.[0]?.type === 'text') {
+                traversalResult.content[0].text += altSection;
+            }
+            return traversalResult;
         }
         catch (error) {
             console.error('Search codebase error:', error);

package/dist/mcp/tools/session-note.tool.js

@@ -3,7 +3,7 @@
  * Save and recall cross-session observations, decisions, and insights
  */
 import { z } from 'zod';
-import { EmbeddingsService } from '../../core/embeddings/embeddings.service.js';
+import { EmbeddingsService, getEmbeddingDimensions } from '../../core/embeddings/embeddings.service.js';
 import { Neo4jService, QUERIES } from '../../storage/neo4j/neo4j.service.js';
 import { TOOL_NAMES, TOOL_METADATA } from '../constants.js';
 import { createErrorResponse, createSuccessResponse, resolveProjectIdOrError, debugLog } from '../utils.js';
@@ -31,17 +31,16 @@ const CREATE_SESSION_NOTE_QUERY = `
   // Link to referenced code nodes (filter out internal coordination nodes)
   WITH n
   UNWIND CASE WHEN size($aboutNodeIds) = 0 THEN [null] ELSE $aboutNodeIds END AS aboutNodeId
-  WITH n, aboutNodeId
-  WHERE aboutNodeId IS NOT NULL
   OPTIONAL MATCH (target)
-  WHERE target.id = aboutNodeId
+  WHERE aboutNodeId IS NOT NULL
+    AND target.id = aboutNodeId
     AND target.projectId = $projectId
     AND NOT target:SessionNote
     AND NOT target:SessionBookmark
     AND NOT target:Pheromone
     AND NOT target:SwarmTask
   WITH n, collect(target) AS targets
-  FOREACH (t IN targets | MERGE (n)-[:ABOUT]->(t))
+  FOREACH (t IN [x IN targets WHERE x IS NOT NULL] | MERGE (n)-[:ABOUT]->(t))
 
   // Link to the latest SessionBookmark for this session (if one exists)
   WITH n
@@ -194,7 +193,7 @@ export const createSaveSessionNoteTool = (server) => {
             // Ensure vector index exists (idempotent — IF NOT EXISTS)
             let hasEmbedding = false;
             try {
-                await neo4jService.run(QUERIES.CREATE_SESSION_NOTES_VECTOR_INDEX);
+                await neo4jService.run(QUERIES.CREATE_SESSION_NOTES_VECTOR_INDEX(getEmbeddingDimensions()));
                 const embeddingsService = new EmbeddingsService();
                 const embeddingText = `${topic}\n\n${content}`;
                 const embedding = await embeddingsService.embedText(embeddingText);

package/dist/storage/neo4j/neo4j.service.js

@@ -136,19 +136,19 @@ export const QUERIES = {
     RETURN labels(n)[0] as nodeType, count(*) as count
     ORDER BY count DESC
   `,
-    CREATE_EMBEDDED_VECTOR_INDEX: `
+    CREATE_EMBEDDED_VECTOR_INDEX: (dimensions) => `
   CREATE VECTOR INDEX embedded_nodes_idx IF NOT EXISTS
   FOR (n:Embedded) ON (n.embedding)
   OPTIONS {indexConfig: {
-    \`vector.dimensions\`: 3072,
+    \`vector.dimensions\`: ${dimensions},
     \`vector.similarity_function\`: 'cosine'
   }}
 `,
-    CREATE_SESSION_NOTES_VECTOR_INDEX: `
+    CREATE_SESSION_NOTES_VECTOR_INDEX: (dimensions) => `
   CREATE VECTOR INDEX session_notes_idx IF NOT EXISTS
   FOR (n:SessionNote) ON (n.embedding)
   OPTIONS {indexConfig: {
-    \`vector.dimensions\`: 3072,
+    \`vector.dimensions\`: ${dimensions},
     \`vector.similarity_function\`: 'cosine'
   }}
 `,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "code-graph-context",
-  "version": "2.9.0",
+  "version": "2.10.0",
   "description": "MCP server that builds code graphs to provide rich context to LLMs",
   "type": "module",
   "homepage": "https://github.com/drewdrewH/code-graph-context#readme",
@@ -34,6 +34,8 @@
   },
   "files": [
     "dist/**/*",
+    "sidecar/embedding_server.py",
+    "sidecar/requirements.txt",
     "README.md",
     "LICENSE",
     ".env.example"

package/sidecar/embedding_server.py

@@ -0,0 +1,147 @@
+"""
+Local embedding server for code-graph-context.
+Uses Qodo-Embed-1-1.5B for high-quality code embeddings without OpenAI dependency.
+Runs as a sidecar process managed by the Node.js MCP server.
+"""
+
+import gc
+import os
+import sys
+import signal
+import logging
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    stream=sys.stderr,
+)
+logger = logging.getLogger("embedding-sidecar")
+
+app = FastAPI(title="code-graph-context embedding sidecar")
+
+model = None
+model_name = os.environ.get("EMBEDDING_MODEL", "Qodo/Qodo-Embed-1-1.5B")
+
+
+class EmbedRequest(BaseModel):
+    texts: list[str]
+    batch_size: int = 8
+
+
+class EmbedResponse(BaseModel):
+    embeddings: list[list[float]]
+    dimensions: int
+    model: str
+
+
+@app.on_event("startup")
+def load_model():
+    global model
+    try:
+        import torch
+        from sentence_transformers import SentenceTransformer
+
+        device = "mps" if torch.backends.mps.is_available() else "cpu"
+        logger.info(f"Loading {model_name} on {device}...")
+        model = SentenceTransformer(model_name, device=device)
+
+        # Warm up with a test embedding
+        test = model.encode(["warmup"], show_progress_bar=False)
+        dims = len(test[0])
+        logger.info(f"Model loaded: {dims} dimensions, device={device}")
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+
+@app.get("/health")
+def health():
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    sample = model.encode(["dim_check"], show_progress_bar=False)
+    return {
+        "status": "ok",
+        "model": model_name,
+        "dimensions": len(sample[0]),
+        "device": str(model.device),
+    }
+
+
+@app.post("/embed", response_model=EmbedResponse)
+async def embed(req: EmbedRequest):
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    if not req.texts:
+        return EmbedResponse(embeddings=[], dimensions=0, model=model_name)
+
+    try:
+        embeddings = _encode_with_oom_fallback(req.texts, req.batch_size)
+        dims = len(embeddings[0])
+        return EmbedResponse(
+            embeddings=embeddings,
+            dimensions=dims,
+            model=model_name,
+        )
+    except Exception as e:
+        logger.error(f"Embedding error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def _encode_with_oom_fallback(texts: list[str], batch_size: int) -> list[list[float]]:
+    """
+    Encode texts, falling back to CPU if MPS runs out of memory.
+    Also retries with smaller batch sizes before giving up.
+    """
+    import torch
+
+    try:
+        result = model.encode(
+            texts,
+            batch_size=batch_size,
+            show_progress_bar=False,
+            normalize_embeddings=True,
+        )
+        return result.tolist()
+    except (torch.mps.OutOfMemoryError, RuntimeError) as e:
+        if "out of memory" not in str(e).lower():
+            raise
+
+        logger.warning(f"MPS OOM with batch_size={batch_size}, len={len(texts)}. Falling back to CPU.")
+
+        # Free MPS memory
+        if hasattr(torch.mps, "empty_cache"):
+            torch.mps.empty_cache()
+        gc.collect()
+
+        # Fall back to CPU for this request
+        original_device = model.device
+        model.to("cpu")
+
+        try:
+            # Use smaller batches on CPU
+            cpu_batch = min(batch_size, 4)
+            result = model.encode(
+                texts,
+                batch_size=cpu_batch,
+                show_progress_bar=False,
+                normalize_embeddings=True,
+            )
+            return result.tolist()
+        finally:
+            # Move back to MPS for future requests
+            try:
+                model.to(original_device)
+            except Exception:
+                logger.warning("Could not move model back to MPS, staying on CPU")
+
+
+def handle_signal(sig, _frame):
+    logger.info(f"Received signal {sig}, shutting down")
+    sys.exit(0)
+
+
+signal.signal(signal.SIGTERM, handle_signal)

package/sidecar/requirements.txt

@@ -0,0 +1,5 @@
+fastapi>=0.104.0
+uvicorn>=0.24.0
+sentence-transformers>=3.0.0
+torch>=2.0.0
+pydantic>=2.0.0