@roll-agent/octopus-agent 0.0.1

package/src/octopus_sponge_skill/roll_server.py

@@ -0,0 +1,391 @@
+from __future__ import annotations
+
+import json
+import sys
+from dataclasses import asdict, is_dataclass
+from typing import Any, Callable
+
+from . import __version__
+from .agent import SpongeAgent
+from .config import AppConfig, load_config_from_env
+from .schema_cache import SchemaCacheError
+
+SERVER_NAME = "octopus-agent"
+SERVER_VERSION = __version__
+PROTOCOL_VERSION = "2024-11-05"
+
+Json = dict[str, Any]
+AgentFactory = Callable[[], SpongeAgent]
+SamplerFactory = Callable[[str, str], str]
+
+
+def main() -> None:
+    serve_stdio()
+
+
+def serve_stdio() -> None:
+    session = StdioSession(sys.stdin.buffer, sys.stdout.buffer)
+    while True:
+        message = session.read_message()
+        if message is None:
+            return
+        response = handle_request(message, sampler=session.sample)
+        if response is not None:
+            session.write_message(response)
+
+
+def handle_request(
+    message: Json,
+    agent_factory: AgentFactory | None = None,
+    sampler: SamplerFactory | None = None,
+) -> Json | None:
+    method = message.get("method")
+    if method is None:
+        return _error_response(message.get("id"), -32600, "Invalid request")
+    if "id" not in message:
+        return None
+
+    request_id = message["id"]
+    try:
+        if method == "initialize":
+            return _result_response(request_id, initialize_result())
+        if method == "ping":
+            return _result_response(request_id, {})
+        if method == "tools/list":
+            return _result_response(request_id, {"tools": tool_definitions()})
+        if method == "tools/call":
+            params = message.get("params") or {}
+            return _result_response(request_id, call_tool(params, agent_factory=agent_factory, sampler=sampler))
+        return _error_response(request_id, -32601, f"Method not found: {method}")
+    except Exception as exc:
+        return _error_response(request_id, -32603, str(exc), {"type": type(exc).__name__})
+
+
+def initialize_result() -> Json:
+    return {
+        "protocolVersion": PROTOCOL_VERSION,
+        "capabilities": {"tools": {}},
+        "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
+    }
+
+
+def tool_definitions() -> list[Json]:
+    return [
+        {
+            "name": "diagnostic_status",
+            "description": "检查 Octopus Sponge Agent 配置、缓存和本地持久化状态",
+            "inputSchema": {
+                "type": "object",
+                "properties": {},
+                "additionalProperties": False,
+            },
+        },
+        {
+            "name": "refresh_schema",
+            "description": "调用 Sponge MCP Server 刷新本地 Schema Cache（表结构缓存）",
+            "inputSchema": {
+                "type": "object",
+                "required": ["traceId"],
+                "properties": {
+                    "traceId": {
+                        "type": "string",
+                        "description": "调用链路 ID，必填，会原样传给 Sponge get_schema",
+                    },
+                    "refreshReason": {
+                        "type": "string",
+                        "description": "刷新原因，默认 manual（人工触发）",
+                    }
+                },
+                "additionalProperties": False,
+            },
+        },
+        {
+            "name": "query_sponge",
+            "description": "把自然语言问题转换为只读 SQL，并通过 Sponge MCP Server 校验和执行",
+            "inputSchema": {
+                "type": "object",
+                "required": ["question", "userId"],
+                "properties": {
+                    "question": {"type": "string", "description": "用户自然语言问题"},
+                    "userId": {"type": "integer", "description": "当前登录用户对应的 Duliri userId"},
+                    "sessionId": {"type": "string", "description": "可选会话 ID"},
+                    "deviceId": {"type": "string", "description": "可选设备 ID"},
+                    "resultFormat": {
+                        "type": "string",
+                        "enum": ["text", "yaml", "md", "markdown", "html"],
+                        "description": "结果输出格式；不传则根据问题自动识别",
+                    },
+                    "includeSql": {
+                        "type": "boolean",
+                        "description": "是否在回参中包含 generatedSql、normalizedSql、repairedSql 等调试信息",
+                    },
+                },
+                "additionalProperties": False,
+            },
+        },
+    ]
+
+
+def call_tool(
+    params: Json,
+    agent_factory: AgentFactory | None = None,
+    sampler: SamplerFactory | None = None,
+) -> Json:
+    tool_name = params.get("name")
+    arguments = params.get("arguments") or {}
+    try:
+        if tool_name == "diagnostic_status":
+            return _tool_json(_diagnostic_status(_load_config()))
+        if tool_name == "refresh_schema":
+            agent = _make_agent(agent_factory)
+            return _tool_json(
+                _refresh_schema(
+                    agent,
+                    trace_id=_required_string(arguments, "traceId"),
+                    refresh_reason=str(arguments.get("refreshReason") or "manual"),
+                )
+            )
+        if tool_name == "query_sponge":
+            agent = _make_agent(agent_factory)
+            answer_kwargs = {
+                "question": _required_string(arguments, "question"),
+                "user_id": _required_int(arguments, "userId"),
+                "session_id": arguments.get("sessionId"),
+                "device_id": arguments.get("deviceId"),
+                "result_format": arguments.get("resultFormat"),
+            }
+            if sampler is not None:
+                answer_kwargs["sampler"] = sampler
+            query_result = agent.query(**answer_kwargs)
+            if arguments.get("includeSql") is True:
+                return _tool_json(
+                    {
+                        "answer": query_result.answer,
+                        "generatedSql": query_result.generated_sql,
+                        "normalizedSql": query_result.normalized_sql,
+                        "repairedSql": query_result.repaired_sql,
+                        "validation": query_result.validation,
+                    }
+                )
+            return _tool_text(query_result.answer)
+        return _tool_error(f"Unknown tool: {tool_name}", code="unknown_tool")
+    except Exception as exc:
+        return _tool_error(str(exc), code=type(exc).__name__)
+
+
+def _diagnostic_status(config: AppConfig) -> Json:
+    cache_exists = config.schema.cache_path.exists()
+    schema_version: str | None = None
+    cache_error: str | None = None
+    try:
+        schema_version = SpongeAgent.from_config(config).schema_cache.version()
+    except SchemaCacheError as exc:
+        cache_error = str(exc)
+
+    return {
+        "server": {"name": SERVER_NAME, "version": SERVER_VERSION},
+        "spongeMcp": {
+            "baseUrl": config.sponge_mcp.base_url,
+            "accessTokenConfigured": bool(config.sponge_mcp.access_token),
+            "timeoutMs": config.sponge_mcp.timeout_ms,
+        },
+        "schema": {
+            "cachePath": str(config.schema.cache_path),
+            "cacheExists": cache_exists,
+            "schemaVersion": schema_version,
+            "cacheError": cache_error,
+        },
+        "audit": {
+            "enabled": config.audit.enabled,
+            "logPath": str(config.audit.log_path),
+        },
+    }
+
+
+def _refresh_schema(agent: SpongeAgent, *, trace_id: str, refresh_reason: str) -> Json:
+    cached_schema = agent.schema_cache.load()
+    current_version = str(cached_schema.get("schemaVersion")) if cached_schema and cached_schema.get("schemaVersion") else None
+    schema = agent.mcp_client.get_schema(
+        trace_id=trace_id,
+        schema_version=current_version,
+        refresh_reason=refresh_reason,
+    )
+    response_changed = schema.get("changed")
+    if schema.get("changed") is False and cached_schema is not None:
+        schema = cached_schema
+    elif schema.get("changed") is False and cached_schema is None:
+        schema = agent.mcp_client.get_schema(
+            trace_id=trace_id,
+            schema_version=None,
+            refresh_reason=refresh_reason,
+        )
+    if "tables" in schema:
+        agent.schema_cache.save(schema)
+    schema_source = "cache" if response_changed is False and cached_schema is not None else "response"
+    agent.audit_logger.log(
+        {
+            "traceId": trace_id,
+            "tool": "get_schema",
+            "schemaVersion": schema.get("schemaVersion"),
+            "refreshReason": refresh_reason,
+            "changed": response_changed,
+        }
+    )
+    return {
+        "changed": response_changed,
+        "request": {
+            "traceId": trace_id,
+            "schemaVersion": current_version,
+            "refreshReason": refresh_reason,
+        },
+        "schema": {
+            "schemaVersion": schema.get("schemaVersion"),
+            "generatedAt": schema.get("generatedAt"),
+            "source": schema_source,
+        },
+        "cachePath": str(agent.schema_cache.path),
+    }
+
+
+def _load_config() -> AppConfig:
+    return load_config_from_env()
+
+
+def _make_agent(agent_factory: AgentFactory | None) -> SpongeAgent:
+    if agent_factory is not None:
+        return agent_factory()
+    return SpongeAgent.from_config(_load_config())
+
+
+def _required_string(arguments: Json, key: str) -> str:
+    value = arguments.get(key)
+    if not isinstance(value, str) or not value:
+        raise ValueError(f"{key} is required")
+    return value
+
+
+def _required_int(arguments: Json, key: str) -> int:
+    value = arguments.get(key)
+    if not isinstance(value, int):
+        raise ValueError(f"{key} is required")
+    return value
+
+
+def _tool_text(text: str) -> Json:
+    return {"content": [{"type": "text", "text": text}]}
+
+
+def _tool_json(data: Any) -> Json:
+    return {
+        "content": [
+            {
+                "type": "text",
+                "text": json.dumps(_jsonable(data), ensure_ascii=False, sort_keys=True),
+            }
+        ]
+    }
+
+
+def _tool_error(message: str, *, code: str) -> Json:
+    return {
+        "isError": True,
+        "content": [
+            {
+                "type": "text",
+                "text": json.dumps({"code": code, "message": message}, ensure_ascii=False, sort_keys=True),
+            }
+        ],
+    }
+
+
+def _jsonable(value: Any) -> Any:
+    if is_dataclass(value):
+        return _jsonable(asdict(value))
+    if isinstance(value, dict):
+        return {str(key): _jsonable(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_jsonable(item) for item in value]
+    return value
+
+
+def _result_response(request_id: Any, result: Json) -> Json:
+    return {"jsonrpc": "2.0", "id": request_id, "result": result}
+
+
+def _error_response(request_id: Any, code: int, message: str, data: Json | None = None) -> Json:
+    error: Json = {"code": code, "message": message}
+    if data is not None:
+        error["data"] = data
+    return {"jsonrpc": "2.0", "id": request_id, "error": error}
+
+
+class StdioSession:
+    def __init__(self, input_stream: Any, output_stream: Any) -> None:
+        self.input_stream = input_stream
+        self.output_stream = output_stream
+        self.next_request_id = 1
+
+    def read_message(self) -> Json | None:
+        return _read_message(self.input_stream)
+
+    def write_message(self, message: Json) -> None:
+        _write_message(self.output_stream, message)
+
+    def sample(self, system_prompt: str, prompt: str) -> str:
+        request_id = f"sampling_{self.next_request_id}"
+        self.next_request_id += 1
+        self.write_message(
+            {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "method": "sampling/createMessage",
+                "params": {
+                    "systemPrompt": system_prompt,
+                    "messages": [{"role": "user", "content": {"type": "text", "text": prompt}}],
+                    "temperature": 0,
+                    "maxTokens": 1600,
+                    "includeContext": "none",
+                },
+            }
+        )
+        while True:
+            message = self.read_message()
+            if message is None:
+                raise RuntimeError("Sampling failed: client closed stdio")
+            if message.get("id") == request_id:
+                if "error" in message:
+                    raise RuntimeError(f"Sampling failed: {message['error']}")
+                return _extract_sampling_text(message.get("result") or {})
+            response = handle_request(message, sampler=self.sample)
+            if response is not None:
+                self.write_message(response)
+
+
+def _extract_sampling_text(result: Json) -> str:
+    content = result.get("content")
+    if isinstance(content, dict) and content.get("type") == "text":
+        return str(content.get("text") or "")
+    if isinstance(content, list):
+        return "\n".join(
+            str(item.get("text") or "") for item in content if isinstance(item, dict) and item.get("type") == "text"
+        )
+    raise RuntimeError("Sampling response did not contain text content")
+
+
+def _read_message(stream: Any) -> Json | None:
+    while True:
+        line = stream.readline()
+        if line == b"":
+            return None
+        if line.strip():
+            return json.loads(line.decode("utf-8"))
+
+
+def _write_message(stream: Any, message: Json) -> None:
+    body = json.dumps(message, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+    stream.write(body + b"\n")
+    stream.flush()
+
+
+if __name__ == "__main__":
+    main()

package/src/octopus_sponge_skill/schema_cache.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+class SchemaCacheError(RuntimeError):
+    """Raised when the persisted schema cache cannot be used."""
+
+
+@dataclass
+class SchemaCache:
+    path: Path
+
+    def load(self) -> dict[str, Any] | None:
+        if not self.path.exists():
+            return None
+        try:
+            data = json.loads(self.path.read_text(encoding="utf-8"))
+        except json.JSONDecodeError as exc:
+            raise SchemaCacheError(f"Invalid schema cache JSON: {self.path}") from exc
+        if not isinstance(data, dict):
+            raise SchemaCacheError(f"Schema cache must be a JSON object: {self.path}")
+        return data
+
+    def save(self, schema: dict[str, Any]) -> None:
+        if not isinstance(schema, dict):
+            raise TypeError("schema must be a dict")
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.path.write_text(
+            json.dumps(schema, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    def version(self) -> str | None:
+        schema = self.load()
+        if schema is None:
+            return None
+        version = schema.get("schemaVersion")
+        return str(version) if version is not None else None
+

package/src/octopus_sponge_skill/schema_context_retriever.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class SchemaContext:
+    tables: list[dict[str, Any]]
+    joins: list[dict[str, Any]]
+    glossary: list[dict[str, Any]]
+    metrics: list[dict[str, Any]]
+    examples: list[dict[str, Any]]
+
+
+def retrieve_schema_context(schema: dict[str, Any], question: str, *, max_tables: int = 8) -> SchemaContext:
+    glossary = _matching_items(schema.get("glossary", []), question, keys=("term", "description"))
+    metrics = _matching_items(schema.get("metrics", []), question, keys=("name", "metric", "description", "definition"))
+    examples = _matching_items(schema.get("examples", []), question, keys=("question", "description"))
+    table_names_from_context = _target_table_names(glossary + metrics + examples)
+    tables = _rank_tables(schema.get("tables", []), question, table_names_from_context)[:max_tables]
+    selected_table_names = {_table_name(table) for table in tables}
+    joins = _filter_joins(schema.get("joins", []), selected_table_names)
+    return SchemaContext(tables=tables, joins=joins, glossary=glossary, metrics=metrics, examples=examples)
+
+
+def _rank_tables(tables: Any, question: str, preferred_names: set[str]) -> list[dict[str, Any]]:
+    if not isinstance(tables, list):
+        return []
+    ranked: list[tuple[int, dict[str, Any]]] = []
+    for table in tables:
+        if not isinstance(table, dict):
+            continue
+        name = _table_name(table)
+        score = 5 if name in preferred_names else 0
+        text = " ".join(
+            [
+                name,
+                str(table.get("comment") or ""),
+                " ".join(_column_text(column) for column in table.get("columns", []) if isinstance(column, dict)),
+            ]
+        ).lower()
+        for token in _question_tokens(question):
+            if token and token in text:
+                score += 1
+        if score > 0 or not preferred_names:
+            ranked.append((score, table))
+    ranked.sort(key=lambda item: item[0], reverse=True)
+    return [table for _, table in ranked]
+
+
+def _matching_items(items: Any, question: str, *, keys: tuple[str, ...]) -> list[dict[str, Any]]:
+    if not isinstance(items, list):
+        return []
+    tokens = _question_tokens(question)
+    matches: list[dict[str, Any]] = []
+    for item in items:
+        if not isinstance(item, dict):
+            continue
+        text = " ".join(str(item.get(key) or "") for key in keys).lower()
+        if any(token and token in text for token in tokens):
+            matches.append(item)
+    return matches[:8]
+
+
+def _target_table_names(items: list[dict[str, Any]]) -> set[str]:
+    names: set[str] = set()
+    for item in items:
+        for key in ("tables", "targetTables"):
+            value = item.get(key)
+            if isinstance(value, list):
+                names.update(str(name) for name in value)
+        targets = item.get("targets")
+        if isinstance(targets, list):
+            for target in targets:
+                if isinstance(target, str) and "." in target:
+                    names.add(target.split(".", 1)[0])
+    return names
+
+
+def _filter_joins(joins: Any, selected_table_names: set[str]) -> list[dict[str, Any]]:
+    if not isinstance(joins, list):
+        return []
+    result: list[dict[str, Any]] = []
+    for join in joins:
+        if not isinstance(join, dict):
+            continue
+        join_tables = {
+            str(join.get("leftTable") or join.get("from") or ""),
+            str(join.get("rightTable") or join.get("to") or ""),
+        }
+        if join_tables & selected_table_names:
+            result.append(join)
+    return result[:12]
+
+
+def _table_name(table: dict[str, Any]) -> str:
+    return str(table.get("tableName") or table.get("name") or "")
+
+
+def _column_text(column: dict[str, Any]) -> str:
+    return " ".join(
+        str(column.get(key) or "") for key in ("columnName", "name", "dataType", "type", "comment")
+    )
+
+
+def _question_tokens(question: str) -> list[str]:
+    lowered = question.lower()
+    tokens = [lowered]
+    tokens.extend(part for part in lowered.replace("_", " ").replace("-", " ").split() if part)
+    for word in ("品牌", "项目", "岗位", "门店", "招聘", "在招", "数量", "列表"):
+        if word in question:
+            tokens.append(word)
+    return list(dict.fromkeys(tokens))
+

package/src/octopus_sponge_skill/sql_generator.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+
+class SqlGenerationError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class GeneratedSql:
+    sql: str
+    tables: list[str]
+    placeholders: list[str]
+
+
+@dataclass
+class RuleBasedSqlGenerator:
+    default_limit: int
+    max_limit: int
+
+    def generate(
+        self,
+        *,
+        question: str,
+        schema: dict[str, Any],
+        scope_hints: dict[str, Any],
+    ) -> GeneratedSql:
+        table_names = _table_names(schema)
+        lowered = question.lower()
+
+        if ("品牌" in question or "brand" in lowered) and {"brand", "brand_project"}.issubset(table_names):
+            sql = (
+                "SELECT b.id AS brand_id, b.name AS brand_name "
+                "FROM brand b "
+                "JOIN brand_project bp ON bp.brand_id = b.id "
+                "WHERE bp.project_id IN (:scope.project_ids) "
+                "AND b.id IN (:scope.brand_ids) "
+                f"LIMIT {self.default_limit}"
+            )
+            return GeneratedSql(sql, ["brand", "brand_project"], _placeholders(scope_hints))
+
+        if ("项目" in question or "project" in lowered) and "sponge_project" in table_names:
+            sql = (
+                "SELECT p.id AS project_id, p.name AS project_name "
+                "FROM sponge_project p "
+                "WHERE p.id IN (:scope.project_ids) "
+                f"LIMIT {self.default_limit}"
+            )
+            return GeneratedSql(sql, ["sponge_project"], [":scope.project_ids"])
+
+        raise SqlGenerationError("No safe SQL template matched the question")
+
+
+def enforce_sql_rules(sql: str, *, max_limit: int) -> None:
+    normalized = sql.strip().lower()
+    if not normalized.startswith("select "):
+        raise SqlGenerationError("SQL must be SELECT only")
+    if ";" in normalized.rstrip(";"):
+        raise SqlGenerationError("SQL must be a single statement")
+    if re.search(r"\b(insert|update|delete|drop|alter|create|truncate)\b", normalized):
+        raise SqlGenerationError("SQL contains a forbidden keyword")
+    if re.search(r"\b(password|token|secret|credential|private_key)\b", normalized):
+        raise SqlGenerationError("SQL contains a sensitive field")
+    limit = re.search(r"\blimit\s+(\d+)\b", normalized)
+    if limit is None:
+        raise SqlGenerationError("SQL must contain LIMIT")
+    if int(limit.group(1)) > max_limit:
+        raise SqlGenerationError(f"SQL LIMIT must not exceed {max_limit}")
+
+
+def _table_names(schema: dict[str, Any]) -> set[str]:
+    tables = schema.get("tables", [])
+    if not isinstance(tables, list):
+        return set()
+    names: set[str] = set()
+    for table in tables:
+        if not isinstance(table, dict):
+            continue
+        name = table.get("name") or table.get("tableName")
+        if name:
+            names.add(str(name))
+    return names
+
+
+def _placeholders(scope_hints: dict[str, Any]) -> list[str]:
+    placeholders: list[str] = []
+    for value in scope_hints.values():
+        if isinstance(value, dict) and value.get("placeholder"):
+            placeholders.append(str(value["placeholder"]))
+    return placeholders