trulens-providers-cortex 1.0.1__tar.gz

trulens_providers_cortex-1.0.1/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.1
+Name: trulens-providers-cortex
+Version: 1.0.1
+Summary: A TruLens extension package adding Snowflake Cortex support for LLM App evaluation.
+Home-page: https://trulens.org/
+License: MIT
+Author: Snowflake Inc.
+Author-email: ml-observability-wg-dl@snowflake.com
+Requires-Python: >=3.9,<3.12
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: snowflake-connector-python (>=3.11,<4.0)
+Requires-Dist: snowflake-snowpark-python (>=1.18,<2.0)
+Requires-Dist: trulens-core (>=1.0.0,<2.0.0)
+Requires-Dist: trulens-feedback (>=1.0.0,<2.0.0)
+Project-URL: Documentation, https://trulens.org/trulens/getting_started/
+Project-URL: Repository, https://github.com/truera/trulens
+Description-Content-Type: text/markdown
+
+# trulens-providers-cortex
+

trulens_providers_cortex-1.0.1/README.md

@@ -0,0 +1 @@
+# trulens-providers-cortex

trulens_providers_cortex-1.0.1/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = [
+  "poetry-core",
+]
+
+[tool.poetry]
+name = "trulens-providers-cortex"
+version = "1.0.1"
+description = "A TruLens extension package adding Snowflake Cortex support for LLM App evaluation."
+authors = [
+  "Snowflake Inc. <ml-observability-wg-dl@snowflake.com>",
+]
+license = "MIT"
+readme = "README.md"
+packages = [
+  { include = "trulens" },
+]
+homepage = "https://trulens.org/"
+documentation = "https://trulens.org/trulens/getting_started/"
+repository = "https://github.com/truera/trulens"
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Operating System :: OS Independent",
+  "Development Status :: 3 - Alpha",
+  "License :: OSI Approved :: MIT License",
+]
+
+[tool.poetry.dependencies]
+python = ">=3.9,<3.12"
+trulens-core = { version = "^1.0.0" }
+trulens-feedback = { version = "^1.0.0" }
+snowflake-connector-python = "^3.11"
+snowflake-snowpark-python = "^1.18"
+
+[tool.poetry.group.dev.dependencies]
+trulens-core = { path = "../../core" }
+trulens-feedback = { path = "../../feedback" }

trulens_providers_cortex-1.0.1/trulens/providers/cortex/__init__.py

@@ -0,0 +1,21 @@
+"""
+!!! note "Additional Dependency Required"
+
+    To use this module, you must have the `trulens-providers-cortex` package installed.
+
+    ```bash
+    pip install trulens-providers-cortex
+    ```
+"""
+
+from importlib.metadata import version
+
+from trulens.core.utils.imports import safe_importlib_package_name
+from trulens.providers.cortex.provider import Cortex
+
+__version__ = version(safe_importlib_package_name(__package__ or __name__))
+
+
+__all__ = [
+    "Cortex",
+]

trulens_providers_cortex-1.0.1/trulens/providers/cortex/config/cortex_model_costs.json

@@ -0,0 +1,16 @@
+{
+    "reka-core": 5.5,
+    "mistral-large": 5.1,
+    "llama3.1-405b": 5,
+    "llama3-70b": 1.21,
+    "llama3.1-70b": 1.21,
+    "snowflake-arctic": 0.84,
+    "jamba-instruct": 0.83,
+    "llama2-chat-70b": 0.45,
+    "reka-flash": 0.45,
+    "mixtral-8x7b": 0.22,
+    "llama3-8b": 0.19,
+    "llama3.1-8b": 0.19,
+    "mistral-7b": 0.12,
+    "gemma-7b": 0.12
+}

trulens_providers_cortex-1.0.1/trulens/providers/cortex/endpoint.py

@@ -0,0 +1,151 @@
+import inspect
+import json
+import logging
+import os
+import pprint
+from typing import Any, Callable, ClassVar, Optional
+
+from snowflake.connector.cursor import SnowflakeCursor
+from snowflake.snowpark import DataFrame
+from snowflake.snowpark import Session
+from trulens.core.feedback import Endpoint
+from trulens.core.feedback import EndpointCallback
+
+logger = logging.getLogger(__name__)
+
+pp = pprint.PrettyPrinter()
+
+
+class CortexCallback(EndpointCallback):
+    model_config: ClassVar[dict] = dict(arbitrary_types_allowed=True)
+    _model_costs: Optional[dict] = None
+    # TODO (Daniel): cost tracking for Cortex finetuned models is not yet implemented.
+
+    def _compute_credits_consumed(
+        self, cortex_model_name: str, n_tokens: int
+    ) -> float:
+        try:
+            if self._model_costs is None:
+                # the credit consumption table needs to be kept up-to-date with
+                # the latest cost information https://www.snowflake.com/legal-files/CreditConsumptionTable.pdf#page=9.
+
+                with open(
+                    os.path.join(
+                        os.path.dirname(os.path.realpath(__file__)),
+                        "config/cortex_model_costs.json",
+                    ),
+                    "r",
+                ) as file:
+                    self._model_costs = json.load(file)
+
+            if cortex_model_name in self._model_costs:
+                return (
+                    self._model_costs[cortex_model_name] * n_tokens / 1e6
+                )  # we maintain config per-1M-token cost
+            else:
+                raise ValueError(
+                    f"Model {cortex_model_name} not valid or not supported yet for cost estimation."
+                )
+        except Exception as e:
+            logger.error(
+                f"Error occurred while computing credits consumed for model {cortex_model_name}: {e}"
+            )
+            return 0.0
+
+    def handle_generation(self, response: dict) -> None:
+        """Get the usage information from Cortex LLM function response's usage field."""
+        usage = response["usage"]
+
+        # Increment number of requests.
+        super().handle_generation(response)
+
+        # Assume a response that had usage field was successful. Note at the time of writing 06/12/2024, the usage
+        # information from Cortex LLM functions is only available when called via snow SQL. It's not fully supported in
+        # Python API such as `from snowflake.cortex import Summarize, Complete, ExtractAnswer, Sentiment, Translate` yet.
+
+        self.cost.n_successful_requests += 1
+
+        for cost_field, cortex_field in [
+            ("n_tokens", "total_tokens"),
+            ("n_cortext_guardrails_tokens", "guardrails_tokens"),
+            ("n_prompt_tokens", "prompt_tokens"),
+            ("n_completion_tokens", "completion_tokens"),
+        ]:
+            setattr(
+                self.cost,
+                cost_field,
+                getattr(self.cost, cost_field, 0) + usage.get(cortex_field, 0),
+            )
+
+        # compute credits consumed in Snowflake account based on tokens processed
+        setattr(
+            self.cost,
+            "cost",
+            getattr(self.cost, "cost", 0)
+            + self._compute_credits_consumed(
+                response["model"], usage.get("total_tokens", 0)
+            ),
+        )
+
+        setattr(self.cost, "cost_currency", "Snowflake credits")
+
+
+class CortexEndpoint(Endpoint):
+    """Snowflake Cortex endpoint."""
+
+    def __init__(self, *args, **kwargs):
+        if hasattr(self, "name"):
+            # singleton already made
+            if len(kwargs) > 0:
+                logger.warning(
+                    "Ignoring additional kwargs for singleton endpoint %s: %s",
+                    self.name,
+                    pp.pformat(kwargs),
+                )
+                self.warning()
+            return
+
+        kwargs["name"] = "cortex"
+        kwargs["callback_class"] = CortexCallback
+
+        super().__init__(*args, **kwargs)
+
+        # Instrument various methods for usage/cost tracking.
+        self._instrument_class(Session, "sql")
+        self._instrument_class(SnowflakeCursor, "fetchall")
+
+    def __new__(cls, *args, **kwargs):
+        return super(Endpoint, cls).__new__(cls, name="cortex")
+
+    def handle_wrapped_call(
+        self,
+        func: Callable,
+        bindings: inspect.BoundArguments,
+        response: Any,
+        callback: Optional[EndpointCallback],
+    ) -> None:
+        counted_something = False
+
+        # response is a snowflake dataframe instance or a list if the response is from cursor.fetchall()
+        try:
+            if isinstance(response, DataFrame):
+                response: dict = json.loads(response.collect()[0][0])
+            elif isinstance(response, list):
+                response: dict = json.loads(response[0][0])
+        except Exception as e:
+            logger.error(f"Error occurred while parsing response: {e}")
+            raise e
+
+        if "usage" in response:
+            counted_something = True
+
+            self.global_callback.handle_generation(response=response)
+
+            if callback is not None:
+                callback.handle_generation(response=response)
+
+        if not counted_something:
+            logger.warning(
+                "Unrecognized Cortex response format. It did not have usage information:\n%s",
+                pp.pformat(response),
+            )

trulens_providers_cortex-1.0.1/trulens/providers/cortex/provider.py

@@ -0,0 +1,175 @@
+import json
+from typing import Any, ClassVar, Dict, Optional, Sequence
+
+from trulens.feedback import LLMProvider
+from trulens.providers.cortex.endpoint import CortexEndpoint
+
+# If this is set, the provider will use this connection. This is useful for server-side evaluations which are done in a stored procedure and must have a single connection throughout the life of the stored procedure.
+# TODO: This is a bit of a hack to pass the connection to the provider. Explore options on how to improve this.
+_SNOWFLAKE_STORED_PROCEDURE_CONNECTION: Any = None
+
+
+class Cortex(
+    LLMProvider
+):  # require `pip install snowflake-snowpark-python` and a active Snowflake account with proper privileges
+    # https://docs.snowflake.com/en/user-guide/snowflake-cortex/llm-functions#availability
+
+    DEFAULT_MODEL_ENGINE: ClassVar[str] = "snowflake-arctic"
+
+    model_engine: str
+    endpoint: CortexEndpoint
+    snowflake_conn: Any
+
+    """Snowflake's Cortex COMPLETE endpoint. Defaults to `snowflake-arctic`.
+       Reference: https://docs.snowflake.com/en/sql-reference/functions/complete-snowflake-cortex
+
+        !!! example
+        === Connecting with user/password:
+
+            ```python
+            connection_parameters = {
+                "account": <account>,
+                "user": <user>,
+                "password": <password>,
+                "role": <role>,
+                "database": <database>,
+                "schema": <schema>,
+                "warehouse": <warehouse>
+            }
+            provider = Cortex(snowflake.connector.connect(
+                **connection_parameters
+            ))
+            ```
+
+        === Connecting with private key:
+
+            ```python
+            connection_parameters = {
+                "account": <account>,
+                "user": <user>,
+                "private_key": <private_key>,
+                "role": <role>,
+                "database": <database>,
+                "schema": <schema>,
+                "warehouse": <warehouse>
+            }
+            provider = Cortex(snowflake.connector.connect(
+                **connection_parameters
+            ))
+
+        === Connecting with a private key file:
+
+            ```python
+            connection_parameters = {
+                "account": <account>,
+                "user": <user>,
+                "private_key_file": <private_key_file>,
+                "private_key_file_pwd": <private_key_file_pwd>,
+                "role": <role>,
+                "database": <database>,
+                "schema": <schema>,
+                "warehouse": <warehouse>
+            }
+            provider = Cortex(snowflake.connector.connect(
+                **connection_parameters
+            ))
+            ```
+
+    Args:
+        snowflake_conn (Any): Snowflake connection.
+
+        model_engine (str, optional): Model engine to use. Defaults to `snowflake-arctic`.
+
+    """
+
+    def __init__(
+        self,
+        snowflake_conn: Any,
+        model_engine: Optional[str] = None,
+        *args,
+        **kwargs: Dict,
+    ):
+        self_kwargs = dict(kwargs)
+
+        self_kwargs["model_engine"] = (
+            self.DEFAULT_MODEL_ENGINE if model_engine is None else model_engine
+        )
+
+        self_kwargs["endpoint"] = CortexEndpoint(*args, **kwargs)
+
+        # Create a Snowflake connector
+        self_kwargs["snowflake_conn"] = _SNOWFLAKE_STORED_PROCEDURE_CONNECTION
+        if _SNOWFLAKE_STORED_PROCEDURE_CONNECTION is None:
+            self_kwargs["snowflake_conn"] = snowflake_conn
+
+        super().__init__(**self_kwargs)
+
+    def _exec_snowsql_complete_command(
+        self,
+        model: str,
+        temperature: float,
+        messages: Optional[Sequence[Dict]] = None,
+    ):
+        # Ensure messages are formatted as a JSON array string
+        if messages is None:
+            messages = []
+
+        messages_json_str = json.dumps(messages)
+
+        options = {"temperature": temperature}
+
+        options_json_str = json.dumps(options)
+
+        completion_input_str = """
+            SELECT SNOWFLAKE.CORTEX.COMPLETE(
+                ?,
+                parse_json(?),
+                parse_json(?)
+            )
+        """
+        if (
+            hasattr(self.snowflake_conn, "_paramstyle")
+            and self.snowflake_conn._paramstyle == "pyformat"
+        ):
+            completion_input_str = completion_input_str.replace("?", "%s")
+
+        # Executing Snow SQL command requires an active snow session
+        cursor = self.snowflake_conn.cursor()
+        try:
+            cursor.execute(
+                completion_input_str,
+                (model, messages_json_str, options_json_str),
+            )
+            result = cursor.fetchall()
+        finally:
+            cursor.close()
+
+        return result
+
+    def _create_chat_completion(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[Sequence[Dict]] = None,
+        **kwargs,
+    ) -> str:
+        if "model" not in kwargs:
+            kwargs["model"] = self.model_engine
+        if "temperature" not in kwargs:
+            kwargs["temperature"] = 0.0
+
+        if messages is not None:
+            kwargs["messages"] = messages
+
+        elif prompt is not None:
+            kwargs["messages"] = [{"role": "system", "content": prompt}]
+        else:
+            raise ValueError("`prompt` or `messages` must be specified.")
+
+        res = self._exec_snowsql_complete_command(**kwargs)
+
+        if len(res) == 0 or len(res[0]) == 0:
+            raise ValueError("No completion returned from Snowflake Cortex.")
+
+        completion = json.loads(res[0][0])["choices"][0]["messages"]
+
+        return completion