"""Wrapper around LLMRails vector database."""
from __future__ import annotations
import json
import logging
import os
import uuid
from typing import Any, Iterable, List, Optional, Tuple
import requests
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
from pydantic import Field
[docs]
class LLMRails(VectorStore):
"""Implementation of Vector Store using LLMRails.
See https://llmrails.com/
Example:
.. code-block:: python
from langchain_community.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:
"""Returns headers that should be attached to each post request."""
return {"X-API-KEY": self._api_key}
[docs]
def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> List[str]:
"""Run more texts through the embeddings and add to the vectorstore.
Args:
texts: Iterable of strings to add to the vectorstore.
Returns:
List of ids from adding the texts into the vectorstore.
"""
names: List[str] = []
for text in texts:
doc_name = str(uuid.uuid4())
response = self._session.post(
f"{self.base_url}/datastores/{self._datastore_id}/text",
json={"name": doc_name, "text": text},
verify=True,
headers=self._get_post_headers(),
)
if response.status_code != 200:
logging.error(
f"Create request failed for doc_name = {doc_name} with status code "
f"{response.status_code}, reason {response.reason}, text "
f"{response.text}"
)
return names
names.append(doc_name)
return names
[docs]
def add_files(
self,
files_list: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> bool:
"""
LLMRails provides a way to add documents directly via our API where
pre-processing and chunking occurs internally in an optimal way
This method provides a way to use that API in LangChain
Args:
files_list: Iterable of strings, each representing a local file path.
Files could be text, HTML, PDF, markdown, doc/docx, ppt/pptx, etc.
see API docs for full list
Returns:
List of ids associated with each of the files indexed
"""
files = []
for file in files_list:
if not os.path.exists(file):
logging.error(f"File {file} does not exist, skipping")
continue
files.append(("file", (os.path.basename(file), open(file, "rb"))))
response = self._session.post(
f"{self.base_url}/datastores/{self._datastore_id}/file",
files=files,
verify=True,
headers=self._get_post_headers(),
)
if response.status_code != 200:
logging.error(
f"Create request failed for datastore = {self._datastore_id} "
f"with status code {response.status_code}, reason {response.reason}, "
f"text {response.text}"
)
return False
return True
[docs]
def similarity_search_with_score(
self, query: str, k: int = 5
) -> List[Tuple[Document, float]]:
"""Return LLMRails documents most similar to query, along with scores.
Args:
query: Text to look up documents similar to.
k: Number of Documents to return. Defaults to 5 Max 10.
alpha: parameter for hybrid search .
Returns:
List of Documents most similar to the query and score for each.
"""
response = self._session.post(
headers=self._get_post_headers(),
url=f"{self.base_url}/datastores/{self._datastore_id}/search",
data=json.dumps({"k": k, "text": query}),
timeout=10,
)
if response.status_code != 200:
logging.error(
"Query failed %s",
f"(code {response.status_code}, reason {response.reason}, details "
f"{response.text})",
)
return []
results = response.json()["results"]
docs = [
(
Document(
page_content=x["text"],
metadata={
key: value
for key, value in x["metadata"].items()
if key != "score"
},
),
x["metadata"]["score"],
)
for x in results
]
return docs
[docs]
def similarity_search(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Document]:
"""Return LLMRails documents most similar to query, along with scores.
Args:
query: Text to look up documents similar to.
k: Number of Documents to return. Defaults to 5.
Returns:
List of Documents most similar to the query
"""
docs_and_scores = self.similarity_search_with_score(query, k=k)
return [doc for doc, _ in docs_and_scores]
[docs]
@classmethod
def from_texts(
cls,
texts: List[str],
embedding: Optional[Embeddings] = None,
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> LLMRails:
"""Construct LLMRails wrapper from raw documents.
This is intended to be a quick way to get started.
Example:
.. code-block:: python
from langchain_community.vectorstores import LLMRails
llm_rails = LLMRails.from_texts(
texts,
datastore_id=datastore_id,
api_key=llm_rails_api_key
)
"""
# Note: LLMRails generates its own embeddings, so we ignore the provided
# embeddings (required by interface)
llm_rails = cls(**kwargs)
llm_rails.add_texts(texts)
return llm_rails
[docs]
def as_retriever(self, **kwargs: Any) -> LLMRailsRetriever:
return LLMRailsRetriever(vectorstore=self, **kwargs)
[docs]
class LLMRailsRetriever(VectorStoreRetriever): # type: ignore[override]
"""Retriever for LLMRails."""
vectorstore: LLMRails
search_kwargs: dict = Field(default_factory=lambda: {"k": 5})
"""Search params.
k: Number of Documents to return. Defaults to 5.
alpha: parameter for hybrid search .
"""
[docs]
def add_texts(self, texts: List[str]) -> None:
"""Add text to the datastore.
Args:
texts (List[str]): The text
"""
self.vectorstore.add_texts(texts)
