[CM-10021] update comet llm to use new endpoint for logging chains and prompts (#128)

* Initial implementation of new endpoint usage

* Draft changes

* Add compatibility for old format of messages in offline mode

* Fix some lint errors

* Add senders dispatching based on backend version

* Add backward compatible send_v2 methods to chain and prompt senders

* Fix lint errors, remove redundant proxy payload class

* Update client

* Update tests

* Update log_chain request function to follow new response format

* Fix lint errors

* Remove exception catching in offline and online message processors

* Add constants.py

* Fix lint errors

* Update outdated tests

* Bump v2 backend version

Author: alexkuzmik

src/comet_llm/chains/api.py

@@ -105,6 +105,7 @@ def log_chain(chain: chain.Chain) -> Optional[llm_result.LLMResult]:
     chain_data = chain.as_dict()
 
     message = messages.ChainMessage(
+        id=messages.generate_id(),
         experiment_info_=chain.experiment_info,
         tags=chain.tags,
         chain_data=chain_data,

src/comet_llm/experiment_api/comet_api_client.py

@@ -13,26 +13,39 @@
 # *******************************************************
 
 import functools
+import logging
 import urllib.parse
 import warnings
-from typing import IO, List, Optional
+from typing import IO, Any, Dict, List, Optional
 
 import requests  # type: ignore
 import urllib3.exceptions
 
-from .. import config
+from .. import config, exceptions, semantic_version
 from ..types import JSONEncodable
-from . import request_exception_wrapper
+from . import error_codes_mapping, payload_constructor
 
 ResponseContent = JSONEncodable
 
+LOGGER = logging.getLogger(__name__)
+
 
 class CometAPIClient:
     def __init__(self, api_key: str, comet_url: str, session: requests.Session):
         self._headers = {"Authorization": api_key}
         self._comet_url = comet_url
         self._session = session
 
+        self.backend_version = semantic_version.SemanticVersion.parse(
+            self.is_alive_ver()["version"]
+        )
+
+    def is_alive_ver(self) -> ResponseContent:
+        return self._request(
+            "GET",
+            "api/isAlive/ver",
+        )
+
     def create_experiment(
         self,
         type_: str,
@@ -122,6 +135,56 @@ def log_experiment_other(
             },
         )
 
+    def log_chain(
+        self,
+        experiment_key: str,
+        chain_asset: Dict[str, JSONEncodable],
+        workspace: Optional[str] = None,
+        project: Optional[str] = None,
+        parameters: Optional[Dict[str, JSONEncodable]] = None,
+        metrics: Optional[Dict[str, JSONEncodable]] = None,
+        tags: Optional[List[str]] = None,
+        others: Optional[Dict[str, JSONEncodable]] = None,
+    ) -> ResponseContent:
+        json = [
+            {
+                "experimentKey": experiment_key,
+                "createExperimentRequest": {
+                    "workspaceName": workspace,
+                    "projectName": project,
+                    "type": "LLM",
+                },
+                "parameters": payload_constructor.chain_parameters_payload(parameters),
+                "metrics": payload_constructor.chain_metrics_payload(metrics),
+                "others": payload_constructor.chain_others_payload(others),
+                "tags": tags,
+                "jsonAsset": {
+                    "extension": "json",
+                    "type": "llm_data",
+                    "fileName": "comet_llm_data.json",
+                    "file": chain_asset,
+                },
+            }
+        ]  # we make a list because endpoint is designed for batches
+
+        batched_response: Dict[str, Dict[str, Any]] = self._request(
+            "POST",
+            "api/rest/v2/write/experiment/llm",
+            json=json,
+        )
+        sub_response = list(batched_response.values())[0]
+        status = sub_response["status"]
+        if status != 200:
+            LOGGER.debug(
+                "Failed to send trace: \nPayload %s, Response %s",
+                str(json),
+                str(batched_response),
+            )
+            error_code = sub_response["content"]["sdk_error_code"]
+            raise exceptions.CometLLMException(error_codes_mapping.MESSAGES[error_code])
+
+        return sub_response["content"]
+
     def _request(self, method: str, path: str, *args, **kwargs) -> ResponseContent:  # type: ignore
         url = urllib.parse.urljoin(self._comet_url, path)
         response = self._session.request(

src/comet_llm/experiment_api/failed_response_handler.py renamed to src/comet_llm/experiment_api/error_codes_mapping.py

@@ -13,24 +13,12 @@
 # *******************************************************
 
 import collections
-import json
-from typing import NoReturn
 
-import requests  # type: ignore
+from .. import backend_error_codes, logging_messages
 
-from .. import backend_error_codes, exceptions, logging_messages
-
-_SDK_ERROR_CODES_LOGGING_MESSAGE = collections.defaultdict(
+MESSAGES = collections.defaultdict(
     lambda: logging_messages.FAILED_TO_SEND_DATA_TO_SERVER,
     {
         backend_error_codes.UNABLE_TO_LOG_TO_NON_LLM_PROJECT: logging_messages.UNABLE_TO_LOG_TO_NON_LLM_PROJECT
     },
 )
-
-
-def handle(exception: requests.RequestException) -> NoReturn:
-    response = exception.response
-    sdk_error_code = json.loads(response.text)["sdk_error_code"]
-    error_message = _SDK_ERROR_CODES_LOGGING_MESSAGE[sdk_error_code]
-
-    raise exceptions.CometLLMException(error_message) from exception

src/comet_llm/experiment_api/payload_constructor.py

@@ -0,0 +1,46 @@
+# -*- 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.
+# *******************************************************
+
+from typing import Dict, List, Optional
+
+from ..types import JSONEncodable
+
+
+def chain_parameters_payload(
+    parameters: Optional[Dict[str, JSONEncodable]]
+) -> List[Dict[str, JSONEncodable]]:
+    return _dict_to_payload_format(parameters, "parameterName", "parameterValue")
+
+
+def chain_metrics_payload(
+    metrics: Optional[Dict[str, JSONEncodable]]
+) -> List[Dict[str, JSONEncodable]]:
+    return _dict_to_payload_format(metrics, "metricName", "metricValue")
+
+
+def chain_others_payload(
+    others: Optional[Dict[str, JSONEncodable]]
+) -> List[Dict[str, JSONEncodable]]:
+    return _dict_to_payload_format(others, "key", "value")
+
+
+def _dict_to_payload_format(
+    source: Optional[Dict[str, JSONEncodable]], key_name: str, value_name: str
+) -> List[Dict[str, JSONEncodable]]:
+    if source is None:
+        return []
+
+    result = [{key_name: key, value_name: value} for key, value in source.items()]
+
+    return result

src/comet_llm/experiment_api/request_exception_wrapper.py

@@ -13,15 +13,16 @@
 # *******************************************************
 
 import functools
+import json
 import logging
 import urllib.parse
 from pprint import pformat
-from typing import Any, Callable, List
+from typing import Any, Callable, List, NoReturn
 
 import requests  # type: ignore
 
 from .. import config, exceptions, logging_messages
-from . import failed_response_handler
+from . import error_codes_mapping
 
 LOGGER = logging.getLogger(__name__)
 
@@ -49,7 +50,7 @@ def wrapper(*args, **kwargs) -> Any:  # type: ignore
                         logging_messages.FAILED_TO_SEND_DATA_TO_SERVER
                     ) from exception
 
-                failed_response_handler.handle(exception)
+                _handle_request_exception(exception)
 
         return wrapper
 
@@ -73,3 +74,11 @@ def _debug_log(exception: requests.RequestException) -> None:
         # Make sure we won't fail on attempt to debug.
         # It's mainly for tests when response object can be mocked
         pass
+
+
+def _handle_request_exception(exception: requests.RequestException) -> NoReturn:
+    response = exception.response
+    sdk_error_code = json.loads(response.text)["sdk_error_code"]
+    error_message = error_codes_mapping.MESSAGES[sdk_error_code]
+
+    raise exceptions.CometLLMException(error_message) from exception

src/comet_llm/message_processing/messages.py

@@ -13,24 +13,32 @@
 # *******************************************************
 
 import dataclasses
-import inspect
+import uuid
 from typing import Any, ClassVar, Dict, List, Optional, Union
 
 from comet_llm.types import JSONEncodable
 
 from .. import experiment_info, logging_messages
 
 
+def generate_id() -> str:
+    return uuid.uuid4().hex
+
+
 @dataclasses.dataclass
 class BaseMessage:
     experiment_info_: experiment_info.ExperimentInfo
+    id: str
     VERSION: ClassVar[int]
 
     @classmethod
     def from_dict(
         cls, d: Dict[str, Any], api_key: Optional[str] = None
     ) -> "BaseMessage":
-        d.pop("VERSION")  #
+        version = d.pop("VERSION")
+        if version == 1:
+            # Message was dumped before id was introduced. We can generate it now.
+            d["id"] = generate_id()
 
         experiment_info_dict: Dict[str, Optional[str]] = d.pop("experiment_info_")
         experiment_info_ = experiment_info.get(
@@ -57,7 +65,7 @@ class PromptMessage(BaseMessage):
     metadata: Optional[Dict[str, Union[str, bool, float, None]]]
     tags: Optional[List[str]]
 
-    VERSION: ClassVar[int] = 1
+    VERSION: ClassVar[int] = 2
 
 
 @dataclasses.dataclass
@@ -69,4 +77,4 @@ class ChainMessage(BaseMessage):
     others: Dict[str, JSONEncodable]
     # 'other' - is a name of an attribute of experiment, logged via log_other
 
-    VERSION: ClassVar[int] = 1
+    VERSION: ClassVar[int] = 2

src/comet_llm/message_processing/offline_message_processor.py

@@ -44,15 +44,9 @@ def process(self, message: messages.BaseMessage) -> None:
             file_path = pathlib.Path(self._offline_directory, self._current_file_name)
 
             if isinstance(message, messages.PromptMessage):
-                try:
-                    return prompt.send(message, str(file_path))
-                except Exception:
-                    LOGGER.error("Failed to log prompt", exc_info=True)
+                return prompt.send(message, str(file_path))
             elif isinstance(message, messages.ChainMessage):
-                try:
-                    return chain.send(message, str(file_path))
-                except Exception:
-                    LOGGER.error("Failed to log chain", exc_info=True)
+                return chain.send(message, str(file_path))
 
         LOGGER.debug(f"Unsupported message type {message}")
         return None

src/comet_llm/message_processing/online_message_processor.py

@@ -28,15 +28,9 @@ def __init__(self) -> None:
 
     def process(self, message: messages.BaseMessage) -> Optional[llm_result.LLMResult]:
         if isinstance(message, messages.PromptMessage):
-            try:
-                return prompt.send(message)
-            except Exception:
-                LOGGER.error("Failed to log prompt", exc_info=True)
+            return prompt.send(message)
         elif isinstance(message, messages.ChainMessage):
-            try:
-                return chain.send(message)
-            except Exception:
-                LOGGER.error("Failed to log chain", exc_info=True)
+            return chain.send(message)
 
         LOGGER.debug(f"Unsupported message type {message}")
         return None

src/comet_llm/message_processing/online_senders/chain.py

@@ -15,12 +15,23 @@
 import io
 import json
 
-from comet_llm import app, convert, experiment_api, llm_result
+from comet_llm import app, convert, experiment_api, llm_result, url_helpers
+from comet_llm.experiment_api import comet_api_client
 
 from .. import messages
+from . import constants
 
 
 def send(message: messages.ChainMessage) -> llm_result.LLMResult:
+    client = comet_api_client.get(message.experiment_info_.api_key)
+
+    if client.backend_version >= constants.V2_BACKEND_VERSION:
+        return _send_v2(message, client)
+
+    return _send_v1(message)
+
+
+def _send_v1(message: messages.ChainMessage) -> llm_result.LLMResult:
     experiment_api_ = experiment_api.ExperimentAPI.create_new(
         api_key=message.experiment_info_.api_key,
         workspace=message.experiment_info_.workspace,
@@ -48,3 +59,24 @@ def send(message: messages.ChainMessage) -> llm_result.LLMResult:
     return llm_result.LLMResult(
         id=experiment_api_.id, project_url=experiment_api_.project_url
     )
+
+
+def _send_v2(
+    message: messages.ChainMessage, client: comet_api_client.CometAPIClient
+) -> llm_result.LLMResult:
+    metrics = {"chain_duration": message.duration}
+    parameters = convert.chain_metadata_to_flat_parameters(message.metadata)
+
+    response = client.log_chain(
+        experiment_key=message.id,
+        chain_asset=message.chain_data,
+        workspace=message.experiment_info_.workspace,
+        project=message.experiment_info_.project_name,
+        tags=message.tags,
+        metrics=metrics,
+        parameters=parameters,
+        others=message.others,
+    )
+    project_url: str = url_helpers.experiment_to_project_url(response["link"])
+
+    return llm_result.LLMResult(id=message.id, project_url=project_url)

src/comet_llm/message_processing/online_senders/constants.py

@@ -0,0 +1,15 @@
+# -*- 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.
+# *******************************************************
+
+V2_BACKEND_VERSION = "3.25.80"