From 193bbe2c59ad726870a052baeb0900efe36e3d25 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Mon, 22 Sep 2025 11:15:26 -0700
Subject: [PATCH 01/31] GML-1994 GML-1992

---
 chat-history/Dockerfile                       |   2 +-
 common/config.py                              |   7 +
 common/metrics/tg_proxy.py                    |  13 +-
 common/utils.py                               | 183 ++++++++++++++++++
 configs/server_config.json                    |   2 +
 docker-compose.yml                            |   2 +-
 ecc/app/graphrag/graph_rag.py                 |  70 ++++++-
 ecc/app/graphrag/workers.py                   |   2 +-
 graphrag-ui/package.json                      |   1 +
 .../src/components/CustomChatMessage.tsx      |   4 +-
 graphrag-ui/tailwind.config.js                |   5 +-
 graphrag/app/agent/agent_generation.py        |  22 ++-
 graphrag/app/agent/agent_graph.py             |  19 +-
 .../app/supportai/retrievers/BaseRetriever.py |  15 ++
 .../retrievers/CommunityRetriever.py          |   1 +
 .../supportai/retrievers/HybridRetriever.py   |   1 +
 16 files changed, 317 insertions(+), 32 deletions(-)
 create mode 100644 common/utils.py

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 1f3e72c..ad6ae59 100644
--- a/common/config.py
+++ b/common/config.py
@@ -91,6 +91,13 @@
     raise Exception("embedding_service is not found in llm_config")
 embedding_dimension = embedding_config.get("dimensions", 1536)
 
+# Get context window size from llm_config
+# -1 means unlimited tokens (no truncation), otherwise use the specified limit
+token_limit = llm_config.get("token_limit", 1000000)    # 1000000 tokens is the default context window size for GPT-4
+
+# Get encoding name from llm_config
+encoding_name = llm_config.get("encoding_name", "cl100k_base")    # Default to GPT-4 encoding
+
 if graphrag_config is None:
     graphrag_config = {"reuse_embedding": True}
 if "chunker" not in graphrag_config:
diff --git a/common/metrics/tg_proxy.py b/common/metrics/tg_proxy.py
index 51ce7b5..132aca7 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,9 @@ def _runInstalledQuery(self, query_name, params, usePost=False):
             result = None
             while not result:
                 ret = self._tg_connection.checkQueryStatus(restppid)
-                LogWriter.info(f"Ret: {ret}")
+                #LogWriter.info(f"Ret: {ret}")
                 if not ret:
-                    time.sleep(0.1)
+                    time.sleep(1)
                     continue
                 if ret[0]["status"].lower() == "success":
                     LogWriter.info(
@@ -72,7 +73,7 @@ def _runInstalledQuery(self, query_name, params, usePost=False):
                             result = None
                             if i >= 9:
                                 raise e
-                            time.sleep(0.1)
+                            time.sleep(1)
                 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 +88,8 @@ def _runInstalledQuery(self, query_name, params, usePost=False):
                     raise Exception(
                         f"Query {query_name} with restppid {restppid} timed out"
                     )
+                else:
+                    time.sleep(1) # Still running, wait for 1 second
             success = True
         except Exception as e:
             LogWriter.error(f"Error running query {query_name}: {str(e)}")
diff --git a/common/utils.py b/common/utils.py
new file mode 100644
index 0000000..a7e481d
--- /dev/null
+++ b/common/utils.py
@@ -0,0 +1,183 @@
+# 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, config: dict = {}):
+        """
+        Initialize the token calculator.
+        
+        Args:
+            config: Configuration dictionary containing token_limit and model_name
+                               Use <= 0 for unlimited tokens (no truncation).
+            encoding_name: Tiktoken encoding name (default: cl100k_base for GPT-4)
+        """
+        self.max_context_tokens = config.get("token_limit", 1000000)
+        self.model_name = config.get("model_name", "gpt-4")
+        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_context_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict:
+        """
+        Truncate retrieved sources to fit within the token limit.
+        
+        Args:
+            sources_dict: Dictionary of retrieved source documents
+            max_tokens: Maximum number of tokens allowed (defaults to self.max_context_tokens)
+            
+        Returns:
+            Dictionary of truncated 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:
+            logger.info(f"Unlimited tokens enabled - returning all {len(sources_dict)} sources without truncation")
+            return sources_dict
+            
+        # Calculate how much to truncate
+        truncation_ratio = max_tokens / total_tokens
+        logger.info(f"Truncating context from {total_tokens} to {max_tokens} tokens (ratio: {truncation_ratio:.2f})")
+        
+        # Truncate each string value in the dictionary
+        truncated_sources = {}
+        for key, value in sources_dict.items():
+            if isinstance(value, str):
+                # Calculate how many characters to keep based on token ratio
+                char_limit = int(len(value) * truncation_ratio)
+                truncated_value = value[:char_limit]
+                if len(value) > char_limit:
+                    truncated_value += "..."
+                truncated_sources[key] = truncated_value
+            elif isinstance(value, list):
+                # Handle list of strings
+                truncated_list = []
+                for item in value:
+                    if isinstance(item, str):
+                        char_limit = int(len(item) * truncation_ratio)
+                        truncated_item = item[:char_limit]
+                        if len(item) > char_limit:
+                            truncated_item += "..."
+                        truncated_list.append(truncated_item)
+                    else:
+                        truncated_list.append(item)
+                truncated_sources[key] = truncated_list
+            elif isinstance(value, dict):
+                logger.info(f"Truncating sub-dictionary: {key}")
+                partial_tokens = self.count_tokens(value)
+                partial_ratio = partial_tokens / total_tokens
+                truncated_sources[key] = self.truncate_context_to_token_limit(value, int(max_tokens * partial_ratio))
+            else:
+                # Keep non-string values as-is
+                truncated_sources[key] = value
+                
+        # Verify the truncated result is within limits
+        final_tokens = self.count_tokens(truncated_sources)
+        logger.info(f"Final truncated context tokens: {final_tokens}")
+        
+        return truncated_sources
+
+    def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> 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
+        """
+        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_tokens(self, text: str | dict, max_tokens: int) -> 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_context_to_token_limit(text, max_tokens)
+        else:
+            return self.truncate_text_to_token_limit(text, max_tokens)
\ No newline at end of file
diff --git a/configs/server_config.json b/configs/server_config.json
index 41764e4..baa4592 100644
--- a/configs/server_config.json
+++ b/configs/server_config.json
@@ -14,6 +14,8 @@
     "chat_history_api": "http://chat-history:8002"
   },
   "llm_config": {
+    "token_limit": -1,
+    "encoding_name": "cl100k_base",
     "embedding_service": {
       "model_name": "text-embedding-3-small",
       "embedding_model_service": "openai",
diff --git a/docker-compose.yml b/docker-compose.yml
index b22745a..8be754b 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -74,7 +74,7 @@ services:
       - graphrag
 
   tigergraph:
-    image: tigergraph/community:4.2.0
+    image: tigergraph/community:4.2.1
     container_name: tigergraph
     platform: linux/amd64
     ports:
diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py
index 0e06e75..fc9068d 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,
@@ -109,11 +147,9 @@ async def chunk_docs(
                 raise
 
     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/workers.py b/ecc/app/graphrag/workers.py
index e439304..570211c 100644
--- a/ecc/app/graphrag/workers.py
+++ b/ecc/app/graphrag/workers.py
@@ -144,7 +144,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/graphrag-ui/package.json b/graphrag-ui/package.json
index b112a73..9983a70 100755
--- a/graphrag-ui/package.json
+++ b/graphrag-ui/package.json
@@ -37,6 +37,7 @@
     "reagraph": "^4.19.1",
     "tailwind-merge": "^2.3.0",
     "tailwindcss-animate": "^1.0.7",
+    "@tailwindcss/typography": "^0.5.18",
     "zod": "^3.23.8"
   },
   "devDependencies": {
diff --git a/graphrag-ui/src/components/CustomChatMessage.tsx b/graphrag-ui/src/components/CustomChatMessage.tsx
index 732f1cc..9fa02a0 100755
--- a/graphrag-ui/src/components/CustomChatMessage.tsx
+++ b/graphrag-ui/src/components/CustomChatMessage.tsx
@@ -85,14 +85,14 @@ export const CustomChatMessage: FC<IChatbotMessageProps> = ({
   return (
     <>
       {typeof message === "string" ? (
-        <div className="text-sm max-w-[230px] md:max-w-[80%] mt-7 mb-7">
+        <div className="prose dark:prose-invert text-sm max-w-[230px] md:max-w-[80%] mt-7 mb-7">
           <Markdown className="typewriter">{message}</Markdown>
         </div>
       ) : message.key === null ? (
         message
       ) : (
         <div className="flex flex-col w-full relative">
-          <div className="text-sm w-full mt-7 mb-7">
+          <div className="prose dark:prose-invert text-sm w-full mt-7 mb-7">
             {message.response_type === "progress" ? (
               <p className="graphrag-thinking typewriter">{message.content}</p>
             ) : (
diff --git a/graphrag-ui/tailwind.config.js b/graphrag-ui/tailwind.config.js
index 9128517..45cea7e 100644
--- a/graphrag-ui/tailwind.config.js
+++ b/graphrag-ui/tailwind.config.js
@@ -82,5 +82,8 @@ module.exports = {
       },
     },
   },
-  plugins: [require("tailwindcss-animate")],
+  plugins: [
+    require("tailwindcss-animate"),
+    require('@tailwindcss/typography'),
+  ],
 };
diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py
index 5391746..baf86af 100644
--- a/graphrag/app/agent/agent_generation.py
+++ b/graphrag/app/agent/agent_generation.py
@@ -12,6 +12,7 @@
 # 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
@@ -19,19 +20,22 @@
 from pydantic import BaseModel, Field
 from common.logs.logwriter import LogWriter
 from common.logs.log import req_id_cv
+from common.utils 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.")
+    generated_answer: str = Field(description="The generated answer to the question in markdown format. Make sure maintain a professional tone.")
     citation: list[str] = Field(description="The citation for the answer. List the information used.")
 
 class TigerGraphAgentGenerator:
     def __init__(self, llm_model):
         self.llm = llm_model
+        self.token_calculator = TokenCalculator(config=completion_config)
 
-    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 +46,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_tokens(context, max_context_tokens)
+                    logger.info(f"Truncated context from {context_tokens} to {self.token_calculator.count_tokens(context)} tokens")
+
         answer_parser = PydanticOutputParser(pydantic_object=GraphRAGAnswerOutput)
 
         prompt = PromptTemplate(
@@ -62,6 +77,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 b73b7aa..54b34ea 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -31,6 +31,7 @@
 
 from common.logs.log import req_id_cv
 from common.py_schemas import GraphRAGResponse, MapQuestionToSchemaResponse
+from common.config import graphrag_config
 
 logger = logging.getLogger(__name__)
 
@@ -228,9 +229,9 @@ def hybrid_search(self, state):
         step = retriever.search(
             state["question"],
             indices=["Document", "DocumentChunk", "Entity", "Relationship"],
-            top_k=5,
-            num_seen_min=2,
-            num_hops=3,
+            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),
         )
 
         query_name = "GraphRAG_Hybrid_Search"
@@ -259,7 +260,7 @@ 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"
@@ -287,7 +288,7 @@ 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"
@@ -314,9 +315,9 @@ 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"
@@ -372,7 +373,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(
diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py
index 647fb70..ba92379 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 import TokenCalculator
+from common.config import completion_config
 
 from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
@@ -27,6 +29,7 @@ def __init__(
         self.embedding_store = embedding_store
         self.embedding_store.set_graphname(connection.graphname)
         self.logger = logging.getLogger(__name__)
+        self.token_calculator = TokenCalculator(config=completion_config)
 
     def _install_query(self, query_name):
         with open(f"common/gsql/supportai/retrievers/{query_name}.gsql", "r") as f:
@@ -125,6 +128,18 @@ 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_context_to_token_limit(retrieved)
+                    self.logger.info(f"Truncated context from {retrieved_tokens} to {self.token_calculator.count_tokens(retrieved)} 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..ace15c5 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=1073741800,
                 usePost=True
             )  
         if len(res) > 1 and "verbose" in res[1]:

From 64882a2fab21b1a98738422059fd1088492fc814 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 23 Sep 2025 12:06:00 -0700
Subject: [PATCH 02/31] fix format

---
 common/utils.py                               | 34 +++++++++----------
 ecc/app/graphrag/graph_rag.py                 |  6 ++--
 .../app/supportai/retrievers/BaseRetriever.py |  1 -
 3 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/common/utils.py b/common/utils.py
index a7e481d..c84f209 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -21,11 +21,11 @@
 
 class TokenCalculator:
     """Utility class for token counting and text truncation operations."""
-    
+
     def __init__(self, config: dict = {}):
         """
         Initialize the token calculator.
-        
+
         Args:
             config: Configuration dictionary containing token_limit and model_name
                                Use <= 0 for unlimited tokens (no truncation).
@@ -70,31 +70,31 @@ def count_tokens(self, text: str | dict) -> int:
     def truncate_context_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict:
         """
         Truncate retrieved sources to fit within the token limit.
-        
+
         Args:
             sources_dict: Dictionary of retrieved source documents
             max_tokens: Maximum number of tokens allowed (defaults to self.max_context_tokens)
-            
+
         Returns:
             Dictionary of truncated 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:
             logger.info(f"Unlimited tokens enabled - returning all {len(sources_dict)} sources without truncation")
             return sources_dict
-            
+
         # Calculate how much to truncate
         truncation_ratio = max_tokens / total_tokens
         logger.info(f"Truncating context from {total_tokens} to {max_tokens} tokens (ratio: {truncation_ratio:.2f})")
-        
+
         # Truncate each string value in the dictionary
         truncated_sources = {}
         for key, value in sources_dict.items():
@@ -126,21 +126,21 @@ def truncate_context_to_token_limit(self, sources_dict: dict, max_tokens: Option
             else:
                 # Keep non-string values as-is
                 truncated_sources[key] = value
-                
+
         # Verify the truncated result is within limits
         final_tokens = self.count_tokens(truncated_sources)
         logger.info(f"Final truncated context tokens: {final_tokens}")
-        
+
         return truncated_sources
 
     def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> 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
         """
@@ -148,15 +148,15 @@ def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> str:
             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")
@@ -169,11 +169,11 @@ def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> str:
     def truncate_to_tokens(self, text: str | dict, max_tokens: int) -> 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
         """
diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py
index fc9068d..cfbab3d 100644
--- a/ecc/app/graphrag/graph_rag.py
+++ b/ecc/app/graphrag/graph_rag.py
@@ -95,7 +95,7 @@ async def stream_chunks(
         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:
@@ -115,11 +115,11 @@ async def stream_chunks(
                 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,
     docs_chan: Channel,
@@ -147,7 +147,7 @@ async def chunk_docs(
                 raise
 
     logger.info("Chunk Processing End")
-    
+
     logger.info("closing extract_chan")
     await extract_chan.put(None)
 
diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py
index ba92379..f8250eb 100644
--- a/graphrag/app/supportai/retrievers/BaseRetriever.py
+++ b/graphrag/app/supportai/retrievers/BaseRetriever.py
@@ -139,7 +139,6 @@ def _generate_response(self, question, retrieved, verbose):
                     retrieved = self.token_calculator.truncate_context_to_token_limit(retrieved)
                     self.logger.info(f"Truncated context from {retrieved_tokens} to {self.token_calculator.count_tokens(retrieved)} tokens")
 
-        
         model = self.llm_service.llm
         prompt = self.llm_service.supportai_response_prompt
 

From 5c73866db52af954d9ddab97bf6f297bbcd3b89a Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Wed, 24 Sep 2025 12:56:59 -0700
Subject: [PATCH 03/31] Update token_limit config

---
 common/config.py                              |  9 ++---
 common/embeddings/embedding_services.py       | 35 +++++++++++++++----
 common/utils.py                               | 26 +++++++-------
 graphrag/app/agent/agent_generation.py        | 10 +++---
 .../app/supportai/retrievers/BaseRetriever.py |  6 ++--
 5 files changed, 55 insertions(+), 31 deletions(-)

diff --git a/common/config.py b/common/config.py
index 42db679..0062091 100644
--- a/common/config.py
+++ b/common/config.py
@@ -93,10 +93,11 @@
 
 # Get context window size from llm_config
 # <=0 means unlimited tokens (no truncation), otherwise use the specified limit
-token_limit = llm_config.get("token_limit", 1000000)    # 1000000 tokens is the default context window size for GPT-4
-
-# Get encoding name from llm_config
-encoding_name = llm_config.get("encoding_name", "cl100k_base")    # Default to GPT-4 encoding
+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"]
 
 if graphrag_config is None:
     graphrag_config = {"reuse_embedding": True}
diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py
index 12cb9f2..75ba9c1 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 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,10 +162,10 @@ 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.embeddings import AzureOpenAIEmbeddings
 
-        self.embeddings = AzureOpenAIEmbeddings(deployment=config["azure_deployment"])
+        self.embeddings = AzureOpenAIEmbeddings(model=self.model_name, deployment=config["azure_deployment"])
 
 
 class OpenAI_Embedding(EmbeddingModel):
@@ -183,7 +205,7 @@ def __init__(self, config):
         import boto3
         from langchain_aws import BedrockEmbeddings
 
-        super().__init__(config=config, model_name=config["model_name"])
+        super().__init__(config=config, model_name=config.get("model_name", "amazon.titan-embed-text-v1"))
 
         client = boto3.client(
             "bedrock-runtime",
@@ -195,7 +217,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):
@@ -204,13 +226,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/utils.py b/common/utils.py
index c84f209..e3784f5 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -22,17 +22,17 @@
 class TokenCalculator:
     """Utility class for token counting and text truncation operations."""
 
-    def __init__(self, config: dict = {}):
+    def __init__(self, token_limit: int = 0, model_name: str = "gpt-4"):
         """
         Initialize the token calculator.
 
         Args:
-            config: Configuration dictionary containing token_limit and model_name
+            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).
-            encoding_name: Tiktoken encoding name (default: cl100k_base for GPT-4)
         """
-        self.max_context_tokens = config.get("token_limit", 1000000)
-        self.model_name = config.get("model_name", "gpt-4")
+        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:
@@ -67,12 +67,12 @@ def count_tokens(self, text: str | dict) -> int:
             # Fallback: rough estimation (1 token ≈ 4 characters for English text)
             return len(text) // 4
 
-    def truncate_context_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict:
+    def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict:
         """
-        Truncate retrieved sources to fit within the token limit.
+        Truncate dictionary to fit within the token limit.
 
         Args:
-            sources_dict: Dictionary of retrieved source documents
+            sources_dict: Dictionary to truncate
             max_tokens: Maximum number of tokens allowed (defaults to self.max_context_tokens)
 
         Returns:
@@ -88,7 +88,6 @@ def truncate_context_to_token_limit(self, sources_dict: dict, max_tokens: Option
 
         # If unlimited tokens is enabled, return all sources without truncation
         if self.is_unlimited_tokens() or max_tokens <= 0 or total_tokens <= max_tokens:
-            logger.info(f"Unlimited tokens enabled - returning all {len(sources_dict)} sources without truncation")
             return sources_dict
 
         # Calculate how much to truncate
@@ -133,7 +132,7 @@ def truncate_context_to_token_limit(self, sources_dict: dict, max_tokens: Option
 
         return truncated_sources
 
-    def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> str:
+    def truncate_text_to_token_limit(self, text: str, max_tokens: Optional[int] = None) -> str:
         """
         Truncate text to fit within the specified token limit.
 
@@ -144,6 +143,9 @@ def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> str:
         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:
@@ -166,7 +168,7 @@ def truncate_text_to_token_limit(self, text: str, max_tokens: int) -> str:
                 return text
             return text[:max_chars] + "..."
 
-    def truncate_to_tokens(self, text: str | dict, max_tokens: int) -> str:
+    def truncate_to_token_limit(self, text: str | dict, max_tokens: Optional[int] = None) -> str:
         """
         Truncate text to fit within the specified token limit.
 
@@ -178,6 +180,6 @@ def truncate_to_tokens(self, text: str | dict, max_tokens: int) -> str:
             Truncated text
         """
         if isinstance(text, dict):
-            return self.truncate_context_to_token_limit(text, max_tokens)
+            return self.truncate_dict_to_token_limit(text, max_tokens)
         else:
             return self.truncate_text_to_token_limit(text, max_tokens)
\ No newline at end of file
diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py
index baf86af..7634a58 100644
--- a/graphrag/app/agent/agent_generation.py
+++ b/graphrag/app/agent/agent_generation.py
@@ -27,13 +27,13 @@
 logger = logging.getLogger(__name__)
 
 class GraphRAGAnswerOutput(BaseModel):
-    generated_answer: str = Field(description="The generated answer to the question in markdown format. Make sure maintain a professional tone.")
+    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.")
 
 class TigerGraphAgentGenerator:
     def __init__(self, llm_model):
         self.llm = llm_model
-        self.token_calculator = TokenCalculator(config=completion_config)
+        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 | dict, query: str = "") -> dict:
         """Generate an answer based on the question and context.
@@ -51,11 +51,11 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") -
             # 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:
+            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_tokens(context, max_context_tokens)
-                    logger.info(f"Truncated context from {context_tokens} to {self.token_calculator.count_tokens(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)
 
diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py
index f8250eb..52ba61f 100644
--- a/graphrag/app/supportai/retrievers/BaseRetriever.py
+++ b/graphrag/app/supportai/retrievers/BaseRetriever.py
@@ -29,7 +29,7 @@ def __init__(
         self.embedding_store = embedding_store
         self.embedding_store.set_graphname(connection.graphname)
         self.logger = logging.getLogger(__name__)
-        self.token_calculator = TokenCalculator(config=completion_config)
+        self.token_calculator = TokenCalculator(token_limit=completion_config.get("token_limit"), model_name=completion_config.get("llm_model"))
 
     def _install_query(self, query_name):
         with open(f"common/gsql/supportai/retrievers/{query_name}.gsql", "r") as f:
@@ -136,8 +136,8 @@ def _generate_response(self, question, retrieved, verbose):
             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_context_to_token_limit(retrieved)
-                    self.logger.info(f"Truncated context from {retrieved_tokens} to {self.token_calculator.count_tokens(retrieved)} 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

From 0ecb0a0168bd1b937bac1159b9501278b15b133b Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Wed, 24 Sep 2025 13:28:28 -0700
Subject: [PATCH 04/31] Fix typo

---
 graphrag/app/agent/agent_generation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py
index 7634a58..394a6f8 100644
--- a/graphrag/app/agent/agent_generation.py
+++ b/graphrag/app/agent/agent_generation.py
@@ -51,7 +51,7 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") -
             # 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:
+            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)

From 68299ca95c73453a9140cb458f314ccf5aa4684d Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Wed, 24 Sep 2025 14:51:54 -0700
Subject: [PATCH 05/31] Remove hard-coded default graph

---
 graphrag-ui/src/components/Bot.tsx      | 11 +++++-
 graphrag-ui/src/components/Contexts.tsx |  2 +-
 graphrag-ui/src/components/Start.tsx    | 51 ++++++++++++++-----------
 3 files changed, 38 insertions(+), 26 deletions(-)

diff --git a/graphrag-ui/src/components/Bot.tsx b/graphrag-ui/src/components/Bot.tsx
index f80dd80..ed033dd 100644
--- a/graphrag-ui/src/components/Bot.tsx
+++ b/graphrag-ui/src/components/Bot.tsx
@@ -22,7 +22,7 @@ import {
 const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getConversationId?:any }) => {
   const [store, setStore] = useState<any>();
   const [currentDate, setCurrentDate] = useState('');
-  const [selectedGraph, setSelectedGraph] = useState(localStorage.getItem("selectedGraph") || 'TigerGraphRAG');
+  const [selectedGraph, setSelectedGraph] = useState(localStorage.getItem("selectedGraph") || '');
   const [ragPattern, setRagPattern] = useState(localStorage.getItem("ragPattern") || 'Hybrid Search');
   const navigate = useNavigate();
 
@@ -30,6 +30,13 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
     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);
+    }
+
     const date = new Date();
     const options: Intl.DateTimeFormatOptions = { year: 'numeric', month: 'long', day: 'numeric', weekday: 'long' };
     const formattedDate = date.toLocaleDateString('en-US', options);
@@ -98,7 +105,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
               <DropdownMenuGroup>
                 {store?.graphs.map((f, i) => (
                   <DropdownMenuItem key={i} onSelect={() => handleSelect(f)}>
-                    <span>{f}</span>
+                    <span>{f.replace("BarClays", "Barclays")}</span>
                   </DropdownMenuItem>
                 ))}
               </DropdownMenuGroup>
diff --git a/graphrag-ui/src/components/Contexts.tsx b/graphrag-ui/src/components/Contexts.tsx
index 142fc08..0d03f88 100644
--- a/graphrag-ui/src/components/Contexts.tsx
+++ b/graphrag-ui/src/components/Contexts.tsx
@@ -1,3 +1,3 @@
 import React, { createContext } from "react";
 
-export const SelectedGraphContext = createContext("pyTigerGraphRAG");
\ No newline at end of file
+export const SelectedGraphContext = createContext("");
\ No newline at end of file
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;

From 3bda3948cce72d948d28fce5b9460446638ed508 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Wed, 24 Sep 2025 15:42:17 -0700
Subject: [PATCH 06/31] Remove hard-coded default graph

---
 graphrag-ui/src/components/Bot.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphrag-ui/src/components/Bot.tsx b/graphrag-ui/src/components/Bot.tsx
index ed033dd..6b0ba56 100644
--- a/graphrag-ui/src/components/Bot.tsx
+++ b/graphrag-ui/src/components/Bot.tsx
@@ -95,7 +95,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
                   className="!h-[48px] !outline-b !outline-gray-300 dark:!outline-[#3D3D3D] h-[70px] flex justify-end items-center bg-white dark:bg-background z-50 rounded-tr-lg"
                 >
                   <img src="/graph-icon.svg" alt="" className="mr-2" />
-                  {selectedGraph} <MdKeyboardArrowDown className="text-2xl" />
+                  {selectedGraph.replace("BarClays", "Barclays")} <MdKeyboardArrowDown className="text-2xl" />
                 </Button>
               </DropdownMenuTrigger>
 

From 061c4a67f7f001b2e1ab27574fae0da069d29f33 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Wed, 24 Sep 2025 20:58:39 -0700
Subject: [PATCH 07/31] Fix bug

---
 common/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/common/utils.py b/common/utils.py
index e3784f5..db1d745 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -121,7 +121,7 @@ def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[
                 logger.info(f"Truncating sub-dictionary: {key}")
                 partial_tokens = self.count_tokens(value)
                 partial_ratio = partial_tokens / total_tokens
-                truncated_sources[key] = self.truncate_context_to_token_limit(value, int(max_tokens * partial_ratio))
+                truncated_sources[key] = self.truncate_dict_to_token_limit(value, int(max_tokens * partial_ratio))
             else:
                 # Keep non-string values as-is
                 truncated_sources[key] = value

From c14ab8764929cb39f8bf22183916e00da9c73826 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 25 Sep 2025 13:37:29 -0700
Subject: [PATCH 08/31] Revise token truncation logic

---
 chat-history/config/config_test.go            |   4 +-
 common/utils.py                               | 122 ++++++++++++------
 graphrag/app/routers/ui.py                    |   5 +-
 .../app/supportai/retrievers/BaseRetriever.py |   1 +
 graphrag/app/supportai/supportai.py           |   8 +-
 5 files changed, 89 insertions(+), 51 deletions(-)

diff --git a/chat-history/config/config_test.go b/chat-history/config/config_test.go
index 8caeef9..76cdcb0 100644
--- a/chat-history/config/config_test.go
+++ b/chat-history/config/config_test.go
@@ -10,7 +10,7 @@ func TestLoadConfig(t *testing.T) {
 	tgConfigPath := setup(t)
 
 	cfg, err := LoadConfig(map[string]string{
-		"tgconfig":   tgConfigPath,
+		"tgconfig": tgConfigPath,
 	})
 	if err != nil {
 		t.Fatal(err)
@@ -29,7 +29,7 @@ func TestLoadConfig(t *testing.T) {
 	}
 }
 
-func setup(t *testing.T) (string, string) {
+func setup(t *testing.T) string {
 	tmp := t.TempDir()
 
 	tgConfigPath := fmt.Sprintf("%s/%s", tmp, "server_config.json")
diff --git a/common/utils.py b/common/utils.py
index db1d745..603219f 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -69,14 +69,15 @@ def count_tokens(self, text: str | dict) -> int:
 
     def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict:
         """
-        Truncate dictionary to fit within the token limit.
+        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 truncated sources that fit within the token limit
+            Dictionary of sources that fit within the token limit
         """
         if max_tokens is None:
             max_tokens = self.max_context_tokens
@@ -90,46 +91,83 @@ def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[
         if self.is_unlimited_tokens() or max_tokens <= 0 or total_tokens <= max_tokens:
             return sources_dict
 
-        # Calculate how much to truncate
-        truncation_ratio = max_tokens / total_tokens
-        logger.info(f"Truncating context from {total_tokens} to {max_tokens} tokens (ratio: {truncation_ratio:.2f})")
-
-        # Truncate each string value in the dictionary
+        # Convert dict to list of (key, value) pairs for processing
+        items = list(sources_dict.items())
         truncated_sources = {}
-        for key, value in sources_dict.items():
-            if isinstance(value, str):
-                # Calculate how many characters to keep based on token ratio
-                char_limit = int(len(value) * truncation_ratio)
-                truncated_value = value[:char_limit]
-                if len(value) > char_limit:
-                    truncated_value += "..."
-                truncated_sources[key] = truncated_value
-            elif isinstance(value, list):
-                # Handle list of strings
-                truncated_list = []
-                for item in value:
-                    if isinstance(item, str):
-                        char_limit = int(len(item) * truncation_ratio)
-                        truncated_item = item[:char_limit]
-                        if len(item) > char_limit:
-                            truncated_item += "..."
-                        truncated_list.append(truncated_item)
-                    else:
-                        truncated_list.append(item)
-                truncated_sources[key] = truncated_list
-            elif isinstance(value, dict):
-                logger.info(f"Truncating sub-dictionary: {key}")
-                partial_tokens = self.count_tokens(value)
-                partial_ratio = partial_tokens / total_tokens
-                truncated_sources[key] = self.truncate_dict_to_token_limit(value, int(max_tokens * partial_ratio))
-            else:
-                # Keep non-string values as-is
-                truncated_sources[key] = value
+        current_tokens = 0
 
-        # Verify the truncated result is within limits
-        final_tokens = self.count_tokens(truncated_sources)
-        logger.info(f"Final truncated context tokens: {final_tokens}")
+        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:
@@ -156,8 +194,8 @@ def truncate_text_to_token_limit(self, text: str, max_tokens: Optional[int] = No
             truncated_text = self.token_encoding.decode(truncated_tokens)
 
             # Add ellipsis to indicate truncation
-            if len(tokens) > max_tokens:
-                truncated_text += "..."
+            #if len(tokens) > max_tokens:
+            #    truncated_text += "..."
 
             return truncated_text
         except Exception as e:
@@ -166,7 +204,7 @@ def truncate_text_to_token_limit(self, text: str, max_tokens: Optional[int] = No
             max_chars = max_tokens * 4
             if len(text) <= max_chars:
                 return text
-            return text[:max_chars] + "..."
+            return text[:max_chars] #+ "..."
 
     def truncate_to_token_limit(self, text: str | dict, max_tokens: Optional[int] = None) -> str:
         """
diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py
index eb8e05d..1db67aa 100644
--- a/graphrag/app/routers/ui.py
+++ b/graphrag/app/routers/ui.py
@@ -338,13 +338,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({
diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py
index 52ba61f..b6f6583 100644
--- a/graphrag/app/supportai/retrievers/BaseRetriever.py
+++ b/graphrag/app/supportai/retrievers/BaseRetriever.py
@@ -32,6 +32,7 @@ def __init__(
         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(
diff --git a/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py
index 0b5fb70..c2a8a06 100644
--- a/graphrag/app/supportai/supportai.py
+++ b/graphrag/app/supportai/supportai.py
@@ -32,10 +32,10 @@ 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",
+        #"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:

From 9d3dfb2e0efe00ed4b79c5b04fb47379c46cf4d2 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 25 Sep 2025 17:07:07 -0700
Subject: [PATCH 09/31] Revise token truncation logic

---
 README.md       | 89 ++++++++++++++++++++++++++++++++++++-------------
 common/utils.py |  4 +--
 2 files changed, 68 insertions(+), 25 deletions(-)

diff --git a/README.md b/README.md
index 115bad4..c204eff 100644
--- a/README.md
+++ b/README.md
@@ -48,24 +48,6 @@ The quickest way to access TigerGraph GraphRAG is to deploy its docker image wit
 * TigerGraph DB 4.2+.
 * API key of your LLM provider. (An LLM provider refers to a company or organization that offers Large Language Models (LLMs) as a service. The API key verifies the identity of the requester, ensuring that the request is coming from a registered and authorized user or application.) Currently, GraphRAG supports the following LLM providers: OpenAI, Azure OpenAI, GCP, AWS Bedrock.
 
-#### Deploy with Kubernetes
-* Step 1: Get kubernetes deployment file
-  - Download the [graphrag-k8s.yml](https://raw.githubusercontent.com/tigergraph/ecosys/refs/heads/master/tutorials/graphrag/graphrag-k8s.yml) file directly
-
-* Step 2: Set up configurations
-  Next, in the same directory as the Kubernetes deployment file is in, create a `configs` directory and download the following configuration files:
-  * [configs/server_config.json](https://raw.githubusercontent.com/tigergraph/ecosys/refs/heads/master/tutorials/graphrag/configs/server_config.json)
-
-  Update the TigerGraph database information, LLM API keys and other configs accordingly.
-
-* Step 3: Start all services
-  Replace `/path/to/graphrag/configs` with the absolute path of the `configs` folder inside `graphrag-k8s.yml`, and update the TigerGraph database information and other configs accordingly.
-
-  Now, simply run `kubectl apply -f graphrag-k8s.yml` and wait for all the services to start.
-
-> Note: Nginx Ingress should be installed using `kubectl apply -f https://raw.githubusercontent.com/kubernetes/ingress-nginx/controller-v1.2.1/deploy/static/provider/cloud/deploy.yaml`
-
-
 #### Deploy with Docker Compose
 * Step 1: Get docker-compose file
   - Download the [docker-compose.yml](https://raw.githubusercontent.com/tigergraph/ecosys/refs/heads/master/tutorials/graphrag/docker-compose.yml) file directly
@@ -87,7 +69,13 @@ The quickest way to access TigerGraph GraphRAG is to deploy its docker image wit
     └── docker-compose.yml
 ```
 
-* Step 3 (Optional): Configure Logging Level in Dockerfile
+* Step 3: Adjust configurations
+
+Edit `llm_config` section of `configs/server_config.json` and replace `<YOUR_OPENAI_API_KEY>` to your own OPENAI_API_KEY. 
+
+> If desired, you can also change the model to be used for the embedding service and completion service to your preferred models to adjust the output from the LLM service.
+
+* Step 4 (Optional): Configure Logging Level in Dockerfile
 
 To configure the logging level of the service, edit the Docker Compose file.
 
@@ -111,18 +99,73 @@ This line can be changed to support different logging levels.
 | `DEBUG_PII` | Finer-grained information that could potentially include `PII`, such as a user’s question, the complete function call (with parameters), and the LLM’s natural language response. |
 | NOTSET | All messages are processed. |
 
+* Step 5: Start all services
+
+Uncomment `tigergraph` section from `docker-compose.yml` if it's commented out. Please follow the [instructions](https://github.com/tigergraph/ecosys/blob/master/tutorials/GSQL.md#set-up-environment) to download TigerGraph docker image.
+
+Now, simply run `docker compose up -d` and wait for all the services to start.
+
+> Note: `graphrag` container will be down if TigerGraph service is not ready. Log into the `tigergraph` container, bring up tigergraph services and rerun `docker compose up -d` should resolve the issue.
+
+* Step 6: Stop all services (when needed)
+Run command `docker compose down` and wait for all the service containers to stopped and removed.
+
+#### Use Standalone TigerGraph instance (Optional)
+
+> **_Note:_** Vector feature is available in both TigerGraph Community Edition 4.2.0+ and Enterprise Edition 4.2.0+.
+
+If you prefer to start a TigerGraph Community Edition instance without a license key, please make sure the container can be accessed from the GraphRAG containers by add `--network graphrag_default`:
+```
+docker run -d -p 14240:14240 --name tigergraph --ulimit nofile=1000000:1000000 --init --network graphrag_default -t tigergraph/community:4.2.1
+```
+
+> Use **tigergraph/tigergraph:4.2.1** if Enterprise Edition is preferred.
+> Setting up **DNS** or `/etc/hosts` properly is an alternative solution to ensure contains can connect to each other.
+> Or modify`hostname` in `db_config` section of `configs/server_config.json` and replace `http://tigergraph` to your tigergraph container IP address, e.g., `http://172.19.0.2`. 
+
+Check the service status with the following commands:
+```
+docker exec -it tigergraph /bin/bash
+gadmin status
+gadmin start all
+```
+
+After using the database, and you want to shutdown it, use the following shell commmand
+```
+gadmin stop all
+```
+
+#### Deploy with Kubernetes
+> Note: TigerGraph instance should be deployed separately in advance and accessible from the kubernetes env
+
+* Step 1: Get kubernetes deployment file
+  - Download the [graphrag-k8s.yml](https://raw.githubusercontent.com/tigergraph/ecosys/refs/heads/master/tutorials/graphrag/graphrag-k8s.yml) file directly
+
+* Step 2: Set up configurations
+  Next, in the same directory as the Kubernetes deployment file is in, create a `configs` directory and download the following configuration files:
+  * [configs/server_config.json](https://raw.githubusercontent.com/tigergraph/ecosys/refs/heads/master/tutorials/graphrag/configs/server_config.json)
+
+* Step 3: Update the TigerGraph database information, LLM API keys and other configs accordingly in `configs/server_config.json`.
 
 * Step 4: Start all services
+  Replace `/path/to/graphrag/configs` with the absolute path of the `configs` folder inside `graphrag-k8s.yml`, and update the TigerGraph database information and other configs accordingly.
 
-  Uncomment `tigergraph` section from `docker-compose.yml` if it's commented out. Please follow the [instructions](https://github.com/tigergraph/ecosys/blob/master/tutorials/GSQL.md#set-up-environment) to download TigerGraph docker image.
+  Now, simply run `kubectl apply -f graphrag-k8s.yml` and wait for all the services in the deployment to be started.
 
-  Now, simply run `docker compose up -d` and wait for all the services to start.
+> Note: Nginx Ingress should be installed using `kubectl apply -f https://raw.githubusercontent.com/kubernetes/ingress-nginx/controller-v1.2.1/deploy/static/provider/cloud/deploy.yaml`
 
-> Note: `graphrag` container will be down if TigerGraph service is not ready. Log into the `tigergraph` container, bring up tigergraph services and rerun `docker compose up -d` should resolve the issue.
+* Step 5: Stop all services (when needed)
+  Run `kubectl delete -f graphrag-k8s.yml` and wait for all the services in the deployment to be deleted.
 
-## Data Ingestion
+> Note: Nginx Ingress should be deleted using `kubectl delete -f https://raw.githubusercontent.com/kubernetes/ingress-nginx/controller-v1.2.1/deploy/static/provider/cloud/deploy.yaml` if port 80 needs to be released
+
+### Data Ingestion
 For data ingestion, please follow the ![GraphRAG Demo Notebook](./docs/notebooks/GraphRAGDemo.ipynb)
 
+### Access GraphRAG ChatBot from UI
+Open your browser to access `http://localhost` to access GraphRAG Chat. If you're accessing GraphRAG UI remotely, use the machine's hostname of IP address instead.
+
+
 ## Detailed Configurations
 
 ### LLM provider configuration
diff --git a/common/utils.py b/common/utils.py
index 603219f..a810fb5 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -92,7 +92,7 @@ def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[
             return sources_dict
 
         # Convert dict to list of (key, value) pairs for processing
-        items = list(sources_dict.items())
+        items = sorted(source_dict.items(), key=lambda x: (".png" not in x, -len(v)))
         truncated_sources = {}
         current_tokens = 0
 
@@ -220,4 +220,4 @@ def truncate_to_token_limit(self, text: str | dict, max_tokens: Optional[int] =
         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)
\ No newline at end of file
+            return self.truncate_text_to_token_limit(text, max_tokens)

From c10bd02d0164483669510cdbcaa2737129e61951 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 25 Sep 2025 17:33:57 -0700
Subject: [PATCH 10/31] fix bug

---
 common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt | 4 ++--
 common/utils.py                                              | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
index 108047b..4003b66 100644
--- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
+++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
@@ -3,8 +3,8 @@ Given the following context in JSON format, 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 to extract and include the image links in markdown syntax in the generated_answer if its summary is used.
-Use markdown syntax to geneate the answer, including title, bulleted or numbered list, images and tables if any.
-Always only return a JSON contains the answer and otheh required fields. Assign an empty value to the field if you cannot determine it.
+Use markdown syntax to geneate the answer, including title, bulleted or numbered list, images and tables if any, and keep it compact.
+Always only return a JSON contains the answer and other required fields. Assign an empty value to the field if you cannot determine it.
 
 
 Question: {question}
diff --git a/common/utils.py b/common/utils.py
index a810fb5..d828cbf 100644
--- a/common/utils.py
+++ b/common/utils.py
@@ -92,7 +92,7 @@ def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[
             return sources_dict
 
         # Convert dict to list of (key, value) pairs for processing
-        items = sorted(source_dict.items(), key=lambda x: (".png" not in x, -len(v)))
+        items = sorted(sources_dict.items(), key=lambda x: (".png" not in x, -len(x)))
         truncated_sources = {}
         current_tokens = 0
 

From 0112e2f5f8232aad49a3cd409982ba8807650139 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 25 Sep 2025 20:33:31 -0700
Subject: [PATCH 11/31] Revise prompt and support separate chat_model for
 chatbot

---
 common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt | 4 ++--
 graphrag/app/agent/agent.py                                  | 3 +++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
index 4003b66..31d9d5c 100644
--- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
+++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
@@ -3,8 +3,8 @@ Given the following context in JSON format, 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 to extract and include the image links in markdown syntax in the generated_answer if its summary is used.
-Use markdown syntax to geneate the answer, including title, bulleted or numbered list, images and tables if any, and keep it compact.
-Always only return a JSON contains the answer and other required fields. Assign an empty value to the field if you cannot determine it.
+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.
+Always make sure return a JSON contains the answer and other required fields after validation. Assign an empty value to the field if you cannot determine it.
 
 
 Question: {question}
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"])

From f26175aad0fec237d917bbaa6d0fbcbc7008dc28 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 25 Sep 2025 21:03:27 -0700
Subject: [PATCH 12/31] Make citation optional as some models cannot handle it
 well

---
 graphrag/app/agent/agent_generation.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py
index 394a6f8..db7cbce 100644
--- a/graphrag/app/agent/agent_generation.py
+++ b/graphrag/app/agent/agent_generation.py
@@ -17,6 +17,7 @@
 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
@@ -28,7 +29,7 @@
 
 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=None)
 
 class TigerGraphAgentGenerator:
     def __init__(self, llm_model):

From 749dc1a73663092ad58a148dd003ad85b4600665 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 25 Sep 2025 21:24:47 -0700
Subject: [PATCH 13/31] Fix bug

---
 graphrag/app/agent/agent_generation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py
index db7cbce..31f7a2f 100644
--- a/graphrag/app/agent/agent_generation.py
+++ b/graphrag/app/agent/agent_generation.py
@@ -29,7 +29,7 @@
 
 class GraphRAGAnswerOutput(BaseModel):
     generated_answer: str = Field(description="The generated answer to the question. Make sure maintain a professional tone.")
-    citation: Optional[list[str]] = Field(description="The citation for the answer. List the metadata of the parts of the context used.", default=None)
+    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):

From 9b774bff0181650afc69af0aba76e2956571570e Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 26 Sep 2025 18:20:32 -0700
Subject: [PATCH 14/31] Adjust prompt

---
 .../aws_bedrock_claude3haiku/chatbot_response.txt    |  6 ++++--
 common/prompts/google_gemini/chatbot_response.txt    |  1 +
 common/prompts/openai_gpt4/chatbot_response.txt      |  3 +++
 graphrag-ui/src/index.css                            |  4 ++--
 graphrag/app/agent/agent_graph.py                    | 12 ++++++++----
 5 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
index 31d9d5c..4267b1f 100644
--- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
+++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
@@ -2,12 +2,14 @@ You are a highly efficient and empathetic AI-powered assistant in JSON parsing a
 Given the following context in JSON format, 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 to extract and include the image links in markdown syntax in the generated_answer if its summary is used.
+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.
-Always make sure return a JSON contains the answer and other required fields after validation. Assign an empty value to the field if you cannot determine it.
+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 after validation. Assign an empty value to the field if you cannot determine it.
 
 
 Question: {question}
 Answer: {context}
+Query: {query}
 Format: {format_instructions}
 
diff --git a/common/prompts/google_gemini/chatbot_response.txt b/common/prompts/google_gemini/chatbot_response.txt
index 5865304..c83cbd5 100644
--- a/common/prompts/google_gemini/chatbot_response.txt
+++ b/common/prompts/google_gemini/chatbot_response.txt
@@ -29,4 +29,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 7d02dff..c2b8333 100644
--- a/common/prompts/openai_gpt4/chatbot_response.txt
+++ b/common/prompts/openai_gpt4/chatbot_response.txt
@@ -19,6 +19,9 @@ Format your answer using Markdown. Organize the content into paragraphs, bullete
 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 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/graphrag-ui/src/index.css b/graphrag-ui/src/index.css
index f4c43ee..d2dc878 100755
--- a/graphrag-ui/src/index.css
+++ b/graphrag-ui/src/index.css
@@ -563,12 +563,12 @@
 .typewriter {
   --lines: 500;
   --lineHeight: 1.5rem;
-  --timePerLine: 1s;
+  --timePerLine: 0.2s;
   --widthCh: 22;
   /* --width: calc(var(--widthCh) * 1ch);  */
   /* do not touch the time property!!! */
   --time: calc(var(--lines) * var(--timePerLine));
-  animation: grow var(--time) steps(var(--lines));
+  animation: grow var(--time) steps(var(--lines), end);
   animation-fill-mode: forwards;
   background: var(--bgColor);
   line-height: var(--lineHeight);
diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index 77f1e42..8c00b39 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -129,7 +129,7 @@ def apologize(self, state):
         """
         self.emit_progress(DONE)
         state["answer"] = GraphRAGResponse(
-            natural_language_response="I'm sorry, I don't know the answer to that question. Please try rephrasing your question.",
+            natural_language_response="I'm sorry, there isn't enough context to answer your question. Please try rephrasing it.",
             answered_question=False,
             response_type="error",
             query_sources={"error": True, "error_history": state["error_history"]},
@@ -390,7 +390,11 @@ def generate_answer(self, state):
         if state["lookup_source"] == "supportai":
             import re
 
-            citations = [re.sub(r"_chunk_\d+", "", x) for x in answer.citation]
+            if answer.citation:
+                citations = [re.sub(r"_chunk_\d+", "", x) for x in answer.citation]
+            else:
+                citations = [re.sub(r"_chunk_\d+", "", x) for x in set(state["context"]["result"]["final_retrieval"].keys())]
+
             state["context"]["reasoning"] = list(set(citations))
 
         try:
@@ -427,7 +431,7 @@ def replace_s3_urls_with_presigned(self, content, expires_in=3600):
             Any: Content with S3 URLs replaced by presigned URLs (same type as input).
         """
 
-        s3_url_pattern = r's3://([\w\-.]+)/([\w\-\./]+)'
+        s3_url_pattern = r'\(s3://([^/]+)/([^\)]+)\)'
         s3 = boto3.client('s3')
 
         def presign(match):
@@ -438,7 +442,7 @@ def presign(match):
                     Params={'Bucket': bucket, 'Key': key},
                     ExpiresIn=expires_in
                 )
-                return url
+                return f"({url})"
             except Exception as e:
                 logger.error(f"Failed to presign S3 url for s3://{bucket}/{key}: {e}")
                 return match.group(0)

From 7ff4c4e1980ca4eaee115258f13d77558ce8ad67 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 26 Sep 2025 23:45:26 -0700
Subject: [PATCH 15/31] Add StreamChunkContent.gsql

---
 common/gsql/graphrag/StreamChunkContent.gsql | 8 ++++++++
 ecc/app/graphrag/util.py                     | 1 +
 ecc/app/supportai/util.py                    | 1 +
 3 files changed, 10 insertions(+)
 create mode 100644 common/gsql/graphrag/StreamChunkContent.gsql

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<DocumentChunk> 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/ecc/app/graphrag/util.py b/ecc/app/graphrag/util.py
index 6a7534a..a053f1b 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",
diff --git a/ecc/app/supportai/util.py b/ecc/app/supportai/util.py
index 24d653d..6de5fc3 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)
 

From 2f64336b5cec879c9e85c0f8728601e205e2827d Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 30 Sep 2025 17:11:44 -0700
Subject: [PATCH 16/31] Update default search to vector search

---
 graphrag-ui/src/components/Bot.tsx |  4 ++--
 graphrag/app/agent/agent_graph.py  | 12 ++++++------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/graphrag-ui/src/components/Bot.tsx b/graphrag-ui/src/components/Bot.tsx
index 6b0ba56..e800b42 100644
--- a/graphrag-ui/src/components/Bot.tsx
+++ b/graphrag-ui/src/components/Bot.tsx
@@ -95,7 +95,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
                   className="!h-[48px] !outline-b !outline-gray-300 dark:!outline-[#3D3D3D] h-[70px] flex justify-end items-center bg-white dark:bg-background z-50 rounded-tr-lg"
                 >
                   <img src="/graph-icon.svg" alt="" className="mr-2" />
-                  {selectedGraph.replace("BarClays", "Barclays")} <MdKeyboardArrowDown className="text-2xl" />
+                  {selectedGraph} <MdKeyboardArrowDown className="text-2xl" />
                 </Button>
               </DropdownMenuTrigger>
 
@@ -105,7 +105,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
               <DropdownMenuGroup>
                 {store?.graphs.map((f, i) => (
                   <DropdownMenuItem key={i} onSelect={() => handleSelect(f)}>
-                    <span>{f.replace("BarClays", "Barclays")}</span>
+                    <span>{f}</span>
                   </DropdownMenuItem>
                 ))}
               </DropdownMenuGroup>
diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index 8c00b39..d5e58dc 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -84,9 +84,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):
@@ -237,7 +237,7 @@ def hybrid_search(self, state):
             num_hops=graphrag_config.get("num_hops", 2),
         )
 
-        query_name = "GraphRAG_Hybrid_Search"
+        query_name = "GraphRAG_Hybrid_Vector_Search"
         state["context"] = {
             "function_call": query_name,
             "result": step[0],
@@ -266,7 +266,7 @@ def similarity_search(self, state):
             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],
@@ -294,7 +294,7 @@ def sibling_search(self, state):
             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],
@@ -323,7 +323,7 @@ def community_search(self, state):
             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],

From ef47297bfad46f6c2137e459e230902513521e18 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 30 Sep 2025 19:55:38 -0700
Subject: [PATCH 17/31] Upgrade fastapi

---
 common/requirements.txt                    |  4 ++--
 graphrag-ui/src/actions/ActionProvider.tsx | 17 +++++++++++++++--
 graphrag/app/routers/ui.py                 | 19 +++++++++++++++----
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/common/requirements.txt b/common/requirements.txt
index 77be241..92dbb7c 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
@@ -151,7 +151,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/graphrag-ui/src/actions/ActionProvider.tsx b/graphrag-ui/src/actions/ActionProvider.tsx
index a268b5a..9307964 100644
--- a/graphrag-ui/src/actions/ActionProvider.tsx
+++ b/graphrag-ui/src/actions/ActionProvider.tsx
@@ -84,10 +84,12 @@ const ActionProvider: React.FC<ActionProviderProps> = ({
   const [messageHistory, setMessageHistory] = useState<MessageEvent<Message>[]>(
     [],
   );
-  const { sendMessage, lastMessage, readyState } = useWebSocket(WS_URL, {
+  const { sendMessage, lastMessage, readyState, getWebSocket } = 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");
@@ -98,6 +100,17 @@ const ActionProvider: React.FC<ActionProviderProps> = ({
       console.log("WebSocket connection " + conversationIdToSend + " established to " + WS_URL);
       sendMessage(conversationIdToSend);
     },
+    onError: (error) => {
+      console.error("WebSocket error:", error);
+    },
+    onClose: (event) => {
+      console.error("WebSocket closed:", event.code, event.reason);
+      console.log("WebSocket state:", getWebSocket()?.readyState);
+    },
+    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/app/routers/ui.py b/graphrag/app/routers/ui.py
index 1db67aa..13fdaf3 100644
--- a/graphrag/app/routers/ui.py
+++ b/graphrag/app/routers/ui.py
@@ -476,10 +476,21 @@ 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()

From 84a8ca7c4d9146b4d6187abc007e25c7225335b1 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Wed, 1 Oct 2025 21:33:17 -0700
Subject: [PATCH 18/31] Bug fix

---
 .../embeddings/tigergraph_embedding_store.py  |  2 +-
 common/metrics/tg_proxy.py                    |  2 +-
 graphrag-ui/src/actions/ActionProvider.tsx    | 12 +++----
 graphrag-ui/src/components/Bot.tsx            | 31 ++++++++++++-------
 graphrag-ui/src/components/Contexts.tsx       |  3 +-
 graphrag/app/agent/agent_graph.py             |  2 ++
 graphrag/app/routers/ui.py                    |  8 +++--
 7 files changed, 37 insertions(+), 23 deletions(-)

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/metrics/tg_proxy.py b/common/metrics/tg_proxy.py
index 132aca7..b5aa75d 100644
--- a/common/metrics/tg_proxy.py
+++ b/common/metrics/tg_proxy.py
@@ -73,7 +73,7 @@ def _runInstalledQuery(self, query_name, params, sizeLimit=None, usePost=False):
                             result = None
                             if i >= 9:
                                 raise e
-                            time.sleep(1)
+                        time.sleep(1)
                 elif ret[0]["status"].lower() == "aborted":
                     LogWriter.error(
                         f"request_id={req_id_cv.get()} query {query_name} with RESTPP ID {restppid} aborted"
diff --git a/graphrag-ui/src/actions/ActionProvider.tsx b/graphrag-ui/src/actions/ActionProvider.tsx
index 9307964..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,11 +80,12 @@ const ActionProvider: React.FC<ActionProviderProps> = ({
   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<MessageEvent<Message>[]>(
     [],
   );
-  const { sendMessage, lastMessage, readyState, getWebSocket } = useWebSocket(WS_URL, {
+  const { sendMessage, lastMessage, readyState } = useWebSocket(WS_URL, {
     onOpen: () => {
       // Send authentication credentials
       const creds = localStorage.getItem("creds");
@@ -92,7 +93,7 @@ const ActionProvider: React.FC<ActionProviderProps> = ({
       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();
@@ -104,8 +105,7 @@ const ActionProvider: React.FC<ActionProviderProps> = ({
       console.error("WebSocket error:", error);
     },
     onClose: (event) => {
-      console.error("WebSocket closed:", event.code, event.reason);
-      console.log("WebSocket state:", getWebSocket()?.readyState);
+      console.log("WebSocket closed:", event.code, event.reason);
     },
     shouldReconnect: (closeEvent) => {
       console.log("WebSocket should reconnect:", closeEvent.code !== 1000);
diff --git a/graphrag-ui/src/components/Bot.tsx b/graphrag-ui/src/components/Bot.tsx
index e800b42..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 {
@@ -23,7 +23,7 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
   const [store, setStore] = useState<any>();
   const [currentDate, setCurrentDate] = useState('');
   const [selectedGraph, setSelectedGraph] = useState(localStorage.getItem("selectedGraph") || '');
-  const [ragPattern, setRagPattern] = useState(localStorage.getItem("ragPattern") || 'Hybrid Search');
+  const [ragPattern, setRagPattern] = useState(localStorage.getItem("ragPattern") || '');
   const navigate = useNavigate();
 
   useEffect(() => {
@@ -37,6 +37,12 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
       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);
@@ -53,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();
   };
 
@@ -115,15 +122,17 @@ const Bot = ({ layout, getConversationId }: { layout?: string | undefined, getCo
         </div>
       
       <SelectedGraphContext.Provider value={selectedGraph}>
-        <Chatbot
-          // eslint-disable-next-line
-          // @ts-ignore
-          config={config}
-          fullPage={layout}
-          getConversationId={getConversationId}
-          messageParser={MessageParser}
-          actionProvider={ActionProvider}
-        />
+        <RagPatternContext.Provider value={ragPattern}>
+          <Chatbot
+            // eslint-disable-next-line
+            // @ts-ignore
+            config={config}
+            fullPage={layout}
+            getConversationId={getConversationId}
+            messageParser={MessageParser}
+            actionProvider={ActionProvider}
+          />
+        </RagPatternContext.Provider>
       </SelectedGraphContext.Provider>
     </div>
   );
diff --git a/graphrag-ui/src/components/Contexts.tsx b/graphrag-ui/src/components/Contexts.tsx
index 0d03f88..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("");
\ No newline at end of file
+export const SelectedGraphContext = createContext("");
+export const RagPatternContext = createContext("");
\ No newline at end of file
diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index d5e58dc..d6c0c3d 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -235,6 +235,8 @@ def hybrid_search(self, state):
             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=graphrag_config.get("chunk_only", True),
+            doc_only=graphrag_config.get("doc_only", False),
         )
 
         query_name = "GraphRAG_Hybrid_Vector_Search"
diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py
index 13fdaf3..a6b6cb5 100644
--- a/graphrag/app/routers/ui.py
+++ b/graphrag/app/routers/ui.py
@@ -384,6 +384,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]
@@ -403,7 +404,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
@@ -457,7 +458,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.
@@ -493,7 +495,7 @@ async def chat(
         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()

From c652e8be0797d505be0e570031954826b5dbe756 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 2 Oct 2025 20:25:51 -0700
Subject: [PATCH 19/31] Add edge info to retrieved results

---
 .../Chunk_Sibling_Vector_Search.gsql           |  6 +++++-
 .../GraphRAG_Community_Vector_Search.gsql      |  8 +++++---
 .../GraphRAG_Hybrid_Vector_Search.gsql         | 18 +++++++++++-------
 .../src/components/CustomChatMessage.tsx       |  2 +-
 graphrag/app/agent/agent_graph.py              |  7 ++++---
 5 files changed, 26 insertions(+), 15 deletions(-)

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<FLOAT> query_vector, UINT top_k=5, UINT lookback=3, INT lookahead=3, BOOL verbose = False) { 
   TYPEDEF TUPLE<VERTEX v, STRING t> VertexTypes;
+  TYPEDEF TUPLE<VERTEX s, VERTEX t> EdgeTypes;
   TYPEDEF tuple<Vertex v, Float score> Similarity_Results;
   HeapAccum<Similarity_Results>(top_k, score DESC) @@topk_set;
   MapAccum<STRING, SetAccum<VertexTypes>> @@verbose_info;
@@ -27,6 +28,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type,
   SetAccum<VERTEX> @@start_set;
   SetAccum<VertexTypes> @@start_set_type;
   SetAccum<VertexTypes> @@sibling_set_type;
+  SetAccum<EdgeTypes> @@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<FLOAT> query_vector, INT community_level=2, INT top_k = 3, BOOL with_chunk = true, BOOL with_doc = false, BOOL verbose = false) { 
+  TYPEDEF TUPLE<VERTEX s, VERTEX t> EdgeTypes;
   MapAccum<Vertex, SetAccum<String>> @@final_retrieval;
   MapAccum<STRING, SetAccum<Vertex>> @@verbose_info;
   SetAccum<STRING> @context;
   SetAccum<Vertex> @children;
   SetAccum<Vertex> @@start_set;
+  SetAccum<EdgeTypes> @@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<FLOAT>
 
       IF with_doc THEN
           related_chunks = SELECT c FROM Content:c -(<HAS_CONTENT)- Document:d -(HAS_CHILD>)- 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 -(<HAS_CONTENT)- DocumentChunk:d -(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);
       END;
   END;
@@ -56,7 +58,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Community_Vector_Search(LIST<FLOAT>
       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<STRING> v_types,
   LIST<FLOAT> 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<VERTEX v, STRING t> VertexTypes;
+  TYPEDEF TUPLE<VERTEX s, VERTEX t> EdgeTypes;
   TYPEDEF tuple<Vertex v, Float score> Similarity_Results;
   HeapAccum<Similarity_Results>(top_k, score DESC) @@topk_set;
   SetAccum<VERTEX> @@start_set;
   SetAccum<VertexTypes> @@start_set_type;
   SetAccum<VERTEX> @@tmp_set;
   SetAccum<VertexTypes> @@selected_set_type;
+  SetAccum<EdgeTypes> @@edges;
   SumAccum<INT> @visited;
   SumAccum<INT> @num_times_seen;
   MapAccum<Vertex, SetAccum<String>> @@result_set;
@@ -89,7 +91,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY GraphRAG_Hybrid_Vector_Search(Set<STRING> 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<STRING> 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<STRING> 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/graphrag-ui/src/components/CustomChatMessage.tsx b/graphrag-ui/src/components/CustomChatMessage.tsx
index 9fa02a0..971feeb 100755
--- a/graphrag-ui/src/components/CustomChatMessage.tsx
+++ b/graphrag-ui/src/components/CustomChatMessage.tsx
@@ -112,7 +112,7 @@ export const CustomChatMessage: FC<IChatbotMessageProps> = ({
               {/* <pre>{JSON.stringify(message, null, 2)}</pre> */}
               <div className="relative w-full h-[550px] my-10 border border-solid border-[#000]">
                 {message.query_sources?.result ? (
-                  <KnowledgeGraphPro data={message.query_sources?.result} />
+                  <KnowledgeGraphPro data={message.query_sources?.result.edges || {}} />
                 ) : (
                   <div className="flex items-center justify-center h-full text-gray-500">
                     No graph data available
diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index d6c0c3d..6fa36bc 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -229,13 +229,14 @@ 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"],
+            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=graphrag_config.get("chunk_only", True),
+            chunk_only=chunk_only,
             doc_only=graphrag_config.get("doc_only", False),
         )
 
@@ -366,7 +367,7 @@ def generate_answer(self, state):
                 f"""request_id={req_id_cv.get()} Got result: {state["context"]["result"]}"""
             )
             answer = step.generate_answer(
-                state["question"], state["context"]["result"]
+                state["question"], state["context"]["result"]["final_retrieval"]
             )
         elif state["lookup_source"] == "inquiryai":
             logger.debug_pii(

From a93b704866338d5bcf651e19768ffe8ad2b549b8 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 3 Oct 2025 16:10:47 -0700
Subject: [PATCH 20/31] GML-2002 fix visual graph generation issue

---
 graphrag-ui/package.json                      | 22 +++++-----
 .../src/components/CustomChatMessage.tsx      |  4 +-
 .../components/graphs/KnowledgeGraphPro.tsx   | 43 +++++++++++++------
 graphrag/app/agent/agent_graph.py             | 15 +++----
 4 files changed, 48 insertions(+), 36 deletions(-)

diff --git a/graphrag-ui/package.json b/graphrag-ui/package.json
index 9983a70..08fdd57 100755
--- a/graphrag-ui/package.json
+++ b/graphrag-ui/package.json
@@ -25,33 +25,35 @@
     "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",
     "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/components/CustomChatMessage.tsx b/graphrag-ui/src/components/CustomChatMessage.tsx
index 971feeb..9404394 100755
--- a/graphrag-ui/src/components/CustomChatMessage.tsx
+++ b/graphrag-ui/src/components/CustomChatMessage.tsx
@@ -111,8 +111,8 @@ export const CustomChatMessage: FC<IChatbotMessageProps> = ({
               {/* {message.query_sources?.result ? <pre>{JSON.stringify(message.query_sources?.result, null, 2)}</pre> : null} */}
               {/* <pre>{JSON.stringify(message, null, 2)}</pre> */}
               <div className="relative w-full h-[550px] my-10 border border-solid border-[#000]">
-                {message.query_sources?.result ? (
-                  <KnowledgeGraphPro data={message.query_sources?.result.edges || {}} />
+                {message.query_sources?.result.edges ? (
+                  <KnowledgeGraphPro data={message.query_sources?.result.edges} />
                 ) : (
                   <div className="flex items-center justify-center h-full text-gray-500">
                     No graph data available
diff --git a/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx b/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx
index b5ba6d3..028bafc 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<any | null>(null);
   // const [sdata, setsdata] = useState(JSON.parse(data));
-  const [edges, setEdges] = useState([]);
+  const [edges, setEdges] = useState<any[]>([]);
+  const [nodes, setNodes] = useState<any[]>([]);
   const [dataArray, setdataArray] = useState<any>();
   const [vId, setvId] = useState<any>();
 
   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<string>();
+
+    parsedData.forEach((e, i) => {
+      if (e.v && e.t) {
+        edgesList.push({ id: `e${i}`, source: e.v, target: e.t });
+        nodesSet.add(e.v);
+        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 && <pre>{edges}</pre>} */}
-    {/* {typeof sdata !== 'number' && typeof sdata !== 'string' && dataArray?.edgez && dataArray?.nodes ? (
+    { edges && nodes ? (
       <GraphCanvas
         ref={ref} 
-        nodes={dataArray?.nodes} 
-        edges={dataArray?.edgez}
+        nodes={nodes} 
+        edges={edges}
         labelType="all"
         theme={darkTheme}
         sizingType="centrality"
         draggable
       />
     ) : <div className='m-10'>Sorry no graph or table available</div> }
-    {typeof sdata !== 'number' && typeof sdata !== 'string' ? (<div className='absolute top-[10px] right-[10px] w-[170px]'>
+    {edges && nodes ? (<div className='absolute top-[10px] right-[10px] w-[170px]'>
       <button className='block w-full text-right text-sm' onClick={() => ref.current?.centerGraph()}>Home</button>
-    </div>) : null} */}
+    </div>) : null}
   </>
  )
 }
\ No newline at end of file
diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index 6fa36bc..a41a5eb 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -369,6 +369,11 @@ def generate_answer(self, state):
             answer = step.generate_answer(
                 state["question"], state["context"]["result"]["final_retrieval"]
             )
+
+            if not answer.citation:
+                answer.citation = [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"]}"""
@@ -390,16 +395,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
-
-            if answer.citation:
-                citations = [re.sub(r"_chunk_\d+", "", x) for x in answer.citation]
-            else:
-                citations = [re.sub(r"_chunk_\d+", "", x) for x in set(state["context"]["result"]["final_retrieval"].keys())]
-
-            state["context"]["reasoning"] = list(set(citations))
-
         try:
             if isinstance(self.llm_provider, AWSBedrock):
                 answer.generated_answer = self.replace_s3_urls_with_presigned(answer.generated_answer)

From fb4987c1103b0c263ae8ad6c728ecafb70053490 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 3 Oct 2025 17:19:45 -0700
Subject: [PATCH 21/31] Fix bug

---
 graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx | 8 ++++----
 graphrag/app/agent/agent_graph.py                       | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx b/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx
index 028bafc..4ebb2dc 100644
--- a/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx
+++ b/graphrag-ui/src/components/graphs/KnowledgeGraphPro.tsx
@@ -54,9 +54,9 @@ export const KnowledgeGraphPro = ({ data }) => {
     let nodesSet = new Set<string>();
 
     parsedData.forEach((e, i) => {
-      if (e.v && e.t) {
-        edgesList.push({ id: `e${i}`, source: e.v, target: e.t });
-        nodesSet.add(e.v);
+      if (e.s && e.t) {
+        edgesList.push({ id: `e${i}`, source: e.s, target: e.t });
+        nodesSet.add(e.s);
         nodesSet.add(e.t);
       }
     });
@@ -96,7 +96,7 @@ export const KnowledgeGraphPro = ({ data }) => {
   <>{/* edges ? JSON.stringify(edges) : 'no data' */}
     {/* {edges} */}
     {/* {edges && <pre>{edges}</pre>} */}
-    { edges && nodes ? (
+    { edges && nodes && nodes.length > 0 && edges.length > 0 ? (
       <GraphCanvas
         ref={ref} 
         nodes={nodes} 
diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index a41a5eb..fca5097 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -371,7 +371,7 @@ def generate_answer(self, state):
             )
 
             if not answer.citation:
-                answer.citation = [state["context"]["result"]["final_retrieval"].keys()]
+                answer.citation = list(state["context"]["result"]["final_retrieval"].keys())
             state["context"]["reasoning"] = list(set(answer.citation))
 
         elif state["lookup_source"] == "inquiryai":

From d41b59fa8212b4bedc818536edf5797841bd1ab8 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 7 Oct 2025 12:04:39 -0700
Subject: [PATCH 22/31] Revise prompts for AML

---
 .../chatbot_response.txt                      | 15 ++++++
 .../prompts/openai_gpt4/chatbot_response.txt  | 49 +++++++++----------
 .../prompts/openai_gpt4/generate_cypher.txt   |  1 +
 3 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
index 4267b1f..063360f 100644
--- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
+++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
@@ -7,6 +7,21 @@ Use compact markdown syntax to geneate the answer, including title, bulleted or
 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 after validation. Assign an empty value to the field if you cannot determine it.
 
+For questions related to financial graph or transaction graph, create a suspicious activity report.
+- The narrative should be clear, comprehensive, and avoid institution-specific acronyms.
+- The narrative should include paragraphs for Summary of **Investigation**, **Suspicious Activity Overview**, **Details of Suspicious Activities**, **Investigation Conducted**, and **Conclusion**
+- The narrative must provide information about the subject to include phone numbers, email addresses, addresses and government IDs.
+- The narrative must include any suspicious transactions and the start of the suspicious transactions.
+- The narrative must specify the suspicious activity observed (types of transactions, amount, frequency).
+- The narrative must tell the complete story. A reviewer should understand the full picture without needing additional context.
+
+Narrative Writing Guidelines:
+- Write in clear, complete sentences
+- Spell out all acronyms (no institution-specific jargon)
+- Include specific dates, amounts, and account numbers
+- Describe the investigation conducted
+- Note any customer explanations received and why they were insufficient
+- Include any supporting documentation references
 
 Question: {question}
 Answer: {context}
diff --git a/common/prompts/openai_gpt4/chatbot_response.txt b/common/prompts/openai_gpt4/chatbot_response.txt
index c2b8333..a31f669 100644
--- a/common/prompts/openai_gpt4/chatbot_response.txt
+++ b/common/prompts/openai_gpt4/chatbot_response.txt
@@ -1,31 +1,30 @@
-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:
+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.
+- 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, and prioritize to use images, tables and context appear repeatedly.
+- 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 after validation. Assign an empty value to the field if you cannot determine it.
 
-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].
+For questions related to financial graph or transaction graph, create a suspicious activity report.
+- The narrative should be clear, comprehensive, and avoid institution-specific acronyms.
+- The narrative should include paragraphs for Summary of **Investigation**, **Suspicious Activity Overview**, **Details of Suspicious Activities**, **Investigation Conducted**, and **Conclusion**
+- The narrative must provide information about the subject to include phone numbers, email addresses, addresses and government IDs.
+- The narrative must include any suspicious transactions and the start of the suspicious transactions.
+- The narrative must specify the suspicious activity observed (types of transactions, amount, frequency).
+- The narrative must tell the complete story. A reviewer should understand the full picture without needing additional context.
 
-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.
-
-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 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.
+Narrative Writing Guidelines:
+- Write in clear, complete sentences
+- Spell out all acronyms (no institution-specific jargon)
+- Include specific dates, amounts, and account numbers
+- Describe the investigation conducted
+- Note any customer explanations received and why they were insufficient
+- Include any supporting documentation references
 
 Question: {question}
-Context: {context}
+Answer: {context}
 Query: {query}
 Format: {format_instructions}
+
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.
 

From c7b4dd8e11ddc15a3aa25bc0af9eab1220fbe99e Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 7 Oct 2025 12:15:44 -0700
Subject: [PATCH 23/31] Revise prompts for AML

---
 common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt | 3 +--
 common/prompts/openai_gpt4/chatbot_response.txt              | 3 +--
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
index 063360f..05532c4 100644
--- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
+++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
@@ -24,7 +24,6 @@ Narrative Writing Guidelines:
 - Include any supporting documentation references
 
 Question: {question}
-Answer: {context}
+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 a31f669..408e22c 100644
--- a/common/prompts/openai_gpt4/chatbot_response.txt
+++ b/common/prompts/openai_gpt4/chatbot_response.txt
@@ -24,7 +24,6 @@ Narrative Writing Guidelines:
 - Include any supporting documentation references
 
 Question: {question}
-Answer: {context}
+Context: {context}
 Query: {query}
 Format: {format_instructions}
-

From 00b1a18863d03a7a0b2b6406c41e838e1ce1c996 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 7 Oct 2025 18:51:38 -0700
Subject: [PATCH 24/31] Add remark-gfm plugin

---
 graphrag-ui/package.json                         | 1 +
 graphrag-ui/src/components/CustomChatMessage.tsx | 7 ++++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/graphrag-ui/package.json b/graphrag-ui/package.json
index 9983a70..aa7ad2e 100755
--- a/graphrag-ui/package.json
+++ b/graphrag-ui/package.json
@@ -35,6 +35,7 @@
     "react-router-dom": "^6.23.1",
     "react-use-websocket": "^4.8.1",
     "reagraph": "^4.19.1",
+    "remark-gfm": "^4.0.0",
     "tailwind-merge": "^2.3.0",
     "tailwindcss-animate": "^1.0.7",
     "@tailwindcss/typography": "^0.5.18",
diff --git a/graphrag-ui/src/components/CustomChatMessage.tsx b/graphrag-ui/src/components/CustomChatMessage.tsx
index 9fa02a0..4cff5bd 100755
--- a/graphrag-ui/src/components/CustomChatMessage.tsx
+++ b/graphrag-ui/src/components/CustomChatMessage.tsx
@@ -1,5 +1,6 @@
 import { FC, useState } from "react";
-import Markdown from 'react-markdown'
+import ReactMarkdown from 'react-markdown'
+import remarkGfm from 'remark-gfm'
 import {
   Dialog,
   DialogContent,
@@ -86,7 +87,7 @@ export const CustomChatMessage: FC<IChatbotMessageProps> = ({
     <>
       {typeof message === "string" ? (
         <div className="prose dark:prose-invert text-sm max-w-[230px] md:max-w-[80%] mt-7 mb-7">
-          <Markdown className="typewriter">{message}</Markdown>
+          <ReactMarkdown remarkPlugins={[remarkGfm]} className="typewriter">{message}</ReactMarkdown>
         </div>
       ) : message.key === null ? (
         message
@@ -96,7 +97,7 @@ export const CustomChatMessage: FC<IChatbotMessageProps> = ({
             {message.response_type === "progress" ? (
               <p className="graphrag-thinking typewriter">{message.content}</p>
             ) : (
-              <Markdown className="typewriter">{message.content}</Markdown>
+              <ReactMarkdown remarkPlugins={[remarkGfm]} className="typewriter">{message.content}</ReactMarkdown>
             )}
             <Interactions
               message={message} 

From 1b2df586918ca71a73983606c266c0fc93cb84e2 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 24 Oct 2025 13:16:07 -0700
Subject: [PATCH 25/31] Support adjust boto3 pool size

---
 common/embeddings/embedding_services.py | 10 +++++++++-
 ecc/app/graphrag/util.py                |  4 +---
 ecc/app/supportai/util.py               |  3 +--
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py
index 75ba9c1..4688989 100644
--- a/common/embeddings/embedding_services.py
+++ b/common/embeddings/embedding_services.py
@@ -202,14 +202,22 @@ class AWS_Bedrock_Embedding(EmbeddingModel):
     """AWS Bedrock Embedding Model"""
 
     def __init__(self, config):
-        import boto3
+        import boto3, botocore
         from langchain_aws import BedrockEmbeddings
 
         super().__init__(config=config, model_name=config.get("model_name", "amazon.titan-embed-text-v1"))
 
+        boto3_config = config.get("boto3_config", {})
+        client_config = botocore.config.Config(
+            max_pool_connections=boto3_config.get("max_pool_connections", 50),
+            read_timeout=boto3_config.get("read_timeout", 300),
+            retries={"max_attempts": boto3_config.get("retries", 10)},
+        )
+
         client = boto3.client(
             "bedrock-runtime",
             region_name="us-east-1",
+            config=client_config,
             aws_access_key_id=config["authentication_configuration"][
                 "AWS_ACCESS_KEY_ID"
             ],
diff --git a/ecc/app/graphrag/util.py b/ecc/app/graphrag/util.py
index a053f1b..8fd79cd 100644
--- a/ecc/app/graphrag/util.py
+++ b/ecc/app/graphrag/util.py
@@ -170,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]
     if v_id == "''" or v_id == '""':
         return ""
-    v_id = v_id.replace("(", "").replace(")", "")
+    v_id = v_id.replace(" ", "-").lower().replace("(", "").replace(")", "")
 
     return v_id
 
diff --git a/ecc/app/supportai/util.py b/ecc/app/supportai/util.py
index 6de5fc3..c12068e 100644
--- a/ecc/app/supportai/util.py
+++ b/ecc/app/supportai/util.py
@@ -153,13 +153,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]
     if v_id == "''" or v_id == '""':
         return ""
+    v_id = v_id.replace(" ", "-").lower().replace("(", "").replace(")", "")
 
     return v_id
 

From f19644f9167b81a825a0b2c77302db9b4a8be3aa Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 24 Oct 2025 14:06:10 -0700
Subject: [PATCH 26/31] GML-2008 support dimensions for Azure

---
 common/embeddings/embedding_services.py    |  6 +++---
 common/llm_services/aws_bedrock_service.py | 11 ++++++++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py
index 4688989..c0c574d 100644
--- a/common/embeddings/embedding_services.py
+++ b/common/embeddings/embedding_services.py
@@ -165,7 +165,7 @@ def __init__(self, config):
         super().__init__(config, model_name=config.get("model_name", "text-embedding-3-small"))
         from langchain.embeddings import AzureOpenAIEmbeddings
 
-        self.embeddings = AzureOpenAIEmbeddings(model=self.model_name, deployment=config["azure_deployment"])
+        self.embeddings = AzureOpenAIEmbeddings(model=self.model_name, dimensions=self.dimensions, deployment=config["azure_deployment"])
 
 
 class OpenAI_Embedding(EmbeddingModel):
@@ -211,12 +211,12 @@ def __init__(self, config):
         client_config = botocore.config.Config(
             max_pool_connections=boto3_config.get("max_pool_connections", 50),
             read_timeout=boto3_config.get("read_timeout", 300),
-            retries={"max_attempts": boto3_config.get("retries", 10)},
+            retries={"max_attempts": boto3_config.get("retries", 5)},
         )
 
         client = boto3.client(
             "bedrock-runtime",
-            region_name="us-east-1",
+            region_name=config.get("region_name", "us-east-1"),
             config=client_config,
             aws_access_key_id=config["authentication_configuration"][
                 "AWS_ACCESS_KEY_ID"
diff --git a/common/llm_services/aws_bedrock_service.py b/common/llm_services/aws_bedrock_service.py
index 29348ea..1693e82 100644
--- a/common/llm_services/aws_bedrock_service.py
+++ b/common/llm_services/aws_bedrock_service.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import os
-import boto3
+import boto3, botocore
 from langchain_aws import ChatBedrock
 import logging
 from common.llm_services import LLM_Model
@@ -27,9 +27,18 @@ class AWSBedrock(LLM_Model):
     def __init__(self, config):
         super().__init__(config)
         model_name = config["llm_model"]
+
+        boto3_config = config.get("boto3_config", {})
+        client_config = botocore.config.Config(
+            max_pool_connections=boto3_config.get("max_pool_connections", 50),
+            read_timeout=boto3_config.get("read_timeout", 300),
+            retries={"max_attempts": boto3_config.get("retries", 5)},
+        )
+
         client = boto3.client(
             "bedrock-runtime",
             region_name=config.get("region_name", "us-east-1"),
+            config=client_config,
             aws_access_key_id=config["authentication_configuration"][
                 "AWS_ACCESS_KEY_ID"
             ],

From b0a0ef8750cfad81511bc7a2f2903deeb5ddf101 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Tue, 28 Oct 2025 21:43:56 -0700
Subject: [PATCH 27/31] Integrate Bedrock Change

---
 graphrag/app/supportai/supportai.py | 73 +++++++++++++++++++----------
 1 file changed, 49 insertions(+), 24 deletions(-)

diff --git a/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py
index 4e6410f..e455fd6 100644
--- a/graphrag/app/supportai/supportai.py
+++ b/graphrag/app/supportai/supportai.py
@@ -136,7 +136,9 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre
     # 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
@@ -147,30 +149,52 @@ 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):
@@ -211,7 +235,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'
             )
@@ -265,7 +289,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)
@@ -419,6 +443,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("]")}
 

From a81f5282e56ee829bfe9189015b1f67c45044a3c Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Fri, 31 Oct 2025 15:07:15 -0700
Subject: [PATCH 28/31] revert tg version

---
 docker-compose.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker-compose.yml b/docker-compose.yml
index 7d38202..8be754b 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -74,7 +74,7 @@ services:
       - graphrag
 
   tigergraph:
-    image: tigergraph/tigergraph:4.3.0
+    image: tigergraph/community:4.2.1
     container_name: tigergraph
     platform: linux/amd64
     ports:

From dec1375d86db131f423ed6d8e0025f3b5eed566d Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Thu, 6 Nov 2025 15:59:36 -0800
Subject: [PATCH 29/31] Revert some changes

---
 common/embeddings/embedding_services.py       |  2 +-
 common/metrics/tg_proxy.py                    |  7 ++-
 .../chatbot_response.txt                      | 43 ++++++++++---------
 .../prompts/openai_gpt4/generate_function.txt |  5 ++-
 .../{utils.py => utils/token_calculator.py}   |  1 +
 ecc/app/graphrag/util.py                      |  2 +-
 ecc/app/supportai/util.py                     |  2 +-
 graphrag/app/agent/agent_generation.py        |  2 +-
 .../app/supportai/retrievers/BaseRetriever.py |  2 +-
 graphrag/app/supportai/supportai.py           |  8 +---
 10 files changed, 38 insertions(+), 36 deletions(-)
 rename common/{utils.py => utils/token_calculator.py} (99%)

diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py
index ad73784..8a24265 100644
--- a/common/embeddings/embedding_services.py
+++ b/common/embeddings/embedding_services.py
@@ -11,7 +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 import TokenCalculator
+from common.utils.token_calculator import TokenCalculator
 
 logger = logging.getLogger(__name__)
 
diff --git a/common/metrics/tg_proxy.py b/common/metrics/tg_proxy.py
index b5aa75d..804d66f 100644
--- a/common/metrics/tg_proxy.py
+++ b/common/metrics/tg_proxy.py
@@ -55,9 +55,8 @@ def _runInstalledQuery(self, query_name, params, sizeLimit=None, usePost=False):
             result = None
             while not result:
                 ret = self._tg_connection.checkQueryStatus(restppid)
-                #LogWriter.info(f"Ret: {ret}")
                 if not ret:
-                    time.sleep(1)
+                    time.sleep(0.5)
                     continue
                 if ret[0]["status"].lower() == "success":
                     LogWriter.info(
@@ -73,7 +72,7 @@ def _runInstalledQuery(self, query_name, params, sizeLimit=None, usePost=False):
                             result = None
                             if i >= 9:
                                 raise e
-                        time.sleep(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"
@@ -89,7 +88,7 @@ def _runInstalledQuery(self, query_name, params, sizeLimit=None, usePost=False):
                         f"Query {query_name} with restppid {restppid} timed out"
                     )
                 else:
-                    time.sleep(1) # Still running, wait for 1 second
+                    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 05532c4..11c9c46 100644
--- a/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
+++ b/common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt
@@ -1,27 +1,30 @@
-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 after validation. Assign an empty value to the field if you cannot determine it.
-
-For questions related to financial graph or transaction graph, create a suspicious activity report.
-- The narrative should be clear, comprehensive, and avoid institution-specific acronyms.
-- The narrative should include paragraphs for Summary of **Investigation**, **Suspicious Activity Overview**, **Details of Suspicious Activities**, **Investigation Conducted**, and **Conclusion**
-- The narrative must provide information about the subject to include phone numbers, email addresses, addresses and government IDs.
-- The narrative must include any suspicious transactions and the start of the suspicious transactions.
-- The narrative must specify the suspicious activity observed (types of transactions, amount, frequency).
-- The narrative must tell the complete story. A reviewer should understand the full picture without needing additional context.
-
-Narrative Writing Guidelines:
-- Write in clear, complete sentences
-- Spell out all acronyms (no institution-specific jargon)
-- Include specific dates, amounts, and account numbers
-- Describe the investigation conducted
-- Note any customer explanations received and why they were insufficient
-- Include any supporting documentation references
+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}
 Context: {context}
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/utils.py b/common/utils/token_calculator.py
similarity index 99%
rename from common/utils.py
rename to common/utils/token_calculator.py
index d828cbf..53d96ce 100644
--- a/common/utils.py
+++ b/common/utils/token_calculator.py
@@ -221,3 +221,4 @@ def truncate_to_token_limit(self, text: str | dict, max_tokens: Optional[int] =
             return self.truncate_dict_to_token_limit(text, max_tokens)
         else:
             return self.truncate_text_to_token_limit(text, max_tokens)
+
diff --git a/ecc/app/graphrag/util.py b/ecc/app/graphrag/util.py
index 8fd79cd..25f4cff 100644
--- a/ecc/app/graphrag/util.py
+++ b/ecc/app/graphrag/util.py
@@ -173,9 +173,9 @@ def process_id(v_id: str):
     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(" ", "-").lower().replace("(", "").replace(")", "")
 
     return v_id
 
diff --git a/ecc/app/supportai/util.py b/ecc/app/supportai/util.py
index c12068e..2e670c7 100644
--- a/ecc/app/supportai/util.py
+++ b/ecc/app/supportai/util.py
@@ -156,9 +156,9 @@ def process_id(v_id: str):
     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(" ", "-").lower().replace("(", "").replace(")", "")
 
     return v_id
 
diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py
index 31f7a2f..d38bd8b 100644
--- a/graphrag/app/agent/agent_generation.py
+++ b/graphrag/app/agent/agent_generation.py
@@ -21,7 +21,7 @@
 from pydantic import BaseModel, Field
 from common.logs.logwriter import LogWriter
 from common.logs.log import req_id_cv
-from common.utils import TokenCalculator
+from common.utils.token_calculator import TokenCalculator
 from common.config import completion_config
 
 
diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py
index b6f6583..6c1b16a 100644
--- a/graphrag/app/supportai/retrievers/BaseRetriever.py
+++ b/graphrag/app/supportai/retrievers/BaseRetriever.py
@@ -3,7 +3,7 @@
 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 import TokenCalculator
+from common.utils.token_calculator import TokenCalculator
 from common.config import completion_config
 
 from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser
diff --git a/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py
index b373bf0..33b3fb0 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,7 +141,7 @@ 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"])
@@ -677,4 +673,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")

From 41531db0433e85f0b0127a5cd759469031b75491 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Mon, 10 Nov 2025 09:10:56 -0800
Subject: [PATCH 30/31] fix minor issue

---
 docker-compose.yml                  | 1 -
 graphrag/app/supportai/supportai.py | 8 ++++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

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/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py
index 33b3fb0..91805f9 100644
--- a/graphrag/app/supportai/supportai.py
+++ b/graphrag/app/supportai/supportai.py
@@ -209,12 +209,12 @@ def trigger_bedrock_bda(input_uri, output_uri, region, aws_access_key, aws_secre
         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://<bucket_name>/<prefix>")
         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):
@@ -276,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:
@@ -578,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

From 38dc7dfd4b5609b842a691cd525cd50a659643b9 Mon Sep 17 00:00:00 2001
From: Chengbiao Jin <chengbiao.jin@tigergraph.com>
Date: Mon, 10 Nov 2025 09:28:04 -0800
Subject: [PATCH 31/31] revert citation changes

---
 graphrag/app/agent/agent_graph.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py
index a808cf8..0874217 100644
--- a/graphrag/app/agent/agent_graph.py
+++ b/graphrag/app/agent/agent_graph.py
@@ -370,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"]}"""
@@ -391,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):