switch from | to Union

Author: swilly22

graphrag_sdk/attribute.py

@@ -1,6 +1,7 @@
 import re
 import json
 import logging
+from typing import Union
 from graphrag_sdk.fixtures.regex import *
 
 logger = logging.getLogger(__name__)
@@ -89,12 +90,12 @@ def __init__(
         self.required = required
 
     @staticmethod
-    def from_json(txt: str | dict):
+    def from_json(txt: Union[str, dict]):
         """
         Creates an Attribute object from a JSON string or dictionary.
 
         Args:
-            txt (str | dict): The JSON string or dictionary representing the Attribute.
+            txt (Union[str, dict]): The JSON string or dictionary representing the Attribute.
 
         Returns:
             Attribute: The created Attribute object.

graphrag_sdk/entity.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from typing import Union
 from .attribute import Attribute
 from falkordb import Node as GraphNode
 import re
@@ -20,7 +21,7 @@ class Entity:
 
     Methods:
         from_graph(entity: GraphNode) -> Entity: Creates an Entity object from a GraphNode object.
-        from_json(txt: dict | str) -> Entity: Creates an Entity object from a JSON string or dictionary.
+        from_json(txt: Union[dict, str]) -> Entity: Creates an Entity object from a JSON string or dictionary.
         to_json() -> dict: Converts the Entity object to a JSON dictionary.
         merge(entity2: Entity) -> Entity: Overwrites attributes of self with attributes of entity2.
         get_unique_attributes() -> list[Attribute]: Returns a list of unique attributes of the entity.
@@ -64,12 +65,12 @@ def from_graph(entity: GraphNode):
         )
 
     @staticmethod
-    def from_json(txt: dict | str):
+    def from_json(txt: Union[dict, str]):
         """
         Create an Entity object from a JSON representation.
 
         Args:
-            txt (dict | str): The JSON representation of the Entity. It can be either a dictionary or a string.
+            txt Union[dict, str]: The JSON representation of the Entity. It can be either a dictionary or a string.
 
         Returns:
             Entity: The Entity object created from the JSON representation.

graphrag_sdk/helpers.py

@@ -1,12 +1,13 @@
 import re
 import logging
 import graphrag_sdk
+from typing import Union
 from fix_busted_json import repair_json
 
 logger = logging.getLogger(__name__)
 
 
-def extract_json(text: str | dict, skip_repair=False) -> str:
+def extract_json(text: Union[str, dict], skip_repair=False) -> str:
     if not isinstance(text, str):
         text = str(text)
     regex = r"(?:```)?(?:json)?([^`]*)(?:\\n)?(?:```)?"
@@ -79,7 +80,7 @@ def extract_cypher(text: str):
 
 def validate_cypher(
     cypher: str, ontology: graphrag_sdk.Ontology
-) -> list[str] | None:
+) -> Union[list[str], None]:
     try:
         if not cypher or len(cypher) == 0:
             return ["Cypher statement is empty"]
@@ -215,4 +216,4 @@ def validate_cypher_relation_directions(
         except Exception:
             continue
 
-    return errors
+    return errors

graphrag_sdk/kg.py

@@ -1,7 +1,7 @@
 import logging
 import warnings
-from typing import Optional
 from falkordb import FalkorDB
+from typing import Optional, Union
 from graphrag_sdk.ontology import Ontology
 from graphrag_sdk.source import AbstractSource
 from graphrag_sdk.chat_session import ChatSession
@@ -45,14 +45,14 @@ def __init__(
             model (GenerativeModel): The Google GenerativeModel to use.
             host (str): FalkorDB hostname.
             port (int): FalkorDB port number.
-            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.
+            username (Union[str, None]): FalkorDB username.
+            password (Union[str, None]): FalkorDB password.
+            ontology (Union[Ontology, None]): Ontology to use.
+            cypher_system_instruction (Union[str, None]): Cypher system instruction. Make sure you have {ontology} in the instruction.
+            qa_system_instruction (Union[str, None]): QA system instruction.
+            cypher_gen_prompt (Union[str, None]): Cypher generation prompt. Make sure you have {question} in the prompt.
+            qa_prompt (Union[str, None]): QA prompt. Make sure you have {question}, {context} and {cypher} in the prompt.
+            cypher_gen_prompt_history (Union[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 == "":
@@ -165,7 +165,7 @@ def process_sources(
 
 
     def _create_graph_with_sources(
-        self, sources: list[AbstractSource] | None = None, instructions: str = None, hide_progress: bool = False
+        self, sources: Union[list[AbstractSource], None] = None, instructions: str = None, hide_progress: bool = False
 ) -> None:
 
         step = ExtractDataStep(

graphrag_sdk/models/gemini.py

@@ -1,4 +1,6 @@
 import os
+from typing import Union
+
 from .model import (
     OutputMethod,
     GenerativeModel,
@@ -23,8 +25,8 @@ class GeminiGenerativeModel(GenerativeModel):
     def __init__(
         self,
         model_name: str,
-        generation_config: GoogleGenerationConfig | None = None,
-        system_instruction: str | None = None,
+        generation_config: Union[GoogleGenerationConfig, None] = None,
+        system_instruction: Union[str, None] = None,
     ):
         self._model_name = model_name
         self._generation_config = generation_config
@@ -59,7 +61,7 @@ def with_system_instruction(self, system_instruction: str) -> "GenerativeModel":
 
         return self
 
-    def start_chat(self, args: dict | None = None) -> GenerativeModelChatSession:
+    def start_chat(self, args: Union[dict, None] = None) -> GenerativeModelChatSession:
         return GeminiChatSession(self, args)
 
     def parse_generate_content_response(
@@ -99,7 +101,7 @@ def from_json(json: dict) -> "GenerativeModel":
 
 class GeminiChatSession(GenerativeModelChatSession):
 
-    def __init__(self, model: GeminiGenerativeModel, args: dict | None = None):
+    def __init__(self, model: GeminiGenerativeModel, args: Union[dict, None] = None):
         self._model = model
         self._chat_session = self._model._model.start_chat(
             history=args.get("history", []) if args is not None else [],

graphrag_sdk/models/model.py

@@ -1,5 +1,6 @@
-from abc import ABC, abstractmethod
 from enum import Enum
+from typing import Union
+from abc import ABC, abstractmethod
 
 
 class FinishReason:
@@ -30,12 +31,12 @@ class GenerativeModelConfig:
 
     def __init__(
         self,
-        temperature: float | None = None,
-        top_p: float | None = None,
-        top_k: int | None = None,
-        max_output_tokens: int | None = None,
-        stop_sequences: list[str] | None = None,
-        response_format: dict | None = None,
+        temperature: Union[float, None] = None,
+        top_p: Union[float, None] = None,
+        top_k: Union[int, None] = None,
+        max_output_tokens: Union[int, None] = None,
+        stop_sequences: Union[list[str], None] = None,
+        response_format: Union[dict, None] = None,
     ):
         self.temperature = temperature
         self.top_p = top_p
@@ -103,7 +104,7 @@ def with_system_instruction(self, system_instruction: str) -> "GenerativeModel":
         pass
 
     @abstractmethod
-    def start_chat(self, args: dict | None) -> GenerativeModelChatSession:
+    def start_chat(self, args: Union[dict, None]) -> GenerativeModelChatSession:
         pass
 
     @staticmethod

graphrag_sdk/models/openai.py

@@ -6,6 +6,7 @@
     FinishReason,
     GenerativeModelChatSession,
 )
+from typing import Union
 from openai import OpenAI
 
 
@@ -16,8 +17,8 @@ class OpenAiGenerativeModel(GenerativeModel):
     def __init__(
         self,
         model_name: str,
-        generation_config: GenerativeModelConfig | None = None,
-        system_instruction: str | None = None,
+        generation_config: Union[GenerativeModelConfig, None] = None,
+        system_instruction: Union[str, None] = None,
     ):
         self.model_name = model_name
         self.generation_config = generation_config or GenerativeModelConfig()
@@ -36,7 +37,7 @@ def with_system_instruction(self, system_instruction: str) -> "GenerativeModel":
 
         return self
 
-    def start_chat(self, args: dict | None = None) -> GenerativeModelChatSession:
+    def start_chat(self, args: Union[dict, None] = None) -> GenerativeModelChatSession:
         return OpenAiChatSession(self, args)
 
     def parse_generate_content_response(self, response: any) -> GenerationResponse:
@@ -75,7 +76,7 @@ class OpenAiChatSession(GenerativeModelChatSession):
 
     _history = []
 
-    def __init__(self, model: OpenAiGenerativeModel, args: dict | None = None):
+    def __init__(self, model: OpenAiGenerativeModel, args: Union[dict, None] = None):
         self._model = model
         self._args = args
         self._history = (

graphrag_sdk/ontology.py

@@ -3,7 +3,7 @@
 import graphrag_sdk
 from .entity import Entity
 from falkordb import Graph
-from typing import Optional
+from typing import Optional, Union
 from .relation import Relation
 from graphrag_sdk.source import AbstractSource
 from graphrag_sdk.models import GenerativeModel
@@ -81,12 +81,12 @@ def from_sources(
         return step.run(boundaries=boundaries)
 
     @staticmethod
-    def from_json(txt: dict | str):
+    def from_json(txt: Union[dict, str]):
         """
         Creates an Ontology object from a JSON representation.
 
         Args:
-            txt (dict | str): The JSON representation of the ontology. It can be either a dictionary or a string.
+            txt (Union[dict, str]): The JSON representation of the ontology. It can be either a dictionary or a string.
 
         Returns:
             The Ontology object created from the JSON representation.
@@ -401,4 +401,4 @@ def save_to_graph(self, graph: Graph):
         for relation in self.relations:
             query = relation.to_graph_query()
             logger.debug(f"Query: {query}")
-            graph.query(query)
+            graph.query(query)

graphrag_sdk/orchestrator/execution_plan.py

@@ -1,4 +1,5 @@
 from json import loads
+from typing import Union
 from graphrag_sdk.orchestrator.step import PlanStep, StepBlockType
 
 
@@ -10,7 +11,7 @@ def __init__(self, steps: list[PlanStep]):
         self.steps = steps
 
     @staticmethod
-    def from_json(json: str | dict) -> "ExecutionPlan":
+    def from_json(json: Union[str, dict]) -> "ExecutionPlan":
         if isinstance(json, str):
             json = loads(json)
         return ExecutionPlan([PlanStep.from_json(step) for step in json])

graphrag_sdk/orchestrator/orchestrator_decision.py

@@ -1,5 +1,6 @@
-from graphrag_sdk.orchestrator.step import PlanStep
 from json import loads
+from typing import Union
+from graphrag_sdk.orchestrator.step import PlanStep
 
 
 class OrchestratorDecisionCode:
@@ -21,7 +22,7 @@ def from_str(code: str) -> str:
 
 class OrchestratorDecision:
     def __init__(
-        self, code: OrchestratorDecisionCode, new_step: PlanStep | None = None
+        self, code: OrchestratorDecisionCode, new_step: Union[PlanStep, None] = None
     ):
         self.code = code
         self.new_step = new_step
@@ -33,7 +34,7 @@ def to_json(self) -> dict:
         }
 
     @staticmethod
-    def from_json(json: dict | str) -> "OrchestratorDecision":
+    def from_json(json: Union[dict, str]) -> "OrchestratorDecision":
         json = json if isinstance(json, dict) else loads(json)
         return OrchestratorDecision(
             OrchestratorDecisionCode.from_str(json["code"]),