Update app.py

app.py

@@ -1,24 +1,18 @@
 import os
 import sys
 import json
-import shutil
 import logging
 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 torch
 from datetime import datetime
 from pydantic import BaseModel
 
-# 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_api.log')
-    ]
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger("TxAgentAPI")
 
@@ -34,36 +28,6 @@ except ImportError as e:
     logger.error(f"Failed to import TxAgent: {str(e)}")
     raise
 
-# Configuration
-persistent_dir = "/data/hf_cache"
-os.makedirs(persistent_dir, exist_ok=True)
-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
-for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir]:
-    os.makedirs(directory, exist_ok=True)
-    logger.info(f"Created directory: {directory}")
-
-# Set environment variables
-os.environ["HF_HOME"] = model_cache_dir
-os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
-
-# Request models
-class ChatRequest(BaseModel):
-    message: str
-    temperature: float = 0.7
-    max_new_tokens: int = 512
-    history: Optional[List[Dict]] = None
-
-class MultistepRequest(BaseModel):
-    message: str
-    temperature: float = 0.7
-    max_new_tokens: int = 512
-    max_round: int = 5
-
 # Initialize FastAPI app
 app = FastAPI(
     title="TxAgent API",
@@ -88,29 +52,10 @@ async def startup_event():
     global agent
     try:
         logger.info("Initializing TxAgent...")
-        agent = init_agent()
-        logger.info("TxAgent initialized successfully")
-    except Exception as e:
-        logger.error(f"Failed to initialize agent: {str(e)}", exc_info=True)
-        raise RuntimeError(f"Failed to initialize agent: {str(e)}")
-
-def init_agent():
-    """Initialize and return the TxAgent instance"""
-    try:
-        tool_path = os.path.join(tool_cache_dir, "new_tool.json")
-        if not os.path.exists(tool_path):
-            logger.info(f"Copying tool file to {tool_path}")
-            default_tool = os.path.abspath("data/new_tool.json")
-            if os.path.exists(default_tool):
-                shutil.copy(default_tool, tool_path)
-            else:
-                raise FileNotFoundError(f"Default tool file not found at {default_tool}")
-        
-        logger.info("Creating TxAgent instance")
         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},
+            tool_files_dict={},  # No tool files in this example
             enable_finish=True,
             enable_rag=False,
             force_finish=True,
@@ -119,10 +64,16 @@ def init_agent():
             seed=100
         )
         agent.init_model()
-        return agent
+        logger.info("TxAgent initialized successfully")
     except Exception as e:
-        logger.error(f"Error in init_agent: {str(e)}", exc_info=True)
-        raise
+        logger.error(f"Failed to initialize agent: {str(e)}")
+        raise RuntimeError(f"Failed to initialize agent: {str(e)}")
+
+class ChatRequest(BaseModel):
+    message: str
+    temperature: float = 0.7
+    max_new_tokens: int = 512
+    history: Optional[List[Dict]] = None
 
 @app.post("/chat")
 async def chat_endpoint(request: ChatRequest):
@@ -135,97 +86,21 @@ async def chat_endpoint(request: ChatRequest):
             temperature=request.temperature,
             max_new_tokens=request.max_new_tokens
         )
-        logger.info("Chat response generated successfully")
-        return JSONResponse({
-            "status": "success",
-            "response": response,
-            "timestamp": datetime.now().isoformat()
-        })
-    except Exception as e:
-        logger.error(f"Chat error: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/multistep")
-async def multistep_endpoint(request: MultistepRequest):
-    """Run multi-step reasoning"""
-    try:
-        logger.info(f"Multistep request received: {request.message[:50]}...")
-        response = agent.run_multistep_agent(
-            message=request.message,
-            temperature=request.temperature,
-            max_new_tokens=request.max_new_tokens,
-            max_round=request.max_round
-        )
-        logger.info("Multistep reasoning completed successfully")
         return JSONResponse({
             "status": "success",
             "response": response,
             "timestamp": datetime.now().isoformat()
         })
     except Exception as e:
-        logger.error(f"Multistep error: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/analyze")
-async def analyze_document(file: UploadFile = File(...)):
-    """Analyze a medical document"""
-    try:
-        logger.info(f"Document analysis request received for: {file.filename}")
-        
-        # 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())
-        logger.info(f"File saved temporarily at {temp_path}")
-        
-        # Process the document
-        text = agent.extract_text_from_file(temp_path)
-        analysis = agent.analyze_text(text)
-        logger.info("Document analysis completed successfully")
-        
-        # Generate report
-        report_filename = f"{os.path.splitext(file.filename)[0]}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
-        report_path = os.path.join(report_dir, report_filename)
-        with open(report_path, "w") as f:
-            json.dump({
-                "filename": file.filename,
-                "analysis": analysis,
-                "timestamp": datetime.now().isoformat()
-            }, f, indent=2)
-        logger.info(f"Report generated at {report_path}")
-        
-        # Clean up
-        os.remove(temp_path)
-        logger.info(f"Temporary file {temp_path} removed")
-        
-        return JSONResponse({
-            "status": "success",
-            "analysis": analysis,
-            "report_path": report_path,
-            "timestamp": datetime.now().isoformat()
-        })
-    except Exception as e:
-        logger.error(f"Document analysis error: {str(e)}", exc_info=True)
+        logger.error(f"Chat error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/status")
 async def service_status():
     """Check service status"""
-    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",
         "timestamp": datetime.now().isoformat()
-    }
-    logger.info(f"Status check: {status}")
-    return status
-
-if __name__ == "__main__":
-    try:
-        logger.info("Starting TxAgent API server")
-        import uvicorn
-        uvicorn.run(app, host="0.0.0.0", port=8000, log_config=None)
-    except Exception as e:
-        logger.error(f"Failed to start server: {str(e)}", exc_info=True)
-        raise
+    }