diff --git a/.sisyphus/plans/graph-rag-spec.md b/.sisyphus/plans/graph-rag-spec.md new file mode 100644 index 0000000..bb4e9e6 --- /dev/null +++ b/.sisyphus/plans/graph-rag-spec.md @@ -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, + name: String, + path: String, + embedding_model: Model, + hnsw: Hnsw<'static, f32, DistCosine>, // ephemeral, rebuilt on load + bm25: SearchEngine, // ephemeral, rebuilt on load + data: RagData, // serialized to YAML + last_sources: RwLock>, +} +``` + +### `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, + pub top_k: usize, + pub batch_size: Option, + pub next_file_id: FileId, + pub document_paths: Vec, + pub files: IndexMap, + #[serde(with = "serde_vectors")] + pub vectors: IndexMap>, +} +``` + +### `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 { + 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> +``` +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>, + list_of_weights: Vec, + top_k: usize, +) -> Vec +``` + +### `RagData::del` (`src/rag/mod.rs:953`): +```rust +pub fn del(&mut self, file_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); + } + } + } +} +``` + +### `RagNode` (`src/graph/types.rs:331`): +```rust +pub struct RagNode { + pub documents: Vec, + pub query: Option, + pub top_k: Option, + pub embedding_model: Option, + pub chunk_size: Option, + pub chunk_overlap: Option, + pub reranker_model: Option, + pub batch_size: Option, + pub state_updates: Option>, + pub timeout: Option, +} +``` + +### `Client` trait (`src/client/common.rs:40`): +- `async fn chat_completions(&self, input: Input) -> Result` — needs `Input` +- `async fn chat_completions_inner(&self, client: &ReqwestClient, data: ChatCompletionsData) -> Result` — accessible on `Box` via vtable +- `async fn embeddings(&self, data: &EmbeddingsData) -> Result>>` +- `async fn rerank(&self, data: &RerankData) -> Result` +- `fn build_client(&self) -> Result` +- `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`. + +### `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, +pub rag_reranker_model: Option, +pub rag_top_k: usize, // default: 5 +pub rag_chunk_size: Option, +pub rag_chunk_overlap: Option, +pub rag_template: Option, +``` + +### `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>` 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 }` +- `Relationship { relation_type: String, weight: f32 }` +- `ExtractionResult { entities: Vec, relationships: Vec }` +- `ExtractedEntity { name: String, r#type: String, description: Option }` +- `ExtractedRelationship { from: String, to: String, r#type: String, weight: Option }` +- `KnowledgeGraph { graph: StableGraph, entity_index: IndexMap, document_entities: IndexMap> }` + +### 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>` — ephemeral reverse map + +### `extract_entities(client: &dyn Client, chunk: &str) -> Result`: +- 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>, // 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, +#[serde(default)] +pub knowledge_graph: KnowledgeGraph, +``` + +### `RagData::new` — two new params: `graph_enabled: bool, extractor_model: Option` + +### `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 = 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 +``` +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, +pub extractor_model: Option, +``` + +--- + +## 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, // 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, +#[serde(default, skip_serializing_if = "Option::is_none")] +pub extractor_model: Option, +``` + +--- + +## 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) diff --git a/Cargo.lock b/Cargo.lock index 63cdc4a..6685978 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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]] diff --git a/Cargo.toml b/Cargo.toml index 434b578..0552831 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/config.agent.example.yaml b/config.agent.example.yaml index e723938..a6bc3a6 100644 --- a/config.agent.example.yaml +++ b/config.agent.example.yaml @@ -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 diff --git a/config.example.yaml b/config.example.yaml index eab018b..9c16ff4 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -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) diff --git a/graph.example.yaml b/graph.example.yaml index 21e5ab3..14e021b 100644 --- a/graph.example.yaml +++ b/graph.example.yaml @@ -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: , sources: [, ...] } context: "{{output.context}}" # writes `context` -> `reducers.context = concat` diff --git a/src/config/agent.rs b/src/config/agent.rs index 5950b44..a7c0b4b 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -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() diff --git a/src/config/app_config.rs b/src/config/app_config.rs index 8da9548..e892e37 100644 --- a/src/config/app_config.rs +++ b/src/config/app_config.rs @@ -74,6 +74,9 @@ pub struct AppConfig { pub rag_chunk_size: Option, pub rag_chunk_overlap: Option, pub rag_template: Option, + pub rag_extractor_model: Option, + pub rag_extractor_prompt: Option, + pub rag_graph_hops: usize, #[serde(default)] pub document_loaders: HashMap, @@ -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::(&get_env_name("rag_template")) { self.rag_template = v; } + if let Some(v) = super::read_env_value::(&get_env_name("rag_extractor_model")) { + self.rag_extractor_model = v; + } + if let Some(v) = super::read_env_value::(&get_env_name("rag_extractor_prompt")) { + self.rag_extractor_prompt = v; + } + if let Some(v) = super::read_env_value::(&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) diff --git a/src/config/mod.rs b/src/config/mod.rs index 3a0e05a..b1ea413 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -250,6 +250,9 @@ pub struct Config { pub rag_chunk_size: Option, pub rag_chunk_overlap: Option, pub rag_template: Option, + pub rag_extractor_model: Option, + pub rag_extractor_prompt: Option, + pub rag_graph_hops: usize, #[serde(default)] pub document_loaders: HashMap, @@ -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(), diff --git a/src/graph/types.rs b/src/graph/types.rs index b06aabf..437ddd7 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -352,6 +352,15 @@ pub struct RagNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub batch_size: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extractor_model: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extractor_prompt: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub graph_hops: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub state_updates: Option>, diff --git a/src/graph/validator.rs b/src/graph/validator.rs index c04309c..7d5c370 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -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, }), diff --git a/src/rag/graph.rs b/src/rag/graph.rs new file mode 100644 index 0000000..26f4151 --- /dev/null +++ b/src/rag/graph.rs @@ -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, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Relationship { + pub relation_type: String, + pub weight: f32, +} + +#[derive(Debug, Deserialize)] +pub struct ExtractionResult { + pub entities: Vec, + pub relationships: Vec, +} + +#[derive(Debug, Deserialize)] +pub struct ExtractedEntity { + pub name: String, + #[serde(rename = "type")] + pub entity_type: String, + pub description: Option, +} + +#[derive(Debug, Deserialize)] +pub struct ExtractedRelationship { + pub from: String, + pub to: String, + #[serde(rename = "type")] + pub relation_type: String, + pub weight: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KnowledgeGraph { + pub graph: StableGraph, + /// Lowercased entity name → raw node index + pub entity_index: IndexMap, + /// DocumentId inner value → raw node indices for entities in that chunk + pub document_entities: IndexMap>, +} + +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 = 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 = doc_ids.iter().map(|d| d.0).collect(); + for raw_id in &removing { + self.document_entities.swap_remove(raw_id); + } + + let still_used: HashSet = self + .document_entities + .values() + .flat_map(|v| v.iter().copied()) + .collect(); + + let to_remove: Vec = 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> { + let mut map: IndexMap> = 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 { + let mut expanded: indexmap::IndexSet = seed_nodes.iter().copied().collect(); + let mut frontier: Vec = seed_nodes.to_vec(); + for _ in 0..hops { + let mut next_frontier: Vec = 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 { + 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::>() + .join("\n") + } else { + text.to_string() + }; + + serde_json::from_str::(&json) + .context("Failed to parse entity extraction JSON") +} diff --git a/src/rag/mod.rs b/src/rag/mod.rs index 05c13ef..c22fa30 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -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, data: RagData, last_sources: RwLock>, + node_to_docs: IndexMap>, } 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, pub top_k: Option, pub batch_size: Option, + pub extractor_model: Option, + pub extractor_prompt: Option, + pub graph_hops: Option, +} + +#[derive(Debug, Clone, Default)] +pub struct GraphRagConfig { + pub extractor_model: Option, + pub extractor_prompt: Option, + pub graph_hops: Option, } 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 { 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( - vec![vector_search_ids, keyword_search_ids], - vec![1.125, 1.0], - top_k, - ); + 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 { + 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 = 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 = 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, #[serde(with = "serde_vectors")] pub vectors: IndexMap>, + #[serde(default)] + pub extractor_model: Option, + #[serde(default)] + pub extractor_prompt: Option, + #[serde(default)] + pub graph_hops: Option, + #[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, top_k: usize, batch_size: Option, + 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) { + 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 { + 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> { + 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:) -> fmt::Result { - write!(f, "{} ({})", self.value, self.description) + write!(f, "{}", self.display) } } @@ -1103,6 +1346,21 @@ fn set_chunk_size(model: &Model) -> Result { value.parse().map_err(|_| anyhow!("Invalid chunk_size")) } +fn set_graph_hops(default_value: usize) -> Result { + 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::() { + 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 { 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(),