From b40dbf3c68a2def746cd83a78e6d61408ef7c502 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 13 Aug 2025 11:50:24 +0200
Subject: [PATCH] feat: adds new error classes to litellm instructor

---
 cognee/infrastructure/llm/exceptions.py       | 21 +++++++++++++++++--
 .../litellm_instructor/llm/get_llm_client.py  |  5 ++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/cognee/infrastructure/llm/exceptions.py b/cognee/infrastructure/llm/exceptions.py
index af3aa5832..287820448 100644
--- a/cognee/infrastructure/llm/exceptions.py
+++ b/cognee/infrastructure/llm/exceptions.py
@@ -1,5 +1,22 @@
-from cognee.exceptions.exceptions import CriticalError
+from cognee.exceptions.exceptions import CogneeValidationError
 
 
-class ContentPolicyFilterError(CriticalError):
+class ContentPolicyFilterError(CogneeValidationError):
     pass
+
+
+class LLMAPIKeyNotSetError(CogneeValidationError):
+    """
+    Raised when the LLM API key is not set in the configuration.
+    """
+    def __init__(self, message: str = "LLM API key is not set."):
+        super().__init__(message=message, name="LLMAPIKeyNotSetError")
+
+
+class UnsupportedLLMProviderError(CogneeValidationError):
+    """
+    Raised when an unsupported LLM provider is specified in the configuration.
+    """
+    def __init__(self, provider: str):
+        message = f"Unsupported LLM provider: {provider}"
+        super().__init__(message=message, name="UnsupportedLLMProviderError")
diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py
index 22d101077..b486fa089 100644
--- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py
+++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py
@@ -2,11 +2,14 @@
 
 from enum import Enum
 
-from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.llm import get_llm_config
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.ollama.adapter import (
     OllamaAPIAdapter,
 )
+from cognee.infrastructure.llm.exceptions import (
+    LLMAPIKeyNotSetError,
+    UnsupportedLLMProviderError,
+)
 
 
 # Define an Enum for LLM Providers