"""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, prompt_preview=json.dumps(messages[:3], ensure_ascii=False)[:500], response_preview=( json.dumps(response, ensure_ascii=False)[:1000] 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, } if max_tokens: payload["max_tokens"] = 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"), } 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