graphrag_sdk 0.2.2__tar.gz → 0.3.1__tar.gz

{graphrag_sdk-0.2.2 → graphrag_sdk-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphrag_sdk
-Version: 0.2.2
+Version: 0.3.1
 Summary: 
 Home-page: http://RAG-SDK.readthedocs.io
 License: MIT
@@ -15,20 +15,26 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Provides-Extra: all
+Provides-Extra: google-generativeai
+Provides-Extra: ollama
+Provides-Extra: openai
+Provides-Extra: vertexai
 Requires-Dist: backoff (>=2.2.1,<3.0.0)
 Requires-Dist: bs4 (>=0.0.2,<0.0.3)
 Requires-Dist: falkordb (>=1.0.9,<2.0.0)
 Requires-Dist: fix-busted-json (>=0.0.18,<0.0.19)
-Requires-Dist: google-generativeai (>=0.8.1,<0.9.0)
+Requires-Dist: google-generativeai (>=0.8.1,<0.9.0) ; extra == "all" or extra == "google-generativeai"
 Requires-Dist: ipykernel (>=6.29.5,<7.0.0)
-Requires-Dist: ollama (>=0.2.1,<0.3.0)
-Requires-Dist: openai (>=1.35.9,<2.0.0)
+Requires-Dist: ollama (>=0.2.1,<0.3.0) ; extra == "all" or extra == "ollama"
+Requires-Dist: openai (>=1.35.9,<2.0.0) ; extra == "all" or extra == "openai"
 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)
+Requires-Dist: vertexai (>=1.49.0,<2.0.0) ; extra == "all" or extra == "vertexai"
 Project-URL: Repository, https://github.com/FalkorDB/GraphRAG-SDK
 Description-Content-Type: text/markdown
 
@@ -55,7 +61,13 @@ GraphRAG-SDK is a comprehensive solution for building Graph Retrieval-Augmented
 ### Install
 
 ```sh
-pip install graphrag_sdk
+pip install graphrag_sdk[all]
+```
+
+or (if only OpenAI is required)
+
+```sh
+pip install graphrag_sdk[openai]
 ```
 
 ### Prerequisites
@@ -75,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
 

{graphrag_sdk-0.2.2 → graphrag_sdk-0.3.1}/README.md

@@ -21,7 +21,13 @@ GraphRAG-SDK is a comprehensive solution for building Graph Retrieval-Augmented
 ### Install
 
 ```sh
-pip install graphrag_sdk
+pip install graphrag_sdk[all]
+```
+
+or (if only OpenAI is required)
+
+```sh
+pip install graphrag_sdk[openai]
 ```
 
 ### Prerequisites
@@ -41,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
 

{graphrag_sdk-0.2.2 → graphrag_sdk-0.3.1}/graphrag_sdk/fixtures/prompts.py

@@ -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:
+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.
 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}
 """

{graphrag_sdk-0.2.2 → graphrag_sdk-0.3.1}/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.2.2 → graphrag_sdk-0.3.1}/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.2.2 → graphrag_sdk-0.3.1}/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.2.2 → graphrag_sdk-0.3.1}/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.2.2 → graphrag_sdk-0.3.1}/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.2.2 → graphrag_sdk-0.3.1}/graphrag_sdk/steps/graph_query_step.py

@@ -40,20 +40,14 @@ class GraphQueryGenerationStep(Step):
         self.chat_session = chat_session
         self.last_answer = last_answer
 
-    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) 
                     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
-                    )
                 )
                 logger.debug(f"Cypher Prompt: {cypher_prompt}")
                 cypher_statement_response = self.chat_session.send_message(
@@ -67,7 +61,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 +75,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.2.2 → graphrag_sdk-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphrag_sdk"
-version = "0.2.2"
+version = "0.3.1"
 license = "MIT"
 description = ""
 authors = ["FalkorDB <info@falkordb.com>"]
@@ -15,16 +15,24 @@ falkordb = "^1.0.9"
 typing-extensions = "^4.12.1"
 bs4 = "^0.0.2"
 pypdf = "^4.2.0"
-vertexai = "^1.49.0"
 backoff = "^2.2.1"
+requests = "^2.32.3"
 python-abc = "^0.2.0"
 ratelimit = "^2.2.1"
 python-dotenv = "^1.0.1"
-openai = "^1.35.9"
 fix-busted-json = "^0.0.18"
-ollama = "^0.2.1"
 ipykernel = "^6.29.5"
-google-generativeai = "^0.8.1"
+vertexai = { version = "^1.49.0", optional = true } 
+ollama = { version = "^0.2.1", optional = true }
+openai = { version = "^1.35.9", optional = true }
+google-generativeai = { version = "^0.8.1", optional = true }
+
+[tool.poetry.extras]
+all = ["vertexai", "ollama", "openai", "google-generativeai"]
+ollama = ["ollama"]
+openai = ["openai"]
+vertexai = ["vertexai"]
+google-generativeai = ["google-generativeai"]
 
 [tool.poetry.group.test.dependencies]
 pytest = "^8.2.1"