diff --git a/chat-history/Dockerfile b/chat-history/Dockerfile index b1fbc35..aa4f5ed 100644 --- a/chat-history/Dockerfile +++ b/chat-history/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.22.3 as builder +FROM golang:1.22.3 AS builder WORKDIR /app COPY . . diff --git a/common/config.py b/common/config.py index 010c319..703d3f8 100644 --- a/common/config.py +++ b/common/config.py @@ -91,6 +91,14 @@ raise Exception("embedding_service is not found in llm_config") embedding_dimension = embedding_config.get("dimensions", 1536) +# Get context window size from llm_config +# <=0 means unlimited tokens (no truncation), otherwise use the specified limit +if "token_limit" in llm_config: + if "token_limit" not in completion_config: + completion_config["token_limit"] = llm_config["token_limit"] + if "token_limit" not in embedding_config: + embedding_config["token_limit"] = llm_config["token_limit"] + # Get multimodal_service config (optional, for vision/image tasks) multimodal_config = llm_config.get("multimodal_service") diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py index 15900dd..8a24265 100644 --- a/common/embeddings/embedding_services.py +++ b/common/embeddings/embedding_services.py @@ -11,6 +11,7 @@ from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter from common.metrics.prometheus_metrics import metrics +from common.utils.token_calculator import TokenCalculator logger = logging.getLogger(__name__) @@ -32,6 +33,7 @@ def __init__(self, config: dict, model_name: str): self.embeddings = None self.model_name = model_name self.dimensions = config.get("dimensions", 1536) + self.token_calculator = TokenCalculator(token_limit=config.get("token_limit", 8192), model_name=model_name) LogWriter.info( f"request_id={req_id_cv.get()} instantiated AI model_name={model_name} with dimensions={self.dimensions}" ) @@ -51,6 +53,13 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]: try: LogWriter.info(f"request_id={req_id_cv.get()} ENTRY embed_documents()") + + if not self.token_calculator.is_unlimited_tokens(): + max_context_tokens = self.token_calculator.get_max_context_tokens() + if any(len(text) > max_context_tokens for text in texts): + if any(self.token_calculator.count_tokens(text) > max_context_tokens for text in texts): + texts = [self.token_calculator.truncate_to_token_limit(text, max_context_tokens) for text in texts] + if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings): docs = self.embeddings.embed_documents(texts, output_dimensionality=self.dimensions) else: @@ -85,6 +94,13 @@ def embed_query(self, question: str) -> List[float]: logger.debug_pii( f"request_id={req_id_cv.get()} embed_query() embedding question={question}" ) + + if not self.token_calculator.is_unlimited_tokens(): + max_context_tokens = self.token_calculator.get_max_context_tokens() + if len(question) > max_context_tokens: + if self.token_calculator.count_tokens(question) > max_context_tokens: + question = self.token_calculator.truncate_to_token_limit(question, max_context_tokens) + if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings): query_embedding = self.embeddings.embed_query(question, output_dimensionality=self.dimensions) else: @@ -117,6 +133,12 @@ async def aembed_query(self, question: str) -> List[float]: # try: LogWriter.info(f"request_id={req_id_cv.get()} ENTRY aembed_query()") logger.debug_pii(f"aembed_query() embedding question={question}") + if not self.token_calculator.is_unlimited_tokens(): + max_context_tokens = self.token_calculator.get_max_context_tokens() + if len(question) > max_context_tokens: + if self.token_calculator.count_tokens(question) > max_context_tokens: + question = self.token_calculator.truncate_to_token_limit(question, max_content_tokens) + if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings): query_embedding = await self.embeddings.aembed_query(question, output_dimensionality=self.dimensions) else: @@ -140,7 +162,7 @@ class AzureOpenAI_Ada002(EmbeddingModel): """Azure OpenAI Ada-002 Embedding Model""" def __init__(self, config): - super().__init__(config, model_name=config.get("model_name", "OpenAI ada-002")) + super().__init__(config, model_name=config.get("model_name", "text-embedding-3-small")) from langchain_openai import AzureOpenAIEmbeddings self.embeddings = AzureOpenAIEmbeddings(model=self.model_name, dimensions=self.dimensions, deployment=config["azure_deployment"]) @@ -203,7 +225,7 @@ def __init__(self, config): "AWS_SECRET_ACCESS_KEY" ], ) - self.embeddings = BedrockEmbeddings(client=client) + self.embeddings = BedrockEmbeddings(client=client, model_id=self.model_name) class Ollama_Embedding(EmbeddingModel): @@ -212,13 +234,12 @@ class Ollama_Embedding(EmbeddingModel): def __init__(self, config): from langchain_ollama import OllamaEmbeddings - super().__init__(config=config, model_name=config.get("model_name", "llama2")) + super().__init__(config=config, model_name=config.get("model_name", "llama3")) # Get Ollama configuration from config base_url = config.get("base_url", "http://localhost:11434") - model_name = config.get("model_name", "llama3") self.embeddings = OllamaEmbeddings( - model=model_name, + model=self.model_name, base_url=base_url ) diff --git a/common/embeddings/tigergraph_embedding_store.py b/common/embeddings/tigergraph_embedding_store.py index a560233..0baf79b 100644 --- a/common/embeddings/tigergraph_embedding_store.py +++ b/common/embeddings/tigergraph_embedding_store.py @@ -103,7 +103,7 @@ def install_vector_queries(self): need_install = True logger.info(f"Done creating vector query {q_name} with status {q_res}") #TBD - if need_install and False: + if need_install: logger.info(f"Installing supportai queries all together") query_res = self.conn.gsql( """USE GRAPH {}\nINSTALL QUERY ALL\n""".format( diff --git a/common/gsql/graphrag/StreamChunkContent.gsql b/common/gsql/graphrag/StreamChunkContent.gsql new file mode 100644 index 0000000..6b488eb --- /dev/null +++ b/common/gsql/graphrag/StreamChunkContent.gsql @@ -0,0 +1,8 @@ +CREATE OR REPLACE DISTRIBUTED QUERY StreamChunkContent(Vertex chunk) { + DocChunk = {chunk}; + + // Get the document's content and mark it as processed + ChunkContent = SELECT c FROM DocChunk:d -(HAS_CONTENT)-> Content:c + POST-ACCUM d.epoch_processed = datetime_to_epoch(now()); + PRINT ChunkContent; +} diff --git a/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql b/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql index 99525e7..dc8d520 100644 --- a/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql +++ b/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql @@ -17,6 +17,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, LIST query_vector, UINT top_k=5, UINT lookback=3, INT lookahead=3, BOOL verbose = False) { TYPEDEF TUPLE VertexTypes; + TYPEDEF TUPLE EdgeTypes; TYPEDEF tuple Similarity_Results; HeapAccum(top_k, score DESC) @@topk_set; MapAccum> @@verbose_info; @@ -27,6 +28,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, SetAccum @@start_set; SetAccum @@start_set_type; SetAccum @@sibling_set_type; + SetAccum @@edges; all_chunks = {v_type}; 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")); @@ -44,6 +46,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, start = SELECT tgt FROM start:s -(IS_AFTER)-> :tgt ACCUM tgt.@touched += TRUE, @@sibling_set_type += VertexTypes(tgt, tgt.type), + @@edges += EdgeTypes(s, tgt), FOREACH (key, val) IN s.@distances DO tgt.@distances += (key -> -1*i) END; @@ -54,6 +57,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, start = SELECT tgt FROM start:s -(reverse_IS_AFTER)-> :tgt ACCUM tgt.@touched += TRUE, @@sibling_set_type += VertexTypes(tgt, tgt.type), + @@edges += EdgeTypes(s, tgt), FOREACH (key, val) IN s.@distances DO tgt.@distances += (key -> i) END; @@ -71,7 +75,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, @@verbose_info += ("selected_set" -> @@sibling_set_type); - PRINT @@final_retrieval as final_retrieval; + PRINT @@final_retrieval as final_retrieval, @@edges as edges; IF verbose THEN PRINT @@verbose_info as verbose; diff --git a/common/gsql/supportai/retrievers/GraphRAG_Community_Vector_Search.gsql b/common/gsql/supportai/retrievers/GraphRAG_Community_Vector_Search.gsql index 46a433f..7777bca 100644 --- a/common/gsql/supportai/retrievers/GraphRAG_Community_Vector_Search.gsql +++ b/common/gsql/supportai/retrievers/GraphRAG_Community_Vector_Search.gsql @@ -15,11 +15,13 @@ */ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Community_Vector_Search(LIST query_vector, INT community_level=2, INT top_k = 3, BOOL with_chunk = true, BOOL with_doc = false, BOOL verbose = false) { + TYPEDEF TUPLE EdgeTypes; MapAccum> @@final_retrieval; MapAccum> @@verbose_info; SetAccum @context; SetAccum @children; SetAccum @@start_set; + SetAccum @@edges; filtered_comms = SELECT c FROM Community:c WHERE c.iteration == community_level and length(c.description) > 0; start_comms = vectorSearch({Community.embedding}, query_vector, top_k, {candidate_set: filtered_comms}); @@ -43,11 +45,11 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Community_Vector_Search(LIST IF with_doc THEN related_chunks = SELECT c FROM Content:c -()- DocumentChunk:dc -(CONTAINS_ENTITY>)- Entity:v -(RESOLVES_TO>)- ResolvedEntity:r -(IN_COMMUNITY>)- selected_comms:m - ACCUM m.@context += c.text, m.@children += d + ACCUM m.@context += c.text, m.@children += d, @@edges += EdgeTypes(m, d) POST-ACCUM @@verbose_info += ("related_chunks" -> m.@children); ELSE related_chunks = SELECT c FROM Content:c -()- Entity:v -(RESOLVES_TO>)- ResolvedEntity:r -(IN_COMMUNITY>)- selected_comms:m - ACCUM m.@context += c.text, m.@children += d + ACCUM m.@context += c.text, m.@children += d, @@edges += EdgeTypes(m, d) POST-ACCUM @@verbose_info += ("related_chunks" -> m.@children); END; END; @@ -56,7 +58,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Community_Vector_Search(LIST ACCUM s.@context += s.description, s.@context += c.@context POST-ACCUM(s) @@final_retrieval += (s -> s.@context); - PRINT @@final_retrieval as final_retrieval; + PRINT @@final_retrieval as final_retrieval, @@edges as edges; IF verbose THEN PRINT @@verbose_info as verbose; diff --git a/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql b/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql index d8c483d..b2af498 100644 --- a/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql +++ b/common/gsql/supportai/retrievers/GraphRAG_Hybrid_Vector_Search.gsql @@ -17,12 +17,14 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set v_types, LIST query_vector, UINT top_k=5, UINT num_hops=3, UINT num_seen_min=1, BOOL chunk_only = False, BOOL doc_only = False, BOOL verbose = False) { TYPEDEF TUPLE VertexTypes; + TYPEDEF TUPLE EdgeTypes; TYPEDEF tuple Similarity_Results; HeapAccum(top_k, score DESC) @@topk_set; SetAccum @@start_set; SetAccum @@start_set_type; SetAccum @@tmp_set; SetAccum @@selected_set_type; + SetAccum @@edges; SumAccum @visited; SumAccum @num_times_seen; MapAccum> @@result_set; @@ -89,7 +91,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set v_ IF NOT (chunk_only OR doc_only) THEN FOREACH v IN s.@parents DO IF v IN @@start_set THEN - @@result_set += (v -> s.@context) + @@result_set += (v -> s.@context), @@edges += EdgeTypes(v, s) END END END; @@ -99,11 +101,12 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set v_ IF doc_only THEN @@selected_set_type.clear(); res = SELECT s FROM doc_chunks:s -(reverse_HAS_CHILD>:e1)- Document:d -(HAS_CONTENT>:e2)- Content:t - ACCUM s.@context += t.text, @@selected_set_type += VertexTypes(d, d.type) - POST-ACCUM(s) - FOREACH v IN s.@parents DO + ACCUM d.@context += t.text, d.@parents += s.@parents, @@selected_set_type += VertexTypes(d, d.type) + POST-ACCUM(d) + @@result_set += (d -> d.@context), + FOREACH v IN d.@parents DO IF v IN @@start_set THEN - @@result_set += (v -> s.@context) + @@edges += EdgeTypes(v, d) END END; @@verbose_info += ("selected_set" -> @@selected_set_type); @@ -111,15 +114,16 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set v_ res = SELECT s FROM doc_chunks:s -(HAS_CONTENT>:e)- Content:t ACCUM s.@context += t.text POST-ACCUM(s) + @@result_set += (s -> s.@context), FOREACH v IN s.@parents DO IF v IN @@start_set THEN - @@result_set += (v -> s.@context) + @@edges += EdgeTypes(v, s) END END; @@verbose_info += ("selected_set" -> @@selected_set_type); END; - PRINT @@result_set as final_retrieval; + PRINT @@result_set as final_retrieval, @@edges as edges; IF verbose THEN PRINT @@verbose_info as verbose; diff --git a/common/metrics/tg_proxy.py b/common/metrics/tg_proxy.py index 51ce7b5..804d66f 100644 --- a/common/metrics/tg_proxy.py +++ b/common/metrics/tg_proxy.py @@ -36,17 +36,18 @@ def hooked(*args, **kwargs): def _req(self, method: str, url: str, authMode: str = "token", *args, **kwargs): # we always use token auth # always use proxy endpoint in GUI for restpp and gsql + # logger.info(f"Requesting {method} {url} with authMode {authMode} and args {args} and kwargs {kwargs}") if self.auth_mode == "pwd": return self.original_req(method, url, authMode, *args, **kwargs) else: return self.original_req(method, url, "token", *args, **kwargs) - def _runInstalledQuery(self, query_name, params, usePost=False): + def _runInstalledQuery(self, query_name, params, sizeLimit=None, usePost=False): start_time = time.time() metrics.tg_inprogress_requests.labels(query_name=query_name).inc() try: restppid = self._tg_connection.runInstalledQuery( - query_name, params, runAsync=True, usePost=usePost + query_name, params, runAsync=True, usePost=usePost, sizeLimit=sizeLimit ) LogWriter.info( f"request_id={req_id_cv.get()} query {query_name} started with RESTPP ID {restppid}" @@ -54,9 +55,8 @@ def _runInstalledQuery(self, query_name, params, usePost=False): result = None while not result: ret = self._tg_connection.checkQueryStatus(restppid) - LogWriter.info(f"Ret: {ret}") if not ret: - time.sleep(0.1) + time.sleep(0.5) continue if ret[0]["status"].lower() == "success": LogWriter.info( @@ -72,7 +72,7 @@ def _runInstalledQuery(self, query_name, params, usePost=False): result = None if i >= 9: raise e - time.sleep(0.1) + time.sleep(0.5) elif ret[0]["status"].lower() == "aborted": LogWriter.error( f"request_id={req_id_cv.get()} query {query_name} with RESTPP ID {restppid} aborted" @@ -87,6 +87,8 @@ def _runInstalledQuery(self, query_name, params, usePost=False): raise Exception( f"Query {query_name} with restppid {restppid} timed out" ) + else: + time.sleep(0.5) # Still running, wait for 0.5 second success = True except Exception as e: LogWriter.error(f"Error running query {query_name}: {str(e)}") diff --git a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt index 7874480..11c9c46 100644 --- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt +++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt @@ -1,14 +1,32 @@ -You are a highly efficient and empathetic AI-powered assistant in JSON parsing and generating. -Given the following context in JSON format, rephrase it to answer the question. +You are a highly efficient and empathetic AI-powered customer support agent. Your goal is to provide accurate, helpful, and friendly support while maintaining professionalism. Follow these guidelines: + +1. **Understand the User's Needs**: Carefully interpret the user's question or issue. Ask clarifying questions if necessary. +2. **Provide Accurate Answers**: Use the given context in JSON format as well as the data structure of the original query used to fecth the context to deliver reliable and precise responses. +3. **Provide Clean Code Examples**: Provide necessary code examples if there is code snippet in the given context to answer the question. Use less code whenever possible. +4. **Be Polite and Friendly**: Maintain a tone that is warm, respectful, and supportive. +5. **Stay Within Scope**: Respond only with information relevant to the user's question. If the information is unavailable, apologize and suggest alternative actions. +6. **Encourage Next Steps**: Proactively suggest solutions or next steps, such as contacting a human agent if the issue cannot be resolved by the AI. +7. **Escalate When Needed**: Recognize complex or sensitive queries that require human attention and politely guide the user to contact human support. +8. **Reference Relevant Knowledge**: Ground your responses based on the context and resources, and referencing the data struction from the query providedy. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID]. + +Example format for responses: +- **Solution**: "Here's the information you requested: [details]." +- **Clarification**: "Could you provide more details about [specific aspect]?" +- **Escalation**: "This issue may require human assistance. Please reach out via [Support Portal](https://www.tigergraph.com/support/) or email us at [support@tigergraph.com](mailto:support@tigergraph.com)." + +Format your answer using Markdown. Organize the content into paragraphs, bulleted or numbered lists, and include links to images where relevant. +Make sure to extract and include the image references in [IMAGE_REF:image_id] format in your generated_answer if they are present in the context. Images are critical visual information that must be included in your response. Do NOT modify or omit these image references. + +Give the context in JSON format, combine and rephrase it to answer the question. Use only the provided information in context without adding any reasoning or additional logic. -Make sure all information in the context are covered in the generated answer. +Make sure all information in the context are covered in the generated answer, including any image markdown references. Make sure to extract and include the image links in markdown syntax in the generated answer when their summaries are referenced, and preserve the link URLs in their original format. Use compact markdown syntax to geneate the answer, including title, bulleted or numbered list, images and tables if any, and place images or tables below the related text section. Ensure that each row of every table, including the header row, starts on a new line. -Always only return a JSON contains the answer and otheh required fields. Assign an empty value to the field if you cannot determine it. - +Generate the answer in JSON format, make sure to escape necessary characters in order to return a valid JSON response only. +Make sure all the fields required by the format instructions are included, set a field to empty if you don't have that information. Question: {question} -Answer: {context} +Context: {context} +Query: {query} Format: {format_instructions} - diff --git a/common/prompts/aws_bedrock_claude3haiku/generate_cypher.txt b/common/prompts/aws_bedrock_claude3haiku/generate_cypher.txt index 66a24e1..732ed49 100644 --- a/common/prompts/aws_bedrock_claude3haiku/generate_cypher.txt +++ b/common/prompts/aws_bedrock_claude3haiku/generate_cypher.txt @@ -12,6 +12,7 @@ Always use double quotes for strings instead of single quotes. Always convert strings to lower case using toLower() function for string comparision in WHERE clause. Use alias for ORDER BY if any, avoid using short alias names especially single letter alias, always use meaningful words connected by underscore. Always make sure the alias or attributes used in ORDER BY is the same type in RETURN. Always add ASC or DESC for ORDER BY based on data type. +For questions like "summarize" or "write a summary" about something, fetch all information on its neighbour nodes and edges. Avoid to generate invalid OpenCypher queries based on the errors from history below. diff --git a/common/prompts/google_gemini/chatbot_response.txt b/common/prompts/google_gemini/chatbot_response.txt index cb23c53..8088bac 100644 --- a/common/prompts/google_gemini/chatbot_response.txt +++ b/common/prompts/google_gemini/chatbot_response.txt @@ -30,4 +30,5 @@ Make sure all the fields required by the format instructions are included, set a Question: {question} Context: {context} +Query: {query} Format: {format_instructions} diff --git a/common/prompts/openai_gpt4/chatbot_response.txt b/common/prompts/openai_gpt4/chatbot_response.txt index b9cab1a..11c9c46 100644 --- a/common/prompts/openai_gpt4/chatbot_response.txt +++ b/common/prompts/openai_gpt4/chatbot_response.txt @@ -20,6 +20,9 @@ Make sure to extract and include the image references in [IMAGE_REF:image_id] fo Give the context in JSON format, combine and rephrase it to answer the question. Use only the provided information in context without adding any reasoning or additional logic. Make sure all information in the context are covered in the generated answer, including any image markdown references. +Make sure to extract and include the image links in markdown syntax in the generated answer when their summaries are referenced, and preserve the link URLs in their original format. +Use compact markdown syntax to geneate the answer, including title, bulleted or numbered list, images and tables if any, and place images or tables below the related text section. +Ensure that each row of every table, including the header row, starts on a new line. Generate the answer in JSON format, make sure to escape necessary characters in order to return a valid JSON response only. Make sure all the fields required by the format instructions are included, set a field to empty if you don't have that information. diff --git a/common/prompts/openai_gpt4/generate_cypher.txt b/common/prompts/openai_gpt4/generate_cypher.txt index 66a24e1..732ed49 100644 --- a/common/prompts/openai_gpt4/generate_cypher.txt +++ b/common/prompts/openai_gpt4/generate_cypher.txt @@ -12,6 +12,7 @@ Always use double quotes for strings instead of single quotes. Always convert strings to lower case using toLower() function for string comparision in WHERE clause. Use alias for ORDER BY if any, avoid using short alias names especially single letter alias, always use meaningful words connected by underscore. Always make sure the alias or attributes used in ORDER BY is the same type in RETURN. Always add ASC or DESC for ORDER BY based on data type. +For questions like "summarize" or "write a summary" about something, fetch all information on its neighbour nodes and edges. Avoid to generate invalid OpenCypher queries based on the errors from history below. diff --git a/common/prompts/openai_gpt4/generate_function.txt b/common/prompts/openai_gpt4/generate_function.txt index 781d63d..00ae7f0 100644 --- a/common/prompts/openai_gpt4/generate_function.txt +++ b/common/prompts/openai_gpt4/generate_function.txt @@ -20,5 +20,8 @@ Sixth Docstring: {doc6} Seventh Docstring: {doc7} Eighth Docstring: {doc8} -Follow the output directions below on how to structure your response: +If the output of this function answers the user's question, immediately return that answer. + +Follow the output directions below on how to structure your response +Only include valid JSON do not include any other texts which would render the response invalid JSON. {format_instructions} diff --git a/common/requirements.txt b/common/requirements.txt index 42c7c48..1491f56 100644 --- a/common/requirements.txt +++ b/common/requirements.txt @@ -33,7 +33,7 @@ docstring_parser==0.16 emoji==2.14.1 environs==14.2.0 exceptiongroup==1.3.0 -fastapi==0.115.14 +fastapi==0.118.0 filelock==3.18.0 filetype==1.2.0 fonttools==4.58.4 @@ -154,7 +154,7 @@ smmap==5.0.2 sniffio==1.3.1 soupsieve==2.7 SQLAlchemy==2.0.41 -starlette==0.46.2 +starlette==0.48.0 tabulate==0.9.0 tenacity==9.1.2 threadpoolctl==3.6.0 diff --git a/common/utils/token_calculator.py b/common/utils/token_calculator.py new file mode 100644 index 0000000..53d96ce --- /dev/null +++ b/common/utils/token_calculator.py @@ -0,0 +1,224 @@ +# Copyright (c) 2025 TigerGraph, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import tiktoken +import sys +from typing import List, Any, Optional + +logger = logging.getLogger(__name__) + +class TokenCalculator: + """Utility class for token counting and text truncation operations.""" + + def __init__(self, token_limit: int = 0, model_name: str = "gpt-4"): + """ + Initialize the token calculator. + + Args: + token_limit: Maximum number of tokens allowed for retrieved context + model_name: Name of the model to use for token counting + Use <= 0 for unlimited tokens (no truncation). + """ + self.max_context_tokens = token_limit + self.model_name = model_name + try: + self.token_encoding = tiktoken.encoding_for_model(self.model_name) + except Exception as e: + self.token_encoding = tiktoken.get_encoding("cl100k_base") + logger.warning(f"Error getting encoding for model {self.model_name}, using cl100k_base: {e}") + logger.info(f"Initialized TokenCalculator with max_context_tokens: {self.max_context_tokens} and encoding: {self.token_encoding}") + + def set_max_context_tokens(self, max_tokens: int): + """Set the maximum number of tokens allowed for retrieved context.""" + self.max_context_tokens = max_tokens + if self.is_unlimited_tokens(): + logger.info("Set token limit to unlimited (no truncation)") + else: + logger.info(f"Set max context tokens to: {max_tokens}") + + def get_max_context_tokens(self) -> int: + """Get the current maximum number of tokens allowed for retrieved context.""" + return self.max_context_tokens if not self.is_unlimited_tokens() else sys.maxsize + + def is_unlimited_tokens(self) -> bool: + """Check if token limit is set to unlimited.""" + return (self.max_context_tokens <= 0) + + def count_tokens(self, text: str | dict) -> int: + """Count the number of tokens in the given text.""" + try: + if not isinstance(text, str): + text = str(text) + return len(self.token_encoding.encode(text)) + except Exception as e: + logger.warning(f"Error counting tokens: {e}, using character-based estimation") + # Fallback: rough estimation (1 token ≈ 4 characters for English text) + return len(text) // 4 + + def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict: + """ + Truncate dictionary to fit within the token limit by keeping original values + until hitting the max_tokens limit, then ignoring remaining values. + + Args: + sources_dict: Dictionary to truncate + max_tokens: Maximum number of tokens allowed (defaults to self.max_context_tokens) + + Returns: + Dictionary of sources that fit within the token limit + """ + if max_tokens is None: + max_tokens = self.max_context_tokens + + if not sources_dict: + return sources_dict + + total_tokens = self.count_tokens(sources_dict) + + # If unlimited tokens is enabled, return all sources without truncation + if self.is_unlimited_tokens() or max_tokens <= 0 or total_tokens <= max_tokens: + return sources_dict + + # Convert dict to list of (key, value) pairs for processing + items = sorted(sources_dict.items(), key=lambda x: (".png" not in x, -len(x))) + truncated_sources = {} + current_tokens = 0 + + for key, value in items: + # Calculate tokens for this key-value pair + item_tokens = self.count_tokens({key: value}) + + # Check if adding this item would exceed the limit + if current_tokens + item_tokens <= max_tokens: + # Add the complete item + truncated_sources[key] = value + current_tokens += item_tokens + logger.debug(f"Added complete item '{key}' ({item_tokens} tokens, total: {current_tokens})") + else: + # Check if we can add a partial version of this item + remaining_tokens = max_tokens - current_tokens + if remaining_tokens > 0: + # Try to add a truncated version of this item + if isinstance(value, str): + # Truncate string to fit remaining tokens + truncated_value = self.truncate_text_to_token_limit(value, remaining_tokens) + if truncated_value: + truncated_sources[key] = truncated_value + current_tokens += self.count_tokens(truncated_value) + logger.debug(f"Added truncated string '{key}' ({self.count_tokens(truncated_value)} tokens, total: {current_tokens})") + elif isinstance(value, list): + # Add as many list items as possible + truncated_list = [] + for item in value: + if isinstance(item, str): + item_tokens = self.count_tokens(item) + if current_tokens + item_tokens <= max_tokens: + truncated_list.append(item) + current_tokens += item_tokens + else: + # Try to add a truncated version of this item + remaining = max_tokens - current_tokens + if remaining > 0: + truncated_item = self.truncate_text_to_token_limit(item, remaining) + if truncated_item: + truncated_list.append(truncated_item) + current_tokens += self.count_tokens(truncated_item) + break + else: + # For non-string items, add if there's space + item_tokens = self.count_tokens(item) + if current_tokens + item_tokens <= max_tokens: + truncated_list.append(item) + current_tokens += item_tokens + else: + break + if truncated_list: + truncated_sources[key] = truncated_list + logger.debug(f"Added truncated list '{key}' ({len(truncated_list)} items, total: {current_tokens})") + elif isinstance(value, dict): + # Recursively truncate sub-dictionary + remaining = max_tokens - current_tokens + if remaining > 0: + truncated_subdict = self.truncate_dict_to_token_limit(value, remaining) + if truncated_subdict: + truncated_sources[key] = truncated_subdict + current_tokens += self.count_tokens(truncated_subdict) + logger.debug(f"Added truncated sub-dict '{key}' ({self.count_tokens(truncated_subdict)} tokens, total: {current_tokens})") + else: + # For other types, add if there's space + if current_tokens + item_tokens <= max_tokens: + truncated_sources[key] = value + current_tokens += item_tokens + logger.debug(f"Added complete non-string item '{key}' ({item_tokens} tokens, total: {current_tokens})") + else: + # No more space, stop processing + logger.debug(f"Stopping truncation - no more space for item '{key}'") + break + + logger.info(f"Final truncated context tokens: {current_tokens} (limit: {max_tokens})") + return truncated_sources + + def truncate_text_to_token_limit(self, text: str, max_tokens: Optional[int] = None) -> str: + """ + Truncate text to fit within the specified token limit. + + Args: + text: Text to truncate + max_tokens: Maximum number of tokens allowed + + Returns: + Truncated text + """ + if max_tokens is None: + max_tokens = self.max_context_tokens + + try: + tokens = self.token_encoding.encode(text) + if len(tokens) <= max_tokens: + return text + + # Truncate to max_tokens and decode back to text + truncated_tokens = tokens[:max_tokens] + truncated_text = self.token_encoding.decode(truncated_tokens) + + # Add ellipsis to indicate truncation + #if len(tokens) > max_tokens: + # truncated_text += "..." + + return truncated_text + except Exception as e: + logger.warning(f"Error truncating text: {e}, using character-based truncation") + # Fallback: rough estimation (1 token ≈ 4 characters) + max_chars = max_tokens * 4 + if len(text) <= max_chars: + return text + return text[:max_chars] #+ "..." + + def truncate_to_token_limit(self, text: str | dict, max_tokens: Optional[int] = None) -> str: + """ + Truncate text to fit within the specified token limit. + + Args: + text: Text to truncate + max_tokens: Maximum number of tokens allowed + + Returns: + Truncated text + """ + if isinstance(text, dict): + return self.truncate_dict_to_token_limit(text, max_tokens) + else: + return self.truncate_text_to_token_limit(text, max_tokens) + diff --git a/configs/server_config.json b/configs/server_config.json index d607a10..7c1d468 100644 --- a/configs/server_config.json +++ b/configs/server_config.json @@ -14,6 +14,7 @@ "chat_history_api": "http://chat-history:8002" }, "llm_config": { + "token_limit": 0, "authentication_configuration": { "OPENAI_API_KEY": "YOUR_OPENAI_API_KEY_HERE" }, diff --git a/docker-compose.yml b/docker-compose.yml index 29d4cad..8be754b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -18,7 +18,6 @@ services: USE_CYPHER: "true" volumes: - ./configs/:/code/configs - - YOUR_DATA_PATH_HERE:/data graphrag-ecc: image: tigergraph/graphrag-ecc:latest diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py index 0e06e75..cfbab3d 100644 --- a/ecc/app/graphrag/graph_rag.py +++ b/ecc/app/graphrag/graph_rag.py @@ -81,6 +81,44 @@ async def stream_docs( logger.info("closing docs chan") docs_chan.close() +async def stream_chunks( + conn: AsyncTigerGraphConnection, + extract_chan: Channel, + embed_chan: Channel, + ttl_batches: int = 10, +): + """ + Streams the chunk contents into the extract_chan and embed_chan + """ + logger.info("streaming chunks") + for i in range(ttl_batches): + chunk_ids = await stream_ids(conn, "DocumentChunk", i, ttl_batches) + if chunk_ids["error"]: + continue + + for c in chunk_ids["ids"]: + try: + async with tg_sem: + res = await conn.runInstalledQuery( + "StreamChunkContent", + params={"chunk": c}, + ) + content = res[0]["ChunkContent"][0]["attributes"]["text"].encode('utf-8').decode('unicode_escape') + logger.info("chunk writes to extract_chan") + await extract_chan.put((content, c)) + + # send chunks to be embedded + logger.info("chunk writes to embed_chan") + await embed_chan.put((c, content, "DocumentChunk")) + except Exception as e: + exc = traceback.format_exc() + logger.error(f"Error retrieving chunk: {c} --> {e}\n{exc}") + continue # try retrieving the next doc + + logger.info("stream_chunks done") + logger.info("closing extract_chan") + await extract_chan.put(None) + async def chunk_docs( conn: AsyncTigerGraphConnection, @@ -110,10 +148,8 @@ async def chunk_docs( logger.info("Chunk Processing End") - # close the extract chan -- chunk_doc is the only sender - # and chunk_doc calls are kicked off from here logger.info("closing extract_chan") - extract_chan.close() + await extract_chan.put(None) async def upsert(upsert_chan: Channel): @@ -248,6 +284,7 @@ async def extract( embed_chan: Channel, extractor: BaseExtractor, conn: AsyncTigerGraphConnection, + num_senders: int, ): """ Creates and starts one worker for each extract job @@ -257,13 +294,20 @@ async def extract( logger.info("Entity Extration Start") # consume task queue async with asyncio.TaskGroup() as grp: + done_count = 0 while True: try: item = await extract_chan.get() - if entity_extraction_switch: - grp.create_task( - workers.extract(upsert_chan, embed_chan, extractor, conn, *item) - ) + if item is None: # sender finished + done_count += 1 + if done_count == num_senders: + logger.info("All senders finished, exiting extract.") + break + else: + if entity_extraction_switch: + grp.create_task( + workers.extract(upsert_chan, embed_chan, extractor, conn, *item) + ) except ChannelClosed: break except Exception: @@ -271,7 +315,8 @@ async def extract( logger.info("Entity Extration End") - logger.info("closing upsert and embed chan") + logger.info("closing extract, upsert and embed chan") + extract_chan.close() upsert_chan.close() embed_chan.close() @@ -497,6 +542,8 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection): embed_chan = Channel() upsert_chan = Channel() extract_chan = Channel() + num_chunk_senders = 2 + async with asyncio.TaskGroup() as grp: # get docs grp.create_task(stream_docs(conn, docs_chan, 100)) @@ -504,6 +551,9 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection): grp.create_task( chunk_docs(conn, docs_chan, embed_chan, upsert_chan, extract_chan) ) + # process existing chunks + grp.create_task(stream_chunks(conn, extract_chan, embed_chan, 100)) + # upsert chunks grp.create_task(upsert(upsert_chan)) grp.create_task(load(conn)) @@ -511,7 +561,7 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection): grp.create_task(embed(embed_chan, embedding_store, graphname)) # extract entities grp.create_task( - extract(extract_chan, upsert_chan, embed_chan, extractor, conn) + extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders) ) logger.info("Join docs_chan") await docs_chan.join() diff --git a/ecc/app/graphrag/util.py b/ecc/app/graphrag/util.py index 6a7534a..25f4cff 100644 --- a/ecc/app/graphrag/util.py +++ b/ecc/app/graphrag/util.py @@ -83,6 +83,7 @@ async def init( requried_queries = [ "common/gsql/graphrag/StreamIds", "common/gsql/graphrag/StreamDocContent", + "common/gsql/graphrag/StreamChunkContent", "common/gsql/graphrag/SetEpochProcessing", "common/gsql/graphrag/ResolveRelationships", "common/gsql/graphrag/get_community_children", @@ -169,14 +170,12 @@ def map_attrs(attributes: dict): def process_id(v_id: str): - v_id = v_id.replace(" ", "_").replace("/", "").replace("%", "percent").lower() - has_func = re.compile(r"(.*)\(").findall(v_id) if len(has_func) > 0: v_id = has_func[0] + v_id = v_id.replace(" ", "-").lower().replace("(", "").replace(")", "") if v_id == "''" or v_id == '""': return "" - v_id = v_id.replace("(", "").replace(")", "") return v_id diff --git a/ecc/app/graphrag/workers.py b/ecc/app/graphrag/workers.py index 1317322..1537885 100644 --- a/ecc/app/graphrag/workers.py +++ b/ecc/app/graphrag/workers.py @@ -148,7 +148,7 @@ async def upsert_chunk(conn: AsyncTigerGraphConnection, doc_id, chunk_id, chunk) conn, "DocumentChunk", chunk_id, - attributes={"epoch_added": date_added, "idx": int(chunk_id.split("_")[-1])}, + attributes={"epoch_added": date_added, "epoch_processed": date_added, "idx": int(chunk_id.split("_")[-1])}, ) await util.upsert_vertex( conn, diff --git a/ecc/app/supportai/util.py b/ecc/app/supportai/util.py index 24d653d..2e670c7 100644 --- a/ecc/app/supportai/util.py +++ b/ecc/app/supportai/util.py @@ -80,6 +80,7 @@ async def init( "common/gsql/supportai/Check_Nonexistent_Vertices", "common/gsql/graphRAG/StreamIds", "common/gsql/graphRAG/StreamDocContent" + "common/gsql/graphRAG/StreamChunkContent" ] await install_queries(requried_queries, conn) @@ -152,11 +153,10 @@ def map_attrs(attributes: dict): def process_id(v_id: str): - v_id = v_id.replace(" ", "_").replace("/", "").replace("%", "percent").lower() - has_func = re.compile(r"(.*)\(").findall(v_id) if len(has_func) > 0: v_id = has_func[0] + v_id = v_id.replace(" ", "-").lower().replace("(", "").replace(")", "") if v_id == "''" or v_id == '""': return "" diff --git a/graphrag-ui/package.json b/graphrag-ui/package.json index 9983a70..4ef113e 100755 --- a/graphrag-ui/package.json +++ b/graphrag-ui/package.json @@ -25,33 +25,36 @@ "install": "^0.13.0", "lucide-react": "^0.390.0", "npm": "^10.8.1", - "react": "^18.2.0", + "react": "^18.3.1", "react-chatbot-kit": "^2.2.2", - "react-dom": "^18.2.0", + "react-dom": "^18.3.1", "react-hook-form": "^7.51.5", "react-i18next": "^14.1.2", "react-icons": "^5.2.1", "react-markdown": "^9.0.1", "react-router-dom": "^6.23.1", "react-use-websocket": "^4.8.1", - "reagraph": "^4.19.1", + "reagraph": "4.15.19", + "@react-three/fiber": "8.13.3", + "@react-three/drei": "9.56.1", + "remark-gfm": "^4.0.0", "tailwind-merge": "^2.3.0", "tailwindcss-animate": "^1.0.7", "@tailwindcss/typography": "^0.5.18", "zod": "^3.23.8" }, "devDependencies": { - "@types/react": "^18.2.66", - "@types/react-dom": "^18.2.22", - "@typescript-eslint/eslint-plugin": "^7.2.0", - "@typescript-eslint/parser": "^7.2.0", - "@vitejs/plugin-react-swc": "^3.5.0", + "@types/react": "^18.2.0", + "@types/react-dom": "^18.2.0", + "@typescript-eslint/eslint-plugin": "^7.5.0", + "@typescript-eslint/parser": "^7.5.0", + "@vitejs/plugin-react-swc": "^3.6.0", "autoprefixer": "^10.4.19", "eslint": "^8.57.0", - "eslint-plugin-react-hooks": "^4.6.0", + "eslint-plugin-react-hooks": "^6.1.1", "eslint-plugin-react-refresh": "^0.4.6", "postcss": "^8.4.38", - "tailwindcss": "^3.4.3", + "tailwindcss": "^3.4.18", "typescript": "^5.2.2", "vite": "^5.2.0" } diff --git a/graphrag-ui/src/actions/ActionProvider.tsx b/graphrag-ui/src/actions/ActionProvider.tsx index a268b5a..968635f 100644 --- a/graphrag-ui/src/actions/ActionProvider.tsx +++ b/graphrag-ui/src/actions/ActionProvider.tsx @@ -2,7 +2,7 @@ import React, {useState, useCallback, useEffect, useContext} from 'react'; import {createClientMessage} from 'react-chatbot-kit'; import useWebSocket, {ReadyState} from 'react-use-websocket'; import Loader from '../components/Loader'; -import { SelectedGraphContext } from '../components/Contexts'; +import { SelectedGraphContext, RagPatternContext } from '../components/Contexts'; interface ActionProviderProps { createChatBotMessage: any; @@ -80,17 +80,20 @@ const ActionProvider: React.FC = ({ children, }) => { const selectedGraph = useContext(SelectedGraphContext); - const WS_URL = "/ui/" + selectedGraph + "/chat"; + const selectedRagPattern = useContext(RagPatternContext); + const WS_URL = "/ui/" + selectedGraph + "/chat" + "?rag_pattern=" + selectedRagPattern; const [messageHistory, setMessageHistory] = useState[]>( [], ); const { sendMessage, lastMessage, readyState } = useWebSocket(WS_URL, { onOpen: () => { // Send authentication credentials - queryGraphragWs2(localStorage.getItem("creds")!); + const creds = localStorage.getItem("creds"); + console.log("Sending credentials, length:", creds ? creds.length : 0); + queryGraphragWs2(creds!); // Send RAG pattern - sendMessage(localStorage.getItem("ragPattern") || "Hybrid Search"); + //sendMessage(selectedRagPattern); // Send conversation ID (or "new" for new conversation) const conversationId = conversationManager.getCurrentConversationId(); @@ -98,6 +101,16 @@ const ActionProvider: React.FC = ({ console.log("WebSocket connection " + conversationIdToSend + " established to " + WS_URL); sendMessage(conversationIdToSend); }, + onError: (error) => { + console.error("WebSocket error:", error); + }, + onClose: (event) => { + console.log("WebSocket closed:", event.code, event.reason); + }, + shouldReconnect: (closeEvent) => { + console.log("WebSocket should reconnect:", closeEvent.code !== 1000); + return closeEvent.code !== 1000; // Don't reconnect on normal closure + }, }); // Initialize conversation manager with any existing conversation data diff --git a/graphrag-ui/src/components/Bot.tsx b/graphrag-ui/src/components/Bot.tsx index f80dd80..95b2009 100644 --- a/graphrag-ui/src/components/Bot.tsx +++ b/graphrag-ui/src/components/Bot.tsx @@ -6,7 +6,7 @@ import ActionProvider from "../actions/ActionProvider.js"; import config from "../actions/config.js"; import MessageParser from "../actions/MessageParser.js"; import { MdKeyboardArrowDown } from "react-icons/md"; -import { SelectedGraphContext } from './Contexts.js'; +import { SelectedGraphContext, RagPatternContext } from './Contexts.js'; import { Button } from "@/components/ui/button"; import { @@ -22,14 +22,27 @@ import { const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getConversationId?:any }) => { const [store, setStore] = useState(); const [currentDate, setCurrentDate] = useState(''); - const [selectedGraph, setSelectedGraph] = useState(localStorage.getItem("selectedGraph") || 'TigerGraphRAG'); - const [ragPattern, setRagPattern] = useState(localStorage.getItem("ragPattern") || 'Hybrid Search'); + const [selectedGraph, setSelectedGraph] = useState(localStorage.getItem("selectedGraph") || ''); + const [ragPattern, setRagPattern] = useState(localStorage.getItem("ragPattern") || ''); const navigate = useNavigate(); useEffect(() => { const parseStore = JSON.parse(localStorage.getItem("site") || "{}"); setStore(parseStore); + // Set default selectedGraph to first graph if no value in localStorage + if (!localStorage.getItem("selectedGraph") && parseStore?.graphs?.length > 0) { + const firstGraph = parseStore.graphs[0]; + setSelectedGraph(firstGraph); + localStorage.setItem("selectedGraph", firstGraph); + } + + // Set default ragPattern if no value in localStorage + if (!localStorage.getItem("ragPattern")) { + setRagPattern("Hybrid Search"); + localStorage.setItem("ragPattern", "Hybrid Search"); + } + const date = new Date(); const options: Intl.DateTimeFormatOptions = { year: 'numeric', month: 'long', day: 'numeric', weekday: 'long' }; const formattedDate = date.toLocaleDateString('en-US', options); @@ -46,6 +59,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo const handleSelectRag = (value) => { setRagPattern(value); localStorage.setItem("ragPattern", value); + navigate("/chat"); //window.location.reload(); }; @@ -108,15 +122,17 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo - + + + ); diff --git a/graphrag-ui/src/components/Contexts.tsx b/graphrag-ui/src/components/Contexts.tsx index 142fc08..9493c26 100644 --- a/graphrag-ui/src/components/Contexts.tsx +++ b/graphrag-ui/src/components/Contexts.tsx @@ -1,3 +1,4 @@ import React, { createContext } from "react"; -export const SelectedGraphContext = createContext("pyTigerGraphRAG"); \ No newline at end of file +export const SelectedGraphContext = createContext(""); +export const RagPatternContext = createContext(""); \ No newline at end of file diff --git a/graphrag-ui/src/components/CustomChatMessage.tsx b/graphrag-ui/src/components/CustomChatMessage.tsx index 36d2ebd..1b7c3fc 100755 --- a/graphrag-ui/src/components/CustomChatMessage.tsx +++ b/graphrag-ui/src/components/CustomChatMessage.tsx @@ -1,5 +1,6 @@ import { FC, useState, useEffect } from "react"; -import Markdown from 'react-markdown' +import ReactMarkdown from 'react-markdown' +import remarkGfm from 'remark-gfm' import { Dialog, DialogContent, @@ -162,7 +163,7 @@ export const CustomChatMessage: FC = ({ <> {typeof message === "string" ? (
- {message} + {message}
) : message.key === null ? ( message @@ -172,7 +173,7 @@ export const CustomChatMessage: FC = ({ {message.response_type === "progress" ? (

{message.content}

) : ( - {message.content} + {message.content} )} = ({ {/* {message.query_sources?.result ? 
{JSON.stringify(message.query_sources?.result, null, 2)}
: null} */} {/* 
{JSON.stringify(message, null, 2)}
*/}
- {message.query_sources?.result ? ( - + {message.query_sources?.result.edges ? ( + ) : (
No graph data available diff --git a/graphrag-ui/src/components/Start.tsx b/graphrag-ui/src/components/Start.tsx index 7802f85..6ddd927 100644 --- a/graphrag-ui/src/components/Start.tsx +++ b/graphrag-ui/src/components/Start.tsx @@ -3,29 +3,34 @@ import { HiOutlineChatBubbleOvalLeft } from "react-icons/hi2"; import { useTheme } from "@/components/ThemeProvider"; -const questions = - localStorage.getItem('selectedGraph') === 'pyTigerGraphRAG' - ? [ - {title: 'What can pyTG do?'}, - {title: 'How do I authenticate with the DB?'}, - {title: 'How do I run PageRank?'}, - {title: 'Help me initialize a NodePiece model'}, - {title: 'Can I load data with pyTG?'}, - {title: 'Write a short python snippet to run a loading job'}, - ] - : localStorage.getItem('selectedGraph') === 'Transaction_Fraud' ? - [ - {title: 'How many transactions are there?'}, - {title: 'Tell me about transaction fraud.'}, - {title: 'Describe flow of one transaction.'}, - {title: 'How TigerGraph can help me?'} - // { title: "How to use visualization correctly?" }, - // { title: "How to detect fraud in transactions?" }, - // { title: "What is William Torres' ID?" }, - // { title: "What's his email?" }, - // {title:"How do I get a count of vertices in Python?"} - ] - : []; +const questions = (() => { + const selectedGraph = localStorage.getItem('selectedGraph'); + + if (selectedGraph?.includes('pyTigerGraphRAG') || selectedGraph?.includes('pyTG')) { + return [ + {title: 'What can pyTG do?'}, + {title: 'How do I authenticate with the DB?'}, + {title: 'How do I run PageRank?'}, + {title: 'Help me initialize a NodePiece model'}, + {title: 'Can I load data with pyTG?'}, + {title: 'Write a short python snippet to run a loading job'}, + ]; + } else if (selectedGraph?.includes('Transaction_Fraud') || selectedGraph?.includes('fraud')) { + return [ + {title: 'How many transactions are there?'}, + {title: 'Tell me about transaction fraud.'}, + {title: 'Describe flow of one transaction.'}, + {title: 'How TigerGraph can help me?'} + // { title: "How to use visualization correctly?" }, + // { title: "How to detect fraud in transactions?" }, + // { title: "What is William Torres' ID?" }, + // { title: "What's his email?" }, + // {title:"How do I get a count of vertices in Python?"} + ]; + } + + return []; +})(); interface Start { props: any; diff --git a/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx b/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx index b5ba6d3..4ebb2dc 100644 --- a/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx +++ b/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx @@ -6,15 +6,16 @@ export const KnowledgeGraphPro = ({ data }) => { const [theme, setTheme] = useState(localStorage.getItem("vite-ui-theme")); const ref = useRef(null); // const [sdata, setsdata] = useState(JSON.parse(data)); - const [edges, setEdges] = useState([]); + const [edges, setEdges] = useState([]); + const [nodes, setNodes] = useState([]); const [dataArray, setdataArray] = useState(); const [vId, setvId] = useState(); useEffect(() => { + let parsedData: any[] = []; if (typeof data === 'string') { - const parseData = JSON.parse(data); - setEdges(parseData); + parsedData = JSON.parse(data); // do i need to parse for question 'show me 5 transacitons with details' // { @@ -41,17 +42,31 @@ export const KnowledgeGraphPro = ({ data }) => { // } else null } else if (Array.isArray(data)) { // YES THERE ARE 5 from question 'show me 5 transacitons with details' - const setresults = data[1]["@@edges"]; + // const setresults = data[1]["@@edges"]; + parsedData = data; // ^ this is valid for question 'what cards have more than 800 transactions between april 1 2021 to august 1 2021' - setEdges(setresults); - } else { + } else { // Handle object data directly - setEdges(data); + parsedData = data.edges; } + + let edgesList: any[] = []; + let nodesSet = new Set(); + + parsedData.forEach((e, i) => { + if (e.s && e.t) { + edgesList.push({ id: `e${i}`, source: e.s, target: e.t }); + nodesSet.add(e.s); + nodesSet.add(e.t); + } + }); + + setEdges(edgesList); + setNodes(Array.from(nodesSet).map(id => ({ id, label: id }))); }, [data]); useEffect(() => { - console.log('\n\n\n\n\n\n\n\n\n\n PARSED edges', edges); + console.log('PARSED edges', edges); },[]) // const getNodes = edges.map((d:any) => ( @@ -78,23 +93,23 @@ export const KnowledgeGraphPro = ({ data }) => { // ] return ( - <>{edges ? JSON.stringify(edges) : 'no data'} + <>{/* edges ? JSON.stringify(edges) : 'no data' */} {/* {edges} */} {/* {edges && 
{edges}
} */} - {/* {typeof sdata !== 'number' && typeof sdata !== 'string' && dataArray?.edgez && dataArray?.nodes ? ( + { edges && nodes && nodes.length > 0 && edges.length > 0 ? ( ) :
Sorry no graph or table available
} - {typeof sdata !== 'number' && typeof sdata !== 'string' ? (
+ {edges && nodes ? (
-
) : null} */} +
) : null} ) } \ No newline at end of file diff --git a/graphrag/app/agent/agent.py b/graphrag/app/agent/agent.py index 14f5df9..fb4dcb6 100644 --- a/graphrag/app/agent/agent.py +++ b/graphrag/app/agent/agent.py @@ -166,6 +166,9 @@ def question_for_agent( def make_agent(graphname, conn, use_cypher, ws: WebSocket = None, supportai_retriever="hybridsearch") -> TigerGraphAgent: + if "chat_model" in llm_config["completion_service"]: + llm_config["completion_service"]["llm_model"] = llm_config["completion_service"]["chat_model"] + if llm_config["completion_service"]["llm_service"].lower() == "openai": llm_service_name = "openai" llm_provider = OpenAI(llm_config["completion_service"]) diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py index 5391746..d38bd8b 100644 --- a/graphrag/app/agent/agent_generation.py +++ b/graphrag/app/agent/agent_generation.py @@ -12,26 +12,31 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import logging from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser from langchain_community.callbacks.manager import get_openai_callback +from typing import Optional from pydantic import BaseModel, Field from common.logs.logwriter import LogWriter from common.logs.log import req_id_cv +from common.utils.token_calculator import TokenCalculator +from common.config import completion_config logger = logging.getLogger(__name__) class GraphRAGAnswerOutput(BaseModel): generated_answer: str = Field(description="The generated answer to the question. Make sure maintain a professional tone.") - citation: list[str] = Field(description="The citation for the answer. List the information used.") + citation: Optional[list[str]] = Field(description="The citation for the answer. List the metadata of the parts of the context used.", default=[]) class TigerGraphAgentGenerator: def __init__(self, llm_model): self.llm = llm_model + self.token_calculator = TokenCalculator(token_limit=completion_config.get("token_limit"), model_name=completion_config.get("llm_model")) - def generate_answer(self, question: str, context: str, query: str = "") -> dict: + def generate_answer(self, question: str, context: str | dict, query: str = "") -> dict: """Generate an answer based on the question and context. Args: question: str: The question to generate an answer for. @@ -42,6 +47,17 @@ def generate_answer(self, question: str, context: str, query: str = "") -> dict: """ LogWriter.info(f"request_id={req_id_cv.get()} ENTRY generate_answer") + # Truncate context to fit within token limit + if not self.token_calculator.is_unlimited_tokens(): + # Reserve tokens for question, query, and format instructions (approximately 1000 tokens) + max_context_tokens = self.token_calculator.get_max_context_tokens() - 1000 + + if len(str(context)) > max_context_tokens: + context_tokens = self.token_calculator.count_tokens(context) + if context_tokens > max_context_tokens: + context = self.token_calculator.truncate_to_token_limit(context, max_context_tokens) + logger.info(f"Truncated context from {context_tokens} to {max_context_tokens} tokens") + answer_parser = PydanticOutputParser(pydantic_object=GraphRAGAnswerOutput) prompt = PromptTemplate( @@ -62,6 +78,9 @@ def generate_answer(self, question: str, context: str, query: str = "") -> dict: # Chain rag_chain = prompt | self.llm.model | answer_parser + if isinstance(context, dict): + context = json.dumps(context) + usage_data = {} with get_openai_callback() as cb: generation = rag_chain.invoke({"question": question, "context": context, "query": query}) diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index 2fc19d5..0874217 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -35,6 +35,7 @@ from common.logs.log import req_id_cv from common.py_schemas import GraphRAGResponse, MapQuestionToSchemaResponse from common.llm_services.aws_bedrock_service import AWSBedrock +from common.config import graphrag_config logger = logging.getLogger(__name__) @@ -84,9 +85,9 @@ def __init__( self.supportai_enabled = True self.supportai_retriever = supportai_retriever.lower().replace(" ", "") try: - self.db_connection.getQueryMetadata("GraphRAG_Hybrid_Search") + self.db_connection.getQueryMetadata("StreamDocContent") except TigerGraphException as e: - logger.info(f"GraphRAG_Hybrid_Search not found in the graph {self.db_connection.graphname}. Disabling supportai.") + logger.info(f"StreamDocContent not found in the graph {self.db_connection.graphname}. Disabling supportai.") self.supportai_enabled = False def emit_progress(self, msg): @@ -229,15 +230,18 @@ def hybrid_search(self, state): self.llm_provider.model, self.db_connection, ) + chunk_only=graphrag_config.get("chunk_only", True) step = retriever.search( state["question"], - indices=["Document", "DocumentChunk", "Entity", "Relationship"], - top_k=3, - num_seen_min=2, - num_hops=2, + indices=(["DocumentChunk"] if chunk_only else ["Document", "DocumentChunk", "Entity"]), + top_k=graphrag_config.get("top_k", 5), + num_seen_min=graphrag_config.get("num_seen_min", 2), + num_hops=graphrag_config.get("num_hops", 2), + chunk_only=chunk_only, + doc_only=graphrag_config.get("doc_only", False), ) - query_name = "GraphRAG_Hybrid_Search" + query_name = "GraphRAG_Hybrid_Vector_Search" state["context"] = { "function_call": query_name, "result": step[0], @@ -263,10 +267,10 @@ def similarity_search(self, state): step = retriever.search( state["question"], index="DocumentChunk", - top_k=5 + top_k=graphrag_config.get("top_k", 5) ) - query_name = "Content_Similarity_Search" + query_name = "Content_Similarity_Vector_Search" state["context"] = { "function_call": query_name, "result": step[0], @@ -291,10 +295,10 @@ def sibling_search(self, state): step = retriever.search( state["question"], index="DocumentChunk", - top_k=3 + top_k=graphrag_config.get("top_k", 5) ) - query_name = "Chunk_Sibling_Search" + query_name = "Chunk_Sibling_Vector_Search" state["context"] = { "function_call": query_name, "result": step[0], @@ -318,12 +322,12 @@ def community_search(self, state): ) step = retriever.search( state["question"], - community_level=2, - top_k=5, - with_chunk=True, + community_level=graphrag_config.get("community_level", 2), + top_k=graphrag_config.get("top_k", 5), + with_chunk=graphrag_config.get("with_chunk", True), ) - query_name = "GraphRAG_Community_Search" + query_name = "GraphRAG_Community_Vector_Search" state["context"] = { "function_call": query_name, "result": step[0], @@ -366,6 +370,11 @@ def generate_answer(self, state): answer = step.generate_answer( state["question"], state["context"]["result"]["final_retrieval"] ) + + if not answer.citation: + answer.citation = list(state["context"]["result"]["final_retrieval"].keys()) + state["context"]["reasoning"] = list(set(answer.citation)) + elif state["lookup_source"] == "inquiryai": logger.debug_pii( f"""request_id={req_id_cv.get()} Got result: {state["context"]["result"]}""" @@ -376,7 +385,7 @@ def generate_answer(self, state): logger.error(f"Failed to serialize context to JSON: {e}") raise ValueError("Invalid context data format. Unable to convert to JSON.") - answer = step.generate_answer(state["question"], context_data_str) + answer = step.generate_answer(state["question"], state["context"]["result"]) elif state["lookup_source"] == "cypher": logger.debug_pii( @@ -387,12 +396,6 @@ def generate_answer(self, state): f"request_id={req_id_cv.get()} Generated answer: {answer.generated_answer}" ) - if state["lookup_source"] == "supportai": - import re - - citations = [re.sub(r"_chunk_\d+", "", x) for x in answer.citation] - state["context"]["reasoning"] = list(set(citations)) - try: # Replace S3 URLs with presigned URLs (for AWS Bedrock BDA processing) if isinstance(self.llm_provider, AWSBedrock): diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index b8c7f69..aba80eb 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -403,13 +403,12 @@ async def load_conversation_history(conversation_id: str, usr_auth: str) -> list ) res.raise_for_status() conversation_data = res.json() - # Convert conversation messages to the format expected by the agent history = [] - for msg in conversation_data.get("messages", []): + for msg in conversation_data: if msg.get("role") == "user": # Find the corresponding system response - for response_msg in conversation_data.get("messages", []): + for response_msg in conversation_data: if (response_msg.get("role") == "system" and response_msg.get("parent_id") == msg.get("message_id")): history.append({ @@ -450,6 +449,7 @@ async def graph_query( graphname: str, creds: Annotated[tuple[list[str], HTTPBasicCredentials], Depends(ui_basic_auth)], q: str | None = None, + rag_pattern: str | None = None, conversation_id: str | None = None, ): creds = creds[1] @@ -469,7 +469,7 @@ async def graph_query( # create agent # get retrieval pattern to use - rag_pattern = "hybridsearch" + rag_pattern = rag_pattern or "hybridsearch" agent = make_agent(graphname, conn, use_cypher, supportai_retriever=rag_pattern) prev_id = None @@ -523,7 +523,8 @@ async def graph_query( @router.websocket(route_prefix + "/{graphname}/chat") async def chat( graphname: str, - websocket: WebSocket + websocket: WebSocket, + rag_pattern: str | None = None, ): """ WebSocket endpoint for chat functionality with conversation history support. @@ -542,13 +543,24 @@ async def chat( await websocket.accept() - # AUTH - # this will error if auth does not pass. FastAPI will correctly respond depending on error - usr_auth = await websocket.receive_text() - _, conn = ws_basic_auth(usr_auth, graphname) + # AUTH with proper error handling and timeout + try: + logger.info(f"WebSocket connected, waiting for authentication for graph: {graphname}") + usr_auth = await asyncio.wait_for(websocket.receive_text(), timeout=10.0) + logger.info(f"Received authentication data, length: {len(usr_auth)}") + _, conn = ws_basic_auth(usr_auth, graphname) + logger.info("Authentication successful") + except asyncio.TimeoutError: + logger.error("WebSocket authentication timeout - no credentials received") + await websocket.close(code=1008, reason="Authentication timeout") + return + except Exception as e: + logger.error(f"Authentication failed: {e}") + await websocket.close(code=1008, reason=f"Authentication failed") + return # Get RAG pattern - rag_pattern = await websocket.receive_text() + rag_pattern = rag_pattern or "hybridsearch" # Get conversation ID conversation_id = await websocket.receive_text() diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py index 647fb70..6c1b16a 100644 --- a/graphrag/app/supportai/retrievers/BaseRetriever.py +++ b/graphrag/app/supportai/retrievers/BaseRetriever.py @@ -3,6 +3,8 @@ from common.metrics.tg_proxy import TigerGraphConnectionProxy from common.llm_services.base_llm import LLM_Model from common.py_schemas import CandidateScore, CandidateGenerator +from common.utils.token_calculator import TokenCalculator +from common.config import completion_config from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser from langchain_core.prompts import ChatPromptTemplate, PromptTemplate @@ -27,8 +29,10 @@ def __init__( self.embedding_store = embedding_store self.embedding_store.set_graphname(connection.graphname) self.logger = logging.getLogger(__name__) + self.token_calculator = TokenCalculator(token_limit=completion_config.get("token_limit"), model_name=completion_config.get("llm_model")) def _install_query(self, query_name): + self.logger.info(f"Installing query {query_name}") with open(f"common/gsql/supportai/retrievers/{query_name}.gsql", "r") as f: query = f.read() res = self.conn.gsql( @@ -125,6 +129,17 @@ def _expand_question(self, question, top_k, verbose): return questions def _generate_response(self, question, retrieved, verbose): + # Truncate retrieved sources to fit within token limit + if not self.token_calculator.is_unlimited_tokens(): + # Reserve tokens for question, query, and format instructions (approximately 1000 tokens) + max_context_tokens = self.token_calculator.get_max_context_tokens() - 1000 + + if len(retrieved) > max_context_tokens: + retrieved_tokens = self.token_calculator.count_tokens(retrieved) + if retrieved_tokens > max_context_tokens: + retrieved = self.token_calculator.truncate_to_token_limit(retrieved, max_context_tokens) + self.logger.info(f"Truncated retrieved text from {retrieved_tokens} to {max_context_tokens} tokens") + model = self.llm_service.llm prompt = self.llm_service.supportai_response_prompt diff --git a/graphrag/app/supportai/retrievers/CommunityRetriever.py b/graphrag/app/supportai/retrievers/CommunityRetriever.py index 549c58d..8fe16c1 100644 --- a/graphrag/app/supportai/retrievers/CommunityRetriever.py +++ b/graphrag/app/supportai/retrievers/CommunityRetriever.py @@ -60,6 +60,7 @@ def search(self, question, community_level: int, top_k: int = 5, similarity_thre "with_doc": with_doc, "verbose": verbose, }, + sizeLimit=1000000000, usePost=True ) diff --git a/graphrag/app/supportai/retrievers/HybridRetriever.py b/graphrag/app/supportai/retrievers/HybridRetriever.py index 9caf922..da3afd8 100644 --- a/graphrag/app/supportai/retrievers/HybridRetriever.py +++ b/graphrag/app/supportai/retrievers/HybridRetriever.py @@ -75,6 +75,7 @@ def search(self, question, indices, top_k=1, similarity_threshold=0.90, num_hops "doc_only": doc_only, "verbose": verbose, }, + sizeLimit=1000000000, usePost=True ) if len(res) > 1 and "verbose" in res[1]: diff --git a/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py index 9a5b978..91805f9 100644 --- a/graphrag/app/supportai/supportai.py +++ b/graphrag/app/supportai/supportai.py @@ -32,10 +32,6 @@ def init_supportai(conn: TigerGraphConnection, graphname: str) -> tuple[dict, di "common/gsql/supportai/Selected_Set_Display.gsql", "common/gsql/supportai/retrievers/GraphRAG_Hybrid_Search_Display.gsql", "common/gsql/supportai/retrievers/GraphRAG_Community_Search_Display.gsql", - "common/gsql/supportai/retrievers/Chunk_Sibling_Search.gsql", - "common/gsql/supportai/retrievers/Content_Similarity_Search.gsql", - "common/gsql/supportai/retrievers/GraphRAG_Hybrid_Search.gsql", - "common/gsql/supportai/retrievers/GraphRAG_Community_Search.gsql", ] if "- VERTEX ResolvedEntity" in current_schema: @@ -145,12 +141,14 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre # Set default configuration values # Configure granularity options - granularity_types = data_source_config.get('granularity', ["DOCUMENT", "ELEMENT"]) + granularity_types = data_source_config.get('granularity', ["DOCUMENT", "PAGE", "ELEMENT", "WORD"]) # Configure text format options text_format_types = data_source_config.get('text_format', ["MARKDOWN"]) - project_arn = None + # If project ARN is provided, use it, otherwise create a new project + project_arn = data_source_config.get('project_arn', None) + try: # there is a bug in AWS bedrock, it does not delete projects properly, so here # we generate random project name each time below @@ -161,40 +159,62 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre # bda_client.delete_data_automation_project(projectArn=project["projectArn"]) # time.sleep(2) # break - - # Create BDA project - logger.info(f"Creating BDA project") - project_name = f"bda-preprocessing-{uuid.uuid4().hex[:6]}" - project_response = bda_client.create_data_automation_project( - projectName=project_name, - projectDescription='Preprocessing multi data formats using bedrock data automation', - projectStage='DEVELOPMENT', - standardOutputConfiguration={ - "document": { - "extraction": { - "granularity": {"types": granularity_types}, - "boundingBox": {"state": "ENABLED"} + project_stage = "LIVE" + if not project_arn: + # Create BDA project + logger.info(f"Creating BDA project") + project_name = f"bda-preprocessing-{uuid.uuid4().hex[:6]}" + project_response = bda_client.create_data_automation_project( + projectName=project_name, + projectDescription='Preprocessing multi data formats using bedrock data automation', + projectStage='DEVELOPMENT', + standardOutputConfiguration={ + "document": { + "extraction": { + "granularity": {"types": granularity_types}, + "boundingBox": {"state": "ENABLED"} + }, + "generativeField": {"state": "ENABLED"}, + "outputFormat": { + "textFormat": {"types": text_format_types}, + "additionalFileFormat": {"state": "ENABLED"} + } }, - "generativeField": {"state": "ENABLED"}, - "outputFormat": { - "textFormat": {"types": text_format_types}, - "additionalFileFormat": {"state": "ENABLED"} + 'image': { + 'extraction': { + 'category': { + 'state': 'ENABLED', + 'types': [ + 'TEXT_DETECTION', + 'LOGOS', + ] + }, + 'boundingBox': {'state': 'ENABLED'} + }, + 'generativeField': { + 'state': 'ENABLED', + 'types': [ + 'IMAGE_SUMMARY' + ] + } } - }} - ) + } + ) + + project_arn = project_response['projectArn'] + project_stage = "DEVELOPMENT" + logger.info(f"Created BDA project {project_name} with ARN {project_arn}") - project_arn = project_response['projectArn'] - logger.info(f"Created BDA project {project_name} with ARN {project_arn}") # Get file names from S3 s3_pattern = re.compile(r'^s3://[a-z0-9\.-]{3,63}/.+$') if s3_pattern.match(input_uri): input_bucket = input_uri[5:].split("/")[0] - input_prefix = "/".join(input_uri[5:].split("/")[1:]) + input_prefix = "/".join(input_uri[5:].split("/")[1:])+"/" elif "//" in input_uri or input_uri.startswith("/") or input_uri.startswith("."): raise Exception("Input URI is not in the format of s3:///") else: input_bucket = input_uri.split("/")[0] - input_prefix = "/".join(input_uri.split("/")[1:]) + input_prefix = "/".join(input_uri.split("/")[1:])+"/" input_uri = "s3://" + input_uri if not s3_pattern.match(output_uri): @@ -225,7 +245,7 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre }, dataAutomationConfiguration={ 'dataAutomationProjectArn': project_arn, - 'stage': 'DEVELOPMENT' + 'stage': project_stage }, dataAutomationProfileArn=f'arn:aws:bedrock:{region}:{account_id}:data-automation-profile/us.data-automation-v1' ) @@ -256,6 +276,8 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre job['status'] = status time.sleep(1) max_timeout -= 1 + if max_timeout % 60 == 0: + logger.info(f"Waiting for BDA job {job_arn} on file {file_key} to complete") if all(job['status'] for job in job_results): break if max_timeout <= 0: @@ -279,7 +301,7 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre finally: # Clean up: Delete the BDA project after all jobs are completed - if project_arn: + if not data_source_config.get('project_arn', None): try: logger.info(f"Cleaning up BDA project: {project_arn}") delete_response = bda_client.delete_data_automation_project(projectArn=project_arn) @@ -454,6 +476,7 @@ def create_ingest( else: raise Exception("Data source not implemented") + logger.debug(f"Ingest template: USE GRAPH {graphname}\nBEGIN\n{ingest_template}\nEND\n") load_job_created = conn.gsql(f"USE GRAPH {graphname}\nBEGIN\n{ingest_template}\nEND\n") res = {"load_job_id": load_job_created.split(":")[1].strip(" [").strip(" ").strip(".").strip("]")} @@ -557,6 +580,8 @@ def ingest( else: s3_bucket = loader_info.file_path.split("/")[0] s3_prefix = "/".join(loader_info.file_path.split("/")[1:]) + if not s3_prefix.endswith("/"): + s3_prefix += "/" # --- Begin: S3 markdown extraction and TigerGraph loading --- # Possiblely we can download the files locally and then process them with next conn.runDocumentIngest() after it supports folder @@ -652,4 +677,4 @@ def ingest( } else: - raise Exception("Data source and file format combination not implemented") \ No newline at end of file + raise Exception("Data source and file format combination not implemented")