Update app.py

app.py

@@ -87,8 +87,12 @@ async def startup_event():
     global agent
     try:
         logger.info("Initializing TxAgent...")
+        # Get Hugging Face token from environment variable
+        hf_token = os.getenv("HUGGINGFACE_TOKEN")
+        model_kwargs = {"token": hf_token} if hf_token else {}
+        
         agent = TxAgent(
-            model_name="mims-harvard/TxAgent-T1-Llama出租3.1-8B",
+            model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",  # Corrected model name
             rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
             tool_files_dict={},
             enable_finish=True,
@@ -96,17 +100,23 @@ async def startup_event():
             force_finish=True,
             enable_checker=True,
             step_rag_num=4,
-            seed=100
+            seed=100,
+            **model_kwargs
         )
         agent.init_model()
         logger.info("TxAgent initialized successfully")
     except Exception as e:
         logger.error(f"Failed to initialize agent: {str(e)}")
-        raise RuntimeError(f"Failed to initialize agent: {str(e)}")
+        # Allow app to start, but endpoints will return errors if agent is None
+        agent = None
 
 @app.post("/chat")
 async def chat_endpoint(request: ChatRequest):
     """Handle chat conversations with formatting options"""
+    if agent is None:
+        logger.error("TxAgent not initialized")
+        raise HTTPException(status_code=503, detail="TxAgent service unavailable")
+    
     try:
         logger.info(f"Chat request received (format: {request.format})")
         raw_response = agent.chat(
@@ -136,32 +146,28 @@ async def chat_endpoint(request: ChatRequest):
 @app.post("/upload")
 async def upload_file(file: UploadFile = File(...)):
     """Handle file uploads and process with TxAgent"""
+    if agent is None:
+        logger.error("TxAgent not initialized")
+        raise HTTPException(status_code=503, detail="TxAgent service unavailable")
+    
     try:
         logger.info(f"File upload received: {file.filename}")
-        
-        # Read file content
         content = ""
         if file.filename.endswith('.pdf'):
             pdf_reader = PyPDF2.PdfReader(file.file)
             for page in pdf_reader.pages:
                 content += page.extract_text() or ""
         else:
-            # Assume text-based file
             content = await file.read()
             content = content.decode('utf-8', errors='ignore')
         
-        # Create a message to analyze the file content
-        message = f"Analyze the following medical document content:\n\n{content[:10000]}"  # Limit to 10k chars for brevity
-        
-        # Process with TxAgent
+        message = f"Analyze the following medical document content:\n\n{content[:10000]}"
         raw_response = agent.chat(
             message=message,
             history=[],
             temperature=0.7,
             max_new_tokens=512
         )
-        
-        # Format response (using clean format for consistency)
         formatted_response = {
             "raw": raw_response,
             "clean": clean_text_response(raw_response),
@@ -169,7 +175,6 @@ async def upload_file(file: UploadFile = File(...)):
             "html": markdown.markdown(raw_response)
         }
         response_content = formatted_response["clean"]
-        
         return JSONResponse({
             "status": "success",
             "format": "clean",
@@ -187,7 +192,7 @@ async def upload_file(file: UploadFile = File(...)):
 async def service_status():
     """Check service status"""
     return {
-        "status": "running",
+        "status": "running" if agent else "failed",
         "version": "2.0.0",
         "model": agent.model_name if agent else "not loaded",
         "formats_available": ["raw", "clean", "structured", "html"],