mirror of https://github.com/hwchase17/langchain
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!!pull/10865/head
Sharath Rajasekar
10 months ago
committed by
GitHub
parent
957956ba6d
commit
96023f94d9
@ -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
|
||||
}
|
||||
@ -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))
|
||||
|
||||
```
|
||||
|
||||
@ -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="<javelin-ai-gateway-uri>",
|
||||
route="<javelin-ai-gateway-chat-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)
|
||||
@ -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="<javelin-ai-gateway-uri>",
|
||||
route="<your-javelin-gateway-embeddings-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]
|
||||
@ -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="<your-javelin-ai-gateway-uri>",
|
||||
route="<your-javelin-ai-gateway-completions-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"
|
||||
Loading…
Reference in New Issue