backend/main.py

"""Streaming chat orchestration utilities for the frontend voicebot."""

from __future__ import annotations

import asyncio
import logging
import os
from queue import Queue
from threading import Lock, Thread
from typing import AsyncGenerator, Dict, Iterator, List, Optional

from dotenv import load_dotenv
from langfuse import Langfuse
from langfuse.decorators import langfuse_context, observe
import sys
sys.path.append(os.path.abspath('./backend'))
from models import LLMFinanceAnalyzer
from functions import MongoHybridSearch


load_dotenv(override=True)

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


langfuse = Langfuse(
    secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
    public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
    host=os.getenv("LANGFUSE_HOST"),
)
langfuse_context.configure(environment="development")


try:
    llm_analyzer = LLMFinanceAnalyzer()
    search_engine = MongoHybridSearch()
    logger.info("Initialized LLM analyzer and Mongo hybrid search for streaming chat.")
except Exception as exc:
    logger.critical("Failed to initialise backend components: %s", exc, exc_info=True)
    raise


_stream_loop: Optional[asyncio.AbstractEventLoop] = None
_stream_thread: Optional[Thread] = None
_stream_loop_lock: "Lock" = Lock()


def _loop_worker(loop: asyncio.AbstractEventLoop) -> None:
    asyncio.set_event_loop(loop)
    loop.run_forever()


def _ensure_stream_loop() -> asyncio.AbstractEventLoop:
    global _stream_loop, _stream_thread

    with _stream_loop_lock:
        if _stream_loop is None or _stream_loop.is_closed():
            _stream_loop = asyncio.new_event_loop()
            _stream_thread = Thread(target=_loop_worker, args=(_stream_loop,), daemon=True)
            _stream_thread.start()

    return _stream_loop


def _create_truncated_history(
    full_conversation: List[Dict[str, str]],
    max_assistant_length: int,
) -> List[Dict[str, str]]:
    truncated = []
    for msg in full_conversation:
        processed = msg.copy()
        if processed.get("role") == "assistant" and len(processed.get("content", "")) > max_assistant_length:
            processed["content"] = processed["content"][:max_assistant_length] + "..."
        truncated.append(processed)
    return truncated


def _generate_pseudo_conversation(conversation: List[Dict[str, str]]) -> List[Dict[str, str]]:
    pseudo = "".join(f"{msg.get('role', 'unknown')}: {msg.get('content', '')}\n" for msg in conversation)
    return [{"role": "user", "content": pseudo.strip()}]


@observe()
async def _stream_chat_async(history: List[Dict[str, str]], message: str) -> AsyncGenerator[str, None]:
    full_conversation = [msg.copy() for msg in history] + [{"role": "user", "content": message}]
    truncated_history = _create_truncated_history(full_conversation, 300)
    pseudo_conversation = _generate_pseudo_conversation(truncated_history)

    rag_decision = "yes"
    logger.info("RAG decision: %s", rag_decision)

    if rag_decision == "yes":
        query = await llm_analyzer.generate_subquery(pseudo_conversation)
        if query is None:
            yield "ขออภัยค่ะ ไม่สามารถวิเคราะห์คำถามเพื่อดึงข้อมูลได้"
            return

        retrieved_data = ""
        if query:
            try:
                docs = await search_engine.search_documents(query)
                retrieved_data = "\n-------\n".join(docs)
                logger.info("Retrieved %d documents for streaming response.", len(docs))
            except Exception as search_err:
                logger.error("Error during document search: %s", search_err, exc_info=True)
                yield "ขออภัยค่ะ เกิดข้อผิดพลาดขณะค้นหาข้อมูล"
                return

        limited_conversation = full_conversation[-7:] if len(full_conversation) > 7 else full_conversation
        response_generator = llm_analyzer.generate_normal_response(retrieved_data, limited_conversation)

        async for chunk in response_generator:
            if chunk:
                yield chunk
                await asyncio.sleep(0.05)
    else:
        limited_conversation = full_conversation[-9:] if len(full_conversation) > 9 else full_conversation
        final_response = await llm_analyzer.generate_non_rag_response(limited_conversation)
        if final_response:
            yield final_response
        else:
            yield "ขออภัยค่ะ เกิดข้อผิดพลาดในการประมวลผลคำถามของคุณ"


def stream_chat_response(history: List[Dict[str, str]], message: str) -> Iterator[str]:
    """Synchronously iterate over streaming LLM chunks."""

    loop = _ensure_stream_loop()
    output_queue: "Queue[Optional[str]]" = Queue()

    async def runner() -> None:
        try:
            async for chunk in _stream_chat_async(history, message):
                output_queue.put_nowait(str(chunk))
        except Exception as exc:  # noqa: BLE001
            logger.error("Unhandled error in async chat stream: %s", exc, exc_info=True)
            output_queue.put_nowait(f"[Error: {exc}]")
        finally:
            output_queue.put_nowait(None)

    future = asyncio.run_coroutine_threadsafe(runner(), loop)

    while True:
        chunk = output_queue.get()
        if chunk is None:
            break
        yield chunk

    # Propagate any exception that was not handled in runner().
    future.result()


__all__ = ["stream_chat_response"]