mirror of
https://github.com/katanemo/plano.git
synced 2026-05-06 06:12:59 +02:00
307 lines
11 KiB
Python
307 lines
11 KiB
Python
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.12"
|
|
# dependencies = ["openai>=1.0.0"]
|
|
# ///
|
|
"""
|
|
Model Affinity Demo — Agentic Tool-Calling Loop
|
|
|
|
Runs the same agentic loop twice through Plano:
|
|
1. Without model affinity — the router may pick different models per turn
|
|
2. With model affinity — all turns use the model selected on turn 1
|
|
|
|
Each loop is a real tool-calling agent: the LLM decides which tools to call,
|
|
we provide simulated results, and the LLM continues until it has enough
|
|
information to produce a final answer. Each turn is a separate request to
|
|
Plano, so the router classifies intent independently every time.
|
|
|
|
Usage:
|
|
planoai up config.yaml # start Plano
|
|
uv run demo.py # run this demo
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import uuid
|
|
|
|
from openai import AsyncOpenAI
|
|
from openai.types.chat import ChatCompletionMessageParam
|
|
|
|
PLANO_URL = os.environ.get("PLANO_URL", "http://localhost:12000")
|
|
|
|
SYSTEM_PROMPT = (
|
|
"You are a database selection analyst. Use the provided tools to gather "
|
|
"benchmark data and case studies, then recommend PostgreSQL or MongoDB "
|
|
"for a high-traffic e-commerce backend. Be concise."
|
|
)
|
|
|
|
USER_QUERY = (
|
|
"Should we use PostgreSQL or MongoDB for our e-commerce platform? "
|
|
"We need strong consistency for orders but flexible schemas for products. "
|
|
"Use the tools to research both options, then give a recommendation."
|
|
)
|
|
|
|
TOOLS = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_db_benchmarks",
|
|
"description": "Fetch performance benchmarks for a database under a given workload.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"database": {
|
|
"type": "string",
|
|
"enum": ["postgresql", "mongodb"],
|
|
},
|
|
"workload": {
|
|
"type": "string",
|
|
"enum": ["read_heavy", "write_heavy", "mixed"],
|
|
},
|
|
},
|
|
"required": ["database", "workload"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_case_studies",
|
|
"description": "Retrieve real-world e-commerce case studies for a database.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"database": {
|
|
"type": "string",
|
|
"enum": ["postgresql", "mongodb"],
|
|
},
|
|
},
|
|
"required": ["database"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "check_feature_support",
|
|
"description": "Check if a database supports a specific feature.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"database": {
|
|
"type": "string",
|
|
"enum": ["postgresql", "mongodb"],
|
|
},
|
|
"feature": {"type": "string"},
|
|
},
|
|
"required": ["database", "feature"],
|
|
},
|
|
},
|
|
},
|
|
]
|
|
|
|
# Simulated tool responses
|
|
_BENCHMARKS = {
|
|
("postgresql", "mixed"): {
|
|
"read_qps": 42000,
|
|
"write_qps": 21000,
|
|
"p99_ms": 6,
|
|
"notes": "Solid all-round; MVCC keeps reads non-blocking",
|
|
},
|
|
("mongodb", "mixed"): {
|
|
"read_qps": 60000,
|
|
"write_qps": 50000,
|
|
"p99_ms": 3,
|
|
"notes": "Flexible schema accelerates feature iteration",
|
|
},
|
|
}
|
|
|
|
_CASE_STUDIES = {
|
|
"postgresql": [
|
|
{"company": "Shopify", "notes": "Moved orders back to Postgres for ACID"},
|
|
{
|
|
"company": "Zalando",
|
|
"notes": "Postgres + Citus for sharded order processing",
|
|
},
|
|
],
|
|
"mongodb": [
|
|
{"company": "eBay", "notes": "Product catalogue — flexible attribute schemas"},
|
|
{"company": "Alibaba", "notes": "Session/cart data — high write throughput"},
|
|
],
|
|
}
|
|
|
|
_FEATURES = {
|
|
("postgresql", "acid transactions"): {"supported": True, "notes": "Full ACID"},
|
|
("mongodb", "acid transactions"): {
|
|
"supported": True,
|
|
"notes": "Multi-doc ACID since v4.0",
|
|
},
|
|
("postgresql", "horizontal sharding"): {
|
|
"supported": True,
|
|
"notes": "Via Citus extension",
|
|
},
|
|
("mongodb", "horizontal sharding"): {
|
|
"supported": True,
|
|
"notes": "Native auto-balancing",
|
|
},
|
|
}
|
|
|
|
|
|
def dispatch_tool(name: str, args: dict) -> str:
|
|
if name == "get_db_benchmarks":
|
|
key = (args["database"], args["workload"])
|
|
return json.dumps(_BENCHMARKS.get(key, {"error": f"no data for {key}"}))
|
|
if name == "get_case_studies":
|
|
return json.dumps(_CASE_STUDIES.get(args["database"], {"error": "unknown db"}))
|
|
if name == "check_feature_support":
|
|
key = (args["database"], args["feature"].lower())
|
|
for k, v in _FEATURES.items():
|
|
if k[0] == key[0] and k[1] in key[1]:
|
|
return json.dumps(v)
|
|
return json.dumps({"error": f"no data for {key}"})
|
|
return json.dumps({"error": f"unknown tool {name}"})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Agentic loop — runs tool calls until the LLM produces a final answer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def run_agent_loop(
|
|
affinity_id: str | None = None,
|
|
max_turns: int = 10,
|
|
) -> tuple[str, list[dict]]:
|
|
"""
|
|
Run a tool-calling agent loop against Plano.
|
|
|
|
Returns (final_answer, trace) where trace is a list of
|
|
{"turn": int, "model": str, "tool_calls": [...]} dicts.
|
|
"""
|
|
client = AsyncOpenAI(base_url=f"{PLANO_URL}/v1", api_key="EMPTY")
|
|
headers = {"X-Model-Affinity": affinity_id} if affinity_id else None
|
|
|
|
messages: list[ChatCompletionMessageParam] = [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": USER_QUERY},
|
|
]
|
|
trace: list[dict] = []
|
|
|
|
for turn in range(1, max_turns + 1):
|
|
resp = await client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=messages,
|
|
tools=TOOLS,
|
|
tool_choice="auto",
|
|
max_completion_tokens=800,
|
|
extra_headers=headers,
|
|
)
|
|
|
|
choice = resp.choices[0]
|
|
turn_info: dict = {"turn": turn, "model": resp.model}
|
|
|
|
if choice.finish_reason == "tool_calls" and choice.message.tool_calls:
|
|
tool_names = [tc.function.name for tc in choice.message.tool_calls]
|
|
turn_info["tool_calls"] = tool_names
|
|
trace.append(turn_info)
|
|
|
|
messages.append(choice.message)
|
|
for tc in choice.message.tool_calls:
|
|
args = json.loads(tc.function.arguments or "{}")
|
|
result = dispatch_tool(tc.function.name, args)
|
|
messages.append(
|
|
{"role": "tool", "content": result, "tool_call_id": tc.id}
|
|
)
|
|
else:
|
|
turn_info["tool_calls"] = []
|
|
trace.append(turn_info)
|
|
return (choice.message.content or "").strip(), trace
|
|
|
|
return "(max turns reached)", trace
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Display helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def short_model(model: str) -> str:
|
|
return model.split("/")[-1] if "/" in model else model
|
|
|
|
|
|
def print_trace(trace: list[dict]) -> None:
|
|
for t in trace:
|
|
model = short_model(t["model"])
|
|
tools = ", ".join(t["tool_calls"]) if t["tool_calls"] else "final answer"
|
|
print(f" turn {t['turn']} [{model:<30}] {tools}")
|
|
|
|
|
|
def print_summary(label: str, trace: list[dict]) -> None:
|
|
models = [t["model"] for t in trace]
|
|
unique = set(models)
|
|
if len(unique) == 1:
|
|
print(
|
|
f" ✓ {label}: {short_model(next(iter(unique)))} "
|
|
f"for all {len(models)} turns"
|
|
)
|
|
else:
|
|
switches = sum(1 for a, b in zip(models, models[1:]) if a != b)
|
|
names = ", ".join(sorted(short_model(m) for m in unique))
|
|
print(f" ✗ {label}: model switched {switches} time(s) — {names}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def main() -> None:
|
|
print()
|
|
print(" ╔══════════════════════════════════════════════════════════╗")
|
|
print(" ║ Model Affinity Demo — Agentic Loop ║")
|
|
print(" ╚══════════════════════════════════════════════════════════╝")
|
|
print()
|
|
print(f" Plano : {PLANO_URL}")
|
|
print(f' Query : "{USER_QUERY[:65]}…"')
|
|
print()
|
|
print(" The agent calls tools (get_db_benchmarks, get_case_studies,")
|
|
print(" check_feature_support) across multiple turns. Each turn is")
|
|
print(" a separate request to Plano — the router classifies intent")
|
|
print(" independently, so different turns may get different models.")
|
|
print()
|
|
|
|
# --- Run 1: without affinity ---
|
|
print(" ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
|
print(" Run 1: WITHOUT Model Affinity")
|
|
print(" ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
|
print()
|
|
answer1, trace1 = await run_agent_loop(affinity_id=None)
|
|
print_trace(trace1)
|
|
print()
|
|
print_summary("Without affinity", trace1)
|
|
print()
|
|
|
|
# --- Run 2: with affinity ---
|
|
aid = str(uuid.uuid4())
|
|
print(" ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
|
print(f" Run 2: WITH Model Affinity (X-Model-Affinity: {aid[:8]}…)")
|
|
print(" ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
|
print()
|
|
answer2, trace2 = await run_agent_loop(affinity_id=aid)
|
|
print_trace(trace2)
|
|
print()
|
|
print_summary("With affinity ", trace2)
|
|
print()
|
|
|
|
# --- Final answer ---
|
|
print(" ══ Agent recommendation (affinity session) ════════════════")
|
|
print()
|
|
for line in answer2.splitlines():
|
|
print(f" {line}")
|
|
print()
|
|
print(" ═══════════════════════════════════════════════════════════")
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|