mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-10 22:32:16 +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