Add configurable prompts #35
Conversation
WalkthroughThe changes involve modifications to several classes within the Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant KnowledgeGraph
participant ChatSession
participant GraphQueryGenerationStep
participant QAStep
User->>KnowledgeGraph: chat_session(cypher_system_instruction, qa_system_instruction, cypher_gen_prompt, qa_prompt)
KnowledgeGraph->>ChatSession: Create with provided parameters
ChatSession->>GraphQueryGenerationStep: Run with cypher_prompt
ChatSession->>QAStep: Run with qa_prompt
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
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: 2
🧹 Outside diff range and nitpick comments (6)
graphrag_sdk/steps/graph_query_step.py (1)
54-62: Fix indentation consistency.The indentation in the ternary expressions is inconsistent, which affects readability. Ensure consistent indentation throughout the nested expressions.
- cypher_prompt = ( - (CYPHER_GEN_PROMPT.format(question=question) - if self.last_answer is None - else CYPHER_GEN_PROMPT_WITH_HISTORY.format(question=question, last_answer=self.last_answer)) - if error is False - else CYPHER_GEN_PROMPT_WITH_ERROR.format( - question=question, error=error - ) - ) + cypher_prompt = ( + (CYPHER_GEN_PROMPT.format(question=question) + if self.last_answer is None + else CYPHER_GEN_PROMPT_WITH_HISTORY.format( + question=question, + last_answer=self.last_answer + )) + if error is False + else CYPHER_GEN_PROMPT_WITH_ERROR.format( + question=question, + error=error + ) + )graphrag_sdk/chat_session.py (3)
54-55: Add type hints and documentation for new attributes.The new prompt attributes would benefit from type hints and docstring documentation to improve code maintainability.
Add type hints and update the class docstring:
class ChatSession: + cypher_prompt: str | None + qa_prompt: str | None """ Represents a chat session with a Knowledge Graph. Args: model_config (KnowledgeGraphModelConfig): The model configuration to use. ontology (Ontology): The ontology to use. graph (Graph): The graph to query. + cypher_system_instruction (str, optional): Custom system instructions for Cypher generation. + qa_system_instruction (str, optional): Custom system instructions for QA. + cypher_gen_prompt (str, optional): Custom prompt template for Cypher generation. + qa_prompt (str, optional): Custom prompt template for QA.
63-63: Consider consistent instruction handling approaches.The QA system instruction handling uses a different approach compared to the Cypher system instruction handling. Consider using the same pattern for consistency.
- qa_system_instruction or GRAPH_QA_SYSTEM + if qa_system_instruction is None: + qa_system_instruction = GRAPH_QA_SYSTEM + self.qa_chat_session = model_config.qa.with_system_instruction(qa_system_instruction).start_chat()
82-82: Consider caching step instances.The steps are recreated for each message, which could be inefficient for long chat sessions since most parameters remain constant.
Consider creating the steps once in the constructor:
def __init__(self, ...): # ... existing code ... + self.cypher_step = GraphQueryGenerationStep( + graph=self.graph, + chat_session=self.cypher_chat_session, + ontology=self.ontology, + cypher_prompt=self.cypher_prompt, + ) + self.qa_step = QAStep( + chat_session=self.qa_chat_session, + qa_prompt=self.qa_prompt, + ) def send_message(self, message: str): - cypher_step = GraphQueryGenerationStep(...) + self.cypher_step.last_answer = self.last_answer + (context, cypher) = self.cypher_step.run(message) # ... rest of the method ... - qa_step = QAStep(...) + answer = self.qa_step.run(message, cypher, context)Also applies to: 92-92
graphrag_sdk/kg.py (2)
137-141: Update method docstring to document new parametersThe method signature has been updated with new parameters, but the docstring hasn't been updated to reflect these changes. Please add parameter descriptions to help users understand the purpose of each new parameter.
Apply this diff:
def chat_session(self, cypher_system_instruction: str = None, qa_system_instruction: str = None, cypher_gen_prompt: str = None, qa_prompt: str = None) -> ChatSession: + """ + Create a new chat session with optional custom instructions and prompts. + + Parameters: + cypher_system_instruction (str, optional): Custom system instruction for Cypher query generation + qa_system_instruction (str, optional): Custom system instruction for question answering + cypher_gen_prompt (str, optional): Custom prompt template for Cypher query generation + qa_prompt (str, optional): Custom prompt template for question answering + + Returns: + ChatSession: A new chat session instance + """ chat_session = ChatSession(self._model_config, self.ontology, self.graph, cypher_system_instruction, qa_system_instruction, cypher_gen_prompt, qa_prompt) return chat_session
137-141: Consider adding parameter validation and default valuesThe method accepts optional string parameters but doesn't validate them or provide default values. Consider:
- Validating that provided strings are not empty
- Using the imported
GRAPH_QA_SYSTEMandCYPHER_GEN_SYSTEMas default valuesHere's a suggested implementation:
def chat_session(self, cypher_system_instruction: str = None, qa_system_instruction: str = None, cypher_gen_prompt: str = None, qa_prompt: str = None) -> ChatSession: + # Use default system instructions if not provided + cypher_system_instruction = cypher_system_instruction or CYPHER_GEN_SYSTEM + qa_system_instruction = qa_system_instruction or GRAPH_QA_SYSTEM + + # Validate non-empty strings if provided + if cypher_gen_prompt is not None and not cypher_gen_prompt.strip(): + raise ValueError("cypher_gen_prompt cannot be empty") + if qa_prompt is not None and not qa_prompt.strip(): + raise ValueError("qa_prompt cannot be empty") + chat_session = ChatSession(self._model_config, self.ontology, self.graph, cypher_system_instruction, qa_system_instruction, cypher_gen_prompt, qa_prompt) return chat_session
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
graphrag_sdk/chat_session.py(4 hunks)graphrag_sdk/kg.py(1 hunks)graphrag_sdk/steps/graph_query_step.py(1 hunks)graphrag_sdk/steps/qa_step.py(1 hunks)
🔇 Additional comments (5)
graphrag_sdk/steps/qa_step.py (2)
20-20: LGTM! Clean implementation of configurable prompts
The addition of the optional qa_prompt parameter with proper type hinting and None default maintains backward compatibility while enabling prompt customization.
Also applies to: 24-24
27-32: Verify format string compatibility with custom prompts
The code assumes that any custom qa_prompt will contain the format placeholders {context}, {cypher}, and {question}. If a custom prompt is provided without these placeholders, it will raise a KeyError at runtime.
Let's verify the format string requirements are documented:
Consider adding validation for required placeholders:
def run(self, question: str, cypher: str, context: str):
graph_qa_prompt = self.qa_prompt or GRAPH_QA_SYSTEM
+ required_placeholders = ["{context}", "{cypher}", "{question}"]
+ if self.qa_prompt and not all(p in self.qa_prompt for p in required_placeholders):
+ raise ValueError(
+ f"Custom qa_prompt must contain all placeholders: {required_placeholders}"
+ )
qa_prompt = graph_qa_prompt.format(
context=context, cypher=cypher, question=question
)graphrag_sdk/steps/graph_query_step.py (1)
36-36: LGTM! Clean implementation of the configurable prompt.
The addition of the optional cypher_prompt parameter and its corresponding instance variable is well-implemented, following Python's best practices.
Also applies to: 43-43
graphrag_sdk/chat_session.py (2)
28-30: LGTM! Well-structured parameter additions.
The new optional parameters are well-named and properly typed, maintaining backward compatibility while adding configurability.
28-30: Verify consistent parameter usage across the codebase.
The new parameters should be consistently handled in related classes and methods.
✅ Verification successful
Let me gather more specific information about the parameter usage.
Parameters are correctly propagated and used throughout the codebase
The verification shows that:
- The new parameters are properly handled in
KnowledgeGraph.chat_session()method which correctly forwards them toChatSession ChatSessionhas proper fallback logic forcypher_system_instructionandqa_system_instruction- The parameters are used in the expected way with
GraphQueryGenerationStepandQAStepclasses - No inconsistencies found in parameter usage across different model implementations (Gemini, OpenAI, Azure, Ollama)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for related class instantiations and verify parameter handling
echo "Searching for ChatSession instantiations..."
rg -l "ChatSession\("
echo "\nSearching for related step classes..."
rg -l "class.*Step"
echo "\nChecking GraphQueryGenerationStep and QAStep parameter handling..."
ast-grep --pattern 'class GraphQueryGenerationStep {
$$$
def __init__($$$) {
$$$
}
$$$
}'
ast-grep --pattern 'class QAStep {
$$$
def __init__($$$) {
$$$
}
$$$
}'
Length of output: 1564
Script:
#!/bin/bash
# Check ChatSession instantiations and their parameters
echo "Checking ChatSession instantiations..."
rg "ChatSession\(" -A 5
echo "\nChecking GraphQueryGenerationStep implementation..."
rg "class GraphQueryGenerationStep" -A 10 graphrag_sdk/steps/graph_query_step.py
echo "\nChecking QAStep implementation..."
rg "class QAStep" -A 10 graphrag_sdk/steps/qa_step.py
echo "\nChecking where cypher_system_instruction and qa_system_instruction are used..."
rg "cypher_system_instruction|qa_system_instruction"
Length of output: 6099
| if self.cypher_prompt is not None: | ||
| cypher_prompt = self.cypher_prompt | ||
| else: | ||
| cypher_prompt = ( | ||
| (CYPHER_GEN_PROMPT.format(question=question) | ||
| if self.last_answer is None | ||
| else CYPHER_GEN_PROMPT_WITH_HISTORY.format(question=question, last_answer=self.last_answer)) | ||
| if error is False | ||
| else CYPHER_GEN_PROMPT_WITH_ERROR.format( | ||
| question=question, error=error | ||
| ) | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Refactor nested ternary expressions for better readability.
The current implementation uses nested ternary expressions which can be hard to read and maintain. Consider extracting the logic into a separate method for better clarity.
Here's a suggested refactor:
- if self.cypher_prompt is not None:
- cypher_prompt = self.cypher_prompt
- else:
- cypher_prompt = (
- (CYPHER_GEN_PROMPT.format(question=question)
- if self.last_answer is None
- else CYPHER_GEN_PROMPT_WITH_HISTORY.format(question=question, last_answer=self.last_answer))
- if error is False
- else CYPHER_GEN_PROMPT_WITH_ERROR.format(
- question=question, error=error
- )
- )
+ cypher_prompt = self._get_cypher_prompt(question, error)
+
+ def _get_cypher_prompt(self, question: str, error: Exception | bool) -> str:
+ if self.cypher_prompt is not None:
+ return self.cypher_prompt
+
+ if error:
+ return CYPHER_GEN_PROMPT_WITH_ERROR.format(
+ question=question,
+ error=error
+ )
+
+ if self.last_answer is None:
+ return CYPHER_GEN_PROMPT.format(question=question)
+
+ return CYPHER_GEN_PROMPT_WITH_HISTORY.format(
+ question=question,
+ last_answer=self.last_answer
+ )Committable suggestion skipped: line range outside the PR's diff.
| if cypher_system_instruction is None: | ||
| cypher_system_instruction = CYPHER_GEN_SYSTEM.replace("#ONTOLOGY", str(ontology.to_json())) | ||
| else: | ||
| cypher_system_instruction = cypher_system_instruction + "\nOntology:\n" + str(ontology.to_json()) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Standardize ontology injection method.
The code uses two different approaches to inject the ontology:
- Using
replace("#ONTOLOGY", str(ontology.to_json()))for the default case - Using string concatenation with
"\nOntology:\n"for custom instructions
This inconsistency could lead to formatting differences and maintenance issues.
Consider standardizing the approach:
- if cypher_system_instruction is None:
- cypher_system_instruction = CYPHER_GEN_SYSTEM.replace("#ONTOLOGY", str(ontology.to_json()))
- else:
- cypher_system_instruction = cypher_system_instruction + "\nOntology:\n" + str(ontology.to_json())
+ base_instruction = cypher_system_instruction or CYPHER_GEN_SYSTEM
+ ontology_json = str(ontology.to_json())
+ cypher_system_instruction = base_instruction.replace("#ONTOLOGY", ontology_json)Committable suggestion skipped: line range outside the PR's diff.
Summary by CodeRabbit
New Features
ChatSessionclass to support custom instructions and prompts for Cypher and QA sessions.KnowledgeGraphclass to allow additional parameters for initiating chat sessions.GraphQueryGenerationStepandQAStepclasses for improved flexibility.Bug Fixes
GraphQueryGenerationStepandQAStepto streamline prompt generation.Documentation