agent-management/backend/rag-ai-bridge/server.py

import logging
import json
import os
import re
from pathlib import Path
from typing import Any

import yaml
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

BRIDGE_DIR = Path(__file__).resolve().parent
CONFIG_PATH = BRIDGE_DIR / "config.yaml"
ENV_PATH = BRIDGE_DIR / ".env"

DEFAULTS = {
    "server": {"host": "127.0.0.1", "port": 18733},
    "llm": {
        "provider": "openai",
        "model": "gpt-4o-mini",
        "base_url": "https://api.openai.com/v1",
        "temperature": 1.0,
        "timeout": 60,
        "generation_max_tokens": 8192,
        "answer_max_tokens": 4096,
    },
    "vanna": {"enabled": True},
    "text2cypher": {"enabled": True},
    "neo4j": {"uri": "bolt://127.0.0.1:7687", "username": "neo4j"},
    "security": {"readonly_sql": True, "readonly_cypher": True, "default_limit": 50},
}

ENV_MAPPINGS = {
    "RAG_AI_BRIDGE_HOST": ("server", "host", str),
    "RAG_AI_BRIDGE_PORT": ("server", "port", int),
    "LLM_PROVIDER": ("llm", "provider", str),
    "OPENAI_MODEL": ("llm", "model", str),
    "OPENAI_BASE_URL": ("llm", "base_url", str),
    "LLM_TEMPERATURE": ("llm", "temperature", float),
    "LLM_TIMEOUT": ("llm", "timeout", int),
    "LLM_GENERATION_MAX_TOKENS": ("llm", "generation_max_tokens", int),
    "LLM_ANSWER_MAX_TOKENS": ("llm", "answer_max_tokens", int),
    "VANNA_ENABLED": ("vanna", "enabled", "bool"),
    "TEXT2CYPHER_ENABLED": ("text2cypher", "enabled", "bool"),
    "NEO4J_URI": ("neo4j", "uri", str),
    "NEO4J_USERNAME": ("neo4j", "username", str),
    "READONLY_SQL": ("security", "readonly_sql", "bool"),
    "READONLY_CYPHER": ("security", "readonly_cypher", "bool"),
    "DEFAULT_LIMIT": ("security", "default_limit", int),
}


def _merge(base: dict, override: dict) -> dict:
    result = {key: value.copy() if isinstance(value, dict) else value for key, value in base.items()}
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(result.get(key), dict):
            result[key] = _merge(result[key], value)
        else:
            result[key] = value
    return result


def _convert(value: str, converter):
    if converter == "bool":
        return value.strip().lower() in {"1", "true", "yes", "on"}
    return converter(value)


def load_settings() -> dict:
    file_config = {}
    if CONFIG_PATH.exists():
        with CONFIG_PATH.open("r", encoding="utf-8") as stream:
            file_config = yaml.safe_load(stream) or {}
        if not isinstance(file_config, dict):
            raise ValueError("config.yaml root must be a mapping")
    else:
        logging.warning("config.yaml not found; using defaults. Copy config.example.yaml to config.yaml to customize.")

    # override=False 保证操作系统环境变量优先于 .env。
    load_dotenv(ENV_PATH, override=False)
    settings = _merge(DEFAULTS, file_config)
    for env_name, (section, key, converter) in ENV_MAPPINGS.items():
        value = os.getenv(env_name)
        if value is not None and value != "":
            settings.setdefault(section, {})[key] = _convert(value, converter)
    return settings


SETTINGS = load_settings()
app = FastAPI(title="RAG AI Bridge", version="1.1.0")


class Text2SqlRequest(BaseModel):
    query: str = Field(min_length=1)
    datasourceId: int
    dialect: str = "mysql"
    ddl: str = ""
    schemaVersion: str = ""
    documentation: str = ""
    examples: list[Any] = Field(default_factory=list)
    tableWhitelist: list[str] = Field(default_factory=list)
    selectedTables: list[str] = Field(default_factory=list)
    foreignKeys: list[dict[str, Any]] = Field(default_factory=list)
    relationships: list[dict[str, Any]] = Field(default_factory=list)
    sampledValues: dict[str, list[Any]] = Field(default_factory=dict)
    referenceRows: dict[str, list[dict[str, Any]]] = Field(default_factory=dict)
    contextWarnings: list[str] = Field(default_factory=list)
    maxRows: int = Field(default=SETTINGS["security"]["default_limit"], ge=1, le=1000)
    entityMentions: list[str] = Field(default_factory=list)


class Text2CypherRequest(BaseModel):
    query: str = Field(min_length=1)
    graphSourceId: int
    schema: str = ""
    examples: list[Any] = Field(default_factory=list)
    allowedLabels: list[str] = Field(default_factory=list)
    allowedRelationships: list[str] = Field(default_factory=list)
    allowedProperties: dict[str, list[str]] = Field(default_factory=dict)
    businessRules: str = ""
    maxDepth: int = Field(default=3, ge=1, le=10)
    entityMentions: list[str] = Field(default_factory=list)


class AnswerRequest(BaseModel):
    question: str = Field(min_length=1)
    evidences: list[dict[str, Any]] = Field(default_factory=list)


class GovernanceRequest(BaseModel):
    sourceType: str
    schemaText: str
    sourceDescription: str = ""
    maxExamples: int = Field(default=8, ge=1, le=20)


class PlanRequest(BaseModel):
    question: str
    capabilities: dict[str, str] = Field(default_factory=dict)


class RepairRequest(BaseModel):
    language: str
    question: str
    query: str
    error: str
    schemaText: str
    dialect: str = ""
    selectedTables: list[str] = Field(default_factory=list)
    foreignKeys: list[dict[str, Any]] = Field(default_factory=list)
    relationships: list[dict[str, Any]] = Field(default_factory=list)
    sampledValues: dict[str, list[Any]] = Field(default_factory=dict)
    referenceRows: dict[str, list[dict[str, Any]]] = Field(default_factory=dict)
    maxRows: int = 50
    maxDepth: int = 3


def _read_only(text: str, language: str) -> str:
    match = re.search(r"```(?:sql|cypher)?\s*(.*?)```", text, re.I | re.S)
    query = (match.group(1) if match else text).strip()
    if query.endswith(";"):
        query = query[:-1].rstrip()
    if ";" in query:
        raise ValueError(f"generated {language} contains multiple statements")
    readonly_enabled = SETTINGS["security"]["readonly_sql" if language == "SQL" else "readonly_cypher"]
    if not readonly_enabled:
        return query
    forbidden = r"\b(insert|update|delete|merge|create|drop|alter|truncate|grant|revoke|load\s+csv|call)\b"
    if re.search(forbidden, query, re.I):
        raise ValueError(f"generated {language} contains a write or unsafe operation")
    if language == "SQL" and not re.match(r"^(select|with)\b", query, re.I):
        raise ValueError(f"generated SQL is not a SELECT/CTE: {query[:200]!r}")
    if language == "Cypher" and not re.match(r"^(match|optional\s+match|with|unwind)\b", query, re.I):
        raise ValueError("generated Cypher is not read-only")
    return query


def _extract_code(text: str, language: str) -> str:
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.I | re.S).strip()
    match = re.search(rf"```(?:{language.lower()})?\s*(.*?)```", text, re.I | re.S)
    if match:
        return match.group(1).strip()
    start = r"\b(?:SELECT|WITH)\b" if language == "SQL" else r"\b(?:MATCH|OPTIONAL\s+MATCH|WITH|UNWIND)\b"
    statement = re.search(start + r"[\s\S]*", text, re.I)
    extracted = (statement.group(0) if statement else text).strip()
    return re.split(r"\n\s*(?:Explanation|Reasoning|说明|解释)\s*[:：]", extracted, maxsplit=1, flags=re.I)[0].strip()


def _llm_configured() -> bool:
    return SETTINGS["llm"]["provider"].lower() == "openai" and bool(os.getenv("OPENAI_API_KEY"))


def _neo4j_configured() -> bool:
    neo4j = SETTINGS["neo4j"]
    return bool(neo4j.get("uri") and neo4j.get("username") and os.getenv("NEO4J_PASSWORD"))


def _openai_client(timeout: int | None = None):
    from openai import OpenAI
    return OpenAI(api_key=os.environ["OPENAI_API_KEY"], base_url=SETTINGS["llm"]["base_url"],
                  timeout=_client_timeout(timeout), max_retries=0)


def _client_timeout(timeout: int | None = None):
    configured = SETTINGS["llm"]["timeout"] if timeout is None else timeout
    return None if configured is None or configured <= 0 else configured


def _completion_content(response) -> str:
    choice = response.choices[0]
    content = choice.message.content or ""
    if content.strip():
        return content
    finish_reason = getattr(choice, "finish_reason", None) or "unknown"
    raise ValueError(f"model returned empty content (finish_reason={finish_reason})")


def _json_object(text: str) -> dict:
    cleaned = re.sub(r"<think>.*?</think>", "", text or "", flags=re.I | re.S).strip()
    fenced = re.search(r"```(?:json)?\s*(.*?)```", cleaned, re.I | re.S)
    candidate = fenced.group(1) if fenced else cleaned
    start, end = candidate.find("{"), candidate.rfind("}")
    if start < 0 or end < start:
        raise ValueError("LLM did not return a JSON object")
    return json.loads(candidate[start:end + 1])


def _short_error(error: str) -> str:
    return re.sub(r"(?i)(password|token|api[_ -]?key)\s*[:=]\s*\S+", r"\1=[redacted]", error)[:800]


def _text2sql_messages(req: Text2SqlRequest) -> list[dict[str, str]]:
    if req.dialect.lower() != "mysql":
        raise ValueError("Text2SQL generation currently supports MySQL only")
    context = {
        "dialect": "mysql",
        "schemaVersion": req.schemaVersion,
        "selectedTables": req.selectedTables,
        "minimalDdl": req.ddl,
        "foreignKeys": req.foreignKeys,
        "relationships": req.relationships,
        "sampledValues": req.sampledValues,
        "referenceRows": req.referenceRows,
        "entityMentions": req.entityMentions,
        "generationRules": req.documentation,
    }
    return [{
        "role": "system",
        "content": (
            "You generate exactly one final read-only MySQL SQL query. "
            "Return SQL only, without comments, explanation or intermediate queries. "
            "Use only tables and columns in minimalDdl and JOIN only through explicitly provided relationships. "
            "A CONDITIONAL relationship requires every listed condition in the JOIN or WHERE clause. "
            "Filter literals must be grounded in sampledValues/referenceRows. "
            "Never invent business mappings, identifiers, columns, tables, values, joins or answers. "
            "Preserve conjunctions: a request for both A and B must require both, normally with GROUP BY/HAVING or equivalent logic. "
            "For highest/best by a stated score, ORDER BY that score DESC; never replace it with nearest distance. "
            "Use MySQL syntax only; DISTINCT ON is forbidden. Never generate DML or DDL. "
            f"Limit results to at most {req.maxRows} rows."
        ),
    }, {
        "role": "user",
        "content": (
            "Verified query context (JSON):\n"
            + json.dumps(context, ensure_ascii=False, default=str)
            + f"\nQuestion: {req.query}"
        ),
    }]


@app.get("/health")
def health():
    return {
        "ok": True,
        "llmConfigured": _llm_configured(),
        "neo4jConfigured": _neo4j_configured(),
        "vannaEnabled": bool(SETTINGS["vanna"]["enabled"]),
        "text2cypherEnabled": bool(SETTINGS["text2cypher"]["enabled"]),
    }


@app.post("/text2sql")
def text2sql(req: Text2SqlRequest):
    if not SETTINGS["vanna"]["enabled"]:
        raise HTTPException(status_code=503, detail="Vanna SQL generation is disabled")
    if not _llm_configured():
        raise HTTPException(status_code=503, detail="OpenAI provider or OPENAI_API_KEY is not configured")
    try:
        response = _openai_client().chat.completions.create(
            model=SETTINGS["llm"]["model"], messages=_text2sql_messages(req),
            temperature=SETTINGS["llm"]["temperature"],
            max_tokens=SETTINGS["llm"]["generation_max_tokens"],
        )
        sql = _read_only(_extract_code(_completion_content(response), "SQL"), "SQL")
        return {"sql": sql, "confidence": 0.0,
                "usedContext": ["minimalDdl", "relationships", "sampledValues", "referenceRows"],
                "warnings": req.contextWarnings}
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"Vanna SQL generation unavailable: {exc}") from exc


@app.post("/text2cypher")
def text2cypher(req: Text2CypherRequest):
    if not SETTINGS["text2cypher"]["enabled"]:
        raise HTTPException(status_code=503, detail="Text-to-Cypher generation is disabled")
    if not _llm_configured():
        raise HTTPException(status_code=503, detail="OpenAI provider or OPENAI_API_KEY is not configured")
    try:
        from openai import OpenAI

        llm = OpenAI(
            api_key=os.environ["OPENAI_API_KEY"],
            base_url=SETTINGS["llm"]["base_url"],
            timeout=_client_timeout(),
            max_retries=0,
        )
        prompt = (
            "Generate one read-only Cypher query for the question. "
            "Return only Cypher, without explanation. Never use CREATE, MERGE, DELETE, SET, REMOVE, DROP, LOAD CSV or CALL.\n"
            "Return nodes, relationships, or paths that carry the requested properties; do not return only scalar projections.\n"
            "Every node pattern MUST declare exactly one label copied EXACTLY from allowedLabels (preserve uppercase, underscore, spelling). "
            "NEVER invent, translate, pluralize, or guess labels from the question text — every label in the Cypher MUST appear verbatim in allowedLabels. "
            "If a concept has no exact match, pick the closest allowedLabels entry and reuse its exact spelling. "
            "Repeating a previously-bound variable is fine, but any new node variable must carry a label.\n"
            f"Schema:\n{req.schema}\n"
            f"Allowed labels: {req.allowedLabels}\n"
            f"Allowed relationships: {req.allowedRelationships}\n"
            f"Allowed properties: {req.allowedProperties}\n"
            f"Business rules:\n{req.businessRules}\n"
            f"Maximum path depth: {req.maxDepth}\n"
            f"Entity mentions that must be grounded to an allowed property (prefer name/title/code): {req.entityMentions}. User names may be abbreviations; unless an exact stored value is known from examples, use CONTAINS instead of equality.\n"
            f"Examples:\n{req.examples}\n"
            f"Question: {req.query}"
        )
        response = llm.chat.completions.create(
            model=SETTINGS["llm"]["model"],
            messages=[{"role": "user", "content": prompt}],
            temperature=SETTINGS["llm"]["temperature"],
            max_tokens=SETTINGS["llm"]["generation_max_tokens"],
        )
        candidate = _extract_code(response.choices[0].message.content or "", "Cypher")
        warnings = []
        try:
            cypher = _read_only(candidate, "Cypher")
        except ValueError as exc:
            # The bridge only generates candidates. Java performs the authoritative
            # schema validation, read-only Guard and the single allowed repair before execution.
            cypher = candidate
            warnings.append(str(exc))
        return {"cypher": cypher, "confidence": 0.0, "usedSchema": req.allowedLabels, "warnings": warnings}
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"Neo4j GraphRAG Cypher generation unavailable: {exc}") from exc


@app.post("/answer")
def answer(req: AnswerRequest):
    if not _llm_configured():
        raise HTTPException(status_code=503, detail="OpenAI provider or OPENAI_API_KEY is not configured")
    try:
        from openai import OpenAI

        client = OpenAI(
            api_key=os.environ["OPENAI_API_KEY"],
            base_url=SETTINGS["llm"]["base_url"],
            timeout=_client_timeout(),
            max_retries=0,
        )
        evidence_json = json.dumps(req.evidences[:20], ensure_ascii=False, default=str)
        response = client.chat.completions.create(
            model=SETTINGS["llm"]["model"],
            messages=[
                {"role": "system", "content": "你是多源RAG回答助手。只能依据给定证据作答，明确区分文档、结构化数据和图谱依据；证据不足时直说，不得编造。回答使用简体中文。"},
                {"role": "user", "content": f"问题：{req.question}\n\n证据：{evidence_json}"},
            ],
            temperature=SETTINGS["llm"]["temperature"],
            max_tokens=SETTINGS["llm"]["answer_max_tokens"],
        )
        return {"answer": response.choices[0].message.content or "", "warnings": []}
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"Answer generation unavailable: {exc}") from exc


@app.post("/governance/suggest")
def governance_suggest(req: GovernanceRequest):
    if not _llm_configured(): raise HTTPException(status_code=503, detail="LLM is not configured")
    prompt = f"""Analyze the real {req.sourceType} schema and propose a safe business-facing RAG policy.
Return JSON only with keys: summary, allowedLabels, allowedRelationships, allowedProperties,
tableWhitelist, documentation, businessRules, examples. examples must contain question plus sql or cypher.
Never invent schema names. Prefer a coherent business subgraph over opening unrelated technical metadata.
Generate at most {req.maxExamples} representative reviewed-candidate examples covering lookup, filtering,
aggregation or graph traversal as applicable.
Source description: {req.sourceDescription}
Schema:\n{req.schemaText}"""
    try:
        response = _openai_client(12).chat.completions.create(model=SETTINGS["llm"]["model"],
            messages=[{"role":"user","content":prompt}], temperature=SETTINGS["llm"]["temperature"],
            max_tokens=SETTINGS["llm"]["answer_max_tokens"])
        return _json_object(response.choices[0].message.content or "")
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"governance suggestion unavailable: {exc}") from exc


@app.post("/plan")
def plan(req: PlanRequest):
    if not _llm_configured(): raise HTTPException(status_code=503, detail="LLM is not configured")
    prompt = f"""Split the user question into source-specific retrieval questions. Return JSON only:
{{"DOCUMENT":"...","STRUCTURED_DATA":"...","GRAPH":"..."}}.
Omit sources that cannot contribute according to their capability description. Preserve the user's intent.
Question: {req.question}\nCapabilities: {json.dumps(req.capabilities, ensure_ascii=False)}"""
    try:
        response = _openai_client(12).chat.completions.create(model=SETTINGS["llm"]["model"],
            messages=[{"role":"user","content":prompt}], temperature=SETTINGS["llm"]["temperature"], max_tokens=1024)
        return _json_object(response.choices[0].message.content or "")
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"question planning unavailable: {exc}") from exc


@app.post("/repair")
def repair(req: RepairRequest):
    language = req.language.upper()
    if language == "SQL":
        context = json.dumps({
            "dialect": req.dialect or "mysql",
            "selectedTables": req.selectedTables,
            "minimalDdl": req.schemaText,
            "foreignKeys": req.foreignKeys,
            "relationships": req.relationships,
            "sampledValues": req.sampledValues,
            "referenceRows": req.referenceRows,
        }, ensure_ascii=False, default=str)
        prompt = f"""Repair this read-only MySQL query after EXPLAIN failed. Return only one corrected SQL statement.
Use only tables and columns in minimalDdl, only explicitly provided relationships, and only sampled/reference values.
Every CONDITIONAL relationship must include all listed conditions. Avoid cartesian products.
Never invent business mappings, schema identifiers, values or answers. Preserve AND semantics for simultaneous conditions.
For highest/best by a stated score use ORDER BY that score DESC. Do not use PostgreSQL DISTINCT ON.
Question: {req.question}
Verified context: {context}
Failed query: {req.query}
Error: {_short_error(req.error)}
Maximum rows: {req.maxRows}."""
    else:
        prompt = f"""Repair this read-only {language} query after EXPLAIN failed. Return only the corrected query.
Use only the supplied schema. Make exactly one statement. Do not use write operations.
The failed query used a label or relationship type that is NOT in the Schema below — open the Schema, find the closest matching entry, and copy its name EXACTLY (preserve uppercase, underscore, spelling).
NEVER invent, translate, pluralize, or guess labels from the question text; every label/relationship MUST appear verbatim in the Schema.
Every node pattern MUST declare exactly one label from the Schema — write (p:Person), never (p) or ().
Question: {req.question}\nSchema: {req.schemaText}\nFailed query: {req.query}\nError: {_short_error(req.error)}
Maximum rows: {req.maxRows}; maximum graph depth: {req.maxDepth}."""
    try:
        response = _openai_client(SETTINGS["llm"]["timeout"]).chat.completions.create(model=SETTINGS["llm"]["model"],
            messages=[{"role":"user","content":prompt}], temperature=SETTINGS["llm"]["temperature"],
            max_tokens=SETTINGS["llm"]["generation_max_tokens"])
        fixed = _read_only(_extract_code(response.choices[0].message.content or "", language), language)
        return {"query": fixed}
    except Exception as exc:
        raise HTTPException(status_code=503, detail=f"query repair unavailable: {exc}") from exc


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host=str(SETTINGS["server"]["host"]), port=int(SETTINGS["server"]["port"]))