mirror of
https://github.com/hwchase17/langchain
synced 2024-10-31 15:20:26 +00:00
117 lines
2.7 KiB
Plaintext
117 lines
2.7 KiB
Plaintext
# MLflow AI Gateway
|
|
|
|
The MLflow AI Gateway service is a powerful tool designed to streamline the usage and management of various large language model (LLM) providers, such as OpenAI and Anthropic, within an organization. It offers a high-level interface that simplifies the interaction with these services by providing a unified endpoint to handle specific LLM related requests. See [the MLflow AI Gateway documentation](https://mlflow.org/docs/latest/gateway/index.html) for more details.
|
|
|
|
## Installation and Setup
|
|
|
|
Install `mlflow` with MLflow AI Gateway dependencies:
|
|
|
|
```sh
|
|
pip install 'mlflow[gateway]'
|
|
```
|
|
|
|
Set the OpenAI API key as an environment variable:
|
|
|
|
```sh
|
|
export OPENAI_API_KEY=...
|
|
```
|
|
|
|
Create a configuration file:
|
|
|
|
```yaml
|
|
routes:
|
|
- name: completions
|
|
route_type: llm/v1/completions
|
|
model:
|
|
provider: openai
|
|
name: text-davinci-003
|
|
config:
|
|
openai_api_key: $OPENAI_API_KEY
|
|
|
|
- name: embeddings
|
|
route_type: llm/v1/embeddings
|
|
model:
|
|
provider: openai
|
|
name: text-embedding-ada-002
|
|
config:
|
|
openai_api_key: $OPENAI_API_KEY
|
|
```
|
|
|
|
Start the Gateway server:
|
|
|
|
```sh
|
|
mlflow gateway start --config-path /path/to/config.yaml
|
|
```
|
|
|
|
## Completions Example
|
|
|
|
```python
|
|
import mlflow
|
|
from langchain import LLMChain, PromptTemplate
|
|
from langchain.llms import MlflowAIGateway
|
|
|
|
gateway = MlflowAIGateway(
|
|
gateway_uri="http://127.0.0.1:5000",
|
|
route="completions",
|
|
params={
|
|
"temperature": 0.0,
|
|
"top_p": 0.1,
|
|
},
|
|
)
|
|
|
|
llm_chain = LLMChain(
|
|
llm=gateway,
|
|
prompt=PromptTemplate(
|
|
input_variables=["adjective"],
|
|
template="Tell me a {adjective} joke",
|
|
),
|
|
)
|
|
result = llm_chain.run(adjective="funny")
|
|
print(result)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.langchain.log_model(chain, "model")
|
|
|
|
model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
print(model.predict([{"adjective": "funny"}]))
|
|
```
|
|
|
|
## Embeddings Example
|
|
|
|
```python
|
|
from langchain.embeddings import MlflowAIGatewayEmbeddings
|
|
|
|
embeddings = MlflowAIGatewayEmbeddings(
|
|
gateway_uri="http://127.0.0.1:5000",
|
|
route="embeddings",
|
|
)
|
|
|
|
print(embeddings.embed_query("hello"))
|
|
print(embeddings.embed_documents(["hello"]))
|
|
```
|
|
|
|
## Databricks MLflow AI Gateway
|
|
|
|
Databricks MLflow AI Gateway is in private preview.
|
|
Please contact a Databricks representative to enroll in the preview.
|
|
|
|
```python
|
|
from langchain import LLMChain, PromptTemplate
|
|
from langchain.llms import MlflowAIGateway
|
|
|
|
gateway = MlflowAIGateway(
|
|
gateway_uri="databricks",
|
|
route="completions",
|
|
)
|
|
|
|
llm_chain = LLMChain(
|
|
llm=gateway,
|
|
prompt=PromptTemplate(
|
|
input_variables=["adjective"],
|
|
template="Tell me a {adjective} joke",
|
|
),
|
|
)
|
|
result = llm_chain.run(adjective="funny")
|
|
print(result)
|
|
```
|