llama-index

1
from typing import Any, Dict, List, Optional
2

3
from llama_index.legacy.readers.base import BaseReader
4
from llama_index.legacy.schema import Document
5

6

7
class MetalReader(BaseReader):
8
    """Metal reader.
9

10
    Args:
11
        api_key (str): Metal API key.
12
        client_id (str): Metal client ID.
13
        index_id (str): Metal index ID.
14
    """
15

16
    def __init__(self, api_key: str, client_id: str, index_id: str):
17
        import_err_msg = (
18
            "`metal_sdk` package not found, please run `pip install metal_sdk`"
19
        )
20
        try:
21
            import metal_sdk  # noqa
22
        except ImportError:
23
            raise ImportError(import_err_msg)
24
        from metal_sdk.metal import Metal
25

26
        """Initialize with parameters."""
27
        self._api_key = api_key
28
        self._client_id = client_id
29
        self._index_id = index_id
30
        self.metal_client = Metal(api_key, client_id, index_id)
31

32
    def load_data(
33
        self,
34
        limit: int,
35
        query_embedding: Optional[List[float]] = None,
36
        filters: Optional[Dict[str, Any]] = None,
37
        separate_documents: bool = True,
38
        **query_kwargs: Any
39
    ) -> List[Document]:
40
        """Load data from Metal.
41

42
        Args:
43
            query_embedding (Optional[List[float]]): Query embedding for search.
44
            limit (int): Number of results to return.
45
            filters (Optional[Dict[str, Any]]): Filters to apply to the search.
46
            separate_documents (Optional[bool]): Whether to return separate
47
                documents per retrieved entry. Defaults to True.
48
            **query_kwargs: Keyword arguments to pass to the search.
49

50
        Returns:
51
            List[Document]: A list of documents.
52
        """
53
        payload = {
54
            "embedding": query_embedding,
55
            "filters": filters,
56
        }
57
        response = self.metal_client.search(payload, limit=limit, **query_kwargs)
58

59
        documents = []
60
        for item in response["data"]:
61
            text = item["text"] or (item["metadata"] and item["metadata"]["text"])
62
            documents.append(Document(text=text))
63

64
        if not separate_documents:
65
            text_list = [doc.get_content() for doc in documents]
66
            text = "\n\n".join(text_list)
67
            documents = [Document(text=text)]
68

69
        return documents
70