feat: dispatch the tool loop over web_search and load_skill

2 files changed, 259 insertions, 18 deletions

diff --git a/llamachat/backend.py b/llamachat/backend.py
index 338f396..686aa6f 100644
--- a/llamachat/backend.py
+++ b/llamachat/backend.py
@@ -28,6 +28,7 @@ from PIL import Image
 
 from . import search
 from . import providers as providers_mod
+from . import skills
 
 THUMB_SIZE = (128, 128)
 
@@ -84,6 +85,14 @@ class SearchConfig:
 
 
 @dataclass
+class SkillsConfig:
+    """What stream_chat needs to offer and run skill loading."""
+    enabled: bool = False
+    store: skills.SkillStore | None = None
+    max_searches: int = 2  # tool rounds per turn, same cap as search
+
+
+@dataclass
 class Attachment:
     """A file the user attached, ready for both the API and the database."""
     path: Path
@@ -276,7 +285,11 @@ class Client:
         return str((choices[0].get("message") or {}).get("content") or "")
 
     def stream_chat(
-        self, model: str, messages: list[dict], search_cfg: "SearchConfig | None" = None
+        self,
+        model: str,
+        messages: list[dict],
+        search_cfg: "SearchConfig | None" = None,
+        skills_cfg: "SkillsConfig | None" = None,
     ):
         """Yield (kind, text) pairs as the model produces them.
 
@@ -285,25 +298,30 @@ class Client:
         separate `reasoning_content` delta field, which is what lets the UI
         keep the two apart.
 
-        With `search_cfg`, the model is offered a `web_search` tool and this
-        becomes a loop: a round that ends in a tool call is searched for, the
-        result appended, and the request re-sent. Two further kinds are
-        yielded around each search, 'search_start' and 'search_done'. After
-        `max_searches` rounds the tool is withdrawn, which forces an answer.
+        With `search_cfg`, the model is offered a `web_search` tool; with
+        `skills_cfg`, a `load_skill` tool. Either turns this into a loop: a
+        round that ends in a tool call is handled, the result appended, and
+        the request re-sent. Two further kinds are yielded around each search,
+        'search_start' and 'search_done', and one around each loaded skill,
+        'skill_loaded'. After the round cap the tools are withdrawn, which
+        forces an answer.
 
         Loading a different model makes the router unload the previous one,
         so the first chunk can take several seconds. That wait happens
         inside the initial `stream()` call.
         """
-        if search_cfg is None or not search_cfg.enabled:
+        search_on = search_cfg is not None and search_cfg.enabled
+        skills_on = skills_cfg is not None and skills_cfg.enabled
+        if not (search_on or skills_on):
             yield from self._stream_once(model, messages, tools=None)
             return
 
         # The loop mutates its own copy; the caller's list is history.
         convo = list(messages)
-        for round_number in range(search_cfg.max_searches + 1):
-            last_round = round_number == search_cfg.max_searches
-            tools = None if last_round else [search.TOOL_SCHEMA]
+        cap = search_cfg.max_searches if search_on else skills_cfg.max_searches
+        for round_number in range(cap + 1):
+            last_round = round_number == cap
+            tools = None if last_round else self._tool_schemas(search_cfg, skills_cfg)
             calls: dict = {}
             saw_tool_finish = False
             reasoning = ""
@@ -318,11 +336,7 @@ class Client:
                         reasoning += piece
                     yield (kind, piece)
 
-            wanted = [
-                call
-                for call in calls.values()
-                if call["name"] == search.TOOL_SCHEMA["function"]["name"]
-            ]
+            wanted = [call for call in calls.values() if call.get("name")]
             if not (saw_tool_finish and wanted) or last_round:
                 return
 
@@ -352,10 +366,12 @@ class Client:
             convo.append(assistant)
             # The round after this one is the last, so its results are the
             # last the model will get. Saying so is what makes it answer
-            # instead of asking for a search it can no longer make.
-            final = round_number + 1 == search_cfg.max_searches
+            # instead of asking for a tool it can no longer use.
+            final = round_number + 1 == cap
             for call in wanted:
-                yield from self._run_search(convo, call, search_cfg, final)
+                yield from self._run_tool(
+                    convo, call, search_cfg, skills_cfg, final
+                )
 
     def _run_search(
         self, convo: list[dict], call: dict, cfg: "SearchConfig", last: bool = False
@@ -396,6 +412,70 @@ class Client:
             json.dumps({"query": query, "results": results, "error": ""}),
         )
 
+    def _tool_schemas(self, search_cfg, skills_cfg):
+        """The tools to offer this round, or None when there are none."""
+        tools = []
+        if search_cfg is not None and search_cfg.enabled:
+            tools.append(search.TOOL_SCHEMA)
+        if (
+            skills_cfg is not None and skills_cfg.enabled
+            and skills_cfg.store is not None
+        ):
+            tools.append(skills_cfg.store.tool_schema())
+        return tools or None
+
+    def _run_tool(self, convo, call, search_cfg, skills_cfg, last=False):
+        """Route one tool call to its handler, or answer that it is unknown."""
+        name = call["name"]
+        if (
+            search_cfg is not None and search_cfg.enabled
+            and name == search.TOOL_SCHEMA["function"]["name"]
+        ):
+            yield from self._run_search(convo, call, search_cfg, last)
+            return
+        if (
+            skills_cfg is not None and skills_cfg.enabled
+            and skills_cfg.store is not None
+            and name == skills.LOAD_SKILL_NAME
+        ):
+            yield from self._run_skill(convo, call, skills_cfg.store, last)
+            return
+        # A tool call for a name that exists nowhere still needs an answer,
+        # or the next request is rejected for an unanswered call.
+        convo.append(
+            {
+                "role": "tool",
+                "tool_call_id": call["id"],
+                "content": f"error: unknown tool '{name}'",
+            }
+        )
+
+    def _run_skill(self, convo, call, store, last=False):
+        """Load a skill's text and append it to `convo`.
+
+        Every failure still appends a tool message, so the turn completes
+        either way. Success additionally yields ("skill_loaded", name) so
+        the UI can persist the load for later turns.
+        """
+        name = skills.parse_name(call["arguments"])
+        if not name:
+            convo.append(
+                skills.tool_message(
+                    call["id"], None, error="Malformed skill request", last=last
+                )
+            )
+            return
+        skill = store.load(name)
+        if skill is None:
+            convo.append(
+                skills.tool_message(
+                    call["id"], None, error=f"No such skill: {name}", last=last
+                )
+            )
+            return
+        convo.append(skills.tool_message(call["id"], skill, last=last))
+        yield ("skill_loaded", name)
+
     def _stream_once(self, model: str, messages: list[dict], tools: list | None):
         """One request/response round, yielding parsed SSE pairs."""
         body = {
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()