ag2 0.9.9__py3-none-any.whl → 0.9.10__py3-none-any.whl

autogen/agentchat/chat.py

@@ -7,11 +7,12 @@
 import asyncio
 import datetime
 import logging
+import uuid
 import warnings
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import partial
-from typing import Any
+from typing import Any, TypedDict
 
 from ..doc_utils import export_module
 from ..events.agent_events import PostCarryoverProcessingEvent
@@ -24,26 +25,38 @@ Prerequisite = tuple[int, int]
 __all__ = ["ChatResult", "a_initiate_chats", "initiate_chats"]
 
 
+class CostDict(TypedDict):
+    usage_including_cached_inference: dict[str, Any]
+    usage_excluding_cached_inference: dict[str, Any]
+
+
 @dataclass
 @export_module("autogen")
 class ChatResult:
     """(Experimental) The result of a chat. Almost certain to be changed."""
 
-    chat_id: int = None
+    chat_id: int = field(default_factory=lambda: uuid.uuid4().int)
     """chat id"""
-    chat_history: list[dict[str, Any]] = None
+
+    chat_history: list[dict[str, Any]] = field(default_factory=list)
     """The chat history."""
-    summary: str = None
+
+    summary: str = ""
     """A summary obtained from the chat."""
-    cost: dict[str, dict[str, Any]] = (
-        None  # keys: "usage_including_cached_inference", "usage_excluding_cached_inference"
+
+    cost: CostDict = field(
+        default_factory=lambda: {
+            "usage_including_cached_inference": {},
+            "usage_excluding_cached_inference": {},
+        }
     )
     """The cost of the chat.
        The value for each usage type is a dictionary containing cost information for that specific type.
            - "usage_including_cached_inference": Cost information on the total usage, including the tokens in cached inference.
            - "usage_excluding_cached_inference": Cost information on the usage of tokens, excluding the tokens in cache. No larger than "usage_including_cached_inference".
     """
-    human_input: list[str] = None
+
+    human_input: list[str] = field(default_factory=list)
     """A list of human input solicited during the chat."""
 
 
@@ -141,37 +154,36 @@ def initiate_chats(chat_queue: list[dict[str, Any]]) -> list[ChatResult]:
     """Initiate a list of chats.
 
     Args:
-        chat_queue (List[Dict]): A list of dictionaries containing the information about the chats.
-
+        chat_queue (List[Dict]): A list of dictionaries containing the information about the chats.\n
             Each dictionary should contain the input arguments for
-            [`ConversableAgent.initiate_chat`](../ConversableAgent#initiate-chat).
-            For example:
-                - `"sender"` - the sender agent.
-                - `"recipient"` - the recipient agent.
-                - `"clear_history"` (bool) - whether to clear the chat history with the agent.
-                Default is True.
-                - `"silent"` (bool or None) - (Experimental) whether to print the messages in this
-                conversation. Default is False.
-                - `"cache"` (Cache or None) - the cache client to use for this conversation.
-                Default is None.
-                - `"max_turns"` (int or None) - maximum number of turns for the chat. If None, the chat
-                will continue until a termination condition is met. Default is None.
-                - `"summary_method"` (str or callable) - a string or callable specifying the method to get
-                a summary from the chat. Default is DEFAULT_summary_method, i.e., "last_msg".
-                - `"summary_args"` (dict) - a dictionary of arguments to be passed to the summary_method.
-                Default is {}.
-                - `"message"` (str, callable or None) - if None, input() will be called to get the
-                initial message.
-                - `**context` - additional context information to be passed to the chat.
-                - `"carryover"` - It can be used to specify the carryover information to be passed
-                to this chat. If provided, we will combine this carryover with the "message" content when
-                generating the initial chat message in `generate_init_message`.
-                - `"finished_chat_indexes_to_exclude_from_carryover"` - It can be used by specifying a list of indexes of the finished_chats list,
-                from which to exclude the summaries for carryover. If 'finished_chat_indexes_to_exclude_from_carryover' is not provided or an empty list,
-                then summary from all the finished chats will be taken.
+            [`ConversableAgent.initiate_chat`](../ConversableAgent#initiate-chat).\n
+            For example:\n
+                - `"sender"` - the sender agent.\n
+                - `"recipient"` - the recipient agent.\n
+                - `"clear_history"` (bool) - whether to clear the chat history with the agent.\n
+                  Default is True.\n
+                - `"silent"` (bool or None) - (Experimental) whether to print the messages in this\n
+                  conversation. Default is False.\n
+                - `"cache"` (Cache or None) - the cache client to use for this conversation.\n
+                  Default is None.\n
+                - `"max_turns"` (int or None) - maximum number of turns for the chat. If None, the chat\n
+                  will continue until a termination condition is met. Default is None.\n
+                - `"summary_method"` (str or callable) - a string or callable specifying the method to get\n
+                  a summary from the chat. Default is DEFAULT_summary_method, i.e., "last_msg".\n
+                - `"summary_args"` (dict) - a dictionary of arguments to be passed to the summary_method.\n
+                  Default is {}.\n
+                - `"message"` (str, callable or None) - if None, input() will be called to get the\n
+                  initial message.\n
+                - `**context` - additional context information to be passed to the chat.\n
+                - `"carryover"` - It can be used to specify the carryover information to be passed\n
+                  to this chat. If provided, we will combine this carryover with the "message" content when\n
+                  generating the initial chat message in `generate_init_message`.\n
+                - `"finished_chat_indexes_to_exclude_from_carryover"` - It can be used by specifying a list of indexes of the finished_chats list,\n
+                  from which to exclude the summaries for carryover. If 'finished_chat_indexes_to_exclude_from_carryover' is not provided or an empty list,\n
+                  then summary from all the finished chats will be taken.\n
 
     Returns:
-        (list): a list of ChatResult objects corresponding to the finished chats in the chat_queue.
+        (list): a list of ChatResult objects corresponding to the finished chats in the chat_queue.\n
     """
     consolidate_chat_info(chat_queue)
     _validate_recipients(chat_queue)
@@ -220,7 +232,7 @@ async def _dependent_chat_future(
     finished_chat_indexes_to_exclude_from_carryover = chat_info.get(
         "finished_chat_indexes_to_exclude_from_carryover", []
     )
-    finished_chats = dict()
+    finished_chats = {}
     for chat in prerequisite_chat_futures:
         chat_future = prerequisite_chat_futures[chat]
         if chat_future.cancelled():
@@ -291,18 +303,18 @@ async def a_initiate_chats(chat_queue: list[dict[str, Any]]) -> dict[int, ChatRe
     num_chats = chat_book.keys()
     prerequisites = __create_async_prerequisites(chat_queue)
     chat_order_by_id = __find_async_chat_order(num_chats, prerequisites)
-    finished_chat_futures = dict()
+    finished_chat_futures = {}
     for chat_id in chat_order_by_id:
         chat_info = chat_book[chat_id]
         prerequisite_chat_ids = chat_info.get("prerequisites", [])
-        pre_chat_futures = dict()
+        pre_chat_futures = {}
         for pre_chat_id in prerequisite_chat_ids:
             pre_chat_future = finished_chat_futures[pre_chat_id]
             pre_chat_futures[pre_chat_id] = pre_chat_future
         current_chat_future = await _dependent_chat_future(chat_id, chat_info, pre_chat_futures)
         finished_chat_futures[chat_id] = current_chat_future
     await asyncio.gather(*list(finished_chat_futures.values()))
-    finished_chats = dict()
+    finished_chats = {}
     for chat in finished_chat_futures:
         chat_result = finished_chat_futures[chat].result()
         finished_chats[chat] = chat_result

autogen/agentchat/contrib/agent_eval/criterion.py

@@ -17,7 +17,7 @@ class Criterion(BaseModel):
     name: str
     description: str
     accepted_values: list[str]
-    sub_criteria: list[Criterion] = list()
+    sub_criteria: list[Criterion] = []
 
     @staticmethod
     def parse_json_str(criteria: str):

autogen/agentchat/contrib/capabilities/text_compressors.py

@@ -34,11 +34,11 @@ class LLMLingua:
 
     def __init__(
         self,
-        prompt_compressor_kwargs: dict = dict(
-            model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",
-            use_llmlingua2=True,
-            device_map="cpu",
-        ),
+        prompt_compressor_kwargs: dict = {
+            "model_name": "microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",
+            "use_llmlingua2": True,
+            "device_map": "cpu",
+        },
         structured_compression: bool = False,
     ) -> None:
         """Args:

autogen/agentchat/contrib/capabilities/tools_capability.py

@@ -14,7 +14,7 @@ class ToolsCapability:
     """
 
     def __init__(self, tool_list: list[Tool]):
-        self.tools = [tool for tool in tool_list]
+        self.tools = list(tool_list)
 
     def add_to_agent(self, agent: ConversableAgent):
         """Add tools to the given agent."""

autogen/agentchat/contrib/capabilities/transforms.py

@@ -337,7 +337,7 @@ class TextMessageCompressor:
         self,
         text_compressor: TextCompressor | None = None,
         min_tokens: int | None = None,
-        compression_params: dict = dict(),
+        compression_params: dict = {},
         cache: AbstractCache | None = None,
         filter_dict: dict[str, Any] | None = None,
         exclude_filter: bool = True,

autogen/agentchat/contrib/captainagent/agent_builder.py

@@ -380,7 +380,7 @@ Match roles in the role set to each expert in expert set.
 
     def clear_all_agents(self, recycle_endpoint: bool | None = True):
         """Clear all cached agents."""
-        for agent_name in [agent_name for agent_name in self.agent_procs_assign]:
+        for agent_name in list(self.agent_procs_assign):
             self.clear_agent(agent_name, recycle_endpoint)
         print(colored("All agents have been cleared.", "yellow"), flush=True)
 

autogen/agentchat/contrib/captainagent/captainagent.py

@@ -149,25 +149,26 @@ Note that the previous experts will forget everything after you obtain the respo
         description: str | None = DEFAULT_DESCRIPTION,
         **kwargs: Any,
     ):
-        """Args:
-        name (str): agent name.
-        system_message (str): system message for the ChatCompletion inference.
-            Please override this attribute if you want to reprogram the agent.
-        llm_config (LLMConfig or dict or False): llm inference configuration.
-            Please refer to [OpenAIWrapper.create](https://docs.ag2.ai/latest/docs/api-reference/autogen/OpenAIWrapper/#autogen.OpenAIWrapper.create) for available options.
-        is_termination_msg (function): a function that takes a message in the form of a dictionary
-            and returns a boolean value indicating if this received message is a termination message.
-            The dict can contain the following keys: "content", "role", "name", "function_call".
-        max_consecutive_auto_reply (int): the maximum number of consecutive auto replies.
-            default to None (no limit provided, class attribute MAX_CONSECUTIVE_AUTO_REPLY will be used as the limit in this case).
-            The limit only plays a role when human_input_mode is not "ALWAYS".
-        agent_lib (str): the path or a JSON file of the agent library for retrieving the nested chat instantiated by CaptainAgent.
-        tool_lib (str): the path to the tool library for retrieving the tools used in the nested chat instantiated by CaptainAgent.
-        nested_config (dict): the configuration for the nested chat instantiated by CaptainAgent.
-            A full list of keys and their functionalities can be found in [docs](https://docs.ag2.ai/latest/docs/user-guide/reference-agents/captainagent).
-        agent_config_save_path (str): the path to save the generated or retrieved agent configuration.
-        **kwargs (dict): Please refer to other kwargs in
-            [ConversableAgent](https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/conversable_agent.py#L74).
+        """
+        Args:\n
+        name (str): agent name.\n
+        system_message (str): system message for the ChatCompletion inference.\n
+            Please override this attribute if you want to reprogram the agent.\n
+        llm_config (LLMConfig or dict or False): llm inference configuration.\n
+            Please refer to [OpenAIWrapper.create](https://docs.ag2.ai/latest/docs/api-reference/autogen/OpenAIWrapper/#autogen.OpenAIWrapper.create) for available options.\n
+        is_termination_msg (function): a function that takes a message in the form of a dictionary\n
+            and returns a boolean value indicating if this received message is a termination message.\n
+            The dict can contain the following keys: "content", "role", "name", "function_call".\n
+        max_consecutive_auto_reply (int): the maximum number of consecutive auto replies.\n
+            default to None (no limit provided, class attribute MAX_CONSECUTIVE_AUTO_REPLY will be used as the limit in this case).\n
+            The limit only plays a role when human_input_mode is not "ALWAYS".\n
+        agent_lib (str): the path or a JSON file of the agent library for retrieving the nested chat instantiated by CaptainAgent.\n
+        tool_lib (str): the path to the tool library for retrieving the tools used in the nested chat instantiated by CaptainAgent.\n
+        nested_config (dict): the configuration for the nested chat instantiated by CaptainAgent.\n
+            A full list of keys and their functionalities can be found in [docs](https://docs.ag2.ai/latest/docs/user-guide/reference-agents/captainagent).\n
+        agent_config_save_path (str): the path to save the generated or retrieved agent configuration.\n
+        **kwargs (dict): Please refer to other kwargs in\n
+            [ConversableAgent](https://github.com/ag2ai/ag2/blob/main/autogen/agentchat/conversable_agent.py#L74).\n
         """
         super().__init__(
             name,

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -55,7 +55,7 @@ class FalkorGraphQueryEngine:
         self.username = username
         self.password = password
         self.model = model or OpenAiGenerativeModel("gpt-4o")
-        self.model_config = KnowledgeGraphModelConfig.with_model(model)
+        self.model_config = KnowledgeGraphModelConfig.with_model(self.model)
         self.ontology = ontology
         self.knowledge_graph: KnowledgeGraph | None = None  # type: ignore[no-any-unimported]
         self.falkordb = FalkorDB(host=self.host, port=self.port, username=self.username, password=self.password)
@@ -139,9 +139,6 @@ class FalkorGraphQueryEngine:
 
         response = self._chat_session.send_message(question)
 
-        # History will be considered when querying by setting the last_answer
-        self._chat_session.last_answer = response["response"]
-
         return GraphStoreQueryResult(answer=response["response"], results=[])
 
     def delete(self) -> bool:
@@ -167,4 +164,4 @@ class FalkorGraphQueryEngine:
         if self.ontology_table_name not in self.falkordb.list_graphs():
             raise ValueError(f"Knowledge graph {self.name} has not been created.")
         graph = self.__get_ontology_storage_graph()
-        return Ontology.from_graph(graph)
+        return Ontology.from_schema_graph(graph)

autogen/agentchat/contrib/graph_rag/graph_rag_capability.py

@@ -14,10 +14,10 @@ __all__ = ["GraphRagCapability"]
 class GraphRagCapability(AgentCapability):
     """A graph-based RAG capability uses a graph query engine to give a conversable agent the graph-based RAG ability.
 
-    An agent class with graph-based RAG capability could
-    1. create a graph in the underlying database with input documents.
-    2. retrieved relevant information based on messages received by the agent.
-    3. generate answers from retrieved information and send messages back.
+    An agent class with graph-based RAG capability could:\n
+    1. create a graph in the underlying database with input documents.\n
+    2. retrieved relevant information based on messages received by the agent.\n
+    3. generate answers from retrieved information and send messages back.\n
 
     For example,
     ```python
@@ -41,7 +41,7 @@ class GraphRagCapability(AgentCapability):
     user_proxy.initiate_chat(graph_rag_agent, message="Name a few actors who've played in 'The Matrix'")
 
     # ChatResult(
-        # chat_id=None,
+        # chat_id=uuid.uuid4().int,
         # chat_history=[
             # {'content': 'Name a few actors who've played in \'The Matrix\'', 'role': 'graph_rag_agent'},
             # {'content': 'A few actors who have played in The Matrix are:

autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py

@@ -28,23 +28,24 @@ with optional_import_block():
 
 @require_optional_import("llama_index", "neo4j")
 class Neo4jGraphQueryEngine:
-    """This class serves as a wrapper for a property graph query engine backed by LlamaIndex and Neo4j,
-    facilitating the creating, connecting, updating, and querying of LlamaIndex property graphs.
-
-    It builds a property graph Index from input documents,
-    storing and retrieving data from the property graph in the Neo4j database.
-
-    It extracts triplets, i.e., [entity] -> [relationship] -> [entity] sets,
-    from the input documents using llamIndex extractors.
-
-    Users can provide custom entities, relationships, and schema to guide the extraction process.
-
-    If strict is True, the engine will extract triplets following the schema
-    of allowed relationships for each entity specified in the schema.
-
-    It also leverages LlamaIndex's chat engine which has a conversation history internally to provide context-aware responses.
-
-    For usage, please refer to example notebook/agentchat_graph_rag_neo4j.ipynb
+    """
+    This class serves as a wrapper for a property graph query engine backed by LlamaIndex and Neo4j,\n
+    facilitating the creating, connecting, updating, and querying of LlamaIndex property graphs.\n
+    \n
+    It builds a property graph Index from input documents,\n
+    storing and retrieving data from the property graph in the Neo4j database.\n
+    \n
+    It extracts triplets, i.e., [entity] -> [relationship] -> [entity] sets,\n
+    from the input documents using llamIndex extractors.\n
+    \n
+    Users can provide custom entities, relationships, and schema to guide the extraction process.\n
+    \n
+    If strict is True, the engine will extract triplets following the schema\n
+    of allowed relationships for each entity specified in the schema.\n
+    \n
+    It also leverages LlamaIndex's chat engine which has a conversation history internally to provide context-aware responses.\n
+    \n
+    For usage, please refer to example notebook/agentchat_graph_rag_neo4j.ipynb\n
     """
 
     def __init__(  # type: ignore[no-any-unimported]

autogen/agentchat/contrib/rag/mongodb_query_engine.py

@@ -17,7 +17,7 @@ with optional_import_block():
     from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
     from llama_index.core.embeddings import BaseEmbedding
     from llama_index.core.schema import Document as LlamaDocument
-    from llama_index.llms.langchain.base import LLM
+    from llama_index.llms.langchain.base import LLM  # type: ignore[attr-defined]
     from llama_index.llms.openai import OpenAI
     from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
     from pymongo import MongoClient
@@ -114,7 +114,7 @@ class MongoDBQueryEngine:
         logger.info("Vector database created.")
         self.vector_search_engine = MongoDBAtlasVectorSearch(
             mongodb_client=self.vector_db.client,  # type: ignore[union-attr]
-            db_name=self.database_name,
+            db_name=self.database_name,  # type: ignore[arg-type]
             collection_name=self.collection_name,
         )
         logger.info("Vector search engine created.")

autogen/agentchat/contrib/rag/query_engine.py

@@ -29,18 +29,18 @@ class RAGQueryEngine(Protocol):
         """Initialize the database with the input documents or records.
 
         This method initializes database with the input documents or records.
-        Usually, it takes the following steps:
-        1. connecting to a database.
-        2. insert records
-        3. build indexes etc.
+        Usually, it takes the following steps:\n
+        1. connecting to a database.\n
+        2. insert records.\n
+        3. build indexes etc.\n
 
-        Args:
-            new_doc_dir (Optional[Union[Path, str]]): A directory containing documents to be ingested.
-            new_doc_paths_or_urls (Optional[Sequence[Union[Path, str]]]): A list of paths or URLs to documents to be ingested.
-            *args: Any additional arguments
-            **kwargs: Any additional keyword arguments
-        Returns:
-            bool: True if initialization is successful, False otherwise
+        Args:\n
+            new_doc_dir (Optional[Union[Path, str]]): A directory containing documents to be ingested.\n
+            new_doc_paths_or_urls (Optional[Sequence[Union[Path, str]]]): A list of paths or URLs to documents to be ingested.\n
+            *args: Any additional arguments\n
+            **kwargs: Any additional keyword arguments\n
+        Returns:\n
+            bool: True if initialization is successful, False otherwise\n
         """
         ...
 

autogen/agentchat/contrib/retrieve_assistant_agent.py

@@ -7,10 +7,13 @@
 import warnings
 from typing import Any
 
+from typing_extensions import deprecated
+
 from ..agent import Agent
 from ..assistant_agent import AssistantAgent
 
 
+@deprecated("The RetrieveAssistantAgent is deprecated. Please use the AssistantAgent instead.")
 class RetrieveAssistantAgent(AssistantAgent):
     """(Experimental) Retrieve Assistant agent, designed to solve a task with LLM.
 

autogen/agentchat/contrib/swarm_agent.py

@@ -1163,8 +1163,9 @@ async def a_run_swarm(
             except Exception as e:
                 response.iostream.send(ErrorEvent(error=e))  # type: ignore[call-arg]
 
-    asyncio.create_task(stream_run())
-
+    task = asyncio.create_task(stream_run())
+    # prevent the task from being garbage collected
+    response._task_ref = task  # type: ignore[attr-defined]
     return response
 
 

autogen/agentchat/contrib/vectordb/couchbase.py

@@ -279,7 +279,7 @@ class CouchbaseVectorDB(VectorDB):
 
         for i in range(0, len(docs), batch_size):
             batch = docs[i : i + batch_size]
-            docs_to_upsert = dict()
+            docs_to_upsert = {}
             for doc in batch:
                 doc_id = doc["id"]
                 embedding = self.embedding_function([

autogen/agentchat/contrib/vectordb/mongodb.py

@@ -437,7 +437,7 @@ class MongoDBAtlasVectorDB(VectorDB):
         if include is None:
             include_fields = {"_id": 1, "content": 1, "metadata": 1}
         else:
-            include_fields = {k: 1 for k in set(include).union({"_id"})}
+            include_fields = dict.fromkeys(set(include).union({"_id"}), 1)
         collection = self.get_collection(collection_name)
         if ids is not None:
             docs = collection.find({"_id": {"$in": ids}}, include_fields)

autogen/agentchat/contrib/web_surfer.py

@@ -277,7 +277,7 @@ class WebSurferAgent(ConversableAgent):
         self._assistant.reset()  # type: ignore[no-untyped-call]
 
         # Clone the messages to give context
-        self._assistant.chat_messages[self._user_proxy] = list()
+        self._assistant.chat_messages[self._user_proxy] = []
         history = messages[0 : len(messages) - 1]
         for message in history:
             self._assistant.chat_messages[self._user_proxy].append(message)