@roll-agent/octopus-agent 0.0.1 → 0.0.2

package/src/octopus_sponge_skill/agent.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import re
 from dataclasses import dataclass, field
 from typing import Any
@@ -9,10 +10,20 @@ from .config import AppConfig
 from .context import UserContext, new_session_id, new_trace_id
 from .llm_sql_generator import Sampler, SamplingSqlGenerator
 from .mcp_client import SpongeMcpClient
-from .result_renderer import render_result
+from .query_sql_cache import QuerySqlCache
+from .result_renderer import normalize_result_format, render_result
 from .schema_cache import SchemaCache
 from .schema_context_retriever import retrieve_schema_context
-from .sql_generator import RuleBasedSqlGenerator, enforce_sql_rules
+from .sql_generator import (
+    GeneratedSql,
+    RuleBasedSqlGenerator,
+    SqlGenerationError,
+    enforce_sql_rules,
+    is_low_quality_sql_for_question,
+    is_recruiting_match_effect_query,
+    sql_satisfies_time_range,
+    validate_sql_against_schema,
+)
 
 
 @dataclass(frozen=True)
@@ -32,6 +43,7 @@ class SpongeAgent:
     audit_logger: AuditLogger
     _last_validation_error: dict[str, Any] = field(default_factory=dict, init=False)
     _last_validation_result: dict[str, Any] = field(default_factory=dict, init=False)
+    _last_validation_trace_id: str | None = field(default=None, init=False)
 
     @classmethod
     def from_config(cls, config: AppConfig) -> "SpongeAgent":
@@ -50,82 +62,236 @@ class SpongeAgent:
         self,
         *,
         question: str,
-        user_id: int,
+        original_question: str | None = None,
         session_id: str | None = None,
-        device_id: str | None = None,
         sampler: Sampler | None = None,
         result_format: str | None = None,
+        fast_mode: bool = True,
     ) -> str:
         return self.query(
             question=question,
-            user_id=user_id,
+            original_question=original_question,
             session_id=session_id,
-            device_id=device_id,
             sampler=sampler,
             result_format=result_format,
+            fast_mode=fast_mode,
         ).answer
 
     def query(
         self,
         *,
         question: str,
-        user_id: int,
+        original_question: str | None = None,
         session_id: str | None = None,
-        device_id: str | None = None,
         sampler: Sampler | None = None,
         result_format: str | None = None,
+        fast_mode: bool = True,
     ) -> QueryResult:
+        user_question = original_question or question
+        if _is_mutation_request(user_question):
+            return _query_only_result()
         context = UserContext(
-            user_id=user_id,
             session_id=session_id or new_session_id(),
-            device_id=device_id,
         )
-        sql_question = _strip_result_format_instruction(question)
+        sql_question = _strip_result_format_instruction(user_question)
         schema = self._ensure_schema()
-        allowed_context = self._get_allowed_context(context, sql_question)
-        scope_hints = allowed_context.get("scopeHints", {})
+        schema_version = str(schema.get("schemaVersion")) if schema.get("schemaVersion") is not None else None
+        repaired_sql: str | None = None
+        prefer_rule_based = is_recruiting_match_effect_query(sql_question)
+        sql_cache = _make_query_sql_cache(self.config)
+        cached_sql = sql_cache.get(question=sql_question, schema_version=schema_version) if sql_cache is not None else None
+        if cached_sql is not None and not sql_satisfies_time_range(cached_sql, sql_question, schema):
+            cached_sql = None
+        if cached_sql is not None and is_low_quality_sql_for_question(sql_question, cached_sql):
+            sql_cache.delete(question=sql_question, schema_version=schema_version)
+            cached_sql = None
+        if cached_sql is not None and validate_sql_against_schema(cached_sql, schema):
+            sql_cache.delete(question=sql_question, schema_version=schema_version)
+            cached_sql = None
+        if cached_sql is not None:
+            generated = GeneratedSql(sql=cached_sql, tables=[], placeholders=[])
         if sampler is not None:
-            schema_context = retrieve_schema_context(schema, sql_question)
-            generator = SamplingSqlGenerator(
-                sampler=sampler,
+            schema_context = None
+            generator: SamplingSqlGenerator | None = None
+            if cached_sql is None:
+                schema_context = retrieve_schema_context(schema, sql_question)
+                generator = SamplingSqlGenerator(
+                    sampler=sampler,
+                    default_limit=self.config.sql.default_limit,
+                    max_limit=self.config.sql.max_limit,
+                )
+                if fast_mode or prefer_rule_based:
+                    try:
+                        generated = _generate_rule_based_sql(
+                            question=sql_question,
+                            schema=schema,
+                            default_limit=self.config.sql.default_limit,
+                            max_limit=self.config.sql.max_limit,
+                        )
+                    except SqlGenerationError:
+                        generated, repaired_sql = _generate_with_sampler(
+                            generator=generator,
+                            question=sql_question,
+                            schema_context=schema_context,
+                        )
+                    else:
+                        repaired_sql = None
+                else:
+                    generated, repaired_sql = _generate_with_sampler(
+                        generator=generator,
+                        question=sql_question,
+                        schema_context=schema_context,
+                    )
+                    if generated.sql == "" and repaired_sql is None:
+                        return _query_only_result()
+                if fast_mode and generated.sql == "" and repaired_sql is None:
+                    return _query_only_result()
+                if not fast_mode and generated.sql == "" and repaired_sql is None:
+                    return _query_only_result()
+        elif cached_sql is None:
+            schema_context = None
+            generated = _generate_rule_based_sql(
+                question=sql_question,
+                schema=schema,
                 default_limit=self.config.sql.default_limit,
                 max_limit=self.config.sql.max_limit,
             )
-            generated = generator.generate(
-                question=sql_question,
-                schema_context=schema_context,
-                scope_hints=scope_hints,
-            )
         else:
             schema_context = None
-            generator = RuleBasedSqlGenerator(
-                default_limit=self.config.sql.default_limit,
-                max_limit=self.config.sql.max_limit,
-            )
-            generated = generator.generate(
-                question=sql_question,
-                schema=schema,
-                scope_hints=scope_hints,
-            )
-        enforce_sql_rules(generated.sql, max_limit=self.config.sql.max_limit)
-        normalized_sql = self._validate_sql(context, question, generated.sql)
-        repaired_sql: str | None = None
-        if sampler is not None and normalized_sql is None and schema_context is not None:
-            repaired = generator.repair(
-                question=sql_question,
-                original_sql=generated.sql,
-                validation_error=self._last_validation_error,
-                schema_context=schema_context,
-                scope_hints=scope_hints,
-            )
-            repaired_sql = repaired.sql
-            enforce_sql_rules(repaired.sql, max_limit=self.config.sql.max_limit)
-            normalized_sql = self._validate_sql(context, question, repaired.sql, raise_on_error=True)
+            generator = None
+        candidate_sql = repaired_sql or _apply_default_limit(generated.sql, sql_question, self.config.sql.default_limit)
+        if not sql_satisfies_time_range(candidate_sql, sql_question, schema):
+            if sampler is None or schema_context is None or generator is None:
+                return _query_only_result()
+            try:
+                repaired = generator.repair(
+                    question=sql_question,
+                    original_sql=candidate_sql,
+                    validation_error={
+                        "code": "MISSING_TIME_RANGE_FILTER",
+                        "message": "用户问题包含时间范围，SQL 必须在 WHERE 条件中直接包含对应时间字段过滤。",
+                    },
+                    schema_context=schema_context,
+                )
+            except SqlGenerationError:
+                return _query_only_result()
+            generated = GeneratedSql(sql="", tables=[], placeholders=[])
+            repaired_sql = _apply_default_limit(repaired.sql, sql_question, self.config.sql.default_limit)
+            candidate_sql = repaired_sql
+            if not sql_satisfies_time_range(candidate_sql, sql_question, schema):
+                return _query_only_result()
+        if is_low_quality_sql_for_question(sql_question, candidate_sql):
+            if sampler is None or schema_context is None or generator is None:
+                return _query_only_result()
+            try:
+                repaired = generator.repair(
+                    question=sql_question,
+                    original_sql=candidate_sql,
+                    validation_error={
+                        "code": "LOW_QUALITY_SQL",
+                        "message": "SQL 与问题意图不匹配，报名/工单/候选人明细不能退化成品牌列表，且必须使用正确工单表、状态和岗位关联。",
+                    },
+                    schema_context=schema_context,
+                )
+            except SqlGenerationError:
+                return _query_only_result()
+            generated = GeneratedSql(sql="", tables=[], placeholders=[])
+            repaired_sql = _apply_default_limit(repaired.sql, sql_question, self.config.sql.default_limit)
+            candidate_sql = repaired_sql
+            if is_low_quality_sql_for_question(sql_question, candidate_sql):
+                return _query_only_result()
+        schema_issues = validate_sql_against_schema(candidate_sql, schema)
+        if schema_issues:
+            if sampler is None or schema_context is None or generator is None:
+                return _query_only_result()
+            try:
+                repaired = generator.repair(
+                    question=sql_question,
+                    original_sql=candidate_sql,
+                    validation_error={
+                        "code": schema_issues[0].code,
+                        "message": schema_issues[0].message,
+                    },
+                    schema_context=schema_context,
+                )
+            except SqlGenerationError:
+                return _query_only_result()
+            generated = GeneratedSql(sql="", tables=[], placeholders=[])
+            repaired_sql = _apply_default_limit(repaired.sql, sql_question, self.config.sql.default_limit)
+            candidate_sql = repaired_sql
+            if validate_sql_against_schema(candidate_sql, schema):
+                return _query_only_result()
+        try:
+            enforce_sql_rules(candidate_sql, max_limit=self.config.sql.max_limit)
+        except SqlGenerationError:
+            if sampler is None:
+                return _query_only_result()
+            if not fast_mode:
+                return _query_only_result()
+            try:
+                generated = _generate_rule_based_sql(
+                    question=sql_question,
+                    schema=schema,
+                    default_limit=self.config.sql.default_limit,
+                    max_limit=self.config.sql.max_limit,
+                )
+                repaired_sql = None
+                candidate_sql = _apply_default_limit(generated.sql, sql_question, self.config.sql.default_limit)
+                enforce_sql_rules(candidate_sql, max_limit=self.config.sql.max_limit)
+            except SqlGenerationError:
+                return _query_only_result()
+        normalized_sql = self._validate_sql(context, user_question, candidate_sql)
+        if sampler is not None and normalized_sql is None and schema_context is not None and generator is not None:
+            try:
+                repaired = generator.repair(
+                    question=sql_question,
+                    original_sql=generated.sql,
+                    validation_error=self._last_validation_error,
+                    schema_context=schema_context,
+                )
+            except SqlGenerationError:
+                try:
+                    if not fast_mode:
+                        return _query_only_result()
+                    generated = _generate_rule_based_sql(
+                        question=sql_question,
+                        schema=schema,
+                        default_limit=self.config.sql.default_limit,
+                        max_limit=self.config.sql.max_limit,
+                    )
+                    repaired_sql = None
+                    candidate_sql = _apply_default_limit(generated.sql, sql_question, self.config.sql.default_limit)
+                    enforce_sql_rules(candidate_sql, max_limit=self.config.sql.max_limit)
+                    normalized_sql = self._validate_sql(context, user_question, candidate_sql, raise_on_error=True)
+                except SqlGenerationError:
+                    return _query_only_result()
+            else:
+                generated = GeneratedSql(sql="", tables=[], placeholders=[])
+                repaired_sql = _apply_default_limit(repaired.sql, sql_question, self.config.sql.default_limit)
+                try:
+                    enforce_sql_rules(repaired_sql, max_limit=self.config.sql.max_limit)
+                except SqlGenerationError:
+                    return _query_only_result()
+                normalized_sql = self._validate_sql(context, user_question, repaired_sql, raise_on_error=True)
         if normalized_sql is None:
             raise RuntimeError("SQL validation failed")
-        result = self._execute_sql(context, question, normalized_sql)
+        if sql_cache is not None and cached_sql is None and not is_low_quality_sql_for_question(sql_question, normalized_sql):
+            sql_cache.set(question=sql_question, schema_version=schema_version, sql=normalized_sql)
+        result = self._execute_sql(context, user_question, normalized_sql)
+        answer = render_result(user_question, result, result_format)
+        if _should_append_default_limit_notice(
+            question=sql_question,
+            result=result,
+            default_limit=self.config.sql.default_limit,
+        ):
+            answer = _append_default_limit_notice(
+                answer,
+                user_question,
+                result_format,
+                _default_limit_notice(self.config.sql.default_limit),
+            )
         return QueryResult(
-            answer=render_result(question, result, result_format),
+            answer=answer,
             generated_sql=generated.sql,
             normalized_sql=normalized_sql,
             repaired_sql=repaired_sql,
@@ -140,7 +306,6 @@ class SpongeAgent:
             schema_response = self.mcp_client.get_schema(
                 trace_id=trace_id,
                 schema_version=cached_version,
-                refresh_reason="startup",
             )
         except Exception:
             if cached_schema is not None:
@@ -166,13 +331,6 @@ class SpongeAgent:
             )
             return cached_schema
 
-        if schema_response.get("changed") is False and cached_schema is None:
-            schema_response = self.mcp_client.get_schema(
-                trace_id=new_trace_id("trace_schema"),
-                schema_version=None,
-                refresh_reason="startup",
-            )
-
         if "tables" in schema_response:
             self.schema_cache.save(schema_response)
             self.audit_logger.log(
@@ -189,26 +347,6 @@ class SpongeAgent:
             return cached_schema
         raise RuntimeError("Sponge get_schema did not return tables and no local schema cache exists")
 
-    def _get_allowed_context(self, context: UserContext, question: str) -> dict[str, Any]:
-        trace_id = new_trace_id()
-        result = self.mcp_client.get_allowed_context(
-            user_id=context.user_id,
-            trace_id=trace_id,
-            session_id=context.session_id,
-            device_id=context.device_id,
-            question=question,
-        )
-        self.audit_logger.log(
-            {
-                "traceId": trace_id,
-                "sessionId": context.session_id,
-                "userId": context.user_id,
-                "question": question,
-                "tool": "get_allowed_context",
-            }
-        )
-        return result
-
     def _validate_sql(
         self,
         context: UserContext,
@@ -218,10 +356,11 @@ class SpongeAgent:
         raise_on_error: bool = False,
     ) -> str | None:
         trace_id = new_trace_id()
+        self._last_validation_trace_id = None
         result = self.mcp_client.validate_sql(
-            user_id=context.user_id,
             trace_id=trace_id,
             session_id=context.session_id,
+            question=question,
             sql=sql,
         )
         self._last_validation_result = result
@@ -229,7 +368,6 @@ class SpongeAgent:
             {
                 "traceId": trace_id,
                 "sessionId": context.session_id,
-                "userId": context.user_id,
                 "question": question,
                 "tool": "validate_sql",
                 "sqlHash": sql_hash(sql),
@@ -244,14 +382,17 @@ class SpongeAgent:
             if raise_on_error:
                 raise RuntimeError(f"SQL validation failed: {code}")
             return None
+        self._last_validation_trace_id = trace_id
         return str(result.get("normalizedSql") or sql)
 
     def _execute_sql(self, context: UserContext, question: str, sql: str) -> dict[str, Any]:
-        trace_id = new_trace_id()
+        if self._last_validation_trace_id is None:
+            raise RuntimeError("SQL execution requires a prior successful validate_sql trace")
+        trace_id = self._last_validation_trace_id
         result = self.mcp_client.execute_sql(
-            user_id=context.user_id,
             trace_id=trace_id,
             session_id=context.session_id,
+            question=question,
             sql=sql,
             timeout_ms=self.config.sponge_mcp.timeout_ms,
         )
@@ -259,7 +400,6 @@ class SpongeAgent:
             {
                 "traceId": trace_id,
                 "sessionId": context.session_id,
-                "userId": context.user_id,
                 "question": question,
                 "tool": "execute_sql",
                 "sqlHash": sql_hash(sql),
@@ -278,11 +418,159 @@ class SpongeAgent:
 def _strip_result_format_instruction(question: str) -> str:
     cleaned = question
     patterns = [
-        r"[，,。；;]?\s*输出\s*(?:yaml|html|markdown|md)\s*(?:格式)?",
+        r"^\s*返回\s*(?:json|yaml|html|markdown|md)\s*格式\s*的?",
+        r"[，,。；;]?\s*输出\s*(?:json|yaml|html|markdown|md)\s*(?:格式)?",
         r"[，,。；;]?\s*用\s*表格\s*展示",
-        r"[，,。；;]?\s*以\s*(?:yaml|html|markdown|md)\s*(?:格式)?\s*输出",
-        r"[，,。；;]?\s*return\s+(?:yaml|html|markdown|md)",
+        r"[，,。；;]?\s*以\s*(?:json|yaml|html|markdown|md)\s*(?:格式)?\s*输出",
+        r"[，,。；;]?\s*return\s+(?:json|yaml|html|markdown|md)",
     ]
     for pattern in patterns:
         cleaned = re.sub(pattern, "", cleaned, flags=re.IGNORECASE)
     return cleaned.strip() or question
+
+
+def _query_only_result() -> QueryResult:
+    return QueryResult(
+        answer="目前只能查询",
+        generated_sql="",
+        normalized_sql="",
+        repaired_sql=None,
+        validation={},
+    )
+
+
+def _generate_with_sampler(
+    *,
+    generator: SamplingSqlGenerator,
+    question: str,
+    schema_context: Any,
+) -> tuple[GeneratedSql, str | None]:
+    try:
+        return (
+            generator.generate(
+                question=question,
+                schema_context=schema_context,
+            ),
+            None,
+        )
+    except SqlGenerationError as sampling_exc:
+        try:
+            repaired = generator.repair(
+                question=question,
+                original_sql="",
+                validation_error={"code": type(sampling_exc).__name__, "message": str(sampling_exc)},
+                schema_context=schema_context,
+            )
+        except SqlGenerationError:
+            return GeneratedSql(sql="", tables=[], placeholders=[]), None
+        return (
+            GeneratedSql(sql="", tables=[], placeholders=[]),
+            _apply_default_limit(repaired.sql, question, generator.default_limit),
+        )
+
+
+def _generate_rule_based_sql(
+    *,
+    question: str,
+    schema: dict[str, Any],
+    default_limit: int,
+    max_limit: int,
+) -> GeneratedSql:
+    generator = RuleBasedSqlGenerator(
+        default_limit=default_limit,
+        max_limit=max_limit,
+    )
+    return generator.generate(
+        question=question,
+        schema=schema,
+    )
+
+
+def _generate_sampling_fallback_sql(
+    *,
+    question: str,
+    schema: dict[str, Any],
+    default_limit: int,
+    max_limit: int,
+) -> GeneratedSql:
+    if re.search(r"(岗位|职位|招聘|在招)", question) is None:
+        raise SqlGenerationError("No sampling fallback matched the question")
+    return _generate_rule_based_sql(
+        question=question,
+        schema=schema,
+        default_limit=default_limit,
+        max_limit=max_limit,
+    )
+
+
+def _make_query_sql_cache(config: AppConfig) -> QuerySqlCache | None:
+    if config.sql.query_cache_path is None:
+        return None
+    return QuerySqlCache(config.sql.query_cache_path, config.sql.query_cache_ttl_minutes)
+
+
+def _is_mutation_request(question: str) -> bool:
+    mutation_match = re.search(
+        r"(新增|新建|创建|添加|插入|录入|修改|更新|编辑|变更|调整|删除|移除|作废|发布|下架|保存|提交|审批|导入|同步|insert|update|delete|create|drop|alter|truncate)",
+        question,
+        flags=re.IGNORECASE,
+    )
+    if mutation_match is None:
+        return False
+    if _is_mutation_analytics_request(question):
+        return False
+    return True
+
+
+def _is_mutation_analytics_request(question: str) -> bool:
+    return bool(
+        re.search(r"(统计|查询|查看|看看|多少|几次|次数|数量|count|列表|明细|记录|日志|历史|趋势|分组)", question, flags=re.IGNORECASE)
+        and re.search(r"(新增|新建|创建|添加|插入|录入|修改|更新|编辑|变更|调整|删除|移除|作废|发布|下架|保存|提交|审批|导入|同步)", question)
+        and re.search(r"(次数|数量|count|列表|明细|记录|日志|历史|趋势|分组|显示|展示|返回|包含|字段)", question, flags=re.IGNORECASE)
+    )
+
+
+def _apply_default_limit(sql: str, question: str, default_limit: int) -> str:
+    if _question_requests_limit(question):
+        return sql
+    return re.sub(r"\blimit\s+\d+\b", f"LIMIT {default_limit}", sql, count=1, flags=re.IGNORECASE)
+
+
+def _question_requests_limit(question: str) -> bool:
+    lowered = question.lower()
+    if re.search(r"\b(?:limit|top)\s+\d+\b", lowered):
+        return True
+    if re.search(r"(?:返回|展示|查询|列出|给我|取)\s*\d+\s*(?:条|行|个|项|家|名)", question):
+        return True
+    return bool(re.search(r"前\s*\d+\s*(?:条|行|个|项|家|名)?", question))
+
+
+def _default_limit_notice(default_limit: int) -> str:
+    return f"本次为默认查询 {default_limit} 条数据；如果需要查询更多数据，可以修改查询口径。"
+
+
+def _should_append_default_limit_notice(*, question: str, result: dict[str, Any], default_limit: int) -> bool:
+    if _question_requests_limit(question):
+        return False
+    try:
+        row_count = int(result.get("rowCount", len(result.get("rows", []))))
+    except (TypeError, ValueError):
+        return False
+    return row_count == default_limit
+
+
+def _append_default_limit_notice(answer: str, question: str, result_format: str | None, notice: str) -> str:
+    output_format = normalize_result_format(question, result_format)
+    if output_format == "json":
+        try:
+            payload = json.loads(answer)
+        except json.JSONDecodeError:
+            return f"{answer}\n\n{notice}"
+        if isinstance(payload, dict):
+            payload["notice"] = notice
+            return json.dumps(payload, ensure_ascii=False, indent=2) + "\n"
+    if output_format == "yaml":
+        return f"{answer.rstrip()}\nnotice: {notice}\n"
+    if output_format == "html":
+        return f"{answer}<p>{notice}</p>"
+    return f"{answer}\n\n{notice}"

package/src/octopus_sponge_skill/config.py

@@ -29,6 +29,8 @@ 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)
@@ -59,8 +61,10 @@ def load_config() -> AppConfig:
         ),
         sql=SqlConfig(
             max_repair_attempts=2,
-            default_limit=100,
+            default_limit=50,
             max_limit=500,
+            query_cache_path=PACKAGE_ROOT / "data/query-sql-cache.json",
+            query_cache_ttl_minutes=30,
         ),
         audit=AuditConfig(
             enabled=True,

package/src/octopus_sponge_skill/context.py

@@ -6,9 +6,7 @@ from uuid import uuid4
 
 @dataclass(frozen=True)
 class UserContext:
-    user_id: int
     session_id: str
-    device_id: str | None = None
 
 
 def new_trace_id(prefix: str = "trace") -> str:
@@ -17,4 +15,3 @@ def new_trace_id(prefix: str = "trace") -> str:
 
 def new_session_id() -> str:
     return f"session_{uuid4().hex}"
-

package/src/octopus_sponge_skill/device_info.py

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