tools/rag.py

"""
RAG search tool — tìm kiếm hybrid (dense + BM25) trên knowledge base PDF của conversation.
"""

import logging
from .base import register_tool

logger = logging.getLogger(__name__)


@register_tool(
    name="rag_search",
    description=(
        "Tìm kiếm trong tài liệu PDF đã được index cho conversation/room này. "
        "Dùng khi người dùng hỏi về nội dung bài học, tài liệu đã chia sẻ, "
        "slides, hoặc kiến thức chuyên ngành liên quan đến lớp học. "
        "Trả về các đoạn văn bản liên quan nhất từ knowledge base."
    ),
    parameters=[
        {
            "name": "query",
            "type": "string",
            "description": "Câu hỏi hoặc chủ đề cần tìm kiếm trong tài liệu",
            "required": True,
        },
        {
            "name": "conversation_id",
            "type": "string",
            "description": "ID cuộc hội thoại hoặc room để giới hạn phạm vi tìm kiếm",
            "required": True,
        },
    ],
)
def tool_rag_search(query: str, conversation_id: str) -> str:
    try:
        from src.pdf_rag import hybrid_search
        chunks = hybrid_search(query, conversation_id, top_k=5)
    except RuntimeError as e:
        return f"(Lỗi kết nối knowledge base: {e})"
    except Exception:
        logger.exception("rag_search thất bại cho conversation '%s'", conversation_id)
        return "(Lỗi tìm kiếm tài liệu. Vui lòng thử lại.)"

    if not chunks:
        return "(Không tìm thấy nội dung liên quan trong knowledge base của conversation này.)"

    return "\n\n---\n\n".join(chunks)