Spaces:
Sleeping
Sleeping
File size: 5,880 Bytes
0d5ac71 dc9d63b 0d5ac71 dc9d63b 0d5ac71 9fa019c dc9d63b 9fa019c dc9d63b 9fa019c dc9d63b 0d5ac71 dc9d63b 0d5ac71 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import asyncio
import json
import os
from typing import Any
import mcp.types as types
from mcp import CreateMessageResult
from mcp.server import Server
from mcp.server.stdio import stdio_server
from ourllm import genratequestionnaire, gradeanswers
DATA_DIR = "data"
os.makedirs(DATA_DIR, exist_ok=True)
app = Server("mcp-drift-server")
registered_models = {}
def get_all_models():
"""Retrieve all registered models."""
return list(registered_models.keys())
def search_models(query: str):
"""Search registered models by name."""
return [model for model in registered_models if query.lower() in model.lower()]
def get_model_details(model_name: str):
"""Get details of a specific model."""
return registered_models.get(model_name, None)
def save_model(model_name: str, model_details: dict):
"""Save a new model or update an existing one."""
registered_models[model_name] = model_details
with open(os.path.join(DATA_DIR, "models.json"), "w") as f:
json.dump(registered_models, f, indent=2)
@app.list_tools()
async def list_tools() -> list[types.Tool]:
return [
types.Tool(
name="run_initial_diagnostics",
description="Generate and store baseline diagnostics for a connected LLM.",
inputSchema={"type":"object",
"properties": {
"model": {
"type": "string",
"description": "The name of the model to run diagnostics on"
},
"model_capabilities": {
"type": "string",
"description": "Full description of the model's capabilities, including any special features"
}
},
"required": ["model", "model_capabilities"]},
),
types.Tool(
name="check_drift",
description="Re-run diagnostics and compare to baseline for drift scoring.",
inputSchema={"type":"object",
"properties": {
"model": {
"type": "string",
"description": "The name of the model to run diagnostics on"
},
},
"required": ["model"]},
),
]
# === Sampling Wrapper ===
async def sample(messages: list[types.SamplingMessage], max_tokens=300) -> CreateMessageResult:
return await app.request_context.session.create_message(
messages=messages,
max_tokens=max_tokens,
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]:
if arguments and "model" in arguments:
model = arguments["model"]
else:
raise(ValueError("Model details is required"))
# 1. Ask the server's internal LLM to generate a questionnaire
questions = await genratequestionnaire(model, arguments["model_capabilities"]) # Server-side trusted LLM
# 2. Send questionnaire to target LLM (i.e., the client)
answers = await sample(questions) # Client model's answers
# 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)
return [types.TextContent(type="text", text="Baseline stored for model: " + model)]
async def check_drift(arguments: dict[str, str]) -> list[types.TextContent]:
if arguments and "model" in arguments:
model = arguments["model"]
else:
raise (ValueError("Model details is required"))
baseline_path = get_baseline_path(model)
if not os.path.exists(baseline_path):
return [types.TextContent(type="text", text="No baseline exists for model: " + model)]
with open(baseline_path) as f:
data = json.load(f)
questions = [types.SamplingMessage(role="user", content=types.TextContent(type="text", text=q)) for q in
data["questions"]]
old_answers = data["answers"]
# 1. Ask the model again
new_answers_msgs = await sample(questions)
new_answers = [m.content.text for m in new_answers_msgs]
grading_response = await 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)
# 4. Optionally alert if high drift
alert = "🚨 Significant drift detected!" if float(drift_score) > 50 else "✅ Drift within acceptable limits."
return [
types.TextContent(type="text", text=f"Drift score for {model}: {drift_score}"),
types.TextContent(type="text", text=alert)
]
@app.call_tool()
async def call_tool(name: str, arguments: dict[str, Any] | None = None):
if name == "run_initial_diagnostics":
return await run_initial_diagnostics(arguments)
elif name == "check_drift":
return await check_drift(arguments)
else:
raise ValueError(f"Unknown tool: {name}")
# === Entrypoint ===
async def main():
async with stdio_server() as streams:
await app.run(streams[0], streams[1], app.create_initialization_options())
if __name__ == "__main__":
asyncio.run(main())
|