Skip to content

refactor(llm): unify invoke/ainvoke dispatch across providers - #626

Open
matteomedioli wants to merge 3 commits into
mainfrom
matteo/gemini-llm-sync-async-core-split-GENKGB-1699
Open

matteomedioli wants to merge 3 commits into
mainfrom
matteo/gemini-llm-sync-async-core-split-GENKGB-1699

Conversation

@matteomedioli

@matteomedioli matteomedioli commented Sep 13, 2026

Copy link
Copy Markdown
Contributor

Description

Refactors LLMBase.invoke/ainvoke from an abstract dispatcher every subclass had to reimplement into a concrete dispatcher: it branches on input type (str vs list[LLMMessage]) and delegates to four narrow abstract hooks (_invoke_v1, _invoke_v2, _ainvoke_v1, _ainvoke_v2). All eight providers (AnthropicLLM, OpenAILLM, GeminiLLM, CohereLLM, MistralAILLM, OllamaLLM, BedrockLLM, VertexAILLM) migrated onto this contract, each now sharing one request builder and response parser per input version, plus a pair of _call_sync/_call_async transport hooks that differ only in whether the SDK call is awaited. This removes the duplicated sync/async dispatch and request/response logic that previously lived four times per provider.

Along the way, this fixes bugs the duplication had let drift out of sync:

  • BedrockLLM now inherits from LLMBase (previously LLMInterface/LLMInterfaceV2 directly), aligning its hierarchy with every other provider. Its async path no longer stacks a sync-side rate-limit retry inside the thread-pool-executed call on top of its own async retry layer.
  • OllamaLLM's async path now spreads model_params into the client.chat call the same way the sync path always did, instead of passing it verbatim as the options value — this fixed a double-nested options bug when model_params carried a sibling key (e.g. {"options": {...}, "format": "json"}).
  • CohereLLM's v2 error paths now raise LLMGenerationError(e) from e consistently, preserving the original SDK exception (previously a generic "Error calling cohere" message on v2 only).

Breaking (targets 2.0): a subclass overriding invoke/ainvoke directly (a supported extension point for BaseAnthropicLLM/BaseOpenAILLM/BaseGeminiLLM) must migrate to overriding the four hooks instead. The public invoke/ainvoke/invoke_with_tools/ainvoke_with_tools calling contract for callers is unchanged.

Note: GeminiLLM's usage-tracking fix (response.usage_metadata -> LLMResponse.usage) landed separately in #621 and is already on main; this branch is rebased on top of it, not re-introducing it.

Type of Change

  • Refactoring (no functional changes) — plus incidental bug fixes surfaced by the consolidation (see above)

Complexity

Complexity: Medium

How Has This Been Tested?

  • Unit 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
  • CLA (https://neo4j.com/developer/cla/) has been signed
  • CHANGELOG.md updated if appropriate

@matteomedioli matteomedioli changed the title fix(llm): populate LLMUsage from Gemini's usage_metadata (#621) refactor(llm): unify invoke/ainvoke dispatch across providers, fix Gemini usage tracking Sep 13, 2026
@matteomedioli
matteomedioli force-pushed the matteo/gemini-llm-sync-async-core-split-GENKGB-1699 branch 2 times, most recently from 7ce315b to 73ec23b Compare September 14, 2026 13:40
LLMBase.invoke/ainvoke were abstract dispatchers every subclass had to
reimplement. Makes them concrete: they branch on input type (str vs
list[LLMMessage]) and delegate to four narrow abstract hooks
(_invoke_v1, _invoke_v2, _ainvoke_v1, _ainvoke_v2). All eight providers
(AnthropicLLM, OpenAILLM, GeminiLLM, CohereLLM, MistralAILLM, OllamaLLM,
BedrockLLM, VertexAILLM) migrate onto this contract, each sharing one
request builder and response parser per input version plus a
_call_sync/_call_async transport pair that differs only in whether the
SDK call is awaited. Removes the duplicated sync/async dispatch and
request/response logic that previously lived four times per provider.

Also fixes bugs the duplication let drift out of sync:
- BedrockLLM now inherits from LLMBase (previously LLMInterface/
  LLMInterfaceV2 directly); its async path no longer stacks a
  sync-side retry inside the thread-pool call on top of its own async
  retry layer.
- OllamaLLM's async path spreads model_params into client.chat the
  same way the sync path always did, instead of passing it verbatim
  as the options value.
- CohereLLM's v2 error paths raise LLMGenerationError(e) from e
  consistently, preserving the original SDK exception.

Breaking (targets 2.0): a subclass overriding invoke/ainvoke directly
must migrate to overriding the four hooks instead. The public
invoke/ainvoke/invoke_with_tools/ainvoke_with_tools calling contract
for callers is unchanged.
@matteomedioli
matteomedioli force-pushed the matteo/gemini-llm-sync-async-core-split-GENKGB-1699 branch from 73ec23b to ce86eae Compare September 14, 2026 13:44
@matteomedioli matteomedioli changed the title refactor(llm): unify invoke/ainvoke dispatch across providers, fix Gemini usage tracking refactor(llm): unify invoke/ainvoke dispatch across providers Sep 14, 2026
@matteomedioli
matteomedioli deleted the matteo/gemini-llm-sync-async-core-split-GENKGB-1699 branch September 14, 2026 13:57
@matteomedioli
matteomedioli restored the matteo/gemini-llm-sync-async-core-split-GENKGB-1699 branch September 14, 2026 14:04
@matteomedioli matteomedioli reopened this Sep 14, 2026
@matteomedioli
matteomedioli marked this pull request as ready for review September 15, 2026 07:52
@matteomedioli
matteomedioli requested a review from a team as a code owner September 15, 2026 07:52
@matteomedioli
matteomedioli force-pushed the matteo/gemini-llm-sync-async-core-split-GENKGB-1699 branch from a82c87d to ee58fbd Compare September 15, 2026 12:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant