Merge branch 'main' into CM-8548_fix_error_message_for_invalid_timestamp

src/comet_llm/chains/api.py

@@ -16,11 +16,20 @@
 import json
 from typing import Dict, List, Optional
 
-from .. import app, convert, experiment_api, experiment_info, llm_result
+from .. import (
+    app,
+    config,
+    convert,
+    exceptions,
+    experiment_api,
+    experiment_info,
+    llm_result,
+)
 from ..types import JSONEncodable
 from . import chain, state
 
 
+@exceptions.filter(allow_raising=config.raising_enabled(), summary=app.SUMMARY)
 def start_chain(
     inputs: Dict[str, JSONEncodable],
     api_key: Optional[str] = None,

src/comet_llm/config.py

@@ -37,6 +37,7 @@ def _extend_comet_ml_config() -> None:
     CONFIG_MAP_EXTENSION = {
         "comet.disable": {"type": int, "default": 0},
         "comet.logging.console": {"type": str, "default": "INFO"},
+        "comet.raise_exceptions_on_error": {"type": int, "default": 1},
     }
 
     comet_ml_config.CONFIG_MAP.update(CONFIG_MAP_EXTENSION)
@@ -74,6 +75,10 @@ def comet_disabled() -> bool:
     return bool(_COMET_ML_CONFIG["comet.disable"])
 
 
+def raising_enabled() -> bool:
+    return bool(_COMET_ML_CONFIG["comet.raise_exceptions_on_error"])
+
+
 def logging_available() -> bool:
     if api_key() is None:
         return False

src/comet_llm/logs_registry.py → src/comet_llm/exceptions/__init__.py

@@ -12,19 +12,5 @@
 #  LICENSE file in the root directory of this package.
 # *******************************************************
 
-import collections
-from typing import DefaultDict, Dict
-
-
-class LogsRegistry:
-    def __init__(self) -> None:
-        self._registry: DefaultDict[str, int] = collections.defaultdict(lambda: 0)
-
-    def register_log(self, project_url: str) -> None:
-        self._registry[project_url] += 1
-
-    def as_dict(self) -> Dict[str, int]:
-        return self._registry.copy()
-
-    def empty(self) -> bool:
-        return len(self._registry) == 0
+from .exceptions import CometLLMException
+from .filter_decorator import filter

src/comet_llm/exceptions.py → src/comet_llm/exceptions/exceptions.py

@@ -14,4 +14,6 @@
 
 
 class CometLLMException(Exception):
-    pass
+    def __init__(self, *args, log_message_once: bool = False) -> None:  # type: ignore
+        super().__init__(*args)
+        self.log_message_once = log_message_once

src/comet_llm/exceptions/filter_decorator.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# *******************************************************
+#   ____                     _               _
+#  / ___|___  _ __ ___   ___| |_   _ __ ___ | |
+# | |   / _ \| '_ ` _ \ / _ \ __| | '_ ` _ \| |
+# | |__| (_) | | | | | |  __/ |_ _| | | | | | |
+#  \____\___/|_| |_| |_|\___|\__(_)_| |_| |_|_|
+#
+#  Sign up for free at https://www.comet.com
+#  Copyright (C) 2015-2023 Comet ML INC
+#  This source code is licensed under the MIT license found in the
+#  LICENSE file in the root directory of this package.
+# *******************************************************
+
+import functools
+import logging
+from typing import TYPE_CHECKING, Any, Callable
+
+from comet_llm import logging as comet_logging
+
+if TYPE_CHECKING:
+    from comet_llm import summary
+
+LOGGER = logging.getLogger(__name__)
+
+
+def filter(allow_raising: bool, summary: "summary.Summary") -> Callable:
+    def decorator(function: Callable) -> Callable:
+        @functools.wraps(function)
+        def wrapper(*args, **kwargs) -> Any:  # type: ignore
+            try:
+                return function(*args, **kwargs)
+            except Exception as exception:
+                summary.increment_failed()
+
+                if allow_raising:
+                    raise
+
+                if getattr(exception, "log_message_once", False):
+                    comet_logging.log_once_at_level(
+                        LOGGER,
+                        logging.ERROR,
+                        str(exception),
+                        exc_info=True,
+                        extra={"show_traceback": True},
+                    )
+                else:
+                    LOGGER.error(
+                        str(exception),
+                        exc_info=True,
+                        extra={"show_traceback": True},
+                    )
+
+        return wrapper
+
+    return decorator

src/comet_llm/experiment_api/request_exception_wrapper.py

@@ -21,7 +21,7 @@
 import requests  # type: ignore
 
 from .. import config, exceptions
-from ..handlers import failed_response
+from . import failed_response_handler
 
 LOGGER = logging.getLogger(__name__)
 
@@ -44,7 +44,9 @@ def wrapper(*args, **kwargs) -> Any:  # type: ignore
                             f"installation is up-to-date and check the traceback for more details."
                         )
                 if exception.response is not None:
-                    exception_args.append(failed_response.handle(exception.response))
+                    exception_args.append(
+                        failed_response_handler.handle(exception.response)
+                    )
 
                 _debug_log(exception)
 

src/comet_llm/experiment_info.py

@@ -34,8 +34,11 @@ def get(
     api_key_not_found_message: Optional[str] = None,
 ) -> ExperimentInfo:
     api_key = api_key if api_key else config.api_key()
+
     if api_key is None and api_key_not_found_message is not None:
-        raise exceptions.CometLLMException(api_key_not_found_message)
+        raise exceptions.CometLLMException(
+            api_key_not_found_message, log_message_once=True
+        )
 
     workspace = workspace if workspace else config.workspace()
     project_name = project_name if project_name else config.project_name()

src/comet_llm/logging.py

@@ -17,10 +17,10 @@
 import sys
 from typing import Any, Callable
 
-_LOG_ONCE_CACHE = set()
-
 from . import config
 
+_LOG_ONCE_CACHE = set()
+
 
 def setup() -> None:
     root = logging.getLogger("comet_llm")

src/comet_llm/prompts/api.py

@@ -18,11 +18,12 @@
 
 import comet_llm.convert
 
-from .. import app, experiment_api, experiment_info, llm_result
+from .. import app, config, exceptions, experiment_api, experiment_info, llm_result
 from ..chains import version
 from . import convert, preprocess
 
 
+@exceptions.filter(allow_raising=config.raising_enabled(), summary=app.SUMMARY)
 def log_prompt(
     prompt: str,
     output: str,

src/comet_llm/summary.py

@@ -12,25 +12,36 @@
 #  LICENSE file in the root directory of this package.
 # *******************************************************
 
+import collections
 import logging
-
-from . import logs_registry
+import threading
+from typing import DefaultDict
 
 LOGGER = logging.getLogger(__name__)
 
 
 class Summary:
     def __init__(self) -> None:
-        self._registry = logs_registry.LogsRegistry()
+        self._logs_registry: DefaultDict[str, int] = collections.defaultdict(lambda: 0)
+        self._failed = 0
+        self._lock = threading.Lock()
 
     def add_log(self, project_url: str, name: str) -> None:
-        if self._registry.empty():
-            LOGGER.info("%s logged to %s", name.capitalize(), project_url)
+        with self._lock:
+            if len(self._logs_registry) == 0:
+                LOGGER.info("%s logged to %s", name.capitalize(), project_url)
 
-        self._registry.register_log(project_url)
+            self._logs_registry[project_url] += 1
 
-    def print(self) -> None:
-        registry_items = self._registry.as_dict().items()
+    def increment_failed(self) -> None:
+        with self._lock:
+            self._failed += 1
 
-        for project, logs_amount in registry_items:
+    def print(self) -> None:
+        for project, logs_amount in self._logs_registry.items():
             LOGGER.info("%d prompts and chains logged to %s", logs_amount, project)
+
+        if self._failed > 0:
+            LOGGER.info(
+                "%d prompts and chains were not logged because of errors", self._failed
+            )

tests/unit/exceptions/test_filter_decorator.py

@@ -0,0 +1,63 @@
+import logging
+
+import pytest
+from testix import *
+
+from comet_llm.exceptions import exceptions, filter_decorator
+
+
+@pytest.fixture(autouse=True)
+def mock_imports(patch_module):
+    patch_module(filter_decorator, "LOGGER")
+    patch_module(filter_decorator, "comet_logging")
+
+
+def test_filter__no_exceptions_raised__nothing_done():
+    NOT_USED = None
+    @filter_decorator.filter(allow_raising=True, summary=NOT_USED)
+    def f():
+        return 42
+
+    assert f() == 42
+
+
+def test_filter__upraising_allowed__function_raised_exception__exception_raised_to_user():
+    @filter_decorator.filter(allow_raising=True, summary=Fake("summary"))
+    def f():
+        raise Exception("some-message")
+
+    with Scenario() as s:
+        s.summary.increment_failed()
+        with pytest.raises(Exception):
+            f()
+
+
+def test_filter__upraising_not_allowed__function_raised_exception__exception_info_logged():
+    @filter_decorator.filter(allow_raising=False, summary=Fake("summary"))
+    def f():
+        raise Exception("some-message")
+
+    with Scenario() as s:
+        s.summary.increment_failed()
+        s.LOGGER.error(
+            "some-message",
+            exc_info=True,
+            extra={"show_traceback": True}
+        )
+        assert f() is None
+
+
+def test_filter__upraising_not_allowed__function_raised_exception__exception_has_log_message_once_attribute_True__exception_info_logged_once():
+    @filter_decorator.filter(allow_raising=False, summary=Fake("summary"))
+    def f():
+        raise exceptions.CometLLMException("some-message", log_message_once=True)
+    with Scenario() as s:
+        s.summary.increment_failed()
+        s.comet_logging.log_once_at_level(
+            filter_decorator.LOGGER,
+            logging.ERROR,
+            "some-message",
+            exc_info=True,
+            extra={"show_traceback": True}
+        )
+        assert f() is None

tests/unit/experiment_api/test_request_exception_wrapper.py

@@ -1,17 +1,18 @@
 import json
 
+import box
 import pytest
 import requests
 from testix import *
 
-from comet_llm import exceptions
+from comet_llm.exceptions import exceptions
 from comet_llm.experiment_api import request_exception_wrapper
 
 
 @pytest.fixture(autouse=True)
 def mock_imports(patch_module):
     patch_module(request_exception_wrapper, "config")
-    patch_module(request_exception_wrapper, "failed_response")
+    patch_module(request_exception_wrapper, "failed_response_handler")
 
 
 def test_wrap_no_exceptions():
@@ -56,18 +57,18 @@ def f():
 
 
 def test_wrap__request_exception_non_llm_project_sdk_code__log_specifc_message_in_exception():
+    response = box.Box(text=json.dumps({"sdk_error_code": 34323}))
+
     @request_exception_wrapper.wrap()
     def f():
         exception = requests.RequestException()
-        response = requests.Response
-        response.text = json.dumps({"sdk_error_code": 34323})
         exception.response = response
         raise exception
 
     expected_log_message = "Failed to send prompt to the specified project as it is not an LLM project, please specify a different project name."
 
     with Scenario() as s:
-        s.failed_response.handle(requests.Response) >> expected_log_message
+        s.failed_response_handler.handle(response) >> expected_log_message
         with pytest.raises(exceptions.CometLLMException) as excinfo:
             f()
 

tests/unit/prompts/test_preprocess.py

@@ -1,7 +1,7 @@
 import pytest
 from testix import *
 
-from comet_llm import exceptions
+from comet_llm.exceptions import exceptions
 from comet_llm.prompts import preprocess
 
 

tests/unit/test_experiment_info.py

@@ -1,7 +1,8 @@
 import pytest
 from testix import *
 
-from comet_llm import exceptions, experiment_info
+from comet_llm import experiment_info
+from comet_llm.exceptions import exceptions
 
 
 @pytest.fixture(autouse=True)
@@ -13,7 +14,7 @@ def test_get__api_key_not_specified_and_not_found_in_config__exception_raised():
     with Scenario() as s:
         s.config.api_key() >> None
 
-        with pytest.raises(exceptions.CometLLMException, match="error-message"):
+        with pytest.raises(exceptions.CometLLMException, match="error-message") as exc_info:
             experiment_info.get(
                 None,
                 "the-workspace",