src/logging/logger.py

"""Logger implementation for interaction tracking."""

import os
import json
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Optional, List, Dict, Any
from contextlib import contextmanager
from threading import Lock
import time

from src.logging.schema import (
    InteractionLog, ToolCall, Citation, QualityRatings, Labels, ErrorDetail
)


class InteractionTracker:
    """Helper class to track an interaction and its components."""
    
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.response: Optional[str] = None
        self.tool_calls: List[ToolCall] = []
        self.retrieved_ctx_ids: List[str] = []
        self.citations: List[Citation] = []
        self.error: Optional[ErrorDetail] = None
        self.start_time: float = time.time()
        
    def set_response(self, response: str) -> None:
        """Set the model response."""
        self.response = response
        
    def add_tool_call(self, tool_call: ToolCall) -> None:
        """Add a tool call record."""
        self.tool_calls.append(tool_call)
        
    def add_retrieved_context(self, doc_id: str) -> None:
        """Add a retrieved document ID."""
        if doc_id not in self.retrieved_ctx_ids:
            self.retrieved_ctx_ids.append(doc_id)
            
    def add_citation(self, citation: Citation) -> None:
        """Add a citation."""
        self.citations.append(citation)
        
    def set_error(self, error: ErrorDetail) -> None:
        """Set error information."""
        self.error = error
        
    def get_latency_ms(self) -> int:
        """Get elapsed time in milliseconds."""
        return int((time.time() - self.start_time) * 1000)


class InteractionLogger:
    """Singleton logger for interaction tracking."""
    
    _instance: Optional['InteractionLogger'] = None
    _lock = Lock()
    
    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
            
        self.log_dir = Path("_out/logs")
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.feedback_dir = self.log_dir
        self._initialized = True
        
    def _get_log_file_path(self) -> Path:
        """Get path to today's log file."""
        date_str = datetime.now(timezone.utc).strftime("%Y%m%d")
        return self.log_dir / f"interactions_{date_str}.jsonl"

    def _get_feedback_file_path(self) -> Path:
        """Get path to today's feedback file."""
        date_str = datetime.now(timezone.utc).strftime("%Y%m%d")
        return self.feedback_dir / f"feedback_{date_str}.jsonl"
    
    def log_interaction(
        self,
        session_id: str,
        prompt: Optional[str] = None,
        response: Optional[str] = None,
        model_version: Optional[str] = None,
        prompt_version: Optional[str] = None,
        tools_schema_version: Optional[str] = None,
        retrieved_ctx_ids: Optional[List[str]] = None,
        citations: Optional[List[Citation]] = None,
        tool_calls: Optional[List[ToolCall]] = None,
        quality: Optional[QualityRatings] = None,
        labels: Optional[Labels] = None,
        correction: Optional[str] = None,
        competition_id: Optional[str] = None,
        snapshot_id: Optional[str] = None,
        lat_ms: Optional[int] = None,
        error: Optional[ErrorDetail] = None,
        user_id: Optional[str] = None,
        environment: Optional[str] = None,
        api_endpoint: Optional[str] = None,
    ) -> None:
        """
        Log an interaction.
        
        Args:
            session_id: Unique session identifier
            prompt: Input prompt/question
            response: Model response
            model_version: Version of model used
            prompt_version: Version of prompt used
            tools_schema_version: Version of tools schema
            retrieved_ctx_ids: List of retrieved document IDs
            citations: List of citations in response
            tool_calls: List of tool calls made
            quality: Quality ratings
            labels: Review labels
            correction: Corrected response if provided
            competition_id: Competition ID
            snapshot_id: Content snapshot ID
            lat_ms: Latency in milliseconds
            error: Error details if any
            user_id: User ID
            environment: Environment name
            api_endpoint: API endpoint called
        """
        log = InteractionLog(
            session_id=session_id,
            prompt=prompt,
            response=response,
            model_version=model_version,
            prompt_version=prompt_version,
            tools_schema_version=tools_schema_version,
            retrieved_ctx_ids=retrieved_ctx_ids or [],
            citations=citations or [],
            tool_calls=tool_calls or [],
            quality=quality,
            labels=labels,
            correction=correction,
            competition_id=competition_id,
            snapshot_id=snapshot_id,
            lat_ms=lat_ms,
            error=error,
            user_id=user_id,
            environment=environment,
            api_endpoint=api_endpoint,
            citations_present=bool(citations),
            rag_empty=not (retrieved_ctx_ids and len(retrieved_ctx_ids) > 0),
        )
        
        # Write to JSONL file
        log_path = self._get_log_file_path()
        with open(log_path, "a") as f:
            f.write(log.to_jsonl() + "\n")
    
    def log_feedback(
        self,
        session_id: str,
        rating_type: str,
        quality_ratings: Optional[QualityRatings] = None,
        correction: Optional[str] = None,
        reason: Optional[str] = None,
    ) -> None:
        """
        Log user feedback.
        
        Args:
            session_id: Session ID
            rating_type: "thumbs_up" or "thumbs_down"
            quality_ratings: Quality ratings if provided
            correction: Correction if provided
            reason: Reason for feedback
        """
        feedback = {
            "ts": datetime.now(timezone.utc).isoformat(),
            "session_id": session_id,
            "rating_type": rating_type,
            "quality_ratings": quality_ratings.model_dump() if quality_ratings else None,
            "correction": correction,
            "reason": reason,
        }
        
        feedback_path = self._get_feedback_file_path()
        with open(feedback_path, "a") as f:
            f.write(json.dumps(feedback) + "\n")
    
    @contextmanager
    def track_interaction(self, session_id: str, **log_kwargs):
        """
        Context manager for automatic interaction tracking.
        
        Args:
            session_id: Session ID
            **log_kwargs: Additional kwargs to pass to log_interaction
            
        Yields:
            InteractionTracker instance
        """
        tracker = InteractionTracker(session_id)
        try:
            yield tracker
        finally:
            # Log the interaction with collected data
            self.log_interaction(
                session_id=session_id,
                response=tracker.response,
                retrieved_ctx_ids=tracker.retrieved_ctx_ids,
                citations=tracker.citations,
                tool_calls=tracker.tool_calls,
                lat_ms=tracker.get_latency_ms(),
                error=tracker.error,
                **log_kwargs
            )


# Global logger instance
_logger: Optional[InteractionLogger] = None


def get_logger() -> InteractionLogger:
    """
    Get or create the global logger instance.
    
    Returns:
        InteractionLogger instance
    """
    global _logger
    if _logger is None:
        _logger = InteractionLogger()
    return _logger