Skip to main content

max / makenotwork

4.8 KB · 189 lines History Blame Raw
1 //! Ollama [`InferenceProvider`].
2 //!
3 //! Talks to a local Ollama server over its `/api/chat` endpoint and
4 //! translates the neutral [`Message`] / [`ToolSpec`] types to and from
5 //! Ollama's OpenAI-shaped wire format.
6
7 use async_trait::async_trait;
8 use serde::{Deserialize, Serialize};
9 use serde_json::{Value, json};
10
11 use crate::agent::{InferenceProvider, Message, Role, ToolCall};
12 use crate::error::Result;
13 use crate::tool::ToolSpec;
14
15 const DEFAULT_OLLAMA_URL: &str = "http://localhost:11434";
16
17 pub struct OllamaProvider {
18 client: reqwest::Client,
19 base_url: String,
20 model: String,
21 }
22
23 impl OllamaProvider {
24 pub fn new(model: impl Into<String>) -> Self {
25 Self {
26 client: reqwest::Client::new(),
27 base_url: DEFAULT_OLLAMA_URL.to_string(),
28 model: model.into(),
29 }
30 }
31
32 #[must_use]
33 pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
34 self.base_url = url.into();
35 self
36 }
37 }
38
39 #[async_trait]
40 impl InferenceProvider for OllamaProvider {
41 async fn chat(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Message> {
42 let wire_messages: Vec<WireMessage> =
43 messages.iter().map(WireMessage::from_neutral).collect();
44 let wire_tools: Vec<WireTool> = tools
45 .iter()
46 .map(|s| WireTool {
47 r#type: "function",
48 function: WireFunction {
49 name: s.name.clone(),
50 description: s.description.clone(),
51 parameters: s.input_schema.clone(),
52 },
53 })
54 .collect();
55
56 let req = ChatRequest {
57 model: &self.model,
58 messages: &wire_messages,
59 tools: &wire_tools,
60 stream: false,
61 };
62
63 let resp: ChatResponse = self
64 .client
65 .post(format!("{}/api/chat", self.base_url))
66 .json(&req)
67 .send()
68 .await?
69 .error_for_status()?
70 .json()
71 .await?;
72
73 Ok(resp.message.into_neutral())
74 }
75 }
76
77 // ---- wire types ----
78
79 #[derive(Debug, Serialize)]
80 struct ChatRequest<'a> {
81 model: &'a str,
82 messages: &'a [WireMessage],
83 tools: &'a [WireTool],
84 stream: bool,
85 }
86
87 #[derive(Debug, Deserialize)]
88 struct ChatResponse {
89 message: WireMessage,
90 }
91
92 #[derive(Debug, Serialize, Deserialize)]
93 struct WireMessage {
94 role: String,
95 #[serde(default, skip_serializing_if = "Option::is_none")]
96 content: Option<String>,
97 #[serde(default, skip_serializing_if = "Option::is_none")]
98 tool_calls: Option<Vec<WireToolCall>>,
99 #[serde(default, skip_serializing_if = "Option::is_none")]
100 tool_name: Option<String>,
101 }
102
103 impl WireMessage {
104 fn from_neutral(m: &Message) -> Self {
105 let role = match m.role {
106 Role::System => "system",
107 Role::User => "user",
108 Role::Assistant => "assistant",
109 Role::Tool => "tool",
110 }
111 .to_string();
112 let tool_calls = if m.tool_calls.is_empty() {
113 None
114 } else {
115 Some(
116 m.tool_calls
117 .iter()
118 .map(|c| WireToolCall {
119 function: WireToolCallFunction {
120 name: c.name.clone(),
121 arguments: c.arguments.clone(),
122 },
123 })
124 .collect(),
125 )
126 };
127 Self {
128 role,
129 content: m.content.clone(),
130 tool_calls,
131 tool_name: m.tool_name.clone(),
132 }
133 }
134
135 fn into_neutral(self) -> Message {
136 let role = match self.role.as_str() {
137 "system" => Role::System,
138 "user" => Role::User,
139 "tool" => Role::Tool,
140 _ => Role::Assistant,
141 };
142 let tool_calls = self
143 .tool_calls
144 .unwrap_or_default()
145 .into_iter()
146 .map(|c| ToolCall {
147 id: String::new(),
148 name: c.function.name,
149 arguments: c.function.arguments,
150 })
151 .collect();
152 Message {
153 role,
154 content: self.content,
155 tool_calls,
156 tool_name: self.tool_name,
157 }
158 }
159 }
160
161 #[derive(Debug, Serialize, Deserialize)]
162 struct WireToolCall {
163 function: WireToolCallFunction,
164 }
165
166 #[derive(Debug, Serialize, Deserialize)]
167 struct WireToolCallFunction {
168 name: String,
169 #[serde(default = "empty_object")]
170 arguments: Value,
171 }
172
173 fn empty_object() -> Value {
174 json!({})
175 }
176
177 #[derive(Debug, Serialize)]
178 struct WireTool {
179 r#type: &'static str,
180 function: WireFunction,
181 }
182
183 #[derive(Debug, Serialize)]
184 struct WireFunction {
185 name: String,
186 description: String,
187 parameters: Value,
188 }
189