Update app.py
Browse files
app.py
CHANGED
@@ -3,13 +3,14 @@ import sys
|
|
3 |
import json
|
4 |
import re
|
5 |
import logging
|
6 |
-
from fastapi import FastAPI, HTTPException
|
7 |
from fastapi.responses import JSONResponse
|
8 |
from fastapi.middleware.cors import CORSMiddleware
|
9 |
from typing import List, Dict, Optional
|
10 |
from datetime import datetime
|
11 |
from pydantic import BaseModel
|
12 |
import markdown
|
|
|
13 |
|
14 |
# Configure logging
|
15 |
logging.basicConfig(
|
@@ -57,26 +58,18 @@ class ChatRequest(BaseModel):
|
|
57 |
# Response cleaning functions
|
58 |
def clean_text_response(text: str) -> str:
|
59 |
"""Basic text cleaning"""
|
60 |
-
# Remove excessive newlines and whitespace
|
61 |
text = re.sub(r'\n\s*\n', '\n\n', text)
|
62 |
-
# Remove redundant spaces
|
63 |
text = re.sub(r'[ ]+', ' ', text)
|
64 |
-
# Clean markdown formatting
|
65 |
text = text.replace("**", "").replace("__", "")
|
66 |
return text.strip()
|
67 |
|
68 |
def structure_medical_response(text: str) -> Dict:
|
69 |
"""Structure medical content into categories"""
|
70 |
result = {"overview": "", "symptoms": [], "types": {}, "notes": ""}
|
71 |
-
|
72 |
-
# Extract overview
|
73 |
overview_end = text.find("Type 1 Diabetes:")
|
74 |
result["overview"] = clean_text_response(text[:overview_end])
|
75 |
-
|
76 |
-
# Extract by type
|
77 |
type_sections = re.split(r'(Type \d Diabetes:)', text[overview_end:])
|
78 |
current_type = None
|
79 |
-
|
80 |
for section in type_sections:
|
81 |
if section.startswith("Type "):
|
82 |
current_type = section.replace(":", "").strip().lower()
|
@@ -84,7 +77,6 @@ def structure_medical_response(text: str) -> Dict:
|
|
84 |
elif current_type:
|
85 |
points = [p.strip() for p in section.split('\n') if p.strip()]
|
86 |
result["types"][current_type] = points
|
87 |
-
|
88 |
return result
|
89 |
|
90 |
# Initialize agent at startup
|
@@ -96,7 +88,7 @@ async def startup_event():
|
|
96 |
try:
|
97 |
logger.info("Initializing TxAgent...")
|
98 |
agent = TxAgent(
|
99 |
-
model_name="mims-harvard/TxAgent-T1-Llama
|
100 |
rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
|
101 |
tool_files_dict={},
|
102 |
enable_finish=True,
|
@@ -117,26 +109,19 @@ async def chat_endpoint(request: ChatRequest):
|
|
117 |
"""Handle chat conversations with formatting options"""
|
118 |
try:
|
119 |
logger.info(f"Chat request received (format: {request.format})")
|
120 |
-
|
121 |
-
# Get raw response from agent
|
122 |
raw_response = agent.chat(
|
123 |
message=request.message,
|
124 |
history=request.history,
|
125 |
temperature=request.temperature,
|
126 |
max_new_tokens=request.max_new_tokens
|
127 |
)
|
128 |
-
|
129 |
-
# Format response based on request
|
130 |
formatted_response = {
|
131 |
"raw": raw_response,
|
132 |
"clean": clean_text_response(raw_response),
|
133 |
"structured": structure_medical_response(raw_response),
|
134 |
"html": markdown.markdown(raw_response)
|
135 |
}
|
136 |
-
|
137 |
-
# Select the requested format
|
138 |
response_content = formatted_response.get(request.format, formatted_response["clean"])
|
139 |
-
|
140 |
return JSONResponse({
|
141 |
"status": "success",
|
142 |
"format": request.format,
|
@@ -144,11 +129,60 @@ async def chat_endpoint(request: ChatRequest):
|
|
144 |
"timestamp": datetime.now().isoformat(),
|
145 |
"available_formats": list(formatted_response.keys())
|
146 |
})
|
147 |
-
|
148 |
except Exception as e:
|
149 |
logger.error(f"Chat error: {str(e)}")
|
150 |
raise HTTPException(status_code=500, detail=str(e))
|
151 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
@app.get("/status")
|
153 |
async def service_status():
|
154 |
"""Check service status"""
|
|
|
3 |
import json
|
4 |
import re
|
5 |
import logging
|
6 |
+
from fastapi import FastAPI, HTTPException, UploadFile, File
|
7 |
from fastapi.responses import JSONResponse
|
8 |
from fastapi.middleware.cors import CORSMiddleware
|
9 |
from typing import List, Dict, Optional
|
10 |
from datetime import datetime
|
11 |
from pydantic import BaseModel
|
12 |
import markdown
|
13 |
+
import PyPDF2
|
14 |
|
15 |
# Configure logging
|
16 |
logging.basicConfig(
|
|
|
58 |
# Response cleaning functions
|
59 |
def clean_text_response(text: str) -> str:
|
60 |
"""Basic text cleaning"""
|
|
|
61 |
text = re.sub(r'\n\s*\n', '\n\n', text)
|
|
|
62 |
text = re.sub(r'[ ]+', ' ', text)
|
|
|
63 |
text = text.replace("**", "").replace("__", "")
|
64 |
return text.strip()
|
65 |
|
66 |
def structure_medical_response(text: str) -> Dict:
|
67 |
"""Structure medical content into categories"""
|
68 |
result = {"overview": "", "symptoms": [], "types": {}, "notes": ""}
|
|
|
|
|
69 |
overview_end = text.find("Type 1 Diabetes:")
|
70 |
result["overview"] = clean_text_response(text[:overview_end])
|
|
|
|
|
71 |
type_sections = re.split(r'(Type \d Diabetes:)', text[overview_end:])
|
72 |
current_type = None
|
|
|
73 |
for section in type_sections:
|
74 |
if section.startswith("Type "):
|
75 |
current_type = section.replace(":", "").strip().lower()
|
|
|
77 |
elif current_type:
|
78 |
points = [p.strip() for p in section.split('\n') if p.strip()]
|
79 |
result["types"][current_type] = points
|
|
|
80 |
return result
|
81 |
|
82 |
# Initialize agent at startup
|
|
|
88 |
try:
|
89 |
logger.info("Initializing TxAgent...")
|
90 |
agent = TxAgent(
|
91 |
+
model_name="mims-harvard/TxAgent-T1-Llama出租3.1-8B",
|
92 |
rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
|
93 |
tool_files_dict={},
|
94 |
enable_finish=True,
|
|
|
109 |
"""Handle chat conversations with formatting options"""
|
110 |
try:
|
111 |
logger.info(f"Chat request received (format: {request.format})")
|
|
|
|
|
112 |
raw_response = agent.chat(
|
113 |
message=request.message,
|
114 |
history=request.history,
|
115 |
temperature=request.temperature,
|
116 |
max_new_tokens=request.max_new_tokens
|
117 |
)
|
|
|
|
|
118 |
formatted_response = {
|
119 |
"raw": raw_response,
|
120 |
"clean": clean_text_response(raw_response),
|
121 |
"structured": structure_medical_response(raw_response),
|
122 |
"html": markdown.markdown(raw_response)
|
123 |
}
|
|
|
|
|
124 |
response_content = formatted_response.get(request.format, formatted_response["clean"])
|
|
|
125 |
return JSONResponse({
|
126 |
"status": "success",
|
127 |
"format": request.format,
|
|
|
129 |
"timestamp": datetime.now().isoformat(),
|
130 |
"available_formats": list(formatted_response.keys())
|
131 |
})
|
|
|
132 |
except Exception as e:
|
133 |
logger.error(f"Chat error: {str(e)}")
|
134 |
raise HTTPException(status_code=500, detail=str(e))
|
135 |
|
136 |
+
@app.post("/upload")
|
137 |
+
async def upload_file(file: UploadFile = File(...)):
|
138 |
+
"""Handle file uploads and process with TxAgent"""
|
139 |
+
try:
|
140 |
+
logger.info(f"File upload received: {file.filename}")
|
141 |
+
|
142 |
+
# Read file content
|
143 |
+
content = ""
|
144 |
+
if file.filename.endswith('.pdf'):
|
145 |
+
pdf_reader = PyPDF2.PdfReader(file.file)
|
146 |
+
for page in pdf_reader.pages:
|
147 |
+
content += page.extract_text() or ""
|
148 |
+
else:
|
149 |
+
# Assume text-based file
|
150 |
+
content = await file.read()
|
151 |
+
content = content.decode('utf-8', errors='ignore')
|
152 |
+
|
153 |
+
# Create a message to analyze the file content
|
154 |
+
message = f"Analyze the following medical document content:\n\n{content[:10000]}" # Limit to 10k chars for brevity
|
155 |
+
|
156 |
+
# Process with TxAgent
|
157 |
+
raw_response = agent.chat(
|
158 |
+
message=message,
|
159 |
+
history=[],
|
160 |
+
temperature=0.7,
|
161 |
+
max_new_tokens=512
|
162 |
+
)
|
163 |
+
|
164 |
+
# Format response (using clean format for consistency)
|
165 |
+
formatted_response = {
|
166 |
+
"raw": raw_response,
|
167 |
+
"clean": clean_text_response(raw_response),
|
168 |
+
"structured": structure_medical_response(raw_response),
|
169 |
+
"html": markdown.markdown(raw_response)
|
170 |
+
}
|
171 |
+
response_content = formatted_response["clean"]
|
172 |
+
|
173 |
+
return JSONResponse({
|
174 |
+
"status": "success",
|
175 |
+
"format": "clean",
|
176 |
+
"response": response_content,
|
177 |
+
"timestamp": datetime.now().isoformat(),
|
178 |
+
"available_formats": list(formatted_response.keys())
|
179 |
+
})
|
180 |
+
except Exception as e:
|
181 |
+
logger.error(f"File upload error: {str(e)}")
|
182 |
+
raise HTTPException(status_code=500, detail=str(e))
|
183 |
+
finally:
|
184 |
+
file.file.close()
|
185 |
+
|
186 |
@app.get("/status")
|
187 |
async def service_status():
|
188 |
"""Check service status"""
|