NgEhddlmZddlZddlZddlmZddlmZmZm Z m Z m Z m Z m Z mZddlmZddlmZerddlZGddeZGd d eZdS) ) annotationsN)md5) TYPE_CHECKINGAnyDictList NamedTuplePatternTupleUnion) GraphDocument) GraphStorec*eZdZdZd dZd dZdd Zd S)AGEQueryExceptionzException for the AGE queries. exceptionUnion[str, Dict]returnNonect|tr(d|vr|dnd|_d|vr|dnd|_dS||_d|_dS)Nmessageunknowndetails) isinstancedictrr)selfrs `/var/www/html/ai-engine/env/lib/python3.11/site-packages/langchain_community/graphs/age_graph.py__init__zAGEQueryException.__init__s` i & & %3< 3I3I9Y//yDL3< 3I3I9Y//yDLLL$DL$DLLLstrc|jSN)rrs r get_messagezAGEQueryException.get_message |rrc|jSr!)rr"s r get_detailszAGEQueryException.get_detailsr$rN)rrrrrr)rr)__name__ __module__ __qualname____doc__rr#r&rrrrsV((%%%%rrcpeZdZUdZdddddddZejd Zd ed < d@dAdZ dBdZ dCdZ dDdZ dEdZ edFd ZdGd#ZdHd$ZdId%ZedJd&ZedKd'ZedLd+ZedMd-ZedNd0ZifdOd3Ze dPdQd8ZedRd:Z dSdTd?Zd4S)UAGEGrapha 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. STRINGDOUBLEINTEGERLISTMAPBOOLEAN)rfloatintlistrboolz [^0-9a-zA-Z]+r label_regexT graph_namerconfDict[str, Any]creater8rrc||_ ddl}n#t$rtdwxYw|jd i||_|5}d|}|||}||rvd|} |||j nS#|j $r$} tdt| dd} ~ wwxYwtd||||}|j|_|ddddS#1swxYwYdS) zCreate a new AGEGraph instance.rNzXCould not import psycopg2 python package. Please install it with `pip install psycopg2`.z9SELECT graphid FROM ag_catalog.ag_graph WHERE name = '{}'zS SELECT ag_catalog.create_graph('{}'); zCould not create the graphrdetailzFGraph "{}" does not exist in the database and "create" is set to Falser,)r:psycopg2 ImportErrorconnect connection _get_cursorformatexecutefetchonecommitErrorrr Exceptiongraphidrefresh_schema) rr:r;r=rAcursgraph_id_querydatacreate_statementes rrzAGEGraph.__init__Ds %  OOOO   A   +(*22T22     * "4PVV  LL ( ( (==??D|(z**%  %5666..0000#>/+G*-a&&$= &,,  ^,,,}}
, 'type':, 'end':}" rNrUa( 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); r:e_label)starttypeendzError fetching triplesr?)rArBrErFr:rGfetchallappendjsonloadsfedgetrJrr) rr`rArR triple_query triple_schemarNlabelqrPds r _get_tripleszAGEGraph._get_tripless  OOOO   -          4!   ''4?E'RRLLOOO==??D!  &,,)-AC);(, 16(:(:'+z!#q'9  ~+'?&)!ff                0sG &!&!D:#BC<;D:< D*D%%D**D::D>D>cV||}||S)a} 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 "(:``)-[:``]->(:``)" )rx_format_triples)rr`tripless r_get_triples_strzAGEGraph._get_triples_strs+##H--##G,,,rr{c(dfd|D}|S)a 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':, 'type':, 'end':} Returns: List[str]: a list of relationships in the form "(:``)-[:``]->(:``)" z$(:`{start}`)-[:`{type}`]->(:`{end}`)c*g|]}jdi|S)r,)rF).0tripletriple_templates r z,AGEGraph._format_triples..s-PPPf//99&99PPPrr,)r{rtrs @rrzzAGEGraph._format_tripless*APPPPPPP rrbList[Dict[str, Any]]c  ddl}n"#t$r}td|d}~wwxYwd}g}|5}|D]}||j|} ||n1#|j$r$}tdt|dd}~wwxYw| } ti} | D]g} tj | j D]9\} } | | |jt#| jf:hd| D|d }|| dddn #1swxYwY|S) a 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': , 'properties': [ { 'property': , 'type': },... ] }" rNrUz SELECT * FROM ag_catalog.cypher('{graph_name}', $$ MATCH (a:`{n_label}`) RETURN properties(a) AS props LIMIT 100 $$) AS (props agtype); )r:n_labelzError fetching node propertiesr?cg|] \}}||d S)propertyrjr,rkvs rrz1AGEGraph._get_node_properties..I$"L"L"L$!Q1#=#="L"L"Lr) propertiesr]rArBrErFr:rGrJrrrlsetrnropropsitemsaddtypesrjr(rm)rrbrArRnode_properties_querynode_propertiesrNrurvrPsrwrrnps r_get_node_propertieszAGEGraph._get_node_propertiess$*  OOOO   -  !       +4! + +)00#1LLOOOO~+'G&)!ff}}GGAAA!% 17 3 3 9 9 ; ;AA1q$*T!WW-=">?@@@@A#M"L!"L"L"L# &&r****9 + + + + + + + + + + + + + + + +>G &!&"E+$A:9E+: B(B##B((B6E++E/2E/c  ddl}n"#t$r}td|d}~wwxYwd}g}|5}|D]}||j|} ||n1#|j$r$}tdt|dd}~wwxYw| } ti} | D]g} tj | j D]9\} } | | |jt#| jf:hd| D|d }|| dddn #1swxYwY|S) a 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': , 'properties': [ { 'property': , 'type': },... ] }" rNrUz SELECT * FROM ag_catalog.cypher('{graph_name}', $$ MATCH ()-[e:`{e_label}`]->() RETURN properties(e) AS props LIMIT 100 $$) AS (props agtype); rgzError fetching edge propertiesr?cg|] \}}||d Srr,rs rrz1AGEGraph._get_edge_properties..rr)rrjr)rr`rArRedge_properties_queryedge_propertiesrNrurvrPrrwrrrs r_get_edge_propertieszAGEGraph._get_edge_propertiesPs$,  OOOO   -   !       +4! + +)00#1LLOOOO~+'G&)!ff}}GGAAA!% 17 3 3 9 9 ; ;AA1q$*T!WW-=">?@@@@A#M"L!"L"L"L! &&r****9 + + + + + + + + + + + + + + + +>rc:|\}}||}||}||}d|d|d||d|_d|Dd|D|id|_dS) z~ Refresh the graph schema information by updating the available labels, relationships, and properties z4 Node properties are the following: z< Relationship properties are the following: z6 The relationships are the following: z c,i|]}|d|dS)r]rr,rels r z+AGEGraph.refresh_schema..s#TTTb2h<L)9TTTrc,i|]}|d|dS)rjrr,rs rrz+AGEGraph.refresh_schema..s#QQQ2"V*b&6QQQr) node_props rel_props relationshipsmetadataN)rcrxrrrzschemastructured_schema)rrbr`rtrrs rrMzAGEGraph.refresh_schemas"--//())(33 33H==33H==         m , ,     UTOTTTQQQQQ* " " rc|jS)zReturns the schema of the Graph)rr"s r get_schemazAGEGraph.get_schemas {rc|jS)z*Returns the structured schema of the Graph)rr"s rget_structured_schemazAGEGraph.get_structured_schemas %%rfieldidxr6c|}d|vr-|ddS|s|dvrd|S|ddddS) a 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 z as )truefalsenullcolumn_(_))stripsplit isnumericreplace)rrs r _get_col_namezAGEGraph._get_col_names  U??;;v&&r*0022 2 __   <%+D"D"D"S?? "==c**223;; ;rr^c`d}d|vr|ddddddddddd }d d |Dvrtd d t|D}dd|D}nd}d }|||||S)a 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 zlSELECT {projection} FROM ag_catalog.cypher('{graph_name}', $$ {query} $$) AS ({fields});rrdistinctzorder byrskiplimit,*c6g|]}|Sr,)r)rxs rrz(AGEGraph._wrap_query..s 111Qqwwyy111rzBAGE graph does not support 'RETURN *' statements in Cypher queriescJg|] \}}t||!Sr,)r.r)rrrs rrz(AGEGraph._wrap_query..s97AsE&&uc22r, cJg|] }|dddz!S).rz agtype)r)rrs rrz(AGEGraph._wrap_query..s.FFFeS!!"% 1FFFrza agtype)r:r^fields projection)lowerr ValueError enumeratejoinrF)r^r:templater fields_str select_strs r _wrap_queryzAGEGraph._wrap_querysU  u{{}} $ $ x%z""2'z""1&vq " w # s 11&11111 6 ENvEVEVF FFvFFFJJ $J !!    rrecordr ci}i}|jD]}t||}t|trrd|vrn|dd}|dd}|dkr2t j|}|d||d<|jD]!}t||}t|tr;d|vr7|dd}|dd}nd}|dkr+t j|d||<|dkrWt j|}||d i|d ||d if||<t|trt j|n|||<#|S) an 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 z::rrvertexridrrqstart_idruend_id)_fieldsgetattrrrrrnroget)rrwverticesrrdtyperrqs r_record_to_dictzAGEGraph._record_to_dicts  F FA""A!S!! Fdaii b)GGDMM!$H$$!Z]]F-3ZZ -E-EHVD\* B BA""A!S!! daii b)GGDMM!$  z!}}((66!&z!}}LLj!1266MLLh44! )31c(:(:Atz!}}}!rparamsrcR ddl}n"#t$r}td|d}~wwxYw|j}5} ||jn]#|j$rP}j td |t|dd}~wwxYw| }|g}nfd|D}|cdddS#1swxYwYdS)a 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 rNrUzError executing graph query: {}r?c:g|]}|Sr,)r)rrwrs rrz"AGEGraph.query..ys'@@@a$..q11@@@r)rArBrr:rErGrDrIrJrollbackrrFrrl) rr^rrArR wrapped_queryrNrPresults ` rr^zAGEGraph.queryOs  OOOO   -  ((@@      4  ]+++&&((((>   ((***'#D#K#KE#R#R"%a&& ==??D|A@@@4@@@)                  sG '"'D.B  D C$A CC$$+DD #D NrrUnion[str, None]cbg}|D]4\}}d|dtj|}||5|Hd|vrD|t |t rdtj|nd|dd|zdzS) a 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 `z`: Nrzid: {r})rrndumpsrmrrr)rrrrrprops r_format_propertieszAGEGraph._format_properties}s$$&&  DAq,q,,TZ]],,D LL     >d*44 LL+5b#+>+>O'tz"~~'''K2KK   TYYu%%%++rrucBtjtjd|S)z remove any disallowed characters from a label and replace with '_' Args: label (str): the original label Returns: str: the sanitized version of the label r)resubr.r9)rus rclean_graph_labelszAGEGraph.clean_graph_labelssvh*C777rFgraph_documentsList[GraphDocument]include_sourcec |sdnd}d}|D]l}|rj|jjdsKt|jjd|jjd<|jD]}|j|j d<|rS| |j | |j | |jj}nL| t|j | |j }|||jD]%}|jj|jj d<|jj|jj d<t|jj | |jj t|jj | |jj t|j | |j d} |j d i| }||'nd S) aQ 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 z6 MERGE (n:`{label}` {properties}) z MERGE (n:`{label}` {properties}) MERGE (d:Document {d_properties}) MERGE (d)-[:MENTIONS]->(n) z MERGE (from:`{f_label}` {f_properties}) MERGE (to:`{t_label}` {t_properties}) MERGE (from)-[:`{r_label}` {r_properties}]->(to) rzutf-8)rur d_properties)rur)f_label f_propertiest_label t_propertiesr_label r_propertiesNr,)sourcerrr page_contentencode hexdigestnodesrrrFrjrr.rr^rtargetupper) rrrnode_insert_queryedge_insert_querydocnoder^rqinputss radd_graph_documentszAGEGraph.add_graph_documentssr&"      #' "' "C "z*..t44"03 /66w??11ikkJ'-   " "(,%! -44"i#'#:#:4?#K#K%)%<%X&&&X&<<<\<0< < < \< |444\4l02,,,,,\;?,,,,\,2 8 8 8\ 8LQI"I"I"I"I"I"I"rr.) __future__rrnrhashlibrtypingrrrrr r r r )langchain_community.graphs.graph_documentr &langchain_community.graphs.graph_storerrWrArKrr.r,rrrs"""""" TTTTTTTTTTTTTTTTTTTTCCCCCC====== $L "L "L "L "L "zL "L "L "L "L "r