Spaces:
Sleeping
Sleeping
Delete api.py
Browse files
api.py
DELETED
@@ -1,346 +0,0 @@
|
|
1 |
-
from fastapi import FastAPI, HTTPException, Request
|
2 |
-
from fastapi.middleware.cors import CORSMiddleware
|
3 |
-
from typing import Dict, List, Optional, Union, Any
|
4 |
-
from pydantic import BaseModel, Field
|
5 |
-
from datetime import datetime
|
6 |
-
import logging
|
7 |
-
import json
|
8 |
-
import os
|
9 |
-
from dotenv import load_dotenv
|
10 |
-
from dify_client_python.dify_client import models
|
11 |
-
from sse_starlette.sse import EventSourceResponse
|
12 |
-
import httpx
|
13 |
-
from json_parser import SSEParser
|
14 |
-
from logger_config import setup_logger
|
15 |
-
from fastapi.responses import StreamingResponse
|
16 |
-
from fastapi.responses import JSONResponse
|
17 |
-
from response_formatter import ResponseFormatter
|
18 |
-
import traceback
|
19 |
-
|
20 |
-
# Load environment variables
|
21 |
-
load_dotenv()
|
22 |
-
|
23 |
-
# Configure logging
|
24 |
-
logging.basicConfig(
|
25 |
-
level=logging.INFO,
|
26 |
-
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
27 |
-
)
|
28 |
-
logger = logging.getLogger(__name__)
|
29 |
-
|
30 |
-
class AgentOutput(BaseModel):
|
31 |
-
"""Structured output from agent processing"""
|
32 |
-
thought_content: str
|
33 |
-
observation: Optional[str]
|
34 |
-
tool_outputs: List[Dict]
|
35 |
-
citations: List[Dict]
|
36 |
-
metadata: Dict
|
37 |
-
raw_response: str
|
38 |
-
|
39 |
-
class AgentRequest(BaseModel):
|
40 |
-
"""Enhanced request model with additional parameters"""
|
41 |
-
query: str
|
42 |
-
conversation_id: Optional[str] = None
|
43 |
-
stream: bool = True
|
44 |
-
inputs: Dict = {}
|
45 |
-
files: List = []
|
46 |
-
user: str = "default_user"
|
47 |
-
response_mode: str = "streaming"
|
48 |
-
|
49 |
-
class AgentProcessor:
|
50 |
-
def __init__(self, api_key: str):
|
51 |
-
self.api_key = api_key
|
52 |
-
self.api_base = "https://rag-engine.go-yamamoto.com/v1"
|
53 |
-
self.formatter = ResponseFormatter()
|
54 |
-
self.client = httpx.AsyncClient(timeout=60.0)
|
55 |
-
self.logger = setup_logger("agent_processor")
|
56 |
-
|
57 |
-
async def log_request_details(
|
58 |
-
self,
|
59 |
-
request: AgentRequest,
|
60 |
-
start_time: datetime
|
61 |
-
) -> None:
|
62 |
-
"""Log detailed request information"""
|
63 |
-
self.logger.debug(
|
64 |
-
"Request details: \n"
|
65 |
-
f"Query: {request.query}\n"
|
66 |
-
f"User: {request.user}\n"
|
67 |
-
f"Conversation ID: {request.conversation_id}\n"
|
68 |
-
f"Stream mode: {request.stream}\n"
|
69 |
-
f"Start time: {start_time}\n"
|
70 |
-
f"Inputs: {request.inputs}\n"
|
71 |
-
f"Files: {len(request.files)} files attached"
|
72 |
-
)
|
73 |
-
|
74 |
-
async def log_error(
|
75 |
-
self,
|
76 |
-
error: Exception,
|
77 |
-
context: Optional[Dict] = None
|
78 |
-
) -> None:
|
79 |
-
"""Log detailed error information"""
|
80 |
-
error_msg = (
|
81 |
-
f"Error type: {type(error).__name__}\n"
|
82 |
-
f"Error message: {str(error)}\n"
|
83 |
-
f"Stack trace:\n{traceback.format_exc()}\n"
|
84 |
-
)
|
85 |
-
if context:
|
86 |
-
error_msg += f"Context:\n{json.dumps(context, indent=2)}"
|
87 |
-
|
88 |
-
self.logger.error(error_msg)
|
89 |
-
|
90 |
-
async def cleanup(self):
|
91 |
-
"""Cleanup method to properly close client"""
|
92 |
-
await self.client.aclose()
|
93 |
-
|
94 |
-
async def process_stream(self, request: AgentRequest):
|
95 |
-
start_time = datetime.now()
|
96 |
-
await self.log_request_details(request, start_time)
|
97 |
-
|
98 |
-
headers = {
|
99 |
-
"Authorization": f"Bearer {self.api_key}",
|
100 |
-
"Content-Type": "application/json",
|
101 |
-
"Accept": "text/event-stream"
|
102 |
-
}
|
103 |
-
|
104 |
-
chat_request = {
|
105 |
-
"query": request.query,
|
106 |
-
"inputs": request.inputs,
|
107 |
-
"response_mode": "streaming" if request.stream else "blocking",
|
108 |
-
"user": request.user,
|
109 |
-
"conversation_id": request.conversation_id,
|
110 |
-
"files": request.files
|
111 |
-
}
|
112 |
-
|
113 |
-
async def event_generator():
|
114 |
-
parser = SSEParser()
|
115 |
-
citations = []
|
116 |
-
metadata = {}
|
117 |
-
|
118 |
-
try:
|
119 |
-
async with self.client.stream(
|
120 |
-
"POST",
|
121 |
-
f"{self.api_base}/chat-messages",
|
122 |
-
headers=headers,
|
123 |
-
json=chat_request
|
124 |
-
) as response:
|
125 |
-
self.logger.debug(
|
126 |
-
f"Stream connection established\n"
|
127 |
-
f"Status: {response.status_code}\n"
|
128 |
-
f"Headers: {dict(response.headers)}"
|
129 |
-
)
|
130 |
-
|
131 |
-
buffer = ""
|
132 |
-
async for line in response.aiter_lines():
|
133 |
-
if not line.strip():
|
134 |
-
continue
|
135 |
-
|
136 |
-
self.logger.debug(f"Raw SSE line: {line}")
|
137 |
-
|
138 |
-
if "data:" in line:
|
139 |
-
try:
|
140 |
-
data = line.split("data:", 1)[1].strip()
|
141 |
-
parsed = json.loads(data)
|
142 |
-
|
143 |
-
if parsed.get("event") == "message_end":
|
144 |
-
citations = parsed.get("retriever_resources", [])
|
145 |
-
metadata = parsed.get("metadata", {})
|
146 |
-
self.logger.debug(
|
147 |
-
f"Message end event:\n"
|
148 |
-
f"Citations: {citations}\n"
|
149 |
-
f"Metadata: {metadata}"
|
150 |
-
)
|
151 |
-
|
152 |
-
formatted = self.format_terminal_output(
|
153 |
-
parsed,
|
154 |
-
citations=citations,
|
155 |
-
metadata=metadata
|
156 |
-
)
|
157 |
-
if formatted:
|
158 |
-
self.logger.info(formatted)
|
159 |
-
except Exception as e:
|
160 |
-
await self.log_error(
|
161 |
-
e,
|
162 |
-
{"line": line, "event": "parse_data"}
|
163 |
-
)
|
164 |
-
|
165 |
-
buffer += line + "\n"
|
166 |
-
|
167 |
-
if line.startswith("data:") or buffer.strip().endswith("}"):
|
168 |
-
try:
|
169 |
-
processed_response = parser.parse_sse_event(buffer)
|
170 |
-
if processed_response and isinstance(processed_response, dict):
|
171 |
-
cleaned_response = self.clean_response(processed_response)
|
172 |
-
if cleaned_response:
|
173 |
-
xml_content = cleaned_response.get("content", "")
|
174 |
-
yield f"data: {xml_content}\n\n"
|
175 |
-
except Exception as parse_error:
|
176 |
-
await self.log_error(
|
177 |
-
parse_error,
|
178 |
-
{"buffer": buffer, "event": "process_buffer"}
|
179 |
-
)
|
180 |
-
error_xml = (
|
181 |
-
f"<agent_response>"
|
182 |
-
f"<error>{str(parse_error)}</error>"
|
183 |
-
f"</agent_response>"
|
184 |
-
)
|
185 |
-
yield f"data: {error_xml}\n\n"
|
186 |
-
finally:
|
187 |
-
buffer = ""
|
188 |
-
|
189 |
-
except httpx.ConnectError as e:
|
190 |
-
await self.log_error(e, {"event": "connection_error"})
|
191 |
-
error_xml = (
|
192 |
-
f"<agent_response>"
|
193 |
-
f"<error>Connection error: {str(e)}</error>"
|
194 |
-
f"</agent_response>"
|
195 |
-
)
|
196 |
-
yield f"data: {error_xml}\n\n"
|
197 |
-
except Exception as e:
|
198 |
-
await self.log_error(e, {"event": "stream_error"})
|
199 |
-
error_xml = (
|
200 |
-
f"<agent_response>"
|
201 |
-
f"<error>Streaming error: {str(e)}</error>"
|
202 |
-
f"</agent_response>"
|
203 |
-
)
|
204 |
-
yield f"data: {error_xml}\n\n"
|
205 |
-
finally:
|
206 |
-
end_time = datetime.now()
|
207 |
-
duration = (end_time - start_time).total_seconds()
|
208 |
-
self.logger.info(f"Request completed in {duration:.2f} seconds")
|
209 |
-
|
210 |
-
return StreamingResponse(
|
211 |
-
event_generator(),
|
212 |
-
media_type="text/event-stream",
|
213 |
-
headers={
|
214 |
-
"Cache-Control": "no-cache",
|
215 |
-
"Connection": "keep-alive",
|
216 |
-
"X-Accel-Buffering": "no",
|
217 |
-
"Access-Control-Allow-Origin": "*"
|
218 |
-
}
|
219 |
-
)
|
220 |
-
|
221 |
-
def format_terminal_output(
|
222 |
-
self,
|
223 |
-
response: Dict,
|
224 |
-
citations: List[Dict] = None,
|
225 |
-
metadata: Dict = None
|
226 |
-
) -> Optional[str]:
|
227 |
-
"""Format response for terminal output"""
|
228 |
-
event_type = response.get("event")
|
229 |
-
|
230 |
-
if event_type == "agent_thought":
|
231 |
-
thought = response.get("thought", "")
|
232 |
-
observation = response.get("observation", "")
|
233 |
-
terminal_output, _ = self.formatter.format_thought(
|
234 |
-
thought,
|
235 |
-
observation,
|
236 |
-
citations=citations,
|
237 |
-
metadata=metadata
|
238 |
-
)
|
239 |
-
return terminal_output
|
240 |
-
|
241 |
-
elif event_type == "agent_message":
|
242 |
-
message = response.get("answer", "")
|
243 |
-
terminal_output, _ = self.formatter.format_message(message)
|
244 |
-
return terminal_output
|
245 |
-
|
246 |
-
elif event_type == "error":
|
247 |
-
error = response.get("error", "Unknown error")
|
248 |
-
terminal_output, _ = self.formatter.format_error(error)
|
249 |
-
return terminal_output
|
250 |
-
|
251 |
-
return None
|
252 |
-
|
253 |
-
def clean_response(self, response: Dict) -> Optional[Dict]:
|
254 |
-
"""Clean and transform the response for frontend consumption"""
|
255 |
-
try:
|
256 |
-
event_type = response.get("event")
|
257 |
-
if not event_type:
|
258 |
-
return None
|
259 |
-
|
260 |
-
# Handle different event types
|
261 |
-
if event_type == "agent_thought":
|
262 |
-
thought = response.get("thought", "")
|
263 |
-
observation = response.get("observation", "")
|
264 |
-
_, xml_output = self.formatter.format_thought(thought, observation)
|
265 |
-
return {
|
266 |
-
"type": "thought",
|
267 |
-
"content": xml_output
|
268 |
-
}
|
269 |
-
|
270 |
-
elif event_type == "agent_message":
|
271 |
-
message = response.get("answer", "")
|
272 |
-
_, xml_output = self.formatter.format_message(message)
|
273 |
-
return {
|
274 |
-
"type": "message",
|
275 |
-
"content": xml_output
|
276 |
-
}
|
277 |
-
|
278 |
-
elif event_type == "error":
|
279 |
-
error = response.get("error", "Unknown error")
|
280 |
-
_, xml_output = self.formatter.format_error(error)
|
281 |
-
return {
|
282 |
-
"type": "error",
|
283 |
-
"content": xml_output
|
284 |
-
}
|
285 |
-
|
286 |
-
return None
|
287 |
-
except Exception as e:
|
288 |
-
logger.error(f"Error cleaning response: {str(e)}")
|
289 |
-
return None
|
290 |
-
|
291 |
-
# Initialize FastAPI app
|
292 |
-
app = FastAPI(docs_url="/", redoc_url=None)
|
293 |
-
agent_processor = None
|
294 |
-
|
295 |
-
# Add CORS middleware
|
296 |
-
app.add_middleware(
|
297 |
-
CORSMiddleware,
|
298 |
-
allow_origins=["*"],
|
299 |
-
allow_credentials=True,
|
300 |
-
allow_methods=["*"],
|
301 |
-
allow_headers=["*"],
|
302 |
-
)
|
303 |
-
|
304 |
-
@app.on_event("startup")
|
305 |
-
async def startup_event():
|
306 |
-
global agent_processor
|
307 |
-
api_key = os.getenv("DIFY_API_KEY", "app-kVHTrZzEmFXEBfyXOi4rro7M")
|
308 |
-
agent_processor = AgentProcessor(api_key=api_key)
|
309 |
-
|
310 |
-
@app.on_event("shutdown")
|
311 |
-
async def shutdown_event():
|
312 |
-
global agent_processor
|
313 |
-
if agent_processor:
|
314 |
-
await agent_processor.cleanup()
|
315 |
-
|
316 |
-
@app.post("/v1/agent")
|
317 |
-
async def process_agent_request(request: AgentRequest):
|
318 |
-
try:
|
319 |
-
logger.info(f"Processing agent request: {request.query}")
|
320 |
-
return await agent_processor.process_stream(request)
|
321 |
-
|
322 |
-
except Exception as e:
|
323 |
-
logger.error(f"Error in agent request processing: {e}", exc_info=True)
|
324 |
-
raise HTTPException(status_code=500, detail=str(e))
|
325 |
-
|
326 |
-
@app.middleware("http")
|
327 |
-
async def error_handling_middleware(request: Request, call_next):
|
328 |
-
try:
|
329 |
-
response = await call_next(request)
|
330 |
-
return response
|
331 |
-
except Exception as e:
|
332 |
-
logger.error(f"Unhandled error: {str(e)}", exc_info=True)
|
333 |
-
return JSONResponse(
|
334 |
-
status_code=500,
|
335 |
-
content={"error": "Internal server error occurred"}
|
336 |
-
)
|
337 |
-
|
338 |
-
# Add host and port parameters to the launch
|
339 |
-
if __name__ == "__main__":
|
340 |
-
import uvicorn
|
341 |
-
uvicorn.run(
|
342 |
-
"api:app",
|
343 |
-
host="0.0.0.0",
|
344 |
-
port=7860,
|
345 |
-
reload=True
|
346 |
-
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|