From bddcd1cbeb41a2657ad23f39967b289c00adfbd0 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 9 Jul 2025 18:49:04 -0700 Subject: [PATCH] Fix bugs in Async mode with bearer authentication --- common/config.py | 10 +++- common/db/connections.py | 19 ++++-- .../Chunk_Sibling_Vector_Search.gsql | 2 +- .../Content_Similarity_Vector_Search.gsql | 2 +- .../GraphRAG_Hybrid_Vector_Search.gsql | 2 +- common/metrics/tg_proxy.py | 2 - .../prompts/openai_gpt4/generate_cypher.txt | 7 --- ecc/app/graphrag/graph_rag.py | 35 +++++------ ecc/app/main.py | 60 ++++++++++++++----- ecc/app/supportai/supportai_init.py | 11 ++-- 10 files changed, 90 insertions(+), 60 deletions(-) diff --git a/common/config.py b/common/config.py index 4346bd5..a088500 100644 --- a/common/config.py +++ b/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() diff --git a/common/db/connections.py b/common/db/connections.py index f72ee32..9e4f208 100644 --- a/common/db/connections.py +++ b/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. diff --git a/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql b/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql index c3b57db..7f84287 100644 --- a/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql +++ b/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 @@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; diff --git a/common/gsql/supportai/retrievers/Content_Similarity_Vector_Search.gsql b/common/gsql/supportai/retrievers/Content_Similarity_Vector_Search.gsql index 622e43e..bd02e34 100644 --- a/common/gsql/supportai/retrievers/Content_Similarity_Vector_Search.gsql +++ b/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 @@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; diff --git a/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql b/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql index 294e1b0..d6c8e16 100644 --- a/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql +++ b/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql @@ -22,7 +22,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set 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 diff --git a/common/metrics/tg_proxy.py b/common/metrics/tg_proxy.py index 46782c0..51ce7b5 100644 --- a/common/metrics/tg_proxy.py +++ b/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): diff --git a/common/prompts/openai_gpt4/generate_cypher.txt b/common/prompts/openai_gpt4/generate_cypher.txt index a4d920c..b186a01 100644 --- a/common/prompts/openai_gpt4/generate_cypher.txt +++ b/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. diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py index 1e1f81c..75c321e 100644 --- a/ecc/app/graphrag/graph_rag.py +++ b/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") diff --git a/ecc/app/main.py b/ecc/app/main.py index 58281c3..becdde5 100644 --- a/ecc/app/main.py +++ b/ecc/app/main.py @@ -6,14 +6,15 @@ 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, @@ -21,9 +22,8 @@ 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 diff --git a/ecc/app/supportai/supportai_init.py b/ecc/app/supportai/supportai_init.py index 4a43490..1356a8e 100644 --- a/ecc/app/supportai/supportai_init.py +++ b/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)