116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
"""LLM Gateway: unified interface for multiple model providers.
|
|
|
|
Supports OpenAI-compatible APIs (OpenAI, DeepSeek, vLLM, etc.)
|
|
"""
|
|
import json
|
|
import time
|
|
import httpx
|
|
from typing import Optional, Tuple
|
|
from app.config import settings
|
|
from app.models.worker import AIWorker
|
|
from app.models.cost import CostLog
|
|
from sqlalchemy.orm import Session
|
|
|
|
|
|
class LLMGateway:
|
|
"""Unified LLM gateway with token counting and cost tracking."""
|
|
|
|
@staticmethod
|
|
def _get_api_key(worker: AIWorker) -> str:
|
|
return worker.api_key or settings.DEFAULT_LLM_API_KEY
|
|
|
|
@staticmethod
|
|
def _get_base_url(worker: AIWorker) -> str:
|
|
return worker.base_url or settings.DEFAULT_LLM_BASE_URL
|
|
|
|
@staticmethod
|
|
def call(
|
|
worker: AIWorker,
|
|
messages: list[dict],
|
|
db: Session,
|
|
tenant_id: int,
|
|
project_id: Optional[int] = None,
|
|
task_id: Optional[int] = None,
|
|
) -> Tuple[str, CostLog]:
|
|
"""Make an LLM API call and record cost.
|
|
|
|
Returns (response_content, cost_log).
|
|
"""
|
|
api_key = LLMGateway._get_api_key(worker)
|
|
base_url = LLMGateway._get_base_url(worker)
|
|
url = f"{base_url.rstrip('/')}/chat/completions"
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
payload = {
|
|
"model": worker.model_name,
|
|
"messages": messages,
|
|
"temperature": worker.temperature,
|
|
}
|
|
|
|
start = time.time()
|
|
cost_log = CostLog(
|
|
tenant_id=tenant_id,
|
|
project_id=project_id,
|
|
task_id=task_id,
|
|
worker_id=worker.id,
|
|
provider=worker.provider,
|
|
model_name=worker.model_name,
|
|
)
|
|
|
|
try:
|
|
with httpx.Client(timeout=120) as client:
|
|
resp = client.post(url, json=payload, headers=headers)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
content = data["choices"][0]["message"]["content"]
|
|
usage = data.get("usage", {})
|
|
|
|
cost_log.prompt_tokens = usage.get("prompt_tokens", 0)
|
|
cost_log.completion_tokens = usage.get("completion_tokens", 0)
|
|
cost_log.total_tokens = usage.get("total_tokens", 0)
|
|
cost_log.cost_cents = LLMGateway._estimate_cost(
|
|
worker.model_name,
|
|
cost_log.prompt_tokens,
|
|
cost_log.completion_tokens,
|
|
)
|
|
cost_log.duration_ms = int((time.time() - start) * 1000)
|
|
cost_log.status = "success"
|
|
|
|
db.add(cost_log)
|
|
db.commit()
|
|
db.refresh(cost_log)
|
|
|
|
return content, cost_log
|
|
|
|
except Exception as e:
|
|
cost_log.status = "error"
|
|
cost_log.error_message = str(e)
|
|
cost_log.duration_ms = int((time.time() - start) * 1000)
|
|
db.add(cost_log)
|
|
db.commit()
|
|
raise
|
|
|
|
@staticmethod
|
|
def _estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> int:
|
|
"""Estimate cost in cents. Rough pricing, adjustable."""
|
|
# Per 1M tokens, in cents (1 USD = 100 cents)
|
|
pricing = {
|
|
"gpt-4o": {"input": 250, "output": 1000},
|
|
"gpt-4o-mini": {"input": 15, "output": 60},
|
|
"gpt-4-turbo": {"input": 1000, "output": 3000},
|
|
"deepseek-chat": {"input": 14, "output": 28},
|
|
"deepseek-reasoner": {"input": 55, "output": 219},
|
|
"deepseek-v4-flash": {"input": 14, "output": 28},
|
|
"qwen3.6-plus": {"input": 40, "output": 120},
|
|
}
|
|
rates = pricing.get(model, {"input": 15, "output": 60}) # default: cheap
|
|
cost = (
|
|
prompt_tokens * rates["input"] / 1_000_000
|
|
+ completion_tokens * rates["output"] / 1_000_000
|
|
)
|
|
return int(cost * 100)
|