diff --git a/application/defs/candidate_cre_defs.py b/application/defs/candidate_cre_defs.py new file mode 100644 index 000000000..0c8218858 --- /dev/null +++ b/application/defs/candidate_cre_defs.py @@ -0,0 +1,40 @@ +import re +from dataclasses import dataclass + + +@dataclass +class CandidateCRE: + """Represents a single CRE candidate match for a CheatsheetRecord.""" + + cre_id: str + name: str + description: str + score: float + + def __post_init__(self): + for field_name, value in [("cre_id", self.cre_id), ("name", self.name)]: + if not isinstance(value, str) or not value.strip(): + raise ValueError( + f"CandidateCRE: field '{field_name}' must be a non-empty string" + ) + + # description is optional in CRE DB model, empty string allowed. + if not isinstance(self.description, str): + raise ValueError("CandidateCRE: field 'description' must be a string") + + if not re.match(r"\d{3}-\d{3}", self.cre_id): + raise ValueError( + f"CandidateCRE: field 'cre_id' must match pattern 'NNN-NNN', got {self.cre_id!r}" + ) + + if not isinstance(self.score, float): + raise ValueError("CandidateCRE: field 'score' must be a float") + + if not (0.0 <= self.score <= 1.0): + raise ValueError( + f"CandidateCRE: field 'score' must be between 0.0 and 1.0, got {self.score}" + ) + + self.cre_id = self.cre_id.strip() + self.name = self.name.strip() + self.description = self.description.strip() diff --git a/application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py b/application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py new file mode 100644 index 000000000..cf6388f6d --- /dev/null +++ b/application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py @@ -0,0 +1,69 @@ +import numpy as np + +from scipy import sparse +from sklearn.metrics.pairwise import cosine_similarity +from typing import List + +from application.database import db +from application.defs.candidate_cre_defs import CandidateCRE +from application.defs.cheatsheet_defs import CheatsheetRecord +from application.defs import cre_defs +from application.prompt_client import prompt_client + + +def retrieve_candidate_cres( + record: CheatsheetRecord, + cache: db.Node_collection, + ph: prompt_client.PromptHandler, + top_k: int = 20, +) -> List[CandidateCRE]: + """Retrieve top-k CRE candidates for a CheatsheetRecord using embedding similarity. + + Args: + record: the cheatsheet record to find CRE candidates for. + cache: database instance for retrieving CRE embeddings. + ph: PromptHandler instance for generating text embeddings. + top_k: number of top candidates to return, defaults to 20. + + Returns: + List of CandidateCRE objects sorted descending by similarity score. + """ + # build query text from record summary and headings + query_text = record.summary + " " + " ".join(record.headings) + + # embed query text via PromptHandler + query_vector = ph.get_text_embeddings(query_text) + + # fetch all pre-stored CRE vectors from DB + db_cre_embeddings = cache.get_embeddings_by_doc_type(cre_defs.Credoctypes.CRE.value) + + # build ordered lists of internal ids and their vectors + internal_ids = list(db_cre_embeddings.keys()) + cre_matrix = sparse.csr_matrix( + np.array(list(db_cre_embeddings.values())).astype(np.float64) + ) + + # compute cosine similarity between query and all CRE vectors + query_array = sparse.csr_matrix( + np.array(query_vector).reshape(1, -1).astype(np.float64) + ) + similarities = cosine_similarity(query_array, cre_matrix)[0] + + # get top-k indices sorted descending by score + top_k_indices = np.argsort(similarities)[::-1][:top_k] + + # fetch full CRE objects and build CandidateCRE list + candidates = [] + for idx in top_k_indices: + internal_id = internal_ids[idx] + cre = cache.get_cre_by_db_id(internal_id) + candidates.append( + CandidateCRE( + cre_id=cre.id, + name=cre.name, + description=cre.description or "", + score=float(similarities[idx]), + ) + ) + + return candidates