Update app.py
Browse files
app.py
CHANGED
@@ -87,8 +87,12 @@ async def startup_event():
|
|
87 |
global agent
|
88 |
try:
|
89 |
logger.info("Initializing TxAgent...")
|
|
|
|
|
|
|
|
|
90 |
agent = TxAgent(
|
91 |
-
model_name="mims-harvard/TxAgent-T1-Llama
|
92 |
rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
|
93 |
tool_files_dict={},
|
94 |
enable_finish=True,
|
@@ -96,17 +100,23 @@ async def startup_event():
|
|
96 |
force_finish=True,
|
97 |
enable_checker=True,
|
98 |
step_rag_num=4,
|
99 |
-
seed=100
|
|
|
100 |
)
|
101 |
agent.init_model()
|
102 |
logger.info("TxAgent initialized successfully")
|
103 |
except Exception as e:
|
104 |
logger.error(f"Failed to initialize agent: {str(e)}")
|
105 |
-
|
|
|
106 |
|
107 |
@app.post("/chat")
|
108 |
async def chat_endpoint(request: ChatRequest):
|
109 |
"""Handle chat conversations with formatting options"""
|
|
|
|
|
|
|
|
|
110 |
try:
|
111 |
logger.info(f"Chat request received (format: {request.format})")
|
112 |
raw_response = agent.chat(
|
@@ -136,32 +146,28 @@ async def chat_endpoint(request: ChatRequest):
|
|
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 |
-
|
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),
|
@@ -169,7 +175,6 @@ async def upload_file(file: UploadFile = File(...)):
|
|
169 |
"html": markdown.markdown(raw_response)
|
170 |
}
|
171 |
response_content = formatted_response["clean"]
|
172 |
-
|
173 |
return JSONResponse({
|
174 |
"status": "success",
|
175 |
"format": "clean",
|
@@ -187,7 +192,7 @@ async def upload_file(file: UploadFile = File(...)):
|
|
187 |
async def service_status():
|
188 |
"""Check service status"""
|
189 |
return {
|
190 |
-
"status": "running",
|
191 |
"version": "2.0.0",
|
192 |
"model": agent.model_name if agent else "not loaded",
|
193 |
"formats_available": ["raw", "clean", "structured", "html"],
|
|
|
87 |
global agent
|
88 |
try:
|
89 |
logger.info("Initializing TxAgent...")
|
90 |
+
# Get Hugging Face token from environment variable
|
91 |
+
hf_token = os.getenv("HUGGINGFACE_TOKEN")
|
92 |
+
model_kwargs = {"token": hf_token} if hf_token else {}
|
93 |
+
|
94 |
agent = TxAgent(
|
95 |
+
model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B", # Corrected model name
|
96 |
rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
|
97 |
tool_files_dict={},
|
98 |
enable_finish=True,
|
|
|
100 |
force_finish=True,
|
101 |
enable_checker=True,
|
102 |
step_rag_num=4,
|
103 |
+
seed=100,
|
104 |
+
**model_kwargs
|
105 |
)
|
106 |
agent.init_model()
|
107 |
logger.info("TxAgent initialized successfully")
|
108 |
except Exception as e:
|
109 |
logger.error(f"Failed to initialize agent: {str(e)}")
|
110 |
+
# Allow app to start, but endpoints will return errors if agent is None
|
111 |
+
agent = None
|
112 |
|
113 |
@app.post("/chat")
|
114 |
async def chat_endpoint(request: ChatRequest):
|
115 |
"""Handle chat conversations with formatting options"""
|
116 |
+
if agent is None:
|
117 |
+
logger.error("TxAgent not initialized")
|
118 |
+
raise HTTPException(status_code=503, detail="TxAgent service unavailable")
|
119 |
+
|
120 |
try:
|
121 |
logger.info(f"Chat request received (format: {request.format})")
|
122 |
raw_response = agent.chat(
|
|
|
146 |
@app.post("/upload")
|
147 |
async def upload_file(file: UploadFile = File(...)):
|
148 |
"""Handle file uploads and process with TxAgent"""
|
149 |
+
if agent is None:
|
150 |
+
logger.error("TxAgent not initialized")
|
151 |
+
raise HTTPException(status_code=503, detail="TxAgent service unavailable")
|
152 |
+
|
153 |
try:
|
154 |
logger.info(f"File upload received: {file.filename}")
|
|
|
|
|
155 |
content = ""
|
156 |
if file.filename.endswith('.pdf'):
|
157 |
pdf_reader = PyPDF2.PdfReader(file.file)
|
158 |
for page in pdf_reader.pages:
|
159 |
content += page.extract_text() or ""
|
160 |
else:
|
|
|
161 |
content = await file.read()
|
162 |
content = content.decode('utf-8', errors='ignore')
|
163 |
|
164 |
+
message = f"Analyze the following medical document content:\n\n{content[:10000]}"
|
|
|
|
|
|
|
165 |
raw_response = agent.chat(
|
166 |
message=message,
|
167 |
history=[],
|
168 |
temperature=0.7,
|
169 |
max_new_tokens=512
|
170 |
)
|
|
|
|
|
171 |
formatted_response = {
|
172 |
"raw": raw_response,
|
173 |
"clean": clean_text_response(raw_response),
|
|
|
175 |
"html": markdown.markdown(raw_response)
|
176 |
}
|
177 |
response_content = formatted_response["clean"]
|
|
|
178 |
return JSONResponse({
|
179 |
"status": "success",
|
180 |
"format": "clean",
|
|
|
192 |
async def service_status():
|
193 |
"""Check service status"""
|
194 |
return {
|
195 |
+
"status": "running" if agent else "failed",
|
196 |
"version": "2.0.0",
|
197 |
"model": agent.model_name if agent else "not loaded",
|
198 |
"formats_available": ["raw", "clean", "structured", "html"],
|