Update app.py

app.py

@@ -35,10 +35,10 @@ def encode_image(image):
         logger.error(f"Error encoding image: {str(e)}")
         raise
 
-def analyze_construction_image(image, follow_up_question=""):
+def analyze_construction_image(image):
     if image is None:
         logger.warning("No image provided")
-        return "Error: No image uploaded", "", ""
+        return "Error: No image uploaded", "", "", []
 
     try:
         logger.info("Starting image analysis")
@@ -62,12 +62,6 @@ def analyze_construction_image(image, follow_up_question=""):
             }
         ]
 
-        if follow_up_question:
-            messages.append({
-                "role": "user",
-                "content": follow_up_question
-            })
-
         logger.info("Sending request to Groq API")
         completion = client.chat.completions.create(
             model="llama-3.2-90b-vision-preview",
@@ -89,39 +83,84 @@ def analyze_construction_image(image, follow_up_question=""):
             parsed_result = json.loads(result)
         except json.JSONDecodeError:
             logger.error("Failed to parse API response as JSON")
-            return "Error: Invalid response format", "", ""
+            return "Error: Invalid response format", "", "", []
 
         snag_category = parsed_result.get('snag_category', 'N/A')
         snag_description = parsed_result.get('snag_description', 'N/A')
         desnag_steps = '\n'.join(parsed_result.get('desnag_steps', ['N/A']))
 
         logger.info("Analysis completed successfully")
-        return snag_category, snag_description, desnag_steps
+        
+        # Initialize chat history with analysis results
+        chat_history = [
+            (None, f"Snag Category: {snag_category}\n\nSnag Description: {snag_description}\n\nSteps to Desnag:\n{desnag_steps}")
+        ]
+        
+        return snag_category, snag_description, desnag_steps, chat_history
     except Exception as e:
         logger.error(f"Error during image analysis: {str(e)}")
-        return f"Error: {str(e)}", "", ""
+        return f"Error: {str(e)}", "", "", []
+
+def chat_about_image(message, chat_history):
+    try:
+        # Prepare the conversation history for the API
+        messages = [
+            {"role": "system", "content": "You are an AI assistant specialized in analyzing construction site images and answering questions about them. Use the information from the initial analysis to answer user queries."},
+        ]
+        
+        # Add chat history to messages
+        for human, ai in chat_history:
+            if human:
+                messages.append({"role": "user", "content": human})
+            if ai:
+                messages.append({"role": "assistant", "content": ai})
+        
+        # Add the new user message
+        messages.append({"role": "user", "content": message})
+        
+        # Make API call
+        completion = client.chat.completions.create(
+            model="llama-3.2-90b-chat",  # Using the chat model here
+            messages=messages,
+            temperature=0.7,
+            max_tokens=500,
+            top_p=1,
+            stream=False,
+            stop=None
+        )
+        
+        response = completion.choices[0].message.content
+        chat_history.append((message, response))
+        
+        return "", chat_history
+    except Exception as e:
+        logger.error(f"Error during chat: {str(e)}")
+        return "", chat_history + [(message, f"Error: {str(e)}")]
 
 # Create the Gradio interface
-iface = gr.Interface(
-    fn=analyze_construction_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Construction Image"),
-        gr.Textbox(label="Follow-up Question (Optional)")
-    ],
-    outputs=[
-        gr.Textbox(label="Snag Category"),
-        gr.Textbox(label="Snag Description"),
-        gr.Textbox(label="Steps to Desnag")
-    ],
-    title="Construction Image Analyzer (Llama 3.2 90B Vision via Groq)",
-    description="Upload a construction site image to identify issues and get desnag steps using Llama 3.2 90B Vision technology through Groq API. You can also ask follow-up questions about the image.",
-    examples=[
-        ["example_image1.jpg", "What safety concerns do you see?"],
-        ["example_image2.jpg", "Is there any visible structural damage?"]
-    ],
-    cache_examples=False,
-    theme="default"
-)
+with gr.Blocks() as iface:
+    gr.Markdown("# Construction Image Analyzer with Chat")
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload Construction Image")
+            analyze_button = gr.Button("Analyze Image")
+        with gr.Column(scale=2):
+            snag_category = gr.Textbox(label="Snag Category")
+            snag_description = gr.Textbox(label="Snag Description")
+            desnag_steps = gr.Textbox(label="Steps to Desnag")
+    
+    chatbot = gr.Chatbot(label="Chat about the Image")
+    msg = gr.Textbox(label="Ask a question about the image")
+    clear = gr.Button("Clear")
+
+    analyze_button.click(
+        analyze_construction_image,
+        inputs=[image_input],
+        outputs=[snag_category, snag_description, desnag_steps, chatbot]
+    )
+
+    msg.submit(chat_about_image, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 # Launch the app
 if __name__ == "__main__":