DB integrated + Working on the README.md

.idea/dataSources.xml

@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="DataSourceManagerImpl" format="xml" multifile-model="true">
+    <data-source source="LOCAL" name="drift_detector" uuid="88b3f45d-5575-4285-8690-21bff284a779">
+      <driver-ref>sqlite.xerial</driver-ref>
+      <synchronize>true</synchronize>
+      <jdbc-driver>org.sqlite.JDBC</jdbc-driver>
+      <jdbc-url>jdbc:sqlite:$PROJECT_DIR$/drift_detector.sqlite3</jdbc-url>
+      <working-dir>$ProjectFileDir$</working-dir>
+      <libraries>
+        <library>
+          <url>file://$APPLICATION_CONFIG_DIR$/jdbc-drivers/Xerial SQLiteJDBC/3.45.1/org/xerial/sqlite-jdbc/3.45.1.0/sqlite-jdbc-3.45.1.0.jar</url>
+        </library>
+        <library>
+          <url>file://$APPLICATION_CONFIG_DIR$/jdbc-drivers/Xerial SQLiteJDBC/3.45.1/org/slf4j/slf4j-api/1.7.36/slf4j-api-1.7.36.jar</url>
+        </library>
+      </libraries>
+    </data-source>
+  </component>
+</project>

README.md

@@ -11,8 +11,14 @@ license: mit
 ---
 
 # Drift Detector
-Drift Detector is an MCP server, designed to detect drift in LLM performance over time. 
-This implementation is intended as a proof of concept and is not intended for production use.
+Drift Detector is an MCP server, designed to detect drift in LLM performance over time by using the power of the **sampling** functionality of MCP. 
+This implementation is intended as a **proof of concept** and is **NOT intended** for production use without significant changes.
+
+## The Idea
+
+The drift detector is a server that can be connected to any LLM client that supports the MCP sampling functionality. 
+It allows you to monitor the performance of your LLM models over time, detecting any drift in their behavior.
+This is particularly useful for applications where the model's performance may change due to various factors, such as changes in the data distribution, model updates, or other external influences.
 
 ## How to run
 
@@ -64,5 +70,48 @@ The gradio interface in [app.py](app.py) is an example dashboard which allows us
 
 ### Drift Detector Server
 
-The Drift Detector server is implemented using the MCP python SDK
-    
+The Drift Detector server is implemented using the MCP python SDK.
+It exposes the following tools:
+
+1. **run_initial_diagnostics**
+   - **Purpose**: Establishes a baseline for model behavior using adaptive sampling techniques
+   - **Parameters**:
+     - `model`: The name of the model to run diagnostics on
+     - `model_capabilities`: Full description of the model's capabilities and special features
+   - **Sampling Process**:
+     - First generates a tailored questionnaire based on model-specific capabilities
+     - Collects responses by sampling the target model with controlled parameters (temperature=0.7)
+     - Each question is processed individually to ensure proper context isolation
+     - Baseline samples are stored as paired question-answer JSON records for future comparison
+   - **Output**: Confirmation message indicating successful baseline creation
+
+2. **check_drift**
+   - **Purpose**: Measures potential drift by comparative sampling against the baseline
+   - **Parameters**:
+     - `model`: The name of the model to check for drift
+   - **Sampling Process**:
+     - Retrieves the original questions from the baseline
+     - Re-samples the model with identical questions using the same sampling parameters
+     - Maintains consistent context conditions to ensure fair comparison
+     - Uses differential analysis to compare semantic and functional differences between sample sets
+   - **Drift Evaluation**:
+     - Calculates a numerical drift score based on answer divergence
+     - Provides threshold-based alerts when drift exceeds acceptable limits (score > 50)
+     - Stores the latest sample responses for audit and trend analysis
+
+## Flow
+
+The intended flow is as follows:
+1. When the client contacts the server for the first time, it will run the `run_initial_diagnostics` tool.
+2. The server will generate a tailored questionnaire based on the model's capabilities.
+3. This questionnaire will be used to collect responses from the model, establishing a baseline for future comparisons.
+4. Once the baseline is established, the server will store the paired question-answer JSON records.
+5. The client can then use the `check_drift` tool to measure potential drift in the model's performance.
+6. The server will retrieve the original questions from the baseline and re-sample the model with identical questions.
+7. The server will maintain consistent context conditions to ensure fair comparison.
+8. If significant drift is detected (score > 50), the server will provide an alert and store the latest sample responses for audit and trend analysis.
+9. The client can visualize the drift data through the Gradio interface, allowing users to track changes in model performance over time.
+
+
+
+

database_module/__init__.py

@@ -1,3 +1,10 @@
 #database_module/__init__.py
 from .db import init_db
-from .mcp_tools import get_all_models_handler, search_models_handler
+from .mcp_tools import (
+    get_all_models_handler,
+    search_models_handler,
+    save_diagnostic_data,
+    get_baseline_diagnostics,
+    save_drift_score,
+    register_model_with_capabilities
+)

database_module/db.py

@@ -1,6 +1,6 @@
 #database_module/db.py
 import os
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, inspect, text
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
 
@@ -26,9 +26,25 @@ SessionLocal = sessionmaker(
 )
 Base = declarative_base()
 
+def apply_migrations():
+    """
+    Apply any necessary migrations to existing tables.
+    """
+    with engine.connect() as conn:
+        # Check if the models table exists and has the capabilities column
+        inspector = inspect(engine)
+        if "models" in inspector.get_table_names():
+            columns = [col['name'] for col in inspector.get_columns('models')]
+            if "capabilities" not in columns:
+                # Add capabilities column to models table
+                conn.execute(text("ALTER TABLE models ADD COLUMN capabilities TEXT"))
+                conn.commit()
+                print("Migration: Added capabilities column to models table")
+
 def init_db():
     """
     Create tables if they don't exist.
     Call this once at application startup.
     """
     Base.metadata.create_all(bind=engine)
+    apply_migrations()

database_module/mcp_tools.py

@@ -2,9 +2,10 @@
 from sqlalchemy.orm import Session
 from sqlalchemy import or_
 from .db import SessionLocal
-from .models import ModelEntry, DriftEntry
+from .models import ModelEntry, DriftEntry, DiagnosticData
 from datetime import datetime
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
+import json
 
 
 def get_all_models_handler(_: Dict[str, Any]) -> List[Dict[str, Any]]:
@@ -111,3 +112,98 @@ def get_drift_history_handler(params: Dict[str, Any]) -> List[Dict[str, Any]]:
             {"date": e.date.isoformat(), "drift_score": e.drift_score}
             for e in entries
         ]
+
+
+# === New functions for drift detection database operations ===
+
+def save_diagnostic_data(
+    model_name: str,
+    questions: list,
+    answers: list,
+    is_baseline: bool = False
+) -> None:
+    """
+    Save diagnostic questions and answers to the database
+    """
+    with SessionLocal() as session:
+        # Check if model exists, create if not
+        model = session.query(ModelEntry).filter_by(name=model_name).first()
+        if not model:
+            model = ModelEntry(
+                name=model_name,
+                created=datetime.utcnow().date(),
+                description=""
+            )
+            session.add(model)
+
+        # Create new diagnostic entry
+        diagnostic = DiagnosticData(
+            model_name=model_name,
+            is_baseline=1 if is_baseline else 0,
+            questions=questions,
+            answers=answers,
+            created=datetime.utcnow()
+        )
+        session.add(diagnostic)
+        session.commit()
+
+
+def get_baseline_diagnostics(model_name: str) -> Optional[Dict[str, Any]]:
+    """
+    Retrieve baseline diagnostics for a model
+    """
+    with SessionLocal() as session:
+        baseline = session.query(DiagnosticData)\
+            .filter_by(model_name=model_name, is_baseline=1)\
+            .order_by(DiagnosticData.created.desc())\
+            .first()
+
+        if not baseline:
+            return None
+
+        return {
+            "questions": baseline.questions,
+            "answers": baseline.answers,
+            "created": baseline.created.isoformat()
+        }
+
+
+def save_drift_score(model_name: str, drift_score: str) -> None:
+    """
+    Save drift score to database
+    """
+    # Try to convert score to float if possible
+    try:
+        score_float = float(drift_score)
+    except ValueError:
+        score_float = None
+
+    with SessionLocal() as session:
+        entry = DriftEntry(
+            model_name=model_name,
+            date=datetime.utcnow(),
+            drift_score=score_float
+        )
+        session.add(entry)
+        session.commit()
+
+
+def register_model_with_capabilities(model_name: str, capabilities: str) -> None:
+    """
+    Register a model with capabilities or update if already exists
+    """
+    with SessionLocal() as session:
+        model = session.query(ModelEntry).filter_by(name=model_name).first()
+
+        if model:
+            model.capabilities = capabilities
+        else:
+            model = ModelEntry(
+                name=model_name,
+                created=datetime.utcnow().date(),
+                capabilities=capabilities,
+                description=""
+            )
+            session.add(model)
+
+        session.commit()

database_module/models.py

@@ -1,5 +1,6 @@
 # database_module/models.py
-from sqlalchemy import Column, String, Date, Integer, Float, Text
+from sqlalchemy import Column, String, Date, Integer, Float, Text, JSON, DateTime
+from datetime import datetime
 from .db import Base
 
 class ModelEntry(Base):
@@ -9,11 +10,22 @@ class ModelEntry(Base):
     name = Column(String, unique=True, nullable=False, index=True)
     created = Column(Date, nullable=False)
     description = Column(Text, nullable=True)
+    capabilities = Column(Text, nullable=True)  # Added to store model_capabilities
 
 class DriftEntry(Base):
     __tablename__ = "drift_history"
 
     id = Column(Integer, primary_key=True, index=True)
     model_name = Column(String, nullable=False, index=True)
-    date = Column(Date, nullable=False)
-    drift_score = Column(Float, nullable=False)
+    date = Column(DateTime, nullable=False, default=datetime.utcnow)
+    drift_score = Column(Float, nullable=True)
+
+class DiagnosticData(Base):
+    __tablename__ = "diagnostic_data"
+
+    id = Column(Integer, primary_key=True, index=True)
+    model_name = Column(String, nullable=False, index=True)
+    created = Column(DateTime, nullable=False, default=datetime.utcnow)
+    is_baseline = Column(Integer, nullable=False, default=0)  # 0=latest, 1=baseline
+    questions = Column(JSON, nullable=True)
+    answers = Column(JSON, nullable=True)

requirements.txt

@@ -4,5 +4,5 @@ plotly
 asyncio
 typing
 sqlalchemy
-psycopg2
+psycopg2-binary
 fast-agent-mcp

server.py

@@ -10,7 +10,14 @@ from mcp.server.stdio import stdio_server
 
 from ourllm import genratequestionnaire, gradeanswers
 from database_module import init_db
-from database_module import get_all_models_handler, search_models_handler
+from database_module import (
+    get_all_models_handler,
+    search_models_handler,
+    save_diagnostic_data,
+    get_baseline_diagnostics,
+    save_drift_score,
+    register_model_with_capabilities
+)
 
 # Initialize data directory and database
 DATA_DIR = "data"
@@ -19,25 +26,6 @@ init_db()
 
 app = Server("mcp-drift-server")
 
-
-# === Sampling Helper ===
-async def sample(messages: List[types.SamplingMessage], max_tokens: int = 300) -> CreateMessageResult:
-    return await app.request_context.session.create_message(
-        messages=messages,
-        max_tokens=max_tokens,
-        temperature=0.7,
-    )
-
-
-# === Baseline File Helpers ===
-def get_baseline_path(model_name: str) -> str:
-    return os.path.join(DATA_DIR, f"{model_name}_baseline.json")
-
-
-def get_response_path(model_name: str) -> str:
-    return os.path.join(DATA_DIR, f"{model_name}_latest.json")
-
-
 # === Tool Manifest ===
 @app.list_tools()
 async def list_tools() -> List[types.Tool]:
@@ -79,8 +67,6 @@ async def list_tools() -> List[types.Tool]:
         ),
     ]
 
-
-
 # === Sampling Wrapper ===
 async def sample(messages: list[types.SamplingMessage], max_tokens=600) -> CreateMessageResult:
     return await app.request_context.session.create_message(
@@ -89,101 +75,70 @@ async def sample(messages: list[types.SamplingMessage], max_tokens=600) -> Creat
         temperature=0.7
     )
 
-
-# === Baseline File Paths ===
-def get_baseline_path(model_name):
-    return os.path.join(DATA_DIR, f"{model_name}_baseline.json")
-
-
-def get_response_path(model_name):
-    return os.path.join(DATA_DIR, f"{model_name}_latest.json")
-
-
-
 # === Core Logic ===
 async def run_initial_diagnostics(arguments: Dict[str, Any]) -> List[types.TextContent]:
     model = arguments["model"]
-    caps  = arguments["model_capabilities"]
-
-    # 1. Generate questionnaire
-    questions = await genratequestionnaire(model, caps)
-
-    # 2. Ask the target LLM (client)
-    answers = await sample(questions)
-
-
-    # 3. Persist baseline
+    capabilities = arguments["model_capabilities"]
 
     # 1. Ask the server's internal LLM to generate a questionnaire
-
-    questions = genratequestionnaire(model, arguments["model_capabilities"])  # Server-side trusted LLM
+    questions = genratequestionnaire(model, capabilities)  # Server-side trusted LLM
     answers = []
     for q in questions:
         a = await sample([q])
         answers.append(a)
 
-    # 3. Save Q/A pair
-
-    with open(get_baseline_path(model), "w") as f:
-        json.dump({
-            "questions": [m.content.text for m in questions],
-            "answers":   [m.content.text for m in answers]
-        }, f, indent=2)
+    # 2. Save the model capabilities and questions/answers to database
+    register_model_with_capabilities(model, capabilities)
+    save_diagnostic_data(
+        model_name=model,
+        questions=[m.content.text for m in questions],
+        answers=[m.content.text for m in answers],
+        is_baseline=True
+    )
 
     return [types.TextContent(type="text", text=f"✅ Baseline stored for model: {model}")]
 
-
 async def check_drift(arguments: Dict[str, Any]) -> List[types.TextContent]:
-    model     = arguments["model"]
-    base_path = get_baseline_path(model)
+    model = arguments["model"]
+
+    # Get baseline from database
+    baseline = get_baseline_diagnostics(model)
 
     # Ensure baseline exists
-    if not os.path.exists(base_path):
+    if not baseline:
         return [types.TextContent(type="text", text=f"❌ No baseline for model: {model}")]
 
-    # Load questions + old answers
-    with open(base_path) as f:
-        data = json.load(f)
-    questions   = [
+    # Convert questions to sampling messages
+    questions = [
         types.SamplingMessage(role="user", content=types.TextContent(type="text", text=q))
-        for q in data["questions"]
+        for q in baseline["questions"]
     ]
-    old_answers = data["answers"]
+    old_answers = baseline["answers"]
 
-
-    # 1. Get fresh answers
-    new_msgs    = await sample(questions)
-    new_answers = [m.content.text for m in new_msgs]
-
-    # 1. Ask the model again
+    # Ask the model again
     new_answers_msgs = []
     for q in questions:
         a = await sample([q])
         new_answers_msgs.append(a)
     new_answers = [m.content.text for m in new_answers_msgs]
 
-
-    # 2. Grade for drift
-    grading     = await gradeanswers(old_answers, new_answers)
-    drift_score = grading[0].content.text.strip()
-
-
-    # 3. Save latest
+    # Grade the answers and get a drift score
     grading_response = gradeanswers(old_answers, new_answers)
     drift_score = grading_response[0].content.text.strip()
 
-    # 3. Save the response
-
-    with open(get_response_path(model), "w") as f:
-        json.dump({
-            "new_answers": new_answers,
-            "drift_score": drift_score
-        }, f, indent=2)
+    # Save the latest responses and drift score to database
+    save_diagnostic_data(
+        model_name=model,
+        questions=baseline["questions"],
+        answers=new_answers,
+        is_baseline=False
+    )
+    save_drift_score(model, drift_score)
 
-    # 4. Alert threshold
+    # Alert threshold
     try:
         score_val = float(drift_score)
-        alert     = "🚨 Significant drift!" if score_val > 50 else "✅ Drift OK"
+        alert = "🚨 Significant drift!" if score_val > 50 else "✅ Drift OK"
     except ValueError:
         alert = "⚠️ Drift score not numeric"
 
@@ -192,26 +147,52 @@ async def check_drift(arguments: Dict[str, Any]) -> List[types.TextContent]:
         types.TextContent(type="text", text=alert)
     ]
 
+# Database tool handlers
+async def get_all_models_handler_async(_: Dict[str, Any]) -> List[types.TextContent]:
+    models = get_all_models_handler({})
+    if not models:
+        return [types.TextContent(type="text", text="No models registered.")]
+
+    model_list = "\n".join([f"• {m['name']} - {m['description']}" for m in models])
+    return [types.TextContent(
+        type="text",
+        text=f"Registered models:\n{model_list}"
+    )]
+
+async def search_models_handler_async(arguments: Dict[str, Any]) -> List[types.TextContent]:
+    query = arguments.get("query", "")
+    models = search_models_handler({"search_term": query})
+
+    if not models:
+        return [types.TextContent(
+            type="text",
+            text=f"No models found matching '{query}'."
+        )]
+
+    model_list = "\n".join([f"• {m['name']} - {m['description']}" for m in models])
+    return [types.TextContent(
+        type="text",
+        text=f"Models matching '{query}':\n{model_list}"
+    )]
 
 # === Dispatcher ===
 @app.call_tool()
 async def dispatch_tool(name: str, arguments: Dict[str, Any] | None = None):
     if name == "run_initial_diagnostics":
-        return await run_initial_diagnostics(arguments or {})
-    if name == "check_drift":
-        return await check_drift(arguments or {})
-    if name == "get_all_models":
-        return await get_all_models_handler()
-    if name == "search_models":
-        return await search_models_handler(arguments or {})
-    raise ValueError(f"Unknown tool: {name}")
-
+        return await run_initial_diagnostics(arguments)
+    elif name == "check_drift":
+        return await check_drift(arguments)
+    elif name == "get_all_models":
+        return await get_all_models_handler_async(arguments or {})
+    elif name == "search_models":
+        return await search_models_handler_async(arguments or {})
+    else:
+        raise ValueError(f"Unknown tool: {name}")
 
 # === Entrypoint ===
 async def main():
     async with stdio_server() as (reader, writer):
         await app.run(reader, writer, app.create_initialization_options())
 
-
 if __name__ == "__main__":
     asyncio.run(main())