salim0986/graph-bug-ai
0
1"""2M7 — Multi-provider LLM abstraction backed by LiteLLM.3 4`LLMClient` is the single entry point for all AI generation in Graph Bug.5It translates Graph Bug's internal tier names (flash / pro / thinking) to6provider-specific LiteLLM model strings, so the rest of the codebase stays7provider-agnostic.8 9Supported providers: gemini · anthropic · openai · mistral · ollama10 11Usage:12 config = LLMConfig(provider="anthropic", api_key="sk-ant-...")13 client = LLMClient(config)14 text = await client.generate(tier="pro", prompt="...")15 review = await client.generate_structured(tier="pro", prompt="...", schema=ReviewOutput)16"""17 18from __future__ import annotations19 20import asyncio21import json22import re23from dataclasses import dataclass, field24from typing import Dict, Optional, Type25 26import litellm27from pydantic import BaseModel28 29from .logger import setup_logger30from .review_schema import ReviewOutput31 32logger = setup_logger(__name__)33 34# Suppress verbose litellm logging unless DEBUG is explicitly requested.35litellm.set_verbose = False # type: ignore[attr-defined]36 37 38# ---------------------------------------------------------------------------39# Tier → provider-specific model map40# ---------------------------------------------------------------------------41 42DEFAULT_TIER_MAP: Dict[str, Dict[str, str]] = {43 "gemini": {44 "flash": "gemini/gemini-2.5-flash",45 "pro": "gemini/gemini-2.5-pro",46 "thinking": "gemini/gemini-2.5-flash",47 },48 "anthropic": {49 "flash": "anthropic/claude-haiku-4-5-20251001",50 "pro": "anthropic/claude-sonnet-4-6",51 "thinking": "anthropic/claude-opus-4-8",52 },53 "openai": {54 "flash": "openai/gpt-4o-mini",55 "pro": "openai/gpt-4o",56 "thinking": "openai/o1",57 },58 "mistral": {59 "flash": "mistral/mistral-small-latest",60 "pro": "mistral/mistral-large-latest",61 "thinking": "mistral/mistral-large-latest",62 },63 "ollama": {64 "flash": "ollama/qwen2.5-coder:7b",65 "pro": "ollama/qwen2.5-coder:32b",66 "thinking": "ollama/qwen2.5-coder:32b",67 },68}69 70# Environment variable names for each provider's API key.71_PROVIDER_ENV: Dict[str, str] = {72 "gemini": "GEMINI_API_KEY",73 "anthropic": "ANTHROPIC_API_KEY",74 "openai": "OPENAI_API_KEY",75 "mistral": "MISTRAL_API_KEY",76}77 78 79# ---------------------------------------------------------------------------80# Configuration81# ---------------------------------------------------------------------------82 83@dataclass84class LLMConfig:85 provider: str = "gemini"86 api_key: Optional[str] = None87 tier_map: Dict[str, Dict[str, str]] = field(default_factory=lambda: DEFAULT_TIER_MAP)88 temperature: float = 0.789 max_tokens: int = 819290 retry_attempts: int = 391 retry_delay: float = 2.092 93 94# ---------------------------------------------------------------------------95# Client96# ---------------------------------------------------------------------------97 98class LLMClient:99 """100 Multi-provider LLM client backed by LiteLLM.101 102 All generation goes through `generate()` (raw text) or103 `generate_structured()` (Pydantic model). The model selected depends on104 the `tier` argument and the provider configured in `LLMConfig`.105 """106 107 def __init__(self, config: LLMConfig) -> None:108 self.config = config109 # Do NOT write to os.environ here — that is a global mutation and races110 # when multiple users trigger concurrent reviews with different providers.111 # Instead, pass api_key directly to each litellm.completion() call.112 # M10: cumulative usage stats for this client instance.113 self.tokens_in: int = 0114 self.tokens_out: int = 0115 self.total_cost: float = 0.0116 117 # ----------------------------------------------------------------118 # Public API119 # ----------------------------------------------------------------120 121 def get_usage_stats(self) -> Dict[str, object]:122 """Return cumulative token and cost stats for this client instance."""123 return {124 "tokens_in": self.tokens_in,125 "tokens_out": self.tokens_out,126 "total_tokens": self.tokens_in + self.tokens_out,127 "total_cost": self.total_cost,128 }129 130 def resolve_model(self, tier: str) -> str:131 """Return the LiteLLM model string for the given tier."""132 provider_map = self.config.tier_map.get(self.config.provider, {})133 model = provider_map.get(tier)134 if not model:135 model = provider_map.get("pro") or next(iter(provider_map.values()), "gemini/gemini-2.5-flash")136 return model137 138 async def generate(self, tier: str, prompt: str) -> str:139 """140 Generate raw text for the given tier.141 Retries with exponential back-off on transient errors.142 Records token usage, cost, and a LangFuse trace on success.143 """144 model = self.resolve_model(tier)145 for attempt in range(self.config.retry_attempts):146 try:147 logger.info(148 f"[M7] LLMClient.generate provider={self.config.provider} "149 f"model={model} attempt={attempt + 1}"150 )151 call_kwargs: Dict[str, object] = dict(152 model=model,153 messages=[{"role": "user", "content": prompt}],154 temperature=self.config.temperature,155 max_tokens=self.config.max_tokens,156 )157 if self.config.api_key:158 call_kwargs["api_key"] = self.config.api_key159 response = await asyncio.to_thread(litellm.completion, **call_kwargs)160 text: str = response.choices[0].message.content or ""161 logger.info(f"[M7] Generated {len(text)} chars")162 # M10: track tokens/cost and record in LangFuse163 in_tok, out_tok, cost = self._extract_usage(response)164 self.tokens_in += in_tok165 self.tokens_out += out_tok166 self.total_cost += cost167 self._log_usage(model, in_tok + out_tok)168 self._observe("generate", model, prompt, text, in_tok, out_tok,169 tier=tier, provider=self.config.provider)170 return text171 except Exception as e:172 logger.warning(f"[M7] generate attempt {attempt + 1} failed: {e}")173 if attempt < self.config.retry_attempts - 1:174 await asyncio.sleep(self.config.retry_delay * (2 ** attempt))175 else:176 logger.error(f"[M7] Max retries exceeded for generate: {e}")177 raise178 return ""179 180 async def generate_structured(181 self,182 tier: str,183 prompt: str,184 schema: Type[BaseModel] = ReviewOutput,185 ) -> BaseModel:186 """187 Generate a structured JSON response that conforms to `schema`.188 189 Uses LiteLLM's `response_format={"type": "json_object"}` where190 supported (OpenAI, Anthropic tool-use mode, etc.). For providers191 that don't support JSON mode, we generate text and parse it.192 Falls back to a minimal schema instance on any error so callers193 always receive a usable object.194 """195 model = self.resolve_model(tier)196 for attempt in range(self.config.retry_attempts):197 try:198 logger.info(199 f"[M7] LLMClient.generate_structured provider={self.config.provider} "200 f"model={model} schema={schema.__name__} attempt={attempt + 1}"201 )202 call_kwargs_s: Dict[str, object] = dict(203 model=model,204 messages=[{"role": "user", "content": prompt}],205 temperature=self.config.temperature,206 max_tokens=self.config.max_tokens,207 response_format={"type": "json_object"},208 )209 if self.config.api_key:210 call_kwargs_s["api_key"] = self.config.api_key211 response = await asyncio.to_thread(litellm.completion, **call_kwargs_s)212 text = response.choices[0].message.content or ""213 logger.info(f"[M7] Structured response: {len(text)} chars")214 # M10: track tokens/cost and record in LangFuse215 in_tok, out_tok, cost = self._extract_usage(response)216 self.tokens_in += in_tok217 self.tokens_out += out_tok218 self.total_cost += cost219 self._log_usage(model, in_tok + out_tok)220 self._observe("generate_structured", model, prompt, text[:500], in_tok, out_tok,221 tier=tier, provider=self.config.provider, schema=schema.__name__)222 return self._parse_json(text, schema)223 except Exception as e:224 logger.warning(f"[M7] generate_structured attempt {attempt + 1} failed: {e}")225 if attempt < self.config.retry_attempts - 1:226 await asyncio.sleep(self.config.retry_delay * (2 ** attempt))227 else:228 logger.error(f"[M7] Max retries exceeded for generate_structured: {e}")229 return self._fallback(schema, str(e))230 return self._fallback(schema, "Max retries exceeded")231 232 # ----------------------------------------------------------------233 # Internal helpers234 # ----------------------------------------------------------------235 236 @staticmethod237 def _extract_usage(response: object) -> tuple[int, int, float]:238 """Return (input_tokens, output_tokens, cost_usd) from a LiteLLM response."""239 usage = getattr(response, "usage", None)240 in_tok = int(getattr(usage, "prompt_tokens", 0) or 0)241 out_tok = int(getattr(usage, "completion_tokens", 0) or 0)242 cost = 0.0243 try:244 cost = float(litellm.completion_cost(completion_response=response) or 0.0)245 except Exception:246 pass247 return in_tok, out_tok, cost248 249 @staticmethod250 def _log_usage(model: str, total_tokens: int) -> None:251 """Wire token counts into the global TokenBudget."""252 if total_tokens <= 0:253 return254 try:255 from .cost_optimizer import token_budget # local import to avoid circular256 token_budget.log_usage(model, total_tokens)257 except Exception:258 pass259 260 @staticmethod261 def _observe(262 name: str,263 model: str,264 prompt: str,265 output: str,266 in_tok: int,267 out_tok: int,268 **meta: object,269 ) -> None:270 """Record the call in LangFuse (no-op when LangFuse not configured)."""271 try:272 from .observability import observe_llm_call273 observe_llm_call(name, model, prompt, output, in_tok, out_tok, **meta)274 except Exception:275 pass276 277 @staticmethod278 def _parse_json(text: str, schema: Type[BaseModel]) -> BaseModel:279 """Strip fences, parse JSON, validate against schema. Falls back gracefully."""280 text = text.strip()281 text = re.sub(r"^```(?:json)?\s*\n?", "", text)282 text = re.sub(r"\n?```\s*$", "", text)283 text = text.strip()284 try:285 data = json.loads(text)286 return schema.model_validate(data)287 except Exception as e:288 logger.warning(f"[M7] JSON parse failed: {e}; returning fallback")289 return LLMClient._fallback(schema, text[:500] or "Unparseable output")290 291 @staticmethod292 def _fallback(schema: Type[BaseModel], reason: str) -> BaseModel:293 """Return a minimal valid instance for schemas that share ReviewOutput fields."""294 try:295 return schema( # type: ignore[call-arg]296 summary=reason,297 overall_assessment="comment",298 risk_level="low",299 )300 except Exception:301 return schema() # type: ignore[call-arg]302