adding memory layer

Author: dori

config.json

@@ -0,0 +1,9 @@
+{
+  "database": {
+    "provider": "in_memory",
+    "url": "",
+    "max_context_records": 20,
+    "context_enrichment_count": 5
+  },
+  "enable_llm_fallback": true
+}

judge_mcp_flow.md

@@ -32,6 +32,7 @@ dependencies = [
     "pydantic>=2.0.0",
     "jinja2>=3.1.0",
     "litellm>=1.0.0",
+    "aiosqlite>=0.19.0",
 ]
 
 [project.urls]

pyproject.toml

@@ -0,0 +1,124 @@
+"""
+Configuration management for MCP as a Judge.
+
+This module handles loading and managing configuration from various sources
+including config files, environment variables, and defaults.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from .models import DatabaseConfig
+
+
+class Config(BaseModel):
+    """Main configuration model for the application."""
+    
+    database: DatabaseConfig = Field(default_factory=DatabaseConfig)
+    enable_llm_fallback: bool = Field(
+        default=True,
+        description="Whether to enable LLM fallback when MCP sampling is not available"
+    )
+
+
+def load_config(config_path: Optional[str] = None) -> Config:
+    """
+    Load configuration from file and environment variables.
+    
+    Args:
+        config_path: Path to config file. If None, looks for config.json in current directory
+        
+    Returns:
+        Config object with loaded settings
+    """
+    # Default configuration
+    config_data = {}
+    
+    # Try to load from config file
+    if config_path is None:
+        config_path = "config.json"
+    
+    config_file = Path(config_path)
+    if config_file.exists():
+        try:
+            with open(config_file, 'r') as f:
+                config_data = json.load(f)
+        except (json.JSONDecodeError, IOError) as e:
+            print(f"Warning: Could not load config file {config_path}: {e}")
+    
+    # Override with environment variables if present
+    db_provider = os.getenv("MCP_JUDGE_DB_PROVIDER")
+    if db_provider:
+        if "database" not in config_data:
+            config_data["database"] = {}
+        config_data["database"]["provider"] = db_provider
+    
+    db_url = os.getenv("MCP_JUDGE_DB_URL")
+    if db_url:
+        if "database" not in config_data:
+            config_data["database"] = {}
+        config_data["database"]["url"] = db_url
+    
+    max_context = os.getenv("MCP_JUDGE_MAX_CONTEXT_RECORDS")
+    if max_context:
+        if "database" not in config_data:
+            config_data["database"] = {}
+        try:
+            config_data["database"]["max_context_records"] = int(max_context)
+        except ValueError:
+            print(f"Warning: Invalid value for MCP_JUDGE_MAX_CONTEXT_RECORDS: {max_context}")
+    
+    llm_fallback = os.getenv("MCP_JUDGE_ENABLE_LLM_FALLBACK")
+    if llm_fallback:
+        config_data["enable_llm_fallback"] = llm_fallback.lower() in ("true", "1", "yes", "on")
+    
+    return Config(**config_data)
+
+
+def create_default_config_file(config_path: str = "config.json") -> None:
+    """
+    Create a default configuration file.
+    
+    Args:
+        config_path: Path where to create the config file
+    """
+    default_config = {
+        "database": {
+            "provider": "in_memory",
+            "url": "",
+            "max_context_records": 10
+        },
+        "enable_llm_fallback": True
+    }
+    
+    with open(config_path, 'w') as f:
+        json.dump(default_config, f, indent=2)
+    
+    print(f"Created default configuration file: {config_path}")
+
+
+def get_database_provider_from_url(url: str) -> str:
+    """
+    Determine database provider from URL.
+    
+    Args:
+        url: Database connection URL
+        
+    Returns:
+        Provider name: 'sqlite', 'postgresql', or 'in_memory'
+    """
+    if not url:
+        return "in_memory"
+    
+    url_lower = url.lower()
+    if url_lower.startswith("sqlite://") or url_lower.endswith(".db"):
+        return "sqlite"
+    elif url_lower.startswith("postgresql://") or url_lower.startswith("postgres://"):
+        return "postgresql"
+    else:
+        # Default to in_memory for unknown URLs
+        return "in_memory"

src/mcp_as_a_judge/config.py

@@ -0,0 +1,197 @@
+"""
+Conversation History Service for MCP Judge Tools.
+
+This service handles:
+1. Loading historical context for LLM enrichment
+2. Saving tool interactions as conversation records
+3. Managing session-based conversation history
+"""
+
+from typing import List, Optional
+
+from .config import Config
+from .db import ConversationHistoryDB, ConversationRecord, create_database_provider
+
+
+class ConversationHistoryService:
+    """Service for managing conversation history in judge tools."""
+    
+    def __init__(self, config: Config, db_provider: Optional[ConversationHistoryDB] = None):
+        """
+        Initialize the conversation history service.
+        
+        Args:
+            config: Application configuration
+            db_provider: Optional database provider (will create one if not provided)
+        """
+        self.config = config
+        self.db = db_provider or create_database_provider(config)
+    
+    async def load_context_for_enrichment(self, session_id: str) -> tuple[List[ConversationRecord], List[str]]:
+        """
+        Load recent conversation records for LLM context enrichment.
+        
+        Args:
+            session_id: Session identifier
+            
+        Returns:
+            Tuple of (conversation_records, conversation_ids)
+            - conversation_records: Full records for LLM context
+            - conversation_ids: Just the IDs for saving in new record's context field
+        """
+        count = self.config.database.context_enrichment_count
+        
+        # Load recent conversations for this session
+        recent_records = await self.db.get_session_conversations(
+            session_id=session_id,
+            limit=count
+        )
+        
+        # Extract just the IDs for context reference
+        context_ids = [record.id for record in recent_records]
+        
+        return recent_records, context_ids
+    
+    async def save_tool_interaction(
+        self,
+        session_id: str,
+        tool_name: str,
+        tool_input: str,
+        tool_output: str,
+        context_ids: List[str]
+    ) -> str:
+        """
+        Save a tool interaction as a conversation record.
+        
+        Args:
+            session_id: Session identifier from AI agent
+            tool_name: Name of the judge tool (e.g., 'judge_coding_plan')
+            tool_input: Input that was passed to the tool
+            tool_output: Output/result from the tool
+            context_ids: IDs of conversation records that were used for context enrichment
+            
+        Returns:
+            ID of the created conversation record
+        """
+        record_id = await self.db.save_conversation(
+            session_id=session_id,
+            source=tool_name,
+            input_data=tool_input,
+            context=context_ids,
+            output=tool_output
+        )
+        
+        return record_id
+    
+    def format_context_for_llm(self, context_records: List[ConversationRecord]) -> str:
+        """
+        Format conversation history for LLM context enrichment.
+        
+        Args:
+            context_records: Recent conversation records
+            
+        Returns:
+            Formatted context string for LLM
+        """
+        if not context_records:
+            return "No previous conversation history available."
+        
+        context_lines = ["## Previous Conversation History"]
+        context_lines.append("Here are the recent interactions in this session for context:")
+        context_lines.append("")
+        
+        # Format records (most recent first)
+        for i, record in enumerate(context_records, 1):
+            context_lines.append(f"### {i}. {record.source} ({record.timestamp.strftime('%Y-%m-%d %H:%M:%S')})")
+            context_lines.append(f"**Input:** {record.input}")
+            context_lines.append(f"**Output:** {record.output}")
+            context_lines.append("")
+        
+        context_lines.append("---")
+        context_lines.append("Use this context to make more informed decisions.")
+        context_lines.append("")
+        
+        return "\n".join(context_lines)
+    
+    async def get_session_summary(self, session_id: str) -> dict:
+        """
+        Get a summary of the session's conversation history.
+        
+        Args:
+            session_id: Session identifier
+            
+        Returns:
+            Dictionary with session statistics
+        """
+        all_records = await self.db.get_session_conversations(session_id)
+        
+        # Count by tool type
+        tool_counts = {}
+        for record in all_records:
+            tool_counts[record.source] = tool_counts.get(record.source, 0) + 1
+        
+        return {
+            "session_id": session_id,
+            "total_interactions": len(all_records),
+            "tool_usage": tool_counts,
+            "latest_interaction": all_records[0].timestamp.isoformat() if all_records else None,
+            "context_enrichment_count": self.config.database.context_enrichment_count,
+            "max_context_records": self.config.database.max_context_records
+        }
+
+
+# Convenience functions for easy integration with existing tools
+
+async def enrich_with_context(
+    service: ConversationHistoryService,
+    session_id: str,
+    base_prompt: str
+) -> tuple[str, List[str]]:
+    """
+    Enrich a base prompt with conversation history context.
+    
+    Args:
+        service: ConversationHistoryService instance
+        session_id: Session identifier
+        base_prompt: Original prompt to enrich
+        
+    Returns:
+        Tuple of (enriched_prompt, context_ids)
+    """
+    context_records, context_ids = await service.load_context_for_enrichment(session_id)
+    context_text = service.format_context_for_llm(context_records)
+    
+    enriched_prompt = f"{context_text}\n## Current Request\n{base_prompt}"
+    
+    return enriched_prompt, context_ids
+
+
+async def save_tool_result(
+    service: ConversationHistoryService,
+    session_id: str,
+    tool_name: str,
+    original_input: str,
+    tool_result: str,
+    context_ids: List[str]
+) -> str:
+    """
+    Save a tool's result to conversation history.
+    
+    Args:
+        service: ConversationHistoryService instance
+        session_id: Session identifier
+        tool_name: Name of the tool
+        original_input: Original input to the tool
+        tool_result: Result from the tool
+        context_ids: Context IDs that were used for enrichment
+        
+    Returns:
+        ID of the saved conversation record
+    """
+    return await service.save_tool_interaction(
+        session_id=session_id,
+        tool_name=tool_name,
+        tool_input=original_input,
+        tool_output=tool_result,
+        context_ids=context_ids
+    )

src/mcp_as_a_judge/conversation_history_service.py

@@ -0,0 +1,18 @@
+"""
+Database abstraction layer for MCP as a Judge.
+
+This module provides database interfaces and providers for storing
+conversation history and tool interactions.
+"""
+
+from .factory import DatabaseFactory, create_database_provider
+from .interface import ConversationHistoryDB, ConversationRecord
+from .providers import InMemoryProvider
+
+__all__ = [
+    "ConversationHistoryDB",
+    "ConversationRecord",
+    "InMemoryProvider",
+    "DatabaseFactory",
+    "create_database_provider"
+]