diff --git a/crates/terraphim_grep/Cargo.toml b/crates/terraphim_grep/Cargo.toml index 7dc6fd3..1ea9e65 100644 --- a/crates/terraphim_grep/Cargo.toml +++ b/crates/terraphim_grep/Cargo.toml @@ -37,6 +37,7 @@ terraphim_config = { version = "1.15.0" } fff-search = { version = "0.8.4", optional = true } async-trait.workspace = true +reqwest = { workspace = true, features = ["json", "rustls-tls"], optional = true } parking_lot = "0.12" tracing = { workspace = true } tracing-subscriber = { version = "0.3", features = ["env-filter"] } @@ -49,7 +50,7 @@ clap = { version = "4", features = ["derive"] } # compiles to a no-op stub that returns no matches (see hybrid_searcher.rs). # It is also required for the no-KG failover path. default = ["llm", "code-search"] -llm = ["dep:terraphim_service"] +llm = ["dep:terraphim_service", "dep:reqwest"] # Enable fast file-finder code search. This is enabled by default so that # terraphim-grep can fall back to plain enhanced grep when no knowledge graph # thesaurus is configured. diff --git a/crates/terraphim_grep/src/lib.rs b/crates/terraphim_grep/src/lib.rs index fdf0057..2f658f9 100644 --- a/crates/terraphim_grep/src/lib.rs +++ b/crates/terraphim_grep/src/lib.rs @@ -23,6 +23,8 @@ pub mod error; pub mod hybrid_searcher; pub mod kg_curation; +#[cfg(feature = "llm")] +pub mod openrouter_client; pub mod rlm_context; pub mod signatures; pub mod sufficiency_judge; @@ -191,17 +193,19 @@ impl TerraphimGrep { let prompt = ctx.build_prompt(); + let task_instruction = if options.include_answer { + format!( + "{}\n\nSynthesise an answer based on the context above.", + signatures::AnswerSignature {}.instructions() + ) + } else { + "List the relevant findings.\n\nProvide a brief answer based on the context above." + .to_string() + }; + let messages = vec![serde_json::json!({ "role": "user", - "content": format!( - "{}\n\n{}\n\nProvide a brief answer based on the context above.", - prompt, - if options.include_answer { - "Synthesise an answer." - } else { - "List the relevant findings." - } - ) + "content": format!("{}\n\n{}", prompt, task_instruction) })]; let llm_response = if let Some(ref client) = self.llm_client { @@ -433,4 +437,73 @@ mod tests { ); assert_eq!(result.stats.kg_hits, 0); } + + /// The RLM prompt for `include_answer` must embed the `AnswerSignature` + /// JSON instructions so the model knows it must return structured output. + #[test] + fn test_answer_prompt_includes_json_instructions() { + let prompt = "Query: test\n## Retrieved Context\nchunk".to_string(); + let include_answer = true; + + let task_instruction = if include_answer { + format!( + "{}\n\nSynthesise an answer based on the context above.", + signatures::AnswerSignature {}.instructions() + ) + } else { + "List the relevant findings.\n\nProvide a brief answer based on the context above." + .to_string() + }; + + let message = serde_json::json!({ + "role": "user", + "content": format!("{}\n\n{}", prompt, task_instruction) + }); + + let content = message["content"].as_str().expect("content string"); + assert!( + content.contains("\"answer\": the synthesised answer"), + "prompt must embed AnswerSignature instructions" + ); + assert!( + content.contains("\"citations\": array of {source, line (optional), excerpt}"), + "prompt must embed citation format" + ); + assert!( + content.contains("\"confidence\": a number between 0 and 1"), + "prompt must embed confidence format" + ); + } + + /// The non-answer path must NOT embed AnswerSignature instructions. + #[test] + fn test_list_prompt_excludes_json_instructions() { + let prompt = "Query: test\n## Retrieved Context\nchunk".to_string(); + let include_answer = false; + + let task_instruction = if include_answer { + format!( + "{}\n\nSynthesise an answer based on the context above.", + signatures::AnswerSignature {}.instructions() + ) + } else { + "List the relevant findings.\n\nProvide a brief answer based on the context above." + .to_string() + }; + + let message = serde_json::json!({ + "role": "user", + "content": format!("{}\n\n{}", prompt, task_instruction) + }); + + let content = message["content"].as_str().expect("content string"); + assert!( + !content.contains("\"answer\": the synthesised answer"), + "list prompt must not embed AnswerSignature instructions" + ); + assert!( + content.contains("List the relevant findings"), + "list prompt must contain the list instruction" + ); + } } diff --git a/crates/terraphim_grep/src/main.rs b/crates/terraphim_grep/src/main.rs index 938024d..a796383 100644 --- a/crates/terraphim_grep/src/main.rs +++ b/crates/terraphim_grep/src/main.rs @@ -4,6 +4,8 @@ use std::sync::Arc; use anyhow::{Context, Result}; use clap::{Parser, Subcommand, ValueEnum}; use terraphim_automata::AutomataPath; +#[cfg(feature = "llm")] +use terraphim_grep::openrouter_client; use terraphim_grep::{ GrepOptions, GrepResult, Haystack, HybridSearcher, SufficiencyJudge, TerraphimGrep, }; @@ -318,6 +320,37 @@ fn build_llm_for_role( } } } + + // Prefer a long-timeout OpenRouter client built directly in grep so + // slow LLM providers do not hit the shared 10-second API timeout. + // This takes precedence over `role_from_env` because `OPENROUTER_API_KEY` + // is the most common way to configure grep. If the direct client cannot + // be built, fall back to `role_from_env` rather than disabling the LLM. + if let Some(key) = std::env::var("OPENROUTER_API_KEY") + .ok() + .filter(|s| !s.is_empty()) + { + let model = std::env::var("OPENROUTER_MODEL") + .unwrap_or_else(|_| "qwen/qwen3-coder:free".to_string()); + if model.ends_with(":free") { + tracing::warn!( + "Using free-tier default model `{}`. RLM prompts may contain \ + repository content excerpts; verify the provider's data-handling \ + terms are acceptable for your codebase.", + model + ); + } + match openrouter_client::OpenRouterClient::new(&key, &model) { + Ok(client) => return Some(openrouter_client::into_llm_client(client)), + Err(e) => { + tracing::warn!( + "Failed to build direct OpenRouter client ({}); falling back to role_from_env", + e + ); + } + } + } + role_from_env(role_name)? } }; diff --git a/crates/terraphim_grep/src/openrouter_client.rs b/crates/terraphim_grep/src/openrouter_client.rs new file mode 100644 index 0000000..bace19d --- /dev/null +++ b/crates/terraphim_grep/src/openrouter_client.rs @@ -0,0 +1,210 @@ +//! Thin OpenRouter client for `terraphim-grep` with a long request timeout. +//! +//! The shared `terraphim_service` API client uses a 10-second timeout, which is +//! too short for many LLM providers when the RLM prompt is large. This client +//! keeps the same request shape but uses a 120-second timeout so RLM synthesis +//! can complete without aborting the connection mid-response. +//! +//! **Maintainability note**: this module intentionally duplicates the shared +//! `terraphim_service` OpenRouter client rather than modifying it, because +//! `terraphim_service` is consumed as an external registry dependency and its +//! timeout is not currently configurable. A TODO item is to upstream a +//! configurable timeout into `terraphim_service` and remove this duplication. +//! +//! **Safety note**: `summarize()` returns a hard error. This is safe only as +//! long as no `terraphim-grep` code path calls `summarize()` through the +//! `LlmClient` trait object; all current grep synthesis uses +//! `chat_completion()`. + +use std::sync::Arc; +use std::time::Duration; + +use terraphim_service::llm::{ChatOptions, LlmClient, SummarizeOptions}; + +/// OpenRouter-backed LLM client with an extended timeout. +pub struct OpenRouterClient { + client: reqwest::Client, + api_key: String, + model: String, + base_url: String, +} + +impl OpenRouterClient { + /// Create a new client for the given OpenRouter API key and model id. + pub fn new(api_key: &str, model: &str) -> anyhow::Result { + anyhow::ensure!(!api_key.is_empty(), "OpenRouter API key cannot be empty"); + anyhow::ensure!(!model.is_empty(), "OpenRouter model cannot be empty"); + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(120)) + .user_agent(concat!( + "terraphim-grep/", + env!("CARGO_PKG_VERSION"), + " (https://terraphim.ai)" + )) + .build()?; + + Ok(Self { + client, + api_key: api_key.to_string(), + model: model.to_string(), + base_url: std::env::var("OPENROUTER_BASE_URL") + .unwrap_or_else(|_| "https://openrouter.ai/api/v1".to_string()), + }) + } +} + +#[async_trait::async_trait] +impl LlmClient for OpenRouterClient { + fn name(&self) -> &'static str { + "openrouter" + } + + async fn summarize( + &self, + _content: &str, + _opts: SummarizeOptions, + ) -> terraphim_service::Result { + Err(terraphim_service::ServiceError::Config( + "summarize not supported by terraphim-grep openrouter client".to_string(), + )) + } + + async fn chat_completion( + &self, + messages: Vec, + opts: ChatOptions, + ) -> terraphim_service::Result { + let request_body = serde_json::json!({ + "model": self.model, + "messages": messages, + "max_tokens": opts.max_tokens.unwrap_or(512), + "temperature": opts.temperature.unwrap_or(0.2), + "stream": false + }); + + let response = self + .client + .post(format!("{}/chat/completions", self.base_url)) + .header("Authorization", format!("Bearer {}", self.api_key)) + .header("Content-Type", "application/json") + .header("HTTP-Referer", "https://terraphim.ai") + .header("X-Title", "Terraphim AI") + .json(&request_body) + .send() + .await + .map_err(|e| { + terraphim_service::ServiceError::Common( + terraphim_service::error::CommonError::Network { + message: format!("OpenRouter request failed: {e}"), + source: Some(Box::new(e)), + }, + ) + })?; + + if response.status() == 429 { + return Err(terraphim_service::ServiceError::Common( + terraphim_service::error::CommonError::Network { + message: "OpenRouter rate limit exceeded".to_string(), + source: None, + }, + )); + } + if !response.status().is_success() { + let error_text = response + .text() + .await + .unwrap_or_else(|_| "Unknown error".to_string()); + return Err(terraphim_service::ServiceError::Common( + terraphim_service::error::CommonError::Network { + message: format!("OpenRouter API error: {error_text}"), + source: None, + }, + )); + } + + let response_json: serde_json::Value = response.json().await.map_err(|e| { + terraphim_service::ServiceError::Common( + terraphim_service::error::CommonError::Network { + message: format!("OpenRouter response JSON decode failed: {e}"), + source: Some(Box::new(e)), + }, + ) + })?; + + let content = response_json + .get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("message")) + .and_then(|m| m.get("content")) + .and_then(|t| t.as_str()) + .unwrap_or("") + .to_string(); + + Ok(content) + } +} + +/// Convenience wrapper that returns the client as a trait object. +pub fn into_llm_client(client: OpenRouterClient) -> Arc { + Arc::new(client) as Arc +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_openrouter_client_new_validates_inputs() { + assert!(OpenRouterClient::new("", "model").is_err()); + assert!(OpenRouterClient::new("key", "").is_err()); + assert!(OpenRouterClient::new("key", "model").is_ok()); + } + + #[test] + fn test_openrouter_client_base_url_resolution() { + // SAFETY: test code, single-threaded, no concurrent env access. + let saved = std::env::var("OPENROUTER_BASE_URL").ok(); + + // Default when unset. + unsafe { + std::env::remove_var("OPENROUTER_BASE_URL"); + } + let client = OpenRouterClient::new("key", "model").expect("build client"); + assert_eq!(client.base_url, "https://openrouter.ai/api/v1"); + + // Custom when set. + unsafe { + std::env::set_var("OPENROUTER_BASE_URL", "https://example.com/v1"); + } + let client = OpenRouterClient::new("key", "model").expect("build client"); + assert_eq!(client.base_url, "https://example.com/v1"); + + // Restore. + if let Some(v) = saved { + unsafe { + std::env::set_var("OPENROUTER_BASE_URL", v); + } + } else { + unsafe { + std::env::remove_var("OPENROUTER_BASE_URL"); + } + } + } + + #[test] + fn test_openrouter_client_timeout_is_120_seconds() { + let client = OpenRouterClient::new("key", "model").expect("build client"); + // reqwest::Client does not expose its timeout directly, but we can + // verify the client was built successfully with the expected config. + assert_eq!(client.model, "model"); + assert_eq!(client.api_key, "key"); + } + + #[test] + fn test_into_llm_client_returns_trait_object() { + let client = OpenRouterClient::new("key", "model").expect("build client"); + let llm: Arc = into_llm_client(client); + assert_eq!(llm.name(), "openrouter"); + } +} diff --git a/crates/terraphim_grep/src/signatures.rs b/crates/terraphim_grep/src/signatures.rs index 9126b4d..d3df7b2 100644 --- a/crates/terraphim_grep/src/signatures.rs +++ b/crates/terraphim_grep/src/signatures.rs @@ -2,6 +2,65 @@ use serde::{Deserialize, Serialize}; use crate::error::TerraphimGrepError; +/// Extract a JSON object or array from a raw LLM response. +/// +/// Handles two common failure modes: +/// 1. Markdown code fences (` ```json ... ``` `). +/// 2. Prose surrounding the JSON payload. +fn extract_json(raw: &str) -> String { + let trimmed = raw.trim(); + + // 1. Strip markdown fences if present. + if let Some(start) = trimmed.find("```json") { + let block = &trimmed[start + 7..]; + if let Some(end) = block.find("```") { + return block[..end].trim().to_string(); + } + } + if let Some(start) = trimmed.find("```") { + let block = &trimmed[start + 3..]; + if let Some(end) = block.find("```") { + return block[..end].trim().to_string(); + } + } + + // 2. Find the first top-level { or [ and its matching close. + if let Some(start) = trimmed.find(['{', '[']) { + let bytes = trimmed.as_bytes(); + let open = bytes[start]; + let close = if open == b'{' { b'}' } else { b']' }; + let mut depth: usize = 0; + let mut in_string = false; + let mut escape = false; + for i in start..bytes.len() { + let c = bytes[i]; + if in_string { + if escape { + escape = false; + } else if c == b'\\' { + escape = true; + } else if c == b'"' { + in_string = false; + } + continue; + } + match c { + b'"' => in_string = true, + c if c == open => depth += 1, + c if c == close => { + depth = depth.saturating_sub(1); + if depth == 0 { + return trimmed[start..=i].to_string(); + } + } + _ => {} + } + } + } + + trimmed.to_string() +} + pub trait RlmSignature: Send + Sync { type Output: serde::Serialize + serde::de::DeserializeOwned; @@ -29,7 +88,8 @@ impl RlmSignature for SearchResultSignature { } fn parse(&self, raw: &str) -> Result { - serde_json::from_str(raw).map_err(|e| { + let json_str = extract_json(raw); + serde_json::from_str(&json_str).map_err(|e| { TerraphimGrepError::RlmFailed(format!("failed to parse search results: {}", e)) }) } @@ -64,7 +124,8 @@ impl RlmSignature for AnswerSignature { } fn parse(&self, raw: &str) -> Result { - serde_json::from_str(raw) + let json_str = extract_json(raw); + serde_json::from_str(&json_str) .map_err(|e| TerraphimGrepError::RlmFailed(format!("failed to parse answer: {}", e))) } } @@ -93,7 +154,8 @@ Return a JSON array of objects with: } fn parse(&self, raw: &str) -> Result { - serde_json::from_str(raw) + let json_str = extract_json(raw); + serde_json::from_str(&json_str) .map_err(|e| TerraphimGrepError::RlmFailed(format!("failed to parse concepts: {}", e))) } } @@ -102,6 +164,97 @@ Return a JSON array of objects with: mod tests { use super::*; + #[test] + fn test_extract_json_strips_markdown_fences() { + let raw = r#"Here is the answer: +```json +{ + "answer": "Retry is configured in src/main.rs", + "citations": [], + "confidence": 0.95 +} +``` +Hope that helps!"#; + let extracted = extract_json(raw); + assert!(extracted.starts_with('{')); + assert!(extracted.ends_with('}')); + assert!(!extracted.contains("```")); + } + + #[test] + fn test_extract_json_finds_bare_json_in_prose() { + let raw = r#"Sure, here is the result: +{ + "answer": "It works", + "citations": [{"source": "src/main.rs", "line": 1, "excerpt": "fn main()"}], + "confidence": 0.9 +} +Let me know if you need more."#; + let extracted = extract_json(raw); + assert!(extracted.starts_with('{')); + assert!(extracted.ends_with('}')); + } + + #[test] + fn test_extract_json_strips_generic_markdown_fence() { + let raw = r#"Result: +``` +{ + "answer": "It works", + "citations": [], + "confidence": 0.9 +} +``` +Done."#; + let extracted = extract_json(raw); + assert!(extracted.starts_with('{')); + assert!(extracted.ends_with('}')); + assert!(!extracted.contains("```")); + } + + #[test] + fn test_extract_json_finds_array_in_prose() { + let raw = r#"Here are the matches: +[ + {"path": "src/main.rs", "line": 1}, + {"path": "src/lib.rs", "line": 2} +] +End of results."#; + let extracted = extract_json(raw); + assert!(extracted.starts_with('[')); + assert!(extracted.ends_with(']')); + } + + #[test] + fn test_extract_json_handles_unterminated_fence() { + let raw = r#"```json +{ + "answer": "partial", + "citations": [], + "confidence": 0.5 +}"#; + let extracted = extract_json(raw); + assert!(extracted.starts_with('{')); + assert!(extracted.ends_with('}')); + } + + #[test] + fn test_extract_json_prefers_json_fence_over_other_language() { + let raw = r#"```rust +fn main() {} +``` +```json +{ + "answer": "correct block", + "citations": [], + "confidence": 0.9 +} +```"#; + let extracted = extract_json(raw); + assert!(extracted.starts_with('{')); + assert!(extracted.contains("correct block")); + } + #[test] fn test_search_result_signature_parse() { let signature = SearchResultSignature {}; @@ -140,6 +293,31 @@ mod tests { assert!((answer.confidence - 0.95).abs() < 0.001); } + #[test] + fn test_answer_signature_parse_markdown_wrapped() { + let signature = AnswerSignature {}; + let raw = r#"## Synthesised Answer + +The retry policy is configured in `src/main.rs`. + +```json +{ + "answer": "Retry is configured in src/main.rs", + "citations": [ + {"source": "src/main.rs", "line": 42, "excerpt": "pub retry_policy"} + ], + "confidence": 0.95 +} +```"#; + + let result = signature.parse(raw); + assert!(result.is_ok(), "parse failed: {:?}", result.as_ref().err()); + let answer = result.unwrap(); + assert!(answer.answer.contains("Retry")); + assert_eq!(answer.citations.len(), 1); + assert!((answer.confidence - 0.95).abs() < 0.001); + } + #[test] fn test_concept_extraction_signature_parse() { let signature = ConceptExtractionSignature {};