Israelbliz/User-Modeling-Agent

1"""LLM client — provider-agnostic wrapper for OpenAI and Gemini.2 3Why a wrapper:4  - Two-tier model selection (reasoning vs bulk) without scattering model names5  - Two-provider support (OpenAI / Gemini), switchable via LLM_PROVIDER env var6  - Built-in retry on transient errors7  - Pydantic-validated structured outputs8  - Single chokepoint for logging / token accounting9 10The provider is chosen at construction time from settings.llm_provider:11  - 'openai' (default) → gpt-4o + gpt-4o-mini via langchain-openai12  - 'gemini'           → gemini-2.5-flash + gemini-2.5-flash-lite via langchain-google-genai13 14Both providers share the same interface, so calling code never needs to15care which one is active.16 17Usage:18    llm = LLMClient()19    answer = llm.complete("Why is the sky blue?", model="bulk")20    parsed = llm.structured(prompt, ReviewOutput, model="reasoning")21"""22from __future__ import annotations23 24import logging25import time26from typing import Any, Type, TypeVar27 28from langchain_core.language_models import BaseChatModel29from langchain_core.output_parsers import PydanticOutputParser30from langchain_core.prompts import ChatPromptTemplate31from pydantic import BaseModel32from tenacity import (retry, stop_after_attempt, wait_exponential,33                      retry_if_exception)34 35from core.config import settings36 37log = logging.getLogger(__name__)38 39T = TypeVar("T", bound=BaseModel)40 41 42def _build_openai_models(temp_reasoning: float, temp_bulk: float) -> tuple[BaseChatModel, BaseChatModel]:43    """Construct OpenAI reasoning + bulk models."""44    from langchain_openai import ChatOpenAI45    if not settings.openai_api_key:46        raise RuntimeError(47            "LLM_PROVIDER=openai but OPENAI_API_KEY not set. "48            "Add it to .env or switch LLM_PROVIDER to 'gemini'."49        )50    reasoning = ChatOpenAI(51        model=settings.openai_reasoning_model,52        temperature=temp_reasoning,53        api_key=settings.openai_api_key,54    )55    bulk = ChatOpenAI(56        model=settings.openai_bulk_model,57        temperature=temp_bulk,58        api_key=settings.openai_api_key,59    )60    return reasoning, bulk61 62 63def _build_gemini_models(temp_reasoning: float, temp_bulk: float) -> tuple[BaseChatModel, BaseChatModel]:64    """Construct Gemini reasoning + bulk models."""65    try:66        from langchain_google_genai import ChatGoogleGenerativeAI67    except ImportError as e:68        raise ImportError(69            "LLM_PROVIDER=gemini but langchain-google-genai is not installed. "70            "Run: pip install langchain-google-genai"71        ) from e72 73    if not settings.gemini_api_key:74        raise RuntimeError(75            "LLM_PROVIDER=gemini but GEMINI_API_KEY not set. "76            "Get a key at https://aistudio.google.com/apikey and add it to .env."77        )78    reasoning = ChatGoogleGenerativeAI(79        model=settings.gemini_reasoning_model,80        temperature=temp_reasoning,81        google_api_key=settings.gemini_api_key,82    )83    bulk = ChatGoogleGenerativeAI(84        model=settings.gemini_bulk_model,85        temperature=temp_bulk,86        google_api_key=settings.gemini_api_key,87    )88    return reasoning, bulk89 90 91def _should_failover(exc: Exception) -> bool:92    """Decide whether an exception warrants trying the fallback provider.93 94    Triggers on quota / rate-limit errors AND on transient service errors95    (5xx, timeouts, connection failures) — i.e. any sign the primary96    provider is currently unable to serve the request. Does NOT trigger on97    clear client-side mistakes (bad request, malformed schema), which the98    fallback could not fix either.99    """100    text = f"{type(exc).__name__} {exc}".lower()101    quota = ("429", "quota", "rate limit", "ratelimit", "resource exhausted",102             "resource_exhausted", "exceeded", "too many requests")103    transient = ("500", "502", "503", "504", "overloaded", "unavailable",104                 "timeout", "timed out", "connection", "internal error",105                 "service")106    return any(s in text for s in quota + transient)107 108 109def _is_quota_error(exc: Exception) -> bool:110    """True only for rate-limit / quota-exhausted errors (used to skip the111    slow retry-backoff so failover happens fast on quota limits)."""112    text = f"{type(exc).__name__} {exc}".lower()113    signals = ("429", "quota", "rate limit", "ratelimit",114               "resource exhausted", "resource_exhausted",115               "exceeded", "too many requests")116    return any(s in text for s in signals)117 118 119class LLMClient:120    """Two-tier, two-provider LLM client with automatic failover.121 122    Tier 'reasoning' → flagship model (gpt-4o / gemini-2.5-flash).123    Tier 'bulk' → cheap/fast model (gpt-4o-mini / gemini-2.5-flash-lite).124 125    Failover: the primary provider is chosen from settings.llm_provider.126    If the other provider's API key is also present, it is built as a127    fallback. When a call to the primary fails with a quota / rate-limit128    error, the identical call is retried on the fallback provider — so a129    judge hitting the free Gemini tier's limit mid-demo never sees an130    error. If no fallback key is configured, the client behaves exactly131    as a single-provider client.132    """133 134    def __init__(self, temperature_reasoning: float = 0.7,135                 temperature_bulk: float = 0.3,136                 provider: str | None = None):137        self.provider = (provider or settings.llm_provider).lower()138        log.info(f"LLMClient initializing with primary provider={self.provider!r}")139 140        if self.provider == "openai":141            self._reasoning, self._bulk = _build_openai_models(142                temperature_reasoning, temperature_bulk)143        elif self.provider == "gemini":144            self._reasoning, self._bulk = _build_gemini_models(145                temperature_reasoning, temperature_bulk)146        else:147            raise ValueError(148                f"Unknown LLM_PROVIDER={self.provider!r}; expected 'openai' or 'gemini'")149 150        # Build the OTHER provider as a fallback, if its key is available.151        self.fallback_provider: str | None = None152        self._fb_reasoning: BaseChatModel | None = None153        self._fb_bulk: BaseChatModel | None = None154        try:155            if self.provider == "gemini" and settings.openai_api_key:156                self._fb_reasoning, self._fb_bulk = _build_openai_models(157                    temperature_reasoning, temperature_bulk)158                self.fallback_provider = "openai"159            elif self.provider == "openai" and settings.gemini_api_key:160                self._fb_reasoning, self._fb_bulk = _build_gemini_models(161                    temperature_reasoning, temperature_bulk)162                self.fallback_provider = "gemini"163        except Exception as e:  # fallback is best-effort; never block startup164            log.warning(f"Fallback provider unavailable, continuing without it: {e}")165            self.fallback_provider = None166 167        if self.fallback_provider:168            log.info(f"Failover enabled: {self.provider} → {self.fallback_provider} "169                     f"on quota errors")170        else:171            log.info("No fallback provider configured; running single-provider")172 173    def _model(self, tier: str) -> BaseChatModel:174        if tier == "reasoning":175            return self._reasoning176        if tier == "bulk":177            return self._bulk178        raise ValueError(f"Unknown tier {tier!r}; expected 'reasoning' or 'bulk'")179 180    def _fb_model(self, tier: str) -> BaseChatModel | None:181        if tier == "reasoning":182            return self._fb_reasoning183        if tier == "bulk":184            return self._fb_bulk185        return None186 187    # ──────────────────────────────────────────────────────────────────188    # Free-form completion189    # ──────────────────────────────────────────────────────────────────190    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10),191           retry=retry_if_exception(lambda e: not _is_quota_error(e)))192    def complete(self, prompt: str, model: str = "bulk",193                 system: str | None = None) -> str:194        messages: list[Any] = []195        if system:196            messages.append(("system", system))197        messages.append(("human", "{input}"))198        template = ChatPromptTemplate.from_messages(messages)199 200        def _run(model_obj: BaseChatModel) -> str:201            t0 = time.time()202            result = (template | model_obj).invoke({"input": prompt})203            content = result.content204            if isinstance(content, list):205                content = "".join(206                    p.get("text", "") if isinstance(p, dict) else str(p)207                    for p in content)208            log.info(f"LLM complete [{model}] {time.time() - t0:.2f}s · "209                     f"prompt {len(prompt)} chars · output {len(content)} chars")210            return content211 212        try:213            return _run(self._model(model))214        except Exception as e:215            fb = self._fb_model(model)216            if fb is not None and _should_failover(e):217                log.warning(f"Primary provider {self.provider} failed "218                            f"({type(e).__name__}); failing over to "219                            f"{self.fallback_provider}")220                return _run(fb)221            raise222 223    # ──────────────────────────────────────────────────────────────────224    # Structured output — pydantic-validated225    # ──────────────────────────────────────────────────────────────────226    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10),227           retry=retry_if_exception(lambda e: not _is_quota_error(e)))228    def structured(self, prompt: str, schema: Type[T], model: str = "reasoning",229                   system: str | None = None) -> T:230        """Run prompt, parse output into the given Pydantic schema.231 232        Uses LangChain's PydanticOutputParser. On a quota / rate-limit error233        from the primary provider, the same call is retried on the fallback.234        """235        parser = PydanticOutputParser(pydantic_object=schema)236        format_instructions = parser.get_format_instructions()237 238        messages: list[Any] = []239        if system:240            messages.append(("system", system))241        messages.append(("human", "{input}\n\n{format_instructions}"))242        template = ChatPromptTemplate.from_messages(messages)243 244        def _run(model_obj: BaseChatModel) -> T:245            t0 = time.time()246            chain = template | model_obj | parser247            out = chain.invoke({248                "input": prompt,249                "format_instructions": format_instructions,250            })251            log.info(f"LLM structured [{model}] {time.time() - t0:.2f}s · "252                     f"schema {schema.__name__} · prompt {len(prompt)} chars")253            return out254 255        try:256            return _run(self._model(model))257        except Exception as e:258            fb = self._fb_model(model)259            if fb is not None and _should_failover(e):260                log.warning(f"Primary provider {self.provider} failed "261                            f"({type(e).__name__}); failing over to "262                            f"{self.fallback_provider}")263                return _run(fb)264            raise265