@roll-agent/octopus-agent 0.0.1 → 0.0.2

package/src/octopus_sponge_skill/schema_cache.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -41,3 +42,8 @@ class SchemaCache:
         version = schema.get("schemaVersion")
         return str(version) if version is not None else None
 
+    def is_fresh(self, max_age_minutes: int) -> bool:
+        if max_age_minutes <= 0 or not self.path.exists():
+            return False
+        age_seconds = time.time() - self.path.stat().st_mtime
+        return age_seconds <= max_age_minutes * 60

package/src/octopus_sponge_skill/schema_context_retriever.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 from typing import Any
 
@@ -13,17 +14,69 @@ class SchemaContext:
     examples: list[dict[str, Any]]
 
 
+LOCATION_RECRUITING_TABLES = (
+    "province",
+    "city",
+    "region",
+    "job_basic_info",
+    "job_address",
+    "job_store",
+    "store",
+)
+
+WORK_ORDER_DETAIL_TABLES = (
+    "mvp_applied_jobs_by_helped_account",
+    "brand",
+    "job_basic_info",
+    "sponge_project",
+)
+
+
 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]
+    tables = _select_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 _select_tables(tables: Any, question: str, preferred_names: set[str], max_tables: int) -> list[dict[str, Any]]:
+    ranked = _rank_tables(tables, question, preferred_names)
+    if _needs_work_order_detail_tables(question) and isinstance(tables, list):
+        return _prepend_named_tables(tables, ranked, WORK_ORDER_DETAIL_TABLES, max_tables)
+    if not _needs_location_recruiting_tables(question) or not isinstance(tables, list):
+        return ranked[:max_tables]
+
+    return _prepend_named_tables(tables, ranked, LOCATION_RECRUITING_TABLES, max_tables)
+
+
+def _prepend_named_tables(
+    tables: Any,
+    ranked: list[dict[str, Any]],
+    names: tuple[str, ...],
+    max_tables: int,
+) -> list[dict[str, Any]]:
+    by_name = {_table_name(table): table for table in tables if isinstance(table, dict)}
+    selected: list[dict[str, Any]] = []
+    seen: set[str] = set()
+    for name in names:
+        table = by_name.get(name)
+        if table is not None:
+            selected.append(table)
+            seen.add(name)
+    for table in ranked:
+        name = _table_name(table)
+        if name not in seen:
+            selected.append(table)
+            seen.add(name)
+        if len(selected) >= max_tables:
+            break
+    return selected[:max_tables]
+
+
 def _rank_tables(tables: Any, question: str, preferred_names: set[str]) -> list[dict[str, Any]]:
     if not isinstance(tables, list):
         return []
@@ -33,6 +86,8 @@ def _rank_tables(tables: Any, question: str, preferred_names: set[str]) -> list[
             continue
         name = _table_name(table)
         score = 5 if name in preferred_names else 0
+        if _needs_work_order_detail_tables(question) and name in WORK_ORDER_DETAIL_TABLES:
+            score += 10
         text = " ".join(
             [
                 name,
@@ -53,14 +108,29 @@ def _matching_items(items: Any, question: str, *, keys: tuple[str, ...]) -> list
     if not isinstance(items, list):
         return []
     tokens = _question_tokens(question)
-    matches: list[dict[str, Any]] = []
+    matches: list[tuple[int, 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]
+        text = " ".join(str(item.get(key) or "") for key in keys + ("sql",)).lower()
+        score = sum(1 for token in tokens if token and token in text)
+        if _needs_work_order_detail_tables(question):
+            for term in (
+                "mvp_applied_jobs_by_helped_account",
+                "work_order_status",
+                "current_helper_status_id",
+                "job_basic_info.job_id",
+                "报名",
+                "工单",
+                "候选人状态",
+                "面试成功",
+            ):
+                if term.lower() in text:
+                    score += 5
+        if score > 0:
+            matches.append((score, item))
+    matches.sort(key=lambda match: match[0], reverse=True)
+    return [item for _, item in matches[:8]]
 
 
 def _target_table_names(items: list[dict[str, Any]]) -> set[str]:
@@ -70,6 +140,9 @@ def _target_table_names(items: list[dict[str, Any]]) -> set[str]:
             value = item.get(key)
             if isinstance(value, list):
                 names.update(str(name) for name in value)
+        sql = item.get("sql")
+        if isinstance(sql, str):
+            names.update(referenced_table for referenced_table in _referenced_tables(sql))
         targets = item.get("targets")
         if isinstance(targets, list):
             for target in targets:
@@ -108,8 +181,40 @@ 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 ("品牌", "项目", "岗位", "门店", "招聘", "在招", "数量", "列表"):
+    for word in (
+        "品牌",
+        "项目",
+        "岗位",
+        "门店",
+        "招聘",
+        "在招",
+        "数量",
+        "列表",
+        "报名",
+        "报名情况",
+        "报名工单",
+        "工单",
+        "候选人",
+        "候选人状态",
+        "面试成功",
+        "进行中",
+        "姓名",
+    ):
         if word in question:
             tokens.append(word)
     return list(dict.fromkeys(tokens))
 
+
+def _referenced_tables(sql: str) -> set[str]:
+    return set(match.group(1) for match in re.finditer(r"\b(?:from|join)\s+([a-zA-Z_][a-zA-Z0-9_]*)\b", sql, re.IGNORECASE))
+
+
+def _needs_location_recruiting_tables(question: str) -> bool:
+    return any(word in question for word in ("岗位", "门店", "招聘", "在招"))
+
+
+def _needs_work_order_detail_tables(question: str) -> bool:
+    return bool(
+        any(word in question for word in ("报名", "工单", "候选人"))
+        and any(word in question for word in ("姓名", "岗位", "项目", "候选人状态"))
+    )