wuyiqunLu commited on
Commit
d4e2332
·
unverified ·
1 Parent(s): 4a72a82

feat: add timeout and remove png from code result (#82)

Browse files

tested locally and see the final result:
<img width="1441" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/2c015adc-9793-4325-ab13-31eeb65c652b">

Files changed (1) hide show
  1. app/api/vision-agent/route.ts +44 -2
app/api/vision-agent/route.ts CHANGED
@@ -11,6 +11,7 @@ import { getPresignedUrl } from '@/lib/aws';
11
  // export const runtime = 'edge';
12
  export const dynamic = 'force-dynamic';
13
  export const maxDuration = 300; // This function can run for a maximum of 5 minutes
 
14
 
15
  const uploadBase64 = async (
16
  base64: string,
@@ -91,7 +92,7 @@ export const POST = withLogging(
91
 
92
  const fetchResponse = await fetch(
93
  `https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
94
- // `http://localhost:5001/v1/agent/chat?agent_class=vision_agent&self_reflection=${enableSelfReflection}`,
95
  {
96
  method: 'POST',
97
  headers: {
@@ -148,6 +149,7 @@ export const POST = withLogging(
148
  const decoder = new TextDecoder('utf-8');
149
  let maxChunkSize = 0;
150
  let buffer = '';
 
151
  const stream = new ReadableStream({
152
  async start(controller) {
153
  // const parser = createParser(streamParser);
@@ -158,6 +160,23 @@ export const POST = withLogging(
158
  const lines = buffer
159
  .split('\n')
160
  .filter(line => line.trim().length > 0);
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
162
  let done = false;
163
  const parseLine = async (
@@ -166,9 +185,32 @@ export const POST = withLogging(
166
  ) => {
167
  try {
168
  const msg = JSON.parse(line);
169
- if (msg.type !== 'final_code') {
 
 
 
 
 
170
  return line;
171
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
172
  const result = JSON.parse(msg.payload.result) as ResultPayload;
173
  for (let index = 0; index < result.results.length; index++) {
174
  const png = result.results[index].png ?? '';
 
11
  // export const runtime = 'edge';
12
  export const dynamic = 'force-dynamic';
13
  export const maxDuration = 300; // This function can run for a maximum of 5 minutes
14
+ const TIMEOUT_MILI_SECONDS = 5 * 60 * 1000;
15
 
16
  const uploadBase64 = async (
17
  base64: string,
 
92
 
93
  const fetchResponse = await fetch(
94
  `https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
95
+ // `http://localhost:5001/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
96
  {
97
  method: 'POST',
98
  headers: {
 
149
  const decoder = new TextDecoder('utf-8');
150
  let maxChunkSize = 0;
151
  let buffer = '';
152
+ let time = Date.now();
153
  const stream = new ReadableStream({
154
  async start(controller) {
155
  // const parser = createParser(streamParser);
 
160
  const lines = buffer
161
  .split('\n')
162
  .filter(line => line.trim().length > 0);
163
+ if (lines.length === 0) {
164
+ if (Date.now() - time > TIMEOUT_MILI_SECONDS) {
165
+ logger.info(
166
+ session,
167
+ {
168
+ message: 'Agent timed out',
169
+ },
170
+ request,
171
+ '__Agent_timeout__',
172
+ );
173
+ controller.error(
174
+ `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
175
+ );
176
+ }
177
+ } else {
178
+ time = Date.now();
179
+ }
180
  buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
181
  let done = false;
182
  const parseLine = async (
 
185
  ) => {
186
  try {
187
  const msg = JSON.parse(line);
188
+ if (
189
+ msg.type !== 'final_code' &&
190
+ (msg.type !== 'code' ||
191
+ msg.status === 'started' ||
192
+ msg.status === 'running')
193
+ ) {
194
  return line;
195
  }
196
+ if (msg.type === 'code') {
197
+ const result = JSON.parse(
198
+ msg.payload.result,
199
+ ) as ResultPayload;
200
+ if (result && result.results) {
201
+ msg.payload.result = JSON.stringify({
202
+ ...result,
203
+ results: result.results.map((_result: any) => {
204
+ return {
205
+ ..._result,
206
+ png: undefined,
207
+ mp4: undefined,
208
+ };
209
+ }),
210
+ });
211
+ }
212
+ return JSON.stringify(msg);
213
+ }
214
  const result = JSON.parse(msg.payload.result) as ResultPayload;
215
  for (let index = 0; index < result.results.length; index++) {
216
  const png = result.results[index].png ?? '';