graphrag_sdk 0.3.0__tar.gz → 0.3.2__tar.gz

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphrag_sdk
-Version: 0.3.0
+Version: 0.3.2
 Summary: 
 Home-page: http://RAG-SDK.readthedocs.io
 License: MIT
@@ -32,6 +32,7 @@ Requires-Dist: pypdf (>=4.2.0,<5.0.0)
 Requires-Dist: python-abc (>=0.2.0,<0.3.0)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: ratelimit (>=2.2.1,<3.0.0)
+Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: typing-extensions (>=4.12.1,<5.0.0)
 Requires-Dist: vertexai (>=1.49.0,<2.0.0) ; extra == "all" or extra == "vertexai"
 Project-URL: Repository, https://github.com/FalkorDB/GraphRAG-SDK
@@ -86,19 +87,9 @@ docker run -p 6379:6379 -p 3000:3000 -it --rm  -v ./data:/data falkordb/falkordb
 Currently, this SDK supports the following LLMs API:
 
 * [OpenAI](https://openai.com/index/openai-api) Recommended model:`gpt-4o`
-* [google](https://makersuite.google.com/app/apikey) Recommended model:`gemini-1.5-flash-001`
-
-Make sure that a `.env` file is present with all required credentials.
-
-   <details>
-     <summary>.env</summary>
-  
-   ```
-   OPENAI_API_KEY="OPENAI_API_KEY"
-   GOOGLE_API_KEY="GOOGLE_API_KEY"
-   ```
-  
-   </details>
+* [Google](https://makersuite.google.com/app/apikey) Recommended model:`gemini-1.5-flash-001`
+* [Azure-OpenAI](https://ai.azure.com) Recommended model:`gpt-4o`
+* [Ollama](https://ollama.com/) Available only to the Q&A step. Recommended models: `llama3`. Ollama models are suitable for the Q&A step only (after the knowledge graph (KG) created).
 
 ## Basic Usage
 
@@ -219,9 +210,9 @@ After generating the initial ontology, you can review it and make any necessary
 
 Once you are satisfied with the ontology, you can proceed to use it for creating and managing your Knowledge Graph (KG).
 
-### Knowledge Graph
+### Knowledge Graph Agent
 
-Now, you can use the SDK to create a Knowledge Graph (KG) from your sources and ontology.
+Now, you can use the SDK to create a Knowledge Graph (KG) from your sources and ontology for Q&A.
 
 ```python
 # After approving the ontology, load it from disk.
@@ -240,6 +231,45 @@ kg.process_sources(sources)
 
 You can update the KG at any time by processing more sources with the `process_sources` method.
 
+#### Configurable Prompts
+When creating your Knowledge Graph (KG) agent, you can customize the prompts to tailor its behavior. This step is optional but can enhance functionality.
+
+There are five types of prompts:
+
+1. **`cypher_system_instruction`**  
+   - System instructions for the Cypher generation step.  
+   - **Note:** Ensure your prompt includes `{ontology}`.
+
+2. **`qa_system_instruction`**  
+   - System instructions for the Q&A step.
+
+3. **`cypher_gen_prompt`**  
+   - The prompt used during the Cypher generation step.  
+   - **Note:** Include `{question}` in your prompt.
+
+4. **`cypher_gen_prompt_history`**  
+   - The prompt for Cypher generation when history needs to be considered.  
+   - **Note:** Include `{question}` and `{last_answer}` in your prompt.
+
+5. **`qa_prompt`**  
+   - The prompt used during the Q&A step.  
+   - **Note:** Include `{question}`, `{context}`, and `{cypher}` in your prompt.
+
+Here’s an example configuration:
+
+```python
+kg = KnowledgeGraph(
+    name="kg_name",
+    model_config=KnowledgeGraphModelConfig.with_model(model),
+    ontology=ontology,
+    cypher_system_instruction=cypher_system_instruction,
+    qa_system_instruction=qa_system_instruction,
+    cypher_gen_prompt=cypher_gen_prompt,
+    cypher_gen_prompt_history=cypher_gen_prompt_history,
+    qa_prompt=qa_prompt
+)
+```
+
 ### Graph RAG
 
 At this point, you have a Knowledge Graph that can be queried using this SDK. Use the method `chat_session` for start a conversation.

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/README.md

@@ -47,19 +47,9 @@ docker run -p 6379:6379 -p 3000:3000 -it --rm  -v ./data:/data falkordb/falkordb
 Currently, this SDK supports the following LLMs API:
 
 * [OpenAI](https://openai.com/index/openai-api) Recommended model:`gpt-4o`
-* [google](https://makersuite.google.com/app/apikey) Recommended model:`gemini-1.5-flash-001`
-
-Make sure that a `.env` file is present with all required credentials.
-
-   <details>
-     <summary>.env</summary>
-  
-   ```
-   OPENAI_API_KEY="OPENAI_API_KEY"
-   GOOGLE_API_KEY="GOOGLE_API_KEY"
-   ```
-  
-   </details>
+* [Google](https://makersuite.google.com/app/apikey) Recommended model:`gemini-1.5-flash-001`
+* [Azure-OpenAI](https://ai.azure.com) Recommended model:`gpt-4o`
+* [Ollama](https://ollama.com/) Available only to the Q&A step. Recommended models: `llama3`. Ollama models are suitable for the Q&A step only (after the knowledge graph (KG) created).
 
 ## Basic Usage
 
@@ -180,9 +170,9 @@ After generating the initial ontology, you can review it and make any necessary
 
 Once you are satisfied with the ontology, you can proceed to use it for creating and managing your Knowledge Graph (KG).
 
-### Knowledge Graph
+### Knowledge Graph Agent
 
-Now, you can use the SDK to create a Knowledge Graph (KG) from your sources and ontology.
+Now, you can use the SDK to create a Knowledge Graph (KG) from your sources and ontology for Q&A.
 
 ```python
 # After approving the ontology, load it from disk.
@@ -201,6 +191,45 @@ kg.process_sources(sources)
 
 You can update the KG at any time by processing more sources with the `process_sources` method.
 
+#### Configurable Prompts
+When creating your Knowledge Graph (KG) agent, you can customize the prompts to tailor its behavior. This step is optional but can enhance functionality.
+
+There are five types of prompts:
+
+1. **`cypher_system_instruction`**  
+   - System instructions for the Cypher generation step.  
+   - **Note:** Ensure your prompt includes `{ontology}`.
+
+2. **`qa_system_instruction`**  
+   - System instructions for the Q&A step.
+
+3. **`cypher_gen_prompt`**  
+   - The prompt used during the Cypher generation step.  
+   - **Note:** Include `{question}` in your prompt.
+
+4. **`cypher_gen_prompt_history`**  
+   - The prompt for Cypher generation when history needs to be considered.  
+   - **Note:** Include `{question}` and `{last_answer}` in your prompt.
+
+5. **`qa_prompt`**  
+   - The prompt used during the Q&A step.  
+   - **Note:** Include `{question}`, `{context}`, and `{cypher}` in your prompt.
+
+Here’s an example configuration:
+
+```python
+kg = KnowledgeGraph(
+    name="kg_name",
+    model_config=KnowledgeGraphModelConfig.with_model(model),
+    ontology=ontology,
+    cypher_system_instruction=cypher_system_instruction,
+    qa_system_instruction=qa_system_instruction,
+    cypher_gen_prompt=cypher_gen_prompt,
+    cypher_gen_prompt_history=cypher_gen_prompt_history,
+    qa_prompt=qa_prompt
+)
+```
+
 ### Graph RAG
 
 At this point, you have a Knowledge Graph that can be queried using this SDK. Use the method `chat_session` for start a conversation.

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/agents/kg_agent.py

@@ -136,7 +136,7 @@ class KGAgent(Agent):
 
         """
         output = self.chat_session.send_message(params["prompt"])
-        return output
+        return output['response']
 
     def __repr__(self):
         """

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/chat_session.py

@@ -1,9 +1,8 @@
+from falkordb import Graph
 from graphrag_sdk.ontology import Ontology
+from graphrag_sdk.steps.qa_step import QAStep
 from graphrag_sdk.model_config import KnowledgeGraphModelConfig
 from graphrag_sdk.steps.graph_query_step import GraphQueryGenerationStep
-from graphrag_sdk.steps.qa_step import QAStep
-from graphrag_sdk.fixtures.prompts import GRAPH_QA_SYSTEM, CYPHER_GEN_SYSTEM
-from falkordb import Graph
 
 
 class ChatSession:
@@ -25,7 +24,9 @@ class ChatSession:
         >>> chat_session.send_message("What is the capital of France?")
     """
 
-    def __init__(self, model_config: KnowledgeGraphModelConfig, ontology: Ontology, graph: Graph):
+    def __init__(self, model_config: KnowledgeGraphModelConfig, ontology: Ontology, graph: Graph,
+                cypher_system_instruction: str, qa_system_instruction: str,
+                cypher_gen_prompt: str, qa_prompt: str, cypher_gen_prompt_history: str):
         """
         Initializes a new ChatSession object.
 
@@ -44,14 +45,21 @@ class ChatSession:
         self.model_config = model_config
         self.graph = graph
         self.ontology = ontology
+        cypher_system_instruction = cypher_system_instruction.format(ontology=str(ontology.to_json()))
+
+        
+        self.cypher_prompt = cypher_gen_prompt
+        self.qa_prompt = qa_prompt
+        self.cypher_prompt_with_history = cypher_gen_prompt_history
+        
         self.cypher_chat_session = (
             model_config.cypher_generation.with_system_instruction(
-                CYPHER_GEN_SYSTEM.replace("#ONTOLOGY", str(ontology.to_json()))
+                cypher_system_instruction
             ).start_chat()
         )
         self.qa_chat_session = model_config.qa.with_system_instruction(
-            GRAPH_QA_SYSTEM
-        ).start_chat()
+                qa_system_instruction
+            ).start_chat()
         self.last_answer = None
 
     def send_message(self, message: str):
@@ -62,24 +70,42 @@ class ChatSession:
             message (str): The message to send.
 
         Returns:
-            str: The response to the message.
+            dict: The response to the message in the following format:
+                    {"question": message, 
+                    "response": answer, 
+                    "context": context, 
+                    "cypher": cypher}
         """
         cypher_step = GraphQueryGenerationStep(
             graph=self.graph,
             chat_session=self.cypher_chat_session,
             ontology=self.ontology,
             last_answer=self.last_answer,
+            cypher_prompt=self.cypher_prompt,
+            cypher_prompt_with_history=self.cypher_prompt_with_history
         )
 
         (context, cypher) = cypher_step.run(message)
 
         if not cypher or len(cypher) == 0:
-            return "I am sorry, I could not find the answer to your question"
+            return {
+                "question": message,
+                "response": "I am sorry, I could not find the answer to your question",
+                "context": None,
+                "cypher": None
+                }
 
         qa_step = QAStep(
             chat_session=self.qa_chat_session,
+            qa_prompt=self.qa_prompt,
         )
 
         answer = qa_step.run(message, cypher, context)
         self.last_answer = answer
-        return answer
+        
+        return {
+            "question": message, 
+            "response": answer, 
+            "context": context, 
+            "cypher": cypher
+            }

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/fixtures/prompts.py

@@ -429,7 +429,7 @@ Single-Pair minimal-weight paths: Find minimal-weight paths between a pair of en
 Single-Source minimal-weight paths: Find minimal-weight paths from a given source entity using algo.SSpaths().
 
 Ontology:
-#ONTOLOGY
+{ontology}
 
 
 For example, given the question "Which managers own Neo4j stocks?", the OpenCypher statement should look like this:
@@ -441,9 +441,10 @@ RETURN m, s
 """
 
 CYPHER_GEN_PROMPT = """
-Using the ontology provided, generate an OpenCypher statement to query the graph database returning all relevant entities, relationships, and attributes to answer the question below:
-If you cannot generate a OpenCypher statement for any reason, return an empty string.
+Using the ontology provided, generate an OpenCypher statement to query the graph database returning all relevant entities, relationships, and attributes to answer the question below.
+If you cannot generate a OpenCypher statement for any reason, return an empty response.
 Respect the order of the relationships, the arrows should always point from the "source" to the "target".
+Please think if your answer is a valid Cypher query, and correct it if it is not.
 
 Question: {question}
 """
@@ -481,8 +482,8 @@ The information part contains the provided information that you must use to cons
 The provided information is authoritative, you must never doubt it or try to use your internal knowledge to correct it.
 Make the answer sound as a response to the question. Do not mention that you based the result on the given information.
 Do not answer more than the question asks for.
-Here is an example:
 
+Here is an example:
 Question: Which managers own Neo4j stocks?
 Context:[manager:CTL LLC, manager:JANE STREET GROUP LLC]
 Helpful Answer: CTL LLC, JANE STREET GROUP LLC owns Neo4j stocks.

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/kg.py

@@ -1,16 +1,15 @@
 import logging
-from graphrag_sdk.ontology import Ontology
+import warnings
 from falkordb import FalkorDB
+from graphrag_sdk.ontology import Ontology
 from graphrag_sdk.source import AbstractSource
-from graphrag_sdk.model_config import KnowledgeGraphModelConfig
-from graphrag_sdk.steps.extract_data_step import ExtractDataStep
-from graphrag_sdk.steps.graph_query_step import GraphQueryGenerationStep
-from graphrag_sdk.fixtures.prompts import GRAPH_QA_SYSTEM, CYPHER_GEN_SYSTEM
-from graphrag_sdk.steps.qa_step import QAStep
 from graphrag_sdk.chat_session import ChatSession
-from graphrag_sdk.helpers import map_dict_to_cypher_properties
 from graphrag_sdk.attribute import AttributeType, Attribute
-from graphrag_sdk.models import GenerativeModelChatSession
+from graphrag_sdk.helpers import map_dict_to_cypher_properties
+from graphrag_sdk.model_config import KnowledgeGraphModelConfig
+from graphrag_sdk.steps.extract_data_step import ExtractDataStep
+from graphrag_sdk.fixtures.prompts import (GRAPH_QA_SYSTEM, CYPHER_GEN_SYSTEM,
+                                CYPHER_GEN_PROMPT, GRAPH_QA_PROMPT, CYPHER_GEN_PROMPT_WITH_HISTORY)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -31,6 +30,11 @@ class KnowledgeGraph:
         port: int = 6379,
         username: str | None = None,
         password: str | None = None,
+        cypher_system_instruction: str = None,
+        qa_system_instruction: str = None,
+        cypher_gen_prompt: str = None,
+        qa_prompt: str = None,
+        cypher_gen_prompt_history: str = None,
     ):
         """
         Initialize Knowledge Graph
@@ -43,6 +47,11 @@ class KnowledgeGraph:
             username (str|None): FalkorDB username.
             password (str|None): FalkorDB password.
             ontology (Ontology|None): Ontology to use.
+            cypher_system_instruction (str|None): Cypher system instruction. Make sure you have {ontology} in the instruction.
+            qa_system_instruction (str|None): QA system instruction.
+            cypher_gen_prompt (str|None): Cypher generation prompt. Make sure you have {question} in the prompt.
+            qa_prompt (str|None): QA prompt. Make sure you have {question}, {context} and {cypher} in the prompt.
+            cypher_gen_prompt_history (str|None): Cypher generation prompt with history. Make sure you have {question} and {last_answer} in the prompt.
         """
 
         if not isinstance(name, str) or name == "":
@@ -57,6 +66,44 @@ class KnowledgeGraph:
         self._model_config = model_config
         self.sources = set([])
 
+        if cypher_system_instruction is None:
+            cypher_system_instruction = CYPHER_GEN_SYSTEM
+        else:
+            if "{ontology}" not in cypher_system_instruction:
+                warnings.warn("Cypher system instruction should contain {ontology}", category=UserWarning)
+
+        if qa_system_instruction is None:
+            qa_system_instruction = GRAPH_QA_SYSTEM
+
+        if cypher_gen_prompt is None:
+            cypher_gen_prompt = CYPHER_GEN_PROMPT
+        else:
+            if "{question}" not in cypher_gen_prompt:
+                raise Exception("Cypher generation prompt should contain {question}")
+
+        if qa_prompt is None:
+            qa_prompt = GRAPH_QA_PROMPT
+        else:
+            if "{question}" not in qa_prompt or "{context}" not in qa_prompt:
+                raise Exception("QA prompt should contain {question} and {context}")
+            if "{cypher}" not in qa_prompt:
+                warnings.warn("QA prompt should contain {cypher}", category=UserWarning)
+
+        if cypher_gen_prompt_history is None:
+            cypher_gen_prompt_history = CYPHER_GEN_PROMPT_WITH_HISTORY
+        else:
+            if "{question}" not in cypher_gen_prompt_history:
+                raise Exception("Cypher generation prompt with history should contain {question}")
+            if "{last_answer}" not in cypher_gen_prompt_history:
+                warnings.warn("Cypher generation prompt with history should contain {last_answer}", category=UserWarning)
+
+        # Assign the validated values
+        self.cypher_system_instruction = cypher_system_instruction
+        self.qa_system_instruction = qa_system_instruction
+        self.cypher_gen_prompt = cypher_gen_prompt
+        self.qa_prompt = qa_prompt
+        self.cypher_gen_prompt_history = cypher_gen_prompt_history
+
     # Attributes
 
     @property
@@ -135,8 +182,9 @@ class KnowledgeGraph:
             setattr(self, key, None)
 
     def chat_session(self) -> ChatSession:
-        return ChatSession(self._model_config, self.ontology, self.graph)
-
+        chat_session = ChatSession(self._model_config, self.ontology, self.graph, self.cypher_system_instruction,
+                                   self.qa_system_instruction, self.cypher_gen_prompt, self.qa_prompt, self.cypher_gen_prompt_history)
+        return chat_session
     def add_node(self, entity: str, attributes: dict):
         """
         Add a node to the knowledge graph, checking if it matches the ontology

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/models/azure_openai.py

@@ -227,3 +227,35 @@ class AzureOpenAiChatSession(GenerativeModelChatSession):
             config['response_format'] = { "type": "json_object" }
         
         return config
+    
+    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.
+        
+        Example:
+            Before:
+            [
+                {"role": "system", "content": "System message"},
+                {"role": "user", "content": "User message"},
+                {"role": "assistant", "content": "Assistant response"},
+            ]
+            After:
+            [
+                {"role": "system", "content": "System message"},
+            ]
+        
+        Note: Does nothing if the chat history is empty or contains only a system message.
+        """
+        # Keep at least the system message if present
+        min_length = 1 if self._model.system_instruction else 0
+        if len(self._history) - 2 >= min_length:
+            self._history.pop()
+            self._history.pop()
+        else:
+            # Reset to initial state with just system message if present
+            self._history = (
+            [{"role": "system", "content": self._model.system_instruction}]
+            if self._model.system_instruction is not None
+            else []
+        )

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/models/gemini.py

@@ -117,3 +117,26 @@ class GeminiChatSession(GenerativeModelChatSession):
                 "temperature": 0
             }
         return self._model._generation_config
+    
+    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.
+        
+        Example:
+            Before:
+            [
+                {"role": "user", "content": "User message"},
+                {"role": "assistant", "content": "Assistant response"},
+            ]
+            After:
+            []
+
+        Note: Does nothing if the chat history is empty or contains only a system message.
+        """
+        if len(self._chat_session.history) >= 2:
+            self._chat_session.history.pop()
+            self._chat_session.history.pop()
+        else:
+            self._chat_session.history = []
+

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/models/ollama.py

@@ -216,3 +216,35 @@ class OllamaChatSession(GenerativeModelChatSession):
             config['format'] = 'json'
         
         return config
+    
+    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.
+        
+        Example:
+            Before:
+            [
+                {"role": "system", "content": "System message"},
+                {"role": "user", "content": "User message"},
+                {"role": "assistant", "content": "Assistant response"},
+            ]
+            After:
+            [
+                {"role": "system", "content": "System message"},
+            ]
+
+        Note: Does nothing if the chat history is empty or contains only a system message.
+        """
+        # 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:
+            self._chat_history.pop()
+            self._chat_history.pop()
+        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 []
+        )

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/models/openai.py

@@ -106,3 +106,35 @@ class OpenAiChatSession(GenerativeModelChatSession):
             config['response_format'] = { "type": "json_object" }
         
         return config
+    
+    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.
+        
+        Example:
+            Before:
+            [
+                {"role": "system", "content": "System message"},
+                {"role": "user", "content": "User message"},
+                {"role": "assistant", "content": "Assistant response"},
+            ]
+            After:
+            [
+                {"role": "system", "content": "System message"},
+            ]
+
+        Note: Does nothing if the chat history is empty or contains only a system message.
+        """
+        # Keep at least the system message if present
+        min_length = 1 if self._model.system_instruction else 0
+        if len(self._history) - 2 >= min_length:
+            self._history.pop()
+            self._history.pop()
+        else:
+            # Reset to initial state with just system message if present
+            self._history = (
+            [{"role": "system", "content": self._model.system_instruction}]
+            if self._model.system_instruction is not None
+            else []
+        )

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/orchestrator/orchestrator.py

@@ -1,5 +1,6 @@
-from graphrag_sdk.models import GenerativeModel
 from graphrag_sdk.agents import Agent
+from graphrag_sdk.models import GenerativeModel
+from graphrag_sdk.models.model import OutputMethod
 from .orchestrator_runner import OrchestratorRunner
 from graphrag_sdk.fixtures.prompts import (
     ORCHESTRATOR_SYSTEM,
@@ -16,13 +17,12 @@ logger = logging.getLogger(__name__)
 
 class Orchestrator:
 
-    _agents = []
-    _chat = None
-
     def __init__(self, model: GenerativeModel, backstory: str = ""):
         self._model = model
         self._backstory = backstory
-
+        self._agents = []
+        self._chat = None
+        
     def _get_chat(self):
         if self._chat is None:
             self._chat = self._model.with_system_instruction(

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/steps/graph_query_step.py

@@ -1,21 +1,15 @@
+import logging
+from falkordb import Graph
 from graphrag_sdk.steps.Step import Step
 from graphrag_sdk.ontology import Ontology
 from graphrag_sdk.models import (
     GenerativeModelChatSession,
 )
-from graphrag_sdk.fixtures.prompts import (
-    CYPHER_GEN_SYSTEM,
-    CYPHER_GEN_PROMPT,
-    CYPHER_GEN_PROMPT_WITH_ERROR,
-    CYPHER_GEN_PROMPT_WITH_HISTORY,
-)
-import logging
 from graphrag_sdk.helpers import (
     extract_cypher,
     validate_cypher,
     stringify_falkordb_response,
 )
-from falkordb import Graph
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -33,28 +27,26 @@ class GraphQueryGenerationStep(Step):
         chat_session: GenerativeModelChatSession,
         config: dict = None,
         last_answer: str = None,
+        cypher_prompt: str = None,
+        cypher_prompt_with_history: str = None,
     ) -> None:
         self.ontology = ontology
         self.config = config or {}
         self.graph = graph
         self.chat_session = chat_session
         self.last_answer = last_answer
+        self.cypher_prompt = cypher_prompt
+        self.cypher_prompt_with_history = cypher_prompt_with_history
 
-    def run(self, question: str, retries: int = 5):
-        error = False
-
+    def run(self, question: str, retries: int = 10):
         cypher = ""
-        while error is not None and retries > 0:
+        for i in range(retries):
             try:
                 cypher_prompt = (
-                    (CYPHER_GEN_PROMPT.format(question=question) 
+                    (self.cypher_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
-                    )
-                )
+                    else self.cypher_prompt_with_history.format(question=question, last_answer=self.last_answer))
+                )   
                 logger.debug(f"Cypher Prompt: {cypher_prompt}")
                 cypher_statement_response = self.chat_session.send_message(
                     cypher_prompt,
@@ -67,7 +59,6 @@ class GraphQueryGenerationStep(Step):
                     return (None, None)
 
                 validation_errors = validate_cypher(cypher, self.ontology)
-                # print(f"Is valid: {is_valid}")
                 if validation_errors is not None:
                     raise Exception("\n".join(validation_errors))
 
@@ -82,6 +73,6 @@ class GraphQueryGenerationStep(Step):
             except Exception as e:
                 logger.debug(f"Error: {e}")
                 error = e
-                retries -= 1
+                self.chat_session.delete_last_message()
 
         raise Exception("Failed to generate Cypher query: " + str(error))

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/graphrag_sdk/steps/qa_step.py

@@ -1,9 +1,7 @@
+import logging
 from graphrag_sdk.steps.Step import Step
 from graphrag_sdk.models import GenerativeModelChatSession
 
-from graphrag_sdk.fixtures.prompts import GRAPH_QA_SYSTEM, GRAPH_QA_PROMPT
-import logging
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
@@ -17,17 +15,18 @@ class QAStep(Step):
         self,
         chat_session: GenerativeModelChatSession,
         config: dict = None,
+        qa_prompt: str = None,
     ) -> None:
         self.config = config or {}
         self.chat_session = chat_session
+        self.qa_prompt = qa_prompt
 
     def run(self, question: str, cypher: str, context: str):
-
-        qa_prompt = GRAPH_QA_PROMPT.format(
+        qa_prompt = self.qa_prompt.format(
             context=context, cypher=cypher, question=question
         )
 
-        # logger.debug(f"QA Prompt: {qa_prompt}")
+        logger.debug(f"QA Prompt: {qa_prompt}")
         qa_response = self.chat_session.send_message(qa_prompt)
 
         return qa_response.text

{graphrag_sdk-0.3.0 → graphrag_sdk-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphrag_sdk"
-version = "0.3.0"
+version = "0.3.2"
 license = "MIT"
 description = ""
 authors = ["FalkorDB <info@falkordb.com>"]
@@ -16,6 +16,7 @@ typing-extensions = "^4.12.1"
 bs4 = "^0.0.2"
 pypdf = "^4.2.0"
 backoff = "^2.2.1"
+requests = "^2.32.3"
 python-abc = "^0.2.0"
 ratelimit = "^2.2.1"
 python-dotenv = "^1.0.1"