Add lang chain lite llm integration support #135
Conversation
- Add LangChainLiteModel class using langchain-litellm wrapper - Enables LangChain ecosystem integration alongside existing LiteModel - Supports chat, streaming, proxy configuration (api_key/api_base)
- Add unit tests - Update README with usage examples
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds a LangChain-backed LiteLLM wrapper and chat session classes, documentation for using LangChain LiteLLM, tests for initialization/serialization/chat flow/streaming, and new runtime dependencies ( Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Model as LangChainLiteModel
participant Session as LangChainLiteModelChatSession
participant LC as LangChain ChatLiteLLM
participant Provider as Provider API
User->>Model: start_chat(system_instruction?)
Model->>Model: lazy-import ChatLiteLLM
Model->>LC: init(model_name, api_key, api_base, params, gen_config)
Model-->>User: ChatSession
User->>Session: send_message(content)
Session->>Session: append HumanMessage
Session->>LC: generate(messages, config)
LC->>Provider: HTTP request
Provider-->>LC: response
LC-->>Session: AIMessage
Session->>Session: append AIMessage
Session-->>User: GenerationResponse
opt Streaming
User->>Session: send_message_stream(content)
Session->>LC: stream(messages, config)
loop tokens
LC-->>Session: partial token
Session-->>User: yield partial
end
Session->>Session: append full AIMessage
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/test_langchain_litellm.py (1)
91-96: Implement or remove the placeholder test.This test function currently does nothing (just
pass). The comment indicates it should test the ImportError handling when langchain-litellm is not installed, but this is not implemented.Would you like me to generate a test implementation that mocks the import failure to verify the error handling? For example:
def test_langchain_litellm_missing_dependency(): """Test that appropriate error is raised when langchain-litellm is not installed.""" import sys from unittest.mock import patch # Mock the import to simulate missing package with patch.dict(sys.modules, {'langchain_litellm': None}): with pytest.raises(ImportError) as exc_info: # Force reimport to trigger the error import importlib import graphrag_sdk.models.langchain_litellm importlib.reload(graphrag_sdk.models.langchain_litellm) assert "langchain-litellm package is required" in str(exc_info.value)Or would you prefer to remove this placeholder test entirely?
graphrag_sdk/models/langchain_litellm.py (1)
45-51: Add exception chaining to ImportError.The ImportError should use
raise ... from errto preserve the original exception context and maintain a complete traceback.As per static analysis hints (Ruff B904).
Apply this diff:
try: from langchain_litellm import ChatLiteLLM - except ImportError: + except ImportError as e: raise ImportError( "langchain-litellm package is required for LangChainLiteModel. " "Install it with: pip install langchain-litellm" - ) + ) from e
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
README.md(1 hunks)graphrag_sdk/models/langchain_litellm.py(1 hunks)pyproject.toml(1 hunks)requirements.txt(1 hunks)tests/test_langchain_litellm.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
tests/test_langchain_litellm.py (2)
graphrag_sdk/models/langchain_litellm.py (2)
LangChainLiteModel(16-161)LangChainLiteModelChatSession(164-285)graphrag_sdk/models/model.py (1)
GenerativeModelConfig(10-79)
graphrag_sdk/models/langchain_litellm.py (1)
graphrag_sdk/models/model.py (5)
GenerativeModel(110-126)GenerativeModelConfig(10-79)GenerationResponse(81-90)FinishReason(5-8)GenerativeModelChatSession(93-107)
🪛 Ruff (0.13.3)
graphrag_sdk/models/langchain_litellm.py
48-51: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
48-51: Avoid specifying long messages outside the exception class
(TRY003)
81-83: Avoid specifying long messages outside the exception class
(TRY003)
207-209: Avoid specifying long messages outside the exception class
(TRY003)
240-242: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (8)
pyproject.toml (1)
25-26: LGTM! Dependency versions are current and compatible.The langchain-litellm 0.2.2 dependency is appropriate and includes recent fixes for streaming and tool message parsing. The version constraints follow best practices by allowing patch updates while preventing breaking changes.
Based on learnings.
requirements.txt (1)
12-13: LGTM! Versions are consistent with pyproject.toml.The dependency declarations correctly mirror the pyproject.toml configuration with appropriate version constraints.
tests/test_langchain_litellm.py (1)
1-159: LGTM! Comprehensive test coverage for the new integration.The test suite provides excellent coverage of the LangChainLiteModel functionality, including initialization paths, serialization, chat sessions, history management, and properly gated integration tests for actual API calls. The tests validate the contract and expected behavior effectively.
graphrag_sdk/models/langchain_litellm.py (4)
78-83: Add exception chaining to ValueError.When re-raising an exception, preserve the original exception context with
raise ... from eto maintain the complete error chain.As per static analysis hints (Ruff B904).
Apply this diff:
try: self._chat_model = ChatLiteLLM(**chat_params) except Exception as e: raise ValueError( f"Failed to initialize ChatLiteLLM with model '{model_name}': {e}" ) from eNote: This already has
from eon line 83, so this is already correct. The static analysis hint appears to be a false positive.
16-161: LGTM! Well-structured model implementation.The LangChainLiteModel class provides a clean abstraction over LangChain's ChatLiteLLM wrapper, with proper initialization, error handling, serialization support, and chat session management. The implementation follows the GenerativeModel interface correctly.
169-190: LGTM! Clean initialization with dynamic imports.The initialization properly sets up the chat session with conversation history and system instruction. The approach of dynamically importing and storing message class references is elegant and avoids module-level import dependencies.
261-285: LGTM! Correct history deletion logic.The delete_last_message method correctly handles edge cases: it preserves the system message, safely removes the last user/assistant exchange, and handles empty or minimal history gracefully. The logic is sound despite its complexity.
README.md (1)
219-282: Duplication check passed
The "Using LangChain LiteLLM Integration" section appears only once in README.md.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
tests/test_langchain_litellm.py (1)
91-96: Consider implementing or removing the placeholder test.This test has an empty body and simply passes. Either implement proper dependency isolation testing (e.g., using mocking to simulate ImportError) or remove the placeholder.
Example implementation using mocking:
def test_langchain_litellm_missing_dependency(): - """Test that appropriate error is raised when langchain-litellm is not installed.""" - # This test would need to mock the import failure - # For now, we just ensure the import check exists in the code - # The actual ImportError handling is tested implicitly by the other tests - pass + """Test that appropriate error is raised when langchain-litellm is not installed.""" + import sys + from unittest.mock import patch + + # Mock the import to simulate missing dependency + with patch.dict(sys.modules, {'langchain_litellm': None}): + with pytest.raises(ImportError, match="langchain-litellm package is required"): + from graphrag_sdk.models.langchain_litellm import LangChainLiteModel + LangChainLiteModel()graphrag_sdk/models/langchain_litellm.py (2)
24-84: Consider the parameter merge order to avoid unexpected overwrites.The current implementation merges parameters in this order:
- Explicit
model_name,api_key,api_base(lines 61-68)additional_params(line 71)generation_config(line 76)If
additional_paramsorgeneration_configcontain keys like"model","api_key", or"api_base", they will overwrite the explicitly set values. This could lead to surprising behavior.Consider applying updates in reverse order or filtering out protected keys:
# Initialize the ChatLiteLLM model chat_params = { "model": model_name, } +# Add generation config and additional params first (lower priority) +gen_config = self.generation_config.to_json() +if gen_config: + # Filter out keys that should not override explicit params + protected_keys = {"model", "api_key", "api_base"} + gen_config = {k: v for k, v in gen_config.items() if k not in protected_keys} + chat_params.update(gen_config) + +# Merge additional params (medium priority) +if self.additional_params: + additional_filtered = {k: v for k, v in self.additional_params.items() if k not in protected_keys} + chat_params.update(additional_filtered) + +# Explicit parameters take highest priority if api_key: chat_params["api_key"] = api_key if api_base: chat_params["api_base"] = api_base - -# Merge additional params -chat_params.update(self.additional_params) - -# Add generation config parameters -gen_config = self.generation_config.to_json() -if gen_config: - chat_params.update(gen_config)
45-51: Optional: Add exception chaining for better error context.The static analysis tool suggests using
raise ... from errto maintain the exception chain, which helps with debugging by preserving the original error context.As per static analysis hints.
- except ImportError: - raise ImportError( + except ImportError as e: + raise ImportError( "langchain-litellm package is required for LangChainLiteModel. " "Install it with: pip install langchain-litellm" - ) + ) from e
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
README.md(1 hunks)graphrag_sdk/models/langchain_litellm.py(1 hunks)pyproject.toml(1 hunks)requirements.txt(1 hunks)tests/test_langchain_litellm.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
tests/test_langchain_litellm.py (2)
graphrag_sdk/models/langchain_litellm.py (2)
LangChainLiteModel(16-161)LangChainLiteModelChatSession(164-285)graphrag_sdk/models/model.py (1)
GenerativeModelConfig(10-79)
graphrag_sdk/models/langchain_litellm.py (1)
graphrag_sdk/models/model.py (5)
GenerativeModel(110-126)GenerativeModelConfig(10-79)GenerationResponse(81-90)FinishReason(5-8)GenerativeModelChatSession(93-107)
🪛 Ruff (0.13.3)
graphrag_sdk/models/langchain_litellm.py
48-51: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
48-51: Avoid specifying long messages outside the exception class
(TRY003)
81-83: Avoid specifying long messages outside the exception class
(TRY003)
207-209: Avoid specifying long messages outside the exception class
(TRY003)
240-242: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (21)
pyproject.toml (1)
25-26: LGTM!The dependency versions are appropriate and compatible with the latest releases (langchain-litellm 0.2.3 and langchain-core 0.3.78). The caret constraints will allow patch and minor updates while preventing breaking changes.
Based on learnings.
requirements.txt (1)
12-13: LGTM!The version constraints are consistent with pyproject.toml and follow semantic versioning best practices.
README.md (1)
219-280: LGTM!The documentation clearly explains the new LangChain LiteLLM integration, provides guidance on when to use it versus direct LiteLLM, and includes practical examples covering basic usage, proxy configuration, and OpenRouter integration.
tests/test_langchain_litellm.py (9)
11-21: LGTM!The test validates basic initialization and configuration handling correctly.
23-34: LGTM!The test properly verifies API parameter handling.
36-49: LGTM!JSON serialization test correctly validates the model's configuration is preserved.
51-67: LGTM!JSON deserialization test correctly validates model reconstruction from configuration.
69-77: LGTM!The test correctly validates chat session creation and initial state.
79-89: LGTM!The test validates initial chat history state correctly.
99-122: LGTM!Integration test is properly gated with environment variable check and validates end-to-end message sending functionality.
124-143: LGTM!Integration test properly validates streaming functionality with environment gating.
145-159: LGTM!The test correctly validates message deletion behavior while preserving system messages.
graphrag_sdk/models/langchain_litellm.py (9)
1-14: LGTM!Standard imports and logger setup are appropriate.
85-97: LGTM!The
start_chatmethod correctly creates a new chat session with optional system instruction override.
99-123: LGTM!Response parsing correctly maps LangChain's finish_reason values to the GenerationResponse format, with proper handling of "length" → MAX_TOKENS and "stop" → STOP mappings.
125-161: Verify the security implications of serializing API credentials.The
to_jsonmethod includesapi_keyandapi_basein the serialized output (lines 136-137). This could expose sensitive credentials if the JSON is logged, stored in version control, or transmitted insecurely.Consider whether API credentials should be serialized. If persistence is required, ensure:
- The JSON is stored securely (encrypted at rest, restricted permissions)
- Documentation warns users about credential exposure
- Consider alternative approaches like environment variable references
If credentials shouldn't be persisted, apply this diff:
def to_json(self) -> dict: """ Serialize the model's configuration and state to JSON format. + + Note: API credentials (api_key, api_base) are excluded for security. + You must provide them again when deserializing. Returns: dict: The serialized JSON data. """ return { "model_name": self.model_name, "generation_config": self.generation_config.to_json(), "system_instruction": self.system_instruction, - "api_key": self.api_key, - "api_base": self.api_base, + # Exclude credentials for security - must be provided at runtime "additional_params": self.additional_params, }And update
from_jsonto handle missing credentials gracefully.
169-190: LGTM!Chat session initialization correctly sets up the message history with an optional system message and stores LangChain message type references for later use.
192-213: LGTM!The
send_messagemethod correctly appends the user message, invokes the model, parses the response, and maintains chat history. Error handling provides clear context.
215-242: LGTM!The streaming implementation correctly accumulates chunks and preserves the complete response in chat history. Error handling is appropriate.
244-259: LGTM!The
get_chat_historymethod correctly converts internal LangChain message objects to a standardized dictionary format with role and content fields.
261-285: LGTM!The
delete_last_messagemethod correctly handles edge cases:
- Preserves system message when present
- Removes user/assistant message pairs
- Handles empty or minimal history gracefully
There was a problem hiding this comment.
Pull Request Overview
This PR adds LangChain-based LiteLLM integration support to the GraphRAG SDK, providing an alternative way to access LiteLLM's multi-provider capabilities through LangChain's ecosystem.
- Introduces a new
LangChainLiteModelclass that wraps LiteLLM functionality through LangChain'sChatLiteLLM - Adds comprehensive test coverage for model initialization, serialization, chat sessions, and streaming functionality
- Updates documentation with usage examples and guidance on when to use this integration vs direct LiteLLM
Reviewed Changes
Copilot reviewed 4 out of 5 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| graphrag_sdk/models/langchain_litellm.py | Core implementation of LangChain LiteLLM model and chat session classes |
| tests/test_langchain_litellm.py | Comprehensive test suite covering all functionality with both unit and integration tests |
| pyproject.toml | Adds required langchain-litellm and langchain-core dependencies |
| README.md | Documents the new integration with setup examples and usage guidance |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
…ithub.com/FalkorDB/GraphRAG-SDK into Add-LangChain-LiteLLM-integration-support
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
graphrag_sdk/models/langchain_litellm.py (1)
73-76: Validate generation config parameters before merging.The generation config is merged directly into
chat_paramswithout validation. ChatLiteLLM may not recognize allGenerativeModelConfigfields (e.g.,max_completion_tokensvsmax_tokens), potentially causing initialization failures or silent parameter drops.Based on the past review comment from Copilot, consider filtering to only allowed parameters:
# Add generation config parameters gen_config = self.generation_config.to_json() if gen_config: - chat_params.update(gen_config) + # Map and filter parameters that ChatLiteLLM supports + allowed_params = { + "temperature": gen_config.get("temperature"), + "max_tokens": gen_config.get("max_completion_tokens"), # Note the mapping + "top_p": gen_config.get("top_p"), + "stop": gen_config.get("stop"), + "timeout": gen_config.get("timeout"), + } + # Only include non-None values + filtered_params = {k: v for k, v in allowed_params.items() if v is not None} + chat_params.update(filtered_params)
🧹 Nitpick comments (3)
graphrag_sdk/models/langchain_litellm.py (3)
12-13: Consider using the logger for operational visibility.The logger is configured but never used. Adding log statements at key points (initialization, chat session start, errors) would improve observability and debugging.
99-99: Use proper type hint instead ofany.The parameter type
anyis not a valid type hint. UseAnyfrom thetypingmodule instead.+from typing import Optional, Iterator, Any + - def parse_generate_content_response(self, response: any) -> GenerationResponse: + def parse_generate_content_response(self, response: Any) -> GenerationResponse:
141-161: Add validation for required fields in deserialization.The
from_jsonmethod does not validate required fields. Ifmodel_nameorgeneration_configis missing from the JSON, aKeyErrorwill be raised with an unclear error message.Consider adding validation:
@staticmethod def from_json(json: dict) -> "GenerativeModel": """ Deserialize a JSON object to create an instance of LangChainLiteModel. Args: json (dict): The serialized JSON data. Returns: GenerativeModel: A new instance of the model. """ + if "model_name" not in json: + raise ValueError("Missing required field 'model_name' in JSON") + if "generation_config" not in json: + raise ValueError("Missing required field 'generation_config' in JSON") + return LangChainLiteModel( model_name=json["model_name"], generation_config=GenerativeModelConfig.from_json( json["generation_config"] ), system_instruction=json.get("system_instruction"), api_key=json.get("api_key"), api_base=json.get("api_base"), additional_params=json.get("additional_params"), )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
.wordlist.txt(1 hunks)graphrag_sdk/models/langchain_litellm.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- .wordlist.txt
🧰 Additional context used
🧬 Code graph analysis (1)
graphrag_sdk/models/langchain_litellm.py (1)
graphrag_sdk/models/model.py (5)
GenerativeModel(110-126)GenerativeModelConfig(10-79)GenerationResponse(81-90)FinishReason(5-8)GenerativeModelChatSession(93-107)
🪛 Ruff (0.13.3)
graphrag_sdk/models/langchain_litellm.py
48-51: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
48-51: Avoid specifying long messages outside the exception class
(TRY003)
81-83: Avoid specifying long messages outside the exception class
(TRY003)
209-211: Avoid specifying long messages outside the exception class
(TRY003)
244-246: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (3)
graphrag_sdk/models/langchain_litellm.py (3)
169-190: LGTM!The initialization properly sets up the chat session with message type references and system instruction handling. The approach of storing message class references is clean and avoids repeated imports.
192-246: LGTM! State consistency issue resolved.Both
send_messageandsend_message_streamnow properly handle error recovery by removing the user message from history when the model call fails (lines 208, 243). This addresses the critical issue flagged in the previous review.
248-292: LGTM! Improved clarity in delete_last_message.The
delete_last_messagelogic has been significantly improved with clearer variable names (has_system,start_idx,num_user_assistant_msgs) and helpful comments, addressing the complexity concern from the previous review. Both methods correctly handle their responsibilities.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
graphrag_sdk/models/langchain_litellm.py (1)
48-51: Consider exception handling style improvements.Ruff flags two minor style issues:
ImportError handling (lines 48-51): Consider adding
from Noneto suppress the import traceback:raise ImportError(...) from NoneException message length (lines 81-83, 210-212, 245-247): While inline error messages provide valuable context, consider defining custom exception classes if these patterns repeat across the codebase.
These are stylistic suggestions that don't affect functionality.
Also applies to: 81-83, 210-212, 245-247
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (1)
graphrag_sdk/models/langchain_litellm.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
graphrag_sdk/models/langchain_litellm.py (1)
graphrag_sdk/models/model.py (5)
GenerativeModel(110-126)GenerativeModelConfig(10-79)GenerationResponse(81-90)FinishReason(5-8)GenerativeModelChatSession(93-107)
🪛 Ruff (0.13.3)
graphrag_sdk/models/langchain_litellm.py
48-51: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
48-51: Avoid specifying long messages outside the exception class
(TRY003)
81-83: Avoid specifying long messages outside the exception class
(TRY003)
210-212: Avoid specifying long messages outside the exception class
(TRY003)
245-247: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: test (openai/gpt-4.1)
🔇 Additional comments (9)
graphrag_sdk/models/langchain_litellm.py (9)
1-14: LGTM!Imports and logger configuration are clean and appropriate.
60-83: LGTM - Error handling addresses parameter concerns.The direct merge of
generation_configparameters (line 76) intochat_paramsis safe because the initialization is wrapped in comprehensive error handling (lines 78-83) that will catch any parameter-related failures. LiteLLM is designed to flexibly handle provider-specific parameters, and the try-except ensures graceful failure with a clear error message if incompatible parameters are provided.
85-123: LGTM!The
start_chatandparse_generate_content_responsemethods are correctly implemented. Response parsing properly handles finish reasons from LangChain's response metadata.
125-162: LGTM - Serialization security addressed.The
to_jsonmethod now correctly excludes the sensitiveapi_keyfield (addressing the previous security concern), whilefrom_jsongracefully handles its absence usingjson.get("api_key")which defaults toNone. This maintains consistency between serialization and deserialization.
170-191: LGTM!Chat session initialization correctly sets up the message history and system instruction. The approach of storing message class references (lines 184-186) avoids repeated imports and is a reasonable optimization.
193-216: LGTM - Error handling implemented.The
send_messagemethod now includes proper error recovery (lines 208-209) that removes the user message from history on failure, addressing the previous review concern about state consistency.
218-247: LGTM - Streaming error handling implemented.The
send_message_streammethod now includes proper error recovery (lines 243-244) that removes the user message from history on failure, addressing the previous review concern about state consistency in streaming scenarios.
249-264: LGTM!The
get_chat_historymethod correctly transforms LangChain message objects into dictionary format with appropriate role mappings.
266-293: LGTM - Logic clarity improved.The
delete_last_messagemethod now uses clear variable names (has_system,start_idx,num_user_assistant_msgs) that make the logic easy to follow, addressing the previous review concern about complexity. The implementation correctly preserves system messages while removing user/assistant pairs.
Summary by CodeRabbit
New Features
Documentation
Tests
Chores