main.py

"""
FastAPI application entry point.
"""
import os
import logging
from contextlib import asynccontextmanager
from dotenv import load_dotenv

from fastapi import FastAPI, Depends, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from sqlalchemy import text
from sqlalchemy.orm import Session
from pydantic import BaseModel
from typing import Dict, Any, Optional
import asyncio


# Load environment variables from .env file
load_dotenv()

from config import settings
from database import get_db

# Configure logging
logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Lifespan context manager for startup and shutdown events.
    
    Yields:
        None: Application is ready
    """
    # Startup
    logger.info("Starting application...")
    yield
    # Shutdown
    logger.info("Shutting down application...")


# Create FastAPI app
app = FastAPI(
    title="Self-Improving Engine API",
    description="FastAPI application with PostgreSQL database and ACE system",
    version="1.0.0",
    debug=settings.debug,
    lifespan=lifespan,
)


@app.get("/")
async def root():
    """
    Root endpoint.
    
    Returns:
        dict: Welcome message
    """
    return {"message": "Welcome to Self-Improving Engine API"}


@app.get("/health")
async def health_check(db: Session = Depends(get_db)):
    """
    Health check endpoint.
    
    Checks database connectivity and returns service status.
    
    Args:
        db: Database session dependency
    
    Returns:
        JSONResponse: Health status with database connection status
    """
    try:
        # Test database connection
        db.execute(text("SELECT 1"))
        db_status = "connected"
    except Exception as e:
        logger.error(f"Database health check failed: {e}")
        db_status = "disconnected"
    
    return JSONResponse(
        status_code=200 if db_status == "connected" else 503,
        content={
            "status": "healthy" if db_status == "connected" else "unhealthy",
            "database": db_status,
        }
    )


# Lazy initialization of ACE components
_ace_initialized = False
fraud_agent = None
selector = None
reflector = None
curator = None
training_pipeline = None
playbook = None


def init_ace_components():
    """Initialize ACE components lazily."""
    global _ace_initialized, fraud_agent, selector, reflector, curator, training_pipeline, playbook
    
    if _ace_initialized:
        return
    
    try:
        from ace.bullet_playbook import BulletPlaybook
        from ace.hybrid_selector import HybridSelector
        from ace.reflector import Reflector
        from ace.curator import Curator
        from ace.training_pipeline import TrainingPipeline
        from agents.mock_fraud_agent import MockFraudAgent
        
        fraud_agent = MockFraudAgent()
        selector = HybridSelector()
        reflector = Reflector()
        curator = Curator()
        training_pipeline = TrainingPipeline(fraud_agent, selector, reflector, curator)
        playbook = BulletPlaybook()
        
        _ace_initialized = True
        logger.info("ACE components initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize ACE components: {e}")
        raise


# Request models
class AnalyzeRequest(BaseModel):
    transaction: Dict[str, Any]
    mode: str = "hybrid"


class EvaluateRequest(BaseModel):
    input_text: str
    node: str
    output: str
    ground_truth: Optional[str] = None


class GetBulletsRequest(BaseModel):
    query: str
    node: str
    mode: str = "hybrid"


class TrainRequest(BaseModel):
    """Request model for offline training."""
    dataset: list[Dict[str, Any]]
    node: str = "fraud_detection"
    max_samples: int = 10


class TraceRequest(BaseModel):
    """Request model for tracing (online learning)."""
    model_config = {"protected_namespaces": ()}
    
    input_text: str
    node: str
    output: str
    model_type: Optional[str] = None
    session_id: Optional[str] = None
    run_id: Optional[str] = None
    ground_truth: Optional[str] = None
    agent_reasoning: Optional[str] = None
    bullet_ids: Optional[Dict[str, list[str]]] = None  # {"full": [...], "online": [...]}


class ContextRequest(BaseModel):
    """Request model for getting context."""
    input_text: str
    node: str
    max_bullets_per_evaluator: int = 10


@app.post("/api/v1/train")
async def train_offline(request: TrainRequest, db: Session = Depends(get_db)):
    """
    Train offline mode with provided dataset.
    
    Receives dataset with inputs and outputs, trains bullets on it with Darwin-Gödel evolution.
    
    Args:
        request: Training request with dataset
        db: Database session
    
    Returns:
        dict: Training results with statistics
    """
    init_ace_components()
    
    try:
        from ace.training_pipeline import TrainingPipeline
        from ace.bullet_playbook import BulletPlaybook
        from ace.hybrid_selector import HybridSelector
        from ace.reflector import Reflector
        from ace.curator import Curator
        from ace.pattern_manager import PatternManager
        from ace.darwin_bullet_evolver import DarwinBulletEvolver
        
        # Initialize components with database
        pattern_manager = PatternManager(db_session=db)
        selector = HybridSelector(db_session=db, pattern_manager=pattern_manager)
        reflector = Reflector()
        curator = Curator()
        
        # Initialize Darwin-Gödel evolution if enabled
        darwin_evolver = None
        if settings.enable_darwin_evolution:
            darwin_evolver = DarwinBulletEvolver(db_session=db)
            logger.info("Darwin-Gödel evolution enabled")
        else:
            logger.info("Darwin-Gödel evolution disabled")
        
        # Create dummy agent (we don't actually use it in training)
        class DummyAgent:
            pass
        
        training_pipeline = TrainingPipeline(DummyAgent(), selector, reflector, curator, darwin_evolver=darwin_evolver)
        playbook = BulletPlaybook(db_session=db)
        
        # Process dataset
        dataset = request.dataset[:request.max_samples]
        logger.info(f"Training on {len(dataset)} samples for node: {request.node}")
        
        bullets_generated = []
        for idx, item in enumerate(dataset):
            query = item.get('query', str(item))
            ground_truth = item.get('answer', item.get('output', 'DECLINE'))
            predicted = item.get('predicted', ground_truth)  # Use provided prediction or ground truth
            
            # Get pattern classification
            pattern_id, confidence = pattern_manager.classify_input_to_category(
                input_summary=query,
                node=request.node
            )
            
            # Generate bullet
            bullet_id = await training_pipeline.add_bullet_from_reflection(
                query=query,
                predicted=predicted,
                correct=ground_truth,
                node=request.node,
                agent_reasoning="",
                playbook=playbook,
                source="offline",
                evaluator=request.node
            )
            
            if bullet_id:
                bullets_generated.append(bullet_id)
        
        # Get final stats
        final_bullets = playbook.get_bullets_for_node(request.node)
        
        return {
            "status": "success",
            "node": request.node,
            "samples_processed": len(dataset),
            "bullets_generated": len(bullets_generated),
            "total_bullets": len(final_bullets),
            "unique_bullets": len(set(b.content for b in final_bullets))
        }
    
    except Exception as e:
        logger.error(f"Error in train: {e}")
        raise HTTPException(status_code=500, detail=str(e))



@app.post("/api/v1/trace")
async def trace_and_learn(request: TraceRequest, db: Session = Depends(get_db)):
    """
    Trace transaction and learn online (synchronous).
    
    Saves transaction to database and executes online training.
    Uses LLM judge for evaluation if ground_truth is not provided.
    Returns only after all processing is complete.
    
    Args:
        request: Trace request with input, output, and ground truth
        db: Database session
    
    Returns:
        dict: Transaction info with pattern_id and is_correct
    """
    try:
        from models import Transaction, LLMJudge
        from ace.pattern_manager import PatternManager
        from ace.training_pipeline import TrainingPipeline
        from ace.bullet_playbook import BulletPlaybook
        from ace.hybrid_selector import HybridSelector
        from ace.reflector import Reflector
        from ace.curator import Curator
        from ace.darwin_bullet_evolver import DarwinBulletEvolver
        from openai import OpenAI
        import json
        import os
        
        logger.info(f"Processing trace for node={request.node}, session_id={request.session_id}, run_id={request.run_id}")
        
        # Step 1: Determine mode from model_type (map "full" to "offline_online")
        mode = request.model_type or "online"
        if mode == "full":
            mode = "offline_online"
        
        # Step 2: Pattern classification
        pattern_manager = PatternManager(db_session=db)
        pattern_id, confidence = pattern_manager.classify_input_to_category(
            input_summary=request.input_text,
            node=request.node
        )
        logger.info(f"Pattern classified: pattern_id={pattern_id}, confidence={confidence:.2f}")
        
        # Step 3: Determine correctness (will be evaluated in Step 6)
        is_correct = False
        correct_decision = request.output
        
        if request.ground_truth:
            is_correct = (request.output == request.ground_truth)
            correct_decision = request.ground_truth
        else:
            # Will be determined based on evaluator consensus in Step 6
            is_correct = False
        
        # Step 4: Save transaction
        system_prompt = request.input_text.split("System Prompt:")[1].strip().split("User Prompt:")[0].strip()
        user_prompt = request.input_text.split("User Prompt:")[1].strip()
        transaction_data = {
            "systemprompt": system_prompt,
            "userprompt": user_prompt,
            "output": request.output,
            "reasoning": request.agent_reasoning or ""
        }
        
        txn = Transaction(
            transaction_data=transaction_data,
            mode=mode,
            node=request.node,
            session_id=request.session_id,
            run_id=request.run_id,
            predicted_decision=request.output,
            correct_decision=correct_decision or request.output,
            is_correct=is_correct,
            input_pattern_id=pattern_id
        )
        db.add(txn)
        db.commit()
        db.flush()
        
        # Step 5: Get evaluators for this node from LLM judges table
        from models import LLMJudge
        evaluators = db.query(LLMJudge.evaluator).filter(
            LLMJudge.node == request.node,
            LLMJudge.is_active == True
        ).distinct().all()
        
        evaluator_names = [e[0] for e in evaluators] if evaluators else []
        
        # Only generate bullets if evaluators exist for this node
        if not evaluator_names:
            logger.info(f"No LLM judges found for node {request.node}. Skipping bullet generation.")
            return {
                "status": "success",
                "node": request.node,
                "transaction_id": txn.id,
                "pattern_id": pattern_id,
                "is_correct": is_correct,
                "generated_bullets": [],
                "message": "Transaction saved, but no LLM judges found for this node"
            }
        # Process improvement (bullet generation) for nodes with evaluators
        logger.info(f"Processing with evaluators: {evaluator_names}")
        
        # Initialize components for online learning
        selector = HybridSelector(db_session=db, pattern_manager=pattern_manager)
        reflector = Reflector()
        curator = Curator()
        
        # Initialize Darwin-Gödel evolution if enabled
        darwin_evolver = None
        enable_evolution = os.getenv("ENABLE_DARWIN_EVOLUTION", "true").lower() == "true"
        if enable_evolution:
            darwin_evolver = DarwinBulletEvolver(db_session=db)
            logger.info("Darwin-Gödel evolution enabled")
        
        class DummyAgent:
            pass
        
        training_pipeline = TrainingPipeline(DummyAgent(), selector, reflector, curator, darwin_evolver=darwin_evolver)
        playbook = BulletPlaybook(db_session=db)
        
        # Step 6: Evaluate with all evaluators (for ALL modes including vanilla)
        generated_bullets = []
        evaluator_results = {}  # Cache evaluator results for reuse in Step 8
        
        for evaluator_name in evaluator_names:
            # Get the specific judge for this evaluator
            evaluator_judge = db.query(LLMJudge).filter(
                LLMJudge.node == request.node,
                LLMJudge.evaluator == evaluator_name,
                LLMJudge.is_active == True
            ).first()

            input_text = txn.transaction_data["userprompt"]
            
            # Get judge reasoning for this specific evaluator
            evaluator_judge_reasoning = ""
            evaluator_is_correct = False
            if evaluator_judge:
                client = OpenAI()
                
                if request.ground_truth:
                    evaluation_prompt = f"""{evaluator_judge.system_prompt}

Input: {input_text}

Output: {request.output}

Ground Truth: {request.ground_truth}

Evaluate the output based on these criteria:
{json.dumps(evaluator_judge.evaluation_criteria, indent=2) if evaluator_judge.evaluation_criteria else "Use your best judgment"}

Respond in JSON format:
{{
    "is_correct": true/false,
    "confidence": 0.0-1.0,
    "reasoning": "brief explanation"
}}"""
                else:
                    evaluation_prompt = f"""{evaluator_judge.system_prompt}

Input: {input_text}

Output: {request.output}

Evaluate the output based on these criteria:
{json.dumps(evaluator_judge.evaluation_criteria, indent=2) if evaluator_judge.evaluation_criteria else "Use your best judgment"}

Respond in JSON format:
{{
    "is_correct": true/false,
    "confidence": 0.0-1.0,
    "reasoning": "brief explanation"
}}"""
                
                response = client.chat.completions.create(
                    model=evaluator_judge.model,
                    messages=[
                        {"role": "system", "content": evaluator_judge.system_prompt},
                        {"role": "user", "content": evaluation_prompt}
                    ],
                    response_format={"type": "json_object"},
                )
                
                result = json.loads(response.choices[0].message.content)
                evaluator_judge_reasoning = result.get("reasoning", "")
                evaluator_is_correct = result.get("is_correct", False)
                
                # Cache the result for reuse in Step 8
                evaluator_results[evaluator_name] = evaluator_is_correct
            elif request.ground_truth:
                # If no judge but ground truth exists, compare
                evaluator_is_correct = (request.output == request.ground_truth)
                evaluator_results[evaluator_name] = evaluator_is_correct

            # Step 6b: Generate bullet ONLY if not vanilla mode
            if mode != "vanilla":
                bullet_id = await training_pipeline.add_bullet_from_reflection(
                    query=request.input_text,
                    predicted=request.output,
                    correct=request.ground_truth or request.output,
                    node=request.node,
                    agent_reasoning=request.agent_reasoning or "",
                    playbook=playbook,
                    source="online",
                    evaluator=evaluator_name,
                    judge_reasoning=evaluator_judge_reasoning
                )
                
                if bullet_id:
                    logger.info(f"Generated bullet for evaluator {evaluator_name}: {bullet_id}")
                    generated_bullets.append(bullet_id)
                else:
                    logger.info(f"No bullet generated for evaluator {evaluator_name} (likely duplicate)")
        
        if mode != "vanilla":
            logger.info(f"Generated {len(generated_bullets)} bullets out of {len(evaluator_names)} evaluators")
            
            # Step 7: Record bullet effectiveness
            if request.bullet_ids and pattern_id:
                if "full" in request.bullet_ids:
                    for bullet_id_used in request.bullet_ids["full"]:
                        pattern_manager.record_bullet_effectiveness(
                            pattern_id=pattern_id,
                            bullet_id=bullet_id_used,
                            node=request.node,
                            is_helpful=is_correct
                        )
                
                if "online" in request.bullet_ids:
                    for bullet_id_used in request.bullet_ids["online"]:
                        pattern_manager.record_bullet_effectiveness(
                            pattern_id=pattern_id,
                            bullet_id=bullet_id_used,
                            node=request.node,
                            is_helpful=is_correct
                        )
            
        # Step 8: Update session run metrics (only if session_id and run_id provided)
        if request.session_id and request.run_id:
            from models import SessionRunMetrics
            
            # Track metrics for each evaluator
            for evaluator_name in evaluator_names:
                # Use cached result from Step 6 if available, otherwise evaluate
                evaluator_is_correct = evaluator_results.get(evaluator_name)
                
                if evaluator_is_correct is None:
                    continue

                metrics = db.query(SessionRunMetrics).filter(
                    SessionRunMetrics.session_id == request.session_id,
                    SessionRunMetrics.run_id == request.run_id,
                    SessionRunMetrics.node == request.node,
                    SessionRunMetrics.evaluator == evaluator_name,
                    SessionRunMetrics.mode == mode
                ).first()
                
                if not metrics:
                    metrics = SessionRunMetrics(
                        session_id=request.session_id,
                        run_id=request.run_id,
                        node=request.node,
                        evaluator=evaluator_name,
                        mode=mode,
                        correct_count=0,
                        total_count=0,
                        accuracy=0.0
                    )
                    db.add(metrics)
                
                metrics.total_count += 1
                if evaluator_is_correct:
                    metrics.correct_count += 1
                metrics.accuracy = metrics.correct_count / metrics.total_count if metrics.total_count > 0 else 0.0
                
                logger.info(f"Updated metrics for evaluator {evaluator_name}: is_correct={evaluator_is_correct}, accuracy={metrics.accuracy:.3f}")
            
            db.commit()
            logger.info(f"Updated metrics for {len(evaluator_names)} evaluators (mode={mode})")
        
        logger.info(f"Trace processing completed for transaction {txn.id}")
        
        # Get bullet details for all generated bullets
        bullet_infos = []
        if generated_bullets:
            from models import Bullet as BulletModel
            for bullet_id in generated_bullets:
                bullet = db.query(BulletModel).filter(BulletModel.id == bullet_id).first()
                if bullet:
                    bullet_infos.append({
                        "bullet_id": bullet.id,
                        "content": bullet.content,
                        "evaluator": bullet.evaluator,
                        "source": bullet.source
                    })
        
        return {
            "status": "success",
            "node": request.node,
            "transaction_id": txn.id,
            "pattern_id": pattern_id,
            "is_correct": is_correct,
            "generated_bullets": bullet_infos,
            "message": "Processing completed"
        }
    
    except Exception as e:
        logger.error(f"Error in trace: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/v1/context")
async def get_context(request: ContextRequest, db: Session = Depends(get_db)):
    """
    Get context for agent prompt.
    
    Returns full text context with bullets organized by evaluator.
    Returns both full context (offline + online) and online-only context.
    
    Args:
        request: Context request with input and node
        db: Database session
    
    Returns:
        dict: Context with 'full' and 'online' keys containing full text
    """
    init_ace_components()
    
    try:
        from ace.bullet_playbook import BulletPlaybook
        from ace.hybrid_selector import HybridSelector
        from ace.pattern_manager import PatternManager
        
        # Initialize components with database
        pattern_manager = PatternManager(db_session=db)
        selector = HybridSelector(db_session=db, pattern_manager=pattern_manager)
        playbook = BulletPlaybook(db_session=db)
        
        # Get pattern classification
        pattern_id, confidence = pattern_manager.classify_input_to_category(
            input_summary=request.input_text,
            node=request.node
        )
        
        # Get all evaluators from LLM judges table
        from models import LLMJudge
        all_evaluators = db.query(LLMJudge.evaluator).filter(
            LLMJudge.node == request.node,
            LLMJudge.is_active == True
        ).distinct().all()
        
        evaluator_names = [e[0] for e in all_evaluators] if all_evaluators else []
        
        # If no evaluators, return empty context
        if not evaluator_names:
            logger.info(f"No evaluators found for node {request.node}. Returning empty context.")
            return {
                "status": "success",
                "node": request.node,
                "pattern_id": None,
                "bullet_ids": {
                    "full": [],
                    "online": []
                },
                "context": {
                    "full": "",
                    "online": ""
                }
            }
        
        # Build full context (offline + online bullets)
        full_context = ""
        online_context = ""
        full_bullet_ids = []  # Track bullets for full context
        online_bullet_ids = []  # Track bullets for online context
        
        for evaluator_name in evaluator_names:
            # Get ALL bullets for this evaluator
            evaluator_bullets = playbook.get_bullets_for_node(request.node, evaluator=evaluator_name)
            
            # Select top bullets for this evaluator using intelligent selection
            if evaluator_bullets:
                # Temporarily create a filtered playbook for this evaluator
                temp_bullets = [b for b in evaluator_bullets]
                temp_playbook = BulletPlaybook()
                temp_playbook.bullets = temp_bullets
                temp_playbook._node_index[request.node] = temp_bullets
                
                selected, _ = selector.select_bullets(
                    query=request.input_text,
                    node=request.node,
                    playbook=temp_playbook,
                    n_bullets=min(request.max_bullets_per_evaluator, len(temp_bullets), 10),
                    iteration=0,
                    pattern_id=pattern_id
                )
                
                # Add to full context
                if selected:
                    full_context += f"\n\n{evaluator_name.upper()} Rules:\n"
                    for bullet in selected[:request.max_bullets_per_evaluator]:
                        full_context += f"- {bullet.content}\n"
                        full_bullet_ids.append(bullet.id)  # Track bullet IDs for full context
                    
                    # Add online and evolved bullets to online context
                    online_bullets = [b for b in selected if b.source in ("online", "evolution")]
                    if online_bullets:
                        online_context += f"\n\n{evaluator_name.upper()} Rules:\n"
                        for bullet in online_bullets[:request.max_bullets_per_evaluator]:
                            online_context += f"- {bullet.content}\n"
                            online_bullet_ids.append(bullet.id)  # Track bullet IDs for online context
        
        return {
            "status": "success",
            "node": request.node,
            "pattern_id": pattern_id,
            "bullet_ids": {
                "full": full_bullet_ids,      # Bullets for offline + online context
                "online": online_bullet_ids   # Bullets for online-only context
            },
            "context": {
                "full": full_context.strip(),
                "online": online_context.strip()
            }
        }
    
    except Exception as e:
        logger.error(f"Error in context: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/v1/metrics/{session_id}")
async def get_session_metrics(session_id: str, db: Session = Depends(get_db)):
    """
    Get metrics for a specific session.
    
    Returns metrics grouped by run_id and evaluator.
    
    Args:
        session_id: Session identifier
        db: Database session
    
    Returns:
        dict: Metrics organized by run_id and evaluator
    """
    try:
        from models import SessionRunMetrics
        
        # Get all metrics for this session
        metrics = db.query(SessionRunMetrics).filter(
            SessionRunMetrics.session_id == session_id
        ).all()
        
        if not metrics:
            return {
                "status": "success",
                "session_id": session_id,
                "message": "No metrics found for this session",
                "metrics": {}
            }
        
        # Organize metrics by run_id, evaluator, and mode (aggregate across all nodes)
        organized_metrics = {}
        
        for metric in metrics:
            run_id = metric.run_id
            evaluator = metric.evaluator
            mode = metric.mode
            
            if run_id not in organized_metrics:
                organized_metrics[run_id] = {}
            
            if evaluator not in organized_metrics[run_id]:
                organized_metrics[run_id][evaluator] = {}
            
            if mode not in organized_metrics[run_id][evaluator]:
                organized_metrics[run_id][evaluator][mode] = {
                    "correct_count": 0,
                    "total_count": 0,
                    "accuracy": 0.0
                }
            
            # Accumulate metrics for this mode (aggregate across all nodes)
            organized_metrics[run_id][evaluator][mode]["correct_count"] += metric.correct_count
            organized_metrics[run_id][evaluator][mode]["total_count"] += metric.total_count
        
        # Calculate accuracy for each mode
        for run_id, evaluators in organized_metrics.items():
            for evaluator, modes in evaluators.items():
                for mode, data in modes.items():
                    if data["total_count"] > 0:
                        data["accuracy"] = data["correct_count"] / data["total_count"]
                    else:
                        data["accuracy"] = 0.0
        
        return {
            "status": "success",
            "session_id": session_id,
            "metrics": organized_metrics
        }
    
    except Exception as e:
        logger.error(f"Error getting metrics: {e}")
        raise HTTPException(status_code=500, detail=str(e))