mostlyai-mock 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

mostlyai/mock/__init__.py

@@ -15,4 +15,4 @@
 from mostlyai.mock.core import sample
 
 __all__ = ["sample"]
-__version__ = "0.1.3"  # Do not set this manually. Use poetry version [params].
+__version__ = "0.1.5"  # Do not set this manually. Use poetry version [params].

mostlyai/mock/core.py

@@ -18,7 +18,7 @@ import json
 from collections import deque
 from collections.abc import Generator
 from enum import Enum
-from typing import Any, Literal, Type
+from typing import Any, Literal
 
 import litellm
 import pandas as pd
@@ -28,7 +28,7 @@ from tqdm import tqdm
 litellm.suppress_debug_info = True
 
 SYSTEM_PROMPT = """
-You are a specialized mock data generator designed to create highly realistic, contextually appropriate data based on schema definitions. 
+You are a specialized mock data generator designed to create highly realistic, contextually appropriate data based on schema definitions.
 
 Your task is to:
 
@@ -58,7 +58,7 @@ class LLMConfig(BaseModel):
 
 
 class MockConfig(RootModel[dict[str, "TableConfig"]]):
-    root: dict[str, TableConfig] = Field(..., min_items=1)
+    root: dict[str, TableConfig] = Field(..., min_length=1)
 
     @field_validator("root")
     @classmethod
@@ -127,7 +127,7 @@ class MockConfig(RootModel[dict[str, "TableConfig"]]):
 
 class TableConfig(BaseModel):
     prompt: str = ""
-    columns: dict[str, ColumnConfig] = Field(..., min_items=1)
+    columns: dict[str, ColumnConfig] = Field(..., min_length=1)
     primary_key: str | None = None
     foreign_keys: list[ForeignKeyConfig] = Field(default_factory=list)
 
@@ -261,7 +261,7 @@ def _create_table_prompt(
 
     # add existing data to augment
     if existing_data is not None:
-        prompt += f"\n## Existing Data to Augment:\n\n"
+        prompt += "\n## Existing Data to Augment:\n\n"
         prompt += f"{existing_data.to_json(orient='records', date_format='iso', indent=2)}\n\n"
 
     # define foreign keys
@@ -314,11 +314,11 @@ def _create_table_prompt(
 
     if foreign_keys:
         prompt += (
-            f"The first Foreign Key column from Foreign Keys section may only contain values from Context Table Data. "
-            f"The following Foreign Key columns from Foreign Keys section (if exists) may only contain values from Non-Context Table Data sections. "
-            f"If either relevant Context Table Data or Non-Context Table Data is not present, this means that table has self-dependency. "
-            f"In this case, ensure that the foreign keys are consistent with primary keys of the table. "
-            f"Pay attention to prompt of the Foreign Key column to understand the relationship.\n\n"
+            "The first Foreign Key column from Foreign Keys section may only contain values from Context Table Data. "
+            "The following Foreign Key columns from Foreign Keys section (if exists) may only contain values from Non-Context Table Data sections. "
+            "If either relevant Context Table Data or Non-Context Table Data is not present, this means that table has self-dependency. "
+            "In this case, ensure that the foreign keys are consistent with primary keys of the table. "
+            "Pay attention to prompt of the Foreign Key column to understand the relationship.\n\n"
         )
 
     if existing_data is not None:
@@ -358,7 +358,7 @@ def _create_table_rows_generator(
     llm_config: LLMConfig,
 ) -> Generator[dict]:
     def create_table_response_format(columns: dict[str, ColumnConfig]) -> BaseModel:
-        def create_annotation(column_config: ColumnConfig) -> Type:
+        def create_annotation(column_config: ColumnConfig) -> type:
             if column_config.values or column_config.dtype is DType.CATEGORY:
                 return Literal[tuple(column_config.values)]
             return {
@@ -488,6 +488,12 @@ def _create_table_rows_generator(
                 table_name: df.sample(frac=1.0).head(non_context_size) for table_name, df in non_context_data.items()
             }
 
+        if context_batch is None:
+            # for root tables, scale down batch size in order to prevent excessive generations
+            remaining_rows = sample_size - yielded_sequences
+            if batch_size >= remaining_rows:
+                batch_size = remaining_rows + 2  # +2 because LLM may not always count the rows correctly
+
         llm_prompt = _create_table_prompt(
             name=name,
             prompt=prompt,
@@ -553,6 +559,36 @@ def _convert_table_rows_generator_to_df(
     return df
 
 
+def _harmonize_tables(tables: dict[str, dict], existing_data: dict[str, pd.DataFrame] | None) -> dict[str, dict]:
+    def _infer_dtype(series: pd.Series) -> DType:
+        if pd.api.types.is_integer_dtype(series):
+            return DType.INTEGER
+        elif pd.api.types.is_float_dtype(series):
+            return DType.FLOAT
+        elif pd.api.types.is_datetime64_dtype(series):
+            return DType.DATETIME
+        elif pd.api.types.is_bool_dtype(series):
+            return DType.BOOLEAN
+        else:
+            return DType.STRING
+
+    if existing_data is None:
+        return tables
+
+    tables = tables.copy()
+    for table_name, existing_table in existing_data.items():
+        table_config = tables.setdefault(table_name, {})
+        column_configs = table_config.setdefault("columns", {})
+        existing_column_configs = {
+            existing_column: {"dtype": _infer_dtype(existing_table[existing_column])}
+            for existing_column in existing_table.columns
+            if existing_column not in column_configs
+        }
+        column_configs = {**existing_column_configs, **column_configs}
+        table_config["columns"] = column_configs
+    return tables
+
+
 def _harmonize_sample_size(sample_size: int | dict[str, int], config: MockConfig) -> dict[str, int]:
     if isinstance(sample_size, int):
         return {table_name: sample_size for table_name in config.root}
@@ -626,24 +662,25 @@ def sample(
     return_type: Literal["auto", "dict"] = "auto",
 ) -> pd.DataFrame | dict[str, pd.DataFrame]:
     """
-    Generate mock data by prompting an LLM.
+    Generate mock data from scratch or enrich existing data by prompting an LLM.
 
     Args:
         tables (dict[str, dict]): The table specifications to generate mock data for. See examples for usage.
         sample_size (int | dict[str, int]): The number of rows to generate for each subject table.
             If a single integer is provided, the same number of rows will be generated for each subject table.
-            If a dictionary is provided, the number of rows to generate for each subject table can be specified
-            individually.
+            If a dictionary is provided, the number of rows to generate for each subject table can be specified individually.
             Default is 10. Ignored if existing_data is provided.
+            If a table has a foreign key, the sample size is determined by the corresponding foreign key prompt. If nothing specified, a few rows per parent record are generated.
         existing_data (dict[str, pd.DataFrame] | None): Existing data to augment. If provided, the sample_size argument is ignored.
             Default is None.
-        model (str): The LiteLLM chat completion model to be used. Requires support for structured output / JSON mode.
+        model (str): The LiteLLM chat completion model to be used. Model needs to support structured output / JSON mode.
             Examples include:
-            - `openai/gpt-4.1-nano` (default; fastest)
+            - `openai/gpt-4.1-nano` (default; fast, and smart)
             - `openai/gpt-4.1-mini` (slower, but smarter)
             - `openai/gpt-4.1` (slowest, but smartest)
             - `gemini/gemini-2.0-flash`
             - `gemini/gemini-2.5-flash-preview-04-17`
+            - 'groq/gemma2-9b-it`
             - `groq/llama-3.3-70b-versatile`
             - `anthropic/claude-3-7-sonnet-latest`
             See https://docs.litellm.ai/docs/providers/ for more options.
@@ -656,7 +693,7 @@ def sample(
         - pd.DataFrame: A single DataFrame containing the generated mock data, if only one table is provided.
         - dict[str, pd.DataFrame]: A dictionary containing the generated mock data for each table, if multiple tables are provided.
 
-    Example of single table (without PK):
+    Example of generating mock data for a single table (without PK):
     ```python
     from mostlyai import mock
 
@@ -679,7 +716,7 @@ def sample(
     df = mock.sample(tables=tables, sample_size=10, model="openai/gpt-4.1-nano")
     ```
 
-    Example of multiple tables (with PK/FK relationships):
+    Example of generating mock data for multiple tables (with PK/FK relationships):
     ```python
     from mostlyai import mock
 
@@ -746,7 +783,34 @@ def sample(
     df_items = data["items"]
     ```
 
-    Example of data augmentation:
+    Example of enriching a single dataframe:
+    ```python
+    from mostlyai import mock
+    import pandas as pd
+
+    tables = {
+        "patients": {
+            "prompt": "Patients of a hospital in Finland",
+            "columns": {
+                "full_name": {"prompt": "first name and last name of the patient", "dtype": "string"},
+                "date_of_birth": {"prompt": "date of birth", "dtype": "date"},
+                "place_of_birth": {"prompt": "place of birth", "dtype": "string"},
+            },
+        },
+    }
+    existing_df = pd.DataFrame({
+        "age": [25, 30, 35, 40],
+        "gender": ["male", "male", "female", "female"],
+    })
+    enriched_df = mock.sample(
+        tables=tables,
+        existing_data={"patients": existing_df},
+        model="openai/gpt-4.1-nano"
+    )
+    enriched_df
+    ```
+
+    Example of enriching / augmenting an existing dataset:
     ```python
     from mostlyai import mock
     import pandas as pd
@@ -803,7 +867,9 @@ def sample(
     ```
     """
 
+    tables: dict[str, TableConfig] = _harmonize_tables(tables, existing_data)
     config = MockConfig(tables)
+
     llm_config = LLMConfig(model=model, api_key=api_key, temperature=temperature, top_p=top_p)
 
     sample_size: dict[str, int] = _harmonize_sample_size(sample_size, config)

mostlyai/mock/mcp_server.py

@@ -1,3 +1,17 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 import tempfile
 

{mostlyai_mock-0.1.3.dist-info → mostlyai_mock-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-mock
-Version: 0.1.3
+Version: 0.1.5
 Summary: LLM-generated Mock Data
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-mock
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-mock
@@ -169,7 +169,7 @@ tables = {
 }
 data = mock.sample(
     tables=tables,
-    sample_size=2, 
+    sample_size=2,
     model="openai/gpt-4.1"
 )
 print(data["customers"])
@@ -250,9 +250,6 @@ tables = {
     "guests": {
         "prompt": "Guests of an Alpine ski hotel in Austria",
         "columns": {
-            "guest_id": {"prompt": "the unique id of the guest", "dtype": "integer"},
-            "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
-            "nationality": {"prompt": "2-letter code for the nationality", "dtype": "string"},
             "gender": {"dtype": "category", "values": ["male", "female"]},
             "age": {"prompt": "age in years; min: 18, max: 80; avg: 25", "dtype": "integer"},
             "room_number": {"prompt": "room number", "dtype": "integer"},

mostlyai_mock-0.1.5.dist-info/RECORD

@@ -0,0 +1,8 @@
+mostlyai/mock/__init__.py,sha256=-bfsVZJQ0OkN5b3IRP3F9aUCiA8Eq1-RmAqBmTg0O0g,714
+mostlyai/mock/core.py,sha256=V7KG7nOQPU95v6lRoSIfJuYivS0pNZ3rbiNC6SqDZSc,38075
+mostlyai/mock/mcp_server.py,sha256=kWMIjKCwnvYfjY8B2IdP4JNs8ik_8jA6ISCDqrG9utc,2137
+mostlyai_mock-0.1.5.dist-info/METADATA,sha256=LfugCsu7ANDZk2ozNFHDxgCqY42etJIdkXcfc-S-cUE,13887
+mostlyai_mock-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mostlyai_mock-0.1.5.dist-info/entry_points.txt,sha256=XDbppUIAaCWW0nresVep8zb71pkzZuFA16jCBHq8CU8,61
+mostlyai_mock-0.1.5.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mostlyai_mock-0.1.5.dist-info/RECORD,,

mostlyai_mock-0.1.3.dist-info/RECORD

@@ -1,8 +0,0 @@
-mostlyai/mock/__init__.py,sha256=38sp2aKJVtPa3koRxanlBS6fe_ccVQvIieILlKb-xuw,714
-mostlyai/mock/core.py,sha256=lO5OzuOz7bvjaLHpfiN-wyjFBPD0oSHSqEA4v8q436Y,35318
-mostlyai/mock/mcp_server.py,sha256=Vp0bWzE8wUyA6k4PHLa0TbkuI9s07E48xPrAUgf_5qU,1563
-mostlyai_mock-0.1.3.dist-info/METADATA,sha256=rkHeGDlNUM2cqSxWY_R47FWXsOLktpdl_COja8zYz28,14161
-mostlyai_mock-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mostlyai_mock-0.1.3.dist-info/entry_points.txt,sha256=XDbppUIAaCWW0nresVep8zb71pkzZuFA16jCBHq8CU8,61
-mostlyai_mock-0.1.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mostlyai_mock-0.1.3.dist-info/RECORD,,