Connect a local model to the MCP tool

from typing import Literal

from mcp.server.fastmcp import FastMCP
from mcp.types import ToolAnnotations

server = FastMCP("Fictional library")
GUIDE_URI = "library://guide/visiting"
HOURS = {"central": "09:00–17:00", "west": "10:00–16:00"}


@server.resource(GUIDE_URI, mime_type="text/plain")
def visiting_guide() -> str:
    return "Ask at the front desk for a visitor card. Equipment bookings go to the media desk."


@server.tool(annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False))
def lookup_hours(branch: Literal["central", "west"]) -> dict[str, str]:
    """Return the opening hours for one fictional library branch."""
    return {"branch": branch, "hours": HOURS[branch]}


if __name__ == "__main__":
    server.run(transport="stdio")

import asyncio
from datetime import timedelta
from pathlib import Path
import sys

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from library_mcp import GUIDE_URI


def library_server():
    return StdioServerParameters(
        command=sys.executable,
        args=[str(Path(__file__).with_name("library_mcp.py").resolve())],
    )


async def main():
    async with stdio_client(library_server()) as (read, write):
        async with ClientSession(read, write, read_timeout_seconds=timedelta(seconds=15)) as session:
            await session.initialize()
            tools = await session.list_tools()
            resources = await session.list_resources()
            print("Tools:", ", ".join(tool.name for tool in tools.tools))
            print("Resources:", ", ".join(str(item.uri) for item in resources.resources))
            guide = await session.read_resource(GUIDE_URI)
            print("Visiting guide:", guide.contents[0].text)

            result = await session.call_tool("lookup_hours", {"branch": "central"})
            if result.isError or result.structuredContent is None:
                raise ValueError("The server could not return the branch hours.")
            print("Central hours:", result.structuredContent["hours"])

            invalid = await session.call_tool("lookup_hours", {"branch": "north"})
            if not invalid.isError:
                raise ValueError("Expected the server to reject an unknown branch.")
            print("Unknown branch rejected.")


if __name__ == "__main__":
    asyncio.run(main())

import argparse
import asyncio
from datetime import timedelta
import json

from mcp import ClientSession
from mcp.client.stdio import stdio_client

from inspect_library_mcp import library_server
from library_mcp import GUIDE_URI, HOURS
from semantic_notes import api

TOOL = "lookup_hours"


def approved_arguments(calls, selected_branch):
    if not isinstance(calls, list) or len(calls) != 1 or not isinstance(calls[0], dict):
        return None
    function = calls[0].get("function")
    if not isinstance(function, dict) or function.get("name") != TOOL:
        return None
    arguments = function.get("arguments")
    if arguments == {"branch": selected_branch} and selected_branch in HOURS:
        return arguments
    return None


def chat(messages, tools=None):
    payload = {
        "model": "qwen2.5:1.5b",
        "messages": messages,
        "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 140},
        "stream": False,
    }
    if tools:
        payload["tools"] = tools
    return api("/api/chat", payload)["message"]


async def get_answer(branch):
    async with stdio_client(library_server()) as (read, write):
        async with ClientSession(read, write, read_timeout_seconds=timedelta(seconds=15)) as session:
            await session.initialize()
            tools = (await session.list_tools()).tools
            tool = next((item for item in tools if item.name == TOOL), None)
            if tool is None:
                raise ValueError("The library server does not provide lookup_hours.")
            tool_schema = {
                "type": "function",
                "function": {"name": tool.name, "description": tool.description,
                             "parameters": tool.inputSchema},
            }
            messages = [{
                "role": "user",
                "content": f"Use the lookup_hours tool for the {branch} branch.",
            }]
            reply = await asyncio.to_thread(chat, messages, [tool_schema])
            arguments = approved_arguments(reply.get("tool_calls"), branch)
            if arguments is None:
                print(f"Please confirm the branch. No lookup ran for {branch}.")
                return

            result = await session.call_tool(TOOL, arguments)
            data = result.structuredContent
            if result.isError or not isinstance(data, dict):
                raise ValueError("The library server could not complete the lookup.")
            hours = data.get("hours")
            if data.get("branch") != branch or not isinstance(hours, str) or not hours:
                raise ValueError("The tool result did not match the selected branch.")
            print(f"Verified lookup: {branch} {hours}")
            guide = (await session.read_resource(GUIDE_URI)).contents[0].text
            messages.extend([
                reply,
                {"role": "tool", "tool_name": TOOL, "content": json.dumps(data)},
                {"role": "user", "content": (
                    f"State the hours from the tool and one visitor tip from this reference: {guide}"
                )},
            ])
            final = await asyncio.to_thread(chat, messages)
            print("Model reply:", final.get("content") or "(empty)")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Ask a local model to look up one branch.")
    parser.add_argument("branch", choices=HOURS)
    args = parser.parse_args()
    try:
        asyncio.run(get_answer(args.branch))
    except (OSError, KeyError, ValueError) as error:
        raise SystemExit(str(error)) from None

import argparse
from contextlib import closing
from itertools import islice
import json
import math
from pathlib import Path
import sqlite3
from urllib.error import HTTPError
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

BASE = "http://127.0.0.1:11434"
MODEL = "all-minilm:latest"
MAX_PASSAGES = 200


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def api(endpoint, payload=None):
    request = Request(BASE + endpoint)
    if payload is not None:
        request.data = json.dumps(payload).encode("utf-8")
        request.add_header("Content-Type", "application/json")
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(request, timeout=120) as response:
        return json.load(response)


def model_digest():
    for model in api("/api/tags")["models"]:
        if model["name"] == MODEL:
            return model["digest"]
    raise ValueError(f"Pull {MODEL} with Ollama first.")


def embed(texts):
    result = api("/api/embed", {
        "model": MODEL,
        "input": texts,
        "truncate": False,
    })["embeddings"]
    if len(result) != len(texts):
        raise ValueError("Ollama returned the wrong number of embeddings.")
    return result


def cosine(left, right):
    if len(left) != len(right):
        raise ValueError("Vector dimensions changed. Rebuild the index.")
    length = math.hypot(*left) * math.hypot(*right)
    return sum(x * y for x, y in zip(left, right)) / length if length else 0


def build_index(folder, database):
    from index_notes import passages

    rows = list(islice(passages(folder), MAX_PASSAGES + 1))
    if not rows or len(rows) > MAX_PASSAGES:
        raise ValueError(f"Use 1–{MAX_PASSAGES} nonempty passages for this example.")
    digest = model_digest()
    stored = []
    for start in range(0, len(rows), 16):
        batch = rows[start:start + 16]
        vectors = embed([text for _, _, text in batch])
        stored.extend((*row, json.dumps(vector)) for row, vector in zip(batch, vectors))
    if model_digest() != digest:
        raise ValueError("The embedding model changed during indexing. Build the index again.")
    with closing(sqlite3.connect(database)) as connection, connection:
        connection.execute("CREATE TABLE IF NOT EXISTS model (name TEXT, digest TEXT)")
        connection.execute(
            "CREATE TABLE IF NOT EXISTS passages "
            "(source TEXT, number INTEGER, text TEXT, vector TEXT)"
        )
        connection.execute("DELETE FROM passages")
        connection.execute("DELETE FROM model")
        connection.execute("INSERT INTO model VALUES (?, ?)", (MODEL, digest))
        connection.executemany("INSERT INTO passages VALUES (?, ?, ?, ?)", stored)
    print(f"Saved {len(stored)} embeddings in {database}.")


def find_notes(database, query):
    if not query.strip() or not database.is_file():
        raise ValueError("Build the index and enter a nonempty question.")
    with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as connection:
        connection.execute("BEGIN")
        identity = connection.execute("SELECT name, digest FROM model").fetchone()
        rows = connection.execute("SELECT source, number, text, vector FROM passages").fetchall()
    digest = model_digest()
    if identity != (MODEL, digest):
        raise ValueError("The embedding model differs from the one in the index. Rebuild it.")
    vector, = embed([query])
    if model_digest() != digest:
        raise ValueError("The embedding model changed during the query. Rebuild the index.")
    matches = [
        (cosine(json.loads(embedding), vector), source, number, text)
        for source, number, text, embedding in rows
    ]
    ranked = sorted(matches, key=lambda match: match[0], reverse=True)
    for score, source, number, text in ranked[:3]:
        print(f"{score:.3f} | {source}, passage {number}\n{text}\n")


def main():
    parser = argparse.ArgumentParser(description="Store local embeddings and search them later.")
    parser.add_argument("--db", type=Path, default=Path("notes-vectors.sqlite3"))
    commands = parser.add_subparsers(dest="command", required=True)
    build = commands.add_parser("build")
    build.add_argument("folder", type=Path)
    find = commands.add_parser("find")
    find.add_argument("query")
    args = parser.parse_args()

    try:
        if args.command == "build":
            build_index(args.folder, args.db)
        else:
            find_notes(args.db, args.query)
    except HTTPError as error:
        raise SystemExit(f"Ollama returned HTTP {error.code}: {error.read().decode()}") from None
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()