Skip to content

Commit

Permalink
Code cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevinmessiaen
kevinmessiaen committed Nov 7, 2024
1 parent f897262 commit 2657b19
Show file tree
Hide file tree
Showing 11 changed files with 157 additions and 783 deletions.
199 changes: 70 additions & 129 deletions docs/open_source/scan/scan_llm/index.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,194 +57,135 @@ like this:
```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"] = "your-api-key"

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

# 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"] = "your-org-id"
os.environ["OPENAI_API_BASE"] = "openaiai-api-base"
```

More information on [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"

# 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.set_llm_model("azure/<your_deployment_name>")
giskard.llm.set_embedding_model("azure/<your_deployment_name>")

# optional
os.environ["AZURE_AD_TOKEN"] = ""
os.environ["AZURE_API_TYPE"] = ""
giskard.llm.set_embedding_model('my-embedding-model')
```

More information on [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"] = "…"
os.environ['MISTRAL_API_KEY'] = ""

# Create a giskard Mistral client
mistral_client = MistralClient()

# Set the default client
giskard.llm.set_default_client(mistral_client)

# 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-tiny")
giskard.llm.set_embedding_model("mistral/mistral-embed")
```

More information on [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")
giskard.llm.set_llm_model("ollama/llama2", api_base="http://localhost:11434") # See supported models here: https://docs.litellm.ai/docs/providers/ollama#ollama-models
```

# 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)
More information on [LiteLLM documentation](https://docs.litellm.ai/docs/providers/ollama)

# Set the default client and embedding
giskard.llm.set_default_client(llm_client)
giskard.llm.embeddings.set_default_embedding(embed_client)
```

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

More information on [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"])
os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

# 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")

# 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")
giskard.llm.set_embedding_model("bedrock/amazon.titan-embed-text-v1")
```

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

More information on [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"])
os.environ["GEMINI_API_KEY"] = "your-api-key"

# Create a giskard Gemini client
gemini_client = GeminiClient()

# Set the default client
giskard.llm.set_default_client(gemini_client)

# 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-pro")
```

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

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

```python
import requests
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",
)
data = json.loads(response.get("body").read())

return ChatMessage(role="assistant", message=data["completion"])

# Create an instance of the custom client
llm_client = MyLLMClient()

# Set the default client
set_default_client(llm_client)

# 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)
import litellm
import os
from typing import Optional


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 is not provided")

response = requests.post('https://www.my-fake-llm.ai/chat/completion', json={
'messages': messages
}, headers={'Authorization': api_key})

return litellm.ModelResponse(**response.json())


my_custom_llm = MyCustomLLM()

litellm.custom_provider_map = [ # 👈 KEY STEP - REGISTER HANDLER
{"provider": "my-custom-llm", "custom_handler": my_custom_llm}
]

api_key = os.environ['MY_SECRET_KEY']

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

::::::
Expand Down
Loading

0 comments on commit 2657b19

Please sign in to comment.