src/llm_client.py

"""
src/llm_client.py — Ollama / OpenAI-compatible text chat client.

Uses only text chat completions; no image generation, no g4f, no GPT-4 models.
Model and endpoint are loaded from environment variables via src/config.py.
"""
import json
import logging
from typing import Dict, Iterator, List, Optional

import requests

from src.config import (
    CHAT_COMPLETION_URL,
    MAX_TOKENS,
    MODEL_NAME,
    OLLAMA_API_KEY,
    REQUEST_TIMEOUT,
    TEMPERATURE,
)

logger = logging.getLogger(__name__)


class LLMError(Exception):
    """Raised when the LLM call fails for any reason."""


class OllamaChatClient:
    """
    Thin wrapper around any OpenAI-compatible /chat/completions endpoint.
    Tested with Ollama's built-in OpenAI-compatibility layer.
    """

    def __init__(self) -> None:
        self.url = CHAT_COMPLETION_URL
        self.model = MODEL_NAME
        self.api_key = OLLAMA_API_KEY

        if not self.url:
            raise LLMError(
                "MODEL_BASE_URL is not configured. "
                "Set it in .env or as a Hugging Face Space secret."
            )
        if not self.model:
            raise LLMError(
                "MODEL_NAME is not configured. "
                "Set it in .env or as a Hugging Face Space secret."
            )

    def _headers(self) -> Dict[str, str]:
        headers: Dict[str, str] = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers

    def chat(
        self,
        messages: List[Dict[str, str]],
        system_prompt: Optional[str] = None,
    ) -> str:
        """
        Send a chat completion request and return the assistant's reply as a string.

        Args:
            messages:      Ordered list of {"role": ..., "content": ...} dicts.
            system_prompt: Optional system instruction prepended before messages.

        Returns:
            The assistant's text response.

        Raises:
            LLMError: On any connectivity, HTTP, or parsing failure.
        """
        all_messages: List[Dict[str, str]] = []
        if system_prompt:
            all_messages.append({"role": "system", "content": system_prompt})
        all_messages.extend(messages)

        payload = {
            "model": self.model,
            "messages": all_messages,
            "max_tokens": MAX_TOKENS,
            "temperature": TEMPERATURE,
            "stream": False,
        }

        logger.debug(
            "LLM request | model=%s | url=%s | messages=%d",
            self.model,
            self.url,
            len(all_messages),
        )

        try:
            resp = requests.post(
                self.url,
                headers=self._headers(),
                json=payload,
                timeout=REQUEST_TIMEOUT,
            )
            resp.raise_for_status()
            data = resp.json()
            return data["choices"][0]["message"]["content"].strip()

        except requests.exceptions.ConnectionError:
            raise LLMError(
                f"Cannot connect to the Ollama endpoint: {self.url}\n"
                "Check that MODEL_BASE_URL is correct and the service is running."
            )
        except requests.exceptions.Timeout:
            raise LLMError(
                f"Request timed out after {REQUEST_TIMEOUT}s. "
                "Try increasing REQUEST_TIMEOUT or check model availability."
            )
        except requests.exceptions.HTTPError as exc:
            body = exc.response.text[:300] if exc.response is not None else ""
            raise LLMError(
                f"Ollama API returned HTTP {exc.response.status_code}: {body}"
            )
        except (KeyError, IndexError, json.JSONDecodeError) as exc:
            raise LLMError(
                f"Unexpected response format from Ollama endpoint: {exc}"
            )

    def stream_chat(
        self,
        messages: List[Dict[str, str]],
        system_prompt: Optional[str] = None,
    ) -> Iterator[str]:
        """
        Stream a chat completion and yield text chunks as they arrive.

        Uses the Ollama/OpenAI SSE streaming format (stream=True).
        Raises LLMError on connectivity, HTTP, or parsing failure.
        """
        all_messages: List[Dict[str, str]] = []
        if system_prompt:
            all_messages.append({"role": "system", "content": system_prompt})
        all_messages.extend(messages)

        payload = {
            "model": self.model,
            "messages": all_messages,
            "max_tokens": MAX_TOKENS,
            "temperature": TEMPERATURE,
            "stream": True,
        }

        logger.debug(
            "LLM stream request | model=%s | messages=%d",
            self.model,
            len(all_messages),
        )

        try:
            with requests.post(
                self.url,
                headers=self._headers(),
                json=payload,
                timeout=REQUEST_TIMEOUT,
                stream=True,
            ) as resp:
                resp.raise_for_status()
                for raw_line in resp.iter_lines():
                    if not raw_line:
                        continue
                    line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
                    if not line.startswith("data:"):
                        continue
                    data = line[5:].strip()
                    if data == "[DONE]":
                        break
                    try:
                        obj = json.loads(data)
                        chunk = obj["choices"][0]["delta"].get("content", "")
                        if chunk:
                            yield chunk
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue

        except requests.exceptions.ConnectionError:
            raise LLMError(
                f"Cannot connect to the Ollama endpoint: {self.url}\n"
                "Check that MODEL_BASE_URL is correct and the service is running."
            )
        except requests.exceptions.Timeout:
            raise LLMError(
                f"Request timed out after {REQUEST_TIMEOUT}s. "
                "Try increasing REQUEST_TIMEOUT or check model availability."
            )
        except requests.exceptions.HTTPError as exc:
            body = exc.response.text[:300] if exc.response is not None else ""
            raise LLMError(
                f"Ollama API returned HTTP {exc.response.status_code}: {body}"
            )