Update tools.py

BlueBrain · Dec 11, 2024 · ee4b97b · ee4b97b
1 parent 8aad2eb
commit ee4b97b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/neuroagent/app/routers/database/tools.py b/src/neuroagent/app/routers/database/tools.py
@@ -24,7 +24,7 @@ async def get_tool_calls(
     memory: Annotated[AsyncSqliteSaver | None, Depends(get_agent_memory)],
     thread_id: str,
     message_id: str,
-) -> list[dict[str, Any]]:
+) -> list[ToolCallSchema]:
     """Get tool calls of a specific message.
     \f
 
@@ -93,14 +93,14 @@ async def get_tool_calls(
     )
 
     # From sub list, extract tool calls
-    tool_calls: list[dict[str, Any]] = []
+    tool_calls: list[ToolCallSchema] = []
     for message in message_list[previous_content_message + 1 : relevant_message]:
         if isinstance(message, AIMessage):
             tool_calls.extend(
                 [
                     ToolCallSchema(
                         call_id=tool["id"], name=tool["name"], arguments=tool["args"]
-                    ).model_dump()
+                    )
                     for tool in message.tool_calls
                 ]
             )