Ali2206 commited on
Commit
064da3f
·
verified ·
1 Parent(s): ddc3dc1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +53 -19
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-3.1-8B",
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"""