modules/retriever.py

from tavily import TavilyClient
import logging

class Retriever:
    def __init__(self, api_key):
        self.client = TavilyClient(api_key=api_key)
        
    def search(self, query, max_results=5):
        """
        Search for relevant content using Tavily API
        """
        try:
            response = self.client.search(
                query=query,
                search_depth="advanced",
                max_results=max_results,
                include_answer=False,
                include_raw_content=False
            )
            return response.get('results', [])
        except Exception as e:
            logging.error(f"Search failed: {str(e)}")
            return []
    
    def get_related_queries(self, query):
        """
        Generate related search queries
        """
        # This could be enhanced with LLM-based query expansion
        return [
            f"{query} research paper",
            f"{query} latest developments",
            f"{query} pros and cons"
        ]