agent/chat.py

"""Main chat agent via gemini-3.1-flash-lite (LangChain; Streamlit-safe)."""

from __future__ import annotations

import time
from collections.abc import Iterator

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage

from agent.llm import get_chat_llm, message_text
from agent.state import ChatTurn
from config import CHAT_SYSTEM_PROMPT

# Soft-chunk size so Streamlit write_stream paints gradually even when the
# model returns large deltas.
_STREAM_CHUNK_CHARS = 8
_STREAM_CHUNK_PAUSE_S = 0.015


def _to_lc_messages(turns: list[ChatTurn]) -> list[BaseMessage]:
    messages: list[BaseMessage] = [
        SystemMessage(content=CHAT_SYSTEM_PROMPT),
    ]
    for turn in turns:
        if turn["role"] == "user":
            messages.append(HumanMessage(content=turn["content"]))
        else:
            messages.append(AIMessage(content=turn["content"]))
    return messages


def _soft_chunks(text: str) -> Iterator[str]:
    step = _STREAM_CHUNK_CHARS
    for i in range(0, len(text), step):
        yield text[i : i + step]
        if _STREAM_CHUNK_PAUSE_S:
            time.sleep(_STREAM_CHUNK_PAUSE_S)


def chat_reply_stream(turns: list[ChatTurn]) -> Iterator[str]:
    """Yield reply chunks. Uses LangChain streaming (avoids genai Client close bugs)."""
    if not turns:
        raise ValueError("empty chat history")

    llm = get_chat_llm(temperature=0.7)
    try:
        for chunk in llm.stream(_to_lc_messages(turns)):
            text = message_text(chunk)
            if text:
                yield from _soft_chunks(text)
    except Exception:
        # Fallback: one-shot call, then soft-chunk for write_stream UX
        full = message_text(llm.invoke(_to_lc_messages(turns))).strip()
        if not full:
            raise
        yield from _soft_chunks(full)


def chat_reply(turns: list[ChatTurn]) -> str:
    """Non-streaming helper."""
    if not turns:
        raise ValueError("empty chat history")
    llm = get_chat_llm(temperature=0.7)
    text = message_text(llm.invoke(_to_lc_messages(turns))).strip()
    if not text:
        raise RuntimeError("empty model response")
    return text