|
35 | 35 | from data_formulator.agents.agent_report_gen import ReportGenAgent |
36 | 36 | from data_formulator.agents.client_utils import Client |
37 | 37 |
|
38 | | -from data_formulator.workflows.exploration_flow import run_exploration_flow_streaming |
39 | 38 | from data_formulator.agents.data_agent import DataAgent |
40 | 39 |
|
41 | 40 | # Get logger for this module (logging config done in app.py) |
@@ -411,104 +410,6 @@ def derive_data(): |
411 | 410 | response.headers.add('Access-Control-Allow-Origin', '*') |
412 | 411 | return response |
413 | 412 |
|
414 | | -@agent_bp.route('/explore-data-streaming', methods=['GET', 'POST']) |
415 | | -def explore_data_streaming(): |
416 | | - def generate(): |
417 | | - if request.is_json: |
418 | | - logger.setLevel(logging.INFO) |
419 | | - |
420 | | - logger.info("# explore-data-streaming request") |
421 | | - content = request.get_json() |
422 | | - token = content["token"] |
423 | | - |
424 | | - # each table is a dict with {"name": xxx, "rows": [...]} |
425 | | - input_tables = content["input_tables"] |
426 | | - initial_plan = content["initial_plan"] # The exploration question |
427 | | - max_iterations = content.get("max_iterations", 3) # Number of exploration iterations |
428 | | - max_repair_attempts = content.get("max_repair_attempts", 1) |
429 | | - agent_exploration_rules = content.get("agent_exploration_rules", "") |
430 | | - agent_coding_rules = content.get("agent_coding_rules", "") |
431 | | - conversation_history = content.get("conversation_history", None) |
432 | | - |
433 | | - logger.debug("== input tables ===>") |
434 | | - for table in input_tables: |
435 | | - logger.debug(f"===> Table: {table['name']} (first 5 rows)") |
436 | | - logger.debug(table['rows'][:5]) |
437 | | - |
438 | | - logger.debug("== exploration question ===") |
439 | | - logger.debug(initial_plan) |
440 | | - |
441 | | - # Model config for the exploration flow |
442 | | - model_config = { |
443 | | - "endpoint": content['model']['endpoint'], |
444 | | - "model": content['model']['model'], |
445 | | - "api_key": content['model']['api_key'], |
446 | | - "api_base": content['model'].get('api_base', ''), |
447 | | - "api_version": content['model'].get('api_version', '') |
448 | | - } |
449 | | - |
450 | | - # Get identity for workspace (used for both SQL and Python with WorkspaceWithTempData) |
451 | | - identity_id = get_identity_id() |
452 | | - |
453 | | - try: |
454 | | - for result in run_exploration_flow_streaming( |
455 | | - model_config=model_config, |
456 | | - input_tables=input_tables, |
457 | | - initial_plan=initial_plan, |
458 | | - identity_id=identity_id, |
459 | | - max_iterations=max_iterations, |
460 | | - max_repair_attempts=max_repair_attempts, |
461 | | - agent_exploration_rules=agent_exploration_rules, |
462 | | - agent_coding_rules=agent_coding_rules, |
463 | | - conversation_history=conversation_history |
464 | | - ): |
465 | | - response_data = { |
466 | | - "token": token, |
467 | | - "status": "ok", |
468 | | - "result": result |
469 | | - } |
470 | | - |
471 | | - yield json.dumps(response_data) + '\n' |
472 | | - |
473 | | - # Break if we get a completion result |
474 | | - if result.get("type") == "completion": |
475 | | - break |
476 | | - |
477 | | - except Exception as e: |
478 | | - logger.setLevel(logging.WARNING) |
479 | | - logger.error(f"Error in exploration flow: {e}") |
480 | | - logger.error(traceback.format_exc()) |
481 | | - error_data = { |
482 | | - "token": token, |
483 | | - "status": "error", |
484 | | - "result": None, |
485 | | - "error_message": str(e) |
486 | | - } |
487 | | - yield json.dumps(error_data) + '\n' |
488 | | - |
489 | | - logger.setLevel(logging.WARNING) |
490 | | - |
491 | | - else: |
492 | | - error_data = { |
493 | | - "token": "", |
494 | | - "status": "error", |
495 | | - "result": None, |
496 | | - "error_message": "Invalid request format" |
497 | | - } |
498 | | - yield json.dumps(error_data) + '\n' |
499 | | - |
500 | | - response = Response( |
501 | | - stream_with_context(generate()), |
502 | | - mimetype='application/json', |
503 | | - headers={ |
504 | | - 'Access-Control-Allow-Origin': '*', |
505 | | - 'Access-Control-Allow-Methods': 'GET, POST, OPTIONS', |
506 | | - 'Access-Control-Allow-Headers': 'Content-Type' |
507 | | - } |
508 | | - ) |
509 | | - return response |
510 | | - |
511 | | - |
512 | 413 | @agent_bp.route('/data-agent-streaming', methods=['GET', 'POST']) |
513 | 414 | def data_agent_streaming(): |
514 | 415 | """Autonomous data exploration agent endpoint (SWE-agent style). |
|
0 commit comments