Update appStore/rag.py

appStore/rag.py

@@ -32,40 +32,49 @@ def get_prompt(context, label):
 # def completion_with_backoff(**kwargs):
 #     return openai.ChatCompletion.create(**kwargs)
 
-# construct query, send to HF API and process response
+class ChatCompletionResult:
+    def __init__(self):
+        self.content = ""
+
+    def add_content(self, text):
+        self.content += text
+
+    def get_full_content(self):
+        return self.content.strip()
+
 def run_query(context, label, model_sel_name):
     '''
-    For non-streamed completion, enable the following 2 lines and comment out the code below
+    Summarize provided test
     '''
     chatbot_role = """You are an analyst specializing in climate change impact assessments and producing insights from policy documents."""
     messages = [{"role": "system", "content": chatbot_role},{"role": "user", "content": get_prompt(context, label)}]
     
     # Initialize the client, pointing it to one of the available models
-    client = InferenceClient(model_sel_name, token = hf_token)
+    client = InferenceClient(model_sel_name, token=hf_token)
     
-    # instantiate ChatCompletion as a generator object (stream is set to True)
-    chat_completion = client.chat_completion(
+    # Instantiate ChatCompletion as a generator object (stream is set to True)
+    chat_completion = client.chat.completions.create(
         messages=messages,
         stream=True
     )
-    # chat_completion = completion_with_backoff(messages=messages, stream=True)
 
-    # iterate through the streamed output
-    report = []
+    # Create an object to store the full chat completion
+    completion_result = ChatCompletionResult()
     res_box = st.empty()
+
+    # Iterate through the streamed output
     for chunk in chat_completion:
-        # extract the object containing the text (totally different structure when streaming)
-        if chunk.choices is not None: # sometimes returns None - probably the prompt needs work
+        # Extract the object containing the text
+        if chunk.choices is not None:
             chunk_message = chunk.choices[0].delta
-            # test to make sure there is text in the object (some don't have)
             if 'content' in chunk_message:
-                report.append(chunk_message['content']) # extract the message
-                # add the latest text and merge it with all previous
-                result = "".join(report).strip()
-                # res_box.success(result) # output to response text box
-                res_box.success(result)
-
+                completion_result.add_content(chunk_message['content'])  # Store the message
+                # Add the latest text and merge it with all previous
+                result = completion_result.get_full_content()
+                res_box.success(result)  # Output to response text box
 
+    # Return the stored chat completion object for later use
+    return completion_result