Update src/streamlit_app.py

src/streamlit_app.py

@@ -28,7 +28,9 @@ from datasets import load_dataset, get_dataset_split_names
 from PIL import Image
 import openai
 import comet_llm
-from opik import track
+from functools import wraps
+from opik import track as opik_track
+
 
 
 
@@ -109,7 +111,13 @@ st.title("🩺 Multimodal Medical Chatbot")
 query = st.text_input("Enter your medical question or symptom description:")
 uploaded_file = st.file_uploader("Upload an image to find similar medical cases:", type=["png", "jpg", "jpeg"])
 
-@track
+def safe_track(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        return opik_track(func)(*args, **kwargs)
+    return wrapper
+
+@safe_track
 def get_chat_completion_openai(client, prompt: str):
     return client.chat.completions.create(
         model="gpt-4o",  # or "gpt-4" if you need the older GPT-4
@@ -118,7 +126,7 @@ def get_chat_completion_openai(client, prompt: str):
         max_tokens=150
     )
 
-@track
+@safe_track
 def get_similar_prompt(query):
     text_embeddings = embed_dataset_texts(combined_texts)  # cached
     query_embedding = embed_query_text(query)  # recalculated each time