diff --git a/backend/app/services/llm_service.py b/backend/app/services/llm_service.py new file mode 100644 index 0000000..3c37f3e --- /dev/null +++ b/backend/app/services/llm_service.py @@ -0,0 +1,113 @@ +"""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}, # $2.50 / $10 per 1M + "gpt-4o-mini": {"input": 15, "output": 60}, # $0.15 / $0.60 per 1M + "gpt-4-turbo": {"input": 1000, "output": 3000}, + "deepseek-chat": {"input": 14, "output": 28}, + "deepseek-reasoner": {"input": 55, "output": 219}, + } + 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)