tools/document_retriever.py

from langchain_core.tools import tool
from langchain_community.document_loaders import TextLoader, CSVLoader, PyPDFLoader
import logging
import os

logger = logging.getLogger(__name__)

@tool
async def document_retriever_tool(task_id: str, query: str, file_type: str) -> str:
    """Retrieve content from a document"""
    try:
        file_path = f"temp_{task_id}.{file_type}"
        if not os.path.exists(file_path):
            logger.warning(f"Document not found: {file_path}")
            return "Document not found"
        
        if file_type == "txt":
            loader = TextLoader(file_path)
        elif file_type == "csv":
            loader = CSVLoader(file_path)
        elif file_type == "pdf":
            loader = PyPDFLoader(file_path)
        else:
            return f"Unsupported file type: {file_type}"
        
        docs = loader.load()
        return "\n".join(doc.page_content for doc in docs)
    except Exception as e:
        logger.error(f"Error retrieving document for task {task_id}: {e}")
        return f"Error: {str(e)}"