| 147 | return all_results[:top_number] |
| 148 | |
| 149 | def query( |
| 150 | self, |
| 151 | query_text: str, |
| 152 | query_embedding: List[float], |
| 153 | knowledge_id_list: list[str], |
| 154 | document_id_list: list[str], |
| 155 | exclude_document_id_list: list[str], |
| 156 | exclude_paragraph_list: list[str], |
| 157 | is_active: bool, |
| 158 | top_n: int, |
| 159 | similarity: float, |
| 160 | search_mode: SearchMode, |
| 161 | ): |
| 162 | exclude_dict = {} |
| 163 | if knowledge_id_list is None or len(knowledge_id_list) == 0: |
| 164 | return [] |
| 165 | for search_handle in search_handle_list: |
| 166 | if search_handle.support(search_mode): |
| 167 | # Query per knowledge base to leverage per-KB partial HNSW indexes |
| 168 | # (WHERE knowledge_id = '{k_id}'), which won't be used with knowledge_id__in |
| 169 | def build_query_set(kid): |
| 170 | qs = QuerySet(Embedding).filter(knowledge_id=kid, is_active=is_active) |
| 171 | if document_id_list is not None and len(document_id_list) > 0: |
| 172 | qs = qs.filter(document_id__in=document_id_list) |
| 173 | if exclude_document_id_list is not None and len(exclude_document_id_list) > 0: |
| 174 | qs = qs.exclude(document_id__in=exclude_document_id_list) |
| 175 | if exclude_paragraph_list is not None and len(exclude_paragraph_list) > 0: |
| 176 | qs = qs.exclude(paragraph_id__in=exclude_paragraph_list) |
| 177 | qs = qs.exclude(**exclude_dict) |
| 178 | return qs |
| 179 | if len(knowledge_id_list) == 1: |
| 180 | query_set = build_query_set(knowledge_id_list[0]) |
| 181 | return search_handle.handle( |
| 182 | query_set, query_text, query_embedding, top_n, similarity, search_mode, knowledge_id_list |
| 183 | ) |
| 184 | else: |
| 185 | all_results = [] |
| 186 | for kid in knowledge_id_list: |
| 187 | query_set = build_query_set(kid) |
| 188 | results = search_handle.handle( |
| 189 | query_set, query_text, query_embedding, top_n, similarity, search_mode, knowledge_id_list |
| 190 | ) |
| 191 | all_results.extend(results) |
| 192 | all_results.sort(key=lambda x: x.get("similarity", x.get("comprehensive_score", 0)), reverse=True) |
| 193 | return all_results[:top_n] |
| 194 | |
| 195 | def update_by_source_id(self, source_id: str, instance: Dict): |
| 196 | QuerySet(Embedding).filter(source_id=source_id).update(**instance) |