@retrivora-ai/rag-engine 0.4.3 → 0.4.5
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/{DocumentChunker-BICIjSuG.d.mts → DocumentChunker-3yElxTO3.d.mts} +9 -2
- package/dist/{DocumentChunker-BICIjSuG.d.ts → DocumentChunker-3yElxTO3.d.ts} +9 -2
- package/dist/{MongoDBProvider-ZKW34AEL.mjs → MongoDBProvider-RE3Q5S5B.mjs} +1 -1
- package/dist/{RagConfig-CVt24lbC.d.ts → RagConfig-BgRDL9Vy.d.mts} +37 -1
- package/dist/{RagConfig-CVt24lbC.d.mts → RagConfig-BgRDL9Vy.d.ts} +37 -1
- package/dist/SimpleGraphProvider-M6T7SE7D.mjs +62 -0
- package/dist/{chunk-IWHCAQEA.mjs → chunk-PQKTC73Y.mjs} +1 -1
- package/dist/{chunk-UUZ3F4WK.mjs → chunk-PRC5CZIZ.mjs} +197 -42
- package/dist/handlers/index.d.mts +2 -2
- package/dist/handlers/index.d.ts +2 -2
- package/dist/handlers/index.js +272 -42
- package/dist/handlers/index.mjs +1 -1
- package/dist/{index-OI2--lvT.d.mts → index-7qeLTPBL.d.mts} +1 -1
- package/dist/{index-D1hoNXMT.d.ts → index-DowY4_K0.d.ts} +1 -1
- package/dist/index.d.mts +15 -5
- package/dist/index.d.ts +15 -5
- package/dist/index.js +161 -1
- package/dist/index.mjs +158 -1
- package/dist/server.d.mts +52 -12
- package/dist/server.d.ts +52 -12
- package/dist/server.js +272 -42
- package/dist/server.mjs +2 -2
- package/package.json +1 -1
- package/src/components/ConfigProvider.tsx +1 -0
- package/src/components/DocumentUpload.tsx +192 -0
- package/src/config/RagConfig.ts +27 -0
- package/src/config/constants.ts +7 -0
- package/src/config/serverConfig.ts +1 -0
- package/src/core/Pipeline.ts +89 -10
- package/src/core/ProviderRegistry.ts +40 -8
- package/src/index.ts +1 -0
- package/src/providers/graphdb/BaseGraphProvider.ts +43 -0
- package/src/providers/graphdb/SimpleGraphProvider.ts +66 -0
- package/src/providers/vectordb/MongoDBProvider.ts +3 -3
- package/src/rag/DocumentChunker.ts +77 -34
- package/src/rag/EntityExtractor.ts +43 -0
- package/src/types/index.ts +19 -0
- package/src/utils/DocumentParser.ts +1 -1
package/src/config/RagConfig.ts
CHANGED
|
@@ -12,6 +12,7 @@ import {
|
|
|
12
12
|
VECTOR_DB_PROVIDERS,
|
|
13
13
|
LLM_PROVIDERS,
|
|
14
14
|
EMBEDDING_PROVIDERS,
|
|
15
|
+
GRAPH_DB_PROVIDERS,
|
|
15
16
|
} from './constants';
|
|
16
17
|
|
|
17
18
|
export type VectorDBProvider = typeof VECTOR_DB_PROVIDERS[number];
|
|
@@ -42,6 +43,19 @@ export interface VectorDBConfig {
|
|
|
42
43
|
options: Record<string, unknown>;
|
|
43
44
|
}
|
|
44
45
|
|
|
46
|
+
// ---------------------------------------------------------------------------
|
|
47
|
+
// Graph DB
|
|
48
|
+
// ---------------------------------------------------------------------------
|
|
49
|
+
|
|
50
|
+
export type GraphDBProvider = typeof GRAPH_DB_PROVIDERS[number];
|
|
51
|
+
|
|
52
|
+
export interface GraphDBConfig {
|
|
53
|
+
/** Which graph database to use */
|
|
54
|
+
provider: GraphDBProvider;
|
|
55
|
+
/** Provider-specific options (URI, credentials, etc.) */
|
|
56
|
+
options: Record<string, unknown>;
|
|
57
|
+
}
|
|
58
|
+
|
|
45
59
|
// ---------------------------------------------------------------------------
|
|
46
60
|
// LLM Provider
|
|
47
61
|
// ---------------------------------------------------------------------------
|
|
@@ -128,6 +142,8 @@ export interface UIConfig {
|
|
|
128
142
|
visualStyle?: 'glass' | 'solid';
|
|
129
143
|
/** Border radius: 'none', 'sm', 'md', 'lg', 'xl', 'full' */
|
|
130
144
|
borderRadius?: 'none' | 'sm' | 'md' | 'lg' | 'xl' | 'full';
|
|
145
|
+
/** Whether to allow file uploads directly from the UI */
|
|
146
|
+
allowUpload?: boolean;
|
|
131
147
|
}
|
|
132
148
|
|
|
133
149
|
// ---------------------------------------------------------------------------
|
|
@@ -143,6 +159,14 @@ export interface RAGConfig {
|
|
|
143
159
|
chunkSize?: number;
|
|
144
160
|
/** Overlap between adjacent chunks in tokens */
|
|
145
161
|
chunkOverlap?: number;
|
|
162
|
+
/** Characters used to split text, in order of priority */
|
|
163
|
+
separators?: string[];
|
|
164
|
+
/** Whether to use query transformation (e.g. HyDE) */
|
|
165
|
+
useQueryTransformation?: boolean;
|
|
166
|
+
/** Whether to use graph-based retrieval */
|
|
167
|
+
useGraphRetrieval?: boolean;
|
|
168
|
+
/** Whether to perform reranking on retrieved results */
|
|
169
|
+
useReranking?: boolean;
|
|
146
170
|
}
|
|
147
171
|
|
|
148
172
|
// ---------------------------------------------------------------------------
|
|
@@ -170,4 +194,7 @@ export interface RagConfig {
|
|
|
170
194
|
|
|
171
195
|
/** Optional RAG pipeline tuning knobs */
|
|
172
196
|
rag?: RAGConfig;
|
|
197
|
+
|
|
198
|
+
/** Optional Graph database configuration */
|
|
199
|
+
graphDb?: GraphDBConfig;
|
|
173
200
|
}
|
package/src/config/constants.ts
CHANGED
|
@@ -39,6 +39,13 @@ export const EMBEDDING_PROVIDERS = [
|
|
|
39
39
|
'custom',
|
|
40
40
|
] as const;
|
|
41
41
|
|
|
42
|
+
export const GRAPH_DB_PROVIDERS = [
|
|
43
|
+
'neo4j',
|
|
44
|
+
'memgraph',
|
|
45
|
+
'simple',
|
|
46
|
+
'custom',
|
|
47
|
+
] as const;
|
|
48
|
+
|
|
42
49
|
export const PROVIDERS_WITH_EMBEDDINGS = [
|
|
43
50
|
'openai',
|
|
44
51
|
'ollama',
|
|
@@ -169,6 +169,7 @@ export function getRagConfig(env: Record<string, string | undefined> = process.e
|
|
|
169
169
|
"Hello! I'm your AI assistant. Ask me anything about your documents.",
|
|
170
170
|
visualStyle: (readString(env, 'NEXT_PUBLIC_UI_VISUAL_STYLE') ?? readString(env, 'UI_VISUAL_STYLE') ?? 'glass') as 'glass' | 'solid',
|
|
171
171
|
borderRadius: (readString(env, 'NEXT_PUBLIC_UI_BORDER_RADIUS') ?? readString(env, 'UI_BORDER_RADIUS') ?? 'xl') as 'none' | 'sm' | 'md' | 'lg' | 'xl' | 'full',
|
|
172
|
+
allowUpload: (readString(env, 'NEXT_PUBLIC_ALLOW_UPLOAD') ?? readString(env, 'UI_ALLOW_UPLOAD') ?? 'false') === 'true',
|
|
172
173
|
},
|
|
173
174
|
rag: {
|
|
174
175
|
topK: readNumber(env, 'RAG_TOP_K', 5),
|
package/src/core/Pipeline.ts
CHANGED
|
@@ -1,11 +1,13 @@
|
|
|
1
1
|
import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
|
|
2
|
+
import { BaseGraphProvider } from '../providers/graphdb/BaseGraphProvider';
|
|
2
3
|
import { ILLMProvider, ChatMessage } from '../llm/ILLMProvider';
|
|
3
4
|
import { DocumentChunker } from '../rag/DocumentChunker';
|
|
5
|
+
import { EntityExtractor } from '../rag/EntityExtractor';
|
|
4
6
|
import { RagConfig } from '../config/RagConfig';
|
|
5
7
|
import { ProviderRegistry } from './ProviderRegistry';
|
|
6
8
|
import { BatchProcessor, BatchOptions } from './BatchProcessor';
|
|
7
9
|
import { EmbeddingStrategyResolver } from '../config/EmbeddingStrategy';
|
|
8
|
-
import { IngestDocument, ChatResponse } from '../types';
|
|
10
|
+
import { IngestDocument, ChatResponse, GraphSearchResult } from '../types';
|
|
9
11
|
|
|
10
12
|
interface QueryFieldHint {
|
|
11
13
|
field?: string;
|
|
@@ -171,9 +173,11 @@ function buildQueryFilter(question: string, hints: QueryFieldHint[]): QueryFilte
|
|
|
171
173
|
*/
|
|
172
174
|
export class Pipeline {
|
|
173
175
|
private vectorDB!: BaseVectorProvider;
|
|
176
|
+
private graphDB?: BaseGraphProvider;
|
|
174
177
|
private llmProvider!: ILLMProvider;
|
|
175
178
|
private embeddingProvider!: ILLMProvider;
|
|
176
179
|
private chunker: DocumentChunker;
|
|
180
|
+
private entityExtractor?: EntityExtractor;
|
|
177
181
|
private config: RagConfig;
|
|
178
182
|
private initialised = false;
|
|
179
183
|
|
|
@@ -202,6 +206,12 @@ export class Pipeline {
|
|
|
202
206
|
this.llmProvider = llmProvider;
|
|
203
207
|
this.embeddingProvider = embeddingProvider;
|
|
204
208
|
|
|
209
|
+
if (this.config.graphDb) {
|
|
210
|
+
this.graphDB = await ProviderRegistry.createGraphProvider(this.config.graphDb);
|
|
211
|
+
await this.graphDB.initialize();
|
|
212
|
+
this.entityExtractor = new EntityExtractor(this.llmProvider);
|
|
213
|
+
}
|
|
214
|
+
|
|
205
215
|
await this.vectorDB.initialize();
|
|
206
216
|
this.initialised = true;
|
|
207
217
|
}
|
|
@@ -270,6 +280,34 @@ export class Pipeline {
|
|
|
270
280
|
docId: doc.docId,
|
|
271
281
|
chunksIngested: upsertResult.totalProcessed,
|
|
272
282
|
});
|
|
283
|
+
|
|
284
|
+
// Graph ingestion
|
|
285
|
+
// Graph ingestion - Batched to prevent timeouts and LLM overload
|
|
286
|
+
if (this.graphDB && this.entityExtractor) {
|
|
287
|
+
console.log(`[Pipeline] Extracting entities for doc ${doc.docId} (${chunks.length} chunks)...`);
|
|
288
|
+
|
|
289
|
+
const extractionOptions: BatchOptions = {
|
|
290
|
+
batchSize: 2, // Low concurrency for LLM extraction
|
|
291
|
+
maxRetries: 1,
|
|
292
|
+
initialDelayMs: 500,
|
|
293
|
+
};
|
|
294
|
+
|
|
295
|
+
await BatchProcessor.processBatch(
|
|
296
|
+
chunks,
|
|
297
|
+
async (batch) => {
|
|
298
|
+
for (const chunk of batch) {
|
|
299
|
+
try {
|
|
300
|
+
const { nodes, edges } = await this.entityExtractor!.extract(chunk.content);
|
|
301
|
+
if (nodes.length > 0) await this.graphDB!.addNodes(nodes);
|
|
302
|
+
if (edges.length > 0) await this.graphDB!.addEdges(edges);
|
|
303
|
+
} catch (err) {
|
|
304
|
+
console.warn(`[Pipeline] Entity extraction failed for chunk:`, err);
|
|
305
|
+
}
|
|
306
|
+
}
|
|
307
|
+
},
|
|
308
|
+
extractionOptions
|
|
309
|
+
);
|
|
310
|
+
}
|
|
273
311
|
} catch (error) {
|
|
274
312
|
console.error(`[Pipeline] Failed to ingest document ${doc.docId}:`, error);
|
|
275
313
|
results.push({ docId: doc.docId, chunksIngested: 0 });
|
|
@@ -286,29 +324,70 @@ export class Pipeline {
|
|
|
286
324
|
const scoreThreshold = this.config.rag?.scoreThreshold ?? 0.0;
|
|
287
325
|
|
|
288
326
|
try {
|
|
289
|
-
|
|
327
|
+
let searchQuery = question;
|
|
328
|
+
|
|
329
|
+
// 1. Query Transformation
|
|
330
|
+
if (this.config.rag?.useQueryTransformation) {
|
|
331
|
+
searchQuery = await this.rewriteQuery(question, history);
|
|
332
|
+
}
|
|
333
|
+
|
|
334
|
+
// 2. Vector Retrieval
|
|
335
|
+
const queryVector = await this.embeddingProvider.embed(searchQuery, { taskType: 'query' });
|
|
290
336
|
const fieldHints = extractQueryFieldHints(question);
|
|
291
|
-
// Build a provider-agnostic filter from hints + original query text. Providers
|
|
292
|
-
// may choose to use `queryText` for hybrid keyword matching and `metadata`
|
|
293
|
-
// for exact/filtered metadata queries.
|
|
294
337
|
const filter = buildQueryFilter(question, fieldHints) as Record<string, unknown>;
|
|
295
|
-
|
|
296
|
-
// Keep the raw hints available as well for providers that want richer parsing
|
|
297
338
|
filter.__entityHints = fieldHints;
|
|
298
339
|
|
|
299
340
|
const rawMatches = await this.vectorDB.query(queryVector, topK, ns, filter);
|
|
300
|
-
|
|
301
341
|
const sources = rawMatches.filter((m) => m.score >= scoreThreshold);
|
|
302
|
-
|
|
342
|
+
|
|
343
|
+
// 3. Graph Retrieval
|
|
344
|
+
let graphData: GraphSearchResult | undefined;
|
|
345
|
+
if (this.graphDB && this.config.rag?.useGraphRetrieval) {
|
|
346
|
+
graphData = await this.graphDB.query(searchQuery);
|
|
347
|
+
}
|
|
348
|
+
|
|
349
|
+
// 4. Context Augmentation
|
|
350
|
+
let context = sources.length
|
|
303
351
|
? sources.map((m, i) => `[Source ${i + 1}]\n${m.content}`).join('\n\n---\n\n')
|
|
304
352
|
: 'No relevant context found.';
|
|
305
353
|
|
|
354
|
+
if (graphData && graphData.nodes.length > 0) {
|
|
355
|
+
const graphContext = graphData.nodes.map(n =>
|
|
356
|
+
`Entity: ${n.label} (${n.id})${n.properties ? ' - ' + JSON.stringify(n.properties) : ''}`
|
|
357
|
+
).join('\n');
|
|
358
|
+
context = `GRAPH KNOWLEDGE:\n${graphContext}\n\nVECTOR CONTEXT:\n${context}`;
|
|
359
|
+
}
|
|
360
|
+
|
|
361
|
+
// 5. Generation
|
|
306
362
|
const messages: ChatMessage[] = [...history, { role: 'user', content: question }];
|
|
307
363
|
const reply = await this.llmProvider.chat(messages, context);
|
|
308
364
|
|
|
309
|
-
return { reply, sources };
|
|
365
|
+
return { reply, sources, graphData };
|
|
310
366
|
} catch (error) {
|
|
311
367
|
throw new Error(`[Pipeline] Chat failed: ${error instanceof Error ? error.message : String(error)}`);
|
|
312
368
|
}
|
|
313
369
|
}
|
|
370
|
+
|
|
371
|
+
/**
|
|
372
|
+
* Rewrite the user query for better retrieval performance.
|
|
373
|
+
*/
|
|
374
|
+
private async rewriteQuery(question: string, history: ChatMessage[]): Promise<string> {
|
|
375
|
+
const prompt = `
|
|
376
|
+
Given the following conversation history and a new question, rewrite the question to be a better search query for a vector database.
|
|
377
|
+
Focus on extracting the core intent and entities. Do not answer the question, just rewrite it.
|
|
378
|
+
|
|
379
|
+
History:
|
|
380
|
+
${history.map(m => `${m.role}: ${m.content}`).join('\n')}
|
|
381
|
+
|
|
382
|
+
New Question: ${question}
|
|
383
|
+
|
|
384
|
+
Optimized Search Query:`;
|
|
385
|
+
|
|
386
|
+
const rewrite = await this.llmProvider.chat([
|
|
387
|
+
{ role: 'system', content: 'You are an assistant that optimizes search queries for RAG systems.' },
|
|
388
|
+
{ role: 'user', content: prompt }
|
|
389
|
+
], '');
|
|
390
|
+
|
|
391
|
+
return rewrite.trim() || question;
|
|
392
|
+
}
|
|
314
393
|
}
|
|
@@ -1,5 +1,6 @@
|
|
|
1
1
|
import { BaseVectorProvider } from '../providers/vectordb/BaseVectorProvider';
|
|
2
|
-
import {
|
|
2
|
+
import { BaseGraphProvider } from '@/providers/graphdb/BaseGraphProvider';
|
|
3
|
+
import { VectorDBConfig, GraphDBConfig, LLMConfig, EmbeddingConfig } from '../config/RagConfig';
|
|
3
4
|
import { LLMFactory } from '../llm/LLMFactory';
|
|
4
5
|
import { ILLMProvider } from '../llm/ILLMProvider';
|
|
5
6
|
|
|
@@ -12,18 +13,19 @@ import { ILLMProvider } from '../llm/ILLMProvider';
|
|
|
12
13
|
*/
|
|
13
14
|
export class ProviderRegistry {
|
|
14
15
|
private static vectorProviders: Record<string, new (config: VectorDBConfig) => BaseVectorProvider> = {};
|
|
16
|
+
private static graphProviders: Record<string, new (config: GraphDBConfig) => BaseGraphProvider> = {};
|
|
15
17
|
|
|
16
|
-
/**
|
|
17
|
-
* Register a custom vector provider class by name.
|
|
18
|
-
* The name must match the provider value used in VectorDBConfig.provider.
|
|
19
|
-
*
|
|
20
|
-
* @example
|
|
21
|
-
* ProviderRegistry.registerVectorProvider('my-db', MyCustomProvider);
|
|
22
|
-
*/
|
|
23
18
|
static registerVectorProvider(name: string, providerClass: new (config: VectorDBConfig) => BaseVectorProvider) {
|
|
24
19
|
this.vectorProviders[name] = providerClass;
|
|
25
20
|
}
|
|
26
21
|
|
|
22
|
+
/**
|
|
23
|
+
* Register a custom graph provider class by name.
|
|
24
|
+
*/
|
|
25
|
+
static registerGraphProvider(name: string, providerClass: new (config: GraphDBConfig) => BaseGraphProvider) {
|
|
26
|
+
this.graphProviders[name] = providerClass;
|
|
27
|
+
}
|
|
28
|
+
|
|
27
29
|
/**
|
|
28
30
|
* Creates a vector database provider based on the configuration.
|
|
29
31
|
* Built-in providers are dynamically imported to avoid bundling all SDKs.
|
|
@@ -85,6 +87,36 @@ export class ProviderRegistry {
|
|
|
85
87
|
}
|
|
86
88
|
}
|
|
87
89
|
|
|
90
|
+
/**
|
|
91
|
+
* Creates a graph database provider based on the configuration.
|
|
92
|
+
*/
|
|
93
|
+
static async createGraphProvider(config: GraphDBConfig): Promise<BaseGraphProvider> {
|
|
94
|
+
const { provider } = config;
|
|
95
|
+
|
|
96
|
+
// Custom registered provider takes priority
|
|
97
|
+
if (this.graphProviders[provider]) {
|
|
98
|
+
return new this.graphProviders[provider](config);
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
// Built-in providers — lazy-loaded
|
|
102
|
+
switch (provider) {
|
|
103
|
+
case 'neo4j': {
|
|
104
|
+
// Placeholder for real Neo4j provider
|
|
105
|
+
throw new Error('[ProviderRegistry] Neo4j provider not implemented yet.');
|
|
106
|
+
}
|
|
107
|
+
case 'simple': {
|
|
108
|
+
const { SimpleGraphProvider } = await import('@/providers/graphdb/SimpleGraphProvider');
|
|
109
|
+
return new SimpleGraphProvider(config);
|
|
110
|
+
}
|
|
111
|
+
default:
|
|
112
|
+
throw new Error(
|
|
113
|
+
`[ProviderRegistry] Unsupported graph provider: "${provider}". ` +
|
|
114
|
+
`Built-in providers: simple. ` +
|
|
115
|
+
`For custom providers, call ProviderRegistry.registerGraphProvider("${provider}", YourClass).`
|
|
116
|
+
);
|
|
117
|
+
}
|
|
118
|
+
}
|
|
119
|
+
|
|
88
120
|
/**
|
|
89
121
|
* Creates an LLM provider based on the configuration.
|
|
90
122
|
*/
|
package/src/index.ts
CHANGED
|
@@ -6,6 +6,7 @@
|
|
|
6
6
|
// ── React Components ─────────────────────────────────────────
|
|
7
7
|
export { ChatWidget } from './components/ChatWidget';
|
|
8
8
|
export { ChatWindow } from './components/ChatWindow';
|
|
9
|
+
export { DocumentUpload } from './components/DocumentUpload';
|
|
9
10
|
export { MessageBubble } from './components/MessageBubble';
|
|
10
11
|
export { SourceCard } from './components/SourceCard';
|
|
11
12
|
export { ConfigProvider, useConfig } from './components/ConfigProvider';
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
import { GraphDBConfig } from '../../config/RagConfig';
|
|
2
|
+
import { GraphNode, Edge, GraphSearchResult } from '../../types';
|
|
3
|
+
|
|
4
|
+
/**
|
|
5
|
+
* BaseGraphProvider — Abstract base class for Graph DB providers.
|
|
6
|
+
*/
|
|
7
|
+
export abstract class BaseGraphProvider {
|
|
8
|
+
protected config: GraphDBConfig;
|
|
9
|
+
|
|
10
|
+
constructor(config: GraphDBConfig) {
|
|
11
|
+
this.config = config;
|
|
12
|
+
}
|
|
13
|
+
|
|
14
|
+
/**
|
|
15
|
+
* Initialise connection to the graph database.
|
|
16
|
+
*/
|
|
17
|
+
abstract initialize(): Promise<void>;
|
|
18
|
+
|
|
19
|
+
/**
|
|
20
|
+
* Add nodes to the graph.
|
|
21
|
+
*/
|
|
22
|
+
abstract addNodes(nodes: GraphNode[]): Promise<void>;
|
|
23
|
+
|
|
24
|
+
/**
|
|
25
|
+
* Add edges (relationships) to the graph.
|
|
26
|
+
*/
|
|
27
|
+
abstract addEdges(edges: Edge[]): Promise<void>;
|
|
28
|
+
|
|
29
|
+
/**
|
|
30
|
+
* Query the graph for relevant nodes and edges.
|
|
31
|
+
*/
|
|
32
|
+
abstract query(queryText: string, limit?: number): Promise<GraphSearchResult>;
|
|
33
|
+
|
|
34
|
+
/**
|
|
35
|
+
* Check if the database is reachable.
|
|
36
|
+
*/
|
|
37
|
+
abstract ping(): Promise<boolean>;
|
|
38
|
+
|
|
39
|
+
/**
|
|
40
|
+
* Close connection.
|
|
41
|
+
*/
|
|
42
|
+
abstract disconnect(): Promise<void>;
|
|
43
|
+
}
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
import { BaseGraphProvider } from './BaseGraphProvider';
|
|
2
|
+
import { GraphNode, Edge, GraphSearchResult } from '../../types';
|
|
3
|
+
|
|
4
|
+
/**
|
|
5
|
+
* SimpleGraphProvider — In-memory graph storage for demonstration and local dev.
|
|
6
|
+
*/
|
|
7
|
+
export class SimpleGraphProvider extends BaseGraphProvider {
|
|
8
|
+
private nodes: Map<string, GraphNode> = new Map();
|
|
9
|
+
private edges: Edge[] = [];
|
|
10
|
+
|
|
11
|
+
async initialize(): Promise<void> {
|
|
12
|
+
console.log('[SimpleGraphProvider] Initialised in-memory graph store.');
|
|
13
|
+
}
|
|
14
|
+
|
|
15
|
+
async addNodes(nodes: GraphNode[]): Promise<void> {
|
|
16
|
+
for (const node of nodes) {
|
|
17
|
+
this.nodes.set(node.id, node);
|
|
18
|
+
}
|
|
19
|
+
}
|
|
20
|
+
|
|
21
|
+
async addEdges(edges: Edge[]): Promise<void> {
|
|
22
|
+
this.edges.push(...edges);
|
|
23
|
+
}
|
|
24
|
+
|
|
25
|
+
async query(queryText: string, limit = 5): Promise<GraphSearchResult> {
|
|
26
|
+
const q = queryText.toLowerCase();
|
|
27
|
+
|
|
28
|
+
// Simple keyword matching for demonstration
|
|
29
|
+
const matchedNodes = Array.from(this.nodes.values()).filter(node =>
|
|
30
|
+
node.id.toLowerCase().includes(q) ||
|
|
31
|
+
node.label.toLowerCase().includes(q) ||
|
|
32
|
+
JSON.stringify(node.properties).toLowerCase().includes(q)
|
|
33
|
+
).slice(0, limit);
|
|
34
|
+
|
|
35
|
+
const matchedNodeIds = new Set(matchedNodes.map(n => n.id));
|
|
36
|
+
const matchedEdges = this.edges.filter(edge =>
|
|
37
|
+
matchedNodeIds.has(edge.source) || matchedNodeIds.has(edge.target)
|
|
38
|
+
);
|
|
39
|
+
|
|
40
|
+
// Also include connected nodes
|
|
41
|
+
for (const edge of matchedEdges) {
|
|
42
|
+
if (!matchedNodeIds.has(edge.source)) {
|
|
43
|
+
const source = this.nodes.get(edge.source);
|
|
44
|
+
if (source) matchedNodes.push(source);
|
|
45
|
+
}
|
|
46
|
+
if (!matchedNodeIds.has(edge.target)) {
|
|
47
|
+
const target = this.nodes.get(edge.target);
|
|
48
|
+
if (target) matchedNodes.push(target);
|
|
49
|
+
}
|
|
50
|
+
}
|
|
51
|
+
|
|
52
|
+
return {
|
|
53
|
+
nodes: Array.from(new Set(matchedNodes)),
|
|
54
|
+
edges: matchedEdges
|
|
55
|
+
};
|
|
56
|
+
}
|
|
57
|
+
|
|
58
|
+
async ping(): Promise<boolean> {
|
|
59
|
+
return true;
|
|
60
|
+
}
|
|
61
|
+
|
|
62
|
+
async disconnect(): Promise<void> {
|
|
63
|
+
this.nodes.clear();
|
|
64
|
+
this.edges = [];
|
|
65
|
+
}
|
|
66
|
+
}
|
|
@@ -94,10 +94,10 @@ export class MongoDBProvider extends BaseVectorProvider {
|
|
|
94
94
|
|
|
95
95
|
const results = await this.collection!.aggregate(pipeline).toArray();
|
|
96
96
|
|
|
97
|
-
return results.map((res
|
|
97
|
+
return (results as unknown as Array<{ _id: string; score: number; [key: string]: unknown }>).map((res) => ({
|
|
98
98
|
id: res._id,
|
|
99
|
-
content: res[this.contentKey],
|
|
100
|
-
metadata: res[this.metadataKey],
|
|
99
|
+
content: res[this.contentKey] as string,
|
|
100
|
+
metadata: (res[this.metadataKey] as Record<string, unknown>) || {},
|
|
101
101
|
score: res.score,
|
|
102
102
|
}));
|
|
103
103
|
}
|
|
@@ -23,19 +23,23 @@ export interface ChunkOptions {
|
|
|
23
23
|
docId?: string | number;
|
|
24
24
|
/** Extra metadata to attach to every chunk */
|
|
25
25
|
metadata?: Record<string, unknown>;
|
|
26
|
+
/** Characters used to split text, in order of priority */
|
|
27
|
+
separators?: string[];
|
|
26
28
|
}
|
|
27
29
|
|
|
28
30
|
export class DocumentChunker {
|
|
29
31
|
private readonly chunkSize: number;
|
|
30
32
|
private readonly chunkOverlap: number;
|
|
33
|
+
private readonly separators: string[];
|
|
31
34
|
|
|
32
|
-
constructor(chunkSize = 1000, chunkOverlap = 200) {
|
|
35
|
+
constructor(chunkSize = 1000, chunkOverlap = 200, separators = ['\n\n', '\n', ' ', '']) {
|
|
33
36
|
this.chunkSize = chunkSize;
|
|
34
37
|
this.chunkOverlap = chunkOverlap;
|
|
38
|
+
this.separators = separators;
|
|
35
39
|
}
|
|
36
40
|
|
|
37
41
|
/**
|
|
38
|
-
* Split a single text string into overlapping chunks.
|
|
42
|
+
* Split a single text string into overlapping chunks using a recursive strategy.
|
|
39
43
|
*/
|
|
40
44
|
chunk(text: string, options: ChunkOptions = {}): Chunk[] {
|
|
41
45
|
const {
|
|
@@ -43,53 +47,92 @@ export class DocumentChunker {
|
|
|
43
47
|
chunkOverlap = this.chunkOverlap,
|
|
44
48
|
docId = `doc_${Date.now()}`,
|
|
45
49
|
metadata = {},
|
|
50
|
+
separators = this.separators,
|
|
46
51
|
} = options;
|
|
47
52
|
|
|
48
|
-
|
|
49
|
-
const
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
const sentences = cleaned
|
|
54
|
-
.split(/(?<=[.!?])\s+|\n{2,}/)
|
|
55
|
-
.map((s) => s.trim())
|
|
56
|
-
.filter(Boolean);
|
|
57
|
-
|
|
58
|
-
const chunks: Chunk[] = [];
|
|
59
|
-
let current = '';
|
|
53
|
+
const finalChunks: Chunk[] = [];
|
|
54
|
+
const splits = this.recursiveSplit(text, separators, chunkSize);
|
|
55
|
+
|
|
56
|
+
let currentChunk: string[] = [];
|
|
57
|
+
let currentLength = 0;
|
|
60
58
|
let chunkIndex = 0;
|
|
61
59
|
|
|
62
|
-
for (const
|
|
63
|
-
if (
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
metadata: { ...metadata, docId, chunkIndex: chunkIndex - 1 },
|
|
71
|
-
});
|
|
72
|
-
}
|
|
60
|
+
for (const split of splits) {
|
|
61
|
+
if (currentLength + split.length > chunkSize && currentChunk.length > 0) {
|
|
62
|
+
// Current chunk is full, push it
|
|
63
|
+
finalChunks.push({
|
|
64
|
+
id: `${docId}_chunk_${chunkIndex++}`,
|
|
65
|
+
content: currentChunk.join('').trim(),
|
|
66
|
+
metadata: { ...metadata, docId, chunkIndex: chunkIndex - 1 },
|
|
67
|
+
});
|
|
73
68
|
|
|
74
|
-
// Start new chunk with overlap
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
69
|
+
// Start new chunk with overlap
|
|
70
|
+
// This is a simplified overlap: we keep items from the end of the previous chunk
|
|
71
|
+
// until we reach the overlap size.
|
|
72
|
+
const overlapItems = [];
|
|
73
|
+
let overlapLen = 0;
|
|
74
|
+
for (let i = currentChunk.length - 1; i >= 0; i--) {
|
|
75
|
+
if (overlapLen + currentChunk[i].length <= chunkOverlap) {
|
|
76
|
+
overlapItems.unshift(currentChunk[i]);
|
|
77
|
+
overlapLen += currentChunk[i].length;
|
|
78
|
+
} else {
|
|
79
|
+
break;
|
|
80
|
+
}
|
|
79
81
|
}
|
|
82
|
+
currentChunk = overlapItems;
|
|
83
|
+
currentLength = overlapLen;
|
|
80
84
|
}
|
|
85
|
+
|
|
86
|
+
currentChunk.push(split);
|
|
87
|
+
currentLength += split.length;
|
|
81
88
|
}
|
|
82
89
|
|
|
83
90
|
// Push the last remaining chunk
|
|
84
|
-
if (
|
|
85
|
-
|
|
91
|
+
if (currentChunk.length > 0) {
|
|
92
|
+
finalChunks.push({
|
|
86
93
|
id: `${docId}_chunk_${chunkIndex}`,
|
|
87
|
-
content:
|
|
94
|
+
content: currentChunk.join('').trim(),
|
|
88
95
|
metadata: { ...metadata, docId, chunkIndex },
|
|
89
96
|
});
|
|
90
97
|
}
|
|
91
98
|
|
|
92
|
-
return
|
|
99
|
+
return finalChunks;
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
/**
|
|
103
|
+
* Recursively split text based on separators.
|
|
104
|
+
*/
|
|
105
|
+
private recursiveSplit(text: string, separators: string[], chunkSize: number): string[] {
|
|
106
|
+
const finalSplits: string[] = [];
|
|
107
|
+
|
|
108
|
+
// Find the first separator that exists in the text
|
|
109
|
+
let separator = separators[separators.length - 1]; // default to last (usually empty string)
|
|
110
|
+
let nextSeparators: string[] = [];
|
|
111
|
+
|
|
112
|
+
for (let i = 0; i < separators.length; i++) {
|
|
113
|
+
if (text.includes(separators[i])) {
|
|
114
|
+
separator = separators[i];
|
|
115
|
+
nextSeparators = separators.slice(i + 1);
|
|
116
|
+
break;
|
|
117
|
+
}
|
|
118
|
+
}
|
|
119
|
+
|
|
120
|
+
// Split the text
|
|
121
|
+
const parts = text.split(separator);
|
|
122
|
+
|
|
123
|
+
for (const part of parts) {
|
|
124
|
+
if (part.length <= chunkSize) {
|
|
125
|
+
finalSplits.push(part + separator);
|
|
126
|
+
} else if (nextSeparators.length > 0) {
|
|
127
|
+
// Part is too long, try splitting with next level separators
|
|
128
|
+
finalSplits.push(...this.recursiveSplit(part, nextSeparators, chunkSize));
|
|
129
|
+
} else {
|
|
130
|
+
// No more separators, just hard split (should not happen if empty string is the last separator)
|
|
131
|
+
finalSplits.push(part);
|
|
132
|
+
}
|
|
133
|
+
}
|
|
134
|
+
|
|
135
|
+
return finalSplits;
|
|
93
136
|
}
|
|
94
137
|
|
|
95
138
|
/**
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
import { ILLMProvider } from '../llm/ILLMProvider';
|
|
2
|
+
import { GraphNode, Edge } from '../types';
|
|
3
|
+
|
|
4
|
+
/**
|
|
5
|
+
* EntityExtractor — Uses an LLM to extract entities and relations from text.
|
|
6
|
+
*/
|
|
7
|
+
export class EntityExtractor {
|
|
8
|
+
private llm: ILLMProvider;
|
|
9
|
+
|
|
10
|
+
constructor(llm: ILLMProvider) {
|
|
11
|
+
this.llm = llm;
|
|
12
|
+
}
|
|
13
|
+
|
|
14
|
+
/**
|
|
15
|
+
* Extract nodes and edges from a text chunk.
|
|
16
|
+
*/
|
|
17
|
+
async extract(text: string): Promise<{ nodes: GraphNode[]; edges: Edge[] }> {
|
|
18
|
+
const prompt = `
|
|
19
|
+
Extract entities and relationships from the following text.
|
|
20
|
+
Format the output as JSON with two keys: "nodes" (array of {id, label, properties}) and "edges" (array of {source, target, type, properties}).
|
|
21
|
+
Use the same ID for the same entity.
|
|
22
|
+
|
|
23
|
+
Text:
|
|
24
|
+
"${text}"
|
|
25
|
+
|
|
26
|
+
Output JSON:
|
|
27
|
+
`;
|
|
28
|
+
|
|
29
|
+
const response = await this.llm.chat([
|
|
30
|
+
{ role: 'system', content: 'You are an expert at knowledge graph extraction. Respond only with valid JSON.' },
|
|
31
|
+
{ role: 'user', content: prompt }
|
|
32
|
+
], '');
|
|
33
|
+
|
|
34
|
+
try {
|
|
35
|
+
// Clean up LLM response (sometimes it adds markdown blocks)
|
|
36
|
+
const cleanJson = response.replace(/```json\n?|\n?```/g, '').trim();
|
|
37
|
+
return JSON.parse(cleanJson);
|
|
38
|
+
} catch {
|
|
39
|
+
console.warn('[EntityExtractor] Failed to parse LLM response as JSON:', response);
|
|
40
|
+
return { nodes: [], edges: [] };
|
|
41
|
+
}
|
|
42
|
+
}
|
|
43
|
+
}
|
package/src/types/index.ts
CHANGED
|
@@ -21,4 +21,23 @@ export interface IngestDocument {
|
|
|
21
21
|
export interface ChatResponse {
|
|
22
22
|
reply: string;
|
|
23
23
|
sources: VectorMatch[];
|
|
24
|
+
graphData?: GraphSearchResult;
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
export interface GraphNode {
|
|
28
|
+
id: string;
|
|
29
|
+
label: string;
|
|
30
|
+
properties?: Record<string, unknown>;
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
export interface Edge {
|
|
34
|
+
source: string;
|
|
35
|
+
target: string;
|
|
36
|
+
type: string;
|
|
37
|
+
properties?: Record<string, unknown>;
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
export interface GraphSearchResult {
|
|
41
|
+
nodes: GraphNode[];
|
|
42
|
+
edges: Edge[];
|
|
24
43
|
}
|
|
@@ -35,7 +35,7 @@ export class DocumentParser {
|
|
|
35
35
|
if (extension === 'pdf' || mimeType === 'application/pdf') {
|
|
36
36
|
try {
|
|
37
37
|
// Dynamic import to avoid errors if not installed
|
|
38
|
-
const pdf = await import('pdf-parse
|
|
38
|
+
const pdf = await import('pdf-parse');
|
|
39
39
|
const buffer = Buffer.isBuffer(file) ? file : Buffer.from(await (file as File).arrayBuffer());
|
|
40
40
|
const data = await pdf.default(buffer);
|
|
41
41
|
return data.text;
|