fix security issue and reformat

Author: dori

src/mcp_as_a_judge/config.py

@@ -20,7 +20,7 @@ class Config(BaseModel):
     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"
+        description="Whether to enable LLM fallback when MCP sampling is not available",
     )
 
 
@@ -69,11 +69,18 @@ def load_config(config_path: str | None = None) -> Config:
         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}")
+            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")
+        config_data["enable_llm_fallback"] = llm_fallback.lower() in (
+            "true",
+            "1",
+            "yes",
+            "on",
+        )
 
     return Config(**config_data)
 
@@ -86,15 +93,11 @@ def create_default_config_file(config_path: str = "config.json") -> None:
         config_path: Path where to create the config file
     """
     default_config = {
-        "database": {
-            "provider": "in_memory",
-            "url": "",
-            "max_context_records": 10
-        },
-        "enable_llm_fallback": True
+        "database": {"provider": "in_memory", "url": "", "max_context_records": 10},
+        "enable_llm_fallback": True,
     }
 
-    with open(config_path, 'w') as f:
+    with open(config_path, "w") as f:
         json.dump(default_config, f, indent=2)
 
     print(f"Created default configuration file: {config_path}")

src/mcp_as_a_judge/db/__init__.py

@@ -14,5 +14,5 @@
     "ConversationRecord",
     "DatabaseFactory",
     "SQLiteProvider",
-    "create_database_provider"
+    "create_database_provider",
 ]

src/mcp_as_a_judge/db/conversation_history_service.py

@@ -19,7 +19,9 @@
 class ConversationHistoryService:
     """Service for managing conversation history in judge tools."""
 
-    def __init__(self, config: Config, db_provider: ConversationHistoryDB | None = None):
+    def __init__(
+        self, config: Config, db_provider: ConversationHistoryDB | None = None
+    ):
         """
         Initialize the conversation history service.
 
@@ -30,7 +32,9 @@ def __init__(self, config: Config, db_provider: ConversationHistoryDB | None = N
         self.config = config
         self.db = db_provider or create_database_provider(config)
 
-    async def load_context_for_enrichment(self, session_id: str) -> list[ConversationRecord]:
+    async def load_context_for_enrichment(
+        self, session_id: str
+    ) -> list[ConversationRecord]:
         """
         Load recent conversation records for LLM context enrichment.
 
@@ -45,18 +49,14 @@ async def load_context_for_enrichment(self, session_id: str) -> list[Conversatio
         # Load recent conversations for this session
         recent_records = await self.db.get_session_conversations(
             session_id=session_id,
-            limit=self.config.database.context_enrichment_count # load last X records
+            limit=self.config.database.context_enrichment_count,  # load last X records
         )
 
         logger.info(f"📚 Retrieved {len(recent_records)} conversation records from DB")
         return recent_records
 
     async def save_tool_interaction(
-        self,
-        session_id: str,
-        tool_name: str,
-        tool_input: str,
-        tool_output: str
+        self, session_id: str, tool_name: str, tool_input: str, tool_output: str
     ) -> str:
         """
         Save a tool interaction as a conversation record.
@@ -70,13 +70,15 @@ async def save_tool_interaction(
         Returns:
             ID of the created conversation record
         """
-        logger.info(f"💾 Saving tool interaction to SQLite DB for session: {session_id}, tool: {tool_name}")
+        logger.info(
+            f"💾 Saving tool interaction to SQLite DB for session: {session_id}, tool: {tool_name}"
+        )
 
         record_id = await self.db.save_conversation(
             session_id=session_id,
             source=tool_name,
             input_data=tool_input,
-            output=tool_output
+            output=tool_output,
         )
 
         logger.info(f"✅ Saved conversation record with ID: {record_id}")
@@ -96,19 +98,27 @@ def format_context_for_llm(self, context_records: list[ConversationRecord]) -> s
             logger.info("📝 No conversation history to format for LLM context")
             return "No previous conversation history available."
 
-        logger.info(f"📝 Formatting {len(context_records)} conversation records for LLM context enrichment")
+        logger.info(
+            f"📝 Formatting {len(context_records)} conversation records for LLM context enrichment"
+        )
 
         context_lines = ["## Previous Conversation History"]
-        context_lines.append("Here are the recent interactions in this session for context:")
+        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"### {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("")
-            logger.info(f"   Formatted record {i}: {record.source} from {record.timestamp}")
+            logger.info(
+                f"   Formatted record {i}: {record.source} from {record.timestamp}"
+            )
 
         context_lines.append("---")
         context_lines.append("Use this context to make more informed decisions.")
@@ -141,18 +151,19 @@ async def get_session_summary(self, session_id: str) -> dict:
             "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,
+            "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
+            "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
+    service: ConversationHistoryService, session_id: str, base_prompt: str
 ) -> str:
     """
     Enrich a base prompt with conversation history context.
@@ -165,16 +176,17 @@ async def enrich_with_context(
     Returns:
         Enriched prompt with conversation history context
     """
-    logger.info(f"🔄 Starting context enrichment for session {session_id}, base_prompt: {base_prompt}")
+    logger.info(
+        f"🔄 Starting context enrichment for session {session_id}, base_prompt: {base_prompt}"
+    )
 
     context_records = 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}"
 
-    logger.info(f"🎯 Context enrichment completed for session {session_id}, enriched_prompt: {enriched_prompt}")
+    logger.info(
+        f"🎯 Context enrichment completed for session {session_id}, enriched_prompt: {enriched_prompt}"
+    )
 
     return enriched_prompt
-
-
-

src/mcp_as_a_judge/db/factory.py

@@ -17,7 +17,7 @@ class DatabaseFactory:
 
     _providers: ClassVar[dict[str, type[ConversationHistoryDB]]] = {
         "in_memory": SQLiteProvider,  # SQLite in-memory (:memory: or empty URL)
-        "sqlite": SQLiteProvider,     # SQLite file-based storage
+        "sqlite": SQLiteProvider,  # SQLite file-based storage
         # Future providers can be added here:
         # "postgresql": PostgreSQLProvider,
         # "mysql": MySQLProvider,
@@ -56,14 +56,14 @@ def create_provider(cls, config: Config) -> ConversationHistoryDB:
             return provider_class(
                 max_context_records=config.database.max_context_records,
                 retention_days=config.database.record_retention_days,
-                url=config.database.url
+                url=config.database.url,
             )
         else:
             # For future network database providers (PostgreSQL, MySQL, etc.)
             return provider_class(
                 url=config.database.url,
                 max_context_records=config.database.max_context_records,
-                retention_days=config.database.record_retention_days
+                retention_days=config.database.record_retention_days,
             )
 
     @classmethod
@@ -73,7 +73,9 @@ def get_available_providers(cls) -> list[str]:
 
     # Not in use - option to register additional providers
     @classmethod
-    def register_provider(cls, name: str, provider_class: type[ConversationHistoryDB]) -> None:
+    def register_provider(
+        cls, name: str, provider_class: type[ConversationHistoryDB]
+    ) -> None:
         """
         Register a new database provider.
 
@@ -84,7 +86,6 @@ def register_provider(cls, name: str, provider_class: type[ConversationHistoryDB
         cls._providers[name] = provider_class
 
 
-
 # Convenience function
 def create_database_provider(config: Config) -> ConversationHistoryDB:
     """

src/mcp_as_a_judge/db/interface.py

@@ -19,14 +19,18 @@ class ConversationRecord(BaseModel):
     source: str  # tool name
     input: str  # tool input query
     output: str  # tool output string
-    timestamp: datetime = Field(default_factory=datetime.utcnow)  # when the record was created
+    timestamp: datetime = Field(
+        default_factory=datetime.utcnow
+    )  # when the record was created
 
 
 class ConversationHistoryDB(ABC):
     """Abstract interface for conversation history database operations."""
 
     @abstractmethod
-    def __init__(self, max_context_records: int = 20, retention_days: int = 1, url: str = "") -> None:
+    def __init__(
+        self, max_context_records: int = 20, retention_days: int = 1, url: str = ""
+    ) -> None:
         """
         Initialize the database provider.
 
@@ -39,11 +43,7 @@ def __init__(self, max_context_records: int = 20, retention_days: int = 1, url:
 
     @abstractmethod
     async def save_conversation(
-        self,
-        session_id: str,
-        source: str,
-        input_data: str,
-        output: str
+        self, session_id: str, source: str, input_data: str, output: str
     ) -> str:
         """
         Save a conversation record to the database.
@@ -61,9 +61,7 @@ async def save_conversation(
 
     @abstractmethod
     async def get_session_conversations(
-        self,
-        session_id: str,
-        limit: int | None = None
+        self, session_id: str, limit: int | None = None
     ) -> list[ConversationRecord]:
         """
         Retrieve all conversation records for a session.