KGPython/controller/QAController.py

import json
import traceback
import uuid

from util.neo4j_utils import Neo4jUtil, neo4j_client
import httpx
from robyn import jsonify, Response

from app import app
from controller.client import client


# --- 核心工具函数：解决元组返回错误及中文乱码 ---
def create_response(status_code, data_dict):
    """
    统一响应格式封装。
    1. 确保返回的是 Robyn 预期的 Response 对象。
    2. description 必须是字符串（json.dumps 结果）。
    3. 强制使用 UTF-8 防止中文乱码。
    """
    return Response(
        status_code=status_code,
        description=json.dumps(data_dict, ensure_ascii=False),
        headers={"Content-Type": "application/json; charset=utf-8"}
    )


def convert_to_g6_format(data):
    entities = data.get("entities", [])
    relations = data.get("relations", [])

    # 创建实体名称到唯一ID的映射
    name_to_id = {}
    nodes = []

    for ent in entities:
        name = ent.get("n")
        if name and name not in name_to_id:
            node_id = str(uuid.uuid4())
            name_to_id[name] = node_id
            nodes.append({
                "id": node_id,
                "label": name,
                "data": {
                    "type": ent.get("t")  # 用于 G6 的节点样式区分
                }
            })

    # 构建边，并为每条边生成唯一 ID
    edges = []
    for rel in relations:
        e1 = rel.get("e1")
        e2 = rel.get("e2")
        r = rel.get("r")

        source_id = name_to_id.get(e1)
        target_id = name_to_id.get(e2)

        if source_id and target_id:
            edge_id = str(uuid.uuid4())
            edges.append({
                "id": edge_id,
                "source": source_id,
                "target": target_id,
                "label": r,
                "data": {
                    "label": r
                }
            })
        else:
            print(f"Warning: Entity not found for relation: {rel}")

    return {
        "nodes": nodes,
        "edges": edges
    }


@app.post("/api/qa/analyze")
async def analyze(request):
    body = request.json()
    input_text = body.get("text", "").strip()

    if not input_text:
        # 使用 create_response 统一返回格式，避免使用 jsonify 可能带来的元组嵌套问题
        return create_response(400, {"error": "缺少 text 字段"})

    try:
        # 1. 提取实体
        resp = await client.post(
            "/getEntity",
            json={"text": input_text},
            timeout=1800.0
        )

        qaList = []
        if resp.status_code in (200, 202):
            resp_json = resp.json()
            # 处理字符串形式的 data 字段
            resp_json_data = resp_json.get("data", "{}")
            if isinstance(resp_json_data, str):
                resp_json_data = json.loads(resp_json_data)

            entities = resp_json_data.get("entities", [])
            print(f"提取到的实体: {entities}")

            # 查询 Neo4j 邻居（此逻辑保留，虽目前未直接放入 qaList，可能用于后续扩展）
            for name in entities:
                neo4j_client.find_neighbors_with_relationshipsAI(
                    node_label=None,
                    direction="both",
                    node_properties={"name": name},
                    rel_type=None
                )

            # 2. 问答代理获取答案列表
            resp_agent = await client.post(
                "/question_agent",
                json={"neo4j_data": [], "text": input_text},
                timeout=1800.0
            )

            resp_data = resp_agent.json()
            inner_data = json.loads(resp_data["data"])
            items = inner_data.get("json", [])

            # 按权重排序取前5
            sorted_items = sorted(items, key=lambda x: x.get("sort", 0), reverse=True)
            top5 = sorted_items[:5]

            # 3. 对每个答案提取关系图谱
            for item in top5:
                resp_ext = await client.post(
                    "/extract_entities_and_relations",
                    json={"text": item['answer']},
                    timeout=1800.0
                )

                if resp_ext.status_code in (200, 202):
                    result = resp_ext.json()
                    g6_data = convert_to_g6_format(result)
                    qaList.append({
                        "answer": item["answer"],
                        "result": g6_data,
                    })
                    print(f"处理成功 xh: {item.get('xh')}, sort: {item.get('sort')}")

            # --- 修复点：使用 create_response 返回解析后的数组 ---
            return create_response(200, qaList)

        else:
            return create_response(resp.status_code, {
                "error": "提交失败",
                "detail": resp.text
            })

    except Exception as e:
        error_trace = traceback.format_exc()
        print("❌ 发生异常:")
        print(error_trace)
        return create_response(500, {"error": str(e), "traceback": error_trace})