Spaces:
Sleeping
Sleeping
Update api.py
Browse files
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 |
-
|
101 |
-
|
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=
|
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 |
-
|
143 |
-
|
|
|
|
|
144 |
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
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
|
235 |
-
|
236 |
-
|
237 |
-
|
238 |
-
|
239 |
-
|
240 |
-
|
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,
|