Entity details shown for entity mode

backend/src/QA_integration.py

@@ -177,7 +177,7 @@ def get_rag_chain(llm, system_template=CHAT_SYSTEM_TEMPLATE):
         logging.error(f"Error creating RAG chain: {e}")
         raise
 
-def format_documents(documents, model):
+def format_documents(documents, model,chat_mode_settings):
     prompt_token_cutoff = 4
     for model_names, value in CHAT_TOKEN_CUT_OFF.items():
         if model in model_names:
@@ -197,12 +197,16 @@ def format_documents(documents, model):
         try:
             source = doc.metadata.get('source', "unknown")
             sources.add(source)
-
             if 'entities' in doc.metadata:
-                if 'entityids' in doc.metadata['entities']:
-                    entities.setdefault('entityids', set()).update(doc.metadata['entities']['entityids'])
-                if 'relationshipids' in doc.metadata['entities']:
-                    entities.setdefault('relationshipids', set()).update(doc.metadata['entities']['relationshipids'])
+                if chat_mode_settings["mode"] == CHAT_ENTITY_VECTOR_MODE:
+                    entity_ids = [entry['entityids'] for entry in doc.metadata['entities'] if 'entityids' in entry]
+                    entities.setdefault('entityids', set()).update(entity_ids)
+                else:
+                    if 'entityids' in doc.metadata['entities']:
+                        entities.setdefault('entityids', set()).update(doc.metadata['entities']['entityids'])
+                    if 'relationshipids' in doc.metadata['entities']:
+                        entities.setdefault('relationshipids', set()).update(doc.metadata['entities']['relationshipids'])
+
             if 'communitydetails' in doc.metadata:
                 existing_ids = {entry['id'] for entry in global_communities}
                 new_entries = [entry for entry in doc.metadata["communitydetails"] if entry['id'] not in existing_ids]
@@ -225,7 +229,7 @@ def process_documents(docs, question, messages, llm, model,chat_mode_settings):
     start_time = time.time()
 
     try:
-        formatted_docs, sources, entitydetails, communities = format_documents(docs, model)
+        formatted_docs, sources, entitydetails, communities = format_documents(docs, model,chat_mode_settings)
 
         rag_chain = get_rag_chain(llm=llm)
 

backend/src/chunkid_entities.py

@@ -190,7 +190,7 @@ def get_entities_from_chunkids(uri, username, password, database ,nodedetails,en
         elif mode == CHAT_ENTITY_VECTOR_MODE:
 
             if "entitydetails" in nodedetails and nodedetails["entitydetails"]:
-                entity_ids = [item["id"] for item in nodedetails["entitydetails"]]
+                entity_ids = [item for item in nodedetails["entitydetails"]["entityids"]]
                 logging.info(f"chunkid_entities module: Starting for entity ids: {entity_ids}")
                 result = process_entityids(driver, entity_ids)
                 if "chunk_data" in result.keys():

backend/src/shared/constants.py

@@ -521,7 +521,7 @@
 LOCAL_COMMUNITY_SEARCH_QUERY = """
 WITH collect(node) AS nodes, 
      avg(score) AS score, 
-     collect({{id: elementId(node), score: score}}) AS metadata
+     collect({{entityids: elementId(node), score: score}}) AS metadata
 
 WITH score, nodes, metadata,