Fix tool output endpoint

BlueBrain · Nov 13, 2024 · e4a3cec · e4a3cec
1 parent 0a2fbd3
commit e4a3cec
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/swarm_copy/app/routers/tools.py b/swarm_copy/app/routers/tools.py
@@ -96,15 +96,16 @@ async def get_tool_returns(
         select(Messages)
         .where(
             Messages.thread_id == thread_id,
-            Messages.entity == Entity.AI_TOOL,
+            Messages.entity == Entity.TOOL,
         )
         .order_by(Messages.order)
     )
     tool_messages = messages_result.scalars().all()
 
     tool_output = []
     for msg in tool_messages:
-        if json.loads(msg.content).get("tool_call_id") == tool_call_id:
-            tool_output.append(json.loads(msg.content))
+        msg_content = json.loads(msg.content)
+        if msg_content.get("tool_call_id") == tool_call_id:
+            tool_output.append(msg_content)
 
     return tool_output