feat: AI Worker 平台 MVP v1.0.0 - 多租户/项目管理/AI Worker/任务编排/HITL审核/成本治理

2026-08-12 13:31:01 +08:00
parent 230a2cea16
commit 7caa52e64f
1 changed files with 113 additions and 0 deletions
@@ -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)