assistant/llm.py

from __future__ import annotations

from functools import lru_cache
from typing import Iterator

from llama_cpp import Llama

from config import CONFIG


@lru_cache(maxsize=1)
def _llm():
    """Build the llama.cpp model once, applying the LoRA adapter if set."""
    kwargs = dict(
        model_path=CONFIG.model_path(),
        n_ctx=CONFIG.llm.n_ctx,
        n_threads=CONFIG.n_threads,
        verbose=False,
    )
    if CONFIG.llm.chat_format:
        kwargs["chat_format"] = CONFIG.llm.chat_format
    lora = CONFIG.lora_path()
    if lora:
        kwargs["lora_path"] = lora
    return Llama(**kwargs)


def warmup() -> None:
    """Load the model and run one token so the first real turn isn't cold."""
    _llm().create_chat_completion(
        messages=[{"role": "user", "content": "hola"}], max_tokens=1
    )


def complete(prompt: str, max_tokens: int = 48) -> str:
    """One short, non-streamed completion. Used to build a search query."""
    out = _llm().create_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens, temperature=0.2,
    )
    return (out["choices"][0]["message"].get("content") or "").strip()


def stream_reply(
    messages: list[dict], temperature: float | None = None
) -> Iterator[str]:
    """Stream the answer token by token."""
    for part in _llm().create_chat_completion(
        messages=messages, stream=True,
        temperature=CONFIG.temperature if temperature is None else temperature,
        max_tokens=CONFIG.max_tokens,
    ):
        delta = part["choices"][0]["delta"].get("content")
        if delta:
            yield delta