diff --git a/docs/extras/integrations/text_embedding/llm_rails.ipynb b/docs/extras/integrations/text_embedding/llm_rails.ipynb new file mode 100644 index 0000000000..e92d3ec09a --- /dev/null +++ b/docs/extras/integrations/text_embedding/llm_rails.ipynb @@ -0,0 +1,133 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "278b6c63", + "metadata": {}, + "source": [ + "# LLMRails\n", + "\n", + "Let's load the LLMRails Embeddings class.\n", + "\n", + "To use LLMRails embedding you need to pass api key by argument or set it in environment with `LLM_RAILS_API_KEY` key.\n", + "To gey API Key you need to sign up in https://console.llmrails.com/signup and then go to https://console.llmrails.com/api-keys and copy key from there after creating one key in platform." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "0be1af71", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.embeddings import LLMRailsEmbeddings" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "2c66e5da", + "metadata": {}, + "outputs": [], + "source": [ + "embeddings = LLMRailsEmbeddings(model='embedding-english-v1') # or embedding-multi-v1" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "01370375", + "metadata": {}, + "outputs": [], + "source": [ + "text = \"This is a test document.\"" + ] + }, + { + "cell_type": "markdown", + "id": "a42e4035", + "metadata": {}, + "source": [ + "To generate embeddings, you can either query an invidivual text, or you can query a list of texts." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "91bc875d-829b-4c3d-8e6f-fc2dda30a3bd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[-0.09996652603149414,\n", + " 0.015568195842206478,\n", + " 0.17670190334320068,\n", + " 0.16521021723747253,\n", + " 0.21193109452724457]" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "query_result = embeddings.embed_query(text)\n", + "query_result[:5]" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "a4b0d49e-0c73-44b6-aed5-5b426564e085", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[-0.04242777079343796,\n", + " 0.016536075621843338,\n", + " 0.10052520781755447,\n", + " 0.18272875249385834,\n", + " 0.2079043835401535]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "doc_result = embeddings.embed_documents([text])\n", + "doc_result[0][:5]" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + }, + "vscode": { + "interpreter": { + "hash": "e971737741ff4ec9aff7dc6155a1060a59a8a6d52c757dbbe66bf8ee389494b1" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/libs/langchain/langchain/embeddings/llm_rails.py b/libs/langchain/langchain/embeddings/llm_rails.py new file mode 100644 index 0000000000..60312384af --- /dev/null +++ b/libs/langchain/langchain/embeddings/llm_rails.py @@ -0,0 +1,72 @@ +""" This file is for LLMRails Embedding """ +import logging +import os +from typing import List, Optional + +import requests + +from langchain.pydantic_v1 import BaseModel, Extra +from langchain.schema.embeddings import Embeddings + + +class LLMRailsEmbeddings(BaseModel, Embeddings): + """LLMRails embedding models. + + To use, you should have the environment + variable ``LLM_RAILS_API_KEY`` set with your API key or pass it + as a named parameter to the constructor. + + Model can be one of ["embedding-english-v1","embedding-multi-v1"] + + Example: + .. code-block:: python + + from langchain.embeddings import LLMRailsEmbeddings + cohere = LLMRailsEmbeddings( + model="embedding-english-v1", api_key="my-api-key" + ) + """ + + model: str = "embedding-english-v1" + """Model name to use.""" + + api_key: Optional[str] = None + """LLMRails API key.""" + + class Config: + """Configuration for this pydantic object.""" + + extra = Extra.forbid + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Call out to Cohere's embedding endpoint. + + Args: + texts: The list of texts to embed. + + Returns: + List of embeddings, one for each text. + """ + api_key = self.api_key or os.environ.get("LLM_RAILS_API_KEY") + if api_key is None: + logging.warning("Can't find LLMRails credentials in environment.") + raise ValueError("LLM_RAILS_API_KEY is not set") + + response = requests.post( + "https://api.llmrails.com/v1/embeddings", + headers={"X-API-KEY": api_key}, + json={"input": texts, "model": self.model}, + timeout=60, + ) + return [item["embedding"] for item in response.json()["data"]] + + def embed_query(self, text: str) -> List[float]: + """Call out to Cohere's embedding endpoint. + + Args: + text: The text to embed. + + Returns: + Embeddings for the text. + """ + return self.embed_documents([text])[0]