Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -7,7 +7,6 @@ from typing import Dict, Optional, Union
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 
-# Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger("TxAgent")
 
@@ -20,18 +19,7 @@ class TxAgent:
                  enable_checker: bool = True,
                  step_rag_num: int = 4,
                  seed: Optional[int] = None):
-        """
-        Initialize the TxAgent with specified configuration.
-        
-        Args:
-            model_name: Name/path of the main LLM model
-            rag_model_name: Name/path of the RAG model
-            tool_files_dict: Dictionary of tool files
-            force_finish: Whether to force finish when max tokens reached
-            enable_checker: Whether to enable reasoning trace checker
-            step_rag_num: Number of RAG tools to retrieve per step
-            seed: Random seed for reproducibility
-        """
+        """Initialize TxAgent without vLLM dependencies."""
         self.model_name = model_name
         self.rag_model_name = rag_model_name
         self.tool_files_dict = tool_files_dict or {}
@@ -48,24 +36,24 @@ class TxAgent:
         logger.info(f"Initialized TxAgent with model: {model_name} on device: {self.device}")
 
     def init_model(self):
-        """Initialize both the main model and RAG model."""
+        """Initialize models using transformers only."""
         self.load_llm_model()
         self.load_rag_model()
         logger.info("Model initialization complete")
 
     def load_llm_model(self):
-        """Load the main LLM model."""
+        """Load the main LLM model using transformers."""
         try:
             logger.info(f"Loading LLM model: {self.model_name}")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_name,
-                cache_dir=os.getenv("TRANSFORMERS_CACHE")
+                cache_dir=os.getenv("HF_HOME")
             )
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_name,
                 torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
                 device_map="auto",
-                cache_dir=os.getenv("TRANSFORMERS_CACHE")
+                cache_dir=os.getenv("HF_HOME")
             )
             logger.info(f"LLM model loaded on {self.device}")
         except Exception as e:
@@ -86,29 +74,16 @@ class TxAgent:
             raise RuntimeError(f"Failed to load RAG model: {str(e)}")
 
     def process_document(self, file_path: str) -> Dict[str, Union[str, Dict]]:
-        """
-        Process a medical document and return analysis results.
-        
-        Args:
-            file_path: Path to the document file (PDF, CSV, or Excel)
-            
-        Returns:
-            Dictionary containing:
-            - status: "success" or "error"
-            - analysis: Detailed analysis results or error message
-            - model: Model used for analysis
-        """
+        """Process a document and return real analysis results."""
         try:
-            # 1. Extract text from document
             text = self.extract_text_from_file(file_path)
             if not text:
                 return {
                     "status": "error",
-                    "message": "Failed to extract text - unsupported file type or empty document",
+                    "message": "Failed to extract text",
                     "model": self.model_name
                 }
 
-            # 2. Analyze with LLM
             analysis = self.analyze_text(text)
             
             return {
@@ -118,23 +93,15 @@ class TxAgent:
             }
             
         except Exception as e:
-            logger.error(f"Document processing failed: {str(e)}", exc_info=True)
+            logger.error(f"Document processing failed: {str(e)}")
             return {
                 "status": "error",
-                "message": f"Processing error: {str(e)}",
+                "message": str(e),
                 "model": self.model_name
             }
 
     def extract_text_from_file(self, file_path: str) -> Optional[str]:
-        """
-        Extract text from supported file types (PDF, CSV, Excel).
-        
-        Args:
-            file_path: Path to the input file
-            
-        Returns:
-            Extracted text as string, or None if extraction fails
-        """
+        """Extract text from PDF, CSV, or Excel files."""
         try:
             if file_path.endswith('.pdf'):
                 with pdfplumber.open(file_path) as pdf:
@@ -160,18 +127,9 @@ class TxAgent:
             raise RuntimeError(f"Text extraction failed: {str(e)}")
 
     def analyze_text(self, text: str, max_tokens: int = 1000) -> str:
-        """
-        Analyze extracted text using the LLM.
-        
-        Args:
-            text: Text to analyze
-            max_tokens: Maximum tokens to generate
-            
-        Returns:
-            Analysis results as string
-        """
+        """Analyze extracted text using the LLM."""
         try:
-            prompt = f"""Analyze this medical document and provide:
+            prompt = f"""Analyze this medical document:
 1. Diagnostic patterns
 2. Medication issues
 3. Recommended follow-ups
@@ -192,7 +150,7 @@ Document:
             raise RuntimeError(f"Analysis failed: {str(e)}")
 
     def cleanup(self):
-        """Clean up resources and clear memory."""
+        """Clean up resources."""
         if hasattr(self, 'model'):
             del self.model
         if hasattr(self, 'rag_model'):