Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d8eec1d427
|
||
|
|
382916c3ee
|
||
|
|
bc3cc10a7b
|
||
|
|
b91f738209
|
||
|
|
4f0dae9b49
|
||
|
|
deb673ebc9
|
+5
-1
@@ -134,6 +134,10 @@ enabled_mcp_servers: null # Which MCP servers to enable by default.
|
||||
# - slack
|
||||
# Example (comma-separated form):
|
||||
# enabled_mcp_servers: github,slack,ddg-search
|
||||
no_workspace_mcp: false # Disable loading workspace-local MCP servers from .coyote/mcp.json (default: false).
|
||||
# When false (the default), Coyote merges .coyote/mcp.json from the current directory
|
||||
# into the global MCP registry at startup. Workspace entries shadow global ones on
|
||||
# name collision. Set to true (or pass --no-workspace-mcp) to skip this entirely.
|
||||
|
||||
# ---- Skills ----
|
||||
# Skills are modular knowledge or capability packs the LLM can load and unload mid-conversation.
|
||||
@@ -199,7 +203,7 @@ rag_chunk_size: null # Defines the size of chunks for document proce
|
||||
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)
|
||||
rag_graph_hops: 1 # Number of hops to expand from matched entities at query time (0 = seed nodes only; 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)
|
||||
|
||||
+1
-1
@@ -227,7 +227,7 @@ nodes:
|
||||
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)
|
||||
graph_hops: 1 # Graph expansion depth at query time (0 = seed nodes only; 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`
|
||||
|
||||
@@ -195,6 +195,9 @@ pub struct Cli {
|
||||
/// Skip discovery and application of all sbx mixins (user and built-in)
|
||||
#[arg(long, requires = "sandbox")]
|
||||
pub no_mixins: bool,
|
||||
/// Disable loading workspace MCP servers from .coyote/mcp.json
|
||||
#[arg(long)]
|
||||
pub no_workspace_mcp: bool,
|
||||
}
|
||||
|
||||
impl Cli {
|
||||
|
||||
@@ -50,7 +50,7 @@ fn prepare_chat_completions(
|
||||
|
||||
let url = format!(
|
||||
"{}/openai/deployments/{}/chat/completions?api-version=2024-12-01-preview",
|
||||
&api_base,
|
||||
api_base,
|
||||
self_.model.real_name()
|
||||
);
|
||||
|
||||
@@ -69,7 +69,7 @@ fn prepare_embeddings(self_: &AzureOpenAIClient, data: &EmbeddingsData) -> Resul
|
||||
|
||||
let url = format!(
|
||||
"{}/openai/deployments/{}/embeddings?api-version=2024-10-21",
|
||||
&api_base,
|
||||
api_base,
|
||||
self_.model.real_name()
|
||||
);
|
||||
|
||||
|
||||
@@ -88,6 +88,7 @@ pub struct AppConfig {
|
||||
|
||||
pub user_agent: Option<String>,
|
||||
pub save_shell_history: bool,
|
||||
pub no_workspace_mcp: bool,
|
||||
pub sync_models_url: Option<String>,
|
||||
|
||||
pub clients: Vec<ClientConfig>,
|
||||
@@ -162,6 +163,7 @@ impl Default for AppConfig {
|
||||
|
||||
user_agent: None,
|
||||
save_shell_history: true,
|
||||
no_workspace_mcp: false,
|
||||
sync_models_url: None,
|
||||
|
||||
clients: vec![],
|
||||
@@ -238,6 +240,7 @@ impl AppConfig {
|
||||
|
||||
user_agent: config.user_agent,
|
||||
save_shell_history: config.save_shell_history,
|
||||
no_workspace_mcp: false,
|
||||
sync_models_url: config.sync_models_url,
|
||||
|
||||
clients: config.clients,
|
||||
|
||||
+10
-10
@@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::config::{
|
||||
GIT_DIR_NAME, GITIGNORE_FILE_NAME, MEMORY_DIR_NAME, MEMORY_INDEX_FILE_NAME,
|
||||
WORKSPACE_MEMORY_DIR_NAME, WORKSPACE_MEMORY_FILE_NAME, paths,
|
||||
WORKSPACE_COYOTE_DIR_NAME, WORKSPACE_MEMORY_FILE_NAME, paths,
|
||||
};
|
||||
|
||||
pub const DEFAULT_MEMORY_CAP_WITH_TOOLS: usize = 6_000;
|
||||
@@ -27,7 +27,7 @@ pub enum WorkspaceMemory {
|
||||
|
||||
pub fn discover_workspace_memory(start: &Path) -> Option<WorkspaceMemory> {
|
||||
for dir in start.ancestors() {
|
||||
let structured = dir.join(WORKSPACE_MEMORY_DIR_NAME).join(MEMORY_DIR_NAME);
|
||||
let structured = dir.join(WORKSPACE_COYOTE_DIR_NAME).join(MEMORY_DIR_NAME);
|
||||
if structured.join(MEMORY_INDEX_FILE_NAME).exists() {
|
||||
return Some(WorkspaceMemory::Structured {
|
||||
workspace_root: dir.to_path_buf(),
|
||||
@@ -84,8 +84,8 @@ pub fn bootstrap_workspace_memory(git_root: &Path) -> Result<PathBuf> {
|
||||
|
||||
fn append_gitignore_entry(git_root: &Path) -> Result<bool> {
|
||||
let gitignore = git_root.join(GITIGNORE_FILE_NAME);
|
||||
let entry = format!("{WORKSPACE_MEMORY_DIR_NAME}/{MEMORY_DIR_NAME}/");
|
||||
let entry_no_slash = format!("{WORKSPACE_MEMORY_DIR_NAME}/{MEMORY_DIR_NAME}");
|
||||
let entry = format!("{WORKSPACE_COYOTE_DIR_NAME}/{MEMORY_DIR_NAME}/");
|
||||
let entry_no_slash = format!("{WORKSPACE_COYOTE_DIR_NAME}/{MEMORY_DIR_NAME}");
|
||||
|
||||
let existing = fs::read_to_string(&gitignore).unwrap_or_default();
|
||||
let already_present = existing.lines().any(|line| {
|
||||
@@ -347,7 +347,7 @@ mod tests {
|
||||
let root = temp_root("phase1");
|
||||
let workspace = root.join("workspace");
|
||||
let workspace_memory_dir = workspace
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME);
|
||||
fs::create_dir_all(&workspace_memory_dir).unwrap();
|
||||
fs::write(
|
||||
@@ -382,7 +382,7 @@ mod tests {
|
||||
let root = temp_root("prefer");
|
||||
let workspace = root.join("ws");
|
||||
let structured = workspace
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME);
|
||||
fs::create_dir_all(&structured).unwrap();
|
||||
fs::write(structured.join(MEMORY_INDEX_FILE_NAME), "s").unwrap();
|
||||
@@ -415,7 +415,7 @@ mod tests {
|
||||
let root = temp_root("indexes_only");
|
||||
let workspace = root.join("ws");
|
||||
let structured = workspace
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME);
|
||||
fs::create_dir_all(&structured).unwrap();
|
||||
fs::write(
|
||||
@@ -450,7 +450,7 @@ mod tests {
|
||||
let root = temp_root("drill_bodies");
|
||||
let workspace = root.join("ws");
|
||||
let structured = workspace
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME);
|
||||
fs::create_dir_all(&structured).unwrap();
|
||||
fs::write(structured.join(MEMORY_INDEX_FILE_NAME), "idx").unwrap();
|
||||
@@ -485,7 +485,7 @@ mod tests {
|
||||
let root = temp_root("cap");
|
||||
let workspace = root.join("ws");
|
||||
let structured = workspace
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME);
|
||||
fs::create_dir_all(&structured).unwrap();
|
||||
fs::write(structured.join(MEMORY_INDEX_FILE_NAME), "idx").unwrap();
|
||||
@@ -575,7 +575,7 @@ mod tests {
|
||||
let root = temp_root("walk_up");
|
||||
let workspace = root.join("ws");
|
||||
let mem_dir = workspace
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME);
|
||||
fs::create_dir_all(&mem_dir).unwrap();
|
||||
fs::write(mem_dir.join(MEMORY_INDEX_FILE_NAME), "idx").unwrap();
|
||||
|
||||
+1
-1
@@ -143,7 +143,7 @@ const MCP_FILE_NAME: &str = "mcp.json";
|
||||
const MEMORY_DIR_NAME: &str = "memory";
|
||||
const MEMORY_INDEX_FILE_NAME: &str = "MEMORY.md";
|
||||
const WORKSPACE_MEMORY_FILE_NAME: &str = "COYOTE.md";
|
||||
const WORKSPACE_MEMORY_DIR_NAME: &str = ".coyote";
|
||||
const WORKSPACE_COYOTE_DIR_NAME: &str = ".coyote";
|
||||
const SBX_KIT_DIR_NAME: &str = "sbx-kit";
|
||||
const SBX_KIT_HASH_FILE: &str = "kit.sha256";
|
||||
const SBX_MIXIN_FILE_NAME: &str = "sbx-mixin.yaml";
|
||||
|
||||
+37
-13
@@ -5,7 +5,7 @@ use super::{
|
||||
GLOBAL_TOOLS_UTILS_DIR_NAME, MACROS_DIR_NAME, MCP_FILE_NAME, MEMORY_DIR_NAME,
|
||||
MEMORY_INDEX_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME, SBX_KIT_DIR_NAME,
|
||||
SBX_KIT_HASH_FILE, SBX_MIXIN_FILE_NAME, SBX_MIXIN_KITS_DIR_NAME, SBX_VAULT_MIXINS_DIR_NAME,
|
||||
SKILLS_DIR_NAME, WORKSPACE_MEMORY_DIR_NAME,
|
||||
SKILLS_DIR_NAME, WORKSPACE_COYOTE_DIR_NAME,
|
||||
};
|
||||
use crate::client::ProviderModels;
|
||||
use crate::config::REPL_HISTORY_DIR_NAME;
|
||||
@@ -118,7 +118,7 @@ pub fn global_tools_sbx_mixin_file() -> PathBuf {
|
||||
pub fn find_workspace_sbx_mixin(start: &Path) -> Option<PathBuf> {
|
||||
for dir in start.ancestors() {
|
||||
let candidate = dir
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(SBX_MIXIN_FILE_NAME);
|
||||
if candidate.exists() {
|
||||
return Some(candidate);
|
||||
@@ -193,6 +193,24 @@ pub fn skill_file(name: &str) -> PathBuf {
|
||||
skill_dir(name).join("SKILL.md")
|
||||
}
|
||||
|
||||
pub fn workspace_skills_dir() -> PathBuf {
|
||||
env::current_dir()
|
||||
.unwrap_or_default()
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(SKILLS_DIR_NAME)
|
||||
}
|
||||
|
||||
pub fn workspace_skill_file(name: &str) -> PathBuf {
|
||||
workspace_skills_dir().join(name).join("SKILL.md")
|
||||
}
|
||||
|
||||
pub fn workspace_mcp_config_file() -> PathBuf {
|
||||
env::current_dir()
|
||||
.unwrap_or_default()
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MCP_FILE_NAME)
|
||||
}
|
||||
|
||||
pub fn validate_skill_name(name: &str) -> Result<()> {
|
||||
if name.is_empty() {
|
||||
bail!("Skill name cannot be empty");
|
||||
@@ -318,7 +336,7 @@ pub fn global_memory_index_path() -> PathBuf {
|
||||
|
||||
pub fn workspace_memory_dir_for(workspace_root: &Path) -> PathBuf {
|
||||
workspace_root
|
||||
.join(WORKSPACE_MEMORY_DIR_NAME)
|
||||
.join(WORKSPACE_COYOTE_DIR_NAME)
|
||||
.join(MEMORY_DIR_NAME)
|
||||
}
|
||||
|
||||
@@ -405,15 +423,21 @@ pub fn has_macro(name: &str) -> bool {
|
||||
|
||||
pub fn list_skills() -> Vec<String> {
|
||||
let mut names = Vec::new();
|
||||
if let Ok(rd) = read_dir(skills_dir()) {
|
||||
for entry in rd.flatten() {
|
||||
if let Ok(file_type) = entry.file_type()
|
||||
&& file_type.is_dir()
|
||||
&& let Some(name) = entry.file_name().to_str()
|
||||
&& entry.path().join("SKILL.md").is_file()
|
||||
&& validate_skill_name(name).is_ok()
|
||||
{
|
||||
names.push(name.to_string());
|
||||
let mut seen = HashSet::new();
|
||||
|
||||
for dir in [workspace_skills_dir(), skills_dir()] {
|
||||
if let Ok(rd) = read_dir(dir) {
|
||||
for entry in rd.flatten() {
|
||||
if let Ok(file_type) = entry.file_type()
|
||||
&& file_type.is_dir()
|
||||
&& let Some(name) = entry.file_name().to_str()
|
||||
&& !seen.contains(name)
|
||||
&& entry.path().join("SKILL.md").is_file()
|
||||
&& validate_skill_name(name).is_ok()
|
||||
{
|
||||
seen.insert(name.to_string());
|
||||
names.push(name.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -423,7 +447,7 @@ pub fn list_skills() -> Vec<String> {
|
||||
}
|
||||
|
||||
pub fn has_skill(name: &str) -> bool {
|
||||
skill_file(name).is_file()
|
||||
workspace_skill_file(name).is_file() || skill_file(name).is_file()
|
||||
}
|
||||
|
||||
pub fn local_models_override() -> Result<Vec<ProviderModels>> {
|
||||
|
||||
@@ -261,7 +261,7 @@ impl Session {
|
||||
data["messages"] = json!(self.messages);
|
||||
|
||||
let output = serde_yaml::to_string(&data)
|
||||
.with_context(|| format!("Unable to show info about session '{}'", &self.name))?;
|
||||
.with_context(|| format!("Unable to show info about session '{}'", self.name))?;
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
|
||||
+5
-1
@@ -117,7 +117,11 @@ impl Skill {
|
||||
|
||||
pub fn load(name: &str) -> Result<Self> {
|
||||
paths::validate_skill_name(name)?;
|
||||
let path = paths::skill_file(name);
|
||||
let path = if paths::workspace_skill_file(name).is_file() {
|
||||
paths::workspace_skill_file(name)
|
||||
} else {
|
||||
paths::skill_file(name)
|
||||
};
|
||||
let content = read_to_string(&path)
|
||||
.with_context(|| format!("Failed to read skill '{name}' at {}", path.display()))?;
|
||||
Ok(Skill::new(name, &content))
|
||||
|
||||
+4
-4
@@ -147,10 +147,10 @@ pub async fn eval_tool_calls(
|
||||
let mut is_all_null = true;
|
||||
for call in calls {
|
||||
if let Some(msg) = ctx.tool_scope.tool_tracker.check_loop(&call.clone()) {
|
||||
let dup_msg = format!("{{\"tool_call_loop_alert\":{}}}", &msg.trim());
|
||||
let dup_msg = format!("{{\"tool_call_loop_alert\":{}}}", msg.trim());
|
||||
println!(
|
||||
"{}",
|
||||
warning_text(format!("{}: ⚠️ Tool-call loop detected! ⚠️", &call.name).as_str())
|
||||
warning_text(format!("{}: ⚠️ Tool-call loop detected! ⚠️", call.name).as_str())
|
||||
);
|
||||
let val = json!(dup_msg);
|
||||
output.push(ToolResult::new(call, val));
|
||||
@@ -870,7 +870,7 @@ impl Functions {
|
||||
let root_dir = paths::functions_dir();
|
||||
let tool_path = format!(
|
||||
"{}/{binary_name}",
|
||||
&paths::global_tools_dir().to_string_lossy()
|
||||
paths::global_tools_dir().to_string_lossy()
|
||||
);
|
||||
content_template
|
||||
.replace("{function_name}", binary_name)
|
||||
@@ -881,7 +881,7 @@ impl Functions {
|
||||
let root_dir = paths::agent_data_dir(agent_name);
|
||||
let tool_path = format!(
|
||||
"{}/{binary_name}",
|
||||
&paths::global_tools_dir().to_string_lossy()
|
||||
paths::global_tools_dir().to_string_lossy()
|
||||
);
|
||||
content_template
|
||||
.replace("{function_name}", binary_name)
|
||||
|
||||
+6
-2
@@ -187,7 +187,11 @@ async fn main() -> Result<()> {
|
||||
let abort_signal = create_abort_signal();
|
||||
let start_mcp_servers = cli.agent.is_none() && cli.role.is_none();
|
||||
let cfg = Config::load_with_interpolation(info_flag).await?;
|
||||
let app_config: Arc<AppConfig> = Arc::new(AppConfig::from_config(cfg)?);
|
||||
let mut app_config = AppConfig::from_config(cfg)?;
|
||||
if cli.no_workspace_mcp {
|
||||
app_config.no_workspace_mcp = true;
|
||||
}
|
||||
let app_config: Arc<AppConfig> = Arc::new(app_config);
|
||||
let app_state: Arc<AppState> = Arc::new(
|
||||
AppState::init(
|
||||
app_config,
|
||||
@@ -559,7 +563,7 @@ async fn shell_execute(
|
||||
|
||||
match answer_char {
|
||||
'e' => {
|
||||
debug!("{} {:?}", shell.cmd, &[&shell.arg, &eval_str]);
|
||||
debug!("{} {:?}", shell.cmd, [&shell.arg, &eval_str]);
|
||||
let code = run_command(&shell.cmd, &[&shell.arg, &eval_str], None)?;
|
||||
if code == 0 && app.save_shell_history {
|
||||
let _ = append_to_shell_history(&shell.name, &eval_str, code);
|
||||
|
||||
+47
-1
@@ -214,7 +214,53 @@ impl McpRegistry {
|
||||
spec.validate(name)?;
|
||||
}
|
||||
|
||||
registry.config = Some(mcp_servers_config);
|
||||
let mut merged = mcp_servers_config;
|
||||
if !app_config.no_workspace_mcp {
|
||||
let ws_path = paths::workspace_mcp_config_file();
|
||||
if ws_path.try_exists().unwrap_or(false) {
|
||||
match tokio::fs::read_to_string(&ws_path).await {
|
||||
Ok(ws_content) if !ws_content.trim().is_empty() => {
|
||||
match interpolate_secrets(&ws_content, vault) {
|
||||
Ok((parsed, missing)) if missing.is_empty() => {
|
||||
match serde_json::from_str::<McpServersConfig>(&parsed) {
|
||||
Ok(ws_config) => {
|
||||
let mut loaded = Vec::new();
|
||||
for (name, spec) in ws_config.mcp_servers {
|
||||
match spec.validate(&name) {
|
||||
Ok(_) => {
|
||||
loaded.push(name.clone());
|
||||
merged.mcp_servers.insert(name, spec);
|
||||
}
|
||||
Err(e) => warn!(
|
||||
"Invalid workspace MCP server '{name}': {e}. Skipping."
|
||||
),
|
||||
}
|
||||
}
|
||||
if !loaded.is_empty() {
|
||||
eprintln!(
|
||||
"Loading workspace MCP servers: {}",
|
||||
loaded.join(", ")
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => warn!(
|
||||
"Failed to parse workspace MCP config: {e}. Skipping."
|
||||
),
|
||||
}
|
||||
}
|
||||
Ok((_, missing)) => warn!(
|
||||
"Workspace MCP config references missing vault secrets: {missing:?}. Skipping."
|
||||
),
|
||||
Err(e) => {
|
||||
warn!("Failed to process workspace MCP config: {e}. Skipping.")
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
registry.config = Some(merged);
|
||||
|
||||
if start_mcp_servers && app_config.mcp_server_support {
|
||||
abortable_run_with_spinner(
|
||||
|
||||
+503
-21
@@ -2,12 +2,21 @@ use super::DocumentId;
|
||||
use crate::client::*;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use indexmap::IndexMap;
|
||||
use indexmap::{IndexMap, IndexSet};
|
||||
use petgraph::Direction;
|
||||
use petgraph::graph::NodeIndex;
|
||||
use petgraph::stable_graph::StableGraph;
|
||||
use petgraph::visit::EdgeRef;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashSet;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
/// Heuristic upper bound on chunk size before warning the user that the
|
||||
/// extraction LLM call may be truncated. Not a hard limit.
|
||||
const MAX_CHUNK_CHARS: usize = 24_000;
|
||||
|
||||
/// Maximum number of nodes the BFS may visit during a single graph_search.
|
||||
/// Keeps the synchronous traversal bounded on dense graphs.
|
||||
pub const MAX_GRAPH_NODES: usize = 500;
|
||||
|
||||
const EXTRACTION_PROMPT: &str = r#"Extract entities and relationships from the following text chunk.
|
||||
|
||||
@@ -89,16 +98,27 @@ impl Default for KnowledgeGraph {
|
||||
|
||||
impl KnowledgeGraph {
|
||||
pub fn merge(&mut self, doc_id: DocumentId, result: ExtractionResult) {
|
||||
let mut chunk_nodes: Vec<u32> = vec![];
|
||||
let mut chunk_nodes: IndexSet<u32> = IndexSet::new();
|
||||
|
||||
for extracted in &result.entities {
|
||||
let key = extracted.name.to_lowercase();
|
||||
let normalized_type = extracted.entity_type.to_uppercase();
|
||||
let node_raw = if let Some(&existing) = self.entity_index.get(&key) {
|
||||
let idx = NodeIndex::new(existing as usize);
|
||||
if self.graph.contains_node(idx) {
|
||||
let node = &mut self.graph[idx];
|
||||
if node.entity_type == "OTHER" && normalized_type != "OTHER" {
|
||||
node.entity_type = normalized_type;
|
||||
}
|
||||
if node.description.is_none() {
|
||||
node.description = extracted.description.clone();
|
||||
}
|
||||
}
|
||||
existing
|
||||
} else {
|
||||
let entity = Entity {
|
||||
name: extracted.name.clone(),
|
||||
entity_type: extracted.entity_type.clone(),
|
||||
entity_type: normalized_type,
|
||||
description: extracted.description.clone(),
|
||||
};
|
||||
let idx = self.graph.add_node(entity);
|
||||
@@ -106,7 +126,7 @@ impl KnowledgeGraph {
|
||||
self.entity_index.insert(key, raw);
|
||||
raw
|
||||
};
|
||||
chunk_nodes.push(node_raw);
|
||||
chunk_nodes.insert(node_raw);
|
||||
}
|
||||
|
||||
for extracted in &result.relationships {
|
||||
@@ -118,11 +138,14 @@ impl KnowledgeGraph {
|
||||
) {
|
||||
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 already_exists = self
|
||||
.graph
|
||||
.edges_connecting(from_idx, to_idx)
|
||||
.any(|e| e.weight().relation_type == extracted.relation_type);
|
||||
if !already_exists {
|
||||
let rel = Relationship {
|
||||
relation_type: extracted.relation_type.clone(),
|
||||
weight: extracted.weight.unwrap_or(1.0),
|
||||
weight: extracted.weight.unwrap_or(1.0).clamp(0.0, 1.0),
|
||||
};
|
||||
self.graph.add_edge(from_idx, to_idx, rel);
|
||||
}
|
||||
@@ -158,6 +181,10 @@ impl KnowledgeGraph {
|
||||
.filter(|raw| !still_used.contains(raw))
|
||||
.collect();
|
||||
|
||||
if to_remove.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
for raw in to_remove {
|
||||
let idx = NodeIndex::new(raw as usize);
|
||||
if self.graph.contains_node(idx) {
|
||||
@@ -166,6 +193,57 @@ impl KnowledgeGraph {
|
||||
self.entity_index.swap_remove(&name);
|
||||
}
|
||||
}
|
||||
|
||||
self.compact();
|
||||
}
|
||||
|
||||
/// Rebuild the internal graph with consecutive node indices. Eliminates
|
||||
/// the null tombstone slots that petgraph's StableGraph accumulates after
|
||||
/// repeated `remove_node` calls, keeping serialized YAML size in check.
|
||||
fn compact(&mut self) {
|
||||
let mut new_graph: StableGraph<Entity, Relationship> = StableGraph::new();
|
||||
let mut old_to_new: HashMap<u32, u32> = HashMap::new();
|
||||
|
||||
for &old_raw in self.entity_index.values() {
|
||||
let old_idx = NodeIndex::new(old_raw as usize);
|
||||
if self.graph.contains_node(old_idx) {
|
||||
let entity = self.graph[old_idx].clone();
|
||||
let new_idx = new_graph.add_node(entity);
|
||||
old_to_new.insert(old_raw, new_idx.index() as u32);
|
||||
}
|
||||
}
|
||||
|
||||
for edge_idx in self.graph.edge_indices() {
|
||||
if let Some((from, to)) = self.graph.edge_endpoints(edge_idx) {
|
||||
let from_raw = from.index() as u32;
|
||||
let to_raw = to.index() as u32;
|
||||
if let (Some(&new_from), Some(&new_to)) =
|
||||
(old_to_new.get(&from_raw), old_to_new.get(&to_raw))
|
||||
{
|
||||
let rel = self.graph[edge_idx].clone();
|
||||
new_graph.add_edge(
|
||||
NodeIndex::new(new_from as usize),
|
||||
NodeIndex::new(new_to as usize),
|
||||
rel,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for raw in self.entity_index.values_mut() {
|
||||
if let Some(&new_raw) = old_to_new.get(raw) {
|
||||
*raw = new_raw;
|
||||
}
|
||||
}
|
||||
|
||||
for node_raws in self.document_entities.values_mut() {
|
||||
*node_raws = node_raws
|
||||
.iter()
|
||||
.filter_map(|raw| old_to_new.get(raw).copied())
|
||||
.collect();
|
||||
}
|
||||
|
||||
self.graph = new_graph;
|
||||
}
|
||||
|
||||
pub fn build_node_to_docs(&self) -> IndexMap<u32, Vec<DocumentId>> {
|
||||
@@ -179,30 +257,79 @@ impl KnowledgeGraph {
|
||||
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();
|
||||
/// BFS from seed nodes with weight-decayed scoring.
|
||||
///
|
||||
/// Seed node scores are provided by the caller (typically token-overlap
|
||||
/// ratios). Each neighbor's score is `edge_weight * parent_score`, so
|
||||
/// strongly-connected neighbors rank higher and weakly-connected ones
|
||||
/// naturally contribute less. Traversal is capped at `MAX_GRAPH_NODES`
|
||||
/// total nodes; the highest-scored frontier nodes are expanded first so
|
||||
/// the budget is spent on the most relevant entities.
|
||||
///
|
||||
/// Returns a map of raw node index → score (includes seed nodes).
|
||||
pub fn expand_neighbors_scored(
|
||||
&self,
|
||||
seed_scores: &[(u32, f32)],
|
||||
hops: usize,
|
||||
) -> IndexMap<u32, f32> {
|
||||
let mut node_scores: IndexMap<u32, f32> = IndexMap::new();
|
||||
for &(raw, score) in seed_scores {
|
||||
node_scores.insert(raw, score);
|
||||
}
|
||||
|
||||
let mut frontier: Vec<(u32, f32)> = seed_scores.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);
|
||||
if node_scores.len() >= MAX_GRAPH_NODES {
|
||||
break;
|
||||
}
|
||||
|
||||
frontier.sort_unstable_by(|a, b| {
|
||||
b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
let mut next_frontier: Vec<(u32, f32)> = vec![];
|
||||
|
||||
'nodes: for (raw, parent_score) in &frontier {
|
||||
let idx = NodeIndex::new(*raw as usize);
|
||||
if !self.graph.contains_node(idx) {
|
||||
continue;
|
||||
}
|
||||
for dir in [Direction::Outgoing, Direction::Incoming] {
|
||||
for edge_ref in self.graph.edges_directed(idx, dir) {
|
||||
let neighbor_idx = match dir {
|
||||
Direction::Outgoing => edge_ref.target(),
|
||||
Direction::Incoming => edge_ref.source(),
|
||||
};
|
||||
let neighbor_raw = neighbor_idx.index() as u32;
|
||||
let candidate = edge_ref.weight().weight * parent_score;
|
||||
|
||||
match node_scores.entry(neighbor_raw) {
|
||||
indexmap::map::Entry::Vacant(e) => {
|
||||
e.insert(candidate);
|
||||
next_frontier.push((neighbor_raw, candidate));
|
||||
}
|
||||
indexmap::map::Entry::Occupied(mut e) => {
|
||||
if candidate > *e.get() {
|
||||
*e.get_mut() = candidate;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if node_scores.len() >= MAX_GRAPH_NODES {
|
||||
break 'nodes;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
frontier = next_frontier;
|
||||
if frontier.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
expanded.into_iter().collect()
|
||||
|
||||
node_scores
|
||||
}
|
||||
}
|
||||
|
||||
@@ -213,6 +340,14 @@ pub async fn extract_entities(
|
||||
chunk: &str,
|
||||
prompt_template: Option<&str>,
|
||||
) -> Result<ExtractionResult> {
|
||||
if chunk.len() > MAX_CHUNK_CHARS {
|
||||
warn!(
|
||||
"Entity extraction chunk is {} chars (heuristic limit: {}); \
|
||||
the LLM response may be truncated",
|
||||
chunk.len(),
|
||||
MAX_CHUNK_CHARS
|
||||
);
|
||||
}
|
||||
let template = prompt_template.unwrap_or(EXTRACTION_PROMPT);
|
||||
let prompt = template.replace("__CHUNK__", chunk);
|
||||
let mut messages = vec![Message::new(
|
||||
@@ -250,3 +385,350 @@ pub async fn extract_entities(
|
||||
serde_json::from_str::<ExtractionResult>(&json)
|
||||
.context("Failed to parse entity extraction JSON")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn entity(name: &str, entity_type: &str) -> ExtractedEntity {
|
||||
ExtractedEntity {
|
||||
name: name.to_string(),
|
||||
entity_type: entity_type.to_string(),
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn rel(from: &str, to: &str, rel_type: &str, weight: f32) -> ExtractedRelationship {
|
||||
ExtractedRelationship {
|
||||
from: from.to_string(),
|
||||
to: to.to_string(),
|
||||
relation_type: rel_type.to_string(),
|
||||
weight: Some(weight),
|
||||
}
|
||||
}
|
||||
|
||||
fn doc(id: usize) -> DocumentId {
|
||||
DocumentId(id)
|
||||
}
|
||||
|
||||
fn extraction(
|
||||
entities: Vec<ExtractedEntity>,
|
||||
rels: Vec<ExtractedRelationship>,
|
||||
) -> ExtractionResult {
|
||||
ExtractionResult {
|
||||
entities,
|
||||
relationships: rels,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_deduplicates_by_lowercase_name() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![
|
||||
entity("Python", "TECHNOLOGY"),
|
||||
entity("python", "TECHNOLOGY"),
|
||||
],
|
||||
vec![],
|
||||
),
|
||||
);
|
||||
assert_eq!(kg.entity_index.len(), 1);
|
||||
assert_eq!(kg.graph.node_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_chunk_nodes_no_duplicate_doc_entries() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(1),
|
||||
extraction(
|
||||
vec![
|
||||
entity("Python", "TECHNOLOGY"),
|
||||
entity("python", "TECHNOLOGY"),
|
||||
],
|
||||
vec![],
|
||||
),
|
||||
);
|
||||
let count = kg.document_entities.get(&1).map(|v| v.len()).unwrap_or(0);
|
||||
assert_eq!(
|
||||
count, 1,
|
||||
"duplicate entity in one chunk should produce one doc_entity entry"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_normalizes_entity_type_to_uppercase() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(vec![entity("Django", "technology")], vec![]),
|
||||
);
|
||||
let raw = kg.entity_index["django"];
|
||||
assert_eq!(
|
||||
kg.graph[NodeIndex::new(raw as usize)].entity_type,
|
||||
"TECHNOLOGY"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_promotes_type_from_other_to_specific() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(doc(0), extraction(vec![entity("Python", "OTHER")], vec![]));
|
||||
kg.merge(
|
||||
doc(1),
|
||||
extraction(vec![entity("Python", "TECHNOLOGY")], vec![]),
|
||||
);
|
||||
let raw = kg.entity_index["python"];
|
||||
assert_eq!(
|
||||
kg.graph[NodeIndex::new(raw as usize)].entity_type,
|
||||
"TECHNOLOGY"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_does_not_demote_specific_type_to_other() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(vec![entity("Python", "TECHNOLOGY")], vec![]),
|
||||
);
|
||||
kg.merge(doc(1), extraction(vec![entity("Python", "OTHER")], vec![]));
|
||||
let raw = kg.entity_index["python"];
|
||||
assert_eq!(
|
||||
kg.graph[NodeIndex::new(raw as usize)].entity_type,
|
||||
"TECHNOLOGY"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_allows_multiple_relation_types_between_same_pair() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![
|
||||
entity("Python", "TECHNOLOGY"),
|
||||
entity("Django", "TECHNOLOGY"),
|
||||
],
|
||||
vec![rel("Python", "Django", "implements", 0.9)],
|
||||
),
|
||||
);
|
||||
kg.merge(
|
||||
doc(1),
|
||||
extraction(
|
||||
vec![
|
||||
entity("Python", "TECHNOLOGY"),
|
||||
entity("Django", "TECHNOLOGY"),
|
||||
],
|
||||
vec![rel("Python", "Django", "uses", 0.8)],
|
||||
),
|
||||
);
|
||||
let from_idx = NodeIndex::new(kg.entity_index["python"] as usize);
|
||||
let to_idx = NodeIndex::new(kg.entity_index["django"] as usize);
|
||||
let count = kg.graph.edges_connecting(from_idx, to_idx).count();
|
||||
assert_eq!(
|
||||
count, 2,
|
||||
"two different relation types should produce two edges"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_deduplicates_same_relation_type() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
|
||||
vec![rel("A", "B", "uses", 1.0)],
|
||||
),
|
||||
);
|
||||
kg.merge(
|
||||
doc(1),
|
||||
extraction(
|
||||
vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
|
||||
vec![rel("A", "B", "uses", 0.5)],
|
||||
),
|
||||
);
|
||||
let from_idx = NodeIndex::new(kg.entity_index["a"] as usize);
|
||||
let to_idx = NodeIndex::new(kg.entity_index["b"] as usize);
|
||||
let count = kg.graph.edges_connecting(from_idx, to_idx).count();
|
||||
assert_eq!(
|
||||
count, 1,
|
||||
"same relation type should not create a duplicate edge"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_documents_preserves_entity_shared_across_docs() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![entity("Python", "TECHNOLOGY"), entity("A", "CONCEPT")],
|
||||
vec![],
|
||||
),
|
||||
);
|
||||
kg.merge(
|
||||
doc(1),
|
||||
extraction(
|
||||
vec![entity("Python", "TECHNOLOGY"), entity("B", "CONCEPT")],
|
||||
vec![],
|
||||
),
|
||||
);
|
||||
kg.remove_documents(&[doc(0)]);
|
||||
assert!(
|
||||
kg.entity_index.contains_key("python"),
|
||||
"shared entity should survive"
|
||||
);
|
||||
assert!(
|
||||
!kg.entity_index.contains_key("a"),
|
||||
"exclusive entity should be removed"
|
||||
);
|
||||
assert!(
|
||||
kg.entity_index.contains_key("b"),
|
||||
"other doc's entity should survive"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_documents_noop_on_empty_slice() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(doc(0), extraction(vec![entity("X", "CONCEPT")], vec![]));
|
||||
kg.remove_documents(&[]);
|
||||
assert_eq!(kg.entity_index.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_documents_compacts_graph() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
// doc 0: A, B with an edge
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
|
||||
vec![rel("A", "B", "uses", 1.0)],
|
||||
),
|
||||
);
|
||||
// doc 1: C only
|
||||
kg.merge(doc(1), extraction(vec![entity("C", "CONCEPT")], vec![]));
|
||||
|
||||
kg.remove_documents(&[doc(0)]);
|
||||
|
||||
assert_eq!(kg.graph.node_count(), 1);
|
||||
let c_raw = kg.entity_index["c"];
|
||||
assert_eq!(
|
||||
c_raw, 0,
|
||||
"compacted graph should give surviving node index 0"
|
||||
);
|
||||
let refs = kg.document_entities.get(&1).cloned().unwrap_or_default();
|
||||
assert_eq!(refs, vec![0u32]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_zero_hops_returns_seeds_only() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
|
||||
vec![rel("A", "B", "uses", 0.9)],
|
||||
),
|
||||
);
|
||||
let a_raw = kg.entity_index["a"];
|
||||
let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 0);
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[&a_raw], 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_one_hop_decays_score_by_edge_weight() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
|
||||
vec![rel("A", "B", "uses", 0.8)],
|
||||
),
|
||||
);
|
||||
let a_raw = kg.entity_index["a"];
|
||||
let b_raw = kg.entity_index["b"];
|
||||
let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1);
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[&a_raw], 1.0);
|
||||
let b_score = result[&b_raw];
|
||||
assert!(
|
||||
(b_score - 0.8).abs() < 1e-6,
|
||||
"neighbor score should be edge_weight * parent_score = 0.8, got {b_score}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_incoming_edges_also_traversed() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
// Edge goes B → A; seeding A should still discover B via incoming edge
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")],
|
||||
vec![rel("B", "A", "uses", 0.7)],
|
||||
),
|
||||
);
|
||||
let a_raw = kg.entity_index["a"];
|
||||
let b_raw = kg.entity_index["b"];
|
||||
let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1);
|
||||
assert!(
|
||||
result.contains_key(&b_raw),
|
||||
"B should be reachable via incoming edge from A"
|
||||
);
|
||||
let b_score = result[&b_raw];
|
||||
assert!((b_score - 0.7).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_picks_best_path_score() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
// A(0.5) → C(0.9): score 0.45; B(1.0) → C(0.4): score 0.40 — A→C path wins.
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(
|
||||
vec![
|
||||
entity("A", "CONCEPT"),
|
||||
entity("B", "CONCEPT"),
|
||||
entity("C", "CONCEPT"),
|
||||
],
|
||||
vec![rel("A", "C", "uses", 0.9), rel("B", "C", "uses", 0.4)],
|
||||
),
|
||||
);
|
||||
let a_raw = kg.entity_index["a"];
|
||||
let b_raw = kg.entity_index["b"];
|
||||
let c_raw = kg.entity_index["c"];
|
||||
let seeds = vec![(a_raw, 0.5f32), (b_raw, 1.0f32)];
|
||||
let result = kg.expand_neighbors_scored(&seeds, 1);
|
||||
let c_score = result[&c_raw];
|
||||
// Best path: B(1.0) * 0.4 = 0.4, A(0.5) * 0.9 = 0.45 → should be 0.45
|
||||
assert!(
|
||||
(c_score - 0.45).abs() < 1e-6,
|
||||
"C score should reflect best path (0.45), got {c_score}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_node_to_docs_maps_shared_entity_to_multiple_docs() {
|
||||
let mut kg = KnowledgeGraph::default();
|
||||
kg.merge(
|
||||
doc(0),
|
||||
extraction(vec![entity("Python", "TECHNOLOGY")], vec![]),
|
||||
);
|
||||
kg.merge(
|
||||
doc(1),
|
||||
extraction(vec![entity("Python", "TECHNOLOGY")], vec![]),
|
||||
);
|
||||
let n2d = kg.build_node_to_docs();
|
||||
let raw = kg.entity_index["python"];
|
||||
let docs = &n2d[&raw];
|
||||
assert!(docs.contains(&DocumentId(0)));
|
||||
assert!(docs.contains(&DocumentId(1)));
|
||||
}
|
||||
}
|
||||
|
||||
+51
-41
@@ -25,6 +25,8 @@ use std::{
|
||||
};
|
||||
use tokio::time::sleep;
|
||||
|
||||
const BM25_SEED_SCORE: f32 = 0.5;
|
||||
|
||||
const RAG_TEMPLATE: &str = r#"Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags)
|
||||
|
||||
<context>
|
||||
@@ -752,14 +754,14 @@ impl Rag {
|
||||
bail!("No RAG files");
|
||||
}
|
||||
|
||||
if self.data.extractor_model.is_some()
|
||||
&& !new_doc_contents.is_empty()
|
||||
if !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();
|
||||
let mut failures = 0usize;
|
||||
for (i, (doc_id, content)) in new_doc_contents.into_iter().enumerate() {
|
||||
progress(
|
||||
&spinner,
|
||||
@@ -774,14 +776,21 @@ impl Rag {
|
||||
{
|
||||
Ok(result) => self.data.knowledge_graph.merge(doc_id, result),
|
||||
Err(e) => {
|
||||
debug!("Entity extraction failed for doc {doc_id:?}: {e}")
|
||||
warn!("Entity extraction failed for doc {doc_id:?}: {e}");
|
||||
failures += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if failures > 0 {
|
||||
progress(
|
||||
&spinner,
|
||||
format!("Entity extraction: {failures}/{total} chunks failed"),
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => debug!("Failed to create extractor client: {e}"),
|
||||
Err(e) => warn!("Failed to create extractor client: {e}"),
|
||||
},
|
||||
Err(e) => debug!("Extractor model not found: {e}"),
|
||||
Err(e) => warn!("Extractor model not found: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -930,9 +939,31 @@ impl Rag {
|
||||
if kg.entity_index.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
let query_lower = query.to_lowercase();
|
||||
|
||||
let mut seed_nodes: Vec<u32> = kg
|
||||
let query_lower = query.to_lowercase();
|
||||
let query_tokens: Vec<&str> = query_lower.split_whitespace().collect();
|
||||
let token_count = query_tokens.len().max(1);
|
||||
|
||||
let score_node = |raw: u32| -> f32 {
|
||||
let idx = NodeIndex::new(raw as usize);
|
||||
if !kg.graph.contains_node(idx) {
|
||||
return 0.0;
|
||||
}
|
||||
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
|
||||
};
|
||||
|
||||
let mut seed_scores: Vec<(u32, f32)> = kg
|
||||
.entity_index
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
@@ -946,52 +977,31 @@ impl Rag {
|
||||
.any(|token| token.trim_matches(|c: char| !c.is_alphanumeric()) == name_str)
|
||||
}
|
||||
})
|
||||
.map(|(_, &raw)| raw)
|
||||
.map(|(_, &raw)| (raw, score_node(raw).max(BM25_SEED_SCORE)))
|
||||
.collect();
|
||||
|
||||
if seed_nodes.is_empty() {
|
||||
if seed_scores.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;
|
||||
for &raw in node_raws {
|
||||
seed_scores.push((raw, BM25_SEED_SCORE));
|
||||
if seed_scores.len() >= top_k {
|
||||
break 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if seed_nodes.is_empty() {
|
||||
if seed_scores.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
|
||||
let mut scored: Vec<(u32, f32)> = kg
|
||||
.expand_neighbors_scored(&seed_scores, hops)
|
||||
.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));
|
||||
|
||||
@@ -1349,11 +1359,11 @@ fn set_chunk_size(model: &Model) -> Result<usize> {
|
||||
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_help_message("Number of hops to expand from matched entities (0 = seed nodes only, 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(_) => Validation::Valid,
|
||||
_ => Validation::Invalid("Must be a non-negative integer".into()),
|
||||
};
|
||||
Ok(out)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user