tiny_agent_rs/core/
agent.rs

1use crate::{
2    error::{AgentError, Result},
3    schemas::{CompletionSchema, SchemaHandle},
4    services::openai_client::OpenAIClient,
5    tools::FunctionFactory,
6};
7use serde_json::{json, Value};
8use std::time::Duration;
9
10/// Main agent
11#[derive(Debug)]
12pub struct Agent {
13    openai_client: OpenAIClient,
14    function_factory: FunctionFactory,
15    model: String,
16    max_iterations: usize,
17    max_tokens: Option<u32>,
18    timeout: Duration,
19    completion_schema: Option<SchemaHandle>,
20}
21
22impl Agent {
23    pub fn new(api_key: String, function_factory: FunctionFactory) -> Self {
24        Self {
25            openai_client: OpenAIClient::new(api_key),
26            function_factory,
27            model: "openai/gpt-4.1-mini".to_string(),
28            max_iterations: 10,
29            max_tokens: Some(1000),
30            timeout: Duration::from_secs(120),
31            completion_schema: None,
32        }
33    }
34
35    pub fn with_model(mut self, model: impl Into<String>) -> Self {
36        self.model = model.into();
37        self
38    }
39
40    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
41        self.openai_client.set_base_url(base_url);
42        self
43    }
44
45    pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
46        self.max_iterations = max_iterations;
47        self
48    }
49
50    pub fn with_max_tokens(mut self, max_tokens: Option<u32>) -> Self {
51        self.max_tokens = max_tokens;
52        self
53    }
54
55    pub fn with_timeout(mut self, timeout: Duration) -> Self {
56        self.timeout = timeout;
57        self
58    }
59
60    pub fn with_completion_schema<T: CompletionSchema>(mut self) -> Self {
61        self.completion_schema = Some(T::schema().clone());
62        self
63    }
64
65    pub(crate) fn max_iterations(&self) -> usize {
66        self.max_iterations
67    }
68
69    pub(crate) fn completion_schema(&self) -> Option<&SchemaHandle> {
70        self.completion_schema.as_ref()
71    }
72
73    pub(crate) fn function_factory(&self) -> &FunctionFactory {
74        &self.function_factory
75    }
76
77    pub(crate) fn model(&self) -> &str {
78        &self.model
79    }
80
81    pub(crate) fn max_tokens(&self) -> Option<u32> {
82        self.max_tokens
83    }
84
85    pub(crate) fn timeout(&self) -> Duration {
86        self.timeout
87    }
88
89    pub fn clear_completion_schema(mut self) -> Self {
90        self.completion_schema = None;
91        self
92    }
93
94    pub async fn run(&self, prompt: &str) -> Result<String> {
95        let messages = vec![
96            json!({
97                "role": "system",
98                "content": "You are a helpful assistant with access to tools. Use tools when necessary to provide accurate information. Be concise and helpful. When you are ready to give the final response, you MUST call the `final_answer` tool with an `answer` string instead of replying directly."
99            }),
100            json!({
101                "role": "user",
102                "content": prompt
103            }),
104        ];
105
106        self.run_with_messages(messages).await
107    }
108
109    pub(crate) async fn make_raw_request(&self, request_body: &Value) -> Result<Value> {
110        self.openai_client
111            .chat_completion(request_body, self.timeout)
112            .await
113    }
114
115    pub fn from_env() -> Result<Self> {
116        let api_key = std::env::var("OPENAI_API_KEY").map_err(|_| {
117            AgentError::Config(
118                "OPENAI_API_KEY environment variable must be set before creating an Agent"
119                    .to_string(),
120            )
121        })?;
122        let function_factory = FunctionFactory::new();
123        let mut agent = Self::new(api_key, function_factory);
124        if let Ok(base_url) =
125            std::env::var("OPENAI_BASE_URL").or_else(|_| std::env::var("OPENROUTER_BASE_URL"))
126        {
127            agent.openai_client.set_base_url(base_url);
128        }
129        Ok(agent)
130    }
131}
tiny_agent_rs/core/agent.rs

tiny_agent_rs/core/
agent.rs
