@roll-agent/octopus-agent 0.0.9 → 0.0.11

package/src/octopus_skill/schema_context_retriever.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import re
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 
@@ -12,50 +12,133 @@ class SchemaContext:
     glossary: list[dict[str, Any]]
     metrics: list[dict[str, Any]]
     examples: list[dict[str, Any]]
+    table_defaults: list[dict[str, Any]] = field(default_factory=list)
+    enums: list[dict[str, Any]] = field(default_factory=list)
 
 
-LOCATION_RECRUITING_TABLES = (
-    "province",
-    "city",
-    "region",
-    "job_basic_info",
-    "job_address",
-    "job_store",
-    "store",
-)
-
-WORK_ORDER_DETAIL_TABLES = (
-    "mvp_applied_jobs_by_helped_account",
-    "c_helped_account",
-    "resume",
-    "job",
-    "job_basic_info",
-    "brand",
-    "supplier",
-    "sponge_project",
-)
-
-
-def retrieve_schema_context(schema: dict[str, Any], question: str, *, max_tables: int = 12) -> SchemaContext:
+def retrieve_schema_context(
+    schema: dict[str, Any],
+    question: str,
+    *,
+    max_tables: int = 20,
+    required_tables: set[str] | None = None,
+) -> SchemaContext:
+    """Assemble the schema slice handed to the LLM.
+
+    Everything is driven by the MCP ``get_schema`` payload: tables are ranked by
+    overlap with the question and with the glossary/metrics/examples that match
+    the question, then the schema's own ``joins`` graph is walked to pull in any
+    bridge tables needed to connect them. The agent stores no per-business table
+    lists, keyword triggers or hand-coded joins.
+    """
     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 = _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)
-    joins = _append_derived_joins(schema, joins, selected_table_names)
-    return SchemaContext(tables=tables, joins=joins, glossary=glossary, metrics=metrics, examples=examples)
+    example_table_names = _ordered_target_table_names(examples)
+    glossary_metric_table_names = _target_table_names(glossary + metrics)
+    required_tables = required_tables or set()
+    preferred = table_names_from_context | required_tables
+
+    tables = _rank_tables(schema.get("tables", []), question, preferred)[:max_tables]
+    if isinstance(schema.get("tables"), list):
+        # Pin the must-have tables to the front in a deterministic order:
+        # required (resolved-entity) tables first, then tables referenced by the
+        # best matched examples, then glossary / metrics. Example table order is
+        # preserved because examples often encode the intended join path; this
+        # keeps late-declared fact tables from being evicted by weaker matches.
+        ordered_names = _ordered_preferred_names(
+            schema.get("tables", []),
+            required_tables,
+            example_table_names,
+            glossary_metric_table_names,
+        )
+        if ordered_names:
+            tables = _prepend_named_tables(schema.get("tables", []), tables, ordered_names, max_tables)
+        tables = _expand_via_join_graph(schema, tables, max_tables)
 
+    selected_table_names = {_table_name(table) for table in tables}
+    joins = _filter_joins(schema.get("joins", []), selected_table_names, required_tables)
+    table_defaults = _filter_table_defaults(schema.get("tableDefaults", []), selected_table_names)
+    enums = _filter_enums(schema.get("enums", []), selected_table_names)
+    return SchemaContext(
+        tables=tables,
+        joins=joins,
+        glossary=glossary,
+        metrics=metrics,
+        examples=examples,
+        table_defaults=table_defaults,
+        enums=enums,
+    )
 
-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 _ordered_preferred_names(
+    tables: Any,
+    required_tables: set[str],
+    example_tables: tuple[str, ...],
+    context_tables: set[str],
+) -> tuple[str, ...]:
+    """Deterministic priority order for must-have tables.
+
+    Required (resolved-entity) tables come first and follow schema declaration
+    order. Matched example tables come next in the example SQL/table order, then
+    glossary / metric tables follow schema order. Required tables always win the
+    max_tables budget, while example join paths stay intact.
+    """
+    if not isinstance(tables, list):
+        return ()
+    schema_order = [_table_name(table) for table in tables if isinstance(table, dict)]
+    schema_names = set(schema_order)
+    required_names = [name for name in schema_order if name in required_tables]
+    example_names = [
+        name
+        for name in example_tables
+        if name in schema_names and name not in required_tables
+    ]
+    context_names = [
+        name
+        for name in schema_order
+        if name in context_tables and name not in required_tables and name not in example_names
+    ]
+    return tuple(dict.fromkeys(required_names + example_names + context_names))
+
+
+def _expand_via_join_graph(
+    schema: dict[str, Any],
+    tables: list[dict[str, Any]],
+    max_tables: int,
+) -> list[dict[str, Any]]:
+    """Add bridge tables that connect two or more already-selected tables.
+
+    Relationship knowledge lives entirely in ``schema.joins``; this walks that
+    graph so indirect paths (e.g. brand -> brand_project -> sponge_project) are
+    complete without hard-coding any business-specific table set.
+    """
+    all_tables = schema.get("tables", [])
+    joins = schema.get("joins", [])
+    if not isinstance(all_tables, list) or not isinstance(joins, list):
+        return tables[:max_tables]
+    by_name = {_table_name(table): table for table in all_tables if isinstance(table, dict)}
+    selected = {_table_name(table) for table in tables}
+    adjacency: dict[str, set[str]] = {}
+    for join in joins:
+        if not isinstance(join, dict):
+            continue
+        left = str(join.get("leftTable") or join.get("from") or "")
+        right = str(join.get("rightTable") or join.get("to") or "")
+        if left and right:
+            adjacency.setdefault(left, set()).add(right)
+            adjacency.setdefault(right, set()).add(left)
+    result = list(tables)
+    for candidate, neighbours in adjacency.items():
+        if len(result) >= max_tables:
+            break
+        if candidate in selected or candidate not in by_name:
+            continue
+        if len(neighbours & selected) >= 2:
+            result.append(by_name[candidate])
+            selected.add(candidate)
+    return result[:max_tables]
 
 
 def _prepend_named_tables(
@@ -85,14 +168,13 @@ def _prepend_named_tables(
 def _rank_tables(tables: Any, question: str, preferred_names: set[str]) -> list[dict[str, Any]]:
     if not isinstance(tables, list):
         return []
+    tokens = _question_tokens(question)
     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
-        if _needs_work_order_detail_tables(question) and name in WORK_ORDER_DETAIL_TABLES:
-            score += 10
         text = " ".join(
             [
                 name,
@@ -100,7 +182,7 @@ def _rank_tables(tables: Any, question: str, preferred_names: set[str]) -> list[
                 " ".join(_column_text(column) for column in table.get("columns", []) if isinstance(column, dict)),
             ]
         ).lower()
-        for token in _question_tokens(question):
+        for token in tokens:
             if token and token in text:
                 score += 1
         if score > 0 or not preferred_names:
@@ -119,19 +201,6 @@ def _matching_items(items: Any, question: str, *, keys: tuple[str, ...]) -> list
             continue
         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)
@@ -139,71 +208,89 @@ def _matching_items(items: Any, question: str, *, keys: tuple[str, ...]) -> list
 
 
 def _target_table_names(items: list[dict[str, Any]]) -> set[str]:
+    return set(_ordered_target_table_names(items))
+
+
+def _ordered_target_table_names(items: list[dict[str, Any]]) -> tuple[str, ...]:
     names: set[str] = set()
+    ordered: list[str] = []
     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)
+                for name in value:
+                    _append_unique_table(ordered, names, str(name))
         sql = item.get("sql")
         if isinstance(sql, str):
-            names.update(referenced_table for referenced_table in _referenced_tables(sql))
+            for referenced_table in _referenced_tables(sql):
+                _append_unique_table(ordered, names, referenced_table)
         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
+                    _append_unique_table(ordered, names, target.split(".", 1)[0])
+    return tuple(ordered)
 
 
-def _filter_joins(joins: Any, selected_table_names: set[str]) -> list[dict[str, Any]]:
+def _append_unique_table(ordered: list[str], seen: set[str], name: str) -> None:
+    if not name or name in seen:
+        return
+    ordered.append(name)
+    seen.add(name)
+
+
+def _filter_joins(
+    joins: Any,
+    selected_table_names: set[str],
+    priority_table_names: set[str] | None = None,
+) -> list[dict[str, Any]]:
     if not isinstance(joins, list):
         return []
-    result: list[dict[str, Any]] = []
-    for join in joins:
+    priority_table_names = priority_table_names or set()
+    direct: list[tuple[int, int, dict[str, Any]]] = []
+    adjacent: list[tuple[int, int, dict[str, Any]]] = []
+    for index, join in enumerate(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]
+        matched = join_tables & selected_table_names
+        priority_score = len(join_tables & priority_table_names)
+        if join_tables <= selected_table_names:
+            direct.append((-priority_score, index, join))
+        elif matched:
+            adjacent.append((-priority_score, index, join))
+    direct.sort()
+    adjacent.sort()
+    return [join for _, _, join in direct + adjacent][:16]
+
+
+def _filter_table_defaults(table_defaults: Any, selected_table_names: set[str]) -> list[dict[str, Any]]:
+    """Pass through default filters for selected tables (data from get_schema)."""
+    if not isinstance(table_defaults, list):
+        return []
+    result: list[dict[str, Any]] = []
+    for entry in table_defaults:
+        if isinstance(entry, dict) and str(entry.get("table") or "") in selected_table_names:
+            result.append(entry)
+    return result
 
 
-def _append_derived_joins(
-    schema: dict[str, Any],
-    joins: list[dict[str, Any]],
-    selected_table_names: set[str],
-) -> list[dict[str, Any]]:
-    result = list(joins)
-    existing = {
-        (
-            str(join.get("leftTable") or ""),
-            str(join.get("leftColumn") or ""),
-            str(join.get("rightTable") or ""),
-            str(join.get("rightColumn") or ""),
-        )
-        for join in result
-    }
-    if (
-        {"resume", "supplier"} <= selected_table_names
-        and _table_has_column(schema, "resume", "supplier_id")
-        and _table_has_column(schema, "supplier", "id")
-        and ("supplier", "id", "resume", "supplier_id") not in existing
-        and ("resume", "supplier_id", "supplier", "id") not in existing
-    ):
-        result.append(
-            {
-                "description": "简历和供应商关系；供应商名称可使用 supplier.nick_name",
-                "leftTable": "supplier",
-                "leftColumn": "id",
-                "rightTable": "resume",
-                "rightColumn": "supplier_id",
-            }
-        )
-    return result[:16]
+def _filter_enums(enums: Any, selected_table_names: set[str]) -> list[dict[str, Any]]:
+    """Pass through enum definitions whose field belongs to a selected table."""
+    if not isinstance(enums, list):
+        return []
+    result: list[dict[str, Any]] = []
+    for entry in enums:
+        if not isinstance(entry, dict):
+            continue
+        field = str(entry.get("field") or "")
+        table = field.split(".", 1)[0] if "." in field else ""
+        if table in selected_table_names:
+            result.append(entry)
+    return result
 
 
 def _table_name(table: dict[str, Any]) -> str:
@@ -216,65 +303,27 @@ def _column_text(column: dict[str, Any]) -> str:
     )
 
 
-def _table_has_column(schema: dict[str, Any], table_name: str, column_name: str) -> bool:
-    tables = schema.get("tables", [])
-    if not isinstance(tables, list):
-        return False
-    for table in tables:
-        if not isinstance(table, dict) or _table_name(table) != table_name:
-            continue
-        columns = table.get("columns", [])
-        if not isinstance(columns, list):
-            return False
-        return any(
-            isinstance(column, dict)
-            and str(column.get("columnName") or column.get("name") or "") == column_name
-            for column in columns
-        )
-    return False
-
-
 def _question_tokens(question: str) -> list[str]:
+    """Generic tokenization with no business vocabulary.
+
+    Produces the lowered question, its ascii word splits, and CJK character
+    bigrams so Chinese questions can still overlap Chinese table comments /
+    glossary text without any hand-maintained keyword list.
+    """
     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)
+    tokens.extend(part for part in re.split(r"[\s_\-]+", lowered) if part)
+    for run in re.findall(r"[\u4e00-\u9fff]+", question):
+        if len(run) == 1:
+            tokens.append(run)
+        for index in range(len(run) - 1):
+            tokens.append(run[index : index + 2])
     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 ("姓名", "岗位", "项目", "候选人状态"))
-    )
+def _referenced_tables(sql: str) -> tuple[str, ...]:
+    names: list[str] = []
+    seen: set[str] = set()
+    for match in re.finditer(r"\b(?:from|join)\s+([a-zA-Z_][a-zA-Z0-9_]*)\b", sql, re.IGNORECASE):
+        _append_unique_table(names, seen, match.group(1))
+    return tuple(names)