app/chatbot.py

"""
chatbot.py
----------
Orchestrates the full RAG pipeline for a single user query.

Step 3 Enhancements:
- chat() now accepts score_threshold and passes it to retriever
- chat_stream() added: does retrieval first, then streams LLM tokens
  Returns (generator, sources, query) tuple for UI streaming support
"""

import logging
from dataclasses import dataclass, field
from typing import Generator, List, Tuple

from components.llm_handler import LLMHandler
from components.prompt_template import build_prompt, format_sources
from components.retriever import Retriever
from components.vector_store import VectorStore

logger = logging.getLogger(__name__)


@dataclass
class ChatResponse:
    """Container for a single chatbot turn (non-streaming)."""
    answer:  str
    sources: List[str] = field(default_factory=list)
    query:   str = ""


class Chatbot:
    """
    End-to-end RAG chatbot.

    Args:
        vector_store: Pre-built or loaded VectorStore instance.
        retriever:    Retriever wrapping the vector store.
        llm:          LLMHandler for text generation.
    """

    def __init__(
        self,
        vector_store: VectorStore,
        retriever: Retriever | None = None,
        llm: LLMHandler | None = None,
    ) -> None:
        self.vector_store = vector_store
        self.retriever    = retriever or Retriever(vector_store)
        self.llm          = llm or LLMHandler()

    # ── Non-streaming (kept for backward compatibility) ───────────────────────

    def chat(
        self,
        query: str,
        top_k: int | None = None,
        history: List[dict] | None = None,
        score_threshold: float | None = None,
    ) -> ChatResponse:
        """
        Process a user query through the full RAG pipeline (non-streaming).

        Args:
            query:           User's natural-language question.
            top_k:           Number of chunks to retrieve.
            history:         Previous chat turns for conversation memory.
            score_threshold: Minimum similarity score for retrieved chunks.

        Returns:
            ChatResponse with answer and source references.
        """
        query = query.strip()
        if not query:
            return ChatResponse(answer="Please enter a question.", sources=[], query=query)

        if not self.vector_store.is_ready:
            return ChatResponse(
                answer="No documents have been ingested yet. Please upload documents first.",
                sources=[],
                query=query,
            )

        logger.info("Processing query: '%s'", query[:100])

        results      = self.retriever.retrieve(query, k=top_k, score_threshold=score_threshold)
        context_docs = [doc for doc, _ in results]

        if not context_docs:
            return ChatResponse(
                answer="I couldn't find any relevant information in the knowledge base.",
                sources=[],
                query=query,
            )

        prompt  = build_prompt(query, context_docs, history=history)
        answer  = self.llm.generate(prompt)
        sources = format_sources(results)

        logger.info("Query answered. Sources: %s | Answer length: %d chars", sources, len(answer))
        return ChatResponse(answer=answer, sources=sources, query=query)

    # ── Streaming ─────────────────────────────────────────────────────────────

    def chat_stream(
        self,
        query: str,
        top_k: int | None = None,
        history: List[dict] | None = None,
        score_threshold: float | None = None,
    ) -> Tuple[Generator | None, List[str], str]:
        """
        Process a query and stream the LLM response token by token.

        Retrieval happens first (fast — milliseconds) so sources are known
        immediately. The token generator is returned alongside sources so
        the UI can stream the answer while sources are already available.

        Args:
            query:           User's natural-language question.
            top_k:           Number of chunks to retrieve.
            history:         Previous chat turns for conversation memory.
            score_threshold: Minimum similarity score for retrieved chunks.

        Returns:
            Tuple of (token_generator, sources, cleaned_query).
            Returns (None, [], query) on early-exit conditions.
        """
        query = query.strip()
        if not query:
            return None, [], query

        if not self.vector_store.is_ready:
            return None, [], query

        logger.info("Processing streaming query: '%s'", query[:100])

        results      = self.retriever.retrieve(query, k=top_k, score_threshold=score_threshold)
        context_docs = [doc for doc, _ in results]

        if not context_docs:
            return None, [], query

        prompt       = build_prompt(query, context_docs, history=history)
        sources      = format_sources(results)
        token_stream = self.llm.generate_stream(prompt)

        # Return generator + sources — generation starts when UI iterates the generator
        return token_stream, sources, query