dashboard.py

"""
Dashboard for the Bellingcat Python project using Smolagents.
"""
import os
from dotenv import load_dotenv
import hydra
from omegaconf import DictConfig

from smolagents import CodeAgent, OpenAIServerModel, tool, GradioUI

from config.hydra_config import AppConfig
from clients.ollama_client import OllamaClient
from services.embedding_service import EmbeddingService
from utils.db_utils import DatabaseManager
from utils.helpers import format_similarity_results

# Load environment variables
load_dotenv()


@tool
def retrieve(query: str) -> str:
    """
    Retrieves and formats similar content based on the provided query.

    This function performs a similarity search using the embedding service
    to find up to 5 most relevant content pieces matching the input query. It then formats
    the results into a structured string for easy reading and further processing by the agent.

    The search leverages embeddings to find semantically similar content, making it useful
    for retrieving relevant information from a knowledge base or document collection.

    Args:
        query (str): The search query string. Should be a descriptive phrase, sentence,
                     or keywords describing the desired content. For best results, use
                     natural language queries that clearly express the information need.
                     Examples: "machine learning algorithms", "history of artificial intelligence".

    Returns:
        str: A formatted string containing the search results. Results are separated by
             double newlines ('\n\n'). Each result block follows this exact format:
             "Source: {url}\nSimilarity: {sim}\n{content}"
             where:
             - url: The source URL of the content
             - sim: Similarity score (float between 0 and 1, higher is more similar)
             - content: The actual content snippet or text

             If no results are found, returns an empty string.

    Raises:
        Any exceptions propagated from the embedding service, such as:
        - ConnectionError: If unable to connect to the search backend
        - ValueError: If the query format is invalid
        - RuntimeError: If the search service is unavailable

    Note:
        This function is designed to be used as a tool in smolagents workflows.
        The limit of 5 results is fixed to balance information richness with response length.
        For debugging, the function relies on any logging within the embedding service.

    Example:
        >>> result = retrieve("quantum computing basics")
        >>> print(result)
        Source: https://example.com/quantum
        Similarity: 0.92
        Quantum computing uses quantum mechanics principles...

        Source: https://example.com/physics
        Similarity: 0.87
        Quantum mechanics fundamentals applied to computing...
    """
    # This will be initialized when the dashboard is configured
    global embedding_service, dashboard_config
    
    try:
        results = embedding_service.search_similar_content(query, limit=dashboard_config.search_limit)
        return format_similarity_results(results, limit=dashboard_config.search_limit)
    except Exception as e:
        return f"Error retrieving content: {str(e)}"


def create_dashboard_services(cfg: AppConfig):
    """
    Create the services needed for the dashboard.
    
    Args:
        cfg: Application configuration
        
    Returns:
        Tuple of (embedding_service, db_manager)
    """
    # Create database manager
    db_manager = DatabaseManager(cfg.database)
    
    # Create Ollama client for embeddings
    ollama_client = OllamaClient(cfg.embedding)
    
    # Create embedding service
    embedding_service = EmbeddingService(cfg.embedding, ollama_client, db_manager)
    
    return embedding_service, db_manager


@hydra.main(version_base=None, config_path="config", config_name="config")
def main(cfg: DictConfig) -> None:
    """
    Main entry point for the dashboard using Hydra configuration.
    
    Args:
        cfg: Hydra configuration object
    """
    global embedding_service, dashboard_config
    
    dashboard_config = cfg.dashboard
    
    # Create services
    embedding_service, db_manager = create_dashboard_services(cfg)
    
    # Initialize database to ensure it's available
    try:
        db_manager.init_db()
        print("Database initialized for dashboard")
    except Exception as e:
        print(f"Warning: Could not initialize database: {e}")
    
    # Check if embedding service is healthy
    if not embedding_service.health_check():
        print("Warning: Embedding service health check failed")
    
    # Create the model for the agent
    model = OpenAIServerModel(
        model_id=dashboard_config.model,
        api_key=cfg.clients.openrouter.api_key,
        api_base=cfg.clients.openrouter.base_url
    )
    
    # Create the agent with the retrieve tool
    agent = CodeAgent(
        tools=[retrieve],
        model=model,
    )
    
    # Create and launch the UI
    ui = GradioUI(agent)
    
    print(f"Launching dashboard on port {dashboard_config.server_port}")
    print(f"Using model: {dashboard_config.model}")
    print(f"Search limit: {dashboard_config.search_limit}")
    
    ui.launch(share=False, server_port=dashboard_config.server_port)


if __name__ == "__main__":
    main()