Graphrag integration

[lspinheiro]

Why are these changes needed?

This PR adds initial integration between graphrag and autogen by exposing local and global search as tools that can be used in autogen-agentchat. To be followed up with a user-guide/cookbook. I I added no tests because the test data I used was fairly large and I'm not sure we have a stablished way to add tests for those more complex integrations but there is a script below that I used. The indexing needs to be done in graphrag first, the goal is to illustrate the e2e steps in a notebook.

Would appreciate some initial feedback, hoping to gradually extend with more flexible configuration, integration of drift search and examples.

Related issue number

Checks

• I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[rysweet]

hi @lspinheiro - this is exciting. its also marked as DRAFT in the subject line but not marked as such in the PR - I'm marking as draft and please set it back by clicking Ready to Review when you are ready.

[ekzhu]

Exciting to see this!! I love the tool idea. The tool itself can also be stateful and shared by multiple agents.

[lspinheiro]

Thanks @ekzhu and @rysweet . This should be ready for review now. Still needs improvements as mentioned in the description, but the tools can be used. I used the following test script.

import asyncio
from autogen_core import CancellationToken
from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
from autogen_ext.tools.graphrag import (
    GlobalSearchTool,
    LocalSearchTool,
    GlobalDataConfig,
    LocalDataConfig,
    EmbeddingConfig,
)
from azure.identity import DefaultAzureCredential, get_bearer_token_provider


async def main():
    openai_client = AzureOpenAIChatCompletionClient(
        model="gpt-4o-mini",
        azure_endpoint="https://<resource-name>.openai.azure.com", 
        azure_deployment="gpt-4o-mini",
        api_version="2024-08-01-preview",
        azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
    )

    # Global search example
    global_config = GlobalDataConfig(
        input_dir="./autogen-test/ragtest/output"
    )
    
    global_tool = GlobalSearchTool.from_config(
        openai_client=openai_client,
        data_config=global_config
    )

    global_args = {
        "query": "What does the station-master says about Dr. Becher?"
    }

    global_result = await global_tool.run_json(global_args, CancellationToken())
    print("\nGlobal Search Result:")
    print(global_result)
    
    # Local search example
    local_config = LocalDataConfig(
        input_dir="./autogen-test/ragtest/output"
    )

    embedding_config = EmbeddingConfig(
        model="text-embedding-3-small",
        api_base="https://<resource-name>.openai.azure.com", 
        deployment_name="text-embedding-3-small",
        api_version="2023-05-15",
        api_type="azure",
        azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"),
        max_retries=10,
        request_timeout=180.0,
    )

    local_tool = LocalSearchTool.from_config(
        openai_client=openai_client,
        data_config=local_config,
        embedding_config=embedding_config
    )

    local_args = {
        "query": "What does the station-master says about Dr. Becher?"
    }

    local_result = await local_tool.run_json(local_args, CancellationToken())
    print("\nLocal Search Result:")
    print(local_result)


if __name__ == "__main__":
    asyncio.run(main())

[lspinheiro]

@jackgerrits , I had to add verride-dependencies for pydantic and tenacity because the current version of pydantic is below their minimum requirement and there is a conflict with llamaindex which requires a lower version of tenancity, but it is a dev dependency for us. Let me know if you have any concerns with the approach

[gagb]

Thank you! More documentation would help me review this PR. I would like to be able to build the docs page on this PR and see the example.

[gagb]

Related #4438

[lspinheiro]

Thank you! More documentation would help me review this PR. I would like to be able to build the docs page on this PR and see the example.

@gagb , I added a sample with a readme and some docstrings that should help with the review.

[ekzhu]

Is there a minimum example for the settings.yaml that can work with a few text documents?

[ekzhu]

We would like to have a dev to just copy and run the example in a script and it will work out of box.

So let's also include the installation instructions:

pip install "autogen-agentchat==0.4.0.dev12" "autogen-ext[graphrag]==0.4.0.dev12"

And have a text file downloaded from an online URL. e.g., the README.md file in our repo, with GraphRAG commands that build an index.

Then show an example of a YAML file for setttings.

Then the Python script that runs the agent with the tool.

[lspinheiro]

Is there a minimum example for the settings.yaml that can work with a few text documents?

@ekzhu I wanted to avoid bringing the graphrag setup steps over in an intentional way. The set up process is reasonably complex and requires initialization (which generates the settings), setting up data, tuning prompts and then indexing. My thought was that it would be best to direct the users to the graphrag docs to understand how to go through the initial steps and then, once indexing is done, they can easily integrate with autogen with the tools.

The reasoning was that it is quite easy for things to go wrong in those steps and if I think users should look for the solutions to problems they may during setup in the graphrag repo and docs and adding steps that may not work out-of-the-box may encourage them to ask for help in autogen's issue. What do you think of that? Maybe we can expand the note to make this clearer?