backend/dependencies.py

"""
Dependencies for FastAPI routes
"""

import logging
from typing import Generator, Optional, Any
from fastapi import Depends, HTTPException, status

# Initialize global testers
knowledge_graph_tester = None
prompt_reconstructor = None

# Import knowledge graph tester conditionally
try:
    from agentgraph.prompt_tester import KnowledgeGraphTester
    from agentgraph.prompt_reconstructor import PromptReconstructor
    TESTER_AVAILABLE = True
    # Update type annotations after imports
    knowledge_graph_tester: Optional[KnowledgeGraphTester] = None
    prompt_reconstructor: Optional[PromptReconstructor] = None
except ImportError:
    TESTER_AVAILABLE = False

# Import database utilities
from backend.database import get_db
from sqlalchemy.orm import Session

# Setup logger
logger = logging.getLogger("agent_monitoring_server")


def get_db_session() -> Generator[Session, None, None]:
    """Get a database session"""
    session = next(get_db())
    try:
        yield session
    finally:
        session.close()


def get_knowledge_graph_tester(kg_path: Optional[str] = None) -> Any:
    """Get or create a KnowledgeGraphTester instance"""
    global knowledge_graph_tester

    if not TESTER_AVAILABLE:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Prompt tester module not available"
        )

    if knowledge_graph_tester is None:
        try:
            from backend.server_config import DEFAULT_KNOWLEDGE_GRAPH
            
            # Use provided path or default
            knowledge_graph_path = kg_path or DEFAULT_KNOWLEDGE_GRAPH
            knowledge_graph_tester = KnowledgeGraphTester(
                knowledge_graph_path=knowledge_graph_path,
                model="gpt-4o"  # Default model
            )
            logger.info(f"Initialized KnowledgeGraphTester with {knowledge_graph_path}")
        except Exception as e:
            logger.error(f"Failed to initialize KnowledgeGraphTester: {str(e)}")
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Error initializing tester: {str(e)}"
            )

    return knowledge_graph_tester


def get_prompt_reconstructor() -> Any:
    """Get or create a PromptReconstructor instance"""
    global prompt_reconstructor

    if not TESTER_AVAILABLE:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Prompt reconstructor module not available"
        )

    if prompt_reconstructor is None:
        try:
            prompt_reconstructor = PromptReconstructor()
            logger.info("Initialized PromptReconstructor")
        except Exception as e:
            logger.error(f"Failed to initialize PromptReconstructor: {str(e)}")
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Error initializing reconstructor: {str(e)}"
            )

    return prompt_reconstructor