fix invoke response

Author: ilackarms

python/packages/autogen-studio/autogenstudio/web/routes/invoke.py

@@ -3,9 +3,11 @@
 from fastapi import APIRouter
 from autogen_core import ComponentModel
 from autogenstudio.teammanager import TeamManager
-from autogenstudio.datamodel import Response
+from autogenstudio.datamodel import Response, TeamResult
+from autogen_agentchat.base import TaskResult
 from typing import Union
 from pydantic import BaseModel
+from autogen_agentchat.messages import BaseTextChatMessage
 
 router = APIRouter()
 team_manager = TeamManager()
@@ -19,8 +21,44 @@ async def invoke(request: InvokeTaskRequest):
     response = Response(message="Task successfully completed", status=True, data=None)
     try:
         result_message = await team_manager.run(task=request.task, team_config=request.team_config)
-        response.data = result_message
+        response.data = _format_team_result(result_message)
     except Exception as e:
         response.message = str(e)
         response.status = False
     return response
+
+def _format_team_result(team_result: TeamResult) -> dict:
+    """
+    Format the result from TeamResult to a dictionary.
+    """
+    formatted_result = {
+        "task_result": _format_task_result(team_result.task_result),
+        "usage": team_result.usage,
+        "duration": team_result.duration,
+    }
+    return formatted_result
+
+def _format_task_result(task_result: TaskResult) -> dict:
+    """
+    Format the result from TeamResult to a dictionary.
+    """
+    for message in task_result.messages:
+        print(f"Message: {message}, Type: {type(message)}")
+    formatted_result = {
+        "messages": [
+            _format_message(message) for message in task_result.messages
+        ],
+        "stop_reason": task_result.stop_reason,
+    }
+    return formatted_result
+
+def _format_message(message: BaseTextChatMessage) -> dict:
+    """
+    Format the message to a dictionary.
+    """
+    return {
+        "source": message.source,
+        "models_usage": message.models_usage,
+        "metadata": message.metadata,
+        "content": message.content,
+    }