Update app.py

app.py

@@ -3,13 +3,14 @@ import sys
 import json
 import re
 import logging
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, UploadFile, File
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from typing import List, Dict, Optional
 from datetime import datetime
 from pydantic import BaseModel
 import markdown
+import PyPDF2
 
 # Configure logging
 logging.basicConfig(
@@ -57,26 +58,18 @@ class ChatRequest(BaseModel):
 # Response cleaning functions
 def clean_text_response(text: str) -> str:
     """Basic text cleaning"""
-    # Remove excessive newlines and whitespace
     text = re.sub(r'\n\s*\n', '\n\n', text)
-    # Remove redundant spaces
     text = re.sub(r'[ ]+', ' ', text)
-    # Clean markdown formatting
     text = text.replace("**", "").replace("__", "")
     return text.strip()
 
 def structure_medical_response(text: str) -> Dict:
     """Structure medical content into categories"""
     result = {"overview": "", "symptoms": [], "types": {}, "notes": ""}
-    
-    # Extract overview
     overview_end = text.find("Type 1 Diabetes:")
     result["overview"] = clean_text_response(text[:overview_end])
-    
-    # Extract by type
     type_sections = re.split(r'(Type \d Diabetes:)', text[overview_end:])
     current_type = None
-    
     for section in type_sections:
         if section.startswith("Type "):
             current_type = section.replace(":", "").strip().lower()
@@ -84,7 +77,6 @@ def structure_medical_response(text: str) -> Dict:
         elif current_type:
             points = [p.strip() for p in section.split('\n') if p.strip()]
             result["types"][current_type] = points
-    
     return result
 
 # Initialize agent at startup
@@ -96,7 +88,7 @@ async def startup_event():
     try:
         logger.info("Initializing TxAgent...")
         agent = TxAgent(
-            model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+            model_name="mims-harvard/TxAgent-T1-Llama出租3.1-8B",
             rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
             tool_files_dict={},
             enable_finish=True,
@@ -117,26 +109,19 @@ async def chat_endpoint(request: ChatRequest):
     """Handle chat conversations with formatting options"""
     try:
         logger.info(f"Chat request received (format: {request.format})")
-        
-        # Get raw response from agent
         raw_response = agent.chat(
             message=request.message,
             history=request.history,
             temperature=request.temperature,
             max_new_tokens=request.max_new_tokens
         )
-        
-        # Format response based on request
         formatted_response = {
             "raw": raw_response,
             "clean": clean_text_response(raw_response),
             "structured": structure_medical_response(raw_response),
             "html": markdown.markdown(raw_response)
         }
-        
-        # Select the requested format
         response_content = formatted_response.get(request.format, formatted_response["clean"])
-        
         return JSONResponse({
             "status": "success",
             "format": request.format,
@@ -144,11 +129,60 @@ async def chat_endpoint(request: ChatRequest):
             "timestamp": datetime.now().isoformat(),
             "available_formats": list(formatted_response.keys())
         })
-        
     except Exception as e:
         logger.error(f"Chat error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/upload")
+async def upload_file(file: UploadFile = File(...)):
+    """Handle file uploads and process with TxAgent"""
+    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
+        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),
+            "structured": structure_medical_response(raw_response),
+            "html": markdown.markdown(raw_response)
+        }
+        response_content = formatted_response["clean"]
+        
+        return JSONResponse({
+            "status": "success",
+            "format": "clean",
+            "response": response_content,
+            "timestamp": datetime.now().isoformat(),
+            "available_formats": list(formatted_response.keys())
+        })
+    except Exception as e:
+        logger.error(f"File upload error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        file.file.close()
+
 @app.get("/status")
 async def service_status():
     """Check service status"""