@roll-agent/octopus-agent 0.0.10 → 0.0.11

package/src/octopus_skill/entity_resolver.py

@@ -0,0 +1,446 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Iterable
+
+
+# Scorer returns a similarity in [0.0, 1.0]. Backed by rapidfuzz when available
+# (better CJK fuzzy matching), with a stdlib difflib fallback so the package
+# still works if rapidfuzz is not installed.
+Scorer = Callable[[str, str], float]
+
+try:
+    from rapidfuzz import fuzz as _rapidfuzz_fuzz
+
+    def default_scorer(a: str, b: str) -> float:
+        return _rapidfuzz_fuzz.ratio(a, b) / 100.0
+
+except ImportError:  # pragma: no cover - exercised only without rapidfuzz
+    from difflib import SequenceMatcher
+
+    def default_scorer(a: str, b: str) -> float:
+        return SequenceMatcher(None, a, b).ratio()
+
+
+GEO_TYPES = ("province", "city", "region")
+# Trailing tokens dropped when comparing geo names so "上海市" matches "上海".
+_GEO_SUFFIXES = ("特别行政区", "自治区", "新区", "地区", "盟", "省", "市", "区", "县", "旗")
+
+
+@dataclass(frozen=True)
+class EntityItem:
+    entity_type: str
+    id: Any
+    name: str
+    aliases: tuple[str, ...]
+    parent_id: Any | None
+
+
+@dataclass(frozen=True)
+class ResolvedEntity:
+    entity_type: str
+    id: Any
+    canonical: str
+    table: str
+    column: str
+    score: float
+    matched_via: str  # exact | alias | partial | fuzzy
+    parent_id: Any | None = None
+
+
+@dataclass(frozen=True)
+class Mention:
+    """A dictionary entity detected as a substring of a larger text."""
+
+    entity: ResolvedEntity
+    start: int
+    end: int
+    text: str  # the matched dictionary name/alias (normalized)
+
+
+@dataclass
+class _TypeIndex:
+    entity_type: str
+    table: str
+    id_column: str
+    name_column: str
+    parent_type: str | None
+    items: list[EntityItem] = field(default_factory=list)
+    _exact: dict[str, EntityItem] = field(default_factory=dict)
+    _alias: dict[str, EntityItem] = field(default_factory=dict)
+    # (normalized_match_text, item) pairs used for substring and fuzzy matching
+    _candidates: list[tuple[str, EntityItem]] = field(default_factory=list)
+    by_id: dict[Any, EntityItem] = field(default_factory=dict)
+
+    def add(self, item: EntityItem, extra_aliases: Iterable[str] = ()) -> None:
+        self.items.append(item)
+        self.by_id.setdefault(item.id, item)
+        name_key = _norm(item.name)
+        if name_key:
+            self._exact.setdefault(name_key, item)
+            self._candidates.append((name_key, item))
+            core = _geo_core(name_key)
+            if core != name_key and core:
+                self._alias.setdefault(core, item)
+                self._candidates.append((core, item))
+        for alias in tuple(item.aliases) + tuple(extra_aliases):
+            alias_key = _norm(alias)
+            if alias_key:
+                self._alias.setdefault(alias_key, item)
+                self._candidates.append((alias_key, item))
+
+    def lookup(self, key: str) -> tuple[EntityItem, str] | None:
+        if key in self._exact:
+            return self._exact[key], "exact"
+        if key in self._alias:
+            return self._alias[key], "alias"
+        core = _geo_core(key)
+        if core and core in self._exact:
+            return self._exact[core], "exact"
+        if core and core in self._alias:
+            return self._alias[core], "alias"
+        return None
+
+
+@dataclass
+class EntityResolver:
+    indexes: dict[str, _TypeIndex]
+    scorer: Scorer = default_scorer
+
+    @classmethod
+    def from_payload(
+        cls,
+        dictionary: dict[str, Any],
+        *,
+        geo_aliases: dict[str, dict[str, str]] | None = None,
+        scorer: Scorer = default_scorer,
+    ) -> "EntityResolver":
+        geo_aliases = geo_aliases or {}
+        entities = dictionary.get("entities") if isinstance(dictionary, dict) else None
+        indexes: dict[str, _TypeIndex] = {}
+        if isinstance(entities, dict):
+            for entity_type, payload in entities.items():
+                if not isinstance(payload, dict):
+                    continue
+                index = _TypeIndex(
+                    entity_type=entity_type,
+                    table=str(payload.get("table") or entity_type),
+                    id_column=str(payload.get("idColumn") or "id"),
+                    name_column=str(payload.get("nameColumn") or "name"),
+                    parent_type=payload.get("parentType"),
+                )
+                # canonical name -> [aliases] coming from the geo alias config
+                alias_by_canonical: dict[str, list[str]] = {}
+                for alias, canonical in (geo_aliases.get(entity_type) or {}).items():
+                    if alias.startswith("_"):
+                        continue
+                    alias_by_canonical.setdefault(_norm(canonical), []).append(alias)
+                for raw in payload.get("items") or []:
+                    if not isinstance(raw, dict):
+                        continue
+                    name = str(raw.get("name") or "")
+                    if not name:
+                        continue
+                    aliases = tuple(
+                        str(value)
+                        for value in (raw.get("aliases") or [])
+                        if str(value).strip()
+                    )
+                    item = EntityItem(
+                        entity_type=entity_type,
+                        id=raw.get(index.id_column) if index.id_column in raw else raw.get("id"),
+                        name=name,
+                        aliases=aliases,
+                        parent_id=raw.get("parentId"),
+                    )
+                    # Config canonical names may omit geo suffixes (e.g. the
+                    # config says "南京" while the table stores "南京市"), so
+                    # match on both the full name and its geo core.
+                    name_key = _norm(name)
+                    extra = list(alias_by_canonical.get(name_key, ()))
+                    core_key = _geo_core(name_key)
+                    if core_key != name_key:
+                        extra.extend(alias_by_canonical.get(core_key, ()))
+                    index.add(item, extra_aliases=extra)
+                indexes[entity_type] = index
+        return cls(indexes=indexes, scorer=scorer)
+
+    def resolve(
+        self,
+        span: str,
+        entity_type: str,
+        *,
+        limit: int = 5,
+        min_score: float = 0.6,
+    ) -> list[ResolvedEntity]:
+        index = self.indexes.get(entity_type)
+        if index is None:
+            return []
+        key = _norm(span)
+        if not key:
+            return []
+
+        direct = index.lookup(key)
+        if direct is not None:
+            item, via = direct
+            score = 1.0 if via == "exact" else 0.95
+            return [_to_resolved(item, index, score, via)]
+
+        scored: dict[Any, ResolvedEntity] = {}
+        for candidate_text, item in index._candidates:
+            score = self.scorer(key, candidate_text)
+            if score < min_score:
+                continue
+            existing = scored.get(item.id)
+            if existing is None or score > existing.score:
+                scored[item.id] = _to_resolved(item, index, round(score, 4), "fuzzy")
+        ranked = sorted(scored.values(), key=lambda entity: entity.score, reverse=True)
+        return ranked[:limit]
+
+    def find_mentions(
+        self,
+        text: str,
+        entity_type: str,
+        *,
+        min_len: int = 2,
+    ) -> list[Mention]:
+        """Detect dictionary entities that appear as substrings of ``text``.
+
+        This is the reverse of :meth:`resolve`: instead of resolving a single
+        pre-extracted span, it scans ``text`` for any known name/alias. Each
+        item is reported once, using its longest matching candidate. Matching
+        is exact/alias only (no fuzzy) to keep substring scanning precise.
+        """
+        index = self.indexes.get(entity_type)
+        if index is None:
+            return []
+        haystack = _norm(text)
+        if not haystack:
+            return []
+        best: dict[Any, Mention] = {}
+        for candidate_text, item in index._candidates:
+            if len(candidate_text) < min_len:
+                continue
+            start = haystack.find(candidate_text)
+            if start < 0:
+                continue
+            via = "exact" if candidate_text == _norm(item.name) else "alias"
+            score = 1.0 if via == "exact" else 0.95
+            mention = Mention(
+                entity=_to_resolved(item, index, score, via),
+                start=start,
+                end=start + len(candidate_text),
+                text=candidate_text,
+            )
+            existing = best.get(item.id)
+            if existing is None or (mention.end - mention.start) > (existing.end - existing.start):
+                best[item.id] = mention
+        return sorted(best.values(), key=lambda mention: mention.start)
+
+    def find_containing(
+        self,
+        span: str,
+        entity_type: str,
+        *,
+        limit: int | None = None,
+        min_len: int = 2,
+    ) -> list[ResolvedEntity]:
+        """Find entities whose known name/alias contains ``span``.
+
+        This covers user shorthand such as "必胜客-1" matching projects named
+        "上海必胜客-1" and "南京必胜客-1".
+        """
+        index = self.indexes.get(entity_type)
+        if index is None:
+            return []
+        key = _norm(span)
+        if len(key) < min_len:
+            return []
+
+        best: dict[Any, ResolvedEntity] = {}
+        order: dict[Any, int] = {}
+        for candidate_text, item in index._candidates:
+            if len(candidate_text) < min_len or key not in candidate_text:
+                continue
+            if candidate_text == _norm(item.name):
+                via = "exact" if key == candidate_text else "partial"
+                score = 1.0 if key == candidate_text else round(0.9 * len(key) / len(candidate_text), 4)
+            else:
+                via = "alias" if key == candidate_text else "partial"
+                score = 0.95 if key == candidate_text else round(0.85 * len(key) / len(candidate_text), 4)
+            resolved = _to_resolved(item, index, score, via)
+            existing = best.get(item.id)
+            if existing is None:
+                order[item.id] = len(order)
+            if existing is None or resolved.score > existing.score:
+                best[item.id] = resolved
+
+        ranked = sorted(
+            best.values(),
+            key=lambda entity: (-entity.score, order.get(entity.id, 0)),
+        )
+        return ranked if limit is None else ranked[:limit]
+
+    def resolve_location(
+        self,
+        span: str,
+        *,
+        min_score: float = 0.6,
+    ) -> dict[str, ResolvedEntity]:
+        """Decompose a location span into province / city / region.
+
+        Handles compounds like "上海宝山" and disambiguates duplicate names
+        (e.g. 宝山区 exists in both 上海 and 黑龙江) by preferring the geo entity
+        whose parent chain is consistent with the other levels mentioned in the
+        same span.
+        """
+        key = _norm(span)
+        if not key:
+            return {}
+
+        # Collect every geo item whose name/alias is a substring, keeping ALL
+        # items per matched span (duplicate names must all be considered).
+        spans: dict[tuple[int, int], dict[str, list[EntityItem]]] = {}
+        for geo_type in GEO_TYPES:
+            index = self.indexes.get(geo_type)
+            if index is None:
+                continue
+            for candidate_text, item in index._candidates:
+                if len(candidate_text) < 2:
+                    continue
+                start = key.find(candidate_text)
+                if start < 0:
+                    continue
+                bucket = spans.setdefault((start, start + len(candidate_text)), {})
+                items = bucket.setdefault(geo_type, [])
+                if all(existing.id != item.id for existing in items):
+                    items.append(item)
+
+        if not spans:
+            return self._fuzzy_location_fallback(span, min_score)
+
+        # Accept non-overlapping spans, longest first.
+        accepted = self._accept_geo_spans(spans)
+        provinces = [item for bucket in accepted for item in bucket.get("province", [])]
+        cities = [item for bucket in accepted for item in bucket.get("city", [])]
+        regions = [item for bucket in accepted for item in bucket.get("region", [])]
+
+        return self._assemble_location(provinces, cities, regions)
+
+    def _accept_geo_spans(
+        self,
+        spans: dict[tuple[int, int], dict[str, list[EntityItem]]],
+    ) -> list[dict[str, list[EntityItem]]]:
+        ordered = sorted(spans.items(), key=lambda item: item[0][1] - item[0][0], reverse=True)
+        accepted: list[tuple[tuple[int, int], dict[str, list[EntityItem]]]] = []
+        for (start, end), bucket in ordered:
+            if any(start < other_end and other_start < end for (other_start, other_end), _ in accepted):
+                continue
+            accepted.append(((start, end), bucket))
+        return [bucket for _, bucket in accepted]
+
+    def _assemble_location(
+        self,
+        provinces: list[EntityItem],
+        cities: list[EntityItem],
+        regions: list[EntityItem],
+    ) -> dict[str, ResolvedEntity]:
+        city_index = self.indexes.get("city")
+        province_ids = {item.id for item in provinces}
+        city_ids = {item.id for item in cities}
+
+        chosen_region: EntityItem | None = None
+        chosen_city: EntityItem | None = None
+        if regions:
+            # Prefer a region whose parent city is mentioned, else whose parent
+            # city rolls up to a mentioned province; fall back to the first.
+            for region in regions:
+                if region.parent_id in city_ids:
+                    chosen_region = region
+                    break
+            if chosen_region is None and province_ids and city_index is not None:
+                for region in regions:
+                    parent_city = city_index.by_id.get(region.parent_id)
+                    if parent_city is not None and parent_city.parent_id in province_ids:
+                        chosen_region = region
+                        chosen_city = parent_city
+                        break
+            if chosen_region is None:
+                chosen_region = regions[0]
+
+        if chosen_city is None and chosen_region is not None and city_index is not None:
+            # Anchor the city to the region's declared parent when available.
+            parent_city = city_index.by_id.get(chosen_region.parent_id)
+            if parent_city is not None and (not cities or parent_city.id in city_ids):
+                chosen_city = parent_city
+        if chosen_city is None and cities:
+            chosen_city = cities[0]
+
+        chosen_province: EntityItem | None = None
+        if chosen_city is not None and provinces:
+            chosen_province = next((p for p in provinces if p.id == chosen_city.parent_id), None)
+        if chosen_province is None and provinces:
+            chosen_province = provinces[0]
+
+        result: dict[str, ResolvedEntity] = {}
+        if chosen_province is not None:
+            result["province"] = _to_resolved(chosen_province, self.indexes["province"], 1.0, "exact")
+        if chosen_city is not None:
+            result["city"] = _to_resolved(chosen_city, self.indexes["city"], 1.0, "exact")
+        if chosen_region is not None:
+            result["region"] = _to_resolved(chosen_region, self.indexes["region"], 1.0, "exact")
+        return result
+
+    def _fuzzy_location_fallback(self, span: str, min_score: float) -> dict[str, ResolvedEntity]:
+        best: ResolvedEntity | None = None
+        for geo_type in GEO_TYPES:
+            for resolved in self.resolve(span, geo_type, limit=1, min_score=min_score):
+                if best is None or resolved.score > best.score:
+                    best = resolved
+        return {best.entity_type: best} if best is not None else {}
+
+
+def load_geo_aliases(path: Path) -> dict[str, dict[str, str]]:
+    if not path.exists():
+        return {}
+    data = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(data, dict):
+        return {}
+    return {
+        key: value
+        for key, value in data.items()
+        if isinstance(value, dict) and not key.startswith("_")
+    }
+
+
+def _to_resolved(
+    item: EntityItem,
+    index: _TypeIndex,
+    score: float,
+    via: str,
+) -> ResolvedEntity:
+    return ResolvedEntity(
+        entity_type=item.entity_type,
+        id=item.id,
+        canonical=item.name,
+        table=index.table,
+        column=index.name_column,
+        score=score,
+        matched_via=via,
+        parent_id=item.parent_id,
+    )
+
+
+def _norm(value: str) -> str:
+    return value.strip().lower()
+
+
+def _geo_core(value: str) -> str:
+    core = value
+    for suffix in _GEO_SUFFIXES:
+        if core.endswith(suffix) and len(core) > len(suffix):
+            core = core[: -len(suffix)]
+            break
+    return core

package/src/octopus_skill/llm_sql_generator.py

@@ -18,6 +18,7 @@ class SamplingSqlGenerator:
     sampler: Sampler
     default_limit: int
     max_limit: int
+    resolved_entities: 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,
+            resolved_entities=self.resolved_entities,
         )
         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,
+            resolved_entities=self.resolved_entities,
         )
         return _parse_generated_sql(self.sampler(NL2SQL_SYSTEM_PROMPT, prompt))
 

package/src/octopus_skill/mcp_client.py

@@ -31,6 +31,22 @@ class SpongeMcpClient:
         }
         return self._post("get_schema", payload)
 
+    def get_entity_dictionary(
+        self,
+        *,
+        trace_id: str,
+        session_id: str | None = None,
+        dict_version: str | None = None,
+        entity_types: list[str] | tuple[str, ...] | None = None,
+    ) -> dict[str, Any]:
+        payload = {
+            "traceId": trace_id,
+            "sessionId": session_id,
+            "dictVersion": dict_version,
+            "entityTypes": list(entity_types) if entity_types else None,
+        }
+        return self._post("get_entity_dictionary", payload)
+
     def validate_sql(
         self,
         *,