mirror of
https://github.com/cactus-compute/needle.git
synced 2026-10-21 06:09:38 +08:00
Raj Ojha
c11ccba41e
agent.extract drops the system facts the agent was initialized with because it forwarded only weights and strict to the module-level extract function. This prevented relative language from resolving and caused temporal validation to fail on licensed relative dates. Forward self._system_text and self._generation to extract, and test that custom facts, auto-injected date lines, and weights/strict/generation are all preserved. Signed-off-by: Rajojha-1 <rajojha047@gmail.com>
403 lines
13 KiB
Python
403 lines
13 KiB
Python
import json
|
|
import os
|
|
import warnings
|
|
|
|
import pytest
|
|
|
|
ENVELOPE = json.dumps({"type": "call", "confidence": 0.9,
|
|
"function_calls": [{"name": "City",
|
|
"arguments": {"city": "Paris"}}]}).encode("utf-8")
|
|
|
|
|
|
class _Stub:
|
|
def __init__(self, calls):
|
|
self.calls = calls
|
|
|
|
def needle_init(self, system, tools, index):
|
|
self.calls.append("init")
|
|
return 0
|
|
|
|
def needle_load(self, blob, size):
|
|
self.calls.append("load")
|
|
return 0
|
|
|
|
def needle_complete(self, text, *args):
|
|
self.calls.append("complete")
|
|
buffer = args[-2]
|
|
buffer.value = ENVELOPE
|
|
return 0
|
|
|
|
def needle_reset(self):
|
|
self.calls.append("reset")
|
|
|
|
|
|
class _WorkerStub:
|
|
next_pid = 1000
|
|
|
|
def __init__(self, calls, library, weights, system, tools, tool_index,
|
|
buffer_size, generation=2):
|
|
self.calls = calls
|
|
self.pid = type(self).next_pid
|
|
type(self).next_pid += 1
|
|
self.calls.append(("worker_start", self.pid, library, weights))
|
|
|
|
def complete(self, text, max_new_tokens):
|
|
self.calls.append(("worker_complete", self.pid, text))
|
|
return ENVELOPE.decode("utf-8")
|
|
|
|
def embed(self, text):
|
|
self.calls.append(("worker_embed", self.pid, text))
|
|
return [0.6, 0.8]
|
|
|
|
def reset(self):
|
|
self.calls.append(("worker_reset", self.pid))
|
|
|
|
def close(self):
|
|
self.calls.append(("worker_close", self.pid))
|
|
|
|
|
|
@pytest.fixture
|
|
def engine(monkeypatch, tmp_path):
|
|
import needle
|
|
|
|
calls = []
|
|
stubs = {2: _Stub(calls), 3: _Stub(calls)}
|
|
monkeypatch.setattr(needle, "_lib", lambda generation=2: stubs[generation])
|
|
monkeypatch.setattr(needle, "_library_path",
|
|
lambda generation=2: f"/tmp/libneedle{generation}")
|
|
monkeypatch.setattr(
|
|
needle, "FineTuneWorker",
|
|
lambda library, weights, system, tools, tool_index, buffer_size,
|
|
generation=2:
|
|
_WorkerStub(calls, library, weights, system, tools, tool_index,
|
|
buffer_size, generation))
|
|
monkeypatch.setattr(needle, "_active", {})
|
|
monkeypatch.setattr(needle, "_loaded_base", {})
|
|
base = tmp_path / "needle3.cact"
|
|
base.write_bytes((0x05E12A84).to_bytes(4, "little") + b"base weights")
|
|
monkeypatch.setattr(needle, "_base_weights_path", lambda generation: str(base))
|
|
return calls
|
|
|
|
|
|
@pytest.fixture
|
|
def tuned(tmp_path):
|
|
path = tmp_path / "tuned.cact"
|
|
path.write_bytes((0x05E12A83).to_bytes(4, "little") + b"tuned weights")
|
|
return str(path)
|
|
|
|
|
|
@pytest.fixture
|
|
def tuned_v3(tmp_path):
|
|
path = tmp_path / "tuned-v3.cact"
|
|
path.write_bytes((0x05E12A84).to_bytes(4, "little") + b"tuned weights")
|
|
return str(path)
|
|
|
|
|
|
def _tuned_agent(path):
|
|
import needle
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
return needle.Needle(tools="[]", weights=path)
|
|
|
|
|
|
def test_generation_3_without_weights_loads_the_base_archive_in_process(engine):
|
|
import needle
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
agent = needle.Needle(tools="[]", generation=3)
|
|
assert agent._generation == 3 and agent._tuned is False and agent._weights is None
|
|
response = agent.complete("hello")
|
|
assert response.get("confidence", "kept") is not None
|
|
assert engine.count("load") == 1 and "complete" in engine
|
|
assert not any(call[0] == "worker_start" for call in engine if isinstance(call, tuple))
|
|
needle.Needle(tools="[]", generation=3).complete("again")
|
|
assert engine.count("load") == 1
|
|
assert agent._track_props()["tuned"] is False
|
|
|
|
|
|
def test_default_agent_is_needle_3(engine):
|
|
import needle
|
|
|
|
agent = needle.Needle(tools="[]")
|
|
assert agent._generation == 3 and agent._weights is None and agent._tuned is False
|
|
agent.complete("hello")
|
|
assert engine.count("load") == 1 and "complete" in engine
|
|
|
|
|
|
def test_generation_2_runs_its_embedded_engine(engine):
|
|
import needle
|
|
|
|
agent = needle.Needle(tools="[]", generation=2)
|
|
assert agent._generation == 2 and agent._weights is None
|
|
agent.complete("hello")
|
|
assert engine.count("load") == 0 and "complete" in engine
|
|
|
|
|
|
def test_base_agent_and_tuned_worker_coexist(engine, tuned):
|
|
import needle
|
|
|
|
tuned_agent = _tuned_agent(tuned)
|
|
base = needle.Needle(tools="[]", generation=2)
|
|
tuned_agent.complete("tuned")
|
|
base.complete("base")
|
|
assert any(call[0] == "worker_complete" for call in engine if isinstance(call, tuple))
|
|
assert "complete" in engine
|
|
|
|
|
|
def test_tuned_worker_does_not_rebind_existing_base_agent(engine, tuned):
|
|
import needle
|
|
|
|
base = needle.Needle(tools="[]", generation=2)
|
|
base.complete("hello")
|
|
tuned_agent = _tuned_agent(tuned)
|
|
tuned_agent.complete("hello")
|
|
base.complete("again")
|
|
assert engine.count("load") == 0
|
|
assert engine.count("complete") == 2
|
|
|
|
|
|
def test_tuned_agent_loads_one_worker_and_reuses_it(engine, tuned):
|
|
agent = _tuned_agent(tuned)
|
|
agent.complete("one")
|
|
agent.complete("two")
|
|
starts = [call for call in engine
|
|
if isinstance(call, tuple) and call[0] == "worker_start"]
|
|
assert len(starts) == 1
|
|
assert len([call for call in engine
|
|
if isinstance(call, tuple) and call[0] == "worker_complete"]) == 2
|
|
|
|
|
|
def test_extract_uses_an_explicit_tuned_worker(engine, tuned):
|
|
import needle
|
|
|
|
schema = {"name": "City", "parameters": {"type": "object",
|
|
"properties": {"city": {"type": "string"}}}}
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
assert needle.extract("the weather in Paris", schema, weights=tuned) == {
|
|
"city": "Paris"}
|
|
assert any(call[0] == "worker_start" and call[3] == tuned
|
|
for call in engine if isinstance(call, tuple))
|
|
|
|
|
|
def test_v2_and_v3_weights_route_to_independent_engines(engine, tuned, tuned_v3,
|
|
monkeypatch):
|
|
import needle
|
|
|
|
routed = []
|
|
monkeypatch.setattr(needle, "_library_path",
|
|
lambda generation=2: routed.append(generation) or
|
|
f"/tmp/libneedle{generation}")
|
|
|
|
v2 = _tuned_agent(tuned)
|
|
v3 = _tuned_agent(tuned_v3)
|
|
assert v2._generation == 2
|
|
assert v3._generation == 3
|
|
assert routed == [2, 3]
|
|
|
|
|
|
def test_v3_routes_embeddings_to_its_worker(engine, tuned_v3):
|
|
agent = _tuned_agent(tuned_v3)
|
|
assert agent.embed("schema") == [0.6, 0.8]
|
|
assert ("worker_embed", agent._worker.pid, "schema") in engine
|
|
|
|
|
|
def test_unknown_archive_tag_is_rejected_before_loading(engine, tmp_path):
|
|
path = tmp_path / "unknown.cact"
|
|
path.write_bytes(b"NOPE")
|
|
with pytest.raises(RuntimeError, match="unknown .cact format tag"):
|
|
_tuned_agent(str(path))
|
|
|
|
|
|
def test_two_finetunes_use_independent_workers(engine, tuned, tmp_path):
|
|
other = tmp_path / "other.cact"
|
|
other.write_bytes((0x05E12A83).to_bytes(4, "little") + b"other weights")
|
|
|
|
first = _tuned_agent(tuned)
|
|
second = _tuned_agent(str(other))
|
|
assert first._worker.pid != second._worker.pid
|
|
first.complete("one")
|
|
second.complete("two")
|
|
assert ("worker_complete", first._worker.pid, "one") in engine
|
|
assert ("worker_complete", second._worker.pid, "two") in engine
|
|
first.close()
|
|
second.close()
|
|
|
|
|
|
def test_extraction_rejects_fabricated_temporal_year():
|
|
import datetime
|
|
import needle
|
|
import pydantic
|
|
|
|
class Invoice(pydantic.BaseModel):
|
|
due_date: datetime.datetime
|
|
|
|
response = {"validation": {"ungrounded": ["Invoice.due_date"]}}
|
|
correct = {"due_date": "2034-09-05T00:00:00Z"}
|
|
needle._validate_extraction("due on 5th September 2034", Invoice,
|
|
correct, response)
|
|
|
|
fabricated = {"due_date": "2024-09-05T00:00:00Z"}
|
|
with pytest.raises(needle.ExtractionValidationError, match="due_date"):
|
|
needle._validate_extraction("due on 5th September 2034", Invoice,
|
|
fabricated, response)
|
|
|
|
with pytest.raises(needle.ExtractionValidationError, match="due_date"):
|
|
needle._validate_extraction(
|
|
"due on 5th September 42", Invoice,
|
|
{"due_date": "2026-09-05T00:00:00Z"}, response)
|
|
|
|
assert needle._source_years("due September 5, 2034") == {2034}
|
|
assert needle._source_years("due September 5") == set()
|
|
assert needle._source_years("due 5th September 42") == {42}
|
|
assert needle._source_years("Invoice 42 is due tomorrow at 5") == set()
|
|
|
|
|
|
def test_tool_index_path_takes_a_path_object(engine, tmp_path):
|
|
import needle
|
|
|
|
index = tmp_path / "tools.idx"
|
|
agent = needle.Needle(tools="[]", tool_index_path=index)
|
|
|
|
assert agent._tool_index_path == os.fspath(index).encode("utf-8")
|
|
|
|
|
|
def test_tool_index_path_agrees_between_the_base_and_tuned_paths(engine, tuned, tmp_path):
|
|
import needle
|
|
|
|
index = tmp_path / "tools.idx"
|
|
base = needle.Needle(tools="[]", tool_index_path=index)
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
needle.Needle(tools="[]", weights=tuned, tool_index_path=index)
|
|
|
|
started = [call for call in engine
|
|
if isinstance(call, tuple) and call[0] == "worker_start"]
|
|
assert started
|
|
assert base._tool_index_path.decode("utf-8") == os.fspath(index)
|
|
|
|
|
|
def test_platform_tuned_archive_keeps_confidence(engine, tmp_path, fake_cact):
|
|
import needle
|
|
|
|
fake_cact(tmp_path / "platform.cact", [1, 2])
|
|
fake_cact(tmp_path / "local.cact", [1])
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
calibrated = needle.Needle(tools="[]", weights=str(tmp_path / "platform.cact"))
|
|
assert calibrated.complete("hello").get("confidence") is not None
|
|
with pytest.warns(UserWarning, match="confidence as None"):
|
|
uncalibrated = needle.Needle(tools="[]", weights=str(tmp_path / "local.cact"))
|
|
assert uncalibrated.complete("hello").get("confidence", "kept") is None
|
|
|
|
|
|
def test_stateless_agent_resets_before_every_query(engine):
|
|
import needle
|
|
|
|
agent = needle.Needle(tools="[]", generation=3, stateless=True)
|
|
agent.complete("one")
|
|
agent.complete("two")
|
|
resets = [c for c in engine if c == "reset"]
|
|
assert len(resets) == 2
|
|
|
|
|
|
def test_unreset_queries_warn_once(engine):
|
|
import needle
|
|
|
|
agent = needle.Needle(tools="[]", generation=3)
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
for _ in range(needle.UNRESET_TURNS):
|
|
agent.complete("q")
|
|
with pytest.warns(UserWarning, match="without reset"):
|
|
agent.complete("q")
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
agent.complete("q")
|
|
agent.reset()
|
|
for _ in range(needle.UNRESET_TURNS):
|
|
agent.complete("q")
|
|
|
|
|
|
def test_tool_result_turns_do_not_count_as_queries(engine):
|
|
import needle
|
|
|
|
calls = []
|
|
|
|
@needle.tool
|
|
def ping(x: str):
|
|
"Ping."
|
|
calls.append(x)
|
|
return {"ok": True}
|
|
|
|
agent = needle.Needle(tools=[ping], generation=3)
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error")
|
|
for _ in range(needle.UNRESET_TURNS):
|
|
agent.run("go")
|
|
|
|
|
|
def test_agent_extract_carries_its_own_system_facts(engine, monkeypatch):
|
|
import needle
|
|
|
|
seen = {}
|
|
|
|
def spy(text, schema, system=None, max_new_tokens=512, weights=None, strict=True, generation=None):
|
|
seen.update(text=text, system=system, weights=weights, strict=strict, generation=generation)
|
|
return None
|
|
|
|
monkeypatch.setattr(needle, "extract", spy)
|
|
facts = "date: 2026-07-21 Tue 14:30; locale: en-US"
|
|
agent = needle.Needle(tools="[]", system=facts)
|
|
|
|
assert agent.extract("dinner tomorrow at 7", {"type": "object"}) is None
|
|
assert seen["system"] == facts
|
|
assert seen["text"] == "dinner tomorrow at 7"
|
|
|
|
|
|
def test_agent_extract_without_system_facts_sends_none_when_autodate_disabled(engine, monkeypatch):
|
|
import needle
|
|
|
|
seen = {}
|
|
monkeypatch.setattr(
|
|
needle, "extract",
|
|
lambda text, schema, system=None, **kwargs: seen.update(system=system))
|
|
needle.Needle(tools="[]", auto_date=False).extract("anything", {"type": "object"})
|
|
|
|
assert seen["system"] is None
|
|
|
|
|
|
def test_agent_extract_without_system_facts_carries_date_line(engine, monkeypatch):
|
|
import needle
|
|
|
|
seen = {}
|
|
monkeypatch.setattr(
|
|
needle, "extract",
|
|
lambda text, schema, system=None, **kwargs: seen.update(system=system))
|
|
needle.Needle(tools="[]").extract("anything", {"type": "object"})
|
|
|
|
assert seen["system"] is not None
|
|
assert seen["system"].startswith("date: ")
|
|
|
|
|
|
def test_agent_extract_still_carries_weights_strict_and_generation(engine, tuned, monkeypatch):
|
|
import needle
|
|
|
|
seen = {}
|
|
monkeypatch.setattr(
|
|
needle, "extract",
|
|
lambda text, schema, system=None, max_new_tokens=512, weights=None, strict=True, generation=None:
|
|
seen.update(system=system, weights=weights, strict=strict, generation=generation))
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
agent = needle.Needle(tools="[]", weights=tuned, system="device: phone", generation=3)
|
|
agent.extract("anything", {"type": "object"}, strict=False)
|
|
|
|
assert seen["weights"] == tuned
|
|
assert seen["strict"] is False
|
|
assert seen["generation"] == agent._generation
|
|
assert "device: phone" in seen["system"]
|
|
assert seen["system"].startswith("date: ")
|