feat: refactor ai code

Author: dori

src/mcp_as_a_judge/constants.py

@@ -15,5 +15,4 @@
 DATABASE_URL = "sqlite://:memory:"
 MAX_SESSION_RECORDS = 20  # Maximum records to keep per session (FIFO)
 MAX_TOTAL_SESSIONS = 50  # Maximum total sessions to keep (LRU cleanup)
-MAX_CONTEXT_TOKENS = 50000  # Maximum tokens for conversation history context (1 token ≈ 4 characters)
-
+MAX_CONTEXT_TOKENS = 50000  # Maximum tokens for session token (1 token ≈ 4 characters)

src/mcp_as_a_judge/db/conversation_history_service.py

@@ -13,8 +13,8 @@
     create_database_provider,
 )
 from mcp_as_a_judge.db.db_config import Config
-from mcp_as_a_judge.logging_config import get_logger
 from mcp_as_a_judge.db.token_utils import filter_records_by_token_limit
+from mcp_as_a_judge.logging_config import get_logger
 
 # Set up logger
 logger = get_logger(__name__)
@@ -36,7 +36,9 @@ def __init__(
         self.config = config
         self.db = db_provider or create_database_provider(config)
 
-    async def load_context_for_enrichment(self, session_id: str, current_prompt: str = "") -> list[ConversationRecord]:
+    async def load_filtered_context_for_enrichment(
+        self, session_id: str, current_prompt: str = ""
+    ) -> list[ConversationRecord]:
         """
         Load recent conversation records for LLM context enrichment.
 
@@ -61,7 +63,9 @@ async def load_context_for_enrichment(self, session_id: str, current_prompt: str
 
         # Apply LLM context filtering: ensure history + current prompt will fit within token limit
         # This filters the list without modifying the database (only token limit matters for LLM)
-        filtered_records = filter_records_by_token_limit(recent_records, current_prompt=current_prompt)
+        filtered_records = filter_records_by_token_limit(
+            recent_records, current_prompt=current_prompt
+        )
 
         logger.info(
             f"✅ Returning {len(filtered_records)} conversation records for LLM context"
@@ -102,23 +106,9 @@ async def save_tool_interaction_and_cleanup(
         logger.info(f"✅ Saved conversation record with ID: {record_id}")
         return record_id
 
-    async def save_tool_interaction(
-        self, session_id: str, tool_name: str, tool_input: str, tool_output: str
-    ) -> str:
-        """
-        Save a tool interaction as a conversation record.
-
-        DEPRECATED: Use save_tool_interaction_and_cleanup() instead.
-        This method is kept for backward compatibility.
-        """
-        logger.warning(
-            "save_tool_interaction() is deprecated. Use save_tool_interaction_and_cleanup() instead."
-        )
-        return await self.save_tool_interaction_and_cleanup(
-            session_id, tool_name, tool_input, tool_output
-        )
-
-    def format_conversation_history_as_json_array( self, conversation_history: list[ConversationRecord]) -> list[dict]:
+    def format_conversation_history_as_json_array(
+        self, conversation_history: list[ConversationRecord]
+    ) -> list[dict]:
         """
         Convert conversation history list to JSON array for prompt injection.
 

src/mcp_as_a_judge/db/providers/sqlite_provider.py

@@ -9,13 +9,13 @@
 from datetime import UTC, datetime
 
 from sqlalchemy import create_engine
-from sqlmodel import Session, SQLModel, asc, desc, select
+from sqlmodel import Session, SQLModel, desc, select
 
 from mcp_as_a_judge.constants import MAX_CONTEXT_TOKENS
 from mcp_as_a_judge.db.cleanup_service import ConversationCleanupService
 from mcp_as_a_judge.db.interface import ConversationHistoryDB, ConversationRecord
-from mcp_as_a_judge.logging_config import get_logger
 from mcp_as_a_judge.db.token_utils import calculate_record_tokens
+from mcp_as_a_judge.logging_config import get_logger
 
 # Set up logger
 logger = get_logger(__name__)
@@ -107,9 +107,11 @@ def _cleanup_old_messages(self, session_id: str) -> int:
         """
         with Session(self.engine) as session:
             # Get current records ordered by timestamp DESC (newest first for token calculation)
-            count_stmt = select(ConversationRecord).where(
-                ConversationRecord.session_id == session_id
-            ).order_by(desc(ConversationRecord.timestamp))
+            count_stmt = (
+                select(ConversationRecord)
+                .where(ConversationRecord.session_id == session_id)
+                .order_by(desc(ConversationRecord.timestamp))
+            )
             current_records = session.exec(count_stmt).all()
             current_count = len(current_records)
 

src/mcp_as_a_judge/db/token_utils.py

@@ -7,8 +7,6 @@
 
 from mcp_as_a_judge.constants import MAX_CONTEXT_TOKENS
 
-from mcp_as_a_judge.db.interface import ConversationRecord
-
 
 def calculate_tokens(text: str) -> int:
     """
@@ -60,9 +58,7 @@ def calculate_total_tokens(records: list) -> int:
     return sum(record.tokens for record in records if hasattr(record, "tokens"))
 
 
-def filter_records_by_token_limit(
-    records: list, current_prompt: str = ""
-) -> list:
+def filter_records_by_token_limit(records: list, current_prompt: str = "") -> list:
     """
     Filter conversation records to stay within token and record limits.
 
@@ -81,7 +77,9 @@ def filter_records_by_token_limit(
         return []
 
     # Calculate current prompt tokens
-    current_prompt_tokens = calculate_record_tokens(current_prompt, "") if current_prompt else 0
+    current_prompt_tokens = (
+        calculate_record_tokens(current_prompt, "") if current_prompt else 0
+    )
 
     # Calculate total tokens including current prompt
     history_tokens = calculate_total_tokens(records)
@@ -96,7 +94,9 @@ def filter_records_by_token_limit(
     filtered_records = records.copy()
     current_history_tokens = history_tokens
 
-    while (current_history_tokens + current_prompt_tokens) > MAX_CONTEXT_TOKENS and len(filtered_records) > 1:
+    while (current_history_tokens + current_prompt_tokens) > MAX_CONTEXT_TOKENS and len(
+        filtered_records
+    ) > 1:
         # Remove the oldest record (last in the list)
         removed_record = filtered_records.pop()
         current_history_tokens -= getattr(removed_record, "tokens", 0)

src/mcp_as_a_judge/server.py

@@ -12,6 +12,7 @@
 from mcp.server.fastmcp import Context, FastMCP
 from pydantic import ValidationError
 
+from mcp_as_a_judge.constants import MAX_CONTEXT_TOKENS
 from mcp_as_a_judge.db.conversation_history_service import ConversationHistoryService
 from mcp_as_a_judge.db.db_config import load_config
 from mcp_as_a_judge.elicitation_provider import elicitation_provider
@@ -48,8 +49,6 @@
     tool_description_provider,
 )
 
-from src.mcp_as_a_judge.constants import MAX_CONTEXT_TOKENS
-
 # Initialize centralized logging
 setup_logging()
 
@@ -90,9 +89,16 @@ async def build_workflow(
 
     try:
         # STEP 1: Load conversation history and format as JSON array
-        conversation_history = await conversation_service.load_filtered_context_for_enrichment(session_id, json.dumps(original_input))
-        history_json_array = conversation_service.format_conversation_history_as_json_array(conversation_history)
-        
+        conversation_history = (
+            await conversation_service.load_filtered_context_for_enrichment(
+                session_id, json.dumps(original_input)
+            )
+        )
+        history_json_array = (
+            conversation_service.format_conversation_history_as_json_array(
+                conversation_history
+            )
+        )
 
         # STEP 2: Create system and user messages with separate context and conversation history
         system_vars = WorkflowGuidanceSystemVars(
@@ -135,7 +141,7 @@ async def build_workflow(
         log_error(e, "build_workflow")
         # Return a default workflow guidance in case of error
         return WorkflowGuidance(
-            next_tool="elicit_missing_requirements",
+            next_tool="raise_missing_requirements",
             reasoning="An error occurred during workflow generation. Please provide more details.",
             preparation_needed=[
                 "Review the error and provide more specific requirements"
@@ -558,9 +564,16 @@ async def judge_coding_plan(
 
     try:
         # STEP 1: Load conversation history and format as JSON array
-        conversation_history = await conversation_service.load_filtered_context_for_enrichment(session_id, json.dumps(original_input))
-        history_json_array = conversation_service.format_conversation_history_as_json_array(conversation_history)
-        
+        conversation_history = (
+            await conversation_service.load_filtered_context_for_enrichment(
+                session_id, json.dumps(original_input)
+            )
+        )
+        history_json_array = (
+            conversation_service.format_conversation_history_as_json_array(
+                conversation_history
+            )
+        )
 
         # STEP 2: Use helper function for main evaluation with JSON array conversation history
         evaluation_result = await _evaluate_coding_plan(
@@ -639,8 +652,16 @@ async def judge_code_change(
 
     try:
         # STEP 1: Load conversation history and format as JSON array
-        conversation_history = await conversation_service.load_filtered_context_for_enrichment(session_id,json.dumps(original_input))
-        history_json_array = conversation_service.format_conversation_history_as_json_array(conversation_history)
+        conversation_history = (
+            await conversation_service.load_filtered_context_for_enrichment(
+                session_id, json.dumps(original_input)
+            )
+        )
+        history_json_array = (
+            conversation_service.format_conversation_history_as_json_array(
+                conversation_history
+            )
+        )
 
         # STEP 2: Create system and user messages with separate context and conversation history
         system_vars = JudgeCodeChangeSystemVars(

tests/test_conversation_history_service_integration.py

@@ -53,7 +53,9 @@ async def test_service_save_and_retrieve_lifecycle(self, service):
         # PHASE 2: Retrieve conversation history
         print("\n📖 PHASE 2: Retrieving conversation history...")
 
-        conversation_history = await service.load_context_for_enrichment(session_id)
+        conversation_history = await service.load_filtered_context_for_enrichment(
+            session_id
+        )
         assert len(conversation_history) == 2, (
             f"Expected 2 records, got {len(conversation_history)}"
         )
@@ -101,7 +103,7 @@ async def test_service_save_and_retrieve_lifecycle(self, service):
             )
 
         # Should only get max_session_records (20) records
-        limited_history = await service.load_context_for_enrichment(session_id)
+        limited_history = await service.load_filtered_context_for_enrichment(session_id)
         expected_count = service.config.database.max_session_records
         assert len(limited_history) == expected_count, (
             f"Expected {expected_count} records, got {len(limited_history)}"
@@ -151,7 +153,7 @@ async def test_service_with_context_ids(self, service):
         )
 
         # Retrieve and verify
-        history = await service.load_context_for_enrichment(session_id)
+        history = await service.load_filtered_context_for_enrichment(session_id)
         assert len(history) == 3
 
         # Verify the conversation flow makes sense
@@ -168,7 +170,9 @@ async def test_service_empty_and_error_cases(self, service):
         print("=" * 60)
 
         # Test empty session
-        empty_history = await service.load_context_for_enrichment("nonexistent_session")
+        empty_history = await service.load_filtered_context_for_enrichment(
+            "nonexistent_session"
+        )
         assert len(empty_history) == 0
         print("✅ Empty session handled correctly")
 
@@ -180,14 +184,16 @@ async def test_service_empty_and_error_cases(self, service):
 
         # Test with special characters in data
         special_session = "special_chars_session"
-        await service.save_tool_interaction(
+        await service.save_tool_interaction_and_cleanup(
             session_id=special_session,
             tool_name="test_tool",
             tool_input="Input with 'quotes' and \"double quotes\" and \n newlines",
             tool_output="Result with émojis 🎉 and unicode ñ characters",
         )
 
-        special_history = await service.load_context_for_enrichment(special_session)
+        special_history = await service.load_filtered_context_for_enrichment(
+            special_session
+        )
         assert len(special_history) == 1
 
         special_json = service.format_conversation_history_as_json_array(
@@ -223,7 +229,7 @@ async def test_service_performance_with_large_dataset(self, service):
 
         # Retrieve records
         start_time = datetime.now()
-        history = await service.load_context_for_enrichment(session_id)
+        history = await service.load_filtered_context_for_enrichment(session_id)
         retrieve_time = datetime.now() - start_time
 
         print(