utils/llm_chat.py

"""Chat LLM integration using NVIDIA NIM (llama-3.1-8b-instruct) with streaming."""

import logging
import os
from typing import Generator

from openai import OpenAI

logger = logging.getLogger(__name__)

_NIM_BASE_URL = "https://integrate.api.nvidia.com/v1"
_CHAT_MODEL = "meta/llama-3.1-8b-instruct"


def _get_client() -> OpenAI:
    api_key = os.environ.get("NVIDIA_API_KEY")
    if not api_key:
        raise EnvironmentError("NVIDIA_API_KEY environment variable is not set.")
    return OpenAI(base_url=_NIM_BASE_URL, api_key=api_key)


def chat_with_video(
    message: str,
    chat_history: list,
    system_prompt: str,
) -> Generator[str, None, None]:
    """Stream a response from the chat LLM given a user message and conversation history.

    chat_history is expected in Gradio's list-of-[user, assistant] format.
    Yields response text chunks for Gradio streaming.
    """
    client = _get_client()

    messages = [{"role": "system", "content": system_prompt}]
    for pair in chat_history:
        if pair[0]:
            messages.append({"role": "user", "content": pair[0]})
        if pair[1]:
            messages.append({"role": "assistant", "content": pair[1]})
    messages.append({"role": "user", "content": message})

    logger.info("Sending chat message to %s", _CHAT_MODEL)
    stream = client.chat.completions.create(
        model=_CHAT_MODEL,
        messages=messages,
        stream=True,
        max_tokens=1024,
        temperature=0.3,
    )

    for chunk in stream:
        delta = chunk.choices[0].delta.content
        if delta:
            yield delta