Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[GSK-3846] Add support to LiteLLM #2069

Merged
merged 42 commits into from
Nov 21, 2024
Merged
Show file tree
Hide file tree
Changes from 41 commits
Commits
Show all changes
42 commits
Select commit Hold shift + click to select a range
f61cc21
Created litellm client
kevinmessiaen Nov 5, 2024
e3844ff
Updated documentation
kevinmessiaen Nov 5, 2024
c524217
Added litellm embedding
kevinmessiaen Nov 5, 2024
d6b032f
Code improvement
kevinmessiaen Nov 5, 2024
e31cdfa
Added deprecated warnings
kevinmessiaen Nov 5, 2024
0f5ade7
Fixed typo
kevinmessiaen Nov 5, 2024
3045060
Improved documentation and llm setup
kevinmessiaen Nov 7, 2024
f49a2bc
Added back fastembed as default
kevinmessiaen Nov 7, 2024
1157bda
Removed todo: LiteLLM does not support embeddings for Gemini and Ollama
kevinmessiaen Nov 7, 2024
10e4113
Typo
kevinmessiaen Nov 7, 2024
5fb6a78
Fixed embeddings
kevinmessiaen Nov 7, 2024
f897262
Default model to gpt-4o
kevinmessiaen Nov 7, 2024
2657b19
Code cleanup
kevinmessiaen Nov 7, 2024
b04f126
Code cleanup
kevinmessiaen Nov 7, 2024
4633aa4
Skip LiteLLM tests with pydantic < 2
kevinmessiaen Nov 8, 2024
63ace19
Added test for custom client
kevinmessiaen Nov 8, 2024
1b382ee
Added test for embedding
kevinmessiaen Nov 8, 2024
713f0b0
Fixed tests
kevinmessiaen Nov 8, 2024
deca09a
Merge branch 'main' into feature/litellm
henchaves Nov 14, 2024
e54c414
Merge branch 'main' into feature/litellm
henchaves Nov 14, 2024
dee0e83
Reintroduced old way to set LLM models
kevinmessiaen Nov 15, 2024
7703d51
Reintroduced old way to set LLM models
kevinmessiaen Nov 15, 2024
5349fc2
Reintroduced old clients
kevinmessiaen Nov 15, 2024
6458f97
Merge branch 'main' into feature/litellm
kevinmessiaen Nov 15, 2024
2756e27
Fixed OpenAI embeddings
kevinmessiaen Nov 15, 2024
2b88ed3
Update Setting up the LLM client docs
henchaves Nov 15, 2024
1dc73d9
Update Setting up the LLM client docs pt 2
henchaves Nov 15, 2024
b39731e
Update testset generation docs
henchaves Nov 15, 2024
7eaf007
Update scan llm docs
henchaves Nov 15, 2024
5f51327
Merge branch 'main' into feature/litellm
henchaves Nov 18, 2024
39c4fa9
Removed response_format with ollama models due to issue in litellm
kevinmessiaen Nov 19, 2024
b09d266
Added dumb trim
kevinmessiaen Nov 19, 2024
911d6e5
Fixed output
kevinmessiaen Nov 19, 2024
40bede9
Add _parse_json_output to LiteLLM
henchaves Nov 19, 2024
77e6a4f
Added way to disable structured output
kevinmessiaen Nov 20, 2024
cab45a1
Fix test_litellm_client
henchaves Nov 21, 2024
5f39da1
Merge branch 'main' into feature/litellm
henchaves Nov 21, 2024
78dd03e
Check if format is json before calling _parse_json_output
henchaves Nov 21, 2024
82712c7
Set LITELLM_LOG as error level
henchaves Nov 21, 2024
a61e4b2
Add `disable_structured_output` to bedrock examples
henchaves Nov 21, 2024
3d33028
Format files
henchaves Nov 21, 2024
a571312
Fix sonar issues
henchaves Nov 21, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
209 changes: 75 additions & 134 deletions docs/open_source/scan/scan_llm/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -38,220 +38,161 @@ processed.

## Before starting

In the following example, we illustrate the procedure using **OpenAI** and **Azure OpenAI**; however, please note that
our platform supports a variety of language models. For details on configuring different models, visit
our [🤖 Setting up the LLM Client page](../../setting_up/index.md)

Before starting, make sure you have installed the LLM flavor of Giskard:
First of all, make sure you have installed the LLM flavor of Giskard:

```bash
pip install "giskard[llm]"
```

For the LLM-assisted detectors to work, you need to have an OpenAI API key. You can set it in your notebook
like this:
For the LLM-assisted detectors to work, you need to set up a LLM client. Our platform supports a variety of language models, and you can find the details on configuring different models in our [🤖 Setting up the LLM Client page](../../setting_up/index.md) or follow the instructions below for each provider:

:::::::{tab-set}
::::::{tab-item} OpenAI

```python
import os
import giskard
from giskard.llm.client.openai import OpenAIClient

# Set the OpenAI API key
os.environ["OPENAI_API_KEY"] = "sk-…"
os.environ["OPENAI_API_KEY"] = "" # "my-openai-api-key"

# Create a giskard OpenAI client
openai_client = OpenAIClient(model="gpt-4o")
# Optional, setup a model (default LLM is gpt-4o, default embedding model is text-embedding-3-small)
giskard.llm.set_llm_model("gpt-4o")
giskard.llm.set_embedding_model("text-embedding-3-small")

# Set the default client
giskard.llm.set_llm_api("openai")
giskard.llm.set_default_client(openai_client)
# Optional Keys - OpenAI Organization, OpenAI API Base
os.environ["OPENAI_ORGANIZATION"] = "" # "my-openai-organization"
os.environ["OPENAI_API_BASE"] = "" # "https://api.openai.com"
```

More information on [OpenAI LiteLLM documentation](https://docs.litellm.ai/docs/providers/openai)

::::::
::::::{tab-item} Azure OpenAI

Require `openai>=1.0.0`

```python
import os
import giskard

# Set the Azure OpenAI API key and endpoint
os.environ['AZURE_OPENAI_API_KEY'] = '...'
os.environ['AZURE_OPENAI_ENDPOINT'] = 'https://xxx.openai.azure.com'
os.environ['OPENAI_API_VERSION'] = '2023-07-01-preview'
os.environ["AZURE_API_KEY"] = "" # "my-azure-api-key"
os.environ["AZURE_API_BASE"] = "" # "https://example-endpoint.openai.azure.com"
os.environ["AZURE_API_VERSION"] = "" # "2023-05-15"

giskard.llm.set_llm_model("azure/<your_llm_name>")
giskard.llm.set_embedding_model("azure/<your_embed_model_name>")

# You'll need to provide the name of the model that you've deployed
# Beware, the model provided must be capable of using function calls
giskard.llm.set_llm_model('my-gpt-4-model')
giskard.llm.embeddings.openai.set_embedding_model('my-embedding-model')
# Optional Keys - Azure AD Token, Azure API Type
os.environ["AZURE_AD_TOKEN"] = ""
os.environ["AZURE_API_TYPE"] = ""
```

More information on [Azure LiteLLM documentation](https://docs.litellm.ai/docs/providers/azure)

::::::
::::::{tab-item} Mistral

```python
import os
import giskard
from giskard.llm.client.mistral import MistralClient

# Set the Mistral API key
os.environ["MISTRAL_API_KEY"] = "…"

# Create a giskard Mistral client
mistral_client = MistralClient()

# Set the default client
giskard.llm.set_default_client(mistral_client)
os.environ["MISTRAL_API_KEY"] = "" # "my-mistral-api-key"

# You may also want to set the default embedding model
# Check the Custom Client code snippet for more details
giskard.llm.set_llm_model("mistral/mistral-large-latest")
giskard.llm.set_embedding_model("mistral/mistral-embed")
```

More information on [Mistral LiteLLM documentation](https://docs.litellm.ai/docs/providers/mistral)

::::::
::::::{tab-item} Ollama

```python
import giskard
from openai import OpenAI
from giskard.llm.client.openai import OpenAIClient
from giskard.llm.embeddings.openai import OpenAIEmbedding

# Setup the OpenAI client with API key and base URL for Ollama
_client = OpenAI(base_url="http://localhost:11434/v1/", api_key="ollama")

# Wrap the original OpenAI client with giskard OpenAI client and embedding
llm_client = OpenAIClient(model="llama3.2", client=_client)
embed_client = OpenAIEmbedding(model="nomic-embed-text", client=_client)
api_base = "http://localhost:11434" # default api_base for local Ollama

# Set the default client and embedding
giskard.llm.set_default_client(llm_client)
giskard.llm.embeddings.set_default_embedding(embed_client)
# See supported models here: https://docs.litellm.ai/docs/providers/ollama#ollama-models
giskard.llm.set_llm_model("ollama/llama3.1", disable_structured_output=True, api_base=api_base)
giskard.llm.set_embedding_model("ollama/nomic-embed-text", api_base=api_base)
```

More information on [Ollama LiteLLM documentation](https://docs.litellm.ai/docs/providers/ollama)

::::::
::::::{tab-item} Claude 3
::::::{tab-item} AWS Bedrock

More information on [Bedrock LiteLLM documentation](https://docs.litellm.ai/docs/providers/bedrock)

```python
import os
import boto3
import giskard

from giskard.llm.client.bedrock import ClaudeBedrockClient
from giskard.llm.embeddings.bedrock import BedrockEmbedding

# Create a Bedrock client
bedrock_runtime = boto3.client("bedrock-runtime", region_name=os.environ["AWS_DEFAULT_REGION"])

# Wrap the Beddock client with giskard Bedrock client and embedding
claude_client = ClaudeBedrockClient(bedrock_runtime, model="anthropic.claude-3-haiku-20240307-v1:0")
embed_client = BedrockEmbedding(bedrock_runtime, model="amazon.titan-embed-text-v1")
os.environ["AWS_ACCESS_KEY_ID"] = "" # "my-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # "my-aws-secret-access-key"
os.environ["AWS_REGION_NAME"] = "" # "us-west-2"

# Set the default client and embedding
giskard.llm.set_default_client(claude_client)
giskard.llm.embeddings.set_default_embedding(embed_client)
giskard.llm.set_llm_model("bedrock/anthropic.claude-3-sonnet-20240229-v1:0", disable_structured_output=True)
giskard.llm.set_embedding_model("bedrock/amazon.titan-embed-image-v1")
```

::::::
::::::{tab-item} Gemini

More information on [Gemini LiteLLM documentation](https://docs.litellm.ai/docs/providers/gemini)

```python
import os
import giskard
import google.generativeai as genai
from giskard.llm.client.gemini import GeminiClient

# Set the Gemini API key
os.environ["GEMINI_API_KEY"] = "…"

# Configure the Gemini API
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# Create a giskard Gemini client
gemini_client = GeminiClient()

# Set the default client
giskard.llm.set_default_client(gemini_client)
os.environ["GEMINI_API_KEY"] = "" # "my-gemini-api-key"

# You may also want to set the default embedding model
# Check the Custom Client code snippet for more details
giskard.llm.set_llm_model("gemini/gemini-1.5-pro")
giskard.llm.set_embedding_model("gemini/text-embedding-004")
```

::::::
::::::{tab-item} Custom Client

More information on [Custom Format LiteLLM documentation](https://docs.litellm.ai/docs/providers/custom_llm_server)

```python
import os
import requests
from typing import Optional

import litellm
import giskard
from typing import Sequence, Optional
from giskard.llm.client import set_default_client
from giskard.llm.client.base import LLMClient, ChatMessage

# Create a custom client by extending the LLMClient class
class MyLLMClient(LLMClient):
def __init__(self, my_client):
self._client = my_client

def complete(
self,
messages: Sequence[ChatMessage],
temperature: float = 1,
max_tokens: Optional[int] = None,
caller_id: Optional[str] = None,
seed: Optional[int] = None,
format=None,
) -> ChatMessage:
# Create the prompt
prompt = ""
for msg in messages:
if msg.role.lower() == "assistant":
prefix = "\n\nAssistant: "
else:
prefix = "\n\nHuman: "

prompt += prefix + msg.content

prompt += "\n\nAssistant: "

# Create the body
params = {
"prompt": prompt,
"max_tokens_to_sample": max_tokens or 1000,
"temperature": temperature,
"top_p": 0.9,
}
body = json.dumps(params)

response = self._client.invoke_model(
body=body,
modelId=self._model_id,
accept="application/json",
contentType="application/json",


class MyCustomLLM(litellm.CustomLLM):
def completion(self, messages: str, api_key: Optional[str] = None, **kwargs) -> litellm.ModelResponse:
api_key = api_key or os.environ.get("MY_SECRET_KEY")
if api_key is None:
raise litellm.AuthenticationError("`api_key` was not provided")

response = requests.post(
"https://www.my-custom-llm.ai/chat/completion",
json={"messages": messages},
headers={"Authorization": api_key},
)
data = json.loads(response.get("body").read())

return ChatMessage(role="assistant", message=data["completion"])
return litellm.ModelResponse(**response.json())

# Create an instance of the custom client
llm_client = MyLLMClient()
os.eviron["MY_SECRET_KEY"] = "" # "my-secret-key"

# Set the default client
set_default_client(llm_client)
my_custom_llm = MyCustomLLM()

# It's also possible to create a custom embedding class extending BaseEmbedding
# Or you can use FastEmbed for a pre-built embedding model:
from giskard.llm.embeddings.fastembed import try_get_fastembed_embeddings
embed_client = try_get_fastembed_embeddings()
giskard.llm.embeddings.set_default_embedding(embed_client)
litellm.custom_provider_map = [ # 👈 KEY STEP - REGISTER HANDLER
{"provider": "my-custom-llm-endpoint", "custom_handler": my_custom_llm}
]

api_key = os.environ["MY_SECRET_KEY"]

giskard.llm.set_llm_model("my-custom-llm-endpoint/my-custom-model", api_key=api_key)
```

::::::
:::::::

We are now ready to start.

(model-wrapping)=

## Step 1: Wrap your model
Expand Down
Loading
Loading