Add LiteLLM AI gateway provider for LLM and embeddings#545
Open
RheagalFire wants to merge 2 commits into
Open
Conversation
Signed-off-by: RheagalFire <arishalam121@gmail.com>
ali-sedaghatbaf
requested changes
Jun 29, 2026
| ): | ||
| call_kwargs["response_format"] = { | ||
| "type": "json_schema", | ||
| "json_schema": { |
Contributor
There was a problem hiding this comment.
"strict": True is missing in json_schema
| llm = LiteLLMChat(model_name="ollama/llama3") | ||
| """ | ||
|
|
||
| supports_structured_output = True |
Contributor
There was a problem hiding this comment.
Not all providers support structured output.
| else: | ||
| raise ValueError(f"Invalid input type for ainvoke method - {type(input)}") | ||
|
|
||
| def invoke_with_tools( |
Contributor
There was a problem hiding this comment.
lacks rate limit handling
| content=res.choices[0].message.content, | ||
| ) | ||
|
|
||
| async def ainvoke_with_tools( |
Contributor
There was a problem hiding this comment.
lacks rate limit handling
|
|
||
|
|
||
| # pylint: disable=redefined-builtin, arguments-differ, raise-missing-from, no-else-return, import-outside-toplevel | ||
| class LiteLLMChat(LLMBase): |
Contributor
There was a problem hiding this comment.
Docs page is missing for this provider
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| # Generated by Claude Code |
Contributor
There was a problem hiding this comment.
Shouldn't be committed!
- Add "strict": True to json_schema response_format in _invoke_v2 and _ainvoke_v2, matching the OpenAI provider pattern - Set supports_structured_output = False since not all LiteLLM-supported providers handle structured output - Add @rate_limit_handler_decorator to invoke_with_tools and @async_rate_limit_handler_decorator to ainvoke_with_tools for consistent rate limit handling across all methods - Add LiteLLM provider documentation to user guide and API reference - Remove "Generated by Claude Code" comments from test files
RheagalFire
force-pushed
the
feat/add-litellm-provider
branch
from
dd1905f to
e17683f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
LiteLLMChatLLM provider (src/neo4j_graphrag/llm/litellm_llm.py) - supports both V1 (string input) and V2 (message list input) interfaces, structured output viaresponse_format, and tool callingLiteLLMEmbeddingsembeddings provider (src/neo4j_graphrag/embeddings/litellm.py) - sync and async embedding generationpip install "neo4j-graphrag[litellm]"(litellm>=1.80.0,<1.87.0)Usage:
Type of Change
Complexity
Complexity: Low
How Has This Been Tested?
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):