Add provider implementations (OpenAI, Anthropic, Ollama)

This commit is contained in:
2026-02-04 12:32:11 +00:00
parent de3ffc114e
commit 8e3a8716fa

View File

@@ -0,0 +1,106 @@
import asyncio
import time
from typing import Any, AsyncIterator, Dict, Optional
import httpx
from .base import ProviderBase, ProviderResponse
from ..core.exceptions import ProviderError
class OllamaProvider(ProviderBase):
def __init__(
self,
model: str = "llama2",
temperature: float = 0.7,
base_url: str = "http://localhost:11434",
**kwargs,
):
super().__init__(None, model, temperature, **kwargs)
self.base_url = base_url.rstrip('/')
@property
def name(self) -> str:
return "ollama"
async def complete(
self,
prompt: str,
system_prompt: Optional[str] = None,
max_tokens: Optional[int] = None,
**kwargs,
) -> ProviderResponse:
start_time = time.time()
try:
async with httpx.AsyncClient() as client:
payload = {
"model": self.model,
"prompt": prompt,
"stream": False,
"options": {
"temperature": self.temperature,
}
}
if max_tokens:
payload["options"]["num_predict"] = max_tokens
response = await client.post(
f"{self.base_url}/api/generate",
json=payload,
timeout=120.0
)
response.raise_for_status()
data = response.json()
latency_ms = (time.time() - start_time) * 1000
return ProviderResponse(
content=data.get("response", ""),
model=self.model,
provider=self.name,
latency_ms=latency_ms,
)
except httpx.HTTPStatusError as e:
raise ProviderError(f"Ollama HTTP error: {e}")
except httpx.RequestError as e:
raise ProviderError(f"Ollama connection error: {e}")
async def stream_complete(
self,
prompt: str,
system_prompt: Optional[str] = None,
max_tokens: Optional[int] = None,
**kwargs,
) -> AsyncIterator[str]:
try:
async with httpx.AsyncClient() as client:
payload = {
"model": self.model,
"prompt": prompt,
"stream": True,
"options": {
"temperature": self.temperature,
}
}
if max_tokens:
payload["options"]["num_predict"] = max_tokens
async with client.stream(
"POST",
f"{self.base_url}/api/generate",
json=payload,
timeout=120.0
) as response:
async for line in response.aiter_lines():
import json
data = json.loads(line)
if "response" in data:
yield data["response"]
except httpx.HTTPStatusError as e:
raise ProviderError(f"Ollama HTTP error: {e}")
except httpx.RequestError as e:
raise ProviderError(f"Ollama connection error: {e}")
def validate_api_key(self) -> bool:
return True