code.py

"""Model Context Protocol — talk to tools through a contract, not an import.

Compare with tool_use/: there the agent imports Python functions directly, so
the agent and its tools must live in the same process, same language, same
deployment. MCP puts a protocol between them. The agent discovers what exists
at runtime by asking, instead of having tools hardcoded at import time.

That indirection is the whole value: swap the server, keep the agent.

ponytail: this runs a JSON-RPC-shaped server in-process rather than over
stdio/HTTP, so the example needs no extra dependency or subprocess. The wire
format and the discover -> call flow are real; only the transport is faked.
Point `send()` at a real transport and nothing above it changes.
"""

import json
from datetime import date

from langchain_core.messages import HumanMessage, ToolMessage

from resources.agent import llm
from resources.helper import show_response

# ---------------------------------------------------------------------------
# SERVER SIDE — knows nothing about LLMs. It just publishes a tool catalogue
# and executes calls against it. In production this is a separate process,
# possibly written by someone else, possibly not in Python.
# ---------------------------------------------------------------------------

INVENTORY = {"SKU-101": 42, "SKU-202": 0, "SKU-303": 7}

# The catalogue IS the contract. Schemas are declared as data, which is what
# lets a client discover them without importing any server code.
SERVER_TOOLS = [
    {
        "name": "check_stock",
        "description": "Return the units in stock for a product SKU.",
        "inputSchema": {
            "type": "object",
            "properties": {"sku": {"type": "string"}},
            "required": ["sku"],
        },
    },
    {
        "name": "restock_eta",
        "description": "Return the restock date for an out-of-stock SKU.",
        "inputSchema": {
            "type": "object",
            "properties": {"sku": {"type": "string"}},
            "required": ["sku"],
        },
    },
]


def server(request):
    """Handle one JSON-RPC request. The only entry point the client gets."""
    method, params = request["method"], request.get("params", {})

    if method == "tools/list":
        return {"result": {"tools": SERVER_TOOLS}}

    if method == "tools/call":
        name, args = params["name"], params["arguments"]
        if name == "check_stock":
            units = INVENTORY.get(args["sku"])
            text = (
                f"{units} units" if units is not None else f"unknown SKU {args['sku']}"
            )
        elif name == "restock_eta":
            text = f"{date(2026, 9, 1)}" if INVENTORY.get(args["sku"]) == 0 else "in stock"
        else:
            # Unknown method is a protocol error, not a crash.
            return {"error": {"code": -32601, "message": f"no tool {name}"}}
        return {"result": {"content": [{"type": "text", "text": text}]}}

    return {"error": {"code": -32601, "message": f"unknown method {method}"}}


# ---------------------------------------------------------------------------
# CLIENT SIDE — knows nothing about inventory. It discovers whatever the
# server offers and adapts it to the model's tool format.
# ---------------------------------------------------------------------------


def send(method, **params):
    """The transport seam. Swap this for stdio or HTTP; nothing else changes."""
    request = {"jsonrpc": "2.0", "id": 1, "method": method, "params": params}
    print(f"  -> {json.dumps(request)}")
    response = server(request)
    print(f"  <- {json.dumps(response)}")
    if "error" in response:
        raise RuntimeError(response["error"]["message"])
    return response["result"]


def discover():
    """Ask the server what it can do, then translate to the model's schema.

    Nothing here names check_stock or restock_eta. Add a tool to the server
    and this client picks it up on the next run with no code change.
    """
    tools = send("tools/list")["tools"]
    return [
        {
            "type": "function",
            "function": {
                "name": t["name"],
                "description": t["description"],
                "parameters": t["inputSchema"],
            },
        }
        for t in tools
    ]


MAX_STEPS = 5


def ask(question, tool_schemas):
    """Standard tool loop - but every execution goes over the protocol."""
    messages = [HumanMessage(question)]
    bound = llm.bind_tools(tool_schemas)

    for _ in range(MAX_STEPS):
        reply = bound.invoke(messages)
        messages.append(reply)
        if not reply.tool_calls:
            return reply

        for call in reply.tool_calls:
            result = send("tools/call", name=call["name"], arguments=call["args"])
            text = result["content"][0]["text"]
            messages.append(ToolMessage(text, tool_call_id=call["id"]))

    return reply


if __name__ == "__main__":
    print(f"{'=' * 60}\nDISCOVERY\n{'=' * 60}")
    schemas = discover()
    print(f"\n  discovered {len(schemas)} tool(s): "
          f"{[s['function']['name'] for s in schemas]}")

    for question in [
        "How many units of SKU-101 do we have?",
        "Is SKU-202 available, and if not when does it come back?",
    ]:
        print(f"\n{'=' * 60}\nQUESTION: {question}\n{'=' * 60}")
        show_response(ask(question, schemas))