Skip to main content

ds_api/agent/
agent_core.rs

1use std::collections::HashMap;
2
3use crate::api::ApiClient;
4use crate::conversation::{Conversation, Summarizer};
5use crate::raw::request::message::{Message, Role};
6use crate::tool_trait::Tool;
7use serde_json::Value;
8
9/// Information about a tool call requested by the model.
10///
11/// Yielded as `AgentEvent::ToolCall` when the model requests a tool invocation.
12/// At this point the tool has not yet been executed.
13#[derive(Debug, Clone)]
14pub struct ToolCallInfo {
15    pub id: String,
16    pub name: String,
17    pub args: Value,
18}
19
20/// The result of a completed tool invocation.
21///
22/// Yielded as `AgentEvent::ToolResult` after the tool has finished executing.
23#[derive(Debug, Clone)]
24pub struct ToolCallResult {
25    pub id: String,
26    pub name: String,
27    pub args: Value,
28    pub result: Value,
29}
30
31/// Events emitted by [`AgentStream`][crate::agent::AgentStream].
32///
33/// Each variant represents a distinct, self-contained event in the agent lifecycle:
34///
35/// - `Token(String)` — a text fragment from the assistant.  In streaming mode each
36///   `Token` is a single SSE delta; in non-streaming mode the full response text
37///   arrives as one `Token`.
38/// - `ToolCall(ToolCallInfo)` — the model has requested a tool invocation.  One event
39///   is emitted per call, before execution begins.
40/// - `ToolResult(ToolCallResult)` — a tool has finished executing.  One event is
41///   emitted per call, in the same order as the corresponding `ToolCall` events.
42#[derive(Debug, Clone)]
43pub enum AgentEvent {
44    Token(String),
45    ToolCall(ToolCallInfo),
46    ToolResult(ToolCallResult),
47}
48
49/// An agent that combines a [`Conversation`] with a set of callable tools.
50///
51/// Build one with the fluent builder methods, then call [`chat`][DeepseekAgent::chat]
52/// to start a turn:
53///
54/// ```no_run
55/// use ds_api::{DeepseekAgent, tool};
56/// use serde_json::{Value, json};
57///
58/// struct MyTool;
59///
60/// #[tool]
61/// impl ds_api::Tool for MyTool {
62///     async fn greet(&self, name: String) -> Value {
63///         json!({ "greeting": format!("Hello, {name}!") })
64///     }
65/// }
66///
67/// # #[tokio::main] async fn main() {
68/// let agent = DeepseekAgent::new("sk-...")
69///     .add_tool(MyTool);
70/// # }
71/// ```
72pub struct DeepseekAgent {
73    /// The conversation manages history, the API client, and context-window compression.
74    pub(crate) conversation: Conversation,
75    pub(crate) tools: Vec<Box<dyn Tool>>,
76    pub(crate) tool_index: HashMap<String, usize>,
77    /// When `true` the agent uses SSE streaming for each API turn so `Token` events
78    /// arrive incrementally.  When `false` (default) the full response is awaited.
79    pub(crate) streaming: bool,
80}
81
82impl DeepseekAgent {
83    /// Create a new agent with the given API token.
84    pub fn new(token: impl Into<String>) -> Self {
85        Self {
86            conversation: Conversation::new(ApiClient::new(token)),
87            tools: vec![],
88            tool_index: HashMap::new(),
89            streaming: false,
90        }
91    }
92
93    /// Register a tool (builder-style, supports chaining).
94    ///
95    /// The tool's protocol-level function names are indexed so incoming tool-call
96    /// requests from the model can be dispatched to the correct implementation.
97    pub fn add_tool<TT: Tool + 'static>(mut self, tool: TT) -> Self {
98        let idx = self.tools.len();
99        for raw in tool.raw_tools() {
100            self.tool_index.insert(raw.function.name.clone(), idx);
101        }
102        self.tools.push(Box::new(tool));
103        self
104    }
105
106    /// Push a user message and return an [`AgentStream`][crate::agent::AgentStream]
107    /// that drives the full agent loop (API calls + tool execution).
108    pub fn chat(mut self, user_message: &str) -> crate::agent::stream::AgentStream {
109        self.conversation.push_user_input(user_message);
110        crate::agent::stream::AgentStream::new(self)
111    }
112
113    /// Enable SSE streaming for each API turn (builder-style).
114    pub fn with_streaming(mut self) -> Self {
115        self.streaming = true;
116        self
117    }
118
119    /// Prepend a permanent system prompt to the conversation history (builder-style).
120    ///
121    /// System messages added this way are never removed by the built-in summarizers.
122    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
123        self.conversation
124            .add_message(Message::new(Role::System, &prompt.into()));
125        self
126    }
127
128    /// Replace the summarizer used for context-window management (builder-style).
129    pub fn with_summarizer(mut self, summarizer: impl Summarizer + 'static) -> Self {
130        self.conversation = self.conversation.with_summarizer(summarizer);
131        self
132    }
133}