sdk-seshat-python 0.4.1__tar.gz → 0.4.2__tar.gz

{sdk_seshat_python-0.4.1 → sdk_seshat_python-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sdk-seshat-python
-Version: 0.4.1
+Version: 0.4.2
 Summary: Seshat python SDK is a library to help create ML data pipelines.
 License: Commercial - see LICENSE.txt
 Author: SeshatLabs
@@ -13,29 +13,30 @@ Provides-Extra: flipside-support
 Provides-Extra: postgres-support
 Requires-Dist: backoff (>=2.2.1,<3.0.0)
 Requires-Dist: bokeh (>=3.6.0,<4.0.0)
-Requires-Dist: boto3 (>=1.39.4,<2.0.0)
+Requires-Dist: boto3 (>=1.35.68,<2.0.0)
 Requires-Dist: croniter (>=6.0.0,<7.0.0)
 Requires-Dist: cryptography (>=44.0.0,<45.0.0)
 Requires-Dist: dask[array,complete,dataframe,distributed] (>=2024.10.0,<2025.0.0)
 Requires-Dist: flipside (>=2.1.0,<3.0.0) ; extra == "flipside-support"
 Requires-Dist: langchain (>=0.3.23,<0.4.0)
 Requires-Dist: langchain-community (>=0.3.21,<0.4.0)
-Requires-Dist: langchain-openai (==0.3.28)
+Requires-Dist: langchain-openai (>=0.3.12,<0.4.0)
 Requires-Dist: loguru (>=0.7.3,<0.8.0)
 Requires-Dist: memory-profiler (>=0.61.0,<0.62.0)
 Requires-Dist: openai (>=1.73.0,<2.0.0)
-Requires-Dist: pandas (>=2.3.1,<3.0.0)
+Requires-Dist: pandas (>=2.2.1,<3.0.0)
 Requires-Dist: psycopg2-binary (>=2.9,<3.0) ; extra == "postgres-support"
 Requires-Dist: pyarmor (>=8.5.1,<9.0.0)
 Requires-Dist: pydantic (>=2.7.4,<3.0.0)
 Requires-Dist: pyspark (>=3.5.1,<4.0.0)
 Requires-Dist: python-logstash-async (>=4.0.2,<5.0.0)
-Requires-Dist: requests (>=2.32.0,<3.0.0)
-Requires-Dist: rich (>=14.0.0,<15.0.0)
+Requires-Dist: requests (==2.32.0)
+Requires-Dist: rich (>=13.9.4,<14.0.0)
 Requires-Dist: scikit-learn (>=1.4.1.post1,<2.0.0)
+Requires-Dist: setuptools (>=80.9.0,<81.0.0)
 Requires-Dist: sqlalchemy (>=2.0.29,<3.0.0)
 Requires-Dist: toml (>=0.10.2,<0.11.0)
-Requires-Dist: typer (>=0.16.0,<0.17.0)
+Requires-Dist: typer (>=0.12.3,<0.13.0)
 Description-Content-Type: text/markdown
 
 

{sdk_seshat_python-0.4.1 → sdk_seshat_python-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sdk-seshat-python"
-version = "0.4.1"
+version = "0.4.2"
 description = "Seshat python SDK is a library to help create ML data pipelines."
 authors =  ["SeshatLabs <info@seshatlabs.xyz>"]
 packages = [{ include = "seshat", from = "." }]
@@ -10,19 +10,19 @@ license = "Commercial - see LICENSE.txt"
 
 [tool.poetry.dependencies]
 python = "^3.11"
-pandas = "^2.3.1"
+pandas = "^2.2.1"
 scikit-learn = "^1.4.1.post1"
 pyspark = "^3.5.1"
 flipside = "^2.1.0"
 sqlalchemy = "^2.0.29"
 memory-profiler = "^0.61.0"
-typer = "^0.16.0"
+typer = "^0.12.3"
 dask = {extras = ["array", "complete", "dataframe", "distributed"], version = "^2024.10.0"}
 bokeh = "^3.6.0"
 toml = "^0.10.2"
-rich = "^14.0.0"
-boto3 = "^1.39.4"
-requests = "^2.32.0"
+rich = "^13.9.4"
+boto3 = "^1.35.68"
+requests = "2.32.0"
 backoff = "^2.2.1"
 cryptography = "^44.0.0"
 loguru = "^0.7.3"
@@ -30,11 +30,12 @@ openai = "^1.73.0"
 pydantic = "^2.7.4"
 langchain = "^0.3.23"
 langchain-community = "^0.3.21"
-langchain-openai = "0.3.28"
+langchain-openai = "^0.3.12"
 pyarmor = "^8.5.1"
+python-logstash-async = "^4.0.2"
 croniter = "^6.0.0"
 psycopg2-binary = { version = "^2.9", optional = true }
-python-logstash-async = "^4.0.2"
+setuptools = "^80.9.0"
 
 [tool.poetry.extras]
 flipside_support = ["flipside"]

{sdk_seshat_python-0.4.1 → sdk_seshat_python-0.4.2}/seshat/transformer/reducer/base.py

@@ -1,4 +1,5 @@
 import logging
+import math
 from typing import Any, Callable, Dict, List, TypeAlias
 
 import pandas as pd
@@ -18,7 +19,17 @@ from seshat.utils.clean_json import JSONCleaner
 InputType: TypeAlias = List[Dict[str, Any]]
 OutputType: TypeAlias = List[Dict[str, Any]]
 ProcessResponseFn: TypeAlias = Callable[[InputType, InputType], OutputType]
-ProcessBatchFn: TypeAlias = Callable[[list[dict[str, Any]]], dict]
+GetExtraCtxFn: TypeAlias = Callable[[List[Dict[str, Any]]], Dict[str, Any]]
+
+
+def math_nan_to_none(row):
+    d = row.asDict()
+    for key, value in d.items():
+        if value is None:
+            continue
+        if isinstance(value, float) and math.isnan(value):
+            d[key] = None
+    return Row(**d)
 
 
 class SFrameReducer(Transformer):
@@ -39,12 +50,14 @@ class LLMInsightExtractor(SFrameReducer):
 
     Parameters
     ----------
-    get_llm_client : BaseChatModel
+    get_llm_client : () ->  BaseChatModel
         The LLM client used to generate insights.
     template_prompt : str
         The template prompt to send to the LLM. Should include placeholders for data.
     id_column : str, optional
         The column name to use as an identifier when expanding results. Required if expand_on_id is True.
+    join_cols : list[str], optional
+        The columns to use for joining the extracted insights back to the original DataFrame. If not provided, defaults to [id_column].
     template_context : str, optional
         The system context to provide to the LLM. Defaults to a basic data scientist role.
     llm_input_columns : List[str], optional
@@ -52,7 +65,7 @@ class LLMInsightExtractor(SFrameReducer):
     process_llm_json_response_fn : ProcessResponseFn, optional
         Function to process the JSON response from the LLM.
     get_extra_context : ProcessBatchFn, optional
-        Function to process batches of data before sending to the LLM. Receives the current format_args dict and the list of batch_responses, and should return a dict to update format_args.
+        Function to process data before sending to the LLM. Receives the current data and should return a dict to update format_args.
     process_llm_response : Callable, optional
         Function to process the raw LLM response before JSON parsing.
     retry : int, default=3
@@ -71,6 +84,8 @@ class LLMInsightExtractor(SFrameReducer):
         Whether to expand results based on ID column. Requires id_column.
     inject_keys : dict[str, str], optional
         Additional keys to inject into the template prompt.
+    merge_result : bool, default=True
+        Whether to merge the extracted insights back to the original DataFrame.
 
     Raises
     ------
@@ -137,10 +152,11 @@ class LLMInsightExtractor(SFrameReducer):
         get_llm_client: Callable[[], "BaseChatModel"],
         template_prompt: str,
         id_column: str = None,
+        join_cols: list[str] = None,
         template_context: str = None,
         llm_input_columns: List[str] = None,
         process_llm_json_response_fn: ProcessResponseFn = None,
-        get_extra_context: ProcessBatchFn = None,
+        get_extra_context: GetExtraCtxFn = None,
         process_llm_response: Callable = None,
         retry: int = 3,
         llm_result_cleaner: Callable = JSONCleaner().clean,
@@ -151,6 +167,7 @@ class LLMInsightExtractor(SFrameReducer):
         groupby_inject_key: str = None,
         expand_on_id: bool = False,
         inject_keys: dict[str, str] = None,
+        merge_result: bool = True,
     ):
 
         super().__init__(group_keys)
@@ -172,9 +189,11 @@ class LLMInsightExtractor(SFrameReducer):
             Your task is to analyze and provide insights about the given dataset.
         """
         )
+        self.merge_result = merge_result
 
         self.group_by_columns = group_by_columns
         self.id_column = id_column
+        self.join_cols = join_cols or [self.id_column]
         self.expand_on_id = expand_on_id
         self.static_injected_data = inject_keys
         self.process_llm_response = process_llm_response
@@ -277,11 +296,19 @@ class LLMInsightExtractor(SFrameReducer):
             prompt_kwargs = (
                 self.get_extra_context(batch_result) if self.get_extra_context else {}
             )
-            batch_result += self.extract_insight(batch_data, prompt_kwargs, **kwargs)
+            res = self.perform_extract(batch_data, prompt_kwargs, **kwargs)
+            if res:
+                batch_result += res
         return batch_result
 
+    def extract_insight_one_shot(
+        self, data: List[Dict[str, Any]], **kwargs
+    ) -> List[Dict[str, Any]]:
+        prompt_kwargs = self.get_extra_context(data) if self.get_extra_context else {}
+        return self.perform_extract(data, prompt_kwargs, **kwargs)
+
     @track
-    def extract_insight(
+    def perform_extract(
         self, data: List[Dict[str, Any]], prompt_kwargs=None, **kwargs
     ) -> List[Dict[str, Any]]:
         """
@@ -329,11 +356,17 @@ class LLMInsightExtractor(SFrameReducer):
                     mask &= default[col] == val
                 group_df = default[mask]
 
-            selected = (
-                group_df[[*self.llm_input_columns]]
-                if self.llm_input_columns
-                else group_df
-            )
+            llm_input_columns = self.llm_input_columns or group_df.columns
+            if set(llm_input_columns) - set(group_df.columns):
+                continue
+
+            selected = group_df[[*llm_input_columns]]
+
+            if (
+                self.id_column in selected.columns
+                and selected[self.id_column].isnull().all()
+            ):
+                continue
 
             # If batch mode, create chunks otherwise use whole data
             if self.batch_mode:
@@ -343,11 +376,16 @@ class LLMInsightExtractor(SFrameReducer):
                 ]
             else:
                 data = selected.to_dict("records")
-            group_name = "-".join([group.get(c) for c in self.group_by_columns])
+            group_name = (
+                "-".join([group.get(c) for c in self.group_by_columns]) if group else ""
+            )
             inputs.append({"group_name": group_name, "data": data})
 
     def reduce_df(self, default: DataFrame, **kwargs) -> Dict[str, DataFrame]:
         # Find the groups if group_by_columns set
+        if default.empty:
+            return {"default": default}
+
         groups = (
             default[[*self.group_by_columns]].drop_duplicates().to_dict("records")
             if self.group_by_columns
@@ -359,28 +397,38 @@ class LLMInsightExtractor(SFrameReducer):
         self._find_extract_inputs(default, groups, inputs)
 
         extract_func = (
-            self.extract_insight_batch if self.batch_mode else self.extract_insight
+            self.extract_insight_batch
+            if self.batch_mode
+            else self.extract_insight_one_shot
         )
         results = []
         for d in inputs:
             results.extend(extract_func(**d))
         results = pd.DataFrame(results)
 
-        if self.expand_on_id:
+        if self.expand_on_id and not results.empty:
             redundant_cols = [
                 col
                 for col in results.columns
-                if col in default.columns and col != self.id_column
+                if col in default.columns
+                and col != self.id_column
+                and col not in set(self.join_cols)
             ]
             results = (
                 results.explode(self.id_column)
                 .set_index(self.id_column)
                 .drop(columns=redundant_cols, axis=1)
-                .join(default.set_index(self.id_column), how="right")
+            )
+        if not self.merge_result:
+            return {"default": results.drop_duplicates()}
+        if not results.empty:
+            default = (
+                pd.merge(default, results, on=self.join_cols, how="left")
                 .reset_index()
+                .drop_duplicates()
             )
 
-        return {"default": results}
+        return {"default": default}
 
     def reduce_spf(
         self, default: PySparkDataFrame, **kwargs
@@ -407,6 +455,9 @@ class LLMInsightExtractor(SFrameReducer):
         # To avoid calling process group again, cache rdd and
         # use count() to trigger running process_group and cache the result.
         rdd = default.rdd.mapPartitions(process_group)
+        # Because pandas result maybe contains math nan values and
+        # these values are not valid for spark
+        rdd = rdd.map(math_nan_to_none)
         rdd.cache()
         rdd.count()