Severian commited on
Commit
565b90d
·
verified ·
1 Parent(s): 46149cf

Update api.py

Browse files
Files changed (1) hide show
  1. api.py +65 -134
api.py CHANGED
@@ -10,7 +10,7 @@ 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
@@ -57,7 +57,24 @@ class AgentProcessor:
57
  self.formatter = ResponseFormatter()
58
  self.client = httpx.AsyncClient(timeout=60.0)
59
  self.logger = setup_logger("agent_processor")
 
 
 
 
 
 
 
60
 
 
 
 
 
 
 
 
 
 
 
61
 
62
  async def log_request_details(
63
  self,
@@ -97,153 +114,67 @@ class AgentProcessor:
97
  await self.client.aclose()
98
 
99
  async def process_stream(self, request: AgentRequest):
100
- start_time = datetime.now()
101
- await self.log_request_details(request, start_time)
102
-
103
- headers = {
104
- "Authorization": f"Bearer {self.api_key}",
105
- "Content-Type": "application/json",
106
- "Accept": "text/event-stream"
107
- }
108
 
109
- chat_request = {
110
- "query": request.query,
111
- "inputs": request.inputs,
112
- "response_mode": "streaming" if request.stream else "blocking",
113
- "user": request.user,
114
- "conversation_id": request.conversation_id,
115
- "files": request.files
116
- }
117
-
118
  async def event_generator():
119
- parser = SSEParser()
120
- citations = []
121
- metadata = {}
122
- tool_outputs = []
123
- buffer = ""
124
-
125
  try:
126
  async with self.client.stream(
127
  "POST",
128
  f"{self.api_base}/chat-messages",
129
- headers=headers,
130
- json=chat_request
131
  ) as response:
132
- self.logger.debug(
133
- f"Stream connection established\n"
134
- f"Status: {response.status_code}\n"
135
- f"Headers: {dict(response.headers)}"
136
- )
137
-
138
  async for line in response.aiter_lines():
139
  if not line.strip():
140
  continue
141
 
142
- self.logger.debug(f"Raw SSE line: {line}")
143
- buffer += line + "\n"
 
 
144
 
145
- if "data:" in line:
146
- try:
147
- data = line.split("data:", 1)[1].strip()
148
- parsed = json.loads(data)
149
-
150
- # Enhanced mermaid diagram handling
151
- if parsed.get("observation"):
152
- try:
153
- observation = parsed["observation"]
154
- if isinstance(observation, str):
155
- if "mermaid_diagram" in observation:
156
- try:
157
- # Clean and extract diagram content
158
- cleaned_content = parser.clean_mermaid_content(
159
- observation
160
- )
161
-
162
- # Create tool output without extra wrapping
163
- tool_output = {
164
- "type": "mermaid_diagram",
165
- "content": cleaned_content
166
- }
167
- tool_outputs.append(tool_output)
168
-
169
- # Send clean event
170
- yield (
171
- "event: tool_output\n"
172
- f"data: {json.dumps(tool_output)}\n\n"
173
- )
174
- except Exception as e:
175
- self.logger.error(
176
- f"Failed to process mermaid diagram: {e}"
177
- )
178
- except Exception as e:
179
- self.logger.error(
180
- f"Error processing observation: {e}"
181
- )
182
-
183
- if parsed.get("event") == "message_end":
184
- citations = parsed.get("retriever_resources", [])
185
- metadata = parsed.get("metadata", {})
186
- metadata["tool_outputs"] = tool_outputs
187
- self.logger.debug(
188
- f"Message end event:\n"
189
- f"Citations: {citations}\n"
190
- f"Metadata: {metadata}\n"
191
- f"Tool outputs: {tool_outputs}"
192
- )
193
-
194
- formatted = self.format_terminal_output(
195
- parsed,
196
- citations=citations,
197
- metadata=metadata,
198
- tool_outputs=tool_outputs
199
- )
200
- if formatted:
201
- self.logger.info(formatted)
202
- yield f"data: {formatted}\n\n"
203
-
204
- except Exception as e:
205
- await self.log_error(
206
- e,
207
- {"line": line, "event": "parse_data"}
208
- )
209
 
210
- # Process buffer when we have a complete message
211
- if buffer.strip().endswith("}"):
212
- try:
213
- processed_response = parser.parse_sse_event(buffer)
214
- if processed_response and isinstance(processed_response, dict):
215
- cleaned_response = self.clean_response(processed_response)
216
- if cleaned_response:
217
- xml_content = cleaned_response.get("content", "")
218
- yield f"data: {xml_content}\n\n"
219
- except Exception as parse_error:
220
- await self.log_error(
221
- parse_error,
222
- {"buffer": buffer, "event": "process_buffer"}
223
- )
224
- error_xml = (
225
- f"<agent_response>"
226
- f"<error>{str(parse_error)}</error>"
227
- f"</agent_response>"
228
- )
229
- yield f"data: {error_xml}\n\n"
230
- finally:
231
- buffer = ""
232
-
233
  except Exception as e:
234
- self.logger.error(f"Stream processing error: {str(e)}")
235
- yield f"data: {{'error': '{str(e)}'}}\n\n"
236
-
237
- return StreamingResponse(
238
- event_generator(),
239
- media_type="text/event-stream",
240
- headers={
241
- "Cache-Control": "no-cache",
242
- "Connection": "keep-alive",
243
- "X-Accel-Buffering": "no",
244
- "Access-Control-Allow-Origin": "*"
245
- }
246
- )
247
 
248
  def format_terminal_output(
249
  self,
 
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, MessageState
14
  from logger_config import setup_logger
15
  from fastapi.responses import StreamingResponse
16
  from fastapi.responses import JSONResponse
 
57
  self.formatter = ResponseFormatter()
58
  self.client = httpx.AsyncClient(timeout=60.0)
59
  self.logger = setup_logger("agent_processor")
60
+
61
+ # Initialize headers
62
+ self.headers = {
63
+ "Authorization": f"Bearer {self.api_key}",
64
+ "Content-Type": "application/json",
65
+ "Accept": "text/event-stream"
66
+ }
67
 
68
+ def prepare_request(self, request: AgentRequest) -> Dict:
69
+ """Prepare request payload for API"""
70
+ return {
71
+ "query": request.query,
72
+ "inputs": request.inputs,
73
+ "response_mode": "streaming" if request.stream else "blocking",
74
+ "user": request.user,
75
+ "conversation_id": request.conversation_id,
76
+ "files": request.files
77
+ }
78
 
79
  async def log_request_details(
80
  self,
 
114
  await self.client.aclose()
115
 
116
  async def process_stream(self, request: AgentRequest):
117
+ """Process streaming request and format for frontend"""
118
+ parser = SSEParser()
119
+ formatter = ResponseFormatter()
 
 
 
 
 
120
 
 
 
 
 
 
 
 
 
 
121
  async def event_generator():
 
 
 
 
 
 
122
  try:
123
  async with self.client.stream(
124
  "POST",
125
  f"{self.api_base}/chat-messages",
126
+ headers=self.headers,
127
+ json=self.prepare_request(request)
128
  ) as response:
 
 
 
 
 
 
129
  async for line in response.aiter_lines():
130
  if not line.strip():
131
  continue
132
 
133
+ # Parse the event
134
+ parsed = parser.parse_sse_event(line)
135
+ if not parsed:
136
+ continue
137
 
138
+ event_type = parsed.get("type")
139
+
140
+ # Format based on type
141
+ if event_type == "message":
142
+ formatted = formatter.format_message(
143
+ message=parsed["content"],
144
+ message_id=parsed["message_id"]
145
+ )
146
+ elif event_type == "thought":
147
+ formatted = formatter.format_thought(
148
+ thought=parsed["content"]["thought"],
149
+ observation=parsed["content"]["observation"],
150
+ message_id=parsed["message_id"]
151
+ )
152
+ elif event_type == "tool_output":
153
+ # Special handling for tool outputs
154
+ formatted = formatter.format_thought(
155
+ thought="",
156
+ observation="",
157
+ tool_outputs=[{
158
+ "type": parsed["tool"],
159
+ "content": parsed["content"]
160
+ }],
161
+ message_id=parsed["message_id"]
162
+ )
163
+ else:
164
+ continue
165
+
166
+ if formatted:
167
+ _, xml_output = formatted
168
+ yield f"data: {xml_output}\n\n"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
169
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
170
  except Exception as e:
171
+ self.logger.error(f"Stream error: {str(e)}")
172
+ error_msg = formatter.format_error(str(e))
173
+ if error_msg:
174
+ _, xml_output = error_msg
175
+ yield f"data: {xml_output}\n\n"
176
+
177
+ return EventSourceResponse(event_generator())
 
 
 
 
 
 
178
 
179
  def format_terminal_output(
180
  self,