Update app.py

app.py

@@ -23,71 +23,55 @@ class BasicAgent:
         
     def __call__(self, question: str, file_bytes: bytes = None, file_type: str = None) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        
-        file_info = ""
-        if file_bytes and file_type:
-            file_info = self.process_file(file_bytes, file_type)
-
-        full_prompt = question
-        if file_info:
-            full_prompt += f"\n\n[File Content Below]\n{file_info}"
-            
-        if file_type == "png":
-            try:
-                encoded_image = base64.b64encode(file_bytes).decode("utf-8")
-
-                messages = [HumanMessage(content=[
-                    {"type": "text", "text": question},
-                    {"type": "image_url", "image_url": f"data:image/png;base64,{encoded_image}"}
-                ])]
-
-                response = self.graph.invoke({"messages": messages})
-                answer = response['messages'][-1].content
-                return answer.split("FINAL ANSWER:")[-1].strip()
-
-            except Exception as e:
-                print(f"Error processing image: {e}")
-                return f"Error: {e}"
-
-        elif file_type == "mp3":
-            try:
-                encoded_audio = base64.b64encode(file_bytes).decode("utf-8")
-                audio_mime_type = "audio/mpeg"
-
-                messages = [HumanMessage(content=[
-                    {"type": "text", "text": question},
-                    {
-                        "type": "media",
-                        "data": encoded_audio,
-                        "mime_type": audio_mime_type,
-                    }
-                ])]
-
-                response = self.graph.invoke({"messages": messages})
-                answer = response['messages'][-1].content
-                return answer.split("FINAL ANSWER:")[-1].strip()
-
-            except Exception as e:
-                print(f"Error processing audio: {e}")
-                return f"Error: {e}"
-
-        else:
-            messages = [HumanMessage(content=full_prompt)]
-            messages = self.graph.invoke({"messages": messages})
-            answer = messages['messages'][-1].content
-            return answer.split("FINAL ANSWER:")[-1].strip()
     
-    def process_file(self, file_bytes: bytes, file_type: str) -> str:
+        messages = []
         try:
-            if file_type == "py":
-                return file_bytes.decode("utf-8")
-
-            elif file_type == "xlsx":
-                df = pd.read_excel(BytesIO(file_bytes))
-                return df.to_string()
-
+            # Handle special formats
+            if file_bytes and file_type:
+                if file_type == "png":
+                    encoded_image = base64.b64encode(file_bytes).decode("utf-8")
+                    messages = [HumanMessage(content=[
+                        {"type": "text", "text": question},
+                        {"type": "image_url", "image_url": f"data:image/png;base64,{encoded_image}"}
+                    ])]
+    
+                elif file_type == "mp3":
+                    encoded_audio = base64.b64encode(file_bytes).decode("utf-8")
+                    audio_mime_type = "audio/mpeg"
+                    messages = [HumanMessage(content=[
+                        {"type": "text", "text": question},
+                        {
+                            "type": "media",
+                            "data": encoded_audio,
+                            "mime_type": audio_mime_type,
+                        }
+                    ])]
+    
+                else:
+                    # For general files like .py, .xlsx
+                    file_info = ""
+                    if file_type == "py":
+                        file_info = file_bytes.decode("utf-8")
+                    elif file_type == "xlsx":
+                        df = pd.read_excel(BytesIO(file_bytes))
+                        file_info = df.to_string()
+                    else:
+                        file_info = f"[Unsupported file type: {file_type}]"
+    
+                    full_prompt = f"{question}\n\n[File Content Below]\n{file_info}"
+                    messages = [HumanMessage(content=full_prompt)]
+    
+            else:
+                # No file attached
+                messages = [HumanMessage(content=question)]
+    
+            # Invoke final agent and extract answer
+            response = self.final_agent.invoke({"messages": messages})
+            answer = response['messages'][-1].content
+            return answer.split("FINAL ANSWER:")[-1].strip()
+    
         except Exception as e:
-            return f"Error processing file: {e}"
+            return f"Error during processing: {e}"
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """