Add database integration: SQLite fallback, model and drift handlers

app.py

@@ -1,34 +1,38 @@
+import os
 import gradio as gr
 import asyncio
 from typing import Optional, List, Dict
 from contextlib import AsyncExitStack
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
+from database_module import init_db, get_all_models_handler, search_models_handler
 import json
 from datetime import datetime
 import plotly.graph_objects as go
-import plotly.express as px
+
+# --- Initialize database and MCP tool registration ---
+# Create tables and register MCP handlers
+init_db()
+
+
+# Ensure server.py imports and registers these tools:
+#   app.register_tool("get_all_models", get_all_models_handler)
+#   app.register_tool("search_models", search_models_handler)
 
 class MCPClient:
     def __init__(self):
         self.session: Optional[ClientSession] = None
         self.exit_stack = AsyncExitStack()
-    
-    async def connect_to_server(self, server_script_path: str = "mcp_server.py"):
+
+    async def connect_to_server(self, server_script_path: str = "server.py"):
         """Connect to MCP server"""
         is_python = server_script_path.endswith('.py')
-        is_js = server_script_path.endswith('.js')
-        
-        if not (is_python or is_js):
-            raise ValueError("Server script must be a .py or .js file")
-        
         command = "python" if is_python else "node"
         server_params = StdioServerParameters(
             command=command,
             args=[server_script_path],
             env=None
         )
-        
         stdio_transport = await self.exit_stack.enter_async_context(
             stdio_client(server_params)
         )
@@ -36,87 +40,71 @@ class MCPClient:
         self.session = await self.exit_stack.enter_async_context(
             ClientSession(self.stdio, self.write)
         )
-        
         await self.session.initialize()
-        
-        # List available tools
-        response = await self.session.list_tools()
-        tools = response.tools
-        print("Connected to server with tools:", [tool.name for tool in tools])
-    
+        tools = (await self.session.list_tools()).tools
+        print("Connected to server with tools:", [t.name for t in tools])
+
     async def call_tool(self, tool_name: str, arguments: dict):
         """Call a tool on the MCP server"""
         if not self.session:
             raise RuntimeError("Not connected to server")
-        
-        response = await self.session.call_tool(tool_name, arguments)
-        return response.content
-    
+        return (await self.session.call_tool(tool_name, arguments)).content
+
     async def close(self):
         """Close the MCP client connection"""
         await self.exit_stack.aclose()
 
+
 # Global MCP client instance
 mcp_client = MCPClient()
 
-# Async wrapper functions for Gradio
+
+# Helper to run async functions
 def run_async(coro):
-    """Helper to run async functions in Gradio"""
     try:
-        loop = asyncio.get_event_loop()
+        loop = asyncio.get_running_loop()
     except RuntimeError:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
-    
-    return loop.run_until_complete(coro)
+        return loop.run_until_complete(coro)
+    else:
+        # return result if coroutine returns value, else schedule
+        task = loop.create_task(coro)
+        return loop.run_until_complete(task) if not task.done() else task
+
 
-# Auto-connect to MCP server on startup
+# Initialize MCP connection on startup
 def initialize_mcp_connection():
-    """Initialize MCP connection on startup"""
     try:
         run_async(mcp_client.connect_to_server())
-        print("Successfully connected to MCP server on startup")
+        print("Successfully connected to MCP server")
         return True
     except Exception as e:
-        print(f"Failed to connect to MCP server on startup: {e}")
+        print(f"Failed to connect to MCP server: {e}")
         return False
 
-# MCP client functions
+
+# Wrapper functions remain unchanged but now call real DB-backed MCP tools
 def get_models_from_db():
-    """Get all models from database via MCP"""
     try:
         result = run_async(mcp_client.call_tool("get_all_models", {}))
         return result if isinstance(result, list) else []
     except Exception as e:
         print(f"Error getting models: {e}")
-        # Fallback data for demonstration
-        return [
-            {"name": "llama-3.1-8b-instant", "created": "2025-01-15", "description": "Fast and efficient model for instant responses."},
-            {"name": "llama3-8b-8192", "created": "2025-02-10", "description": "Extended context window model with 8192 tokens."},
-            {"name": "gemini-2.5-pro-preview-06-05", "created": "2025-06-05", "description": "Professional preview version of Gemini 2.5."},
-            {"name": "gemini-2.5-flash-preview-05-20", "created": "2025-05-20", "description": "Flash preview with optimized speed."},
-            {"name": "gemini-1.5-pro", "created": "2024-12-01", "description": "Stable professional release of Gemini 1.5."}
-        ]
+        return []
+
 
 def get_available_model_names():
-    """Get list of available model names for dropdown"""
-    models = get_models_from_db()
-    return [model["name"] for model in models]
+    return [m["name"] for m in get_models_from_db()]
+
 
 def search_models_in_db(search_term: str):
-    """Search models in database via MCP"""
     try:
         result = run_async(mcp_client.call_tool("search_models", {"search_term": search_term}))
         return result if isinstance(result, list) else []
     except Exception as e:
         print(f"Error searching models: {e}")
-        # Fallback search for demonstration
-        all_models = get_models_from_db()
-        if not search_term:
-            return all_models
-        term = search_term.lower()
-        return [model for model in all_models if term in model["name"].lower() or term in model["description"].lower()]
-
+        return [m for m in get_models_from_db() if search_term.lower() in m["name"].lower()]
 def format_dropdown_items(models):
     """Format dropdown items to show model name, creation date, and description preview"""
     formatted_items = []

database_module/__init__.py

@@ -0,0 +1,2 @@
+from .db import init_db
+from .mcp_tools import get_all_models_handler, search_models_handler

database_module/db.py

@@ -0,0 +1,33 @@
+import os
+from sqlalchemy import create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+# Configure your database URL (env var or fallback to SQLite file)
+DATABASE_URL = os.getenv(
+    "DATABASE_URL",
+    "sqlite:///./drift_detector.sqlite3"
+)
+
+# Create engine and session factory
+# For SQLite, disable same-thread check
+connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}
+engine = create_engine(
+    DATABASE_URL,
+    connect_args=connect_args,
+    pool_pre_ping=True
+)
+
+SessionLocal = sessionmaker(
+    autocommit=False,
+    autoflush=False,
+    bind=engine
+)
+Base = declarative_base()
+
+def init_db():
+    """
+    Create tables if they don't exist.
+    Call this once at application startup.
+    """
+    Base.metadata.create_all(bind=engine)

database_module/mcp_tools.py

@@ -0,0 +1,113 @@
+# database_module/mcp_tools.py
+from sqlalchemy.orm import Session
+from sqlalchemy import or_
+from .db import SessionLocal
+from .models import ModelEntry, DriftEntry
+from datetime import datetime
+from typing import Any, Dict, List
+
+
+def get_all_models_handler(_: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """
+    Return all models as list of dicts matching:
+      {name, created (ISO), description}
+    """
+    with SessionLocal() as session:
+        entries = session.query(ModelEntry).all()
+        return [
+            {"name": e.name, "created": e.created.isoformat(), "description": e.description or ""}
+            for e in entries
+        ]
+
+
+def search_models_handler(params: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """
+    Search models by name or description substring (case-insensitive).
+    params: {search_term: str}
+    """
+    term = params.get("search_term", "").strip().lower()
+    with SessionLocal() as session:
+        query = session.query(ModelEntry)
+        if term:
+            like_pattern = f"%{term}%"
+            query = query.filter(
+                or_(
+                    ModelEntry.name.ilike(like_pattern),
+                    ModelEntry.description.ilike(like_pattern)
+                )
+            )
+        entries = query.all()
+        return [
+            {"name": e.name, "created": e.created.isoformat(), "description": e.description or ""}
+            for e in entries
+        ]
+
+
+def get_model_details_handler(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Return a single model's details including system_prompt and description.
+    params: {model_name: str}
+    """
+    model_name = params.get("model_name")
+    with SessionLocal() as session:
+        e = session.query(ModelEntry).filter_by(name=model_name).first()
+        if not e:
+            return {"name": model_name, "system_prompt": "You are a helpful AI assistant.", "description": ""}
+        # You can store system_prompt as a column if desired; here placeholder
+        return {"name": e.name, "system_prompt": "You are a helpful AI assistant.", "description": e.description or ""}
+
+
+def save_model_handler(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Save or update a model's system_prompt.
+    params: {model_name: str, system_prompt: str}
+    """
+    name = params.get("model_name")
+    prompt = params.get("system_prompt", "")
+    with SessionLocal() as session:
+        entry = session.query(ModelEntry).filter_by(name=name).first()
+        if not entry:
+            # New model; created today
+            entry = ModelEntry(
+                name=name,
+                created=datetime.utcnow().date(),
+                description=""
+            )
+            session.add(entry)
+        # Optionally store prompt in another table or JSON field
+        session.commit()
+    return {"message": f"Model '{name}' saved."}
+
+
+def calculate_drift_handler(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Placeholder drift calculation: record a new random drift score today.
+    params: {model_name: str}
+    """
+    import random
+    name = params.get("model_name")
+    score = round(random.uniform(0, 1), 3)
+    today = datetime.utcnow().date()
+    with SessionLocal() as session:
+        entry = DriftEntry(
+            model_name=name,
+            date=today,
+            drift_score=score
+        )
+        session.add(entry)
+        session.commit()
+    return {"drift_score": score, "message": f"Drift recorded for '{name}'."}
+
+
+def get_drift_history_handler(params: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """
+    Return drift history as list of {date, drift_score} for a model.
+    params: {model_name: str}
+    """
+    name = params.get("model_name")
+    with SessionLocal() as session:
+        entries = session.query(DriftEntry).filter_by(model_name=name).order_by(DriftEntry.date).all()
+        return [
+            {"date": e.date.isoformat(), "drift_score": e.drift_score}
+            for e in entries
+        ]

database_module/models.py

@@ -0,0 +1,19 @@
+# database_module/models.py
+from sqlalchemy import Column, String, Date, Integer, Float, Text
+from .db import Base
+
+class ModelEntry(Base):
+    __tablename__ = "models"
+
+    id = Column(Integer, primary_key=True, index=True)
+    name = Column(String, unique=True, nullable=False, index=True)
+    created = Column(Date, nullable=False)
+    description = Column(Text, nullable=True)
+
+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)

ourllm.py

@@ -1,7 +1,47 @@
+import difflib
+from typing import List
+import mcp.types as types
 
-def genratequestionnaire(model, capabilities):
-    return None
+def genratequestionnaire(model: str, capabilities: str) -> List[types.SamplingMessage]:
+    """
+    Generate a baseline questionnaire for the given model.
+    Returns a list of SamplingMessage instances (role="user") with diagnostic questions.
+    """
+    questions = [
+        f"Model Name: {model}\nPlease confirm your model name.",
+        f"Capabilities Overview:\n{capabilities}\nPlease summarize your key capabilities.",
+        "Describe a typical use-case scenario that demonstrates these capabilities.",
+    ]
+    return [
+        types.SamplingMessage(
+            role="user",
+            content=types.TextContent(type="text", text=q)
+        )
+        for q in questions
+    ]
 
 
-def gradeanswers(old_answers, new_answers):
-    return None
+def gradeanswers(old_answers: List[str], new_answers: List[str]) -> List[types.SamplingMessage]:
+    """
+    Compare the old and new answers to compute a drift score.
+    Returns a list with a single SamplingMessage (role="assistant") whose content.text is the drift percentage.
+    """
+    total = len(old_answers)
+    if total == 0:
+        drift_pct = 0.0
+    else:
+        # Count how many answers are sufficiently similar
+        similar_count = 0
+        for old, new in zip(old_answers, new_answers):
+            ratio = difflib.SequenceMatcher(None, old, new).ratio()
+            if ratio >= 0.8:
+                similar_count += 1
+        drift_pct = round((1 - (similar_count / total)) * 100, 2)
+
+    drift_text = f"{drift_pct}"
+    return [
+        types.SamplingMessage(
+            role="assistant",
+            content=types.TextContent(type="text", text=drift_text)
+        )
+    ]

requirements.txt

@@ -2,4 +2,6 @@ gradio
 mcp
 plotly
 asyncio
-typing
+typing
+sqlalchemy
+psycopg2