llm_client.py

"""Self-contained OpenAI-compatible LLM client for GraphAnything.

Targets any `chat.completions`-shaped HTTP endpoint:

  - **Local vLLM** (`vllm serve ...`)
  - **llama.cpp** server (`./server`)
  - **Ollama** (OpenAI-compat bridge on :11434)
  - **LM Studio**
  - **OpenAI** itself, or any commercial OpenAI-compatible host

The client talks plain HTTP via `requests`; no vendor SDK, no model loading
in-process. Point it at a URL, set a model name, optionally a key, and go.

Public surface (the `chat` / `chat_json` interface the rest of the package
expects):

    client = make_client()
    text  = client.chat("Summarise this in one sentence: ...")
    data  = client.chat_json("Return {x: 1, y: 2} as JSON.")

Both methods accept either a plain string (treated as the user message) or
a `[{"role": ..., "content": ...}]` list. `chat_json` requests structured
JSON output via `response_format` when the server supports it, and falls
back to a robust brace-balance scan otherwise (so it also works against
servers that ignore the field, including reasoning models that emit
`<think>...</think>` blocks before the answer).

Configuration (env vars, in priority order):

  - **API base URL**: `GA_API_BASE` → `OPENAI_API_BASE` → `OPENAI_BASE_URL`
                     → `API_BASE`. Default `http://localhost:8000/v1`.
  - **Model name**:   `GA_MODEL` → `OPENAI_MODEL` → `SUMMARY_MODEL_NAME`.
                     Required for chat calls.
  - **API key**:      `GA_API_KEY` → `OPENAI_API_KEY` → `API_KEY`.
                     Optional — many local servers accept any string.
  - **Timeout (s)**:  `GA_HTTP_TIMEOUT`. Default 600.

Aliasing the legacy upstream env vars (`API_KEY`, `API_BASE`,
`SUMMARY_MODEL_NAME`) means an existing setup continues to work without
edits.
"""
from __future__ import annotations

import json as _json
import os
import re
from typing import Any, Optional, Sequence


_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"


def _strip_think(text: str) -> str:
    """Strip reasoning <think>...</think> blocks and any markdown code fences."""
    if not text:
        return text
    if _THINK_CLOSE in text:
        text = text.rsplit(_THINK_CLOSE, 1)[-1]
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    text = re.sub(r"<think>.*$", "", text, flags=re.DOTALL)
    text = re.sub(r"```(?:json)?", "", text)
    return text.strip()


def _extract_json(text: str) -> Optional[str]:
    """Find a parseable top-level JSON object inside `text`."""
    text = _strip_think(text)
    if not text:
        return None
    depth = 0
    start = -1
    in_string = False
    escape = False
    candidates: list[str] = []
    for i, ch in enumerate(text):
        if escape:
            escape = False
            continue
        if ch == "\\":
            escape = True
            continue
        if ch == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}":
            if depth > 0:
                depth -= 1
                if depth == 0 and start != -1:
                    candidates.append(text[start : i + 1])
                    start = -1
    for cand in candidates:
        try:
            parsed = _json.loads(cand)
        except Exception:
            continue
        if parsed:
            return cand
    first, last = text.find("{"), text.rfind("}")
    if first != -1 and last > first:
        cand = text[first : last + 1]
        try:
            if _json.loads(cand):
                return cand
        except Exception:
            pass
    return None


def _env(*names: str, default: Optional[str] = None) -> Optional[str]:
    for n in names:
        v = os.environ.get(n)
        if v:
            return v
    return default


def _normalise_base(base: str) -> str:
    base = base.rstrip("/")
    if not base.endswith("/v1"):
        base = base + "/v1"
    return base


def _to_messages(prompt_or_messages) -> list[dict]:
    if isinstance(prompt_or_messages, list):
        return prompt_or_messages
    return [{"role": "user", "content": str(prompt_or_messages)}]


class LLMClient:
    """OpenAI-compatible chat client.

    Instances are cheap; nothing is loaded until the first call. Thread-safe
    in the sense that `requests.post` is used per-call.
    """

    def __init__(
        self,
        *,
        api_base: Optional[str] = None,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        timeout: Optional[float] = None,
    ) -> None:
        self.api_base = _normalise_base(
            api_base
            or _env(
                "GA_API_BASE",
                "OPENAI_API_BASE",
                "OPENAI_BASE_URL",
                "API_BASE",
                default="http://localhost:8000/v1",
            )
        )
        self.api_key = api_key or _env(
            "GA_API_KEY", "OPENAI_API_KEY", "API_KEY", default=""
        ) or ""
        self.model = model or _env(
            "GA_MODEL", "OPENAI_MODEL", "SUMMARY_MODEL_NAME", default=""
        ) or ""
        self.timeout = float(
            timeout
            if timeout is not None
            else _env("GA_HTTP_TIMEOUT", default="600")  # type: ignore[arg-type]
        )

    # ---------- internals ----------------------------------------------------

    def _post_chat(
        self,
        messages: list[dict],
        *,
        max_tokens: int,
        temperature: float,
        json_mode: bool,
    ) -> str:
        try:
            import requests  # type: ignore
        except ImportError as e:  # pragma: no cover
            raise RuntimeError(
                "graphanything LLM client requires `requests`; "
                "install it with `pip install requests`."
            ) from e

        if not self.model:
            raise RuntimeError(
                "No model configured. Set GA_MODEL (or OPENAI_MODEL / SUMMARY_MODEL_NAME) "
                "to the chat-completions model name your endpoint serves."
            )

        url = self.api_base + "/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key or 'none'}",
        }
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": float(temperature),
            "max_tokens": int(max_tokens),
        }
        if json_mode:
            payload["response_format"] = {"type": "json_object"}

        resp = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
        if resp.status_code == 400 and json_mode:
            # Some servers (e.g. Ollama bridge, older vLLM) reject response_format.
            # Retry without it; we'll regex-extract the JSON ourselves.
            payload.pop("response_format", None)
            resp = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
        if resp.status_code >= 400:
            raise RuntimeError(
                f"chat.completions failed (HTTP {resp.status_code}): {resp.text[:400]}"
            )

        data = resp.json()
        try:
            return data["choices"][0]["message"]["content"] or ""
        except (KeyError, IndexError, TypeError) as e:
            raise RuntimeError(f"unexpected chat.completions response: {data!r}") from e

    # ---------- public API --------------------------------------------------

    def chat(
        self,
        prompt_or_messages,
        *,
        max_tokens: int = 16384,
        temperature: float = 0.7,
        **_unused: Any,
    ) -> str:
        """Free-form chat completion → string. Strips reasoning tags."""
        raw = self._post_chat(
            _to_messages(prompt_or_messages),
            max_tokens=max_tokens,
            temperature=temperature,
            json_mode=False,
        )
        return _strip_think(raw)

    def chat_json(
        self,
        prompt_or_messages,
        *,
        max_tokens: int = 16384,
        temperature: float = 0.0,
        **_unused: Any,
    ) -> Any:
        """Strict-JSON chat completion → parsed Python value.

        Asks the server for JSON via `response_format`; if the server
        ignores the hint or wraps the answer in markdown / `<think>`, falls
        back to a brace-balanced extractor.
        """
        raw = self._post_chat(
            _to_messages(prompt_or_messages),
            max_tokens=max_tokens,
            temperature=temperature,
            json_mode=True,
        )
        # Fast path: server actually obeyed json mode.
        try:
            return _json.loads(raw)
        except Exception:
            pass
        extracted = _extract_json(raw)
        if extracted is None:
            raise RuntimeError(
                f"could not extract JSON from chat output. raw[:300]={raw[:300]!r}"
            )
        return _json.loads(extracted)


def make_client(**kwargs: Any) -> LLMClient:
    """Factory: build a client from kwargs / env, ready to call.

    Usage:
        from GraphAnything.llm_client import make_client
        llm = make_client()                      # all-env
        llm = make_client(model="Qwen3-32B-Instruct", api_base="http://h:8000/v1")
    """
    return LLMClient(**kwargs)