"""Ollama local LLM provider implementation.
This module provides the Ollama provider implementation using the mcp-common
HTTPClientAdapter for connection pooling and aiohttp fallback for HTTP communications.
"""
from __future__ import annotations
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any
from session_buddy.llm.base import LLMProvider
from session_buddy.llm.models import LLMMessage, LLMResponse
if TYPE_CHECKING:
from collections.abc import AsyncGenerator
from session_buddy.di.container import depends
# mcp-common HTTP client adapter (httpx based)
try:
from mcp_common.adapters.http.client import HTTPClientAdapter
from session_buddy.di.container import depends
HTTP_ADAPTER_AVAILABLE = True
except Exception:
HTTPClientAdapter = None # type: ignore[assignment]
HTTP_ADAPTER_AVAILABLE = False
class OllamaProvider(LLMProvider):
"""Ollama local LLM provider using HTTPClientAdapter for connection pooling."""
def __init__(self, config: dict[str, Any]) -> None:
super().__init__(config)
self.base_url = config.get("base_url", "http://localhost:11434")
self.default_model = config.get("default_model", "llama2")
self._available_models: list[str] = []
# Initialize HTTP client adapter if available
self._http_adapter = None
if HTTP_ADAPTER_AVAILABLE and HTTPClientAdapter is not None:
try:
self._http_adapter = depends.get_sync(HTTPClientAdapter)
except Exception:
self._http_adapter = None
async def _make_api_request(
self,
endpoint: str,
data: dict[str, Any],
) -> dict[str, Any]:
"""Make API request to Ollama service with connection pooling."""
url = f"{self.base_url}/{endpoint}"
if self._http_adapter is not None:
try:
async with self._http_adapter as client:
resp = await client.post(url, json=data, timeout=300)
return resp.json() # type: ignore[no-any-return]
except Exception as e:
self.logger.exception(f"HTTP request failed: {e}")
raise
# Fallback to aiohttp (legacy)
try:
import aiohttp
async with (
aiohttp.ClientSession() as session,
session.post(
url,
json=data,
timeout=aiohttp.ClientTimeout(total=300),
) as response,
):
return await response.json() # type: ignore[no-any-return]
except ImportError:
msg = (
"aiohttp package not installed and HTTPClientAdapter not available. "
"Install with: pip install aiohttp or configure mcp-common HTTPClientAdapter"
)
raise ImportError(msg) # type: ignore[no-any-return]
def _convert_messages(self, messages: list[LLMMessage]) -> list[dict[str, str]]:
"""Convert LLMMessage objects to Ollama format."""
return [{"role": msg.role, "content": msg.content} for msg in messages]
async def generate(
self,
messages: list[LLMMessage],
model: str | None = None,
temperature: float = 0.7,
max_tokens: int | None = None,
**kwargs: Any,
) -> LLMResponse:
"""Generate response using Ollama API."""
if not await self.is_available():
msg = "Ollama provider not available"
raise RuntimeError(msg)
model_name = model or self.default_model
try:
data: dict[str, Any] = {
"model": model_name,
"messages": self._convert_messages(messages),
"options": {"temperature": temperature},
}
if max_tokens:
data["options"]["num_predict"] = max_tokens
response = await self._make_api_request("api/chat", data)
return LLMResponse(
content=response.get("message", {}).get("content", ""),
model=model_name,
provider="ollama",
usage={
"prompt_tokens": response.get("prompt_eval_count", 0),
"completion_tokens": response.get("eval_count", 0),
"total_tokens": response.get("prompt_eval_count", 0)
+ response.get("eval_count", 0),
},
finish_reason=response.get("done_reason", "stop"),
timestamp=datetime.now().isoformat(),
)
except Exception as e:
self.logger.exception(f"Ollama generation failed: {e}")
raise
def _prepare_stream_data(
self,
model_name: str,
messages: list[LLMMessage],
temperature: float,
max_tokens: int | None,
) -> dict[str, Any]:
"""Prepare data payload for streaming request."""
data: dict[str, Any] = {
"model": model_name,
"messages": self._convert_messages(messages),
"stream": True,
"options": {"temperature": temperature},
}
if max_tokens:
data["options"]["num_predict"] = max_tokens
return data
def _extract_chunk_content(self, line: bytes) -> str | None:
"""Extract content from a streaming chunk line."""
if not line:
return None
try:
chunk_data = json.loads(line.decode("utf-8"))
if isinstance(chunk_data, dict) and "message" in chunk_data:
message = chunk_data["message"]
if isinstance(message, dict) and "content" in message:
return str(message["content"])
except json.JSONDecodeError:
pass
return None
async def _stream_from_response_aiohttp(self, response: Any) -> AsyncGenerator[str]:
"""Process streaming response from aiohttp and yield content chunks."""
async for line in response.content:
content = self._extract_chunk_content(line)
if content:
yield content
async def _stream_from_response_httpx(self, response: Any) -> AsyncGenerator[str]:
"""Process streaming response from httpx and yield content chunks."""
async for line in response.aiter_bytes():
content = self._extract_chunk_content(line)
if content:
yield content
async def _stream_with_mcp_common(
self,
url: str,
data: dict[str, Any],
) -> AsyncGenerator[str]:
"""Stream using MCP-common HTTP adapter."""
# Note: http_adapter access requires mcp-common integration setup
# This is a placeholder for future mcp-common integration
if False: # Disabled until http_adapter is properly initialized
yield "" # pragma: no cover
else:
# Fallback to aiohttp for now
async for chunk in self._stream_with_aiohttp(url, data):
yield chunk
async def _stream_with_aiohttp(
self,
url: str,
data: dict[str, Any],
) -> AsyncGenerator[str]:
"""Stream using aiohttp fallback."""
try:
import aiohttp
async with (
aiohttp.ClientSession() as session,
session.post(
url,
json=data,
timeout=aiohttp.ClientTimeout(total=300),
) as response,
):
async for chunk in self._stream_from_response_aiohttp(response):
yield chunk
except ImportError:
msg = "aiohttp not installed and mcp-common not available"
raise ImportError(msg)
async def stream_generate( # type: ignore[override]
self,
messages: list[LLMMessage],
model: str | None = None,
temperature: float = 0.7,
max_tokens: int | None = None,
**kwargs: Any,
) -> AsyncGenerator[str]:
"""Stream response using Ollama API with connection pooling."""
if not await self.is_available():
msg = "Ollama provider not available"
raise RuntimeError(msg)
model_name = model or self.default_model
data = self._prepare_stream_data(model_name, messages, temperature, max_tokens)
url = f"{self.base_url}/api/chat"
try:
# Note: mcp-common integration deferred - using aiohttp fallback
# if self._use_mcp_common and self.http_adapter:
# async for chunk in self._stream_with_mcp_common(url, data):
# yield chunk
# else:
async for chunk in self._stream_with_aiohttp(url, data):
yield chunk
except Exception as e:
self.logger.exception(f"Ollama streaming failed: {e}")
raise
async def _check_with_mcp_common(self, url: str) -> bool:
"""Check availability using MCP-common HTTP adapter."""
# Note: http_adapter access requires mcp-common integration setup
# This is a placeholder for future mcp-common integration
return False # Disabled until http_adapter is properly initialized
async def _check_with_aiohttp(self, url: str) -> bool:
"""Check availability using aiohttp fallback."""
try:
import aiohttp
async with (
aiohttp.ClientSession() as session,
session.get(
url,
timeout=aiohttp.ClientTimeout(total=10),
) as response,
):
if response.status == 200:
data = await response.json()
self._available_models = [
model["name"] for model in data.get("models", [])
]
return True
return False
except Exception:
return False
async def is_available(self) -> bool:
"""Check if Ollama is available with connection pooling."""
try:
url = f"{self.base_url}/api/tags"
# Note: mcp-common integration deferred - using aiohttp fallback
# if self._use_mcp_common and self.http_adapter:
# return await self._check_with_mcp_common(url)
return await self._check_with_aiohttp(url)
except Exception:
return False
def get_models(self) -> list[str]:
"""Get available Ollama models."""
return self._available_models or [
"llama2",
"llama2:13b",
"llama2:70b",
"codellama",
"mistral",
"mixtral",
]
