Fix Cypher generation iterative loop#37
Conversation
WalkthroughThis pull request introduces several modifications across multiple files. New entries "openai" and "Ollama" were added to the Changes
Possibly related PRs
Suggested reviewers
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Outside diff range and nitpick comments (7)
.wordlist.txt (1)
28-28: Consider maintaining alphabetical orderWhile the addition of "Ollama" is appropriate, consider maintaining alphabetical order in the wordlist for better maintainability and easier lookups.
html https -www -faq Ollama +www +faqgraphrag_sdk/orchestrator/orchestrator.py (3)
3-3: Remove unused importOutputMethodThe
OutputMethodimport is not used anywhere in this file. If it's not needed for future implementation, it should be removed to keep the codebase clean.-from graphrag_sdk.models.model import OutputMethod🧰 Tools
🪛 Ruff
3-3:
graphrag_sdk.models.model.OutputMethodimported but unusedRemove unused import:
graphrag_sdk.models.model.OutputMethod(F401)
23-24: Add type hints for instance variablesConsider adding type hints to improve code maintainability and IDE support:
- self._agents = [] - self._chat = None + self._agents: list[Agent] = [] + self._chat: Optional[ChatSession] = NoneDon't forget to add the necessary imports:
from typing import Optional from graphrag_sdk.models.chat import ChatSession
Line range hint
52-67: Enhance error handling for Cypher generationGiven that this PR aims to fix Cypher generation iterative loop issues, consider implementing more specific error handling:
- The current generic exception handling doesn't provide much value as it immediately re-raises the error
- Consider catching specific exceptions that might occur during plan creation
def _create_execution_plan(self, question: str): try: response = self._get_chat().send_message( ORCHESTRATOR_EXECUTION_PLAN_PROMPT.replace("#QUESTION", question) ) logger.debug(f"Execution plan response: {response.text}") plan = ExecutionPlan.from_json( extract_json(response.text, skip_repair=True) ) logger.debug(f"Execution plan: {plan}") return plan - except Exception as e: - logger.error(f"Failed to create plan: {e}") - raise e + except ValueError as e: + logger.error(f"Invalid JSON in execution plan response: {e}") + raise + except KeyError as e: + logger.error(f"Missing required field in execution plan: {e}") + raise + except Exception as e: + logger.error(f"Unexpected error creating execution plan: {e}") + raise🧰 Tools
🪛 Ruff
3-3:
graphrag_sdk.models.model.OutputMethodimported but unusedRemove unused import:
graphrag_sdk.models.model.OutputMethod(F401)
graphrag_sdk/steps/graph_query_step.py (1)
45-45: Rename unused loop variable.The loop control variable
iis not used within the loop body.- for i in range(retries): + for _i in range(retries):🧰 Tools
🪛 Ruff
45-45: Loop control variable
inot used within loop bodyRename unused
ito_i(B007)
README.md (1)
62-67: Enhance Azure-OpenAI configuration examplesThe Azure-OpenAI configuration section would be more helpful with concrete example values.
- AZURE_OPENAI_API_KEY="AZURE_OPENAI_API_KEY" - AZURE_ENDPOINT="AZURE_ENDPOINT" - AZURE_API_VERSION="AZURE_API_VERSION" + AZURE_OPENAI_API_KEY="00000000000000000000000000000000" + AZURE_ENDPOINT="https://my-resource.openai.azure.com" + AZURE_API_VERSION="2023-05-15"graphrag_sdk/fixtures/prompts.py (1)
444-447: Correct grammatical errors inCYPHER_GEN_PROMPTfor clarityThere are grammatical issues in the
CYPHER_GEN_PROMPT. Specifically:
- Use "an OpenCypher statement" instead of "a OpenCypher statement".
- Rephrase "Please think if your answer is a valid Cypher query, and correct it if it is not." to "Please ensure that your answer is a valid Cypher query, and correct it if it is not." for better clarity.
Apply this diff to correct the grammatical errors:
-If you cannot generate a OpenCypher statement for any reason, return an empty string. +If you cannot generate an OpenCypher statement for any reason, return an empty string. -Please think if your answer is a valid Cypher query, and correct it if it is not. +Please ensure that your answer is a valid Cypher query, and correct it if it is not.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (9)
.wordlist.txt(2 hunks)README.md(2 hunks)graphrag_sdk/fixtures/prompts.py(1 hunks)graphrag_sdk/models/azure_openai.py(1 hunks)graphrag_sdk/models/gemini.py(1 hunks)graphrag_sdk/models/ollama.py(1 hunks)graphrag_sdk/models/openai.py(1 hunks)graphrag_sdk/orchestrator/orchestrator.py(2 hunks)graphrag_sdk/steps/graph_query_step.py(2 hunks)
🧰 Additional context used
🪛 Ruff
graphrag_sdk/orchestrator/orchestrator.py
3-3: graphrag_sdk.models.model.OutputMethod imported but unused
Remove unused import: graphrag_sdk.models.model.OutputMethod
(F401)
graphrag_sdk/steps/graph_query_step.py
45-45: Loop control variable i not used within loop body
Rename unused i to _i
(B007)
🔇 Additional comments (7)
graphrag_sdk/orchestrator/orchestrator.py (1)
Line range hint 1-67: Verify completeness of Cypher generation fix
While the structural changes in this file look good, I don't see direct changes addressing the Cypher generation iterative loop fix mentioned in the PR title.
Let's verify if there are related changes in other files:
✅ Verification successful
Let me analyze the changes to verify if they address the Cypher generation iterative loop fix. The key changes appear to be in graphrag_sdk/steps/graph_query_step.py where the Cypher generation loop was modified. Let me verify the specific changes to the loop logic.
Let me check the specific changes in the Cypher generation loop implementation by examining the relevant files.
Cypher generation loop has been successfully fixed
The changes in graphrag_sdk/steps/graph_query_step.py show a significant improvement to the Cypher generation loop:
- Changed from a
whileloop with error flag to a cleanerforloop with fixed number of retries - Removed redundant error-based prompt selection logic
- Added
delete_last_message()call to clean up failed attempts - Improved error handling and validation flow
The changes make the Cypher generation more robust by:
- Using a deterministic number of retries (10 attempts)
- Cleaning up the conversation context after failed attempts
- Simplifying the prompt selection logic
- Maintaining proper error propagation
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for Cypher-related code in the codebase
rg -l "cypher" --type py
# Search for any iterative or loop-related changes in recent commits
git log -p --since="1 week ago" | rg -A 5 -B 5 "(?i)(cypher|loop|iterate)"
Length of output: 83162
Script:
#!/bin/bash
# Get the complete implementation of GraphQueryGenerationStep class
ast-grep --pattern 'class GraphQueryGenerationStep {
$$$
}'
# Search for any other files that might contain Cypher generation logic
rg -l "cypher.*generation" --type py
Length of output: 184
Script:
#!/bin/bash
# Get the content of files with Cypher generation logic
cat graphrag_sdk/model_config.py graphrag_sdk/chat_session.py tests/test_kg_ollama.py
# Get the content of graph_query_step.py since it wasn't found by ast-grep
cat graphrag_sdk/steps/graph_query_step.py
Length of output: 11726
graphrag_sdk/steps/graph_query_step.py (2)
Line range hint 45-78: LGTM! Improved error handling and control flow.
The changes improve the code in several ways:
- More predictable loop structure with fixed iterations
- Better error handling with chat session cleanup
- Simplified prompt selection logic based on last_answer
🧰 Tools
🪛 Ruff
45-45: Loop control variable i not used within loop body
Rename unused i to _i
(B007)
43-43: Verify the necessity of increasing default retries to 10.
While increasing retries can help handle transient failures, it could also mask underlying issues and increase latency. Please confirm if this change is based on observed failure patterns in production.
✅ Verification successful
The retry count increase is part of a larger error handling improvement
The change from 5 to 10 retries was made alongside significant improvements to the error handling logic:
- Replaced error-prone
whileloop with a cleanerforloop - Removed the complex error-based prompt selection logic
- Added chat session cleanup on errors to prevent context pollution
- Simplified the overall retry mechanism
These changes make the increased retry count reasonable as the error handling is now more robust and cleaner, with proper cleanup between attempts.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for any related issues or discussions about retry failures
gh issue list --state all --search "retry cypher generation"
# Check commit history for context about retry-related changes
git log -p --grep="retry" -- graphrag_sdk/steps/graph_query_step.py
Length of output: 130
Script:
#!/bin/bash
# Let's try a broader search to understand the context and usage of retries
# Search for retry-related patterns in the file
rg -A 5 -B 5 "retries" graphrag_sdk/steps/graph_query_step.py
# Look for any error handling or failure patterns
rg -A 5 "except|error|fail" graphrag_sdk/steps/graph_query_step.py
# Check the git history with a broader context
git log -p graphrag_sdk/steps/graph_query_step.py
# Look for any test files that might provide insights
fd "test.*graph.*query.*step.*" -t f
Length of output: 6945
graphrag_sdk/models/openai.py (1)
110-111: Verify integration with Cypher generation retry logic
The PR title suggests this is part of fixing Cypher generation loop. Let's verify how this method is used in the retry mechanism.
✅ Verification successful
Let me gather more information about the usage of delete_last_message specifically.
Let me check the implementation of the Cypher generation method to understand the retry logic better.
Integration with Cypher generation retry logic is correctly implemented
The delete_last_message() method is used within the retry loop of Cypher generation. When an exception occurs during Cypher generation, the last message is deleted before retrying, allowing for a clean slate in the next attempt. This is the expected behavior for handling failed attempts in the retry mechanism.
Key implementation details:
- The retry logic uses a for-loop with configurable retries (default 10)
delete_last_message()is called in the exception handler before the next retry attempt- The method is consistently implemented across all LLM providers (OpenAI, Gemini, Azure, Ollama)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for usage of delete_last_message in Cypher generation related code
rg -l "delete_last_message" | xargs rg -A 5 "Cypher|retry"
# Look for retry logic in GraphQueryGenerationStep
ast-grep --pattern 'class GraphQueryGenerationStep {
$$$
retry($$$) {
$$$
}
$$$
}'
Length of output: 1310
Script:
#!/bin/bash
# Search for direct usage of delete_last_message
rg -B 5 -A 5 "delete_last_message"
# Look for retry-related logic in graph_query_step.py
rg -B 5 -A 5 "retry|attempt" graphrag_sdk/steps/graph_query_step.py
Length of output: 2484
Script:
#!/bin/bash
# Look for the main Cypher generation method
ast-grep --pattern 'def generate_cypher($$$) {
$$$
}'
# Also check for any while loops or iteration logic in graph_query_step.py
rg -B 5 -A 10 "while|for" graphrag_sdk/steps/graph_query_step.py
Length of output: 1062
graphrag_sdk/models/ollama.py (2)
220-221: Verify consistent implementation across chat sessions
Let's ensure the delete_last_message implementation is consistent across all chat session classes.
#!/bin/bash
# Description: Check delete_last_message implementation in other chat sessions
# Expected: Similar implementations in other chat session classes
# Find all chat session classes and their delete_last_message implementations
ast-grep --pattern 'class $_ChatSession($_) {
$$$
delete_last_message($$$) {
$$$
}
$$$
}'
# Backup: Search using ripgrep
rg -A 5 "def delete_last_message" graphrag_sdk/models/220-221: Verify usage in Cypher generation retry logic
Let's understand how this method is used in the Cypher generation retry loop mentioned in the PR title.
✅ Verification successful
delete_last_message() is correctly used in Cypher generation retry logic
The method is properly integrated within the retry loop of GraphQueryGenerationStep.run(). When a Cypher query generation attempt fails (due to validation errors or execution issues), delete_last_message() is called to remove the failed attempt from the chat history before the next retry iteration. This ensures a clean slate for subsequent attempts within the maximum retry limit of 10.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find where delete_last_message is used in Cypher generation
# Expected: Usage in retry logic for Cypher generation
# Search for GraphQueryGenerationStep class and its retry logic
ast-grep --pattern 'class GraphQueryGenerationStep {
$$$
$_($$$) {
$$$
delete_last_message
$$$
}
$$$
}'
# Backup: Search using ripgrep
rg -B 10 -A 10 "delete_last_message" graphrag_sdk/steps/
Length of output: 1396
Script:
#!/bin/bash
# Let's get more context around the usage in graph_query_step.py
rg -B 20 -A 10 "delete_last_message" graphrag_sdk/steps/graph_query_step.py
# Also check the class definition and retry mechanism
ast-grep --pattern 'class GraphQueryGenerationStep {
$$$
execute($$$) {
$$$
}
}'
Length of output: 1161
Script:
#!/bin/bash
# Let's get the full context of the retry loop and class
cat graphrag_sdk/steps/graph_query_step.py
Length of output: 2791
graphrag_sdk/models/azure_openai.py (1)
231-232: Verify consistency with other chat session implementations
Let's ensure the delete_last_message implementation is consistent across all chat session classes.
✅ Verification successful
Implementation is consistent across all chat session classes
The verification shows that delete_last_message is implemented consistently across all chat session classes:
- All implementations have identical method signatures:
def delete_last_message(self) - All implementations follow the same pattern of removing the last two items from their respective history arrays:
- OpenAI:
self._history[:-2] - Gemini:
self._chat_session.history[:-2] - Azure OpenAI:
self._history[:-2] - Ollama:
self._chat_history[:-2]
- OpenAI:
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check implementation consistency across chat session classes
# Find all chat session classes and their delete_last_message implementations
ast-grep --pattern 'class $CHAT_CLASS {
$$$
delete_last_message($$$) {
$IMPL
}
$$$
}'
# Check for any differences in method signatures
rg -A 1 "def delete_last_message"
Length of output: 744
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
graphrag_sdk/models/ollama.py (1)
220-237: Implementation looks good with room for minor improvementThe implementation effectively addresses the key concerns:
- Proper documentation with docstring
- System message preservation
- Safe handling of chat history
However, consider adding error handling for edge cases:
def delete_last_message(self): """ Deletes the last message exchange (user message and assistant response) from the chat history. Preserves the system message if present. Note: Does nothing if the chat history is empty or contains only a system message. """ + if not self._chat_history: + return + # Keep at least the system message if present min_length = 1 if self._model.system_instruction else 0 - if len(self._chat_history) - 2 >= min_length: + if (len(self._chat_history) >= 2 and # Ensure we have at least a message pair + self._chat_history[-2]["role"] == "user" and # Validate message pair structure + self._chat_history[-1]["role"] == "assistant" and + len(self._chat_history) - 2 >= min_length): # Check system message preservation self._chat_history = self._chat_history[:-2] else: # Reset to initial state with just system message if present self._chat_history = ( [{"role": "system", "content": self._model.system_instruction}] if self._model.system_instruction is not None else [] )The suggested changes add validation to ensure:
- The chat history exists
- We have a complete message pair (user + assistant)
- The message pair has the correct structure
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
graphrag_sdk/models/azure_openai.py(1 hunks)graphrag_sdk/models/gemini.py(1 hunks)graphrag_sdk/models/ollama.py(1 hunks)graphrag_sdk/models/openai.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- graphrag_sdk/models/azure_openai.py
- graphrag_sdk/models/gemini.py
- graphrag_sdk/models/openai.py
2 participants
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Refactor