johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::{json, Value};

use crate::backend::{
    AgentBackend, BackendError, BackendEvent, BackendStream, RunOutcome, RunRequest,
};
use crate::message::{AgentContent, AgentMessage, Role, RunStatus};
use crate::tools::{content_for_log, ToolDefinition, ToolError, ToolExecutor, ToolObserver};

const MAX_TOOL_ROUNDS: usize = 8;
const TOOL_BUDGET_EXHAUSTED: &str = "Tool budget exhausted: no more tool calls are possible. \
Answer now using only the results you already have, and say briefly what you could not determine.";

/// Backend for OpenAI-compatible `/v1/chat/completions` endpoints.
///
/// Supports a bounded function-calling loop. Tools come from an optional
/// `ToolExecutor` (for example, MCP servers over stdio).
#[derive(Clone)]
pub struct OpenAiCompatibleBackend {
    client: reqwest::Client,
    base_url: String,
    model: String,
    api_key: Option<String>,
    tool_executor: Option<Arc<dyn ToolExecutor>>,
    tool_observer: Option<Arc<dyn ToolObserver>>,
    verbose_logging: bool,
    /// Cap on generated tokens per request (`max_tokens`).
    max_output_tokens: u32,
    /// Wall-clock limit per model request.
    request_timeout: std::time::Duration,
}

/// Default `max_tokens`: generous for code-mode programs and chat replies,
/// small enough that a degenerate generation ends (live: a local model
/// decoded 71 486 tokens on one request with no cap).
pub const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 8192;
/// Default per-request wall-clock limit.
pub const DEFAULT_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);

impl OpenAiCompatibleBackend {
    pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            base_url: base_url.into(),
            model: model.into(),
            api_key: None,
            tool_executor: None,
            tool_observer: None,
            verbose_logging: false,
            max_output_tokens: DEFAULT_MAX_OUTPUT_TOKENS,
            request_timeout: DEFAULT_REQUEST_TIMEOUT,
        }
    }

    pub fn with_max_output_tokens(mut self, tokens: u32) -> Self {
        self.max_output_tokens = tokens;
        self
    }

    pub fn with_request_timeout(mut self, timeout: std::time::Duration) -> Self {
        self.request_timeout = timeout;
        self
    }

    pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
        self.api_key = Some(api_key.into());
        self
    }

    pub fn with_tool_executor(mut self, executor: Arc<dyn ToolExecutor>) -> Self {
        self.tool_executor = Some(executor);
        self
    }

    pub fn with_tool_observer(mut self, observer: Arc<dyn ToolObserver>) -> Self {
        self.tool_observer = Some(observer);
        self
    }

    pub fn with_verbose_logging(mut self, verbose: bool) -> Self {
        self.verbose_logging = verbose;
        self
    }

    async fn post_chat(&self, messages: &[Value], tools: &[Value]) -> Result<Value, BackendError> {
        let url = format!(
            "{}/v1/chat/completions",
            self.base_url.trim_end_matches('/')
        );
        let mut payload = json!({
            "model": self.model,
            "messages": messages,
            "stream": false,
            "max_tokens": self.max_output_tokens,
        });
        if !tools.is_empty() {
            payload["tools"] = Value::Array(tools.to_vec());
            payload["tool_choice"] = json!("auto");
        }

        let mut builder = self
            .client
            .post(url)
            .timeout(self.request_timeout)
            .json(&payload);
        if let Some(api_key) = &self.api_key {
            builder = builder.bearer_auth(api_key);
        }

        let response = builder.send().await.map_err(|error| {
            if error.is_timeout() {
                BackendError::new(format!(
                    "model request timed out after {:?}",
                    self.request_timeout
                ))
            } else {
                BackendError::new(format!("request failed: {error}"))
            }
        })?;
        let status = response.status();
        let body_text = response
            .text()
            .await
            .map_err(|error| BackendError::new(format!("failed to read response: {error}")))?;

        if !status.is_success() {
            return Err(BackendError::new(format!("HTTP {status}: {body_text}")));
        }

        serde_json::from_str(&body_text).map_err(|error| {
            BackendError::new(format!("invalid JSON response: {error}: {body_text}"))
        })
    }
}

fn role_name(role: &Role) -> &'static str {
    match role {
        Role::User => "user",
        Role::Assistant => "assistant",
        Role::System => "system",
    }
}

fn tool_specs(tools: &[ToolDefinition]) -> Vec<Value> {
    tools
        .iter()
        .map(|tool| {
            let mut function = json!({
                "name": tool.name,
                "parameters": tool.input_schema,
            });
            if let Some(description) = &tool.description {
                function["description"] = json!(description);
            }
            json!({ "type": "function", "function": function })
        })
        .collect()
}

const NO_REPLY: &str = "NO_REPLY";

/// A tool call whose only content is the `NO_REPLY` token (quoted or not):
/// `run_code(code="NO_REPLY")`.
fn is_no_reply_call(call: &Value) -> bool {
    let Some(function) = call.get("function") else {
        return false;
    };
    let arguments = match function.get("arguments") {
        Some(Value::String(text)) => serde_json::from_str::<Value>(text).unwrap_or(Value::Null),
        Some(other) => other.clone(),
        None => return false,
    };
    arguments
        .get("code")
        .and_then(Value::as_str)
        .map(|code| {
            code.trim()
                .trim_matches(|c| matches!(c, '"' | '\'' | '`'))
                .trim()
        })
        .is_some_and(|code| code.eq_ignore_ascii_case(NO_REPLY))
}

/// Adds `note` to the newest message when it is a user message.
fn append_to_last_user_message(conversation: &mut [Value], note: &str) {
    if let Some(last) = conversation.last_mut() {
        if last["role"] == "user" {
            if let Some(content) = last["content"].as_str() {
                last["content"] = json!(format!("{content}\n\n{note}"));
            }
        }
    }
}

fn parse_tool_arguments(value: &Value) -> Result<Value, String> {
    match value {
        Value::String(text) => serde_json::from_str(text)
            .map_err(|error| format!("invalid tool arguments JSON: {error}")),
        Value::Null => Ok(Value::Null),
        other => Ok(other.clone()),
    }
}

#[async_trait]
impl AgentBackend for OpenAiCompatibleBackend {
    async fn start_run(&self, request: RunRequest) -> Result<BackendStream, BackendError> {
        let mut conversation = Vec::new();
        if !request.agent.instructions.is_empty() {
            conversation.push(json!({
                "role": "system",
                "content": request.agent.instructions,
            }));
        }
        for message in &request.messages {
            let content = match &message.content {
                AgentContent::Text(text) => text.clone(),
                AgentContent::Json(value) => value.to_string(),
            };
            conversation.push(json!({
                "role": role_name(&message.role),
                "content": content,
            }));
        }

        let tool_executor = self.tool_executor.as_ref().map(|executor| {
            let scope = match &request.thread_id {
                Some(thread_id) => executor.thread_scope(thread_id),
                None => executor.run_scope(),
            };
            scope.unwrap_or_else(|| executor.clone())
        });
        let tools = match &tool_executor {
            Some(executor) => executor
                .list_tools()
                .await
                .map_err(|error| BackendError::new(format!("MCP tools/list failed: {error}")))?,
            None => Vec::new(),
        };
        let tools = tool_specs(&tools);
        if let Some(note) = tool_executor
            .as_ref()
            .and_then(|executor| executor.context_note())
        {
            append_to_last_user_message(&mut conversation, &note);
        }
        let mut summary = TurnSummary::default();

        for round in 0..=MAX_TOOL_ROUNDS {
            summary.rounds += 1;
            // Last round: withdraw the tools so the model must answer with
            // what it has instead of failing the run with no reply. The note
            // is a user message: many chat templates reject a system message
            // after the first position.
            let round_tools: &[Value] = if round == MAX_TOOL_ROUNDS && !tools.is_empty() {
                conversation.push(json!({
                    "role": "user",
                    "content": TOOL_BUDGET_EXHAUSTED,
                }));
                summary.final_round_forced = true;
                &[]
            } else {
                &tools
            };
            let response = self.post_chat(&conversation, round_tools).await?;
            if response
                .pointer("/choices/0/finish_reason")
                .and_then(Value::as_str)
                == Some("length")
            {
                crate::log::emit(
                    "warn",
                    "model.truncated",
                    json!({"round": round, "max_tokens": self.max_output_tokens}),
                );
            }
            let message = response
                .get("choices")
                .and_then(Value::as_array)
                .and_then(|choices| choices.first())
                .and_then(|choice| choice.get("message"))
                .ok_or_else(|| {
                    BackendError::new(format!("response missing choices[0].message: {response}"))
                })?;

            let mut content = message
                .get("content")
                .and_then(Value::as_str)
                .unwrap_or("")
                .to_string();
            let mut tool_calls = message
                .get("tool_calls")
                .and_then(Value::as_array)
                .cloned()
                .unwrap_or_default();
            // Live: models sometimes "say" NO_REPLY by running it as code
            // (`run_code("NO_REPLY")`), which raises NameError and costs
            // rounds. A call that is only the token ends the turn with it.
            if !tool_calls.is_empty() && tool_calls.iter().all(is_no_reply_call) {
                crate::log::emit("info", "model.no_reply_tool_call", json!({"round": round}));
                content = NO_REPLY.to_string();
                tool_calls.clear();
            }

            if !tool_calls.is_empty() {
                if round == MAX_TOOL_ROUNDS {
                    return Err(BackendError::new("tool call round limit exceeded"));
                }
                let executor = tool_executor.as_ref().ok_or_else(|| {
                    BackendError::new("model requested a tool but no tool executor is configured")
                })?;

                conversation.push(json!({
                    "role": "assistant",
                    "content": if content.is_empty() { Value::Null } else { json!(content) },
                    "tool_calls": tool_calls,
                }));

                for call in &tool_calls {
                    let call_id = call
                        .get("id")
                        .and_then(Value::as_str)
                        .unwrap_or("tool-call");
                    let name = call
                        .get("function")
                        .and_then(|function| function.get("name"))
                        .and_then(Value::as_str)
                        .unwrap_or_default();
                    let raw_arguments = call
                        .get("function")
                        .and_then(|function| function.get("arguments"))
                        .cloned()
                        .unwrap_or(Value::Null);

                    let transcript_call_id = if call_id.is_empty() || call_id == "tool-call" {
                        format!("call_{}", uuid::Uuid::new_v4())
                    } else {
                        call_id.to_string()
                    };
                    let parsed_arguments = parse_tool_arguments(&raw_arguments);
                    let observed = parsed_arguments
                        .clone()
                        .unwrap_or_else(|_| raw_arguments.clone());
                    if self.verbose_logging {
                        eprintln!(
                            "{}",
                            json!({
                                "level": "info",
                                "event": "model.tool_call",
                                "tool": name,
                                "arguments": observed,
                            })
                        );
                    }
                    if let Some(observer) = &self.tool_observer {
                        observer.tool_call(&transcript_call_id, name, &observed);
                    }

                    let outcome = match parsed_arguments {
                        Ok(arguments) => executor.call_tool(name, arguments).await,
                        Err(error) => Err(ToolError::new(error)),
                    };
                    let (tool_content, is_error, audit) = match outcome {
                        Ok(output) => (output.content, output.is_error, output.audit),
                        Err(error) => (format!("tool error: {error}"), true, None),
                    };
                    summary.record_tool(&tool_content, audit.as_ref());
                    if self.verbose_logging {
                        eprintln!(
                            "{}",
                            json!({
                                "level": if is_error { "error" } else { "info" },
                                "event": "model.tool_result",
                                "tool": name,
                                "is_error": is_error,
                                "content": content_for_log(&tool_content, 500),
                            })
                        );
                    }
                    if let Some(observer) = &self.tool_observer {
                        observer.tool_result(
                            &transcript_call_id,
                            &tool_content,
                            is_error,
                            audit.as_ref(),
                        );
                    }

                    conversation.push(json!({
                        "role": "tool",
                        "tool_call_id": call_id,
                        "content": tool_content,
                    }));
                }
                continue;
            }

            if self.verbose_logging {
                eprintln!("{}", summary.to_log());
            }
            let assistant = AgentMessage {
                role: Role::Assistant,
                content: AgentContent::Text(content),
            };
            let mut messages = request.messages;
            messages.push(assistant.clone());

            let final_output = assistant.content.clone();
            let events = vec![
                Ok(BackendEvent::Message(assistant)),
                Ok(BackendEvent::Completed(RunOutcome {
                    status: RunStatus::Completed,
                    messages,
                    final_output: Some(final_output),
                    last_error: None,
                })),
            ];

            return Ok(Box::pin(futures::stream::iter(events)));
        }

        Err(BackendError::new("tool call round limit exceeded"))
    }
}

/// Per-run tool metrics, logged as `model.turn_summary` with `--verbose`.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TurnSummary {
    pub rounds: usize,
    pub tool_calls: usize,
    pub nested_calls: usize,
    /// MCP result bytes received inside the sandbox.
    pub bytes_in: usize,
    /// Tool-result bytes sent to the model.
    pub bytes_out: usize,
    pub final_round_forced: bool,
}

impl TurnSummary {
    pub fn record_tool(&mut self, content: &str, audit: Option<&Value>) {
        self.tool_calls += 1;
        self.bytes_out += content.len();
        let nested = audit
            .and_then(|audit| audit.get("nested_calls"))
            .and_then(Value::as_array);
        for call in nested.into_iter().flatten() {
            self.nested_calls += 1;
            self.bytes_in += call.get("bytes").and_then(Value::as_u64).unwrap_or(0) as usize;
        }
    }

    pub fn to_log(&self) -> Value {
        json!({
            "level": "info",
            "event": "model.turn_summary",
            "rounds": self.rounds,
            "run_code_calls": self.tool_calls,
            "nested_calls": self.nested_calls,
            "bytes_in": self.bytes_in,
            "bytes_out": self.bytes_out,
            "final_round_forced": self.final_round_forced,
        })
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn no_reply_as_code_is_recognized() {
        let call = |code: &str| serde_json::json!({"function": {"name": "run_code", "arguments": serde_json::json!({"code": code}).to_string()}});
        assert!(super::is_no_reply_call(&call("NO_REPLY")));
        assert!(super::is_no_reply_call(&call("\"NO_REPLY\"\n")));
        assert!(!super::is_no_reply_call(&call("print('x')\nNO_REPLY")));
        assert!(!super::is_no_reply_call(&call("x = 1")));
    }

    #[test]
    fn context_note_goes_on_the_newest_user_message_only() {
        let mut conversation = vec![
            serde_json::json!({"role": "system", "content": "sys"}),
            serde_json::json!({"role": "user", "content": "hi"}),
        ];
        super::append_to_last_user_message(&mut conversation, "[vars]");
        assert_eq!(conversation[1]["content"], "hi\n\n[vars]");
        assert_eq!(conversation[0]["content"], "sys");

        let mut assistant_last = vec![serde_json::json!({"role": "assistant", "content": "a"})];
        super::append_to_last_user_message(&mut assistant_last, "[vars]");
        assert_eq!(assistant_last[0]["content"], "a");
    }

    #[test]
    fn turn_summary_counts_nested_bytes() {
        let mut summary = TurnSummary {
            rounds: 2,
            ..TurnSummary::default()
        };
        summary.record_tool(
            "small",
            Some(&json!({"nested_calls": [{"tool": "a", "bytes": 1000}, {"tool": "b", "bytes": 24}]})),
        );
        summary.record_tool("err", None);
        let log = summary.to_log();
        assert_eq!(log["event"], "model.turn_summary");
        assert_eq!(log["run_code_calls"], 2);
        assert_eq!(log["nested_calls"], 2);
        assert_eq!(log["bytes_in"], 1024);
        assert_eq!(log["bytes_out"], 8);
        assert_eq!(log["final_round_forced"], false);
    }

    use super::*;

    #[test]
    fn tool_log_content_decodes_json_and_truncates_text() {
        let decoded = content_for_log(r#"{"kcps":[],"errors":[]}"#, 500);
        assert_eq!(decoded["kcps"], json!([]));
        assert_eq!(decoded["errors"], json!([]));

        let truncated = content_for_log("plain text", 4);
        assert_eq!(truncated, json!("plai...[truncated]"));
    }
}