graphrag_sdk 0.2.0__tar.gz → 0.2.2__tar.gz

graphrag_sdk-0.2.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 FalkorDB
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

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

@@ -1,19 +1,23 @@
 Metadata-Version: 2.1
 Name: graphrag_sdk
-Version: 0.2.0
+Version: 0.2.2
 Summary: 
 Home-page: http://RAG-SDK.readthedocs.io
 License: MIT
 Keywords: RAG,Graph-RAG,KnowledgeGraph
 Author: FalkorDB
 Author-email: info@falkordb.com
-Requires-Python: >=3.11.4,<4.0.0
+Requires-Python: >=3.9.0,<4.0.0
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: backoff (>=2.2.1,<3.0.0)
 Requires-Dist: bs4 (>=0.0.2,<0.0.3)
-Requires-Dist: falkordb (>=1.0.4,<2.0.0)
+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: ipykernel (>=6.29.5,<7.0.0)
@@ -25,14 +29,17 @@ Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: ratelimit (>=2.2.1,<3.0.0)
 Requires-Dist: typing-extensions (>=4.12.1,<5.0.0)
 Requires-Dist: vertexai (>=1.49.0,<2.0.0)
-Project-URL: Repository, https://github.com/FalkorDB/falkordb-gemini-kg
+Project-URL: Repository, https://github.com/FalkorDB/GraphRAG-SDK
 Description-Content-Type: text/markdown
 
+[![Dockerhub](https://img.shields.io/docker/pulls/falkordb/falkordb?label=Docker)](https://hub.docker.com/r/falkordb/falkordb/)
+[![pypi](https://badge.fury.io/py/graphrag_sdk.svg)](https://pypi.org/project/graphrag_sdk/)
+[![Discord](https://img.shields.io/discord/1146782921294884966?style=flat-square)](https://discord.gg/6M4QwDXn2w)
+[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-2.1-4baaaa.svg)](code_of_conduct.md) 
+
 # GraphRAG-SDK
 
 [![Try Free](https://img.shields.io/badge/Try%20Free-FalkorDB%20Cloud-FF8101?labelColor=FDE900&style=for-the-badge&link=https://app.falkordb.cloud)](https://app.falkordb.cloud)
-[![Dockerhub](https://img.shields.io/docker/pulls/falkordb/falkordb?label=Docker)](https://hub.docker.com/r/falkordb/falkordb/)
-[![Discord](https://img.shields.io/discord/1146782921294884966?style=flat-square)](https://discord.gg/6M4QwDXn2w)
 
 GraphRAG-SDK is a comprehensive solution for building Graph Retrieval-Augmented Generation (GraphRAG) applications, leveraging [FalkorDB](https://www.falkordb.com/) for optimal performance.
 
@@ -91,10 +98,10 @@ The following example demonstrates the basic usage of this SDK to create a Graph
 ```python
 from dotenv import load_dotenv
 
-from graphrag_sdk.classes.source import URL
+from graphrag_sdk.source import URL
 from graphrag_sdk import KnowledgeGraph, Ontology
 from graphrag_sdk.models.openai import OpenAiGenerativeModel
-from graphrag_sdk.classes.model_config import KnowledgeGraphModelConfig
+from graphrag_sdk.model_config import KnowledgeGraphModelConfig
 load_dotenv()
 
 # Import Data
@@ -150,7 +157,7 @@ The SDK supports the following file formats:
 
 ```python
 import os
-from graphrag_sdk.classes.source import Source
+from graphrag_sdk.source import Source
 
 src_files = "data_folder"
 sources = []
@@ -224,12 +231,9 @@ You can update the KG at any time by processing more sources with the `process_s
 
 ### Graph RAG
 
-At this point, you have a Knowledge Graph that can be queried using this SDK. You can use the `ask` method for single questions or `chat_session` for conversations.
+At this point, you have a Knowledge Graph that can be queried using this SDK. Use the method `chat_session` for start a conversation.
 
 ```python
-# Single question.
-response = kg.ask("What were the last five fights? When were they? How many rounds did they have?")
-print(response)
 
 # Conversation.
 chat = kg.chat_session()

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

@@ -1,8 +1,11 @@
+[![Dockerhub](https://img.shields.io/docker/pulls/falkordb/falkordb?label=Docker)](https://hub.docker.com/r/falkordb/falkordb/)
+[![pypi](https://badge.fury.io/py/graphrag_sdk.svg)](https://pypi.org/project/graphrag_sdk/)
+[![Discord](https://img.shields.io/discord/1146782921294884966?style=flat-square)](https://discord.gg/6M4QwDXn2w)
+[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-2.1-4baaaa.svg)](code_of_conduct.md) 
+
 # GraphRAG-SDK
 
 [![Try Free](https://img.shields.io/badge/Try%20Free-FalkorDB%20Cloud-FF8101?labelColor=FDE900&style=for-the-badge&link=https://app.falkordb.cloud)](https://app.falkordb.cloud)
-[![Dockerhub](https://img.shields.io/docker/pulls/falkordb/falkordb?label=Docker)](https://hub.docker.com/r/falkordb/falkordb/)
-[![Discord](https://img.shields.io/discord/1146782921294884966?style=flat-square)](https://discord.gg/6M4QwDXn2w)
 
 GraphRAG-SDK is a comprehensive solution for building Graph Retrieval-Augmented Generation (GraphRAG) applications, leveraging [FalkorDB](https://www.falkordb.com/) for optimal performance.
 
@@ -61,10 +64,10 @@ The following example demonstrates the basic usage of this SDK to create a Graph
 ```python
 from dotenv import load_dotenv
 
-from graphrag_sdk.classes.source import URL
+from graphrag_sdk.source import URL
 from graphrag_sdk import KnowledgeGraph, Ontology
 from graphrag_sdk.models.openai import OpenAiGenerativeModel
-from graphrag_sdk.classes.model_config import KnowledgeGraphModelConfig
+from graphrag_sdk.model_config import KnowledgeGraphModelConfig
 load_dotenv()
 
 # Import Data
@@ -120,7 +123,7 @@ The SDK supports the following file formats:
 
 ```python
 import os
-from graphrag_sdk.classes.source import Source
+from graphrag_sdk.source import Source
 
 src_files = "data_folder"
 sources = []
@@ -194,12 +197,9 @@ You can update the KG at any time by processing more sources with the `process_s
 
 ### Graph RAG
 
-At this point, you have a Knowledge Graph that can be queried using this SDK. You can use the `ask` method for single questions or `chat_session` for conversations.
+At this point, you have a Knowledge Graph that can be queried using this SDK. Use the method `chat_session` for start a conversation.
 
 ```python
-# Single question.
-response = kg.ask("What were the last five fights? When were they? How many rounds did they have?")
-print(response)
 
 # Conversation.
 chat = kg.chat_session()

{graphrag_sdk-0.2.0 → graphrag_sdk-0.2.2}/graphrag_sdk/agents/agent.py

@@ -144,18 +144,15 @@ class Agent(ABC):
         pass
 
     @abstractmethod
-    def run(
-        self, params: dict, session: GenerativeModelChatSession
-    ) -> tuple[str, GenerativeModelChatSession]:
+    def run(self, params: dict) -> str:
         """
         Run the agent with the given parameters and chat session.
 
         Args:
             params (Dict): The parameters for the agent.
-            session (GenerativeModelChatSession): The chat session.
 
         Returns:
-            Tuple[str, GenerativeModelChatSession]: The agent's response and the updated chat session.
+            str: The agent's response.
         """
         pass
 

{graphrag_sdk-0.2.0 → graphrag_sdk-0.2.2}/graphrag_sdk/agents/kg_agent.py

@@ -43,6 +43,7 @@ class KGAgent(Agent):
         self.agent_id = agent_id
         self.introduction = introduction
         self.kg = kg
+        self.chat_session = self._kg.chat_session()
 
     @property
     def agent_id(self) -> str:
@@ -123,22 +124,19 @@ class KGAgent(Agent):
         """
         self._kg = value
 
-    def run(
-        self, params: dict, session: GenerativeModelChatSession | None = None
-    ) -> tuple[str, GenerativeModelChatSession]:
+    def run(self, params: dict) -> str:
         """
         Ask the agent a question.
 
         Args:
             params (dict): The parameters for the agent.
-            session (GenerativeModelChatSession | None): The chat session to use for the agent. Defaults to None.
 
         Returns:
-            tuple[str, GenerativeModelChatSession]: The agent's response and the updated chat session.
+            str: The agent's response.
 
         """
-        (output, chat_session) = self._kg.ask(params["prompt"], session)
-        return (output, chat_session)
+        output = self.chat_session.send_message(params["prompt"])
+        return output
 
     def __repr__(self):
         """

{graphrag_sdk-0.2.0 → graphrag_sdk-0.2.2}/graphrag_sdk/chat_session.py

@@ -52,6 +52,7 @@ class ChatSession:
         self.qa_chat_session = model_config.qa.with_system_instruction(
             GRAPH_QA_SYSTEM
         ).start_chat()
+        self.last_answer = None
 
     def send_message(self, message: str):
         """
@@ -67,6 +68,7 @@ class ChatSession:
             graph=self.graph,
             chat_session=self.cypher_chat_session,
             ontology=self.ontology,
+            last_answer=self.last_answer,
         )
 
         (context, cypher) = cypher_step.run(message)
@@ -79,5 +81,5 @@ class ChatSession:
         )
 
         answer = qa_step.run(message, cypher, context)
-
-        return answer
+        self.last_answer = answer
+        return answer

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

@@ -188,7 +188,7 @@ Raw text:
 """
 
 BOUNDARIES_PREFIX = """
-Use the following instructions as boundaries for the ontology extraction process.
+Use the following instructions as boundaries for the ontology extraction process:
 {user_boundaries}
 """
 
@@ -330,7 +330,7 @@ Schema:
 }
 ```
 
-For example:
+Output example:
 ```{"entities":[{"label":"Person","attributes":{"name":"John Doe","age":30}},{"label":"Movie","attributes":{"title":"Inception","releaseYear":2010}}],"relations":[{"label":"ACTED_IN","source":{"label":"Person","attributes":{"name":"JohnDoe"}},"target":{"label":"Movie","attributes":{"title":"Inception"}},"attributes":{"role":"Cobb"}}]}```
 
 Ontology:
@@ -338,18 +338,47 @@ Ontology:
 """
 
 EXTRACT_DATA_PROMPT = """
-Extract all possible entities and relations from the text below.
-Use the ontology provided in the system prompt.
-Assign textual IDs whenever required.
-Use double quotes for string values.
-It's imperative that string values are properly escaped.
-All formats should be consistent, for example, dates should be in the format "YYYY-MM-DD".
-If needed, add the correct spacing for text fields, where the text is not properly formatted.
-
-User instructions:
+You are tasked with extracting entities and relations from the text below, using the ontology provided.
+
+**Output Format:**
+
+- Provide the extracted data as a JSON object with two keys: `"entities"` and `"relations"`.
+
+- **Entities**: Represent entities and concepts. Each entity should have a `"label"` and `"attributes"` field.
+
+- **Relations**: Represent relations between entities or concepts. Each relation should have a `"label"`, `"source"`, `"target"`, and `"attributes"` field.
+
+**Guidelines:**
+- **Extract all entities and relations**: Capture all entities and relations mentioned in the text.
+
+- **Use Only the Provided Ontology**: Utilize only the types of entities, relations, and attributes defined in the ontology.
+
+- **Assign IDs Where Required**: Assign textual IDs to entities and relations as specified.
+
+- **Avoid Duplicates**: Ensure each entity and relation is unique; do not include duplicates.
+
+- **Formatting**:
+  - Do not include any introduction or explanation in the response, only the JSON.
+  
+  - Use double quotes for all string values.
+
+  - Properly escape any special characters.
+
+  - Dates should be in the format `"YYYY-MM-DD"`.
+
+  - Correct any spacing or formatting issues in text fields as necessary.
+
+- **Precision**: Be concise and precise in your extraction.
+
+- **Token Limit**: Ensure your response does not exceed **{max_tokens} tokens**.
+
+**User Instructions**:
 {instructions}
 
-Raw Text:
+**Ontology**:
+{ontology}
+
+**Raw Text**:
 {text}
 """
 
@@ -363,6 +392,12 @@ JSON:
 {json}
 """
 
+# This constant is used as a follow-up prompt when the initial data extraction is incomplete or contains duplicates.
+# It instructs the model to complete the answer and ensure uniqueness of entities and relations.
+COMPLETE_DATA_EXTRACTION = """
+Please complete your answer. Ensure that each entity and relations is unique. Do not include duplicates. Please be precise.
+"""
+
 CYPHER_GEN_SYSTEM = """
 Task: Generate OpenCypher statement to query a graph database.
 
@@ -426,6 +461,19 @@ Do not include any apologies or other texts, except the generated OpenCypher sta
 Question: {question}
 """
 
+CYPHER_GEN_PROMPT_WITH_HISTORY = """
+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.
+
+First, determine if the last answers provided to the user over the entire conversation is relevant to the current question. If it is relevant, you may consider incorporating information from it into the query. If it is not relevant, ignore it and generate the query based solely on the question.
+
+If you cannot generate an 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".
+
+Last Answer: {last_answer}
+
+Question: {question}
+"""
 
 GRAPH_QA_SYSTEM = """
 You are an assistant that helps to form nice and human understandable answers.
@@ -540,7 +588,7 @@ Choose between the following steps to create the execution plan:
 
 
 ORCHESTRATOR_SUMMARY_PROMPT = """
-Given the following execution log, generate the final answer to the user's question.
+Given the following execution log and the history of this chat, generate the final answer to the user's question.
 Be very polite and detailed in your response, always providing the reasoning behind the answer.
 
 User question:

{graphrag_sdk-0.2.0 → graphrag_sdk-0.2.2}/graphrag_sdk/kg.py

@@ -118,55 +118,6 @@ class KnowledgeGraph:
 
         step.run(instructions)
 
-    def ask(
-        self, question: str, qa_chat_session: GenerativeModelChatSession | None = None
-    ) -> tuple[str, GenerativeModelChatSession]:
-        """
-        Query the knowledge graph using natural language.
-        Optionally, you can provide a qa_chat_session to use for the query.
-
-        Parameters:
-            question (str): question to ask the knowledge graph
-            qa_chat_session (GenerativeModelChatSession|None): qa_chat_session to use for the query
-
-        Returns:
-            tuple[str, GenerativeModelChatSession]: answer, qa_chat_session
-
-         Example:
-            >>> (ans, qa_chat_session) = kg.ask("List a few movies in which that actor played in")
-            >>> print(ans)
-        """
-
-        cypher_chat_session = (
-            self._model_config.cypher_generation.with_system_instruction(
-                CYPHER_GEN_SYSTEM.replace("#ONTOLOGY", str(self.ontology.to_json())),
-            ).start_chat()
-        )
-        cypher_step = GraphQueryGenerationStep(
-            ontology=self.ontology,
-            chat_session=cypher_chat_session,
-            graph=self.graph,
-        )
-
-        (context, cypher) = cypher_step.run(question)
-
-        if not cypher or len(cypher) == 0:
-            return "I am sorry, I could not find the answer to your question"
-
-        qa_chat_session = (
-            qa_chat_session
-            or self._model_config.qa.with_system_instruction(
-                GRAPH_QA_SYSTEM
-            ).start_chat()
-        )
-        qa_step = QAStep(
-            chat_session=qa_chat_session,
-        )
-
-        answer = qa_step.run(question, cypher, context)
-
-        return (answer, qa_chat_session)
-
     def delete(self) -> None:
         """
         Deletes the knowledge graph and any other related resource

graphrag_sdk-0.2.2/graphrag_sdk/models/azure_openai.py

@@ -0,0 +1,229 @@
+import os
+from typing import Optional
+from openai import AzureOpenAI
+from .model import (
+    OutputMethod,
+    GenerativeModel,
+    GenerativeModelConfig,
+    GenerationResponse,
+    FinishReason,
+    GenerativeModelChatSession,
+)
+
+class AzureOpenAiGenerativeModel(GenerativeModel):
+    """
+    A generative model that interfaces with Azure's OpenAI API for chat completions.
+    """
+
+    client: AzureOpenAI = None
+
+    def __init__(
+        self,
+        model_name: str,
+        generation_config: Optional[GenerativeModelConfig] = None,
+        system_instruction: Optional[str] = None,
+    ):
+        """
+        Initialize the AzureOpenAiGenerativeModel with required parameters.
+
+        Args:
+            model_name (str): Name of the Azure OpenAI model.
+            generation_config (Optional[GenerativeModelConfig]): Configuration settings for generation.
+            system_instruction (Optional[str]): System-level instruction for the model.
+        """
+        self.model_name = model_name
+        self.generation_config = generation_config or GenerativeModelConfig()
+        self.system_instruction = system_instruction
+        
+        # Credentials
+        self.api_key = os.getenv("AZURE_OPENAI_API_KEY")
+        self.azure_endpoint = os.getenv("AZURE_ENDPOINT")
+        self.api_version = os.getenv("AZURE_API_VERSION")
+        
+        if not self.api_key or not self.azure_endpoint or not self.api_version:
+            raise ValueError(
+                "Missing credentials in the environment: AZURE_OPENAI_API_KEY, AZURE_ENDPOINT, or AZURE_API_VERSION."
+            )
+
+
+    def _connect_to_model(self) -> None:
+        """
+        Establish a connection to the Azure OpenAI model by initializing the AzureOpenAI client.
+        """
+        self.client = AzureOpenAI(azure_endpoint=self.azure_endpoint,
+                                api_version=self.api_version,
+                                api_key=self.api_key,
+                                )
+
+    def with_system_instruction(self, system_instruction: str) -> "GenerativeModel":
+        """
+        Set or update the system instruction and connect to the Azure OpenAI model.
+
+        Args:
+            system_instruction (str): System instructions for the model.
+        
+        Returns:
+            GenerativeModel: The updated model instance.
+        """
+        self.system_instruction = system_instruction
+        self._connect_to_model()
+
+        return self
+
+    def start_chat(self, args: Optional[dict] = None) -> GenerativeModelChatSession:
+        """
+        Start a new chat session.
+
+        Args:
+            args (Optional[dict]): Additional arguments for the chat session.
+
+        Returns:
+            GenerativeModelChatSession: A new instance of the chat session.
+        """
+        return AzureOpenAiChatSession(self, args)
+
+    def ask(self, message: str) -> GenerationResponse:
+        """
+        Send a message to the model and receive a response.
+
+        Args:
+            message (str): The user's message input.
+
+        Returns:
+            GenerationResponse: The model's generated response.
+        """
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[
+                {"role": "system", "content": self.system_instruction},
+                {"role": "user", "content": message[:14385]},
+            ],
+            max_tokens=self.generation_config.max_output_tokens,
+            temperature=self.generation_config.temperature,
+            top_p=self.generation_config.top_p,
+            top_k=self.generation_config.top_k,
+            stop=self.generation_config.stop_sequences,
+        )
+        return self._parse_generate_content_response(response)
+
+    def _parse_generate_content_response(self, response: any) -> GenerationResponse:
+        """
+        Parse the model's response and extract content for the user.
+
+        Args:
+            response (any): The raw response from the model.
+
+        Returns:
+            GenerationResponse: Parsed response containing the generated text and finish reason.
+        """
+        return GenerationResponse(
+            text=response.choices[0].message.content,
+            finish_reason=(
+                FinishReason.STOP
+                if response.choices[0].finish_reason == "stop"
+                else (
+                    FinishReason.MAX_TOKENS
+                    if response.choices[0].finish_reason == "length"
+                    else FinishReason.OTHER
+                )
+            ),
+        )
+
+    def to_json(self) -> dict:
+        """
+        Serialize the model's configuration and state to JSON format.
+
+        Returns:
+            dict: The serialized JSON data.
+        """
+        return {
+            "model_name": self.model_name,
+            "generation_config": self.generation_config.to_json(),
+            "system_instruction": self.system_instruction,
+        }
+
+    @staticmethod
+    def from_json(json: dict) -> "GenerativeModel":
+        """
+        Deserialize a JSON object to create an instance of AzureOpenAiGenerativeModel.
+
+        Args:
+            json (dict): The serialized JSON data.
+
+        Returns:
+            GenerativeModel: A new instance of the model.
+        """
+        return AzureOpenAiGenerativeModel(
+            json["model_name"],
+            generation_config=GenerativeModelConfig.from_json(
+                json["generation_config"]
+            ),
+            system_instruction=json["system_instruction"],
+        )
+
+
+class AzureOpenAiChatSession(GenerativeModelChatSession):
+    """
+    A chat session for interacting with the Azure OpenAI model, maintaining conversation history.
+    """
+
+
+    _history = []
+
+    def __init__(self, model: AzureOpenAiGenerativeModel, args: Optional[dict] = None):
+        """
+        Initialize the chat session and set up the conversation history.
+
+        Args:
+            model (AzureOpenAiGenerativeModel): The model instance for the session.
+            args (Optional[dict]): Additional arguments for customization.
+        """
+        self._model = model
+        self._args = args
+        self._history = (
+            [{"role": "system", "content": self._model.system_instruction}]
+            if self._model.system_instruction is not None
+            else []
+        )
+
+    def send_message(self, message: str, output_method: OutputMethod = OutputMethod.DEFAULT) -> GenerationResponse:
+        """
+        Send a message in the chat session and receive the model's response.
+
+        Args:
+            message (str): The message to send.
+            output_method (OutputMethod): Format for the model's output.
+
+        Returns:
+            GenerationResponse: The generated response.
+        """
+        generation_config = self._get_generation_config(output_method)
+        prompt = []
+        prompt.extend(self._history)
+        prompt.append({"role": "user", "content": message[:14385]})
+        response = self._model.client.chat.completions.create(
+            model=self._model.model_name,
+            messages=prompt,
+            **generation_config
+        )
+        content = self._model._parse_generate_content_response(response)
+        self._history.append({"role": "user", "content": message})
+        self._history.append({"role": "assistant", "content": content.text})
+        return content
+    
+    def _get_generation_config(self, output_method: OutputMethod):
+        """
+        Adjust the generation configuration based on the output method.
+
+        Args:
+            output_method (OutputMethod): The desired output method (e.g., default or JSON).
+
+        Returns:
+            dict: The configuration settings for generation.
+        """
+        config = self._model.generation_config.to_json()
+        if output_method == OutputMethod.JSON:
+            config['temperature'] = 0
+            config['response_format'] = { "type": "json_object" }
+        
+        return config

{graphrag_sdk-0.2.0 → graphrag_sdk-0.2.2}/graphrag_sdk/models/gemini.py

@@ -1,5 +1,6 @@
 import os
 from .model import (
+    OutputMethod,
     GenerativeModel,
     GenerativeModelConfig,
     GenerationResponse,
@@ -61,10 +62,6 @@ class GeminiGenerativeModel(GenerativeModel):
     def start_chat(self, args: dict | None = None) -> GenerativeModelChatSession:
         return GeminiChatSession(self, args)
 
-    def ask(self, message: str) -> GenerationResponse:
-        response = self._model.generate_content(message)
-        return self.parse_generate_content_response(response)
-
     def parse_generate_content_response(
         self, response: types.generation_types.GenerateContentResponse
     ) -> GenerationResponse:
@@ -108,6 +105,15 @@ class GeminiChatSession(GenerativeModelChatSession):
             history=args.get("history", []) if args is not None else [],
         )
 
-    def send_message(self, message: str) -> GenerationResponse:
-        response = self._chat_session.send_message(message)
+    def send_message(self, message: str, output_method: OutputMethod = OutputMethod.DEFAULT) -> GenerationResponse:
+        generation_config = self._get_generation_config(output_method)
+        response = self._chat_session.send_message(message, generation_config=generation_config)
         return self._model.parse_generate_content_response(response)
+    
+    def _get_generation_config(self, output_method: OutputMethod):
+        if output_method == OutputMethod.JSON:
+            return {
+                "response_mime_type": "application/json",
+                "temperature": 0
+            }
+        return self._model._generation_config