src/analyzer/utils/query_logger.py

"""
query_logger.py — Logs user queries and AI responses for RLHF analysis

Automatically logs every interaction to CSV and JSONL formats.
Export to RLHF training format with export_logs.py script.
"""
from __future__ import annotations

import csv
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


class QueryLogger:
    """Logs user queries and AI responses to JSON and CSV for RLHF analysis."""

    def __init__(self, log_dir: str = "logs"):
        """
        Initialize query logger.

        Args:
            log_dir: Directory to store log files (default: "logs/")
        """
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)

        # Get today's date for log file naming
        today = datetime.now().strftime("%Y%m%d")
        self.jsonl_path = self.log_dir / f"queries_{today}.jsonl"
        self.csv_path = self.log_dir / f"queries_{today}.csv"

        # Initialize CSV file with headers if it doesn't exist
        if not self.csv_path.exists():
            self._init_csv()

    def _init_csv(self):
        """Initialize CSV file with headers."""
        headers = [
            'timestamp',
            'user_query',
            'ai_response',
            'tools_called',
            'response_time_ms',
            'success',
            'rating',
            'feedback',
            'model',
            'tokens_used'
        ]
        with open(self.csv_path, 'w', newline='', encoding='utf-8') as f:
            writer = csv.DictWriter(f, fieldnames=headers)
            writer.writeheader()

    def log_interaction(
        self,
        user_query: str,
        ai_response: str,
        *,
        tools_called: Optional[List[str]] = None,
        response_time_ms: Optional[int] = None,
        success: bool = True,
        rating: Optional[int] = None,
        feedback: Optional[str] = None,
        model: Optional[str] = None,
        tokens_used: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None
    ):
        """
        Log a user query and AI response.

        Args:
            user_query: The user's input query
            ai_response: The AI's response
            tools_called: List of tool names that were called
            response_time_ms: Response time in milliseconds
            success: Whether the interaction was successful
            rating: Optional 1-5 rating (for RLHF)
            feedback: Optional text feedback (for RLHF)
            model: Model name used (e.g., "gpt-4-mini")
            tokens_used: Number of tokens consumed
            metadata: Additional metadata to log
        """
        timestamp = datetime.utcnow().isoformat()

        # Prepare log entry
        log_entry = {
            'timestamp': timestamp,
            'user_query': user_query,
            'ai_response': ai_response,
            'tools_called': tools_called or [],
            'response_time_ms': response_time_ms,
            'success': success,
            'rating': rating,
            'feedback': feedback,
            'model': model,
            'tokens_used': tokens_used,
            'metadata': metadata or {}
        }

        # Write to JSONL (one JSON object per line)
        try:
            with open(self.jsonl_path, 'a', encoding='utf-8') as f:
                f.write(json.dumps(log_entry) + '\n')
        except Exception as e:
            logger.warning(f"Failed to write to JSONL log: {e}")

        # Write to CSV (Excel-compatible)
        try:
            with open(self.csv_path, 'a', newline='', encoding='utf-8') as f:
                writer = csv.DictWriter(f, fieldnames=[
                    'timestamp', 'user_query', 'ai_response', 'tools_called',
                    'response_time_ms', 'success', 'rating', 'feedback',
                    'model', 'tokens_used'
                ])
                writer.writerow({
                    'timestamp': timestamp,
                    'user_query': user_query,
                    'ai_response': ai_response,
                    'tools_called': ','.join(tools_called or []),
                    'response_time_ms': response_time_ms,
                    'success': success,
                    'rating': rating or '',
                    'feedback': feedback or '',
                    'model': model or '',
                    'tokens_used': tokens_used or ''
                })
        except Exception as e:
            logger.warning(f"Failed to write to CSV log: {e}")

        logger.debug(f"Logged interaction: {user_query[:50]}...")

    def get_stats(self) -> Dict[str, Any]:
        """
        Get statistics about logged queries.

        Returns:
            Dictionary with statistics (total queries, success rate, avg response time, etc.)
        """
        try:
            with open(self.csv_path, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                rows = list(reader)

            if not rows:
                return {
                    'total_queries': 0,
                    'successful_queries': 0,
                    'failed_queries': 0,
                    'avg_response_time_ms': 0,
                    'avg_rating': 0,
                    'rated_queries': 0,
                    'tool_usage': {}
                }

            total = len(rows)
            successful = sum(1 for r in rows if r.get('success') == 'True')
            failed = total - successful

            # Calculate average response time
            response_times = [
                int(r['response_time_ms'])
                for r in rows
                if r.get('response_time_ms') and r['response_time_ms'].isdigit()
            ]
            avg_response_time = sum(response_times) / len(response_times) if response_times else 0

            # Calculate average rating
            ratings = [
                int(r['rating'])
                for r in rows
                if r.get('rating') and r['rating'].isdigit()
            ]
            avg_rating = sum(ratings) / len(ratings) if ratings else 0
            rated_queries = len(ratings)

            # Count tool usage
            tool_usage = {}
            for row in rows:
                tools = row.get('tools_called', '').split(',')
                for tool in tools:
                    tool = tool.strip()
                    if tool:
                        tool_usage[tool] = tool_usage.get(tool, 0) + 1

            return {
                'total_queries': total,
                'successful_queries': successful,
                'failed_queries': failed,
                'avg_response_time_ms': avg_response_time,
                'avg_rating': avg_rating,
                'rated_queries': rated_queries,
                'tool_usage': tool_usage
            }
        except Exception as e:
            logger.warning(f"Failed to calculate stats: {e}")
            return {
                'total_queries': 0,
                'successful_queries': 0,
                'failed_queries': 0,
                'avg_response_time_ms': 0,
                'avg_rating': 0,
                'rated_queries': 0,
                'tool_usage': {}
            }

    def export_for_rlhf(self, output_path: Optional[Path] = None) -> Path:
        """
        Export logs in RLHF training format.

        Format:
        [
            {
                "prompt": "user query",
                "completion": "ai response",
                "rating": 5,
                "feedback": "Great!",
                "tools_used": ["list_grants"],
                "timestamp": "2025-10-23T10:15:30"
            },
            ...
        ]

        Args:
            output_path: Optional custom output path

        Returns:
            Path to exported file
        """
        if output_path is None:
            today = datetime.now().strftime("%Y%m%d")
            output_path = self.log_dir / f"rlhf_data_{today}.json"

        rlhf_data = []

        try:
            # Read from JSONL
            with open(self.jsonl_path, 'r', encoding='utf-8') as f:
                for line in f:
                    entry = json.loads(line)

                    # Only include successful interactions
                    if not entry.get('success', False):
                        continue

                    rlhf_entry = {
                        'prompt': entry['user_query'],
                        'completion': entry['ai_response'],
                        'rating': entry.get('rating'),
                        'feedback': entry.get('feedback'),
                        'tools_used': entry.get('tools_called', []),
                        'timestamp': entry['timestamp'],
                        'response_time_ms': entry.get('response_time_ms'),
                        'model': entry.get('model')
                    }

                    rlhf_data.append(rlhf_entry)

            # Write RLHF format
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(rlhf_data, f, indent=2, ensure_ascii=False)

            logger.info(f"Exported {len(rlhf_data)} interactions to {output_path}")
            return output_path

        except Exception as e:
            logger.error(f"Failed to export RLHF data: {e}")
            raise


# Global singleton instance
_query_logger: Optional[QueryLogger] = None


def get_query_logger(log_dir: str = "logs") -> QueryLogger:
    """
    Get or create the global query logger instance.

    Args:
        log_dir: Directory to store log files

    Returns:
        QueryLogger instance
    """
    global _query_logger
    if _query_logger is None:
        _query_logger = QueryLogger(log_dir=log_dir)
    return _query_logger


# Quick self-test
if __name__ == "__main__":
    logger = get_query_logger(log_dir="_out/test_logs")

    # Log a test interaction
    logger.log_interaction(
        user_query="What grants are available for batteries?",
        ai_response="Here are the battery-related grants: 1. Battery Innovation Grant...",
        tools_called=["list_grants"],
        response_time_ms=1500,
        success=True,
        rating=5,
        feedback="Very helpful!",
        model="gpt-4-mini"
    )

    # Get stats
    stats = logger.get_stats()
    print("Statistics:", json.dumps(stats, indent=2))

    # Export for RLHF
    output = logger.export_for_rlhf()
    print(f"Exported to: {output}")