ZǻiE ddlmZddlZddlmZmZmZddlmZddl m Z ddl m Z m Z mZdZdZd Zd Zd Z dd Z d ddZddZddZdZdZdZdZddZd dZ d!dZe j>ddddddd d ej@df d"dZ!ddde j>f d#dZ"y)$) annotationsN)AnyOptionalUnion)InvalidHybridSearchRankerError)get_metadata_filter) EntityType SearchTypeHybridSearchRankerzCALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) YIELD node, score WITH node, score LIMIT $top_kzCALL db.index.vector.queryRelationships($vector_index_name, $top_k * $effective_search_ratio, $query_vector) YIELD relationship, score WITH relationship, score LIMIT $top_kz~WITH node, vector.similarity.cosine(node.`{embedding_node_property}`, $query_vector) AS score ORDER BY score DESC LIMIT $top_kzMATCH (node:`{node_label}`) WHERE node.`{embedding_node_property}` IS NOT NULL AND size(node.`{embedding_node_property}`) = toInteger($embedding_dimension)zgCALL db.index.fulltext.queryNodes($fulltext_index_name, $query_text, {limit: $top_k}) YIELD node, scorecFdj|}|rd|dS|rd|dSy)z{A helper function to return the CALL subquery syntax: - Either CALL { WITH - or CALL (variables) { ,zCALL (z) { z CALL { WITH  CALL { )join)support_variable_scope_clause variable_list variabless Z/opt/lhia/marcimex/agent/venv/lib/python3.12/site-packages/neo4j_graphrag/neo4j_queries.py_call_subquery_syntaxr8s9'I$ {%((yk++ Fc@t|ddg}|rd}nd}d|d|dS) zBuild the Cypher query to upsert a batch of nodes: - Create the new node - Set its label(s) and properties - Set its embedding properties if any - Return the node elementId nrowrzSET n:$(row.labels) zUWITH n, row CALL apoc.create.addLabels(n, row.labels) YIELD node WITH node as n, row zaUNWIND $rows AS row CREATE (n:__KGBuilder__ {__tmp_internal_id: row.id}) SET n += row.properties z WITH n, row z WITH n, row WHERE row.embedding_properties IS NOT NULL UNWIND keys(row.embedding_properties) as emb CALL db.create.setNodeVectorProperty(n, emb, row.embedding_properties[emb]) RETURN count(*) as nbEmb } RETURN elementId(n) as element_idr)rsupport_dynamic_labels call_prefix set_labelss rupsert_node_queryrGsI(%c5\K+ l  # , -, , rc,t|ddg}d|dS)aBuild the Cypher query to upsert a batch of relationships: - Create the new relationship: only one relationship of a specific type is allowed between the same two nodes - Set its properties - Set its embedding properties if any - Return the node elementId relrraUNWIND $rows as row MATCH (start:__KGBuilder__ {__tmp_internal_id: row.start_node_id}), (end:__KGBuilder__ {__tmp_internal_id: row.end_node_id}) WITH start, end, row CALL apoc.merge.relationship(start, row.type, {}, row.properties, end, row.properties) YIELD rel WITH rel, row z WITH rel, row WHERE row.embedding_properties IS NOT NULL UNWIND keys(row.embedding_properties) as emb CALL db.create.setRelationshipVectorProperty(rel, emb, row.embedding_properties[emb]) } RETURN elementId(rel)r)rrs rupsert_relationship_queryr"gs2(%eU^K  -  rc0t|dg}d|d|dS)z@Removes the temporary __tmp_internal_id property from all nodes.rrz>MATCH (n:__KGBuilder__) WHERE n.__tmp_internal_id IS NOT NULL z9 SET n.__tmp_internal_id = NULL } IN TRANSACTIONS OF z ROWS ON ERROR CONTINUEr)r batch_sizers rdb_cleaning_queryr%s8'%cUK 1 -)\*  rzMATCH (n) WHERE elementId(n) = $node_element_id WITH n CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector) RETURN nzUNWIND $rows AS row MATCH (n) WHERE elementId(n) = row.id WITH n, row CALL db.create.setNodeVectorProperty(n, $embedding_property, row.embedding) RETURN nzMATCH ()-[r]->() WHERE elementId(r) = $rel_element_id WITH r CALL db.create.setRelationshipVectorProperty(r, $embedding_property, $vector) RETURN rzUNWIND $rows AS row MATCH ()-[r]->() WHERE elementId(r) = row.id WITH r, row CALL db.create.setRelationshipVectorProperty(r, $embedding_property, row.embedding) RETURN rcFt|g}tdtd}||zS)a  Construct a cypher query for hybrid search. Args: neo4j_version_is_5_23_or_above (bool): Whether the Neo4j version is 5.23 or above; determines which call syntax is used. Returns: str: The constructed Cypher query string. rz WITH collect({node:node, score:score}) AS nodes, max(score) AS vector_index_max_score UNWIND nodes AS n RETURN n.node AS node, (n.score / vector_index_max_score) AS score UNION z WITH collect({node:node, score:score}) AS nodes, max(score) AS ft_index_max_score UNWIND nodes AS n RETURN n.node AS node, (n.score / ft_index_max_score) AS score } WITH node, max(score) AS score ORDER BY score DESC LIMIT $top_k)rNODE_VECTOR_INDEX_QUERYFULL_TEXT_SEARCH_QUERY)neo4j_version_is_5_23_or_abover query_bodys r_get_hybrid_queryr+sF(&bK # #$ " "#J J  ##rc8|rdnd}tdtd}||zS)a Construct a Cypher query for hybrid search using a linear combination approach with an alpha parameter. This query retrieves normalized scores from both the vector index and full-text index. It then computes the final score as a weighted sum: ``` final_score = alpha * (vector normalized score) + (1 - alpha) * (fulltext normalized score) ``` If a node appears in only one index, the missing score is treated as 0. Args: neo4j_version_is_5_23_or_above (bool): Whether the Neo4j version is 5.23 or above; determines the call syntax. alpha (float): Weight for the vector index normalized score. The full-text score is weighted as (1 - alpha). Returns: str: The constructed Cypher query string. z CALL () { rz WITH collect({node: node, score: score}) AS nodes, max(score) AS vector_index_max_score UNWIND nodes AS n WITH n.node AS node, (n.score / vector_index_max_score) AS rawScore RETURN node, rawScore * $alpha AS score UNION a WITH collect({node: node, score: score}) AS nodes, max(score) AS ft_index_max_score UNWIND nodes AS n WITH n.node AS node, (n.score / ft_index_max_score) AS rawScore RETURN node, rawScore * (1 - $alpha) AS score } WITH node, sum(score) AS score ORDER BY score DESC LIMIT $top_k)r'r()r)alpharr*s r_get_hybrid_query_linearr.sA(#A,iK # #$ " "#J J  ##rc|rdnd}t|d\}}tj||}tj|} ||d<||d|d | z|fS) aKBuild Cypher query for vector search with filters Uses exact KNN. Args: filters (dict[str, Any]): filters used to pre-filter the nodes before vector search node_label (str): node label we want to search for embedding_node_property (str): the name of the property holding the embeddings embedding_dimension (int): the dimension of the embeddings use_parallel_runtime (bool): Whether or not use the parallel runtime to run the query. Defaults to False. Returns: tuple[str, dict[str, Any]]: query and parameters z5CYPHER runtime = parallel parallelRuntimeSupport=all node) node_alias) node_labelembedding_node_property)r4embedding_dimensionz AND (z) )rBASE_VECTOR_EXACT_QUERYformatVECTOR_EXACT_QUERY) filtersr3r4r5use_parallel_runtimeparallel_query where_filters query_params base_query vector_querys r_get_filtered_vector_queryr@s.  @  #6g&"QM<(// 70J&,, 7-L+>L&'J"/+, %ME6$w!8v=k]KL L"#:";<    // / PQ Q *++ +UV V J-- -2B6E./F.GH =k]KL L9+GHH' J WAj\ "F **rc|r|S|rbdj|Dcgc]}d| c}}|tjk(rd|dS|tjk(rd|dSt d||r|SdScc}w) aBuild the RETURN statement after the search is performed Args return_properties (list[str]): list of property names to return. It can't be provided together with retrieval_query. retrieval_query (str): the query to use to retrieve the search results It can't be provided together with return_properties. fallback_return (str): the fallback return statement to use to retrieve the search results Returns: str: the RETURN statement z, .z RETURN node {za} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, scorezRETURN relationship {z} AS relationship, type(relationship) as relationshipType, elementId(relationship) AS elementId, elementId(relationship) AS id, scorerDr0)rr rJrQrP)rVrUrErFpropreturn_properties_cyphers rrSrSs$#'99EV-WT$j-W#X *// ) !9 :; J33 3()A(BC =k]KL L-?525'.Xs A/)rboolrz list[str]returnstr)F)rr`rr`rarb)rr`rarb)rr`r$intrarb)r)r`rarb)r)r`r-floatrarb) r9zdict[str, Any]r3rbr4rbr5rcr:r`ratuple[str, dict[str, Any]])rTr rFr rUOptional[list[str]]rV Optional[str]r3rgr4rgr5z Optional[int]r9zOptional[dict[str, Any]]r)r`r:r`rWzUnion[str, HybridSearchRanker]r-zOptional[float]rare) rVrgrUrfrErgrFr rarb)# __future__rrGtypingrrrneo4j_graphrag.exceptionsrneo4j_graphrag.filtersrneo4j_graphrag.typesr r r r'rRr8r6r(rrr"r%UPSERT_VECTOR_ON_NODE_QUERYUPSERT_VECTORS_ON_NODE_QUERY#UPSERT_VECTOR_ON_RELATIONSHIP_QUERY$UPSERT_VECTORS_ON_RELATIONSHIP_QUERYr+r.r@rJrMr[rSrrrrs#''D6KK$,'S #' 8A   IN#'AE@6 $$%$:$$N& &&!& &  &  &V)oo-1%) $-1)-(,+0!&-?-E-E!t+t+t++t+# t+  t+ + t+'t+&t+%)t+t+ +t+ t+ t+p&*-1%)(oo (6"(6*(6#(6 (6  (6r