sgr_agent/agents/

hybrid.rs

//! HybridAgent — 2-phase agent (reasoning + action).
//!
//! Phase 1: Send a minimal ReasoningTool-only FC call to get the agent's
//! reasoning about what to do next.
//! Phase 2: Send the full toolkit FC call with the reasoning as context,
//! getting back concrete tool calls.
//!
//! Inspired by Python SGRToolCallingAgent — separates "thinking" from "acting"
//! so the model doesn't get overwhelmed by a large tool set during reasoning.

use crate::agent::{Agent, AgentError, Decision};
use crate::client::LlmClient;
use crate::registry::ToolRegistry;
use crate::types::Message;

/// 2-phase hybrid agent.
pub struct HybridAgent<C: LlmClient> {
    client: C,
    system_prompt: String,
}

impl<C: LlmClient> HybridAgent<C> {
    pub fn new(client: C, system_prompt: impl Into<String>) -> Self {
        Self {
            client,
            system_prompt: system_prompt.into(),
        }
    }
}

/// Internal reasoning tool definition for phase 1.
fn reasoning_tool_def() -> crate::tool::ToolDef {
    crate::tool::ToolDef {
        name: "reasoning".to_string(),
        description: "Analyze the situation and decide what tools to use next. Describe your reasoning, the current situation, and which tools you plan to call.".to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "situation": {
                    "type": "string",
                    "description": "Your assessment of the current situation"
                },
                "plan": {
                    "type": "array",
                    "items": { "type": "string" },
                    "description": "Step-by-step plan of what to do next"
                },
                "done": {
                    "type": "boolean",
                    "description": "Set to true if the task is fully complete"
                }
            },
            "required": ["situation", "plan", "done"]
        }),
    }
}

#[async_trait::async_trait]
impl<C: LlmClient> Agent for HybridAgent<C> {
    async fn decide(
        &self,
        messages: &[Message],
        tools: &ToolRegistry,
    ) -> Result<Decision, AgentError> {
        self.decide_stateful(messages, tools, None)
            .await
            .map(|(d, _)| d)
    }

    async fn decide_stateful(
        &self,
        messages: &[Message],
        tools: &ToolRegistry,
        previous_response_id: Option<&str>,
    ) -> Result<(Decision, Option<String>), AgentError> {
        // Prepare messages with system prompt
        let mut msgs = Vec::with_capacity(messages.len() + 1);
        let has_system = messages
            .iter()
            .any(|m| m.role == crate::types::Role::System);
        if !has_system && !self.system_prompt.is_empty() {
            msgs.push(Message::system(&self.system_prompt));
        }
        msgs.extend_from_slice(messages);

        // Phase 1: Reasoning — stateless (fresh context each time)
        let reasoning_defs = vec![reasoning_tool_def()];
        let reasoning_calls = self.client.tools_call(&msgs, &reasoning_defs).await?;

        // Extract reasoning from phase 1
        let (situation, plan, done) = if let Some(rc) = reasoning_calls.first() {
            let sit = rc
                .arguments
                .get("situation")
                .and_then(|s| s.as_str())
                .unwrap_or("")
                .to_string();
            let plan: Vec<String> = rc
                .arguments
                .get("plan")
                .and_then(|p| p.as_array())
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| v.as_str().map(String::from))
                        .collect()
                })
                .unwrap_or_default();
            let done = rc
                .arguments
                .get("done")
                .and_then(|d| d.as_bool())
                .unwrap_or(false);
            (sit, plan, done)
        } else {
            return Ok((
                Decision {
                    situation: String::new(),
                    task: vec![],
                    tool_calls: vec![],
                    completed: true,
                },
                None,
            ));
        };

        // Phase 2: Action — STATEFUL (chain from previous step for token caching)
        let mut action_msgs = msgs.clone();
        let reasoning_context = if done {
            format!(
                "Reasoning: {}\nStatus: Task appears complete. Call the answer/finish tool with the final result.",
                situation
            )
        } else {
            format!("Reasoning: {}\nPlan: {}", situation, plan.join(", "))
        };
        action_msgs.push(Message::assistant(&reasoning_context));
        action_msgs.push(Message::user(
            "Now execute the next step from your plan using the available tools.",
        ));

        // Progressive tool discovery: filter tools by reasoning context.
        // Send only tools mentioned in situation/plan + answer/finish (always needed).
        let context_lower = format!("{} {}", situation, plan.join(" ")).to_lowercase();
        let filtered: Vec<_> = tools
            .to_defs()
            .into_iter()
            .filter(|t| {
                // Always include answer/finish tools
                t.name == "answer"
                    || t.name == "finish_task"
                    || t.name.contains("answer")
                    // Include if tool name appears in reasoning context
                    || context_lower.contains(&t.name.to_lowercase())
                    // Include read/write/search as core tools (almost always needed)
                    || matches!(t.name.as_str(), "read" | "write" | "search")
            })
            .collect();
        let defs = if filtered.is_empty() {
            tools.to_defs()
        } else {
            filtered
        };

        let (tool_calls, new_response_id) = self
            .client
            .tools_call_stateful(&action_msgs, &defs, previous_response_id)
            .await?;

        let completed =
            tool_calls.is_empty() || tool_calls.iter().any(|tc| tc.name == "finish_task");

        Ok((
            Decision {
                situation,
                task: plan,
                tool_calls,
                completed,
            },
            new_response_id,
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent_tool::{Tool, ToolError, ToolOutput};
    use crate::context::AgentContext;
    use crate::tool::ToolDef;
    use crate::types::{SgrError, ToolCall};
    use serde_json::Value;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// Mock client that returns reasoning in phase 1, tool call in phase 2.
    struct MockHybridClient {
        call_count: Arc<AtomicUsize>,
    }

    #[async_trait::async_trait]
    impl LlmClient for MockHybridClient {
        async fn structured_call(
            &self,
            _: &[Message],
            _: &Value,
        ) -> Result<(Option<Value>, Vec<ToolCall>, String), SgrError> {
            Ok((None, vec![], String::new()))
        }
        async fn tools_call(
            &self,
            _: &[Message],
            _tools: &[ToolDef],
        ) -> Result<Vec<ToolCall>, SgrError> {
            let n = self.call_count.fetch_add(1, Ordering::SeqCst);
            if n == 0 {
                // Phase 1: reasoning
                Ok(vec![ToolCall {
                    id: "r1".into(),
                    name: "reasoning".into(),
                    arguments: serde_json::json!({
                        "situation": "Need to read a file",
                        "plan": ["read main.rs", "analyze contents"],
                        "done": false
                    }),
                }])
            } else {
                // Phase 2: action
                Ok(vec![ToolCall {
                    id: "a1".into(),
                    name: "read_file".into(),
                    arguments: serde_json::json!({"path": "main.rs"}),
                }])
            }
        }
        async fn complete(&self, _: &[Message]) -> Result<String, SgrError> {
            Ok(String::new())
        }
    }

    struct DummyTool;
    #[async_trait::async_trait]
    impl Tool for DummyTool {
        fn name(&self) -> &str {
            "read_file"
        }
        fn description(&self) -> &str {
            "read a file"
        }
        fn parameters_schema(&self) -> Value {
            serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
        }
        async fn execute(&self, _: Value, _: &mut AgentContext) -> Result<ToolOutput, ToolError> {
            Ok(ToolOutput::text("file contents"))
        }
    }

    #[tokio::test]
    async fn hybrid_two_phases() {
        let client = MockHybridClient {
            call_count: Arc::new(AtomicUsize::new(0)),
        };
        let agent = HybridAgent::new(client, "test agent");
        let tools = ToolRegistry::new().register(DummyTool);
        let msgs = vec![Message::user("read main.rs")];

        let decision = agent.decide(&msgs, &tools).await.unwrap();
        assert_eq!(decision.situation, "Need to read a file");
        assert_eq!(decision.task.len(), 2);
        assert_eq!(decision.tool_calls.len(), 1);
        assert_eq!(decision.tool_calls[0].name, "read_file");
        assert!(!decision.completed);
    }

    #[tokio::test]
    async fn hybrid_done_still_runs_phase2() {
        // Even when reasoning says done, phase 2 runs to let the model call answer/finish
        struct DoneClient {
            call_count: Arc<AtomicUsize>,
        }
        #[async_trait::async_trait]
        impl LlmClient for DoneClient {
            async fn structured_call(
                &self,
                _: &[Message],
                _: &Value,
            ) -> Result<(Option<Value>, Vec<ToolCall>, String), SgrError> {
                Ok((None, vec![], String::new()))
            }
            async fn tools_call(
                &self,
                _: &[Message],
                _: &[ToolDef],
            ) -> Result<Vec<ToolCall>, SgrError> {
                let n = self.call_count.fetch_add(1, Ordering::SeqCst);
                if n == 0 {
                    Ok(vec![ToolCall {
                        id: "r1".into(),
                        name: "reasoning".into(),
                        arguments: serde_json::json!({
                            "situation": "Task is already complete",
                            "plan": [],
                            "done": true
                        }),
                    }])
                } else {
                    // Phase 2 — model calls finish
                    Ok(vec![ToolCall {
                        id: "a1".into(),
                        name: "finish_task".into(),
                        arguments: serde_json::json!({"summary": "done"}),
                    }])
                }
            }
            async fn complete(&self, _: &[Message]) -> Result<String, SgrError> {
                Ok(String::new())
            }
        }

        let agent = HybridAgent::new(
            DoneClient {
                call_count: Arc::new(AtomicUsize::new(0)),
            },
            "test",
        );
        let tools = ToolRegistry::new().register(DummyTool);
        let msgs = vec![Message::user("done")];

        let decision = agent.decide(&msgs, &tools).await.unwrap();
        // Phase 2 ran and returned finish_task
        assert!(decision.completed);
        assert_eq!(decision.tool_calls.len(), 1);
        assert_eq!(decision.tool_calls[0].name, "finish_task");
    }

    #[tokio::test]
    async fn hybrid_no_reasoning_completes() {
        struct EmptyClient;
        #[async_trait::async_trait]
        impl LlmClient for EmptyClient {
            async fn structured_call(
                &self,
                _: &[Message],
                _: &Value,
            ) -> Result<(Option<Value>, Vec<ToolCall>, String), SgrError> {
                Ok((None, vec![], String::new()))
            }
            async fn tools_call(
                &self,
                _: &[Message],
                _: &[ToolDef],
            ) -> Result<Vec<ToolCall>, SgrError> {
                Ok(vec![])
            }
            async fn complete(&self, _: &[Message]) -> Result<String, SgrError> {
                Ok(String::new())
            }
        }

        let agent = HybridAgent::new(EmptyClient, "test");
        let tools = ToolRegistry::new().register(DummyTool);
        let msgs = vec![Message::user("hello")];

        let decision = agent.decide(&msgs, &tools).await.unwrap();
        assert!(decision.completed);
    }

    #[tokio::test]
    async fn hybrid_two_phases_independent() {
        // Verify that phase 1 and phase 2 don't share state:
        // Both calls use tools_call (stateless), so they are independent.
        // The mock tracks call order and verifies each phase gets separate invocations.
        struct PhaseTrackingClient {
            call_count: Arc<AtomicUsize>,
        }

        #[async_trait::async_trait]
        impl LlmClient for PhaseTrackingClient {
            async fn structured_call(
                &self,
                _: &[Message],
                _: &Value,
            ) -> Result<(Option<Value>, Vec<ToolCall>, String), SgrError> {
                Ok((None, vec![], String::new()))
            }
            async fn tools_call(
                &self,
                msgs: &[Message],
                tools: &[ToolDef],
            ) -> Result<Vec<ToolCall>, SgrError> {
                let n = self.call_count.fetch_add(1, Ordering::SeqCst);
                if n == 0 {
                    // Phase 1: reasoning — only gets reasoning tool
                    assert_eq!(tools.len(), 1, "Phase 1 should only have reasoning tool");
                    assert_eq!(tools[0].name, "reasoning");
                    Ok(vec![ToolCall {
                        id: "r1".into(),
                        name: "reasoning".into(),
                        arguments: serde_json::json!({
                            "situation": "Testing phase independence",
                            "plan": ["call read_file"],
                            "done": false
                        }),
                    }])
                } else {
                    // Phase 2: action — gets full tool registry
                    assert!(
                        tools.len() > 1 || tools[0].name != "reasoning",
                        "Phase 2 should have the real tools, not just reasoning"
                    );
                    // Verify that messages don't contain any implicit state from phase 1
                    // (they will have reasoning context added explicitly as assistant message)
                    let last_msg = msgs.last().unwrap();
                    assert_eq!(
                        last_msg.role,
                        crate::types::Role::User,
                        "Last message in phase 2 should be the action prompt"
                    );
                    Ok(vec![ToolCall {
                        id: "a1".into(),
                        name: "read_file".into(),
                        arguments: serde_json::json!({"path": "test.rs"}),
                    }])
                }
            }
            async fn complete(&self, _: &[Message]) -> Result<String, SgrError> {
                Ok(String::new())
            }
        }

        let call_count = Arc::new(AtomicUsize::new(0));
        let agent = HybridAgent::new(
            PhaseTrackingClient {
                call_count: call_count.clone(),
            },
            "test agent",
        );
        let tools = ToolRegistry::new().register(DummyTool);
        let msgs = vec![Message::user("read test.rs")];

        let decision = agent.decide(&msgs, &tools).await.unwrap();

        // Both phases ran
        assert_eq!(call_count.load(Ordering::SeqCst), 2);
        // Phase 2 returned the action
        assert_eq!(decision.tool_calls.len(), 1);
        assert_eq!(decision.tool_calls[0].name, "read_file");
        assert!(!decision.completed);
    }
}