56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
from skills_learning.config import SkillAuthoringConfig, load_config
|
|
from skills_learning.embeddings import EmbeddingClient, embed_skill_text
|
|
from skills_learning.models import FlowTemplateSkill
|
|
from skills_learning.store import SkillStore, get_default_store
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ScoredSkill:
|
|
skill: FlowTemplateSkill
|
|
score: float
|
|
|
|
|
|
def retrieve_candidate_skills(
|
|
goal: str,
|
|
*,
|
|
store: SkillStore | None = None,
|
|
top_k: int | None = None,
|
|
embedding_client: EmbeddingClient | None = None,
|
|
config: SkillAuthoringConfig | None = None,
|
|
) -> list[ScoredSkill]:
|
|
settings = config or load_config()
|
|
skill_store = store or get_default_store()
|
|
query_vector = embed_skill_text(
|
|
goal,
|
|
client=embedding_client,
|
|
config=settings,
|
|
)
|
|
if query_vector is None:
|
|
return []
|
|
|
|
scored: list[ScoredSkill] = []
|
|
for record in skill_store.list_embeddings():
|
|
skill = skill_store.get_by_id(record.skill_id)
|
|
if skill is None:
|
|
continue
|
|
score = _cosine_similarity(query_vector, record.vector)
|
|
scored.append(ScoredSkill(skill=skill, score=score))
|
|
scored.sort(key=lambda item: item.score, reverse=True)
|
|
return scored[: top_k or settings.top_k]
|
|
|
|
|
|
def _cosine_similarity(left: list[float], right: list[float]) -> float:
|
|
if not left or not right or len(left) != len(right):
|
|
return 0.0
|
|
dot = sum(a * b for a, b in zip(left, right, strict=True))
|
|
left_norm = math.sqrt(sum(value * value for value in left))
|
|
right_norm = math.sqrt(sum(value * value for value in right))
|
|
if left_norm == 0 or right_norm == 0:
|
|
return 0.0
|
|
return dot / (left_norm * right_norm)
|