Graphrag integration

python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py

@@ -5,8 +5,6 @@
 import tiktoken
 from autogen_core import CancellationToken
 from autogen_core.tools import BaseTool
-from pydantic import BaseModel, Field
-
 from graphrag.config.config_file_loader import load_config_from_file
 from graphrag.query.indexer_adapters import (
     read_indexer_communities,
@@ -17,6 +15,7 @@
 from graphrag.query.llm.get_client import get_llm
 from graphrag.query.structured_search.global_search.community_context import GlobalCommunityContext
 from graphrag.query.structured_search.global_search.search import GlobalSearch
+from pydantic import BaseModel, Field
 
 from ._config import GlobalContextConfig as ContextConfig
 from ._config import GlobalDataConfig as DataConfig
@@ -39,12 +38,54 @@ class GlobalSearchTool(BaseTool[GlobalSearchToolArgs, GlobalSearchToolReturn]):
 
     This tool allows you to perform semantic search over a corpus of documents using the GraphRAG framework.
     The search combines graph-based document relationships with semantic embeddings to find relevant information.
+    Example usage with AssistantAgent:
+
+    .. code-block:: python
+
+        import asyncio
+        from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
+        from autogen_ext.tools.graphrag import GlobalSearchTool
+        from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+        from autogen_agentchat.agents import AssistantAgent
+
+
+        async def main():
+            # Initialize the OpenAI client
+            openai_client = OpenAIChatCompletionClient(
+                model="gpt-4o-mini",
+                api_key="<api-key>",
+            )
+
+            # Set up global search tool
+            global_tool = GlobalSearchTool.from_settings(settings_path="./settings.yaml")
+
+            # Create assistant agent with the global search tool
+            assistant_agent = AssistantAgent(
+                name="search_assistant",
+                tools=[global_tool],
+                model_client=openai_client,
+                system_message=(
+                    "You are a tool selector AI assistant using the GraphRAG framework. "
+                    "Your primary task is to determine the appropriate search tool to call based on the user's query. "
+                    "For broader, abstract questions requiring a comprehensive understanding of the dataset, call the 'global_search' function."
+                ),
+            )
+
+            # Run a sample query
+            query = "What is the overall sentiment of the community reports?"
+            response_stream = assistant_agent.run_stream(task=query)
+            async for msg in response_stream:
+                if hasattr(msg, "content"):
+                    print(f"\nAgent response: {msg.content}")
+
+
+        if __name__ == "__main__":
+            asyncio.run(main())
 
     .. note::
 
         This tool requires the :code:`graphrag` extra for the :code:`autogen-ext` package.
-
-        This tool requires indexed data created by the GraphRAG indexing process. See the GraphRAG documentation
+        This tool requires indexed data created by the GraphRAG indexing process. See the [GraphRAG documentation](https://microsoft.github.io/graphrag/)
         for details on how to prepare the required data files.
 
 

python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py

@@ -6,8 +6,6 @@
 import tiktoken
 from autogen_core import CancellationToken
 from autogen_core.tools import BaseTool
-from pydantic import BaseModel, Field
-
 from graphrag.config.config_file_loader import load_config_from_file
 from graphrag.query.indexer_adapters import (
     read_indexer_entities,
@@ -19,6 +17,7 @@
 from graphrag.query.structured_search.local_search.mixed_context import LocalSearchMixedContext
 from graphrag.query.structured_search.local_search.search import LocalSearch
 from graphrag.vector_stores.lancedb import LanceDBVectorStore
+from pydantic import BaseModel, Field
 
 from ._config import LocalContextConfig, SearchConfig
 from ._config import LocalDataConfig as DataConfig
@@ -40,12 +39,54 @@ class LocalSearchTool(BaseTool[LocalSearchToolArgs, LocalSearchToolReturn]):
 
     This tool allows you to perform semantic search over a corpus of documents using the GraphRAG framework.
     The search combines local document context with semantic embeddings to find relevant information.
+    Example usage with AssistantAgent:
+
+    .. code-block:: python
+
+        import asyncio
+        from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
+        from autogen_ext.tools.graphrag import LocalSearchTool
+        from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+        from autogen_agentchat.agents import AssistantAgent
+
+
+        async def main():
+            # Initialize the OpenAI client
+            openai_client = OpenAIChatCompletionClient(
+                model="gpt-4o-mini",
+                api_key="<api-key>",
+            )
+
+            # Set up local search tool
+            local_tool = LocalSearchTool.from_settings(settings_path="./settings.yaml")
+
+            # Create assistant agent with the local search tool
+            assistant_agent = AssistantAgent(
+                name="search_assistant",
+                tools=[local_tool],
+                model_client=openai_client,
+                system_message=(
+                    "You are a tool selector AI assistant using the GraphRAG framework. "
+                    "Your primary task is to determine the appropriate search tool to call based on the user's query. "
+                    "For specific, detailed information about particular entities or relationships, call the 'local_search' function."
+                ),
+            )
+
+            # Run a sample query
+            query = "What does the station-master say about Dr. Becher?"
+            response_stream = assistant_agent.run_stream(task=query)
+            async for msg in response_stream:
+                if hasattr(msg, "content"):
+                    print(f"\nAgent response: {msg.content}")
+
+
+        if __name__ == "__main__":
+            asyncio.run(main())
 
     .. note::
 
         This tool requires the :code:`graphrag` extra for the :code:`autogen-ext` package.
-
-        This tool requires indexed data created by the GraphRAG indexing process. See the GraphRAG documentation
+        This tool requires indexed data created by the GraphRAG indexing process. See the [GraphRAG documentation](https://microsoft.github.io/graphrag/)
         for details on how to prepare the required data files.
 
     Args:

python/pyproject.toml

@@ -22,8 +22,9 @@ dev = [
 
 override-dependencies = [
     "tenacity==9.0.0",
-    "pydantic==2.9.2",
 ]
+[tool.uv.workspace]
+members = ["packages/*"]
 
 [tool.uv.sources]
 agbench = { workspace = true }