agent/tools.py

"""Custom tools for A11y Expert agent."""

from typing import Optional
from database.vector_store_client import VectorStoreClient
from models.embeddings import get_embeddings_client
from loguru import logger

def search_knowledge_base(
    query: str,
    vector_store: VectorStoreClient,
    language: str = "en"
) -> str:
    """
    Search the accessibility knowledge base for relevant information.
    
    Args:
        query: Question or search terms.
        vector_store: The VectorStoreClient instance.
        language: Language filter ('pl' or 'en').
        
    Returns:
        Formatted context from the knowledge base or an error message.
    """
    try:
        logger.info(f"Query: {query} (language: {language})")
        
        embeddings_client = get_embeddings_client()
        query_embedding = embeddings_client.get_embedding(query)
        
        where_clause = f"language = '{language}'"
        results = vector_store.search(
            query_embedding=query_embedding, where=where_clause, top_k=5
        )
        
        if not results:
            return f"No documents found for language '{language}'. Please ensure the knowledge base is loaded."
        
        context_lines = [
            f"[{i}. {r.get('source', 'unknown')} - {r.get('doc_type', 'Document')}]\n"
            f"{r.get('text', '')}\n"
            for i, r in enumerate(results, 1)
        ]
        
        context = "\n".join(context_lines)
        logger.info(f"Found {len(results)} relevant documents")
        
        return context
        
    except Exception as e:
        logger.error(f"Search failed: {e}")
        return f"Error searching knowledge base: {str(e)}"

def get_database_stats(vector_store: VectorStoreClient) -> str:
    """
    Get statistics about the knowledge base.
    
    Args:
        vector_store: The VectorStoreClient instance.
        
    Returns:
        A formatted string of database statistics.
    """
    try:
        count = vector_store.count_documents()
        return f"📊 Knowledge Base Statistics: Total documents: {count}"
        
    except Exception as e:
        logger.error(f"Failed to get stats: {e}")
        return f"Error retrieving statistics: {str(e)}"