Skip to content

Add LiteLLM AI gateway provider for LLM and embeddings#545

Open
RheagalFire wants to merge 1 commit into
neo4j:mainfrom
RheagalFire:feat/add-litellm-provider
Open

Add LiteLLM AI gateway provider for LLM and embeddings#545
RheagalFire wants to merge 1 commit into
neo4j:mainfrom
RheagalFire:feat/add-litellm-provider

Conversation

@RheagalFire

Copy link
Copy Markdown

Description

Add LiteLLM as an LLM and embeddings provider, giving users access to 100+ LLM providers (OpenAI, Anthropic, Google, Azure, AWS Bedrock, Ollama, Cohere, Mistral, and more) through a single unified interface.

What's included:

  • LiteLLMChat LLM provider (src/neo4j_graphrag/llm/litellm_llm.py) - supports both V1 (string input) and V2 (message list input) interfaces, structured output via response_format, and tool calling
  • LiteLLMEmbeddings embeddings provider (src/neo4j_graphrag/embeddings/litellm.py) - sync and async embedding generation
  • Optional dependency: pip install "neo4j-graphrag[litellm]" (litellm>=1.80.0,<1.87.0)
  • 23 unit tests covering both providers

Usage:

from neo4j_graphrag.llm import LiteLLMChat
from neo4j_graphrag.embeddings import LiteLLMEmbeddings

# Use any LiteLLM-supported model
llm = LiteLLMChat(model_name="gpt-4o", api_key="...")
llm = LiteLLMChat(model_name="anthropic/claude-sonnet-4-20250514", api_key="...")
llm = LiteLLMChat(model_name="ollama/llama3")

response = llm.invoke("Hello!")

# Embeddings
embedder = LiteLLMEmbeddings(model="text-embedding-ada-002", api_key="...")
vector = embedder.embed_query("Some text")

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

Manual testing performed with Azure Anthropic Foundry endpoint confirming V1 sync, V2 sync, and async invoke all work correctly with real API calls.

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

Signed-off-by: RheagalFire <arishalam121@gmail.com>
@RheagalFire RheagalFire requested a review from a team as a code owner June 16, 2026 18:56
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