core.py

"""
Core agent orchestration — entry point dùng chung cho API và UI.
"""

import base64
import json
import logging
import mimetypes
import time
from datetime import datetime
from typing import Optional

logger = logging.getLogger(__name__)

from langchain_core.messages import HumanMessage, ToolMessage

from src.conversation_memory import add_turn
from src.graph import run
from src.nodes import final_response_node, image_response_node
from src.pdf_processing import format_chat_history, pdf_to_markdown
from src.qdrant_store import get_custom_prompt
from src.quiz import generate_quiz
from src.redis_client import redis_client
from src.state import MAX_ITERS, AgentState


def final_answer(
    conversation_id: str,
    sender_id: str,
    query: str,
    pdf_path: Optional[str] = None,
    image_path: Optional[str] = None,
    pdf_name: Optional[str] = None,
    gen_quiz: bool = False,
    k_question: Optional[int] = None,
    skip_pdf_indexing: bool = False,
) -> tuple[str, str, str | None, str | None]:
    """
    Khởi tạo AgentState, chạy graph, trả về 4 giá trị.

    Returns:
        (answer, elapsed, chart_type, chart_data)
        - chart_type : "column" | "pie" | None
        - chart_data : JSON string | None (chỉ có khi tool summarize_chart được gọi)

    Raises:
        ValueError: nếu bất kỳ tham số bắt buộc nào rỗng.
    """
    conversation_id = conversation_id.strip()
    sender_id       = sender_id.strip()
    query           = query.strip()

    if not conversation_id:
        raise ValueError("conversation_id không được để trống.")
    if not sender_id:
        raise ValueError("sender_id không được để trống.")
    if not query:
        raise ValueError("query không được để trống.")

    if gen_quiz:
        t0 = time.perf_counter()
        answer = generate_quiz(query, k_question or 10)
        return answer, f"{time.perf_counter() - t0:.2f}s", None, None

    custom_prompt = get_custom_prompt(sender_id)

    if pdf_path is not None and not skip_pdf_indexing:
        # Auto-index vào Qdrant để dùng với rag_search sau này (idempotent)
        try:
            from src.pdf_rag import index_pdf
            index_pdf(pdf_path, pdf_name or "document.pdf", conversation_id)
        except Exception:
            logger.exception("Auto-index PDF thất bại.")

        pdf_content  = pdf_to_markdown(pdf_path)
        chat_history = redis_client.get_chat_history(conversation_id)
        chat_text    = format_chat_history(chat_history)
        tool_content = (
            f"[Nội dung PDF]\n{pdf_content}"
            f"\n\n[Lịch sử trò chuyện]\n{chat_text}"
        )

        state: AgentState = {
            "conversation_id": conversation_id,
            "sender_id":       sender_id,
            "time":            datetime.now().isoformat(),
            "raw_query":       query,
            "query_type":      None,
            "messages":        [
                HumanMessage(content=query),
                ToolMessage(content=tool_content, tool_call_id="pdf_reader", name="pdf_reader"),
            ],
            "iters":           0,
            "max_iters":       MAX_ITERS,
            "final_answer":    None,
            "custom_prompt":   custom_prompt,
        }

        t0 = time.perf_counter()
        result  = final_response_node(state)
        elapsed = f"{time.perf_counter() - t0:.2f}s"
        answer  = result.get("final_answer") or "(Không có kết quả)"
        add_turn(conversation_id, sender_id, query, answer)
        return answer, elapsed, None, None

    if image_path is not None:
        mime_type, _ = mimetypes.guess_type(image_path)
        mime_type = mime_type or "image/jpeg"

        with open(image_path, "rb") as f:
            image_b64 = base64.b64encode(f.read()).decode()

        chat_history = redis_client.get_chat_history(conversation_id)
        chat_text    = format_chat_history(chat_history)

        text_content = f"{query}\n\n[Lịch sử trò chuyện]\n{chat_text}"

        state: AgentState = {
            "conversation_id": conversation_id,
            "sender_id":       sender_id,
            "time":            datetime.now().isoformat(),
            "raw_query":       query,
            "query_type":      None,
            "messages":        [
                HumanMessage(content=[
                    {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_b64}"}},
                    {"type": "text", "text": text_content},
                ]),
            ],
            "iters":           0,
            "max_iters":       MAX_ITERS,
            "final_answer":    None,
            "custom_prompt":   custom_prompt,
        }

        t0 = time.perf_counter()
        result  = image_response_node(state)
        elapsed = f"{time.perf_counter() - t0:.2f}s"
        answer  = result.get("final_answer") or "(Không có kết quả)"
        add_turn(conversation_id, sender_id, query, answer)
        return answer, elapsed, None, None

    initial_state: AgentState = {
        "conversation_id": conversation_id,
        "sender_id":       sender_id,
        "time":            datetime.now().isoformat(),
        "raw_query":       query,
        "query_type":      None,
        "messages":        [],
        "iters":           0,
        "max_iters":       MAX_ITERS,
        "final_answer":    None,
        "custom_prompt":   custom_prompt,
    }

    t0 = time.perf_counter()
    result  = run(initial_state)
    elapsed = f"{time.perf_counter() - t0:.2f}s"

    answer = result.get("final_answer") or "(Không có kết quả)"

    chart_type: str | None = None
    chart_data: str | None = None
    for msg in result.get("messages", []):
        if isinstance(msg, ToolMessage) and getattr(msg, "name", None) == "summarize_chart":
            try:
                parsed = json.loads(msg.content)
                if parsed.get("status") == "success":
                    chart_type = parsed.get("chart_type")
                    chart_data = json.dumps(parsed.get("chart_data", []), ensure_ascii=False)
            except Exception:
                pass
            break

    add_turn(conversation_id, sender_id, query, answer)
    return answer, elapsed, chart_type, chart_data


if __name__ == "__main__":
    answer, elapsed = final_answer(
        conversation_id="04ba40fe-61c7-4906-9f51-5ada0a392dac",
        sender_id="@slavakpa",
        query="tóm tắt nội dung tài liệu này",
        pdf_path="temp/test_doc.pdf",
    )
    print(answer)
    print(f"\n({elapsed})")