Update app.py

app.py

@@ -26,38 +26,27 @@ index.add(embeddings)
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
 
 # --- RAG pipeline ---
-def generate_answer(messages):
-    try:
-        sys.stderr.write("=== generate_answer called ===\n")
-        sys.stderr.flush()
-        
-        user_query = messages[-1]["content"]
-        query_embedding = embedding_model.encode([user_query], convert_to_numpy=True)
-        faiss.normalize_L2(query_embedding)
-        _, indices = index.search(query_embedding, k=5)
-
-        context = "\n".join([f"{codes[i]}: {descriptions[i]}" for i in indices[0]])
-        prompt = f"""Here are some tariff code descriptions:
-{context}
-
-Question: {user_query}
-Answer:"""
-
-        sys.stderr.write(f"Prompt sent to model:\n{prompt}\n")
-        sys.stderr.flush()
-
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=200,
-            temperature=0.7,
-            stop_sequences=["\n\n"]
-        )
-        return {"role": "assistant", "content": response.strip()}
-
-    except Exception as e:
-        sys.stderr.write(f"Error during inference: {type(e).__name__}: {e}\n")
-        sys.stderr.flush()
-        return {"role": "assistant", "content": "An internal error occurred. Please try again."}
+def generate_answer(user_query):
+    sys.stderr.write("=== generate_answer called ===\n")
+    sys.stderr.flush()
+
+    query_embedding = embedding_model.encode([user_query], convert_to_numpy=True)
+    faiss.normalize_L2(query_embedding)
+    _, indices = index.search(query_embedding, k=5)
+
+    context = "\n".join([f"{codes[i]}: {descriptions[i]}" for i in indices[0]])
+    prompt = f"""Here are some tariff code descriptions:\n{context}\n\nQuestion: {user_query}\nAnswer:"""
+
+    sys.stderr.write(f"Prompt sent to model:\n{prompt}\n")
+    sys.stderr.flush()
+
+    response = client.text_generation(
+        prompt,
+        max_new_tokens=200,
+        temperature=0.7,
+        stop_sequences=["\n\n"]
+    )
+    return response.strip()
 
 # --- Gradio Chat Interface ---
 gr.ChatInterface(