Update app.py

app.py

@@ -191,6 +191,8 @@ def create_ui(agent):
                         extracted = "\n".join(results)
                         file_hash_value = file_hash(files[0].name)
 
+                # Truncate extracted content to avoid token limit issues
+                max_content_length = 8000  # Reduced from 12000 to prevent token overflow
                 prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
@@ -198,7 +200,7 @@ def create_ui(agent):
 4. Highlight abnormal results needing follow-up
 
 Medical Records:
-{extracted[:12000]}
+{extracted[:max_content_length]}
 
 ### Potential Oversights:
 """
@@ -210,52 +212,61 @@ Medical Records:
                 full_response = ""
                 response_chunks = []
 
-                # Process streaming response
-                for chunk in agent.run_gradio_chat(
-                    message=prompt,
-                    history=[],
-                    temperature=0.2,
-                    max_new_tokens=2048,
-                    max_token=4096,
-                    call_agent=False,
-                    conversation=[]
-                ):
-                    try:
-                        if chunk is None:
-                            continue
-                        
-                        # Handle different chunk types
-                        if isinstance(chunk, str):
-                            chunk_content = chunk
-                        elif hasattr(chunk, 'content'):
-                            chunk_content = chunk.content
-                        elif isinstance(chunk, list):
-                            chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-                        else:
-                            print("DEBUG: Received unknown type chunk", type(chunk))
+                # Process streaming response with error handling
+                try:
+                    for chunk in agent.run_gradio_chat(
+                        message=prompt,
+                        history=[],
+                        temperature=0.2,
+                        max_new_tokens=2048,
+                        max_token=4096,
+                        call_agent=False,
+                        conversation=[]
+                    ):
+                        try:
+                            if chunk is None:
+                                continue
+                            
+                            # Handle different chunk types
+                            if isinstance(chunk, str):
+                                chunk_content = chunk
+                            elif hasattr(chunk, 'content'):
+                                chunk_content = chunk.content
+                            elif isinstance(chunk, list):
+                                chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
+                            else:
+                                print("DEBUG: Received unknown type chunk", type(chunk))
+                                continue
+
+                            if not chunk_content:
+                                continue
+
+                            response_chunks.append(chunk_content)
+                            full_response = "".join(response_chunks)
+
+                            # Clean the response for display
+                            display_response = full_response.split('[TOOL_CALLS]')[0].strip()
+                            display_response = display_response.replace('[TxAgent]', '').strip()
+
+                            # Update the chat history with the latest response
+                            if len(history) > 0 and history[-1]["role"] == "assistant":
+                                history[-1]["content"] = display_response
+                            else:
+                                history.append({"role": "assistant", "content": display_response})
+
+                            yield history, None
+
+                        except Exception as e:
+                            print("❌ Error processing chunk:", str(e))
+                            traceback.print_exc()
                             continue
 
-                        if not chunk_content:
-                            continue
-
-                        response_chunks.append(chunk_content)
-                        full_response = "".join(response_chunks)
-
-                        # Remove any tool call markers from the displayed response
-                        display_response = full_response.split('[TOOL_CALLS]')[0].strip()
-
-                        # Update the chat history with the latest response
-                        if len(history) > 0 and history[-1]["role"] == "assistant":
-                            history[-1]["content"] = display_response
-                        else:
-                            history.append({"role": "assistant", "content": display_response})
-
-                        yield history, None
-
-                    except Exception as e:
-                        print("❌ Error processing chunk:", str(e))
-                        traceback.print_exc()
-                        continue
+                except Exception as e:
+                    print("❌ Error in model streaming:", str(e))
+                    traceback.print_exc()
+                    history.append({"role": "assistant", "content": f"Error in model response: {str(e)}"})
+                    yield history, None
+                    return
 
                 # Final response handling
                 if not full_response:
@@ -263,6 +274,7 @@ Medical Records:
                 else:
                     # Clean up the final response
                     full_response = full_response.split('[TOOL_CALLS]')[0].strip()
+                    full_response = full_response.replace('[TxAgent]', '').strip()
 
                 # Save report if we have files
                 report_path = None