site/WESP_REL/app/services/local_llm_client.py

"""HTTP-клиент к локальному LLM с OpenAI-совместимым API (например llama-server из llama.cpp)."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

import requests


class LocalLlmError(Exception):
    """Локальный LLM-сервер недоступен или вернул ошибку."""


def _auth_headers(api_key: str) -> Dict[str, str]:
    k = (api_key or "").strip()
    if not k:
        return {}
    return {"Authorization": f"Bearer {k}"}


def openai_list_models(
    *, base_url: str, timeout_sec: float, api_key: str = ""
) -> Dict[str, Any]:
    """GET /v1/models — как у OpenAI; llama-server поддерживает."""
    url = f"{base_url.rstrip('/')}/v1/models"
    try:
        r = requests.get(url, timeout=timeout_sec, headers=_auth_headers(api_key))
    except requests.RequestException as e:
        raise LocalLlmError(f"LLM недоступен: {e}") from e
    if r.status_code != 200:
        raise LocalLlmError(f"LLM GET /v1/models: HTTP {r.status_code}")
    try:
        return r.json()
    except ValueError as e:
        raise LocalLlmError("LLM /v1/models: не JSON") from e


def model_in_openai_payload(payload: Dict[str, Any], want: str) -> bool:
    """Проверяет, есть ли в ответе /v1/models модель с id, совпадающим с want (без учёта регистра)."""
    want_l = (want or "").strip().lower()
    if not want_l:
        return False
    data = payload.get("data")
    if not isinstance(data, list):
        return False
    for item in data:
        if not isinstance(item, dict):
            continue
        mid = str(item.get("id") or "").strip().lower()
        if not mid:
            continue
        if mid == want_l or mid.endswith("/" + want_l) or mid.endswith("\\" + want_l):
            return True
        # имя файла без пути
        if mid.rsplit("/", 1)[-1].rsplit("\\", 1)[-1] == want_l:
            return True
    return False


def _normalize_message_content(content: Any) -> Optional[str]:
    if content is None:
        return None
    if isinstance(content, str):
        return content.strip() or None
    if isinstance(content, list):
        parts: List[str] = []
        for c in content:
            if isinstance(c, dict) and c.get("type") == "text":
                parts.append(str(c.get("text") or ""))
            elif isinstance(c, str):
                parts.append(c)
        joined = "".join(parts).strip()
        return joined or None
    return str(content).strip() or None


def _parse_assistant_message(data: Dict[str, Any]) -> Dict[str, Any]:
    choices = data.get("choices")
    if not isinstance(choices, list) or not choices:
        raise LocalLlmError("LLM /v1/chat/completions: нет choices")
    first = choices[0]
    if not isinstance(first, dict):
        raise LocalLlmError("LLM /v1/chat/completions: choices[0] не объект")
    msg = first.get("message")
    if not isinstance(msg, dict):
        raise LocalLlmError("LLM /v1/chat/completions: нет message")
    content = _normalize_message_content(msg.get("content"))
    tool_calls = msg.get("tool_calls")
    if tool_calls is not None and not isinstance(tool_calls, list):
        tool_calls = None
    return {
        "content": content,
        "tool_calls": tool_calls,
    }


def openai_chat_completion_message(
    *,
    base_url: str,
    model: str,
    messages: List[Dict[str, Any]],
    timeout_sec: float,
    max_tokens: int = 512,
    api_key: str = "",
    tools: Optional[List[Dict[str, Any]]] = None,
    tool_choice: Optional[str] = None,
) -> Dict[str, Any]:
    """POST /v1/chat/completions — возвращает content и опционально tool_calls."""
    url = f"{base_url.rstrip('/')}/v1/chat/completions"
    body: Dict[str, Any] = {
        "model": model,
        "messages": messages,
        "max_tokens": max(16, min(int(max_tokens), 4096)),
    }
    if tools:
        body["tools"] = tools
        body["tool_choice"] = (tool_choice or "auto") if tool_choice else "auto"
    try:
        r = requests.post(url, json=body, timeout=timeout_sec, headers=_auth_headers(api_key))
    except requests.RequestException as e:
        raise LocalLlmError(f"LLM недоступен: {e}") from e
    if r.status_code != 200:
        raise LocalLlmError(f"LLM /v1/chat/completions: HTTP {r.status_code} {r.text[:500]}")
    try:
        data = r.json()
    except ValueError as e:
        raise LocalLlmError("LLM /v1/chat/completions: не JSON") from e
    return _parse_assistant_message(data)


def openai_chat_completion(
    *,
    base_url: str,
    model: str,
    messages: List[Dict[str, Any]],
    timeout_sec: float,
    max_tokens: int = 512,
    api_key: str = "",
) -> str:
    """POST /v1/chat/completions — ответ ассистента из choices[0].message.content (без tool_calls)."""
    msg = openai_chat_completion_message(
        base_url=base_url,
        model=model,
        messages=messages,
        timeout_sec=timeout_sec,
        max_tokens=max_tokens,
        api_key=api_key,
        tools=None,
        tool_choice=None,
    )
    content = msg.get("content")
    if isinstance(content, str) and content.strip():
        return content.strip()
    if msg.get("tool_calls"):
        raise LocalLlmError(
            "LLM вернул вызовы инструментов без финального текста; включите режим инструментов в чате."
        )
    raise LocalLlmError("LLM /v1/chat/completions: нет choices[0].message.content")