pcb-afm-simulation-system/web/backend/app/services/ai_client.py

"""Kimi AI Client - OpenAI compatible API wrapper.

P3-M1: AI service layer foundation.
Supports chat completions, streaming, retries, and call logging.
Endpoint: https://api.kimi.com/coding/v1 (Kimi Code Plan)
Model: k3 (1M context, reasoning enabled)
"""
import json
import time
import logging
from datetime import datetime
from typing import Optional, List, Dict, Any, AsyncIterator

import httpx

from ..config import (
    KIMI_API_KEY, KIMI_BASE_URL, KIMI_MODEL,
    KIMI_MAX_TOKENS, KIMI_TEMPERATURE, KIMI_TIMEOUT, KIMI_MAX_RETRIES
)
from ..database import SessionLocal
from ..models.ai_call_log import AICallLog

logger = logging.getLogger(__name__)


class KimiClient:
    """Kimi AI API client with retry and logging."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
    ):
        self.api_key = api_key or KIMI_API_KEY
        self.base_url = (base_url or KIMI_BASE_URL).rstrip("/")
        self.model = model or KIMI_MODEL
        self._client = httpx.Client(timeout=KIMI_TIMEOUT)

    @property
    def is_configured(self) -> bool:
        """Check if API key is configured."""
        return bool(self.api_key)

    def _headers(self) -> Dict[str, str]:
        return {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }

    def _log_call(
        self,
        endpoint: str,
        model: str,
        messages: List[Dict[str, str]],
        response: Optional[Dict[str, Any]],
        error: Optional[str],
        duration_ms: int,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        total_tokens: int = 0,
    ):
        """Log AI call to database."""
        try:
            db = SessionLocal()
            log = AICallLog(
                endpoint=endpoint,
                model=model,
                # Keep enough of the prompt to debug what the AI was actually
                # asked (BC + experience cases live in the user message); 500
                # chars truncated before the user message, hiding this.
                prompt_preview=json.dumps(messages, ensure_ascii=False)[:8000],
                response_preview=(
                    json.dumps(response, ensure_ascii=False)[:2000]
                    if response else None
                ),
                error=error,
                duration_ms=duration_ms,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=total_tokens,
                status="success" if not error else "failed",
                created_at=datetime.utcnow(),
            )
            db.add(log)
            db.commit()
            db.close()
        except Exception as e:
            logger.warning(f"Failed to log AI call: {e}")

    def chat(
        self,
        messages: List[Dict[str, str]],
        model: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        system_prompt: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Non-streaming chat completion.

        Args:
            messages: List of message dicts with role and content.
            model: Override default model.
            temperature: Sampling temperature.
            max_tokens: Max output tokens.
            system_prompt: Optional system prompt prepended to messages.

        Returns:
            Dict with content, reasoning_content, usage, etc.
        """
        if not self.is_configured:
            raise RuntimeError("KIMI_API_KEY is not configured")

        _model = model or self.model
        _messages = list(messages)
        if system_prompt:
            _messages.insert(0, {"role": "system", "content": system_prompt})

        payload = {
            "model": _model,
            "messages": _messages,
            "temperature": temperature if temperature is not None else KIMI_TEMPERATURE,
        }
        # Default to the configured KIMI_MAX_TOKENS when the caller does not
        # specify one. Reasoning models (k3) spend part of the budget on
        # reasoning_content, so too small a cap (e.g. 2000) can starve the
        # actual content and yield an empty response.
        payload["max_tokens"] = max_tokens if max_tokens else KIMI_MAX_TOKENS

        url = f"{self.base_url}/chat/completions"
        start = time.time()
        last_error = None

        for attempt in range(KIMI_MAX_RETRIES):
            try:
                resp = self._client.post(url, headers=self._headers(), json=payload)
                resp.raise_for_status()
                data = resp.json()
                duration_ms = int((time.time() - start) * 1000)

                choice = data.get("choices", [{}])[0]
                msg = choice.get("message", {})
                usage = data.get("usage", {})

                result = {
                    "content": msg.get("content", ""),
                    "reasoning_content": msg.get("reasoning_content", ""),
                    "model": data.get("model", _model),
                    "usage": usage,
                    "finish_reason": choice.get("finish_reason"),
                }

                # Warn when the output budget was exhausted: reasoning models
                # may return empty content with finish_reason=length. Callers
                # should treat empty content as a failure signal.
                if result["finish_reason"] == "length" and not result["content"]:
                    logger.warning(
                        "Kimi response truncated: finish_reason=length with empty "
                        "content (max_tokens=%s). Consider raising KIMI_MAX_TOKENS.",
                        payload.get("max_tokens"),
                    )

                self._log_call(
                    endpoint="/chat/completions",
                    model=_model,
                    messages=_messages,
                    response=result,
                    error=None,
                    duration_ms=duration_ms,
                    prompt_tokens=usage.get("prompt_tokens", 0),
                    completion_tokens=usage.get("completion_tokens", 0),
                    total_tokens=usage.get("total_tokens", 0),
                )
                return result

            except httpx.HTTPStatusError as e:
                last_error = f"HTTP {e.response.status_code}: {e.response.text[:200]}"
                if e.response.status_code in (401, 403, 404):
                    break  # Don't retry auth/not-found errors
                if attempt < KIMI_MAX_RETRIES - 1:
                    time.sleep(2 ** attempt)
            except Exception as e:
                last_error = str(e)
                if attempt < KIMI_MAX_RETRIES - 1:
                    time.sleep(2 ** attempt)

        duration_ms = int((time.time() - start) * 1000)
        self._log_call(
            endpoint="/chat/completions",
            model=_model,
            messages=_messages,
            response=None,
            error=last_error,
            duration_ms=duration_ms,
        )
        raise RuntimeError(f"Kimi API call failed after {KIMI_MAX_RETRIES} retries: {last_error}")

    def chat_json(
        self,
        messages: List[Dict[str, str]],
        system_prompt: Optional[str] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """Chat completion that parses JSON response.

        Injects instruction to return valid JSON, then parses the response.
        """
        _messages = list(messages)
        _messages.append({
            "role": "user",
            "content": "Return your response as valid JSON only. Do not include markdown code fences or any text outside the JSON object."
        })
        result = self.chat(_messages, system_prompt=system_prompt, **kwargs)
        content = result.get("content", "").strip()

        # Strip code fences if present
        if content.startswith("```"):
            lines = content.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            content = "\n".join(lines).strip()

        try:
            parsed = json.loads(content)
            result["parsed_json"] = parsed
            return result
        except json.JSONDecodeError as e:
            result["json_parse_error"] = str(e)
            result["raw_content"] = content
            logger.warning(f"Failed to parse JSON response: {e}\nContent: {content[:300]}")
            return result

    def list_models(self) -> List[Dict[str, Any]]:
        """List available models."""
        if not self.is_configured:
            raise RuntimeError("KIMI_API_KEY is not configured")
        url = f"{self.base_url}/models"
        resp = self._client.get(url, headers=self._headers())
        resp.raise_for_status()
        return resp.json().get("data", [])

    def health_check(self) -> Dict[str, Any]:
        """Test API key validity and return model info."""
        try:
            models = self.list_models()
            return {
                "status": "ok",
                "configured": self.is_configured,
                "base_url": self.base_url,
                "default_model": self.model,
                "available_models": [
                    {"id": m.get("id"), "display_name": m.get("display_name"),
                     "context_length": m.get("context_length")}
                    for m in models
                ],
            }
        except Exception as e:
            return {
                "status": "error",
                "configured": self.is_configured,
                "base_url": self.base_url,
                "default_model": self.model,
                "error": str(e),
            }


# Global singleton
_kimi_client: Optional[KimiClient] = None


def get_kimi_client() -> KimiClient:
    """Get or create global KimiClient singleton."""
    global _kimi_client
    if _kimi_client is None:
        _kimi_client = KimiClient()
    return _kimi_client