neo4j-labs
diff --git a/‎README.md‎
Lines changed: 64 additions & 53 deletions b/‎README.md‎
Lines changed: 64 additions & 53 deletions
diff --git a/‎backend/README.md‎
Lines changed: 0 additions & 2 deletions b/‎backend/README.md‎
Lines changed: 0 additions & 2 deletions
diff --git a/‎backend/example.env‎
Lines changed: 11 additions & 10 deletions b/‎backend/example.env‎
Lines changed: 11 additions & 10 deletions
diff --git a/‎backend/score.py‎
Lines changed: 128 additions & 23 deletions b/‎backend/score.py‎
Lines changed: 128 additions & 23 deletions
diff --git a/‎backend/src/QA_integration.py‎
Lines changed: 22 additions & 24 deletions b/‎backend/src/QA_integration.py‎
Lines changed: 22 additions & 24 deletions
diff --git a/‎backend/src/communities.py‎
Lines changed: 10 additions & 11 deletions b/‎backend/src/communities.py‎
Lines changed: 10 additions & 11 deletions
diff --git a/‎backend/src/entities/source_extract_params.py‎
Lines changed: 6 additions & 1 deletion b/‎backend/src/entities/source_extract_params.py‎
Lines changed: 6 additions & 1 deletion
diff --git a/‎backend/src/entities/source_node.py‎
Lines changed: 1 addition & 0 deletions b/‎backend/src/entities/source_node.py‎
Lines changed: 1 addition & 0 deletions
@@ -64,8 +64,6 @@ Update the environment variable in `.env` file. Refer example.env in backend fol
 
 `OPENAI_API_KEY`: Open AI key to use incase of openai embeddings
 
-`EMBEDDING_MODEL` : "all-MiniLM-L6-v2" or "openai" or "vertexai"
-
 `NEO4J_URI` : Neo4j URL
 
 `NEO4J_USERNAME` : Neo4J database username
 
@@ -1,6 +1,6 @@
 OPENAI_API_KEY = "" #Mandatory- Default_Value = "openai_api_key"  #This is required if you are using openai embedding model
-EMBEDDING_MODEL = "" #OPTIONAL- Default_Value =""  #this can be openai or vertexai or by default all-MiniLM-L6-v2
-RAGAS_EMBEDDING_MODEL = "" #OPTIONAL- Default_Value ="openai"  #Keep blank if you want to use all-MiniLM-L6-v2 for ragas embeddings
+EMBEDDING_MODEL = "" #OPTIONAL- Default_Value ="all-MiniLM-L6-v2" -- Embedding model to be used for text embedding
+EMBEDDING_PROVIDER = "" #OPTIONAL- Default_Value ="sentence-transformer" -- Embedding provider to be used for text embedding
 IS_EMBEDDING = "" #OPTIONAL- Default_Value ="True" --Flag to enable text embedding
 BUCKET_UPLOAD_FILE = "" #OPTIONAL- Default_Value ="gcs bucket name" -- use the gcs bucket to upload local file to gcs cloud
 BUCKET_FAILED_FILE = "" #OPTIONAL- Default_Value ="gcs bucket name" -- use the gcs bucket for failed file while extraction
@@ -20,7 +20,7 @@ DUPLICATE_SCORE_VALUE = "" #OPTIONAL- Default_Value = "0.97" -- Node score value
 DUPLICATE_TEXT_DISTANCE = "" #OPTIONAL- Default_Value = "3" --This value used to find distance for all node pairs in the graph and calculated based on node properties
 DEFAULT_DIFFBOT_CHAT_MODEL="" #OPTIONAL- Default_Value = "openai_gpt_5_mini"  #whichever model specified here , need to add config for that model in below format)
 GRAPH_CLEANUP_MODEL="" #OPTIONAL- Default_Value = "openai_gpt_5_mini" -- Model name to clean-up graph in post processing
-BEDROCK_EMBEDDING_MODEL="" #OPTIONAL - Default_Value = "model_name,aws_access_key,aws_secret_key,region_name" -- If want to use bedrock embedding #model_name="amazon.titan-embed-text-v1"
+BEDROCK_EMBEDDING_MODEL_KEY="" #OPTIONAL - Default_Value = "aws_access_key,aws_secret_key,region_name" -- If want to use bedrock embedding
 YOUTUBE_TRANSCRIPT_PROXY="" #Mandatory --Proxy key required to process youtube video for getting transcript --Sample Value ="https://user:pass@domain:port"
 EFFECTIVE_SEARCH_RATIO="" #OPTIONAL- Default_Value = "5"
 MAX_TOKEN_CHUNK_SIZE="" #OPTIONAL- Default_Value = "10000" #Max token used to process/extract the file content.
@@ -44,10 +44,11 @@ LLM_MODEL_CONFIG_bedrock_nova_micro_v1="amazon.nova-micro-v1:0,aws_access_key,aw
 LLM_MODEL_CONFIG_bedrock_nova_lite_v1="amazon.nova-lite-v1:0,aws_access_key,aws_secret_key,region_name"
 LLM_MODEL_CONFIG_bedrock_nova_pro_v1="amazon.nova-pro-v1:0,aws_access_key,aws_secret_key,region_name"
 LLM_MODEL_CONFIG_ollama_llama3="llama3_model_name,model_local_url"
-TRACK_TOKEN_USAGE="true" #Add this if you want to track token usage
-DAILY_TOKENS_LIMIT="250000" #Mandatory if TRACK_TOKEN_USAGE is true
-MONTHLY_TOKENS_LIMIT="1000000" #Mandatory if TRACK_TOKEN_USAGE is true
-TOKEN_TRACKER_DB_URI=""  #Mandatory if TRACK_TOKEN_USAGE is true
-TOKEN_TRACKER_DB_USERNAME="" #Mandatory if TRACK_TOKEN_USAGE is true
-TOKEN_TRACKER_DB_PASSWORD="" #Mandatory if TRACK_TOKEN_USAGE is true
-TOKEN_TRACKER_DB_DATABASE="" #Mandatory if TRACK_TOKEN_USAGE is true
+TRACK_USER_USAGE="true" #Add this if you want to track token usage
+DAILY_TOKENS_LIMIT="250000" #Mandatory if TRACK_USER_USAGE is true
+MONTHLY_TOKENS_LIMIT="1000000" #Mandatory if TRACK_USER_USAGE is true
+TOKEN_TRACKER_DB_URI=""  #Mandatory if TRACK_USER_USAGE is true
+TOKEN_TRACKER_DB_USERNAME="" #Mandatory if TRACK_USER_USAGE is true
+TOKEN_TRACKER_DB_PASSWORD="" #Mandatory if TRACK_USER_USAGE is true
+TOKEN_TRACKER_DB_DATABASE="" #Mandatory if TRACK_TOKEN_USAGE is true
+AUTHENTICATION_REQUIRED="" #OPTIONAL- Default_Value = "False" -- Enable authentication for API access
@@ -39,8 +39,6 @@
 )
 load_dotenv() 
 
-EMBEDDING_MODEL = get_value_from_env("EMBEDDING_MODEL", "sentence_transformer")
-
 class SessionChatHistory:
     history_dict = {}
 
@@ -291,7 +289,7 @@ def retrieve_documents(doc_retriever, messages):
 
     return docs,transformed_question
 
-def create_document_retriever_chain(llm, retriever):
+def create_document_retriever_chain(llm, retriever, embedding_provider, embedding_model):
     try:
         logging.info("Starting to create document retriever chain")
 
@@ -305,7 +303,7 @@ def create_document_retriever_chain(llm, retriever):
         output_parser = StrOutputParser()
 
         splitter = TokenTextSplitter(chunk_size=CHAT_DOC_SPLIT_SIZE, chunk_overlap=0)
-        EMBEDDING_FUNCTION , _ = load_embedding_model(EMBEDDING_MODEL) 
+        EMBEDDING_FUNCTION , _ = load_embedding_model(embedding_provider, embedding_model) 
         embeddings_filter = EmbeddingsFilter(
             embeddings=EMBEDDING_FUNCTION,
             similarity_threshold=CHAT_EMBEDDING_FILTER_SCORE_THRESHOLD
@@ -334,7 +332,7 @@ def create_document_retriever_chain(llm, retriever):
         logging.error(f"Error creating document retriever chain: {e}", exc_info=True)
         raise
 
-def initialize_neo4j_vector(graph, chat_mode_settings):
+def initialize_neo4j_vector(graph, chat_mode_settings, embedding_provider, embedding_model):
     try:
         retrieval_query = chat_mode_settings.get("retrieval_query")
         index_name = chat_mode_settings.get("index_name")
@@ -346,7 +344,7 @@ def initialize_neo4j_vector(graph, chat_mode_settings):
 
         if not retrieval_query or not index_name:
             raise ValueError("Required settings 'retrieval_query' or 'index_name' are missing.")
-        EMBEDDING_FUNCTION , _ = load_embedding_model(EMBEDDING_MODEL) 
+        EMBEDDING_FUNCTION , _ = load_embedding_model(embedding_provider, embedding_model) 
         if keyword_index:
             neo_db = Neo4jVector.from_existing_graph(
                 embedding=EMBEDDING_FUNCTION,
@@ -397,10 +395,10 @@ def create_retriever(neo_db, document_names, chat_mode_settings,search_k, score_
         logging.info(f"Successfully created retriever with search_k={search_k}, score_threshold={score_threshold}")
     return retriever
 
-def get_neo4j_retriever(graph, document_names,chat_mode_settings, score_threshold=CHAT_SEARCH_KWARG_SCORE_THRESHOLD):
+def get_neo4j_retriever(graph, document_names,chat_mode_settings, score_threshold=CHAT_SEARCH_KWARG_SCORE_THRESHOLD, embedding_provider=None, embedding_model=None):
     try:
 
-        neo_db = initialize_neo4j_vector(graph, chat_mode_settings)
+        neo_db = initialize_neo4j_vector(graph, chat_mode_settings, embedding_provider, embedding_model)
         # document_names= list(map(str.strip, json.loads(document_names)))
         search_k = chat_mode_settings["top_k"]
         ef_ratio = get_value_from_env("EFFECTIVE_SEARCH_RATIO", 5, "int")
@@ -412,7 +410,7 @@ def get_neo4j_retriever(graph, document_names,chat_mode_settings, score_threshol
         raise Exception(f"An error occurred while retrieving the Neo4jVector index or creating the retriever. Please drop and create a new vector index '{index_name}': {e}") from e 
 
 
-def setup_chat(model, graph, document_names, chat_mode_settings):
+def setup_chat(model, graph, document_names, chat_mode_settings, embedding_provider, embedding_model):
     start_time = time.time()
     try:
         if model == "diffbot":
@@ -421,8 +419,8 @@ def setup_chat(model, graph, document_names, chat_mode_settings):
         llm, model_name, _ = get_llm(model=model)
         logging.info(f"Model called in chat: {model} (version: {model_name})")
 
-        retriever = get_neo4j_retriever(graph=graph, chat_mode_settings=chat_mode_settings, document_names=document_names)
-        doc_retriever = create_document_retriever_chain(llm, retriever)
+        retriever = get_neo4j_retriever(graph=graph, chat_mode_settings=chat_mode_settings, document_names=document_names, embedding_provider=embedding_provider, embedding_model=embedding_model)
+        doc_retriever = create_document_retriever_chain(llm, retriever, embedding_provider, embedding_model)
 
         chat_setup_time = time.time() - start_time
         logging.info(f"Chat setup completed in {chat_setup_time:.2f} seconds")
@@ -433,23 +431,23 @@ def setup_chat(model, graph, document_names, chat_mode_settings):
 
     return llm, doc_retriever, model_name
 
-def process_chat_response(messages, history, question, model, graph, document_names, chat_mode_settings, email=None, uri=None):
+def process_chat_response(messages, history, question, model, graph, document_names, chat_mode_settings, email=None, uri=None, embedding_provider=None, embedding_model=None):
     try:
-        if get_value_from_env("TRACK_TOKEN_USAGE", "false", "bool"):
-            try:
-                track_token_usage(email, uri, 0, model)
-            except LLMGraphBuilderException as e:
-                logging.error(str(e))
-                raise RuntimeError(str(e))
-        llm, doc_retriever, model_version = setup_chat(model, graph, document_names, chat_mode_settings)
+        # if get_value_from_env("TRACK_USER_USAGE", "false", "bool"):
+        #     try:
+        #         track_token_usage(email, uri, 0, model, operation_type="chat")
+        #     except LLMGraphBuilderException as e:
+        #         logging.error(str(e))
+        #         raise RuntimeError(str(e))
+        llm, doc_retriever, model_version = setup_chat(model, graph, document_names, chat_mode_settings,embedding_provider, embedding_model)
 
         docs,transformed_question = retrieve_documents(doc_retriever, messages)  
 
         if docs:
             content, result, total_tokens,formatted_docs = process_documents(docs, question, messages, llm, model, chat_mode_settings)
-            if get_value_from_env("TRACK_TOKEN_USAGE", "false", "bool"):
-                latest_token = track_token_usage(email=email, uri=uri, usage=total_tokens, last_used_model=model)
-                logging.info(f"Total token usage {latest_token} for user {email} ")
+            # if get_value_from_env("TRACK_USER_USAGE", "false", "bool"):
+            #     latest_token = track_token_usage(email=email, uri=uri, usage=total_tokens, last_used_model=model, operation_type="chat")
+            #     logging.info(f"Total token usage {latest_token} for user {email} ")
         else:
             content = "I couldn't find any relevant documents to answer your question."
             result = {"sources": list(), "nodedetails": list(), "entities": list()}
@@ -664,7 +662,7 @@ def get_chat_mode_settings(mode,settings_map=CHAT_MODE_CONFIG_MAP):
 
     return chat_mode_settings
 
-def QA_RAG(graph, model, question, document_names, session_id, mode, write_access=True, email=None, uri=None):
+def QA_RAG(graph, model, question, document_names, session_id, mode, write_access=True, email=None, uri=None, embedding_provider=None, embedding_model=None):
     logging.info(f"Chat Mode: {mode}")
 
     history = create_neo4j_chat_message_history(graph, session_id, write_access)
@@ -695,7 +693,7 @@ def QA_RAG(graph, model, question, document_names, session_id, mode, write_acces
                 "user": "chatbot"
             }
         else:
-            result = process_chat_response(messages,history, question, model, graph, document_names,chat_mode_settings, email, uri)
+            result = process_chat_response(messages,history, question, model, graph, document_names,chat_mode_settings, email, uri, embedding_provider, embedding_model)
 
     result["session_id"] = session_id
 
 
@@ -316,9 +316,9 @@ def create_community_summaries(gds, model, email, uri):
     token_usage = 0
     try:
         #pre check if user allowed to create community summaries
-        if get_value_from_env("TRACK_TOKEN_USAGE", "false", "bool"):
+        if get_value_from_env("TRACK_USER_USAGE", "false", "bool"):
             try:
-                track_token_usage(email, uri, 0, model)
+                track_token_usage(email, uri, 0, model,operation_type="precheck")
             except LLMGraphBuilderException as e:
                 logging.error(str(e))
                 raise RuntimeError(str(e))
@@ -361,21 +361,20 @@ def create_community_summaries(gds, model, email, uri):
 
     finally:
        try:
-           if get_value_from_env("TRACK_TOKEN_USAGE", "false", "bool"):
+           if get_value_from_env("TRACK_USER_USAGE", "false", "bool"):
                if callback_handler:
                    usage = callback_handler.report()
                    token_usage = usage.get("total_tokens", 0)
                    if email and token_usage > 0:
                        email = email.strip().replace('"', '')
-                       latest_token = track_token_usage(email, uri, token_usage, model)
+                       latest_token = track_token_usage(email, uri, token_usage, model,operation_type="community_summary")
                        logging.info(f"In community : Total token usage {latest_token} for user {email} ")
        except Exception as err:
            logging.warning(f"Failed to track token usage: {err}")
 
-def create_community_embeddings(gds):
+def create_community_embeddings(gds, embedding_provider, embedding_model):
     try:
-        embedding_model = get_value_from_env("EMBEDDING_MODEL","sentence_transformer")
-        embeddings, dimension = load_embedding_model(embedding_model)
+        embeddings, dimension = load_embedding_model(embedding_provider, embedding_model)
         logging.info(f"Embedding model '{embedding_model}' loaded successfully.")
 
         logging.info("Fetching community details.")
@@ -467,7 +466,7 @@ def create_fulltext_index(gds, index_type):
         logging.error("An error occurred while creating the full-text index.", exc_info=True)
         logging.error(f"Error details: {str(e)}")
 
-def create_community_properties(gds, model, email, uri):
+def create_community_properties(gds, model, email, uri, embedding_provider, embedding_model):
     commands = [
         (CREATE_COMMUNITY_CONSTRAINT, "created community constraint to the graph."),
         (CREATE_COMMUNITY_LEVELS, "Successfully created community levels."),
@@ -484,7 +483,7 @@ def create_community_properties(gds, model, email, uri):
         create_community_summaries(gds, model, email, uri)
         logging.info("Successfully created community summaries.")
 
-        embedding_dimension = create_community_embeddings(gds)
+        embedding_dimension = create_community_embeddings(gds, embedding_provider, embedding_model)
         logging.info("Successfully created community embeddings.")
 
         create_vector_index(gds=gds,index_type=ENTITY_VECTOR_INDEX_NAME,embedding_dimension=embedding_dimension)
@@ -518,7 +517,7 @@ def clear_communities(gds):
         raise
 
 
-def create_communities(uri, username, password, database,email=None,model=COMMUNITY_CREATION_DEFAULT_MODEL):
+def create_communities(uri, username, password, database,email=None,model=COMMUNITY_CREATION_DEFAULT_MODEL, embedding_provider=None, embedding_model=None):
     try:
         gds = get_gds_driver(uri, username, password, database)
         clear_communities(gds)
@@ -527,7 +526,7 @@ def create_communities(uri, username, password, database,email=None,model=COMMUN
         write_communities_success = write_communities(gds, graph_project)
         if write_communities_success:
             logging.info("Starting Community properties creation process.")
-            create_community_properties(gds,model,email,uri)
+            create_community_properties(gds,model,email,uri, embedding_provider, embedding_model)
             logging.info("Communities creation process completed successfully.")
         else:
             logging.warning("Failed to write communities. Constraint was not applied.")
 
@@ -23,7 +23,8 @@ class SourceScanExtractParams(BaseModel):
     language: Optional[str] = Field(None, description="Language")
     retry_condition: Optional[str] = Field(None, description="Retry condition")
     additional_instructions: Optional[str] = Field(None, description="Additional instructions")
-
+    embedding_provider: Optional[str] = Field(None, description="Embedding provider")
+    embedding_model: Optional[str] = Field(None, description="Embedding model")
 
 def get_source_scan_extract_params(
     source_url: Optional[str] = Form(None),
@@ -46,6 +47,8 @@ def get_source_scan_extract_params(
     language: Optional[str] = Form(None),
     retry_condition: Optional[str] = Form(None),
     additional_instructions: Optional[str] = Form(None),
+    embedding_provider: Optional[str] = Form(None),
+    embedding_model: Optional[str] = Form(None),
 ) -> SourceScanExtractParams:
     return SourceScanExtractParams(
         source_url=source_url,
@@ -68,4 +71,6 @@ def get_source_scan_extract_params(
         language=language,
         retry_condition=retry_condition,
         additional_instructions=additional_instructions,
+        embedding_provider=embedding_provider,
+        embedding_model=embedding_model
     )
@@ -31,3 +31,4 @@ class sourceNode:
     access_token:str=None
     retry_condition:str=None
     token_usage:int=None
+    embedding_model:str=None