Entity embed (#505)

* removed version of pydantic

* updating embedding of entities in post processing

* schema extraction get llm correction

* added the create_entity_embedding task in post processing

---------

Co-authored-by: kartikpersistent <[email protected]>

Author: aashipandya

backend/example.env

@@ -24,6 +24,7 @@ GCS_FILE_CACHE = "" #save the file into GCS or local, SHould be True or False
 NEO4J_USER_AGENT=""
 ENABLE_USER_AGENT = ""
 LLM_MODEL_CONFIG_model_version=""
+ENTITY_EMBEDDING="" True or False
 #examples
 LLM_MODEL_CONFIG_azure-ai-gpt-35="azure_deployment_name,azure_endpoint or base_url,azure_api_key,api_version"
 LLM_MODEL_CONFIG_azure-ai-gpt-4o="gpt-4o,https://YOUR-ENDPOINT.openai.azure.com/,azure_api_key,api_version"

backend/score.py

@@ -18,7 +18,7 @@
 from src.graphDB_dataAccess import graphDBdataAccess
 from src.graph_query import get_graph_results
 from src.chunkid_entities import get_entities_from_chunkids
-from src.post_processing import create_fulltext
+from src.post_processing import create_fulltext, create_entity_embedding
 from sse_starlette.sse import EventSourceResponse
 import json
 from typing import List, Mapping
@@ -262,7 +262,11 @@ async def post_processing(uri=Form(None), userName=Form(None), password=Form(Non
             josn_obj = {'api_name': 'post_processing/create_fulltext_index', 'db_url': uri, 'logging_time': formatted_time(datetime.now(timezone.utc))}
             logger.log_struct(josn_obj)
             logging.info(f'Full Text index created')
-
+        if os.environ.get('ENTITY_EMBEDDING').upper()=="TRUE" and "create_entity_embedding" in tasks:
+            await asyncio.to_thread(create_entity_embedding, graph)
+            josn_obj = {'api_name': 'post_processing/create_entity_embedding', 'db_url': uri, 'logging_time': formatted_time(datetime.now(timezone.utc))}
+            logger.log_struct(josn_obj)
+            logging.info(f'Entity Embeddings created')
         return create_api_response('Success', message='All tasks completed successfully')
 
     except Exception as e:

backend/src/main.py

@@ -1,6 +1,6 @@
 from langchain_community.graphs import Neo4jGraph
 from src.shared.constants import BUCKET_UPLOAD, PROJECT_ID
-from src.shared.schema_extraction import sceham_extraction_from_text
+from src.shared.schema_extraction import schema_extraction_from_text
 from dotenv import load_dotenv
 from datetime import datetime
 import logging
@@ -548,5 +548,5 @@ def populate_graph_schema_from_text(text, model, is_schema_description_cheked):
   Returns:
       data (list): list of lebels and relationTypes
   """
-  result = sceham_extraction_from_text(text, model, is_schema_description_cheked)
+  result = schema_extraction_from_text(text, model, is_schema_description_cheked)
   return {"labels": result.labels, "relationshipTypes": result.relationshipTypes}

backend/src/post_processing.py

@@ -1,7 +1,9 @@
 from neo4j import GraphDatabase
 import logging
 import time
-
+from langchain_community.graphs import Neo4jGraph
+import os
+from src.shared.common_fn import load_embedding_model
 
 DROP_INDEX_QUERY = "DROP INDEX entities IF EXISTS;"
 LABELS_QUERY = "CALL db.labels()"
@@ -55,4 +57,32 @@ def create_fulltext(uri, username, password, database):
     finally:
         driver.close()
         logging.info("Driver closed.")
-        logging.info(f"Process completed in {time.time() - start_time:.2f} seconds.")
+        logging.info(f"Process completed in {time.time() - start_time:.2f} seconds.")
+
+        
+def create_entity_embedding(graph:Neo4jGraph):
+    rows = fetch_entities_for_embedding(graph)
+    for i in range(0, len(rows), 1000):
+        update_embeddings(rows[i:i+1000],graph)
+            
+def fetch_entities_for_embedding(graph):
+    query = """
+                MATCH (e)
+                WHERE NOT (e:Chunk OR e:Document) AND e.embedding IS NULL AND e.id IS NOT NULL
+                RETURN elementId(e) AS elementId, e.id + " " + coalesce(e.description, "") AS text
+                """
+    result = graph.query(query)           
+    return [{"elementId": record["elementId"], "text": record["text"]} for record in result]
+
+def update_embeddings(rows, graph):
+    embedding_model = os.getenv('EMBEDDING_MODEL')
+    embeddings, dimension = load_embedding_model(embedding_model)
+    logging.info(f"update embedding for entities")
+    for row in rows:
+        row['embedding'] = embeddings.embed_query(row['text'])                        
+    query = """
+      UNWIND $rows AS row
+      MATCH (e) WHERE elementId(e) = row.elementId
+      CALL db.create.setNodeVectorProperty(e, "embedding", row.embedding)
+      """  
+    return graph.query(query,params={'rows':rows})          

backend/src/shared/schema_extraction.py

@@ -25,9 +25,9 @@ class Schema(BaseModel):
     "Only return the string types for nodes and relationships, don't return attributes."
 )
 
-def sceham_extraction_from_text(input_text:str, model:str, is_schema_description_cheked:bool):
+def schema_extraction_from_text(input_text:str, model:str, is_schema_description_cheked:bool):
 
-    llm = get_llm(MODEL_VERSIONS[model])
+    llm, model_name = get_llm(model)
     if is_schema_description_cheked:
         schema_prompt = PROMPT_TEMPLATE_WITH_SCHEMA
     else:

frontend/src/utils/Constants.ts

@@ -161,7 +161,7 @@ export const ChatModeOptions = [
   { Icon: 'abc', value: 'vector' },
 ];
 
-export const taskParam: string[] = ['update_similarity_graph', 'create_fulltext_index'];
+export const taskParam: string[] = ['update_similarity_graph', 'create_fulltext_index','create_entity_embedding'];
 
 export const nvlOptions: NvlOptions = {
   allowDynamicMinZoom: true,