Fix LLM configuration usage and Providers API #116
Conversation
## Walkthrough
This update refactors all generative model classes (OpenAI, Azure OpenAI, Gemini, Ollama) to delegate functionality to a unified LiteLLM backend, removing direct SDK dependencies and manual chat management. The `OutputMethod` enum and related logic are eliminated, replaced by flexible parameter passing. Test cases are updated for model naming consistency. Additionally, `KnowledgeGraphModelConfig.with_model` now creates a specialized `extract_data` model with adjusted generation configuration.
## Changes
| File(s) | Change Summary |
|-------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
| graphrag_sdk/models/litellm.py | LiteModel and LiteModelChatSession refactored to accept `additional_params`, remove `output_method`, and simplify config.|
| graphrag_sdk/models/model.py | Removed `OutputMethod` enum; GenerativeModelConfig now supports arbitrary params; updated serialization/deserialization. |
| graphrag_sdk/models/openai.py<br>graphrag_sdk/models/azure_openai.py<br>graphrag_sdk/models/gemini.py<br>graphrag_sdk/models/ollama.py | All model classes now delegate to LiteModel; removed direct SDK usage, manual chat history, and `output_method` logic. |
| graphrag_sdk/model_config.py | `KnowledgeGraphModelConfig.with_model` now creates a specialized `extract_data` model with adjusted config. |
| graphrag_sdk/steps/extract_data_step.py | Removed all `OutputMethod` usage from method signatures and calls. |
| graphrag_sdk/orchestrator/orchestrator.py | Removed unused import of `OutputMethod`. |
| tests/test_auto_create_ontology.py<br>tests/test_kg.py<br>tests/test_ontology_from_kg.py | Updated Gemini model instantiations to use `"gemini-1.5-flash"` instead of `"gemini-1.5-flash-001"`. |
| tests/test_multi_agent.py | Switched from GeminiGenerativeModel to LiteModel with `"azure/o3-mini"` for test setup. |
## Sequence Diagram(s)
```mermaid
sequenceDiagram
participant User
participant GenerativeModel (OpenAI/Gemini/Ollama/Azure)
participant LiteModel
participant LiteModelChatSession
User->>GenerativeModel: start_chat(system_instruction)
GenerativeModel->>LiteModel: start_chat(system_instruction)
LiteModel->>LiteModelChatSession: create session
User->>GenerativeModel: send_message(message)
GenerativeModel->>LiteModel: delegate to chat session
LiteModel->>LiteModelChatSession: send_message(message)
LiteModelChatSession-->>LiteModel: return GenerationResponse
LiteModel-->>GenerativeModel: return GenerationResponse
GenerativeModel-->>User: return GenerationResponseSuggested reviewers
Poem
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
graphrag_sdk/model_config.py (1)
38-73: Excellent enhancement with robust configuration logic.The new
with_modelimplementation provides significant improvements:
- Smart configuration: Sets temperature=0 and JSON response format for extract_data operations, which is ideal for structured data extraction
- Robust fallback: Gracefully handles model creation failures
- Preserves user settings: Respects existing configuration while making necessary adjustments
Consider making the exception handling more specific for better debugging:
- except Exception: + except (TypeError, ValueError, AttributeError) as e: # Fallback: use original model if creation fails + # Log the specific error for debugging if needed extract_data_model = modelgraphrag_sdk/models/gemini.py (1)
19-41: Consider adding**kwargssupport for consistency with other model implementations.The Azure OpenAI implementation accepts
**kwargsand passes them asadditional_paramsto the LiteModel. For consistency across model implementations, consider adding the same support to the Gemini model.def __init__( self, model_name: str, generation_config: Optional[GenerativeModelConfig] = None, system_instruction: Optional[str] = None, + **kwargs: Any, ):And update the LiteModel initialization:
self._lite_model = LiteModel( model_name=lite_model_name, generation_config=generation_config, - system_instruction=system_instruction + system_instruction=system_instruction, + additional_params=kwargs )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
graphrag_sdk/model_config.py(2 hunks)graphrag_sdk/models/azure_openai.py(7 hunks)graphrag_sdk/models/gemini.py(4 hunks)graphrag_sdk/models/litellm.py(6 hunks)graphrag_sdk/models/model.py(4 hunks)graphrag_sdk/models/ollama.py(5 hunks)graphrag_sdk/models/openai.py(6 hunks)graphrag_sdk/orchestrator/orchestrator.py(0 hunks)graphrag_sdk/steps/extract_data_step.py(2 hunks)tests/test_auto_create_ontology.py(1 hunks)tests/test_kg.py(1 hunks)tests/test_multi_agent.py(2 hunks)tests/test_ontology_from_kg.py(1 hunks)
💤 Files with no reviewable changes (1)
- graphrag_sdk/orchestrator/orchestrator.py
🧰 Additional context used
🪛 Ruff (0.11.9)
graphrag_sdk/models/openai.py
3-3: .litellm.LiteModelChatSession imported but unused
Remove unused import: .litellm.LiteModelChatSession
(F401)
22-22: Do not perform function call GenerativeModelConfig in argument defaults; instead, perform the call within the function, or read the default from a module-level singleton variable
(B008)
graphrag_sdk/models/ollama.py
23-23: Do not perform function call GenerativeModelConfig in argument defaults; instead, perform the call within the function, or read the default from a module-level singleton variable
(B008)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: test (openai/gpt-4o)
- GitHub Check: test (gemini/gemini-2.0-flash)
🔇 Additional comments (12)
tests/test_kg.py (1)
71-71: LGTM: Model name updated for consistency.The model name change from
"gemini-1.5-flash-001"to"gemini-1.5-flash"aligns with the refactoring ofGeminiGenerativeModelto use the unified LiteLLM backend.tests/test_auto_create_ontology.py (1)
27-27: LGTM: Consistent model naming update.The model name change matches the pattern established across other test files for the LiteLLM backend refactoring.
tests/test_ontology_from_kg.py (1)
58-58: LGTM: Model naming consistency maintained.The model name update is consistent with the refactoring changes across the test suite.
tests/test_multi_agent.py (2)
9-9: LGTM: Import updated for unified LiteLLM backend.The import change from
GeminiGenerativeModeltoLiteModelaligns with the refactoring to use the unified LiteLLM backend.
188-188: Verify the intentional model provider change.The model has been changed from
GeminiGenerativeModel("gemini-1.5-flash-001")toLiteModel("azure/o3-mini"), which represents a switch from Gemini to Azure OpenAI. While this aligns with testing the LiteModel directly, please confirm this provider change is intentional for this specific test case.#!/bin/bash # Description: Check if other tests also use different model providers or if this is unique # Expected: Find patterns of model provider usage across test files echo "=== Checking model instantiations across test files ===" rg -A 2 -B 2 "(GeminiGenerativeModel|LiteModel|OpenAiGenerativeModel|AzureOpenAiGenerativeModel)" tests/graphrag_sdk/model_config.py (1)
1-2: LGTM: Import statements updated for refactored model interfaces.The imports are correctly updated to use the unified model interfaces.
graphrag_sdk/steps/extract_data_step.py (1)
190-190: Clean API simplificationThe removal of the
output_methodparameter from all model calls aligns well with the unified LiteLLM backend approach and simplifies the API.Also applies to: 197-197, 368-368
graphrag_sdk/models/model.py (1)
35-90: Excellent enhancement for flexible configurationThe addition of
**kwargssupport and dynamic attribute handling makesGenerativeModelConfigmuch more extensible. The serialization methods correctly handle both known and arbitrary parameters, which will be valuable for provider-specific configurations.graphrag_sdk/models/litellm.py (2)
196-212: Well-designed parameter mergingThe approach of merging
generation_configandadditional_paramsusing dictionary unpacking is clean and allows for flexible parameter passing. The order ensures that additional_params can override default configuration values when needed.Also applies to: 232-238
75-75: Verify parameter compatibility with LiteLLMThe change from
max_tokenstomax_completion_tokenscould potentially break compatibility if LiteLLM doesn't support this parameter name across all providers.Please verify that
max_completion_tokensis supported by LiteLLM for all provider types:Does LiteLLM support max_completion_tokens parameter for all providers like OpenAI, Azure OpenAI, Gemini, and Ollama?graphrag_sdk/models/ollama.py (1)
37-69: Excellent refactoring to unified backendThe delegation pattern to LiteLLM is well-implemented:
- Proper model name prefixing with "ollama/"
- Host parameter correctly mapped to
api_base- Clean property accessors maintain API compatibility
- Chat session delegation removes code duplication
Also applies to: 128-177
graphrag_sdk/models/openai.py (1)
33-163: Clean implementation following unified patternAfter addressing the import and default argument issues, this implementation nicely mirrors the pattern used in other model classes, providing a consistent API while delegating to the unified LiteLLM backend.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
tests/test_multi_agent.py (1)
191-191: Fix line length violation.The line exceeds the 79-character limit. Consider breaking it into multiple lines for better readability.
- cls.model = LiteModel("azure/o3-mini", generation_config=GenerativeModelConfig(temperature=None)) + cls.model = LiteModel( + "azure/o3-mini", + generation_config=GenerativeModelConfig(temperature=None) + )graphrag_sdk/model_config.py (1)
39-39: Fix line length violations for better readability.Multiple lines exceed the 79-character limit. Consider breaking them for better code readability.
- The extract_data model will be configured with temperature=0 and JSON response format + The extract_data model will be configured with temperature=0 and JSON + response format for structured data extraction. Args: model (GenerativeModel): The generative model to use. Returns: - KnowledgeGraphModelConfig: The new KnowledgeGraphModelConfig instance. + KnowledgeGraphModelConfig: The new KnowledgeGraphModelConfig + instance. - # User explicitly set temperature (even if to None), only override if not None + # User explicitly set temperature (even if to None), + # only override if not None if model.generation_config.temperature is not None: - # If user explicitly set temperature=None, respect that and don't override + # If user explicitly set temperature=None, respect that and + # don't overrideAlso applies to: 46-46, 54-54, 57-57
graphrag_sdk/models/model.py (3)
38-46: Consider refactoring the constructor to reduce complexity.The constructor has many parameters (7), which makes it complex to use and maintain. Consider using a builder pattern or converting to a dataclass for better ergonomics.
Alternative approach using a dataclass:
from dataclasses import dataclass, field from typing import Any, Dict @dataclass class GenerativeModelConfig: temperature: Any = _TEMP_NOT_SET top_p: Optional[float] = None top_k: Optional[int] = None max_tokens: Optional[int] = None stop: Optional[list] = None response_format: Optional[dict] = None extra_params: Dict[str, Any] = field(default_factory=dict) _temperature_was_set: bool = field(init=False, default=False) def __post_init__(self): if self.temperature is self._TEMP_NOT_SET: self.temperature = None self._temperature_was_set = False else: self._temperature_was_set = True # Set extra parameters as attributes for key, value in self.extra_params.items(): setattr(self, key, value)
14-15: Fix line length violations throughout the file.Multiple lines exceed the 79-character limit. Break long lines for better readability.
- This configuration follows OpenAI-style parameter naming but is designed to be - compatible with other generative models. Supports both predefined and arbitrary parameters. + This configuration follows OpenAI-style parameter naming but is designed to be + compatible with other generative models. Supports both predefined and arbitrary + parameters. - return f"GenerativeModelConfig({', '.join(f'{k}={v}' for k, v in self.to_json().items())})" + items = ', '.join(f'{k}={v}' for k, v in self.to_json().items()) + return f"GenerativeModelConfig({items})" - Serialize the configuration to a dictionary, excluding any fields with None values. + Serialize the configuration to a dictionary, excluding any fields with None + values. - dict: A dictionary containing only the parameters that are explicitly set + dict: A dictionary containing only the parameters that are explicitly set (i.e., not None). - return {k: v for k, v in self.__dict__.items() if v is not None and k != '_temperature_was_set'} + return {k: v for k, v in self.__dict__.items() + if v is not None and k != '_temperature_was_set'} - # For other parameters, only pass if they exist in JSON to maintain default behavior + # For other parameters, only pass if they exist in JSON to maintain + # default behavior - known_keys = {'temperature', 'top_p', 'top_k', 'max_tokens', 'stop', 'response_format'} + known_keys = {'temperature', 'top_p', 'top_k', 'max_tokens', 'stop', + 'response_format'}Also applies to: 68-68, 72-72, 75-75, 83-83, 91-91, 104-104
40-40: Fix parameter spacing formatting issue.Remove unexpected spaces around the equals sign in the parameter definition.
- temperature = _TEMP_NOT_SET, + temperature=_TEMP_NOT_SET,
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
graphrag_sdk/model_config.py(2 hunks)graphrag_sdk/models/model.py(2 hunks)tests/test_multi_agent.py(2 hunks)
🧰 Additional context used
🪛 Flake8 (7.2.0)
graphrag_sdk/model_config.py
[error] 39-39: line too long (93 > 79 characters)
(E501)
[error] 46-46: line too long (82 > 79 characters)
(E501)
[error] 54-54: line too long (90 > 79 characters)
(E501)
[error] 57-57: line too long (86 > 79 characters)
(E501)
graphrag_sdk/models/model.py
[error] 14-14: line too long (82 > 79 characters)
(E501)
[error] 15-15: line too long (95 > 79 characters)
(E501)
[error] 20-20: line too long (81 > 79 characters)
(E501)
[error] 40-40: unexpected spaces around keyword / parameter equals
(E251)
[error] 40-40: unexpected spaces around keyword / parameter equals
(E251)
[error] 67-67: too many blank lines (2)
(E303)
[error] 68-68: line too long (99 > 79 characters)
(E501)
[error] 72-72: line too long (91 > 79 characters)
(E501)
[error] 75-75: line too long (85 > 79 characters)
(E501)
[error] 83-83: line too long (104 > 79 characters)
(E501)
[error] 91-91: line too long (92 > 79 characters)
(E501)
[error] 104-104: line too long (95 > 79 characters)
(E501)
tests/test_multi_agent.py
[error] 191-191: line too long (105 > 79 characters)
(E501)
🪛 Pylint (3.3.7)
graphrag_sdk/models/model.py
[refactor] 38-38: Too many arguments (7/5)
(R0913)
[refactor] 38-38: Too many positional arguments (7/5)
(R0917)
🔇 Additional comments (10)
tests/test_multi_agent.py (1)
9-14: LGTM! Model abstraction refactoring looks good.The migration from
GeminiGenerativeModeltoLiteModelaligns with the unified model abstraction approach described in the PR objectives. The additional imports forGenerativeModelConfigsupport the new configuration approach.graphrag_sdk/model_config.py (3)
38-48: LGTM! Good documentation of the specialized configuration.The docstring clearly explains that the
extract_datamodel will be configured with specific parameters for structured data extraction. This helps users understand the behavior.
66-74: Good error handling with fallback pattern.The try-catch block with fallback to the original model ensures robustness when model instantiation fails. This prevents breaking existing functionality.
49-75: Verify the temperature handling logic works correctly with GenerativeModelConfig.The temperature handling logic is sophisticated and respects user intent, but needs verification that it works correctly with the new
GenerativeModelConfig._temperature_was_setflag implementation.#!/bin/bash # Description: Verify that GenerativeModelConfig properly tracks temperature setting # Expected: Find the _temperature_was_set implementation and sentinel pattern ast-grep --pattern 'class GenerativeModelConfig { $$$ _temperature_was_set = $$$ $$$ }' # Also check for the sentinel object usage rg -A 5 "_TEMP_NOT_SET"graphrag_sdk/models/model.py (6)
14-25: LGTM! Good documentation of the enhanced configuration capabilities.The updated docstring clearly explains the OpenAI-style parameter naming and support for arbitrary parameters. The example demonstrates the usage well.
35-55: Clever sentinel pattern for tracking explicit parameter setting.The use of a sentinel object to differentiate between unset and explicitly set to None is a robust solution. This enables the sophisticated temperature handling logic in
KnowledgeGraphModelConfig.with_model.
62-64: Good use of dynamic attribute storage.Using
setattrto store arbitrary kwargs as instance attributes provides the flexibility needed for supporting various model-specific parameters.
70-83: Verify the serialization logic handles edge cases correctly.The dynamic serialization approach looks good, but ensure it properly handles edge cases and doesn't inadvertently expose internal attributes.
# Test edge cases for the serialization logic config1 = GenerativeModelConfig(temperature=0.5, custom_param="test") print("Basic serialization:", config1.to_json()) config2 = GenerativeModelConfig(temperature=None) print("Explicit None serialization:", config2.to_json()) config3 = GenerativeModelConfig() print("Default serialization:", config3.to_json()) # Test round-trip serialization json_data = config1.to_json() config_restored = GenerativeModelConfig.from_json(json_data) print("Round-trip test:", config_restored.to_json() == json_data)
86-107: Robust deserialization with parameter validation.The conditional parameter passing ensures that only explicitly provided parameters are set, maintaining backward compatibility and default behavior.
132-132: LGTM! OutputMethod removal simplifies the interface.Removing the
output_methodparameter fromsend_messagealigns with the broader refactoring to eliminate the OutputMethod enum and simplifies the API.
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
graphrag_sdk/models/azure_openai.py (1)
78-78: Fix invalid type hint:anyshould beAny.The type hint
anyis not valid in Python. UseAnyfrom thetypingmodule.- def parse_generate_content_response(self, response: any) -> GenerationResponse: + def parse_generate_content_response(self, response: Any) -> GenerationResponse:graphrag_sdk/models/gemini.py (1)
73-73: Fix invalid type hint:anyshould beAny.The type hint
anyis not valid in Python. UseAnyfrom thetypingmodule.- def parse_generate_content_response(self, response: any) -> GenerationResponse: + def parse_generate_content_response(self, response: Any) -> GenerationResponse:
🧹 Nitpick comments (4)
README.md (1)
38-38: Fix punctuation for better readability.Add a comma after "vendor" for improved grammar and readability.
- To choose vendor use the prefix "specific_vendor/your_model", for example "openai/gpt-4.1". + To choose vendor, use the prefix "specific_vendor/your_model", for example "openai/gpt-4.1".graphrag_sdk/model_config.py (1)
51-57: Fix formatting and approve deep copy approach.The deep copy approach correctly prevents mutation of the original model when configuring the extract_data model. However, there are too many blank lines that should be cleaned up.
extract_data_model.generation_config.response_format = {"type": "json_object"} - return KnowledgeGraphModelConfig(graphrag_sdk/models/model.py (2)
41-41: Consider preserving type safety for thestopparameter.Changing from
Optional[list[str]]toOptional[list]reduces type safety. Since stop sequences are expected to be strings, maintaining the more specific type hint would help catch type errors at development time.- stop: Optional[list] = None, + stop: Optional[list[str]] = None,
56-57: Remove extra blank line.- - - def __str__(self) -> str: + + def __str__(self) -> str:
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
.github/workflows/test.yml(2 hunks).wordlist.txt(1 hunks)README.md(3 hunks)examples/movies/demo-movies.ipynb(6 hunks)examples/trip/demo_orchestrator_trip.ipynb(11 hunks)examples/ufc/demo-ufc.ipynb(13 hunks)examples/ufc/ontology.json(5 hunks)graphrag_sdk/model_config.py(4 hunks)graphrag_sdk/models/azure_openai.py(7 hunks)graphrag_sdk/models/gemini.py(4 hunks)graphrag_sdk/models/litellm.py(5 hunks)graphrag_sdk/models/model.py(4 hunks)graphrag_sdk/models/ollama.py(5 hunks)graphrag_sdk/models/openai.py(6 hunks)tests/test_rag.py(1 hunks)tests/test_streaming_response.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- .wordlist.txt
🚧 Files skipped from review as they are similar to previous changes (2)
- .github/workflows/test.yml
- graphrag_sdk/models/openai.py
🧰 Additional context used
🪛 LanguageTool
README.md
[uncategorized] ~38-~38: A comma might be missing here.
Context: ...eployment and use cases.
To choose vendor use the prefix "specific_vendor/your_mo...
(AI_EN_LECTOR_MISSING_PUNCTUATION_COMMA)
🪛 Flake8 (7.2.0)
graphrag_sdk/model_config.py
[error] 56-56: too many blank lines (2)
(E303)
graphrag_sdk/models/model.py
[error] 57-57: too many blank lines (2)
(E303)
🪛 Ruff (0.11.9)
graphrag_sdk/models/litellm.py
57-59: Use a single if statement instead of nested if statements
(SIM102)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: test (openai/gpt-4.1)
🔇 Additional comments (17)
tests/test_streaming_response.py (1)
93-93: LGTM! Model name update aligns with unified LiteLLM backend.The change from "gemini/gemini-2.0-flash" to "openai/gpt-4.1" is consistent with the broader refactoring to use a unified LiteLLM backend across the codebase.
tests/test_rag.py (1)
115-115: LGTM! Consistent model name update.The change from "gemini/gemini-2.0-flash" to "openai/gpt-4.1" maintains consistency with other test files and the unified LiteLLM backend approach.
README.md (1)
39-41: LGTM! Documentation updates align with unified LiteLLM backend.The model name changes from "gemini/gemini-2.0-flash" to "openai/gpt-4.1" throughout the documentation are consistent with the codebase refactoring to use a unified LiteLLM backend.
Also applies to: 71-72, 146-146
examples/movies/demo-movies.ipynb (3)
63-64: LGTM! Environment setup updated for OpenAI integration.The change from GEMINI_API_KEY to OPENAI_API_KEY correctly reflects the model switch to "openai/gpt-4.1" and ensures users configure the appropriate credentials.
66-66: LGTM! Typo correction improves documentation quality.Good catch fixing "Falcor" to "Falkor" for accurate brand representation.
121-121: LGTM! Model name update maintains consistency.The change to "openai/gpt-4.1" aligns with the unified LiteLLM backend approach used throughout the codebase.
graphrag_sdk/model_config.py (2)
1-1: LGTM! Import added for deep copy functionality.The copy module import is necessary for the deep copy operation used to avoid mutating the original model.
30-31: LGTM! Extract data model configured for structured output.Explicitly setting the response_format to JSON ensures structured data extraction, which aligns with the specialized model configuration requirements.
examples/trip/demo_orchestrator_trip.ipynb (1)
18-18: LGTM! Clean migration to LiteModel.The migration from OpenAiGenerativeModel to LiteModel is implemented correctly, with appropriate model naming conventions (e.g., "o3-mini" for the orchestrator).
Also applies to: 51-51, 266-266, 407-407
examples/ufc/demo-ufc.ipynb (2)
44-44: LGTM! Consistent migration to LiteModel.The changes correctly migrate from OpenAiGenerativeModel to LiteModel, maintaining consistency with other files in the PR.
Also applies to: 112-112
144-585: Verify compatibility of the significant ontology schema changes.The ontology structure has changed significantly:
- "Person" entity renamed to "Event"
- "FightEvent" renamed to "Fight"
- New entities added: Referee, FightStatistics, RoundStatistics, SignificantStrikeBreakdown
- Relations have been redefined with new labels and attributes
These changes could be breaking for existing data or dependent systems.
Ensure that:
- Existing UFC data can be migrated to the new schema
- Any dependent code that relies on the old entity names is updated
- The ontology changes are documented in release notes
graphrag_sdk/models/ollama.py (1)
17-178: Excellent refactoring to use LiteModel as the backend.The refactoring cleanly delegates all functionality to the internal LiteModel while maintaining backward compatibility:
- The public API remains unchanged
- Host parameter is correctly mapped to LiteLLM's api_base
- Properties provide access to the original values
- All chat session operations are properly delegated
This implementation successfully unifies the backend while preserving the Ollama-specific interface.
graphrag_sdk/models/azure_openai.py (2)
36-48: LGTM: Clean delegation to LiteModel backend.The refactoring correctly delegates to the internal LiteModel with proper azure/ prefix while maintaining API compatibility through the original model name property.
132-172: LGTM: Consistent chat session delegation pattern.The chat session correctly delegates all operations to the internal LiteLLM chat session, simplifying the implementation and removing manual history management.
graphrag_sdk/models/gemini.py (2)
33-44: LGTM: Consistent LiteModel integration with proper prefix.The implementation correctly creates a LiteModel with the "gemini/" prefix and maintains the original model name for API compatibility.
118-167: LGTM: Clean chat session delegation.The chat session implementation properly delegates all operations to the internal LiteLLM chat session, following the established pattern.
examples/ufc/ontology.json (1)
1-442: LGTM: Comprehensive domain model enhancement.The ontology restructuring significantly improves the UFC domain representation with:
- Semantic clarity (Fighter → Person, more descriptive attribute names)
- Enhanced Fight entity with fight-specific attributes (title_bout, weight_class, etc.)
- New detailed statistics entities (RoundStatistics, SignificantStrikeBreakdown)
- Well-structured relations with appropriate attributes
The JSON structure is valid and the entity/relation definitions are logically consistent for UFC data modeling.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
graphrag_sdk/models/litellm.py (1)
67-70: Simplify nested if statements and document the special case.The nested if statements can be combined for better readability, and the hardcoded temperature setting should be documented.
- if self.model_name == "gpt-4.1": - # Set default temperature to 0 for gpt-4.1 - if generation_config is None: - generation_config = GenerativeModelConfig(temperature=0) + # Set default temperature to 0 for gpt-4.1 for deterministic outputs + if self.model_name == "gpt-4.1" and generation_config is None: + generation_config = GenerativeModelConfig(temperature=0)
🧹 Nitpick comments (1)
graphrag_sdk/models/ollama.py (1)
13-13: Fix formatting: Add blank line before class definition.Static analysis indicates missing blank line before the class definition.
+ class OllamaGenerativeModel(LiteModel):
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
.github/workflows/test.yml(3 hunks)graphrag_sdk/models/azure_openai.py(5 hunks)graphrag_sdk/models/gemini.py(2 hunks)graphrag_sdk/models/litellm.py(6 hunks)graphrag_sdk/models/model.py(2 hunks)graphrag_sdk/models/ollama.py(3 hunks)graphrag_sdk/models/openai.py(5 hunks)pyproject.toml(3 hunks)tests/test_multi_agent.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- .github/workflows/test.yml
🧰 Additional context used
🪛 Flake8 (7.2.0)
graphrag_sdk/models/ollama.py
[error] 13-13: expected 2 blank lines, found 1
(E302)
graphrag_sdk/models/model.py
[error] 58-58: too many blank lines (2)
(E303)
tests/test_multi_agent.py
[error] 12-12: 'graphrag_sdk.GenerativeModelConfig' imported but unused
(F401)
🪛 Ruff (0.11.9)
graphrag_sdk/models/litellm.py
67-69: Use a single if statement instead of nested if statements
(SIM102)
tests/test_multi_agent.py
14-14: graphrag_sdk.GenerativeModelConfig imported but unused
Remove unused import: graphrag_sdk.GenerativeModelConfig
(F401)
🪛 Pylint (3.3.7)
graphrag_sdk/models/model.py
[refactor] 36-36: Too many arguments (7/5)
(R0913)
[refactor] 36-36: Too many positional arguments (7/5)
(R0917)
🔇 Additional comments (21)
tests/test_multi_agent.py (1)
9-9: LGTM! Clean transition to unified LiteLLM backend.The refactoring from
GeminiGenerativeModeltoLiteModelaligns perfectly with the PR objectives to consolidate all generative models under a unified LiteLLM backend.Also applies to: 191-191
pyproject.toml (2)
3-3: Appropriate version bump for the refactoring.The minor version bump to 0.8.0 correctly reflects the significant architectural changes to use unified LiteLLM backend.
26-28: Excellent dependency restructuring for LiteLLM integration.The promotion of
litellmandollamato core dependencies, combined with the fine-grained extras structure, provides users with flexible control over LiteLLM features while supporting the unified backend architecture.Also applies to: 30-63
graphrag_sdk/models/model.py (4)
21-21: Good alignment with OpenAI parameter naming conventions.The change from
max_tokenstomax_completion_tokensaligns with OpenAI's parameter naming and improves consistency across different LLM providers.Also applies to: 41-41, 49-49
44-44: Excellent addition of flexible parameter support.The
**kwargsparameter and dynamic attribute storage allows the configuration to support arbitrary provider-specific parameters while maintaining type safety for common parameters.Also applies to: 53-55
74-74: Clean and maintainable serialization methods.The use of
vars(self)into_json()and the simplifiedfrom_json()method create a robust serialization pattern that automatically handles both predefined and dynamic attributes.Also applies to: 78-79
103-103: Consistent removal of OutputMethod parameter.The removal of the
output_methodparameter from the abstractsend_messagemethod aligns with the elimination of theOutputMethodenum throughout the codebase.graphrag_sdk/models/openai.py (4)
22-22: Good fix for mutable default argument.Changing the default value from
GenerativeModelConfig()toNoneprevents the shared mutable state issue that was identified in past reviews.
33-42: Clever model name handling for LiteLLM compatibility.The approach of internally prefixing with
"openai/"for LiteLLM while exposing the original model name through a property maintains API compatibility while supporting the unified backend.Also applies to: 44-47
68-68: Consistent serialization with clean model names.Returning the original model name without the
openai/prefix in serialization maintains clean external API while using LiteLLM internally.
93-110: Excellent delegation to LiteModelChatSession.The simplified
OpenAiChatSessionthat inherits all functionality fromLiteModelChatSessioneliminates code duplication while maintaining the specific type for API compatibility.graphrag_sdk/models/litellm.py (3)
5-11: Robust handling of optional Ollama dependency.The try-catch import pattern with availability flag provides clean error handling when Ollama client is not installed.
89-89: Consistent parameter naming with OpenAI conventions.The change from
max_tokenstomax_completion_tokensaligns with the parameter name updates throughout the codebase.
225-226: Clean parameter merging for LiteLLM calls.The direct unpacking of
generation_config.to_json()andadditional_paramsprovides a clean way to merge all parameters for the completion call without complex adjustment logic.Also applies to: 251-252
graphrag_sdk/models/ollama.py (2)
13-61: Excellent refactoring to LiteModel architecture.The conversion from direct Ollama client usage to LiteModel inheritance is well-implemented. The approach of converting model names to LiteLLM format (
ollama/{model_name}) while preserving the original name through a property maintains backward compatibility and provides a clean abstraction.
107-124: Clean delegation to LiteModelChatSession.The simplified chat session implementation that delegates all functionality to the parent class is a good approach that eliminates code duplication while maintaining the same interface.
graphrag_sdk/models/azure_openai.py (3)
11-96: Consistent refactoring implementation.The refactoring to inherit from LiteModel follows the same clean pattern as the other model implementations. The Azure-specific model name conversion and parameter handling via **kwargs is well-implemented.
98-115: Proper delegation to parent chat session.The simplified chat session implementation correctly delegates to LiteModelChatSession, maintaining consistency with the other model implementations.
22-22: Verify the default model name "gpt-4.1".The default model name "gpt-4.1" doesn't appear to be a valid Azure OpenAI model name. Standard Azure OpenAI models typically use names like "gpt-4", "gpt-4-turbo", "gpt-35-turbo", etc.
What are the current available model names in Azure OpenAI Service?graphrag_sdk/models/gemini.py (2)
10-91: Consistent and clean LiteModel integration.The refactoring perfectly follows the established pattern from the other model implementations. The conversion to
gemini/{model_name}format, preservation of the original model name via property, and delegation to LiteModel is implemented correctly and maintains backward compatibility.
94-111: Excellent simplification of chat session.The chat session implementation correctly inherits from LiteModelChatSession and delegates all functionality to the parent class, eliminating code duplication while maintaining the same interface.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
README.md (5)
26-29: Clarify installation commandAdd
--upgradeor a version specifier so users with an older cached wheel don’t miss the new LiteLLM integration:-pip install graphrag_sdk +pip install --upgrade "graphrag_sdk>=0.6.0"
31-31: Drop the trailing period in the heading
markdownlint(MD026) flags punctuation at the end of a heading.
### Configure Credentialsis sufficient.
34-38: Grammar & consistency tweaks
- Missing comma: “To choose a vendor, use the prefix …”.
- Keep the prefix pattern consistent across bullets. E.g.:
-* [Google](https://makersuite.google.com/app/apikey) Recommended model:`gemini-2.0-flash` +* [Google](https://makersuite.google.com/app/apikey) Recommended model:`google/gemini-2.0-flash`Please verify the exact strings LiteLLM expects for each provider.
67-68: Streamline the explanatory comment-# Model - vendor: openai, model: gpt-4.1 -> openai/gpt-4.1 +# OpenAI GPT-4.1 modelThe shorter comment conveys the same information without the arrow notation.
142-143: Avoid magic strings by reusing a constantIf the recommended model changes, you’ll need to update multiple snippets. Consider defining once:
MODEL_NAME = "openai/gpt-4.1" model = LiteModel(model_name=MODEL_NAME)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
README.md(3 hunks)pyproject.toml(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- pyproject.toml
🧰 Additional context used
🪛 LanguageTool
README.md
[uncategorized] ~34-~34: A comma might be missing here.
Context: ...eployment and use cases.
To choose vendor use the prefix "specific_vendor/your_mo...
(AI_EN_LECTOR_MISSING_PUNCTUATION_COMMA)
🪛 markdownlint-cli2 (0.17.2)
README.md
31-31: Trailing punctuation in heading
Punctuation: '.'
(MD026, no-trailing-punctuation)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: test (openai/gpt-4.1)
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
graphrag_sdk/models/litellm.py (1)
67-71: Address the nested if statement issue from previous review.This is a duplicate of the previous review comment that wasn't addressed. The nested if statements should be simplified and the special case documented.
Apply this diff to fix the nested if statements:
- if self._internal_model_name == "gpt-4.1": - # Set default temperature to 0 for gpt-4.1 - if generation_config is None: - generation_config = GenerativeModelConfig(temperature=0) + # Set default temperature to 0 for gpt-4.1 for deterministic outputs + if self._internal_model_name == "gpt-4.1" and generation_config is None: + generation_config = GenerativeModelConfig(temperature=0)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
graphrag_sdk/models/litellm.py(6 hunks)tests/test_multi_agent.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_multi_agent.py
🧰 Additional context used
🪛 Ruff (0.11.9)
graphrag_sdk/models/litellm.py
67-69: Use a single if statement instead of nested if statements
(SIM102)
125-125: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🔇 Additional comments (11)
graphrag_sdk/models/litellm.py (11)
5-11: Good defensive programming with optional Ollama import.The conditional import pattern with availability flag is well-implemented, allowing graceful fallback when Ollama is not installed.
34-34: Good addition of flexible parameter passing.The
additional_paramsparameter enables provider-specific configurations while maintaining a clean interface.
56-56: Proper separation of internal and external model names.Using
_internal_model_nameto store the parsed model name while keeping the full provider/model format inself.modelis a good design choice.
60-62: Robust error handling for optional dependency.The availability check with descriptive error message guides users to install the required package.
74-75: Clean initialization of additional parameters.The null-safe initialization of
additional_paramswith empty dict fallback is well-implemented.
76-79: Well-designed property for model name access.The property provides clean access to the internal model name while maintaining encapsulation.
94-94: Correct parameter name update.The change from
max_tokenstomax_completion_tokensaligns with the updatedGenerativeModelConfiginterface.
215-215: Simplified method signature improves usability.Removing the
output_methodparameter simplifies the interface while maintaining functionality through flexible parameter passing.
230-232: Clean parameter merging with proper precedence.The parameter merging allows
additional_paramsto override generation config values, providing flexibility while maintaining backward compatibility.
256-258: Consistent parameter handling in streaming method.The streaming method uses the same parameter merging pattern as the non-streaming version, ensuring consistency.
31-31: [web_search]What GPT-4 model identifiers does the OpenAI API currently support, and is "gpt-4.1" included among them?
Summary by CodeRabbit
Refactor
Bug Fixes
Tests
Chores