Source code for langchain.vectorstores.llm_rails
"""Wrapper around LLMRails vector database."""
from __future__ import annotations
import json
import logging
import os
import uuid
from enum import Enum
from typing import Any, Iterable, List, Optional, Tuple
import requests
from langchain.pydantic_v1 import Field
from langchain.schema import Document
from langchain.schema.embeddings import Embeddings
from langchain.schema.vectorstore import VectorStore, VectorStoreRetriever
[docs]class ModelChoices(str, Enum):
    embedding_english_v1 = "embedding-english-v1"
    embedding_multi_v1 = "embedding-multi-v1"
[docs]class LLMRails(VectorStore):
    """Implementation of Vector Store using LLMRails (https://llmrails.com/).
    Example:
        .. code-block:: python
            from langchain.vectorstores import LLMRails
            vectorstore = LLMRails(
                api_key=llm_rails_api_key,
                datastore_id=datastore_id
            )
    """
[docs]    def __init__(
        self,
        datastore_id: Optional[str] = None,
        api_key: Optional[str] = None,
    ):
        """Initialize with LLMRails API."""
        self._datastore_id = datastore_id or os.environ.get("LLM_RAILS_DATASTORE_ID")
        self._api_key = api_key or os.environ.get("LLM_RAILS_API_KEY")
        if self._api_key is None:
            logging.warning("Can't find Rails credentials in environment.")
        self._session = requests.Session()  # to reuse connections
        self.datastore_id = datastore_id
        self.base_url = "https://api.llmrails.com/v1"
    def _get_post_headers(self) -> dict: