feat: complete production-quality Model Router AI traffic control platform
This commit is contained in:
Vardhman Gupta
@@ -0,0 +1,46 @@
|
||||
import abc
|
||||
import time
|
||||
import asyncio
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.models.schemas import ProviderResponse
|
||||
|
||||
|
||||
class ModelProvider(abc.ABC):
|
||||
"""
|
||||
Abstract Base Class for all Model Providers (Mock, Ollama, OpenAI, Anthropic, Gemini, etc.)
|
||||
Ensures complete provider decoupling from routing logic.
|
||||
"""
|
||||
|
||||
def __init__(self, provider_id: str, name: str, base_url: Optional[str] = None):
|
||||
self.provider_id = provider_id
|
||||
self.name = name
|
||||
self.base_url = base_url
|
||||
|
||||
@abc.abstractmethod
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
"""Execute a text generation call against the model provider."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Stream generated response chunks."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
"""Check provider connectivity, status, and loaded models."""
|
||||
pass
|
||||
@@ -0,0 +1,321 @@
|
||||
import time
|
||||
import httpx
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.models.schemas import ProviderResponse
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OpenAIProvider(ModelProvider):
|
||||
"""
|
||||
OpenAI Provider Adapter.
|
||||
Uses OPENAI_API_KEY from environment variables. Secrets are never exposed.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
super().__init__(provider_id="openai", name="OpenAI", base_url="https://api.openai.com/v1")
|
||||
self.api_key = api_key or settings.OPENAI_API_KEY
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
if not self.api_key:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error="OPENAI_API_KEY not configured in environment.",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="openai",
|
||||
)
|
||||
|
||||
endpoint = f"{self.base_url}/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"model": model_id,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if max_tokens:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload, headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
choice = data["choices"][0]
|
||||
usage = data.get("usage", {})
|
||||
|
||||
return ProviderResponse(
|
||||
content=choice["message"]["content"],
|
||||
finish_reason=choice.get("finish_reason", "stop"),
|
||||
input_tokens=usage.get("prompt_tokens", 0),
|
||||
output_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="openai",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"OpenAI error: {str(exc)}",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="openai",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
if not self.api_key:
|
||||
yield "[OpenAI Error: API key not configured]"
|
||||
return
|
||||
|
||||
endpoint = f"{self.base_url}/chat/completions"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
payload = {
|
||||
"model": model_id,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
async with client.stream("POST", endpoint, json=payload, headers=headers) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if line.startswith("data: ") and not line.endswith("[DONE]"):
|
||||
import json
|
||||
data = json.loads(line[6:])
|
||||
delta = data["choices"][0].get("delta", {}).get("content", "")
|
||||
if delta:
|
||||
yield delta
|
||||
except Exception as exc:
|
||||
yield f"\n[OpenAI Streaming Error: {str(exc)}]"
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
if not self.api_key:
|
||||
return {"provider": "openai", "status": "NOT_CONFIGURED", "message": "API key not set in .env"}
|
||||
return {"provider": "openai", "status": "CONNECTED", "message": "API key configured."}
|
||||
|
||||
|
||||
class AnthropicProvider(ModelProvider):
|
||||
"""
|
||||
Anthropic Provider Adapter.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
super().__init__(provider_id="anthropic", name="Anthropic", base_url="https://api.anthropic.com/v1")
|
||||
self.api_key = api_key or settings.ANTHROPIC_API_KEY
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
if not self.api_key:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error="ANTHROPIC_API_KEY not configured in environment.",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="anthropic",
|
||||
)
|
||||
|
||||
endpoint = f"{self.base_url}/messages"
|
||||
headers = {
|
||||
"x-api-key": self.api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"model": model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": max_tokens or 2048,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload, headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
content = data["content"][0]["text"]
|
||||
usage = data.get("usage", {})
|
||||
|
||||
return ProviderResponse(
|
||||
content=content,
|
||||
finish_reason=data.get("stop_reason", "stop"),
|
||||
input_tokens=usage.get("input_tokens", 0),
|
||||
output_tokens=usage.get("output_tokens", 0),
|
||||
total_tokens=usage.get("input_tokens", 0) + usage.get("output_tokens", 0),
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="anthropic",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"Anthropic error: {str(exc)}",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="anthropic",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
resp = await self.generate(prompt, model_id, system_prompt, temperature, max_tokens)
|
||||
if resp.error:
|
||||
yield f"[{resp.error}]"
|
||||
return
|
||||
for word in resp.content.split(" "):
|
||||
yield word + " "
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
if not self.api_key:
|
||||
return {"provider": "anthropic", "status": "NOT_CONFIGURED", "message": "API key not set in .env"}
|
||||
return {"provider": "anthropic", "status": "CONNECTED", "message": "API key configured."}
|
||||
|
||||
|
||||
class GeminiProvider(ModelProvider):
|
||||
"""
|
||||
Google Gemini Provider Adapter.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
super().__init__(provider_id="gemini", name="Google Gemini", base_url="https://generativelanguage.googleapis.com/v1beta")
|
||||
self.api_key = api_key or settings.GEMINI_API_KEY
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
if not self.api_key:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error="GEMINI_API_KEY not configured in environment.",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="gemini",
|
||||
)
|
||||
|
||||
endpoint = f"{self.base_url}/models/{model_id}:generateContent?key={self.api_key}"
|
||||
contents = []
|
||||
if system_prompt:
|
||||
contents.append({"role": "user", "parts": [{"text": f"System context: {system_prompt}"}]})
|
||||
contents.append({"role": "user", "parts": [{"text": prompt}]})
|
||||
|
||||
payload = {
|
||||
"contents": contents,
|
||||
"generationConfig": {"temperature": temperature},
|
||||
}
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
candidate = data["candidates"][0]
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
usage = data.get("usageMetadata", {})
|
||||
|
||||
return ProviderResponse(
|
||||
content=text,
|
||||
finish_reason=candidate.get("finishReason", "stop"),
|
||||
input_tokens=usage.get("promptTokenCount", 0),
|
||||
output_tokens=usage.get("candidatesTokenCount", 0),
|
||||
total_tokens=usage.get("totalTokenCount", 0),
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="gemini",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"Gemini error: {str(exc)}",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="gemini",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
resp = await self.generate(prompt, model_id, system_prompt, temperature, max_tokens)
|
||||
if resp.error:
|
||||
yield f"[{resp.error}]"
|
||||
return
|
||||
for word in resp.content.split(" "):
|
||||
yield word + " "
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
if not self.api_key:
|
||||
return {"provider": "gemini", "status": "NOT_CONFIGURED", "message": "API key not set in .env"}
|
||||
return {"provider": "gemini", "status": "CONNECTED", "message": "API key configured."}
|
||||
@@ -0,0 +1,108 @@
|
||||
import asyncio
|
||||
import time
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.models.schemas import ProviderResponse
|
||||
|
||||
|
||||
class MockProvider(ModelProvider):
|
||||
"""
|
||||
Mock Provider for development, testing, and zero-cost simulation.
|
||||
All outputs are explicitly labeled as DEMO/MOCK.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(provider_id="mock", name="Mock / Demo Provider")
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
start_time = time.perf_counter()
|
||||
|
||||
# Simulate realistic latency based on tier
|
||||
tier_delay = 0.08 if "fast" in model_id else (0.18 if "balanced" in model_id else 0.35)
|
||||
await asyncio.sleep(tier_delay)
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
|
||||
# Calculate realistic token estimations
|
||||
input_tokens = max(1, len(prompt.split()) * 2)
|
||||
|
||||
# Response tailored to task keywords
|
||||
prompt_lower = prompt.lower()
|
||||
if "python" in prompt_lower or "code" in prompt_lower or "def " in prompt_lower or "function" in prompt_lower:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"```python\n"
|
||||
f"# Mock implementation generated for: {prompt[:60]}...\n"
|
||||
f"def solve_problem(*args, **kwargs):\n"
|
||||
f" \"\"\"Simulated response from Model Router Demo Engine.\"\"\"\n"
|
||||
f" result = [x for x in args if x is not None]\n"
|
||||
f" return {{'status': 'success', 'data': result, 'model': '{model_id}'}}\n"
|
||||
f"```\n\n"
|
||||
f"This code was returned by the simulated {model_id} engine."
|
||||
)
|
||||
elif "debug" in prompt_lower or "error" in prompt_lower or "bug" in prompt_lower:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"### Root Cause Analysis:\n"
|
||||
f"1. **Potential Bottleneck**: Asynchronous task scheduling without proper semaphore limits.\n"
|
||||
f"2. **Mitigation**: Introduce bounded concurrency worker pools.\n"
|
||||
f"3. **Verification**: Checked against {model_id} simulated reasoning engine."
|
||||
)
|
||||
elif "summary" in prompt_lower or "summarize" in prompt_lower:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"**Executive Summary**:\n"
|
||||
f"• Main point: {prompt[:80]}...\n"
|
||||
f"• Key takeaway: High routing accuracy achieved with minimal latency overhead."
|
||||
)
|
||||
else:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"Processed request: \"{prompt[:100]}\"\n\n"
|
||||
f"This is a simulated response generated locally by Model Router Mock Engine without external API calls."
|
||||
)
|
||||
|
||||
output_tokens = max(1, len(content.split()) * 2)
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
return ProviderResponse(
|
||||
content=content,
|
||||
finish_reason="stop",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
time_to_first_token_ms=round(tier_delay * 500, 2),
|
||||
is_mock=True,
|
||||
model=model_id,
|
||||
provider="mock",
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
resp = await self.generate(prompt, model_id, system_prompt, temperature, max_tokens)
|
||||
chunks = resp.content.split(" ")
|
||||
for chunk in chunks:
|
||||
await asyncio.sleep(0.02)
|
||||
yield chunk + " "
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"provider": "mock",
|
||||
"status": "CONNECTED",
|
||||
"message": "Mock provider active (Demo Mode). Zero API keys required.",
|
||||
"models_available": ["mock-fast", "mock-balanced", "mock-power"],
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
import time
|
||||
import httpx
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.models.schemas import ProviderResponse
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OllamaProvider(ModelProvider):
|
||||
"""
|
||||
Local Ollama Provider adapter.
|
||||
Executes real inference against locally running Ollama instance without external API keys.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: Optional[str] = None):
|
||||
url = base_url or settings.OLLAMA_BASE_URL
|
||||
super().__init__(provider_id="ollama", name="Ollama (Local)", base_url=url)
|
||||
|
||||
def _strip_model_prefix(self, model_id: str) -> str:
|
||||
# e.g. "ollama-qwen2.5-coder" -> "qwen2.5-coder"
|
||||
if model_id.startswith("ollama-"):
|
||||
return model_id[len("ollama-"):]
|
||||
return model_id
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
real_model = self._strip_model_prefix(model_id)
|
||||
endpoint = f"{self.base_url.rstrip('/')}/api/generate"
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"model": real_model,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"temperature": temperature,
|
||||
}
|
||||
}
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
if max_tokens:
|
||||
payload["options"]["num_predict"] = max_tokens
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
|
||||
content = data.get("response", "")
|
||||
input_tokens = data.get("prompt_eval_count", max(1, len(prompt.split()) * 2))
|
||||
output_tokens = data.get("eval_count", max(1, len(content.split()) * 2))
|
||||
|
||||
# Nanoseconds to ms if provided by ollama
|
||||
eval_duration_ms = data.get("eval_duration", 0) / 1_000_000.0 if "eval_duration" in data else elapsed_ms
|
||||
|
||||
return ProviderResponse(
|
||||
content=content,
|
||||
finish_reason="stop" if data.get("done") else "length",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
time_to_first_token_ms=round(data.get("prompt_eval_duration", 0) / 1_000_000.0, 2) or None,
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="ollama",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"Ollama connection error ({str(exc)}). Is Ollama running on {self.base_url}?",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="ollama",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
real_model = self._strip_model_prefix(model_id)
|
||||
endpoint = f"{self.base_url.rstrip('/')}/api/generate"
|
||||
payload = {
|
||||
"model": real_model,
|
||||
"prompt": prompt,
|
||||
"stream": True,
|
||||
"options": {"temperature": temperature},
|
||||
}
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
async with client.stream("POST", endpoint, json=payload) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if line:
|
||||
import json
|
||||
data = json.loads(line)
|
||||
chunk = data.get("response", "")
|
||||
if chunk:
|
||||
yield chunk
|
||||
except Exception as exc:
|
||||
yield f"\n[Ollama Streaming Error: {str(exc)}]"
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=3.0) as client:
|
||||
resp = await client.get(f"{self.base_url.rstrip('/')}/api/tags")
|
||||
if resp.status_code == 200:
|
||||
models = [m.get("name") for m in resp.json().get("models", [])]
|
||||
return {
|
||||
"provider": "ollama",
|
||||
"status": "CONNECTED",
|
||||
"base_url": self.base_url,
|
||||
"models_available": models,
|
||||
}
|
||||
except Exception as exc:
|
||||
pass
|
||||
return {
|
||||
"provider": "ollama",
|
||||
"status": "NOT_CONNECTED",
|
||||
"base_url": self.base_url,
|
||||
"message": "Ollama service unreachable on configured port. Start via 'ollama serve'.",
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Dict, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.providers.mock_provider import MockProvider
|
||||
from app.providers.ollama_provider import OllamaProvider
|
||||
from app.providers.external_providers import OpenAIProvider, AnthropicProvider, GeminiProvider
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""
|
||||
Central registry managing provider instances.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._providers: Dict[str, ModelProvider] = {
|
||||
"mock": MockProvider(),
|
||||
"ollama": OllamaProvider(base_url=settings.OLLAMA_BASE_URL),
|
||||
"openai": OpenAIProvider(api_key=settings.OPENAI_API_KEY),
|
||||
"anthropic": AnthropicProvider(api_key=settings.ANTHROPIC_API_KEY),
|
||||
"gemini": GeminiProvider(api_key=settings.GEMINI_API_KEY),
|
||||
}
|
||||
|
||||
def get_provider(self, provider_id: str) -> Optional[ModelProvider]:
|
||||
return self._providers.get(provider_id.lower())
|
||||
|
||||
def register_provider(self, provider_id: str, provider: ModelProvider):
|
||||
self._providers[provider_id.lower()] = provider
|
||||
|
||||
def list_providers(self) -> Dict[str, ModelProvider]:
|
||||
return self._providers
|
||||
|
||||
|
||||
provider_registry = ProviderRegistry()
|
||||
Reference in New Issue
Block a user