mirror of
https://github.com/andrewyng/openworker.git
synced 2026-09-03 13:00:37 +00:00
Rohit C PrasadandDevika
2b45018ffa
Imported from andrewyng/aisuite@1b4bbf303e (contents of its platform/ directory, hoisted to the repo root). Development history prior to this commit lives in that repository. Co-authored-by: Devika <devikaverma11@gmail.com>
155 lines
4.8 KiB
Python
155 lines
4.8 KiB
Python
"""Tests for web search — provider abstraction, the tool, and config resolution.
|
|
|
|
No network: a FakeProvider is injected; third-party key handling and the REST config path
|
|
are exercised without hitting DuckDuckGo/Tavily/Brave.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from coworker.secrets import SecretStore
|
|
from coworker.web import (
|
|
SearchResult,
|
|
build_provider,
|
|
make_web_search_tool,
|
|
provider_names,
|
|
)
|
|
from coworker.web.providers import (
|
|
BraveProvider,
|
|
DuckDuckGoProvider,
|
|
TavilyProvider,
|
|
WebSearchProvider,
|
|
)
|
|
|
|
|
|
class FakeProvider(WebSearchProvider):
|
|
name = "fake"
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def search(self, query, max_results=5):
|
|
self.calls.append((query, max_results))
|
|
return [
|
|
SearchResult(title=f"r{i}", url=f"https://x/{i}", snippet="s")
|
|
for i in range(max_results)
|
|
]
|
|
|
|
|
|
def test_tool_returns_results():
|
|
fake = FakeProvider()
|
|
tool = make_web_search_tool(provider=fake)
|
|
out = tool(query="anthropic", max_results=3)
|
|
assert out["provider"] == "fake"
|
|
assert [r["title"] for r in out["results"]] == ["r0", "r1", "r2"]
|
|
assert fake.calls == [("anthropic", 3)]
|
|
# metadata + schema for the registry
|
|
assert tool.__aisuite_tool_metadata__.category == "web"
|
|
assert tool.__coworker_schema__["function"]["name"] == "web_search"
|
|
|
|
|
|
def test_tool_clamps_max_results():
|
|
fake = FakeProvider()
|
|
make_web_search_tool(provider=fake)(query="q", max_results=99)
|
|
assert fake.calls[0][1] == 10 # clamped to 10
|
|
make_web_search_tool(provider=fake)(query="q", max_results=0)
|
|
assert fake.calls[1][1] == 1 # clamped to >=1
|
|
|
|
|
|
def test_tool_reports_search_errors():
|
|
class Boom(WebSearchProvider):
|
|
name = "boom"
|
|
|
|
def search(self, query, max_results=5):
|
|
raise RuntimeError("network down")
|
|
|
|
out = make_web_search_tool(provider=Boom())(query="q")
|
|
assert "web search failed" in out["error"] and out["provider"] == "boom"
|
|
|
|
|
|
def test_build_provider_default_is_keyless_duckduckgo():
|
|
assert isinstance(build_provider("duckduckgo"), DuckDuckGoProvider)
|
|
assert isinstance(build_provider("unknown-thing"), DuckDuckGoProvider) # falls back
|
|
assert "duckduckgo" in provider_names()
|
|
|
|
|
|
def test_build_provider_third_party_requires_key():
|
|
with pytest.raises(ValueError):
|
|
build_provider("tavily") # no key
|
|
assert isinstance(build_provider("tavily", "tvly-x"), TavilyProvider)
|
|
assert isinstance(build_provider("brave", "brv-x"), BraveProvider)
|
|
|
|
|
|
def test_tool_surfaces_missing_key_error(tmp_path):
|
|
secrets = SecretStore(tmp_path / "secrets.json")
|
|
secrets.put("web_search:default", {"provider": "tavily"}) # no api_key
|
|
out = make_web_search_tool(secrets)(
|
|
query="q"
|
|
) # resolve_provider raises ValueError → error dict
|
|
assert "needs an API key" in out["error"]
|
|
|
|
|
|
def test_resolve_provider_from_secretstore(tmp_path):
|
|
from coworker.web import resolve_provider
|
|
|
|
secrets = SecretStore(tmp_path / "secrets.json")
|
|
secrets.put("web_search:default", {"provider": "tavily", "api_key": "tvly-123"})
|
|
p = resolve_provider(secrets)
|
|
assert p.name == "tavily" and p.api_key == "tvly-123"
|
|
|
|
|
|
def test_web_search_rest(tmp_path, monkeypatch):
|
|
from fastapi.testclient import TestClient
|
|
|
|
from coworker.server.app import create_app
|
|
from coworker.server.manager import SessionManager
|
|
|
|
monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
|
|
client = TestClient(create_app(SessionManager(data_dir=tmp_path / "data")))
|
|
|
|
assert client.get("/v1/web-search").json()["provider"] == "duckduckgo"
|
|
assert (
|
|
client.post(
|
|
"/v1/web-search", json={"provider": "tavily", "api_key": "sk-secret-xyz"}
|
|
).json()["ok"]
|
|
is True
|
|
)
|
|
got = client.get("/v1/web-search").json()
|
|
assert got["provider"] == "tavily" and got["has_key"] is True
|
|
assert (
|
|
"sk-secret-xyz" not in client.get("/v1/web-search").text
|
|
) # key never returned
|
|
assert (
|
|
client.post("/v1/web-search", json={"provider": "nope"}).json()["ok"] is False
|
|
)
|
|
|
|
|
|
def test_engine_registers_web_search(tmp_path):
|
|
from coworker.agent import build_engine
|
|
from coworker.agents import chat_agent
|
|
|
|
eng = build_engine(
|
|
agent=chat_agent(),
|
|
provider=_StubProvider(),
|
|
secrets=SecretStore(tmp_path / "s.json"),
|
|
)
|
|
assert "web_search" in eng.registry.names()
|
|
|
|
|
|
class _StubProvider:
|
|
def complete(self, **_kw):
|
|
from coworker.providers import AssistantTurn
|
|
|
|
return AssistantTurn()
|
|
|
|
def capabilities(self, _model):
|
|
from coworker.providers.base import ModelCapabilities
|
|
|
|
return ModelCapabilities()
|
|
|
|
def stream(self, **_kw):
|
|
from coworker.providers.base import StreamChunk
|
|
|
|
yield StreamChunk(turn=self.complete())
|