From feefa3d2dc8a9e4dc3dec476ab4ddf003f2ecc5d Mon Sep 17 00:00:00 2001 From: Abhijeet Saharan Date: Tue, 23 Jun 2026 16:08:37 +0530 Subject: [PATCH 1/2] feat: add candidate retrieval checkpoints D1 & D2 Signed-off-by: Abhijeet Saharan --- application/defs/candidate_cre_defs.py | 44 +++++++++++ .../parsers/cheatsheet_cre_retriever.py | 74 +++++++++++++++++++ 2 files changed, 118 insertions(+) create mode 100644 application/defs/candidate_cre_defs.py create mode 100644 application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py diff --git a/application/defs/candidate_cre_defs.py b/application/defs/candidate_cre_defs.py new file mode 100644 index 000000000..a359bab0b --- /dev/null +++ b/application/defs/candidate_cre_defs.py @@ -0,0 +1,44 @@ +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() \ No newline at end of file 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..c32e3f21c --- /dev/null +++ b/application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py @@ -0,0 +1,74 @@ +import logging +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 + +logger = logging.getLogger(__name__) + + +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 \ No newline at end of file From 5d098e896afe7ecc1c6234498d9072e9fba9c18f Mon Sep 17 00:00:00 2001 From: Abhijeet Saharan Date: Wed, 24 Jun 2026 10:21:16 +0530 Subject: [PATCH 2/2] fix: address lint test issues Signed-off-by: Abhijeet Saharan --- application/defs/candidate_cre_defs.py | 12 ++++-------- .../parsers/cheatsheet_cre_retriever.py | 9 ++------- 2 files changed, 6 insertions(+), 15 deletions(-) diff --git a/application/defs/candidate_cre_defs.py b/application/defs/candidate_cre_defs.py index a359bab0b..0c8218858 100644 --- a/application/defs/candidate_cre_defs.py +++ b/application/defs/candidate_cre_defs.py @@ -18,11 +18,9 @@ def __post_init__(self): f"CandidateCRE: field '{field_name}' must be a non-empty string" ) - # description is optional in CRE DB model, empty string allowed + # 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" - ) + raise ValueError("CandidateCRE: field 'description' must be a string") if not re.match(r"\d{3}-\d{3}", self.cre_id): raise ValueError( @@ -30,9 +28,7 @@ def __post_init__(self): ) if not isinstance(self.score, float): - raise ValueError( - "CandidateCRE: field 'score' must be a float" - ) + raise ValueError("CandidateCRE: field 'score' must be a float") if not (0.0 <= self.score <= 1.0): raise ValueError( @@ -41,4 +37,4 @@ def __post_init__(self): self.cre_id = self.cre_id.strip() self.name = self.name.strip() - self.description = self.description.strip() \ No newline at end of file + 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 index c32e3f21c..cf6388f6d 100644 --- a/application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py +++ b/application/utils/external_project_parsers/parsers/cheatsheet_cre_retriever.py @@ -1,4 +1,3 @@ -import logging import numpy as np from scipy import sparse @@ -11,8 +10,6 @@ from application.defs import cre_defs from application.prompt_client import prompt_client -logger = logging.getLogger(__name__) - def retrieve_candidate_cres( record: CheatsheetRecord, @@ -38,9 +35,7 @@ def retrieve_candidate_cres( 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 - ) + 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()) @@ -71,4 +66,4 @@ def retrieve_candidate_cres( ) ) - return candidates \ No newline at end of file + return candidates