llama-index

1
"""Pathway reader."""
2

3
from typing import List, Optional, Union
4

5
from llama_index.legacy.readers.base import BaseReader
6
from llama_index.legacy.schema import Document
7

8

9
class PathwayReader(BaseReader):
10
    """Pathway reader.
11

12
    Retrieve documents from Pathway data indexing pipeline.
13

14
    Args:
15
        host (str): The URI where Pathway is currently hosted.
16
        port (str | int): The port number on which Pathway is listening.
17

18
    See Also:
19
        llamaindex.retriever.pathway.PathwayRetriever and,
20
        llamaindex.retriever.pathway.PathwayVectorServer
21
    """
22

23
    def __init__(self, host: str, port: Union[str, int]):
24
        """Initializing the Pathway reader client."""
25
        import_err_msg = "`pathway` package not found, please run `pip install pathway`"
26
        try:
27
            from pathway.xpacks.llm.vector_store import VectorStoreClient
28
        except ImportError:
29
            raise ImportError(import_err_msg)
30
        self.client = VectorStoreClient(host, port)
31

32
    def load_data(
33
        self,
34
        query_text: str,
35
        k: Optional[int] = 4,
36
        metadata_filter: Optional[str] = None,
37
    ) -> List[Document]:
38
        """Load data from Pathway.
39

40
        Args:
41
            query_text (str): The text to get the closest neighbors of.
42
            k (int): Number of results to return.
43
            metadata_filter (str): Filter to be applied.
44

45
        Returns:
46
            List[Document]: A list of documents.
47
        """
48
        results = self.client(query_text, k, metadata_filter)
49
        documents = []
50
        for return_elem in results:
51
            document = Document(
52
                text=return_elem["text"],
53
                extra_info=return_elem["metadata"],
54
            )
55

56
            documents.append(document)
57

58
        return documents
59