""" 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}")