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": 2048,
        "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 = ""
    documentation: str = ""
    examples: list[Any] = Field(default_factory=list)
    tableWhitelist: 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
    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=timeout or SETTINGS["llm"]["timeout"], max_retries=0)


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]


@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:
        import pandas as pd
        from openai import OpenAI
        from vanna.base import VannaBase
        from vanna.openai import OpenAI_Chat

        class RequestContext(VannaBase):
            def __init__(self):
                self.ddl = [req.ddl] if req.ddl else []
                self.documentation = [req.documentation] if req.documentation else []
                self.examples = req.examples

            def get_related_ddl(self, question: str, **kwargs) -> list:
                return self.ddl

            def get_related_documentation(self, question: str, **kwargs) -> list:
                return self.documentation

            def get_similar_question_sql(self, question: str, **kwargs) -> list:
                return self.examples

            def generate_embedding(self, data: str, **kwargs) -> list[float]:
                return []

            def add_ddl(self, ddl: str, **kwargs) -> str:
                self.ddl.append(ddl)
                return str(len(self.ddl))

            def add_documentation(self, documentation: str, **kwargs) -> str:
                self.documentation.append(documentation)
                return str(len(self.documentation))

            def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
                self.examples.append({"question": question, "sql": sql})
                return str(len(self.examples))

            def get_training_data(self, **kwargs) -> pd.DataFrame:
                return pd.DataFrame()

            def remove_training_data(self, id: str, **kwargs) -> bool:
                return False

        class Vanna(RequestContext, OpenAI_Chat):
            def __init__(self):
                RequestContext.__init__(self)
                client = OpenAI(
                    api_key=os.environ["OPENAI_API_KEY"],
                    base_url=SETTINGS["llm"]["base_url"],
                    timeout=SETTINGS["llm"]["timeout"],
                    max_retries=0,
                )
                OpenAI_Chat.__init__(self, client=client, config={
                    "model": SETTINGS["llm"]["model"],
                    "temperature": SETTINGS["llm"]["temperature"],
                })

            def submit_prompt(self, prompt, **kwargs) -> str:
                response = self.client.chat.completions.create(
                    model=SETTINGS["llm"]["model"],
                    messages=prompt,
                    temperature=SETTINGS["llm"]["temperature"],
                    max_tokens=SETTINGS["llm"]["generation_max_tokens"],
                )
                return response.choices[0].message.content or ""

        vn = Vanna()
        prompt = [{
            "role": "system",
            "content": (
                f"You generate exactly one final read-only {req.dialect} SQL query. "
                "Return SQL only, without comments, explanation or intermediate queries. "
                "Use only the supplied schema. Never generate DML or DDL. "
                f"Limit results to at most {req.maxRows} rows."
            ),
        }, {
            "role": "user",
            "content": (
                f"Schema:\n{req.ddl}\nDocumentation:\n{req.documentation}\n"
                f"Allowed tables: {req.tableWhitelist}\nEntity mentions that must be grounded to real columns: {req.entityMentions}. Unless an exact stored value is supplied by examples, use LIKE for abbreviated entity names.\nExamples: {req.examples}\nQuestion: {req.query}"
            ),
        }]
        sql = _read_only(_extract_code(vn.submit_prompt(prompt), "SQL"), "SQL")
        return {"sql": sql, "confidence": 0.0, "usedContext": ["ddl"] if req.ddl else [], "warnings": []}
    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=SETTINGS["llm"]["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"],
        )
        cypher = _read_only(_extract_code(response.choices[0].message.content or "", "Cypher"), "Cypher")
        return {"cypher": cypher, "confidence": 0.0, "usedSchema": req.allowedLabels, "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=SETTINGS["llm"]["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()
    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"]))