Update app.py

app.py

@@ -4,6 +4,7 @@ import pandas as pd
 import gradio as gr
 import re
 import hashlib
+import shutil
 from datetime import datetime
 from collections import defaultdict
 from typing import List, Dict, Tuple
@@ -13,10 +14,38 @@ WORKING_DIR = os.getcwd()
 REPORT_DIR = os.path.join(WORKING_DIR, "reports")
 os.makedirs(REPORT_DIR, exist_ok=True)
 
+# Model configuration
+MODEL_CACHE_DIR = os.path.join(WORKING_DIR, "model_cache")
+os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
+os.environ["HF_HOME"] = MODEL_CACHE_DIR
+os.environ["TRANSFORMERS_CACHE"] = MODEL_CACHE_DIR
+
+# Import TxAgent after setting up environment
+sys.path.append(os.path.join(WORKING_DIR, "src"))
+from txagent.txagent import TxAgent
+
 class PatientHistoryAnalyzer:
     def __init__(self):
         self.max_token_length = 2000
         self.max_text_length = 500
+        self.agent = self._initialize_agent()
+
+    def _initialize_agent(self):
+        """Initialize the TxAgent with proper configuration"""
+        tool_path = os.path.join(WORKING_DIR, "data", "new_tool.json")
+        if not os.path.exists(tool_path):
+            raise FileNotFoundError(f"Tool file not found at {tool_path}")
+        
+        return TxAgent(
+            model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+            rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+            tool_files_dict={"new_tool": tool_path},
+            force_finish=True,
+            enable_checker=True,
+            step_rag_num=4,
+            seed=100,
+            additional_default_tools=[],
+        )
 
     def clean_text(self, text: str) -> str:
         """Clean and normalize text fields"""
@@ -185,6 +214,30 @@ class PatientHistoryAnalyzer:
             "### Recommendations"
         ])
 
+    def _call_agent(self, prompt: str) -> str:
+        """Call TxAgent with proper error handling"""
+        try:
+            response = ""
+            for result in self.agent.run_gradio_chat(
+                message=prompt,
+                history=[],
+                temperature=0.2,
+                max_new_tokens=1024,
+                max_token=2048,
+                call_agent=False,
+                conversation=[],
+            ):
+                if isinstance(result, list):
+                    for r in result:
+                        if hasattr(r, 'content') and r.content:
+                            response += r.content + "\n"
+                elif isinstance(result, str):
+                    response += result + "\n"
+            
+            return response.strip()
+        except Exception as e:
+            return f"Error in model response: {str(e)}"
+
     def generate_report(self, analysis_results: List[str]) -> Tuple[str, str]:
         """Combine analysis results into final report"""
         report = [
@@ -226,14 +279,13 @@ class PatientHistoryAnalyzer:
             patient_data = self.process_excel(file_path)
             prompts = self.generate_analysis_prompt(patient_data)
             
-            # Simulate LLM responses
-            simulated_responses = [
-                "### Summary of Current Status\nPatient shows improvement in blood pressure control...",
-                "### Historical Patterns\nChronic back pain has been a consistent issue...",
-                "### Medication Summary\nCurrent regimen includes 4 medications..."
-            ]
+            # Call TxAgent for each prompt
+            analysis_results = []
+            for prompt in prompts:
+                response = self._call_agent(prompt['content'])
+                analysis_results.append(response)
             
-            return self.generate_report(simulated_responses)
+            return self.generate_report(analysis_results)
             
         except Exception as e:
             return f"Error during analysis: {str(e)}", ""