trulens-providers-cortex 1.0.1a1__tar.gz

trulens_providers_cortex-1.0.1a1/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.1
+Name: trulens-providers-cortex
+Version: 1.0.1a1
+Summary: Library to systematically track and evaluate LLM based applications.
+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.1a1/README.md

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

trulens_providers_cortex-1.0.1a1/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.1a1"
+description = "Library to systematically track and evaluate LLM based applications."
+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", allow-prereleases = true }
+trulens-feedback = { version = "^1.0.0", allow-prereleases = true }
+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.1a1/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.1a1/trulens/providers/cortex/endpoint.py

@@ -0,0 +1,95 @@
+import inspect
+import json
+import logging
+import pprint
+from typing import Any, Callable, ClassVar, Optional
+
+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)
+
+    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_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),
+            )
+
+
+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)
+        self._instrument_class(Session, "sql")
+
+    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
+
+        if isinstance(
+            response, DataFrame
+        ):  # response is a snowflake dataframe instance
+            response: dict = json.loads(response.collect()[0][0])
+
+            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.1a1/trulens/providers/cortex/provider.py

@@ -0,0 +1,105 @@
+import json
+import os
+from typing import ClassVar, Dict, Optional, Sequence
+
+import snowflake
+import snowflake.connector
+from snowflake.connector import SnowflakeConnection
+from trulens.feedback import LLMProvider
+from trulens.providers.cortex.endpoint import CortexEndpoint
+
+
+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
+    """Snowflake's Cortex COMPLETE endpoint. Defaults to `snowflake-arctic`.
+       Reference: https://docs.snowflake.com/en/sql-reference/functions/complete-snowflake-cortex
+    """
+
+    endpoint: CortexEndpoint
+    snowflake_conn: SnowflakeConnection
+
+    def __init__(
+        self, 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.connector.connect(
+            account=os.environ["SNOWFLAKE_ACCOUNT"],
+            user=os.environ["SNOWFLAKE_USER"],
+            password=os.environ["SNOWFLAKE_USER_PASSWORD"],
+            database=os.environ["SNOWFLAKE_DATABASE"],
+            schema=os.environ["SNOWFLAKE_SCHEMA"],
+            warehouse=os.environ["SNOWFLAKE_WAREHOUSE"],
+        )
+        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(
+                %s,
+                parse_json(%s),
+                parse_json(%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)
+
+        completion = json.loads(res[0][0])["choices"][0]["messages"]
+
+        return completion