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

autogen/a2a/utils.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, cast
+from uuid import uuid4
+
+from a2a.types import Artifact, DataPart, Message, Part, Role, TextPart
+from a2a.utils import new_agent_parts_message, new_artifact
+
+from autogen.remote.protocol import RequestMessage, ResponseMessage
+
+CLIENT_TOOLS_KEY = "ag2_client_tools"
+CONTEXT_KEY = "ag2_context_update"
+
+
+def request_message_to_a2a(
+    request_message: RequestMessage,
+    context_id: str,
+) -> Message:
+    metadata: dict[str, Any] = {}
+    if request_message.client_tools:
+        metadata[CLIENT_TOOLS_KEY] = request_message.client_tools
+    if request_message.context:
+        metadata[CONTEXT_KEY] = request_message.context
+
+    return Message(
+        role=Role.user,
+        parts=[message_to_part(message) for message in request_message.messages],
+        message_id=uuid4().hex,
+        context_id=context_id,
+        metadata=metadata,
+    )
+
+
+def request_message_from_a2a(message: Message) -> RequestMessage:
+    metadata = message.metadata or {}
+    return RequestMessage(
+        messages=[message_from_part(part) for part in message.parts],
+        context=metadata.get(CONTEXT_KEY),
+        client_tools=metadata.get(CLIENT_TOOLS_KEY, []),
+    )
+
+
+def response_message_from_a2a_artifacts(artifacts: list[Artifact] | None) -> ResponseMessage | None:
+    if not artifacts:
+        return None
+
+    if len(artifacts) > 1:
+        raise NotImplementedError("Multiple artifacts are not supported")
+
+    artifact = artifacts[-1]
+
+    if not artifact.parts:
+        return None
+
+    if len(artifact.parts) > 1:
+        raise NotImplementedError("Multiple parts are not supported")
+
+    return ResponseMessage(
+        messages=[message_from_part(artifact.parts[-1])],
+        context=(artifact.metadata or {}).get(CONTEXT_KEY),
+    )
+
+
+def response_message_from_a2a_message(message: Message) -> ResponseMessage | None:
+    text_parts: list[Part] = []
+    data_parts: list[Part] = []
+    for part in message.parts:
+        if isinstance(part.root, TextPart):
+            text_parts.append(part)
+        elif isinstance(part.root, DataPart):
+            data_parts.append(part)
+        else:
+            raise NotImplementedError(f"Unsupported part type: {type(part.root)}")
+
+    tpn = len(text_parts)
+    if dpn := len(data_parts):
+        if dpn > 1:
+            raise NotImplementedError("Multiple data parts are not supported")
+
+        if tpn:
+            raise NotImplementedError("Data parts and text parts are not supported together")
+
+        messages = [message_from_part(data_parts[0])]
+    elif tpn == 1:
+        messages = [message_from_part(text_parts[0])]
+    else:
+        messages = [{"content": "\n".join(cast(TextPart, t.root).text for t in text_parts)}]
+
+    return ResponseMessage(
+        messages=messages,
+        context=(message.metadata or {}).get(CONTEXT_KEY),
+    )
+
+
+def response_message_to_a2a(
+    result: ResponseMessage | None,
+    context_id: str | None,
+    task_id: str | None,
+) -> tuple[Artifact, list[Message]]:
+    # mypy ignores could be removed after
+    # https://github.com/a2aproject/a2a-python/pull/503
+
+    if not result:
+        return new_artifact(
+            name="result",
+            parts=[],
+            description=None,  # type: ignore[arg-type]
+        ), []
+
+    artifact = new_artifact(
+        name="result",
+        parts=[message_to_part(result.messages[-1])],
+        description=None,  # type: ignore[arg-type]
+    )
+
+    if result.context:
+        artifact.metadata = {CONTEXT_KEY: result.context}
+
+    return (
+        artifact,
+        [
+            new_agent_parts_message(
+                parts=[message_to_part(m) for m in result.messages],
+                context_id=context_id,
+                task_id=task_id,
+            ),
+        ],
+    )
+
+
+def message_to_part(message: dict[str, Any]) -> Part:
+    message = message.copy()
+    text = message.pop("content", "") or ""
+    return Part(
+        root=TextPart(
+            text=text,
+            metadata=message or None,
+        )
+    )
+
+
+def message_from_part(part: Part) -> dict[str, Any]:
+    root = part.root
+
+    if isinstance(root, TextPart):
+        return {
+            **(root.metadata or {}),
+            "content": root.text,
+        }
+
+    elif isinstance(root, DataPart):
+        if (  # pydantic-ai specific
+            set(root.data.keys()) == {"result"}
+            and root.metadata
+            and "json_schema" in root.metadata
+            and isinstance(data := root.data["result"], dict)
+        ):
+            return data
+
+        return root.data
+
+    else:
+        raise NotImplementedError(f"Unsupported part type: {type(part.root)}")

autogen/agentchat/__init__.py

@@ -15,6 +15,7 @@ from .contrib.swarm_agent import (
     run_swarm,
 )
 from .conversable_agent import ConversableAgent, UpdateSystemMessage, register_function
+from .group import ContextVariables, ReplyResult
 from .group.multi_agent_chat import a_initiate_group_chat, a_run_group_chat, initiate_group_chat, run_group_chat
 from .groupchat import GroupChat, GroupChatManager
 from .user_proxy_agent import UserProxyAgent
@@ -24,10 +25,12 @@ __all__ = [
     "Agent",
     "AssistantAgent",
     "ChatResult",
+    "ContextVariables",
     "ConversableAgent",
     "GroupChat",
     "GroupChatManager",
     "LLMAgent",
+    "ReplyResult",
     "UpdateSystemMessage",
     "UserProxyAgent",
     "a_initiate_chats",

autogen/agentchat/agent.py

@@ -105,7 +105,6 @@ class Agent(Protocol):
         self,
         messages: list[dict[str, Any]] | None = None,
         sender: Optional["Agent"] = None,
-        **kwargs: Any,
     ) -> str | dict[str, Any] | None:
         """Generate a reply based on the received messages.
 
@@ -124,7 +123,6 @@ class Agent(Protocol):
         self,
         messages: list[dict[str, Any]] | None = None,
         sender: Optional["Agent"] = None,
-        **kwargs: Any,
     ) -> str | dict[str, Any] | None:
         """(Async) Generate a reply based on the received messages.
 

autogen/agentchat/assistant_agent.py

@@ -19,10 +19,10 @@ class AssistantAgent(ConversableAgent):
 
     AssistantAgent is a subclass of ConversableAgent configured with a default system message.
     The default system message is designed to solve a task with LLM,
-    including suggesting python code blocks and debugging.
-    `human_input_mode` is default to "NEVER"
-    and `code_execution_config` is default to False.
-    This agent doesn't execute code by default, and expects the user to execute the code.
+    including suggesting python code blocks and debugging. \n
+    `human_input_mode` is default to "NEVER" \n
+    and `code_execution_config` is default to False. \n
+    This agent doesn't execute code by default, and expects the user to execute the code. \n
     """
 
     DEFAULT_SYSTEM_MESSAGE = """You are a helpful AI assistant.
@@ -52,20 +52,20 @@ Reply "TERMINATE" in the end when everything is done.
         **kwargs: Any,
     ):
         """Args:
-        name (str): agent name.
-        system_message (str): system message for the ChatCompletion inference.
+        - 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.
-        llm_config (dict or False or None): 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
+        - llm_config (dict or False or None): llm inference configuration. \n
+            Please refer to [OpenAIWrapper.create](https://docs.ag2.ai/latest/docs/api-reference/autogen/OpenAIWrapper/#autogen.OpenAIWrapper.create) \n
+            for available options. \n
+        - 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.
+            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.
             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".
-        **kwargs (dict): Please refer to other kwargs in
-            [ConversableAgent](https://docs.ag2.ai/latest/docs/api-reference/autogen/ConversableAgent).
+            The limit only plays a role when human_input_mode is not "ALWAYS". \n
+        - **kwargs (dict): Please refer to other kwargs in
+            [ConversableAgent](https://docs.ag2.ai/latest/docs/api-reference/autogen/ConversableAgent). \n
         """
         super().__init__(
             name,

autogen/agentchat/chat.py

@@ -7,12 +7,14 @@
 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 ..code_utils import content_str
 from ..doc_utils import export_module
 from ..events.agent_events import PostCarryoverProcessingEvent
 from ..io.base import IOStream
@@ -24,26 +26,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."""
 
 
@@ -119,7 +133,10 @@ def _post_process_carryover_item(carryover_item):
     if isinstance(carryover_item, str):
         return carryover_item
     elif isinstance(carryover_item, dict) and "content" in carryover_item:
-        return str(carryover_item["content"])
+        content_value = carryover_item.get("content")
+        if isinstance(content_value, (str, list)) or content_value is None:
+            return content_str(content_value)
+        return str(content_value)
     else:
         return str(carryover_item)
 
@@ -141,37 +158,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 +236,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 +307,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]