@roll-agent/octopus-agent 0.0.10 → 0.1.0

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,13 +24,21 @@ class SchemaConfig:
     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
-    query_cache_path: Path | None = None
-    query_cache_ttl_minutes: int = 30
 
 
 @dataclass(frozen=True)
@@ -39,12 +47,19 @@ 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
+    sampling: SamplingConfig = field(default_factory=lambda: SamplingConfig(max_tokens=4096))
+    binding: EntityBindingConfig | None = None
 
 
 def load_config() -> AppConfig:
@@ -63,13 +78,15 @@ def load_config() -> AppConfig:
             max_repair_attempts=2,
             default_limit=50,
             max_limit=500,
-            query_cache_path=_optional_path("OCTOPUS_QUERY_SQL_CACHE_PATH"),
-            query_cache_ttl_minutes=30,
         ),
         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,
     )
 
 
@@ -77,8 +94,36 @@ def load_config_from_env() -> AppConfig:
     return load_config()
 
 
-def _optional_path(env_name: str) -> Path | None:
-    value = os.getenv(env_name)
-    if value is None or not value.strip():
-        return None
-    return Path(value).expanduser()
+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/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

@@ -7,7 +7,7 @@ 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
+from .sql_generator import GeneratedSql, SqlGenerationError, contains_forbidden_dml_keyword
 
 
 Sampler = Callable[[str, str], str]
@@ -18,6 +18,7 @@ class SamplingSqlGenerator:
     sampler: Sampler
     default_limit: int
     max_limit: int
+    confirmed_query_plan: dict[str, Any] | None = None
 
     def generate(
         self,
@@ -30,6 +31,7 @@ class SamplingSqlGenerator:
             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))
 
@@ -48,6 +50,7 @@ class SamplingSqlGenerator:
             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))
 
@@ -63,7 +66,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:
@@ -77,6 +89,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],
@@ -84,6 +97,7 @@ 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,
     )
 
 
@@ -125,5 +139,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,6 +31,22 @@ class SpongeMcpClient:
         }
         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,
         *,