Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -1,8 +1,9 @@
 import os
-import json
 import logging
 import torch
-from typing import List, Dict, Optional, Union
+import pdfplumber
+import pandas as pd
+from typing import Dict, Optional, Union
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 
@@ -15,56 +16,61 @@ class TxAgent:
                  model_name: str,
                  rag_model_name: str,
                  tool_files_dict: Optional[Dict] = None,
-                 use_vllm: bool = False,
                  force_finish: bool = True,
                  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
+        """
         self.model_name = model_name
         self.rag_model_name = rag_model_name
         self.tool_files_dict = tool_files_dict or {}
-        self.use_vllm = use_vllm
         self.force_finish = force_finish
         self.enable_checker = enable_checker
         self.step_rag_num = step_rag_num
         self.seed = seed
         
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model = None
         self.tokenizer = None
         self.rag_model = None
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        logger.info(f"Initializing TxAgent with model: {model_name} on device: {self.device}")
+        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."""
-        self.load_models()
+        self.load_llm_model()
         self.load_rag_model()
         logger.info("Model initialization complete")
 
-    def load_models(self):
+    def load_llm_model(self):
         """Load the main LLM model."""
         try:
-            logger.info(f"Loading model: {self.model_name}")
-            
+            logger.info(f"Loading LLM model: {self.model_name}")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_name,
-                cache_dir=os.environ.get("TRANSFORMERS_CACHE")
+                cache_dir=os.getenv("TRANSFORMERS_CACHE")
             )
-            
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_name,
-                torch_dtype=torch.float16,
+                torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
                 device_map="auto",
-                cache_dir=os.environ.get("TRANSFORMERS_CACHE")
+                cache_dir=os.getenv("TRANSFORMERS_CACHE")
             )
-            
-            logger.info(f"Successfully loaded model on {self.device}")
-            
+            logger.info(f"LLM model loaded on {self.device}")
         except Exception as e:
-            logger.error(f"Failed to load model: {str(e)}")
-            raise RuntimeError(f"Failed to load model: {str(e)}")
+            logger.error(f"Failed to load LLM model: {str(e)}")
+            raise RuntimeError(f"Failed to load LLM model: {str(e)}")
 
     def load_rag_model(self):
         """Load the RAG model."""
@@ -79,37 +85,114 @@ class TxAgent:
             logger.error(f"Failed to load RAG model: {str(e)}")
             raise RuntimeError(f"Failed to load RAG model: {str(e)}")
 
-    def process_document(self, file_path: str) -> Dict:
-        """Process a document and return analysis results."""
-        try:
-            # Extract text (implement your extraction logic)
-            text = self.extract_text(file_path)
+    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)
             
-            # Process with LLM (implement your processing logic)
-            result = self.analyze_text(text)
+        Returns:
+            Dictionary containing:
+            - status: "success" or "error"
+            - analysis: Detailed analysis results or error message
+            - model: Model used for analysis
+        """
+        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",
+                    "model": self.model_name
+                }
+
+            # 2. Analyze with LLM
+            analysis = self.analyze_text(text)
             
             return {
                 "status": "success",
-                "analysis": result,
+                "analysis": analysis,
                 "model": self.model_name
             }
             
         except Exception as e:
-            logger.error(f"Document processing failed: {str(e)}")
-            raise RuntimeError(f"Document processing failed: {str(e)}")
+            logger.error(f"Document processing failed: {str(e)}", exc_info=True)
+            return {
+                "status": "error",
+                "message": f"Processing error: {str(e)}",
+                "model": self.model_name
+            }
 
-    def extract_text(self, file_path: str) -> str:
-        """Extract text from various file formats."""
-        # Implement your text extraction logic here
-        return "Sample extracted text"
+    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
+        """
+        try:
+            if file_path.endswith('.pdf'):
+                with pdfplumber.open(file_path) as pdf:
+                    return "\n".join(
+                        page.extract_text() 
+                        for page in pdf.pages 
+                        if page.extract_text()
+                    )
+            
+            elif file_path.endswith('.csv'):
+                df = pd.read_csv(file_path)
+                return df.to_string()
+                
+            elif file_path.endswith(('.xlsx', '.xls')):
+                df = pd.read_excel(file_path)
+                return df.to_string()
+            
+            logger.warning(f"Unsupported file type: {file_path}")
+            return None
+            
+        except Exception as e:
+            logger.error(f"Text extraction failed: {str(e)}")
+            raise RuntimeError(f"Text extraction failed: {str(e)}")
 
-    def analyze_text(self, text: str) -> str:
-        """Analyze extracted text using the LLM."""
-        # Implement your text analysis logic here
-        return "Sample analysis result"
+    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
+        """
+        try:
+            prompt = f"""Analyze this medical document and provide:
+1. Diagnostic patterns
+2. Medication issues
+3. Recommended follow-ups
+
+Document:
+{text[:8000]}  # Truncate to avoid token limits
+"""
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+        except Exception as e:
+            logger.error(f"Text analysis failed: {str(e)}")
+            raise RuntimeError(f"Analysis failed: {str(e)}")
 
     def cleanup(self):
-        """Clean up resources."""
+        """Clean up resources and clear memory."""
         if hasattr(self, 'model'):
             del self.model
         if hasattr(self, 'rag_model'):