import asyncio import json import time import uuid from fastapi import APIRouter, HTTPException, Depends from fastapi.responses import StreamingResponse from config import Config from schemas.agent_input import AgentInput from schemas.agent_output import AgentOutput from schemas.tool_input import ToolInput from schemas.tool_output import ToolOutput from schemas.chat_message_response import ChatMessageResponseDTO from workflows.workflow_manager import WorkflowType from api.dependencies import get_workflow_manager, get_nacos_manager, get_service_config, get_tool_router, get_prompt_manager from services.app_errors import AppError, ErrorCode from services.ragflow_sync import RagflowSync from services.structured_logger import get_structured_logger from tools.sr_api_tool import SrApiQueryTool router = APIRouter() def _resolve_workflow_type(value: str) -> WorkflowType: try: return WorkflowType(value) except Exception as e: raise AppError( code=ErrorCode.INVALID_WORKFLOW_TYPE, message=f"不支持的工作流类型: {value}", status_code=400, ) from e def _to_http_error(e: Exception) -> HTTPException: if isinstance(e, AppError): return HTTPException(status_code=e.status_code, detail=e.to_dict()) return HTTPException(status_code=500, detail={"code": ErrorCode.INTERNAL_ERROR.value, "message": str(e)}) @router.get("/health") def health_check(service_config=Depends(get_service_config)): return { "status": "ok", "service_name": service_config.service_name, "model_section": service_config.metadata.get("model_section", "") } @router.get("/nacos/status") def nacos_status(nacos_manager=Depends(get_nacos_manager)): return nacos_manager.status() @router.post("/api/workflows", response_model=AgentOutput) def run_workflow(payload: AgentInput, workflow_manager=Depends(get_workflow_manager)): trace_id = uuid.uuid4().hex slog = get_structured_logger() slog.log("INFO", "run_workflow.start", trace_id, {"workflow_type": payload.workflow_type}) try: workflow_type = _resolve_workflow_type(payload.workflow_type) except Exception as e: slog.log("ERROR", "run_workflow.invalid_type", trace_id, error_code=ErrorCode.INVALID_WORKFLOW_TYPE.value, payload={"workflow_type": payload.workflow_type}) raise _to_http_error(e) try: result = workflow_manager.execute_workflow( workflow_type=workflow_type, user_input=payload.input, session_id=payload.session_id, ) slog.log("INFO", "run_workflow.success", trace_id, {"session_id": result.get("session_id")}) except Exception as e: slog.log("ERROR", "run_workflow.failed", trace_id, error_code=ErrorCode.INTERNAL_ERROR.value, payload={"error": str(e)}) raise _to_http_error(e) return AgentOutput( session_id=result["session_id"], workflow_type=result["workflow_type"], result=result["result"], ) @router.post("/api/sql/generate") def generate_sql(payload: AgentInput, workflow_manager=Depends(get_workflow_manager)): """仅生成 SQL，不调用 SR API""" trace_id = uuid.uuid4().hex slog = get_structured_logger() try: workflow_type = _resolve_workflow_type(payload.workflow_type) except Exception as e: slog.log("ERROR", "generate_sql.invalid_type", trace_id, error_code=ErrorCode.INVALID_WORKFLOW_TYPE.value) raise _to_http_error(e) result = workflow_manager.execute_workflow( workflow_type=workflow_type, user_input=payload.input, session_id=payload.session_id, skip_sr_api=True, ) context = (result.get("result") or {}).get("context") or {} sql_text = context.get("final_sql") if not sql_text: e = AppError(code=ErrorCode.SQL_GENERATION_FAILED, message="SQL 生成失败") slog.log("ERROR", "generate_sql.failed", trace_id, error_code=e.code.value, payload={"context_keys": list(context.keys())}) raise _to_http_error(e) slog.log("INFO", "generate_sql.success", trace_id, {"sql_len": len(sql_text)}) return { "session_id": result.get("session_id"), "workflow_type": result.get("workflow_type"), "sql": sql_text, } @router.post("/api/workflows/stream") def run_workflow_stream(payload: AgentInput, workflow_manager=Depends(get_workflow_manager)): trace_id = uuid.uuid4().hex slog = get_structured_logger() try: workflow_type = _resolve_workflow_type(payload.workflow_type) except Exception as e: raise _to_http_error(e) if workflow_type != WorkflowType.CONVERSATION: raise _to_http_error(AppError(code=ErrorCode.INVALID_WORKFLOW_TYPE, message="仅支持对话工作流的流式输出", status_code=400)) stream_cfg = Config.get_section("stream") progress_interval = float(stream_cfg.get("progress_interval", 0.3)) task_id = uuid.uuid4().hex def _build_message(conversation_id: str, answer: str) -> str: dto = ChatMessageResponseDTO( id=uuid.uuid4().hex, event="message", task_id=task_id, message_id=uuid.uuid4().hex, conversation_id=conversation_id, answer=answer, created_at=int(time.time()), ) return f"event: message\ndata: {json.dumps(dto.model_dump(), ensure_ascii=False)}\n\n" async def event_stream(): try: slog.log("INFO", "stream.start", trace_id, {"workflow_type": payload.workflow_type}) # 1) 先仅生成 SQL（不执行 SR API） result = await asyncio.to_thread( workflow_manager.execute_workflow, workflow_type, payload.input, payload.session_id, skip_sr_api=True, ) conversation_id = str(result.get("session_id") or payload.session_id or task_id) context = (result.get("result") or {}).get("context") or {} sql_text = str(context.get("final_sql") or "") if not sql_text: reason = "SQL 生成失败，可能是表未匹配或对应 SQL 提示词不存在" slog.log("ERROR", "stream.sql_generation_failed", trace_id, error_code=ErrorCode.SQL_GENERATION_FAILED.value, payload={"conversation_id": conversation_id}) yield _build_message(conversation_id, reason) yield "event: end\ndata: [DONE]\n\n" return # 2) 先流式返回 SQL yield _build_message(conversation_id, sql_text) # 3) 异步执行 SQL，并及时流式返回执行结果 tool = SrApiQueryTool() task = asyncio.create_task( asyncio.to_thread(tool.run, json.dumps({"sql": sql_text}, ensure_ascii=False)) ) while not task.done(): yield _build_message(conversation_id, "executing_sql") await asyncio.sleep(progress_interval) sql_result = await task slog.log("INFO", "stream.sql_executed", trace_id, {"result_len": len(str(sql_result))}) yield _build_message(conversation_id, str(sql_result)) yield "event: end\ndata: [DONE]\n\n" except Exception as e: slog.log("ERROR", "stream.failed", trace_id, error_code=ErrorCode.INTERNAL_ERROR.value, payload={"error": str(e)}) conversation_id = str(payload.session_id or task_id) yield _build_message(conversation_id, str(e)) yield "event: end\ndata: [DONE]\n\n" return StreamingResponse(event_stream(), media_type="text/event-stream") @router.post("/api/tools/execute", response_model=ToolOutput) def run_tool(payload: ToolInput, tool_router=Depends(get_tool_router)): result = tool_router.call(payload.tool_name, payload.payload) return ToolOutput(**result) @router.post("/api/prompts/reload") def reload_prompts(prompt_manager=Depends(get_prompt_manager)): prompt_manager.reload() return {"ok": True} @router.post("/api/ragflow/table-retrieval/reload") def reload_table_retrieval(): syncer = RagflowSync() result = syncer.upload_table_retrieval() return {"ok": True, "result": result} @router.post("/api/ragflow/table-retrieval/upload") def upload_table_retrieval(): """上传表名检索模板文档""" syncer = RagflowSync() try: result = syncer.upload_table_retrieval() return {"ok": True, "result": result} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @router.put("/api/ragflow/table-retrieval/update") def update_table_retrieval(): """更新表名检索文档（仅文档内容）""" syncer = RagflowSync() try: result = syncer.update_table_retrieval_documents() return {"ok": True, "result": result} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @router.post("/api/ragflow/sql-gen/upload") def upload_sql_gen(): """上传 SQL 生成提示词文档""" syncer = RagflowSync() try: result = syncer.upload_sql_gen() return {"ok": True, "result": result} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @router.put("/api/ragflow/sql-gen/update") def update_sql_gen(): """更新 SQL 生成文档（仅文档内容）""" syncer = RagflowSync() try: result = syncer.update_sql_gen_documents() return {"ok": True, "result": result} except Exception as e: raise HTTPException(status_code=500, detail=str(e))