@roll-agent/octopus-agent 0.0.11 → 0.1.1

package/src/octopus_skill/config.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 
 
@@ -24,15 +24,6 @@ class SchemaConfig:
     refresh_interval_minutes: int
 
 
-@dataclass(frozen=True)
-class DictionaryConfig:
-    cache_path: Path
-    geo_alias_path: Path
-    refresh_on_start: bool
-    refresh_interval_minutes: int
-    entity_types: tuple[str, ...]
-
-
 @dataclass(frozen=True)
 class EntityBindingConfig:
     # Compute bindings and log them, but do not change SQL generation yet.
@@ -56,13 +47,18 @@ class AuditConfig:
     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
-    dictionary: DictionaryConfig | None = None
+    sampling: SamplingConfig = field(default_factory=lambda: SamplingConfig(max_tokens=4096))
     binding: EntityBindingConfig | None = None
 
 
@@ -78,29 +74,19 @@ def load_config() -> AppConfig:
             refresh_on_start=True,
             refresh_interval_minutes=1440,
         ),
-        dictionary=DictionaryConfig(
-            cache_path=PACKAGE_ROOT / "data/sponge-entity-cache.json",
-            geo_alias_path=PACKAGE_ROOT / "data/geo-aliases.json",
-            refresh_on_start=True,
-            refresh_interval_minutes=1440,
-            entity_types=("brand", "province", "city", "region", "project"),
-        ),
         sql=SqlConfig(
             max_repair_attempts=2,
-            default_limit=200,
+            default_limit=50,
             max_limit=500,
         ),
         audit=AuditConfig(
             enabled=True,
             log_path=PACKAGE_ROOT / "logs/octopus-sponge-audit.log",
         ),
-        binding=EntityBindingConfig(
-            shadow_mode=_env_bool("OCTOPUS_BINDING_SHADOW", default=False),
-            enabled_entity_types=_env_csv(
-                "OCTOPUS_BINDING_ENABLED", default=("brand", "location", "project")
-            ),
-            min_fuzzy_score=_env_float("OCTOPUS_BINDING_MIN_SCORE", default=0.6),
+        sampling=SamplingConfig(
+            max_tokens=_env_int("OCTOPUS_SAMPLING_MAX_TOKENS", default=4096),
         ),
+        binding=None,
     )
 
 
@@ -130,3 +116,14 @@ def _env_float(name: str, *, default: float) -> float:
         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/exporter.py

@@ -0,0 +1,66 @@
+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

@@ -0,0 +1,101 @@
+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

@@ -5,9 +5,14 @@ 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 .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
+from .sql_generator import GeneratedSql, SqlGenerationError, contains_forbidden_dml_keyword
 
 
 Sampler = Callable[[str, str], str]
@@ -18,7 +23,7 @@ class SamplingSqlGenerator:
     sampler: Sampler
     default_limit: int
     max_limit: int
-    resolved_entities: dict[str, Any] | None = None
+    confirmed_query_plan: dict[str, Any] | None = None
 
     def generate(
         self,
@@ -31,7 +36,7 @@ class SamplingSqlGenerator:
             schema_context=schema_context,
             default_limit=self.default_limit,
             max_limit=self.max_limit,
-            resolved_entities=self.resolved_entities,
+            confirmed_query_plan=self.confirmed_query_plan,
         )
         return _parse_generated_sql(self.sampler(NL2SQL_SYSTEM_PROMPT, prompt))
 
@@ -50,10 +55,34 @@ class SamplingSqlGenerator:
             schema_context=schema_context,
             default_limit=self.default_limit,
             max_limit=self.max_limit,
-            resolved_entities=self.resolved_entities,
+            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:
@@ -66,7 +95,16 @@ def _parse_generated_sql(text: str) -> GeneratedSql:
         return GeneratedSql(sql=sql, tables=[], placeholders=[])
     sql = payload.get("sql")
     if isinstance(sql, str) and not sql.strip():
-        return GeneratedSql(sql="", tables=[], placeholders=[], used_columns=[], used_joins=[], metric_refs=[])
+        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:
@@ -80,6 +118,7 @@ def _parse_generated_sql(text: str) -> GeneratedSql:
     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],
@@ -87,9 +126,35 @@ def _parse_generated_sql(text: str) -> GeneratedSql:
         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("```"):
@@ -128,5 +193,5 @@ 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 re.search(r"\b(insert|update|delete|drop|alter|create|truncate)\b", normalized):
+    if contains_forbidden_dml_keyword(sql):
         raise SqlGenerationError("LLM SQL contains a forbidden keyword")

package/src/octopus_skill/mcp_client.py

@@ -31,21 +31,21 @@ class SpongeMcpClient:
         }
         return self._post("get_schema", payload)
 
-    def get_entity_dictionary(
+    def search_entities(
         self,
         *,
         trace_id: str,
         session_id: str | None = None,
-        dict_version: 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,
-            "dictVersion": dict_version,
             "entityTypes": list(entity_types) if entity_types else None,
+            "name": name,
         }
-        return self._post("get_entity_dictionary", payload)
+        return self._post("search_entities", payload)
 
     def validate_sql(
         self,