escape special characters

app.py

@@ -1,4 +1,4 @@
-# main.py  
+# app.py  
 import spaces
 from torch.nn import DataParallel  
 from torch import Tensor  
@@ -17,7 +17,7 @@ import gradio as gr
 import torch  
 import torch.nn.functional as F  
 from dotenv import load_dotenv
-from utils import load_env_variables, parse_and_route  
+from utils import load_env_variables, parse_and_route , escape_special_characters 
 from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name , metadata_prompt 
 
 
@@ -49,12 +49,12 @@ class EmbeddingGenerator:
 
     @spaces.GPU  
     def compute_embeddings(self, input_text: str):  
-        # Get the intention  
+        escaped_input_text = escape_special_characters(input_text)  
         intention_completion = self.intention_client.chat.completions.create(  
             model="yi-large",  
             messages=[  
-                {"role": "system", "content": intention_prompt},  
-                {"role": "user", "content": input_text}  
+                {"role": "system", "content": escape_special_characters(intention_prompt)},  
+                {"role": "user", "content": escaped_input_text}  
             ]  
         )  
         intention_output = intention_completion.choices[0].message['content']  
@@ -71,14 +71,14 @@ class EmbeddingGenerator:
             return f"Error: Task '{selected_task}' not found. Please select a valid task."  
   
         query_prefix = f"Instruct: {task_description}\nQuery: "  
-        queries = [input_text]  
+        queries = [escaped_input_text]  
   
         # Get the metadata  
         metadata_completion = self.intention_client.chat.completions.create(  
             model="yi-large",  
             messages=[  
-                {"role": "system", "content": metadata_prompt},  
-                {"role": "user", "content": input_text}  
+                {"role": "system", "content": escape_special_characters(metadata_prompt)},  
+                {"role": "user", "content": escaped_input_text}  
             ]  
         )  
         metadata_output = metadata_completion.choices[0].message['content']  
@@ -93,12 +93,9 @@ class EmbeddingGenerator:
             # Normalize embeddings  
             query_embeddings = F.normalize(query_embeddings, p=2, dim=1)  
             embeddings_list = query_embeddings.detach().cpu().numpy().tolist()  
-        
-            # Include metadata in the embeddings  
-            embeddings_with_metadata = [{"embedding": emb, "metadata": metadata} for emb in embeddings_list]  
   
             self.clear_cuda_cache()  
-            return embeddings_with_metadata  
+            return embeddings_list, metadata   
   
     def extract_metadata(self, metadata_output: str):  
         # Regex pattern to extract key-value pairs  
@@ -143,8 +140,18 @@ def add_documents_to_chroma(client, collection, documents: list, embedding_funct
             )  
 
 def query_chroma(client, collection_name: str, query_text: str, embedding_function: MyEmbeddingFunction):  
+    # Compute query embeddings and metadata  
+    query_embeddings, query_metadata = embedding_function.embedding_generator.compute_embeddings(query_text)  
+  
+    # Initialize Chroma with the collection  
     db = Chroma(client=client, collection_name=collection_name, embedding_function=embedding_function)  
-    result_docs = db.similarity_search(query_text)  
+  
+    # Perform similarity search using the query embeddings and metadata  
+    result_docs = db.similarity_search(  
+        query_embeddings=query_embeddings,  
+        query_metadata=query_metadata  
+    )  
+  
     return result_docs  
 
 
@@ -164,13 +171,13 @@ def respond(
     top_p,  
 ):  
     retrieved_text = query_documents(message)  
-    messages = [{"role": "system", "content": system_message}]  
+    messages = [{"role": "system", "content": escape_special_characters(system_message)}]  
     for val in history:  
         if val[0]:  
             messages.append({"role": "user", "content": val[0]})  
         if val[1]:  
             messages.append({"role": "assistant", "content": val[1]})  
-    messages.append({"role": "user", "content": f"{retrieved_text}\n\n{message}"})  
+    messages.append({"role": "user", "content": f"{retrieved_text}\n\n{escape_special_characters(message)}"})  
     response = ""  
     for message in intention_client.chat_completion(  
         messages,  

utils.py

@@ -30,4 +30,37 @@ def parse_and_route(example_output: str):
         else:  
             return {true_task: "Task description not found"}  
     else:  
-        return "No true task found in the example output" 
+        return "No true task found in the example output" 
+    
+import json  
+  
+def escape_special_characters(text: str) -> str:  
+    """  
+    Escapes special characters in the given text for JSON and cURL compatibility.  
+    """  
+    escaped_text = json.dumps(text)[1:-1]   
+    curl_escaped_text = escaped_text.replace(" ", "\\ ")  
+    curl_escaped_text = curl_escaped_text.replace("&", "\\&")  
+    curl_escaped_text = curl_escaped_text.replace(";", "\\;")  
+    curl_escaped_text = curl_escaped_text.replace("(", "\\(")  
+    curl_escaped_text = curl_escaped_text.replace(")", "\\)")  
+    curl_escaped_text = curl_escaped_text.replace("$", "\\$")  
+    curl_escaped_text = curl_escaped_text.replace("`", "\\`")  
+    curl_escaped_text = curl_escaped_text.replace("|", "\\|")  
+    curl_escaped_text = curl_escaped_text.replace("*", "\\*")  
+    curl_escaped_text = curl_escaped_text.replace("?", "\\?")  
+    curl_escaped_text = curl_escaped_text.replace("<", "\\<")  
+    curl_escaped_text = curl_escaped_text.replace(">", "\\>")  
+    curl_escaped_text = curl_escaped_text.replace("!", "\\!")  
+    curl_escaped_text = curl_escaped_text.replace("{", "\\{")  
+    curl_escaped_text = curl_escaped_text.replace("}", "\\}")  
+    curl_escaped_text = curl_escaped_text.replace("[", "\\[")  
+    curl_escaped_text = curl_escaped_text.replace("]", "\\]")  
+    curl_escaped_text = curl_escaped_text.replace("#", "\\#")  
+    curl_escaped_text = curl_escaped_text.replace("%", "\\%")  
+    curl_escaped_text = curl_escaped_text.replace("^", "\\^")  
+    curl_escaped_text = curl_escaped_text.replace("=", "\\=")  
+    curl_escaped_text = curl_escaped_text.replace("~", "\\~")  
+    curl_escaped_text = curl_escaped_text.replace("'", "\\'")  
+  
+    return curl_escaped_text