Source code for langchain.retrievers.vespa_retriever
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Sequence, Union
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
if TYPE_CHECKING:
    from vespa.application import Vespa
[docs]class VespaRetriever(BaseRetriever):
    """`Vespa` retriever."""
    app: Vespa
    """Vespa application to query."""
    body: Dict
    """Body of the query."""
    content_field: str
    """Name of the content field."""
    metadata_fields: Sequence[str]
    """Names of the metadata fields."""
    def _query(self, body: Dict) -> List[Document]:
        response = self.app.query(body)
        if not str(response.status_code).startswith("2"):
            raise RuntimeError(
                "Could not retrieve data from Vespa. Error code: {}".format(
                    response.status_code
                )
            )
        root = response.json["root"]
        if "errors" in root:
            raise RuntimeError(json.dumps(root["errors"]))
        docs = []
        for child in response.hits:
            page_content = child["fields"].pop(self.content_field, "")
            if self.metadata_fields == "*":
                metadata = child["fields"]
            else:
                metadata = {mf: child["fields"].get(mf) for mf in self.metadata_fields}
            metadata["id"] = child["id"]
            docs.append(Document(page_content=page_content, metadata=metadata))
        return docs
    def _get_relevant_documents(