Local doc cache support for document reconstruct #1066
Conversation
| @@ -126,9 +133,21 @@ def to_docs(self, query_params: "BaseDBReader.QueryParams") -> list[Document]: | |||
| doc.properties[DocumentPropertyTypes.SOURCE] = DocumentSource.DB_QUERY | |||
| assert doc.doc_id, "Retrieved invalid doc with missing doc_id" | |||
| if not doc.parent_id: | |||
| if query_params.document_cache is not None: | |||
| cached_doc = query_params.document_cache.get(f"{index_name}:{doc.doc_id}") | |||
There was a problem hiding this comment.
we probably want to move this get_key logic to a shared method?
| # doc_ids = list(unique_docs.keys()) | ||
| doc_ids = [d for d in list(unique_docs.keys()) if d not in list(cached_docs.keys())] | ||
|
|
||
| # We can't safely exclude embeddings since we might need them for 'rerank', e.g. |
There was a problem hiding this comment.
Do we need to do embedding style reranking locally? Won't the knn result just be ordered? Or do we expect to reorder based on something else after
There was a problem hiding this comment.
No. When I wrote the comment, I thought we use embeddings for reranking. In general, though, I don't know if we can safely drop them in the context of document reconstruct. Maybe document reconstruct should be a set of options not just a boolean.
| doc.elements.sort(key=lambda e: e.element_index if e.element_index is not None else float("inf")) | ||
| if doc.doc_id not in cached_docs: | ||
| doc.elements.sort(key=lambda e: e.element_index if e.element_index is not None else float("inf")) | ||
| if query_params.document_cache is not None: |
There was a problem hiding this comment.
can cached_docs have anything if document_cache is None?
There was a problem hiding this comment.
It'll be an empty dict.
| @@ -224,7 +225,7 @@ def opensearch( | |||
| index_name: str, | |||
| query: Optional[Dict] = None, | |||
| reconstruct_document: bool = False, | |||
| query_kwargs: dict[str, Any] = {}, | |||
| query_kwargs=None, | |||
There was a problem hiding this comment.
can you change to an optional? Just for consistency
No description provided.