max / makenotwork

5.2 KB · 196 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	// reqwest is declared `rustls-no-provider` so this crate does not drag
26	// a C crypto backend into every consumer, which means the provider is
27	// ours to name and `Client::new()` panics without it. `install_default`
28	// returns `Err` only when one is already installed, which is the state
29	// we want, so the result is dropped and the call is safe to repeat.
30	let _ = rustls::crypto::ring::default_provider().install_default();
31
32	Self {
33	client: reqwest::Client::new(),
34	base_url: DEFAULT_OLLAMA_URL.to_string(),
35	model: model.into(),
36	}
37	}
38
39	#[must_use]
40	pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
41	self.base_url = url.into();
42	self
43	}
44	}
45
46	#[async_trait]
47	impl InferenceProvider for OllamaProvider {
48	async fn chat(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Message> {
49	let wire_messages: Vec<WireMessage> =
50	messages.iter().map(WireMessage::from_neutral).collect();
51	let wire_tools: Vec<WireTool> = tools
52	.iter()
53	.map(\|s\| WireTool {
54	r#type: "function",
55	function: WireFunction {
56	name: s.name.clone(),
57	description: s.description.clone(),
58	parameters: s.input_schema.clone(),
59	},
60	})
61	.collect();
62
63	let req = ChatRequest {
64	model: &self.model,
65	messages: &wire_messages,
66	tools: &wire_tools,
67	stream: false,
68	};
69
70	let resp: ChatResponse = self
71	.client
72	.post(format!("{}/api/chat", self.base_url))
73	.json(&req)
74	.send()
75	.await?
76	.error_for_status()?
77	.json()
78	.await?;
79
80	Ok(resp.message.into_neutral())
81	}
82	}
83
84	// ---- wire types ----
85
86	#[derive(Debug, Serialize)]
87	struct ChatRequest<'a> {
88	model: &'a str,
89	messages: &'a [WireMessage],
90	tools: &'a [WireTool],
91	stream: bool,
92	}
93
94	#[derive(Debug, Deserialize)]
95	struct ChatResponse {
96	message: WireMessage,
97	}
98
99	#[derive(Debug, Serialize, Deserialize)]
100	struct WireMessage {
101	role: String,
102	#[serde(default, skip_serializing_if = "Option::is_none")]
103	content: Option<String>,
104	#[serde(default, skip_serializing_if = "Option::is_none")]
105	tool_calls: Option<Vec<WireToolCall>>,
106	#[serde(default, skip_serializing_if = "Option::is_none")]
107	tool_name: Option<String>,
108	}
109
110	impl WireMessage {
111	fn from_neutral(m: &Message) -> Self {
112	let role = match m.role {
113	Role::System => "system",
114	Role::User => "user",
115	Role::Assistant => "assistant",
116	Role::Tool => "tool",
117	}
118	.to_string();
119	let tool_calls = if m.tool_calls.is_empty() {
120	None
121	} else {
122	Some(
123	m.tool_calls
124	.iter()
125	.map(\|c\| WireToolCall {
126	function: WireToolCallFunction {
127	name: c.name.clone(),
128	arguments: c.arguments.clone(),
129	},
130	})
131	.collect(),
132	)
133	};
134	Self {
135	role,
136	content: m.content.clone(),
137	tool_calls,
138	tool_name: m.tool_name.clone(),
139	}
140	}
141
142	fn into_neutral(self) -> Message {
143	let role = match self.role.as_str() {
144	"system" => Role::System,
145	"user" => Role::User,
146	"tool" => Role::Tool,
147	_ => Role::Assistant,
148	};
149	let tool_calls = self
150	.tool_calls
151	.unwrap_or_default()
152	.into_iter()
153	.map(\|c\| ToolCall {
154	id: String::new(),
155	name: c.function.name,
156	arguments: c.function.arguments,
157	})
158	.collect();
159	Message {
160	role,
161	content: self.content,
162	tool_calls,
163	tool_name: self.tool_name,
164	}
165	}
166	}
167
168	#[derive(Debug, Serialize, Deserialize)]
169	struct WireToolCall {
170	function: WireToolCallFunction,
171	}
172
173	#[derive(Debug, Serialize, Deserialize)]
174	struct WireToolCallFunction {
175	name: String,
176	#[serde(default = "empty_object")]
177	arguments: Value,
178	}
179
180	fn empty_object() -> Value {
181	json!({})
182	}
183
184	#[derive(Debug, Serialize)]
185	struct WireTool {
186	r#type: &'static str,
187	function: WireFunction,
188	}
189
190	#[derive(Debug, Serialize)]
191	struct WireFunction {
192	name: String,
193	description: String,
194	parameters: Value,
195	}
196