GML-1970 Fix bugs in Async mode with bearer authentication

common/config.py

@@ -77,12 +77,18 @@
 embedding_dimension = embedding_config.get("dimensions", 1536)
 
 if graphrag_config is None:
-    graphrag_config = {"reuse_embedding", True}
+    graphrag_config = {"reuse_embedding": True}
 if "chunker" not in graphrag_config:
     graphrag_config["chunker"] = "semantic"
 if "extractor" not in graphrag_config:
     graphrag_config["extractor"] = "llm"
 
+reuse_embedding = graphrag_config.get("reuse_embedding", True)
+doc_process_switch = graphrag_config.get("doc_process_switch", True)
+entity_extraction_switch = graphrag_config.get("entity_extraction_switch", doc_process_switch)
+entity_resolution_switch = graphrag_config.get("entity_resolution_switch", entity_extraction_switch)
+community_detection_switch = graphrag_config.get("community_detection_switch", entity_resolution_switch)
+
 if "model_name" not in llm_config or "model_name" not in llm_config["embedding_service"]:
     if "model_name" not in llm_config:
         llm_config["model_name"] = llm_config["embedding_service"]["model_name"]
@@ -133,7 +139,7 @@ def get_llm_service(llm_config) -> LLM_Model:
         password=db_config.get("password", "tigergraph"),
         gsPort=db_config.get("gsPort", "14240"),
         restppPort=db_config.get("restppPort", "9000"),
-        graphname=db_config.get("graphname", "MyGraph"),
+        graphname=db_config.get("graphname", ""),
     )
     if db_config.get("getToken"):
         conn.getToken()

common/db/connections.py

@@ -1,4 +1,5 @@
 import logging
+import asyncio
 from typing import Annotated
 
 from fastapi import Depends, HTTPException, status
@@ -31,6 +32,9 @@ def get_db_connection_id_token(
             tgCloud=True,
             sslPort=14240,
         )
+        asyncio.run(conn.customizeHeader(
+            timeout=db_config["default_timeout"] * 1000, responseSize=5000000
+        ))
     else:
         conn = TigerGraphConnection(
             host=db_config["hostname"],
@@ -39,13 +43,17 @@ def get_db_connection_id_token(
             tgCloud=True,
             sslPort=14240,
         )
-    conn.customizeHeader(
-        timeout=db_config["default_timeout"] * 1000, responseSize=5000000
-    )
-    conn = TigerGraphConnectionProxy(conn, auth_mode="id_token")
+        conn.customizeHeader(
+            timeout=db_config["default_timeout"] * 1000, responseSize=5000000
+        )
+
+        conn = TigerGraphConnectionProxy(conn, auth_mode="token")
 
     try:
-        conn.gsql("USE GRAPH " + graphname)
+        if async_conn:
+            asyncio.run(conn.gsql("USE GRAPH " + graphname))
+        else:
+            conn.gsql("USE GRAPH " + graphname)
     except HTTPError:
         LogWriter.error("Failed to connect to TigerGraph. Incorrect ID Token.")
         raise HTTPException(
@@ -161,6 +169,7 @@ def elevate_db_connection_to_token(host, username, password, graphname, async_co
 
     return conn
 
+
 def get_schema_ver(conn: TigerGraphConnectionProxy) -> int:
     """Retrieves the schema version of the graph by running an interpreted query.
 

common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql

@@ -13,7 +13,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type,
   SetAccum<VertexTypes> @@sibling_set_type;
 
   all_chunks = {v_type};
-  result = SELECT v FROM all_chunks:v POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE"));
+  result = SELECT v FROM all_chunks:v WHERE v.embedding.size() > 0 POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE"));
 
   FOREACH item IN @@topk_set DO
     @@start_set += item.v;

common/gsql/supportai/retrievers/Content_Similarity_Vector_Search.gsql

@@ -8,7 +8,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Content_Similarity_Vector_Search(STRING v_ty
   MapAccum<STRING, STRING> @@final_retrieval;
 
   vset = {v_type};
-  result = SELECT v FROM vset:v POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE"));
+  result = SELECT v FROM vset:v WHERE v.embedding.size() > 0 POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE"));
 
   FOREACH item IN @@topk_set DO
     @@start_set += item.v;

common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql

@@ -22,7 +22,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set<STRING> v_
     END;
 
     vset = {v_type.*};
-    result = SELECT v FROM vset:v POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE"));
+    result = SELECT v FROM vset:v WHERE v.embedding.size() > 0 POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE"));
   END;
 
   WHILE True LIMIT top_k DO

common/metrics/tg_proxy.py

@@ -39,8 +39,6 @@ def _req(self, method: str, url: str, authMode: str = "token", *args, **kwargs):
         if self.auth_mode == "pwd":
             return self.original_req(method, url, authMode, *args, **kwargs)
         else:
-            url = re.sub(r"/gsqlserver/", "/api/gsql-server/", url)
-            url = re.sub(r"/restpp/", "/api/restpp/", url)
             return self.original_req(method, url, "token", *args, **kwargs)
 
     def _runInstalledQuery(self, query_name, params, usePost=False):

common/prompts/openai_gpt4/generate_cypher.txt

@@ -64,13 +64,6 @@ UNION ALL
 UNWIND
 SET
 
-Here's some commonly used abbreviations:
-dt -> date
-pct -> percentage
-qty -> quantity
-lng -> longitude
-cm -> Contract Manufacturer
-
 Always make the cypher query returns the entity in the original question together with the data to be queried.
 Make sure to have correct attribute names in the OpenCypher query and not to name result aliases that are vertex or edge types, operator or function names, and other reserved keywords, always construct alias with multiple words connected with underscore.
 

ecc/app/graphrag/graph_rag.py

@@ -24,7 +24,7 @@
 )
 from pyTigerGraph import AsyncTigerGraphConnection
 
-from common.config import embedding_service, graphrag_config
+from common.config import embedding_service, graphrag_config, entity_extraction_switch, entity_resolution_switch, community_detection_switch, doc_process_switch
 from common.embeddings.base_embedding_store import EmbeddingStore
 from common.extractors.BaseExtractor import BaseExtractor
 
@@ -246,9 +246,10 @@ async def extract(
         while True:
             try:
                 item = await extract_chan.get()
-                grp.create_task(
-                    workers.extract(upsert_chan, embed_chan, extractor, conn, *item)
-                )
+                if entity_extraction_switch:
+                    grp.create_task(
+                        workers.extract(upsert_chan, embed_chan, extractor, conn, *item)
+                    )
             except ChannelClosed:
                 break
             except Exception:
@@ -476,9 +477,6 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection):
     extractor, embedding_store = await init(conn)
     init_start = time.perf_counter()
 
-    doc_process_switch = True
-    entity_resolution_switch = True
-    community_detection_switch = True
     if doc_process_switch:
         logger.info("Doc Processing Start")
         docs_chan = Channel(1)
@@ -514,17 +512,18 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection):
 
     # Type Resolution
     type_start = time.perf_counter()
-    logger.info("Type Processing Start")
-    res = await add_rels_between_types(conn)
-    if res.get("error", False):
-        logger.error(f"Error adding relationships between types: {res}")
-    else:
-        logger.info(f"Added relationships between types: {res}")
+    if entity_extraction_switch:
+        logger.info("Type Processing Start")
+        res = await add_rels_between_types(conn)
+        if res.get("error", False):
+            logger.error(f"Error adding relationships between types: {res}")
+        else:
+            logger.info(f"Added relationships between types: {res}")
     logger.info("Type Processing End")
     type_end = time.perf_counter()
+
     # Entity Resolution
     entity_start = time.perf_counter()
-
     if entity_resolution_switch:
         logger.info("Entity Processing Start")
         while not await check_embedding_rebuilt(conn, "Entity"):
@@ -551,12 +550,11 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection):
         await upsert_chan.join()
         #Resolve relationsihps
         await resolve_relationships(conn)
-
+        while not await check_all_ents_resolved(conn):
+            logger.info(f"Waiting for resolved entites to finish loading")
+            await asyncio.sleep(1)
     entity_end = time.perf_counter()
     logger.info("Entity Processing End")
-    while not await check_all_ents_resolved(conn):
-        logger.info(f"Waiting for resolved entites to finish loading")
-        await asyncio.sleep(1)
 
     # Community Detection
     community_start = time.perf_counter()
@@ -583,7 +581,6 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection):
         await embed_chan.join()
         logger.info("Join upsert_chan")
         await upsert_chan.join()
-
     community_end = time.perf_counter()
     logger.info("Community Processing End")
 

ecc/app/main.py

@@ -6,24 +6,24 @@
 import logging
 from contextlib import asynccontextmanager
 from threading import Thread
-from typing import Annotated, Callable
+from typing import Callable
 
 import asyncio
 import graphrag
 import supportai
 from eventual_consistency_checker import EventualConsistencyChecker
-from fastapi import BackgroundTasks, Depends, FastAPI, Response, status
-from fastapi.security.http import HTTPBase
+from fastapi import BackgroundTasks, Depends, FastAPI, Request, Response, status, HTTPException
+from fastapi.security.http import HTTPBasicCredentials, HTTPAuthorizationCredentials
+from base64 import b64decode
 
 from common.config import (
     db_config,
     graphrag_config,
     embedding_service,
     get_llm_service,
     llm_config,
-    security,
 )
-from common.db.connections import elevate_db_connection_to_token
+from common.db.connections import elevate_db_connection_to_token, get_db_connection_id_token
 from common.embeddings.base_embedding_store import EmbeddingStore
 from common.embeddings.tigergraph_embedding_store import TigerGraphEmbeddingStore
 from common.logs.logwriter import LogWriter
@@ -95,15 +95,15 @@ def initialize_eventual_consistency_checker(
             raise ValueError("Invalid extractor type")
 
         checker = EventualConsistencyChecker(
-            process_interval_seconds,
-            cleanup_interval_seconds,
+            graphrag_config.get("process_interval_seconds", 300),
+            graphrag_config.get("cleanup_interval_seconds", 300),
             graphname,
             embedding_service,
             embedding_store,
             index_names,
             conn,
             extractor,
-            batch_size,
+            graphrag_config.get("batch_size", 100),
         )
         consistency_checkers[graphname] = checker
 
@@ -127,6 +127,25 @@ def start_func_in_thread(f: Callable, *args, **kwargs):
     thread.start()
     LogWriter.info(f'Thread started for function: "{f.__name__}"')
 
+def auth_credentials(
+    request: Request,
+):
+    auth = request.headers.get("Authorization")
+    if not auth:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Authorization header")
+
+    scheme, credentials = auth.split(" ")
+    if scheme == "Bearer":
+        credentials = HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
+        return credentials
+
+    elif scheme == "Basic":
+        username, password = b64decode(credentials).decode().split(":")
+        credentials = HTTPBasicCredentials(username=username, password=password)
+        return credentials
+    else:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Unsupported auth scheme")
+
 
 @app.get("/")
 def root():
@@ -139,16 +158,25 @@ def consistency_status(
     graphname: str,
     ecc_method: str,
     background: BackgroundTasks,
-    credentials: Annotated[HTTPBase, Depends(security)],
     response: Response,
+    credentials = Depends(auth_credentials),
 ):
-    conn = elevate_db_connection_to_token(
-        db_config.get("hostname"),
-        credentials.username,
-        credentials.password,
-        graphname,
-        async_conn=True
-    )
+    if isinstance(credentials, HTTPBasicCredentials):
+        conn = elevate_db_connection_to_token(
+            db_config.get("hostname"),
+            credentials.username,
+            credentials.password,
+            graphname,
+            async_conn=True
+        )
+    elif isinstance(credentials, HTTPAuthorizationCredentials):
+        conn = get_db_connection_id_token(
+            graphname,
+            credentials.credentials,
+            async_conn=True
+        )
+    else:
+        raise HTTPException(status_code=401, detail="Invalid authentication credentials")
 
     asyncio.run(conn.customizeHeader(
         timeout=db_config.get("default_timeout", 300) * 1000, responseSize=5000000

ecc/app/supportai/supportai_init.py

@@ -7,7 +7,7 @@
 from aiochannel import Channel
 from pyTigerGraph import TigerGraphConnection
 
-from common.config import embedding_service
+from common.config import embedding_service, entity_extraction_switch, doc_process_switch
 from common.embeddings.base_embedding_store import EmbeddingStore
 from common.extractors.BaseExtractor import BaseExtractor
 from supportai import workers
@@ -152,9 +152,10 @@ async def extract(
     # consume task queue
     async with asyncio.TaskGroup() as sp:
         async for item in extract_chan:
-            sp.create_task(
-                workers.extract(upsert_chan, embed_chan, extractor, conn, *item)
-            )
+            if entity_extraction_switch:
+                sp.create_task(
+                    workers.extract(upsert_chan, embed_chan, extractor, conn, *item)
+                )
 
     logger.info(f"extract done")
 
@@ -181,8 +182,6 @@ async def run(
     extractor, embedding_store = await init(conn)
     init_start = time.perf_counter()
 
-    doc_process_switch = True
-
     if doc_process_switch:
         logger.info("Doc Processing Start")
         docs_chan = Channel(1)