Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/clients/ollama.py: 22%
145 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Ollama LLM client implementation for local models.
3Production-ready implementation with Ollama API integration.
4"""
6from __future__ import annotations
8from collections.abc import AsyncGenerator, AsyncIterator
9from datetime import UTC, datetime
10from typing import Any
11import uuid
13from lexigram.ai.llm.clients._message_utils import serialize_text_for_ollama
14from lexigram.ai.llm.clients.base import AbstractLLMClient
15from lexigram.ai.llm.config import ClientConfig
16from lexigram.ai.llm.exceptions import LLMError
17from lexigram.ai.llm.model_manager import LLMModelManager
18from lexigram.ai.llm.multimodal.fetcher import fetch_image_as_base64
19from lexigram.ai.llm.types import (
20 AIError,
21 ChatMessage,
22 Completion,
23 StreamChunk,
24 TokenUsage,
25)
26from lexigram.contracts.ai.llm import FunctionCall, ToolCall
27from lexigram.contracts.ai.multimodal import ImageUrlPart
28from lexigram.contracts.core.health import HealthCheckResult, HealthStatus
29from lexigram.result import Ok, Result
30from lexigram.serialization import loads_str
33class OllamaClient(AbstractLLMClient):
34 """Ollama LLM client for local models.
36 Conforms to: :class:`~lexigram.contracts.ai.LLMClientProtocol` protocol via structural typing.
38 Supports running LLMs locally with Ollama:
39 - Llama 3, Mistral, Phi, and other open models
40 - Streaming responses
41 - Zero API costs
42 - Full data privacy
44 Example:
45 >>> from lexigram.ai import ClientConfig
46 >>> config = ClientConfig(
47 ... provider="ollama",
48 ... model="llama3:8b",
49 ... api_base="http://localhost:11434"
50 ... )
51 >>> client = OllamaClient(config)
52 >>> completion = await client.complete([
53 ... ChatMessage(role="user", content="Hello!")
54 ... ])
55 """
57 def __init__(self, config: ClientConfig):
58 """Initialize Ollama client.
60 Args:
61 config: LLM configuration
63 Raises:
64 ImportError: If ollama package is not installed
65 """
66 super().__init__(config=config)
67 self.model_manager: LLMModelManager | None = None
68 self.current_model = config.model
70 try:
71 from ollama import AsyncClient
72 except ImportError as e:
73 raise ImportError(
74 "Ollama client requires 'ollama' package. "
75 "Install with: pip install lexigram-ai-llm[ollama]",
76 ) from e
78 # Ollama typically runs on localhost:11434
79 host = config.api_base or "http://localhost:11434"
80 self.client = AsyncClient(host=host)
81 self._closed = False
83 async def _ensure_model_loaded(self, model_name: str) -> None:
84 """Ensure the specified model is loaded (no-op without an injected model_manager)."""
85 if model_name != self.current_model and self.model_manager is not None:
86 await self.model_manager.switch_provider("ollama")
87 await self.model_manager.load_model(model_name)
88 self.current_model = model_name
90 async def _serialize_messages_for_ollama(
91 self, messages: list[ChatMessage]
92 ) -> list[dict[str, Any]]:
93 """Serialize ChatMessages to Ollama wire format, pre-fetching URL images.
95 Builds Ollama-compatible message dicts with ``content`` (text) and optional
96 ``images`` (list of raw base64 strings, no data URI prefix). ``ImageUrlPart``
97 entries are fetched via ``fetch_image_as_base64`` before serialization.
99 Args:
100 messages: ChatMessages to serialize.
102 Returns:
103 List of Ollama message dicts with optional ``images`` key.
104 """
105 result: list[dict[str, Any]] = []
106 for msg in messages:
107 content = msg.content
108 # Pre-fetch ImageUrlParts before passing to sync serializer
109 if isinstance(content, list):
110 resolved: list[Any] = []
111 for part in content:
112 if isinstance(part, ImageUrlPart):
113 resolved.append(await fetch_image_as_base64(part.url))
114 else:
115 resolved.append(part)
116 content = resolved
118 text, images = serialize_text_for_ollama(content)
119 entry: dict[str, Any] = {
120 "role": msg.role.value,
121 "content": text,
122 }
123 if images:
124 entry["images"] = images
125 if msg.tool_calls:
126 entry["tool_calls"] = self._serialize_tool_calls(msg.tool_calls)
127 result.append(entry)
128 return result
130 async def _do_complete(
131 self,
132 messages: list[ChatMessage],
133 **kwargs: Any,
134 ) -> Result[Completion, LLMError]:
135 """Generate completion from messages.
137 Args:
138 messages: Chat messages
139 **kwargs: Additional Ollama parameters
141 Returns:
142 ``Ok(Completion)`` on success.
144 Raises:
145 AIError: If request fails (all Ollama errors are infrastructure).
146 """
147 try:
148 model_name = kwargs.pop("model", self.config.model)
149 tools = kwargs.pop("tools", None)
150 await self._ensure_model_loaded(model_name)
152 # Convert messages to Ollama format
153 ollama_messages = await self._serialize_messages_for_ollama(messages)
155 # Merge config defaults with kwargs
156 params = {
157 "model": model_name,
158 "messages": ollama_messages,
159 "options": {
160 "temperature": kwargs.pop("temperature", self.config.temperature),
161 "num_predict": kwargs.pop("max_tokens", self.config.max_tokens),
162 },
163 **kwargs,
164 }
166 if tools:
167 params["tools"] = self._convert_tools(tools)
169 # Inject thinking suppress for models that support it
170 if self.config.thinking is not None and self.config.thinking.suppress:
171 params["think"] = False
173 # Make API call
174 response = await self.client.chat(**params)
176 # Convert to our Completion type
177 return Ok(
178 Completion(
179 content=response["message"]["content"],
180 model=response["model"],
181 finish_reason="stop",
182 tool_calls=self._parse_tool_calls(
183 response["message"].get("tool_calls")
184 ),
185 usage=TokenUsage(
186 prompt_tokens=response.get("prompt_eval_count", 0),
187 completion_tokens=response.get("eval_count", 0),
188 total_tokens=response.get("prompt_eval_count", 0)
189 + response.get("eval_count", 0),
190 ),
191 metadata={
192 "total_duration": response.get("total_duration"),
193 "load_duration": response.get("load_duration"),
194 },
195 timestamp=datetime.now(UTC),
196 )
197 )
199 except (LLMError, ValueError, RuntimeError, OSError, ConnectionError) as e:
200 msg = f"Ollama completion failed: {e}"
201 raise AIError(msg) from e
203 async def _do_stream_chat(
204 self,
205 messages: list[ChatMessage],
206 **kwargs: Any,
207 ) -> Result[AsyncIterator[StreamChunk], LLMError]:
208 """Start streaming completion (protocol-aligned method).
210 Args:
211 messages: Chat messages
212 **kwargs: Additional parameters
214 Returns:
215 ``Ok(AsyncIterator[StreamChunk])`` on success.
217 Raises:
218 AIError: If setting up the stream fails.
219 """
220 return Ok(self._stream_impl(messages, **kwargs))
222 async def _stream_impl(
223 self,
224 messages: list[ChatMessage],
225 **kwargs: Any,
226 ) -> AsyncGenerator[StreamChunk, None]:
227 """Internal async generator for streaming.
229 Args:
230 messages: Chat messages
231 **kwargs: Additional parameters
233 Yields:
234 StreamChunk objects
236 Raises:
237 AIError: If streaming fails
238 """
239 try:
240 model_name = kwargs.pop("model", self.config.model)
241 await self._ensure_model_loaded(model_name)
243 # Convert messages to Ollama format
244 ollama_messages = await self._serialize_messages_for_ollama(messages)
246 # Merge config defaults with kwargs
247 params = {
248 "model": model_name,
249 "messages": ollama_messages,
250 "stream": True,
251 "options": {
252 "temperature": kwargs.pop("temperature", self.config.temperature),
253 "num_predict": kwargs.pop("max_tokens", self.config.max_tokens),
254 },
255 **kwargs,
256 }
258 # Inject thinking suppress for models that support it
259 if self.config.thinking is not None and self.config.thinking.suppress:
260 params["think"] = False
262 # Stream API call
263 index = 0
264 stream = await self.client.chat(**params)
265 async for chunk in stream:
266 if "message" in chunk and "content" in chunk["message"]:
267 delta = chunk["message"]["content"]
268 if delta:
269 yield StreamChunk(
270 delta=delta,
271 model=chunk.get("model", self.config.model),
272 finish_reason="stop" if chunk.get("done") else None,
273 index=index,
274 )
275 index += 1
277 except (LLMError, ValueError, RuntimeError, OSError, ConnectionError) as e:
278 msg = f"Ollama streaming failed: {e}"
279 raise AIError(msg) from e
281 async def _do_chat(
282 self,
283 messages: list[ChatMessage],
284 tools: list[ToolCall] | None = None,
285 **kwargs: Any,
286 ) -> Result[Completion, LLMError]:
287 """Generate completion (native tool calling handled in ``_do_complete``).
289 Args:
290 messages: Chat messages
291 tools: Optional tools (exercised via ``complete(..., tools=...)``)
292 **kwargs: Additional parameters
294 Returns:
295 ``Ok(Completion)`` on success. ``Err(LLMError)`` on failure.
296 """
297 # Ollama tool calling is wired through complete(..., tools=...);
298 # delegate directly.
299 return await self._do_complete(messages, **kwargs)
301 # ------------------------------------------------------------------
302 # Native tool calling helpers
303 # ------------------------------------------------------------------
305 @staticmethod
306 def _convert_tools(tools: Any) -> list[dict[str, Any]]:
307 """Convert ``ToolDefinition`` schemas to Ollama tool dicts.
309 Args:
310 tools: Iterable of ToolDefinition (or name/description/parameters
311 duck types) received via ``complete(..., tools=...)``.
313 Returns:
314 List of Ollama-compatible ``{"type": "function", "function": {...}}``
315 dicts.
316 """
317 converted: list[dict[str, Any]] = []
318 for tool in tools:
319 name = getattr(tool, "name", None)
320 if not name:
321 continue
322 description = getattr(tool, "description", None) or ""
323 parameters = getattr(tool, "parameters", None) or {
324 "type": "object",
325 "properties": {},
326 }
327 converted.append(
328 {
329 "type": "function",
330 "function": {
331 "name": name,
332 "description": description,
333 "parameters": parameters,
334 },
335 }
336 )
337 return converted
339 def _serialize_tool_calls(self, tool_calls: list[ToolCall]) -> list[dict[str, Any]]:
340 """Serialize framework ``ToolCall``s back to Ollama wire format.
342 Assistant turns that requested tools must be re-emitted before the
343 matching ``tool`` role responses so the conversation stays consistent.
345 Args:
346 tool_calls: Framework tool calls from a prior assistant turn.
348 Returns:
349 Ollama ``tool_calls`` dicts (``{"function": {"name", "arguments"}}``).
350 """
351 serialized: list[dict[str, Any]] = []
352 for call in tool_calls:
353 if call.function is None:
354 continue
355 arguments = call.function.arguments
356 if isinstance(arguments, str):
357 try:
358 arguments = loads_str(arguments)
359 except (TypeError, ValueError):
360 arguments = {}
361 serialized.append(
362 {
363 "function": {
364 "name": call.function.name,
365 "arguments": dict(arguments)
366 if isinstance(arguments, dict)
367 else {},
368 }
369 }
370 )
371 return serialized
373 @staticmethod
374 def _parse_tool_calls(raw: Any) -> list[ToolCall] | None:
375 """Parse Ollama tool calls into framework ``ToolCall`` objects.
377 Ollama models reply with tool calls as ``Message.ToolCall`` objects
378 (or dicts) carrying ``function.name`` and ``function.arguments``.
379 Ollama does not assign call IDs, so a synthetic ID is generated for
380 round-trip bookkeeping.
382 Args:
383 raw: The raw ``tool_calls`` from the Ollama response message.
385 Returns:
386 List of framework ``ToolCall`` objects, or ``None`` when the
387 response contained no tool calls.
388 """
389 if not raw:
390 return None
391 calls: list[ToolCall] = []
392 for raw_call in raw:
393 function = (
394 raw_call.get("function")
395 if isinstance(raw_call, dict)
396 else getattr(raw_call, "function", None)
397 )
398 if function is None:
399 continue
400 if isinstance(function, dict):
401 name = function.get("name")
402 raw_arguments = function.get("arguments", {})
403 else:
404 name = getattr(function, "name", None)
405 raw_arguments = getattr(function, "arguments", {})
406 if not name:
407 continue
408 calls.append(
409 ToolCall(
410 id=uuid.uuid4().hex,
411 type="function",
412 function=FunctionCall(name=name, arguments=raw_arguments),
413 )
414 )
415 return calls or None
417 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
418 """Perform a lightweight health check against the Ollama daemon.
420 Calls ``list()`` to verify the daemon is running and reachable.
422 Args:
423 timeout: Maximum seconds to wait for the response.
425 Returns:
426 :class:`~lexigram.contracts.core.health.HealthCheckResult`.
427 """
428 try:
429 models_response = await self.client.list()
430 model_names = [m.model for m in models_response.models]
431 return HealthCheckResult(
432 component="llm.ollama",
433 status=HealthStatus.HEALTHY,
434 details={
435 "provider": "ollama",
436 "model": self.current_model,
437 "available_models": model_names,
438 },
439 )
440 except (OSError, ConnectionError, TimeoutError, RuntimeError) as exc:
441 return HealthCheckResult(
442 component="llm.ollama",
443 status=HealthStatus.UNHEALTHY,
444 error=str(exc),
445 )
447 async def close(self) -> None:
448 """Close Ollama client."""
449 if not self._closed:
450 self._closed = True
451 await super().close()