"""A thin, provider-agnostic wrapper over the OpenAI-compatible chat API. Works against self-hosted vLLM, OpenAI, OpenRouter, Berget, Gemini's compatibility endpoint — anything speaking the OpenAI chat protocol. The differences between them that actually bite in production are handled in one place here: * vLLM accepts ``extra_body`` sampler knobs (``repetition_penalty``, ``chat_template_kwargs``) that hosted providers reject with a 4xx. * OpenAI reasoning models (o1/o3/o4/gpt-5) require ``max_completion_tokens`` instead of ``max_tokens``, and refuse any temperature other than 1. * Some models emit ```` blocks inline in ``content`` rather than in ``reasoning_content``; those must never reach a user. Error contract, preserved from the predecessor: on API failure ``generate`` returns an error *string* rather than raising. Call sites branch on ``isinstance(response, str)``. """ from __future__ import annotations import json import os import re import traceback from typing import Any, Dict, Generator, List, Optional, Tuple, Type from openai import OpenAI from openai.types.chat import ChatCompletion from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_message import ( ChatCompletionMessage as _OpenAIChatCompletionMessage, ) from pydantic import BaseModel from .tools import execute_tool, parse_function_call_arguments _THINK_RE = re.compile(r".*?\s*", re.DOTALL | re.IGNORECASE) # Models that reject `max_tokens` and any temperature but 1. _REASONING_PREFIXES = ("o1", "o3", "o4", "gpt-5") # Sampler fields only self-hosted vLLM understands. _VLLM_ONLY_EXTRA_BODY = ("repetition_penalty", "chat_template_kwargs") class ChatCompletionMessage(_OpenAIChatCompletionMessage): """Assistant message, extended with the structured-output fields. When ``generate(format=...)`` is used, ``content`` holds the *parsed model instance* rather than text, and ``parsed`` / ``content_text`` carry the instance and its raw JSON respectively. """ model_config = {"extra": "allow"} def _strip_think(text: str) -> str: return _THINK_RE.sub("", text).strip() def _is_reasoning_model(model: str) -> bool: name = (model or "").split("/")[-1].lower() return name.startswith(_REASONING_PREFIXES) class LLM: """One configured connection to a chat model. Args: system_message: Seeds ``self.messages`` when no explicit history is passed. temperature: Default sampling temperature; overridable per call. model: Model identifier sent to the provider. base_url: OpenAI-compatible endpoint, including the ``/v1`` suffix. api_key: Provider key. Its presence also marks this as an *external* provider, which suppresses the vLLM-only ``extra_body`` fields. think: Default reasoning mode. When false on self-hosted vLLM, chain-of-thought is disabled at the template level so no reasoning tokens are generated. """ def __init__( self, system_message: str = "You are an assistant.", temperature: float = 0.01, model: Optional[str] = None, max_length_answer: int = 3000, messages: Optional[List[dict]] = None, chat: bool = True, tools: Optional[list] = None, think: bool = False, timeout: int = 240, silent: bool = False, presence_penalty: float = 0.3, top_p: float = 0.9, extra_body: Optional[Dict[str, Any]] = None, base_url: Optional[str] = None, api_key: Optional[str] = None, max_retries: int = 4, ) -> None: self.model = model or os.getenv("LLM_MODEL", "smart") self.system_message = system_message self.messages = messages or [{"role": "system", "content": system_message}] self.max_length_answer = max_length_answer self.chat = chat self.tools = tools or [] self.think = think self.silent = silent self.options = { "temperature": temperature, "presence_penalty": presence_penalty, "top_p": top_p, } # repetition_penalty > 1.0 damps already-seen tokens; 1.2 breaks generation # loops without measurably hurting quality. self.extra_body = extra_body if extra_body is not None else {"repetition_penalty": 1.2} self._api_key = api_key self.base_url = base_url or os.getenv("LLM_DIRECT_URL") or "" if not self.base_url: raise ValueError( "No LLM endpoint configured. Pass base_url= or set LLM_DIRECT_URL." ) # max_retries covers transient 408/409/429/5xx with exponential backoff. # Above the SDK default of 2 because a self-hosted vLLM engine crash returns # 500 and usually restarts within seconds — the extra attempts ride it out. self.client = OpenAI( base_url=self.base_url, api_key=self._api_key or os.getenv("LLM_BEARER", "NONE"), timeout=timeout, max_retries=max_retries, ) # -- request assembly ----------------------------------------------------- @property def _is_external_provider(self) -> bool: """A caller-supplied key means a hosted provider, not our own vLLM.""" return bool(self._api_key) def _build_extra_body(self, think: Optional[bool]) -> Optional[dict]: body = dict(self.extra_body or {}) effective_think = self.think if think is None else think if not effective_think and not self._is_external_provider: ctk = dict(body.get("chat_template_kwargs") or {}) ctk["enable_thinking"] = False body["chat_template_kwargs"] = ctk if self._is_external_provider: for key in _VLLM_ONLY_EXTRA_BODY: body.pop(key, None) return body or None def _sampling_kwargs(self, model: str, temperature: Optional[float], max_tokens: Optional[int]) -> dict: temp = self.options["temperature"] if temperature is None else temperature limit = max_tokens or self.max_length_answer if _is_reasoning_model(model): # These models accept only the default temperature and a different token field. return {"max_completion_tokens": limit} return { "temperature": temp, "top_p": self.options["top_p"], "max_tokens": limit, } def _create(self, **kwargs) -> ChatCompletion: """Call the API, retrying once if the token-limit field name is rejected.""" try: return self.client.chat.completions.create(**kwargs) except Exception as exc: swapped = _swap_token_param(kwargs, exc) if swapped is None: raise _friendly_error(exc, self.base_url) from exc return self.client.chat.completions.create(**swapped) # -- public API ----------------------------------------------------------- def generate( self, query: Optional[str] = None, *, messages: Optional[List[dict]] = None, tools: Optional[list] = None, model: Optional[str] = None, temperature: Optional[float] = None, format: Optional[Type[BaseModel]] = None, stream: bool = False, think: Optional[bool] = None, max_tokens: Optional[int] = None, extra_body: Optional[Dict[str, Any]] = None, auto_execute_tools: bool = True, ): """Run one completion. Returns a :class:`ChatCompletionMessage`, a generator of ``(kind, chunk)`` pairs when ``stream`` is set, or an error string when the call fails. Args: query: Convenience for a single user turn; ignored if ``messages`` is given. messages: Full conversation to send. Also becomes this instance's history. tools: Tool schemas to advertise. Defaults to the instance's tools. format: Pydantic model requesting structured output via JSON schema. auto_execute_tools: Execute returned tool calls and append their results to the history. It does *not* continue the conversation — the caller decides whether to call again. """ if messages is not None: self.messages = list(messages) elif query is not None: self.messages.append({"role": "user", "content": query}) resolved_model = model or self.model if extra_body is not None: self.extra_body = extra_body try: if format is not None: return self._generate_structured( resolved_model, format, temperature, max_tokens, think ) request = { "model": resolved_model, "messages": self.messages, "extra_body": self._build_extra_body(think), **self._sampling_kwargs(resolved_model, temperature, max_tokens), } tools_to_use = self.tools if tools is None else tools if tools_to_use: request["tools"] = tools_to_use if stream: request["stream"] = True return self._read_stream(self.client.chat.completions.create(**request)) message = self._create(**request).choices[0].message if auto_execute_tools and getattr(message, "tool_calls", None): self._run_tool_calls(message.tool_calls) if isinstance(message.content, str): message.content = _strip_think(message.content) self.messages.append({"role": "assistant", "content": message.content}) if not self.chat: self.messages = self.messages[:1] return message except Exception as exc: if not self.silent: traceback.print_exc() return f"LLM request failed: {exc}" # -- structured output ---------------------------------------------------- def _generate_structured( self, model: str, format: Type[BaseModel], temperature: Optional[float], max_tokens: Optional[int], think: Optional[bool], ) -> ChatCompletionMessage: # vLLM's json_schema mode rejects `role: tool` turns, so fold them into # user turns. Done on a copy — the caller's history stays intact. messages = [ {"role": "user", "content": f"Tool output:\n{m.get('content', '')}"} if m.get("role") == "tool" else m for m in self.messages ] response = self._create( model=model, messages=messages, extra_body=self._build_extra_body(think), response_format={ "type": "json_schema", "json_schema": {"name": format.__name__, "schema": format.model_json_schema()}, }, **self._sampling_kwargs(model, temperature, max_tokens), ) content_text = response.choices[0].message.content or "" parsed = format.model_validate_json(_extract_json(content_text)) message = ChatCompletionMessage.model_construct(role="assistant", content=parsed) message.parsed = parsed message.parsed_dict = parsed.model_dump() message.content_text = content_text return message # -- tool execution ------------------------------------------------------- def _run_tool_calls(self, tool_calls) -> None: """Execute each returned tool call, appending results as `tool` messages. A failing tool appends its error rather than raising, so the model can see what went wrong and correct itself on the next turn. """ for call in tool_calls: fn = getattr(call, "function", None) if fn is None: continue name = getattr(fn, "name", None) try: args = parse_function_call_arguments(getattr(fn, "arguments", None)) result = execute_tool(name, args) content = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False) except Exception as exc: if not self.silent: print(f"[llm] tool {name} failed: {exc}") content = json.dumps({"error": str(exc)}, ensure_ascii=False) self.messages.append({"role": "tool", "name": name or "unknown", "content": content}) # -- streaming ------------------------------------------------------------ def _read_stream(self, response) -> Generator[Tuple[str, str], None, None]: """Yield ``("thinking" | "content", text)`` pairs as they arrive. Reasoning arrives either in ``reasoning_content`` or inline as ```` blocks; both are surfaced as "thinking" so callers can render or drop them. """ in_think_block = False for chunk in response: if not chunk.choices: continue delta = chunk.choices[0].delta reasoning = getattr(delta, "reasoning_content", None) if reasoning: yield "thinking", reasoning text = getattr(delta, "content", None) if not text: continue if "" in text: in_think_block = True text = text.split("", 1)[0] if "" in text: in_think_block = False text = text.split("", 1)[1] if not text: continue yield ("thinking" if in_think_block else "content"), text # -- module helpers ----------------------------------------------------------- def _extract_json(text: str) -> str: """Pull the JSON object out of a response that may carry stray prose.""" text = _strip_think(text).strip() if text.startswith("{"): return text match = re.search(r"\{.*\}", text, re.DOTALL) return match.group(0) if match else text def _swap_token_param(kwargs: dict, exc: Exception) -> Optional[dict]: """Retry payload with the other token-limit field, if that's what was rejected. Providers disagree about `max_tokens` vs `max_completion_tokens`, and the model lists that require each keep changing. Rather than track them, react to the error. """ text = str(exc).lower() if "max_tokens" not in text and "max_completion_tokens" not in text: return None alternates = {"max_tokens": "max_completion_tokens", "max_completion_tokens": "max_tokens"} for current, replacement in alternates.items(): if current in kwargs: swapped = dict(kwargs) swapped[replacement] = swapped.pop(current) swapped.pop("temperature", None) # reasoning models reject it too return swapped return None def _friendly_error(exc: Exception, base_url: str) -> Exception: """Turn provider errors into something a deployer can act on.""" text = str(exc).lower() if any(s in text for s in ("connection", "timeout", "refused", "unreachable")): return RuntimeError(f"LLM endpoint unreachable at {base_url}. Is the server running? ({exc})") if any(s in text for s in ("401", "403", "invalid api key", "unauthorized")): return RuntimeError(f"LLM provider rejected the API key. ({exc})") return exc