filter-to-session

Author: galshubeli

graphrag_sdk/attribute.py

@@ -57,15 +57,15 @@ class Attribute:
     """
 
     def __init__(
-        self, name: str, attr_type: AttributeType, unique: bool, required: bool = False
+        self, name: str, attr_type: AttributeType, unique: bool = False, required: bool = False
     ):
         """
         Initialize a new Attribute object.
 
         Args:
             name (str): The name of the attribute.
             attr_type (AttributeType): The type of the attribute.
-            unique (bool): Indicates whether the attribute should be unique.
+            unique (bool, optional): Indicates whether the attribute should be unique. Defaults to False.
             required (bool, optional): Indicates whether the attribute is required. Defaults to False.
         """
         self.name = re.sub(r"([^a-zA-Z0-9_])", "_", name)
@@ -130,29 +130,24 @@ def from_string(txt: str):
 
         return Attribute(name, AttributeType.from_string(attr_type), unique, required)
 
-    def to_json(self, include_all: bool = True):
+    def to_json(self):
         """
         Converts the attribute object to a JSON representation.
 
-        Args:
-            include_all (bool): Whether to include both "unique" and "required" fields in the output. Default is True.
-
         Returns:
             dict: A dictionary representing the attribute object in JSON format.
                 The dictionary contains the following keys:
                 - "name": The name of the attribute.
                 - "type": The type of the attribute.
-                Optionally includes:
-                - "unique": A boolean indicating whether the attribute is unique (if include_all is True).
-                - "required": A boolean indicating whether the attribute is required (if include_all is True).
+                - "unique": A boolean indicating whether the attribute is unique.
+                - "required": A boolean indicating whether the attribute is required.
         """
         json_data = {
             "name": self.name,
             "type": self.type,
+            "unique": self.unique,
+            "required": self.required,
         }
-        if include_all:
-            json_data["unique"] = self.unique
-            json_data["required"] = self.required
 
         return json_data
 

graphrag_sdk/chat_session.py

@@ -1,3 +1,4 @@
+import json
 from falkordb import Graph
 from graphrag_sdk.ontology import Ontology
 from graphrag_sdk.steps.qa_step import QAStep
@@ -45,8 +46,11 @@ def __init__(self, model_config: KnowledgeGraphModelConfig, ontology: Ontology,
         self.model_config = model_config
         self.graph = graph
         self.ontology = ontology
-        cypher_system_instruction = cypher_system_instruction.format(ontology=str(ontology.to_json(include_all=False)))
-
+        
+        # Filter the ontology to remove unique and required attributes that are not needed for Q&A. 
+        ontology_prompt = self.clean_ontology_for_prompt(ontology)
+                
+        cypher_system_instruction = cypher_system_instruction.format(ontology=ontology_prompt)
 
         self.cypher_prompt = cypher_gen_prompt
         self.qa_prompt = qa_prompt
@@ -108,4 +112,31 @@ def send_message(self, message: str):
             "response": answer, 
             "context": context, 
             "cypher": cypher
-            }
+            }
+        
+    def clean_ontology_for_prompt(self, ontology: dict):
+        """
+        Cleans the ontology by removing 'unique' and 'required' keys and prepares it for use in a prompt.
+
+        Args:
+            ontology (dict): The ontology to clean and transform.
+
+        Returns:
+            str: The cleaned ontology as a JSON string.
+        """
+        # Convert the ontology object to a JSON.
+        ontology = ontology.to_json()
+        
+        # Remove unique and required attributes from the ontology.
+        for entity in ontology["entities"]:
+            for attribute in entity["attributes"]:
+                del attribute['unique']
+                del attribute['required']
+        
+        for relation in ontology["relations"]:
+            for attribute in relation["attributes"]:
+                del attribute['unique']
+                del attribute['required']
+        
+        # Return the transformed ontology as a JSON string
+        return json.dumps(ontology)

graphrag_sdk/entity.py

@@ -82,7 +82,7 @@ def from_json(txt: dict | str):
             txt.get("description", ""),
         )
 
-    def to_json(self, include_all: bool = True) -> dict:
+    def to_json(self) -> dict:
         """
         Convert the entity object to a JSON representation.
 
@@ -95,7 +95,7 @@ def to_json(self, include_all: bool = True) -> dict:
         """
         return {
             "label": self.label,
-            "attributes": [attr.to_json(include_all=include_all) for attr in self.attributes],
+            "attributes": [attr.to_json() for attr in self.attributes],
             "description": self.description,
         }
 

graphrag_sdk/ontology.py

@@ -132,7 +132,7 @@ def from_kg_graph(graph: Graph, node_limit: int = 100,):
             attributes = graph.query(
                 f"""MATCH (a:{label[0]}) call {{ with a return [k in keys(a) | [k, typeof(a[k])]] as types }} 
                 WITH types limit {node_limit} unwind types as kt RETURN kt, count(1)""").result_set
-            ontology.add_entity(Entity(label[0], [Attribute(attr[0][0], attr[0][1], False, False) for attr in attributes]))
+            ontology.add_entity(Entity(label[0], [Attribute(attr[0][0], attr[0][1]) for attr in attributes]))
 
         # Process each relationship type and extract attributes, limited to the specified number of nodes
         for label in r_labels:
@@ -143,7 +143,7 @@ def from_kg_graph(graph: Graph, node_limit: int = 100,):
                         attributes = graph.query(
                             f"""MATCH ()-[a:{label[0]}]->() call {{ with a return [k in keys(a) | [k, typeof(a[k])]] as types }} 
                             WITH types limit {node_limit} unwind types as kt RETURN kt, count(1)""").result_set
-                        ontology.add_relation(Relation(label[0], label_s[0], label_t[0], [Attribute(attr[0][0], attr[0][1], False, False) for attr in attributes]))
+                        ontology.add_relation(Relation(label[0], label_s[0], label_t[0], [Attribute(attr[0][0], attr[0][1]) for attr in attributes]))
 
         return ontology
 
@@ -165,16 +165,16 @@ def add_relation(self, relation: Relation):
         """
         self.relations.append(relation)
 
-    def to_json(self, include_all: bool = True) -> dict:
+    def to_json(self) -> dict:
         """
         Converts the ontology object to a JSON representation.
 
         Returns:
             A dictionary representing the ontology object in JSON format.
         """
         return {
-            "entities": [entity.to_json(include_all) for entity in self.entities],
-            "relations": [relation.to_json(include_all) for relation in self.relations],
+            "entities": [entity.to_json() for entity in self.entities],
+            "relations": [relation.to_json() for relation in self.relations],
         }
 
     def merge_with(self, o: "Ontology"):

graphrag_sdk/relation.py

@@ -205,7 +205,7 @@ def from_string(txt: str):
             [Attribute.from_string(attr) for attr in attributes],
         )
 
-    def to_json(self, include_all: bool = True) -> dict:
+    def to_json(self) -> dict:
         """
         Converts the Relation object to a JSON dictionary.
 
@@ -216,7 +216,7 @@ def to_json(self, include_all: bool = True) -> dict:
             "label": self.label,
             "source": self.source.to_json(),
             "target": self.target.to_json(),
-            "attributes": [attr.to_json(include_all=include_all) for attr in self.attributes],
+            "attributes": [attr.to_json() for attr in self.attributes],
         }
 
     def combine(self, relation2: "Relation"):