106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
"""Exercise the instrumentation as a base install would, then print JSON.
|
|
|
|
Run in a fresh interpreter. greenlet and aiosqlite are made unimportable
|
|
before anything else loads, as on ``pip install fastmcp-feedback`` without
|
|
the ``instrumentation`` extra. Argument: a directory for the feedback database.
|
|
"""
|
|
|
|
import sys
|
|
from importlib.metadata import version
|
|
|
|
# SQLAlchemy 2.1 refuses to import sqlalchemy.ext.asyncio without greenlet,
|
|
# once: a second import succeeds. 2.0 imports it and fails on first use.
|
|
# On 2.0 the module is blocked too, so both versions fail at import.
|
|
_blocked = ["greenlet", "aiosqlite"]
|
|
if tuple(int(p) for p in version("sqlalchemy").split(".")[:2]) < (2, 1):
|
|
_blocked.append("sqlalchemy.ext.asyncio")
|
|
for _name in _blocked:
|
|
sys.modules[_name] = None # import raises ModuleNotFoundError
|
|
|
|
import asyncio # noqa: E402
|
|
import io # noqa: E402
|
|
import json # noqa: E402
|
|
from pathlib import Path # noqa: E402
|
|
|
|
STORAGE = "fastmcp_feedback.instrumentation.storage"
|
|
EMBEDDINGS = "fastmcp_feedback.instrumentation.embeddings"
|
|
|
|
|
|
async def run(db_dir: Path) -> dict:
|
|
from fastmcp import Client, FastMCP
|
|
|
|
import fastmcp_feedback.instrumentation as inst
|
|
from fastmcp_feedback import add_feedback_tools
|
|
from fastmcp_feedback.instrumentation import (
|
|
DEFAULT_TABLE_PREFIX,
|
|
JsonLinesSink,
|
|
MemorySink,
|
|
instrument,
|
|
)
|
|
|
|
out: dict = {"prefix": DEFAULT_TABLE_PREFIX}
|
|
app = FastMCP("base-install")
|
|
|
|
@app.tool
|
|
def add(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
memory = MemorySink()
|
|
lines = io.StringIO()
|
|
mw = instrument(app, [memory, JsonLinesSink(lines)], mode="full")
|
|
add_feedback_tools(
|
|
app, database_url=f"sqlite:///{db_dir / 'fb.db'}", instrumentation=mw
|
|
)
|
|
|
|
async with Client(app) as client:
|
|
result = await client.call_tool("add", {"a": 2, "b": 3})
|
|
out["result"] = result.content[0].text
|
|
fb = await client.call_tool(
|
|
"submit_feedback",
|
|
{
|
|
"request": {
|
|
"type": "bug",
|
|
"title": "base install",
|
|
"description": "submitted without the extra",
|
|
"submitter": "worker",
|
|
}
|
|
},
|
|
)
|
|
out["feedback"] = json.loads(fb.content[0].text)
|
|
|
|
mw.record_event("job.done", key="j1", attrs={"n": 1})
|
|
await mw.flush()
|
|
out["tools"] = [r.tool for r in memory.records if hasattr(r, "tool")]
|
|
out["event_kinds"] = [r.kind for r in memory.records if hasattr(r, "kind")]
|
|
out["jsonl_lines"] = len(lines.getvalue().splitlines())
|
|
out["capture"] = await mw.capture_recent_calls()
|
|
out["similar"] = await mw.similar("anything")
|
|
out["similar_feedback"] = await mw.similar_feedback("bug-1")
|
|
out["feedback_context"] = await mw.feedback_context("bug-1")
|
|
out["link_feedback"] = await mw.link_feedback("bug-1", [])
|
|
out["events_for"] = len(await mw.events_for("j1"))
|
|
out["database_sink"] = mw.database_sink is None
|
|
out["embedding_sink"] = mw.embedding_sink is None
|
|
await mw.aclose()
|
|
|
|
errors = {}
|
|
for name in (
|
|
"DatabaseSink",
|
|
"EmbeddingSink",
|
|
"build_metadata",
|
|
"backfill_feedback",
|
|
):
|
|
for _attempt in range(2): # SQLAlchemy only complains the first time
|
|
try:
|
|
getattr(inst, name)
|
|
errors.setdefault(name, []).append(None)
|
|
except ImportError as exc:
|
|
errors.setdefault(name, []).append(str(exc))
|
|
out["errors"] = errors
|
|
out["loaded"] = sorted(m for m in (STORAGE, EMBEDDINGS) if m in sys.modules)
|
|
return out
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(json.dumps(asyncio.run(run(Path(sys.argv[1]))), default=str))
|