@steno-ai/engine 0.1.11 → 0.1.12

package/src/extraction/pipeline.ts

@@ -11,6 +11,7 @@ import type {
 } from './types.js';
 import { extractHeuristic } from './heuristic.js';
 import { extractWithLLM, normalizeEntityName } from './llm-extractor.js';
+import { extractCodebaseFacts } from './codebase-extractor.js';
 import { deduplicateFacts } from './dedup.js';
 import { processContradictions } from './contradiction.js';
 import { buildEntityIdMap, persistEdges } from './entity-extractor.js';
@@ -30,6 +31,8 @@ export interface PipelineConfig {
   decayNormalizationK?: number;
   entityTypes?: string[];
   domainEntityTypes?: DomainEntityType[];
+  /** Pipeline mode: 'conversation' (default) for chat/document memory, 'codebase' for code understanding */
+  mode?: 'conversation' | 'codebase';
 }
 
 // ---------------------------------------------------------------------------
@@ -89,87 +92,127 @@ async function executeExtraction(
 ): Promise<PipelineResult> {
   const tier = config.extractionTier ?? 'auto';
   const textContent = inputToText(input);
+  const isCodebaseMode = config.mode === 'codebase';
 
   let totalTokensInput = 0;
   let totalTokensOutput = 0;
   const tiersUsed: string[] = [];
 
-  // Run heuristic extraction (always)
-  const heuristicResult = extractHeuristic(textContent);
-  tiersUsed.push('heuristic');
+  let mergedFacts: ExtractedFact[];
+  let mergedEntities: ExtractedEntity[];
+  let mergedEdges: ExtractedEdge[];
 
-  let mergedFacts: ExtractedFact[] = heuristicResult.facts;
-  let mergedEntities: ExtractedEntity[] = heuristicResult.entities;
-  let mergedEdges: ExtractedEdge[] = heuristicResult.edges;
-
-  // Run LLM extraction (unless heuristic_only)
-  if (tier !== 'heuristic_only') {
+  if (isCodebaseMode) {
+    // ── CODEBASE MODE: Skip heuristic, use codebase-specific extraction ──
     const llmToUse = tier === 'smart_only' ? (config.smartLLM ?? config.cheapLLM) : config.cheapLLM;
     const llmTier = tier === 'smart_only' ? 'smart_llm' : 'cheap_llm';
 
-    // Fetch existing facts for the LLM to compare against (enables contradiction detection)
-    let existingFactsForLLM = input.existingFacts?.map(f => ({
-      lineageId: f.lineageId,
-      content: f.content,
-    }));
-
-    if (!existingFactsForLLM || existingFactsForLLM.length === 0) {
-      try {
-        const queryEmbedding = await config.embedding.embed(textContent.slice(0, 1000));
-        const similar = await config.storage.vectorSearch({
-          embedding: queryEmbedding,
-          tenantId: input.tenantId,
-          scope: input.scope,
-          scopeId: input.scopeId,
-          limit: 20,
-          minSimilarity: 0.3,
-          validOnly: true,
-        });
-        if (similar.length > 0) {
-          existingFactsForLLM = similar.map(s => ({
-            lineageId: s.fact.lineageId ?? s.fact.id,
-            content: s.fact.content,
-          }));
-        }
-      } catch {
-        // Non-fatal — extraction continues without existing facts
-      }
-    }
-
-    const llmResult = await extractWithLLM(
-      { llm: llmToUse, tier: llmTier, entityTypes: config.entityTypes, domainEntityTypes: config.domainEntityTypes },
+    const codebaseResult = await extractCodebaseFacts(
+      { llm: llmToUse, tier: llmTier },
       textContent,
-      existingFactsForLLM,
     );
 
-    totalTokensInput += llmResult.tokensInput;
-    totalTokensOutput += llmResult.tokensOutput;
+    totalTokensInput += codebaseResult.tokensInput;
+    totalTokensOutput += codebaseResult.tokensOutput;
     tiersUsed.push(llmTier);
 
     // Escalate to smart LLM if confidence is low
-    let finalLLMResult = llmResult;
-    if (
-      tier === 'auto' &&
-      llmResult.confidence < 0.6 &&
-      config.smartLLM
-    ) {
-      const smartResult = await extractWithLLM(
-        { llm: config.smartLLM, tier: 'smart_llm', entityTypes: config.entityTypes, domainEntityTypes: config.domainEntityTypes },
+    if (tier === 'auto' && codebaseResult.confidence < 0.6 && config.smartLLM) {
+      const smartResult = await extractCodebaseFacts(
+        { llm: config.smartLLM, tier: 'smart_llm' },
         textContent,
-        existingFactsForLLM,
       );
       totalTokensInput += smartResult.tokensInput;
       totalTokensOutput += smartResult.tokensOutput;
       tiersUsed.push('smart_llm');
-      finalLLMResult = smartResult;
+      mergedFacts = smartResult.facts;
+      mergedEntities = smartResult.entities;
+      mergedEdges = smartResult.edges;
+    } else {
+      mergedFacts = codebaseResult.facts;
+      mergedEntities = codebaseResult.entities;
+      mergedEdges = codebaseResult.edges;
     }
+  } else {
+    // ── CONVERSATION MODE: Original extraction logic ──
+
+    // Run heuristic extraction (always)
+    const heuristicResult = extractHeuristic(textContent);
+    tiersUsed.push('heuristic');
+
+    mergedFacts = heuristicResult.facts;
+    mergedEntities = heuristicResult.entities;
+    mergedEdges = heuristicResult.edges;
+
+    // Run LLM extraction (unless heuristic_only)
+    if (tier !== 'heuristic_only') {
+      const llmToUse = tier === 'smart_only' ? (config.smartLLM ?? config.cheapLLM) : config.cheapLLM;
+      const llmTier = tier === 'smart_only' ? 'smart_llm' : 'cheap_llm';
+
+      // Fetch existing facts for the LLM to compare against (enables contradiction detection)
+      let existingFactsForLLM = input.existingFacts?.map(f => ({
+        lineageId: f.lineageId,
+        content: f.content,
+      }));
+
+      if (!existingFactsForLLM || existingFactsForLLM.length === 0) {
+        try {
+          const queryEmbedding = await config.embedding.embed(textContent.slice(0, 1000));
+          const similar = await config.storage.vectorSearch({
+            embedding: queryEmbedding,
+            tenantId: input.tenantId,
+            scope: input.scope,
+            scopeId: input.scopeId,
+            limit: 20,
+            minSimilarity: 0.3,
+            validOnly: true,
+          });
+          if (similar.length > 0) {
+            existingFactsForLLM = similar.map(s => ({
+              lineageId: s.fact.lineageId ?? s.fact.id,
+              content: s.fact.content,
+            }));
+          }
+        } catch {
+          // Non-fatal — extraction continues without existing facts
+        }
+      }
+
+      const llmResult = await extractWithLLM(
+        { llm: llmToUse, tier: llmTier, entityTypes: config.entityTypes, domainEntityTypes: config.domainEntityTypes },
+        textContent,
+        existingFactsForLLM,
+      );
+
+      totalTokensInput += llmResult.tokensInput;
+      totalTokensOutput += llmResult.tokensOutput;
+      tiersUsed.push(llmTier);
+
+      // Escalate to smart LLM if confidence is low
+      let finalLLMResult = llmResult;
+      if (
+        tier === 'auto' &&
+        llmResult.confidence < 0.6 &&
+        config.smartLLM
+      ) {
+        const smartResult = await extractWithLLM(
+          { llm: config.smartLLM, tier: 'smart_llm', entityTypes: config.entityTypes, domainEntityTypes: config.domainEntityTypes },
+          textContent,
+          existingFactsForLLM,
+        );
+        totalTokensInput += smartResult.tokensInput;
+        totalTokensOutput += smartResult.tokensOutput;
+        tiersUsed.push('smart_llm');
+        finalLLMResult = smartResult;
+      }
 
-    // Merge results: LLM takes priority
-    mergedFacts = mergeFacts(heuristicResult.facts, finalLLMResult.facts);
+      // Merge results: LLM takes priority
+      mergedFacts = mergeFacts(heuristicResult.facts, finalLLMResult.facts);
 
-    // Merge entities: dedup by canonicalName, LLM overwrites heuristic
-    mergedEntities = mergeEntities(heuristicResult.entities, finalLLMResult.entities);
-    mergedEdges = [...heuristicResult.edges, ...finalLLMResult.edges];
+      // Merge entities: dedup by canonicalName, LLM overwrites heuristic
+      mergedEntities = mergeEntities(heuristicResult.entities, finalLLMResult.entities);
+      mergedEdges = [...heuristicResult.edges, ...finalLLMResult.edges];
+    }
   }
 
   // Run deduplication against existing memories
@@ -203,7 +246,8 @@ async function executeExtraction(
 
   // Ensure "User" entity exists — LLM creates edges like "user → shops_at → target"
   // but never extracts "user" as an entity, so those edges get silently dropped.
-  if (!mergedEntities.some(e => e.canonicalName === 'user')) {
+  // Skip in codebase mode — there is no "user" entity in code analysis.
+  if (!isCodebaseMode && !mergedEntities.some(e => e.canonicalName === 'user')) {
     mergedEntities.push({
       name: 'User',
       entityType: 'person',

package/src/extraction/types.ts

@@ -60,7 +60,7 @@ export interface ExtractionInput {
   scope: Scope;
   scopeId: string;
   sessionId?: string;
-  inputType: 'conversation' | 'document' | 'url' | 'raw_text' | 'image' | 'audio' | 'code';
+  inputType: 'conversation' | 'document' | 'url' | 'raw_text' | 'image' | 'audio' | 'code' | 'codebase_scan' | 'file_change' | 'architecture_doc';
   data: unknown;
   existingFacts?: Array<{ id: string; lineageId: string; content: string; embedding?: number[] }>;
 }