llm_client.py

"""
llm_client.py — Smart LLM client router for bootstrap-basil.

Routes API calls to the local vLLM instance for Tutor and Sophie models,
and to the standard OpenAI API for all other models (Grader, TaskAgent, etc.).

Usage:
    from llm_client import create_smart_client
    client = create_smart_client()
    # Use exactly like an openai.OpenAI client:
    client.chat.completions.create(model=TUTOR_MODEL, ...)
"""

import os
from openai import OpenAI

# Local vLLM endpoint (runs on the remote box at 127.0.0.1:8000)
_LOCAL_BASE_URL = "http://127.0.0.1:8000/v1"
_LOCAL_API_KEY = "localtoken"

# Models that should be routed to the local vLLM instance.
# Populated at import time from config to stay in sync.
_LOCAL_MODELS: frozenset = frozenset()


def _get_local_models() -> frozenset:
    """Import TUTOR_MODEL and SOPHIE_MODEL from config (lazy, avoids circular imports)."""
    global _LOCAL_MODELS
    if not _LOCAL_MODELS:
        try:
            from config import TUTOR_MODEL, SOPHIE_MODEL
            _LOCAL_MODELS = frozenset({TUTOR_MODEL, SOPHIE_MODEL})
        except ImportError:
            pass
    return _LOCAL_MODELS


class _Completions:
    """Proxy for client.chat.completions — routes to correct underlying client."""

    def __init__(self, openai_client: OpenAI, local_client: OpenAI):
        self._openai = openai_client
        self._local = local_client

    def create(self, *, model: str, **kwargs):
        local_models = _get_local_models()
        # Also route any Qwen-prefixed model to local vLLM
        is_local = model in local_models or model.startswith("Qwen/")
        target = self._local if is_local else self._openai
        return target.chat.completions.create(model=model, **kwargs)


class _Chat:
    def __init__(self, openai_client: OpenAI, local_client: OpenAI):
        self.completions = _Completions(openai_client, local_client)


class SmartClient:
    """
    Drop-in replacement for openai.OpenAI that automatically routes
    TUTOR_MODEL / SOPHIE_MODEL calls to the local vLLM endpoint and
    everything else to the standard OpenAI API.
    """

    def __init__(self):
        self._openai_client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY", "not-needed-all-local"),
            timeout=60.0,
        )
        self._local_client = OpenAI(
            base_url=_LOCAL_BASE_URL,
            api_key=_LOCAL_API_KEY,
            timeout=120.0,
        )
        self.chat = _Chat(self._openai_client, self._local_client)


def create_smart_client() -> SmartClient:
    """Return a SmartClient instance (creates both underlying clients)."""
    return SmartClient()