GML-1997 token calculator, GML-1998 Extra Chat model

chat-history/Dockerfile

@@ -1,4 +1,4 @@
-FROM golang:1.22.3 as builder
+FROM golang:1.22.3 AS builder
 
 WORKDIR /app 
 COPY . .

common/config.py

@@ -91,6 +91,14 @@
     raise Exception("embedding_service is not found in llm_config")
 embedding_dimension = embedding_config.get("dimensions", 1536)
 
+# Get context window size from llm_config
+# <=0 means unlimited tokens (no truncation), otherwise use the specified limit
+if "token_limit" in llm_config:
+    if "token_limit" not in completion_config:
+        completion_config["token_limit"] = llm_config["token_limit"]
+    if "token_limit" not in embedding_config:
+        embedding_config["token_limit"] = llm_config["token_limit"]
+
 # Get multimodal_service config (optional, for vision/image tasks)
 multimodal_config = llm_config.get("multimodal_service")
 

common/embeddings/embedding_services.py

@@ -11,6 +11,7 @@
 from common.logs.log import req_id_cv
 from common.logs.logwriter import LogWriter
 from common.metrics.prometheus_metrics import metrics
+from common.utils.token_calculator import TokenCalculator
 
 logger = logging.getLogger(__name__)
 
@@ -32,6 +33,7 @@ def __init__(self, config: dict, model_name: str):
         self.embeddings = None
         self.model_name = model_name
         self.dimensions = config.get("dimensions", 1536)
+        self.token_calculator = TokenCalculator(token_limit=config.get("token_limit", 8192), model_name=model_name)
         LogWriter.info(
             f"request_id={req_id_cv.get()} instantiated AI model_name={model_name} with dimensions={self.dimensions}"
         )
@@ -51,6 +53,13 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:
 
         try:
             LogWriter.info(f"request_id={req_id_cv.get()} ENTRY embed_documents()")
+
+            if not self.token_calculator.is_unlimited_tokens():
+                max_context_tokens = self.token_calculator.get_max_context_tokens()
+                if any(len(text) > max_context_tokens for text in texts):
+                    if any(self.token_calculator.count_tokens(text) > max_context_tokens for text in texts):
+                        texts = [self.token_calculator.truncate_to_token_limit(text, max_context_tokens) for text in texts]
+
             if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings):
                 docs = self.embeddings.embed_documents(texts, output_dimensionality=self.dimensions)
             else:
@@ -85,6 +94,13 @@ def embed_query(self, question: str) -> List[float]:
             logger.debug_pii(
                 f"request_id={req_id_cv.get()} embed_query() embedding question={question}"
             )
+
+            if not self.token_calculator.is_unlimited_tokens():
+                max_context_tokens = self.token_calculator.get_max_context_tokens()
+                if len(question) > max_context_tokens:
+                    if self.token_calculator.count_tokens(question) > max_context_tokens:
+                        question = self.token_calculator.truncate_to_token_limit(question, max_context_tokens)
+
             if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings):
                 query_embedding = self.embeddings.embed_query(question, output_dimensionality=self.dimensions)
             else:
@@ -117,6 +133,12 @@ async def aembed_query(self, question: str) -> List[float]:
         # try:
         LogWriter.info(f"request_id={req_id_cv.get()} ENTRY aembed_query()")
         logger.debug_pii(f"aembed_query() embedding question={question}")
+        if not self.token_calculator.is_unlimited_tokens():
+            max_context_tokens = self.token_calculator.get_max_context_tokens()
+            if len(question) > max_context_tokens:
+                if self.token_calculator.count_tokens(question) > max_context_tokens:
+                    question = self.token_calculator.truncate_to_token_limit(question, max_content_tokens)
+
         if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings):
             query_embedding = await self.embeddings.aembed_query(question, output_dimensionality=self.dimensions)
         else:
@@ -140,7 +162,7 @@ class AzureOpenAI_Ada002(EmbeddingModel):
     """Azure OpenAI Ada-002 Embedding Model"""
 
     def __init__(self, config):
-        super().__init__(config, model_name=config.get("model_name", "OpenAI ada-002"))
+        super().__init__(config, model_name=config.get("model_name", "text-embedding-3-small"))
         from langchain_openai import AzureOpenAIEmbeddings
 
         self.embeddings = AzureOpenAIEmbeddings(model=self.model_name, dimensions=self.dimensions, deployment=config["azure_deployment"])
@@ -203,7 +225,7 @@ def __init__(self, config):
                 "AWS_SECRET_ACCESS_KEY"
             ],
         )
-        self.embeddings = BedrockEmbeddings(client=client)
+        self.embeddings = BedrockEmbeddings(client=client, model_id=self.model_name)
 
 
 class Ollama_Embedding(EmbeddingModel):
@@ -212,13 +234,12 @@ class Ollama_Embedding(EmbeddingModel):
     def __init__(self, config):
         from langchain_ollama import OllamaEmbeddings
 
-        super().__init__(config=config, model_name=config.get("model_name", "llama2"))
+        super().__init__(config=config, model_name=config.get("model_name", "llama3"))
 
         # Get Ollama configuration from config
         base_url = config.get("base_url", "http://localhost:11434")
-        model_name = config.get("model_name", "llama3")
 
         self.embeddings = OllamaEmbeddings(
-            model=model_name,
+            model=self.model_name,
             base_url=base_url
         )

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(

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;
+}

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;

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;

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,27 +101,29 @@ 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);
   ELSE
       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;

common/metrics/tg_proxy.py

@@ -36,27 +36,27 @@ 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}"
             )
             result = None
             while not result:
                 ret = self._tg_connection.checkQueryStatus(restppid)
-                LogWriter.info(f"Ret: {ret}")
                 if not ret:
-                    time.sleep(0.1)
+                    time.sleep(0.5)
                     continue
                 if ret[0]["status"].lower() == "success":
                     LogWriter.info(
@@ -72,7 +72,7 @@ def _runInstalledQuery(self, query_name, params, usePost=False):
                             result = None
                             if i >= 9:
                                 raise e
-                            time.sleep(0.1)
+                        time.sleep(0.5)
                 elif ret[0]["status"].lower() == "aborted":
                     LogWriter.error(
                         f"request_id={req_id_cv.get()} query {query_name} with RESTPP ID {restppid} aborted"
@@ -87,6 +87,8 @@ def _runInstalledQuery(self, query_name, params, usePost=False):
                     raise Exception(
                         f"Query {query_name} with restppid {restppid} timed out"
                     )
+                else:
+                    time.sleep(0.5) # Still running, wait for 0.5 second
             success = True
         except Exception as e:
             LogWriter.error(f"Error running query {query_name}: {str(e)}")

common/prompts/aws_bedrock_claude3haiku/chatbot_response.txt

@@ -1,14 +1,32 @@
-You are a highly efficient and empathetic AI-powered assistant in JSON parsing and generating.
-Given the following context in JSON format, rephrase it to answer the question.
+You are a highly efficient and empathetic AI-powered customer support agent. Your goal is to provide accurate, helpful, and friendly support while maintaining professionalism. Follow these guidelines:
+
+1. **Understand the User's Needs**: Carefully interpret the user's question or issue. Ask clarifying questions if necessary.
+2. **Provide Accurate Answers**: Use the given context in JSON format as well as the data structure of the original query used to fecth the context to deliver reliable and precise responses.
+3. **Provide Clean Code Examples**: Provide necessary code examples if there is code snippet in the given context to answer the question. Use less code whenever possible.
+4. **Be Polite and Friendly**: Maintain a tone that is warm, respectful, and supportive.
+5. **Stay Within Scope**: Respond only with information relevant to the user's question. If the information is unavailable, apologize and suggest alternative actions.
+6. **Encourage Next Steps**: Proactively suggest solutions or next steps, such as contacting a human agent if the issue cannot be resolved by the AI.
+7. **Escalate When Needed**: Recognize complex or sensitive queries that require human attention and politely guide the user to contact human support.
+8. **Reference Relevant Knowledge**: Ground your responses based on the context and resources, and referencing the data struction from the query providedy. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID].
+
+Example format for responses:
+- **Solution**: "Here's the information you requested: [details]."
+- **Clarification**: "Could you provide more details about [specific aspect]?"
+- **Escalation**: "This issue may require human assistance. Please reach out via [Support Portal](https://www.tigergraph.com/support/) or email us at [support@tigergraph.com](mailto:support@tigergraph.com)."
+
+Format your answer using Markdown. Organize the content into paragraphs, bulleted or numbered lists, and include links to images where relevant.
+Make sure to extract and include the image references in [IMAGE_REF:image_id] format in your generated_answer if they are present in the context. Images are critical visual information that must be included in your response. Do NOT modify or omit these image references.
+
+Give the context in JSON format, combine and rephrase it to answer the question.
 Use only the provided information in context without adding any reasoning or additional logic.
-Make sure all information in the context are covered in the generated answer.
+Make sure all information in the context are covered in the generated answer, including any image markdown references.
 Make sure to extract and include the image links in markdown syntax in the generated answer when their summaries are referenced, and preserve the link URLs in their original format.
 Use compact markdown syntax to geneate the answer, including title, bulleted or numbered list, images and tables if any, and place images or tables below the related text section.
 Ensure that each row of every table, including the header row, starts on a new line.
-Always only return a JSON contains the answer and otheh required fields. Assign an empty value to the field if you cannot determine it.
-
+Generate the answer in JSON format, make sure to escape necessary characters in order to return a valid JSON response only.
+Make sure all the fields required by the format instructions are included, set a field to empty if you don't have that information.
 
 Question: {question}
-Answer: {context}
+Context: {context}
+Query: {query}
 Format: {format_instructions}
-

common/prompts/aws_bedrock_claude3haiku/generate_cypher.txt

@@ -12,6 +12,7 @@ Always use double quotes for strings instead of single quotes.
 Always convert strings to lower case using toLower() function for string comparision in WHERE clause.
 Use alias for ORDER BY if any, avoid using short alias names especially single letter alias, always use meaningful words connected by underscore.
 Always make sure the alias or attributes used in ORDER BY is the same type in RETURN. Always add ASC or DESC for ORDER BY based on data type.
+For questions like "summarize" or "write a summary" about something, fetch all information on its neighbour nodes and edges.
 
 Avoid to generate invalid OpenCypher queries based on the errors from history below.