Skip to main content

ds_api/agent/
agent_core.rs

1use std::collections::HashMap;
2
3use crate::api::ApiClient;
4use crate::conversation::{Conversation, DeepseekConversation, Summarizer};
5use crate::raw::request::message::{Message, Role};
6use crate::tool_trait::Tool;
7use serde_json::Value;
8
9/// Tool call event (result).
10///
11/// Represents a single tool invocation result produced by the agent.
12#[derive(Debug, Clone)]
13pub struct ToolCallEvent {
14    pub id: String,
15    pub name: String,
16    pub args: Value,
17    pub result: Value,
18}
19
20/// Single agent response exposed to callers.
21///
22/// May contain assistant text content or a list of tool call events.
23#[derive(Debug, Clone)]
24pub struct AgentResponse {
25    pub content: Option<String>,
26    pub tool_calls: Vec<ToolCallEvent>,
27}
28
29/// DeepseekAgent: encapsulates a conversation (`Conversation`) and a collection of tools.
30///
31/// Responsible for coordinating API calls and executing tools as requested by the model.
32/// Fields use `pub(crate)` visibility so the sibling `stream` submodule can access internal
33/// state (for example `tools` and `tool_index`).
34pub struct DeepseekAgent {
35    pub(crate) client: ApiClient,
36    pub(crate) conversation: DeepseekConversation,
37    pub(crate) tools: Vec<Box<dyn Tool>>,
38    pub(crate) tool_index: HashMap<String, usize>,
39    /// If true, the agent uses streaming for each API turn and yields text fragments
40    /// one by one as they arrive.  If false (default), it waits for the full response.
41    pub(crate) streaming: bool,
42}
43
44impl DeepseekAgent {
45    /// Create an Agent using the provided token.
46    ///
47    /// This internally constructs an `ApiClient` and attaches a `DeepseekConversation`.
48    pub fn new(token: impl Into<String>) -> Self {
49        let client = ApiClient::new(token.into());
50        let conversation = DeepseekConversation::new(client.clone());
51        Self {
52            client,
53            conversation,
54            tools: vec![],
55            tool_index: HashMap::new(),
56            streaming: false,
57        }
58    }
59
60    /// Add a tool (supports method chaining).
61    ///
62    /// Registers the tool and indexes its raw (protocol) function names so incoming
63    /// tool call requests can be routed to the correct implementation.
64    pub fn add_tool<TT: Tool + 'static>(mut self, tool: TT) -> Self {
65        let idx = self.tools.len();
66        for raw in tool.raw_tools() {
67            // raw.function.name is the protocol-level name used to match tool call requests
68            self.tool_index.insert(raw.function.name.clone(), idx);
69        }
70        self.tools.push(Box::new(tool));
71        self
72    }
73
74    /// Push a user message into the conversation and return an `AgentStream`.
75    ///
76    /// The returned `AgentStream` drives the API request and any subsequent tool execution.
77    /// The return type uses a fully-qualified path and depends on the sibling `stream`
78    /// submodule providing `AgentStream`.
79    pub fn chat(mut self, user_message: &str) -> crate::agent::stream::AgentStream {
80        self.conversation.push_user_input(user_message.to_string());
81        crate::agent::stream::AgentStream::new(self)
82    }
83
84    /// Enable streaming text output for each API turn.
85    ///
86    /// When set, the agent uses `send_stream` internally so text fragments are
87    /// yielded to the caller as they arrive instead of waiting for a full response.
88    pub fn with_streaming(mut self) -> Self {
89        self.streaming = true;
90        self
91    }
92
93    /// Set a custom system prompt to inject at the start of the conversation.
94    ///
95    /// Builder-style: returns `self` so the call can be chained.
96    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
97        let p = prompt.into();
98        // Add a system message to the conversation history
99        self.conversation
100            .add_message(Message::new(Role::System, p.as_str()));
101        self
102    }
103
104    /// Set a summarizer to use for conversation summarization.
105    ///
106    /// Builder-style: returns `self` so the call can be chained.
107    pub fn with_summarizer(mut self, summarizer: impl Summarizer + 'static) -> Self {
108        self.conversation = self.conversation.with_summarizer(summarizer);
109        self
110    }
111}