diff options
Diffstat (limited to 'test_llamachat.py')
| -rwxr-xr-x | test_llamachat.py | 161 |
1 files changed, 161 insertions, 0 deletions
diff --git a/test_llamachat.py b/test_llamachat.py index 81a430b..0c7d2d0 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -2763,6 +2763,28 @@ def _tool_round(query: str = "q"): ] +def _skill_round(name: str = "firefly-cli"): + """One streamed round that asks to load a skill.""" + return [ + ( + "tool_calls", + json.dumps( + [ + { + "index": 0, + "id": "call_2", + "function": { + "name": "load_skill", + "arguments": json.dumps({"name": name}), + }, + } + ] + ), + ), + ("tool_finish", ""), + ] + + def test_search_loop_cap(): import urllib.request @@ -3265,6 +3287,142 @@ def test_final_round_note(): print("ok final round note") +def test_skill_tool_round(): + from llamachat import skills + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write_skills(root, {"firefly-cli/SKILL.md": SKILLS_SAMPLE["firefly-cli/SKILL.md"]}) + store = skills.SkillStore(root) + + client = backend.Client("http://x") + seen_messages = [] + sent_tools = [] + + def scripted(model, messages, tools): + sent_tools.append(tools) + seen_messages.append([dict(m) for m in messages]) + if len(seen_messages) == 1: + yield from _skill_round() + else: + yield ("content", "answered") + + client._stream_once = scripted + cfg = backend.SkillsConfig(enabled=True, store=store, max_searches=1) + out = list(client.stream_chat( + "m", [{"role": "user", "content": "hi"}], skills_cfg=cfg + )) + + # The load is reported and the answer flows through. + assert ("skill_loaded", "firefly-cli") in out, out + assert ("content", "answered") in out, out + + # The offered schema is the load_skill one; the final round has none. + assert sent_tools[0] == [store.tool_schema()], sent_tools[0] + assert sent_tools[-1] is None, sent_tools + + # The second request carries the assistant tool call and a tool + # reply with the skill body. + second = seen_messages[1] + assistant = next( + m for m in second if m["role"] == "assistant" and m.get("tool_calls") + ) + assert assistant["tool_calls"][0]["function"]["name"] == "load_skill" + tool_msg = next(m for m in second if m["role"] == "tool") + assert tool_msg["tool_call_id"] == "call_2" + assert "run `firefly auth test` first" in tool_msg["content"].lower() + assert "[skill: firefly-cli]" in tool_msg["content"] + print("ok skill tool round") + + +def test_skill_tool_errors(): + from llamachat import skills + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write_skills(root, {"firefly-cli/SKILL.md": SKILLS_SAMPLE["firefly-cli/SKILL.md"]}) + store = skills.SkillStore(root) + client = backend.Client("http://x") + cfg = backend.SkillsConfig(enabled=True, store=store, max_searches=1) + + # Unknown skill name -> error tool message, no skill_loaded yield. + rounds = [_skill_round("vault-librarian"), [("content", "done")]] + seen = [] + + def scripted(model, messages, tools): + seen.append(messages) + yield from rounds[min(len(seen) - 1, len(rounds) - 1)] + + client._stream_once = scripted + out = list(client.stream_chat( + "m", [{"role": "user", "content": "hi"}], skills_cfg=cfg + )) + assert ("content", "done") in out, out + assert not any(k == "skill_loaded" for k, _ in out), out + tool_msg = next(m for m in seen[1] if m["role"] == "tool") + assert "No such skill: vault-librarian" in tool_msg["content"] + + # Malformed arguments are answered rather than left dangling. + bad = {"id": "call_9", "name": "load_skill", "arguments": "{not json"} + convo: list = [] + emitted = list(client._run_skill(convo, bad, store)) + assert emitted == [], emitted + assert convo[-1]["role"] == "tool" + assert "error" in convo[-1]["content"] + + # A tool call naming a tool that does not exist is answered too. + rounds2 = [[ + ("tool_calls", json.dumps([{ + "index": 0, "id": "call_7", + "function": {"name": "no_such_tool", "arguments": "{}"}, + }])), + ("tool_finish", ""), + ], [("content", "done")]] + seen2 = [] + + def scripted2(model, messages, tools): + seen2.append(messages) + yield from rounds2[min(len(seen2) - 1, len(rounds2) - 1)] + + client._stream_once = scripted2 + out2 = list(client.stream_chat( + "m", [{"role": "user", "content": "hi"}], skills_cfg=cfg + )) + assert ("content", "done") in out2, out2 + unknown = next(m for m in seen2[1] if m["role"] == "tool") + assert "no_such_tool" in unknown["content"] + print("ok skill tool errors") + + +def test_skill_round_cap(): + from llamachat import skills + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write_skills(root, {"firefly-cli/SKILL.md": SKILLS_SAMPLE["firefly-cli/SKILL.md"]}) + store = skills.SkillStore(root) + + client = backend.Client("http://x") + sent_tools = [] + + def always_skill_calls(model, messages, tools): + sent_tools.append(tools) + yield from _skill_round() + + client._stream_once = always_skill_calls + cfg = backend.SkillsConfig(enabled=True, store=store, max_searches=1) + list(client.stream_chat( + "m", [{"role": "user", "content": "hi"}], skills_cfg=cfg + )) + + # max_searches rounds offer the tool, then a final round without it, + # even when search is entirely disabled (skills-only loop path). + assert len(sent_tools) == 2, sent_tools + assert sent_tools[0] == [store.tool_schema()], sent_tools[0] + assert sent_tools[-1] is None, "the final round must withdraw the tool" + print("ok skill round cap") + + def test_title_cleaning(): clean = backend.clean_title @@ -3582,6 +3740,9 @@ if __name__ == "__main__": test_date_note() test_search_html() test_final_round_note() + test_skill_tool_round() + test_skill_tool_errors() + test_skill_round_cap() test_title_cleaning() test_title_request() test_needs_title() |