autogen/python/packages/autogen-studio/autogenstudio/utils/utils.py

72 lines
2.9 KiB
Python

import base64
from typing import Sequence
from autogen_agentchat.messages import ChatMessage, MultiModalMessage, TextMessage
from autogen_core import Image
from autogen_core.models import UserMessage
from loguru import logger
def construct_task(query: str, files: list[dict] | None = None) -> Sequence[ChatMessage]:
"""
Construct a task from a query string and list of files.
Returns a list of ChatMessage objects suitable for processing by the agent system.
Args:
query: The text query from the user
files: List of file objects with properties name, content, and type
Returns:
List of BaseChatMessage objects (TextMessage, MultiModalMessage)
"""
if files is None:
files = []
messages = []
# Add the user's text query as a TextMessage
if query:
messages.append(TextMessage(source="user", content=query))
# Process each file based on its type
for file in files:
try:
if file.get("type", "").startswith("image/"):
# Handle image file using from_base64 method
# The content is already base64 encoded according to the convertFilesToBase64 function
image = Image.from_base64(file["content"])
messages.append(
MultiModalMessage(
source="user", content=[image], metadata={"filename": file.get("name", "unknown.img")}
)
)
elif file.get("type", "").startswith("text/"):
# Handle text file as TextMessage
text_content = base64.b64decode(file["content"]).decode("utf-8")
messages.append(
TextMessage(
source="user", content=text_content, metadata={"filename": file.get("name", "unknown.txt")}
)
)
else:
# Log unsupported file types but still try to process based on best guess
logger.warning(f"Potentially unsupported file type: {file.get('type')} for file {file.get('name')}")
if file.get("type", "").startswith("application/"):
# Try to treat as text if it's an application type (like JSON)
text_content = base64.b64decode(file["content"]).decode("utf-8")
messages.append(
TextMessage(
source="user",
content=text_content,
metadata={
"filename": file.get("name", "unknown.file"),
"filetype": file.get("type", "unknown"),
},
)
)
except Exception as e:
logger.error(f"Error processing file {file.get('name')}: {str(e)}")
# Continue processing other files even if one fails
return messages