diff --git a/needle/playground/server.py b/needle/playground/server.py index 6525fe6..019f585 100644 --- a/needle/playground/server.py +++ b/needle/playground/server.py @@ -28,20 +28,19 @@ def load(self): self.tools_json = "[]" def complete(self, tools_json, query): - from .. import Needle, _lib + from .. import Needle with self.lock: if self.agent is None or tools_json != self.tools_json: self.agent = Needle(tools=tools_json, weights=self.weights) self.tools_json = tools_json else: - _lib().needle_reset() + self.agent.reset() return self.agent.complete(query) def reset(self): - from .. import _lib with self.lock: if self.agent is not None: - _lib().needle_reset() + self.agent.close() self.agent = None self.tools_json = None diff --git a/tests/test_playground.py b/tests/test_playground.py new file mode 100644 index 0000000..2f438bf --- /dev/null +++ b/tests/test_playground.py @@ -0,0 +1,84 @@ +import pytest + + +class _Agent: + def __init__(self, tools, weights): + self.tools = tools + self.weights = weights + self.calls = [] + + def complete(self, query, max_new_tokens=256): + self.calls.append(("complete", query)) + return {"type": "text", "text": query} + + def reset(self): + self.calls.append("reset") + + def close(self): + self.calls.append("close") + + +@pytest.fixture +def agents(monkeypatch): + """Stand in for Needle, and fail the test if the global engine is touched.""" + import needle + + made = [] + + def build(tools=None, weights=None, **_kwargs): + agent = _Agent(tools, weights) + made.append(agent) + return agent + + def refuse(*_args, **_kwargs): + raise AssertionError("the playground reached the process-global engine") + + monkeypatch.setattr(needle, "Needle", build) + monkeypatch.setattr(needle, "_lib", refuse) + return made + + +def test_a_second_query_on_the_same_tools_rewinds_that_agent(agents): + from needle.playground.server import Engine + + engine = Engine() + engine.complete("[]", "one") + engine.complete("[]", "two") + + assert len(agents) == 1 + assert agents[0].calls == [("complete", "one"), "reset", ("complete", "two")] + + +def test_new_tools_build_a_fresh_agent_without_rewinding(agents): + from needle.playground.server import Engine + + engine = Engine() + engine.complete("[]", "one") + engine.complete('[{"name": "t"}]', "two") + + assert len(agents) == 2 + assert "reset" not in agents[0].calls + assert agents[1].calls == [("complete", "two")] + + +def test_reset_closes_the_agent_it_drops(agents): + from needle.playground.server import Engine + + engine = Engine() + engine.complete("[]", "one") + engine.reset() + + assert agents[0].calls[-1] == "close" + assert engine.agent is None + assert engine.tools_json is None + + +def test_a_tuned_agent_is_rewound_through_itself(agents): + from needle.playground.server import Engine + + engine = Engine(weights="tuned.cact") + engine.complete("[]", "one") + engine.complete("[]", "two") + + assert agents[0].weights == "tuned.cact" + assert "reset" in agents[0].calls