Add provider implementations (OpenAI, Anthropic, Ollama)
This commit is contained in:
106
src/promptforge/providers/ollama.py
Normal file
106
src/promptforge/providers/ollama.py
Normal 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
|
||||
Reference in New Issue
Block a user