From 96023f94d9db48b4ca3d6849711f030010c2454b Mon Sep 17 00:00:00 2001 From: Sharath Rajasekar Date: Wed, 20 Sep 2023 16:36:39 -0700 Subject: [PATCH] Add Javelin integration (#10275) We are introducing the py integration to Javelin AI Gateway www.getjavelin.io. Javelin is an enterprise-scale fast llm router & gateway. Could you please review and let us know if there is anything missing. Javelin AI Gateway wraps Embedding, Chat and Completion LLMs. Uses javelin_sdk under the covers (pip install javelin_sdk). Author: Sharath Rajasekar, Twitter: @sharathr, @javelinai Thanks!! --- docs/api_reference/guide_imports.json | 11 + docs/extras/integrations/llms/javelin.ipynb | 242 ++++++++++++++++++ .../providers/javelin_ai_gateway.mdx | 92 +++++++ .../langchain/chat_models/__init__.py | 2 + .../chat_models/javelin_ai_gateway.py | 223 ++++++++++++++++ .../langchain/embeddings/__init__.py | 2 + .../embeddings/javelin_ai_gateway.py | 110 ++++++++ libs/langchain/langchain/llms/__init__.py | 3 + .../langchain/llms/javelin_ai_gateway.py | 152 +++++++++++ 9 files changed, 837 insertions(+) create mode 100644 docs/extras/integrations/llms/javelin.ipynb create mode 100644 docs/extras/integrations/providers/javelin_ai_gateway.mdx create mode 100644 libs/langchain/langchain/chat_models/javelin_ai_gateway.py create mode 100644 libs/langchain/langchain/embeddings/javelin_ai_gateway.py create mode 100644 libs/langchain/langchain/llms/javelin_ai_gateway.py diff --git a/docs/api_reference/guide_imports.json b/docs/api_reference/guide_imports.json index 8e4d0fed32..7218665180 100644 --- a/docs/api_reference/guide_imports.json +++ b/docs/api_reference/guide_imports.json @@ -465,6 +465,7 @@ "PromptLayer": "https://python.langchain.com/docs/integrations/callbacks/promptlayer", "Log10": "https://python.langchain.com/docs/integrations/providers/log10", "MLflow AI Gateway": "https://python.langchain.com/docs/integrations/providers/mlflow_ai_gateway", + "Javelin AI Gateway": "https://python.langchain.com/docs/integrations/providers/javelin_ai_gateway", "Flyte": "https://python.langchain.com/docs/integrations/providers/flyte", "Arthur": "https://python.langchain.com/docs/integrations/providers/arthur_tracking", "Chatbots": "https://python.langchain.com/docs/use_cases/chatbots", @@ -1245,6 +1246,7 @@ "Context": "https://python.langchain.com/docs/integrations/callbacks/context", "Label Studio": "https://python.langchain.com/docs/integrations/callbacks/labelstudio", "MLflow AI Gateway": "https://python.langchain.com/docs/integrations/providers/mlflow_ai_gateway", + "Javelin AI Gateway": "https://python.langchain.com/docs/integrations/providers/javelin_ai_gateway", "Chatbots": "https://python.langchain.com/docs/use_cases/chatbots", "Conversational Retrieval Agent": "https://python.langchain.com/docs/use_cases/question_answering/how_to/conversational_retrieval_agents", "Structure answers with OpenAI functions": "https://python.langchain.com/docs/use_cases/question_answering/integrations/openai_functions_retrieval_qa", @@ -1879,6 +1881,15 @@ "ChatMLflowAIGateway": { "MLflow AI Gateway": "https://python.langchain.com/docs/integrations/providers/mlflow_ai_gateway" }, + "JavelinAIGateway": { + "Javelin AI Gateway": "https://python.langchain.com/docs/integrations/providers/javelin_ai_gateway" + }, + "JavelinAIGatewayEmbeddings": { + "Javelin AI Gateway": "https://python.langchain.com/docs/integrations/providers/javelin_ai_gateway" + }, + "ChatJavelinAIGateway": { + "Javelin AI Gateway": "https://python.langchain.com/docs/integrations/providers/javelin_ai_gateway" + }, "SingleStoreDB": { "SingleStoreDB": "https://python.langchain.com/docs/integrations/vectorstores/singlestoredb" }, diff --git a/docs/extras/integrations/llms/javelin.ipynb b/docs/extras/integrations/llms/javelin.ipynb new file mode 100644 index 0000000000..1201e47d6f --- /dev/null +++ b/docs/extras/integrations/llms/javelin.ipynb @@ -0,0 +1,242 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "62bacc68-1976-44eb-9316-d5baf54bf595", + "metadata": {}, + "source": [ + "# Javelin AI Gateway Tutorial\n", + "\n", + "This Jupyter Notebook will explore how to interact with the Javelin AI Gateway using the Python SDK. \n", + "The Javelin AI Gateway facilitates the utilization of large language models (LLMs) like OpenAI, Cohere, Anthropic, and others by \n", + "providing a secure and unified endpoint. The gateway itself provides a centralized mechanism to roll out models systematically, \n", + "provide access security, policy & cost guardrails for enterprises, etc., \n", + "\n", + "For a complete listing of all the features & benefits of Javelin, please visit www.getjavelin.io\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "e52185f8-132b-4585-b73d-6fee928ac199", + "metadata": {}, + "source": [ + "## Step 1: Introduction\n", + "[The Javelin AI Gateway](https://www.getjavelin.io) is an enterprise-grade API Gateway for AI applications. It integrates robust access security, ensuring secure interactions with large language models. Learn more in the [official documentation](https://docs.getjavelin.io).\n" + ] + }, + { + "cell_type": "markdown", + "id": "2e2acdb3-e3b8-422b-b077-7a0d63d18349", + "metadata": {}, + "source": [ + "## Step 2: Installation\n", + "Before we begin, we must install the `javelin_sdk` and set up the Javelin API key as an environment variable. " + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "e91518a4-43ce-443e-b4c0-dbc652eb749f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Requirement already satisfied: javelin_sdk in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (0.1.8)\n", + "Requirement already satisfied: httpx<0.25.0,>=0.24.0 in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from javelin_sdk) (0.24.1)\n", + "Requirement already satisfied: pydantic<2.0.0,>=1.10.7 in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from javelin_sdk) (1.10.12)\n", + "Requirement already satisfied: certifi in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from httpx<0.25.0,>=0.24.0->javelin_sdk) (2023.5.7)\n", + "Requirement already satisfied: httpcore<0.18.0,>=0.15.0 in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from httpx<0.25.0,>=0.24.0->javelin_sdk) (0.17.3)\n", + "Requirement already satisfied: idna in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from httpx<0.25.0,>=0.24.0->javelin_sdk) (3.4)\n", + "Requirement already satisfied: sniffio in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from httpx<0.25.0,>=0.24.0->javelin_sdk) (1.3.0)\n", + "Requirement already satisfied: typing-extensions>=4.2.0 in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from pydantic<2.0.0,>=1.10.7->javelin_sdk) (4.7.1)\n", + "Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from httpcore<0.18.0,>=0.15.0->httpx<0.25.0,>=0.24.0->javelin_sdk) (0.14.0)\n", + "Requirement already satisfied: anyio<5.0,>=3.0 in /usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages (from httpcore<0.18.0,>=0.15.0->httpx<0.25.0,>=0.24.0->javelin_sdk) (3.7.1)\n", + "Note: you may need to restart the kernel to use updated packages.\n" + ] + } + ], + "source": [ + "pip install 'javelin_sdk'" + ] + }, + { + "cell_type": "markdown", + "id": "53b546dc-9ca3-4602-9a7b-d733d99e8e2f", + "metadata": {}, + "source": [ + "## Step 3: Completions Example\n", + "This section will demonstrate how to interact with the Javelin AI Gateway to get completions from a large language model. Here is a Python script that demonstrates this:\n", + "(note) assumes that you have setup a route in the gateway called 'eng_dept03'" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "d36949f0-5354-44ca-9a31-70c769344319", + "metadata": {}, + "outputs": [ + { + "ename": "ImportError", + "evalue": "cannot import name 'JavelinAIGateway' from 'langchain.llms' (/usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages/langchain/llms/__init__.py)", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mImportError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[6], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mchains\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LLMChain\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mllms\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m JavelinAIGateway\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprompts\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m PromptTemplate\n\u001b[1;32m 5\u001b[0m route_completions \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124meng_dept03\u001b[39m\u001b[38;5;124m\"\u001b[39m\n", + "\u001b[0;31mImportError\u001b[0m: cannot import name 'JavelinAIGateway' from 'langchain.llms' (/usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages/langchain/llms/__init__.py)" + ] + } + ], + "source": [ + "from langchain.chains import LLMChain\n", + "from langchain.llms import JavelinAIGateway\n", + "from langchain.prompts import PromptTemplate\n", + "\n", + "route_completions = \"eng_dept03\"\n", + "\n", + "gateway = JavelinAIGateway(\n", + " gateway_uri=\"http://localhost:8000\", # replace with service URL or host/port of Javelin\n", + " route=route_completions,\n", + " model_name=\"text-davinci-003\",\n", + ")\n", + "\n", + "prompt = PromptTemplate(\"Translate the following English text to French: {text}\")\n", + "\n", + "llmchain = LLMChain(llm=gateway, prompt=prompt)\n", + "result = llmchain.run(\"podcast player\")\n", + "\n", + "print(result)\n" + ] + }, + { + "cell_type": "markdown", + "id": "6b63fe93-2e77-4ea9-b8e7-dec2b96b8e95", + "metadata": {}, + "source": [ + "# Step 4: Embeddings Example\n", + "This section demonstrates how to use the Javelin AI Gateway to obtain embeddings for text queries and documents. Here is a Python script that illustrates this:\n", + "(note) assumes that you have setup a route in the gateway called 'embeddings'" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "878e6c1d-be7f-49de-825c-43c266c8714e", + "metadata": {}, + "outputs": [ + { + "ename": "ImportError", + "evalue": "cannot import name 'JavelinAIGatewayEmbeddings' from 'langchain.embeddings' (/usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages/langchain/embeddings/__init__.py)", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mImportError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[9], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01membeddings\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m JavelinAIGatewayEmbeddings\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01membeddings\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mopenai\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m OpenAIEmbeddings\n\u001b[1;32m 4\u001b[0m embeddings \u001b[38;5;241m=\u001b[39m JavelinAIGatewayEmbeddings(\n\u001b[1;32m 5\u001b[0m gateway_uri\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhttp://localhost:8000\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;66;03m# replace with service URL or host/port of Javelin\u001b[39;00m\n\u001b[1;32m 6\u001b[0m route\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124membeddings\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 7\u001b[0m )\n", + "\u001b[0;31mImportError\u001b[0m: cannot import name 'JavelinAIGatewayEmbeddings' from 'langchain.embeddings' (/usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages/langchain/embeddings/__init__.py)" + ] + } + ], + "source": [ + "from langchain.embeddings import JavelinAIGatewayEmbeddings\n", + "from langchain.embeddings.openai import OpenAIEmbeddings\n", + "\n", + "embeddings = JavelinAIGatewayEmbeddings(\n", + " gateway_uri=\"http://localhost:8000\", # replace with service URL or host/port of Javelin\n", + " route=\"embeddings\",\n", + ")\n", + "\n", + "print(embeddings.embed_query(\"hello\"))\n", + "print(embeddings.embed_documents([\"hello\"]))\n" + ] + }, + { + "cell_type": "markdown", + "id": "07c6691b-d333-4598-b2b7-c0933ed75937", + "metadata": {}, + "source": [ + "# Step 5: Chat Example\n", + "This section illustrates how to interact with the Javelin AI Gateway to facilitate a chat with a large language model. Here is a Python script that demonstrates this:\n", + "(note) assumes that you have setup a route in the gateway called 'mychatbot_route'" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "653ef88c-36cd-4730-9c12-43c246b551f1", + "metadata": {}, + "outputs": [ + { + "ename": "ImportError", + "evalue": "cannot import name 'ChatJavelinAIGateway' from 'langchain.chat_models' (/usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages/langchain/chat_models/__init__.py)", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mImportError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mchat_models\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ChatJavelinAIGateway\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HumanMessage, SystemMessage\n\u001b[1;32m 4\u001b[0m messages \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 5\u001b[0m SystemMessage(\n\u001b[1;32m 6\u001b[0m content\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou are a helpful assistant that translates English to French.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 10\u001b[0m ),\n\u001b[1;32m 11\u001b[0m ]\n", + "\u001b[0;31mImportError\u001b[0m: cannot import name 'ChatJavelinAIGateway' from 'langchain.chat_models' (/usr/local/Caskroom/miniconda/base/lib/python3.11/site-packages/langchain/chat_models/__init__.py)" + ] + } + ], + "source": [ + "from langchain.chat_models import ChatJavelinAIGateway\n", + "from langchain.schema import HumanMessage, SystemMessage\n", + "\n", + "messages = [\n", + " SystemMessage(\n", + " content=\"You are a helpful assistant that translates English to French.\"\n", + " ),\n", + " HumanMessage(\n", + " content=\"Artificial Intelligence has the power to transform humanity and make the world a better place\"\n", + " ),\n", + "]\n", + "\n", + "chat = ChatJavelinAIGateway(\n", + " gateway_uri=\"http://localhost:8000\", # replace with service URL or host/port of Javelin\n", + " route=\"mychatbot_route\",\n", + " model_name=\"gpt-3.5-turbo\",\n", + " params={\n", + " \"temperature\": 0.1\n", + " }\n", + ")\n", + "\n", + "print(chat(messages))\n" + ] + }, + { + "cell_type": "markdown", + "id": "6eb9cf33-6505-4e05-808b-645856463a8e", + "metadata": {}, + "source": [ + "Step 6: Conclusion\n", + "This tutorial introduced the Javelin AI Gateway and demonstrated how to interact with it using the Python SDK. \n", + "Remember to check the Javelin [Python SDK](https://www.github.com/getjavelin.io/javelin-python) for more examples and to explore the official documentation for additional details.\n", + "\n", + "Happy coding!" + ] + } + ], + "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.4" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/extras/integrations/providers/javelin_ai_gateway.mdx b/docs/extras/integrations/providers/javelin_ai_gateway.mdx new file mode 100644 index 0000000000..7aae14365b --- /dev/null +++ b/docs/extras/integrations/providers/javelin_ai_gateway.mdx @@ -0,0 +1,92 @@ +# Javelin AI Gateway + +[The Javelin AI Gateway](https://www.getjavelin.io) service is a high-performance, enterprise grade API Gateway for AI applications. +It is designed to streamline the usage and access of various large language model (LLM) providers, +such as OpenAI, Cohere, Anthropic and custom large language models within an organization by incorporating +robust access security for all interactions with LLMs. + +Javelin offers a high-level interface that simplifies the interaction with LLMs by providing a unified endpoint +to handle specific LLM related requests. + +See the Javelin AI Gateway [documentation](https://docs.getjavelin.io) for more details. +[Javelin Python SDK](https://www.github.com/getjavelin/javelin-python) is an easy to use client library meant to be embedded into AI Applications + +## Installation and Setup + +Install `javelin_sdk` to interact with Javelin AI Gateway: + +```sh +pip install 'javelin_sdk' +``` + +Set the Javelin's API key as an environment variable: + +```sh +export JAVELIN_API_KEY=... +``` + +## Completions Example + +```python + +from langchain.chains import LLMChain +from langchain.llms import JavelinAIGateway +from langchain.prompts import PromptTemplate + +route_completions = "eng_dept03" + +gateway = JavelinAIGateway( + gateway_uri="http://localhost:8000", + route=route_completions, + model_name="text-davinci-003", +) + +llmchain = LLMChain(llm=gateway, prompt=prompt) +result = llmchain.run("podcast player") + +print(result) + +``` + +## Embeddings Example + +```python +from langchain.embeddings import JavelinAIGatewayEmbeddings +from langchain.embeddings.openai import OpenAIEmbeddings + +embeddings = JavelinAIGatewayEmbeddings( + gateway_uri="http://localhost:8000", + route="embeddings", +) + +print(embeddings.embed_query("hello")) +print(embeddings.embed_documents(["hello"])) +``` + +## Chat Example +```python +from langchain.chat_models import ChatJavelinAIGateway +from langchain.schema import HumanMessage, SystemMessage + +messages = [ + SystemMessage( + content="You are a helpful assistant that translates English to French." + ), + HumanMessage( + content="Artificial Intelligence has the power to transform humanity and make the world a better place" + ), +] + +chat = ChatJavelinAIGateway( + gateway_uri="http://localhost:8000", + route="mychatbot_route", + model_name="gpt-3.5-turbo" + params={ + "temperature": 0.1 + } +) + +print(chat(messages)) + +``` + diff --git a/libs/langchain/langchain/chat_models/__init__.py b/libs/langchain/langchain/chat_models/__init__.py index c4a8185617..4812e694b8 100644 --- a/libs/langchain/langchain/chat_models/__init__.py +++ b/libs/langchain/langchain/chat_models/__init__.py @@ -26,6 +26,7 @@ from langchain.chat_models.ernie import ErnieBotChat from langchain.chat_models.fake import FakeListChatModel from langchain.chat_models.google_palm import ChatGooglePalm from langchain.chat_models.human import HumanInputChatModel +from langchain.chat_models.javelin_ai_gateway import ChatJavelinAIGateway from langchain.chat_models.jinachat import JinaChat from langchain.chat_models.konko import ChatKonko from langchain.chat_models.litellm import ChatLiteLLM @@ -53,6 +54,7 @@ __all__ = [ "ChatAnyscale", "ChatLiteLLM", "ErnieBotChat", + "ChatJavelinAIGateway", "ChatKonko", "QianfanChatEndpoint", ] diff --git a/libs/langchain/langchain/chat_models/javelin_ai_gateway.py b/libs/langchain/langchain/chat_models/javelin_ai_gateway.py new file mode 100644 index 0000000000..5744f68116 --- /dev/null +++ b/libs/langchain/langchain/chat_models/javelin_ai_gateway.py @@ -0,0 +1,223 @@ +import logging +from typing import Any, Dict, List, Mapping, Optional + +from langchain.callbacks.manager import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) +from langchain.chat_models.base import BaseChatModel +from langchain.pydantic_v1 import BaseModel, Extra +from langchain.schema import ( + ChatGeneration, + ChatResult, +) +from langchain.schema.messages import ( + AIMessage, + BaseMessage, + ChatMessage, + FunctionMessage, + HumanMessage, + SystemMessage, +) + +logger = logging.getLogger(__name__) + + +# Ignoring type because below is valid pydantic code +# Unexpected keyword argument "extra" for "__init_subclass__" of "object" [call-arg] +class ChatParams(BaseModel, extra=Extra.allow): # type: ignore[call-arg] + """Parameters for the `Javelin AI Gateway` LLM.""" + + temperature: float = 0.0 + stop: Optional[List[str]] = None + max_tokens: Optional[int] = None + + +class ChatJavelinAIGateway(BaseChatModel): + """`Javelin AI Gateway` chat models API. + + To use, you should have the ``javelin_sdk`` python package installed. + For more information, see https://docs.getjavelin.io + + Example: + .. code-block:: python + + from langchain.chat_models import ChatJavelinAIGateway + + chat = ChatJavelinAIGateway( + gateway_uri="", + route="", + params={ + "temperature": 0.1 + } + ) + """ + + route: str + """The route to use for the Javelin AI Gateway API.""" + + gateway_uri: Optional[str] = None + """The URI for the Javelin AI Gateway API.""" + + params: Optional[ChatParams] = None + """Parameters for the Javelin AI Gateway LLM.""" + + client: Any + """javelin client.""" + + javelin_api_key: Optional[str] = None + """The API key for the Javelin AI Gateway.""" + + def __init__(self, **kwargs: Any): + try: + from javelin_sdk import ( + JavelinClient, + UnauthorizedError, + ) + except ImportError: + raise ImportError( + "Could not import javelin_sdk python package. " + "Please install it with `pip install javelin_sdk`." + ) + + super().__init__(**kwargs) + if self.gateway_uri: + try: + self.client = JavelinClient( + base_url=self.gateway_uri, api_key=self.javelin_api_key + ) + except UnauthorizedError as e: + raise ValueError("Javelin: Incorrect API Key.") from e + + @property + def _default_params(self) -> Dict[str, Any]: + params: Dict[str, Any] = { + "gateway_uri": self.gateway_uri, + "javelin_api_key": self.javelin_api_key, + "route": self.route, + **(self.params.dict() if self.params else {}), + } + return params + + def _generate( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + message_dicts = [ + ChatJavelinAIGateway._convert_message_to_dict(message) + for message in messages + ] + data: Dict[str, Any] = { + "messages": message_dicts, + **(self.params.dict() if self.params else {}), + } + + resp = self.client.query_route(self.route, query_body=data) + + return ChatJavelinAIGateway._create_chat_result(resp.dict()) + + async def _agenerate( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + message_dicts = [ + ChatJavelinAIGateway._convert_message_to_dict(message) + for message in messages + ] + data: Dict[str, Any] = { + "messages": message_dicts, + **(self.params.dict() if self.params else {}), + } + + resp = await self.client.aquery_route(self.route, query_body=data) + + return ChatJavelinAIGateway._create_chat_result(resp.dict()) + + @property + def _identifying_params(self) -> Dict[str, Any]: + return self._default_params + + def _get_invocation_params( + self, stop: Optional[List[str]] = None, **kwargs: Any + ) -> Dict[str, Any]: + """Get the parameters used to invoke the model FOR THE CALLBACKS.""" + return { + **self._default_params, + **super()._get_invocation_params(stop=stop, **kwargs), + } + + @property + def _llm_type(self) -> str: + """Return type of chat model.""" + return "javelin-ai-gateway-chat" + + @staticmethod + def _convert_dict_to_message(_dict: Mapping[str, Any]) -> BaseMessage: + role = _dict["role"] + content = _dict["content"] + if role == "user": + return HumanMessage(content=content) + elif role == "assistant": + return AIMessage(content=content) + elif role == "system": + return SystemMessage(content=content) + else: + return ChatMessage(content=content, role=role) + + @staticmethod + def _raise_functions_not_supported() -> None: + raise ValueError( + "Function messages are not supported by the Javelin AI Gateway. Please" + " create a feature request at https://docs.getjavelin.io" + ) + + @staticmethod + def _convert_message_to_dict(message: BaseMessage) -> dict: + if isinstance(message, ChatMessage): + message_dict = {"role": message.role, "content": message.content} + elif isinstance(message, HumanMessage): + message_dict = {"role": "user", "content": message.content} + elif isinstance(message, AIMessage): + message_dict = {"role": "assistant", "content": message.content} + elif isinstance(message, SystemMessage): + message_dict = {"role": "system", "content": message.content} + elif isinstance(message, FunctionMessage): + raise ValueError( + "Function messages are not supported by the Javelin AI Gateway. Please" + " create a feature request at https://docs.getjavelin.io" + ) + else: + raise ValueError(f"Got unknown message type: {message}") + + if "function_call" in message.additional_kwargs: + ChatJavelinAIGateway._raise_functions_not_supported() + if message.additional_kwargs: + logger.warning( + "Additional message arguments are unsupported by Javelin AI Gateway " + " and will be ignored: %s", + message.additional_kwargs, + ) + return message_dict + + @staticmethod + def _create_chat_result(response: Mapping[str, Any]) -> ChatResult: + generations = [] + for candidate in response["llm_response"]["choices"]: + message = ChatJavelinAIGateway._convert_dict_to_message( + candidate["message"] + ) + message_metadata = candidate.get("metadata", {}) + gen = ChatGeneration( + message=message, + generation_info=dict(message_metadata), + ) + generations.append(gen) + + response_metadata = response.get("metadata", {}) + return ChatResult(generations=generations, llm_output=response_metadata) diff --git a/libs/langchain/langchain/embeddings/__init__.py b/libs/langchain/langchain/embeddings/__init__.py index 621b88a628..e5cdc38f13 100644 --- a/libs/langchain/langchain/embeddings/__init__.py +++ b/libs/langchain/langchain/embeddings/__init__.py @@ -40,6 +40,7 @@ from langchain.embeddings.huggingface import ( HuggingFaceInstructEmbeddings, ) from langchain.embeddings.huggingface_hub import HuggingFaceHubEmbeddings +from langchain.embeddings.javelin_ai_gateway import JavelinAIGatewayEmbeddings from langchain.embeddings.jina import JinaEmbeddings from langchain.embeddings.llamacpp import LlamaCppEmbeddings from langchain.embeddings.localai import LocalAIEmbeddings @@ -107,6 +108,7 @@ __all__ = [ "AwaEmbeddings", "HuggingFaceBgeEmbeddings", "ErnieEmbeddings", + "JavelinAIGatewayEmbeddings", "OllamaEmbeddings", "QianfanEmbeddingsEndpoint", ] diff --git a/libs/langchain/langchain/embeddings/javelin_ai_gateway.py b/libs/langchain/langchain/embeddings/javelin_ai_gateway.py new file mode 100644 index 0000000000..db97b183c7 --- /dev/null +++ b/libs/langchain/langchain/embeddings/javelin_ai_gateway.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from typing import Any, Iterator, List, Optional + +from langchain.pydantic_v1 import BaseModel +from langchain.schema.embeddings import Embeddings + + +def _chunk(texts: List[str], size: int) -> Iterator[List[str]]: + for i in range(0, len(texts), size): + yield texts[i : i + size] + + +class JavelinAIGatewayEmbeddings(Embeddings, BaseModel): + """ + Wrapper around embeddings LLMs in the Javelin AI Gateway. + + To use, you should have the ``javelin_sdk`` python package installed. + For more information, see https://docs.getjavelin.io + + Example: + .. code-block:: python + + from langchain.embeddings import JavelinAIGatewayEmbeddings + + embeddings = JavelinAIGatewayEmbeddings( + gateway_uri="", + route="" + ) + """ + + client: Any + """javelin client.""" + + route: str + """The route to use for the Javelin AI Gateway API.""" + + gateway_uri: Optional[str] = None + """The URI for the Javelin AI Gateway API.""" + + javelin_api_key: Optional[str] = None + """The API key for the Javelin AI Gateway API.""" + + def __init__(self, **kwargs: Any): + try: + from javelin_sdk import ( + JavelinClient, + UnauthorizedError, + ) + except ImportError: + raise ImportError( + "Could not import javelin_sdk python package. " + "Please install it with `pip install javelin_sdk`." + ) + + super().__init__(**kwargs) + if self.gateway_uri: + try: + self.client = JavelinClient( + base_url=self.gateway_uri, api_key=self.javelin_api_key + ) + except UnauthorizedError as e: + raise ValueError("Javelin: Incorrect API Key.") from e + + def _query(self, texts: List[str]) -> List[List[float]]: + embeddings = [] + for txt in _chunk(texts, 20): + try: + resp = self.client.query_route(self.route, query_body={"input": txt}) + resp_dict = resp.dict() + + embeddings_chunk = resp_dict.get("llm_response", {}).get("data", []) + for item in embeddings_chunk: + if "embedding" in item: + embeddings.append(item["embedding"]) + except ValueError as e: + print("Failed to query route: " + str(e)) + + return embeddings + + async def _aquery(self, texts: List[str]) -> List[List[float]]: + embeddings = [] + for txt in _chunk(texts, 20): + try: + resp = await self.client.aquery_route( + self.route, query_body={"input": txt} + ) + resp_dict = resp.dict() + + embeddings_chunk = resp_dict.get("llm_response", {}).get("data", []) + for item in embeddings_chunk: + if "embedding" in item: + embeddings.append(item["embedding"]) + except ValueError as e: + print("Failed to query route: " + str(e)) + + return embeddings + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + return self._query(texts) + + def embed_query(self, text: str) -> List[float]: + return self._query([text])[0] + + async def aembed_documents(self, texts: List[str]) -> List[List[float]]: + return await self._aquery(texts) + + async def aembed_query(self, text: str) -> List[float]: + result = await self._aquery([text]) + return result[0] diff --git a/libs/langchain/langchain/llms/__init__.py b/libs/langchain/langchain/llms/__init__.py index 8e835ea0a9..f8b2b8a378 100644 --- a/libs/langchain/langchain/llms/__init__.py +++ b/libs/langchain/langchain/llms/__init__.py @@ -54,6 +54,7 @@ from langchain.llms.huggingface_hub import HuggingFaceHub from langchain.llms.huggingface_pipeline import HuggingFacePipeline from langchain.llms.huggingface_text_gen_inference import HuggingFaceTextGenInference from langchain.llms.human import HumanInputLLM +from langchain.llms.javelin_ai_gateway import JavelinAIGateway from langchain.llms.koboldai import KoboldApiLLM from langchain.llms.llamacpp import LlamaCpp from langchain.llms.manifest import ManifestWrapper @@ -161,6 +162,7 @@ __all__ = [ "Writer", "OctoAIEndpoint", "Xinference", + "JavelinAIGateway", "QianfanLLMEndpoint", ] @@ -230,5 +232,6 @@ type_to_cls_dict: Dict[str, Type[BaseLLM]] = { "vllm_openai": VLLMOpenAI, "writer": Writer, "xinference": Xinference, + "javelin-ai-gateway": JavelinAIGateway, "qianfan_endpoint": QianfanLLMEndpoint, } diff --git a/libs/langchain/langchain/llms/javelin_ai_gateway.py b/libs/langchain/langchain/llms/javelin_ai_gateway.py new file mode 100644 index 0000000000..7effebb6b0 --- /dev/null +++ b/libs/langchain/langchain/llms/javelin_ai_gateway.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Mapping, Optional + +from langchain.callbacks.manager import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) +from langchain.llms.base import LLM +from langchain.pydantic_v1 import BaseModel, Extra + + +# Ignoring type because below is valid pydantic code +# Unexpected keyword argument "extra" for "__init_subclass__" of "object" +class Params(BaseModel, extra=Extra.allow): # type: ignore[call-arg] + """Parameters for the Javelin AI Gateway LLM.""" + + temperature: float = 0.0 + stop: Optional[List[str]] = None + max_tokens: Optional[int] = None + + +class JavelinAIGateway(LLM): + """ + Wrapper around completions LLMs in the Javelin AI Gateway. + + To use, you should have the ``javelin_sdk`` python package installed. + For more information, see https://docs.getjavelin.io + + Example: + .. code-block:: python + + from langchain.llms import JavelinAIGateway + + completions = JavelinAIGateway( + gateway_uri="", + route="", + params={ + "temperature": 0.1 + } + ) + """ + + route: str + """The route to use for the Javelin AI Gateway API.""" + + client: Optional[Any] = None + """The Javelin AI Gateway client.""" + + gateway_uri: Optional[str] = None + """The URI of the Javelin AI Gateway API.""" + + params: Optional[Params] = None + """Parameters for the Javelin AI Gateway API.""" + + javelin_api_key: Optional[str] = None + """The API key for the Javelin AI Gateway API.""" + + def __init__(self, **kwargs: Any): + try: + from javelin_sdk import ( + JavelinClient, + UnauthorizedError, + ) + except ImportError: + raise ImportError( + "Could not import javelin_sdk python package. " + "Please install it with `pip install javelin_sdk`." + ) + super().__init__(**kwargs) + if self.gateway_uri: + try: + self.client = JavelinClient( + base_url=self.gateway_uri, api_key=self.javelin_api_key + ) + except UnauthorizedError as e: + raise ValueError("Javelin: Incorrect API Key.") from e + + @property + def _default_params(self) -> Dict[str, Any]: + """Get the default parameters for calling Javelin AI Gateway API.""" + params: Dict[str, Any] = { + "gateway_uri": self.gateway_uri, + "route": self.route, + "javelin_api_key": self.javelin_api_key, + **(self.params.dict() if self.params else {}), + } + return params + + @property + def _identifying_params(self) -> Mapping[str, Any]: + """Get the identifying parameters.""" + return self._default_params + + def _call( + self, + prompt: str, + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> str: + """Call the Javelin AI Gateway API.""" + data: Dict[str, Any] = { + "prompt": prompt, + **(self.params.dict() if self.params else {}), + } + if s := (stop or (self.params.stop if self.params else None)): + data["stop"] = s + + if self.client is not None: + resp = self.client.query_route(self.route, query_body=data) + else: + raise ValueError("Javelin client is not initialized.") + + resp_dict = resp.dict() + + try: + return resp_dict["llm_response"]["choices"][0]["text"] + except KeyError: + return "" + + async def _acall( + self, + prompt: str, + stop: Optional[List[str]] = None, + run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> str: + """Call async the Javelin AI Gateway API.""" + data: Dict[str, Any] = { + "prompt": prompt, + **(self.params.dict() if self.params else {}), + } + if s := (stop or (self.params.stop if self.params else None)): + data["stop"] = s + + if self.client is not None: + resp = await self.client.aquery_route(self.route, query_body=data) + else: + raise ValueError("Javelin client is not initialized.") + + resp_dict = resp.dict() + + try: + return resp_dict["llm_response"]["choices"][0]["text"] + except KeyError: + return "" + + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "javelin-ai-gateway"