working demo DIY

BlueBrain · Dec 2, 2024 · d2f918a · d2f918a
1 parent 8f07f79
commit d2f918a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/swarm_copy/run.py b/swarm_copy/run.py
@@ -371,3 +371,24 @@ async def astream(
             ),
             [],
         )
+
+        # Hack to return a tool message as a user message to have a nice plot.
+        if "scsgetone-tool" == history[-2].get("tool_name"):
+            json_data = json.loads(history[-2]["content"])
+            content_as_str = json.dumps(list(json_data["results"].values()))
+            fake_llm_message = {
+                "content": "PLOT_MESSAGE" + content_as_str,
+                "sender": "Agent",
+                "role": "assistant",
+                "function_call": None,
+                "tool_calls": None,
+            }
+            history.append(fake_llm_message)
+            yield (
+                Response(
+                    messages=history[init_len - 1 :],
+                    agent=active_agent,
+                    context_variables=context_variables,
+                ),
+                "PLOT",
+            )
diff --git a/swarm_copy/stream.py b/swarm_copy/stream.py
@@ -39,7 +39,7 @@ async def stream_agent_response(
             yield chunk
         # Final chunk that contains the whole response
         else:
-            if chunk[1]:
+            if chunk[1] and chunk[1] != "PLOT":
                 time.sleep(0.1)
                 tool_valid_list = json.dumps([tool.model_dump() for tool in chunk[1]])
                 yield "\n<requires_human_approval>\n" + tool_valid_list

diff --git a/swarm_copy/tools/bluenaas_scs_getone.py b/swarm_copy/tools/bluenaas_scs_getone.py
@@ -45,9 +45,11 @@ async def arun(self) -> SimulationDetailsResponse:
             f"Running SCSGetOne tool with inputs {self.input_schema.model_dump()}"
         )
 
+        # list(response.json()["results"].values())[0][0]["x"]
+
         response = await self.metadata.httpx_client.get(
             url=f"{self.metadata.bluenaas_url}/simulation/single-neuron/{self.metadata.vlab_id}/{self.metadata.project_id}/{self.input_schema.simulation_id}",
             headers={"Authorization": f"Bearer {self.metadata.token}"},
         )
 
-        return SimulationDetailsResponse(**response.json())
+        return SimulationDetailsResponse(**response.json()).model_dump_json()