@roll-agent/octopus-agent 0.1.2 → 0.2.0

package/src/octopus_skill/audit_logger.py

@@ -1,29 +0,0 @@
-from __future__ import annotations
-
-import hashlib
-import json
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Any
-
-
-@dataclass
-class AuditLogger:
-    path: Path
-    enabled: bool = True
-
-    def log(self, event: dict[str, Any]) -> None:
-        if not self.enabled:
-            return
-        self.path.parent.mkdir(parents=True, exist_ok=True)
-        record = {
-            "occurredAt": datetime.now(timezone.utc).isoformat(),
-            **event,
-        }
-        with self.path.open("a", encoding="utf-8") as file:
-            file.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
-
-
-def sql_hash(sql: str) -> str:
-    return hashlib.sha256(sql.encode("utf-8")).hexdigest()

package/src/octopus_skill/config.py

@@ -1,129 +0,0 @@
-from __future__ import annotations
-
-import os
-from dataclasses import dataclass, field
-from pathlib import Path
-
-
-PACKAGE_ROOT = Path.cwd()
-DEFAULT_SPONGE_MCP_BASE_URL = "https://sponge-mcp.duliday.com"
-DEFAULT_SPONGE_MCP_ACCESS_TOKEN = ""
-
-
-@dataclass(frozen=True)
-class SpongeMcpConfig:
-    base_url: str
-    access_token: str
-    timeout_ms: int
-
-
-@dataclass(frozen=True)
-class SchemaConfig:
-    cache_path: Path
-    refresh_on_start: bool
-    refresh_interval_minutes: int
-
-
-@dataclass(frozen=True)
-class EntityBindingConfig:
-    # Compute bindings and log them, but do not change SQL generation yet.
-    shadow_mode: bool
-    # Entity types whose bindings are allowed to influence SQL (live rollout).
-    # In shadow_mode this is ignored.
-    enabled_entity_types: tuple[str, ...]
-    min_fuzzy_score: float
-
-
-@dataclass(frozen=True)
-class SqlConfig:
-    max_repair_attempts: int
-    default_limit: int
-    max_limit: int
-
-
-@dataclass(frozen=True)
-class AuditConfig:
-    enabled: bool
-    log_path: Path
-
-
-@dataclass(frozen=True)
-class SamplingConfig:
-    max_tokens: int
-
-
-@dataclass(frozen=True)
-class AppConfig:
-    sponge_mcp: SpongeMcpConfig
-    schema: SchemaConfig
-    sql: SqlConfig
-    audit: AuditConfig
-    sampling: SamplingConfig = field(default_factory=lambda: SamplingConfig(max_tokens=4096))
-    binding: EntityBindingConfig | None = None
-
-
-def load_config() -> AppConfig:
-    return AppConfig(
-        sponge_mcp=SpongeMcpConfig(
-            base_url=os.getenv("SPONGE_MCP_BASE_URL", DEFAULT_SPONGE_MCP_BASE_URL).rstrip("/"),
-            access_token=os.getenv("SPONGE_MCP_ACCESS_TOKEN", DEFAULT_SPONGE_MCP_ACCESS_TOKEN),
-            timeout_ms=30000,
-        ),
-        schema=SchemaConfig(
-            cache_path=PACKAGE_ROOT / "data/sponge-schema-cache.json",
-            refresh_on_start=True,
-            refresh_interval_minutes=1440,
-        ),
-        sql=SqlConfig(
-            max_repair_attempts=2,
-            default_limit=50,
-            max_limit=500,
-        ),
-        audit=AuditConfig(
-            enabled=True,
-            log_path=PACKAGE_ROOT / "logs/octopus-sponge-audit.log",
-        ),
-        sampling=SamplingConfig(
-            max_tokens=_env_int("OCTOPUS_SAMPLING_MAX_TOKENS", default=4096),
-        ),
-        binding=None,
-    )
-
-
-def load_config_from_env() -> AppConfig:
-    return load_config()
-
-
-def _env_bool(name: str, *, default: bool) -> bool:
-    value = os.getenv(name)
-    if value is None or value == "":
-        return default
-    return value.strip().lower() in {"1", "true", "yes", "on"}
-
-
-def _env_csv(name: str, *, default: tuple[str, ...] = ()) -> tuple[str, ...]:
-    value = os.getenv(name)
-    if value is None:
-        return default
-    return tuple(item.strip() for item in value.split(",") if item.strip())
-
-
-def _env_float(name: str, *, default: float) -> float:
-    value = os.getenv(name)
-    if value is None or value == "":
-        return default
-    try:
-        return float(value)
-    except ValueError:
-        return default
-
-
-def _env_int(name: str, *, default: int) -> int:
-    value = os.getenv(name)
-    if value is None or value == "":
-        return default
-    try:
-        parsed = int(value)
-    except ValueError:
-        return default
-    return parsed if parsed > 0 else default

package/src/octopus_skill/context.py

@@ -1,17 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from uuid import uuid4
-
-
-@dataclass(frozen=True)
-class UserContext:
-    session_id: str
-
-
-def new_trace_id(prefix: str = "trace") -> str:
-    return f"{prefix}_{uuid4().hex}"
-
-
-def new_session_id() -> str:
-    return f"session_{uuid4().hex}"

package/src/octopus_skill/device_info.py

@@ -1,7 +0,0 @@
-from __future__ import annotations
-
-import platform
-
-
-def get_device_info() -> str:
-    return platform.platform(aliased=True, terse=True)

package/src/octopus_skill/exporter.py

@@ -1,66 +0,0 @@
-from __future__ import annotations
-
-import csv
-import os
-from dataclasses import dataclass
-from datetime import datetime
-from pathlib import Path
-from typing import Any
-from uuid import uuid4
-
-
-DEFAULT_EXPORT_DIR = Path("/tmp/octopus-agent-exports")
-
-
-@dataclass(frozen=True)
-class CsvExport:
-    path: Path
-    row_count: int
-    column_count: int
-
-
-def export_result_to_csv(result: dict[str, Any], *, prefix: str = "octopus_export") -> CsvExport | None:
-    rows = result.get("rows")
-    if not isinstance(rows, list) or not rows:
-        return None
-    columns = _column_names(result.get("columns"), rows)
-    if not columns:
-        return None
-    export_dir = Path(os.getenv("OCTOPUS_EXPORT_DIR") or DEFAULT_EXPORT_DIR)
-    export_dir.mkdir(parents=True, exist_ok=True)
-    path = export_dir / _csv_filename(prefix)
-    with path.open("w", newline="", encoding="utf-8-sig") as handle:
-        writer = csv.writer(handle)
-        writer.writerow(columns)
-        for row in rows:
-            writer.writerow([_cell_value(row, column) for column in columns])
-    return CsvExport(path=path, row_count=len(rows), column_count=len(columns))
-
-
-def _csv_filename(prefix: str) -> str:
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    suffix = uuid4().hex[:8]
-    safe_prefix = "".join(char if char.isalnum() or char in {"_", "-"} else "_" for char in prefix).strip("_")
-    return f"{safe_prefix or 'octopus_export'}_{timestamp}_{suffix}.csv"
-
-
-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 not isinstance(row, dict):
-        return str(row)
-    value = row.get(column, "")
-    if value is None:
-        return ""
-    return str(value)

package/src/octopus_skill/llm_result_renderer.py

@@ -1,101 +0,0 @@
-from __future__ import annotations
-
-import json
-from typing import Any
-
-from .llm_sql_generator import Sampler
-from .prompt_builder import filter_agent_instructions_by_stage, format_numbered_instruction_lines
-from .result_renderer import RESULT_ROW_DISPLAY_LIMIT, normalize_result_format, render_result
-
-
-RESULT_RENDER_SYSTEM_PROMPT = """你是丸子Agent（Octopus Agent）的查询结果渲染专家。
-业务规则的唯一来源是 payload 中的 mandatoryAgentInstructions；必须逐条遵守，不得使用训练记忆中的业务口径替代 schema 要求。
-根据用户问题、outputFormat 和 result 数据，输出最终用户可见答复。
-直接输出渲染后的文本，不要 JSON 包装，不要解释渲染过程。"""
-
-
-def render_result_with_llm(
-    *,
-    question: str,
-    result: dict[str, Any],
-    agent_instructions: list[dict[str, Any]],
-    sampler: Sampler,
-    result_format: str | None = None,
-) -> str:
-    output_format = normalize_result_format(question, result_format)
-    if not result.get("success", True):
-        return render_result(question, result, result_format)
-
-    if _should_render_detail_rows_locally(result):
-        local_format = "markdown" if output_format == "text" else output_format
-        return render_result(question, result, local_format)
-
-    instruction_lines = format_numbered_instruction_lines(
-        filter_agent_instructions_by_stage(agent_instructions, "after")
-    )
-    row_count = int(result.get("rowCount", len(result.get("rows", []))))
-    rows = result.get("rows", [])
-    sample = rows[:RESULT_ROW_DISPLAY_LIMIT] if isinstance(rows, list) else []
-    payload = {
-        "mandatoryAgentInstructions": instruction_lines,
-        "instruction": "按 mandatoryAgentInstructions 渲染 result，并匹配 outputFormat。",
-        "question": question,
-        "outputFormat": output_format,
-        "result": {
-            "rowCount": row_count,
-            "limited": bool(result.get("limited", False)),
-            "columns": result.get("columns", []),
-            "rows": sample,
-        },
-    }
-    try:
-        rendered = sampler(RESULT_RENDER_SYSTEM_PROMPT, json.dumps(payload, ensure_ascii=False)).strip()
-    except Exception:
-        return render_result(question, result, result_format)
-    if not rendered:
-        return render_result(question, result, result_format)
-    return rendered
-
-
-def _should_render_detail_rows_locally(result: dict[str, Any]) -> bool:
-    try:
-        row_count = int(result.get("rowCount", len(result.get("rows", []))))
-    except (TypeError, ValueError):
-        row_count = 0
-    rows = result.get("rows", [])
-    if row_count <= 0 or not isinstance(rows, list) or not rows:
-        return False
-    return True
-
-
-def _result_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 _is_aggregate_column_name(column: str) -> bool:
-    lowered = column.lower()
-    aggregate_terms = (
-        "count",
-        "sum",
-        "avg",
-        "min",
-        "max",
-        "数量",
-        "总数",
-        "人数",
-        "次数",
-        "平均",
-        "最大",
-        "最小",
-        "合计",
-    )
-    return any(term in lowered or term in column for term in aggregate_terms)

package/src/octopus_skill/llm_sql_generator.py

@@ -1,197 +0,0 @@
-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_expert_review_prompt,
-    build_sql_repair_prompt,
-)
-from .schema_context_retriever import SchemaContext
-from .sql_generator import GeneratedSql, SqlGenerationError, contains_forbidden_dml_keyword
-
-
-Sampler = Callable[[str, str], str]
-
-
-@dataclass
-class SamplingSqlGenerator:
-    sampler: Sampler
-    default_limit: int
-    max_limit: int
-    confirmed_query_plan: dict[str, Any] | None = None
-
-    def generate(
-        self,
-        *,
-        question: str,
-        schema_context: SchemaContext,
-    ) -> GeneratedSql:
-        prompt = build_nl2sql_prompt(
-            question=question,
-            schema_context=schema_context,
-            default_limit=self.default_limit,
-            max_limit=self.max_limit,
-            confirmed_query_plan=self.confirmed_query_plan,
-        )
-        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,
-    ) -> GeneratedSql:
-        prompt = build_sql_repair_prompt(
-            question=question,
-            original_sql=original_sql,
-            validation_error=validation_error,
-            schema_context=schema_context,
-            default_limit=self.default_limit,
-            max_limit=self.max_limit,
-            confirmed_query_plan=self.confirmed_query_plan,
-        )
-        return _parse_generated_sql(self.sampler(NL2SQL_SYSTEM_PROMPT, prompt))
-
-    def review_sql(
-        self,
-        *,
-        question: str,
-        sql: str,
-        schema_context: SchemaContext,
-    ) -> "SqlExpertReview":
-        prompt = build_sql_expert_review_prompt(
-            question=question,
-            sql=sql,
-            schema_context=schema_context,
-            default_limit=self.default_limit,
-            max_limit=self.max_limit,
-            confirmed_query_plan=self.confirmed_query_plan,
-        )
-        return _parse_sql_expert_review(self.sampler(NL2SQL_SYSTEM_PROMPT, prompt))
-
-
-@dataclass(frozen=True)
-class SqlExpertReview:
-    executable: bool
-    issues: list[dict[str, str]]
-    repaired_sql: str | None = None
-
-
-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
-        _require_select_sql(sql)
-        return GeneratedSql(sql=sql, tables=[], placeholders=[])
-    sql = payload.get("sql")
-    if isinstance(sql, str) and not sql.strip():
-        clarification = payload.get("clarification") if isinstance(payload.get("clarification"), dict) else None
-        return GeneratedSql(
-            sql="",
-            tables=[],
-            placeholders=[],
-            used_columns=[],
-            used_joins=[],
-            metric_refs=[],
-            clarification=clarification,
-        )
-    if not isinstance(sql, str):
-        fallback_sql = _extract_sql_text(text)
-        if fallback_sql is None:
-            raise SqlGenerationError("LLM output must contain sql")
-        _require_select_sql(fallback_sql)
-        return GeneratedSql(sql=fallback_sql, tables=[], placeholders=[])
-    sql = sql.strip()
-    _require_select_sql(sql)
-    tables = payload.get("tables") if isinstance(payload.get("tables"), list) else []
-    placeholders = payload.get("placeholders") if isinstance(payload.get("placeholders"), list) else []
-    used_columns = payload.get("usedColumns") if isinstance(payload.get("usedColumns"), list) else []
-    used_joins = payload.get("usedJoins") if isinstance(payload.get("usedJoins"), list) else []
-    metric_refs = payload.get("metricRefs") if isinstance(payload.get("metricRefs"), list) else []
-    clarification = payload.get("clarification") if isinstance(payload.get("clarification"), dict) else None
-    return GeneratedSql(
-        sql=sql,
-        tables=[str(table) for table in tables],
-        placeholders=[str(placeholder) for placeholder in placeholders],
-        used_columns=[str(column) for column in used_columns],
-        used_joins=[str(join) for join in used_joins],
-        metric_refs=[str(metric) for metric in metric_refs],
-        clarification=clarification,
-    )
-
-
-def _parse_sql_expert_review(text: str) -> SqlExpertReview:
-    payload = _load_json_object(text)
-    executable = bool(payload.get("executable", False))
-    issues_payload = payload.get("issues")
-    issues: list[dict[str, str]] = []
-    if isinstance(issues_payload, list):
-        for item in issues_payload:
-            if not isinstance(item, dict):
-                continue
-            issues.append(
-                {
-                    "code": str(item.get("code") or ""),
-                    "message": str(item.get("message") or ""),
-                    "repairHint": str(item.get("repairHint") or ""),
-                }
-            )
-    repaired_sql = payload.get("repairedSql")
-    if isinstance(repaired_sql, str) and repaired_sql.strip():
-        repaired_sql = repaired_sql.strip()
-        _require_select_sql(repaired_sql)
-    else:
-        repaired_sql = None
-    return SqlExpertReview(executable=executable, issues=issues, repaired_sql=repaired_sql)
-
-
-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
-
-
-def _require_select_sql(sql: str) -> None:
-    normalized = sql.strip().lower()
-    if not normalized.startswith("select "):
-        raise SqlGenerationError("LLM SQL must be SELECT only")
-    if contains_forbidden_dml_keyword(sql):
-        raise SqlGenerationError("LLM SQL contains a forbidden keyword")

package/src/octopus_skill/main.py

@@ -1,26 +0,0 @@
-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-agent")
-    parser.add_argument("question")
-    parser.add_argument("--session-id")
-    args = parser.parse_args()
-
-    config = load_config_from_env()
-    agent = SpongeAgent.from_config(config)
-    print(
-        agent.answer(
-            question=args.question,
-            session_id=args.session_id,
-        )
-    )
-
-
-if __name__ == "__main__":
-    main()

package/src/octopus_skill/mcp_client.py

@@ -1,114 +0,0 @@
-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,
-        session_id: str | None = None,
-        schema_version: str | None = None,
-    ) -> dict[str, Any]:
-        payload = {
-            "traceId": trace_id,
-            "sessionId": session_id,
-            "schemaVersion": schema_version,
-        }
-        return self._post("get_schema", payload)
-
-    def search_entities(
-        self,
-        *,
-        trace_id: str,
-        session_id: str | None = None,
-        entity_types: list[str] | tuple[str, ...] | None = None,
-        name: str | None = None,
-    ) -> dict[str, Any]:
-        payload = {
-            "traceId": trace_id,
-            "sessionId": session_id,
-            "entityTypes": list(entity_types) if entity_types else None,
-            "name": name,
-        }
-        return self._post("search_entities", payload)
-
-    def validate_sql(
-        self,
-        *,
-        trace_id: str,
-        session_id: str,
-        question: str,
-        sql: str,
-    ) -> dict[str, Any]:
-        payload = {
-            "traceId": trace_id,
-            "sessionId": session_id,
-            "question": question,
-            "sql": sql,
-        }
-        return self._post("validate_sql", payload)
-
-    def execute_sql(
-        self,
-        *,
-        trace_id: str,
-        session_id: str,
-        question: str,
-        sql: str,
-        timeout_ms: int | None = None,
-    ) -> dict[str, Any]:
-        payload = {
-            "traceId": trace_id,
-            "sessionId": session_id,
-            "question": question,
-            "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