Update app.py

app.py

@@ -26,21 +26,6 @@ if 'FIREWORKS_API_KEY' not in os.environ:
 if 'MISTRAL_API_KEY' not in os.environ:
     os.environ['MISTRAL_API_KEY'] = getpass.getpass('Mistral AI API Key:')
 
-# Create prompt function
-@pxt.udf
-def create_prompt(top_k_list: list[dict], question: str) -> str:
-    concat_top_k = '\n\n'.join(
-        elt['text'] for elt in reversed(top_k_list)
-    )
-    return f'''
-    PASSAGES:
-
-    {concat_top_k}
-
-    QUESTION:
-
-    {question}'''
-
 """Gradio Application"""
 def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, show_question, show_correct_answer, show_gpt4omini, show_llamav3p23b, show_mistralsmall, progress=gr.Progress()):
     # Ensure a clean slate for the demo by removing and recreating the 'rag_demo' directory
@@ -86,6 +71,17 @@ def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, sh
         string_embed=sentence_transformer.using(model_id='sentence-transformers/all-MiniLM-L12-v2')
     )
 
+    # Create prompt function
+    @pxt.udf
+    def create_prompt(top_k_list: list[dict], question: str) -> str:
+        if not top_k_list:
+            return f"QUESTION:\n{question}"
+        
+        concat_top_k = '\n\n'.join(
+            elt['text'] for elt in reversed(top_k_list) if elt and 'text' in elt
+        )
+        return f'''PASSAGES:\n{concat_top_k}\n\nQUESTION:\n{question}'''
+    
     # Define a query function to retrieve the top-k most similar chunks for a given question
     @chunks_t.query
     def top_k(query_text: str):
@@ -96,52 +92,53 @@ def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, sh
               .limit(5)
       )
 
-    # Add computed columns to the queries table for context retrieval and prompt creation
+    # Then modify the messages structure to use a UDF
+    @pxt.udf
+    def create_messages(prompt: str) -> list[dict]:
+        return [
+            {
+                'role': 'system',
+                'content': 'Read the following passages and answer the question based on their contents.'
+            },
+            {
+                'role': 'user',
+                'content': prompt
+            }
+        ]
+    
+    # First add the context and prompt columns
     queries_t.add_computed_column(question_context=chunks_t.queries.top_k(queries_t.question))
     queries_t.add_computed_column(prompt=create_prompt(
-        queries_t.question_context, queries_t.question
+        queries_t.question_context, 
+        queries_t.question
     ))
-
-    # Prepare messages for the OpenAI API, including system instructions and user prompt
-    msgs = [
-        {
-            'role': 'system',
-            'content': 'Read the following passages and answer the question based on their contents.'
-        },
-        {
-            'role': 'user',
-            'content': queries_t.prompt
-        }
-    ]
-
-    progress(0.6, desc="Querying models...")
-
-    # Add OpenAI response column
+    
+    # Add the messages column
+    queries_t.add_computed_column(messages=create_messages(queries_t.prompt))
+    
+    # Then add the response columns using the messages
     queries_t.add_computed_column(response=openai.chat_completions(
         model='gpt-4o-mini-2024-07-18',
-        messages=msgs,
+        messages=queries_t.messages,
         max_tokens=300,
         top_p=0.9,
         temperature=0.7
     ))
-
-    # Create a table in Pixeltable and pick a model hosted on Anthropic with some parameters
+    
     queries_t.add_computed_column(response_2=f_chat_completions(
-      messages=msgs,
-      model='accounts/fireworks/models/llama-v3p2-3b-instruct',
-      # These parameters are optional and can be used to tune model behavior:
-      max_tokens=300,
-      top_p=0.9,
-      temperature=0.7
+        messages=queries_t.messages,
+        model='accounts/fireworks/models/llama-v3p2-3b-instruct',
+        max_tokens=300,
+        top_p=0.9,
+        temperature=0.7
     ))
-
+    
     queries_t.add_computed_column(response_3=chat_completions(
-      messages=msgs,
-      model='mistral-small-latest',
-      # These parameters are optional and can be used to tune model behavior:
-      max_tokens=300,
-      top_p=0.9,
-      temperature=0.7
+        messages=queries_t.messages,
+        model='mistral-small-latest',
+        max_tokens=300,
+        top_p=0.9,
+        temperature=0.7
     ))
 
     # Extract the answer text from the API response