feat: add support to custom embeddings

.github/workflows/build-python.yml

@@ -135,15 +135,20 @@ jobs:
         run: |
           df -h
 
+      - name: Install uv
+        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # v7.2.0
+
       - name: Install dependencies (Windows)
         if: ${{ matrix.os == 'windows-2022' }}
         run: |
+          pdm config use_uv true
           pdm install -G :all --without tensorflow
           pdm list
 
       - name: Install dependencies
         if: ${{ matrix.os != 'windows-2022' }}
         run: |
+          pdm config use_uv true
           pdm install -G :all
           pdm list
 
@@ -357,9 +362,13 @@ jobs:
         uses: r-lib/actions/setup-pandoc@v2
         with:
           pandoc-version: "3.1.7" # https://github.com/jgm/pandoc/releases
+      - name: Install uv
+        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # v7.2.0
 
       - name: Install dependencies
-        run: pdm install -G :all
+        run: |
+          pdm config use_uv true
+          pdm install -G :all
 
       - name: Check new disk space
         run: |

.github/workflows/lock-deps.yml

@@ -58,8 +58,11 @@ jobs:
           version: head # Issue with PDM 2.20.1: https://github.com/pdm-project/pdm/issues/3271
           cache: false
 
+      - name: Install uv
+        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b # v7.2.0
+
       - name: Install dependencies
-        run: rm -rf pdm.lock && pdm lock -G :all
+        run: pdm config use_uv true && rm -rf pdm.lock && pdm lock -G :all
 
       - name: Configure git
         run: |

docs/open_source/scan/scan_llm/index.md

@@ -176,18 +176,39 @@ class MyCustomLLM(litellm.CustomLLM):
         )
 
         return litellm.ModelResponse(**response.json())
+
+    # Custom embeddings (available since LiteLLM v1.71.1)
+    def embedding(
+        self,
+        model: str,
+        input: list,
+        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` was not provided")
+
+        response = requests.post(
+            "https://www.my-custom-llm.ai/embeddings",
+            json={"model": model, "input": input},
+            headers={"Authorization": api_key},
+        )
+
+        return litellm.EmbeddingResponse(**response.json())
 
-os.eviron["MY_SECRET_KEY"] = "" # "my-secret-key"
+os.environ["MY_SECRET_KEY"] = "" # "my-secret-key"
 
 my_custom_llm = MyCustomLLM()
 
 litellm.custom_provider_map = [  # 👈 KEY STEP - REGISTER HANDLER
-    {"provider": "my-custom-llm-endpoint", "custom_handler": my_custom_llm}
+    {"provider": "my-custom-provider", "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)
+giskard.llm.set_llm_model("my-custom-provider/my-custom-model", api_key=api_key)
+giskard.llm.set_embedding_model("my-custom-provider/my-custom-model", api_key=api_key)
 ```
 
 ::::::

docs/open_source/setting_up/index.md

@@ -203,18 +203,39 @@ class MyCustomLLM(litellm.CustomLLM):
         )
 
         return litellm.ModelResponse(**response.json())
+
+    # Custom embeddings (available since LiteLLM v1.71.1)
+    def embedding(
+        self,
+        model: str,
+        input: list,
+        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` was not provided")
+
+        response = requests.post(
+            "https://www.my-custom-llm.ai/embeddings",
+            json={"model": model, "input": input},
+            headers={"Authorization": api_key},
+        )
+
+        return litellm.EmbeddingResponse(**response.json())
 
-os.eviron["MY_SECRET_KEY"] = "" # "my-secret-key"
+os.environ["MY_SECRET_KEY"] = "" # "my-secret-key"
 
 my_custom_llm = MyCustomLLM()
 
 litellm.custom_provider_map = [  # 👈 KEY STEP - REGISTER HANDLER
-    {"provider": "my-custom-llm-endpoint", "custom_handler": my_custom_llm}
+    {"provider": "my-custom-provider", "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)
+giskard.llm.set_llm_model("my-custom-provider/my-custom-model", api_key=api_key)
+giskard.llm.set_embedding_model("my-custom-provider/my-custom-model", api_key=api_key)
 ```
 
 If you run into any issues configuring the LLM client, don't hesitate to [ask us on Discord](https://discord.com/invite/ABvfpbu69R) or open a new issue on [our GitHub repo](https://github.com/Giskard-AI/giskard).

docs/open_source/testset_generation/testset_generation/index.md

@@ -214,18 +214,39 @@ class MyCustomLLM(litellm.CustomLLM):
         )
 
         return litellm.ModelResponse(**response.json())
+
+    # Custom embeddings (available since LiteLLM v1.71.1)
+    def embedding(
+        self,
+        model: str,
+        input: list,
+        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` was not provided")
+
+        response = requests.post(
+            "https://www.my-custom-llm.ai/embeddings",
+            json={"model": model, "input": input},
+            headers={"Authorization": api_key},
+        )
+
+        return litellm.EmbeddingResponse(**response.json())
 
-os.eviron["MY_SECRET_KEY"] = "" # "my-secret-key"
+os.environ["MY_SECRET_KEY"] = "" # "my-secret-key"
 
 my_custom_llm = MyCustomLLM()
 
 litellm.custom_provider_map = [  # 👈 KEY STEP - REGISTER HANDLER
-    {"provider": "my-custom-llm-endpoint", "custom_handler": my_custom_llm}
+    {"provider": "my-custom-provider", "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)
+giskard.llm.set_llm_model("my-custom-provider/my-custom-model", api_key=api_key)
+giskard.llm.set_embedding_model("my-custom-provider/my-custom-model", api_key=api_key)
 ```
 
 ::::::