tools/rag_server.py

from agent.rag_engine import RAGEngine


_rag_engine = None

def get_rag_engine():
    global _rag_engine
    if _rag_engine is None:
        _rag_engine = RAGEngine()
    return _rag_engine


async def search_documents(query: str, k: int = 5) -> dict:
    """
    Search documents using RAG
    
    Args:
        query: Search query
        k: Number of results
        
    Returns:
        Dict with search results
    """
    try:
        rag = get_rag_engine()
        results = await rag.search(query, k=k)
        
        return {
            'success': True,
            'query': query,
            'results': results,
            'count': len(results)
        }
        
    except Exception as e:
        return {'error': str(e), 'success': False}


async def add_document_to_rag(text: str, metadata: dict) -> dict:
    """Add document to RAG index"""
    try:
        rag = get_rag_engine()
        doc_id = await rag.add_document(text, metadata)
        
        return {
            'success': True,
            'document_id': doc_id,
            'text_length': len(text)
        }
        
    except Exception as e:
        return {'error': str(e), 'success': False}