From 987765bbdcdd87d2fa21dbdfd3c94e0a1fd9978d Mon Sep 17 00:00:00 2001
From: Jamie Broomall <88007022+jamie256@users.noreply.github.com>
Date: Fri, 10 Nov 2023 16:08:10 -0600
Subject: [PATCH] Add langkit metadata to schema for inclusion in whylogs
 profiles (#186)

---
 langkit/all_metrics.py         |  6 +++--
 langkit/count_regexes.py       |  2 +-
 langkit/injections.py          |  2 +-
 langkit/input_output.py        |  2 +-
 langkit/light_metrics.py       |  5 ++--
 langkit/llm_metrics.py         |  5 ++--
 langkit/metadata.py            | 41 ++++++++++++++++++++++++++++++
 langkit/nlp_scores.py          |  2 +-
 langkit/pattern_loader.py      |  2 +-
 langkit/regexes.py             |  2 +-
 langkit/sentiment.py           |  2 +-
 langkit/tests/test_metadata.py | 17 +++++++++++++
 langkit/textstat.py            |  7 ++++--
 langkit/themes.py              |  2 +-
 langkit/topics.py              |  2 +-
 langkit/toxicity.py            |  2 +-
 langkit/utils.py               |  2 +-
 poetry.lock                    | 46 +++++-----------------------------
 pyproject.toml                 |  2 +-
 19 files changed, 91 insertions(+), 60 deletions(-)
 create mode 100644 langkit/metadata.py
 create mode 100644 langkit/tests/test_metadata.py

diff --git a/langkit/all_metrics.py b/langkit/all_metrics.py
index cf3198ab..3d66235d 100644
--- a/langkit/all_metrics.py
+++ b/langkit/all_metrics.py
@@ -2,7 +2,9 @@
 from whylogs.experimental.core.udf_schema import udf_schema
 from whylogs.core.schema import DeclarativeSchema
 
-from . import LangKitConfig
+from langkit.metadata import attach_schema_metadata
+
+from langkit import LangKitConfig
 from langkit import injections
 from langkit import topics
 from langkit import regexes
@@ -22,5 +24,5 @@ def init(config: Optional[LangKitConfig] = None) -> DeclarativeSchema:
     themes.init(config=config)
     toxicity.init(config=config)
     input_output.init(config=config)
-    text_schema = udf_schema()
+    text_schema = attach_schema_metadata(udf_schema(), "all_metrics")
     return text_schema
diff --git a/langkit/count_regexes.py b/langkit/count_regexes.py
index c9b1ef53..3e9dc871 100644
--- a/langkit/count_regexes.py
+++ b/langkit/count_regexes.py
@@ -3,7 +3,7 @@
 
 from langkit.pattern_loader import PatternLoader
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 from whylogs.core.stubs import pd
 from typing import Dict, List, Optional, Set, Union
 
diff --git a/langkit/injections.py b/langkit/injections.py
index 26f55111..5168749a 100644
--- a/langkit/injections.py
+++ b/langkit/injections.py
@@ -2,7 +2,7 @@
 from typing import Dict, List, Optional, Union
 from whylogs.core.stubs import pd
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, lang_config, prompt_column
+from langkit import LangKitConfig, lang_config, prompt_column
 from sentence_transformers import SentenceTransformer
 import requests
 from io import BytesIO
diff --git a/langkit/input_output.py b/langkit/input_output.py
index 766798b7..434cb298 100644
--- a/langkit/input_output.py
+++ b/langkit/input_output.py
@@ -4,7 +4,7 @@
 
 from sentence_transformers import util
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 from langkit.transformer import Encoder
 
 _prompt = prompt_column
diff --git a/langkit/light_metrics.py b/langkit/light_metrics.py
index c0779e72..99f884ae 100644
--- a/langkit/light_metrics.py
+++ b/langkit/light_metrics.py
@@ -2,7 +2,8 @@
 from whylogs.experimental.core.udf_schema import udf_schema
 from whylogs.core.schema import DeclarativeSchema
 
-from . import LangKitConfig
+from langkit import LangKitConfig
+from langkit.metadata import attach_schema_metadata
 from langkit import regexes
 from langkit import textstat
 
@@ -11,5 +12,5 @@ def init(config: Optional[LangKitConfig] = None) -> DeclarativeSchema:
     regexes.init(config=config)
     textstat.init(config=config)
 
-    text_schema = udf_schema()
+    text_schema = attach_schema_metadata(udf_schema(), "light_metrics")
     return text_schema
diff --git a/langkit/llm_metrics.py b/langkit/llm_metrics.py
index ea44d5e1..1019d1b5 100644
--- a/langkit/llm_metrics.py
+++ b/langkit/llm_metrics.py
@@ -1,4 +1,5 @@
-from . import LangKitConfig
+from langkit.metadata import attach_schema_metadata
+from langkit import LangKitConfig
 from logging import getLogger
 from typing import Optional
 from whylogs.experimental.core.udf_schema import udf_schema
@@ -27,5 +28,5 @@ def init(config: Optional[LangKitConfig] = None) -> DeclarativeSchema:
     toxicity.init(config=config)
     input_output.init(config=config)
 
-    text_schema = udf_schema()
+    text_schema = attach_schema_metadata(udf_schema(), "llm_metrics")
     return text_schema
diff --git a/langkit/metadata.py b/langkit/metadata.py
new file mode 100644
index 00000000..f7bfae4d
--- /dev/null
+++ b/langkit/metadata.py
@@ -0,0 +1,41 @@
+from langkit import __version__
+from logging import getLogger
+
+from typing import Any, Dict, Optional
+
+
+_LANGKIT_VERSION_METADATA_KEY = "langkit.version"
+_LANGKIT_METRIC_COLLECTION_KEY = "langkit.metric_collection"
+diagnostic_logger = getLogger(__name__)
+
+
+def _check_for_metadata(schema: Any) -> Optional[Dict[str, str]]:
+    if schema is not None and hasattr(schema, "metadata"):
+        metadata = getattr(schema, "metadata")
+        if metadata is not None and isinstance(metadata, dict):
+            return metadata
+    return None
+
+
+def _add_langkit_version_metadata(
+    metadata: Dict[str, str], metric_collection_name: Optional[str]
+) -> Dict[str, str]:
+    if metadata is None:
+        diagnostic_logger.warning("metadata is None, LangKit won't update metadata")
+    else:
+        metadata[_LANGKIT_VERSION_METADATA_KEY] = __version__
+        if metric_collection_name:
+            metadata[_LANGKIT_METRIC_COLLECTION_KEY] = metric_collection_name
+    return metadata
+
+
+def attach_schema_metadata(schema: Any, metric_collection_name: Optional[str]) -> Any:
+    metadata = _check_for_metadata(schema)
+    if metadata is None:
+        diagnostic_logger.warning(
+            "schema does not contain metadata, LangKit won't update metadata"
+        )
+        return schema
+    _add_langkit_version_metadata(metadata, metric_collection_name)
+
+    return schema
diff --git a/langkit/nlp_scores.py b/langkit/nlp_scores.py
index 3bd6fa77..1103d536 100644
--- a/langkit/nlp_scores.py
+++ b/langkit/nlp_scores.py
@@ -2,7 +2,7 @@
 from typing import List, Optional, Set
 from whylogs.experimental.core.udf_schema import register_dataset_udf
 import evaluate
-from . import LangKitConfig, lang_config, response_column
+from langkit import LangKitConfig, lang_config, response_column
 from logging import getLogger
 
 
diff --git a/langkit/pattern_loader.py b/langkit/pattern_loader.py
index 6421c512..c28e6a29 100644
--- a/langkit/pattern_loader.py
+++ b/langkit/pattern_loader.py
@@ -4,7 +4,7 @@
 from logging import getLogger
 from typing import Optional
 
-from . import LangKitConfig, lang_config
+from langkit import LangKitConfig, lang_config
 
 
 diagnostic_logger = getLogger(__name__)
diff --git a/langkit/regexes.py b/langkit/regexes.py
index f68450c5..8f9f6e18 100644
--- a/langkit/regexes.py
+++ b/langkit/regexes.py
@@ -3,7 +3,7 @@
 
 from langkit.pattern_loader import PatternLoader
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 from whylogs.core.metrics.metrics import FrequentItemsMetric
 from whylogs.core.resolvers import MetricSpec
 from typing import Dict, List, Optional
diff --git a/langkit/sentiment.py b/langkit/sentiment.py
index cadf7ffe..a92894f2 100644
--- a/langkit/sentiment.py
+++ b/langkit/sentiment.py
@@ -2,7 +2,7 @@
 from typing import Optional
 
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 
 
 _prompt = prompt_column
diff --git a/langkit/tests/test_metadata.py b/langkit/tests/test_metadata.py
new file mode 100644
index 00000000..161bf773
--- /dev/null
+++ b/langkit/tests/test_metadata.py
@@ -0,0 +1,17 @@
+def test_metadata_langkit_version():
+    import whylogs as why
+    from langkit import __version__
+    from langkit.metadata import (
+        _LANGKIT_VERSION_METADATA_KEY,
+        _LANGKIT_METRIC_COLLECTION_KEY,
+    )
+    from langkit import light_metrics  # noqa
+
+    expected_metric_collection_name = "light_metrics"
+    text_schema = light_metrics.init()
+    results = why.log({"prompt": "hello", "response": "goodbye"}, schema=text_schema)
+    version = results.metadata[_LANGKIT_VERSION_METADATA_KEY]
+    metric_collection_name = results.metadata[_LANGKIT_METRIC_COLLECTION_KEY]
+    assert results.metadata
+    assert version == __version__
+    assert metric_collection_name == expected_metric_collection_name
diff --git a/langkit/textstat.py b/langkit/textstat.py
index 9ca0d678..deca402a 100644
--- a/langkit/textstat.py
+++ b/langkit/textstat.py
@@ -1,9 +1,8 @@
 from logging import getLogger
 from typing import Callable, Dict, List, Optional, Tuple, Union
-import textstat
 from whylogs.core.stubs import pd
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, prompt_column, response_column
+from langkit import LangKitConfig, prompt_column, response_column
 
 
 diagnostic_logger = getLogger(__name__)
@@ -43,6 +42,8 @@
 def wrapper(
     stat_name: str, column: str
 ) -> Callable[[Union[pd.DataFrame, Dict[str, List]]], Union[pd.Series, List]]:
+    import textstat
+
     stat = textstat.textstat.__getattribute__(stat_name)
 
     def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List]:
@@ -54,6 +55,8 @@ def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List
 def aggregate_wrapper(
     column: str,
 ) -> Callable[[Union[pd.DataFrame, Dict[str, List]]], Union[pd.Series, List]]:
+    import textstat
+
     stat = textstat.textstat.text_standard
 
     def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List]:
diff --git a/langkit/themes.py b/langkit/themes.py
index 484b24db..4639081f 100644
--- a/langkit/themes.py
+++ b/langkit/themes.py
@@ -9,7 +9,7 @@
 
 from langkit.transformer import Encoder
 
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 
 diagnostic_logger = getLogger(__name__)
 
diff --git a/langkit/topics.py b/langkit/topics.py
index 598d8992..1c259ccc 100644
--- a/langkit/topics.py
+++ b/langkit/topics.py
@@ -4,7 +4,7 @@
 from transformers import (
     pipeline,
 )
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 
 
 _topics: List[str] = lang_config.topics
diff --git a/langkit/toxicity.py b/langkit/toxicity.py
index 9dd31623..f576b1fd 100644
--- a/langkit/toxicity.py
+++ b/langkit/toxicity.py
@@ -2,7 +2,7 @@
 from typing import Optional
 
 from whylogs.experimental.core.udf_schema import register_dataset_udf
-from . import LangKitConfig, lang_config, prompt_column, response_column
+from langkit import LangKitConfig, lang_config, prompt_column, response_column
 
 
 _prompt = prompt_column
diff --git a/langkit/utils.py b/langkit/utils.py
index 6b5c7c9c..4d0b5f3c 100644
--- a/langkit/utils.py
+++ b/langkit/utils.py
@@ -1,5 +1,5 @@
 import os
-from . import lang_config
+from langkit import lang_config
 
 
 def _get_data_home() -> str:
diff --git a/poetry.lock b/poetry.lock
index ae94f9f8..73796e81 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1521,40 +1521,6 @@ files = [
     {file = "numpy-1.24.4.tar.gz", hash = "sha256:80f5e3a4e498641401868df4208b74581206afbee7cf7b8329daae82676d9463"},
 ]
 
-[[package]]
-name = "numpy"
-version = "1.25.2"
-description = "Fundamental package for array computing in Python"
-optional = false
-python-versions = ">=3.9"
-files = [
-    {file = "numpy-1.25.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:db3ccc4e37a6873045580d413fe79b68e47a681af8db2e046f1dacfa11f86eb3"},
-    {file = "numpy-1.25.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:90319e4f002795ccfc9050110bbbaa16c944b1c37c0baeea43c5fb881693ae1f"},
-    {file = "numpy-1.25.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:dfe4a913e29b418d096e696ddd422d8a5d13ffba4ea91f9f60440a3b759b0187"},
-    {file = "numpy-1.25.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f08f2e037bba04e707eebf4bc934f1972a315c883a9e0ebfa8a7756eabf9e357"},
-    {file = "numpy-1.25.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:bec1e7213c7cb00d67093247f8c4db156fd03075f49876957dca4711306d39c9"},
-    {file = "numpy-1.25.2-cp310-cp310-win32.whl", hash = "sha256:7dc869c0c75988e1c693d0e2d5b26034644399dd929bc049db55395b1379e044"},
-    {file = "numpy-1.25.2-cp310-cp310-win_amd64.whl", hash = "sha256:834b386f2b8210dca38c71a6e0f4fd6922f7d3fcff935dbe3a570945acb1b545"},
-    {file = "numpy-1.25.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:c5462d19336db4560041517dbb7759c21d181a67cb01b36ca109b2ae37d32418"},
-    {file = "numpy-1.25.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c5652ea24d33585ea39eb6a6a15dac87a1206a692719ff45d53c5282e66d4a8f"},
-    {file = "numpy-1.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0d60fbae8e0019865fc4784745814cff1c421df5afee233db6d88ab4f14655a2"},
-    {file = "numpy-1.25.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:60e7f0f7f6d0eee8364b9a6304c2845b9c491ac706048c7e8cf47b83123b8dbf"},
-    {file = "numpy-1.25.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:bb33d5a1cf360304754913a350edda36d5b8c5331a8237268c48f91253c3a364"},
-    {file = "numpy-1.25.2-cp311-cp311-win32.whl", hash = "sha256:5883c06bb92f2e6c8181df7b39971a5fb436288db58b5a1c3967702d4278691d"},
-    {file = "numpy-1.25.2-cp311-cp311-win_amd64.whl", hash = "sha256:5c97325a0ba6f9d041feb9390924614b60b99209a71a69c876f71052521d42a4"},
-    {file = "numpy-1.25.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:b79e513d7aac42ae918db3ad1341a015488530d0bb2a6abcbdd10a3a829ccfd3"},
-    {file = "numpy-1.25.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:eb942bfb6f84df5ce05dbf4b46673ffed0d3da59f13635ea9b926af3deb76926"},
-    {file = "numpy-1.25.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e0746410e73384e70d286f93abf2520035250aad8c5714240b0492a7302fdca"},
-    {file = "numpy-1.25.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d7806500e4f5bdd04095e849265e55de20d8cc4b661b038957354327f6d9b295"},
-    {file = "numpy-1.25.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:8b77775f4b7df768967a7c8b3567e309f617dd5e99aeb886fa14dc1a0791141f"},
-    {file = "numpy-1.25.2-cp39-cp39-win32.whl", hash = "sha256:2792d23d62ec51e50ce4d4b7d73de8f67a2fd3ea710dcbc8563a51a03fb07b01"},
-    {file = "numpy-1.25.2-cp39-cp39-win_amd64.whl", hash = "sha256:76b4115d42a7dfc5d485d358728cdd8719be33cc5ec6ec08632a5d6fca2ed380"},
-    {file = "numpy-1.25.2-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:1a1329e26f46230bf77b02cc19e900db9b52f398d6722ca853349a782d4cff55"},
-    {file = "numpy-1.25.2-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c3abc71e8b6edba80a01a52e66d83c5d14433cbcd26a40c329ec7ed09f37901"},
-    {file = "numpy-1.25.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:1b9735c27cea5d995496f46a8b1cd7b408b3f34b6d50459d9ac8fe3a20cc17bf"},
-    {file = "numpy-1.25.2.tar.gz", hash = "sha256:fd608e19c8d7c55021dffd43bfe5492fab8cc105cc8986f813f8c3c048b38760"},
-]
-
 [[package]]
 name = "nvidia-cublas-cu11"
 version = "11.10.3.66"
@@ -1783,8 +1749,8 @@ files = [
 [package.dependencies]
 numpy = [
     {version = ">=1.20.3", markers = "python_version < \"3.10\""},
-    {version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
     {version = ">=1.23.2", markers = "python_version >= \"3.11\""},
+    {version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
 ]
 python-dateutil = ">=2.8.2"
 pytz = ">=2020.1"
@@ -3442,13 +3408,13 @@ urllib3 = ">=1.25.3"
 
 [[package]]
 name = "whylogs"
-version = "1.3.10"
+version = "1.3.13"
 description = "Profile and monitor your ML data pipeline end-to-end"
 optional = false
 python-versions = ">=3.7.1,<4"
 files = [
-    {file = "whylogs-1.3.10-py3-none-any.whl", hash = "sha256:9aeafad68f5b6fb99006c8910f3863681881e99a2853d2f5a67b1bf27996f776"},
-    {file = "whylogs-1.3.10.tar.gz", hash = "sha256:868f1c2cb6131d7e687e3f2285cdcb019a3d5d7b2db6f762cd44ef3981d8a011"},
+    {file = "whylogs-1.3.13-py3-none-any.whl", hash = "sha256:b2b25519624d39f82094657507c7518a2f8ac8661b153949fb009122ef17d610"},
+    {file = "whylogs-1.3.13.tar.gz", hash = "sha256:ea366c3c2dc00579f7565314fb4b26ac6bad62b9a66defe90427aa6c13fa9fbe"},
 ]
 
 [package.dependencies]
@@ -3766,9 +3732,9 @@ docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.link
 testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy (>=0.9.1)", "pytest-ruff"]
 
 [extras]
-all = ["datasets", "evaluate", "faiss-cpu", "ipywidgets", "nltk", "numpy", "numpy", "openai", "sentence-transformers", "torch"]
+all = ["datasets", "evaluate", "faiss-cpu", "ipywidgets", "nltk", "numpy", "openai", "sentence-transformers", "torch"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8,<4"
-content-hash = "44e72db5b9563952e4b159237480c831bd0ab87ca36fa94a1640835ed421c374"
+content-hash = "2a5a6206e17480c743b5d15a4f8edb451280cf2d3ee1efe1869f78f401d1b539"
diff --git a/pyproject.toml b/pyproject.toml
index 9a169a57..f65a88f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ readme = "DESCRIPTION.md"
 python = ">=3.8,<4"
 textstat = "^0.7.3"
 pandas = "*"
-whylogs = "^1.3.9"
+whylogs = "^1.3.13"
 
 
 # optional dependencies