mirror of https://github.com/hwchase17/langchain
LLMRails (#10796)
### LLMRails Integration This PR provides integration with LLMRails. Implemented here are: langchain/vectorstore/llm_rails.py tests/integration_tests/vectorstores/test_llm_rails.py docs/extras/integrations/vectorstores/llm-rails.ipynb --------- Co-authored-by: Anar Aliyev <aaliyev@mgmt.cloudnet.services> Co-authored-by: Bagatur <baskaryan@gmail.com>pull/10771/head^2
parent
900dbd1cbe
commit
c656a6b966
@ -0,0 +1,203 @@
|
|||||||
|
"""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.vectorstores.base import VectorStore, VectorStoreRetriever
|
||||||
|
|
||||||
|
|
||||||
|
class ModelChoices(str, Enum):
|
||||||
|
embedding_english_v1 = "embedding-english-v1"
|
||||||
|
embedding_multi_v1 = "embedding-multi-v1"
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
|
||||||
|
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,
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
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]
|
||||||
|
|
||||||
|
@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.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
|
||||||
|
|
||||||
|
def as_retriever(self, **kwargs: Any) -> LLMRailsRetriever:
|
||||||
|
return LLMRailsRetriever(vectorstore=self, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class LLMRailsRetriever(VectorStoreRetriever):
|
||||||
|
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 .
|
||||||
|
"""
|
||||||
|
|
||||||
|
def add_texts(self, texts: List[str]) -> None:
|
||||||
|
"""Add text to the datastore.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
texts (List[str]): The text
|
||||||
|
"""
|
||||||
|
self.vectorstore.add_texts(texts)
|
@ -0,0 +1,35 @@
|
|||||||
|
from langchain.vectorstores.llm_rails import LLMRails
|
||||||
|
|
||||||
|
#
|
||||||
|
# For this test to run properly, please setup as follows:
|
||||||
|
# 1. Create a LLMRails account: sign up at https://console.llmrails.com/signup
|
||||||
|
# 2. Create an API_KEY for this corpus with permissions for query and indexing
|
||||||
|
# 3. Create a datastorea and get its id from datastore setting
|
||||||
|
# 3. Setup environment variable:
|
||||||
|
# LLM_RAILS_API_KEY, LLM_RAILS_DATASTORE_ID
|
||||||
|
#
|
||||||
|
|
||||||
|
|
||||||
|
def test_llm_rails_add_documents() -> None:
|
||||||
|
"""Test end to end construction and search."""
|
||||||
|
|
||||||
|
# create a new Vectara instance
|
||||||
|
docsearch: LLMRails = LLMRails()
|
||||||
|
|
||||||
|
# start with some initial texts, added with add_texts
|
||||||
|
texts1 = ["large language model", "information retrieval", "question answering"]
|
||||||
|
docsearch.add_texts(texts1)
|
||||||
|
|
||||||
|
# test without filter
|
||||||
|
output1 = docsearch.similarity_search("large language model", k=1)
|
||||||
|
|
||||||
|
print(output1)
|
||||||
|
assert len(output1) == 1
|
||||||
|
assert output1[0].page_content == "large language model"
|
||||||
|
|
||||||
|
# test without filter but with similarity score
|
||||||
|
output2 = docsearch.similarity_search_with_score("large language model", k=1)
|
||||||
|
|
||||||
|
assert len(output2) == 1
|
||||||
|
assert output2[0][0].page_content == "large language model"
|
||||||
|
assert output2[0][1] > 0
|
Loading…
Reference in New Issue