2025-02-04 08:25:48 -08:00
|
|
|
from typing import List, Union
|
2025-02-04 06:55:04 -08:00
|
|
|
|
2025-02-04 08:25:48 -08:00
|
|
|
from autogen_core import FunctionCall, Image
|
|
|
|
from autogen_core.models import FunctionExecutionResult, LLMMessage, UserMessage
|
2025-03-26 16:19:52 -07:00
|
|
|
from pydantic import BaseModel
|
2025-02-04 06:55:04 -08:00
|
|
|
|
2025-02-04 08:25:48 -08:00
|
|
|
# Type aliases for convenience
|
2025-03-26 16:19:52 -07:00
|
|
|
_StructuredContent = BaseModel
|
2025-02-04 08:25:48 -08:00
|
|
|
_UserContent = Union[str, List[Union[str, Image]]]
|
|
|
|
_AssistantContent = Union[str, List[FunctionCall]]
|
|
|
|
_FunctionExecutionContent = List[FunctionExecutionResult]
|
|
|
|
_SystemContent = str
|
2025-02-04 06:55:04 -08:00
|
|
|
|
2025-02-04 08:25:48 -08:00
|
|
|
|
2025-03-26 16:19:52 -07:00
|
|
|
def content_to_str(
|
|
|
|
content: _UserContent | _AssistantContent | _FunctionExecutionContent | _SystemContent | _StructuredContent,
|
|
|
|
) -> str:
|
2025-02-04 08:25:48 -08:00
|
|
|
"""Convert the content of an LLMMessage to a string."""
|
2025-02-04 06:55:04 -08:00
|
|
|
if isinstance(content, str):
|
|
|
|
return content
|
2025-03-26 16:19:52 -07:00
|
|
|
elif isinstance(content, BaseModel):
|
|
|
|
return content.model_dump_json()
|
2025-02-04 06:55:04 -08:00
|
|
|
else:
|
|
|
|
result: List[str] = []
|
|
|
|
for c in content:
|
|
|
|
if isinstance(c, str):
|
|
|
|
result.append(c)
|
|
|
|
elif isinstance(c, Image):
|
|
|
|
result.append("<image>")
|
|
|
|
else:
|
2025-02-04 08:25:48 -08:00
|
|
|
result.append(str(c))
|
2025-02-04 06:55:04 -08:00
|
|
|
|
|
|
|
return "\n".join(result)
|
|
|
|
|
|
|
|
|
|
|
|
def remove_images(messages: List[LLMMessage]) -> List[LLMMessage]:
|
|
|
|
"""Remove images from a list of LLMMessages"""
|
|
|
|
str_messages: List[LLMMessage] = []
|
|
|
|
for message in messages:
|
|
|
|
if isinstance(message, UserMessage) and isinstance(message.content, list):
|
2025-02-04 08:25:48 -08:00
|
|
|
str_messages.append(UserMessage(content=content_to_str(message.content), source=message.source))
|
2025-02-04 06:55:04 -08:00
|
|
|
else:
|
|
|
|
str_messages.append(message)
|
|
|
|
return str_messages
|