format, add test

Author: victordibia

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py

@@ -32,6 +32,7 @@
 from .. import EVENT_LOGGER_NAME
 from ..base import Handoff as HandoffBase
 from ..base import Response
+from ..memory._base_memory import Memory
 from ..messages import (
     AgentEvent,
     ChatMessage,
@@ -44,7 +45,6 @@
 )
 from ..state import AssistantAgentState
 from ._base_chat_agent import BaseChatAgent
-from ..memory._base_memory import Memory
 
 event_logger = logging.getLogger(EVENT_LOGGER_NAME)
 
@@ -245,8 +245,7 @@ def __init__(
         name: str,
         model_client: ChatCompletionClient,
         *,
-        tools: List[Tool | Callable[..., Any] |
-                    Callable[..., Awaitable[Any]]] | None = None,
+        tools: List[Tool | Callable[..., Any] | Callable[..., Awaitable[Any]]] | None = None,
         handoffs: List[HandoffBase | str] | None = None,
         model_context: ChatCompletionContext | None = None,
         description: str = "An agent that provides assistance with ability to use tools.",
@@ -266,20 +265,19 @@ def __init__(
             elif isinstance(memory, list):
                 self._memory = memory
             else:
-                raise TypeError(
-                    f"Expected Memory, List[Memory], or None, got {type(memory)}")
+                raise TypeError(f"Expected Memory, List[Memory], or None, got {type(memory)}")
 
-        self._system_messages: List[SystemMessage | UserMessage |
-                                    AssistantMessage | FunctionExecutionResultMessage] = []
+        self._system_messages: List[
+            SystemMessage | UserMessage | AssistantMessage | FunctionExecutionResultMessage
+        ] = []
         if system_message is None:
             self._system_messages = []
         else:
             self._system_messages = [SystemMessage(content=system_message)]
         self._tools: List[Tool] = []
         if tools is not None:
             if model_client.model_info["function_calling"] is False:
-                raise ValueError(
-                    "The model does not support function calling.")
+                raise ValueError("The model does not support function calling.")
             for tool in tools:
                 if isinstance(tool, Tool):
                     self._tools.append(tool)
@@ -288,8 +286,7 @@ def __init__(
                         description = tool.__doc__
                     else:
                         description = ""
-                    self._tools.append(FunctionTool(
-                        tool, description=description))
+                    self._tools.append(FunctionTool(tool, description=description))
                 else:
                     raise ValueError(f"Unsupported tool type: {type(tool)}")
         # Check if tool names are unique.
@@ -301,22 +298,19 @@ def __init__(
         self._handoffs: Dict[str, HandoffBase] = {}
         if handoffs is not None:
             if model_client.model_info["function_calling"] is False:
-                raise ValueError(
-                    "The model does not support function calling, which is needed for handoffs.")
+                raise ValueError("The model does not support function calling, which is needed for handoffs.")
             for handoff in handoffs:
                 if isinstance(handoff, str):
                     handoff = HandoffBase(target=handoff)
                 if isinstance(handoff, HandoffBase):
                     self._handoff_tools.append(handoff.handoff_tool)
                     self._handoffs[handoff.name] = handoff
                 else:
-                    raise ValueError(
-                        f"Unsupported handoff type: {type(handoff)}")
+                    raise ValueError(f"Unsupported handoff type: {type(handoff)}")
         # Check if handoff tool names are unique.
         handoff_tool_names = [tool.name for tool in self._handoff_tools]
         if len(handoff_tool_names) != len(set(handoff_tool_names)):
-            raise ValueError(
-                f"Handoff names must be unique: {handoff_tool_names}")
+            raise ValueError(f"Handoff names must be unique: {handoff_tool_names}")
         # Check if handoff tool names not in tool names.
         if any(name in tool_names for name in handoff_tool_names):
             raise ValueError(
@@ -344,8 +338,7 @@ async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token:
         async for message in self.on_messages_stream(messages, cancellation_token):
             if isinstance(message, Response):
                 return message
-        raise AssertionError(
-            "The stream should have returned the final result.")
+        raise AssertionError("The stream should have returned the final result.")
 
     async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
@@ -377,26 +370,22 @@ async def on_messages_stream(
         # Check if the response is a string and return it.
         if isinstance(result.content, str):
             yield Response(
-                chat_message=TextMessage(
-                    content=result.content, source=self.name, models_usage=result.usage),
+                chat_message=TextMessage(content=result.content, source=self.name, models_usage=result.usage),
                 inner_messages=inner_messages,
             )
             return
 
         # Process tool calls.
-        assert isinstance(result.content, list) and all(
-            isinstance(item, FunctionCall) for item in result.content)
-        tool_call_msg = ToolCallRequestEvent(
-            content=result.content, source=self.name, models_usage=result.usage)
+        assert isinstance(result.content, list) and all(isinstance(item, FunctionCall) for item in result.content)
+        tool_call_msg = ToolCallRequestEvent(content=result.content, source=self.name, models_usage=result.usage)
         event_logger.debug(tool_call_msg)
         # Add the tool call message to the output.
         inner_messages.append(tool_call_msg)
         yield tool_call_msg
 
         # Execute the tool calls.
         results = await asyncio.gather(*[self._execute_tool_call(call, cancellation_token) for call in result.content])
-        tool_call_result_msg = ToolCallExecutionEvent(
-            content=results, source=self.name)
+        tool_call_result_msg = ToolCallExecutionEvent(content=results, source=self.name)
         event_logger.debug(tool_call_result_msg)
         await self._model_context.add_message(FunctionExecutionResultMessage(content=results))
         inner_messages.append(tool_call_result_msg)
@@ -416,8 +405,7 @@ async def on_messages_stream(
                 )
             # Return the output messages to signal the handoff.
             yield Response(
-                chat_message=HandoffMessage(
-                    content=handoffs[0].message, target=handoffs[0].target, source=self.name),
+                chat_message=HandoffMessage(content=handoffs[0].message, target=handoffs[0].target, source=self.name),
                 inner_messages=inner_messages,
             )
             return
@@ -431,8 +419,7 @@ async def on_messages_stream(
             await self._model_context.add_message(AssistantMessage(content=result.content, source=self.name))
             # Yield the response.
             yield Response(
-                chat_message=TextMessage(
-                    content=result.content, source=self.name, models_usage=result.usage),
+                chat_message=TextMessage(content=result.content, source=self.name, models_usage=result.usage),
                 inner_messages=inner_messages,
             )
         else:
@@ -448,8 +435,7 @@ async def on_messages_stream(
                 )
             tool_call_summary = "\n".join(tool_call_summaries)
             yield Response(
-                chat_message=ToolCallSummaryMessage(
-                    content=tool_call_summary, source=self.name),
+                chat_message=ToolCallSummaryMessage(content=tool_call_summary, source=self.name),
                 inner_messages=inner_messages,
             )
 
@@ -460,11 +446,9 @@ async def _execute_tool_call(
         try:
             if not self._tools + self._handoff_tools:
                 raise ValueError("No tools are available.")
-            tool = next((t for t in self._tools +
-                        self._handoff_tools if t.name == tool_call.name), None)
+            tool = next((t for t in self._tools + self._handoff_tools if t.name == tool_call.name), None)
             if tool is None:
-                raise ValueError(
-                    f"The tool '{tool_call.name}' is not available.")
+                raise ValueError(f"The tool '{tool_call.name}' is not available.")
             arguments = json.loads(tool_call.arguments)
             result = await tool.run_json(arguments, cancellation_token)
             result_as_str = tool.return_value_as_string(result)

python/packages/autogen-agentchat/src/autogen_agentchat/memory/__init__.py

@@ -0,0 +1,10 @@
+from ._base_memory import Memory, MemoryContent, MemoryMimeType
+from ._list_memory import ListMemory, ListMemoryConfig
+
+__all__ = [
+    "Memory",
+    "MemoryContent",
+    "MemoryMimeType",
+    "ListMemory",
+    "ListMemoryConfig",
+]

python/packages/autogen-agentchat/src/autogen_agentchat/memory/_base_memory.py

@@ -3,8 +3,8 @@
 from typing import Any, Dict, List, Protocol, Union, runtime_checkable
 
 from autogen_core import CancellationToken, Image
-from pydantic import BaseModel, ConfigDict, Field
 from autogen_core.model_context import ChatCompletionContext
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class MemoryMimeType(Enum):
@@ -22,7 +22,7 @@ class MemoryMimeType(Enum):
 
 class MemoryContent(BaseModel):
     content: ContentType
-    mime_type: MemoryMimeType
+    mime_type: MemoryMimeType | str
     metadata: Dict[str, Any] | None = None
     timestamp: datetime | None = None
     source: str | None = None
@@ -35,8 +35,7 @@ class BaseMemoryConfig(BaseModel):
     """Base configuration for memory implementations."""
 
     k: int = Field(default=5, description="Number of results to return")
-    score_threshold: float | None = Field(
-        default=None, description="Minimum relevance score")
+    score_threshold: float | None = Field(default=None, description="Minimum relevance score")
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 

python/packages/autogen-agentchat/src/autogen_agentchat/memory/_list_memory.py

@@ -1,17 +1,16 @@
-from difflib import SequenceMatcher
 import logging
+from difflib import SequenceMatcher
 from typing import Any, List
 
 from autogen_core import CancellationToken, Image
-from pydantic import Field
-
-from ._base_memory import BaseMemoryConfig, MemoryContent, Memory, MemoryMimeType
 from autogen_core.model_context import ChatCompletionContext
 from autogen_core.models import (
     SystemMessage,
 )
+from pydantic import Field
 
 from .. import EVENT_LOGGER_NAME
+from ._base_memory import BaseMemoryConfig, Memory, MemoryContent, MemoryMimeType
 
 event_logger = logging.getLogger(EVENT_LOGGER_NAME)
 
@@ -34,19 +33,10 @@ class ListMemory(Memory):
     Example:
         ```python
         # Initialize memory with custom config
-        memory = ListMemory(
-            name="chat_history",
-            config=ListMemoryConfig(
-                similarity_threshold=0.7,
-                k=3
-            )
-        )
+        memory = ListMemory(name="chat_history", config=ListMemoryConfig(similarity_threshold=0.7, k=3))
 
         # Add memory content
-        content = MemoryContent(
-            content="User prefers formal language",
-            mime_type=MemoryMimeType.TEXT
-        )
+        content = MemoryContent(content="User prefers formal language", mime_type=MemoryMimeType.TEXT)
         await memory.add(content)
 
         # Transform a model context with memory
@@ -106,26 +96,23 @@ async def transform(
 
         # Extract query from last message
         last_message = messages[-1]
-        query_text = last_message.content if isinstance(
-            last_message.content, str) else str(last_message)
-        query = MemoryContent(content=query_text,
-                              mime_type=MemoryMimeType.TEXT)
+        query_text = last_message.content if isinstance(last_message.content, str) else str(last_message)
+        query = MemoryContent(content=query_text, mime_type=MemoryMimeType.TEXT)
 
         # Query memory and format results
         results: List[str] = []
         query_results = await self.query(query)
         for i, result in enumerate(query_results, 1):
             if isinstance(result.content, str):
                 results.append(f"{i}. {result.content}")
-                event_logger.debug(
-                    f"Retrieved memory {i}. {result.content}, score: {result.score}"
-                )
+                event_logger.debug(f"Retrieved memory {i}. {result.content}, score: {result.score}")
 
         # Add memory results to context
         if results:
             memory_context = (
-                "\n The following results were retrieved from memory for this task. You may choose to use them or not. :\n" +
-                "\n".join(results) + "\n"
+                "\n The following results were retrieved from memory for this task. You may choose to use them or not. :\n"
+                + "\n".join(results)
+                + "\n"
             )
             await model_context.add_message(SystemMessage(content=memory_context))
 
@@ -159,10 +146,7 @@ async def query(
         Example:
             ```python
             # Query memories similar to some text
-            query = MemoryContent(
-                content="What's the weather?", 
-                mime_type=MemoryMimeType.TEXT
-            )
+            query = MemoryContent(content="What's the weather?", mime_type=MemoryMimeType.TEXT)
             results = await memory.query(query)
 
             # Check similarity scores
@@ -172,8 +156,8 @@ async def query(
         """
         try:
             query_text = self._extract_text(query)
-        except ValueError:
-            raise ValueError("Query must contain text content")
+        except ValueError as e:
+            raise ValueError("Query must contain text content") from e
 
         results: List[MemoryContent] = []
 
@@ -207,7 +191,7 @@ def _calculate_similarity(self, text1: str, text2: str) -> float:
 
         Note:
             Uses difflib's SequenceMatcher for basic text similarity.
-            For production use cases, consider using more sophisticated 
+            For production use cases, consider using more sophisticated
             similarity metrics or embeddings.
         """
         return SequenceMatcher(None, text1.lower(), text2.lower()).ratio()
@@ -242,14 +226,9 @@ def _extract_text(self, content_item: MemoryContent) -> str:
         elif isinstance(content, Image):
             raise ValueError("Image content cannot be converted to text")
         else:
-            raise ValueError(
-                f"Unsupported content type: {content_item.mime_type}")
+            raise ValueError(f"Unsupported content type: {content_item.mime_type}")
 
-    async def add(
-        self,
-        content: MemoryContent,
-        cancellation_token: CancellationToken | None = None
-    ) -> None:
+    async def add(self, content: MemoryContent, cancellation_token: CancellationToken | None = None) -> None:
         """Add new content to memory.
 
         Args:
@@ -262,3 +241,11 @@ async def add(
             deduplication or content-based filtering.
         """
         self._contents.append(content)
+
+    async def clear(self) -> None:
+        """Clear all memory content."""
+        self._contents = []
+
+    async def cleanup(self) -> None:
+        """Cleanup resources if needed."""
+        pass