Repository: Cinnamon/kotaemon
File: libs/kotaemon/kotaemon/indices/vectorindex.py (Python)

Code before:

        for item in cast(list, text):
            if isinstance(item, str):
                input_.append(Document(text=item, id_=str(uuid.uuid4())))
            elif isinstance(item, Document):
                input_.append(item)
            else:
                raise ValueError(
                    f"Invalid input type {type(item)}, should be str or Document"
                )

        print(f"Getting embeddings for {len(input_)} nodes")
        embeddings = self.embedding(input_)
        print("Adding embeddings to vector store")
        self.vector_store.add(
            embeddings=embeddings,
            ids=[t.doc_id for t in input_],
        )
        if self.doc_store:
            print("Adding documents to doc store")
            self.doc_store.add(input_)
        # save the text content into markdown format
        if self.cache_dir is not None:
            file_name, _ = os.path.splitext(input_[0].metadata['file_name'])
            for i in range(len(input_)):
                with open(self.cache_dir / f"{file_name}_{i}.md", "w") as f:
                    f.write(input_[i].text)
                


class VectorRetrieval(BaseRetrieval):
    """Retrieve list of documents from vector store"""

    vector_store: BaseVectorStore
    doc_store: Optional[BaseDocumentStore] = None
    embedding: BaseEmbeddings
    rerankers: Sequence[BaseReranking] = []
    top_k: int = 5
    retrieval_mode: str = "hybrid"  # vector, text, hybrid

    def run(
        self, text: str | Document, top_k: Optional[int] = None, **kwargs
    ) -> list[RetrievedDocument]:
        """Retrieve a list of documents from vector store

        Args:
            text: the text to retrieve similar documents
            top_k: number of top similar documents to return

        Returns:
            list[RetrievedDocument]: list of retrieved documents

Review comment (performance):
Since `self.cache_dir` can be string and will not work with `/` operator, should use `Path(self.cache_dir)` instead.

Code after:
        input_: list[Document] = []
        if not isinstance(text, list):
            text = [text]

        for item in cast(list, text):
            if isinstance(item, str):
                input_.append(Document(text=item, id_=str(uuid.uuid4())))
            elif isinstance(item, Document):
                input_.append(item)
            else:
                raise ValueError(
                    f"Invalid input type {type(item)}, should be str or Document"
                )

        print(f"Getting embeddings for {len(input_)} nodes")
        embeddings = self.embedding(input_)
        print("Adding embeddings to vector store")
        self.vector_store.add(
            embeddings=embeddings,
            ids=[t.doc_id for t in input_],
        )
        if self.doc_store:
            print("Adding documents to doc store")
            self.doc_store.add(input_)
        # save the chunks content into markdown format
        if self.cache_dir is not None:
            file_name = Path(input_[0].metadata["file_name"])
            for i in range(len(input_)):
                with open(Path(self.cache_dir) / f"{file_name.stem}_{i}.md", "w") as f:
                    f.write(input_[i].text)


class VectorRetrieval(BaseRetrieval):
    """Retrieve list of documents from vector store"""

    vector_store: BaseVectorStore
    doc_store: Optional[BaseDocumentStore] = None
    embedding: BaseEmbeddings
    rerankers: Sequence[BaseReranking] = []
    top_k: int = 5
    retrieval_mode: str = "hybrid"  # vector, text, hybrid

    def run(
        self, text: str | Document, top_k: Optional[int] = None, **kwargs
    ) -> list[RetrievedDocument]:
        """Retrieve a list of documents from vector store

        Args:
            text: the text to retrieve similar documents
            top_k: number of top similar documents to return