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.

[ekzhu]

Since we have already adopted so much of GraphRAG's components here, do you think it makes sense to just use the BaseLLM type instead of the ChatCompletionClient type and creating an adaptor?

[lspinheiro]

One of the things I tried to do is add a factory method where the user doesn't need to be exposed to graphrag components with from_config. This is what I'm using in the example script. The constructor right now accepts the graphrag components to allow more flexibility and to allow users to create the tool even if we forgot to add support to some customisation option in the factory. Still needs to be improved, but that is the goal. In the script I can create the tool just from our llm client + configuration options.

[ekzhu]

I think it makes sense though I am hoping to avoid too much code bloat. Checking with @jackgerrits

[lspinheiro]

Alright, I replaced it with a from_settings method that uses the graphrag configuration file. Should be much clear now.

[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.