@roll-agent/octopus-agent 0.0.1

package/src/octopus_sponge_skill/llm_sql_generator.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from .prompt_builder import NL2SQL_SYSTEM_PROMPT, build_nl2sql_prompt, build_sql_repair_prompt
+from .schema_context_retriever import SchemaContext
+from .sql_generator import GeneratedSql, SqlGenerationError
+
+
+Sampler = Callable[[str, str], str]
+
+
+@dataclass
+class SamplingSqlGenerator:
+    sampler: Sampler
+    default_limit: int
+    max_limit: int
+
+    def generate(
+        self,
+        *,
+        question: str,
+        schema_context: SchemaContext,
+        scope_hints: dict[str, Any],
+    ) -> GeneratedSql:
+        prompt = build_nl2sql_prompt(
+            question=question,
+            schema_context=schema_context,
+            scope_hints=scope_hints,
+            default_limit=self.default_limit,
+            max_limit=self.max_limit,
+        )
+        return _parse_generated_sql(self.sampler(NL2SQL_SYSTEM_PROMPT, prompt))
+
+    def repair(
+        self,
+        *,
+        question: str,
+        original_sql: str,
+        validation_error: dict[str, Any],
+        schema_context: SchemaContext,
+        scope_hints: dict[str, Any],
+    ) -> GeneratedSql:
+        prompt = build_sql_repair_prompt(
+            question=question,
+            original_sql=original_sql,
+            validation_error=validation_error,
+            schema_context=schema_context,
+            scope_hints=scope_hints,
+            max_limit=self.max_limit,
+        )
+        return _parse_generated_sql(self.sampler(NL2SQL_SYSTEM_PROMPT, prompt))
+
+
+def _parse_generated_sql(text: str) -> GeneratedSql:
+    try:
+        payload = _load_json_object(text)
+    except SqlGenerationError:
+        sql = _extract_sql_text(text)
+        if sql is None:
+            raise
+        return GeneratedSql(sql=sql, tables=[], placeholders=[])
+    sql = payload.get("sql")
+    if not isinstance(sql, str) or not sql.strip():
+        fallback_sql = _extract_sql_text(text)
+        if fallback_sql is None:
+            raise SqlGenerationError("LLM output must contain sql")
+        return GeneratedSql(sql=fallback_sql, tables=[], placeholders=[])
+    tables = payload.get("tables") if isinstance(payload.get("tables"), list) else []
+    placeholders = payload.get("placeholders") if isinstance(payload.get("placeholders"), list) else []
+    return GeneratedSql(
+        sql=sql.strip(),
+        tables=[str(table) for table in tables],
+        placeholders=[str(placeholder) for placeholder in placeholders],
+    )
+
+
+def _load_json_object(text: str) -> dict[str, Any]:
+    cleaned = text.strip()
+    if cleaned.startswith("```"):
+        cleaned = cleaned.strip("`")
+        if cleaned.lower().startswith("json"):
+            cleaned = cleaned[4:].strip()
+    try:
+        payload = json.loads(cleaned)
+    except json.JSONDecodeError as exc:
+        start = cleaned.find("{")
+        end = cleaned.rfind("}")
+        if start == -1 or end == -1 or end <= start:
+            raise SqlGenerationError("LLM output is not valid JSON") from exc
+        payload = json.loads(cleaned[start : end + 1])
+    if not isinstance(payload, dict):
+        raise SqlGenerationError("LLM output must be a JSON object")
+    return payload
+
+
+def _extract_sql_text(text: str) -> str | None:
+    cleaned = text.strip()
+    fenced = re.search(r"```(?:sql)?\s*(.*?)```", cleaned, flags=re.IGNORECASE | re.DOTALL)
+    if fenced:
+        cleaned = fenced.group(1).strip()
+    match = re.search(r"\bselect\b.*", cleaned, flags=re.IGNORECASE | re.DOTALL)
+    if not match:
+        return None
+    sql = match.group(0).strip()
+    sql = re.split(r"\n\s*(?:解释|说明|备注)[:：]", sql, maxsplit=1)[0].strip()
+    if sql.endswith(";"):
+        sql = sql[:-1].strip()
+    return sql

package/src/octopus_sponge_skill/main.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import argparse
+
+from .agent import SpongeAgent
+from .config import load_config_from_env
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(prog="octopus-sponge")
+    parser.add_argument("question")
+    parser.add_argument("--user-id", type=int, required=True)
+    parser.add_argument("--session-id")
+    parser.add_argument("--device-id")
+    args = parser.parse_args()
+
+    config = load_config_from_env()
+    agent = SpongeAgent.from_config(config)
+    print(
+        agent.answer(
+            question=args.question,
+            user_id=args.user_id,
+            session_id=args.session_id,
+            device_id=args.device_id,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

package/src/octopus_sponge_skill/mcp_client.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import json
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+from typing import Any
+
+
+class SpongeMcpError(RuntimeError):
+    pass
+
+
+@dataclass
+class SpongeMcpClient:
+    base_url: str
+    access_token: str
+    timeout_ms: int = 30000
+
+    def get_schema(
+        self,
+        *,
+        trace_id: str,
+        schema_version: str | None,
+        refresh_reason: str,
+    ) -> dict[str, Any]:
+        payload = {
+            "traceId": trace_id,
+            "schemaVersion": schema_version,
+            "refreshReason": refresh_reason,
+        }
+        return self._post("get_schema", payload)
+
+    def get_allowed_context(
+        self,
+        *,
+        user_id: int,
+        trace_id: str,
+        session_id: str,
+        question: str,
+        device_id: str | None = None,
+    ) -> dict[str, Any]:
+        payload = {
+            "userId": user_id,
+            "traceId": trace_id,
+            "sessionId": session_id,
+            "deviceId": device_id,
+            "question": question,
+        }
+        return self._post("get_allowed_context", payload)
+
+    def validate_sql(
+        self,
+        *,
+        user_id: int,
+        trace_id: str,
+        session_id: str,
+        sql: str,
+    ) -> dict[str, Any]:
+        payload = {
+            "userId": user_id,
+            "traceId": trace_id,
+            "sessionId": session_id,
+            "sql": sql,
+        }
+        return self._post("validate_sql", payload)
+
+    def execute_sql(
+        self,
+        *,
+        user_id: int,
+        trace_id: str,
+        session_id: str,
+        sql: str,
+        timeout_ms: int | None = None,
+    ) -> dict[str, Any]:
+        payload = {
+            "userId": user_id,
+            "traceId": trace_id,
+            "sessionId": session_id,
+            "sql": sql,
+            "timeoutMs": timeout_ms or self.timeout_ms,
+        }
+        return self._post("execute_sql", payload)
+
+    def _post(self, tool_name: str, payload: dict[str, Any]) -> dict[str, Any]:
+        body = json.dumps({key: value for key, value in payload.items() if value is not None}).encode(
+            "utf-8"
+        )
+        request = urllib.request.Request(
+            url=f"{self.base_url}/mcp/tools/{tool_name}",
+            data=body,
+            method="POST",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.access_token}",
+            },
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=self.timeout_ms / 1000) as response:
+                data = json.loads(response.read().decode("utf-8"))
+        except urllib.error.HTTPError as exc:
+            message = exc.read().decode("utf-8", errors="replace")
+            if tool_name == "validate_sql":
+                try:
+                    error_payload = json.loads(message)
+                except json.JSONDecodeError:
+                    error_payload = None
+                if isinstance(error_payload, dict) and isinstance(error_payload.get("error"), dict):
+                    return {"valid": False, "error": error_payload["error"]}
+            raise SpongeMcpError(f"MCP {tool_name} failed: HTTP {exc.code} {message}") from exc
+        except urllib.error.URLError as exc:
+            raise SpongeMcpError(f"MCP {tool_name} failed: {exc.reason}") from exc
+        if not isinstance(data, dict):
+            raise SpongeMcpError(f"MCP {tool_name} returned non-object JSON")
+        return data

package/src/octopus_sponge_skill/prompt_builder.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from .schema_context_retriever import SchemaContext
+
+
+NL2SQL_SYSTEM_PROMPT = """你是 Octopus 的 Sponge 数据查询 SQL 生成器。
+你只能基于提供的 schema context 生成 MySQL SELECT 查询。
+SQL 必须使用 :scope.project_ids 和 :scope.brand_ids 表示用户可访问范围。
+不得硬编码真实 projectId 或 brandId。
+不得查询 auth 库、密码、令牌、密钥、凭证字段。
+不得生成 INSERT、UPDATE、DELETE、DDL 或多语句 SQL。
+必须包含 LIMIT，且 LIMIT 不超过 500。
+只输出 JSON，不输出解释文本。"""
+
+
+def build_nl2sql_prompt(
+    *,
+    question: str,
+    schema_context: SchemaContext,
+    scope_hints: dict[str, Any],
+    default_limit: int,
+    max_limit: int,
+) -> str:
+    payload = {
+        "question": question,
+        "scopeHints": scope_hints,
+        "limits": {"defaultLimit": default_limit, "maxLimit": max_limit},
+        "schemaContext": {
+            "tables": schema_context.tables,
+            "joins": schema_context.joins,
+            "glossary": schema_context.glossary,
+            "metrics": schema_context.metrics,
+            "examples": schema_context.examples,
+        },
+        "outputSchema": {
+            "sql": "SELECT ... LIMIT 100",
+            "tables": ["table_name"],
+            "placeholders": [":scope.project_ids", ":scope.brand_ids"],
+        },
+    }
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+
+
+def build_sql_repair_prompt(
+    *,
+    question: str,
+    original_sql: str,
+    validation_error: dict[str, Any],
+    schema_context: SchemaContext,
+    scope_hints: dict[str, Any],
+    max_limit: int,
+) -> str:
+    payload = {
+        "task": "repair_sql",
+        "question": question,
+        "originalSql": original_sql,
+        "validationError": validation_error,
+        "scopeHints": scope_hints,
+        "limits": {"maxLimit": max_limit},
+        "schemaContext": {
+            "tables": schema_context.tables,
+            "joins": schema_context.joins,
+            "glossary": schema_context.glossary,
+            "metrics": schema_context.metrics,
+            "examples": schema_context.examples,
+        },
+        "outputSchema": {
+            "sql": "SELECT ... LIMIT 100",
+            "tables": ["table_name"],
+            "placeholders": [":scope.project_ids", ":scope.brand_ids"],
+        },
+    }
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+

package/src/octopus_sponge_skill/result_renderer.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import html
+import json
+from typing import Any
+
+
+SUPPORTED_RESULT_FORMATS = {"text", "yaml", "md", "markdown", "html"}
+
+
+def normalize_result_format(question: str, result_format: str | None = None) -> str:
+    if result_format:
+        normalized = result_format.lower().strip()
+        if normalized not in SUPPORTED_RESULT_FORMATS:
+            raise ValueError(f"Unsupported resultFormat: {result_format}")
+        return "markdown" if normalized == "md" else normalized
+    lowered = question.lower()
+    if "yaml" in lowered:
+        return "yaml"
+    if "html" in lowered:
+        return "html"
+    if "markdown" in lowered or " md" in lowered or "表格" in question:
+        return "markdown"
+    return "text"
+
+
+def render_result(question: str, result: dict[str, Any], result_format: str | None = None) -> str:
+    output_format = normalize_result_format(question, result_format)
+    if not result.get("success", True):
+        return _format_failure(output_format)
+    row_count = int(result.get("rowCount", len(result.get("rows", []))))
+    rows = result.get("rows", [])
+    columns = result.get("columns", [])
+    if row_count == 0:
+        return _format_empty(output_format)
+    sample = rows[:5] if isinstance(rows, list) else []
+    if output_format == "yaml":
+        return _render_yaml(question, row_count, bool(result.get("limited", False)), sample)
+    if output_format == "markdown":
+        return _render_markdown(question, row_count, bool(result.get("limited", False)), sample, columns)
+    if output_format == "html":
+        return _render_html(question, row_count, bool(result.get("limited", False)), sample, columns)
+    return _render_text(question, row_count, bool(result.get("limited", False)), sample, columns)
+
+
+def _format_failure(output_format: str) -> str:
+    if output_format == "yaml":
+        return "success: false\nmessage: 查询失败。\n"
+    if output_format == "markdown":
+        return "查询失败。"
+    if output_format == "html":
+        return "<p>查询失败。</p>"
+    return "查询失败。"
+
+
+def _format_empty(output_format: str) -> str:
+    if output_format == "yaml":
+        return "success: true\nrowCount: 0\nrows: []\nmessage: 没有查询到符合条件的数据。\n"
+    if output_format == "markdown":
+        return "返回 0 行。\n\n没有查询到符合条件的数据。"
+    if output_format == "html":
+        return "<p>没有查询到符合条件的数据。</p>"
+    return "没有查询到符合条件的数据。"
+
+
+def _render_yaml(question: str, row_count: int, limited: bool, rows: list[Any]) -> str:
+    lines = [
+        f"rowCount: {row_count}",
+        f"limited: {_yaml_bool(limited)}",
+        "rows:",
+    ]
+    for row in rows:
+        lines.append("  -")
+        if isinstance(row, dict):
+            for key, value in row.items():
+                lines.append(f"    {key}: {json.dumps(value, ensure_ascii=False)}")
+        else:
+            lines.append(f"    value: {json.dumps(row, ensure_ascii=False)}")
+    return "\n".join(lines) + "\n"
+
+
+def _render_markdown(
+    question: str,
+    row_count: int,
+    limited: bool,
+    rows: list[Any],
+    columns: Any,
+) -> str:
+    table_columns = _column_names(columns, rows)
+    lines = [f"返回 {row_count} 行。", f"limited: {str(limited).lower()}"]
+    if not table_columns:
+        lines.extend(["", "```json", json.dumps(rows, ensure_ascii=False, indent=2), "```"])
+        return "\n".join(lines)
+    lines.extend(["", "|" + "|".join(_escape_markdown(column) for column in table_columns) + "|"])
+    lines.append("|" + "|".join("---" for _ in table_columns) + "|")
+    for row in rows:
+        lines.append("|" + "|".join(_escape_markdown(_cell_value(row, column)) for column in table_columns) + "|")
+    return "\n".join(lines)
+
+
+def _render_html(
+    question: str,
+    row_count: int,
+    limited: bool,
+    rows: list[Any],
+    columns: Any,
+) -> str:
+    table_columns = _column_names(columns, rows)
+    parts = [
+        "<section>",
+        f"<p>返回 {row_count} 行。limited: {str(limited).lower()}</p>",
+    ]
+    if table_columns:
+        parts.append("<table>")
+        parts.append("<thead><tr>" + "".join(f"<th>{html.escape(column)}</th>" for column in table_columns) + "</tr></thead>")
+        parts.append("<tbody>")
+        for row in rows:
+            parts.append("<tr>" + "".join(f"<td>{html.escape(_cell_value(row, column))}</td>" for column in table_columns) + "</tr>")
+        parts.append("</tbody></table>")
+    else:
+        parts.append(f"<pre>{html.escape(json.dumps(rows, ensure_ascii=False, indent=2))}</pre>")
+    parts.append("</section>")
+    return "".join(parts)
+
+
+def _render_text(
+    question: str,
+    row_count: int,
+    limited: bool,
+    rows: list[Any],
+    columns: Any,
+) -> str:
+    visible_columns = _text_columns(question, columns, rows)
+    if not visible_columns:
+        return f"查询到 {row_count} 条结果。"
+
+    noun = _result_noun(question)
+    lines = [f"查询到 {row_count} 个{noun}：", ""]
+    include_id = _question_requests_id(question)
+    for index, row in enumerate(rows, start=1):
+        text = _format_text_row(row, visible_columns, include_id)
+        if text:
+            lines.append(f"{index}. {text}")
+    if limited:
+        lines.append("")
+        lines.append("结果较多，已按当前限制返回部分数据。")
+    return "\n".join(lines)
+
+
+def _text_columns(question: str, columns: Any, rows: list[Any]) -> list[str]:
+    names = _column_names(columns, rows)
+    if not names:
+        return []
+    allow_id = _question_requests_id(question)
+    visible = [name for name in names if allow_id or not _is_id_column(name)]
+    name_columns = [name for name in visible if _is_name_column(name)]
+    if allow_id:
+        return visible
+    return name_columns or visible
+
+
+def _format_text_row(row: Any, columns: list[str], include_id: bool) -> str:
+    if not isinstance(row, dict):
+        return str(row)
+    name_column = next((column for column in columns if _is_name_column(column)), None)
+    id_column = next((column for column in columns if _is_id_column(column)), None)
+    name_value = _cell_value(row, name_column) if name_column else ""
+    id_value = _cell_value(row, id_column) if id_column else ""
+    if include_id and name_value and id_value:
+        return f"{name_value}（ID：{id_value}）"
+    values = [_cell_value(row, column) for column in columns]
+    values = [value for value in values if value]
+    return "，".join(values)
+
+
+def _question_requests_id(question: str) -> bool:
+    lowered = question.lower()
+    return "id" in lowered or "编号" in question or "主键" in question
+
+
+def _is_id_column(column: str) -> bool:
+    lowered = column.lower()
+    return lowered == "id" or lowered.endswith("_id") or lowered.endswith(".id")
+
+
+def _is_name_column(column: str) -> bool:
+    lowered = column.lower()
+    return lowered == "name" or lowered.endswith("_name") or lowered.endswith(".name") or "名称" in column
+
+
+def _result_noun(question: str) -> str:
+    if "品牌" in question:
+        return "品牌"
+    if "岗位" in question:
+        return "岗位"
+    if "门店" in question:
+        return "门店"
+    if "项目" in question:
+        return "项目"
+    return "结果"
+
+
+def _column_names(columns: Any, rows: list[Any]) -> list[str]:
+    names: list[str] = []
+    if isinstance(columns, list):
+        for column in columns:
+            if isinstance(column, dict) and column.get("name"):
+                names.append(str(column["name"]))
+            elif isinstance(column, str):
+                names.append(column)
+    if not names and rows and isinstance(rows[0], dict):
+        names = [str(key) for key in rows[0].keys()]
+    return names
+
+
+def _cell_value(row: Any, column: str) -> str:
+    if isinstance(row, dict):
+        return str(row.get(column, ""))
+    return str(row)
+
+
+def _escape_markdown(value: str) -> str:
+    return value.replace("\\", "\\\\").replace("|", "\\|").replace("\n", " ")
+
+
+def _yaml_bool(value: bool) -> str:
+    return "true" if value else "false"