feat: Implemented graph-based RAG
This commit is contained in:
@@ -0,0 +1,371 @@
|
||||
# Graph RAG Design Spec
|
||||
|
||||
## Status: COMPLETE
|
||||
|
||||
### Verified From Code (all claims backed by actual file reads)
|
||||
|
||||
---
|
||||
|
||||
## Goal
|
||||
|
||||
Extend the existing two-signal hybrid search (vector HNSW + BM25 → RRF) to a three-signal hybrid
|
||||
(vector + BM25 + knowledge graph → RRF). The graph captures entity/relationship knowledge extracted
|
||||
from documents at ingestion time via an LLM call per chunk. At query time, graph traversal expands
|
||||
context beyond semantic similarity.
|
||||
|
||||
---
|
||||
|
||||
## Verified Current Architecture
|
||||
|
||||
### `Rag` struct (`src/rag/mod.rs:48`)
|
||||
```rust
|
||||
pub struct Rag {
|
||||
app_config: Arc<AppConfig>,
|
||||
name: String,
|
||||
path: String,
|
||||
embedding_model: Model,
|
||||
hnsw: Hnsw<'static, f32, DistCosine>, // ephemeral, rebuilt on load
|
||||
bm25: SearchEngine<DocumentId>, // ephemeral, rebuilt on load
|
||||
data: RagData, // serialized to YAML
|
||||
last_sources: RwLock<Option<String>>,
|
||||
}
|
||||
```
|
||||
|
||||
### `RagData` struct (`src/rag/mod.rs:892`)
|
||||
```rust
|
||||
pub struct RagData {
|
||||
pub embedding_model: String,
|
||||
pub chunk_size: usize,
|
||||
pub chunk_overlap: usize,
|
||||
pub reranker_model: Option<String>,
|
||||
pub top_k: usize,
|
||||
pub batch_size: Option<usize>,
|
||||
pub next_file_id: FileId,
|
||||
pub document_paths: Vec<String>,
|
||||
pub files: IndexMap<FileId, RagFile>,
|
||||
#[serde(with = "serde_vectors")]
|
||||
pub vectors: IndexMap<DocumentId, Vec<f32>>,
|
||||
}
|
||||
```
|
||||
|
||||
### `RagData::new` callers (both need updating):
|
||||
1. `Rag::init` (`src/rag/mod.rs:219`) — interactive init path
|
||||
2. `Rag::resolve_init_data` (`src/rag/mod.rs:195`) — config-driven init path
|
||||
|
||||
### `Rag::create` (`src/rag/mod.rs:253`) — all init paths converge here:
|
||||
```rust
|
||||
pub fn create(app: &AppConfig, name: &str, path: &Path, data: RagData) -> Result<Self> {
|
||||
let hnsw = data.build_hnsw();
|
||||
let bm25 = data.build_bm25();
|
||||
let embedding_model = Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?;
|
||||
let rag = Rag { app_config: Arc::new(app.clone()), name: name.to_string(),
|
||||
path: path.display().to_string(), data, embedding_model, hnsw, bm25,
|
||||
last_sources: RwLock::new(None) };
|
||||
Ok(rag)
|
||||
}
|
||||
```
|
||||
|
||||
### `hybrid_search` (`src/rag/mod.rs:710`)
|
||||
```rust
|
||||
async fn hybrid_search(&self, query: &str, top_k: usize, rerank_model: Option<&str>)
|
||||
-> Result<Vec<(DocumentId, String)>>
|
||||
```
|
||||
Runs `vector_search` + `keyword_search` in parallel via `tokio::join!`, then either reranks or
|
||||
applies `reciprocal_rank_fusion(vec![vector_ids, keyword_ids], vec![1.125, 1.0], top_k)`.
|
||||
|
||||
### `reciprocal_rank_fusion` (`src/rag/mod.rs:1186`) — standalone fn, already weight-parameterized:
|
||||
```rust
|
||||
fn reciprocal_rank_fusion(
|
||||
list_of_document_ids: Vec<Vec<DocumentId>>,
|
||||
list_of_weights: Vec<f32>,
|
||||
top_k: usize,
|
||||
) -> Vec<DocumentId>
|
||||
```
|
||||
|
||||
### `RagData::del` (`src/rag/mod.rs:953`):
|
||||
```rust
|
||||
pub fn del(&mut self, file_ids: Vec<FileId>) {
|
||||
for file_id in file_ids {
|
||||
if let Some(file) = self.files.swap_remove(&file_id) {
|
||||
for (document_index, _) in file.documents.iter().enumerate() {
|
||||
let document_id = DocumentId::new(file_id, document_index);
|
||||
self.vectors.swap_remove(&document_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### `RagNode` (`src/graph/types.rs:331`):
|
||||
```rust
|
||||
pub struct RagNode {
|
||||
pub documents: Vec<String>,
|
||||
pub query: Option<String>,
|
||||
pub top_k: Option<usize>,
|
||||
pub embedding_model: Option<String>,
|
||||
pub chunk_size: Option<usize>,
|
||||
pub chunk_overlap: Option<usize>,
|
||||
pub reranker_model: Option<String>,
|
||||
pub batch_size: Option<usize>,
|
||||
pub state_updates: Option<HashMap<String, String>>,
|
||||
pub timeout: Option<u64>,
|
||||
}
|
||||
```
|
||||
|
||||
### `Client` trait (`src/client/common.rs:40`):
|
||||
- `async fn chat_completions(&self, input: Input) -> Result<ChatCompletionsOutput>` — needs `Input`
|
||||
- `async fn chat_completions_inner(&self, client: &ReqwestClient, data: ChatCompletionsData) -> Result<ChatCompletionsOutput>` — accessible on `Box<dyn Client>` via vtable
|
||||
- `async fn embeddings(&self, data: &EmbeddingsData) -> Result<Vec<Vec<f32>>>`
|
||||
- `async fn rerank(&self, data: &RerankData) -> Result<RerankOutput>`
|
||||
- `fn build_client(&self) -> Result<ReqwestClient>`
|
||||
- `fn model(&self) -> &Model`
|
||||
|
||||
**Key finding**: `Input` cannot be constructed without `RequestContext` (which `Rag` doesn't have).
|
||||
Instead, `extract_entities` uses `chat_completions_inner` directly with manually built
|
||||
`ChatCompletionsData`. This is accessible via `Box<dyn Client>`.
|
||||
|
||||
### `Message` (`src/client/message.rs:22`):
|
||||
```rust
|
||||
pub fn new(role: MessageRole, content: MessageContent) -> Self
|
||||
```
|
||||
`MessageRole::User`, `MessageContent::Text(String)` — both confirmed.
|
||||
|
||||
### `AppConfig` RAG fields (`src/config/app_config.rs:71`):
|
||||
```rust
|
||||
pub rag_embedding_model: Option<String>,
|
||||
pub rag_reranker_model: Option<String>,
|
||||
pub rag_top_k: usize, // default: 5
|
||||
pub rag_chunk_size: Option<usize>,
|
||||
pub rag_chunk_overlap: Option<usize>,
|
||||
pub rag_template: Option<String>,
|
||||
```
|
||||
|
||||
### `patch_messages` — confirmed exported from `crate::client::*` (used in `input.rs:5`)
|
||||
|
||||
### `init_client(app_config, model)` — works for any `ModelType`, including `Chat`
|
||||
|
||||
### `ModelType` variants: `Chat`, `Embedding`, `Reranker` (confirmed in `model.rs`)
|
||||
|
||||
### petgraph serde: `NodeIndex` serializes as inner `u32`; `StableGraph` preserves index positions
|
||||
through roundtrip. `IndexMap<DocumentId, Vec<NodeIndex>>` safe for YAML (DocumentId is newtype over
|
||||
usize, serializes as integer key).
|
||||
|
||||
---
|
||||
|
||||
## New Dependency
|
||||
|
||||
```toml
|
||||
petgraph = { version = "0.7", features = ["serde-1"] }
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## New File: `src/rag/graph.rs`
|
||||
|
||||
All graph types and extraction logic. Module declared in `mod.rs` as `mod graph; use self::graph::*;`.
|
||||
|
||||
### Types:
|
||||
- `Entity { name: String, entity_type: String, description: Option<String> }`
|
||||
- `Relationship { relation_type: String, weight: f32 }`
|
||||
- `ExtractionResult { entities: Vec<ExtractedEntity>, relationships: Vec<ExtractedRelationship> }`
|
||||
- `ExtractedEntity { name: String, r#type: String, description: Option<String> }`
|
||||
- `ExtractedRelationship { from: String, to: String, r#type: String, weight: Option<f32> }`
|
||||
- `KnowledgeGraph { graph: StableGraph<Entity, Relationship>, entity_index: IndexMap<String, NodeIndex>, document_entities: IndexMap<DocumentId, Vec<NodeIndex>> }`
|
||||
|
||||
### Key methods on `KnowledgeGraph`:
|
||||
- `merge(doc_id: DocumentId, result: ExtractionResult)` — merges extraction into graph
|
||||
- `remove_documents(ids: &[DocumentId])` — removes entities exclusive to deleted documents
|
||||
- `build_node_to_docs(&self) -> IndexMap<NodeIndex, Vec<DocumentId>>` — ephemeral reverse map
|
||||
|
||||
### `extract_entities(client: &dyn Client, chunk: &str) -> Result<ExtractionResult>`:
|
||||
- Builds `ChatCompletionsData` manually (no `Input` needed)
|
||||
- Calls `patch_messages` then `client.chat_completions_inner(&reqwest_client, data).await`
|
||||
- Strips markdown code fences from response before JSON parse
|
||||
- Temperature: `Some(0.0)` for deterministic extraction
|
||||
|
||||
### Extraction prompt: structured JSON output requesting entities + relationships
|
||||
|
||||
---
|
||||
|
||||
## Changes to `src/rag/mod.rs`
|
||||
|
||||
### `Rag` struct — add one ephemeral field:
|
||||
```rust
|
||||
node_to_docs: IndexMap<NodeIndex, Vec<DocumentId>>, // ephemeral, rebuilt on load
|
||||
```
|
||||
|
||||
### `Rag::create` — build node_to_docs before moving data:
|
||||
```rust
|
||||
let node_to_docs = data.knowledge_graph.build_node_to_docs();
|
||||
// then add to struct literal
|
||||
```
|
||||
|
||||
### `Rag` Clone impl — add:
|
||||
```rust
|
||||
node_to_docs: self.data.knowledge_graph.build_node_to_docs(),
|
||||
```
|
||||
|
||||
### `RagData` struct — three new fields (all `#[serde(default)]` for backward compat):
|
||||
```rust
|
||||
#[serde(default)]
|
||||
pub graph_enabled: bool,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub extractor_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub knowledge_graph: KnowledgeGraph,
|
||||
```
|
||||
|
||||
### `RagData::new` — two new params: `graph_enabled: bool, extractor_model: Option<String>`
|
||||
|
||||
### `RagData::del` — collect doc_ids during existing loop, call `remove_documents` at end:
|
||||
```rust
|
||||
let mut doc_ids_to_remove = vec![];
|
||||
for file_id in file_ids {
|
||||
if let Some(file) = self.files.swap_remove(&file_id) {
|
||||
for (document_index, _) in file.documents.iter().enumerate() {
|
||||
let document_id = DocumentId::new(file_id, document_index);
|
||||
self.vectors.swap_remove(&document_id);
|
||||
doc_ids_to_remove.push(document_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
self.knowledge_graph.remove_documents(&doc_ids_to_remove);
|
||||
```
|
||||
|
||||
### `Rag::init` (line 219) — add two params to `RagData::new`:
|
||||
```rust
|
||||
app.rag_graph_enabled,
|
||||
app.rag_extractor_model.clone(),
|
||||
```
|
||||
|
||||
### `resolve_init_data` — resolve from config+app, pass to `RagData::new`:
|
||||
```rust
|
||||
let graph_enabled = config.graph_enabled.unwrap_or(app.rag_graph_enabled);
|
||||
let extractor_model = config.extractor_model.clone().or_else(|| app.rag_extractor_model.clone());
|
||||
```
|
||||
|
||||
### `sync_documents` — entity extraction block after `rag_files` built, before embedding:
|
||||
```rust
|
||||
if self.data.graph_enabled {
|
||||
if let Some(extractor_model_id) = self.data.extractor_model.clone() {
|
||||
let model = Model::retrieve_model(&self.app_config, &extractor_model_id, ModelType::Chat)?;
|
||||
let client = self.create_embeddings_client(model)?;
|
||||
let total_chunks: usize = rag_files.iter().map(|f| f.documents.len()).sum();
|
||||
let mut chunk_num = 0;
|
||||
let file_offset = next_file_id;
|
||||
for (batch_file_idx, rag_file) in rag_files.iter().enumerate() {
|
||||
let file_id = file_offset + batch_file_idx;
|
||||
for (doc_idx, doc) in rag_file.documents.iter().enumerate() {
|
||||
chunk_num += 1;
|
||||
progress(&spinner, format!("Extracting entities [{chunk_num}/{total_chunks}]"));
|
||||
let doc_id = DocumentId::new(file_id, doc_idx);
|
||||
match extract_entities(client.as_ref(), &doc.page_content).await {
|
||||
Ok(result) => self.data.knowledge_graph.merge(doc_id, result),
|
||||
Err(e) => debug!("Entity extraction failed for {doc_id:?}: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### After line 705 (after hnsw/bm25 rebuild in sync_documents):
|
||||
```rust
|
||||
self.node_to_docs = self.data.knowledge_graph.build_node_to_docs();
|
||||
```
|
||||
|
||||
### `hybrid_search` — add third signal:
|
||||
```rust
|
||||
let graph_search_ids: Vec<DocumentId> = if self.data.graph_enabled
|
||||
&& !self.data.knowledge_graph.entity_index.is_empty()
|
||||
{
|
||||
self.graph_search(query, &keyword_search_ids, top_k)
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
// RRF: extend to 3-way when graph has results, fall back to 2-way otherwise
|
||||
```
|
||||
|
||||
### New `graph_search` method (sync):
|
||||
```rust
|
||||
fn graph_search(&self, query: &str, bm25_anchor_ids: &[DocumentId], top_k: usize) -> Vec<DocumentId>
|
||||
```
|
||||
Phase 1: entity names from query via substring match in `entity_index`.
|
||||
Phase 2: fallback — entities from top BM25 document chunks.
|
||||
Phase 3: expand 1-hop neighbors in `StableGraph`.
|
||||
Phase 4: score docs by entity overlap ratio, return top_k.
|
||||
|
||||
### `RagInitConfig` — two new fields:
|
||||
```rust
|
||||
pub graph_enabled: Option<bool>,
|
||||
pub extractor_model: Option<String>,
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Changes to `src/config/app_config.rs`
|
||||
|
||||
New fields alongside existing `rag_*` block:
|
||||
```rust
|
||||
pub rag_graph_enabled: bool, // default: false
|
||||
pub rag_extractor_model: Option<String>, // default: None
|
||||
```
|
||||
Defaults, env var overrides, and propagation all follow the same pattern as existing `rag_*` fields.
|
||||
|
||||
---
|
||||
|
||||
## Changes to `src/graph/types.rs` — `RagNode`
|
||||
|
||||
```rust
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub graph_enabled: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub extractor_model: Option<String>,
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Changes to `src/config/agent.rs`
|
||||
|
||||
Pass new fields through to `RagInitConfig`:
|
||||
```rust
|
||||
graph_enabled: rag_node.graph_enabled,
|
||||
extractor_model: rag_node.extractor_model.clone(),
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Backward Compatibility
|
||||
|
||||
- All new `RagData` fields have `#[serde(default)]` — old YAML files load without migration
|
||||
- `graph_enabled` defaults `false` — existing RAG instances unchanged
|
||||
- `graph_search_ids` empty → 2-way RRF runs (identical to current behavior)
|
||||
- `node_to_docs` rebuild on `create()` is O(n) over empty map for old instances
|
||||
|
||||
---
|
||||
|
||||
## V1 Scope Exclusions
|
||||
|
||||
- LLM entity extraction from query at search time (V1 uses substring match + BM25 anchoring)
|
||||
- Multi-hop traversal (field reserved, 1-hop only in V1)
|
||||
- Entity embeddings / fuzzy entity lookup
|
||||
- Bincode for large-corpus graph storage
|
||||
- Gleaning / multi-pass extraction
|
||||
|
||||
---
|
||||
|
||||
## Implementation Progress
|
||||
|
||||
- [x] Cargo.toml — petgraph dependency
|
||||
- [x] src/rag/graph.rs — new file
|
||||
- [x] src/rag/mod.rs — mod/use, Rag struct, create, clone
|
||||
- [x] src/rag/mod.rs — RagData fields, new, del
|
||||
- [x] src/rag/mod.rs — Rag::init, resolve_init_data
|
||||
- [x] src/rag/mod.rs — sync_documents extraction block
|
||||
- [x] src/rag/mod.rs — hybrid_search + graph_search
|
||||
- [x] src/rag/mod.rs — RagInitConfig fields
|
||||
- [x] src/config/app_config.rs — new fields
|
||||
- [x] src/config/mod.rs — propagation
|
||||
- [x] src/graph/types.rs — RagNode fields
|
||||
- [x] src/config/agent.rs — propagation
|
||||
- [x] cargo check — clean (0 warnings, 1065 tests passing)
|
||||
Generated
+14
-1
@@ -1455,6 +1455,7 @@ dependencies = [
|
||||
"os_info",
|
||||
"parking_lot",
|
||||
"path-absolutize",
|
||||
"petgraph 0.7.1",
|
||||
"pretty_assertions",
|
||||
"rand 0.10.1",
|
||||
"reedline",
|
||||
@@ -4128,6 +4129,18 @@ version = "2.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
|
||||
|
||||
[[package]]
|
||||
name = "petgraph"
|
||||
version = "0.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772"
|
||||
dependencies = [
|
||||
"fixedbitset",
|
||||
"indexmap 2.14.0",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "petgraph"
|
||||
version = "0.8.3"
|
||||
@@ -6305,7 +6318,7 @@ checksum = "b8765b90061cba6c22b5831f675da109ae5561588290f9fa2317adab2714d5a6"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"nom 8.0.0",
|
||||
"petgraph",
|
||||
"petgraph 0.8.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -74,6 +74,7 @@ html_to_markdown = "0.1.0"
|
||||
rust-embed = "8.5.0"
|
||||
os_info = { version = "3.8.2", default-features = false }
|
||||
bm25 = { version = "2.0.1", features = ["parallelism"] }
|
||||
petgraph = { version = "0.7", features = ["serde-1"] }
|
||||
which = "8.0.0"
|
||||
fuzzy-matcher = "0.3.7"
|
||||
terminal-colorsaurus = "0.4.8"
|
||||
|
||||
@@ -92,6 +92,9 @@ conversation_starters: # Optional conversation starters for the agent
|
||||
- What is the best way to exercise?
|
||||
- How do I manage my time effectively?
|
||||
documents: # Optional documents to load for the agent
|
||||
# To enable graph-based RAG (entity/relationship extraction + knowledge graph retrieval),
|
||||
# set `rag_extractor_model` in your global config.yaml.
|
||||
# See https://github.com/Dark-Alex-17/coyote/wiki/RAG#graph-based-rag
|
||||
- git:/some/repo # Explicitly tell Coyote to use the 'git' document loader using an absolute path
|
||||
- pdf:some-pdf-file.pdf # Explicitly tell Coyote to use the 'pdf' document loader using a relative path
|
||||
- https://some-website.com/some-page
|
||||
|
||||
@@ -197,6 +197,9 @@ rag_reranker_model: null # Specifies the reranker model used for sorting
|
||||
rag_top_k: 5 # Specifies the number of documents to retrieve for answering queries
|
||||
rag_chunk_size: null # Defines the size of chunks for document processing in characters
|
||||
rag_chunk_overlap: null # Defines the overlap between chunks
|
||||
rag_extractor_model: null # LLM model for graph-based entity/relationship extraction; when set, enables a graph RAG signal alongside vector and BM25
|
||||
rag_extractor_prompt: null # Custom extraction prompt template; must contain __CHUNK__ placeholder; defaults to built-in prompt when null
|
||||
rag_graph_hops: 1 # Number of hops to expand from matched entities at query time (1 = direct neighbors; increase for denser graphs)
|
||||
# Defines the query structure using variables like __CONTEXT__, __SOURCES__, and __INPUT__ to tailor searches to specific needs
|
||||
rag_template: |
|
||||
Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags)
|
||||
|
||||
@@ -225,6 +225,9 @@ nodes:
|
||||
chunk_size: 1000
|
||||
chunk_overlap: 100
|
||||
reranker_model: null # Optional reranker for hybrid-search results
|
||||
extractor_model: null # Optional chat model for graph-based entity/relationship extraction; enables graph RAG signal when set
|
||||
extractor_prompt: null # Optional custom extraction prompt; must contain __CHUNK__ placeholder; uses built-in prompt when null
|
||||
graph_hops: 1 # Graph expansion depth at query time (1 = direct neighbors; increase for denser knowledge graphs)
|
||||
batch_size: 100 # Optional embedding-request batch size
|
||||
state_updates: # {{output}} = { context: <str>, sources: [<path>, ...] }
|
||||
context: "{{output.context}}" # writes `context` -> `reducers.context = concat`
|
||||
|
||||
@@ -921,6 +921,9 @@ async fn init_graph_rags(
|
||||
reranker_model: rag_node.reranker_model.clone(),
|
||||
top_k: rag_node.top_k,
|
||||
batch_size: rag_node.batch_size,
|
||||
extractor_model: rag_node.extractor_model.clone(),
|
||||
extractor_prompt: rag_node.extractor_prompt.clone(),
|
||||
graph_hops: rag_node.graph_hops,
|
||||
};
|
||||
let fully_specified = config.embedding_model.is_some()
|
||||
&& config.chunk_size.is_some()
|
||||
|
||||
@@ -74,6 +74,9 @@ pub struct AppConfig {
|
||||
pub rag_chunk_size: Option<usize>,
|
||||
pub rag_chunk_overlap: Option<usize>,
|
||||
pub rag_template: Option<String>,
|
||||
pub rag_extractor_model: Option<String>,
|
||||
pub rag_extractor_prompt: Option<String>,
|
||||
pub rag_graph_hops: usize,
|
||||
|
||||
#[serde(default)]
|
||||
pub document_loaders: HashMap<String, String>,
|
||||
@@ -146,6 +149,9 @@ impl Default for AppConfig {
|
||||
rag_chunk_size: None,
|
||||
rag_chunk_overlap: None,
|
||||
rag_template: None,
|
||||
rag_extractor_model: None,
|
||||
rag_extractor_prompt: None,
|
||||
rag_graph_hops: 1,
|
||||
|
||||
document_loaders: Default::default(),
|
||||
|
||||
@@ -219,6 +225,9 @@ impl AppConfig {
|
||||
rag_chunk_size: config.rag_chunk_size,
|
||||
rag_chunk_overlap: config.rag_chunk_overlap,
|
||||
rag_template: config.rag_template,
|
||||
rag_extractor_model: config.rag_extractor_model,
|
||||
rag_extractor_prompt: config.rag_extractor_prompt,
|
||||
rag_graph_hops: config.rag_graph_hops,
|
||||
|
||||
document_loaders: config.document_loaders,
|
||||
|
||||
@@ -512,6 +521,15 @@ impl AppConfig {
|
||||
if let Some(v) = super::read_env_value::<String>(&get_env_name("rag_template")) {
|
||||
self.rag_template = v;
|
||||
}
|
||||
if let Some(v) = super::read_env_value::<String>(&get_env_name("rag_extractor_model")) {
|
||||
self.rag_extractor_model = v;
|
||||
}
|
||||
if let Some(v) = super::read_env_value::<String>(&get_env_name("rag_extractor_prompt")) {
|
||||
self.rag_extractor_prompt = v;
|
||||
}
|
||||
if let Some(v) = super::read_env_value::<usize>(&get_env_name("rag_graph_hops")) {
|
||||
self.rag_graph_hops = v.unwrap_or(1);
|
||||
}
|
||||
|
||||
if let Ok(v) = env::var(get_env_name("document_loaders"))
|
||||
&& let Ok(v) = serde_json::from_str(&v)
|
||||
|
||||
@@ -250,6 +250,9 @@ pub struct Config {
|
||||
pub rag_chunk_size: Option<usize>,
|
||||
pub rag_chunk_overlap: Option<usize>,
|
||||
pub rag_template: Option<String>,
|
||||
pub rag_extractor_model: Option<String>,
|
||||
pub rag_extractor_prompt: Option<String>,
|
||||
pub rag_graph_hops: usize,
|
||||
|
||||
#[serde(default)]
|
||||
pub document_loaders: HashMap<String, String>,
|
||||
@@ -322,6 +325,9 @@ impl Default for Config {
|
||||
rag_chunk_size: None,
|
||||
rag_chunk_overlap: None,
|
||||
rag_template: None,
|
||||
rag_extractor_model: None,
|
||||
rag_extractor_prompt: None,
|
||||
rag_graph_hops: 1,
|
||||
|
||||
document_loaders: Default::default(),
|
||||
|
||||
|
||||
@@ -352,6 +352,15 @@ pub struct RagNode {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub batch_size: Option<usize>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub extractor_model: Option<String>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub extractor_prompt: Option<String>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub graph_hops: Option<usize>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub state_updates: Option<HashMap<String, String>>,
|
||||
|
||||
|
||||
@@ -1027,6 +1027,9 @@ mod tests {
|
||||
chunk_overlap: None,
|
||||
reranker_model: None,
|
||||
batch_size: None,
|
||||
extractor_model: None,
|
||||
extractor_prompt: None,
|
||||
graph_hops: None,
|
||||
state_updates,
|
||||
timeout: None,
|
||||
}),
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
use super::DocumentId;
|
||||
use crate::client::*;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use indexmap::IndexMap;
|
||||
use petgraph::Direction;
|
||||
use petgraph::graph::NodeIndex;
|
||||
use petgraph::stable_graph::StableGraph;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashSet;
|
||||
|
||||
const EXTRACTION_PROMPT: &str = r#"Extract entities and relationships from the following text chunk.
|
||||
|
||||
Return a JSON object with this exact structure:
|
||||
{
|
||||
"entities": [
|
||||
{"name": "EntityName", "type": "EntityType", "description": "brief description"}
|
||||
],
|
||||
"relationships": [
|
||||
{"from": "EntityA", "to": "EntityB", "type": "relation_verb", "weight": 0.9}
|
||||
]
|
||||
}
|
||||
|
||||
Rules:
|
||||
- Entity types: PERSON, ORGANIZATION, CONCEPT, TECHNOLOGY, LOCATION, EVENT, or OTHER
|
||||
- Relationship types should be short verb phrases (e.g., "uses", "depends_on", "implements", "part_of")
|
||||
- Weight is a float from 0.0 to 1.0 indicating relationship strength (default 1.0)
|
||||
- Only extract entities and relationships clearly stated or strongly implied in the text
|
||||
- Use exact entity names as they appear so relationships can be matched
|
||||
- Return ONLY the JSON object, no markdown fences, no explanation
|
||||
|
||||
Text chunk:
|
||||
__CHUNK__"#;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Entity {
|
||||
pub name: String,
|
||||
pub entity_type: String,
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Relationship {
|
||||
pub relation_type: String,
|
||||
pub weight: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ExtractionResult {
|
||||
pub entities: Vec<ExtractedEntity>,
|
||||
pub relationships: Vec<ExtractedRelationship>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ExtractedEntity {
|
||||
pub name: String,
|
||||
#[serde(rename = "type")]
|
||||
pub entity_type: String,
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ExtractedRelationship {
|
||||
pub from: String,
|
||||
pub to: String,
|
||||
#[serde(rename = "type")]
|
||||
pub relation_type: String,
|
||||
pub weight: Option<f32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KnowledgeGraph {
|
||||
pub graph: StableGraph<Entity, Relationship>,
|
||||
/// Lowercased entity name → raw node index
|
||||
pub entity_index: IndexMap<String, u32>,
|
||||
/// DocumentId inner value → raw node indices for entities in that chunk
|
||||
pub document_entities: IndexMap<usize, Vec<u32>>,
|
||||
}
|
||||
|
||||
impl Default for KnowledgeGraph {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
graph: StableGraph::new(),
|
||||
entity_index: IndexMap::new(),
|
||||
document_entities: IndexMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl KnowledgeGraph {
|
||||
pub fn merge(&mut self, doc_id: DocumentId, result: ExtractionResult) {
|
||||
let mut chunk_nodes: Vec<u32> = vec![];
|
||||
|
||||
for extracted in &result.entities {
|
||||
let key = extracted.name.to_lowercase();
|
||||
let node_raw = if let Some(&existing) = self.entity_index.get(&key) {
|
||||
existing
|
||||
} else {
|
||||
let entity = Entity {
|
||||
name: extracted.name.clone(),
|
||||
entity_type: extracted.entity_type.clone(),
|
||||
description: extracted.description.clone(),
|
||||
};
|
||||
let idx = self.graph.add_node(entity);
|
||||
let raw = idx.index() as u32;
|
||||
self.entity_index.insert(key, raw);
|
||||
raw
|
||||
};
|
||||
chunk_nodes.push(node_raw);
|
||||
}
|
||||
|
||||
for extracted in &result.relationships {
|
||||
let from_key = extracted.from.to_lowercase();
|
||||
let to_key = extracted.to.to_lowercase();
|
||||
if let (Some(&from_raw), Some(&to_raw)) = (
|
||||
self.entity_index.get(&from_key),
|
||||
self.entity_index.get(&to_key),
|
||||
) {
|
||||
let from_idx = NodeIndex::new(from_raw as usize);
|
||||
let to_idx = NodeIndex::new(to_raw as usize);
|
||||
// Avoid duplicate edges
|
||||
if !self.graph.contains_edge(from_idx, to_idx) {
|
||||
let rel = Relationship {
|
||||
relation_type: extracted.relation_type.clone(),
|
||||
weight: extracted.weight.unwrap_or(1.0),
|
||||
};
|
||||
self.graph.add_edge(from_idx, to_idx, rel);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.document_entities
|
||||
.entry(doc_id.0)
|
||||
.or_default()
|
||||
.extend(chunk_nodes);
|
||||
}
|
||||
|
||||
pub fn remove_documents(&mut self, doc_ids: &[DocumentId]) {
|
||||
if doc_ids.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let removing: HashSet<usize> = doc_ids.iter().map(|d| d.0).collect();
|
||||
for raw_id in &removing {
|
||||
self.document_entities.swap_remove(raw_id);
|
||||
}
|
||||
|
||||
let still_used: HashSet<u32> = self
|
||||
.document_entities
|
||||
.values()
|
||||
.flat_map(|v| v.iter().copied())
|
||||
.collect();
|
||||
|
||||
let to_remove: Vec<u32> = self
|
||||
.entity_index
|
||||
.values()
|
||||
.copied()
|
||||
.filter(|raw| !still_used.contains(raw))
|
||||
.collect();
|
||||
|
||||
for raw in to_remove {
|
||||
let idx = NodeIndex::new(raw as usize);
|
||||
if self.graph.contains_node(idx) {
|
||||
let name = self.graph[idx].name.to_lowercase();
|
||||
self.graph.remove_node(idx);
|
||||
self.entity_index.swap_remove(&name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_node_to_docs(&self) -> IndexMap<u32, Vec<DocumentId>> {
|
||||
let mut map: IndexMap<u32, Vec<DocumentId>> = IndexMap::new();
|
||||
for (&doc_raw, node_raws) in &self.document_entities {
|
||||
let doc_id = DocumentId(doc_raw);
|
||||
for &node_raw in node_raws {
|
||||
map.entry(node_raw).or_default().push(doc_id);
|
||||
}
|
||||
}
|
||||
map
|
||||
}
|
||||
|
||||
pub fn expand_neighbors(&self, seed_nodes: &[u32], hops: usize) -> Vec<u32> {
|
||||
let mut expanded: indexmap::IndexSet<u32> = seed_nodes.iter().copied().collect();
|
||||
let mut frontier: Vec<u32> = seed_nodes.to_vec();
|
||||
for _ in 0..hops {
|
||||
let mut next_frontier: Vec<u32> = vec![];
|
||||
for &raw in &frontier {
|
||||
let idx = NodeIndex::new(raw as usize);
|
||||
if self.graph.contains_node(idx) {
|
||||
for dir in [Direction::Outgoing, Direction::Incoming] {
|
||||
for neighbor in self.graph.neighbors_directed(idx, dir) {
|
||||
let n = neighbor.index() as u32;
|
||||
if expanded.insert(n) {
|
||||
next_frontier.push(n);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
frontier = next_frontier;
|
||||
if frontier.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
expanded.into_iter().collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Uses chat_completions_inner directly (bypassing Input) because Rag has no
|
||||
/// RequestContext, which Input::from_str requires.
|
||||
pub async fn extract_entities(
|
||||
client: &dyn Client,
|
||||
chunk: &str,
|
||||
prompt_template: Option<&str>,
|
||||
) -> Result<ExtractionResult> {
|
||||
let template = prompt_template.unwrap_or(EXTRACTION_PROMPT);
|
||||
let prompt = template.replace("__CHUNK__", chunk);
|
||||
let mut messages = vec![Message::new(
|
||||
MessageRole::User,
|
||||
MessageContent::Text(prompt),
|
||||
)];
|
||||
patch_messages(&mut messages, client.model());
|
||||
let reqwest_client = client
|
||||
.build_client()
|
||||
.context("Failed to build HTTP client for entity extraction")?;
|
||||
let data = ChatCompletionsData {
|
||||
messages,
|
||||
temperature: Some(0.0),
|
||||
top_p: None,
|
||||
functions: None,
|
||||
stream: false,
|
||||
};
|
||||
let output = client
|
||||
.chat_completions_inner(&reqwest_client, data)
|
||||
.await
|
||||
.context("Entity extraction LLM call failed")?;
|
||||
|
||||
let text = output.text.trim();
|
||||
// Strip markdown code fences if the model wraps in ```json ... ```
|
||||
let json: String = if text.starts_with("```") {
|
||||
text.lines()
|
||||
.skip(1)
|
||||
.take_while(|l| !l.trim_start().starts_with("```"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
} else {
|
||||
text.to_string()
|
||||
};
|
||||
|
||||
serde_json::from_str::<ExtractionResult>(&json)
|
||||
.context("Failed to parse entity extraction JSON")
|
||||
}
|
||||
+349
-19
@@ -4,15 +4,19 @@ use crate::client::*;
|
||||
use crate::config::*;
|
||||
use crate::utils::*;
|
||||
|
||||
mod graph;
|
||||
mod serde_vectors;
|
||||
mod splitter;
|
||||
|
||||
use self::graph::{KnowledgeGraph, extract_entities};
|
||||
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use bm25::{Language, SearchEngine, SearchEngineBuilder};
|
||||
use hnsw_rs::prelude::*;
|
||||
use indexmap::{IndexMap, IndexSet};
|
||||
use inquire::{Confirm, Select, Text, required, validator::Validation};
|
||||
use parking_lot::RwLock;
|
||||
use petgraph::graph::NodeIndex;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use std::{
|
||||
@@ -54,6 +58,7 @@ pub struct Rag {
|
||||
bm25: SearchEngine<DocumentId>,
|
||||
data: RagData,
|
||||
last_sources: RwLock<Option<String>>,
|
||||
node_to_docs: IndexMap<u32, Vec<DocumentId>>,
|
||||
}
|
||||
|
||||
impl Debug for Rag {
|
||||
@@ -76,6 +81,7 @@ impl Clone for Rag {
|
||||
embedding_model: self.embedding_model.clone(),
|
||||
hnsw: self.data.build_hnsw(),
|
||||
bm25: self.data.build_bm25(),
|
||||
node_to_docs: self.data.knowledge_graph.build_node_to_docs(),
|
||||
data: self.data.clone(),
|
||||
last_sources: RwLock::new(None),
|
||||
}
|
||||
@@ -90,6 +96,16 @@ pub struct RagInitConfig {
|
||||
pub reranker_model: Option<String>,
|
||||
pub top_k: Option<usize>,
|
||||
pub batch_size: Option<usize>,
|
||||
pub extractor_model: Option<String>,
|
||||
pub extractor_prompt: Option<String>,
|
||||
pub graph_hops: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GraphRagConfig {
|
||||
pub extractor_model: Option<String>,
|
||||
pub extractor_prompt: Option<String>,
|
||||
pub graph_hops: Option<usize>,
|
||||
}
|
||||
|
||||
impl Rag {
|
||||
@@ -199,6 +215,17 @@ impl Rag {
|
||||
reranker_model,
|
||||
top_k,
|
||||
batch_size,
|
||||
GraphRagConfig {
|
||||
extractor_model: config
|
||||
.extractor_model
|
||||
.clone()
|
||||
.or_else(|| app.rag_extractor_model.clone()),
|
||||
extractor_prompt: config
|
||||
.extractor_prompt
|
||||
.clone()
|
||||
.or_else(|| app.rag_extractor_prompt.clone()),
|
||||
graph_hops: Some(config.graph_hops.unwrap_or(app.rag_graph_hops)),
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
@@ -216,6 +243,16 @@ impl Rag {
|
||||
let (embedding_model, chunk_size, chunk_overlap) = Self::create_config(app)?;
|
||||
let reranker_model = app.rag_reranker_model.clone();
|
||||
let top_k = app.rag_top_k;
|
||||
let extractor_model = match app.rag_extractor_model.clone() {
|
||||
Some(model) => Some(model),
|
||||
None => select_extractor_model(app)?,
|
||||
};
|
||||
let graph_hops = if extractor_model.is_some() {
|
||||
set_graph_hops(app.rag_graph_hops)?
|
||||
} else {
|
||||
app.rag_graph_hops
|
||||
};
|
||||
let extractor_prompt = app.rag_extractor_prompt.clone();
|
||||
let data = RagData::new(
|
||||
embedding_model.id(),
|
||||
chunk_size,
|
||||
@@ -223,6 +260,11 @@ impl Rag {
|
||||
reranker_model,
|
||||
top_k,
|
||||
embedding_model.max_batch_size(),
|
||||
GraphRagConfig {
|
||||
extractor_model,
|
||||
extractor_prompt,
|
||||
graph_hops: Some(graph_hops),
|
||||
},
|
||||
);
|
||||
let mut rag = Self::create(app, name, save_path, data)?;
|
||||
let mut paths = doc_paths.to_vec();
|
||||
@@ -253,6 +295,7 @@ impl Rag {
|
||||
pub fn create(app: &AppConfig, name: &str, path: &Path, data: RagData) -> Result<Self> {
|
||||
let hnsw = data.build_hnsw();
|
||||
let bm25 = data.build_bm25();
|
||||
let node_to_docs = data.knowledge_graph.build_node_to_docs();
|
||||
let embedding_model =
|
||||
Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?;
|
||||
let rag = Rag {
|
||||
@@ -263,6 +306,7 @@ impl Rag {
|
||||
embedding_model,
|
||||
hnsw,
|
||||
bm25,
|
||||
node_to_docs,
|
||||
last_sources: RwLock::new(None),
|
||||
};
|
||||
Ok(rag)
|
||||
@@ -413,6 +457,9 @@ impl Rag {
|
||||
"chunk_size": self.data.chunk_size,
|
||||
"chunk_overlap": self.data.chunk_overlap,
|
||||
"reranker_model": self.data.reranker_model,
|
||||
"extractor_model": self.data.extractor_model,
|
||||
"extractor_prompt": self.data.extractor_prompt,
|
||||
"graph_hops": self.data.graph_hops.unwrap_or(1),
|
||||
"top_k": self.data.top_k,
|
||||
"batch_size": self.data.batch_size,
|
||||
"document_paths": self.data.document_paths,
|
||||
@@ -673,13 +720,18 @@ impl Rag {
|
||||
let mut files = vec![];
|
||||
let mut document_ids = vec![];
|
||||
let mut embeddings = vec![];
|
||||
let mut new_doc_contents: Vec<(DocumentId, String)> = vec![];
|
||||
|
||||
if !rag_files.is_empty() {
|
||||
let mut texts = vec![];
|
||||
for file in rag_files.into_iter() {
|
||||
for (document_index, document) in file.documents.iter().enumerate() {
|
||||
document_ids.push(DocumentId::new(next_file_id, document_index));
|
||||
texts.push(document.page_content.clone())
|
||||
let doc_id = DocumentId::new(next_file_id, document_index);
|
||||
document_ids.push(doc_id);
|
||||
texts.push(document.page_content.clone());
|
||||
if self.data.extractor_model.is_some() {
|
||||
new_doc_contents.push((doc_id, document.page_content.clone()));
|
||||
}
|
||||
}
|
||||
files.push((next_file_id, file));
|
||||
next_file_id += 1;
|
||||
@@ -700,9 +752,43 @@ impl Rag {
|
||||
bail!("No RAG files");
|
||||
}
|
||||
|
||||
if self.data.extractor_model.is_some()
|
||||
&& !new_doc_contents.is_empty()
|
||||
&& let Some(extractor_model_id) = self.data.extractor_model.clone()
|
||||
{
|
||||
match Model::retrieve_model(&self.app_config, &extractor_model_id, ModelType::Chat) {
|
||||
Ok(model) => match self.create_embeddings_client(model) {
|
||||
Ok(client) => {
|
||||
let total = new_doc_contents.len();
|
||||
for (i, (doc_id, content)) in new_doc_contents.into_iter().enumerate() {
|
||||
progress(
|
||||
&spinner,
|
||||
format!("Extracting entities [{}/{}]", i + 1, total),
|
||||
);
|
||||
match extract_entities(
|
||||
client.as_ref(),
|
||||
&content,
|
||||
self.data.extractor_prompt.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => self.data.knowledge_graph.merge(doc_id, result),
|
||||
Err(e) => {
|
||||
debug!("Entity extraction failed for doc {doc_id:?}: {e}")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => debug!("Failed to create extractor client: {e}"),
|
||||
},
|
||||
Err(e) => debug!("Extractor model not found: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
progress(&spinner, "Building store".into());
|
||||
self.hnsw = self.data.build_hnsw();
|
||||
self.bm25 = self.data.build_bm25();
|
||||
self.node_to_docs = self.data.knowledge_graph.build_node_to_docs();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -755,11 +841,21 @@ impl Rag {
|
||||
ids
|
||||
}
|
||||
None => {
|
||||
let ids = reciprocal_rank_fusion(
|
||||
let ids = if self.data.extractor_model.is_some() {
|
||||
let graph_ids = self.graph_search(query, top_k);
|
||||
debug!("graph_search_ids: {graph_ids:?}");
|
||||
reciprocal_rank_fusion(
|
||||
vec![vector_search_ids, keyword_search_ids, graph_ids],
|
||||
vec![1.125, 1.0, 0.9],
|
||||
top_k,
|
||||
)
|
||||
} else {
|
||||
reciprocal_rank_fusion(
|
||||
vec![vector_search_ids, keyword_search_ids],
|
||||
vec![1.125, 1.0],
|
||||
top_k,
|
||||
);
|
||||
)
|
||||
};
|
||||
debug!("rrf_ids: {ids:?}");
|
||||
ids
|
||||
}
|
||||
@@ -829,6 +925,93 @@ impl Rag {
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn graph_search(&self, query: &str, top_k: usize) -> Vec<DocumentId> {
|
||||
let kg = &self.data.knowledge_graph;
|
||||
if kg.entity_index.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
let query_lower = query.to_lowercase();
|
||||
|
||||
let mut seed_nodes: Vec<u32> = kg
|
||||
.entity_index
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
let name_str = name.as_str();
|
||||
if name_str.contains(' ') {
|
||||
query_lower.contains(name_str)
|
||||
} else {
|
||||
// whole-word match: prevents "go" from seeding on every query containing "Django"
|
||||
query_lower
|
||||
.split_whitespace()
|
||||
.any(|token| token.trim_matches(|c: char| !c.is_alphanumeric()) == name_str)
|
||||
}
|
||||
})
|
||||
.map(|(_, &raw)| raw)
|
||||
.collect();
|
||||
|
||||
if seed_nodes.is_empty() {
|
||||
let bm25_results = self.bm25.search(query, top_k * 2);
|
||||
'outer: for result in bm25_results {
|
||||
if let Some(node_raws) = kg.document_entities.get(&result.document.id.0) {
|
||||
seed_nodes.extend(node_raws.iter().copied());
|
||||
if seed_nodes.len() >= top_k {
|
||||
break 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if seed_nodes.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let hops = self.data.graph_hops.unwrap_or(1);
|
||||
let expanded = kg.expand_neighbors(&seed_nodes, hops);
|
||||
|
||||
let query_tokens: Vec<&str> = query_lower.split_whitespace().collect();
|
||||
let token_count = query_tokens.len().max(1);
|
||||
let mut scored: Vec<(u32, f32)> = expanded
|
||||
.into_iter()
|
||||
.map(|raw| {
|
||||
let idx = NodeIndex::new(raw as usize);
|
||||
let score = if kg.graph.contains_node(idx) {
|
||||
let entity = &kg.graph[idx];
|
||||
let combined = format!(
|
||||
"{} {}",
|
||||
entity.name,
|
||||
entity.description.as_deref().unwrap_or("")
|
||||
)
|
||||
.to_lowercase();
|
||||
query_tokens
|
||||
.iter()
|
||||
.filter(|t| combined.contains(*t))
|
||||
.count() as f32
|
||||
/ token_count as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
(raw, score)
|
||||
})
|
||||
.collect();
|
||||
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));
|
||||
|
||||
let mut result_ids: IndexSet<DocumentId> = IndexSet::new();
|
||||
for (raw, _) in scored {
|
||||
if let Some(doc_ids) = self.node_to_docs.get(&raw) {
|
||||
for &doc_id in doc_ids {
|
||||
result_ids.insert(doc_id);
|
||||
if result_ids.len() >= top_k {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if result_ids.len() >= top_k {
|
||||
break;
|
||||
}
|
||||
}
|
||||
result_ids.into_iter().collect()
|
||||
}
|
||||
|
||||
async fn create_embeddings(
|
||||
&self,
|
||||
data: EmbeddingsData,
|
||||
@@ -902,6 +1085,14 @@ pub struct RagData {
|
||||
pub files: IndexMap<FileId, RagFile>,
|
||||
#[serde(with = "serde_vectors")]
|
||||
pub vectors: IndexMap<DocumentId, Vec<f32>>,
|
||||
#[serde(default)]
|
||||
pub extractor_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub extractor_prompt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub graph_hops: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub knowledge_graph: KnowledgeGraph,
|
||||
}
|
||||
|
||||
impl Debug for RagData {
|
||||
@@ -916,6 +1107,9 @@ impl Debug for RagData {
|
||||
.field("next_file_id", &self.next_file_id)
|
||||
.field("document_paths", &self.document_paths)
|
||||
.field("files", &self.files)
|
||||
.field("extractor_model", &self.extractor_model)
|
||||
.field("extractor_prompt", &self.extractor_prompt)
|
||||
.field("graph_hops", &self.graph_hops)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -928,6 +1122,7 @@ impl RagData {
|
||||
reranker_model: Option<String>,
|
||||
top_k: usize,
|
||||
batch_size: Option<usize>,
|
||||
graph: GraphRagConfig,
|
||||
) -> Self {
|
||||
Self {
|
||||
embedding_model,
|
||||
@@ -940,6 +1135,10 @@ impl RagData {
|
||||
document_paths: Default::default(),
|
||||
files: Default::default(),
|
||||
vectors: Default::default(),
|
||||
extractor_model: graph.extractor_model,
|
||||
extractor_prompt: graph.extractor_prompt,
|
||||
graph_hops: graph.graph_hops,
|
||||
knowledge_graph: KnowledgeGraph::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -951,14 +1150,17 @@ impl RagData {
|
||||
}
|
||||
|
||||
pub fn del(&mut self, file_ids: Vec<FileId>) {
|
||||
let mut graph_doc_ids = vec![];
|
||||
for file_id in file_ids {
|
||||
if let Some(file) = self.files.swap_remove(&file_id) {
|
||||
for (document_index, _) in file.documents.iter().enumerate() {
|
||||
let document_id = DocumentId::new(file_id, document_index);
|
||||
self.vectors.swap_remove(&document_id);
|
||||
graph_doc_ids.push(document_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
self.knowledge_graph.remove_documents(&graph_doc_ids);
|
||||
}
|
||||
|
||||
pub fn add(
|
||||
@@ -1055,29 +1257,70 @@ impl DocumentId {
|
||||
}
|
||||
|
||||
fn select_embedding_model(models: &[&Model]) -> Result<String> {
|
||||
let max_width = models.iter().map(|v| v.id().len()).max().unwrap_or(0);
|
||||
let models: Vec<_> = models
|
||||
.iter()
|
||||
.map(|v| SelectOption::new(v.id(), v.description()))
|
||||
.map(|v| SelectOption::new(v.id(), v.description(), max_width))
|
||||
.collect();
|
||||
let result = Select::new("Select embedding model:", models).prompt()?;
|
||||
let result = Select::new("Select embedding model:", models)
|
||||
.with_formatter(&|opt| opt.value.value.clone())
|
||||
.prompt()?;
|
||||
Ok(result.value)
|
||||
}
|
||||
|
||||
const EXTRACTOR_SKIP: &str = "Skip";
|
||||
|
||||
fn select_extractor_model(app: &AppConfig) -> Result<Option<String>> {
|
||||
let models = list_models(app, ModelType::Chat);
|
||||
if models.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let pad = models
|
||||
.iter()
|
||||
.map(|v| v.id().len())
|
||||
.max()
|
||||
.unwrap_or(0)
|
||||
.max(EXTRACTOR_SKIP.len());
|
||||
let mut options = vec![SelectOption::new(
|
||||
EXTRACTOR_SKIP.to_string(),
|
||||
"vector + full text search only (no graph)".to_string(),
|
||||
pad,
|
||||
)];
|
||||
options.extend(
|
||||
models
|
||||
.iter()
|
||||
.map(|v| SelectOption::new(v.id(), v.description(), pad)),
|
||||
);
|
||||
let result = Select::new("Extractor model for graph-based RAG (optional):", options)
|
||||
.with_formatter(&|opt| opt.value.value.clone())
|
||||
.prompt()?;
|
||||
Ok(if result.value == EXTRACTOR_SKIP {
|
||||
None
|
||||
} else {
|
||||
Some(result.value)
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SelectOption {
|
||||
pub value: String,
|
||||
pub description: String,
|
||||
pub display: String,
|
||||
}
|
||||
|
||||
impl SelectOption {
|
||||
pub fn new(value: String, description: String) -> Self {
|
||||
Self { value, description }
|
||||
pub fn new(value: String, description: String, pad: usize) -> Self {
|
||||
let display = if description.is_empty() {
|
||||
format!("{value:<pad$}")
|
||||
} else {
|
||||
format!("{value:<pad$} ({description})")
|
||||
};
|
||||
Self { value, display }
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for SelectOption {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{} ({})", self.value, self.description)
|
||||
write!(f, "{}", self.display)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1103,6 +1346,21 @@ fn set_chunk_size(model: &Model) -> Result<usize> {
|
||||
value.parse().map_err(|_| anyhow!("Invalid chunk_size"))
|
||||
}
|
||||
|
||||
fn set_graph_hops(default_value: usize) -> Result<usize> {
|
||||
let value = Text::new("Set graph expansion hops:")
|
||||
.with_default(&default_value.to_string())
|
||||
.with_help_message("Number of hops to expand from matched entities (1 = direct neighbors, 2 = neighbors of neighbors)")
|
||||
.with_validator(move |text: &str| {
|
||||
let out = match text.parse::<usize>() {
|
||||
Ok(v) if v >= 1 => Validation::Valid,
|
||||
_ => Validation::Invalid("Must be an integer >= 1".into()),
|
||||
};
|
||||
Ok(out)
|
||||
})
|
||||
.prompt()?;
|
||||
value.parse().map_err(|_| anyhow!("Invalid graph_hops"))
|
||||
}
|
||||
|
||||
fn set_chunk_overlay(default_value: usize) -> Result<usize> {
|
||||
let value = Text::new("Set chunk overlay:")
|
||||
.with_default(&default_value.to_string())
|
||||
@@ -1277,7 +1535,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_new_defaults() {
|
||||
let data = RagData::new("model".into(), 1000, 20, None, 5, None);
|
||||
let data = RagData::new(
|
||||
"model".into(),
|
||||
1000,
|
||||
20,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
assert_eq!(data.embedding_model, "model");
|
||||
assert_eq!(data.chunk_size, 1000);
|
||||
assert_eq!(data.chunk_overlap, 20);
|
||||
@@ -1291,7 +1557,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_get_returns_document() {
|
||||
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let mut data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
let file = RagFile {
|
||||
hash: "abc".into(),
|
||||
path: "test.txt".into(),
|
||||
@@ -1308,13 +1582,29 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_get_returns_none_for_missing_file() {
|
||||
let data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
assert!(data.get(DocumentId::new(99, 0)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_data_get_returns_none_for_missing_document() {
|
||||
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let mut data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
let file = RagFile {
|
||||
hash: "abc".into(),
|
||||
path: "test.txt".into(),
|
||||
@@ -1326,7 +1616,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_del_removes_files_and_vectors() {
|
||||
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let mut data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
let file = RagFile {
|
||||
hash: "abc".into(),
|
||||
path: "test.txt".into(),
|
||||
@@ -1347,14 +1645,30 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_del_nonexistent_is_noop() {
|
||||
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let mut data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
data.del(vec![99]);
|
||||
assert!(data.files.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_data_add_inserts_files_and_vectors() {
|
||||
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let mut data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
let file = RagFile {
|
||||
hash: "xyz".into(),
|
||||
path: "new.txt".into(),
|
||||
@@ -1414,7 +1728,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_build_bm25_empty() {
|
||||
let data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
let engine = data.build_bm25();
|
||||
let results = engine.search("anything", 5);
|
||||
assert!(results.is_empty());
|
||||
@@ -1422,7 +1744,15 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn rag_data_build_bm25_finds_documents() {
|
||||
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
|
||||
let mut data = RagData::new(
|
||||
"m".into(),
|
||||
100,
|
||||
10,
|
||||
None,
|
||||
5,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
let file = RagFile {
|
||||
hash: "h".into(),
|
||||
path: "test.txt".into(),
|
||||
|
||||
Reference in New Issue
Block a user