File size: 5,763 Bytes
116524e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
"""XML skill rendering and skill retrieval helpers."""

from __future__ import annotations

import logging
import re
from collections import defaultdict
from typing import TYPE_CHECKING, Iterable
from xml.sax.saxutils import escape

if TYPE_CHECKING:
    from ace.core.skillbook import Skill, Skillbook
    from ace.deduplication.detector import SimilarityDetector

logger = logging.getLogger(__name__)

RRF_K = 60


def _tokenize(text: str) -> list[str]:
    return re.findall(r"[a-z0-9_]+", text.lower())


def _normalize_keywords(keywords: Iterable[str] | None) -> list[str]:
    if keywords is None:
        return []
    normalized: list[str] = []
    seen: set[str] = set()
    for keyword in keywords:
        text = str(keyword).strip().lower().replace(" ", "_")
        if not text or text in seen:
            continue
        normalized.append(text)
        seen.add(text)
    return normalized


def _keyword_overlap(skill: "Skill", keywords: list[str]) -> int:
    if not keywords:
        return 0
    return sum(1 for keyword in keywords if keyword in skill.keywords)


def render_skills_xml(skills: list["Skill"]) -> str:
    """Render skills as XML ``<strategy>`` elements."""
    if not skills:
        return ""

    parts: list[str] = []
    for skill in skills:
        keyword_attr = ",".join(skill.keywords)
        body = [f"  <issue>{escape(skill.issue)}</issue>"]
        if skill.insight:
            body.append(f"  <insight>{escape(skill.insight)}</insight>")
        body.append(f"  <keywords>{escape(keyword_attr)}</keywords>")
        parts.append(
            f'<strategy id="{escape(skill.id)}" section="{escape(skill.section)}">\n'
            + "\n".join(body)
            + "\n</strategy>"
        )

    strategies_block = "\n".join(parts)
    return (
        f"{strategies_block}\n\n"
        "Adapt these strategies to your current situation; "
        "they are patterns, not rigid rules."
    )


def _lexical_ranking(

    skills: list["Skill"],

    query: str,

) -> list["Skill"]:
    if not skills:
        return []
    query_tokens = _tokenize(query)
    if not query_tokens:
        return skills

    try:
        from rank_bm25 import BM25Okapi

        corpus = [_tokenize(skill.embedding_text()) for skill in skills]
        bm25 = BM25Okapi(corpus)
        scores = bm25.get_scores(query_tokens)
        ranked_pairs = sorted(
            zip(scores, skills),
            key=lambda item: item[0],
            reverse=True,
        )
        return [skill for _, skill in ranked_pairs]
    except Exception as exc:
        logger.debug("BM25 unavailable, falling back to token overlap: %s", exc)
        query_token_set = set(query_tokens)
        ranked_pairs = []
        for skill in skills:
            doc_tokens = set(_tokenize(skill.embedding_text()))
            ranked_pairs.append((len(query_token_set & doc_tokens), skill))
        ranked_pairs.sort(key=lambda item: item[0], reverse=True)
        return [skill for _, skill in ranked_pairs]


def _dense_ranking(

    skills: list["Skill"],

    query: str,

    detector: "SimilarityDetector",

) -> list["Skill"]:
    if not skills:
        return []

    query_embedding = detector.compute_embedding(query)
    if query_embedding is None:
        raise RuntimeError(
            "Failed to embed retrieval query — "
            "check embedding provider credentials / network."
        )

    ranked_pairs: list[tuple[float, Skill]] = []
    for skill in skills:
        if skill.embedding is None:
            continue
        similarity = detector.cosine_similarity(query_embedding, skill.embedding)
        ranked_pairs.append((similarity, skill))

    ranked_pairs.sort(key=lambda item: item[0], reverse=True)
    return [skill for _, skill in ranked_pairs]


def retrieve_top_k(

    skillbook: "Skillbook",

    query: str,

    *,

    top_k: int = 5,

    detector: "SimilarityDetector | None" = None,

    section: str | None = None,

    keywords: list[str] | None = None,

) -> list["Skill"]:
    """Retrieve relevant skills using lexical + dense fusion."""
    if top_k <= 0:
        return []

    candidates = skillbook.skills()
    if section:
        normalized_section = str(section).strip().lower()
        candidates = [
            skill for skill in candidates if skill.section == normalized_section
        ]
    if not candidates:
        return []

    normalized_keywords = _normalize_keywords(keywords)

    if detector is None:
        from ace.deduplication.detector import SimilarityDetector as _Detector
        from ace.protocols.deduplication import DeduplicationConfig

        detector = _Detector(DeduplicationConfig())

    detector.ensure_embeddings(skillbook)
    lexical_ranked = _lexical_ranking(candidates, query)
    dense_ranked = _dense_ranking(candidates, query, detector)

    fused_scores: dict[str, float] = defaultdict(float)
    skills_by_id = {skill.id: skill for skill in candidates}

    for rank, skill in enumerate(lexical_ranked, start=1):
        fused_scores[skill.id] += 1.0 / (RRF_K + rank)
    for rank, skill in enumerate(dense_ranked, start=1):
        fused_scores[skill.id] += 1.0 / (RRF_K + rank)

    if normalized_keywords:
        for skill in candidates:
            overlap = _keyword_overlap(skill, normalized_keywords)
            if overlap:
                fused_scores[skill.id] += 0.25 * overlap

    ranked_ids = sorted(
        fused_scores,
        key=lambda skill_id: fused_scores[skill_id],
        reverse=True,
    )
    return [skills_by_id[skill_id] for skill_id in ranked_ids[:top_k]]