Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -1,22 +1,14 @@
 import os
 import logging
 import torch
-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
 
-# Configure logging
+# Configure logging for Hugging Face Spaces
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler(),
-        logging.FileHandler('txagent_core.log')
-    ]
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger("TxAgent")
 
@@ -27,18 +19,10 @@ class TxAgent:
                  tool_files_dict: Optional[Dict] = None,
                  enable_finish: bool = True,
                  enable_rag: bool = False,
-                 enable_summary: bool = False,
-                 init_rag_num: int = 0,
-                 step_rag_num: int = 0,
-                 summary_mode: str = 'step',
-                 summary_skip_last_k: int = 0,
-                 summary_context_length: Optional[int] = None,
                  force_finish: bool = True,
-                 avoid_repeat: bool = True,
-                 seed: Optional[int] = None,
-                 enable_checker: bool = False,
-                 enable_chat: bool = False,
-                 additional_default_tools: Optional[List] = None):
+                 enable_checker: bool = True,
+                 step_rag_num: int = 4,
+                 seed: Optional[int] = None):
         
         # Initialization parameters
         self.model_name = model_name
@@ -46,18 +30,10 @@ class TxAgent:
         self.tool_files_dict = tool_files_dict or {}
         self.enable_finish = enable_finish
         self.enable_rag = enable_rag
-        self.enable_summary = enable_summary
-        self.summary_mode = summary_mode
-        self.summary_skip_last_k = summary_skip_last_k
-        self.summary_context_length = summary_context_length
-        self.init_rag_num = init_rag_num
-        self.step_rag_num = step_rag_num
         self.force_finish = force_finish
-        self.avoid_repeat = avoid_repeat
-        self.seed = seed
         self.enable_checker = enable_checker
-        self.enable_chat = enable_chat
-        self.additional_default_tools = additional_default_tools or []
+        self.step_rag_num = step_rag_num
+        self.seed = seed
         
         # Device setup
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -66,24 +42,21 @@ class TxAgent:
         self.model = None
         self.tokenizer = None
         self.rag_model = None
-        self.tooluniverse = None
         
         # Prompts
-        self.prompt_multi_step = "You are a helpful assistant that solves problems through step-by-step reasoning."
-        self.self_prompt = "Strictly follow the instruction."
         self.chat_prompt = "You are a helpful assistant for user chat."
         
-        logger.info(f"Initialized TxAgent with model: {model_name} on device: {self.device}")
+        logger.info(f"Initialized TxAgent with model: {model_name}")
 
     def init_model(self):
         """Initialize all models and components"""
         try:
             self.load_llm_model()
-            self.load_rag_model()
-            self.load_tooluniverse()
-            logger.info("All models initialized successfully")
+            if self.enable_rag:
+                self.load_rag_model()
+            logger.info("Models initialized successfully")
         except Exception as e:
-            logger.error(f"Model initialization failed: {str(e)}", exc_info=True)
+            logger.error(f"Model initialization failed: {str(e)}")
             raise
 
     def load_llm_model(self):
@@ -92,7 +65,6 @@ class TxAgent:
             logger.info(f"Loading LLM model: {self.model_name}")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_name,
-                cache_dir=os.getenv("HF_HOME"),
                 trust_remote_code=True
             )
             
@@ -100,39 +72,24 @@ class TxAgent:
                 self.model_name,
                 torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
                 device_map="auto",
-                cache_dir=os.getenv("HF_HOME"),
                 trust_remote_code=True
             )
             logger.info(f"LLM model loaded on {self.device}")
         except Exception as e:
-            logger.error(f"Failed to load LLM model: {str(e)}", exc_info=True)
+            logger.error(f"Failed to load LLM model: {str(e)}")
             raise
 
     def load_rag_model(self):
         """Load the RAG model"""
         try:
             logger.info(f"Loading RAG model: {self.rag_model_name}")
-            self.rag_model = ToolRAGModel(self.rag_model_name)
+            self.rag_model = SentenceTransformer(
+                self.rag_model_name,
+                device=str(self.device)
+            )
             logger.info("RAG model loaded successfully")
         except Exception as 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 with files: %s", self.tool_files_dict)
-            self.tooluniverse = ToolUniverse(tool_files=self.tool_files_dict)
-            self.tooluniverse.load_tools()
-            
-            # Prepare special tools
-            special_tools = self.tooluniverse.prepare_tool_prompts(
-                self.tooluniverse.tool_category_dicts["special_tools"])
-            self.special_tools_name = [tool['name'] for tool in special_tools]
-            
-            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)}", exc_info=True)
+            logger.error(f"Failed to load RAG model: {str(e)}")
             raise
 
     def chat(self, message: str, history: Optional[List[Dict]] = None, 
@@ -176,179 +133,20 @@ class TxAgent:
             return response.strip()
             
         except Exception as e:
-            logger.error(f"Chat failed: {str(e)}", exc_info=True)
+            logger.error(f"Chat failed: {str(e)}")
             raise RuntimeError(f"Chat failed: {str(e)}")
 
-    def run_multistep_agent(self, message: str, temperature: float = 0.7, 
-                           max_new_tokens: int = 512, max_round: int = 5) -> str:
-        """Run multi-step reasoning agent"""
-        try:
-            conversation = [{"role": "system", "content": self.prompt_multi_step}]
-            conversation.append({"role": "user", "content": message})
-            
-            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,
-                    add_generation_prompt=True,
-                    return_tensors="pt"
-                ).to(self.device)
-                
-                generation_config = GenerationConfig(
-                    max_new_tokens=max_new_tokens,
-                    temperature=temperature,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id
-                )
-                
-                outputs = self.model.generate(
-                    inputs,
-                    generation_config=generation_config
-                )
-                
-                response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-                
-                # Check for final answer
-                if "[FinalAnswer]" in response:
-                    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)}", 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",
-                "content": "Provide your final answer now based on all previous reasoning."
-            })
-            
-            inputs = self.tokenizer.apply_chat_template(
-                conversation,
-                add_generation_prompt=True,
-                return_tensors="pt"
-            ).to(self.device)
-            
-            generation_config = GenerationConfig(
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-            
-            outputs = self.model.generate(
-                inputs,
-                generation_config=generation_config
-            )
-            
-            response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-            return response.strip()
-            
-        except Exception as 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"""
         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")
+            logger.info("Resources cleaned up")
         except Exception as e:
-            logger.error(f"Cleanup failed: {str(e)}", exc_info=True)
+            logger.error(f"Cleanup failed: {str(e)}")
             raise
 
     def __del__(self):