257 lines
9.5 KiB
Python
257 lines
9.5 KiB
Python
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))
|