[docs]classAGEQueryException(Exception):"""Exception for the AGE queries."""def__init__(self,exception:Union[str,Dict])->None:ifisinstance(exception,dict):self.message=exception["message"]if"message"inexceptionelse"unknown"self.details=exception["details"]if"details"inexceptionelse"unknown"else:self.message=exceptionself.details="unknown"defget_message(self)->str:returnself.messagedefget_details(self)->Any:returnself.details
[docs]classAGEGraph(GraphStore):""" Apache AGE wrapper for graph operations. Args: graph_name (str): the name of the graph to connect to or create conf (Dict[str, Any]): the pgsql connection config passed directly to psycopg2.connect create (bool): if True and graph doesn't exist, attempt to create it *Security note*: Make sure that the database connection uses credentials that are narrowly-scoped to only include necessary permissions. Failure to do so may result in data corruption or loss, since the calling code may attempt commands that would result in deletion, mutation of data if appropriately prompted or reading sensitive data if such data is present in the database. The best way to guard against such negative outcomes is to (as appropriate) limit the permissions granted to the credentials used with this tool. See https://python.langchain.com/docs/security for more information. """# python type mapping for providing readable types to LLMtypes={"str":"STRING","float":"DOUBLE","int":"INTEGER","list":"LIST","dict":"MAP","bool":"BOOLEAN",}# precompiled regex for checking chars in graph labelslabel_regex:Pattern=re.compile("[^0-9a-zA-Z]+")
[docs]def__init__(self,graph_name:str,conf:Dict[str,Any],create:bool=True)->None:"""Create a new AGEGraph instance."""self.graph_name=graph_name# check that psycopg2 is installedtry:importpsycopg2exceptImportError:raiseImportError("Could not import psycopg2 python package. ""Please install it with `pip install psycopg2`.")self.connection=psycopg2.connect(**conf)withself._get_cursor()ascurs:# check if graph with name graph_name existsgraph_id_query=("""SELECT graphid FROM ag_catalog.ag_graph WHERE name = '{}'""".format(graph_name))curs.execute(graph_id_query)data=curs.fetchone()# if graph doesn't exist and create is True, create itifdataisNone:ifcreate:create_statement=""" SELECT ag_catalog.create_graph('{}'); """.format(graph_name)try:curs.execute(create_statement)self.connection.commit()exceptpsycopg2.Errorase:raiseAGEQueryException({"message":"Could not create the graph","detail":str(e),})else:raiseException(('Graph "{}" does not exist in the database '+'and "create" is set to False').format(graph_name))curs.execute(graph_id_query)data=curs.fetchone()# store graph id and refresh the schemaself.graphid=data.graphidself.refresh_schema()
def_get_cursor(self)->psycopg2.extras.NamedTupleCursor:""" get cursor, load age extension and set search path """try:importpsycopg2.extrasexceptImportErrorase:raiseImportError("Unable to import psycopg2, please install with ""`pip install -U psycopg2`.")fromecursor=self.connection.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor)cursor.execute("""LOAD 'age';""")cursor.execute("""SET search_path = ag_catalog, "$user", public;""")returncursordef_get_labels(self)->Tuple[List[str],List[str]]:""" Get all labels of a graph (for both edges and vertices) by querying the graph metadata table directly Returns Tuple[List[str]]: 2 lists, the first containing vertex labels and the second containing edge labels """e_labels_records=self.query("""MATCH ()-[e]-() RETURN collect(distinct label(e)) as labels""")e_labels=e_labels_records[0]["labels"]ife_labels_recordselse[]n_labels_records=self.query("""MATCH (n) RETURN collect(distinct label(n)) as labels""")n_labels=n_labels_records[0]["labels"]ifn_labels_recordselse[]returnn_labels,e_labelsdef_get_triples(self,e_labels:List[str])->List[Dict[str,str]]:""" Get a set of distinct relationship types (as a list of dicts) in the graph to be used as context by an llm. Args: e_labels (List[str]): a list of edge labels to filter for Returns: List[Dict[str, str]]: relationships as a list of dicts in the format "{'start':<from_label>, 'type':<edge_label>, 'end':<from_label>}" """# age query to get distinct relationship typestry:importpsycopg2exceptImportErrorase:raiseImportError("Unable to import psycopg2, please install with ""`pip install -U psycopg2`.")frometriple_query=""" SELECT * FROM ag_catalog.cypher('{graph_name}', $$ MATCH (a)-[e:`{e_label}`]->(b) WITH a,e,b LIMIT 3000 RETURN DISTINCT labels(a) AS from, type(e) AS edge, labels(b) AS to LIMIT 10 $$) AS (f agtype, edge agtype, t agtype); """triple_schema=[]# iterate desired edge types and add distinct relationship types to resultwithself._get_cursor()ascurs:forlabeline_labels:q=triple_query.format(graph_name=self.graph_name,e_label=label)try:curs.execute(q)data=curs.fetchall()fordindata:# use json.loads to convert returned# strings to python primitivestriple_schema.append({"start":json.loads(d.f)[0],"type":json.loads(d.edge),"end":json.loads(d.t)[0],})exceptpsycopg2.Errorase:raiseAGEQueryException({"message":"Error fetching triples","detail":str(e),})returntriple_schemadef_get_triples_str(self,e_labels:List[str])->List[str]:""" Get a set of distinct relationship types (as a list of strings) in the graph to be used as context by an llm. Args: e_labels (List[str]): a list of edge labels to filter for Returns: List[str]: relationships as a list of strings in the format "(:`<from_label>`)-[:`<edge_label>`]->(:`<to_label>`)" """triples=self._get_triples(e_labels)returnself._format_triples(triples)@staticmethoddef_format_triples(triples:List[Dict[str,str]])->List[str]:""" Convert a list of relationships from dictionaries to formatted strings to be better readable by an llm Args: triples (List[Dict[str,str]]): a list relationships in the form {'start':<from_label>, 'type':<edge_label>, 'end':<from_label>} Returns: List[str]: a list of relationships in the form "(:`<from_label>`)-[:`<edge_label>`]->(:`<to_label>`)" """triple_template="(:`{start}`)-[:`{type}`]->(:`{end}`)"triple_schema=[triple_template.format(**triple)fortripleintriples]returntriple_schemadef_get_node_properties(self,n_labels:List[str])->List[Dict[str,Any]]:""" Fetch a list of available node properties by node label to be used as context for an llm Args: n_labels (List[str]): a list of node labels to filter for Returns: List[Dict[str, Any]]: a list of node labels and their corresponding properties in the form "{ 'labels': <node_label>, 'properties': [ { 'property': <property_name>, 'type': <property_type> },... ] }" """try:importpsycopg2exceptImportErrorase:raiseImportError("Unable to import psycopg2, please install with ""`pip install -U psycopg2`.")frome# cypher query to fetch properties of a given labelnode_properties_query=""" SELECT * FROM ag_catalog.cypher('{graph_name}', $$ MATCH (a:`{n_label}`) RETURN properties(a) AS props LIMIT 100 $$) AS (props agtype); """node_properties=[]withself._get_cursor()ascurs:forlabelinn_labels:q=node_properties_query.format(graph_name=self.graph_name,n_label=label)try:curs.execute(q)exceptpsycopg2.Errorase:raiseAGEQueryException({"message":"Error fetching node properties","detail":str(e),})data=curs.fetchall()# build a set of distinct propertiess=set({})fordindata:# use json.loads to convert to python# primitive and get readable typefork,vinjson.loads(d.props).items():s.add((k,self.types[type(v).__name__]))np={"properties":[{"property":k,"type":v}fork,vins],"labels":label,}node_properties.append(np)returnnode_propertiesdef_get_edge_properties(self,e_labels:List[str])->List[Dict[str,Any]]:""" Fetch a list of available edge properties by edge label to be used as context for an llm Args: e_labels (List[str]): a list of edge labels to filter for Returns: List[Dict[str, Any]]: a list of edge labels and their corresponding properties in the form "{ 'labels': <edge_label>, 'properties': [ { 'property': <property_name>, 'type': <property_type> },... ] }" """try:importpsycopg2exceptImportErrorase:raiseImportError("Unable to import psycopg2, please install with ""`pip install -U psycopg2`.")frome# cypher query to fetch properties of a given labeledge_properties_query=""" SELECT * FROM ag_catalog.cypher('{graph_name}', $$ MATCH ()-[e:`{e_label}`]->() RETURN properties(e) AS props LIMIT 100 $$) AS (props agtype); """edge_properties=[]withself._get_cursor()ascurs:forlabeline_labels:q=edge_properties_query.format(graph_name=self.graph_name,e_label=label)try:curs.execute(q)exceptpsycopg2.Errorase:raiseAGEQueryException({"message":"Error fetching edge properties","detail":str(e),})data=curs.fetchall()# build a set of distinct propertiess=set({})fordindata:# use json.loads to convert to python# primitive and get readable typefork,vinjson.loads(d.props).items():s.add((k,self.types[type(v).__name__]))np={"properties":[{"property":k,"type":v}fork,vins],"type":label,}edge_properties.append(np)returnedge_properties
[docs]defrefresh_schema(self)->None:""" Refresh the graph schema information by updating the available labels, relationships, and properties """# fetch graph schema informationn_labels,e_labels=self._get_labels()triple_schema=self._get_triples(e_labels)node_properties=self._get_node_properties(n_labels)edge_properties=self._get_edge_properties(e_labels)# update the formatted string representationself.schema=f""" Node properties are the following:{node_properties} Relationship properties are the following:{edge_properties} The relationships are the following:{self._format_triples(triple_schema)} """# update the dictionary representationself.structured_schema={"node_props":{el["labels"]:el["properties"]forelinnode_properties},"rel_props":{el["type"]:el["properties"]forelinedge_properties},"relationships":triple_schema,"metadata":{},}
@propertydefget_schema(self)->str:"""Returns the schema of the Graph"""returnself.schema@propertydefget_structured_schema(self)->Dict[str,Any]:"""Returns the structured schema of the Graph"""returnself.structured_schema@staticmethoddef_get_col_name(field:str,idx:int)->str:""" Convert a cypher return field to a pgsql select field If possible keep the cypher column name, but create a generic name if necessary Args: field (str): a return field from a cypher query to be formatted for pgsql idx (int): the position of the field in the return statement Returns: str: the field to be used in the pgsql select statement """# remove white spacefield=field.strip()# if an alias is provided for the field, use itif" as "infield:returnfield.split(" as ")[-1].strip()# if the return value is an unnamed primitive, give it a generic nameeliffield.isnumeric()orfieldin("true","false","null"):returnf"column_{idx}"# otherwise return the value stripping out some common special charselse:returnfield.replace("(","_").replace(")","")@staticmethoddef_wrap_query(query:str,graph_name:str)->str:""" Convert a cypher query to an Apache Age compatible sql query by wrapping the cypher query in ag_catalog.cypher, casting results to agtype and building a select statement Args: query (str): a valid cypher query graph_name (str): the name of the graph to query Returns: str: an equivalent pgsql query """# pgsql templatetemplate="""SELECT {projection} FROM ag_catalog.cypher('{graph_name}', $${query} $$) AS ({fields});"""# if there are any returned fields they must be added to the pgsql queryif"return"inquery.lower():# parse return statement to identify returned fieldsfields=(query.lower().split("return")[-1].split("distinct")[-1].split("order by")[0].split("skip")[0].split("limit")[0].split(","))# raise exception if RETURN * is found as we can't resolve the fieldsif"*"in[x.strip()forxinfields]:raiseValueError("AGE graph does not support 'RETURN *'"+" statements in Cypher queries")# get pgsql formatted field namesfields=[AGEGraph._get_col_name(field,idx)foridx,fieldinenumerate(fields)]# build resulting pgsql relationfields_str=", ".join([field.split(".")[-1]+" agtype"forfieldinfields])# if no return statement we still need to return a single field of type agtypeelse:fields_str="a agtype"select_str="*"returntemplate.format(graph_name=graph_name,query=query,fields=fields_str,projection=select_str,)@staticmethoddef_record_to_dict(record:NamedTuple)->Dict[str,Any]:""" Convert a record returned from an age query to a dictionary Args: record (): a record from an age query result Returns: Dict[str, Any]: a dictionary representation of the record where the dictionary key is the field name and the value is the value converted to a python type """# result holderd={}# prebuild a mapping of vertex_id to vertex mappings to be used# later to build edgesvertices={}forkinrecord._fields:v=getattr(record,k)# agtype comes back '{key: value}::type' which must be parsedifisinstance(v,str)and"::"inv:dtype=v.split("::")[-1]v=v.split("::")[0]ifdtype=="vertex":vertex=json.loads(v)vertices[vertex["id"]]=vertex.get("properties")# iterate returned fields and parse appropriatelyforkinrecord._fields:v=getattr(record,k)ifisinstance(v,str)and"::"inv:dtype=v.split("::")[-1]v=v.split("::")[0]else:dtype=""ifdtype=="vertex":d[k]=json.loads(v).get("properties")# convert edge from id-label->id by replacing id with node information# we only do this if the vertex was also returned in the query# this is an attempt to be consistent with neo4j implementationelifdtype=="edge":edge=json.loads(v)d[k]=(vertices.get(edge["start_id"],{}),edge["label"],vertices.get(edge["end_id"],{}),)else:d[k]=json.loads(v)ifisinstance(v,str)elsevreturnd
[docs]defquery(self,query:str,params:dict={})->List[Dict[str,Any]]:""" Query the graph by taking a cypher query, converting it to an age compatible query, executing it and converting the result Args: query (str): a cypher query to be executed params (dict): parameters for the query (not used in this implementation) Returns: List[Dict[str, Any]]: a list of dictionaries containing the result set """try:importpsycopg2exceptImportErrorase:raiseImportError("Unable to import psycopg2, please install with ""`pip install -U psycopg2`.")frome# convert cypher query to pgsql/age querywrapped_query=self._wrap_query(query,self.graph_name)# execute the query, rolling back on an errorwithself._get_cursor()ascurs:try:curs.execute(wrapped_query)self.connection.commit()exceptpsycopg2.Errorase:self.connection.rollback()raiseAGEQueryException({"message":"Error executing graph query: {}".format(query),"detail":str(e),})data=curs.fetchall()ifdataisNone:result=[]# convert to dictionarieselse:result=[self._record_to_dict(d)fordindata]returnresult
@staticmethoddef_format_properties(properties:Dict[str,Any],id:Union[str,None]=None)->str:""" Convert a dictionary of properties to a string representation that can be used in a cypher query insert/merge statement. Args: properties (Dict[str,str]): a dictionary containing node/edge properties id (Union[str, None]): the id of the node or None if none exists Returns: str: the properties dictionary as a properly formatted string """props=[]# wrap property key in backticks to escapefork,vinproperties.items():prop=f"`{k}`: {json.dumps(v)}"props.append(prop)ifidisnotNoneand"id"notinproperties:props.append(f"id: {json.dumps(id)}"ifisinstance(id,str)elsef"id: {id}")return"{"+", ".join(props)+"}"
[docs]@staticmethoddefclean_graph_labels(label:str)->str:""" remove any disallowed characters from a label and replace with '_' Args: label (str): the original label Returns: str: the sanitized version of the label """returnre.sub(AGEGraph.label_regex,"_",label)
[docs]defadd_graph_documents(self,graph_documents:List[GraphDocument],include_source:bool=False)->None:""" insert a list of graph documents into the graph Args: graph_documents (List[GraphDocument]): the list of documents to be inserted include_source (bool): if True add nodes for the sources with MENTIONS edges to the entities they mention Returns: None """# query for inserting nodesnode_insert_query=(""" MERGE (n:`{label}` {properties}) """ifnotinclude_sourceelse""" MERGE (n:`{label}` {properties}) MERGE (d:Document {d_properties}) MERGE (d)-[:MENTIONS]->(n) """)# query for inserting edgesedge_insert_query=""" MERGE (from:`{f_label}` {f_properties}) MERGE (to:`{t_label}` {t_properties}) MERGE (from)-[:`{r_label}` {r_properties}]->(to) """# iterate docs and insert themfordocingraph_documents:# if we are adding sources, create an id for the sourceifinclude_source:ifnotdoc.source.metadata.get("id"):doc.source.metadata["id"]=md5(doc.source.page_content.encode("utf-8")).hexdigest()# insert entity nodesfornodeindoc.nodes:node.properties["id"]=node.idifinclude_source:query=node_insert_query.format(label=node.type,properties=self._format_properties(node.properties),d_properties=self._format_properties(doc.source.metadata),)else:query=node_insert_query.format(label=AGEGraph.clean_graph_labels(node.type),properties=self._format_properties(node.properties),)self.query(query)# insert relationshipsforedgeindoc.relationships:edge.source.properties["id"]=edge.source.idedge.target.properties["id"]=edge.target.idinputs={"f_label":AGEGraph.clean_graph_labels(edge.source.type),"f_properties":self._format_properties(edge.source.properties),"t_label":AGEGraph.clean_graph_labels(edge.target.type),"t_properties":self._format_properties(edge.target.properties),"r_label":AGEGraph.clean_graph_labels(edge.type).upper(),"r_properties":self._format_properties(edge.properties),}query=edge_insert_query.format(**inputs)self.query(query)