app/services/llm.py

"""
ByteAstra — Local LLM Client.

Connects to any OpenAI-compatible inference server (Ollama, vLLM, llama.cpp server).
Supports both streaming and non-streaming completions.
"""
from __future__ import annotations
import asyncio
import logging
import time
from typing import AsyncGenerator
from openai import AsyncOpenAI
import httpx
from app.config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()

# Single shared async client for the application lifetime
_client: AsyncOpenAI | None = None

# Cache check_connection result to avoid failing when LLM is busy generating
_connection_cache: tuple[bool, float] | None = None  # (result, timestamp)
_CONNECTION_CACHE_TTL = 30.0  # seconds


def get_llm_client() -> AsyncOpenAI:
    global _client
    if _client is None:
        base_url = settings.resolved_llm_base_url
        logger.info("LLM client connecting to: %s", base_url)
        _client = AsyncOpenAI(
            base_url=base_url,
            api_key=settings.llm_api_key,
            # Separate connect vs read timeouts:
            # connect: 5s (fail fast if server is down)
            # read: 300s (CPU 3B @ 600 max_tokens needs up to ~5 min worst case under heavy CPU throttling)
            timeout=httpx.Timeout(connect=5.0, read=300.0, write=10.0, pool=5.0),
        )
    return _client


async def stream_completion(
    messages: list[dict],
    model: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
) -> AsyncGenerator[str, None]:
    """
    Stream a chat completion token by token.
    Yields text delta strings as they arrive from the model.
    """
    client = get_llm_client()
    _model = model or settings.resolved_model_name
    _temp = temperature if temperature is not None else settings.llm_temperature
    _max_tok = max_tokens or settings.llm_max_tokens

    logger.debug("LLM stream_completion | model=%s | messages=%d", _model, len(messages))

    try:
        stream = await client.chat.completions.create(
            model=_model,
            messages=messages,
            temperature=_temp,
            max_tokens=_max_tok,
            stream=True,
            # Explicit stop sequences prevent model from repeating delimiters or going out-of-turn
            stop=["<|im_end|>", "<|im_start|>", "--- STUDENT QUESTION ---", "--- SYLLABUS CONTEXT ---"],
            # Frequency penalty prevents the model from getting stuck in repetition loops (e.g. repeating emojis)
            frequency_penalty=0.1,
            extra_body={"cache_prompt": True},
        )
        async for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                yield delta
    except Exception as exc:
        logger.warning("Local LLM connection failed (%s). Falling back to RAG-context generator.", exc)
        
        # Extract RAG context from the last user message
        user_msg = messages[-1]["content"] if messages else ""
        
        # Let's extract the context blocks
        context_lines = []
        in_context = False
        for line in user_msg.splitlines():
            if "--- CONTEXT ---" in line or "--- SYLLABUS CONTEXT ---" in line:
                in_context = True
                continue
            if "--- STUDENT QUESTION ---" in line:
                break
            if in_context:
                context_lines.append(line)
        
        context_text = "\n".join(context_lines).strip()
        if context_text:
            intro = (
                "*[Fallback Mode: Connected directly to textbook database]*\n\n"
                "According to the classical texts in the BAMS curriculum:\n\n"
            )
            # Stream the intro word-by-word
            intro_words = intro.split(" ")
            for j, word in enumerate(intro_words):
                space = " " if j < len(intro_words) - 1 else ""
                yield word + space
                await asyncio.sleep(0.02)
            
            # Yield the context text word-by-word
            for chunk_part in context_text.split("\n\n"):
                part = chunk_part.strip()
                if not part:
                    continue
                words = part.split(" ")
                for j, word in enumerate(words):
                    space = " " if j < len(words) - 1 else ""
                    yield word + space
                    await asyncio.sleep(0.015)
                yield "\n\n"
        else:
            fallback_msg = (
                "I couldn't establish a connection to the local LLM server, and there is no "
                "relevant context in the database to answer your question."
            )
            for char in fallback_msg:
                yield char
                await asyncio.sleep(0.005)


async def complete(
    messages: list[dict],
    model: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
) -> str:
    """
    Non-streaming completion — returns the full response string.
    Internally collects the stream; useful for evaluation scripts.
    """
    parts: list[str] = []
    async for token in stream_completion(messages, model, temperature, max_tokens):
        parts.append(token)
    return "".join(parts)


async def check_connection() -> bool:
    """Check if the local LLM server is online. Result is cached for 30s.
    This prevents false negatives when the server is busy generating a response."""
    global _connection_cache
    now = time.monotonic()

    # Return cached result if fresh
    if _connection_cache is not None:
        cached_result, cached_at = _connection_cache
        if now - cached_at < _CONNECTION_CACHE_TTL:
            return cached_result

    client = get_llm_client()
    try:
        # Use a dedicated short-timeout client just for health check
        health_client = AsyncOpenAI(
            base_url=settings.resolved_llm_base_url,
            api_key=settings.llm_api_key,
            timeout=httpx.Timeout(connect=2.0, read=2.0, write=2.0, pool=2.0),
        )
        await health_client.models.list()
        result = True
    except Exception as e:
        err_msg = str(e).lower()
        # ConnectionRefusedError or explicit connect errors mean offline.
        # Timeout/read error means server is online but busy.
        if "refused" in err_msg or "connect error" in err_msg or "failed to connect" in err_msg:
            result = False
        else:
            result = True

    _connection_cache = (result, now)
    logger.info("LLM connection check: %s (cached for %ds)", result, int(_CONNECTION_CACHE_TTL))
    return result