agents/llm.py

"""Unified LLM client: uses a single model for both vision and text tasks."""

import asyncio
import base64
import io
import json
import logging
import time

from PIL import Image
from pydantic import BaseModel

logger = logging.getLogger(__name__)


class LLMClient:
    """Single model client for vision + text via llama.cpp."""

    def __init__(self, model_path: str, projection_path: str):
        self.model_path = model_path
        self.projection_path = projection_path
        self.model_id = model_path.split("/")[-1]
        self._model = None

    def _ensure_model(self):
        if self._model is not None:
            return
        from llama_cpp import Llama
        from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler

        logger.info("Loading model: %s", self.model_path)
        chat_handler = MiniCPMv26ChatHandler(
            clip_model_path=self.projection_path
        )
        self._model = Llama(
            model_path=self.model_path,
            chat_handler=chat_handler,
            n_ctx=2048,
            n_threads=2,
            verbose=False,
        )
        logger.info("Model loaded")

    def _prepare_image(self, image: Image.Image) -> str:
        max_dim = 384
        if max(image.size) > max_dim:
            ratio = max_dim / max(image.size)
            new_size = (int(image.width * ratio), int(image.height * ratio))
            image = image.resize(new_size, Image.LANCZOS)
        if image.mode != "RGB":
            image = image.convert("RGB")
        buffer = io.BytesIO()
        image.save(buffer, format="JPEG", quality=85)
        b64 = base64.b64encode(buffer.getvalue()).decode()
        return f"data:image/jpeg;base64,{b64}"

    def _build_example_json(self, schema: type[BaseModel]) -> str:
        example: dict = {}
        for name, field in schema.model_fields.items():
            annotation = field.annotation
            origin = getattr(annotation, "__origin__", None)
            if origin is type(None):
                example[name] = None
                continue
            args = getattr(annotation, "__args__", None)
            if args and type(None) in args:
                annotation = [a for a in args if a is not type(None)][0]
            if annotation == str:
                example[name] = "..."
            elif annotation == int:
                example[name] = 0
            elif annotation == float:
                example[name] = 0.0
            elif annotation == bool:
                example[name] = False
            elif annotation is list or (
                hasattr(annotation, "__origin__")
                and getattr(annotation, "__origin__", None) is list
            ):
                example[name] = ["..."]
            else:
                example[name] = "..."
        return json.dumps(example, indent=2)

    def describe_image(self, image: Image.Image, prompt: str) -> tuple[str, int]:
        """Describe an image. Returns (description, duration_ms)."""
        self._ensure_model()
        data_uri = self._prepare_image(image)

        start = time.monotonic()
        response = self._model.create_chat_completion(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": data_uri}},
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
            max_tokens=256,
        )
        duration_ms = int((time.monotonic() - start) * 1000)
        content = response["choices"][0]["message"]["content"]
        return content, duration_ms

    def generate(
        self,
        system: str,
        prompt: str,
        output_schema: type[BaseModel],
        max_tokens: int = 512,
    ) -> tuple[BaseModel, int]:
        """Generate structured JSON output. Returns (parsed_model, duration_ms)."""
        self._ensure_model()

        example_json = self._build_example_json(output_schema)
        json_instruction = (
            "\n\nRespond ONLY with a valid JSON object. "
            f"Use exactly these keys:\n{example_json}\n"
            "Fill in real values. Do not include any text outside the JSON."
        )

        start = time.monotonic()
        response = self._model.create_chat_completion(
            messages=[
                {"role": "system", "content": system + json_instruction},
                {"role": "user", "content": prompt},
            ],
            max_tokens=max_tokens,
            response_format={"type": "json_object"},
        )
        duration_ms = int((time.monotonic() - start) * 1000)

        content = response["choices"][0]["message"]["content"]
        return output_schema.model_validate_json(content), duration_ms

    async def adescribe_image(
        self, image: Image.Image, prompt: str
    ) -> tuple[str, int]:
        return await asyncio.to_thread(self.describe_image, image, prompt)

    async def agenerate(
        self,
        system: str,
        prompt: str,
        output_schema: type[BaseModel],
        max_tokens: int = 512,
    ) -> tuple[BaseModel, int]:
        return await asyncio.to_thread(
            self.generate, system, prompt, output_schema, max_tokens
        )