mirror of
https://github.com/abhigyanpatwari/GitNexus.git
synced 2026-11-13 08:51:28 +00:00
setting up embeddings pipeline. Facing webgpu directx issue, maybe windows compatibility issue
This commit is contained in:
abhigyanpatwari
parent
f25f1f2827
commit
c7437c3c81
16 changed files with 2318 additions and 27 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -33,3 +33,5 @@ coverage/
|
|||
*.local
|
||||
|
||||
.vercel
|
||||
|
||||
|
||||
|
|
|
|||
153
README.md
153
README.md
|
|
@ -2,7 +2,9 @@
|
|||
|
||||
> Privacy-focused, zero-server knowledge graph generator that runs entirely in your browser.
|
||||
|
||||
Transform codebases into interactive knowledge graphs using AST parsing, Web Workers, and an embedded KuzuDB WASM database. All processing happens locally - your code never leaves your machine. Next step -> settng up AI Layer : An embedings pipeline using a very small embedings model that can run in browser and a Graph RAG tool using LLMs to generate and execute cyfer queries. Aiming to give rich and complete retrieved context enabling Agent to detect unused code, perform security audits, do a BLAST RADIUS analyses of code changes and for overall codebase understanding and explaination.
|
||||
Transform codebases into interactive knowledge graphs using AST parsing, Web Workers, and an embedded KuzuDB WASM database. All processing happens locally - your code never leaves your machine.
|
||||
|
||||
**Next up:** Browser-based embeddings + Graph RAG. The cool part? KuzuDB supports native vector indexing, so we can do semantic search AND graph traversal in a single Cypher query. No separate vector DB needed. See [Work in Progress](#-current-work-in-progress) for the full plan.
|
||||
|
||||
|
||||
|
||||
|
|
@ -19,11 +21,116 @@ https://github.com/user-attachments/assets/f375b00a-78cd-4f93-a96c-9ba924455f49
|
|||
**Actively Building:**
|
||||
|
||||
- [ ] **Graph RAG Agent** - AI chat with Cypher query generation for intelligent code exploration
|
||||
- [ ] **Browser Embeddings** - Small embedding model (e.g., gte-small) for semantic node search + LLM-driven RAG
|
||||
- [ ] **Browser Embeddings** - Small embedding model for semantic node search (see below!)
|
||||
- [ ] **Multi-Worker Pool** - Parallel parsing across multiple Web Workers (currently using single worker)
|
||||
- [ ] **Ollama Support** - Local LLM integration
|
||||
- [ ] **CSV Export** - Export node/relationship tables
|
||||
|
||||
### 🧠 Graph RAG: The Plan
|
||||
|
||||
Here's what we're building for the AI layer. The goal: ask questions in plain English, get answers backed by actual graph traversal + semantic understanding.
|
||||
|
||||
**The Problem:** A regular LLM doesn't know your codebase. It can't tell you what calls `handleAuth` or what breaks if you change `UserService`. You need to give it tools to explore the graph.
|
||||
|
||||
**The Solution:** Combine embeddings (for "find relevant code by meaning") with graph queries (for "trace connections").
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
Q[Your Question] --> EMB[Embed with transformers.js]
|
||||
EMB --> VS[Vector Search in KuzuDB]
|
||||
VS --> ENTRY[Entry Point Nodes]
|
||||
ENTRY --> EXPAND[Graph Traversal via Cypher]
|
||||
EXPAND --> CTX[Rich Context]
|
||||
CTX --> LLM[LLM Generates Answer]
|
||||
```
|
||||
|
||||
**Embedding Model:** We're going with `snowflake-arctic-embed-xs` - a tiny 22M parameter model that runs entirely in the browser via [transformers.js](https://huggingface.co/docs/transformers.js). It outputs 384-dimensional vectors and scores 50.15 on MTEB (comparable to models 5x its size). The model downloads once (~90MB), gets cached, and runs locally forever. Privacy intact. ✅
|
||||
|
||||
**The Pipeline:**
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
subgraph Main["Main Pipeline (Blocking)"]
|
||||
P1[Extract] --> P2[Structure] --> P3[Parse] --> P4[Imports] --> P5[Calls]
|
||||
end
|
||||
|
||||
P5 --> READY[Graph Ready!<br/>User can explore]
|
||||
READY --> BG
|
||||
|
||||
subgraph BG["Background (Non-blocking)"]
|
||||
E1[Load Model] --> E2[Embed Nodes] --> E3[Create Vector Index]
|
||||
end
|
||||
|
||||
E3 --> AI[AI Search Ready!]
|
||||
```
|
||||
|
||||
The idea: you can start exploring the graph immediately after Phase 5. Meanwhile, embeddings are generated in the background. Once done, semantic search unlocks.
|
||||
|
||||
### 💡 A Fun Discovery: Unified Vector + Graph = Superpowers
|
||||
|
||||
While designing this, I stumbled onto something cool. Most Graph RAG systems use **separate databases** - a vector DB (Pinecone, Qdrant) for semantic search and a graph DB (Neo4j) for traversal. This means the LLM has to:
|
||||
|
||||
1. Call vector search → get IDs
|
||||
2. Take those IDs → call graph DB
|
||||
3. Coordinate between two systems
|
||||
|
||||
But KuzuDB WASM supports **native vector indexing** (HNSW). Which means we can do vector search AND graph traversal **in a single Cypher query**:
|
||||
|
||||
```cypher
|
||||
-- Find code similar to "authentication" AND trace what calls it
|
||||
-- ALL IN ONE QUERY! 🤯
|
||||
CALL QUERY_VECTOR_INDEX('CodeNode', 'embedding_idx', $queryVector, 10)
|
||||
WITH node AS match, distance
|
||||
WHERE distance < 0.4
|
||||
MATCH (caller:CodeNode)-[r:CodeRelation {type: 'CALLS'}]->(match)
|
||||
RETURN match.name AS found,
|
||||
caller.name AS called_by,
|
||||
distance AS relevance
|
||||
ORDER BY distance
|
||||
```
|
||||
|
||||
This is kind of a big deal. Here's why:
|
||||
|
||||
**Traditional approach (2 queries, 2 systems):**
|
||||
```
|
||||
semantic_search("auth") → ["id1", "id2", "id3"]
|
||||
↓
|
||||
graph_query("MATCH ... WHERE id IN [...]") → results
|
||||
```
|
||||
|
||||
**Unified KuzuDB approach (1 query, 1 system):**
|
||||
```
|
||||
cypher("CALL QUERY_VECTOR_INDEX(...) WITH node MATCH (node)-[...]->() ...") → results
|
||||
```
|
||||
|
||||
And because `distance` comes back with every result, we get **built-in reranking for free**:
|
||||
|
||||
```cypher
|
||||
-- The LLM can dynamically control relevance thresholds!
|
||||
CALL QUERY_VECTOR_INDEX('CodeNode', 'idx', $vec, 20)
|
||||
WITH node, distance,
|
||||
CASE
|
||||
WHEN distance < 0.15 THEN 'exact_match'
|
||||
WHEN distance < 0.30 THEN 'highly_relevant'
|
||||
ELSE 'related'
|
||||
END AS tier
|
||||
WHERE distance < 0.5
|
||||
MATCH (node)-[*1..2]-(context)
|
||||
RETURN node.name, tier, collect(context.name) AS related
|
||||
ORDER BY distance
|
||||
```
|
||||
|
||||
**What this enables:**
|
||||
- 🎯 **Single query execution** - No round trips between systems
|
||||
- 📊 **Hierarchical relevance** - LLM sees exact matches vs related vs weak
|
||||
- 🌳 **Weighted expansion** - Traverse further from better matches
|
||||
- ⚡ **Dynamic thresholds** - LLM adjusts `WHERE distance < X` per question type
|
||||
- 🔄 **No reranker needed** - Distance IS the relevance score
|
||||
|
||||
Basically, the LLM gets to write one smart query that does semantic search, filters by relevance, expands via graph relationships, and returns ranked results. No separate reranker model, no vector DB API calls, no coordination logic. Just Cypher.
|
||||
|
||||
Still wrapping my head around all the query patterns this unlocks, but I'm pretty excited about it.
|
||||
|
||||
---
|
||||
|
||||
## ⚡ What's New in V2
|
||||
|
|
@ -363,6 +470,7 @@ RETURN f.name
|
|||
- ✅ Polymorphic schema (single node/edge tables)
|
||||
- ✅ CSV generation and bulk loading
|
||||
- ✅ Cypher query execution
|
||||
- 🚧 Vector embeddings + HNSW index (WIP)
|
||||
- 🚧 Graph RAG agent (WIP)
|
||||
|
||||
---
|
||||
|
|
@ -372,9 +480,10 @@ RETURN f.name
|
|||
- **Frontend**: React 18 + TypeScript + Vite + Tailwind CSS v4
|
||||
- **Visualization**: Sigma.js + Graphology + ForceAtlas2 (WebGL)
|
||||
- **Parsing**: Tree-sitter WASM (TypeScript, JavaScript, Python)
|
||||
- **Database**: KuzuDB WASM (in-browser graph database)
|
||||
- **Database**: KuzuDB WASM (in-browser graph database + vector index)
|
||||
- **Concurrency**: Web Worker + Comlink
|
||||
- **Caching**: lru-cache with WASM memory management
|
||||
- **AI (WIP)**: transformers.js for browser embeddings, LangChain for agent orchestration
|
||||
|
||||
---
|
||||
|
||||
|
|
@ -428,27 +537,41 @@ Open http://localhost:5173
|
|||
|
||||
### Graph RAG Agent (WIP)
|
||||
|
||||
The idea: ask questions in plain English, get answers backed by graph queries.
|
||||
The idea: ask questions in plain English, get answers backed by graph queries + semantic understanding.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
flowchart TD
|
||||
USER[Your Question] --> LLM[LLM]
|
||||
LLM --> TOOLS[Pick a Tool]
|
||||
TOOLS --> CYPHER[Run Cypher]
|
||||
TOOLS --> SEARCH[Semantic Search]
|
||||
CYPHER --> CONTEXT[Gather Context]
|
||||
SEARCH --> CONTEXT
|
||||
CONTEXT --> LLM
|
||||
LLM --> |Generates| CYPHER[Unified Cypher Query]
|
||||
|
||||
subgraph KUZU[KuzuDB WASM]
|
||||
CYPHER --> VEC[Vector Search]
|
||||
VEC --> GRAPH[Graph Traversal]
|
||||
GRAPH --> RANK[Ranked Results]
|
||||
end
|
||||
|
||||
RANK --> CTX[Rich Context + Code Snippets]
|
||||
CTX --> LLM
|
||||
LLM --> ANSWER[Your Answer]
|
||||
```
|
||||
|
||||
**Example interactions:**
|
||||
|
||||
- "What functions call `handleAuth`?" → Generates Cypher, returns list
|
||||
- "Show me the blast radius if I change `UserService`" → Traverses dependencies
|
||||
- "Find all files that import from `utils/`" → Pattern matching query
|
||||
- "What functions call `handleAuth`?" → Vector search finds `handleAuth`, Cypher traces callers
|
||||
- "Show me the blast radius if I change `UserService`" → Finds service, traverses 3 hops of dependencies
|
||||
- "How does authentication work in this codebase?" → Semantic search for auth-related code, returns connected components
|
||||
|
||||
**Why pre-built query templates?** LLMs are... creative with Cypher syntax. Instead of letting the LLM generate queries from scratch (and fail half the time), we're building a library of reliable query templates that the LLM can choose from and fill in.
|
||||
**Why dynamic Cypher generation?** Originally we planned to use pre-built query templates (because LLMs can be... creative with syntax). But with the unified vector + graph approach, the LLM just needs to learn one pattern:
|
||||
|
||||
```cypher
|
||||
CALL QUERY_VECTOR_INDEX(...) WITH node, distance
|
||||
WHERE distance < [threshold]
|
||||
MATCH (node)-[relationship pattern]->(connected)
|
||||
RETURN [what you need]
|
||||
ORDER BY distance
|
||||
```
|
||||
|
||||
Give the LLM the schema, a few examples, and let it compose queries. The schema is simple enough that modern LLMs (GPT-4, Claude) handle it well. And if a query fails? The error message is usually clear enough for the LLM to self-correct.
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
892
package-lock.json
generated
892
package-lock.json
generated
File diff suppressed because it is too large
Load diff
|
|
@ -9,6 +9,7 @@
|
|||
"preview": "vite preview"
|
||||
},
|
||||
"dependencies": {
|
||||
"@huggingface/transformers": "^3.0.0",
|
||||
"@isomorphic-git/lightning-fs": "^4.6.2",
|
||||
"@sigma/edge-curve": "^3.1.0",
|
||||
"@tailwindcss/vite": "^4.1.18",
|
||||
|
|
|
|||
125
src/components/EmbeddingStatus.tsx
Normal file
125
src/components/EmbeddingStatus.tsx
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
import { Brain, Loader2, Check, AlertCircle, Zap } from 'lucide-react';
|
||||
import { useAppState } from '../hooks/useAppState';
|
||||
|
||||
/**
|
||||
* Embedding status indicator and trigger button
|
||||
* Shows in header when graph is loaded
|
||||
*/
|
||||
export const EmbeddingStatus = () => {
|
||||
const {
|
||||
embeddingStatus,
|
||||
embeddingProgress,
|
||||
startEmbeddings,
|
||||
graph,
|
||||
viewMode
|
||||
} = useAppState();
|
||||
|
||||
// Only show when exploring a loaded graph
|
||||
if (viewMode !== 'exploring' || !graph) return null;
|
||||
|
||||
const handleStartEmbeddings = async () => {
|
||||
try {
|
||||
await startEmbeddings();
|
||||
} catch (error) {
|
||||
console.error('Embedding failed:', error);
|
||||
}
|
||||
};
|
||||
|
||||
// Idle state - show button to start
|
||||
if (embeddingStatus === 'idle') {
|
||||
return (
|
||||
<button
|
||||
onClick={handleStartEmbeddings}
|
||||
className="flex items-center gap-2 px-3 py-1.5 bg-surface border border-border-subtle rounded-lg text-sm text-text-secondary hover:bg-hover hover:text-text-primary hover:border-accent/50 transition-all group"
|
||||
title="Generate embeddings for semantic search"
|
||||
>
|
||||
<Brain className="w-4 h-4 text-node-interface group-hover:text-accent transition-colors" />
|
||||
<span className="hidden sm:inline">Enable Semantic Search</span>
|
||||
<Zap className="w-3 h-3 text-text-muted" />
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
// Loading model
|
||||
if (embeddingStatus === 'loading') {
|
||||
const downloadPercent = embeddingProgress?.modelDownloadPercent ?? 0;
|
||||
return (
|
||||
<div className="flex items-center gap-2.5 px-3 py-1.5 bg-surface border border-accent/30 rounded-lg text-sm">
|
||||
<Loader2 className="w-4 h-4 text-accent animate-spin" />
|
||||
<div className="flex flex-col gap-0.5">
|
||||
<span className="text-text-secondary text-xs">Loading AI model...</span>
|
||||
<div className="w-24 h-1 bg-elevated rounded-full overflow-hidden">
|
||||
<div
|
||||
className="h-full bg-gradient-to-r from-accent to-node-interface rounded-full transition-all duration-300"
|
||||
style={{ width: `${downloadPercent}%` }}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Embedding in progress
|
||||
if (embeddingStatus === 'embedding') {
|
||||
const processed = embeddingProgress?.nodesProcessed ?? 0;
|
||||
const total = embeddingProgress?.totalNodes ?? 0;
|
||||
const percent = embeddingProgress?.percent ?? 0;
|
||||
|
||||
return (
|
||||
<div className="flex items-center gap-2.5 px-3 py-1.5 bg-surface border border-node-function/30 rounded-lg text-sm">
|
||||
<Loader2 className="w-4 h-4 text-node-function animate-spin" />
|
||||
<div className="flex flex-col gap-0.5">
|
||||
<span className="text-text-secondary text-xs">
|
||||
Embedding {processed}/{total} nodes
|
||||
</span>
|
||||
<div className="w-24 h-1 bg-elevated rounded-full overflow-hidden">
|
||||
<div
|
||||
className="h-full bg-gradient-to-r from-node-function to-accent rounded-full transition-all duration-300"
|
||||
style={{ width: `${percent}%` }}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Indexing
|
||||
if (embeddingStatus === 'indexing') {
|
||||
return (
|
||||
<div className="flex items-center gap-2 px-3 py-1.5 bg-surface border border-node-interface/30 rounded-lg text-sm text-text-secondary">
|
||||
<Loader2 className="w-4 h-4 text-node-interface animate-spin" />
|
||||
<span className="text-xs">Creating vector index...</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Ready
|
||||
if (embeddingStatus === 'ready') {
|
||||
return (
|
||||
<div
|
||||
className="flex items-center gap-2 px-3 py-1.5 bg-node-function/10 border border-node-function/30 rounded-lg text-sm text-node-function"
|
||||
title="Semantic search is ready! Use natural language in the AI chat."
|
||||
>
|
||||
<Check className="w-4 h-4" />
|
||||
<span className="text-xs font-medium">Semantic Ready</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// Error
|
||||
if (embeddingStatus === 'error') {
|
||||
return (
|
||||
<button
|
||||
onClick={handleStartEmbeddings}
|
||||
className="flex items-center gap-2 px-3 py-1.5 bg-red-500/10 border border-red-500/30 rounded-lg text-sm text-red-400 hover:bg-red-500/20 transition-colors"
|
||||
title={embeddingProgress?.error || 'Embedding failed. Click to retry.'}
|
||||
>
|
||||
<AlertCircle className="w-4 h-4" />
|
||||
<span className="text-xs">Failed - Retry</span>
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
return null;
|
||||
};
|
||||
|
||||
|
|
@ -2,6 +2,7 @@ import { Search, Settings, HelpCircle, Sparkles } from 'lucide-react';
|
|||
import { useAppState } from '../hooks/useAppState';
|
||||
import { useState, useMemo, useRef, useEffect } from 'react';
|
||||
import { GraphNode } from '../core/graph/types';
|
||||
import { EmbeddingStatus } from './EmbeddingStatus';
|
||||
|
||||
// Color mapping for node types in search results
|
||||
const NODE_TYPE_COLORS: Record<string, string> = {
|
||||
|
|
@ -184,12 +185,15 @@ export const Header = ({ onFocusNode }: HeaderProps) => {
|
|||
<div className="flex items-center gap-2">
|
||||
{/* Stats */}
|
||||
{graph && (
|
||||
<div className="flex items-center gap-4 mr-4 text-xs text-text-muted">
|
||||
<div className="flex items-center gap-4 mr-2 text-xs text-text-muted">
|
||||
<span>{nodeCount} nodes</span>
|
||||
<span>{edgeCount} edges</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Embedding Status */}
|
||||
<EmbeddingStatus />
|
||||
|
||||
{/* Icon buttons */}
|
||||
<button className="w-9 h-9 flex items-center justify-center rounded-md text-text-secondary hover:bg-hover hover:text-text-primary transition-colors">
|
||||
<Settings className="w-[18px] h-[18px]" />
|
||||
|
|
|
|||
224
src/core/embeddings/embedder.ts
Normal file
224
src/core/embeddings/embedder.ts
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
/**
|
||||
* Embedder Module
|
||||
*
|
||||
* Singleton factory for transformers.js embedding pipeline.
|
||||
* Handles model loading, caching, and both single and batch embedding operations.
|
||||
*
|
||||
* Uses snowflake-arctic-embed-xs by default (22M params, 384 dims, ~90MB)
|
||||
*/
|
||||
|
||||
import { pipeline, env, type FeatureExtractionPipeline } from '@huggingface/transformers';
|
||||
import { DEFAULT_EMBEDDING_CONFIG, type EmbeddingConfig, type ModelProgress } from './types';
|
||||
|
||||
// Module-level state for singleton pattern
|
||||
let embedderInstance: FeatureExtractionPipeline | null = null;
|
||||
let isInitializing = false;
|
||||
let initPromise: Promise<FeatureExtractionPipeline> | null = null;
|
||||
|
||||
/**
|
||||
* Progress callback type for model loading
|
||||
*/
|
||||
export type ModelProgressCallback = (progress: ModelProgress) => void;
|
||||
|
||||
/**
|
||||
* Initialize the embedding model
|
||||
* Uses singleton pattern - only loads once, subsequent calls return cached instance
|
||||
*
|
||||
* @param onProgress - Optional callback for model download progress
|
||||
* @param config - Optional configuration override
|
||||
* @returns Promise resolving to the embedder pipeline
|
||||
*/
|
||||
export const initEmbedder = async (
|
||||
onProgress?: ModelProgressCallback,
|
||||
config: Partial<EmbeddingConfig> = {}
|
||||
): Promise<FeatureExtractionPipeline> => {
|
||||
// Return existing instance if available
|
||||
if (embedderInstance) {
|
||||
return embedderInstance;
|
||||
}
|
||||
|
||||
// If already initializing, wait for that promise
|
||||
if (isInitializing && initPromise) {
|
||||
return initPromise;
|
||||
}
|
||||
|
||||
isInitializing = true;
|
||||
|
||||
const finalConfig = { ...DEFAULT_EMBEDDING_CONFIG, ...config };
|
||||
|
||||
initPromise = (async () => {
|
||||
try {
|
||||
// Configure transformers.js environment
|
||||
env.allowLocalModels = false;
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log(`🧠 Loading embedding model: ${finalConfig.modelId}`);
|
||||
}
|
||||
|
||||
// Create the feature extraction pipeline
|
||||
// Try WebGPU first for speed, fallback to WASM for compatibility
|
||||
const progressCallback = onProgress ? (data: any) => {
|
||||
const progress: ModelProgress = {
|
||||
status: data.status || 'progress',
|
||||
file: data.file,
|
||||
progress: data.progress,
|
||||
loaded: data.loaded,
|
||||
total: data.total,
|
||||
};
|
||||
onProgress(progress);
|
||||
} : undefined;
|
||||
|
||||
// Try preferred device first, fallback to wasm if it fails
|
||||
const devicesToTry: Array<'webgpu' | 'wasm'> =
|
||||
finalConfig.device === 'webgpu' ? ['webgpu', 'wasm'] : ['wasm'];
|
||||
|
||||
let lastError: Error | null = null;
|
||||
|
||||
for (const device of devicesToTry) {
|
||||
try {
|
||||
if (import.meta.env.DEV && devicesToTry.length > 1) {
|
||||
console.log(`🔧 Trying ${device} backend...`);
|
||||
}
|
||||
|
||||
embedderInstance = await pipeline(
|
||||
'feature-extraction',
|
||||
finalConfig.modelId,
|
||||
{
|
||||
device,
|
||||
dtype: 'fp32',
|
||||
progress_callback: progressCallback,
|
||||
}
|
||||
);
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log(`✅ Using ${device} backend`);
|
||||
}
|
||||
break; // Success!
|
||||
|
||||
} catch (err) {
|
||||
lastError = err as Error;
|
||||
if (import.meta.env.DEV) {
|
||||
console.warn(`⚠️ ${device} backend failed:`, err);
|
||||
}
|
||||
// Continue to next device
|
||||
}
|
||||
}
|
||||
|
||||
if (!embedderInstance) {
|
||||
throw lastError || new Error('No backend available');
|
||||
}
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log('✅ Embedding model loaded successfully');
|
||||
}
|
||||
|
||||
return embedderInstance;
|
||||
} catch (error) {
|
||||
isInitializing = false;
|
||||
initPromise = null;
|
||||
embedderInstance = null;
|
||||
throw error;
|
||||
} finally {
|
||||
isInitializing = false;
|
||||
}
|
||||
})();
|
||||
|
||||
return initPromise;
|
||||
};
|
||||
|
||||
/**
|
||||
* Check if the embedder is initialized and ready
|
||||
*/
|
||||
export const isEmbedderReady = (): boolean => {
|
||||
return embedderInstance !== null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Get the embedder instance (throws if not initialized)
|
||||
*/
|
||||
export const getEmbedder = (): FeatureExtractionPipeline => {
|
||||
if (!embedderInstance) {
|
||||
throw new Error('Embedder not initialized. Call initEmbedder() first.');
|
||||
}
|
||||
return embedderInstance;
|
||||
};
|
||||
|
||||
/**
|
||||
* Embed a single text string
|
||||
*
|
||||
* @param text - Text to embed
|
||||
* @returns Float32Array of embedding vector (384 dimensions)
|
||||
*/
|
||||
export const embedText = async (text: string): Promise<Float32Array> => {
|
||||
const embedder = getEmbedder();
|
||||
|
||||
const result = await embedder(text, {
|
||||
pooling: 'mean',
|
||||
normalize: true,
|
||||
});
|
||||
|
||||
// Result is a Tensor, convert to Float32Array
|
||||
return new Float32Array(result.data as ArrayLike<number>);
|
||||
};
|
||||
|
||||
/**
|
||||
* Embed multiple texts in a single batch
|
||||
* More efficient than calling embedText multiple times
|
||||
*
|
||||
* @param texts - Array of texts to embed
|
||||
* @returns Array of Float32Array embedding vectors
|
||||
*/
|
||||
export const embedBatch = async (texts: string[]): Promise<Float32Array[]> => {
|
||||
if (texts.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const embedder = getEmbedder();
|
||||
|
||||
// Process batch
|
||||
const result = await embedder(texts, {
|
||||
pooling: 'mean',
|
||||
normalize: true,
|
||||
});
|
||||
|
||||
// Result shape is [batch_size, dimensions]
|
||||
// Need to split into individual vectors
|
||||
const data = result.data as ArrayLike<number>;
|
||||
const dimensions = DEFAULT_EMBEDDING_CONFIG.dimensions;
|
||||
const embeddings: Float32Array[] = [];
|
||||
|
||||
for (let i = 0; i < texts.length; i++) {
|
||||
const start = i * dimensions;
|
||||
const end = start + dimensions;
|
||||
embeddings.push(new Float32Array(Array.prototype.slice.call(data, start, end)));
|
||||
}
|
||||
|
||||
return embeddings;
|
||||
};
|
||||
|
||||
/**
|
||||
* Convert Float32Array to regular number array (for KuzuDB storage)
|
||||
*/
|
||||
export const embeddingToArray = (embedding: Float32Array): number[] => {
|
||||
return Array.from(embedding);
|
||||
};
|
||||
|
||||
/**
|
||||
* Cleanup the embedder (free memory)
|
||||
* Call this when done with embeddings
|
||||
*/
|
||||
export const disposeEmbedder = async (): Promise<void> => {
|
||||
if (embedderInstance) {
|
||||
// transformers.js pipelines may have a dispose method
|
||||
try {
|
||||
if ('dispose' in embedderInstance && typeof embedderInstance.dispose === 'function') {
|
||||
await embedderInstance.dispose();
|
||||
}
|
||||
} catch {
|
||||
// Ignore disposal errors
|
||||
}
|
||||
embedderInstance = null;
|
||||
initPromise = null;
|
||||
}
|
||||
};
|
||||
|
||||
358
src/core/embeddings/embedding-pipeline.ts
Normal file
358
src/core/embeddings/embedding-pipeline.ts
Normal file
|
|
@ -0,0 +1,358 @@
|
|||
/**
|
||||
* Embedding Pipeline Module
|
||||
*
|
||||
* Orchestrates the background embedding process:
|
||||
* 1. Query embeddable nodes from KuzuDB
|
||||
* 2. Generate text representations
|
||||
* 3. Batch embed using transformers.js
|
||||
* 4. Update KuzuDB with embeddings
|
||||
* 5. Create vector index for semantic search
|
||||
*/
|
||||
|
||||
import { initEmbedder, embedBatch, embedText, embeddingToArray, isEmbedderReady } from './embedder';
|
||||
import { generateBatchEmbeddingTexts, generateEmbeddingText } from './text-generator';
|
||||
import {
|
||||
type EmbeddingProgress,
|
||||
type EmbeddingConfig,
|
||||
type EmbeddableNode,
|
||||
type SemanticSearchResult,
|
||||
type ModelProgress,
|
||||
DEFAULT_EMBEDDING_CONFIG,
|
||||
EMBEDDABLE_LABELS,
|
||||
} from './types';
|
||||
|
||||
/**
|
||||
* Progress callback type
|
||||
*/
|
||||
export type EmbeddingProgressCallback = (progress: EmbeddingProgress) => void;
|
||||
|
||||
/**
|
||||
* Query all embeddable nodes from KuzuDB
|
||||
*/
|
||||
const queryEmbeddableNodes = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>
|
||||
): Promise<EmbeddableNode[]> => {
|
||||
// Build WHERE clause for embeddable labels
|
||||
const labelConditions = EMBEDDABLE_LABELS
|
||||
.map(label => `n.label = '${label}'`)
|
||||
.join(' OR ');
|
||||
|
||||
const cypher = `
|
||||
MATCH (n:CodeNode)
|
||||
WHERE ${labelConditions}
|
||||
RETURN n.id AS id, n.name AS name, n.label AS label,
|
||||
n.filePath AS filePath, n.content AS content,
|
||||
n.startLine AS startLine, n.endLine AS endLine
|
||||
`;
|
||||
|
||||
const rows = await executeQuery(cypher);
|
||||
|
||||
return rows.map(row => ({
|
||||
id: row.id ?? row[0],
|
||||
name: row.name ?? row[1],
|
||||
label: row.label ?? row[2],
|
||||
filePath: row.filePath ?? row[3],
|
||||
content: row.content ?? row[4] ?? '',
|
||||
startLine: row.startLine ?? row[5],
|
||||
endLine: row.endLine ?? row[6],
|
||||
}));
|
||||
};
|
||||
|
||||
/**
|
||||
* Update a single node's embedding in KuzuDB
|
||||
*/
|
||||
const updateNodeEmbedding = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
nodeId: string,
|
||||
embedding: number[]
|
||||
): Promise<void> => {
|
||||
// KuzuDB requires the array to be cast to the correct type
|
||||
const embeddingStr = `[${embedding.join(',')}]`;
|
||||
|
||||
const cypher = `
|
||||
MATCH (n:CodeNode {id: '${nodeId}'})
|
||||
SET n.embedding = CAST(${embeddingStr} AS FLOAT[384])
|
||||
`;
|
||||
|
||||
await executeQuery(cypher);
|
||||
};
|
||||
|
||||
/**
|
||||
* Batch update multiple node embeddings
|
||||
* More efficient than individual updates
|
||||
*/
|
||||
const batchUpdateEmbeddings = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
updates: Array<{ id: string; embedding: number[] }>
|
||||
): Promise<void> => {
|
||||
// Process updates one by one for now
|
||||
// KuzuDB doesn't have great batch update syntax
|
||||
for (const update of updates) {
|
||||
await updateNodeEmbedding(executeQuery, update.id, update.embedding);
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Create the vector index for semantic search
|
||||
*/
|
||||
const createVectorIndex = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>
|
||||
): Promise<void> => {
|
||||
const cypher = `
|
||||
CALL CREATE_VECTOR_INDEX('CodeNode', 'code_embedding_idx', 'embedding', metric := 'cosine')
|
||||
`;
|
||||
|
||||
try {
|
||||
await executeQuery(cypher);
|
||||
} catch (error) {
|
||||
// Index might already exist
|
||||
if (import.meta.env.DEV) {
|
||||
console.warn('Vector index creation warning:', error);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Run the embedding pipeline
|
||||
*
|
||||
* @param executeQuery - Function to execute Cypher queries against KuzuDB
|
||||
* @param onProgress - Callback for progress updates
|
||||
* @param config - Optional configuration override
|
||||
*/
|
||||
export const runEmbeddingPipeline = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
onProgress: EmbeddingProgressCallback,
|
||||
config: Partial<EmbeddingConfig> = {}
|
||||
): Promise<void> => {
|
||||
const finalConfig = { ...DEFAULT_EMBEDDING_CONFIG, ...config };
|
||||
|
||||
try {
|
||||
// Phase 1: Load embedding model
|
||||
onProgress({
|
||||
phase: 'loading-model',
|
||||
percent: 0,
|
||||
modelDownloadPercent: 0,
|
||||
});
|
||||
|
||||
await initEmbedder((modelProgress: ModelProgress) => {
|
||||
// Report model download progress
|
||||
const downloadPercent = modelProgress.progress ?? 0;
|
||||
onProgress({
|
||||
phase: 'loading-model',
|
||||
percent: Math.round(downloadPercent * 0.2), // 0-20% for model loading
|
||||
modelDownloadPercent: downloadPercent,
|
||||
});
|
||||
}, finalConfig);
|
||||
|
||||
onProgress({
|
||||
phase: 'loading-model',
|
||||
percent: 20,
|
||||
modelDownloadPercent: 100,
|
||||
});
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log('🔍 Querying embeddable nodes...');
|
||||
}
|
||||
|
||||
// Phase 2: Query embeddable nodes
|
||||
const nodes = await queryEmbeddableNodes(executeQuery);
|
||||
const totalNodes = nodes.length;
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log(`📊 Found ${totalNodes} embeddable nodes`);
|
||||
}
|
||||
|
||||
if (totalNodes === 0) {
|
||||
onProgress({
|
||||
phase: 'ready',
|
||||
percent: 100,
|
||||
nodesProcessed: 0,
|
||||
totalNodes: 0,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// Phase 3: Batch embed nodes
|
||||
const batchSize = finalConfig.batchSize;
|
||||
const totalBatches = Math.ceil(totalNodes / batchSize);
|
||||
let processedNodes = 0;
|
||||
|
||||
onProgress({
|
||||
phase: 'embedding',
|
||||
percent: 20,
|
||||
nodesProcessed: 0,
|
||||
totalNodes,
|
||||
currentBatch: 0,
|
||||
totalBatches,
|
||||
});
|
||||
|
||||
for (let batchIndex = 0; batchIndex < totalBatches; batchIndex++) {
|
||||
const start = batchIndex * batchSize;
|
||||
const end = Math.min(start + batchSize, totalNodes);
|
||||
const batch = nodes.slice(start, end);
|
||||
|
||||
// Generate texts for this batch
|
||||
const texts = generateBatchEmbeddingTexts(batch, finalConfig);
|
||||
|
||||
// Embed the batch
|
||||
const embeddings = await embedBatch(texts);
|
||||
|
||||
// Update KuzuDB with embeddings
|
||||
const updates = batch.map((node, i) => ({
|
||||
id: node.id,
|
||||
embedding: embeddingToArray(embeddings[i]),
|
||||
}));
|
||||
|
||||
await batchUpdateEmbeddings(executeQuery, updates);
|
||||
|
||||
processedNodes += batch.length;
|
||||
|
||||
// Report progress (20-90% for embedding phase)
|
||||
const embeddingProgress = 20 + ((processedNodes / totalNodes) * 70);
|
||||
onProgress({
|
||||
phase: 'embedding',
|
||||
percent: Math.round(embeddingProgress),
|
||||
nodesProcessed: processedNodes,
|
||||
totalNodes,
|
||||
currentBatch: batchIndex + 1,
|
||||
totalBatches,
|
||||
});
|
||||
}
|
||||
|
||||
// Phase 4: Create vector index
|
||||
onProgress({
|
||||
phase: 'indexing',
|
||||
percent: 90,
|
||||
nodesProcessed: totalNodes,
|
||||
totalNodes,
|
||||
});
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log('📇 Creating vector index...');
|
||||
}
|
||||
|
||||
await createVectorIndex(executeQuery);
|
||||
|
||||
// Complete
|
||||
onProgress({
|
||||
phase: 'ready',
|
||||
percent: 100,
|
||||
nodesProcessed: totalNodes,
|
||||
totalNodes,
|
||||
});
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.log('✅ Embedding pipeline complete!');
|
||||
}
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
console.error('❌ Embedding pipeline error:', error);
|
||||
}
|
||||
|
||||
onProgress({
|
||||
phase: 'error',
|
||||
percent: 0,
|
||||
error: errorMessage,
|
||||
});
|
||||
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Perform semantic search using the vector index
|
||||
*
|
||||
* @param executeQuery - Function to execute Cypher queries
|
||||
* @param query - Search query text
|
||||
* @param k - Number of results to return (default: 10)
|
||||
* @param maxDistance - Maximum distance threshold (default: 0.5)
|
||||
* @returns Array of search results ordered by relevance
|
||||
*/
|
||||
export const semanticSearch = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
query: string,
|
||||
k: number = 10,
|
||||
maxDistance: number = 0.5
|
||||
): Promise<SemanticSearchResult[]> => {
|
||||
if (!isEmbedderReady()) {
|
||||
throw new Error('Embedding model not initialized. Run embedding pipeline first.');
|
||||
}
|
||||
|
||||
// Embed the query
|
||||
const queryEmbedding = await embedText(query);
|
||||
const queryVec = embeddingToArray(queryEmbedding);
|
||||
const queryVecStr = `[${queryVec.join(',')}]`;
|
||||
|
||||
// Query the vector index
|
||||
const cypher = `
|
||||
CALL QUERY_VECTOR_INDEX('CodeNode', 'code_embedding_idx',
|
||||
CAST(${queryVecStr} AS FLOAT[384]), ${k})
|
||||
YIELD node, distance
|
||||
WHERE distance < ${maxDistance}
|
||||
RETURN node.id AS nodeId, node.name AS name, node.label AS label,
|
||||
node.filePath AS filePath, distance,
|
||||
node.startLine AS startLine, node.endLine AS endLine
|
||||
ORDER BY distance
|
||||
`;
|
||||
|
||||
const rows = await executeQuery(cypher);
|
||||
|
||||
return rows.map(row => ({
|
||||
nodeId: row.nodeId ?? row[0],
|
||||
name: row.name ?? row[1],
|
||||
label: row.label ?? row[2],
|
||||
filePath: row.filePath ?? row[3],
|
||||
distance: row.distance ?? row[4],
|
||||
startLine: row.startLine ?? row[5],
|
||||
endLine: row.endLine ?? row[6],
|
||||
}));
|
||||
};
|
||||
|
||||
/**
|
||||
* Semantic search with graph expansion
|
||||
* Finds similar nodes AND their connections
|
||||
*
|
||||
* @param executeQuery - Function to execute Cypher queries
|
||||
* @param query - Search query text
|
||||
* @param k - Number of initial results
|
||||
* @param hops - Number of hops to expand (default: 2)
|
||||
* @returns Search results with connected nodes
|
||||
*/
|
||||
export const semanticSearchWithContext = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
query: string,
|
||||
k: number = 5,
|
||||
hops: number = 2
|
||||
): Promise<any[]> => {
|
||||
if (!isEmbedderReady()) {
|
||||
throw new Error('Embedding model not initialized. Run embedding pipeline first.');
|
||||
}
|
||||
|
||||
// Embed the query
|
||||
const queryEmbedding = await embedText(query);
|
||||
const queryVec = embeddingToArray(queryEmbedding);
|
||||
const queryVecStr = `[${queryVec.join(',')}]`;
|
||||
|
||||
// Query with graph expansion
|
||||
const cypher = `
|
||||
CALL QUERY_VECTOR_INDEX('CodeNode', 'code_embedding_idx',
|
||||
CAST(${queryVecStr} AS FLOAT[384]), ${k})
|
||||
YIELD node AS match, distance
|
||||
WHERE distance < 0.5
|
||||
MATCH (match)-[r:CodeRelation*1..${hops}]-(connected:CodeNode)
|
||||
RETURN match.id AS matchId, match.name AS matchName, match.label AS matchLabel,
|
||||
match.filePath AS matchPath, distance,
|
||||
collect(DISTINCT {
|
||||
id: connected.id,
|
||||
name: connected.name,
|
||||
label: connected.label,
|
||||
relationType: [rel IN r | rel.type]
|
||||
}) AS connections
|
||||
ORDER BY distance
|
||||
`;
|
||||
|
||||
return executeQuery(cypher);
|
||||
};
|
||||
|
||||
11
src/core/embeddings/index.ts
Normal file
11
src/core/embeddings/index.ts
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
/**
|
||||
* Embeddings Module
|
||||
*
|
||||
* Re-exports for the embedding pipeline system.
|
||||
*/
|
||||
|
||||
export * from './types';
|
||||
export * from './embedder';
|
||||
export * from './text-generator';
|
||||
export * from './embedding-pipeline';
|
||||
|
||||
235
src/core/embeddings/text-generator.ts
Normal file
235
src/core/embeddings/text-generator.ts
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
/**
|
||||
* Text Generator Module
|
||||
*
|
||||
* Pure functions to generate embedding text from code nodes.
|
||||
* Combines node metadata with code snippets for semantic matching.
|
||||
*/
|
||||
|
||||
import type { EmbeddableNode, EmbeddingConfig } from './types';
|
||||
import { DEFAULT_EMBEDDING_CONFIG } from './types';
|
||||
|
||||
/**
|
||||
* Extract the filename from a file path
|
||||
*/
|
||||
const getFileName = (filePath: string): string => {
|
||||
const parts = filePath.split('/');
|
||||
return parts[parts.length - 1] || filePath;
|
||||
};
|
||||
|
||||
/**
|
||||
* Extract the directory path from a file path
|
||||
*/
|
||||
const getDirectory = (filePath: string): string => {
|
||||
const parts = filePath.split('/');
|
||||
parts.pop();
|
||||
return parts.join('/') || '';
|
||||
};
|
||||
|
||||
/**
|
||||
* Truncate content to max length, preserving word boundaries
|
||||
*/
|
||||
const truncateContent = (content: string, maxLength: number): string => {
|
||||
if (content.length <= maxLength) {
|
||||
return content;
|
||||
}
|
||||
|
||||
// Find last space before maxLength to avoid cutting words
|
||||
const truncated = content.slice(0, maxLength);
|
||||
const lastSpace = truncated.lastIndexOf(' ');
|
||||
|
||||
if (lastSpace > maxLength * 0.8) {
|
||||
return truncated.slice(0, lastSpace) + '...';
|
||||
}
|
||||
|
||||
return truncated + '...';
|
||||
};
|
||||
|
||||
/**
|
||||
* Clean code content for embedding
|
||||
* Removes excessive whitespace while preserving structure
|
||||
*/
|
||||
const cleanContent = (content: string): string => {
|
||||
return content
|
||||
// Normalize line endings
|
||||
.replace(/\r\n/g, '\n')
|
||||
// Remove excessive blank lines (more than 2)
|
||||
.replace(/\n{3,}/g, '\n\n')
|
||||
// Trim each line
|
||||
.split('\n')
|
||||
.map(line => line.trimEnd())
|
||||
.join('\n')
|
||||
.trim();
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a Function node
|
||||
*/
|
||||
const generateFunctionText = (
|
||||
node: EmbeddableNode,
|
||||
maxSnippetLength: number
|
||||
): string => {
|
||||
const parts: string[] = [
|
||||
`Function: ${node.name}`,
|
||||
`File: ${getFileName(node.filePath)}`,
|
||||
];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a Class node
|
||||
*/
|
||||
const generateClassText = (
|
||||
node: EmbeddableNode,
|
||||
maxSnippetLength: number
|
||||
): string => {
|
||||
const parts: string[] = [
|
||||
`Class: ${node.name}`,
|
||||
`File: ${getFileName(node.filePath)}`,
|
||||
];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a Method node
|
||||
*/
|
||||
const generateMethodText = (
|
||||
node: EmbeddableNode,
|
||||
maxSnippetLength: number
|
||||
): string => {
|
||||
const parts: string[] = [
|
||||
`Method: ${node.name}`,
|
||||
`File: ${getFileName(node.filePath)}`,
|
||||
];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for an Interface node
|
||||
*/
|
||||
const generateInterfaceText = (
|
||||
node: EmbeddableNode,
|
||||
maxSnippetLength: number
|
||||
): string => {
|
||||
const parts: string[] = [
|
||||
`Interface: ${node.name}`,
|
||||
`File: ${getFileName(node.filePath)}`,
|
||||
];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a File node
|
||||
* Uses file name and first N characters of content
|
||||
*/
|
||||
const generateFileText = (
|
||||
node: EmbeddableNode,
|
||||
maxSnippetLength: number
|
||||
): string => {
|
||||
const parts: string[] = [
|
||||
`File: ${node.name}`,
|
||||
`Path: ${node.filePath}`,
|
||||
];
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
// For files, use a shorter snippet since they can be very long
|
||||
const snippet = truncateContent(cleanedContent, Math.min(maxSnippetLength, 300));
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for any embeddable node
|
||||
* Dispatches to the appropriate generator based on node label
|
||||
*
|
||||
* @param node - The node to generate text for
|
||||
* @param config - Optional configuration for max snippet length
|
||||
* @returns Text suitable for embedding
|
||||
*/
|
||||
export const generateEmbeddingText = (
|
||||
node: EmbeddableNode,
|
||||
config: Partial<EmbeddingConfig> = {}
|
||||
): string => {
|
||||
const maxSnippetLength = config.maxSnippetLength ?? DEFAULT_EMBEDDING_CONFIG.maxSnippetLength;
|
||||
|
||||
switch (node.label) {
|
||||
case 'Function':
|
||||
return generateFunctionText(node, maxSnippetLength);
|
||||
case 'Class':
|
||||
return generateClassText(node, maxSnippetLength);
|
||||
case 'Method':
|
||||
return generateMethodText(node, maxSnippetLength);
|
||||
case 'Interface':
|
||||
return generateInterfaceText(node, maxSnippetLength);
|
||||
case 'File':
|
||||
return generateFileText(node, maxSnippetLength);
|
||||
default:
|
||||
// Fallback for any other embeddable type
|
||||
return `${node.label}: ${node.name}\nPath: ${node.filePath}`;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding texts for a batch of nodes
|
||||
*
|
||||
* @param nodes - Array of nodes to generate text for
|
||||
* @param config - Optional configuration
|
||||
* @returns Array of texts in the same order as input nodes
|
||||
*/
|
||||
export const generateBatchEmbeddingTexts = (
|
||||
nodes: EmbeddableNode[],
|
||||
config: Partial<EmbeddingConfig> = {}
|
||||
): string[] => {
|
||||
return nodes.map(node => generateEmbeddingText(node, config));
|
||||
};
|
||||
|
||||
117
src/core/embeddings/types.ts
Normal file
117
src/core/embeddings/types.ts
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
/**
|
||||
* Embedding Pipeline Types
|
||||
*
|
||||
* Type definitions for the embedding generation and semantic search system.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Node labels that should be embedded for semantic search
|
||||
* These are code elements that benefit from semantic matching
|
||||
*/
|
||||
export const EMBEDDABLE_LABELS = [
|
||||
'Function',
|
||||
'Class',
|
||||
'Method',
|
||||
'Interface',
|
||||
'File',
|
||||
] as const;
|
||||
|
||||
export type EmbeddableLabel = typeof EMBEDDABLE_LABELS[number];
|
||||
|
||||
/**
|
||||
* Check if a label should be embedded
|
||||
*/
|
||||
export const isEmbeddableLabel = (label: string): label is EmbeddableLabel =>
|
||||
EMBEDDABLE_LABELS.includes(label as EmbeddableLabel);
|
||||
|
||||
/**
|
||||
* Embedding pipeline phases
|
||||
*/
|
||||
export type EmbeddingPhase =
|
||||
| 'idle'
|
||||
| 'loading-model'
|
||||
| 'embedding'
|
||||
| 'indexing'
|
||||
| 'ready'
|
||||
| 'error';
|
||||
|
||||
/**
|
||||
* Progress information for the embedding pipeline
|
||||
*/
|
||||
export interface EmbeddingProgress {
|
||||
phase: EmbeddingPhase;
|
||||
percent: number;
|
||||
modelDownloadPercent?: number;
|
||||
nodesProcessed?: number;
|
||||
totalNodes?: number;
|
||||
currentBatch?: number;
|
||||
totalBatches?: number;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configuration for the embedding pipeline
|
||||
*/
|
||||
export interface EmbeddingConfig {
|
||||
/** Model identifier for transformers.js */
|
||||
modelId: string;
|
||||
/** Number of nodes to embed in each batch */
|
||||
batchSize: number;
|
||||
/** Embedding vector dimensions */
|
||||
dimensions: number;
|
||||
/** Device to use for inference */
|
||||
device: 'webgpu' | 'wasm' | 'cpu';
|
||||
/** Maximum characters of code snippet to include */
|
||||
maxSnippetLength: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Default embedding configuration
|
||||
* Uses snowflake-arctic-embed-xs for browser efficiency
|
||||
* Tries WebGPU first (fast), auto-fallback to WASM if unavailable
|
||||
*/
|
||||
export const DEFAULT_EMBEDDING_CONFIG: EmbeddingConfig = {
|
||||
modelId: 'Snowflake/snowflake-arctic-embed-xs',
|
||||
batchSize: 16,
|
||||
dimensions: 384,
|
||||
device: 'webgpu', // Try WebGPU first, auto-fallback to WASM
|
||||
maxSnippetLength: 500,
|
||||
};
|
||||
|
||||
/**
|
||||
* Result from semantic search
|
||||
*/
|
||||
export interface SemanticSearchResult {
|
||||
nodeId: string;
|
||||
name: string;
|
||||
label: string;
|
||||
filePath: string;
|
||||
distance: number;
|
||||
startLine?: number;
|
||||
endLine?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Node data for embedding (minimal structure from KuzuDB query)
|
||||
*/
|
||||
export interface EmbeddableNode {
|
||||
id: string;
|
||||
name: string;
|
||||
label: string;
|
||||
filePath: string;
|
||||
content: string;
|
||||
startLine?: number;
|
||||
endLine?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Model download progress from transformers.js
|
||||
*/
|
||||
export interface ModelProgress {
|
||||
status: 'initiate' | 'download' | 'progress' | 'done' | 'ready';
|
||||
file?: string;
|
||||
progress?: number;
|
||||
loaded?: number;
|
||||
total?: number;
|
||||
}
|
||||
|
||||
|
|
@ -149,6 +149,8 @@ const extractContent = (
|
|||
* Headers: id,label,name,filePath,startLine,endLine,content
|
||||
*
|
||||
* All string fields are quoted for RFC 4180 compliance
|
||||
* Note: embedding column is NOT included in CSV - it's populated later via UPDATE queries
|
||||
* by the embedding pipeline after bulk load completes
|
||||
*/
|
||||
export const generateNodeCSV = (
|
||||
graph: KnowledgeGraph,
|
||||
|
|
|
|||
|
|
@ -86,7 +86,8 @@ export const loadGraphToKuzu = async (
|
|||
|
||||
// Use HEADER=true because our CSV generator adds headers
|
||||
// Use PARALLEL=false because content field has quoted newlines
|
||||
await conn.query(`COPY ${NODE_TABLE_NAME} FROM "${nodesPath}" (HEADER=true, PARALLEL=false)`);
|
||||
// Explicitly list columns since CSV doesn't include 'embedding' (populated later via UPDATE)
|
||||
await conn.query(`COPY ${NODE_TABLE_NAME}(id, label, name, filePath, startLine, endLine, content) FROM "${nodesPath}" (HEADER=true, PARALLEL=false)`);
|
||||
await conn.query(`COPY ${EDGE_TABLE_NAME} FROM "${edgesPath}" (HEADER=true, PARALLEL=false)`);
|
||||
|
||||
// Verify results
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ export const EDGE_TABLE_NAME = 'CodeRelation';
|
|||
/**
|
||||
* Node table schema
|
||||
* Stores all code elements: Files, Functions, Classes, etc.
|
||||
* embedding column stores 384-dimensional vectors for semantic search
|
||||
*/
|
||||
export const NODE_SCHEMA = `
|
||||
CREATE NODE TABLE ${NODE_TABLE_NAME} (
|
||||
|
|
@ -24,9 +25,18 @@ CREATE NODE TABLE ${NODE_TABLE_NAME} (
|
|||
startLine INT64,
|
||||
endLine INT64,
|
||||
content STRING,
|
||||
embedding FLOAT[384],
|
||||
PRIMARY KEY (id)
|
||||
)`;
|
||||
|
||||
/**
|
||||
* Create vector index for semantic search
|
||||
* Uses HNSW (Hierarchical Navigable Small World) algorithm with cosine similarity
|
||||
*/
|
||||
export const CREATE_VECTOR_INDEX_QUERY = `
|
||||
CALL CREATE_VECTOR_INDEX('${NODE_TABLE_NAME}', 'code_embedding_idx', 'embedding', metric := 'cosine')
|
||||
`;
|
||||
|
||||
/**
|
||||
* Edge table schema
|
||||
* Stores all relationships: CALLS, IMPORTS, CONTAINS, DEFINES
|
||||
|
|
|
|||
|
|
@ -6,9 +6,11 @@ import { createKnowledgeGraph } from '../core/graph/graph';
|
|||
import { DEFAULT_VISIBLE_LABELS } from '../lib/constants';
|
||||
import type { IngestionWorkerApi } from '../workers/ingestion.worker';
|
||||
import type { FileEntry } from '../services/zip';
|
||||
import type { EmbeddingProgress, SemanticSearchResult } from '../core/embeddings/types';
|
||||
|
||||
export type ViewMode = 'onboarding' | 'loading' | 'exploring';
|
||||
export type RightPanelTab = 'code' | 'chat';
|
||||
export type EmbeddingStatus = 'idle' | 'loading' | 'embedding' | 'indexing' | 'ready' | 'error';
|
||||
|
||||
export interface QueryResult {
|
||||
rows: Record<string, any>[];
|
||||
|
|
@ -67,6 +69,16 @@ interface AppState {
|
|||
runPipelineFromFiles: (files: FileEntry[], onProgress: (p: PipelineProgress) => void) => Promise<PipelineResult>;
|
||||
runQuery: (cypher: string) => Promise<any[]>;
|
||||
isDatabaseReady: () => Promise<boolean>;
|
||||
|
||||
// Embedding state
|
||||
embeddingStatus: EmbeddingStatus;
|
||||
embeddingProgress: EmbeddingProgress | null;
|
||||
|
||||
// Embedding methods
|
||||
startEmbeddings: () => Promise<void>;
|
||||
semanticSearch: (query: string, k?: number) => Promise<SemanticSearchResult[]>;
|
||||
semanticSearchWithContext: (query: string, k?: number, hops?: number) => Promise<any[]>;
|
||||
isEmbeddingReady: boolean;
|
||||
}
|
||||
|
||||
const AppStateContext = createContext<AppState | null>(null);
|
||||
|
|
@ -116,6 +128,10 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
|
|||
|
||||
// Project info
|
||||
const [projectName, setProjectName] = useState<string>('');
|
||||
|
||||
// Embedding state
|
||||
const [embeddingStatus, setEmbeddingStatus] = useState<EmbeddingStatus>('idle');
|
||||
const [embeddingProgress, setEmbeddingProgress] = useState<EmbeddingProgress | null>(null);
|
||||
|
||||
// Worker (single instance shared across app)
|
||||
const workerRef = useRef<Worker | null>(null);
|
||||
|
|
@ -177,6 +193,64 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
|
|||
}
|
||||
}, []);
|
||||
|
||||
// Embedding methods
|
||||
const startEmbeddings = useCallback(async (): Promise<void> => {
|
||||
const api = apiRef.current;
|
||||
if (!api) throw new Error('Worker not initialized');
|
||||
|
||||
setEmbeddingStatus('loading');
|
||||
setEmbeddingProgress(null);
|
||||
|
||||
try {
|
||||
const proxiedOnProgress = Comlink.proxy((progress: EmbeddingProgress) => {
|
||||
setEmbeddingProgress(progress);
|
||||
|
||||
// Update status based on phase
|
||||
switch (progress.phase) {
|
||||
case 'loading-model':
|
||||
setEmbeddingStatus('loading');
|
||||
break;
|
||||
case 'embedding':
|
||||
setEmbeddingStatus('embedding');
|
||||
break;
|
||||
case 'indexing':
|
||||
setEmbeddingStatus('indexing');
|
||||
break;
|
||||
case 'ready':
|
||||
setEmbeddingStatus('ready');
|
||||
break;
|
||||
case 'error':
|
||||
setEmbeddingStatus('error');
|
||||
break;
|
||||
}
|
||||
});
|
||||
|
||||
await api.startEmbeddingPipeline(proxiedOnProgress);
|
||||
} catch (error) {
|
||||
setEmbeddingStatus('error');
|
||||
throw error;
|
||||
}
|
||||
}, []);
|
||||
|
||||
const semanticSearch = useCallback(async (
|
||||
query: string,
|
||||
k: number = 10
|
||||
): Promise<SemanticSearchResult[]> => {
|
||||
const api = apiRef.current;
|
||||
if (!api) throw new Error('Worker not initialized');
|
||||
return api.semanticSearch(query, k);
|
||||
}, []);
|
||||
|
||||
const semanticSearchWithContext = useCallback(async (
|
||||
query: string,
|
||||
k: number = 5,
|
||||
hops: number = 2
|
||||
): Promise<any[]> => {
|
||||
const api = apiRef.current;
|
||||
if (!api) throw new Error('Worker not initialized');
|
||||
return api.semanticSearchWithContext(query, k, hops);
|
||||
}, []);
|
||||
|
||||
const toggleLabelVisibility = useCallback((label: NodeLabel) => {
|
||||
setVisibleLabels(prev => {
|
||||
if (prev.includes(label)) {
|
||||
|
|
@ -219,6 +293,13 @@ export const AppStateProvider = ({ children }: { children: ReactNode }) => {
|
|||
runPipelineFromFiles,
|
||||
runQuery,
|
||||
isDatabaseReady,
|
||||
// Embedding state and methods
|
||||
embeddingStatus,
|
||||
embeddingProgress,
|
||||
startEmbeddings,
|
||||
semanticSearch,
|
||||
semanticSearchWithContext,
|
||||
isEmbeddingReady: embeddingStatus === 'ready',
|
||||
};
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -2,6 +2,14 @@ import * as Comlink from 'comlink';
|
|||
import { runIngestionPipeline, runPipelineFromFiles } from '../core/ingestion/pipeline';
|
||||
import { PipelineProgress, SerializablePipelineResult, serializePipelineResult } from '../types/pipeline';
|
||||
import { FileEntry } from '../services/zip';
|
||||
import {
|
||||
runEmbeddingPipeline,
|
||||
semanticSearch as doSemanticSearch,
|
||||
semanticSearchWithContext as doSemanticSearchWithContext,
|
||||
type EmbeddingProgressCallback,
|
||||
} from '../core/embeddings/embedding-pipeline';
|
||||
import { isEmbedderReady, disposeEmbedder } from '../core/embeddings/embedder';
|
||||
import type { EmbeddingProgress, SemanticSearchResult } from '../core/embeddings/types';
|
||||
|
||||
// Lazy import for Kuzu to avoid breaking worker if SharedArrayBuffer unavailable
|
||||
let kuzuAdapter: typeof import('../core/kuzu/kuzu-adapter') | null = null;
|
||||
|
|
@ -12,6 +20,10 @@ const getKuzuAdapter = async () => {
|
|||
return kuzuAdapter;
|
||||
};
|
||||
|
||||
// Embedding state
|
||||
let embeddingProgress: EmbeddingProgress | null = null;
|
||||
let isEmbeddingComplete = false;
|
||||
|
||||
/**
|
||||
* Worker API exposed via Comlink
|
||||
*
|
||||
|
|
@ -49,8 +61,8 @@ const workerApi = {
|
|||
await kuzu.loadGraphToKuzu(result.graph, result.fileContents);
|
||||
|
||||
if (import.meta.env.DEV) {
|
||||
const stats = await kuzu.getKuzuStats();
|
||||
console.log('KuzuDB loaded:', stats);
|
||||
const stats = await kuzu.getKuzuStats();
|
||||
console.log('KuzuDB loaded:', stats);
|
||||
}
|
||||
} catch {
|
||||
// KuzuDB is optional - silently continue without it
|
||||
|
|
@ -145,6 +157,115 @@ const workerApi = {
|
|||
// Convert to serializable format for transfer back to main thread
|
||||
return serializePipelineResult(result);
|
||||
},
|
||||
|
||||
// ============================================================
|
||||
// Embedding Pipeline Methods
|
||||
// ============================================================
|
||||
|
||||
/**
|
||||
* Start the embedding pipeline in the background
|
||||
* Generates embeddings for all embeddable nodes and creates vector index
|
||||
* @param onProgress - Proxied callback for embedding progress updates
|
||||
*/
|
||||
async startEmbeddingPipeline(
|
||||
onProgress: (progress: EmbeddingProgress) => void
|
||||
): Promise<void> {
|
||||
const kuzu = await getKuzuAdapter();
|
||||
if (!kuzu.isKuzuReady()) {
|
||||
throw new Error('Database not ready. Please load a repository first.');
|
||||
}
|
||||
|
||||
// Reset state
|
||||
embeddingProgress = null;
|
||||
isEmbeddingComplete = false;
|
||||
|
||||
const progressCallback: EmbeddingProgressCallback = (progress) => {
|
||||
embeddingProgress = progress;
|
||||
if (progress.phase === 'ready') {
|
||||
isEmbeddingComplete = true;
|
||||
}
|
||||
onProgress(progress);
|
||||
};
|
||||
|
||||
await runEmbeddingPipeline(kuzu.executeQuery, progressCallback);
|
||||
},
|
||||
|
||||
/**
|
||||
* Perform semantic search on the codebase
|
||||
* @param query - Natural language search query
|
||||
* @param k - Number of results to return (default: 10)
|
||||
* @param maxDistance - Maximum distance threshold (default: 0.5)
|
||||
* @returns Array of search results ordered by relevance
|
||||
*/
|
||||
async semanticSearch(
|
||||
query: string,
|
||||
k: number = 10,
|
||||
maxDistance: number = 0.5
|
||||
): Promise<SemanticSearchResult[]> {
|
||||
const kuzu = await getKuzuAdapter();
|
||||
if (!kuzu.isKuzuReady()) {
|
||||
throw new Error('Database not ready. Please load a repository first.');
|
||||
}
|
||||
if (!isEmbeddingComplete) {
|
||||
throw new Error('Embeddings not ready. Please wait for embedding pipeline to complete.');
|
||||
}
|
||||
|
||||
return doSemanticSearch(kuzu.executeQuery, query, k, maxDistance);
|
||||
},
|
||||
|
||||
/**
|
||||
* Perform semantic search with graph expansion
|
||||
* Finds similar nodes AND their connections
|
||||
* @param query - Natural language search query
|
||||
* @param k - Number of initial results (default: 5)
|
||||
* @param hops - Number of graph hops to expand (default: 2)
|
||||
* @returns Search results with connected nodes
|
||||
*/
|
||||
async semanticSearchWithContext(
|
||||
query: string,
|
||||
k: number = 5,
|
||||
hops: number = 2
|
||||
): Promise<any[]> {
|
||||
const kuzu = await getKuzuAdapter();
|
||||
if (!kuzu.isKuzuReady()) {
|
||||
throw new Error('Database not ready. Please load a repository first.');
|
||||
}
|
||||
if (!isEmbeddingComplete) {
|
||||
throw new Error('Embeddings not ready. Please wait for embedding pipeline to complete.');
|
||||
}
|
||||
|
||||
return doSemanticSearchWithContext(kuzu.executeQuery, query, k, hops);
|
||||
},
|
||||
|
||||
/**
|
||||
* Check if the embedding model is loaded and ready
|
||||
*/
|
||||
isEmbeddingModelReady(): boolean {
|
||||
return isEmbedderReady();
|
||||
},
|
||||
|
||||
/**
|
||||
* Check if embeddings are fully generated and indexed
|
||||
*/
|
||||
isEmbeddingComplete(): boolean {
|
||||
return isEmbeddingComplete;
|
||||
},
|
||||
|
||||
/**
|
||||
* Get current embedding progress
|
||||
*/
|
||||
getEmbeddingProgress(): EmbeddingProgress | null {
|
||||
return embeddingProgress;
|
||||
},
|
||||
|
||||
/**
|
||||
* Cleanup embedding model resources
|
||||
*/
|
||||
async disposeEmbeddingModel(): Promise<void> {
|
||||
await disposeEmbedder();
|
||||
isEmbeddingComplete = false;
|
||||
embeddingProgress = null;
|
||||
},
|
||||
};
|
||||
|
||||
// Expose the worker API to the main thread
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue