Retrieve relevant guidance entries based on query similarity. Args: query: Query text (e.g., subject or question) top_k: Maximum number of entries to retrieve threshold: Minimum similarity threshold Returns:
(
self,
query: str,
top_k: int = 5,
threshold: float = 0.7,
)
| 68 | save_jsonl(self.entries, self.storage_path, append=False) |
| 69 | |
| 70 | def retrieve( |
| 71 | self, |
| 72 | query: str, |
| 73 | top_k: int = 5, |
| 74 | threshold: float = 0.7, |
| 75 | ) -> List[Dict[str, Any]]: |
| 76 | """ |
| 77 | Retrieve relevant guidance entries based on query similarity. |
| 78 | |
| 79 | Args: |
| 80 | query: Query text (e.g., subject or question) |
| 81 | top_k: Maximum number of entries to retrieve |
| 82 | threshold: Minimum similarity threshold |
| 83 | |
| 84 | Returns: |
| 85 | List of relevant entries with similarity scores |
| 86 | """ |
| 87 | if not self.entries: |
| 88 | return [] |
| 89 | |
| 90 | # Get query embedding |
| 91 | query_embedding = self.llm_client.get_embedding(query) |
| 92 | query_emb = np.array(query_embedding) |
| 93 | |
| 94 | # Filter entries with embeddings |
| 95 | valid_entries = [entry for entry in self.entries if "embedding" in entry] |
| 96 | if not valid_entries: |
| 97 | return [] |
| 98 | |
| 99 | # Vectorized similarity calculation using NumPy |
| 100 | embeddings_matrix = np.array([entry["embedding"] for entry in valid_entries]) |
| 101 | |
| 102 | # Calculate cosine similarities in batch |
| 103 | # cos_sim = dot(A, B) / (norm(A) * norm(B)) |
| 104 | query_norm = np.linalg.norm(query_emb) |
| 105 | if query_norm == 0: |
| 106 | return [] |
| 107 | |
| 108 | # Compute dot products |
| 109 | dot_products = np.dot(embeddings_matrix, query_emb) |
| 110 | |
| 111 | # Compute norms of all embeddings |
| 112 | embedding_norms = np.linalg.norm(embeddings_matrix, axis=1) |
| 113 | |
| 114 | # Avoid division by zero |
| 115 | with np.errstate(divide='ignore', invalid='ignore'): |
| 116 | similarities_array = dot_products / (embedding_norms * query_norm) |
| 117 | similarities_array = np.nan_to_num( |
| 118 | similarities_array, nan=0.0, posinf=0.0, neginf=0.0 |
| 119 | ) |
| 120 | |
| 121 | # Filter by threshold and create result list |
| 122 | valid_indices = np.where(similarities_array >= threshold)[0] |
| 123 | |
| 124 | similarities = [ |
| 125 | { |
| 126 | "entry": valid_entries[idx], |
| 127 | "similarity": float(similarities_array[idx]), |
no test coverage detected