Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -1,178 +1,261 @@
-# app.py - FastAPI application
+# txagent.py - Core TxAgent class (simplified but maintains key functionality)
 import os
-import sys
-import json
-import shutil
-from fastapi import FastAPI, HTTPException, UploadFile, File
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from typing import List, Dict, Optional
+import logging
 import torch
-from datetime import datetime
-from pydantic import BaseModel
-
-# Configuration
-persistent_dir = "/data/hf_cache"
-model_cache_dir = os.path.join(persistent_dir, "txagent_models")
-tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
-file_cache_dir = os.path.join(persistent_dir, "cache")
-report_dir = os.path.join(persistent_dir, "reports")
-
-# Create directories if they don't exist
-os.makedirs(model_cache_dir, exist_ok=True)
-os.makedirs(tool_cache_dir, exist_ok=True)
-os.makedirs(file_cache_dir, exist_ok=True)
-os.makedirs(report_dir, exist_ok=True)
-
-# Set environment variables
-os.environ["HF_HOME"] = model_cache_dir
-os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
-
-# Set up Python path
-current_dir = os.path.dirname(os.path.abspath(__file__))
-src_path = os.path.abspath(os.path.join(current_dir, "src"))
-sys.path.insert(0, src_path)
-
-# Request models
-class ChatRequest(BaseModel):
-    message: str
-    temperature: float = 0.7
-    max_new_tokens: int = 512
-    history: Optional[List[Dict]] = None
+import json
+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
 
-class MultistepRequest(BaseModel):
-    message: str
-    temperature: float = 0.7
-    max_new_tokens: int = 512
-    max_round: int = 5
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("TxAgent")
 
-# Initialize FastAPI app
-app = FastAPI(
-    title="TxAgent API",
-    description="API for TxAgent medical document analysis",
-    version="1.0.0"
-)
+class TxAgent:
+    def __init__(self, 
+                 model_name: str,
+                 rag_model_name: str,
+                 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):
+        
+        # Initialization parameters
+        self.model_name = model_name
+        self.rag_model_name = rag_model_name
+        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 []
+        
+        # Device setup
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Models
+        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}")
 
-# CORS configuration
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+    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")
 
-# Initialize agent at startup
-agent = None
+    def load_llm_model(self):
+        """Load the main LLM model"""
+        try:
+            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
+            )
+            
+            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("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)}")
+            raise
 
-@app.on_event("startup")
-async def startup_event():
-    global agent
-    try:
-        agent = init_agent()
-    except Exception as e:
-        raise RuntimeError(f"Failed to initialize agent: {str(e)}")
+    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)
+            logger.info("RAG model loaded successfully")
+        except Exception as e:
+            logger.error(f"Failed to load RAG model: {str(e)}")
+            raise
 
-def init_agent():
-    """Initialize and return the TxAgent instance"""
-    tool_path = os.path.join(tool_cache_dir, "new_tool.json")
-    if not os.path.exists(tool_path):
-        shutil.copy(os.path.abspath("data/new_tool.json"), tool_path)
-    
-    agent = 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},
-        enable_finish=True,
-        enable_rag=False,
-        force_finish=True,
-        enable_checker=True,
-        step_rag_num=4,
-        seed=100
-    )
-    agent.init_model()
-    return agent
+    def load_tooluniverse(self):
+        """Initialize the ToolUniverse"""
+        try:
+            logger.info("Loading ToolUniverse")
+            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)}")
+            raise
 
-@app.post("/chat")
-async def chat_endpoint(request: ChatRequest):
-    """Handle chat conversations"""
-    try:
-        response = agent.chat(
-            message=request.message,
-            history=request.history,
-            temperature=request.temperature,
-            max_new_tokens=request.max_new_tokens
-        )
-        return JSONResponse({
-            "status": "success",
-            "response": response,
-            "timestamp": datetime.now().isoformat()
-        })
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    def chat(self, message: str, history: Optional[List[Dict]] = None, 
+             temperature: float = 0.7, max_new_tokens: int = 512) -> str:
+        """Handle chat conversations"""
+        try:
+            conversation = []
+            
+            # Initialize with system prompt
+            conversation.append({"role": "system", "content": self.chat_prompt})
+            
+            # Add history if provided
+            if history:
+                for msg in history:
+                    conversation.append({"role": msg["role"], "content": msg["content"]})
+            
+            # Add current message
+            conversation.append({"role": "user", "content": message})
+            
+            # Generate response
+            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
+            )
+            
+            # Decode and clean up response
+            response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+            return response.strip()
+            
+        except Exception as e:
+            logger.error(f"Chat failed: {str(e)}")
+            raise RuntimeError(f"Chat failed: {str(e)}")
 
-@app.post("/multistep")
-async def multistep_endpoint(request: MultistepRequest):
-    """Run multi-step reasoning"""
-    try:
-        response = agent.run_multistep_agent(
-            message=request.message,
-            temperature=request.temperature,
-            max_new_tokens=request.max_new_tokens,
-            max_round=request.max_round
-        )
-        return JSONResponse({
-            "status": "success",
-            "response": response,
-            "timestamp": datetime.now().isoformat()
-        })
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=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 _ in range(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:
+                    return response.split("[FinalAnswer]")[-1].strip()
+                
+                # Add to conversation
+                conversation.append({"role": "assistant", "content": response})
+                
+            # If max rounds reached
+            if self.force_finish:
+                return self._force_final_answer(conversation, temperature, max_new_tokens)
+            
+            return "Reasoning rounds exceeded limit without reaching a final answer."
+            
+        except Exception as e:
+            logger.error(f"Multi-step agent failed: {str(e)}")
+            raise RuntimeError(f"Multi-step agent failed: {str(e)}")
 
-@app.post("/analyze")
-async def analyze_document(file: UploadFile = File(...)):
-    """Analyze a medical document"""
-    try:
-        # Save the uploaded file temporarily
-        temp_path = os.path.join(file_cache_dir, file.filename)
-        with open(temp_path, "wb") as f:
-            f.write(await file.read())
-        
-        # Process the document
-        text = agent.extract_text_from_file(temp_path)
-        analysis = agent.analyze_text(text)
-        
-        # Generate report
-        report_path = os.path.join(report_dir, f"{file.filename}.json")
-        with open(report_path, "w") as f:
-            json.dump({
-                "filename": file.filename,
-                "analysis": analysis,
-                "timestamp": datetime.now().isoformat()
-            }, f)
-        
-        # Clean up
-        os.remove(temp_path)
-        
-        return JSONResponse({
-            "status": "success",
-            "analysis": analysis,
-            "report_path": report_path,
-            "timestamp": datetime.now().isoformat()
-        })
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=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:
+            # 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)}")
+            return "Failed to generate final answer."
 
-@app.get("/status")
-async def service_status():
-    """Check service status"""
-    return {
-        "status": "running",
-        "version": "1.0.0",
-        "model": agent.model_name if agent else "not loaded",
-        "device": str(agent.device) if agent else "unknown"
-    }
+    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")
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    def __del__(self):
+        """Destructor to ensure proper cleanup"""
+        self.cleanup()