"""Retriever module."""

from typing import Any, Dict, List, Optional

from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import Document
from langchain.schema.retriever import BaseRetriever


class SearchAPIRetriever(BaseRetriever):
    """Search API retriever."""

    pages: List[Dict] = []

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Get relevant documents."""
        docs = []
        for page in self.pages:
            docs.append(
                Document(
                    page_content=page.get("page_text", ""),
                    metadata={
                        "title": page.get("title", ""),
                        "source": page.get("link", ""),
                    },
                )
            )

        return docs