Coverage for agentos/models/backends/ollama.py: 27%

1"""Ollama backend for AgentOS. 

2 

3Supports local LLM inference via Ollama. 

4Models: llama3, mistral, codellama, phi3, gemma2, deepseek-r1, etc. 

5""" 

6 

7from collections.abc import AsyncIterator 

8from dataclasses import dataclass 

9from typing import Any 

10 

11 

12@dataclass 

13class OllamaConfig: 

14 """Configuration for Ollama backend.""" 

15 

16 base_url: str = "http://localhost:11434" 

17 model: str = "llama3" 

18 temperature: float = 0.7 

19 max_tokens: int = 4096 

20 top_p: float = 0.9 

21 top_k: int = 40 

22 num_ctx: int = 8192 

23 timeout: int = 120 

24 max_retries: int = 3 

25 keep_alive: str = "5m" 

26 

27 

28class OllamaClient: 

29 """Ollama LLM client for local model inference. 

30 

31 Supports: 

32 - Chat completions (streaming and non-streaming) 

33 - Tool calling (function calling) 

34 - Model listing and management 

35 - Custom system prompts 

36 """ 

37 

38 def __init__(self, config: OllamaConfig | None = None): 

39 self.config = config or OllamaConfig() 

40 

41 @property 

42 def _generate_url(self) -> str: 

43 return f"{self.config.base_url}/api/chat" 

44 

45 def _build_payload( 

46 self, 

47 messages: list[dict[str, str]], 

48 system: str | None = None, 

49 **kwargs, 

50 ) -> dict[str, Any]: 

51 msgs = messages.copy() 

52 if system: 

53 msgs.insert(0, {"role": "system", "content": system}) 

54 

55 payload: dict[str, Any] = { 

56 "model": self.config.model, 

57 "messages": msgs, 

58 "stream": False, 

59 "options": { 

60 "temperature": kwargs.get("temperature", self.config.temperature), 

61 "num_predict": kwargs.get("max_tokens", self.config.max_tokens), 

62 "top_p": kwargs.get("top_p", self.config.top_p), 

63 "top_k": self.config.top_k, 

64 "num_ctx": self.config.num_ctx, 

65 }, 

66 } 

67 

68 if kwargs.get("tools"): 

69 payload["tools"] = kwargs["tools"] 

70 

71 return payload 

72 

73 async def _async_request(self, payload: dict) -> dict: 

74 import httpx 

75 

76 timeout = httpx.Timeout(self.config.timeout) 

77 for attempt in range(self.config.max_retries): 

78 try: 

79 async with httpx.AsyncClient() as client: 

80 resp = await client.post( 

81 self._generate_url, 

82 json=payload, 

83 timeout=timeout, 

84 ) 

85 resp.raise_for_status() 

86 return resp.json() 

87 except (httpx.HTTPStatusError, httpx.ConnectError): 

88 if attempt == self.config.max_retries - 1: 

89 raise 

90 import asyncio 

91 

92 await asyncio.sleep(2**attempt) 

93 

94 async def chat( 

95 self, 

96 messages: list[dict[str, str]], 

97 system: str | None = None, 

98 **kwargs, 

99 ) -> dict[str, Any]: 

100 """Send a chat completion request. 

101 

102 Returns dict with 'content', 'role', 'usage', 'model'. 

103 """ 

104 payload = self._build_payload(messages, system, **kwargs) 

105 result = await self._async_request(payload) 

106 

107 message = result.get("message", {}) 

108 tool_calls = [] 

109 if "tool_calls" in message: 

110 for tc in message["tool_calls"]: 

111 func = tc.get("function", {}) 

112 tool_calls.append( 

113 { 

114 "id": tc.get("id", ""), 

115 "name": func.get("name", ""), 

116 "arguments": func.get("arguments", "{}"), 

117 } 

118 ) 

119 

120 return { 

121 "content": message.get("content", ""), 

122 "role": "assistant", 

123 "usage": { 

124 "input_tokens": result.get("prompt_eval_count", 0), 

125 "output_tokens": result.get("eval_count", 0), 

126 "total_duration_ms": result.get("total_duration", 0), 

127 }, 

128 "model": result.get("model", self.config.model), 

129 "done_reason": result.get("done_reason", ""), 

130 "tool_calls": tool_calls, 

131 } 

132 

133 async def chat_stream( 

134 self, 

135 messages: list[dict[str, str]], 

136 system: str | None = None, 

137 **kwargs, 

138 ) -> AsyncIterator[dict[str, Any]]: 

139 """Stream chat completion tokens.""" 

140 payload = self._build_payload(messages, system, **kwargs) 

141 payload["stream"] = True 

142 

143 import httpx 

144 

145 timeout = httpx.Timeout(self.config.timeout * 2) 

146 async with httpx.AsyncClient() as client: 

147 async with client.stream( 

148 "POST", 

149 self._generate_url, 

150 json=payload, 

151 timeout=timeout, 

152 ) as response: 

153 response.raise_for_status() 

154 async for line in response.aiter_lines(): 

155 import json 

156 

157 try: 

158 event = json.loads(line) 

159 message = event.get("message", {}) 

160 yield { 

161 "delta": message.get("content", ""), 

162 "done": event.get("done", False), 

163 "model": event.get("model", self.config.model), 

164 } 

165 if event.get("done"): 

166 break 

167 except json.JSONDecodeError: 

168 continue 

169 

170 def sync_chat( 

171 self, 

172 messages: list[dict[str, str]], 

173 system: str | None = None, 

174 **kwargs, 

175 ) -> dict[str, Any]: 

176 """Synchronous chat completion.""" 

177 import asyncio 

178 

179 loop = asyncio.new_event_loop() 

180 try: 

181 return loop.run_until_complete(self.chat(messages, system, **kwargs)) 

182 finally: 

183 loop.close() 

184 

185 async def list_models(self) -> list[dict[str, Any]]: 

186 """List locally available Ollama models.""" 

187 import httpx 

188 

189 timeout = httpx.Timeout(self.config.timeout) 

190 async with httpx.AsyncClient() as client: 

191 resp = await client.get( 

192 f"{self.config.base_url}/api/tags", 

193 timeout=timeout, 

194 ) 

195 resp.raise_for_status() 

196 data = resp.json() 

197 return [ 

198 { 

199 "name": m.get("name", ""), 

200 "size": m.get("size", 0), 

201 "modified": m.get("modified_at", ""), 

202 "format": m.get("details", {}).get("format", ""), 

203 } 

204 for m in data.get("models", []) 

205 ] 

206 

207 async def pull_model(self, model_name: str) -> AsyncIterator[dict[str, Any]]: 

208 """Pull a model from Ollama registry.""" 

209 import httpx 

210 

211 timeout = httpx.Timeout(600) # 10 min for downloads 

212 async with httpx.AsyncClient() as client: 

213 async with client.stream( 

214 "POST", 

215 f"{self.config.base_url}/api/pull", 

216 json={"name": model_name, "stream": True}, 

217 timeout=timeout, 

218 ) as response: 

219 response.raise_for_status() 

220 async for line in response.aiter_lines(): 

221 import json 

222 

223 try: 

224 event = json.loads(line) 

225 yield event 

226 except json.JSONDecodeError: 

227 continue