Refactor conversation function to improve code clarity and add validation for image data format

app.py

@@ -205,7 +205,7 @@ def conversation(
     model_path,
 ):
     # Initialize LLM
-    if hf_token.strip() != "" and model_path.strip() != "":
+    if hf_token.strip() and model_path.strip():
         llm = HuggingFaceEndpoint(
             repo_id=model_path,
             temperature=temperature,
@@ -220,7 +220,7 @@ def conversation(
             huggingfacehub_api_token=os.getenv("P_HF_TOKEN", "None"),
         )
 
-    # Get vector database collections
+    # Get collections from the vector database
     text_collection = vectordb_client.get_collection(
         "text_db", embedding_function=sentence_transformer_ef
     )
@@ -239,9 +239,10 @@ def conversation(
         n_results=img_context,
     )
 
-    # Process similar images
+    # Process image results
     img_links = similar_images["metadatas"][0] if similar_images["metadatas"] else []
     images_and_locs = []
+
     for distance, link in zip(similar_images["distances"][0], img_links):
         try:
             img = Image.open(io.BytesIO(base64.b64decode(link["image"])))
@@ -250,43 +251,46 @@ def conversation(
         except Exception as e:
             print(f"Error decoding image: {e}")
 
-    # Fallback to placeholder if no images are found
+    # Fallback to placeholder if no images are valid
+    placeholder_path = "assets/placeholder.jpg"
     if not images_and_locs:
-        placeholder_path = "assets/placeholder.jpg"  # Ensure this exists
         if not os.path.exists(placeholder_path):
             raise FileNotFoundError(f"Placeholder image not found at {placeholder_path}")
         images_and_locs = [(placeholder_path, "No images found")]
 
-    # Prepare prompt for the LLM
-    img_desc = "\n".join(similar_images["documents"][0]) if images_and_locs else "No Images Are Provided"
-    template = """
-    Context:
-    {context}
-
-    Included Images:
-    {images}
-    
-    Question:
-    {question}
+    # Validate data format
+    if not all(
+        isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], (str, Image.Image)) and isinstance(item[1], str)
+        for item in images_and_locs
+    ):
+        raise ValueError("images_and_locs must be a list of (image, caption) tuples.")
 
-    Answer:
-
-    """
-    prompt = PromptTemplate(template=template, input_variables=["context", "question"])
+    # Prepare LLM prompt
+    img_desc = "\n".join(similar_images["documents"][0]) if images_and_locs else "No Images Are Provided"
     context = "\n\n".join(results)
-
-    # Generate response
+    prompt = PromptTemplate(
+        template="""
+        Context:
+        {context}
+
+        Included Images:
+        {images}
+        
+        Question:
+        {question}
+
+        Answer:
+        """,
+        input_variables=["context", "question"],
+    )
     response = llm(prompt.format(context=context, question=msg, images=img_desc))
 
-    # Validate `images_and_locs` before returning
-    if not all(isinstance(item, tuple) and len(item) == 2 for item in images_and_locs):
-        raise ValueError("ret_images must be a list of (media, caption) tuples.")
+    print("Images and Locations:", images_and_locs)  # Debug output
 
     return history + [(msg, response)], results, images_and_locs
 
 
 
-
 def check_validity_and_llm(session_states):
     if session_states.get("processed", False) == True:
         return gr.Tabs(selected=2)