Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -1,15 +1,23 @@
-# txagent.py - Core TxAgent class (simplified but maintains key functionality)
 import os
 import logging
 import torch
-import json
+import pdfplumber
+import pandas as pd
 from typing import Dict, Optional, List, Union
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 from sentence_transformers import SentenceTransformer
 from tooluniverse import ToolUniverse
 from .toolrag import ToolRAGModel
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler('txagent_core.log')
+    ]
+)
 logger = logging.getLogger("TxAgent")
 
 class TxAgent:
@@ -69,10 +77,14 @@ class TxAgent:
 
     def init_model(self):
         """Initialize all models and components"""
-        self.load_llm_model()
-        self.load_rag_model()
-        self.load_tooluniverse()
-        logger.info("All models initialized successfully")
+        try:
+            self.load_llm_model()
+            self.load_rag_model()
+            self.load_tooluniverse()
+            logger.info("All models initialized successfully")
+        except Exception as e:
+            logger.error(f"Model initialization failed: {str(e)}", exc_info=True)
+            raise
 
     def load_llm_model(self):
         """Load the main LLM model"""
@@ -93,7 +105,7 @@ class TxAgent:
             )
             logger.info(f"LLM model loaded on {self.device}")
         except Exception as e:
-            logger.error(f"Failed to load LLM model: {str(e)}")
+            logger.error(f"Failed to load LLM model: {str(e)}", exc_info=True)
             raise
 
     def load_rag_model(self):
@@ -103,13 +115,13 @@ class TxAgent:
             self.rag_model = ToolRAGModel(self.rag_model_name)
             logger.info("RAG model loaded successfully")
         except Exception as e:
-            logger.error(f"Failed to load RAG model: {str(e)}")
+            logger.error(f"Failed to load RAG model: {str(e)}", exc_info=True)
             raise
 
     def load_tooluniverse(self):
         """Initialize the ToolUniverse"""
         try:
-            logger.info("Loading ToolUniverse")
+            logger.info("Loading ToolUniverse with files: %s", self.tool_files_dict)
             self.tooluniverse = ToolUniverse(tool_files=self.tool_files_dict)
             self.tooluniverse.load_tools()
             
@@ -120,7 +132,7 @@ class TxAgent:
             
             logger.info(f"ToolUniverse loaded with {len(self.special_tools_name)} special tools")
         except Exception as e:
-            logger.error(f"Failed to load ToolUniverse: {str(e)}")
+            logger.error(f"Failed to load ToolUniverse: {str(e)}", exc_info=True)
             raise
 
     def chat(self, message: str, history: Optional[List[Dict]] = None, 
@@ -164,7 +176,7 @@ class TxAgent:
             return response.strip()
             
         except Exception as e:
-            logger.error(f"Chat failed: {str(e)}")
+            logger.error(f"Chat failed: {str(e)}", exc_info=True)
             raise RuntimeError(f"Chat failed: {str(e)}")
 
     def run_multistep_agent(self, message: str, temperature: float = 0.7, 
@@ -174,7 +186,9 @@ class TxAgent:
             conversation = [{"role": "system", "content": self.prompt_multi_step}]
             conversation.append({"role": "user", "content": message})
             
-            for _ in range(max_round):
+            for round_num in range(1, max_round + 1):
+                logger.info(f"Starting reasoning round {round_num}/{max_round}")
+                
                 # Generate next step
                 inputs = self.tokenizer.apply_chat_template(
                     conversation,
@@ -198,24 +212,31 @@ class TxAgent:
                 
                 # Check for final answer
                 if "[FinalAnswer]" in response:
-                    return response.split("[FinalAnswer]")[-1].strip()
+                    final_answer = response.split("[FinalAnswer]")[-1].strip()
+                    logger.info(f"Final answer found in round {round_num}")
+                    return final_answer
                 
                 # Add to conversation
                 conversation.append({"role": "assistant", "content": response})
+                logger.info(f"Round {round_num} completed without final answer")
                 
             # If max rounds reached
             if self.force_finish:
+                logger.info("Max rounds reached, forcing final answer")
                 return self._force_final_answer(conversation, temperature, max_new_tokens)
             
+            logger.warning("Max rounds reached without final answer")
             return "Reasoning rounds exceeded limit without reaching a final answer."
             
         except Exception as e:
-            logger.error(f"Multi-step agent failed: {str(e)}")
+            logger.error(f"Multi-step agent failed: {str(e)}", exc_info=True)
             raise RuntimeError(f"Multi-step agent failed: {str(e)}")
 
     def _force_final_answer(self, conversation: List[Dict], temperature: float, max_new_tokens: int) -> str:
         """Force a final answer when max rounds reached"""
         try:
+            logger.info("Attempting to force final answer")
+            
             # Add instruction to provide final answer
             conversation.append({
                 "role": "user",
@@ -244,17 +265,91 @@ class TxAgent:
             return response.strip()
             
         except Exception as e:
-            logger.error(f"Failed to force final answer: {str(e)}")
+            logger.error(f"Failed to force final answer: {str(e)}", exc_info=True)
             return "Failed to generate final answer."
 
+    def extract_text_from_file(self, file_path: str) -> Optional[str]:
+        """Extract text from PDF, CSV, or Excel files"""
+        try:
+            logger.info(f"Extracting text from file: {file_path}")
+            
+            if file_path.endswith('.pdf'):
+                with pdfplumber.open(file_path) as pdf:
+                    text = "\n".join(
+                        page.extract_text() 
+                        for page in pdf.pages 
+                        if page.extract_text()
+                    )
+                    logger.info(f"Extracted {len(text)} characters from PDF")
+                    return text
+            
+            elif file_path.endswith('.csv'):
+                df = pd.read_csv(file_path)
+                text = df.to_string()
+                logger.info(f"Extracted {len(text)} characters from CSV")
+                return text
+                
+            elif file_path.endswith(('.xlsx', '.xls')):
+                df = pd.read_excel(file_path)
+                text = df.to_string()
+                logger.info(f"Extracted {len(text)} characters from Excel")
+                return text
+            
+            logger.warning(f"Unsupported file type: {file_path}")
+            return None
+            
+        except Exception as e:
+            logger.error(f"Text extraction failed: {str(e)}", exc_info=True)
+            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"""
+        try:
+            logger.info(f"Analyzing text (first 100 chars): {text[:100]}...")
+            
+            prompt = f"""Analyze this medical document:
+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)
+            
+            generation_config = GenerationConfig(
+                max_new_tokens=max_tokens,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+            
+            outputs = self.model.generate(
+                **inputs,
+                generation_config=generation_config
+            )
+            
+            analysis = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            logger.info("Text analysis completed successfully")
+            return analysis
+            
+        except Exception as e:
+            logger.error(f"Text analysis failed: {str(e)}", exc_info=True)
+            raise RuntimeError(f"Analysis failed: {str(e)}")
+
     def cleanup(self):
         """Clean up resources"""
-        if hasattr(self, 'model'):
-            del self.model
-        if hasattr(self, 'rag_model'):
-            del self.rag_model
-        torch.cuda.empty_cache()
-        logger.info("TxAgent resources cleaned up")
+        try:
+            logger.info("Cleaning up TxAgent resources")
+            if hasattr(self, 'model'):
+                del self.model
+            if hasattr(self, 'rag_model'):
+                del self.rag_model
+            torch.cuda.empty_cache()
+            logger.info("TxAgent resources cleaned up")
+        except Exception as e:
+            logger.error(f"Cleanup failed: {str(e)}", exc_info=True)
+            raise
 
     def __del__(self):
         """Destructor to ensure proper cleanup"""