Slightly improve error handling for external providers (#220)

* [kb] Improve error handling for OpenAIRecord encoder

Catch specific OpenAI errors and re-raise with a clear message

* [kb] Improve the error in case of failing to infer dimensionality

* Added option to ignore specific warnings, or ignore warnings from a specific module

In addition, the `transformers` module is very insisstent, so I added their dedicated mechanism for silencing warnings

* Bug fix - infering dimension incorrectly

* [CLI] Removed dedicated error message for OpenAI auth problem

Instead, individual RecordEncoders and LLMs would need to raise their own errors

* [server] Call health_check() on startup

To detect errors early

* [cli] Bug fix - validate function did not connect to KB

* [LLM] Explicit error when initializing and when using OpenAI components - both LLM and RecordEncoder

* [kb] Slightly better solution for error handling

Catch the error in RecordEncoder base class, then format for each inhertor differently

* linter

* revert accidental commit

* [test] KB test - improved dimension infer testing

1. Needed to change after code was changed.
2. Added a few more test cases
3. Added more assertions on the naive create()

* [test] Fix LLM tests after error handling

Some error types changed

* [cli] Improve error message in case of index already exists

The CLI shouldn't mention `delete_index()`

* [CLI] Improve ChatEngine validate - use less tokens

Sending just "hello" without any limitations can use many tokens

* [LLM] Explicit error for function calling failed

Based on feedback in PR #220

Author: igiloh-pinecone

src/canopy/__init__.py

@@ -1,4 +1,25 @@
 import importlib.metadata
+import warnings
+import logging
+import os
+from typing import List
 
 # Taken from https://stackoverflow.com/a/67097076
 __version__ = importlib.metadata.version("canopy-sdk")
+
+
+IGNORED_WARNINGS: List[str] = [
+]
+
+IGNORED_WARNING_IN_MODULES = [
+    "transformers",
+]
+
+for warning in IGNORED_WARNINGS:
+    warnings.filterwarnings("ignore", message=warning)
+for module in IGNORED_WARNING_IN_MODULES:
+    warnings.filterwarnings("ignore", module=module)
+    logging.getLogger(module).setLevel(logging.ERROR)
+
+# Apparently, `transformers` has its own logging system, and needs to be silenced separately # noqa: E501
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"

src/canopy/knowledge_base/knowledge_base.py

@@ -307,24 +307,28 @@ def create_canopy_index(self,
         if dimension is None:
             try:
                 encoder_dimension = self._encoder.dimension
+                if encoder_dimension is None:
+                    raise RuntimeError(
+                        f"The selected encoder {self._encoder.__class__.__name__} does "
+                        f"not support inferring the vectors' dimensionality."
+                    )
+                dimension = encoder_dimension
             except Exception as e:
                 raise RuntimeError(
-                    f"Failed to infer vectors' dimension from encoder due to error: "
-                    f"{e}. Please fix the error or provide the dimension manually"
+                    f"Canopy has failed to infer vectors' dimensionality using the "
+                    f"selected encoder: {self._encoder.__class__.__name__}. You can "
+                    f"provide the dimension manually, try using a different encoder, or"
+                    f" fix the underlying error:\n{e}"
                 ) from e
-            if encoder_dimension is not None:
-                dimension = encoder_dimension
-            else:
-                raise ValueError("Could not infer dimension from encoder. "
-                                 "Please provide the vectors' dimension")
 
         # connect to pinecone and create index
         connect_to_pinecone()
 
         if self.index_name in list_indexes():
             raise RuntimeError(
-                f"Index {self.index_name} already exists. "
-                "If you wish to delete it, use `delete_index()`. "
+                f"Index {self.index_name} already exists. To connect to an "
+                f"existing index, use `knowledge_base.connect()`. "
+                "If you wish to delete it call `knowledge_base.delete_index()`. "
             )
 
         # create index

src/canopy/knowledge_base/record_encoder/base.py

@@ -100,7 +100,13 @@ def encode_documents(self, documents: List[KBDocChunk]) -> List[KBEncodedDocChun
         """   # noqa: E501
         encoded_docs = []
         for batch in self._batch_iterator(documents, self.batch_size):
-            encoded_docs.extend(self._encode_documents_batch(batch))
+            try:
+                encoded_docs.extend(self._encode_documents_batch(batch))
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to enconde documents using {self.__class__.__name__}. "
+                    f"Error: {self._format_error(e)}"
+                ) from e
 
         return encoded_docs  # TODO: consider yielding a generator
 
@@ -118,7 +124,13 @@ def encode_queries(self, queries: List[Query]) -> List[KBQuery]:
 
         kb_queries = []
         for batch in self._batch_iterator(queries, self.batch_size):
-            kb_queries.extend(self._encode_queries_batch(batch))
+            try:
+                kb_queries.extend(self._encode_queries_batch(batch))
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to enconde queries using {self.__class__.__name__}. "
+                    f"Error: {self._format_error(e)}"
+                ) from e
 
         return kb_queries
 
@@ -137,3 +149,6 @@ async def aencode_queries(self, queries: List[Query]) -> List[KBQuery]:
             kb_queries.extend(await self._aencode_queries_batch(batch))
 
         return kb_queries
+
+    def _format_error(self, err):
+        return f"{err}"

src/canopy/knowledge_base/record_encoder/dense.py

@@ -64,7 +64,8 @@ def dimension(self) -> int:
         Returns:
             dimension(int): the dimension of the encoder
         """  # noqa: E501
-        return len(self._dense_encoder.encode_documents(["hello"])[0])
+        dummy_doc = KBDocChunk(text="hello", id="dummy_doc", document_id="dummy_doc")
+        return len(self.encode_documents([dummy_doc])[0].values)
 
     async def _aencode_documents_batch(self,
                                        documents: List[KBDocChunk]

src/canopy/knowledge_base/record_encoder/openai.py

@@ -1,10 +1,19 @@
 from typing import List
+
+from openai import OpenAIError, RateLimitError, APIConnectionError, AuthenticationError
 from pinecone_text.dense.openai_encoder import OpenAIEncoder
 from canopy.knowledge_base.models import KBDocChunk, KBEncodedDocChunk, KBQuery
 from canopy.knowledge_base.record_encoder.dense import DenseRecordEncoder
 from canopy.models.data_models import Query
 
 
+def _format_openai_error(e):
+    try:
+        return e.response.json()['error']['message']
+    except Exception:
+        return str(e)
+
+
 class OpenAIRecordEncoder(DenseRecordEncoder):
     """
     OpenAIRecordEncoder is a type of DenseRecordEncoder that uses the OpenAI `embeddings` API.
@@ -27,25 +36,31 @@ def __init__(self,
                         Defaults to 400.
             **kwargs: Additional arguments to pass to the underlying `pinecone-text. OpenAIEncoder`.
         """  # noqa: E501
-        encoder = OpenAIEncoder(model_name, **kwargs)
+        try:
+            encoder = OpenAIEncoder(model_name, **kwargs)
+        except OpenAIError as e:
+            raise RuntimeError(
+                "Failed to connect to OpenAI, please make sure that the OPENAI_API_KEY "
+                "environment variable is set correctly.\n"
+                f"Error: {_format_openai_error(e)}"
+            ) from e
         super().__init__(dense_encoder=encoder, batch_size=batch_size)
 
-    def encode_documents(self, documents: List[KBDocChunk]) -> List[KBEncodedDocChunk]:
-        """
-        Encode a list of documents, takes a list of KBDocChunk and returns a list of KBEncodedDocChunk.
-
-        Args:
-            documents: A list of KBDocChunk to encode.
-
-        Returns:
-            encoded chunks: A list of KBEncodedDocChunk, with the `values` field populated by the generated embeddings vector.
-        """  # noqa: E501
-        return super().encode_documents(documents)
-
     async def _aencode_documents_batch(self,
                                        documents: List[KBDocChunk]
                                        ) -> List[KBEncodedDocChunk]:
         raise NotImplementedError
 
     async def _aencode_queries_batch(self, queries: List[Query]) -> List[KBQuery]:
         raise NotImplementedError
+
+    def _format_error(self, err):
+        if isinstance(err, RateLimitError):
+            return (f"Your OpenAI account seem to have reached the rate limit. "
+                    f"Details: {_format_openai_error(err)}")
+        elif isinstance(err, (AuthenticationError, APIConnectionError)):
+            return (f"Failed to connect to OpenAI, please make sure that the "
+                    f"OPENAI_API_KEY environment variable is set correctly. "
+                    f"Details: {_format_openai_error(err)}")
+        else:
+            return _format_openai_error(err)

src/canopy/llm/openai.py

@@ -17,6 +17,13 @@
 from canopy.models.data_models import Messages, Query
 
 
+def _format_openai_error(e):
+    try:
+        return e.response.json()['error']['message']
+    except Exception:
+        return str(e)
+
+
 class OpenAILLM(BaseLLM):
     """
     OpenAI LLM wrapper built on top of the OpenAI Python client.
@@ -48,9 +55,17 @@ def __init__(self,
                     These params can be overridden by passing a `model_params` argument to the `chat_completion` or `enforced_function_call` methods.
         """  # noqa: E501
         super().__init__(model_name)
-        self._client = openai.OpenAI(api_key=api_key,
-                                     organization=organization,
-                                     base_url=base_url)
+        try:
+            self._client = openai.OpenAI(api_key=api_key,
+                                         organization=organization,
+                                         base_url=base_url)
+        except openai.OpenAIError as e:
+            raise RuntimeError(
+                "Failed to connect to OpenAI, please make sure that the OPENAI_API_KEY "
+                "environment variable is set correctly.\n"
+                f"Error: {_format_openai_error(e)}"
+            )
+
         self.default_model_params = kwargs
 
     @property
@@ -96,11 +111,19 @@ def chat_completion(self,
         )
 
         messages = [m.dict() for m in messages]
-        response = self._client.chat.completions.create(model=self.model_name,
-                                                        messages=messages,
-                                                        stream=stream,
-                                                        max_tokens=max_tokens,
-                                                        **model_params_dict)
+        try:
+            response = self._client.chat.completions.create(model=self.model_name,
+                                                            messages=messages,
+                                                            stream=stream,
+                                                            max_tokens=max_tokens,
+                                                            **model_params_dict)
+        except openai.OpenAIError as e:
+            provider_name = self.__class__.__name__.replace("LLM", "")
+            raise RuntimeError(
+                f"Failed to use {provider_name}'s {self.model_name} model for chat "
+                f"completion.\n"
+                f"Error: {_format_openai_error(e)}"
+            )
 
         def streaming_iterator(response):
             for chunk in response:
@@ -175,15 +198,23 @@ def enforced_function_call(self,
         function_dict = cast(ChatCompletionToolParam,
                              {"type": "function", "function": function.dict()})
 
-        chat_completion = self._client.chat.completions.create(
-            messages=[m.dict() for m in messages],
-            model=self.model_name,
-            tools=[function_dict],
-            tool_choice={"type": "function",
-                         "function": {"name": function.name}},
-            max_tokens=max_tokens,
-            **model_params_dict
-        )
+        try:
+            chat_completion = self._client.chat.completions.create(
+                messages=[m.dict() for m in messages],
+                model=self.model_name,
+                tools=[function_dict],
+                tool_choice={"type": "function",
+                             "function": {"name": function.name}},
+                max_tokens=max_tokens,
+                **model_params_dict
+            )
+        except openai.OpenAIError as e:
+            provider_name = self.__class__.__name__.replace("LLM", "")
+            raise RuntimeError(
+                f"Failed to use {provider_name}'s {self.model_name} model for "
+                f"chat completion with enforced function calling.\n"
+                f"Error: {_format_openai_error(e)}"
+            )
 
         result = chat_completion.choices[0].message.tool_calls[0].function.arguments
         arguments = json.loads(result)

src/canopy_cli/cli.py

@@ -22,7 +22,7 @@
 from canopy.knowledge_base import connect_to_pinecone
 from canopy.knowledge_base.chunker import Chunker
 from canopy.chat_engine import ChatEngine
-from canopy.models.data_models import Document
+from canopy.models.data_models import Document, UserMessage
 from canopy.tokenizer import Tokenizer
 from canopy_cli.data_loader import (
     load_from_path,
@@ -44,12 +44,6 @@
 DEFAULT_SERVER_URL = f"http://localhost:8000/{API_VERSION}"
 spinner = Spinner()
 
-OPENAI_AUTH_ERROR_MSG = (
-    "Failed to connect to OpenAI, please make sure that the OPENAI_API_KEY "
-    "environment variable is set correctly.\n"
-    "Please visit https://platform.openai.com/account/api-keys for more details"
-)
-
 
 def check_server_health(url: str):
     try:
@@ -140,9 +134,15 @@ def _validate_chat_engine(config_file: Optional[str]):
     config = _read_config_file(config_file)
     Tokenizer.initialize()
     try:
-        ChatEngine.from_config(config.get("chat_engine", {}))
-    except openai.OpenAIError:
-        raise CLIError(OPENAI_AUTH_ERROR_MSG)
+        # If the server itself will fail, we can't except the error, since it's running
+        # in a different process. Try to load and run the ChatEngine so we can catch
+        # any errors and print a nice message.
+        chat_engine = ChatEngine.from_config(config.get("chat_engine", {}))
+        chat_engine.max_generated_tokens = 5
+        chat_engine.context_engine.knowledge_base.connect()
+        chat_engine.chat(
+            [UserMessage(content="This is a health check. Are you alive? Be concise")]
+        )
     except Exception as e:
         msg = f"Failed to initialize Canopy server. Reason:\n{e}"
         if config_file:
@@ -229,9 +229,15 @@ def new(index_name: str, config: Optional[str]):
     with spinner:
         try:
             kb.create_canopy_index()
-        # TODO: kb should throw a specific exception for each case
+        # TODO: kb should throw a specific exception for failure
         except Exception as e:
-            msg = f"Failed to create a new index. Reason:\n{e}"
+            already_exists_str = f"Index {kb.index_name} already exists"
+            if isinstance(e, RuntimeError) and already_exists_str in str(e):
+                msg = (f"{already_exists_str}, please use a different name."
+                       f"If you wish to delete the index, log in to Pinecone's "
+                       f"Console: https://app.pinecone.io/")
+            else:
+                msg = f"Failed to create a new index. Reason:\n{e}"
             raise CLIError(msg)
     click.echo(click.style("Success!", fg="green"))
     os.environ["INDEX_NAME"] = index_name
@@ -303,8 +309,8 @@ def upsert(index_name: str,
     kb_config = _load_kb_config(config)
     try:
         kb = KnowledgeBase.from_config(kb_config, index_name=index_name)
-    except openai.OpenAIError:
-        raise CLIError(OPENAI_AUTH_ERROR_MSG)
+    except Exception as e:
+        raise CLIError(str(e))
 
     try:
         kb.connect()

src/canopy_server/app.py

@@ -238,7 +238,7 @@ async def health_check() -> HealthStatus:
 
     try:
         msg = UserMessage(content="This is a health check. Are you alive? Be concise")
-        await run_in_threadpool(llm.chat_completion, messages=[msg], max_tokens=50)
+        await run_in_threadpool(llm.chat_completion, messages=[msg], max_tokens=5)
     except Exception as e:
         err_msg = f"Failed to communicate with {llm.__class__.__name__}"
         logger.exception(err_msg)
@@ -281,6 +281,7 @@ async def startup():
     _init_logging()
     _init_engines()
     _init_routes(app)
+    await health_check()
 
 
 def _init_routes(app):