mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-31 04:45:23 +08:00
136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
"""Shared-runtime behaviour tests for the TavilySearch tool.
|
|
|
|
One ``TavilySearchTool`` object is driven through two distinct installation
|
|
bindings. The API key is read from the *current* invocation's installation
|
|
config (via the SDK task-local invocation context) and the vendor client is
|
|
built per call; the vendor's synchronous ``search`` must run off the shared
|
|
event loop.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from langbot_plugin.api.proxies.invocation import bind_invocation # noqa: E402
|
|
from langbot_plugin.entities.io.context import InstallationBinding # noqa: E402
|
|
|
|
import components.tools.tavily_search as tavily_module # noqa: E402
|
|
|
|
|
|
def binding(owner: str) -> InstallationBinding:
|
|
return InstallationBinding(
|
|
instance_uuid=f"instance-{owner}",
|
|
workspace_uuid=f"workspace-{owner}",
|
|
placement_generation=1,
|
|
installation_uuid=f"install-{owner}",
|
|
runtime_revision=1,
|
|
artifact_digest="a" * 64,
|
|
)
|
|
|
|
|
|
class _FakePlugin:
|
|
"""Minimal shared plugin object carrying the runtime handler identity.
|
|
|
|
``get_config`` returns a process-wide fallback (as a dedicated-mode plugin
|
|
would). A stateless component must ignore it while an invocation config is
|
|
bound; the tests assert the invocation config wins.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.plugin_runtime_handler = object()
|
|
|
|
def get_config(self) -> dict:
|
|
return {"tavily_api_key": "key-plugin-default"}
|
|
|
|
|
|
class _FakeTavilyClient:
|
|
"""Synchronous vendor-client stand-in (same shape as TavilyClient.search)."""
|
|
|
|
api_keys: list[str] = []
|
|
searches: list[dict] = []
|
|
delay = 0.0
|
|
|
|
def __init__(self, api_key: str) -> None:
|
|
type(self).api_keys.append(api_key)
|
|
self.api_key = api_key
|
|
|
|
def search(self, **kwargs) -> dict:
|
|
type(self).searches.append(dict(kwargs))
|
|
if self.delay:
|
|
time.sleep(self.delay)
|
|
return {
|
|
"results": [{"title": f"result-{self.api_key}", "url": "u", "content": "c"}],
|
|
}
|
|
|
|
|
|
class TavilySharedRuntimeTests(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self) -> None:
|
|
self.tool = tavily_module.TavilySearchTool()
|
|
self.plugin = _FakePlugin()
|
|
self.tool.plugin = self.plugin
|
|
self.handler = self.plugin.plugin_runtime_handler
|
|
_FakeTavilyClient.api_keys = []
|
|
_FakeTavilyClient.searches = []
|
|
_FakeTavilyClient.delay = 0.0
|
|
|
|
def _drive(self, owner: str, api_key: str):
|
|
return bind_invocation(
|
|
self.handler,
|
|
config={"tavily_api_key": api_key},
|
|
binding=binding(owner),
|
|
)
|
|
|
|
async def test_request_uses_current_invocation_credentials(self) -> None:
|
|
with mock.patch.object(tavily_module, "TavilyClient", _FakeTavilyClient):
|
|
with self._drive("a", "key-owner-a"):
|
|
result_a = await self.tool.call({"query": "query-a"}, None, 1)
|
|
with self._drive("b", "key-owner-b"):
|
|
result_b = await self.tool.call({"query": "query-b"}, None, 2)
|
|
|
|
self.assertEqual(_FakeTavilyClient.api_keys, ["key-owner-a", "key-owner-b"])
|
|
self.assertEqual([s["query"] for s in _FakeTavilyClient.searches],
|
|
["query-a", "query-b"])
|
|
self.assertIn("result-key-owner-a", result_a)
|
|
self.assertIn("result-key-owner-b", result_b)
|
|
|
|
# No per-tenant state retained on the shared component object.
|
|
self.assertLessEqual(set(vars(self.tool)), {"plugin"})
|
|
self.assertNotIn("key-owner-a", repr(vars(self.tool)))
|
|
self.assertNotIn("key-owner-b", repr(vars(self.tool)))
|
|
|
|
async def test_handler_does_not_block_event_loop(self) -> None:
|
|
_FakeTavilyClient.delay = 0.3
|
|
|
|
ticks = [0]
|
|
stop = asyncio.Event()
|
|
|
|
async def ticker() -> None:
|
|
while not stop.is_set():
|
|
ticks[0] += 1
|
|
await asyncio.sleep(0.01)
|
|
|
|
ticker_task = asyncio.create_task(ticker())
|
|
with mock.patch.object(tavily_module, "TavilyClient", _FakeTavilyClient):
|
|
with self._drive("a", "key-owner-a"):
|
|
await self.tool.call({"query": "query-a"}, None, 1)
|
|
|
|
progressed = ticks[0]
|
|
stop.set()
|
|
await ticker_task
|
|
|
|
# Pre-fix assertion: the synchronous client.search ran on the loop, so
|
|
# `progressed` was 0.
|
|
self.assertGreaterEqual(progressed, 3, "event loop stalled during Tavily search")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|