neo4j-graphrag-python icon indicating copy to clipboard operation
neo4j-graphrag-python copied to clipboard

Fix VertexAILLM

Open stellasia opened this issue 6 months ago • 0 comments

Description

The main reason for this PR is that, even if Google GenerativeModel accepts a list of tools as parameter, this is not how we should provide multiple tools. We should instead provide a single tool with multiple function declarations. This PR addresses this behavior.

It also addresses a few other problems related to the way GenerativeModel is instantiated. The usage is:

    # config for chat mode (aka invoke)
    generation_config = GenerationConfig(
        temperature=0.0,
        response_mime_type="application/json",
    )
    # config for tool mode (aka invoke_with_tools)
    # optional
    tool_config = ToolConfig(
        function_calling_config=ToolConfig.FunctionCallingConfig(
            mode=ToolConfig.FunctionCallingConfig.Mode.ANY
        ),
    )

    llm = VertexAILLM(
        model_name="gemini-2.0-flash-001",
        generation_config=generation_config,
        tool_config=tool_config,
    )

Then, internally:

  • if invoke mode: we drop the tool config (it's not useful)
  • if invoke_with_tool mode, we drop the generation config (otherwise the presence of fields like response_mime_type/response_schema raise errors)

Drawback of this approach:

  • We can't have a tool_config per call, which means we can't configure the allowed_function_names parameter per call.

We can't move these parameters to the invoke* methods because then they would be "provider-dependent" and LLMInterfaces wouldn't be interchangeable anymore, which defeats the whole point of the interface.

Type of Change

  • [ ] New feature
  • [x] Bug fix
  • [ ] Breaking change
  • [ ] Documentation update
  • [ ] Project configuration change

Complexity

Complexity: ?

How Has This Been Tested?

  • [x] Unit tests
  • [ ] E2E tests
  • [x] Manual tests

Checklist

The following requirements should have been met (depending on the changes in the branch):

  • [ ] Documentation has been updated
  • [ ] Unit tests have been updated
  • [ ] E2E tests have been updated
  • [ ] Examples have been updated
  • [ ] New files have copyright header
  • [x] CLA (https://neo4j.com/developer/cla/) has been signed
  • [x] CHANGELOG.md updated if appropriate

stellasia avatar May 20 '25 13:05 stellasia