Migrate LLM integration from llm-chain to rig-core #5
Conversation
- Replace llm-chain dependency with rig-core 0.16.0 - Implement unified LLM client using rig's CompletionModel trait - Support OpenAI, Anthropic Claude, and Google Gemini providers - Add automatic temperature parameter handling for newer models - Reduce codebase complexity by ~840 lines through rig abstraction - Maintain existing API compatibility and error handling - All tests pass with new implementation 🦐 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
WalkthroughUnified multiple provider-specific LLM clients into a single Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Caller
participant Analyzer as RigLlmClient
participant Selector as create_provider
participant Provider as RigProvider
participant API as Provider API
Caller->>Analyzer: request_completion(model, prompt, opts)
Analyzer->>Selector: select provider (model, config/env)
Selector-->>Analyzer: RigProvider variant
Note over Analyzer: Build system prompt and messages
Analyzer->>Provider: send_completion_request(messages, opts)
Provider->>API: create completion (temperature, max_tokens)
API-->>Provider: response
Provider-->>Analyzer: extracted text
Analyzer-->>Caller: completion text
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
Cargo.toml(1 hunks)src/analyzer/llm_client.rs(8 hunks)src/analyzer/mod.rs(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
src/analyzer/mod.rs (1)
src/analyzer/llm_client.rs (1)
create_llm_client(656-672)
src/analyzer/llm_client.rs (3)
src/models/analysis.rs (2)
new(123-134)new(257-264)src/analyzer/orchestrator.rs (2)
new(18-31)request(373-373)src/analyzer/prompts.rs (1)
build_system_prompt(537-573)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (rust)
| serde = { version = "1.0", features = ["derive"] } | ||
| serde_json = "1.0" | ||
| thiserror = "1.0" | ||
| rig-core = "0.16.0" |
There was a problem hiding this comment.
Fix the crate alias for rig-core
Line 16 declares the dependency as rig-core, but the Rust crate identifier is rig_core. Because the code in src/analyzer/llm_client.rs imports modules with use rig::..., this currently fails to compile (E0432: unresolved import rig). Rename the dependency to rig (or update the imports to rig_core::) so the crate path resolves.
Apply this diff to align the alias with the use rig::... statements:
-rig-core = "0.16.0"
+rig = { package = "rig-core", version = "0.16.0" }🤖 Prompt for AI Agents
In Cargo.toml around line 16, the dependency is declared as rig-core = "0.16.0"
but the code imports use rig::...; change the crate alias so the package name
maps to the rig identifier — either replace the line with rig = { package =
"rig_core", version = "0.16.0" } or rename the dependency key to rig = "0.16.0"
(and keep package = "rig_core" if necessary) so use rig::... resolves;
alternatively, update all imports from rig::... to rig_core::... if you prefer
keeping the existing dependency name.
- Improve summary extraction logic to avoid 'Summary not provided' messages - Add better error messages for empty responses with debug information - Fix temperature parameter handling for newer model series (o3, o4) - Set provider-specific max_tokens defaults (4096 for Claude, 1000 for others) - Change default model from gpt-5-mini to gpt-4o for better compatibility - Add more robust content extraction patterns for different response formats 🦐 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
src/analyzer/llm_client.rs (5)
517-527: Provider-specific max_tokens fallback is currently unreachableBecause
create_llm_clientsetsmax_tokens=Some(1000), this branch never runs and Anthropic won’t get the intended higher default. Recommend letting the provider-specific defaults apply by default.Apply this diff (also see config change below):
- // Set max_tokens with provider-specific defaults - if let Some(max_tokens) = self.config.max_tokens { - builder = builder.max_tokens(max_tokens as u64); - } else { - // Use provider-specific defaults - let default_tokens = match &self.provider { - RigProvider::Anthropic(_) => 4096, // Claude models typically need explicit max_tokens - _ => 1000, - }; - builder = builder.max_tokens(default_tokens); - } + // Set max_tokens with provider-specific defaults + if let Some(max_tokens) = self.config.max_tokens { + builder = builder.max_tokens(max_tokens as u64); + } else { + // Use provider-specific defaults + let default_tokens = match &self.provider { + RigProvider::Anthropic(_) => 4096, // Claude models typically need explicit max_tokens + _ => 1000, + }; + builder = builder.max_tokens(default_tokens); + }(Combine with the config change at lines 701-708 to expose this path.)
701-708: Let provider defaults drive max_tokensSetting
Some(1000)at construction time forces that cap for all providers. Set toNoneso Anthropic can use 4096 (and others 1000) per your send path.Apply this diff:
let config = LlmConfig { model_name: model.to_string(), api_key, timeout_seconds, max_retries: 3, - max_tokens: Some(1000), + max_tokens: None, // Defer to provider-specific defaults in send_completion_request temperature: Some(0.3), };
499-551: Add retry with backoff for transient provider errors
max_retriesis unused; add a simple retry/backoff around.send()to improve resiliency on rate limits/transient network issues. Rig 0.16.0 emphasizes improved stability; this complements that.Apply this diff:
async fn send_completion_request<M: CompletionModel>( &self, model: M, prompt: &str, system_prompt: String, ) -> Result<String, EbiError> { - let mut builder = model - .completion_request(prompt) - .preamble(system_prompt); - - // Skip temperature for models that don't support it (like GPT-5 series and o1 series) - if let Some(temp) = self.config.temperature { - let model_name = &self.config.model_name; - if !model_name.starts_with("gpt-5") && !model_name.starts_with("o1") && !model_name.starts_with("o3") && !model_name.starts_with("o4") { - builder = builder.temperature(temp as f64); - } - } - - // Set max_tokens with provider-specific defaults - if let Some(max_tokens) = self.config.max_tokens { - builder = builder.max_tokens(max_tokens as u64); - } else { - // Use provider-specific defaults - let default_tokens = match &self.provider { - RigProvider::Anthropic(_) => 4096, // Claude models typically need explicit max_tokens - _ => 1000, - }; - builder = builder.max_tokens(default_tokens); - } - - let response = builder - .send() - .await - .map_err(|e| EbiError::LlmClientError(format!("Request failed for model {}: {}", self.config.model_name, e)))?; + let mut attempt = 0u32; + loop { + let mut builder = model + .completion_request(prompt) + .preamble(system_prompt.clone()); + + // Skip temperature for models that don't support it (like GPT-5 series and o1 series) + if let Some(temp) = self.config.temperature { + let model_name = &self.config.model_name; + if !model_name.starts_with("gpt-5") + && !model_name.starts_with("o1") + && !model_name.starts_with("o3") + && !model_name.starts_with("o4") + { + builder = builder.temperature(temp as f64); + } + } + + // Set max_tokens with provider-specific defaults + if let Some(max_tokens) = self.config.max_tokens { + builder = builder.max_tokens(max_tokens as u64); + } else { + let default_tokens = match &self.provider { + RigProvider::Anthropic(_) => 4096, + _ => 1000, + }; + builder = builder.max_tokens(default_tokens); + } + + match builder.send().await { + Ok(response) => { + // Extract the text content from the response + let mut extracted_text = String::new(); + for content in response.choice.iter() { + if let AssistantContent::Text(text_content) = content { + extracted_text.push_str(&text_content.text); + } + } + + // Handle empty responses + if extracted_text.trim().is_empty() { + return Err(EbiError::LlmClientError(format!( + "Model {} returned empty response. Response had {} choice(s). This may indicate an API issue, authentication problem, or model configuration issue.", + self.config.model_name, response.choice.len() + ))); + } + + return Ok(extracted_text); + } + Err(e) => { + if attempt >= self.config.max_retries { + return Err(EbiError::LlmClientError(format!( + "Request failed for model {} after {} retries: {}", + self.config.model_name, attempt, e + ))); + } + attempt += 1; + // exponential backoff: 200ms, 400ms, 800ms... + let backoff_ms = 200u64.saturating_mul(1u64 << (attempt - 1)); + tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; + } + } + } - - // Extract the text content from the response - let mut extracted_text = String::new(); - for content in response.choice.iter() { - if let AssistantContent::Text(text_content) = content { - extracted_text.push_str(&text_content.text); - } - } - - // Handle empty responses - if extracted_text.trim().is_empty() { - return Err(EbiError::LlmClientError( - format!("Model {} returned empty response. Response had {} choice(s). This may indicate an API issue, authentication problem, or model configuration issue.", - self.config.model_name, response.choice.len()) - )); - } - - Ok(extracted_text) }Based on learnings
725-733: Support fine‑tuned model prefixes for Claude and GeminiOpenAI detection strips
ft:; do the same for Claude/Gemini to accept common fine-tune naming.Apply this diff:
fn is_claude_model(model: &str) -> bool { - let candidate = model.strip_prefix("anthropic/").unwrap_or(model); + let candidate = model.strip_prefix("anthropic/").unwrap_or(model); + let candidate = candidate.strip_prefix("ft:").unwrap_or(candidate); candidate.starts_with("claude-") } fn is_gemini_model(model: &str) -> bool { - let candidate = model.strip_prefix("gemini/").unwrap_or(model); + let candidate = model.strip_prefix("gemini/").unwrap_or(model); + let candidate = candidate.strip_prefix("ft:").unwrap_or(candidate); candidate.starts_with("gemini-") }
466-471: Confirm intended visibility of RigProviderIf external callers or tests need to select/inspect provider variants, consider making this enum public or re‑exporting it; otherwise keeping it private is fine.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
src/analyzer/llm_client.rs(11 hunks)src/analyzer/orchestrator.rs(1 hunks)src/cli/args.rs(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
src/analyzer/orchestrator.rs (1)
src/analyzer/llm_client.rs (1)
new(474-477)
src/analyzer/llm_client.rs (2)
src/analyzer/orchestrator.rs (2)
new(18-31)request(373-373)src/analyzer/prompts.rs (1)
build_system_prompt(537-573)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (rust)
🔇 Additional comments (4)
src/cli/args.rs (2)
17-18: Default model switched to gpt-4o: LGTMChange aligns with orchestrator defaults and unified Rig client.
204-205: Test expectation updated: LGTMAssertion matches the new default model.
src/analyzer/orchestrator.rs (1)
391-396: Orchestrator defaults updated to gpt-4o: LGTMConsistent with CLI and provider selection logic.
src/analyzer/llm_client.rs (1)
509-515: Temperature auto‑handling: LGTMSkips temp for GPT‑5 and o* families, as intended by the migration notes.
Based on learnings
Summary
llm-chaindependency withrig-core 0.16.0for better LLM provider integrationTechnical Changes
llm-chaintorig-core 0.16.0LlmProvidertrait interfaceBenefits
Test Results
Test plan
🦐 Generated with Claude Code
Summary by CodeRabbit