Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7fc06ad9bc
|
||
|
|
e814b9f62d
|
||
|
|
209fbc9e41
|
||
|
|
5eb1daf18d
|
||
|
|
b8990fdfc2
|
||
|
|
7673799d83
|
||
|
|
ce212fe660
|
||
|
|
f855a45493
|
||
|
|
5be4f45671
|
||
|
|
84ce094677
|
||
|
|
fa7eadd08a
|
||
|
|
af91b89cff
|
@@ -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"
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
---
|
||||
name: diagnose
|
||||
temperature: 0.2
|
||||
enabled_tools:
|
||||
- execute_command
|
||||
- fs_cat
|
||||
|
||||
@@ -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`
|
||||
|
||||
@@ -91,6 +91,9 @@ pub struct Cli {
|
||||
/// Disable memory for this invocation
|
||||
#[arg(long)]
|
||||
pub no_memory: bool,
|
||||
/// Skip permission prompts by setting AUTO_CONFIRM for all tools (dangerous!)
|
||||
#[arg(long)]
|
||||
pub dangerously_skip_permissions: bool,
|
||||
/// Bootstrap a memory marker so coyote begins loading memory next run
|
||||
#[arg(long, value_name = "SCOPE", value_enum)]
|
||||
pub init_memory: Option<MemoryScope>,
|
||||
@@ -424,6 +427,18 @@ mod tests {
|
||||
assert!(cli.build_tools);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dangerously_skip_permissions_flag() {
|
||||
let cli = parse(&["--dangerously-skip-permissions"]);
|
||||
assert!(cli.dangerously_skip_permissions);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dangerously_skip_permissions_default_off() {
|
||||
let cli = parse(&[]);
|
||||
assert!(!cli.dangerously_skip_permissions);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_sync_models_flag() {
|
||||
let cli = parse(&["--sync-models"]);
|
||||
|
||||
@@ -4,29 +4,44 @@ use indexmap::IndexMap;
|
||||
use parking_lot::RwLock;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, (String, i64)>>> =
|
||||
type AccessTokenEntry = (String, i64, Option<String>);
|
||||
|
||||
static ACCESS_TOKENS: LazyLock<RwLock<IndexMap<String, AccessTokenEntry>>> =
|
||||
LazyLock::new(|| RwLock::new(IndexMap::new()));
|
||||
|
||||
pub fn get_access_token(client_name: &str) -> Result<String> {
|
||||
ACCESS_TOKENS
|
||||
.read()
|
||||
.get(client_name)
|
||||
.map(|(token, _)| token.clone())
|
||||
.map(|(token, _, _)| token.clone())
|
||||
.ok_or_else(|| anyhow!("Invalid access token"))
|
||||
}
|
||||
|
||||
pub fn get_access_token_account_id(client_name: &str) -> Option<String> {
|
||||
ACCESS_TOKENS
|
||||
.read()
|
||||
.get(client_name)
|
||||
.and_then(|(_, _, account_id)| account_id.clone())
|
||||
}
|
||||
|
||||
pub fn is_valid_access_token(client_name: &str) -> bool {
|
||||
let access_tokens = ACCESS_TOKENS.read();
|
||||
let (token, expires_at) = match access_tokens.get(client_name) {
|
||||
let (token, expires_at, _) = match access_tokens.get(client_name) {
|
||||
Some(v) => v,
|
||||
None => return false,
|
||||
};
|
||||
!token.is_empty() && Utc::now().timestamp() < *expires_at
|
||||
}
|
||||
|
||||
pub fn set_access_token(client_name: &str, token: String, expires_at: i64) {
|
||||
pub fn set_access_token(
|
||||
client_name: &str,
|
||||
token: String,
|
||||
expires_at: i64,
|
||||
account_id: Option<String>,
|
||||
) {
|
||||
let mut access_tokens = ACCESS_TOKENS.write();
|
||||
let entry = access_tokens.entry(client_name.to_string()).or_default();
|
||||
entry.0 = token;
|
||||
entry.1 = expires_at;
|
||||
entry.2 = account_id;
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ mod common;
|
||||
mod gemini_oauth;
|
||||
mod message;
|
||||
pub mod oauth;
|
||||
mod openai_oauth;
|
||||
#[macro_use]
|
||||
mod macros;
|
||||
mod model;
|
||||
|
||||
+50
-19
@@ -57,13 +57,24 @@ pub trait OAuthProvider: Send + Sync {
|
||||
fn fixed_redirect_uri(&self) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn extract_account_id(&self, _response: &Value) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn include_state_in_token_exchange(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OAuthTokens {
|
||||
pub access_token: String,
|
||||
pub refresh_token: String,
|
||||
#[serde(default)]
|
||||
pub refresh_token: Option<String>,
|
||||
pub expires_at: i64,
|
||||
#[serde(default)]
|
||||
pub account_id: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) -> Result<()> {
|
||||
@@ -137,18 +148,17 @@ pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) ->
|
||||
}
|
||||
|
||||
let client = ReqwestClient::new();
|
||||
let request = build_token_request(
|
||||
&client,
|
||||
provider,
|
||||
&[
|
||||
let mut token_params = vec![
|
||||
("grant_type", "authorization_code"),
|
||||
("client_id", provider.client_id()),
|
||||
("code", &code),
|
||||
("code_verifier", &code_verifier),
|
||||
("redirect_uri", &redirect_uri),
|
||||
("state", &state),
|
||||
],
|
||||
);
|
||||
("code", code.as_str()),
|
||||
("code_verifier", code_verifier.as_str()),
|
||||
("redirect_uri", redirect_uri.as_str()),
|
||||
];
|
||||
if provider.include_state_in_token_exchange() {
|
||||
token_params.push(("state", state.as_str()));
|
||||
}
|
||||
let request = build_token_request(&client, provider, &token_params);
|
||||
|
||||
let response: Value = request.send().await?.json().await?;
|
||||
|
||||
@@ -156,20 +166,20 @@ pub async fn run_oauth_flow(provider: &dyn OAuthProvider, client_name: &str) ->
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Missing access_token in response: {response}"))?
|
||||
.to_string();
|
||||
let refresh_token = response["refresh_token"]
|
||||
.as_str()
|
||||
.ok_or_else(|| anyhow!("Missing refresh_token in response: {response}"))?
|
||||
.to_string();
|
||||
let refresh_token = response["refresh_token"].as_str().map(|s| s.to_string());
|
||||
let expires_in = response["expires_in"]
|
||||
.as_i64()
|
||||
.ok_or_else(|| anyhow!("Missing expires_in in response: {response}"))?;
|
||||
|
||||
let expires_at = Utc::now().timestamp() + expires_in;
|
||||
|
||||
let account_id = provider.extract_account_id(&response);
|
||||
|
||||
let tokens = OAuthTokens {
|
||||
access_token,
|
||||
refresh_token,
|
||||
expires_at,
|
||||
account_id,
|
||||
};
|
||||
|
||||
save_oauth_tokens(client_name, &tokens)?;
|
||||
@@ -205,13 +215,19 @@ pub async fn refresh_oauth_token(
|
||||
client_name: &str,
|
||||
tokens: &OAuthTokens,
|
||||
) -> Result<OAuthTokens> {
|
||||
let refresh_token_val = tokens.refresh_token.as_deref().ok_or_else(|| {
|
||||
anyhow!(
|
||||
"No refresh token available for '{}'. Please re-authenticate.",
|
||||
client_name
|
||||
)
|
||||
})?;
|
||||
let request = build_token_request(
|
||||
client,
|
||||
provider,
|
||||
&[
|
||||
("grant_type", "refresh_token"),
|
||||
("client_id", provider.client_id()),
|
||||
("refresh_token", &tokens.refresh_token),
|
||||
("refresh_token", refresh_token_val),
|
||||
],
|
||||
);
|
||||
|
||||
@@ -224,17 +240,22 @@ pub async fn refresh_oauth_token(
|
||||
let refresh_token = response["refresh_token"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| tokens.refresh_token.clone());
|
||||
.or_else(|| tokens.refresh_token.clone());
|
||||
let expires_in = response["expires_in"]
|
||||
.as_i64()
|
||||
.ok_or_else(|| anyhow!("Missing expires_in in refresh response: {response}"))?;
|
||||
|
||||
let expires_at = Utc::now().timestamp() + expires_in;
|
||||
|
||||
let account_id = provider
|
||||
.extract_account_id(&response)
|
||||
.or_else(|| tokens.account_id.clone());
|
||||
|
||||
let new_tokens = OAuthTokens {
|
||||
access_token,
|
||||
refresh_token,
|
||||
expires_at,
|
||||
account_id,
|
||||
};
|
||||
|
||||
save_oauth_tokens(client_name, &new_tokens)?;
|
||||
@@ -262,7 +283,12 @@ pub async fn prepare_oauth_access_token(
|
||||
tokens
|
||||
};
|
||||
|
||||
set_access_token(client_name, tokens.access_token.clone(), tokens.expires_at);
|
||||
set_access_token(
|
||||
client_name,
|
||||
tokens.access_token.clone(),
|
||||
tokens.expires_at,
|
||||
tokens.account_id.clone(),
|
||||
);
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
@@ -371,6 +397,7 @@ pub fn get_oauth_provider(provider_type: &str) -> Option<Box<dyn OAuthProvider>>
|
||||
match provider_type {
|
||||
"claude" => Some(Box::new(super::claude_oauth::ClaudeOAuthProvider)),
|
||||
"gemini" => Some(Box::new(super::gemini_oauth::GeminiOAuthProvider)),
|
||||
"openai" => Some(Box::new(super::openai_oauth::OpenAIOAuthProvider)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -409,7 +436,11 @@ fn client_config_info(client_config: &ClientConfig) -> (&str, &'static str, Opti
|
||||
"claude",
|
||||
c.auth.as_deref(),
|
||||
),
|
||||
ClientConfig::OpenAIConfig(c) => (c.name.as_deref().unwrap_or("openai"), "openai", None),
|
||||
ClientConfig::OpenAIConfig(c) => (
|
||||
c.name.as_deref().unwrap_or("openai"),
|
||||
"openai",
|
||||
c.auth.as_deref(),
|
||||
),
|
||||
ClientConfig::OpenAICompatibleConfig(c) => (
|
||||
c.name.as_deref().unwrap_or("openai-compatible"),
|
||||
"openai-compatible",
|
||||
|
||||
+330
-19
@@ -1,13 +1,17 @@
|
||||
use super::access_token::{get_access_token, get_access_token_account_id};
|
||||
use super::oauth::{self, OAuthProvider};
|
||||
use super::openai_oauth::OpenAIOAuthProvider;
|
||||
use super::*;
|
||||
|
||||
use crate::utils::strip_think_tag;
|
||||
|
||||
use anyhow::{Context, Result, bail};
|
||||
use reqwest::RequestBuilder;
|
||||
use reqwest::{Client as ReqwestClient, RequestBuilder};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
const API_BASE: &str = "https://api.openai.com/v1";
|
||||
const CODEX_API_ENDPOINT: &str = "https://chatgpt.com/backend-api/codex/responses";
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Default)]
|
||||
pub struct OpenAIConfig {
|
||||
@@ -15,6 +19,7 @@ pub struct OpenAIConfig {
|
||||
pub api_key: Option<String>,
|
||||
pub api_base: Option<String>,
|
||||
pub organization_id: Option<String>,
|
||||
pub auth: Option<String>,
|
||||
#[serde(default)]
|
||||
pub models: Vec<ModelData>,
|
||||
pub patch: Option<RequestPatch>,
|
||||
@@ -25,36 +30,131 @@ impl OpenAIClient {
|
||||
config_get_fn!(api_key, get_api_key);
|
||||
config_get_fn!(api_base, get_api_base);
|
||||
|
||||
create_client_config!([("api_key", "API Key", None, true)]);
|
||||
create_oauth_supported_client_config!();
|
||||
}
|
||||
|
||||
impl_client_trait!(
|
||||
OpenAIClient,
|
||||
(
|
||||
prepare_chat_completions,
|
||||
openai_chat_completions,
|
||||
openai_chat_completions_streaming
|
||||
),
|
||||
(prepare_embeddings, openai_embeddings),
|
||||
(noop_prepare_rerank, noop_rerank),
|
||||
);
|
||||
#[async_trait::async_trait]
|
||||
impl Client for OpenAIClient {
|
||||
client_common_fns!();
|
||||
|
||||
fn prepare_chat_completions(
|
||||
fn supports_oauth(&self) -> bool {
|
||||
self.config.auth.as_deref() == Some("oauth")
|
||||
}
|
||||
|
||||
async fn chat_completions_inner(
|
||||
&self,
|
||||
client: &ReqwestClient,
|
||||
data: ChatCompletionsData,
|
||||
) -> Result<ChatCompletionsOutput> {
|
||||
let uses_codex =
|
||||
self.config.auth.as_deref() == Some("oauth") && self.get_api_base().is_err();
|
||||
let request_data = prepare_chat_completions(self, client, data).await?;
|
||||
let builder = self.request_builder(client, request_data);
|
||||
if uses_codex {
|
||||
openai_responses_chat_completions(builder, self.model()).await
|
||||
} else {
|
||||
openai_chat_completions(builder, self.model()).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn chat_completions_streaming_inner(
|
||||
&self,
|
||||
client: &ReqwestClient,
|
||||
handler: &mut SseHandler,
|
||||
data: ChatCompletionsData,
|
||||
) -> Result<()> {
|
||||
let uses_codex =
|
||||
self.config.auth.as_deref() == Some("oauth") && self.get_api_base().is_err();
|
||||
let request_data = prepare_chat_completions(self, client, data).await?;
|
||||
let builder = self.request_builder(client, request_data);
|
||||
|
||||
if uses_codex {
|
||||
openai_responses_streaming(builder, handler).await
|
||||
} else {
|
||||
openai_chat_completions_streaming(builder, handler, self.model()).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn embeddings_inner(
|
||||
&self,
|
||||
client: &ReqwestClient,
|
||||
data: &EmbeddingsData,
|
||||
) -> Result<EmbeddingsOutput> {
|
||||
let request_data = prepare_embeddings(self, client, data).await?;
|
||||
let builder = self.request_builder(client, request_data);
|
||||
openai_embeddings(builder, self.model()).await
|
||||
}
|
||||
|
||||
async fn rerank_inner(
|
||||
&self,
|
||||
client: &ReqwestClient,
|
||||
data: &RerankData,
|
||||
) -> Result<RerankOutput> {
|
||||
let request_data = noop_prepare_rerank(self, data)?;
|
||||
let builder = self.request_builder(client, request_data);
|
||||
noop_rerank(builder, self.model()).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn prepare_chat_completions(
|
||||
self_: &OpenAIClient,
|
||||
client: &ReqwestClient,
|
||||
data: ChatCompletionsData,
|
||||
) -> Result<RequestData> {
|
||||
let api_key = self_.get_api_key()?;
|
||||
let uses_oauth = self_.config.auth.as_deref() == Some("oauth");
|
||||
let has_custom_base = self_.get_api_base().is_ok();
|
||||
|
||||
let uses_codex = uses_oauth && !has_custom_base;
|
||||
|
||||
let url = if uses_codex {
|
||||
CODEX_API_ENDPOINT.to_string()
|
||||
} else {
|
||||
let api_base = self_
|
||||
.get_api_base()
|
||||
.unwrap_or_else(|_| API_BASE.to_string());
|
||||
format!("{}/chat/completions", api_base.trim_end_matches('/'))
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", api_base.trim_end_matches('/'));
|
||||
|
||||
let body = openai_build_chat_completions_body(data, &self_.model);
|
||||
let body = if uses_codex {
|
||||
openai_build_responses_body(data, &self_.model)
|
||||
} else {
|
||||
openai_build_chat_completions_body(data, &self_.model)
|
||||
};
|
||||
|
||||
let mut request_data = RequestData::new(url, body);
|
||||
|
||||
if uses_oauth {
|
||||
let provider = OpenAIOAuthProvider;
|
||||
let ready = oauth::prepare_oauth_access_token(client, &provider, self_.name()).await?;
|
||||
|
||||
if !ready {
|
||||
bail!(
|
||||
"OAuth configured but no tokens found for '{}'. Run: 'coyote --authenticate {}' or '.authenticate' in the REPL",
|
||||
self_.name(),
|
||||
self_.name()
|
||||
);
|
||||
}
|
||||
|
||||
let token = get_access_token(self_.name())?;
|
||||
request_data.bearer_auth(token);
|
||||
|
||||
if let Some(account_id) = get_access_token_account_id(self_.name()) {
|
||||
request_data.header("ChatGPT-Account-Id", account_id);
|
||||
}
|
||||
|
||||
for (key, value) in provider.extra_request_headers() {
|
||||
request_data.header(key, value);
|
||||
}
|
||||
} else if let Ok(api_key) = self_.get_api_key() {
|
||||
request_data.bearer_auth(api_key);
|
||||
} else {
|
||||
bail!(
|
||||
"No authentication configured for '{}'. Set `api_key` or use `auth: oauth` with `coyote --authenticate {}`.",
|
||||
self_.name(),
|
||||
self_.name()
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(organization_id) = &self_.config.organization_id {
|
||||
request_data.header("OpenAI-Organization", organization_id);
|
||||
}
|
||||
@@ -62,8 +162,11 @@ fn prepare_chat_completions(
|
||||
Ok(request_data)
|
||||
}
|
||||
|
||||
fn prepare_embeddings(self_: &OpenAIClient, data: &EmbeddingsData) -> Result<RequestData> {
|
||||
let api_key = self_.get_api_key()?;
|
||||
async fn prepare_embeddings(
|
||||
self_: &OpenAIClient,
|
||||
client: &ReqwestClient,
|
||||
data: &EmbeddingsData,
|
||||
) -> Result<RequestData> {
|
||||
let api_base = self_
|
||||
.get_api_base()
|
||||
.unwrap_or_else(|_| API_BASE.to_string());
|
||||
@@ -74,7 +177,30 @@ fn prepare_embeddings(self_: &OpenAIClient, data: &EmbeddingsData) -> Result<Req
|
||||
|
||||
let mut request_data = RequestData::new(url, body);
|
||||
|
||||
if self_.config.auth.as_deref() == Some("oauth") {
|
||||
let provider = OpenAIOAuthProvider;
|
||||
let ready = oauth::prepare_oauth_access_token(client, &provider, self_.name()).await?;
|
||||
|
||||
if !ready {
|
||||
bail!(
|
||||
"OAuth configured but no tokens found for '{}'. Run: 'coyote --authenticate {}' or '.authenticate' in the REPL",
|
||||
self_.name(),
|
||||
self_.name()
|
||||
);
|
||||
}
|
||||
|
||||
let token = get_access_token(self_.name())?;
|
||||
request_data.bearer_auth(token);
|
||||
} else if let Ok(api_key) = self_.get_api_key() {
|
||||
request_data.bearer_auth(api_key);
|
||||
} else {
|
||||
bail!(
|
||||
"No authentication configured for '{}'. Set `api_key` or use `auth: oauth` with `coyote --authenticate {}`.",
|
||||
self_.name(),
|
||||
self_.name()
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(organization_id) = &self_.config.organization_id {
|
||||
request_data.header("OpenAI-Organization", organization_id);
|
||||
}
|
||||
@@ -402,3 +528,188 @@ fn normalize_function_id(value: &str) -> Option<String> {
|
||||
Some(value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn openai_build_responses_body(data: ChatCompletionsData, model: &Model) -> Value {
|
||||
let ChatCompletionsData {
|
||||
messages,
|
||||
temperature,
|
||||
top_p,
|
||||
functions,
|
||||
stream,
|
||||
} = data;
|
||||
|
||||
let messages_len = messages.len();
|
||||
let input: Vec<Value> = messages
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.flat_map(|(i, message)| {
|
||||
let Message { role, content } = message;
|
||||
match content {
|
||||
MessageContent::ToolCalls(MessageContentToolCalls {
|
||||
tool_results,
|
||||
text: _,
|
||||
sequence: _,
|
||||
}) => tool_results
|
||||
.into_iter()
|
||||
.flat_map(|tool_result| {
|
||||
vec![
|
||||
json!({
|
||||
"type": "function_call",
|
||||
"call_id": tool_result.call.id,
|
||||
"name": tool_result.call.name,
|
||||
"arguments": tool_result.call.arguments.to_string(),
|
||||
}),
|
||||
json!({
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_result.call.id,
|
||||
"output": tool_result.output.to_string(),
|
||||
}),
|
||||
]
|
||||
})
|
||||
.collect(),
|
||||
MessageContent::Text(text) if role.is_assistant() && i != messages_len - 1 => {
|
||||
vec![json!({ "role": role, "content": strip_think_tag(&text) })]
|
||||
}
|
||||
_ => vec![json!({ "role": role, "content": content })],
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut body = json!({
|
||||
"model": &model.real_name(),
|
||||
"input": input,
|
||||
"store": false,
|
||||
});
|
||||
|
||||
if let Some(v) = model.max_tokens_param() {
|
||||
body["max_output_tokens"] = v.into();
|
||||
}
|
||||
if let Some(v) = temperature {
|
||||
body["temperature"] = v.into();
|
||||
}
|
||||
if let Some(v) = top_p {
|
||||
body["top_p"] = v.into();
|
||||
}
|
||||
if stream {
|
||||
body["stream"] = true.into();
|
||||
}
|
||||
if let Some(functions) = functions {
|
||||
body["tools"] = functions
|
||||
.iter()
|
||||
.map(|v| {
|
||||
let mut tool = serde_json::to_value(v).unwrap_or_default();
|
||||
tool["type"] = "function".into();
|
||||
tool
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
pub async fn openai_responses_chat_completions(
|
||||
builder: RequestBuilder,
|
||||
_model: &Model,
|
||||
) -> Result<ChatCompletionsOutput> {
|
||||
let res = builder.send().await?;
|
||||
let status = res.status();
|
||||
let data: Value = res.json().await?;
|
||||
|
||||
if !status.is_success() {
|
||||
catch_error(&data, status.as_u16())?;
|
||||
}
|
||||
|
||||
debug!("non-stream-data: {data}");
|
||||
openai_extract_responses(&data)
|
||||
}
|
||||
|
||||
pub fn openai_extract_responses(data: &Value) -> Result<ChatCompletionsOutput> {
|
||||
let mut text = String::new();
|
||||
let mut tool_calls = vec![];
|
||||
|
||||
if let Some(output) = data["output"].as_array() {
|
||||
for item in output {
|
||||
match item["type"].as_str() {
|
||||
Some("message") => {
|
||||
if let Some(content) = item["content"].as_array() {
|
||||
for part in content {
|
||||
if part["type"].as_str() == Some("output_text")
|
||||
&& let Some(t) = part["text"].as_str()
|
||||
{
|
||||
text.push_str(t);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("function_call") => {
|
||||
if let (Some(name), Some(arguments_str), Some(call_id)) = (
|
||||
item["name"].as_str(),
|
||||
item["arguments"].as_str(),
|
||||
item["call_id"].as_str(),
|
||||
) {
|
||||
let arguments: Value = arguments_str.parse().with_context(|| {
|
||||
format!("Tool call '{name}' has non-JSON arguments '{arguments_str}'")
|
||||
})?;
|
||||
tool_calls.push(ToolCall::new(
|
||||
name.to_string(),
|
||||
arguments,
|
||||
Some(call_id.to_string()),
|
||||
));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if text.is_empty() && tool_calls.is_empty() {
|
||||
bail!("Invalid response data: {data}");
|
||||
}
|
||||
Ok(ChatCompletionsOutput { text, tool_calls })
|
||||
}
|
||||
|
||||
pub async fn openai_responses_streaming(
|
||||
builder: RequestBuilder,
|
||||
handler: &mut SseHandler,
|
||||
) -> Result<()> {
|
||||
let handle = |message: SseMessage| -> Result<bool> {
|
||||
if message.data == "[DONE]" {
|
||||
return Ok(true);
|
||||
}
|
||||
let data: Value = serde_json::from_str(&message.data)?;
|
||||
debug!("stream-data: {data}");
|
||||
|
||||
match data["type"].as_str() {
|
||||
Some("response.output_text.delta") => {
|
||||
if let Some(delta) = data["delta"].as_str().filter(|v| !v.is_empty()) {
|
||||
handler.text(delta)?;
|
||||
}
|
||||
}
|
||||
Some("response.output_item.done") => {
|
||||
let item = &data["item"];
|
||||
if item["type"].as_str() == Some("function_call")
|
||||
&& let (Some(name), Some(arguments_str), Some(call_id)) = (
|
||||
item["name"].as_str(),
|
||||
item["arguments"].as_str(),
|
||||
item["call_id"].as_str(),
|
||||
)
|
||||
{
|
||||
let arguments: Value = arguments_str.parse().with_context(|| {
|
||||
format!("Tool call '{name}' has non-JSON arguments '{arguments_str}'")
|
||||
})?;
|
||||
handler.tool_call(ToolCall::new(
|
||||
name.to_string(),
|
||||
arguments,
|
||||
Some(call_id.to_string()),
|
||||
))?;
|
||||
}
|
||||
}
|
||||
Some("response.completed") => {
|
||||
return Ok(true);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(false)
|
||||
};
|
||||
|
||||
sse_stream(builder, handle).await
|
||||
}
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
use super::oauth::{OAuthProvider, TokenRequestFormat};
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
use serde_json::Value;
|
||||
|
||||
pub struct OpenAIOAuthProvider;
|
||||
|
||||
impl OAuthProvider for OpenAIOAuthProvider {
|
||||
fn provider_name(&self) -> &str {
|
||||
"openai"
|
||||
}
|
||||
|
||||
fn client_id(&self) -> &str {
|
||||
"app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
}
|
||||
|
||||
fn authorize_url(&self) -> &str {
|
||||
"https://auth.openai.com/oauth/authorize"
|
||||
}
|
||||
|
||||
fn token_url(&self) -> &str {
|
||||
"https://auth.openai.com/oauth/token"
|
||||
}
|
||||
|
||||
fn redirect_uri(&self) -> &str {
|
||||
"http://localhost:1455/auth/callback"
|
||||
}
|
||||
|
||||
fn scopes(&self) -> &str {
|
||||
"openid profile email offline_access"
|
||||
}
|
||||
|
||||
fn token_request_format(&self) -> TokenRequestFormat {
|
||||
TokenRequestFormat::FormUrlEncoded
|
||||
}
|
||||
|
||||
fn extra_authorize_params(&self) -> Vec<(&str, &str)> {
|
||||
vec![
|
||||
("id_token_add_organizations", "true"),
|
||||
("codex_cli_simplified_flow", "true"),
|
||||
]
|
||||
}
|
||||
|
||||
fn fixed_redirect_uri(&self) -> Option<String> {
|
||||
Some("http://localhost:1455/auth/callback".to_string())
|
||||
}
|
||||
|
||||
fn include_state_in_token_exchange(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn extract_account_id(&self, response: &Value) -> Option<String> {
|
||||
let id_token = response["id_token"].as_str().unwrap_or_default();
|
||||
let access_token = response["access_token"].as_str().unwrap_or_default();
|
||||
extract_account_id_from_jwt(id_token).or_else(|| extract_account_id_from_jwt(access_token))
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_account_id_from_jwt(token: &str) -> Option<String> {
|
||||
let parts: Vec<&str> = token.splitn(3, '.').collect();
|
||||
if parts.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
let decoded = URL_SAFE_NO_PAD.decode(parts[1]).ok()?;
|
||||
let claims: Value = serde_json::from_slice(&decoded).ok()?;
|
||||
claims["chatgpt_account_id"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
.or_else(|| {
|
||||
claims["https://api.openai.com/auth"]["chatgpt_account_id"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
})
|
||||
.or_else(|| {
|
||||
claims["organizations"][0]["id"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
})
|
||||
}
|
||||
@@ -232,8 +232,8 @@ where
|
||||
.map(|value| value.to_string());
|
||||
let is_event_stream = content_type
|
||||
.as_deref()
|
||||
.map(|ct| ct.starts_with("text/event-stream"))
|
||||
.unwrap_or(false);
|
||||
.map(|ct| ct.is_empty() || ct.starts_with("text/event-stream"))
|
||||
.unwrap_or(true);
|
||||
if !is_event_stream {
|
||||
let header_value = content_type.unwrap_or_default();
|
||||
let text = res.text().await?;
|
||||
|
||||
@@ -483,7 +483,7 @@ pub async fn prepare_gcloud_access_token(
|
||||
let expires_at = Utc::now()
|
||||
+ Duration::try_seconds(expires_in)
|
||||
.ok_or_else(|| anyhow!("Failed to parse expires_in of access_token"))?;
|
||||
set_access_token(client_name, token, expires_at.timestamp())
|
||||
set_access_token(client_name, token, expires_at.timestamp(), None)
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
use crate::mcp::{
|
||||
ConnectedServer, JsonField, McpServer, McpTransportType, oauth, spawn_mcp_server,
|
||||
ConnectedServer, JsonField, McpServer, McpTransportType, is_auth_required_error, oauth,
|
||||
spawn_mcp_server,
|
||||
};
|
||||
|
||||
use anyhow::Result;
|
||||
use anyhow::{Result, anyhow};
|
||||
use parking_lot::Mutex;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
@@ -106,7 +107,18 @@ impl McpFactory {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let handle = spawn_mcp_server(spec, log_path, bearer_token).await?;
|
||||
let handle = spawn_mcp_server(spec, log_path, bearer_token)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
if is_auth_required_error(&e) {
|
||||
anyhow!(
|
||||
"MCP server '{name}' requires OAuth authentication. \
|
||||
Run `coyote --auth-mcp {name}` or `.mcp auth {name}` in the REPL to authenticate."
|
||||
)
|
||||
} else {
|
||||
e
|
||||
}
|
||||
})?;
|
||||
self.insert_active(key, &handle);
|
||||
Ok(handle)
|
||||
}
|
||||
@@ -132,7 +144,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -149,7 +161,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: Some(url.to_string()),
|
||||
headers,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+10
-3
@@ -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(),
|
||||
|
||||
@@ -445,9 +451,10 @@ fn confirm_asset_overwrite(category: AssetCategory, label: &str, target: &Path)
|
||||
}
|
||||
let body = match category {
|
||||
AssetCategory::McpConfig => format!(
|
||||
"This replaces your MCP server configuration at {} with this \
|
||||
build's bundled template. Your configured MCP servers (and any \
|
||||
custom secret references they contain) will be lost.",
|
||||
"This merges the bundled MCP server template into your configuration \
|
||||
at {}. New servers from the bundled template will be added; any \
|
||||
MCP servers you have already configured (including custom secret \
|
||||
references) are left untouched.",
|
||||
target.display()
|
||||
),
|
||||
_ => format!(
|
||||
|
||||
@@ -16,7 +16,7 @@ use super::{MessageContentToolCalls, prompts};
|
||||
use crate::client::{Model, ModelType, list_models};
|
||||
use crate::function::{
|
||||
FunctionDeclaration, Functions, ToolCallTracker, ToolResult, skill::SKILL_FUNCTION_PREFIX,
|
||||
user_interaction::USER_FUNCTION_PREFIX,
|
||||
todo::TODO_FUNCTION_PREFIX, user_interaction::USER_FUNCTION_PREFIX,
|
||||
};
|
||||
use crate::mcp::{
|
||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||
@@ -1466,7 +1466,9 @@ impl RequestContext {
|
||||
.filter(|v| {
|
||||
(v.name.starts_with(USER_FUNCTION_PREFIX)
|
||||
|| (!matches!(role.skills_enabled(), Some(false))
|
||||
&& v.name.starts_with(SKILL_FUNCTION_PREFIX)))
|
||||
&& v.name.starts_with(SKILL_FUNCTION_PREFIX))
|
||||
|| (self.auto_continue_config().enabled
|
||||
&& v.name.starts_with(TODO_FUNCTION_PREFIX)))
|
||||
&& !existing.contains(&v.name)
|
||||
})
|
||||
.cloned()
|
||||
@@ -2682,9 +2684,9 @@ impl RequestContext {
|
||||
None
|
||||
};
|
||||
|
||||
self.use_role_obj(role)?;
|
||||
self.rebuild_tool_scope(app, mcp_servers, abort_signal)
|
||||
.await?;
|
||||
self.use_role_obj(role)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn use_session(
|
||||
@@ -3698,7 +3700,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -5175,19 +5177,26 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn install_mcp_config_overwrites_existing() {
|
||||
fn install_mcp_config_merges_existing() {
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
|
||||
Functions::install_mcp_config().unwrap();
|
||||
let mcp = paths::mcp_config_file();
|
||||
assert!(mcp.exists(), "install_mcp_config should create mcp.json");
|
||||
|
||||
write(&mcp, "USER_MCP_CONFIG").unwrap();
|
||||
let custom_json =
|
||||
r#"{"mcpServers":{"my-custom-server":{"type":"stdio","command":"custom-cmd"}}}"#;
|
||||
write(&mcp, custom_json).unwrap();
|
||||
Functions::install_mcp_config().unwrap();
|
||||
assert_ne!(
|
||||
read_to_string(&mcp).unwrap(),
|
||||
"USER_MCP_CONFIG",
|
||||
"install_mcp_config must overwrite the existing mcp.json"
|
||||
|
||||
let result = read_to_string(&mcp).unwrap();
|
||||
assert!(
|
||||
result.contains("my-custom-server"),
|
||||
"install_mcp_config must preserve user-added MCP servers"
|
||||
);
|
||||
assert!(
|
||||
result.contains("github"),
|
||||
"install_mcp_config must add new bundled servers"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+33
-5
@@ -14,7 +14,7 @@ use crate::config::ensure_parent_exists;
|
||||
use crate::config::paths;
|
||||
use crate::mcp::{
|
||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpServersConfig,
|
||||
};
|
||||
use crate::parsers::{bash, python, typescript};
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
@@ -266,14 +266,42 @@ impl Functions {
|
||||
let file_path = paths::mcp_config_file();
|
||||
let embedded = FunctionAssets::get("mcp.json")
|
||||
.ok_or_else(|| anyhow!("Failed to load embedded mcp.json"))?;
|
||||
let content = unsafe { std::str::from_utf8_unchecked(&embedded.data) };
|
||||
let bundled_content = unsafe { std::str::from_utf8_unchecked(&embedded.data) };
|
||||
let bundled: McpServersConfig =
|
||||
serde_json::from_str(bundled_content).context("failed to parse embedded mcp.json")?;
|
||||
|
||||
ensure_parent_exists(&file_path)?;
|
||||
|
||||
info!("Reinstalling MCP config file: {}", file_path.display());
|
||||
let mut merged = if file_path.exists() {
|
||||
let existing =
|
||||
fs::read_to_string(&file_path).context("failed to read existing mcp.json")?;
|
||||
serde_json::from_str::<McpServersConfig>(&existing)
|
||||
.context("failed to parse existing mcp.json")?
|
||||
} else {
|
||||
McpServersConfig {
|
||||
mcp_servers: IndexMap::new(),
|
||||
}
|
||||
};
|
||||
|
||||
let mut config_file = File::create(&file_path)?;
|
||||
config_file.write_all(content.as_bytes())?;
|
||||
let mut added = Vec::new();
|
||||
for (name, server) in bundled.mcp_servers {
|
||||
if !merged.mcp_servers.contains_key(&name) {
|
||||
merged.mcp_servers.insert(name.clone(), server);
|
||||
added.push(name);
|
||||
}
|
||||
}
|
||||
|
||||
info!("Merging bundled MCP config into: {}", file_path.display());
|
||||
|
||||
let serialized =
|
||||
serde_json::to_string_pretty(&merged).context("failed to serialize merged mcp.json")?;
|
||||
let tmp = file_path.with_extension("json.tmp");
|
||||
fs::write(&tmp, &serialized).context("failed to write temporary mcp.json")?;
|
||||
fs::rename(&tmp, &file_path).context("failed to finalize mcp.json")?;
|
||||
|
||||
if !added.is_empty() {
|
||||
println!(" + new MCP servers: {}", added.join(", "));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}),
|
||||
|
||||
+14
-1
@@ -56,6 +56,12 @@ async fn main() -> Result<()> {
|
||||
CompleteEnv::with_factory(Cli::command).complete();
|
||||
let cli = Cli::parse();
|
||||
|
||||
if cli.dangerously_skip_permissions {
|
||||
unsafe {
|
||||
env::set_var("AUTO_CONFIRM", "true");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(shell) = cli.completions {
|
||||
let mut cmd = Cli::command();
|
||||
shell.generate_completions(&mut cmd);
|
||||
@@ -158,7 +164,14 @@ async fn main() -> Result<()> {
|
||||
}
|
||||
|
||||
let url = spec.url.as_deref().expect("validated: remote spec has url");
|
||||
mcp::oauth::run_mcp_oauth_flow(server_name, url, spec.oauth_client_id.as_deref()).await?;
|
||||
mcp::oauth::run_mcp_oauth_flow(
|
||||
server_name,
|
||||
url,
|
||||
spec.oauth.as_ref().and_then(|o| o.client_id.as_deref()),
|
||||
spec.oauth.as_ref().and_then(|o| o.callback_port),
|
||||
spec.oauth.as_ref().and_then(|o| o.redirect_host.as_deref()),
|
||||
)
|
||||
.await?;
|
||||
println!("Authentication saved. '{server_name}' is now available for use.");
|
||||
|
||||
return Ok(());
|
||||
|
||||
+29
-17
@@ -6,6 +6,7 @@ use crate::config::paths;
|
||||
use crate::utils::{AbortSignal, abortable_run_with_spinner};
|
||||
use crate::vault::Vault;
|
||||
use crate::vault::interpolate_secrets;
|
||||
use anyhow::Error;
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use http::{HeaderName, HeaderValue};
|
||||
@@ -57,6 +58,16 @@ pub(crate) struct McpServersConfig {
|
||||
pub mcp_servers: IndexMap<String, McpServer>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub(crate) struct McpOAuthConfig {
|
||||
#[serde(rename = "clientId", skip_serializing_if = "Option::is_none")]
|
||||
pub client_id: Option<String>,
|
||||
#[serde(rename = "callbackPort", skip_serializing_if = "Option::is_none")]
|
||||
pub callback_port: Option<u16>,
|
||||
#[serde(rename = "redirectHost", skip_serializing_if = "Option::is_none")]
|
||||
pub redirect_host: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub(crate) struct McpServer {
|
||||
@@ -75,7 +86,7 @@ pub(crate) struct McpServer {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub headers: Option<IndexMap<String, String>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub oauth_client_id: Option<String>,
|
||||
pub oauth: Option<McpOAuthConfig>,
|
||||
}
|
||||
|
||||
impl McpServer {
|
||||
@@ -110,10 +121,10 @@ impl McpServer {
|
||||
"MCP server '{name}' is missing a \"command\" field (required for stdio transport)"
|
||||
));
|
||||
}
|
||||
if self.url.is_some() || self.headers.is_some() || self.oauth_client_id.is_some() {
|
||||
if self.url.is_some() || self.headers.is_some() || self.oauth.is_some() {
|
||||
return Err(anyhow!(
|
||||
"MCP server '{name}' has type \"stdio\" but also specifies remote fields \
|
||||
(url/headers/oauth_client_id). Remove the remote fields or change the type to \"http\" or \"sse\"."
|
||||
(url/headers/oauth). Remove the remote fields or change the type to \"http\" or \"sse\"."
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -278,7 +289,7 @@ impl McpRegistry {
|
||||
Err(e) if is_auth_required_error(&e) => {
|
||||
warn!(
|
||||
"MCP server '{id}' requires OAuth authentication. \
|
||||
Run `.mcp auth {id}` in the REPL to authenticate, then restart Coyote."
|
||||
Run `coyote --auth-mcp {id}` or `.mcp auth {id}` in the REPL to authenticate."
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
@@ -399,8 +410,9 @@ fn merge_bearer_token(
|
||||
}
|
||||
}
|
||||
|
||||
fn is_auth_required_error(e: &anyhow::Error) -> bool {
|
||||
e.to_string().contains("Auth required")
|
||||
pub(crate) fn is_auth_required_error(e: &Error) -> bool {
|
||||
e.chain()
|
||||
.any(|cause| cause.to_string().contains("Auth required"))
|
||||
}
|
||||
|
||||
async fn spawn_http_mcp_server(
|
||||
@@ -511,7 +523,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -524,7 +536,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: Some(url.to_string()),
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -537,7 +549,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: Some(url.to_string()),
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -569,7 +581,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -587,7 +599,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: Some("http://localhost".into()),
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -607,7 +619,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: Some(headers),
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -632,7 +644,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -650,7 +662,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: Some("http://localhost".into()),
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -668,7 +680,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: Some("http://localhost".into()),
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -686,7 +698,7 @@ mod tests {
|
||||
cwd: Some("/tmp".into()),
|
||||
url: Some("http://localhost".into()),
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
@@ -711,7 +723,7 @@ mod tests {
|
||||
cwd: None,
|
||||
url: None,
|
||||
headers: None,
|
||||
oauth_client_id: None,
|
||||
oauth: None,
|
||||
};
|
||||
|
||||
let err = spec.validate("test").unwrap_err();
|
||||
|
||||
+6
-2
@@ -80,13 +80,17 @@ pub async fn run_mcp_oauth_flow(
|
||||
server_name: &str,
|
||||
server_url: &str,
|
||||
configured_client_id: Option<&str>,
|
||||
callback_port: Option<u16>,
|
||||
redirect_host: Option<&str>,
|
||||
) -> Result<()> {
|
||||
let metadata = discover_oauth_metadata(server_url).await?;
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0")?;
|
||||
let host = redirect_host.unwrap_or("127.0.0.1");
|
||||
let bind_addr = format!("127.0.0.1:{}", callback_port.unwrap_or(0));
|
||||
let listener = TcpListener::bind(&bind_addr)?;
|
||||
let port = listener.local_addr()?.port();
|
||||
drop(listener);
|
||||
let redirect_uri = format!("http://127.0.0.1:{port}/callback");
|
||||
let redirect_uri = format!("http://{host}:{port}/callback");
|
||||
|
||||
let client_id = if let Some(id) = configured_client_id {
|
||||
id.to_string()
|
||||
|
||||
@@ -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(),
|
||||
|
||||
+18
-6
@@ -640,13 +640,25 @@ pub async fn run_repl_command(
|
||||
.url
|
||||
.as_deref()
|
||||
.expect("validated: remote spec has url");
|
||||
let client_id = spec.oauth_client_id.as_deref();
|
||||
mcp::oauth::run_mcp_oauth_flow(server_name, url, client_id)
|
||||
let client_id = spec
|
||||
.oauth
|
||||
.as_ref()
|
||||
.and_then(|o| o.client_id.as_deref());
|
||||
let callback_port =
|
||||
spec.oauth.as_ref().and_then(|o| o.callback_port);
|
||||
let redirect_host = spec
|
||||
.oauth
|
||||
.as_ref()
|
||||
.and_then(|o| o.redirect_host.as_deref());
|
||||
mcp::oauth::run_mcp_oauth_flow(
|
||||
server_name,
|
||||
url,
|
||||
client_id,
|
||||
callback_port,
|
||||
redirect_host,
|
||||
)
|
||||
.await?;
|
||||
println!(
|
||||
"Authentication saved. \
|
||||
Restart Coyote to connect to '{server_name}'."
|
||||
);
|
||||
println!("Authentication saved.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user