naman1102 commited on
Commit
8fbacb7
·
1 Parent(s): c47b92c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +115 -92
app.py CHANGED
@@ -182,16 +182,17 @@ def discover_attachment(task_id: str, api_url: str) -> Optional[str]:
182
  # -------------------------
183
 
184
  class AgentState(TypedDict):
185
- question: Annotated[str, override]
186
- current_step: Annotated[str, override]
187
- final_answer: Annotated[str, override]
188
- history: Annotated[List[Dict[str, str]], list.__add__]
189
- needs_search: Annotated[bool, override]
190
- search_query: Annotated[str, override]
191
- task_id: Annotated[str, override]
192
- logs: Annotated[Dict[str, Any], merge_dicts]
193
- file_url: Annotated[str, override]
194
- code_blocks: Annotated[List[Dict[str, str]], list.__add__]
 
195
 
196
  # -------------------------
197
  # BasicAgent implementation
@@ -247,7 +248,7 @@ class BasicAgent:
247
 
248
  state: AgentState = {
249
  "question": question,
250
- "current_step": "answer",
251
  "final_answer": "",
252
  "history": [],
253
  "needs_search": False,
@@ -255,7 +256,8 @@ class BasicAgent:
255
  "task_id": task_id,
256
  "logs": {},
257
  "file_url": file_url,
258
- "code_blocks": []
 
259
  }
260
 
261
  print(f"\nProcessing task {task_id}")
@@ -265,7 +267,47 @@ class BasicAgent:
265
  final_state = self.workflow.invoke(state)
266
  return final_state["final_answer"]
267
 
268
- def _process_image(self, state: AgentState) -> AgentState:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
269
  """Process image files using LLaVA."""
270
  try:
271
  print(f"Downloading {state['file_url']} …")
@@ -278,16 +320,20 @@ class BasicAgent:
278
  answer = image_qa_bytes(data, state["question"])
279
 
280
  print(f"Generated answer: {answer}")
281
- state["final_answer"] = answer
282
- state["current_step"] = "done"
283
- return state
 
 
284
  except Exception as e:
285
  print(f"\nError processing image {state['file_url']}: {str(e)}")
286
- state["final_answer"] = f"Error processing image: {str(e)}"
287
- state["current_step"] = "done"
288
- return state
 
 
289
 
290
- def _process_video(self, state: AgentState) -> AgentState:
291
  """Process video files using VideoMAE."""
292
  try:
293
  print(f"Downloading {state['file_url']} …")
@@ -300,16 +346,20 @@ class BasicAgent:
300
  answer = video_label_bytes(data)
301
 
302
  print(f"Generated answer: {answer}")
303
- state["final_answer"] = answer
304
- state["current_step"] = "done"
305
- return state
 
 
306
  except Exception as e:
307
  print(f"\nError processing video {state['file_url']}: {str(e)}")
308
- state["final_answer"] = f"Error processing video: {str(e)}"
309
- state["current_step"] = "done"
310
- return state
 
 
311
 
312
- def _process_spreadsheet(self, state: AgentState) -> AgentState:
313
  """Process spreadsheet files."""
314
  try:
315
  print(f"Downloading {state['file_url']} …")
@@ -322,16 +372,20 @@ class BasicAgent:
322
  answer = sheet_answer_bytes(data)
323
 
324
  print(f"Generated answer: {answer}")
325
- state["final_answer"] = answer
326
- state["current_step"] = "done"
327
- return state
 
 
328
  except Exception as e:
329
  print(f"\nError processing spreadsheet {state['file_url']}: {str(e)}")
330
- state["final_answer"] = f"Error processing spreadsheet: {str(e)}"
331
- state["current_step"] = "done"
332
- return state
 
 
333
 
334
- def _process_python(self, state: AgentState) -> AgentState:
335
  """Process Python files."""
336
  try:
337
  print(f"Downloading {state['file_url']} …")
@@ -344,16 +398,20 @@ class BasicAgent:
344
  answer = run_python(data.decode())
345
 
346
  print(f"Generated answer: {answer}")
347
- state["final_answer"] = answer
348
- state["current_step"] = "done"
349
- return state
 
 
350
  except Exception as e:
351
  print(f"\nError processing Python file {state['file_url']}: {str(e)}")
352
- state["final_answer"] = f"Error processing Python file: {str(e)}"
353
- state["current_step"] = "done"
354
- return state
 
 
355
 
356
- def _process_text(self, state: AgentState) -> AgentState:
357
  """Process text-only questions using LLM."""
358
  print("\nProcessing as text-only question...")
359
  prompt = f"""
@@ -368,53 +426,18 @@ QUESTION:
368
  raw = self._call_llm(prompt, 300)
369
  answer = self._safe_parse(raw)
370
  print(f"Generated answer: {answer}")
371
- state["final_answer"] = answer
 
 
 
 
372
  except Exception as e:
373
  print(f"\nLLM Error in answer generation: {str(e)}")
374
- state["final_answer"] = "I encountered an error while generating the answer."
375
-
376
- state["current_step"] = "done"
377
- return state
378
-
379
- def _route_to_tool(self, state: AgentState) -> str:
380
- """Route the state to the appropriate tool based on file type."""
381
- if not state["file_url"]:
382
- return "process_text"
383
-
384
- try:
385
- response = SESSION.get(state["file_url"], timeout=30)
386
- response.raise_for_status()
387
- data = response.content
388
-
389
- # Get content type from response headers first, fallback to URL-based detection
390
- kind = response.headers.get("Content-Type", "")
391
- if kind in ("application/octet-stream", ""):
392
- # rough sniff: look at the first few bytes
393
- sig = data[:4]
394
- if sig.startswith(b"\x89PNG"):
395
- kind = "image/png"
396
- elif sig.startswith(b"\xFF\xD8"):
397
- kind = "image/jpeg"
398
- elif sig[:2] == b"PK": # XLSX = ZIP
399
- kind = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
400
- elif not kind: # fallback if header missing
401
- kind = mimetypes.guess_type(state["file_url"])[0] or ""
402
-
403
- if "image" in kind:
404
- return "process_image"
405
- elif "video" in kind:
406
- return "process_video"
407
- elif "spreadsheet" in kind or "excel" in kind:
408
- return "process_spreadsheet"
409
- elif state["file_url"].endswith(".py"):
410
- return "process_python"
411
- else:
412
- print(f"Unsupported file type: {kind}")
413
- return "process_text"
414
-
415
- except Exception as e:
416
- print(f"Error determining file type: {str(e)}")
417
- return "process_text"
418
 
419
  def _build_workflow(self) -> Graph:
420
  """Build the workflow graph with conditional edges."""
@@ -431,15 +454,15 @@ QUESTION:
431
  # Set entry point
432
  sg.set_entry_point("route")
433
 
434
- # Add conditional edges
435
  sg.add_conditional_edges(
436
  "route",
437
  {
438
- "process_image": lambda x: x == "process_image",
439
- "process_video": lambda x: x == "process_video",
440
- "process_spreadsheet": lambda x: x == "process_spreadsheet",
441
- "process_python": lambda x: x == "process_python",
442
- "process_text": lambda x: x == "process_text"
443
  }
444
  )
445
 
 
182
  # -------------------------
183
 
184
  class AgentState(TypedDict):
185
+ question: str
186
+ current_step: str
187
+ final_answer: str
188
+ history: List[Dict[str, str]]
189
+ needs_search: bool
190
+ search_query: str
191
+ task_id: str
192
+ logs: Dict[str, Any]
193
+ file_url: str
194
+ code_blocks: List[Dict[str, str]]
195
+ next_step: str # Add this field for routing
196
 
197
  # -------------------------
198
  # BasicAgent implementation
 
248
 
249
  state: AgentState = {
250
  "question": question,
251
+ "current_step": "route",
252
  "final_answer": "",
253
  "history": [],
254
  "needs_search": False,
 
256
  "task_id": task_id,
257
  "logs": {},
258
  "file_url": file_url,
259
+ "code_blocks": [],
260
+ "next_step": "" # Initialize next_step
261
  }
262
 
263
  print(f"\nProcessing task {task_id}")
 
267
  final_state = self.workflow.invoke(state)
268
  return final_state["final_answer"]
269
 
270
+ def _route_to_tool(self, state: AgentState) -> Dict[str, Any]:
271
+ """Route the state to the appropriate tool based on file type."""
272
+ if not state["file_url"]:
273
+ return {"next_step": "process_text"}
274
+
275
+ try:
276
+ response = SESSION.get(state["file_url"], timeout=30)
277
+ response.raise_for_status()
278
+ data = response.content
279
+
280
+ # Get content type from response headers first, fallback to URL-based detection
281
+ kind = response.headers.get("Content-Type", "")
282
+ if kind in ("application/octet-stream", ""):
283
+ # rough sniff: look at the first few bytes
284
+ sig = data[:4]
285
+ if sig.startswith(b"\x89PNG"):
286
+ kind = "image/png"
287
+ elif sig.startswith(b"\xFF\xD8"):
288
+ kind = "image/jpeg"
289
+ elif sig[:2] == b"PK": # XLSX = ZIP
290
+ kind = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
291
+ elif not kind: # fallback if header missing
292
+ kind = mimetypes.guess_type(state["file_url"])[0] or ""
293
+
294
+ if "image" in kind:
295
+ return {"next_step": "process_image"}
296
+ elif "video" in kind:
297
+ return {"next_step": "process_video"}
298
+ elif "spreadsheet" in kind or "excel" in kind:
299
+ return {"next_step": "process_spreadsheet"}
300
+ elif state["file_url"].endswith(".py"):
301
+ return {"next_step": "process_python"}
302
+ else:
303
+ print(f"Unsupported file type: {kind}")
304
+ return {"next_step": "process_text"}
305
+
306
+ except Exception as e:
307
+ print(f"Error determining file type: {str(e)}")
308
+ return {"next_step": "process_text"}
309
+
310
+ def _process_image(self, state: AgentState) -> Dict[str, Any]:
311
  """Process image files using LLaVA."""
312
  try:
313
  print(f"Downloading {state['file_url']} …")
 
320
  answer = image_qa_bytes(data, state["question"])
321
 
322
  print(f"Generated answer: {answer}")
323
+ return {
324
+ "final_answer": answer,
325
+ "current_step": "done",
326
+ "next_step": "done"
327
+ }
328
  except Exception as e:
329
  print(f"\nError processing image {state['file_url']}: {str(e)}")
330
+ return {
331
+ "final_answer": f"Error processing image: {str(e)}",
332
+ "current_step": "done",
333
+ "next_step": "done"
334
+ }
335
 
336
+ def _process_video(self, state: AgentState) -> Dict[str, Any]:
337
  """Process video files using VideoMAE."""
338
  try:
339
  print(f"Downloading {state['file_url']} …")
 
346
  answer = video_label_bytes(data)
347
 
348
  print(f"Generated answer: {answer}")
349
+ return {
350
+ "final_answer": answer,
351
+ "current_step": "done",
352
+ "next_step": "done"
353
+ }
354
  except Exception as e:
355
  print(f"\nError processing video {state['file_url']}: {str(e)}")
356
+ return {
357
+ "final_answer": f"Error processing video: {str(e)}",
358
+ "current_step": "done",
359
+ "next_step": "done"
360
+ }
361
 
362
+ def _process_spreadsheet(self, state: AgentState) -> Dict[str, Any]:
363
  """Process spreadsheet files."""
364
  try:
365
  print(f"Downloading {state['file_url']} …")
 
372
  answer = sheet_answer_bytes(data)
373
 
374
  print(f"Generated answer: {answer}")
375
+ return {
376
+ "final_answer": answer,
377
+ "current_step": "done",
378
+ "next_step": "done"
379
+ }
380
  except Exception as e:
381
  print(f"\nError processing spreadsheet {state['file_url']}: {str(e)}")
382
+ return {
383
+ "final_answer": f"Error processing spreadsheet: {str(e)}",
384
+ "current_step": "done",
385
+ "next_step": "done"
386
+ }
387
 
388
+ def _process_python(self, state: AgentState) -> Dict[str, Any]:
389
  """Process Python files."""
390
  try:
391
  print(f"Downloading {state['file_url']} …")
 
398
  answer = run_python(data.decode())
399
 
400
  print(f"Generated answer: {answer}")
401
+ return {
402
+ "final_answer": answer,
403
+ "current_step": "done",
404
+ "next_step": "done"
405
+ }
406
  except Exception as e:
407
  print(f"\nError processing Python file {state['file_url']}: {str(e)}")
408
+ return {
409
+ "final_answer": f"Error processing Python file: {str(e)}",
410
+ "current_step": "done",
411
+ "next_step": "done"
412
+ }
413
 
414
+ def _process_text(self, state: AgentState) -> Dict[str, Any]:
415
  """Process text-only questions using LLM."""
416
  print("\nProcessing as text-only question...")
417
  prompt = f"""
 
426
  raw = self._call_llm(prompt, 300)
427
  answer = self._safe_parse(raw)
428
  print(f"Generated answer: {answer}")
429
+ return {
430
+ "final_answer": answer,
431
+ "current_step": "done",
432
+ "next_step": "done"
433
+ }
434
  except Exception as e:
435
  print(f"\nLLM Error in answer generation: {str(e)}")
436
+ return {
437
+ "final_answer": "I encountered an error while generating the answer.",
438
+ "current_step": "done",
439
+ "next_step": "done"
440
+ }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
441
 
442
  def _build_workflow(self) -> Graph:
443
  """Build the workflow graph with conditional edges."""
 
454
  # Set entry point
455
  sg.set_entry_point("route")
456
 
457
+ # Add conditional edges based on next_step field
458
  sg.add_conditional_edges(
459
  "route",
460
  {
461
+ "process_image": lambda x: x["next_step"] == "process_image",
462
+ "process_video": lambda x: x["next_step"] == "process_video",
463
+ "process_spreadsheet": lambda x: x["next_step"] == "process_spreadsheet",
464
+ "process_python": lambda x: x["next_step"] == "process_python",
465
+ "process_text": lambda x: x["next_step"] == "process_text"
466
  }
467
  )
468