Skip to content

Commit b3428d4

Browse files
formatted the queries and added logic for empty label (#752)
1 parent eee5e1d commit b3428d4

File tree

1 file changed

+141
-49
lines changed

1 file changed

+141
-49
lines changed

backend/src/shared/constants.py

+141-49
Original file line numberDiff line numberDiff line change
@@ -235,51 +235,97 @@
235235
LOCAL_COMMUNITY_TOP_OUTSIDE_RELS = 10
236236

237237
LOCAL_COMMUNITY_SEARCH_QUERY = """
238-
WITH collect(node) as nodes, avg(score) as score, collect({{id: elementId(node), score:score}}) as metadata
239-
WITH score, nodes,metadata,
240-
collect {{
241-
UNWIND nodes as n
242-
MATCH (n)<-[:HAS_ENTITY]->(c:Chunk)
243-
WITH c, count(distinct n) as freq
244-
RETURN c
245-
ORDER BY freq DESC
246-
LIMIT {topChunks}
247-
}} AS chunks,
248-
collect {{
249-
UNWIND nodes as n
250-
MATCH (n)-[:IN_COMMUNITY]->(c:__Community__)
251-
WITH c, c.community_rank as rank, c.weight AS weight
252-
RETURN c
253-
ORDER BY rank, weight DESC
254-
LIMIT {topCommunities}
255-
}} AS communities,
256-
collect {{
257-
unwind nodes as n
258-
unwind nodes as m
259-
match (n)-[r]->(m)
260-
RETURN distinct r
261-
// todo should we have a limit here?
262-
}} as rels,
263-
collect {{
264-
unwind nodes as n
265-
match path = (n)-[r]-(m:__Entity__)
266-
where not m in nodes
267-
with m, collect(distinct r) as rels, count(*) as freq
268-
ORDER BY freq DESC LIMIT {topOutsideRels}
269-
WITH collect(m) as outsideNodes, apoc.coll.flatten(collect(rels)) as rels
270-
RETURN {{ nodes: outsideNodes, rels: rels }}
271-
}} as outside
238+
WITH collect(node) AS nodes,
239+
avg(score) AS score,
240+
collect({{id: elementId(node), score: score}}) AS metadata
241+
242+
WITH score, nodes, metadata,
243+
244+
collect {{
245+
UNWIND nodes AS n
246+
MATCH (n)<-[:HAS_ENTITY]->(c:Chunk)
247+
WITH c, count(distinct n) AS freq
248+
RETURN c
249+
ORDER BY freq DESC
250+
LIMIT {topChunks}
251+
}} AS chunks,
252+
253+
collect {{
254+
UNWIND nodes AS n
255+
MATCH (n)-[:IN_COMMUNITY]->(c:__Community__)
256+
WITH c, c.community_rank AS rank, c.weight AS weight
257+
RETURN c
258+
ORDER BY rank, weight DESC
259+
LIMIT {topCommunities}
260+
}} AS communities,
261+
262+
collect {{
263+
UNWIND nodes AS n
264+
UNWIND nodes AS m
265+
MATCH (n)-[r]->(m)
266+
RETURN DISTINCT r
267+
// TODO: need to add limit
268+
}} AS rels,
269+
270+
collect {{
271+
UNWIND nodes AS n
272+
MATCH path = (n)-[r]-(m:__Entity__)
273+
WHERE NOT m IN nodes
274+
WITH m, collect(distinct r) AS rels, count(*) AS freq
275+
ORDER BY freq DESC
276+
LIMIT {topOutsideRels}
277+
WITH collect(m) AS outsideNodes, apoc.coll.flatten(collect(rels)) AS rels
278+
RETURN {{ nodes: outsideNodes, rels: rels }}
279+
}} AS outside
272280
"""
273281

274282
LOCAL_COMMUNITY_SEARCH_QUERY_SUFFIX = """
275-
RETURN {chunks: [c in chunks | c.text],
276-
communities: [c in communities | c.summary],
277-
entities: [n in nodes | apoc.coll.removeAll(labels(n),["__Entity__"])[0] + ":"+ n.id + " " + coalesce(n.description, "")],
278-
relationships: [r in rels | startNode(r).id+" "+type(r)+" "+endNode(r).id],
279-
outside: {
280-
nodes: [n in outside[0].nodes | apoc.coll.removeAll(labels(n),["__Entity__"])[0] + ":"+n.id + " " + coalesce(n.description, "")],
281-
relationships: [r in outside[0].rels | apoc.coll.removeAll(labels(startNode(r)),["__Entity__"])[0] + ":"+startNode(r).id+" "+type(r)+" "+apoc.coll.removeAll(labels(startNode(r)),["__Entity__"])[0] + ":"+endNode(r).id]}
282-
} AS text, score, {entities:metadata} as metadata
283+
RETURN {
284+
chunks: [c IN chunks | c.text],
285+
communities: [c IN communities | c.summary],
286+
entities: [
287+
n IN nodes |
288+
CASE
289+
WHEN size(labels(n)) > 1 THEN
290+
apoc.coll.removeAll(labels(n), ["__Entity__"])[0] + ":" + n.id + " " + coalesce(n.description, "")
291+
ELSE
292+
n.id + " " + coalesce(n.description, "")
293+
END
294+
],
295+
relationships: [
296+
r IN rels |
297+
startNode(r).id + " " + type(r) + " " + endNode(r).id
298+
],
299+
outside: {
300+
nodes: [
301+
n IN outside[0].nodes |
302+
CASE
303+
WHEN size(labels(n)) > 1 THEN
304+
apoc.coll.removeAll(labels(n), ["__Entity__"])[0] + ":" + n.id + " " + coalesce(n.description, "")
305+
ELSE
306+
n.id + " " + coalesce(n.description, "")
307+
END
308+
],
309+
relationships: [
310+
r IN outside[0].rels |
311+
CASE
312+
WHEN size(labels(startNode(r))) > 1 THEN
313+
apoc.coll.removeAll(labels(startNode(r)), ["__Entity__"])[0] + ":" + startNode(r).id + " "
314+
ELSE
315+
startNode(r).id + " "
316+
END +
317+
type(r) + " " +
318+
CASE
319+
WHEN size(labels(endNode(r))) > 1 THEN
320+
apoc.coll.removeAll(labels(endNode(r)), ["__Entity__"])[0] + ":" + endNode(r).id
321+
ELSE
322+
endNode(r).id
323+
END
324+
]
325+
}
326+
} AS text,
327+
score,
328+
{entities: metadata} AS metadata
283329
"""
284330

285331
LOCAL_COMMUNITY_DETAILS_QUERY_PREFIX = """
@@ -289,14 +335,60 @@
289335
"""
290336
LOCAL_COMMUNITY_DETAILS_QUERY_SUFFIX = """
291337
WITH *
292-
UNWIND chunks as c
338+
UNWIND chunks AS c
293339
MATCH (c)-[:PART_OF]->(d:Document)
294-
RETURN [c {.*, embedding:null, fileName:d.fileName, fileSource:d.fileSource}] as chunks,
295-
[community in communities | community {.*, embedding:null}] as communities,
296-
[node in nodes+outside[0].nodes | {element_id:elementId(node), labels:labels(node), properties:{id:node.id,description:node.description}}] as nodes,
297-
[r in rels+outside[0].rels | {startNode:{element_id:elementId(startNode(r)), labels:labels(startNode(r)), properties:{id:startNode(r).id,description:startNode(r).description}},
298-
endNode:{element_id:elementId(endNode(r)), labels:labels(endNode(r)), properties:{id:endNode(r).id,description:endNode(r).description}},
299-
relationship: {type:type(r), element_id:elementId(r)}}] as entities
340+
RETURN
341+
[
342+
c {
343+
.*,
344+
embedding: null,
345+
fileName: d.fileName,
346+
fileSource: d.fileSource
347+
}
348+
] AS chunks,
349+
[
350+
community IN communities |
351+
community {
352+
.*,
353+
embedding: null
354+
}
355+
] AS communities,
356+
[
357+
node IN nodes + outside[0].nodes |
358+
{
359+
element_id: elementId(node),
360+
labels: labels(node),
361+
properties: {
362+
id: node.id,
363+
description: node.description
364+
}
365+
}
366+
] AS nodes,
367+
[
368+
r IN rels + outside[0].rels |
369+
{
370+
startNode: {
371+
element_id: elementId(startNode(r)),
372+
labels: labels(startNode(r)),
373+
properties: {
374+
id: startNode(r).id,
375+
description: startNode(r).description
376+
}
377+
},
378+
endNode: {
379+
element_id: elementId(endNode(r)),
380+
labels: labels(endNode(r)),
381+
properties: {
382+
id: endNode(r).id,
383+
description: endNode(r).description
384+
}
385+
},
386+
relationship: {
387+
type: type(r),
388+
element_id: elementId(r)
389+
}
390+
}
391+
] AS entities
300392
"""
301393

302394
CHAT_MODE_CONFIG_MAP= {

0 commit comments

Comments
 (0)