Skip to content

feat(llm): extract BaseVertexAILLM around a model-construction hook - #572

Draft
matteomedioli wants to merge 3 commits into
mainfrom
matteo/base-vertexai-llm
Draft

feat(llm): extract BaseVertexAILLM around a model-construction hook#572
matteomedioli wants to merge 3 commits into
mainfrom
matteo/base-vertexai-llm

Conversation

@matteomedioli

@matteomedioli matteomedioli commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Stack

Merge order: #565#567#566#571#572. Each PR builds on the previous one (this branch is stacked on matteo/base-gemini-llm). Until its base merges, GitHub may show parent commits in the diff — review only the commits unique to this branch.

Description

Same goal as the other base classes — subclass instead of fork — but VertexAI needs a different shape. There is no client object to construct: the SDK works through a global vertexai.init(...) call plus a fresh GenerativeModel created for every request.

So instead of "the subclass creates the clients", BaseVertexAILLM has one abstract method, _get_model(), which returns the model used for a call. That is the only thing a subclass implements — pointing at a different project, region, or endpoint is a few lines. Everything else (message building, config and schema handling, response parsing) is inherited.

  • VertexAILLM implements _get_model() exactly as it did before the extraction — no behavior change.
  • BaseVertexAILLM is exported from neo4j_graphrag.llm; the llm.rst extensibility page now has a Vertex AI section explaining why this hook differs from the Anthropic/OpenAI/Gemini client-injection shape (no http_client/base_url — endpoint config is global via vertexai.init), with a worked subclass example. api.rst cross-references it.
  • The docs note that new projects should prefer GeminiLLM (google-genai SDK), positioning VertexAILLM as legacy support — resolving the maintainer question below in the docs.
  • Tests cover the contract directly: the base class cannot be instantiated without _get_model, and a minimal subclass overriding only _get_model runs invoke() end to end with its custom model.

Question for maintainers: Google is retiring the vertexai.generative_models SDK in favor of google-genai — which GeminiLLM already uses, and which can also reach Vertex-hosted models. This PR still makes sense (it cleanly factors what exists); the docs now recommend GeminiLLM for new work. Shout if you'd rather drop that note or phrase it differently.

Type of Change

  • New feature
  • Bug fix
  • Breaking change
  • Documentation update
  • Project configuration change

Complexity

Complexity: Low

How Has This Been Tested?

  • Unit tests
  • E2E tests
  • Manual tests

Checklist

  • Documentation has been updated
  • Unit tests have been updated
  • Examples have been updated
  • CLA (https://neo4j.com/developer/cla/) has been signed
  • CHANGELOG.md updated

@matteomedioli
matteomedioli requested a review from a team as a code owner July 17, 2026 11:40
@matteomedioli
matteomedioli marked this pull request as draft July 17, 2026 11:41
@matteomedioli
matteomedioli force-pushed the matteo/split-http-client-helper branch from 0b4426b to 4baa924 Compare July 17, 2026 15:02
@matteomedioli
matteomedioli force-pushed the matteo/base-vertexai-llm branch from 5df6253 to e75b657 Compare July 17, 2026 15:04
@matteomedioli
matteomedioli force-pushed the matteo/split-http-client-helper branch from 4baa924 to 2b6a1f6 Compare July 17, 2026 15:07
@matteomedioli
matteomedioli force-pushed the matteo/base-vertexai-llm branch from e75b657 to 7575cd1 Compare July 17, 2026 15:07
@matteomedioli
matteomedioli force-pushed the matteo/split-http-client-helper branch from 2b6a1f6 to 5a6541c Compare July 17, 2026 15:21
@matteomedioli
matteomedioli force-pushed the matteo/base-vertexai-llm branch from 7575cd1 to 0fa27e8 Compare July 17, 2026 15:21
@matteomedioli
matteomedioli force-pushed the matteo/split-http-client-helper branch from 5a6541c to a819760 Compare July 22, 2026 11:42
Base automatically changed from matteo/split-http-client-helper to main July 22, 2026 13:47
@matteomedioli
matteomedioli changed the base branch from main to matteo/base-gemini-llm July 22, 2026 17:34
@matteomedioli
matteomedioli force-pushed the matteo/base-gemini-llm branch from f943b3e to 5ca3ff9 Compare July 23, 2026 12:24
Base automatically changed from matteo/base-gemini-llm to main July 24, 2026 08:39
VertexAILLM has no persistent per-instance SDK client (unlike Anthropic/
OpenAI/Gemini) -- it relies on a global vertexai.init(...) plus a fresh
GenerativeModel built per call. So instead of a client-construction hook,
BaseVertexAILLM declares _get_model(...) as its one abstract method;
VertexAILLM implements it exactly as before, and everything else
(message building, generation-config/schema handling, response parsing)
moves to the base class unchanged.

BaseVertexAILLM is exported from neo4j_graphrag.llm as a documented
extension point, mirroring BaseAnthropicLLM/BaseOpenAILLM/BaseGeminiLLM
but with the mechanism this provider's SDK actually supports.
… to end

A minimal subclass overriding only _get_model now proves its custom
model is the one invoke() calls, with message building and response
parsing inherited — the contract exported to custom subclasses.
Bound methods typed as Callable expose no __func__ under strict mypy
(attr-defined, x4) — this fails CI's 'mypy .' gate. The end-to-end
minimal-subclass test already covers the inheritance contract these
asserts approximated.
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