Refactor conversation function to improve image handling and add descriptive comments

app.py

@@ -204,6 +204,7 @@ def conversation(
     hf_token,
     model_path,
 ):
+    # Initialize the LLM based on inputs
     if hf_token.strip() != "" and model_path.strip() != "":
         llm = HuggingFaceEndpoint(
             repo_id=model_path,
@@ -219,6 +220,7 @@ def conversation(
             huggingfacehub_api_token=os.getenv("P_HF_TOKEN", "None"),
         )
 
+    # Retrieve collections from vector database
     text_collection = vectordb_client.get_collection(
         "text_db", embedding_function=sentence_transformer_ef
     )
@@ -226,23 +228,36 @@ def conversation(
         "image_db", embedding_function=sentence_transformer_ef
     )
 
+    # Query text context
     results = text_collection.query(
         query_texts=[msg], include=["documents"], n_results=num_context
     )["documents"][0]
+
+    # Query image context
     similar_images = image_collection.query(
         query_texts=[msg],
         include=["metadatas", "distances", "documents"],
         n_results=img_context,
     )
-    img_links = [i["image"] for i in similar_images["metadatas"][0]]
-
-    images_and_locs = [
-        Image.open(io.BytesIO(base64.b64decode(i[1])))
-        for i in zip(similar_images["distances"][0], img_links)
-    ]
-    img_desc = "\n".join(similar_images["documents"][0])
-    if len(img_links) == 0:
-        img_desc = "No Images Are Provided"
+
+    # Initialize image links and descriptions
+    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"])))
+            caption = f"Distance: {distance:.2f}"
+            images_and_locs.append((img, caption))
+        except Exception as e:
+            print(f"Error decoding image: {e}")
+
+    # Handle case where no images are found
+    if not images_and_locs:
+        images_and_locs = [("path/to/placeholder/image.jpg", "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}
@@ -258,11 +273,15 @@ def conversation(
     """
     prompt = PromptTemplate(template=template, input_variables=["context", "question"])
     context = "\n\n".join(results)
-    # references = [gr.Textbox(i, visible=True, interactive=False) for i in results]
+
+    # Generate response
     response = llm(prompt.format(context=context, question=msg, images=img_desc))
+
+    # Return updated history, text results, and image locations
     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)