mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-08 22:22:17 +02:00
feat(capabilities): generate agent tool door from registry
This commit is contained in:
parent
23802a74bb
commit
de4b8ab6b0
3 changed files with 180 additions and 1 deletions
|
|
@ -1 +1 @@
|
|||
"""Access doors (05): thin adapters that expose the 04 registry as REST/MCP/chat."""
|
||||
"""Access doors (05): thin adapters that expose the 04 registry as REST/MCP/agent."""
|
||||
|
|
|
|||
52
surfsense_backend/app/capabilities/access/agent.py
Normal file
52
surfsense_backend/app/capabilities/access/agent.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
"""Generate the agent door from the capability registry (05).
|
||||
|
||||
One LangChain tool per verb; each runs the same thin adapter as the REST door
|
||||
(``access/rest.py``): meter-gate -> executor -> charge. The tool returns the
|
||||
verb's serialized output so the model can reason over it; UI cards are the SSE
|
||||
emission handler's job, not this generator's.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from langchain_core.tools import BaseTool, StructuredTool
|
||||
|
||||
from app.capabilities.billing import charge_capability, gate_capability
|
||||
from app.capabilities.store import all_capabilities
|
||||
from app.capabilities.types import Capability, CapabilityContext
|
||||
from app.db import async_session_maker
|
||||
from app.services.web_crawl_credit_service import InsufficientCreditsError
|
||||
|
||||
|
||||
def build_capability_tools(
|
||||
*,
|
||||
workspace_id: int,
|
||||
capabilities: list[Capability] | None = None,
|
||||
) -> list[BaseTool]:
|
||||
"""Emit one tool per verb (defaults to the whole registry)."""
|
||||
caps = capabilities if capabilities is not None else all_capabilities()
|
||||
return [_capability_tool(cap, workspace_id) for cap in caps]
|
||||
|
||||
|
||||
def _capability_tool(capability: Capability, workspace_id: int) -> BaseTool:
|
||||
input_model = capability.input_schema
|
||||
unit = capability.billing_unit
|
||||
executor = capability.executor
|
||||
|
||||
async def _run(**kwargs: object) -> dict | str:
|
||||
payload = input_model(**kwargs)
|
||||
async with async_session_maker() as session:
|
||||
ctx = CapabilityContext(session=session, workspace_id=workspace_id)
|
||||
try:
|
||||
await gate_capability(payload, unit, ctx)
|
||||
except InsufficientCreditsError as exc:
|
||||
return str(exc)
|
||||
output = await executor(payload)
|
||||
await charge_capability(output, unit, ctx)
|
||||
return output.model_dump()
|
||||
|
||||
return StructuredTool.from_function(
|
||||
coroutine=_run,
|
||||
name=capability.name.replace(".", "_"),
|
||||
description=capability.description,
|
||||
args_schema=input_model,
|
||||
)
|
||||
|
|
@ -0,0 +1,127 @@
|
|||
"""The agent door (05): generate one LangChain tool per registry verb."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.capabilities.types import BillingUnit, Capability
|
||||
from app.services.web_crawl_credit_service import InsufficientCreditsError
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
class _EchoInput(BaseModel):
|
||||
text: str = Field(description="The text to echo back.")
|
||||
|
||||
|
||||
class _EchoOutput(BaseModel):
|
||||
echoed: str
|
||||
|
||||
@property
|
||||
def billable_units(self) -> int:
|
||||
return 1
|
||||
|
||||
|
||||
def _capability(
|
||||
*, name: str, output: _EchoOutput, unit=BillingUnit.WEB_CRAWL
|
||||
) -> Capability:
|
||||
async def _executor(payload: _EchoInput) -> _EchoOutput:
|
||||
_executor.seen = payload
|
||||
return output
|
||||
|
||||
cap = Capability(
|
||||
name=name,
|
||||
description=f"{name} does a thing.",
|
||||
input_schema=_EchoInput,
|
||||
output_schema=_EchoOutput,
|
||||
executor=_executor,
|
||||
billing_unit=unit,
|
||||
)
|
||||
cap.executor.seen = None # type: ignore[attr-defined]
|
||||
return cap
|
||||
|
||||
|
||||
class _FakeSessionCtx:
|
||||
async def __aenter__(self):
|
||||
return SimpleNamespace()
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolate(monkeypatch):
|
||||
"""Stub the billing session + charge/gate so tools never hit the DB."""
|
||||
from app.capabilities.access import agent as mod
|
||||
|
||||
monkeypatch.setattr(mod, "async_session_maker", lambda: _FakeSessionCtx())
|
||||
charge = AsyncMock()
|
||||
gate = AsyncMock()
|
||||
monkeypatch.setattr(mod, "charge_capability", charge)
|
||||
monkeypatch.setattr(mod, "gate_capability", gate)
|
||||
return SimpleNamespace(module=mod, charge=charge, gate=gate)
|
||||
|
||||
|
||||
async def test_registry_becomes_one_tool_per_verb(isolate):
|
||||
caps = [
|
||||
_capability(name="web.scrape", output=_EchoOutput(echoed="a")),
|
||||
_capability(name="web.discover", output=_EchoOutput(echoed="b"), unit=None),
|
||||
]
|
||||
|
||||
tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=caps)
|
||||
|
||||
by_name = {t.name: t for t in tools}
|
||||
assert set(by_name) == {"web_scrape", "web_discover"}
|
||||
assert by_name["web_scrape"].description == "web.scrape does a thing."
|
||||
assert by_name["web_scrape"].args_schema is _EchoInput
|
||||
|
||||
|
||||
async def test_input_field_docs_reach_the_model(isolate):
|
||||
"""Per-field descriptions must surface in the tool's args schema (LLM context)."""
|
||||
cap = _capability(name="web.scrape", output=_EchoOutput(echoed="a"))
|
||||
[tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
|
||||
|
||||
assert tool.args["text"]["description"] == "The text to echo back."
|
||||
|
||||
|
||||
async def test_tool_runs_executor_and_returns_serialized_output(isolate):
|
||||
cap = _capability(name="web.scrape", output=_EchoOutput(echoed="hi there"))
|
||||
[tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
|
||||
|
||||
result = await tool.ainvoke({"text": "ping"})
|
||||
|
||||
assert result == {"echoed": "hi there"}
|
||||
assert cap.executor.seen.text == "ping"
|
||||
|
||||
|
||||
async def test_tool_charges_owner(isolate):
|
||||
output = _EchoOutput(echoed="hi")
|
||||
cap = _capability(name="web.scrape", output=output)
|
||||
[tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
|
||||
|
||||
await tool.ainvoke({"text": "ping"})
|
||||
|
||||
isolate.charge.assert_awaited_once()
|
||||
(charged_output, unit, ctx), _ = isolate.charge.call_args
|
||||
assert charged_output is output
|
||||
assert unit is BillingUnit.WEB_CRAWL
|
||||
assert ctx.workspace_id == 7
|
||||
|
||||
|
||||
async def test_over_budget_returns_friendly_message(isolate):
|
||||
cap = _capability(name="web.scrape", output=_EchoOutput(echoed="hi"))
|
||||
isolate.gate.side_effect = InsufficientCreditsError(
|
||||
message="This run would exceed your available credit.",
|
||||
balance_micros=0,
|
||||
required_micros=1_000_000,
|
||||
)
|
||||
[tool] = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
|
||||
|
||||
result = await tool.ainvoke({"text": "ping"})
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "credit" in result.lower()
|
||||
assert cap.executor.seen is None
|
||||
isolate.charge.assert_not_awaited()
|
||||
Loading…
Add table
Add a link
Reference in a new issue