pylance-mcp-server 1.0.0

package/mcp_server/logging_db.py

@@ -0,0 +1,403 @@
+"""
+Conversation Logging Database
+Logs all user prompts, LLM responses, and tool calls for training data collection.
+"""
+
+import sqlite3
+import json
+import time
+from pathlib import Path
+from typing import Optional, Dict, Any, List
+from datetime import datetime
+from contextlib import contextmanager
+
+
+class ConversationLogger:
+    """SQLite database for logging MCP conversations and tool usage."""
+    
+    def __init__(self, db_path: Optional[str] = None):
+        """
+        Initialize conversation logger.
+        
+        Args:
+            db_path: Path to SQLite database file. Defaults to ~/.pylance_mcp/conversations.db
+        """
+        if db_path is None:
+            db_dir = Path.home() / ".pylance_mcp"
+            db_dir.mkdir(parents=True, exist_ok=True)
+            db_path = str(db_dir / "conversations.db")
+        
+        self.db_path = db_path
+        self._init_database()
+    
+    def _init_database(self):
+        """Create database tables if they don't exist."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            
+            # Conversations table - tracks user sessions
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS conversations (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    session_id TEXT UNIQUE NOT NULL,
+                    workspace_root TEXT,
+                    python_version TEXT,
+                    venv_path TEXT,
+                    started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    ended_at TIMESTAMP,
+                    total_messages INTEGER DEFAULT 0,
+                    total_tool_calls INTEGER DEFAULT 0,
+                    metadata TEXT
+                )
+            """)
+            
+            # Messages table - logs every user prompt and LLM response
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    conversation_id INTEGER NOT NULL,
+                    message_type TEXT NOT NULL,
+                    role TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    token_count INTEGER,
+                    model_name TEXT,
+                    metadata TEXT,
+                    FOREIGN KEY (conversation_id) REFERENCES conversations(id)
+                )
+            """)
+            
+            # Tool calls table - logs every MCP tool invocation
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS tool_calls (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    message_id INTEGER,
+                    conversation_id INTEGER NOT NULL,
+                    tool_name TEXT NOT NULL,
+                    arguments TEXT NOT NULL,
+                    result TEXT,
+                    success BOOLEAN,
+                    error_message TEXT,
+                    execution_time_ms REAL,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    FOREIGN KEY (message_id) REFERENCES messages(id),
+                    FOREIGN KEY (conversation_id) REFERENCES conversations(id)
+                )
+            """)
+            
+            # Files modified table - tracks which files were changed
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS files_modified (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    tool_call_id INTEGER NOT NULL,
+                    file_path TEXT NOT NULL,
+                    operation TEXT NOT NULL,
+                    lines_added INTEGER DEFAULT 0,
+                    lines_removed INTEGER DEFAULT 0,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    FOREIGN KEY (tool_call_id) REFERENCES tool_calls(id)
+                )
+            """)
+            
+            # Training data view - optimized for LLM training
+            cursor.execute("""
+                CREATE VIEW IF NOT EXISTS training_data AS
+                SELECT 
+                    c.session_id,
+                    c.workspace_root,
+                    m.id as message_id,
+                    m.role,
+                    m.content,
+                    m.timestamp,
+                    m.model_name,
+                    GROUP_CONCAT(tc.tool_name || '(' || tc.arguments || ')') as tools_used,
+                    GROUP_CONCAT(fm.file_path) as files_modified,
+                    COUNT(DISTINCT tc.id) as tool_call_count
+                FROM conversations c
+                JOIN messages m ON c.id = m.conversation_id
+                LEFT JOIN tool_calls tc ON m.id = tc.message_id
+                LEFT JOIN files_modified fm ON tc.id = fm.tool_call_id
+                GROUP BY m.id
+                ORDER BY m.timestamp
+            """)
+            
+            # Create indexes for performance
+            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_conversation ON messages(conversation_id)")
+            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages(timestamp)")
+            cursor.execute("CREATE INDEX IF NOT EXISTS idx_tool_calls_conversation ON tool_calls(conversation_id)")
+            cursor.execute("CREATE INDEX IF NOT EXISTS idx_tool_calls_tool_name ON tool_calls(tool_name)")
+            
+            conn.commit()
+    
+    @contextmanager
+    def _get_connection(self):
+        """Context manager for database connections."""
+        conn = sqlite3.connect(self.db_path)
+        conn.row_factory = sqlite3.Row
+        try:
+            yield conn
+        finally:
+            conn.close()
+    
+    def start_conversation(self, session_id: str, workspace_root: str, 
+                          python_version: Optional[str] = None,
+                          venv_path: Optional[str] = None,
+                          metadata: Optional[Dict[str, Any]] = None) -> int:
+        """
+        Start a new conversation session.
+        
+        Args:
+            session_id: Unique session identifier
+            workspace_root: Path to the workspace
+            python_version: Python version being used
+            venv_path: Virtual environment path
+            metadata: Additional metadata as dict
+        
+        Returns:
+            conversation_id: Database ID for this conversation
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                INSERT INTO conversations (session_id, workspace_root, python_version, venv_path, metadata)
+                VALUES (?, ?, ?, ?, ?)
+            """, (session_id, workspace_root, python_version, venv_path, 
+                  json.dumps(metadata) if metadata else None))
+            conn.commit()
+            return cursor.lastrowid
+    
+    def end_conversation(self, conversation_id: int):
+        """Mark conversation as ended."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                UPDATE conversations 
+                SET ended_at = CURRENT_TIMESTAMP,
+                    total_messages = (SELECT COUNT(*) FROM messages WHERE conversation_id = ?),
+                    total_tool_calls = (SELECT COUNT(*) FROM tool_calls WHERE conversation_id = ?)
+                WHERE id = ?
+            """, (conversation_id, conversation_id, conversation_id))
+            conn.commit()
+    
+    def log_user_message(self, conversation_id: int, content: str, 
+                        metadata: Optional[Dict[str, Any]] = None) -> int:
+        """
+        Log a user message (prompt).
+        
+        Args:
+            conversation_id: ID of the conversation
+            content: User's message content
+            metadata: Additional metadata
+        
+        Returns:
+            message_id: Database ID for this message
+        """
+        return self._log_message(conversation_id, "user_prompt", "user", content, metadata=metadata)
+    
+    def log_assistant_message(self, conversation_id: int, content: str,
+                             model_name: Optional[str] = None,
+                             token_count: Optional[int] = None,
+                             metadata: Optional[Dict[str, Any]] = None) -> int:
+        """
+        Log an assistant message (LLM response).
+        
+        Args:
+            conversation_id: ID of the conversation
+            content: Assistant's response content
+            model_name: Name of the LLM model used
+            token_count: Number of tokens in response
+            metadata: Additional metadata
+        
+        Returns:
+            message_id: Database ID for this message
+        """
+        return self._log_message(
+            conversation_id, "assistant_response", "assistant", content,
+            model_name=model_name, token_count=token_count, metadata=metadata
+        )
+    
+    def _log_message(self, conversation_id: int, message_type: str, role: str,
+                     content: str, model_name: Optional[str] = None,
+                     token_count: Optional[int] = None,
+                     metadata: Optional[Dict[str, Any]] = None) -> int:
+        """Internal method to log a message."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                INSERT INTO messages (conversation_id, message_type, role, content, 
+                                     model_name, token_count, metadata)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+            """, (conversation_id, message_type, role, content, model_name, token_count,
+                  json.dumps(metadata) if metadata else None))
+            conn.commit()
+            return cursor.lastrowid
+    
+    def log_tool_call(self, conversation_id: int, tool_name: str, arguments: Dict[str, Any],
+                     result: Optional[Any] = None, success: bool = True,
+                     error_message: Optional[str] = None, execution_time_ms: Optional[float] = None,
+                     message_id: Optional[int] = None) -> int:
+        """
+        Log a tool call.
+        
+        Args:
+            conversation_id: ID of the conversation
+            tool_name: Name of the tool called
+            arguments: Tool arguments as dict
+            result: Tool execution result
+            success: Whether the tool call succeeded
+            error_message: Error message if failed
+            execution_time_ms: Execution time in milliseconds
+            message_id: Associated message ID (optional)
+        
+        Returns:
+            tool_call_id: Database ID for this tool call
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                INSERT INTO tool_calls (conversation_id, message_id, tool_name, arguments, 
+                                       result, success, error_message, execution_time_ms)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            """, (conversation_id, message_id, tool_name, json.dumps(arguments),
+                  json.dumps(result) if result else None, success, error_message, execution_time_ms))
+            conn.commit()
+            return cursor.lastrowid
+    
+    def log_file_modification(self, tool_call_id: int, file_path: str, operation: str,
+                             lines_added: int = 0, lines_removed: int = 0):
+        """
+        Log a file modification.
+        
+        Args:
+            tool_call_id: ID of the tool call that modified the file
+            file_path: Path to the modified file
+            operation: Type of operation (edit, create, delete, rename)
+            lines_added: Number of lines added
+            lines_removed: Number of lines removed
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                INSERT INTO files_modified (tool_call_id, file_path, operation, 
+                                           lines_added, lines_removed)
+                VALUES (?, ?, ?, ?, ?)
+            """, (tool_call_id, file_path, operation, lines_added, lines_removed))
+            conn.commit()
+    
+    def get_conversation(self, conversation_id: int) -> Optional[Dict[str, Any]]:
+        """Get conversation details."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT * FROM conversations WHERE id = ?", (conversation_id,))
+            row = cursor.fetchone()
+            return dict(row) if row else None
+    
+    def get_messages(self, conversation_id: int) -> List[Dict[str, Any]]:
+        """Get all messages in a conversation."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT * FROM messages 
+                WHERE conversation_id = ? 
+                ORDER BY timestamp
+            """, (conversation_id,))
+            return [dict(row) for row in cursor.fetchall()]
+    
+    def get_tool_calls(self, conversation_id: int) -> List[Dict[str, Any]]:
+        """Get all tool calls in a conversation."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT * FROM tool_calls 
+                WHERE conversation_id = ? 
+                ORDER BY timestamp
+            """, (conversation_id,))
+            return [dict(row) for row in cursor.fetchall()]
+    
+    def export_training_data(self, output_path: str, format: str = "jsonl"):
+        """
+        Export conversation data for LLM training.
+        
+        Args:
+            output_path: Path to output file
+            format: Export format ("jsonl" or "json")
+        """
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT * FROM training_data")
+            rows = cursor.fetchall()
+            
+            training_examples = []
+            for row in rows:
+                example = dict(row)
+                training_examples.append(example)
+            
+            output_file = Path(output_path)
+            if format == "jsonl":
+                with open(output_file, 'w', encoding='utf-8') as f:
+                    for example in training_examples:
+                        f.write(json.dumps(example) + '\n')
+            else:  # json
+                with open(output_file, 'w', encoding='utf-8') as f:
+                    json.dump(training_examples, f, indent=2)
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """Get overall statistics from the database."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            
+            stats = {}
+            
+            # Total conversations
+            cursor.execute("SELECT COUNT(*) as count FROM conversations")
+            stats['total_conversations'] = cursor.fetchone()['count']
+            
+            # Total messages
+            cursor.execute("SELECT COUNT(*) as count FROM messages")
+            stats['total_messages'] = cursor.fetchone()['count']
+            
+            # Total tool calls
+            cursor.execute("SELECT COUNT(*) as count FROM tool_calls")
+            stats['total_tool_calls'] = cursor.fetchone()['count']
+            
+            # Most used tools
+            cursor.execute("""
+                SELECT tool_name, COUNT(*) as count 
+                FROM tool_calls 
+                GROUP BY tool_name 
+                ORDER BY count DESC 
+                LIMIT 10
+            """)
+            stats['top_tools'] = [dict(row) for row in cursor.fetchall()]
+            
+            # Average tools per conversation
+            cursor.execute("""
+                SELECT AVG(total_tool_calls) as avg_tools 
+                FROM conversations 
+                WHERE total_tool_calls > 0
+            """)
+            stats['avg_tools_per_conversation'] = cursor.fetchone()['avg_tools']
+            
+            # Files modified count
+            cursor.execute("SELECT COUNT(DISTINCT file_path) as count FROM files_modified")
+            stats['unique_files_modified'] = cursor.fetchone()['count']
+            
+            return stats
+
+
+# Global logger instance
+_logger: Optional[ConversationLogger] = None
+
+
+def get_logger() -> ConversationLogger:
+    """Get or create the global conversation logger."""
+    global _logger
+    if _logger is None:
+        _logger = ConversationLogger()
+    return _logger
+
+
+def set_logger(logger: ConversationLogger):
+    """Set a custom logger instance."""
+    global _logger
+    _logger = logger