llama-index

1
"""You Retriever."""
2

3
import logging
4
import os
5
from typing import List, Optional
6

7
import requests
8

9
from llama_index.legacy.callbacks.base import CallbackManager
10
from llama_index.legacy.core.base_retriever import BaseRetriever
11
from llama_index.legacy.indices.query.schema import QueryBundle
12
from llama_index.legacy.schema import NodeWithScore, QueryBundle, TextNode
13

14
logger = logging.getLogger(__name__)
15

16

17
class YouRetriever(BaseRetriever):
18
    """You retriever."""
19

20
    def __init__(
21
        self,
22
        api_key: Optional[str] = None,
23
        callback_manager: Optional[CallbackManager] = None,
24
    ) -> None:
25
        """Init params."""
26
        self._api_key = api_key or os.environ["YOU_API_KEY"]
27
        super().__init__(callback_manager)
28

29
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
30
        """Retrieve."""
31
        headers = {"X-API-Key": self._api_key}
32
        results = requests.get(
33
            f"https://api.ydc-index.io/search?query={query_bundle.query_str}",
34
            headers=headers,
35
        ).json()
36

37
        search_hits = ["\n".join(hit["snippets"]) for hit in results["hits"]]
38
        return [NodeWithScore(node=TextNode(text=s), score=1.0) for s in search_hits]
39