Improved documentation and llm setup

docs/open_source/setting_up/index.md

@@ -1,16 +1,12 @@
 # 🤖 Setting up the LLM Client
 
-This guide focuses primarily on configuring and using various LLM clients supported to run Giskard's LLM-assisted functionalities. These clients include:
+This guide focuses primarily on configuring and using various LLM clients supported to run Giskard's LLM-assisted functionalities. We are using [LiteLLM](https://github.com/BerriAI/litellm) to handle the model calls, you can see the list of supported models in the [LiteLLM documentation](https://docs.litellm.ai/docs/providers).
 
-- OpenAI GPT models (such as GPT-3.5 and GPT-4)
-- Azure OpenAI
-- Mistral
-- Ollama
-- Any Custom Model
+## OpenAI Client Setup
 
-## OpenAI GPT-4 Client Setup
+More information on [LiteLLM documentation](https://docs.litellm.ai/docs/providers/openai)
 
-More information on [litellm documentation](https://docs.litellm.ai/docs/providers/openai)
+### Setup using .env variables
 
 ```python
 import os
@@ -21,12 +17,26 @@ os.environ["OPENAI_API_KEY"] = "your-api-key"
 # 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")
+
+# Optional Keys - OpenAI Organization, OpenAI API Base
+os.environ["OPENAI_ORGANIZATION"] = "your-org-id"
+os.environ["OPENAI_API_BASE"] = "openaiai-api-base"
+```
+
+### Setup using completion params
+
+```python
+import giskard
+
+giskard.llm.set_llm_model("gpt-4", api_key="your-api-key")
+giskard.llm.set_embedding_model("text-embedding-ada-002", api_key="your-api-key")
 ```
 
 ## Azure OpenAI Client Setup
 
 More information on [litellm documentation](https://docs.litellm.ai/docs/providers/azure)
 
+### Setup using .env variables
 
 ```python
 import os
@@ -35,6 +45,7 @@ import giskard
 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_deployment_name>")
 giskard.llm.set_embedding_model("azure/<your_deployment_name>")
 
@@ -43,15 +54,33 @@ os.environ["AZURE_AD_TOKEN"] = ""
 os.environ["AZURE_API_TYPE"] = ""
 ```
 
+### Setup using completion params
+
+```python
+import giskard
+
+# Using api_key, api_base, api_version
+giskard.llm.set_llm_model("azure/<your_deployment_name>", api_base="", api_version="", api_key="")
+giskard.llm.set_embedding_model("azure/<your_deployment_name>", api_base="", api_version="", api_key="")
+
+# Using azure_ad_token, api_base, api_version
+giskard.llm.set_llm_model("azure/<your_deployment_name>", api_base="", api_version="", azure_ad_token="")
+giskard.llm.set_embedding_model("azure/<your_deployment_name>", api_base="", api_version="", azure_ad_token="")
+```
+
+
 ## Mistral Client Setup
 
 More information on [litellm documentation](https://docs.litellm.ai/docs/providers/mistral)
 
+### Setup using .env variables
+
 ```python
 import os
 import giskard
 
 os.environ['MISTRAL_API_KEY'] = ""
+
 giskard.llm.set_llm_model("mistral/mistral-tiny")
 giskard.llm.set_embedding_model("mistral/mistral-embed")
 
@@ -61,27 +90,29 @@ giskard.llm.set_embedding_model("mistral/mistral-embed")
 
 More information on [litellm documentation](https://docs.litellm.ai/docs/providers/ollama)
 
+### Setup using completion params
+
 ```python
-import litellm
 import giskard
 
-giskard.llm.set_llm_model("ollama/llama2") # See supported models here: https://docs.litellm.ai/docs/providers/ollama#ollama-models
-litellm.api_base = "http://localhost:11434"
-
+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
 # TODO: embedding
 ```
 
 ## AWS Bedrock Client Setup
 
 More information on [litellm documentation](https://docs.litellm.ai/docs/providers/bedrock)
 
+### Setup using .env variables
+
 ```python
 import os
 import giskard
 
 os.environ["AWS_ACCESS_KEY_ID"] = ""
 os.environ["AWS_SECRET_ACCESS_KEY"] = ""
 os.environ["AWS_REGION_NAME"] = ""
+
 giskard.llm.set_llm_model("bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
 giskard.llm.set_embedding_model("bedrock/amazon.titan-embed-text-v1")
 ```
@@ -90,6 +121,8 @@ giskard.llm.set_embedding_model("bedrock/amazon.titan-embed-text-v1")
 
 More information on [litellm documentation](https://docs.litellm.ai/docs/providers/gemini)
 
+### Setup using .env variables
+
 ```python
 import os
 import giskard
@@ -105,34 +138,47 @@ giskard.llm.set_llm_model("gemini/gemini-pro")
 More information on [litellm documentation](https://docs.litellm.ai/docs/providers/custom_llm_server    )
 
 ```python
+import requests
 import giskard
 import litellm
-from litellm import CustomLLM, completion, get_llm_provider
+import os
+from typing import Optional
 
 
-class MyCustomLLM(CustomLLM):
-    def completion(self, *args, **kwargs) -> litellm.ModelResponse:
-        return litellm.completion(
-            model="gpt-3.5-turbo",
-            messages=[{"role": "user", "content": "Hello world"}],
-            mock_response="Hi!",
-        )
+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())
+
+    def embedding(self, inputs, api_key: Optional[str] = None, **kwargs) -> litellm.EmbeddingResponse:
+        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/embeddings', json={
+            'inputs': inputs
+        }, headers={'Authorization': api_key})
+
+        return litellm.EmbeddingResponse(**response.json())
 
-    def embedding(self, *args, **kwargs) -> litellm.EmbeddingResponse:
-        return litellm.embedding(
-            model="openai/text-embedding-ada-002",
-            input=["Hello world"]
-
-        )
 
 my_custom_llm = MyCustomLLM()
 
-litellm.custom_provider_map = [ # 👈 KEY STEP - REGISTER HANDLER
+litellm.custom_provider_map = [  # 👈 KEY STEP - REGISTER HANDLER
     {"provider": "my-custom-llm", "custom_handler": my_custom_llm}
 ]
 
-giskard.llm.set_llm_model("my-custom-llm/my-fake-llm-model")
-giskard.llm.set_embedding_model("my-custom-llm/my-fake-embedding-model")
+api_key = os.environ['MY_SECRET_KEY']
+
+giskard.llm.set_llm_model("my-custom-llm/my-fake-llm-model", api_key=api_key)
+giskard.llm.set_embedding_model("my-custom-llm/my-fake-embedding-model", api_key=api_key)
 
 
 ```

giskard/llm/client/__init__.py

@@ -9,7 +9,10 @@
 
 _default_client = None
 _default_llm_api: Optional[str] = None
+
 _default_llm_model = os.getenv("GSK_LLM_MODEL", "gpt-4")
+_default_completion_params = dict()
+
 _default_llm_base_url = os.getenv("GSK_LLM_BASE_URL", None)
 
 
@@ -21,9 +24,6 @@ def set_default_client(client: LLMClient):
     _default_client = client
 
 
-@deprecated(
-    "_unset_default_client is deprecated, check documentation to setup llm: https://docs.giskard.ai/en/latest/open_source/setting_up/index.html"
-)
 def _unset_default_client():
     global _default_client
     _default_client = None
@@ -52,9 +52,13 @@ def set_llm_base_url(llm_base_url: Optional[str]):
     _unset_default_client()
 
 
-def set_llm_model(llm_model: str):
+def set_llm_model(llm_model: str, **kwargs):
     global _default_llm_model
+    global _default_completion_params
+
     _default_llm_model = llm_model
+    _default_completion_params = kwargs
+
     # If the model is set, we unset the default client
     _unset_default_client()
 
@@ -84,7 +88,7 @@ def get_default_client() -> LLMClient:
     try:
         from .litellm import LiteLLMClient
 
-        _default_client = LiteLLMClient(_default_llm_model)
+        _default_client = LiteLLMClient(_default_llm_model, _default_completion_params)
     except ImportError:
         raise ValueError(f"LLM scan using {_default_llm_model} requires litellm")
 

giskard/llm/client/litellm.py

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence
+from typing import Any, Dict, Optional, Sequence
 
 from ...client.python_utils import warning
 from ..errors import LLMImportError
@@ -23,8 +23,18 @@ def _get_response_format(format):
 
 
 class LiteLLMClient(LLMClient):
-    def __init__(self, model: str = "gpt-4o"):
+    def __init__(self, model: str = "gpt-4o", completion_params: Optional[Dict[str, Any]] = None):
+        """Initialize a LiteLLM completion client
+
+        Parameters
+        ----------
+        model : str
+            The name of the language model to use for text completion. see all supported LLMs: https://docs.litellm.ai/docs/providers/
+        completion_params : dict, optional
+            A dictionary containing params for the completion.
+        """
         self.model = model
+        self.completion_params = completion_params or dict()
 
     def _build_supported_completion_params(self, **kwargs):
         supported_params = litellm.get_supported_openai_params(model=self.model)
@@ -46,6 +56,7 @@ def complete(
             **self._build_supported_completion_params(
                 temperature=temperature, max_tokens=max_tokens, seed=seed, response_format=_get_response_format(format)
             ),
+            **self.completion_params,
         )
 
         self.logger.log_call(

giskard/llm/embeddings/__init__.py

@@ -1,18 +1,20 @@
 from typing import Optional
 from typing_extensions import deprecated
 
-from ..client.litellm import LiteLLMClient
 from .base import BaseEmbedding
+from .litellm import LiteLLMEmbedding
 
 _default_embedding = None
+
 _default_embedding_model = "text-embedding-ada-002"
+_default_embedding_params = dict()
 
 
 def get_embedding_model() -> str:
     return _default_embedding_model
 
 
-def set_embedding_model(model: str):
+def set_embedding_model(model: str, **kwargs):
     """
     Set the default embedding model to be used with litellm.
 
@@ -22,7 +24,10 @@ def set_embedding_model(model: str):
         Model name (e.g. 'text-embedding-ada-002' or 'text-embedding-3-large').
     """
     global _default_embedding_model
+    global _default_embedding_params
+
     _default_embedding_model = model
+    _default_embedding_params = kwargs
 
 
 def get_default_embedding():
@@ -34,7 +39,9 @@ def get_default_embedding():
     """
     global _default_embedding
 
-    _default_embedding = _default_embedding or LiteLLMClient(model=get_embedding_model())
+    _default_embedding = _default_embedding or LiteLLMEmbedding(
+        model=get_embedding_model(), embedding_params=_default_embedding_params
+    )
 
     return _default_embedding
 

giskard/llm/embeddings/litellm.py

@@ -1,4 +1,4 @@
-from typing import Sequence
+from typing import Any, Dict, Optional, Sequence
 
 import numpy as np
 
@@ -13,16 +13,20 @@
 
 
 class LiteLLMEmbedding(BaseEmbedding):
-    def __init__(self, model: str, batch_size=40):
-        """
+    def __init__(self, model: str, embedding_params: Optional[Dict[str, Any]] = None, batch_size=40):
+        """Initialize a LiteLLM embedding client
+
         Parameters
         ----------
         model : str
             Model name.
+        embedding_params : dict, optional)
+            A dictionary containing params for the completion.
         batch_size : int, optional
             Batch size for embeddings, by default 40.
         """
         self.model = model
+        self.embedding_params = embedding_params or dict()
         self.batch_size = batch_size
 
     def embed(self, texts: Sequence[str]) -> np.ndarray:
@@ -31,7 +35,7 @@ def embed(self, texts: Sequence[str]) -> np.ndarray:
 
         embeddings = []
         for batch in batched(texts, self.batch_size):
-            response = litellm.embedding(model=self.model, input=batch)
+            response = litellm.embedding(model=self.model, input=batch, **self.embedding_params)
             embeddings.extend([item.embedding for item in response.data])
 
         return np.array(embeddings)