Skip to content

Commit e17683f

Browse files
RheagalFireAarish Irani
authored andcommitted
Address PR review feedback for LiteLLM provider
- 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
1 parent ed935d4 commit e17683f

5 files changed

Lines changed: 61 additions & 5 deletions

File tree

docs/source/api.rst

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,12 @@ CohereEmbeddings
318318
.. autoclass:: neo4j_graphrag.embeddings.cohere.CohereEmbeddings
319319
:members:
320320

321+
LiteLLMEmbeddings
322+
=================
323+
324+
.. autoclass:: neo4j_graphrag.embeddings.litellm.LiteLLMEmbeddings
325+
:members:
326+
321327
**********
322328
Generation
323329
**********
@@ -402,6 +408,13 @@ MistralAILLM
402408
:members:
403409

404410

411+
LiteLLMChat
412+
-----------
413+
414+
.. autoclass:: neo4j_graphrag.llm.litellm_llm.LiteLLMChat
415+
:members:
416+
417+
405418
Rate Limiting
406419
=============
407420

docs/source/user_guide_rag.rst

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ If OpenAI cannot be used directly, there are a few available alternatives:
7979
- Use Anthropic LLM (Claude...).
8080
- Use Mistral LLM
8181
- Use Cohere.
82+
- Use LiteLLM (100+ providers via a unified interface).
8283
- Use a local Ollama model.
8384
- Implement a custom interface.
8485
- Utilize any LangChain chat model.
@@ -214,6 +215,47 @@ To use Cohere, instantiate the `CohereLLM` class:
214215
See :ref:`coherellm`.
215216

216217

218+
Using LiteLLM (100+ Providers)
219+
-------------------------------
220+
221+
LiteLLM provides a unified interface to 100+ LLM providers (OpenAI, Anthropic,
222+
Google, Azure, AWS Bedrock, Ollama, and more) using the OpenAI completion format.
223+
To use LiteLLM, instantiate the `LiteLLMChat` class:
224+
225+
.. code:: python
226+
227+
from neo4j_graphrag.llm import LiteLLMChat
228+
229+
# Use any LiteLLM-supported provider
230+
llm = LiteLLMChat(model_name="gpt-4o", api_key="...")
231+
llm.invoke("say something")
232+
233+
# Anthropic
234+
llm = LiteLLMChat(model_name="anthropic/claude-sonnet-4-20250514", api_key="...")
235+
236+
# Ollama (local)
237+
llm = LiteLLMChat(model_name="ollama/llama3")
238+
239+
# Azure OpenAI
240+
llm = LiteLLMChat(
241+
model_name="azure/my-deployment",
242+
api_key="...",
243+
api_base="https://my-endpoint.openai.azure.com/",
244+
)
245+
246+
# AWS Bedrock
247+
llm = LiteLLMChat(model_name="bedrock/anthropic.claude-v2")
248+
249+
250+
.. note::
251+
252+
In order to run this code, the `litellm` Python package needs to be installed:
253+
`pip install "neo4j_graphrag[litellm]"`
254+
255+
256+
See :ref:`litellmchat`.
257+
258+
217259
Using a Local Model via Ollama
218260
-------------------------------
219261

@@ -620,6 +662,7 @@ Currently, this package supports the following embedders:
620662
- :ref:`cohereembeddings`
621663
- :ref:`azureopenaiembeddings`
622664
- :ref:`ollamaembeddings`
665+
- :ref:`litellmembeddings`
623666

624667
The `OpenAIEmbeddings` was illustrated previously. Here is how to use the `SentenceTransformerEmbeddings`:
625668

src/neo4j_graphrag/llm/litellm_llm.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ class LiteLLMChat(LLMBase):
8888
llm = LiteLLMChat(model_name="ollama/llama3")
8989
"""
9090

91-
supports_structured_output = True
91+
supports_structured_output = False
9292

9393
def __init__(
9494
self,
@@ -145,6 +145,7 @@ async def ainvoke(
145145
else:
146146
raise ValueError(f"Invalid input type for ainvoke method - {type(input)}")
147147

148+
@rate_limit_handler_decorator
148149
def invoke_with_tools(
149150
self,
150151
input: str,
@@ -193,6 +194,7 @@ def invoke_with_tools(
193194
content=res.choices[0].message.content,
194195
)
195196

197+
@async_rate_limit_handler_decorator
196198
async def ainvoke_with_tools(
197199
self,
198200
input: str,
@@ -284,6 +286,7 @@ def _invoke_v2(
284286
"type": "json_schema",
285287
"json_schema": {
286288
"name": response_format.__name__,
289+
"strict": True,
287290
"schema": response_format.model_json_schema(),
288291
},
289292
}
@@ -340,6 +343,7 @@ async def _ainvoke_v2(
340343
"type": "json_schema",
341344
"json_schema": {
342345
"name": response_format.__name__,
346+
"strict": True,
343347
"schema": response_format.model_json_schema(),
344348
},
345349
}

tests/unit/embeddings/test_litellm_embedder.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
# Generated by Claude Code
17-
1816
import sys
1917
from typing import Generator
2018
from unittest.mock import AsyncMock, MagicMock, Mock, patch

tests/unit/llm/test_litellm_llm.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
# Generated by Claude Code
17-
1816
import sys
1917
import warnings
2018
from typing import Generator, List

0 commit comments

Comments
 (0)