WIP Stash

backend/src/agents/expert.py

@@ -1,6 +1,8 @@
 import inference
 from agent import Agent, AgentField
 
+import asyncio
+
 CLASS_DESCRIPTION = """
 Expert AI with customizable system prompt.
 Will get guidelines, documents, and other advanced features in the future.
@@ -22,16 +24,18 @@ class ExpertAgent(Agent):
     async def run(self, run):
         opts = {}
 
-        agent_tools = self.tools
-        print('agent tools ' + repr(agent_tools))
-        run.add_tools(agent_tools)
+        using_tools = False
+        if 'tools' not in run.input:
+            if self.tools:
+                run.input['tools'] = self.tools
+                using_tools = True
 
         if self.config.get("temperature") is not None:
             opts["temperature"] = self.config["temperature"]
         if self.config.get("repetition_penalty") is not None:
             opts["repetition_penalty"] = self.config["repetition_penalty"]
 
-        # 1. Build message list based on input
+        # Build message list based on input
         messages = inference.build_messages(
             system_prompt=self.config.get("system_prompt", ""),
             instruction=run.input.get("instruction"),
@@ -39,14 +43,44 @@ class ExpertAgent(Agent):
             messages=run.input.get("messages"),
         )
 
-        # 2. Call inference, passing the run to feed run.stream and tool calls
-        await inference.chat_completion(
-            self.app,
-            model_id=self.config.get("model"),
-            messages=messages,
-            tools=run.input.get("tools", []),
-            opts=opts,
-            agent_run=run,
-        )
+        # Collect inference output. These are resused across calls
+        tool_calls = []
+        output_stream = None
+        inference_end = asyncio.Event()
+        def on_output(stream):
+            nonlocal output_stream
+            output_stream = stream
+        run.on('output', on_output)
+        run.on('inference-end', lambda event: inference_end.set())
+        run.on('tool-call', lambda tc: tool_calls.append(tc))
 
-        run.get_output_stream().listen(lambda chunk: None, lambda: run.close())
+        while True:
+            # Reset data
+            inference_end.clear()
+            tool_calls = []
+            output_stream = None
+
+            # Call inference, passing the run to get events out
+            await inference.chat_completion(
+                self.app,
+                model_id=self.config.get("model"),
+                messages=messages,
+                tools=run.input.get("tools", []),
+                opts=opts,
+                agent_run=run,
+            )
+
+            await inference_end.wait()
+
+            messages.append(inference.LlamaMessage('assistant', output_stream and await output_stream.content_wait() or ""))
+
+            if using_tools and tool_calls:
+                for tc in tool_calls:
+                    tool = next(t for t in self.tools if t.name == tc.tool_name)
+                    if not tool:
+                        print(f"Bad agent tool name: {tc.name}")
+                        continue
+                    tc.result = await tool.execute_async(**tc.arguments)
+                    messages.append(inference.LlamaMessage('tool', tc))
+            else:
+                break