mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
fix: harden openai-compatible provider, approval replay, and embeddings defaults (#237)
* fix: harden openai-compatible tool flow and local defaults * fix: close approval replay gaps and harden openai-compatible flow * fix: address review feedback and code improvements (takeover #112) - Make ChatCompletionResponse.id Optional<String> to handle providers that omit or null the field - Propagate HTTP client builder errors instead of silently dropping timeout configuration (openai_compatible_chat, nearai_chat) - Add EMBEDDING_DIMENSION env var with smart per-model defaults instead of hardcoding 768/1536 everywhere - Remove duplicated dimension inference logic from main.rs Co-Authored-By: panosAthDBX <[email protected]> Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: harden src/llm/ module from crate audit findings - Replace 9x .expect() on RwLock with graceful poison recovery (nearai.rs: 7, nearai_chat.rs: 2) — eliminates production panics - Propagate HTTP client builder errors in nearai.rs instead of silently dropping timeout config (NearAiProvider::new now returns Result) - Make nearai_chat ChatCompletionResponse.id Optional<String> (mirrors openai_compatible_chat.rs fix for providers that omit id) - Make nearai_chat usage fields optional with defensive parse_usage() helper (was required u32 fields that crash on null/missing) - Truncate error responses to 512 chars in nearai_chat.rs error messages to prevent log bloat and potential data leakage - Delegate 4 missing LlmProvider methods in FailoverProvider (model_metadata, seed_response_chain, get_response_chain_id, calculate_cost) to last-used provider instead of trait defaults Co-Authored-By: Claude Opus 4.6 <[email protected]> * refactor(llm): add RetryProvider, remove openai_compatible_chat, harden decorators - Add composable RetryProvider decorator wrapping any LlmProvider with exponential backoff + jitter, respecting RateLimited retry_after hints - Remove openai_compatible_chat.rs — replaced by rig adapter + RetryProvider - Remove internal retry loop from nearai.rs (was causing double-retry with external RetryProvider, up to 16 attempts instead of 4) - Remove internal retry loop from nearai_chat.rs (same issue) - Wire RetryProvider into main.rs composition chain: each provider gets its own retry wrapper before failover - Move normalize_tool_name to rig_adapter.rs for all rig-based providers - Reconcile is_retryable() vs is_transient() error classification: ModelNotAvailable no longer retryable, Json no longer transient - Fix unchecked Duration subtraction panic in circuit_breaker.rs - Make failover.rs use shared is_retryable() from retry.rs - Remove stale #[allow(dead_code)] on NearAiResponse::id (field is used) Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR review feedback — error handling, dimension validation, libSQL warning - Replace response.text().await.unwrap_or_default() with proper error propagation in nearai.rs and nearai_chat.rs (4 call sites). Failures now return LlmError::RequestFailed with context instead of silently proceeding with an empty string. - Add embedding dimension validation in OllamaEmbeddings::embed_batch(): returns EmbeddingError if Ollama returns embeddings with a dimension that doesn't match the configured value. - Add runtime warning when libSQL backend is used with non-1536 embedding dimension, since the libSQL schema uses F32_BLOB(1536) and cannot store different-dimension vectors. Co-Authored-By: Claude Opus 4.6 <[email protected]> * Apply suggestions from code review Co-authored-by: Copilot <[email protected]> --------- Co-authored-by: panosAthDbx <[email protected]> Co-authored-by: panosAthDBX <[email protected]> Co-authored-by: panosAthDBX <[email protected]> Co-authored-by: Claude Opus 4.6 <[email protected]> Co-authored-by: Copilot <[email protected]>
This commit is contained in:
co-authored by
panosAthDbx
panosAthDBX
panosAthDBX
Claude Opus 4.6
Copilot
parent
e87d7bd066
commit
097a26ace6
@@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
### Added
|
||||||
|
|
||||||
|
- Refactored OpenAI-compatible chat completion routing to use the rig adapter and `RetryProvider` composition for custom base URL usage.
|
||||||
|
- Added Ollama embeddings provider support (`EMBEDDING_PROVIDER=ollama`, `OLLAMA_BASE_URL`) in workspace embeddings.
|
||||||
|
- Added migration `V9__flexible_embedding_dimension.sql` for flexible embedding vector dimensions.
|
||||||
|
|
||||||
|
### Changed
|
||||||
|
|
||||||
|
- Changed default sandbox image to `ironclaw-worker:latest` in config/settings/sandbox defaults.
|
||||||
|
- Improved tool-message sanitization and provider compatibility handling across NEAR AI, rig adapter, and shared LLM provider code.
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- Fixed approval-input aliases (`a`, `/approve`, `/always`, `/deny`, etc.) in submission parsing.
|
||||||
|
- Fixed multi-tool approval resume flow by preserving and replaying deferred tool calls so all prior `tool_use` IDs receive matching `tool_result` messages.
|
||||||
|
- Fixed REPL quit/exit handling to route shutdown through the agent loop for graceful termination.
|
||||||
|
|
||||||
## [0.6.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.5.0...ironclaw-v0.6.0) - 2026-02-19
|
## [0.6.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.5.0...ironclaw-v0.6.0) - 2026-02-19
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
@@ -94,6 +111,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
- Bump MSRV to 1.92, add GCP deployment files ([#40](https://github.com/nearai/ironclaw/pull/40))
|
- Bump MSRV to 1.92, add GCP deployment files ([#40](https://github.com/nearai/ironclaw/pull/40))
|
||||||
- Add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models) ([#31](https://github.com/nearai/ironclaw/pull/31))
|
- Add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models) ([#31](https://github.com/nearai/ironclaw/pull/31))
|
||||||
|
|
||||||
|
|
||||||
## [0.1.3](https://github.com/nearai/ironclaw/compare/v0.1.2...v0.1.3) - 2026-02-12
|
## [0.1.3](https://github.com/nearai/ironclaw/compare/v0.1.2...v0.1.3) - 2026-02-12
|
||||||
|
|
||||||
### Other
|
### Other
|
||||||
|
|||||||
@@ -0,0 +1,43 @@
|
|||||||
|
-- Allow embedding vectors of any dimension (not just 1536).
|
||||||
|
-- This supports Ollama models (768-dim nomic-embed-text, 1024-dim mxbai-embed-large)
|
||||||
|
-- alongside OpenAI models (1536-dim text-embedding-3-small, 3072-dim text-embedding-3-large).
|
||||||
|
--
|
||||||
|
-- NOTE: HNSW indexes require a fixed dimension, so we drop the index.
|
||||||
|
-- Exact (sequential) cosine distance search still works without the index.
|
||||||
|
-- For a personal assistant workspace the dataset is small enough that this
|
||||||
|
-- has negligible impact on query latency.
|
||||||
|
|
||||||
|
-- Drop dependent views first
|
||||||
|
DROP VIEW IF EXISTS chunks_pending_embedding;
|
||||||
|
DROP VIEW IF EXISTS memory_documents_summary;
|
||||||
|
|
||||||
|
DROP INDEX IF EXISTS idx_memory_chunks_embedding;
|
||||||
|
|
||||||
|
ALTER TABLE memory_chunks
|
||||||
|
ALTER COLUMN embedding TYPE vector
|
||||||
|
USING embedding::vector;
|
||||||
|
|
||||||
|
-- Recreate the views
|
||||||
|
CREATE VIEW memory_documents_summary AS
|
||||||
|
SELECT
|
||||||
|
d.id,
|
||||||
|
d.user_id,
|
||||||
|
d.path,
|
||||||
|
d.created_at,
|
||||||
|
d.updated_at,
|
||||||
|
COUNT(c.id) as chunk_count,
|
||||||
|
COUNT(c.embedding) as embedded_chunk_count
|
||||||
|
FROM memory_documents d
|
||||||
|
LEFT JOIN memory_chunks c ON c.document_id = d.id
|
||||||
|
GROUP BY d.id;
|
||||||
|
|
||||||
|
CREATE VIEW chunks_pending_embedding AS
|
||||||
|
SELECT
|
||||||
|
c.id as chunk_id,
|
||||||
|
c.document_id,
|
||||||
|
d.user_id,
|
||||||
|
d.path,
|
||||||
|
LENGTH(c.content) as content_length
|
||||||
|
FROM memory_chunks c
|
||||||
|
JOIN memory_documents d ON d.id = c.document_id
|
||||||
|
WHERE c.embedding IS NULL;
|
||||||
+10
-2
@@ -255,7 +255,10 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Execute each tool (with approval checking and hook interception)
|
// Execute each tool (with approval checking and hook interception)
|
||||||
for mut tc in tool_calls {
|
let mut idx = 0usize;
|
||||||
|
while idx < tool_calls.len() {
|
||||||
|
let mut tc = tool_calls[idx].clone();
|
||||||
|
|
||||||
// Check if tool requires approval
|
// Check if tool requires approval
|
||||||
if let Some(tool) = self.tools().get(&tc.name).await
|
if let Some(tool) = self.tools().get(&tc.name).await
|
||||||
&& tool.requires_approval()
|
&& tool.requires_approval()
|
||||||
@@ -277,7 +280,9 @@ impl Agent {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if !is_auto_approved {
|
if !is_auto_approved {
|
||||||
// Need approval - store pending request and return
|
// Need approval - store pending request and return.
|
||||||
|
// Preserve remaining tool calls so they can be replayed
|
||||||
|
// after approval.
|
||||||
let pending = PendingApproval {
|
let pending = PendingApproval {
|
||||||
request_id: Uuid::new_v4(),
|
request_id: Uuid::new_v4(),
|
||||||
tool_name: tc.name.clone(),
|
tool_name: tc.name.clone(),
|
||||||
@@ -285,6 +290,7 @@ impl Agent {
|
|||||||
description: tool.description().to_string(),
|
description: tool.description().to_string(),
|
||||||
tool_call_id: tc.id.clone(),
|
tool_call_id: tc.id.clone(),
|
||||||
context_messages: context_messages.clone(),
|
context_messages: context_messages.clone(),
|
||||||
|
deferred_tool_calls: tool_calls[idx + 1..].to_vec(),
|
||||||
};
|
};
|
||||||
|
|
||||||
return Ok(AgenticLoopResult::NeedApproval { pending });
|
return Ok(AgenticLoopResult::NeedApproval { pending });
|
||||||
@@ -441,6 +447,8 @@ impl Agent {
|
|||||||
&tc.name,
|
&tc.name,
|
||||||
result_content,
|
result_content,
|
||||||
));
|
));
|
||||||
|
|
||||||
|
idx += 1;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ use chrono::{DateTime, Utc};
|
|||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::llm::ChatMessage;
|
use crate::llm::{ChatMessage, ToolCall};
|
||||||
|
|
||||||
/// A session containing one or more threads.
|
/// A session containing one or more threads.
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
@@ -148,6 +148,10 @@ pub struct PendingApproval {
|
|||||||
pub tool_call_id: String,
|
pub tool_call_id: String,
|
||||||
/// Context messages at the time of the request (to resume from).
|
/// Context messages at the time of the request (to resume from).
|
||||||
pub context_messages: Vec<ChatMessage>,
|
pub context_messages: Vec<ChatMessage>,
|
||||||
|
/// Remaining tool calls from the same assistant message that were not
|
||||||
|
/// executed yet when approval was requested.
|
||||||
|
#[serde(default)]
|
||||||
|
pub deferred_tool_calls: Vec<ToolCall>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A conversation thread within a session.
|
/// A conversation thread within a session.
|
||||||
@@ -946,6 +950,7 @@ mod tests {
|
|||||||
description: "dangerous command".to_string(),
|
description: "dangerous command".to_string(),
|
||||||
tool_call_id: "call_123".to_string(),
|
tool_call_id: "call_123".to_string(),
|
||||||
context_messages: vec![ChatMessage::user("do it")],
|
context_messages: vec![ChatMessage::user("do it")],
|
||||||
|
deferred_tool_calls: vec![],
|
||||||
};
|
};
|
||||||
|
|
||||||
thread.await_approval(approval);
|
thread.await_approval(approval);
|
||||||
@@ -969,6 +974,7 @@ mod tests {
|
|||||||
description: "test".to_string(),
|
description: "test".to_string(),
|
||||||
tool_call_id: "call_456".to_string(),
|
tool_call_id: "call_456".to_string(),
|
||||||
context_messages: vec![],
|
context_messages: vec![],
|
||||||
|
deferred_tool_calls: vec![],
|
||||||
};
|
};
|
||||||
|
|
||||||
thread.await_approval(approval);
|
thread.await_approval(approval);
|
||||||
|
|||||||
+54
-3
@@ -118,19 +118,19 @@ impl SubmissionParser {
|
|||||||
// Approval responses (simple yes/no/always for pending approvals)
|
// Approval responses (simple yes/no/always for pending approvals)
|
||||||
// These are short enough to check explicitly
|
// These are short enough to check explicitly
|
||||||
match lower.as_str() {
|
match lower.as_str() {
|
||||||
"yes" | "y" | "approve" | "ok" => {
|
"yes" | "y" | "approve" | "ok" | "/approve" | "/yes" | "/y" => {
|
||||||
return Submission::ApprovalResponse {
|
return Submission::ApprovalResponse {
|
||||||
approved: true,
|
approved: true,
|
||||||
always: false,
|
always: false,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
"always" | "yes always" | "approve always" => {
|
"always" | "a" | "yes always" | "approve always" | "/always" | "/a" => {
|
||||||
return Submission::ApprovalResponse {
|
return Submission::ApprovalResponse {
|
||||||
approved: true,
|
approved: true,
|
||||||
always: true,
|
always: true,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
"no" | "n" | "deny" | "reject" | "cancel" => {
|
"no" | "n" | "deny" | "reject" | "cancel" | "/deny" | "/no" | "/n" => {
|
||||||
return Submission::ApprovalResponse {
|
return Submission::ApprovalResponse {
|
||||||
approved: false,
|
approved: false,
|
||||||
always: false,
|
always: false,
|
||||||
@@ -475,6 +475,57 @@ mod tests {
|
|||||||
assert!(matches!(submission, Submission::UserInput { content } if content == "/unknown"));
|
assert!(matches!(submission, Submission::UserInput { content } if content == "/unknown"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parser_approval_response_aliases() {
|
||||||
|
// approve once
|
||||||
|
assert!(matches!(
|
||||||
|
SubmissionParser::parse("y"),
|
||||||
|
Submission::ApprovalResponse {
|
||||||
|
approved: true,
|
||||||
|
always: false
|
||||||
|
}
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
SubmissionParser::parse("/approve"),
|
||||||
|
Submission::ApprovalResponse {
|
||||||
|
approved: true,
|
||||||
|
always: false
|
||||||
|
}
|
||||||
|
));
|
||||||
|
|
||||||
|
// approve always
|
||||||
|
assert!(matches!(
|
||||||
|
SubmissionParser::parse("a"),
|
||||||
|
Submission::ApprovalResponse {
|
||||||
|
approved: true,
|
||||||
|
always: true
|
||||||
|
}
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
SubmissionParser::parse("/always"),
|
||||||
|
Submission::ApprovalResponse {
|
||||||
|
approved: true,
|
||||||
|
always: true
|
||||||
|
}
|
||||||
|
));
|
||||||
|
|
||||||
|
// deny
|
||||||
|
assert!(matches!(
|
||||||
|
SubmissionParser::parse("n"),
|
||||||
|
Submission::ApprovalResponse {
|
||||||
|
approved: false,
|
||||||
|
always: false
|
||||||
|
}
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
SubmissionParser::parse("/deny"),
|
||||||
|
Submission::ApprovalResponse {
|
||||||
|
approved: false,
|
||||||
|
always: false
|
||||||
|
}
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_parser_json_exec_approval() {
|
fn test_parser_json_exec_approval() {
|
||||||
let req_id = Uuid::new_v4();
|
let req_id = Uuid::new_v4();
|
||||||
|
|||||||
+174
-1
@@ -11,7 +11,7 @@ use uuid::Uuid;
|
|||||||
use crate::agent::Agent;
|
use crate::agent::Agent;
|
||||||
use crate::agent::compaction::ContextCompactor;
|
use crate::agent::compaction::ContextCompactor;
|
||||||
use crate::agent::dispatcher::{AgenticLoopResult, detect_auth_awaiting, parse_auth_result};
|
use crate::agent::dispatcher::{AgenticLoopResult, detect_auth_awaiting, parse_auth_result};
|
||||||
use crate::agent::session::{Session, ThreadState};
|
use crate::agent::session::{PendingApproval, Session, ThreadState};
|
||||||
use crate::agent::submission::SubmissionResult;
|
use crate::agent::submission::SubmissionResult;
|
||||||
use crate::channels::{IncomingMessage, StatusUpdate};
|
use crate::channels::{IncomingMessage, StatusUpdate};
|
||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
@@ -712,6 +712,7 @@ impl Agent {
|
|||||||
|
|
||||||
// Build context including the tool result
|
// Build context including the tool result
|
||||||
let mut context_messages = pending.context_messages;
|
let mut context_messages = pending.context_messages;
|
||||||
|
let deferred_tool_calls = pending.deferred_tool_calls;
|
||||||
|
|
||||||
// Record result in thread
|
// Record result in thread
|
||||||
{
|
{
|
||||||
@@ -780,6 +781,178 @@ impl Agent {
|
|||||||
result_content,
|
result_content,
|
||||||
));
|
));
|
||||||
|
|
||||||
|
// Replay deferred tool calls from the same assistant message so
|
||||||
|
// every tool_use ID gets a matching tool_result before the next
|
||||||
|
// LLM call.
|
||||||
|
if !deferred_tool_calls.is_empty() {
|
||||||
|
let _ = self
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&message.channel,
|
||||||
|
StatusUpdate::Thinking(format!(
|
||||||
|
"Executing {} deferred tool(s)...",
|
||||||
|
deferred_tool_calls.len()
|
||||||
|
)),
|
||||||
|
&message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut deferred_queue = std::collections::VecDeque::from(deferred_tool_calls);
|
||||||
|
while let Some(tc) = deferred_queue.pop_front() {
|
||||||
|
// Re-check approval for each deferred tool call
|
||||||
|
if let Some(tool) = self.tools().get(&tc.name).await
|
||||||
|
&& tool.requires_approval()
|
||||||
|
{
|
||||||
|
let is_auto_approved = {
|
||||||
|
let sess = session.lock().await;
|
||||||
|
let mut approved = sess.is_tool_auto_approved(&tc.name);
|
||||||
|
if approved && tool.requires_approval_for(&tc.arguments) {
|
||||||
|
approved = false;
|
||||||
|
}
|
||||||
|
approved
|
||||||
|
};
|
||||||
|
|
||||||
|
if !is_auto_approved {
|
||||||
|
let new_pending = PendingApproval {
|
||||||
|
request_id: Uuid::new_v4(),
|
||||||
|
tool_name: tc.name.clone(),
|
||||||
|
parameters: tc.arguments.clone(),
|
||||||
|
description: tool.description().to_string(),
|
||||||
|
tool_call_id: tc.id.clone(),
|
||||||
|
context_messages: context_messages.clone(),
|
||||||
|
deferred_tool_calls: deferred_queue.iter().cloned().collect(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let request_id = new_pending.request_id;
|
||||||
|
let tool_name = new_pending.tool_name.clone();
|
||||||
|
let description = new_pending.description.clone();
|
||||||
|
let parameters = new_pending.parameters.clone();
|
||||||
|
|
||||||
|
{
|
||||||
|
let mut sess = session.lock().await;
|
||||||
|
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||||
|
thread.await_approval(new_pending);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = self
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&message.channel,
|
||||||
|
StatusUpdate::Status("Awaiting approval".into()),
|
||||||
|
&message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
return Ok(SubmissionResult::NeedApproval {
|
||||||
|
request_id,
|
||||||
|
tool_name,
|
||||||
|
description,
|
||||||
|
parameters,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = self
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&message.channel,
|
||||||
|
StatusUpdate::ToolStarted {
|
||||||
|
name: tc.name.clone(),
|
||||||
|
},
|
||||||
|
&message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let deferred_result = self
|
||||||
|
.execute_chat_tool(&tc.name, &tc.arguments, &job_ctx)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let _ = self
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&message.channel,
|
||||||
|
StatusUpdate::ToolCompleted {
|
||||||
|
name: tc.name.clone(),
|
||||||
|
success: deferred_result.is_ok(),
|
||||||
|
},
|
||||||
|
&message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
if let Ok(ref output) = deferred_result
|
||||||
|
&& !output.is_empty()
|
||||||
|
{
|
||||||
|
let _ = self
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&message.channel,
|
||||||
|
StatusUpdate::ToolResult {
|
||||||
|
name: tc.name.clone(),
|
||||||
|
preview: output.clone(),
|
||||||
|
},
|
||||||
|
&message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Record in thread
|
||||||
|
{
|
||||||
|
let mut sess = session.lock().await;
|
||||||
|
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||||
|
&& let Some(turn) = thread.last_turn_mut()
|
||||||
|
{
|
||||||
|
match &deferred_result {
|
||||||
|
Ok(output) => turn.record_tool_result(serde_json::json!(output)),
|
||||||
|
Err(e) => turn.record_tool_error(e.to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Auth detection for deferred tools
|
||||||
|
if let Some((ext_name, instructions)) =
|
||||||
|
detect_auth_awaiting(&tc.name, &deferred_result)
|
||||||
|
{
|
||||||
|
let auth_data = parse_auth_result(&deferred_result);
|
||||||
|
{
|
||||||
|
let mut sess = session.lock().await;
|
||||||
|
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||||
|
thread.enter_auth_mode(ext_name.clone());
|
||||||
|
thread.complete_turn(&instructions);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let _ = self
|
||||||
|
.channels
|
||||||
|
.send_status(
|
||||||
|
&message.channel,
|
||||||
|
StatusUpdate::AuthRequired {
|
||||||
|
extension_name: ext_name,
|
||||||
|
instructions: Some(instructions.clone()),
|
||||||
|
auth_url: auth_data.auth_url,
|
||||||
|
setup_url: auth_data.setup_url,
|
||||||
|
},
|
||||||
|
&message.metadata,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return Ok(SubmissionResult::response(instructions));
|
||||||
|
}
|
||||||
|
|
||||||
|
let deferred_content = match deferred_result {
|
||||||
|
Ok(output) => {
|
||||||
|
let sanitized = self.safety().sanitize_tool_output(&tc.name, &output);
|
||||||
|
self.safety().wrap_for_llm(
|
||||||
|
&tc.name,
|
||||||
|
&sanitized.content,
|
||||||
|
sanitized.was_modified,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
Err(e) => format!("Error: {}", e),
|
||||||
|
};
|
||||||
|
|
||||||
|
context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content));
|
||||||
|
}
|
||||||
|
|
||||||
// Continue the agentic loop (a tool was already executed this turn)
|
// Continue the agentic loop (a tool was already executed this turn)
|
||||||
let result = self
|
let result = self
|
||||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages, true)
|
.run_agentic_loop(message, session.clone(), thread_id, context_messages, true)
|
||||||
|
|||||||
@@ -330,7 +330,13 @@ impl Channel for ReplChannel {
|
|||||||
// Handle local REPL commands (only commands that need
|
// Handle local REPL commands (only commands that need
|
||||||
// immediate local handling stay here)
|
// immediate local handling stay here)
|
||||||
match line.to_lowercase().as_str() {
|
match line.to_lowercase().as_str() {
|
||||||
"/quit" | "/exit" => break,
|
"/quit" | "/exit" => {
|
||||||
|
// Forward shutdown command so the agent loop exits even
|
||||||
|
// when other channels (e.g. web gateway) are still active.
|
||||||
|
let msg = IncomingMessage::new("repl", "default", "/quit");
|
||||||
|
let _ = tx.blocking_send(msg);
|
||||||
|
break;
|
||||||
|
}
|
||||||
"/help" => {
|
"/help" => {
|
||||||
print_help();
|
print_help();
|
||||||
continue;
|
continue;
|
||||||
|
|||||||
@@ -305,6 +305,7 @@ impl Channel for GatewayChannel {
|
|||||||
description,
|
description,
|
||||||
parameters: serde_json::to_string_pretty(¶meters)
|
parameters: serde_json::to_string_pretty(¶meters)
|
||||||
.unwrap_or_else(|_| parameters.to_string()),
|
.unwrap_or_else(|_| parameters.to_string()),
|
||||||
|
thread_id,
|
||||||
},
|
},
|
||||||
StatusUpdate::AuthRequired {
|
StatusUpdate::AuthRequired {
|
||||||
extension_name,
|
extension_name,
|
||||||
|
|||||||
@@ -159,6 +159,7 @@ function connectSSE() {
|
|||||||
|
|
||||||
eventSource.addEventListener('approval_needed', (e) => {
|
eventSource.addEventListener('approval_needed', (e) => {
|
||||||
const data = JSON.parse(e.data);
|
const data = JSON.parse(e.data);
|
||||||
|
if (!isCurrentThread(data.thread_id)) return;
|
||||||
showApproval(data);
|
showApproval(data);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
@@ -137,6 +137,8 @@ pub enum SseEvent {
|
|||||||
tool_name: String,
|
tool_name: String,
|
||||||
description: String,
|
description: String,
|
||||||
parameters: String,
|
parameters: String,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
thread_id: Option<String>,
|
||||||
},
|
},
|
||||||
#[serde(rename = "auth_required")]
|
#[serde(rename = "auth_required")]
|
||||||
AuthRequired {
|
AuthRequired {
|
||||||
@@ -785,12 +787,14 @@ mod tests {
|
|||||||
tool_name: "shell".to_string(),
|
tool_name: "shell".to_string(),
|
||||||
description: "Run ls".to_string(),
|
description: "Run ls".to_string(),
|
||||||
parameters: "{}".to_string(),
|
parameters: "{}".to_string(),
|
||||||
|
thread_id: Some("t1".to_string()),
|
||||||
};
|
};
|
||||||
let ws = WsServerMessage::from_sse_event(&sse);
|
let ws = WsServerMessage::from_sse_event(&sse);
|
||||||
match ws {
|
match ws {
|
||||||
WsServerMessage::Event { event_type, data } => {
|
WsServerMessage::Event { event_type, data } => {
|
||||||
assert_eq!(event_type, "approval_needed");
|
assert_eq!(event_type, "approval_needed");
|
||||||
assert_eq!(data["tool_name"], "shell");
|
assert_eq!(data["tool_name"], "shell");
|
||||||
|
assert_eq!(data["thread_id"], "t1");
|
||||||
}
|
}
|
||||||
_ => panic!("Expected Event variant"),
|
_ => panic!("Expected Event variant"),
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,25 +9,48 @@ use crate::settings::Settings;
|
|||||||
pub struct EmbeddingsConfig {
|
pub struct EmbeddingsConfig {
|
||||||
/// Whether embeddings are enabled.
|
/// Whether embeddings are enabled.
|
||||||
pub enabled: bool,
|
pub enabled: bool,
|
||||||
/// Provider to use: "openai" or "nearai"
|
/// Provider to use: "openai", "nearai", or "ollama"
|
||||||
pub provider: String,
|
pub provider: String,
|
||||||
/// OpenAI API key (for OpenAI provider).
|
/// OpenAI API key (for OpenAI provider).
|
||||||
pub openai_api_key: Option<SecretString>,
|
pub openai_api_key: Option<SecretString>,
|
||||||
/// Model to use for embeddings.
|
/// Model to use for embeddings.
|
||||||
pub model: String,
|
pub model: String,
|
||||||
|
/// Ollama base URL (for Ollama provider). Defaults to http://localhost:11434.
|
||||||
|
pub ollama_base_url: String,
|
||||||
|
/// Embedding vector dimension. Inferred from the model name when not set explicitly.
|
||||||
|
pub dimension: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for EmbeddingsConfig {
|
impl Default for EmbeddingsConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
|
let model = "text-embedding-3-small".to_string();
|
||||||
|
let dimension = default_dimension_for_model(&model);
|
||||||
Self {
|
Self {
|
||||||
enabled: false,
|
enabled: false,
|
||||||
provider: "openai".to_string(),
|
provider: "openai".to_string(),
|
||||||
openai_api_key: None,
|
openai_api_key: None,
|
||||||
model: "text-embedding-3-small".to_string(),
|
model,
|
||||||
|
ollama_base_url: "http://localhost:11434".to_string(),
|
||||||
|
dimension,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Infer the embedding dimension from a well-known model name.
|
||||||
|
///
|
||||||
|
/// Falls back to 1536 (OpenAI text-embedding-3-small default) for unknown models.
|
||||||
|
fn default_dimension_for_model(model: &str) -> usize {
|
||||||
|
match model {
|
||||||
|
"text-embedding-3-small" => 1536,
|
||||||
|
"text-embedding-3-large" => 3072,
|
||||||
|
"text-embedding-ada-002" => 1536,
|
||||||
|
"nomic-embed-text" => 768,
|
||||||
|
"mxbai-embed-large" => 1024,
|
||||||
|
"all-minilm" => 384,
|
||||||
|
_ => 1536,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl EmbeddingsConfig {
|
impl EmbeddingsConfig {
|
||||||
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||||
let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from);
|
let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from);
|
||||||
@@ -38,6 +61,19 @@ impl EmbeddingsConfig {
|
|||||||
let model =
|
let model =
|
||||||
optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone());
|
optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone());
|
||||||
|
|
||||||
|
let ollama_base_url = optional_env("OLLAMA_BASE_URL")?
|
||||||
|
.or_else(|| settings.ollama_base_url.clone())
|
||||||
|
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||||
|
|
||||||
|
let dimension = optional_env("EMBEDDING_DIMENSION")?
|
||||||
|
.map(|s| s.parse::<usize>())
|
||||||
|
.transpose()
|
||||||
|
.map_err(|e| ConfigError::InvalidValue {
|
||||||
|
key: "EMBEDDING_DIMENSION".to_string(),
|
||||||
|
message: format!("must be a positive integer: {e}"),
|
||||||
|
})?
|
||||||
|
.unwrap_or_else(|| default_dimension_for_model(&model));
|
||||||
|
|
||||||
let enabled = optional_env("EMBEDDING_ENABLED")?
|
let enabled = optional_env("EMBEDDING_ENABLED")?
|
||||||
.map(|s| s.parse())
|
.map(|s| s.parse())
|
||||||
.transpose()
|
.transpose()
|
||||||
@@ -52,6 +88,8 @@ impl EmbeddingsConfig {
|
|||||||
provider,
|
provider,
|
||||||
openai_api_key,
|
openai_api_key,
|
||||||
model,
|
model,
|
||||||
|
ollama_base_url,
|
||||||
|
dimension,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+16
-2
@@ -64,6 +64,8 @@ impl std::fmt::Display for LlmBackend {
|
|||||||
pub struct OpenAiDirectConfig {
|
pub struct OpenAiDirectConfig {
|
||||||
pub api_key: SecretString,
|
pub api_key: SecretString,
|
||||||
pub model: String,
|
pub model: String,
|
||||||
|
/// Optional base URL override (e.g. for proxies like VibeProxy).
|
||||||
|
pub base_url: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Configuration for direct Anthropic API access.
|
/// Configuration for direct Anthropic API access.
|
||||||
@@ -71,6 +73,8 @@ pub struct OpenAiDirectConfig {
|
|||||||
pub struct AnthropicDirectConfig {
|
pub struct AnthropicDirectConfig {
|
||||||
pub api_key: SecretString,
|
pub api_key: SecretString,
|
||||||
pub model: String,
|
pub model: String,
|
||||||
|
/// Optional base URL override (e.g. for proxies like VibeProxy).
|
||||||
|
pub base_url: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Configuration for local Ollama.
|
/// Configuration for local Ollama.
|
||||||
@@ -274,7 +278,12 @@ impl LlmConfig {
|
|||||||
hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(),
|
hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(),
|
||||||
})?;
|
})?;
|
||||||
let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string());
|
let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string());
|
||||||
Some(OpenAiDirectConfig { api_key, model })
|
let base_url = optional_env("OPENAI_BASE_URL")?;
|
||||||
|
Some(OpenAiDirectConfig {
|
||||||
|
api_key,
|
||||||
|
model,
|
||||||
|
base_url,
|
||||||
|
})
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
@@ -288,7 +297,12 @@ impl LlmConfig {
|
|||||||
})?;
|
})?;
|
||||||
let model = optional_env("ANTHROPIC_MODEL")?
|
let model = optional_env("ANTHROPIC_MODEL")?
|
||||||
.unwrap_or_else(|| "claude-sonnet-4-20250514".to_string());
|
.unwrap_or_else(|| "claude-sonnet-4-20250514".to_string());
|
||||||
Some(AnthropicDirectConfig { api_key, model })
|
let base_url = optional_env("ANTHROPIC_BASE_URL")?;
|
||||||
|
Some(AnthropicDirectConfig {
|
||||||
|
api_key,
|
||||||
|
model,
|
||||||
|
base_url,
|
||||||
|
})
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ impl Default for SandboxModeConfig {
|
|||||||
timeout_secs: 120,
|
timeout_secs: 120,
|
||||||
memory_limit_mb: 2048,
|
memory_limit_mb: 2048,
|
||||||
cpu_shares: 1024,
|
cpu_shares: 1024,
|
||||||
image: "ghcr.io/nearai/sandbox:latest".to_string(),
|
image: "ironclaw-worker:latest".to_string(),
|
||||||
auto_pull_image: true,
|
auto_pull_image: true,
|
||||||
extra_allowed_domains: Vec::new(),
|
extra_allowed_domains: Vec::new(),
|
||||||
}
|
}
|
||||||
@@ -57,7 +57,7 @@ impl SandboxModeConfig {
|
|||||||
memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?,
|
memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?,
|
||||||
cpu_shares: parse_optional_env("SANDBOX_CPU_SHARES", 1024)?,
|
cpu_shares: parse_optional_env("SANDBOX_CPU_SHARES", 1024)?,
|
||||||
image: optional_env("SANDBOX_IMAGE")?
|
image: optional_env("SANDBOX_IMAGE")?
|
||||||
.unwrap_or_else(|| "ghcr.io/nearai/sandbox:latest".to_string()),
|
.unwrap_or_else(|| "ironclaw-worker:latest".to_string()),
|
||||||
auto_pull_image: optional_env("SANDBOX_AUTO_PULL")?
|
auto_pull_image: optional_env("SANDBOX_AUTO_PULL")?
|
||||||
.map(|s| s.parse())
|
.map(|s| s.parse())
|
||||||
.transpose()
|
.transpose()
|
||||||
|
|||||||
@@ -123,7 +123,11 @@ impl CircuitBreakerProvider {
|
|||||||
);
|
);
|
||||||
Ok(())
|
Ok(())
|
||||||
} else {
|
} else {
|
||||||
let remaining = self.config.recovery_timeout - opened_at.elapsed();
|
let remaining = self
|
||||||
|
.config
|
||||||
|
.recovery_timeout
|
||||||
|
.checked_sub(opened_at.elapsed())
|
||||||
|
.unwrap_or(Duration::ZERO);
|
||||||
Err(LlmError::RequestFailed {
|
Err(LlmError::RequestFailed {
|
||||||
provider: self.inner.model_name().to_string(),
|
provider: self.inner.model_name().to_string(),
|
||||||
reason: format!(
|
reason: format!(
|
||||||
@@ -208,8 +212,16 @@ impl CircuitBreakerProvider {
|
|||||||
/// Returns `true` for errors that indicate the provider is degraded
|
/// Returns `true` for errors that indicate the provider is degraded
|
||||||
/// (server errors, rate limits, network failures, auth infrastructure down).
|
/// (server errors, rate limits, network failures, auth infrastructure down).
|
||||||
///
|
///
|
||||||
/// Client errors (wrong model, bad credentials, context overflow) are NOT
|
/// This answers: "should this error count toward tripping the circuit breaker?"
|
||||||
/// transient: they are the caller's problem, not a sign of backend trouble.
|
///
|
||||||
|
/// Includes `SessionExpired` because repeated session failures signal backend
|
||||||
|
/// auth infrastructure trouble.
|
||||||
|
///
|
||||||
|
/// Excludes client errors that are the caller's problem, not backend trouble:
|
||||||
|
/// `AuthFailed`, `ContextLengthExceeded`, `ModelNotAvailable`, `Json`.
|
||||||
|
///
|
||||||
|
/// See also `retry::is_retryable()` which answers a different question:
|
||||||
|
/// "could retrying this exact request succeed?"
|
||||||
fn is_transient(err: &LlmError) -> bool {
|
fn is_transient(err: &LlmError) -> bool {
|
||||||
matches!(
|
matches!(
|
||||||
err,
|
err,
|
||||||
@@ -219,7 +231,6 @@ fn is_transient(err: &LlmError) -> bool {
|
|||||||
| LlmError::SessionExpired { .. }
|
| LlmError::SessionExpired { .. }
|
||||||
| LlmError::SessionRenewalFailed { .. }
|
| LlmError::SessionRenewalFailed { .. }
|
||||||
| LlmError::Http(_)
|
| LlmError::Http(_)
|
||||||
| LlmError::Json(_)
|
|
||||||
| LlmError::Io(_)
|
| LlmError::Io(_)
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -547,6 +558,9 @@ mod tests {
|
|||||||
provider: "p".into(),
|
provider: "p".into(),
|
||||||
model: "m".into(),
|
model: "m".into(),
|
||||||
}));
|
}));
|
||||||
|
assert!(!is_transient(&LlmError::Json(
|
||||||
|
serde_json::from_str::<String>("bad").unwrap_err()
|
||||||
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
// -- Passthrough delegation tests --
|
// -- Passthrough delegation tests --
|
||||||
|
|||||||
+26
-29
@@ -19,34 +19,11 @@ use rust_decimal::Decimal;
|
|||||||
|
|
||||||
use crate::error::LlmError;
|
use crate::error::LlmError;
|
||||||
use crate::llm::provider::{
|
use crate::llm::provider::{
|
||||||
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
|
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
|
||||||
ToolCompletionResponse,
|
ToolCompletionResponse,
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Returns `true` if the error is transient and the request should be retried
|
use crate::llm::retry::is_retryable;
|
||||||
/// on the next provider in the failover chain.
|
|
||||||
///
|
|
||||||
/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`,
|
|
||||||
/// `SessionRenewalFailed`, `ModelNotAvailable`, `Http`, `Io`.
|
|
||||||
///
|
|
||||||
/// `ModelNotAvailable` is retryable because the next provider in the chain may
|
|
||||||
/// offer a different model, so it's worth trying.
|
|
||||||
///
|
|
||||||
/// Non-retryable errors (`AuthFailed`, `SessionExpired`, `ContextLengthExceeded`)
|
|
||||||
/// propagate immediately because a different provider won't fix them.
|
|
||||||
fn is_retryable(err: &LlmError) -> bool {
|
|
||||||
matches!(
|
|
||||||
err,
|
|
||||||
LlmError::RequestFailed { .. }
|
|
||||||
| LlmError::RateLimited { .. }
|
|
||||||
| LlmError::InvalidResponse { .. }
|
|
||||||
| LlmError::SessionRenewalFailed { .. }
|
|
||||||
// ModelNotAvailable is retryable: the next provider may offer a different model.
|
|
||||||
| LlmError::ModelNotAvailable { .. }
|
|
||||||
| LlmError::Http(_)
|
|
||||||
| LlmError::Io(_)
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Configuration for per-provider cooldown behavior.
|
/// Configuration for per-provider cooldown behavior.
|
||||||
///
|
///
|
||||||
@@ -376,6 +353,26 @@ impl LlmProvider for FailoverProvider {
|
|||||||
Ok(all_models)
|
Ok(all_models)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||||
|
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||||
|
.model_metadata()
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||||
|
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||||
|
.seed_response_chain(thread_id, response_id);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||||
|
self.providers[self.last_used.load(Ordering::Relaxed)].get_response_chain_id(thread_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||||
|
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||||
|
.calculate_cost(input_tokens, output_tokens)
|
||||||
|
}
|
||||||
|
|
||||||
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
|
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
|
||||||
if let Some(provider_idx) = self.take_bound_provider_for_current_task() {
|
if let Some(provider_idx) = self.take_bound_provider_for_current_task() {
|
||||||
return self.providers[provider_idx].effective_model_name(requested_model);
|
return self.providers[provider_idx].effective_model_name(requested_model);
|
||||||
@@ -1111,10 +1108,6 @@ mod tests {
|
|||||||
std::io::ErrorKind::ConnectionReset,
|
std::io::ErrorKind::ConnectionReset,
|
||||||
"reset"
|
"reset"
|
||||||
))));
|
))));
|
||||||
assert!(is_retryable(&LlmError::ModelNotAvailable {
|
|
||||||
provider: "p".into(),
|
|
||||||
model: "m".into(),
|
|
||||||
}));
|
|
||||||
|
|
||||||
// Non-retryable
|
// Non-retryable
|
||||||
assert!(!is_retryable(&LlmError::AuthFailed {
|
assert!(!is_retryable(&LlmError::AuthFailed {
|
||||||
@@ -1127,6 +1120,10 @@ mod tests {
|
|||||||
used: 100_000,
|
used: 100_000,
|
||||||
limit: 50_000,
|
limit: 50_000,
|
||||||
}));
|
}));
|
||||||
|
assert!(!is_retryable(&LlmError::ModelNotAvailable {
|
||||||
|
provider: "p".into(),
|
||||||
|
model: "m".into(),
|
||||||
|
}));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Test: empty providers list returns error (not panic).
|
// Test: empty providers list returns error (not panic).
|
||||||
|
|||||||
+55
-31
@@ -15,7 +15,7 @@ mod nearai_chat;
|
|||||||
mod provider;
|
mod provider;
|
||||||
mod reasoning;
|
mod reasoning;
|
||||||
pub mod response_cache;
|
pub mod response_cache;
|
||||||
mod retry;
|
pub mod retry;
|
||||||
mod rig_adapter;
|
mod rig_adapter;
|
||||||
pub mod session;
|
pub mod session;
|
||||||
|
|
||||||
@@ -32,6 +32,7 @@ pub use reasoning::{
|
|||||||
ToolSelection,
|
ToolSelection,
|
||||||
};
|
};
|
||||||
pub use response_cache::{CachedProvider, ResponseCacheConfig};
|
pub use response_cache::{CachedProvider, ResponseCacheConfig};
|
||||||
|
pub use retry::{RetryConfig, RetryProvider};
|
||||||
pub use rig_adapter::RigAdapter;
|
pub use rig_adapter::RigAdapter;
|
||||||
pub use session::{SessionConfig, SessionManager, create_session_manager};
|
pub use session::{SessionConfig, SessionManager, create_session_manager};
|
||||||
|
|
||||||
@@ -76,7 +77,7 @@ pub fn create_llm_provider_with_config(
|
|||||||
model = %config.model,
|
model = %config.model,
|
||||||
"Using Responses API (chat-api) with session auth"
|
"Using Responses API (chat-api) with session auth"
|
||||||
);
|
);
|
||||||
Ok(Arc::new(NearAiProvider::new(config.clone(), session)))
|
Ok(Arc::new(NearAiProvider::new(config.clone(), session)?))
|
||||||
}
|
}
|
||||||
NearAiApiMode::ChatCompletions => {
|
NearAiApiMode::ChatCompletions => {
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
@@ -99,15 +100,30 @@ fn create_openai_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, Ll
|
|||||||
// (Responses API). The Responses API path in rig-core panics when tool results
|
// (Responses API). The Responses API path in rig-core panics when tool results
|
||||||
// are sent back because ironclaw doesn't thread `call_id` through its ToolCall
|
// are sent back because ironclaw doesn't thread `call_id` through its ToolCall
|
||||||
// type. The Chat Completions API works correctly with the existing code.
|
// type. The Chat Completions API works correctly with the existing code.
|
||||||
let client: openai::CompletionsClient = openai::Client::new(oai.api_key.expose_secret())
|
let client: openai::CompletionsClient = if let Some(ref base_url) = oai.base_url {
|
||||||
.map_err(|e| LlmError::RequestFailed {
|
tracing::info!(
|
||||||
provider: "openai".to_string(),
|
"Using OpenAI direct API (chat completions, model: {}, base_url: {})",
|
||||||
reason: format!("Failed to create OpenAI client: {}", e),
|
oai.model,
|
||||||
})?
|
base_url,
|
||||||
.completions_api();
|
);
|
||||||
|
openai::Client::builder()
|
||||||
|
.base_url(base_url)
|
||||||
|
.api_key(oai.api_key.expose_secret())
|
||||||
|
.build()
|
||||||
|
} else {
|
||||||
|
tracing::info!(
|
||||||
|
"Using OpenAI direct API (chat completions, model: {}, base_url: default)",
|
||||||
|
oai.model,
|
||||||
|
);
|
||||||
|
openai::Client::new(oai.api_key.expose_secret())
|
||||||
|
}
|
||||||
|
.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "openai".to_string(),
|
||||||
|
reason: format!("Failed to create OpenAI client: {}", e),
|
||||||
|
})?
|
||||||
|
.completions_api();
|
||||||
|
|
||||||
let model = client.completion_model(&oai.model);
|
let model = client.completion_model(&oai.model);
|
||||||
tracing::info!("Using OpenAI direct API (model: {})", oai.model);
|
|
||||||
Ok(Arc::new(RigAdapter::new(model, &oai.model)))
|
Ok(Arc::new(RigAdapter::new(model, &oai.model)))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -121,16 +137,25 @@ fn create_anthropic_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>,
|
|||||||
|
|
||||||
use rig::providers::anthropic;
|
use rig::providers::anthropic;
|
||||||
|
|
||||||
let client: anthropic::Client =
|
let client: anthropic::Client = if let Some(ref base_url) = anth.base_url {
|
||||||
anthropic::Client::new(anth.api_key.expose_secret()).map_err(|e| {
|
anthropic::Client::builder()
|
||||||
LlmError::RequestFailed {
|
.api_key(anth.api_key.expose_secret())
|
||||||
provider: "anthropic".to_string(),
|
.base_url(base_url)
|
||||||
reason: format!("Failed to create Anthropic client: {}", e),
|
.build()
|
||||||
}
|
} else {
|
||||||
})?;
|
anthropic::Client::new(anth.api_key.expose_secret())
|
||||||
|
}
|
||||||
|
.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "anthropic".to_string(),
|
||||||
|
reason: format!("Failed to create Anthropic client: {}", e),
|
||||||
|
})?;
|
||||||
|
|
||||||
let model = client.completion_model(&anth.model);
|
let model = client.completion_model(&anth.model);
|
||||||
tracing::info!("Using Anthropic direct API (model: {})", anth.model);
|
tracing::info!(
|
||||||
|
"Using Anthropic direct API (model: {}, base_url: {})",
|
||||||
|
anth.model,
|
||||||
|
anth.base_url.as_deref().unwrap_or("default"),
|
||||||
|
);
|
||||||
Ok(Arc::new(RigAdapter::new(model, &anth.model)))
|
Ok(Arc::new(RigAdapter::new(model, &anth.model)))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -199,26 +224,25 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmPr
|
|||||||
|
|
||||||
use rig::providers::openai;
|
use rig::providers::openai;
|
||||||
|
|
||||||
let api_key = compat
|
let client: openai::CompletionsClient = openai::Client::builder()
|
||||||
.api_key
|
|
||||||
.as_ref()
|
|
||||||
.map(|k| k.expose_secret().to_string())
|
|
||||||
.unwrap_or_else(|| "no-key".to_string());
|
|
||||||
|
|
||||||
let client: openai::Client = openai::Client::builder()
|
|
||||||
.base_url(&compat.base_url)
|
.base_url(&compat.base_url)
|
||||||
.api_key(api_key)
|
.api_key(
|
||||||
|
compat
|
||||||
|
.api_key
|
||||||
|
.as_ref()
|
||||||
|
.map(|k| k.expose_secret().to_string())
|
||||||
|
.unwrap_or_else(|| "no-key".to_string()),
|
||||||
|
)
|
||||||
.build()
|
.build()
|
||||||
.map_err(|e| LlmError::RequestFailed {
|
.map_err(|e| LlmError::RequestFailed {
|
||||||
provider: "openai_compatible".to_string(),
|
provider: "openai_compatible".to_string(),
|
||||||
reason: format!("Failed to create OpenAI-compatible client: {}", e),
|
reason: format!("Failed to create OpenAI-compatible client: {}", e),
|
||||||
})?;
|
})?
|
||||||
|
.completions_api();
|
||||||
|
|
||||||
// OpenAI-compatible providers (e.g. OpenRouter) are most reliable on Chat Completions.
|
let model = client.completion_model(&compat.model);
|
||||||
// This avoids Responses-API-specific assumptions such as required tool call IDs.
|
|
||||||
let model = client.completions_api().completion_model(&compat.model);
|
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
"Using OpenAI-compatible endpoint via Chat Completions API (base_url: {}, model: {})",
|
"Using OpenAI-compatible endpoint (chat completions, base_url: {}, model: {})",
|
||||||
compat.base_url,
|
compat.base_url,
|
||||||
compat.model
|
compat.model
|
||||||
);
|
);
|
||||||
@@ -252,7 +276,7 @@ pub fn create_cheap_llm_provider(
|
|||||||
tracing::info!("Cheap LLM provider: {}", cheap_model);
|
tracing::info!("Cheap LLM provider: {}", cheap_model);
|
||||||
|
|
||||||
match cheap_config.api_mode {
|
match cheap_config.api_mode {
|
||||||
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)))),
|
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)?))),
|
||||||
NearAiApiMode::ChatCompletions => {
|
NearAiApiMode::ChatCompletions => {
|
||||||
Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?)))
|
Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?)))
|
||||||
}
|
}
|
||||||
|
|||||||
+135
-147
@@ -19,7 +19,6 @@ use crate::llm::provider::{
|
|||||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||||
ToolCompletionRequest, ToolCompletionResponse,
|
ToolCompletionRequest, ToolCompletionResponse,
|
||||||
};
|
};
|
||||||
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
|
|
||||||
use crate::llm::session::SessionManager;
|
use crate::llm::session::SessionManager;
|
||||||
|
|
||||||
/// Information about an available model from NEAR AI API.
|
/// Information about an available model from NEAR AI API.
|
||||||
@@ -54,28 +53,34 @@ pub struct NearAiProvider {
|
|||||||
|
|
||||||
impl NearAiProvider {
|
impl NearAiProvider {
|
||||||
/// Create a new NEAR AI provider with a session manager.
|
/// Create a new NEAR AI provider with a session manager.
|
||||||
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Self {
|
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Result<Self, LlmError> {
|
||||||
let client = Client::builder()
|
let client = Client::builder()
|
||||||
.timeout(std::time::Duration::from_secs(120))
|
.timeout(std::time::Duration::from_secs(120))
|
||||||
.build()
|
.build()
|
||||||
.unwrap_or_else(|_| Client::new());
|
.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "nearai".to_string(),
|
||||||
|
reason: format!("Failed to build HTTP client: {}", e),
|
||||||
|
})?;
|
||||||
|
|
||||||
let active_model = std::sync::RwLock::new(config.model.clone());
|
let active_model = std::sync::RwLock::new(config.model.clone());
|
||||||
Self {
|
Ok(Self {
|
||||||
client,
|
client,
|
||||||
config,
|
config,
|
||||||
session,
|
session,
|
||||||
active_model,
|
active_model,
|
||||||
response_chains: std::sync::RwLock::new(HashMap::new()),
|
response_chains: std::sync::RwLock::new(HashMap::new()),
|
||||||
}
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Seed a response chain for a thread (e.g. when restoring from DB).
|
/// Seed a response chain for a thread (e.g. when restoring from DB).
|
||||||
pub fn seed_response_id(&self, thread_id: &str, response_id: String) {
|
pub fn seed_response_id(&self, thread_id: &str, response_id: String) {
|
||||||
let mut chains = self
|
let mut chains = match self.response_chains.write() {
|
||||||
.response_chains
|
Ok(guard) => guard,
|
||||||
.write()
|
Err(poisoned) => {
|
||||||
.expect("response_chains lock poisoned");
|
tracing::warn!("response_chains lock poisoned in seed; recovering");
|
||||||
|
poisoned.into_inner()
|
||||||
|
}
|
||||||
|
};
|
||||||
chains.insert(
|
chains.insert(
|
||||||
thread_id.to_string(),
|
thread_id.to_string(),
|
||||||
ChainState {
|
ChainState {
|
||||||
@@ -87,19 +92,25 @@ impl NearAiProvider {
|
|||||||
|
|
||||||
/// Get the last response ID for a thread (for persistence).
|
/// Get the last response ID for a thread (for persistence).
|
||||||
pub fn get_response_id(&self, thread_id: &str) -> Option<String> {
|
pub fn get_response_id(&self, thread_id: &str) -> Option<String> {
|
||||||
let chains = self
|
let chains = match self.response_chains.read() {
|
||||||
.response_chains
|
Ok(guard) => guard,
|
||||||
.read()
|
Err(poisoned) => {
|
||||||
.expect("response_chains lock poisoned");
|
tracing::warn!("response_chains lock poisoned in get; recovering");
|
||||||
|
poisoned.into_inner()
|
||||||
|
}
|
||||||
|
};
|
||||||
chains.get(thread_id).map(|c| c.response_id.clone())
|
chains.get(thread_id).map(|c| c.response_id.clone())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Store a response chain state after a successful call.
|
/// Store a response chain state after a successful call.
|
||||||
fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) {
|
fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) {
|
||||||
let mut chains = self
|
let mut chains = match self.response_chains.write() {
|
||||||
.response_chains
|
Ok(guard) => guard,
|
||||||
.write()
|
Err(poisoned) => {
|
||||||
.expect("response_chains lock poisoned");
|
tracing::warn!("response_chains lock poisoned in store; recovering");
|
||||||
|
poisoned.into_inner()
|
||||||
|
}
|
||||||
|
};
|
||||||
chains.insert(
|
chains.insert(
|
||||||
thread_id.to_string(),
|
thread_id.to_string(),
|
||||||
ChainState {
|
ChainState {
|
||||||
@@ -111,10 +122,13 @@ impl NearAiProvider {
|
|||||||
|
|
||||||
/// Clear the chain for a thread (on error / fallback).
|
/// Clear the chain for a thread (on error / fallback).
|
||||||
fn clear_chain(&self, thread_id: &str) {
|
fn clear_chain(&self, thread_id: &str) {
|
||||||
let mut chains = self
|
let mut chains = match self.response_chains.write() {
|
||||||
.response_chains
|
Ok(guard) => guard,
|
||||||
.write()
|
Err(poisoned) => {
|
||||||
.expect("response_chains lock poisoned");
|
tracing::warn!("response_chains lock poisoned in clear; recovering");
|
||||||
|
poisoned.into_inner()
|
||||||
|
}
|
||||||
|
};
|
||||||
chains.remove(thread_id);
|
chains.remove(thread_id);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -160,7 +174,10 @@ impl NearAiProvider {
|
|||||||
})?;
|
})?;
|
||||||
|
|
||||||
let status = response.status();
|
let status = response.status();
|
||||||
let response_text = response.text().await.unwrap_or_default();
|
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "nearai".to_string(),
|
||||||
|
reason: format!("Failed to read response body: {}", e),
|
||||||
|
})?;
|
||||||
|
|
||||||
if !status.is_success() {
|
if !status.is_success() {
|
||||||
if status.as_u16() == 401 {
|
if status.as_u16() == 401 {
|
||||||
@@ -283,139 +300,95 @@ impl NearAiProvider {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Inner request implementation with retry logic for transient errors.
|
/// Inner request implementation (single attempt).
|
||||||
///
|
///
|
||||||
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
|
/// Does not retry internally — retries are handled by the external
|
||||||
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
|
/// `RetryProvider` wrapper in the composition chain.
|
||||||
async fn send_request_inner<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
|
async fn send_request_inner<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
|
||||||
&self,
|
&self,
|
||||||
path: &str,
|
path: &str,
|
||||||
body: &T,
|
body: &T,
|
||||||
) -> Result<R, LlmError> {
|
) -> Result<R, LlmError> {
|
||||||
let url = self.api_url(path);
|
let url = self.api_url(path);
|
||||||
let max_retries = self.config.max_retries;
|
let token = self.session.get_token().await?;
|
||||||
|
|
||||||
for attempt in 0..=max_retries {
|
tracing::debug!("Sending request to NEAR AI: {}", url);
|
||||||
let token = self.session.get_token().await?;
|
tracing::debug!("Request body: {:?}", body);
|
||||||
|
|
||||||
tracing::debug!(
|
let response = self
|
||||||
"Sending request to NEAR AI: {} (attempt {})",
|
.client
|
||||||
url,
|
.post(&url)
|
||||||
attempt + 1
|
.header("Authorization", format!("Bearer {}", token.expose_secret()))
|
||||||
);
|
.header("Content-Type", "application/json")
|
||||||
tracing::debug!("Request body: {:?}", body);
|
.json(body)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| {
|
||||||
|
tracing::error!("NEAR AI request failed: {}", e);
|
||||||
|
LlmError::Http(e)
|
||||||
|
})?;
|
||||||
|
|
||||||
let response = self
|
let status = response.status();
|
||||||
.client
|
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||||
.post(&url)
|
provider: "nearai".to_string(),
|
||||||
.header("Authorization", format!("Bearer {}", token.expose_secret()))
|
reason: format!("Failed to read response body: {}", e),
|
||||||
.header("Content-Type", "application/json")
|
})?;
|
||||||
.json(body)
|
|
||||||
.send()
|
|
||||||
.await;
|
|
||||||
|
|
||||||
let response = match response {
|
tracing::debug!("NEAR AI response status: {}", status);
|
||||||
Ok(r) => r,
|
tracing::debug!("NEAR AI response body: {}", response_text);
|
||||||
Err(e) => {
|
|
||||||
tracing::error!("NEAR AI request failed: {}", e);
|
|
||||||
// Network errors (timeout, connection refused) are transient
|
|
||||||
if attempt < max_retries {
|
|
||||||
let delay = retry_backoff_delay(attempt);
|
|
||||||
tracing::warn!(
|
|
||||||
"NEAR AI request error (attempt {}/{}), retrying in {:?}: {}",
|
|
||||||
attempt + 1,
|
|
||||||
max_retries + 1,
|
|
||||||
delay,
|
|
||||||
e,
|
|
||||||
);
|
|
||||||
tokio::time::sleep(delay).await;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
return Err(e.into());
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let status = response.status();
|
if !status.is_success() {
|
||||||
let response_text = response.text().await.unwrap_or_default();
|
let status_code = status.as_u16();
|
||||||
|
|
||||||
tracing::debug!("NEAR AI response status: {}", status);
|
// Check for session expiration (401 with specific message patterns)
|
||||||
tracing::debug!("NEAR AI response body: {}", response_text);
|
if status_code == 401 {
|
||||||
|
let lower = response_text.to_lowercase();
|
||||||
|
let is_session_expired = lower.contains("session")
|
||||||
|
&& (lower.contains("expired") || lower.contains("invalid"));
|
||||||
|
|
||||||
if !status.is_success() {
|
if is_session_expired {
|
||||||
let status_code = status.as_u16();
|
return Err(LlmError::SessionExpired {
|
||||||
|
|
||||||
// Check for session expiration (401 with specific message patterns)
|
|
||||||
if status_code == 401 {
|
|
||||||
let lower = response_text.to_lowercase();
|
|
||||||
let is_session_expired = lower.contains("session")
|
|
||||||
&& (lower.contains("expired") || lower.contains("invalid"));
|
|
||||||
|
|
||||||
if is_session_expired {
|
|
||||||
return Err(LlmError::SessionExpired {
|
|
||||||
provider: "nearai".to_string(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
// Generic 401 -- not retryable
|
|
||||||
return Err(LlmError::AuthFailed {
|
|
||||||
provider: "nearai".to_string(),
|
provider: "nearai".to_string(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check if this is a transient error worth retrying
|
return Err(LlmError::AuthFailed {
|
||||||
if is_retryable_status(status_code) && attempt < max_retries {
|
|
||||||
let delay = retry_backoff_delay(attempt);
|
|
||||||
tracing::warn!(
|
|
||||||
"NEAR AI returned HTTP {} (attempt {}/{}), retrying in {:?}",
|
|
||||||
status_code,
|
|
||||||
attempt + 1,
|
|
||||||
max_retries + 1,
|
|
||||||
delay,
|
|
||||||
);
|
|
||||||
tokio::time::sleep(delay).await;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Non-retryable error or exhausted retries
|
|
||||||
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
|
|
||||||
if status_code == 429 {
|
|
||||||
return Err(LlmError::RateLimited {
|
|
||||||
provider: "nearai".to_string(),
|
|
||||||
retry_after: None,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
return Err(LlmError::RequestFailed {
|
|
||||||
provider: "nearai".to_string(),
|
|
||||||
reason: error.error,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
return Err(LlmError::RequestFailed {
|
|
||||||
provider: "nearai".to_string(),
|
provider: "nearai".to_string(),
|
||||||
reason: format!("HTTP {}: {}", status, response_text),
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Success -- parse the response
|
if status_code == 429 {
|
||||||
return match serde_json::from_str::<R>(&response_text) {
|
return Err(LlmError::RateLimited {
|
||||||
Ok(parsed) => Ok(parsed),
|
provider: "nearai".to_string(),
|
||||||
Err(e) => {
|
retry_after: None,
|
||||||
tracing::debug!("Response is not expected JSON format: {}", e);
|
});
|
||||||
tracing::debug!("Will try alternative parsing in caller");
|
}
|
||||||
Err(LlmError::InvalidResponse {
|
|
||||||
provider: "nearai".to_string(),
|
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
|
||||||
reason: format!("Parse error: {}. Raw: {}", e, response_text),
|
return Err(LlmError::RequestFailed {
|
||||||
})
|
provider: "nearai".to_string(),
|
||||||
}
|
reason: error.error,
|
||||||
};
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
return Err(LlmError::RequestFailed {
|
||||||
|
provider: "nearai".to_string(),
|
||||||
|
reason: format!("HTTP {}: {}", status, response_text),
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// This is unreachable because the loop always returns, but the compiler
|
// Success -- parse the response
|
||||||
// cannot prove that. Return a generic error as a safety net.
|
match serde_json::from_str::<R>(&response_text) {
|
||||||
Err(LlmError::RequestFailed {
|
Ok(parsed) => Ok(parsed),
|
||||||
provider: "nearai".to_string(),
|
Err(e) => {
|
||||||
reason: "retry loop exited unexpectedly".to_string(),
|
tracing::debug!("Response is not expected JSON format: {}", e);
|
||||||
})
|
tracing::debug!("Will try alternative parsing in caller");
|
||||||
|
Err(LlmError::InvalidResponse {
|
||||||
|
provider: "nearai".to_string(),
|
||||||
|
reason: format!("Parse error: {}. Raw: {}", e, response_text),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -464,7 +437,9 @@ impl LlmProvider for NearAiProvider {
|
|||||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||||
let thread_id = req.metadata.get("thread_id").cloned();
|
let thread_id = req.metadata.get("thread_id").cloned();
|
||||||
let (instructions, input) = split_messages(req.messages, false);
|
let mut messages = req.messages;
|
||||||
|
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||||
|
let (instructions, input) = split_messages(messages, false);
|
||||||
|
|
||||||
let request = NearAiRequest {
|
let request = NearAiRequest {
|
||||||
model,
|
model,
|
||||||
@@ -582,13 +557,20 @@ impl LlmProvider for NearAiProvider {
|
|||||||
) -> Result<ToolCompletionResponse, LlmError> {
|
) -> Result<ToolCompletionResponse, LlmError> {
|
||||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||||
let thread_id = req.metadata.get("thread_id").cloned();
|
let thread_id = req.metadata.get("thread_id").cloned();
|
||||||
|
let mut messages = req.messages;
|
||||||
|
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||||
|
|
||||||
// Look up chaining state for this thread
|
// Look up chaining state for this thread
|
||||||
let chain_state = thread_id.as_ref().and_then(|tid| {
|
let chain_state = thread_id.as_ref().and_then(|tid| {
|
||||||
let chains = self
|
let chains = match self.response_chains.read() {
|
||||||
.response_chains
|
Ok(guard) => guard,
|
||||||
.read()
|
Err(poisoned) => {
|
||||||
.expect("response_chains lock poisoned");
|
tracing::warn!(
|
||||||
|
"response_chains lock poisoned in complete_with_tools; recovering"
|
||||||
|
);
|
||||||
|
poisoned.into_inner()
|
||||||
|
}
|
||||||
|
};
|
||||||
chains
|
chains
|
||||||
.get(tid)
|
.get(tid)
|
||||||
.map(|c| (c.response_id.clone(), c.input_count))
|
.map(|c| (c.response_id.clone(), c.input_count))
|
||||||
@@ -601,7 +583,7 @@ impl LlmProvider for NearAiProvider {
|
|||||||
|
|
||||||
// When chaining, only send new messages (the delta since last call).
|
// When chaining, only send new messages (the delta since last call).
|
||||||
// Tool results are converted to function_call_output items.
|
// Tool results are converted to function_call_output items.
|
||||||
let (instructions, all_input) = split_messages(req.messages, chaining);
|
let (instructions, all_input) = split_messages(messages, chaining);
|
||||||
let input = if chaining && all_input.len() > prev_input_count {
|
let input = if chaining && all_input.len() > prev_input_count {
|
||||||
all_input[prev_input_count..].to_vec()
|
all_input[prev_input_count..].to_vec()
|
||||||
} else {
|
} else {
|
||||||
@@ -806,18 +788,25 @@ impl LlmProvider for NearAiProvider {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn active_model_name(&self) -> String {
|
fn active_model_name(&self) -> String {
|
||||||
self.active_model
|
match self.active_model.read() {
|
||||||
.read()
|
Ok(guard) => guard.clone(),
|
||||||
.expect("active_model lock poisoned")
|
Err(poisoned) => {
|
||||||
.clone()
|
tracing::warn!("active_model lock poisoned while reading; continuing");
|
||||||
|
poisoned.into_inner().clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
||||||
let mut guard = self
|
match self.active_model.write() {
|
||||||
.active_model
|
Ok(mut guard) => {
|
||||||
.write()
|
*guard = model.to_string();
|
||||||
.expect("active_model lock poisoned");
|
}
|
||||||
*guard = model.to_string();
|
Err(poisoned) => {
|
||||||
|
tracing::warn!("active_model lock poisoned while writing; continuing");
|
||||||
|
*poisoned.into_inner() = model.to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -952,7 +941,6 @@ struct NearAiTool {
|
|||||||
/// Primary response format (output array style)
|
/// Primary response format (output array style)
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct NearAiResponse {
|
struct NearAiResponse {
|
||||||
#[allow(dead_code)]
|
|
||||||
id: String,
|
id: String,
|
||||||
output: Vec<NearAiOutputItem>,
|
output: Vec<NearAiOutputItem>,
|
||||||
usage: NearAiUsage,
|
usage: NearAiUsage,
|
||||||
|
|||||||
+197
-125
@@ -16,18 +16,29 @@ use crate::llm::provider::{
|
|||||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||||||
};
|
};
|
||||||
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
|
|
||||||
|
|
||||||
/// NEAR AI Chat Completions API provider.
|
/// NEAR AI Chat Completions API provider.
|
||||||
pub struct NearAiChatProvider {
|
pub struct NearAiChatProvider {
|
||||||
client: Client,
|
client: Client,
|
||||||
config: NearAiConfig,
|
config: NearAiConfig,
|
||||||
active_model: std::sync::RwLock<String>,
|
active_model: std::sync::RwLock<String>,
|
||||||
|
flatten_tool_messages: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl NearAiChatProvider {
|
impl NearAiChatProvider {
|
||||||
/// Create a new NEAR AI chat completions provider with API key auth.
|
/// Create a new NEAR AI chat completions provider with API key auth.
|
||||||
|
///
|
||||||
|
/// By default this enables tool-message flattening for compatibility with
|
||||||
|
/// providers that reject `role: "tool"` messages (e.g. NEAR cloud-api).
|
||||||
pub fn new(config: NearAiConfig) -> Result<Self, LlmError> {
|
pub fn new(config: NearAiConfig) -> Result<Self, LlmError> {
|
||||||
|
Self::new_with_flatten(config, true)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a chat completions provider with configurable tool-message flattening.
|
||||||
|
pub fn new_with_flatten(
|
||||||
|
config: NearAiConfig,
|
||||||
|
flatten_tool_messages: bool,
|
||||||
|
) -> Result<Self, LlmError> {
|
||||||
if config.api_key.is_none() {
|
if config.api_key.is_none() {
|
||||||
return Err(LlmError::AuthFailed {
|
return Err(LlmError::AuthFailed {
|
||||||
provider: "nearai_chat".to_string(),
|
provider: "nearai_chat".to_string(),
|
||||||
@@ -37,22 +48,29 @@ impl NearAiChatProvider {
|
|||||||
let client = Client::builder()
|
let client = Client::builder()
|
||||||
.timeout(std::time::Duration::from_secs(120))
|
.timeout(std::time::Duration::from_secs(120))
|
||||||
.build()
|
.build()
|
||||||
.unwrap_or_else(|_| Client::new());
|
.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "nearai_chat".to_string(),
|
||||||
|
reason: format!("Failed to build HTTP client: {}", e),
|
||||||
|
})?;
|
||||||
|
|
||||||
let active_model = std::sync::RwLock::new(config.model.clone());
|
let active_model = std::sync::RwLock::new(config.model.clone());
|
||||||
Ok(Self {
|
Ok(Self {
|
||||||
client,
|
client,
|
||||||
config,
|
config,
|
||||||
active_model,
|
active_model,
|
||||||
|
flatten_tool_messages,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn api_url(&self, path: &str) -> String {
|
fn api_url(&self, path: &str) -> String {
|
||||||
format!(
|
let base = self.config.base_url.trim_end_matches('/');
|
||||||
"{}/v1/{}",
|
let path = path.trim_start_matches('/');
|
||||||
self.config.base_url,
|
|
||||||
path.trim_start_matches('/')
|
if base.ends_with("/v1") {
|
||||||
)
|
format!("{}/{}", base, path)
|
||||||
|
} else {
|
||||||
|
format!("{}/v1/{}", base, path)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn api_key(&self) -> String {
|
fn api_key(&self) -> String {
|
||||||
@@ -63,116 +81,75 @@ impl NearAiChatProvider {
|
|||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Send a request to the chat completions API with retry on transient errors.
|
/// Send a single request to the chat completions API.
|
||||||
///
|
///
|
||||||
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
|
/// Does not retry internally — retries are handled by the external
|
||||||
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
|
/// `RetryProvider` wrapper in the composition chain.
|
||||||
async fn send_request<T: Serialize, R: for<'de> Deserialize<'de>>(
|
async fn send_request<T: Serialize, R: for<'de> Deserialize<'de>>(
|
||||||
&self,
|
&self,
|
||||||
body: &T,
|
body: &T,
|
||||||
) -> Result<R, LlmError> {
|
) -> Result<R, LlmError> {
|
||||||
let url = self.api_url("chat/completions");
|
let url = self.api_url("chat/completions");
|
||||||
let max_retries = self.config.max_retries;
|
|
||||||
|
|
||||||
for attempt in 0..=max_retries {
|
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
|
||||||
tracing::debug!(
|
|
||||||
"Sending request to NEAR AI Chat: {} (attempt {})",
|
|
||||||
url,
|
|
||||||
attempt + 1,
|
|
||||||
);
|
|
||||||
|
|
||||||
if tracing::enabled!(tracing::Level::DEBUG)
|
if tracing::enabled!(tracing::Level::DEBUG)
|
||||||
&& let Ok(json) = serde_json::to_string(body)
|
&& let Ok(json) = serde_json::to_string(body)
|
||||||
{
|
{
|
||||||
tracing::debug!("NEAR AI Chat request body: {}", json);
|
tracing::debug!("NEAR AI Chat request body: {}", json);
|
||||||
}
|
}
|
||||||
|
|
||||||
let response = self
|
let response = self
|
||||||
.client
|
.client
|
||||||
.post(&url)
|
.post(&url)
|
||||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||||
.header("Content-Type", "application/json")
|
.header("Content-Type", "application/json")
|
||||||
.json(body)
|
.json(body)
|
||||||
.send()
|
.send()
|
||||||
.await;
|
.await
|
||||||
|
.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "nearai_chat".to_string(),
|
||||||
|
reason: e.to_string(),
|
||||||
|
})?;
|
||||||
|
|
||||||
let response = match response {
|
let status = response.status();
|
||||||
Ok(r) => r,
|
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||||
Err(e) => {
|
provider: "nearai_chat".to_string(),
|
||||||
tracing::error!("NEAR AI Chat request failed: {}", e);
|
reason: format!("Failed to read response body: {}", e),
|
||||||
if attempt < max_retries {
|
})?;
|
||||||
let delay = retry_backoff_delay(attempt);
|
|
||||||
tracing::warn!(
|
|
||||||
"NEAR AI Chat request error (attempt {}/{}), retrying in {:?}: {}",
|
|
||||||
attempt + 1,
|
|
||||||
max_retries + 1,
|
|
||||||
delay,
|
|
||||||
e,
|
|
||||||
);
|
|
||||||
tokio::time::sleep(delay).await;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
return Err(LlmError::RequestFailed {
|
|
||||||
provider: "nearai_chat".to_string(),
|
|
||||||
reason: e.to_string(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let status = response.status();
|
tracing::debug!("NEAR AI Chat response status: {}", status);
|
||||||
let response_text = response.text().await.unwrap_or_default();
|
tracing::debug!("NEAR AI Chat response body: {}", response_text);
|
||||||
|
|
||||||
tracing::debug!("NEAR AI Chat response status: {}", status);
|
if !status.is_success() {
|
||||||
tracing::debug!("NEAR AI Chat response body: {}", response_text);
|
let status_code = status.as_u16();
|
||||||
|
|
||||||
if !status.is_success() {
|
if status_code == 401 {
|
||||||
let status_code = status.as_u16();
|
return Err(LlmError::AuthFailed {
|
||||||
|
|
||||||
// Auth errors are not retryable
|
|
||||||
if status_code == 401 {
|
|
||||||
return Err(LlmError::AuthFailed {
|
|
||||||
provider: "nearai_chat".to_string(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
// Transient errors: retry with backoff
|
|
||||||
if is_retryable_status(status_code) && attempt < max_retries {
|
|
||||||
let delay = retry_backoff_delay(attempt);
|
|
||||||
tracing::warn!(
|
|
||||||
"NEAR AI Chat returned HTTP {} (attempt {}/{}), retrying in {:?}",
|
|
||||||
status_code,
|
|
||||||
attempt + 1,
|
|
||||||
max_retries + 1,
|
|
||||||
delay,
|
|
||||||
);
|
|
||||||
tokio::time::sleep(delay).await;
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
// Non-retryable or exhausted retries
|
|
||||||
if status_code == 429 {
|
|
||||||
return Err(LlmError::RateLimited {
|
|
||||||
provider: "nearai_chat".to_string(),
|
|
||||||
retry_after: None,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
return Err(LlmError::RequestFailed {
|
|
||||||
provider: "nearai_chat".to_string(),
|
provider: "nearai_chat".to_string(),
|
||||||
reason: format!("HTTP {}: {}", status, response_text),
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Success — parse the response
|
if status_code == 429 {
|
||||||
return serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
|
return Err(LlmError::RateLimited {
|
||||||
|
provider: "nearai_chat".to_string(),
|
||||||
|
retry_after: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||||
|
return Err(LlmError::RequestFailed {
|
||||||
provider: "nearai_chat".to_string(),
|
provider: "nearai_chat".to_string(),
|
||||||
reason: format!("JSON parse error: {}. Raw: {}", e, response_text),
|
reason: format!("HTTP {}: {}", status, truncated),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
// Safety net: unreachable because the loop always returns
|
serde_json::from_str(&response_text).map_err(|e| {
|
||||||
Err(LlmError::RequestFailed {
|
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||||
provider: "nearai_chat".to_string(),
|
LlmError::InvalidResponse {
|
||||||
reason: "retry loop exited unexpectedly".to_string(),
|
provider: "nearai_chat".to_string(),
|
||||||
|
reason: format!("JSON parse error: {}. Raw: {}", e, truncated),
|
||||||
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -192,12 +169,16 @@ impl NearAiChatProvider {
|
|||||||
})?;
|
})?;
|
||||||
|
|
||||||
let status = response.status();
|
let status = response.status();
|
||||||
let response_text = response.text().await.unwrap_or_default();
|
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||||
|
provider: "nearai_chat".to_string(),
|
||||||
|
reason: format!("Failed to read response body: {}", e),
|
||||||
|
})?;
|
||||||
|
|
||||||
if !status.is_success() {
|
if !status.is_success() {
|
||||||
|
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||||
return Err(LlmError::RequestFailed {
|
return Err(LlmError::RequestFailed {
|
||||||
provider: "nearai_chat".to_string(),
|
provider: "nearai_chat".to_string(),
|
||||||
reason: format!("HTTP {}: {}", status, response_text),
|
reason: format!("HTTP {}: {}", status, truncated),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -228,8 +209,10 @@ struct ApiModelEntry {
|
|||||||
impl LlmProvider for NearAiChatProvider {
|
impl LlmProvider for NearAiChatProvider {
|
||||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||||
|
let mut raw_messages = req.messages;
|
||||||
|
crate::llm::provider::sanitize_tool_messages(&mut raw_messages);
|
||||||
let messages: Vec<ChatCompletionMessage> =
|
let messages: Vec<ChatCompletionMessage> =
|
||||||
req.messages.into_iter().map(|m| m.into()).collect();
|
raw_messages.into_iter().map(|m| m.into()).collect();
|
||||||
|
|
||||||
let request = ChatCompletionRequest {
|
let request = ChatCompletionRequest {
|
||||||
model,
|
model,
|
||||||
@@ -261,11 +244,13 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
_ => FinishReason::Unknown,
|
_ => FinishReason::Unknown,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let (input_tokens, output_tokens) = parse_usage(response.usage.as_ref());
|
||||||
|
|
||||||
Ok(CompletionResponse {
|
Ok(CompletionResponse {
|
||||||
content,
|
content,
|
||||||
finish_reason,
|
finish_reason,
|
||||||
input_tokens: response.usage.prompt_tokens,
|
input_tokens,
|
||||||
output_tokens: response.usage.completion_tokens,
|
output_tokens,
|
||||||
response_id: None,
|
response_id: None,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -275,14 +260,18 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
req: ToolCompletionRequest,
|
req: ToolCompletionRequest,
|
||||||
) -> Result<ToolCompletionResponse, LlmError> {
|
) -> Result<ToolCompletionResponse, LlmError> {
|
||||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||||
|
let mut raw_messages = req.messages;
|
||||||
|
crate::llm::provider::sanitize_tool_messages(&mut raw_messages);
|
||||||
let messages: Vec<ChatCompletionMessage> =
|
let messages: Vec<ChatCompletionMessage> =
|
||||||
req.messages.into_iter().map(|m| m.into()).collect();
|
raw_messages.into_iter().map(|m| m.into()).collect();
|
||||||
|
|
||||||
// NEAR AI cloud-api does not support multi-turn tool calling (rejects
|
// Some OpenAI-compatible providers reject `role:"tool"` messages.
|
||||||
// any request containing role:"tool" messages with HTTP 400). Rewrite
|
// When enabled, rewrite tool-call / tool-result pairs into plain text.
|
||||||
// tool-call / tool-result pairs into plain text so the conversation
|
let messages = if self.flatten_tool_messages {
|
||||||
// history is preserved without using unsupported message roles.
|
flatten_tool_messages(messages)
|
||||||
let messages = flatten_tool_messages(messages);
|
} else {
|
||||||
|
messages
|
||||||
|
};
|
||||||
|
|
||||||
let tools: Vec<ChatCompletionTool> = req
|
let tools: Vec<ChatCompletionTool> = req
|
||||||
.tools
|
.tools
|
||||||
@@ -349,12 +338,14 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let (input_tokens, output_tokens) = parse_usage(response.usage.as_ref());
|
||||||
|
|
||||||
Ok(ToolCompletionResponse {
|
Ok(ToolCompletionResponse {
|
||||||
content,
|
content,
|
||||||
tool_calls,
|
tool_calls,
|
||||||
finish_reason,
|
finish_reason,
|
||||||
input_tokens: response.usage.prompt_tokens,
|
input_tokens,
|
||||||
output_tokens: response.usage.completion_tokens,
|
output_tokens,
|
||||||
response_id: None,
|
response_id: None,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -384,18 +375,25 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn active_model_name(&self) -> String {
|
fn active_model_name(&self) -> String {
|
||||||
self.active_model
|
match self.active_model.read() {
|
||||||
.read()
|
Ok(guard) => guard.clone(),
|
||||||
.expect("active_model lock poisoned")
|
Err(poisoned) => {
|
||||||
.clone()
|
tracing::warn!("active_model lock poisoned while reading; continuing");
|
||||||
|
poisoned.into_inner().clone()
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> {
|
fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> {
|
||||||
let mut guard = self
|
match self.active_model.write() {
|
||||||
.active_model
|
Ok(mut guard) => {
|
||||||
.write()
|
*guard = model.to_string();
|
||||||
.expect("active_model lock poisoned");
|
}
|
||||||
*guard = model.to_string();
|
Err(poisoned) => {
|
||||||
|
tracing::warn!("active_model lock poisoned while writing; continuing");
|
||||||
|
*poisoned.into_inner() = model.to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -545,9 +543,11 @@ struct ChatCompletionFunction {
|
|||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct ChatCompletionResponse {
|
struct ChatCompletionResponse {
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
id: String,
|
#[serde(default)]
|
||||||
|
id: Option<String>,
|
||||||
choices: Vec<ChatCompletionChoice>,
|
choices: Vec<ChatCompletionChoice>,
|
||||||
usage: ChatCompletionUsage,
|
#[serde(default)]
|
||||||
|
usage: Option<ChatCompletionUsage>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
@@ -579,18 +579,90 @@ struct ChatCompletionToolCallFunction {
|
|||||||
arguments: String,
|
arguments: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize, Default)]
|
||||||
struct ChatCompletionUsage {
|
struct ChatCompletionUsage {
|
||||||
prompt_tokens: u32,
|
#[serde(default)]
|
||||||
completion_tokens: u32,
|
prompt_tokens: Option<u64>,
|
||||||
#[allow(dead_code)]
|
#[serde(default)]
|
||||||
total_tokens: u32,
|
completion_tokens: Option<u64>,
|
||||||
|
#[serde(default)]
|
||||||
|
total_tokens: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn saturate_u32(val: u64) -> u32 {
|
||||||
|
val.min(u32::MAX as u64) as u32
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_usage(usage: Option<&ChatCompletionUsage>) -> (u32, u32) {
|
||||||
|
let Some(u) = usage else {
|
||||||
|
return (0, 0);
|
||||||
|
};
|
||||||
|
let input = u.prompt_tokens.map(saturate_u32).unwrap_or(0);
|
||||||
|
let output = u.completion_tokens.map(saturate_u32).unwrap_or_else(|| {
|
||||||
|
// Fall back to total - prompt if completion is missing.
|
||||||
|
match (u.total_tokens, u.prompt_tokens) {
|
||||||
|
(Some(total), Some(prompt)) => saturate_u32(total.saturating_sub(prompt)),
|
||||||
|
(Some(total), None) => saturate_u32(total),
|
||||||
|
_ => 0,
|
||||||
|
}
|
||||||
|
});
|
||||||
|
(input, output)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
fn test_nearai_config(base_url: &str) -> NearAiConfig {
|
||||||
|
NearAiConfig {
|
||||||
|
model: "test-model".to_string(),
|
||||||
|
base_url: base_url.to_string(),
|
||||||
|
auth_base_url: "https://private.near.ai".to_string(),
|
||||||
|
session_path: std::path::PathBuf::from("/tmp/session.json"),
|
||||||
|
api_mode: crate::config::NearAiApiMode::ChatCompletions,
|
||||||
|
api_key: Some(secrecy::SecretString::from("test-key".to_string())),
|
||||||
|
cheap_model: None,
|
||||||
|
fallback_model: None,
|
||||||
|
max_retries: 0,
|
||||||
|
circuit_breaker_threshold: None,
|
||||||
|
circuit_breaker_recovery_secs: 30,
|
||||||
|
response_cache_enabled: false,
|
||||||
|
response_cache_ttl_secs: 3600,
|
||||||
|
response_cache_max_entries: 1000,
|
||||||
|
failover_cooldown_secs: 300,
|
||||||
|
failover_cooldown_threshold: 3,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_api_url_with_base_without_v1() {
|
||||||
|
let mut cfg = test_nearai_config("http://127.0.0.1:8318");
|
||||||
|
|
||||||
|
let provider = NearAiChatProvider::new(cfg.clone()).expect("provider");
|
||||||
|
assert_eq!(
|
||||||
|
provider.api_url("chat/completions"),
|
||||||
|
"http://127.0.0.1:8318/v1/chat/completions"
|
||||||
|
);
|
||||||
|
|
||||||
|
cfg.base_url = "http://127.0.0.1:8318/".to_string();
|
||||||
|
let provider = NearAiChatProvider::new(cfg).expect("provider");
|
||||||
|
assert_eq!(
|
||||||
|
provider.api_url("/chat/completions"),
|
||||||
|
"http://127.0.0.1:8318/v1/chat/completions"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_api_url_with_base_already_v1() {
|
||||||
|
let cfg = test_nearai_config("http://127.0.0.1:8318/v1");
|
||||||
|
|
||||||
|
let provider = NearAiChatProvider::new(cfg).expect("provider");
|
||||||
|
assert_eq!(
|
||||||
|
provider.api_url("chat/completions"),
|
||||||
|
"http://127.0.0.1:8318/v1/chat/completions"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_message_conversion() {
|
fn test_message_conversion() {
|
||||||
let msg = ChatMessage::user("Hello");
|
let msg = ChatMessage::user("Hello");
|
||||||
|
|||||||
@@ -347,3 +347,124 @@ pub trait LlmProvider: Send + Sync {
|
|||||||
input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens)
|
input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Sanitize a message list to ensure tool_use / tool_result integrity.
|
||||||
|
///
|
||||||
|
/// LLM APIs (especially Anthropic) require every tool_result to reference a
|
||||||
|
/// tool_call_id that exists in an immediately preceding assistant message's
|
||||||
|
/// tool_calls. Orphaned tool_results cause HTTP 400 errors.
|
||||||
|
///
|
||||||
|
/// This function:
|
||||||
|
/// 1. Tracks all tool_call_ids emitted by assistant messages.
|
||||||
|
/// 2. Rewrites orphaned tool_result messages (whose tool_call_id has no
|
||||||
|
/// matching assistant tool_call) as user messages so the content is
|
||||||
|
/// preserved without violating the protocol.
|
||||||
|
///
|
||||||
|
/// Call this before sending messages to any LLM provider.
|
||||||
|
pub fn sanitize_tool_messages(messages: &mut [ChatMessage]) {
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
|
// Collect all tool_call_ids from assistant messages with tool_calls.
|
||||||
|
let mut known_ids: HashSet<String> = HashSet::new();
|
||||||
|
for msg in messages.iter() {
|
||||||
|
if msg.role == Role::Assistant
|
||||||
|
&& let Some(ref calls) = msg.tool_calls
|
||||||
|
{
|
||||||
|
for tc in calls {
|
||||||
|
known_ids.insert(tc.id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rewrite orphaned tool_result messages as user messages.
|
||||||
|
for msg in messages.iter_mut() {
|
||||||
|
if msg.role != Role::Tool {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let is_orphaned = match &msg.tool_call_id {
|
||||||
|
Some(id) => !known_ids.contains(id),
|
||||||
|
None => true,
|
||||||
|
};
|
||||||
|
if is_orphaned {
|
||||||
|
let tool_name = msg.name.as_deref().unwrap_or("unknown");
|
||||||
|
tracing::debug!(
|
||||||
|
tool_call_id = ?msg.tool_call_id,
|
||||||
|
tool_name,
|
||||||
|
"Rewriting orphaned tool_result as user message",
|
||||||
|
);
|
||||||
|
msg.role = Role::User;
|
||||||
|
msg.content = format!("[Tool `{}` returned: {}]", tool_name, msg.content);
|
||||||
|
msg.tool_call_id = None;
|
||||||
|
msg.name = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sanitize_preserves_valid_pairs() {
|
||||||
|
let tc = ToolCall {
|
||||||
|
id: "call_1".to_string(),
|
||||||
|
name: "echo".to_string(),
|
||||||
|
arguments: serde_json::json!({}),
|
||||||
|
};
|
||||||
|
let mut messages = vec![
|
||||||
|
ChatMessage::user("hello"),
|
||||||
|
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
|
||||||
|
ChatMessage::tool_result("call_1", "echo", "result"),
|
||||||
|
];
|
||||||
|
sanitize_tool_messages(&mut messages);
|
||||||
|
assert_eq!(messages[2].role, Role::Tool);
|
||||||
|
assert_eq!(messages[2].tool_call_id, Some("call_1".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sanitize_rewrites_orphaned_tool_result() {
|
||||||
|
let mut messages = vec![
|
||||||
|
ChatMessage::user("hello"),
|
||||||
|
ChatMessage::assistant("I'll use a tool"),
|
||||||
|
ChatMessage::tool_result("call_missing", "search", "some result"),
|
||||||
|
];
|
||||||
|
sanitize_tool_messages(&mut messages);
|
||||||
|
assert_eq!(messages[2].role, Role::User);
|
||||||
|
assert!(messages[2].content.contains("[Tool `search` returned:"));
|
||||||
|
assert!(messages[2].tool_call_id.is_none());
|
||||||
|
assert!(messages[2].name.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sanitize_handles_no_tool_messages() {
|
||||||
|
let mut messages = vec![
|
||||||
|
ChatMessage::system("prompt"),
|
||||||
|
ChatMessage::user("hello"),
|
||||||
|
ChatMessage::assistant("hi"),
|
||||||
|
];
|
||||||
|
let original_len = messages.len();
|
||||||
|
sanitize_tool_messages(&mut messages);
|
||||||
|
assert_eq!(messages.len(), original_len);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_sanitize_multiple_orphaned() {
|
||||||
|
let tc = ToolCall {
|
||||||
|
id: "call_1".to_string(),
|
||||||
|
name: "echo".to_string(),
|
||||||
|
arguments: serde_json::json!({}),
|
||||||
|
};
|
||||||
|
let mut messages = vec![
|
||||||
|
ChatMessage::user("test"),
|
||||||
|
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
|
||||||
|
ChatMessage::tool_result("call_1", "echo", "ok"),
|
||||||
|
// These are orphaned (call_2 and call_3 have no matching assistant message)
|
||||||
|
ChatMessage::tool_result("call_2", "search", "orphan 1"),
|
||||||
|
ChatMessage::tool_result("call_3", "http", "orphan 2"),
|
||||||
|
];
|
||||||
|
sanitize_tool_messages(&mut messages);
|
||||||
|
assert_eq!(messages[2].role, Role::Tool); // call_1 is valid
|
||||||
|
assert_eq!(messages[3].role, Role::User); // call_2 orphaned
|
||||||
|
assert_eq!(messages[4].role, Role::User); // call_3 orphaned
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+334
-24
@@ -1,15 +1,50 @@
|
|||||||
//! Shared retry helpers for LLM providers.
|
//! Shared retry helpers and composable `RetryProvider` decorator for LLM providers.
|
||||||
//!
|
//!
|
||||||
//! Provides exponential backoff with jitter and retryable status classification
|
//! Provides:
|
||||||
//! used by both `NearAiProvider` and `NearAiChatProvider`.
|
//! - `is_retryable()` — `LlmError`-level retryability classification (shared with `failover.rs`)
|
||||||
|
//! - `retry_backoff_delay()` — exponential backoff with jitter
|
||||||
|
//! - `RetryProvider` — decorator that wraps any `LlmProvider` with automatic retries
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
use rand::Rng;
|
use rand::Rng;
|
||||||
|
use rust_decimal::Decimal;
|
||||||
|
|
||||||
/// Returns `true` if the HTTP status code is transient and worth retrying.
|
use crate::error::LlmError;
|
||||||
pub(crate) fn is_retryable_status(status: u16) -> bool {
|
use crate::llm::provider::{
|
||||||
matches!(status, 429 | 500 | 502 | 503 | 504)
|
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
|
||||||
|
ToolCompletionResponse,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// Returns `true` if the `LlmError` is transient and the request should be retried.
|
||||||
|
///
|
||||||
|
/// Used by `RetryProvider` (retry the same provider) and `FailoverProvider`
|
||||||
|
/// (try the next provider). The question is: "could this exact same request
|
||||||
|
/// succeed if we try again?"
|
||||||
|
///
|
||||||
|
/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`,
|
||||||
|
/// `SessionRenewalFailed`, `Http`, `Io`.
|
||||||
|
///
|
||||||
|
/// Non-retryable: `AuthFailed`, `SessionExpired`, `ContextLengthExceeded`,
|
||||||
|
/// `ModelNotAvailable`, `Json`.
|
||||||
|
/// - `SessionExpired` — handled by session renewal layer, not by retry
|
||||||
|
/// - `ModelNotAvailable` — the model won't appear between attempts
|
||||||
|
/// - `Json` — a serde parse bug, not a transient failure
|
||||||
|
///
|
||||||
|
/// See also `circuit_breaker::is_transient()` which answers a different
|
||||||
|
/// question: "does this error indicate the backend is degraded?"
|
||||||
|
pub(crate) fn is_retryable(err: &LlmError) -> bool {
|
||||||
|
matches!(
|
||||||
|
err,
|
||||||
|
LlmError::RequestFailed { .. }
|
||||||
|
| LlmError::RateLimited { .. }
|
||||||
|
| LlmError::InvalidResponse { .. }
|
||||||
|
| LlmError::SessionRenewalFailed { .. }
|
||||||
|
| LlmError::Http(_)
|
||||||
|
| LlmError::Io(_)
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Calculate exponential backoff delay with random jitter.
|
/// Calculate exponential backoff delay with random jitter.
|
||||||
@@ -31,31 +66,183 @@ pub(crate) fn retry_backoff_delay(attempt: u32) -> Duration {
|
|||||||
Duration::from_millis(delay_ms)
|
Duration::from_millis(delay_ms)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Configuration for the retry decorator.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct RetryConfig {
|
||||||
|
/// Maximum number of retry attempts (not counting the initial attempt).
|
||||||
|
/// Default: 3.
|
||||||
|
pub max_retries: u32,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for RetryConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self { max_retries: 3 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Composable decorator that wraps any `LlmProvider` with automatic retries.
|
||||||
|
///
|
||||||
|
/// On transient errors, sleeps using exponential backoff and retries.
|
||||||
|
/// On non-transient errors (`AuthFailed`, `ContextLengthExceeded`, `SessionExpired`),
|
||||||
|
/// returns immediately.
|
||||||
|
///
|
||||||
|
/// Special handling for `RateLimited { retry_after }`: uses the provider-suggested
|
||||||
|
/// duration if available, otherwise falls back to standard backoff.
|
||||||
|
pub struct RetryProvider {
|
||||||
|
inner: Arc<dyn LlmProvider>,
|
||||||
|
config: RetryConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RetryProvider {
|
||||||
|
pub fn new(inner: Arc<dyn LlmProvider>, config: RetryConfig) -> Self {
|
||||||
|
Self { inner, config }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl LlmProvider for RetryProvider {
|
||||||
|
fn model_name(&self) -> &str {
|
||||||
|
self.inner.model_name()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||||
|
self.inner.cost_per_token()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
|
let mut last_error: Option<LlmError> = None;
|
||||||
|
|
||||||
|
for attempt in 0..=self.config.max_retries {
|
||||||
|
let req = request.clone();
|
||||||
|
match self.inner.complete(req).await {
|
||||||
|
Ok(resp) => return Ok(resp),
|
||||||
|
Err(err) => {
|
||||||
|
if !is_retryable(&err) || attempt == self.config.max_retries {
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
|
||||||
|
let delay = match &err {
|
||||||
|
LlmError::RateLimited {
|
||||||
|
retry_after: Some(duration),
|
||||||
|
..
|
||||||
|
} => *duration,
|
||||||
|
_ => retry_backoff_delay(attempt),
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::warn!(
|
||||||
|
provider = %self.inner.model_name(),
|
||||||
|
attempt = attempt + 1,
|
||||||
|
max_retries = self.config.max_retries,
|
||||||
|
delay_ms = delay.as_millis() as u64,
|
||||||
|
error = %err,
|
||||||
|
"Retrying after transient error"
|
||||||
|
);
|
||||||
|
|
||||||
|
last_error = Some(err);
|
||||||
|
tokio::time::sleep(delay).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Err(last_error.unwrap_or_else(|| LlmError::RequestFailed {
|
||||||
|
provider: self.inner.model_name().to_string(),
|
||||||
|
reason: "retry loop exited unexpectedly".to_string(),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn complete_with_tools(
|
||||||
|
&self,
|
||||||
|
request: ToolCompletionRequest,
|
||||||
|
) -> Result<ToolCompletionResponse, LlmError> {
|
||||||
|
let mut last_error: Option<LlmError> = None;
|
||||||
|
|
||||||
|
for attempt in 0..=self.config.max_retries {
|
||||||
|
let req = request.clone();
|
||||||
|
match self.inner.complete_with_tools(req).await {
|
||||||
|
Ok(resp) => return Ok(resp),
|
||||||
|
Err(err) => {
|
||||||
|
if !is_retryable(&err) || attempt == self.config.max_retries {
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
|
||||||
|
let delay = match &err {
|
||||||
|
LlmError::RateLimited {
|
||||||
|
retry_after: Some(duration),
|
||||||
|
..
|
||||||
|
} => *duration,
|
||||||
|
_ => retry_backoff_delay(attempt),
|
||||||
|
};
|
||||||
|
|
||||||
|
tracing::warn!(
|
||||||
|
provider = %self.inner.model_name(),
|
||||||
|
attempt = attempt + 1,
|
||||||
|
max_retries = self.config.max_retries,
|
||||||
|
delay_ms = delay.as_millis() as u64,
|
||||||
|
error = %err,
|
||||||
|
"Retrying after transient error (tools)"
|
||||||
|
);
|
||||||
|
|
||||||
|
last_error = Some(err);
|
||||||
|
tokio::time::sleep(delay).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Err(last_error.unwrap_or_else(|| LlmError::RequestFailed {
|
||||||
|
provider: self.inner.model_name().to_string(),
|
||||||
|
reason: "retry loop exited unexpectedly".to_string(),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||||
|
self.inner.list_models().await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||||
|
self.inner.model_metadata().await
|
||||||
|
}
|
||||||
|
|
||||||
|
fn active_model_name(&self) -> String {
|
||||||
|
self.inner.active_model_name()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
||||||
|
self.inner.set_model(model)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||||
|
self.inner.seed_response_chain(thread_id, response_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||||
|
self.inner.get_response_chain_id(thread_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||||
|
self.inner.calculate_cost(input_tokens, output_tokens)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
#[test]
|
use crate::testing::StubLlm;
|
||||||
fn test_is_retryable_status() {
|
|
||||||
// Transient errors should be retryable
|
|
||||||
assert!(is_retryable_status(429));
|
|
||||||
assert!(is_retryable_status(500));
|
|
||||||
assert!(is_retryable_status(502));
|
|
||||||
assert!(is_retryable_status(503));
|
|
||||||
assert!(is_retryable_status(504));
|
|
||||||
|
|
||||||
// Client errors should not be retryable
|
fn make_request() -> CompletionRequest {
|
||||||
assert!(!is_retryable_status(400));
|
CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")])
|
||||||
assert!(!is_retryable_status(401));
|
|
||||||
assert!(!is_retryable_status(403));
|
|
||||||
assert!(!is_retryable_status(404));
|
|
||||||
assert!(!is_retryable_status(422));
|
|
||||||
|
|
||||||
// Success codes should not be retryable
|
|
||||||
assert!(!is_retryable_status(200));
|
|
||||||
assert!(!is_retryable_status(201));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn make_tool_request() -> ToolCompletionRequest {
|
||||||
|
ToolCompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")], vec![])
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fast_config(max_retries: u32) -> RetryConfig {
|
||||||
|
RetryConfig { max_retries }
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- Backoff delay tests --
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_retry_backoff_delay_exponential_growth() {
|
fn test_retry_backoff_delay_exponential_growth() {
|
||||||
// Run multiple samples to verify the range, accounting for jitter
|
// Run multiple samples to verify the range, accounting for jitter
|
||||||
@@ -93,4 +280,127 @@ mod tests {
|
|||||||
let delay = retry_backoff_delay(30);
|
let delay = retry_backoff_delay(30);
|
||||||
assert!(delay.as_millis() >= 100);
|
assert!(delay.as_millis() >= 100);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// -- is_retryable() classification tests --
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_is_retryable_classification() {
|
||||||
|
// Retryable
|
||||||
|
assert!(is_retryable(&LlmError::RequestFailed {
|
||||||
|
provider: "p".into(),
|
||||||
|
reason: "err".into(),
|
||||||
|
}));
|
||||||
|
assert!(is_retryable(&LlmError::RateLimited {
|
||||||
|
provider: "p".into(),
|
||||||
|
retry_after: None,
|
||||||
|
}));
|
||||||
|
assert!(is_retryable(&LlmError::InvalidResponse {
|
||||||
|
provider: "p".into(),
|
||||||
|
reason: "bad".into(),
|
||||||
|
}));
|
||||||
|
assert!(is_retryable(&LlmError::SessionRenewalFailed {
|
||||||
|
provider: "p".into(),
|
||||||
|
reason: "timeout".into(),
|
||||||
|
}));
|
||||||
|
assert!(is_retryable(&LlmError::Io(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::ConnectionReset,
|
||||||
|
"reset"
|
||||||
|
))));
|
||||||
|
|
||||||
|
// NOT retryable
|
||||||
|
assert!(!is_retryable(&LlmError::AuthFailed {
|
||||||
|
provider: "p".into(),
|
||||||
|
}));
|
||||||
|
assert!(!is_retryable(&LlmError::SessionExpired {
|
||||||
|
provider: "p".into(),
|
||||||
|
}));
|
||||||
|
assert!(!is_retryable(&LlmError::ContextLengthExceeded {
|
||||||
|
used: 100_000,
|
||||||
|
limit: 50_000,
|
||||||
|
}));
|
||||||
|
assert!(!is_retryable(&LlmError::ModelNotAvailable {
|
||||||
|
provider: "p".into(),
|
||||||
|
model: "m".into(),
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- RetryProvider tests --
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn success_on_first_attempt() {
|
||||||
|
let stub = Arc::new(StubLlm::new("ok").with_model_name("test"));
|
||||||
|
let retry = RetryProvider::new(stub.clone(), fast_config(3));
|
||||||
|
|
||||||
|
let resp = retry.complete(make_request()).await;
|
||||||
|
assert!(resp.is_ok());
|
||||||
|
assert_eq!(resp.unwrap().content, "ok");
|
||||||
|
assert_eq!(stub.calls(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn retries_transient_errors_then_succeeds() {
|
||||||
|
// StubLlm starts failing, then we flip it to succeed.
|
||||||
|
// With max_retries=2, it will try 3 times total.
|
||||||
|
let stub = Arc::new(StubLlm::failing("test"));
|
||||||
|
let retry = RetryProvider::new(stub.clone(), fast_config(2));
|
||||||
|
|
||||||
|
// Spawn a task that flips the stub to succeed after a short delay
|
||||||
|
let stub_clone = stub.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
// Wait for at least 1 retry attempt (backoff is ~1s, so 1.5s should be enough)
|
||||||
|
tokio::time::sleep(Duration::from_millis(1500)).await;
|
||||||
|
stub_clone.set_failing(false);
|
||||||
|
});
|
||||||
|
|
||||||
|
let resp = retry.complete(make_request()).await;
|
||||||
|
assert!(resp.is_ok());
|
||||||
|
// Should have called at least twice (first fail, then succeed after flip)
|
||||||
|
assert!(stub.calls() >= 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn non_transient_error_fails_immediately() {
|
||||||
|
let stub = Arc::new(StubLlm::failing_non_transient("test"));
|
||||||
|
let retry = RetryProvider::new(stub.clone(), fast_config(3));
|
||||||
|
|
||||||
|
let err = retry.complete(make_request()).await.unwrap_err();
|
||||||
|
assert!(matches!(err, LlmError::ContextLengthExceeded { .. }));
|
||||||
|
// Should only be called once — no retries for non-transient errors
|
||||||
|
assert_eq!(stub.calls(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn exhausts_retries_then_returns_error() {
|
||||||
|
let stub = Arc::new(StubLlm::failing("test"));
|
||||||
|
// max_retries=0 means only the initial attempt, no retries
|
||||||
|
let retry = RetryProvider::new(stub.clone(), fast_config(0));
|
||||||
|
|
||||||
|
let err = retry.complete(make_request()).await.unwrap_err();
|
||||||
|
assert!(matches!(err, LlmError::RequestFailed { .. }));
|
||||||
|
assert_eq!(stub.calls(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn complete_with_tools_retries_same_as_complete() {
|
||||||
|
let stub = Arc::new(StubLlm::failing_non_transient("test"));
|
||||||
|
let retry = RetryProvider::new(stub.clone(), fast_config(3));
|
||||||
|
|
||||||
|
let err = retry
|
||||||
|
.complete_with_tools(make_tool_request())
|
||||||
|
.await
|
||||||
|
.unwrap_err();
|
||||||
|
assert!(matches!(err, LlmError::ContextLengthExceeded { .. }));
|
||||||
|
assert_eq!(stub.calls(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn passthrough_methods_delegate_to_inner() {
|
||||||
|
let stub = Arc::new(StubLlm::new("ok").with_model_name("my-model"));
|
||||||
|
let retry = RetryProvider::new(stub, fast_config(3));
|
||||||
|
|
||||||
|
assert_eq!(retry.model_name(), "my-model");
|
||||||
|
assert_eq!(retry.active_model_name(), "my-model");
|
||||||
|
assert_eq!(retry.cost_per_token(), (Decimal::ZERO, Decimal::ZERO));
|
||||||
|
assert_eq!(retry.calculate_cost(100, 50), Decimal::ZERO);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+73
-3
@@ -18,6 +18,8 @@ use serde::Serialize;
|
|||||||
use serde::de::DeserializeOwned;
|
use serde::de::DeserializeOwned;
|
||||||
use serde_json::Value as JsonValue;
|
use serde_json::Value as JsonValue;
|
||||||
|
|
||||||
|
use std::collections::HashSet;
|
||||||
|
|
||||||
use crate::error::LlmError;
|
use crate::error::LlmError;
|
||||||
use crate::llm::costs;
|
use crate::llm::costs;
|
||||||
use crate::llm::provider::{
|
use crate::llm::provider::{
|
||||||
@@ -414,7 +416,9 @@ where
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let (preamble, history) = convert_messages(&request.messages);
|
let mut messages = request.messages;
|
||||||
|
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||||
|
let (preamble, history) = convert_messages(&messages);
|
||||||
|
|
||||||
let rig_req = build_rig_request(
|
let rig_req = build_rig_request(
|
||||||
preamble,
|
preamble,
|
||||||
@@ -459,7 +463,12 @@ where
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let (preamble, history) = convert_messages(&request.messages);
|
let known_tool_names: HashSet<String> =
|
||||||
|
request.tools.iter().map(|t| t.name.clone()).collect();
|
||||||
|
|
||||||
|
let mut messages = request.messages;
|
||||||
|
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||||
|
let (preamble, history) = convert_messages(&messages);
|
||||||
let tools = convert_tools(&request.tools);
|
let tools = convert_tools(&request.tools);
|
||||||
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
||||||
|
|
||||||
@@ -481,7 +490,20 @@ where
|
|||||||
reason: e.to_string(),
|
reason: e.to_string(),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let (text, tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
let (text, mut tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||||||
|
|
||||||
|
// Normalize tool call names: some proxies prepend "proxy_" prefixes.
|
||||||
|
for tc in &mut tool_calls {
|
||||||
|
let normalized = normalize_tool_name(&tc.name, &known_tool_names);
|
||||||
|
if normalized != tc.name {
|
||||||
|
tracing::debug!(
|
||||||
|
original = %tc.name,
|
||||||
|
normalized = %normalized,
|
||||||
|
"Normalized tool call name from provider",
|
||||||
|
);
|
||||||
|
tc.name = normalized;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Ok(ToolCompletionResponse {
|
Ok(ToolCompletionResponse {
|
||||||
content: text,
|
content: text,
|
||||||
@@ -513,6 +535,25 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Normalize a tool call name returned by an OpenAI-compatible provider.
|
||||||
|
///
|
||||||
|
/// Some proxies (e.g. VibeProxy) prepend `proxy_` to tool names.
|
||||||
|
/// If the returned name doesn't match any known tool but stripping a
|
||||||
|
/// `proxy_` prefix yields a match, use the stripped version.
|
||||||
|
fn normalize_tool_name(name: &str, known_tools: &HashSet<String>) -> String {
|
||||||
|
if known_tools.contains(name) {
|
||||||
|
return name.to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(stripped) = name.strip_prefix("proxy_")
|
||||||
|
&& known_tools.contains(stripped)
|
||||||
|
{
|
||||||
|
return stripped.to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
name.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -801,4 +842,33 @@ mod tests {
|
|||||||
assert_eq!(saturate_u32(u64::MAX), u32::MAX);
|
assert_eq!(saturate_u32(u64::MAX), u32::MAX);
|
||||||
assert_eq!(saturate_u32(u32::MAX as u64), u32::MAX);
|
assert_eq!(saturate_u32(u32::MAX as u64), u32::MAX);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// -- normalize_tool_name tests --
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_normalize_tool_name_exact_match() {
|
||||||
|
let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]);
|
||||||
|
assert_eq!(normalize_tool_name("echo", &known), "echo");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_normalize_tool_name_proxy_prefix_match() {
|
||||||
|
let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]);
|
||||||
|
assert_eq!(normalize_tool_name("proxy_echo", &known), "echo");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_normalize_tool_name_proxy_prefix_no_match_kept() {
|
||||||
|
let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]);
|
||||||
|
assert_eq!(
|
||||||
|
normalize_tool_name("proxy_unknown", &known),
|
||||||
|
"proxy_unknown"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_normalize_tool_name_unknown_passthrough() {
|
||||||
|
let known = HashSet::from(["echo".to_string()]);
|
||||||
|
assert_eq!(normalize_tool_name("other_tool", &known), "other_tool");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+88
-19
@@ -26,9 +26,9 @@ use ironclaw::{
|
|||||||
hooks::{HookRegistry, bootstrap_hooks},
|
hooks::{HookRegistry, bootstrap_hooks},
|
||||||
llm::{
|
llm::{
|
||||||
CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig,
|
CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig,
|
||||||
FailoverProvider, LlmProvider, ResponseCacheConfig, SessionConfig,
|
FailoverProvider, LlmProvider, ResponseCacheConfig, RetryConfig, RetryProvider,
|
||||||
create_cheap_llm_provider, create_llm_provider, create_llm_provider_with_config,
|
SessionConfig, create_cheap_llm_provider, create_llm_provider,
|
||||||
create_session_manager,
|
create_llm_provider_with_config, create_session_manager,
|
||||||
},
|
},
|
||||||
orchestrator::{
|
orchestrator::{
|
||||||
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
|
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
|
||||||
@@ -42,7 +42,9 @@ use ironclaw::{
|
|||||||
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
|
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
|
||||||
wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools},
|
wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools},
|
||||||
},
|
},
|
||||||
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
|
workspace::{
|
||||||
|
EmbeddingProvider, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings, Workspace,
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
#[cfg(feature = "libsql")]
|
#[cfg(feature = "libsql")]
|
||||||
@@ -115,18 +117,20 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
&config.llm.nearai.base_url,
|
&config.llm.nearai.base_url,
|
||||||
session,
|
session,
|
||||||
)
|
)
|
||||||
.with_model(&config.embeddings.model, 1536),
|
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||||
|
)),
|
||||||
|
"ollama" => Some(Arc::new(
|
||||||
|
ironclaw::workspace::OllamaEmbeddings::new(
|
||||||
|
&config.embeddings.ollama_base_url,
|
||||||
|
)
|
||||||
|
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||||
)),
|
)),
|
||||||
_ => {
|
_ => {
|
||||||
if let Some(api_key) = config.embeddings.openai_api_key() {
|
if let Some(api_key) = config.embeddings.openai_api_key() {
|
||||||
let dim = match config.embeddings.model.as_str() {
|
|
||||||
"text-embedding-3-large" => 3072,
|
|
||||||
_ => 1536,
|
|
||||||
};
|
|
||||||
Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model(
|
Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model(
|
||||||
api_key,
|
api_key,
|
||||||
&config.embeddings.model,
|
&config.embeddings.model,
|
||||||
dim,
|
config.embeddings.dimension,
|
||||||
)))
|
)))
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
@@ -137,6 +141,23 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Warn if libSQL backend is used with non-1536 embedding dimension.
|
||||||
|
// libSQL schema uses F32_BLOB(1536) which cannot be altered without a
|
||||||
|
// table rebuild, so non-1536 embeddings will cause storage failures.
|
||||||
|
if config.database.backend == ironclaw::config::DatabaseBackend::LibSql
|
||||||
|
&& config.embeddings.enabled
|
||||||
|
&& config.embeddings.dimension != 1536
|
||||||
|
{
|
||||||
|
tracing::warn!(
|
||||||
|
configured_dimension = config.embeddings.dimension,
|
||||||
|
"Embedding dimension {} is not 1536. The libSQL schema uses \
|
||||||
|
F32_BLOB(1536) which requires exactly 1536 dimensions. \
|
||||||
|
Embedding storage will fail. Use PostgreSQL or set \
|
||||||
|
EMBEDDING_DIMENSION=1536.",
|
||||||
|
config.embeddings.dimension
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
// Create a Database-trait-backed workspace for the memory command
|
// Create a Database-trait-backed workspace for the memory command
|
||||||
let db: Arc<dyn ironclaw::db::Database> =
|
let db: Arc<dyn ironclaw::db::Database> =
|
||||||
ironclaw::db::connect_from_config(&config.database)
|
ironclaw::db::connect_from_config(&config.database)
|
||||||
@@ -599,6 +620,22 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
let llm = create_llm_provider(&config.llm, session.clone())?;
|
let llm = create_llm_provider(&config.llm, session.clone())?;
|
||||||
tracing::info!("LLM provider initialized: {}", llm.model_name());
|
tracing::info!("LLM provider initialized: {}", llm.model_name());
|
||||||
|
|
||||||
|
// Wrap each provider with RetryProvider for automatic retries on transient errors.
|
||||||
|
// RetryProvider sits inside FailoverProvider so each provider in the failover chain
|
||||||
|
// gets its own retry attempts before the failover moves to the next provider.
|
||||||
|
let retry_config = RetryConfig {
|
||||||
|
max_retries: config.llm.nearai.max_retries,
|
||||||
|
};
|
||||||
|
let llm: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
|
||||||
|
tracing::info!(
|
||||||
|
max_retries = retry_config.max_retries,
|
||||||
|
"LLM retry wrapper enabled"
|
||||||
|
);
|
||||||
|
Arc::new(RetryProvider::new(llm, retry_config.clone()))
|
||||||
|
} else {
|
||||||
|
llm
|
||||||
|
};
|
||||||
|
|
||||||
// Wrap in failover if a fallback model is configured
|
// Wrap in failover if a fallback model is configured
|
||||||
let llm: Arc<dyn LlmProvider> =
|
let llm: Arc<dyn LlmProvider> =
|
||||||
if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() {
|
if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() {
|
||||||
@@ -615,6 +652,12 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
fallback = %fallback.model_name(),
|
fallback = %fallback.model_name(),
|
||||||
"LLM failover enabled"
|
"LLM failover enabled"
|
||||||
);
|
);
|
||||||
|
// Wrap fallback with retry too
|
||||||
|
let fallback: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
|
||||||
|
Arc::new(RetryProvider::new(fallback, retry_config.clone()))
|
||||||
|
} else {
|
||||||
|
fallback
|
||||||
|
};
|
||||||
let cooldown_config = CooldownConfig {
|
let cooldown_config = CooldownConfig {
|
||||||
cooldown_duration: std::time::Duration::from_secs(
|
cooldown_duration: std::time::Duration::from_secs(
|
||||||
config.llm.nearai.failover_cooldown_secs,
|
config.llm.nearai.failover_cooldown_secs,
|
||||||
@@ -685,28 +728,39 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
match config.embeddings.provider.as_str() {
|
match config.embeddings.provider.as_str() {
|
||||||
"nearai" => {
|
"nearai" => {
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
"Embeddings enabled via NEAR AI (model: {})",
|
"Embeddings enabled via NEAR AI (model: {}, dim: {})",
|
||||||
config.embeddings.model
|
config.embeddings.model,
|
||||||
|
config.embeddings.dimension,
|
||||||
);
|
);
|
||||||
Some(Arc::new(
|
Some(Arc::new(
|
||||||
NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone())
|
NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone())
|
||||||
.with_model(&config.embeddings.model, 1536),
|
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
"ollama" => {
|
||||||
|
tracing::info!(
|
||||||
|
"Embeddings enabled via Ollama (model: {}, url: {}, dim: {})",
|
||||||
|
config.embeddings.model,
|
||||||
|
config.embeddings.ollama_base_url,
|
||||||
|
config.embeddings.dimension,
|
||||||
|
);
|
||||||
|
Some(Arc::new(
|
||||||
|
OllamaEmbeddings::new(&config.embeddings.ollama_base_url)
|
||||||
|
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
// Default to OpenAI for unknown providers
|
// Default to OpenAI for unknown providers
|
||||||
if let Some(api_key) = config.embeddings.openai_api_key() {
|
if let Some(api_key) = config.embeddings.openai_api_key() {
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
"Embeddings enabled via OpenAI (model: {})",
|
"Embeddings enabled via OpenAI (model: {}, dim: {})",
|
||||||
config.embeddings.model
|
config.embeddings.model,
|
||||||
|
config.embeddings.dimension,
|
||||||
);
|
);
|
||||||
Some(Arc::new(OpenAiEmbeddings::with_model(
|
Some(Arc::new(OpenAiEmbeddings::with_model(
|
||||||
api_key,
|
api_key,
|
||||||
&config.embeddings.model,
|
&config.embeddings.model,
|
||||||
match config.embeddings.model.as_str() {
|
config.embeddings.dimension,
|
||||||
"text-embedding-3-large" => 3072,
|
|
||||||
_ => 1536, // text-embedding-3-small and ada-002
|
|
||||||
},
|
|
||||||
)))
|
)))
|
||||||
} else {
|
} else {
|
||||||
tracing::warn!("Embeddings configured but OPENAI_API_KEY not set");
|
tracing::warn!("Embeddings configured but OPENAI_API_KEY not set");
|
||||||
@@ -719,6 +773,21 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Warn if libSQL backend is used with non-1536 embedding dimension.
|
||||||
|
if config.database.backend == ironclaw::config::DatabaseBackend::LibSql
|
||||||
|
&& config.embeddings.enabled
|
||||||
|
&& config.embeddings.dimension != 1536
|
||||||
|
{
|
||||||
|
tracing::warn!(
|
||||||
|
configured_dimension = config.embeddings.dimension,
|
||||||
|
"Embedding dimension {} is not 1536. The libSQL schema uses \
|
||||||
|
F32_BLOB(1536) which requires exactly 1536 dimensions. \
|
||||||
|
Embedding storage will fail. Use PostgreSQL or set \
|
||||||
|
EMBEDDING_DIMENSION=1536.",
|
||||||
|
config.embeddings.dimension
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
// Register memory tools if database is available
|
// Register memory tools if database is available
|
||||||
if let Some(ref db) = db {
|
if let Some(ref db) = db {
|
||||||
let mut workspace = Workspace::new_with_db("default", Arc::clone(db));
|
let mut workspace = Workspace::new_with_db("default", Arc::clone(db));
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ impl Default for SandboxConfig {
|
|||||||
memory_limit_mb: 2048,
|
memory_limit_mb: 2048,
|
||||||
cpu_shares: 1024,
|
cpu_shares: 1024,
|
||||||
network_allowlist: default_allowlist(),
|
network_allowlist: default_allowlist(),
|
||||||
image: "ghcr.io/nearai/sandbox:latest".to_string(),
|
image: "ironclaw-worker:latest".to_string(),
|
||||||
auto_pull_image: true,
|
auto_pull_image: true,
|
||||||
proxy_port: 0,
|
proxy_port: 0,
|
||||||
}
|
}
|
||||||
|
|||||||
+1
-1
@@ -462,7 +462,7 @@ fn default_sandbox_cpu_shares() -> u32 {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn default_sandbox_image() -> String {
|
fn default_sandbox_image() -> String {
|
||||||
"ghcr.io/nearai/sandbox:latest".to_string()
|
"ironclaw-worker:latest".to_string()
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for SandboxSettings {
|
impl Default for SandboxSettings {
|
||||||
|
|||||||
@@ -354,6 +354,123 @@ impl EmbeddingProvider for NearAiEmbeddings {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Ollama embedding provider using a local Ollama instance.
|
||||||
|
///
|
||||||
|
/// Ollama serves embedding models (e.g. `nomic-embed-text`, `mxbai-embed-large`)
|
||||||
|
/// via a REST API, typically at `http://localhost:11434`.
|
||||||
|
pub struct OllamaEmbeddings {
|
||||||
|
client: reqwest::Client,
|
||||||
|
base_url: String,
|
||||||
|
model: String,
|
||||||
|
dimension: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OllamaEmbeddings {
|
||||||
|
/// Create a new Ollama embedding provider.
|
||||||
|
///
|
||||||
|
/// Defaults to `nomic-embed-text` (768 dimensions).
|
||||||
|
pub fn new(base_url: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
client: reqwest::Client::new(),
|
||||||
|
base_url: base_url.into(),
|
||||||
|
model: "nomic-embed-text".to_string(),
|
||||||
|
dimension: 768,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Use a specific model with a given dimension.
|
||||||
|
pub fn with_model(mut self, model: impl Into<String>, dimension: usize) -> Self {
|
||||||
|
self.model = model.into();
|
||||||
|
self.dimension = dimension;
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct OllamaEmbedRequest<'a> {
|
||||||
|
model: &'a str,
|
||||||
|
input: &'a [String],
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct OllamaEmbedResponse {
|
||||||
|
embeddings: Vec<Vec<f32>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl EmbeddingProvider for OllamaEmbeddings {
|
||||||
|
fn dimension(&self) -> usize {
|
||||||
|
self.dimension
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model_name(&self) -> &str {
|
||||||
|
&self.model
|
||||||
|
}
|
||||||
|
|
||||||
|
fn max_input_length(&self) -> usize {
|
||||||
|
// Most Ollama embedding models support 8192 tokens (~32k chars)
|
||||||
|
32_000
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||||
|
if text.len() > self.max_input_length() {
|
||||||
|
return Err(EmbeddingError::TextTooLong {
|
||||||
|
length: text.len(),
|
||||||
|
max: self.max_input_length(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let embeddings = self.embed_batch(&[text.to_string()]).await?;
|
||||||
|
embeddings
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
.ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||||
|
if texts.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
|
||||||
|
let request = OllamaEmbedRequest {
|
||||||
|
model: &self.model,
|
||||||
|
input: texts,
|
||||||
|
};
|
||||||
|
|
||||||
|
let url = format!("{}/api/embed", self.base_url);
|
||||||
|
|
||||||
|
let response = self.client.post(&url).json(&request).send().await?;
|
||||||
|
|
||||||
|
let status = response.status();
|
||||||
|
|
||||||
|
if !status.is_success() {
|
||||||
|
let error_text = response.text().await.unwrap_or_default();
|
||||||
|
return Err(EmbeddingError::HttpError(format!(
|
||||||
|
"Ollama returned HTTP {}: {}",
|
||||||
|
status, error_text
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let result: OllamaEmbedResponse = response.json().await.map_err(|e| {
|
||||||
|
EmbeddingError::InvalidResponse(format!("Failed to parse Ollama response: {}", e))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Validate that returned embeddings match the configured dimension.
|
||||||
|
for (i, emb) in result.embeddings.iter().enumerate() {
|
||||||
|
if emb.len() != self.dimension {
|
||||||
|
return Err(EmbeddingError::InvalidResponse(format!(
|
||||||
|
"Ollama returned embedding of dimension {}, expected {} at index {}",
|
||||||
|
emb.len(),
|
||||||
|
self.dimension,
|
||||||
|
i
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(result.embeddings)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// A mock embedding provider for testing.
|
/// A mock embedding provider for testing.
|
||||||
///
|
///
|
||||||
/// Generates deterministic embeddings based on text hash.
|
/// Generates deterministic embeddings based on text hash.
|
||||||
|
|||||||
@@ -50,7 +50,9 @@ mod search;
|
|||||||
|
|
||||||
pub use chunker::{ChunkConfig, chunk_document};
|
pub use chunker::{ChunkConfig, chunk_document};
|
||||||
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
|
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
|
||||||
pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings};
|
pub use embeddings::{
|
||||||
|
EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings,
|
||||||
|
};
|
||||||
#[cfg(feature = "postgres")]
|
#[cfg(feature = "postgres")]
|
||||||
pub use repository::Repository;
|
pub use repository::Repository;
|
||||||
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
|
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
|
||||||
|
|||||||
Reference in New Issue
Block a user