Add langkit metadata to schema for inclusion in whylogs profiles (#186)

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

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
 

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

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

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

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

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

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
 
 

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__)

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

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

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

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]:

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__)
 

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

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

langkit/utils.py

@@ -1,5 +1,5 @@
 import os
-from . import lang_config
+from langkit import lang_config
 
 
 def _get_data_home() -> str:

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