Update app.py

app.py

@@ -1,15 +1,17 @@
 import os
 import sys
 import json
+import re
 import logging
-from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi import FastAPI, HTTPException
 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
 
-# Configure logging for Hugging Face Spaces
+# Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
@@ -32,7 +34,7 @@ except ImportError as e:
 app = FastAPI(
     title="TxAgent API",
     description="API for TxAgent medical document analysis",
-    version="1.0.0"
+    version="2.0.0"
 )
 
 # CORS configuration
@@ -44,6 +46,47 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Request models
+class ChatRequest(BaseModel):
+    message: str
+    temperature: float = 0.7
+    max_new_tokens: int = 512
+    history: Optional[List[Dict]] = None
+    format: Optional[str] = "clean"  # Options: raw, clean, structured, html
+
+# 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()
+            result["types"][current_type] = []
+        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
 agent = None
 
@@ -55,7 +98,7 @@ async def startup_event():
         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={},  # No tool files in this example
+            tool_files_dict={},
             enable_finish=True,
             enable_rag=False,
             force_finish=True,
@@ -69,28 +112,39 @@ async def startup_event():
         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):
-    """Handle chat conversations"""
+    """Handle chat conversations with formatting options"""
     try:
-        logger.info(f"Chat request received: {request.message[:50]}...")
-        response = agent.chat(
+        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",
-            "response": response,
-            "timestamp": datetime.now().isoformat()
+            "format": request.format,
+            "response": response_content,
+            "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))
@@ -100,7 +154,8 @@ async def service_status():
     """Check service status"""
     return {
         "status": "running",
-        "version": "1.0.0",
+        "version": "2.0.0",
         "model": agent.model_name if agent else "not loaded",
+        "formats_available": ["raw", "clean", "structured", "html"],
         "timestamp": datetime.now().isoformat()
     }