mostlyai-mock 0.0.1__tar.gz

mostlyai_mock-0.0.1/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+.idea/
+.vscode/
+.ipynb_checkpoints/
+.DS_Store
+dist/
+examples/ws-*/
+LICENSE_HEADER
+/site/

mostlyai_mock-0.0.1/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

mostlyai_mock-0.0.1/LICENSE_HEADER

@@ -0,0 +1,13 @@
+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.

mostlyai_mock-0.0.1/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.4
+Name: mostlyai-mock
+Version: 0.0.1
+Summary: Synthetic Mock Data
+Project-URL: homepage, https://github.com/mostly-ai/mostlyai-mock
+Project-URL: repository, https://github.com/mostly-ai/mostlyai-mock
+Project-URL: documentation, https://mostly-ai.github.io/mostlyai-mock/
+Author-email: MOSTLY AI <dev@mostly.ai>
+License-Expression: Apache-2.0
+License-File: LICENSE
+License-File: LICENSE_HEADER
+Requires-Python: >=3.10
+Requires-Dist: litellm>=1.67.0
+Requires-Dist: numpy>=1.26.3
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: pyarrow>=14.0.0
+Requires-Dist: pydantic<3.0.0,>=2.0.0
+Description-Content-Type: text/markdown
+
+# Synthetic Mock Data 🔮
+
+[![Documentation](https://img.shields.io/badge/docs-latest-green)](https://mostly-ai.github.io/mostlyai-mock/) [![stats](https://pepy.tech/badge/mostlyai-mock)](https://pypi.org/project/mostlyai-mock/) ![license](https://img.shields.io/github/license/mostly-ai/mostlyai-mock) ![GitHub Release](https://img.shields.io/github/v/release/mostly-ai/mostlyai-mock) ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mostlyai-mock)
+
+Create data out of nothing. Prompt LLMs for Tabular Data.
+
+## Installation
+
+The latest release of `mostlyai-mock` can be installed via pip:
+
+```bash
+pip install -U mostlyai-mock
+```
+
+Note: An API key to a LLM endpoint, with structured response, is required. It is recommended to set such a key as an environment variable (e.g. `OPENAI_API_KEY`, `GEMINI_API_KEY`, etc.). Alternatively, the key needs to be passed to every call to the library iteself via the parameter `api_key`.
+
+## Quick Start
+
+### Single Table
+
+```python
+from mostlyai import mock
+
+tables = {
+    "guests": {
+        "description": "Guests of an Alpine ski hotel in Austria",
+        "columns": {
+            "nationality": {"prompt": "2-letter code for the nationality", "dtype": "string"},
+            "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+            "gender": {"prompt": "gender of the guest; male or female", "dtype": "string"},
+            "age": {"prompt": "age in years; min: 18, max: 80; avg: 25", "dtype": "integer"},
+            "date_of_birth": {"prompt": "date of birth", "dtype": "date"},
+            "checkin_time": {"prompt": "the check in timestamp of the guest; may 2025", "dtype": "datetime"},
+            "is_vip": {"prompt": "is the guest a VIP", "dtype": "boolean"},
+            "price_per_night": {"prompt": "price paid per night, in EUR", "dtype": "float"},
+        },
+    }
+}
+df = mock.sample(tables=tables, sample_size=10, model="openai/gpt-4.1-nano")
+print(df)
+```
+
+### Multiple Tables
+
+```python
+from mostlyai import mock
+
+tables = {
+    "guests": {
+        "description": "Guests of an Alpine ski hotel in Austria",
+        "columns": {
+            "id": {"prompt": "the unique id of the guest", "dtype": "integer"},
+            "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+        },
+        "primary_key": "id",
+    },
+    "purchases": {
+        "description": "Purchases of a Guest during their stay",
+        "columns": {
+            "guest_id": {"prompt": "the guest id for that purchase", "dtype": "integer"},
+            "purchase_id": {"prompt": "the unique id of the purchase", "dtype": "string"},
+            "text": {"prompt": "purchase text description", "dtype": "string"},
+            "amount": {"prompt": "purchase amount in EUR", "dtype": "float"},
+        },
+        "foreign_keys": [
+            {
+                "column": "guest_id",
+                "referenced_table": "guests",
+                "description": "each guest has anywhere between 1 and 10 purchases",
+            }
+        ],
+    },
+}
+data = mock.sample(tables=tables, sample_size=5, model="openai/gpt-4.1-nano")
+df_guests = data["guests"]
+df_purchases = data["purchases"]
+print(df_guests)
+print(df_purchases)
+```

mostlyai_mock-0.0.1/README.md

@@ -0,0 +1,79 @@
+# Synthetic Mock Data 🔮
+
+[![Documentation](https://img.shields.io/badge/docs-latest-green)](https://mostly-ai.github.io/mostlyai-mock/) [![stats](https://pepy.tech/badge/mostlyai-mock)](https://pypi.org/project/mostlyai-mock/) ![license](https://img.shields.io/github/license/mostly-ai/mostlyai-mock) ![GitHub Release](https://img.shields.io/github/v/release/mostly-ai/mostlyai-mock) ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/mostlyai-mock)
+
+Create data out of nothing. Prompt LLMs for Tabular Data.
+
+## Installation
+
+The latest release of `mostlyai-mock` can be installed via pip:
+
+```bash
+pip install -U mostlyai-mock
+```
+
+Note: An API key to a LLM endpoint, with structured response, is required. It is recommended to set such a key as an environment variable (e.g. `OPENAI_API_KEY`, `GEMINI_API_KEY`, etc.). Alternatively, the key needs to be passed to every call to the library iteself via the parameter `api_key`.
+
+## Quick Start
+
+### Single Table
+
+```python
+from mostlyai import mock
+
+tables = {
+    "guests": {
+        "description": "Guests of an Alpine ski hotel in Austria",
+        "columns": {
+            "nationality": {"prompt": "2-letter code for the nationality", "dtype": "string"},
+            "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+            "gender": {"prompt": "gender of the guest; male or female", "dtype": "string"},
+            "age": {"prompt": "age in years; min: 18, max: 80; avg: 25", "dtype": "integer"},
+            "date_of_birth": {"prompt": "date of birth", "dtype": "date"},
+            "checkin_time": {"prompt": "the check in timestamp of the guest; may 2025", "dtype": "datetime"},
+            "is_vip": {"prompt": "is the guest a VIP", "dtype": "boolean"},
+            "price_per_night": {"prompt": "price paid per night, in EUR", "dtype": "float"},
+        },
+    }
+}
+df = mock.sample(tables=tables, sample_size=10, model="openai/gpt-4.1-nano")
+print(df)
+```
+
+### Multiple Tables
+
+```python
+from mostlyai import mock
+
+tables = {
+    "guests": {
+        "description": "Guests of an Alpine ski hotel in Austria",
+        "columns": {
+            "id": {"prompt": "the unique id of the guest", "dtype": "integer"},
+            "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+        },
+        "primary_key": "id",
+    },
+    "purchases": {
+        "description": "Purchases of a Guest during their stay",
+        "columns": {
+            "guest_id": {"prompt": "the guest id for that purchase", "dtype": "integer"},
+            "purchase_id": {"prompt": "the unique id of the purchase", "dtype": "string"},
+            "text": {"prompt": "purchase text description", "dtype": "string"},
+            "amount": {"prompt": "purchase amount in EUR", "dtype": "float"},
+        },
+        "foreign_keys": [
+            {
+                "column": "guest_id",
+                "referenced_table": "guests",
+                "description": "each guest has anywhere between 1 and 10 purchases",
+            }
+        ],
+    },
+}
+data = mock.sample(tables=tables, sample_size=5, model="openai/gpt-4.1-nano")
+df_guests = data["guests"]
+df_purchases = data["purchases"]
+print(df_guests)
+print(df_purchases)
+```

mostlyai_mock-0.0.1/mostlyai/mock/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from mostlyai.mock.core import sample
+
+__all__ = ["sample"]
+__version__ = "0.0.1"  # Do not set this manually. Use poetry version [params].

mostlyai_mock-0.0.1/mostlyai/mock/core.py

@@ -0,0 +1,512 @@
+# 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.
+
+from __future__ import annotations
+
+import json
+from collections import deque
+from collections.abc import Generator
+from enum import Enum
+
+import litellm
+import pandas as pd
+from pydantic import BaseModel, Field, RootModel, create_model, field_validator
+from tqdm import tqdm
+
+SYSTEM_PROMPT = f"""
+You are a specialized synthetic data generator designed to create
+highly realistic, contextually appropriate data based on schema definitions. Your task is to:
+
+1. Generate data that strictly adheres to the provided schema constraints (data types, ranges, formats)
+2. Ensure logical consistency across related tables and foreign key relationships
+3. Create contextually appropriate values that reflect real-world patterns and distributions
+4. Produce diverse, non-repetitive data that avoids obvious patterns
+5. Respect uniqueness constraints and other data integrity rules
+6. Return well-formatted JSON output that can be directly parsed.
+7. Don't use markdown formatting.
+
+For numeric fields, generate realistic distributions rather than random values. For text fields, create contextually \
+appropriate content. For dates and timestamps, ensure logical chronology. Always maintain referential integrity \
+across tables.
+"""
+
+
+class LLMConfig(BaseModel):
+    model: str
+    api_key: str | None = None
+
+
+class MockConfig(RootModel[dict[str, "TableConfig"]]):
+    root: dict[str, TableConfig] = Field(..., min_items=1)
+
+    @field_validator("root")
+    @classmethod
+    def validate_consistency_of_relationships(cls, tables: dict[str, TableConfig]) -> dict[str, TableConfig]:
+        for table_name, table_config in tables.items():
+            if not table_config.foreign_keys:
+                continue
+
+            for fk in table_config.foreign_keys:
+                if fk.referenced_table not in tables:
+                    raise ValueError(
+                        f"Foreign key violation in table '{table_name}': "
+                        f"Referenced table '{fk.referenced_table}' does not exist"
+                    )
+
+                referenced_config = tables[fk.referenced_table]
+                if not referenced_config.primary_key:
+                    raise ValueError(
+                        f"Foreign key violation in table '{table_name}': "
+                        f"Referenced table '{fk.referenced_table}' has no primary key defined"
+                    )
+
+                if fk.column not in table_config.columns:
+                    raise ValueError(
+                        f"Foreign key violation in table '{table_name}': "
+                        f"Column '{fk.column}' does not exist in the schema"
+                    )
+
+                fk_field = table_config.columns[fk.column]
+                pk_field = referenced_config.columns[referenced_config.primary_key]
+                if fk_field.dtype != pk_field.dtype:
+                    raise ValueError(
+                        f"Foreign key violation in table '{table_name}': "
+                        f"Column '{fk.column}' type '{fk_field.dtype}' does not match "
+                        f"referenced primary key '{referenced_config.primary_key}' type '{pk_field.dtype}'"
+                    )
+
+        return tables
+
+
+class TableConfig(BaseModel):
+    description: str = ""
+    columns: dict[str, ColumnConfig] = Field(..., min_items=1)
+    primary_key: str | None = None
+    foreign_keys: list[ForeignKeyConfig] = Field(default_factory=list, min_length=0, max_length=1)
+
+
+class ColumnConfig(BaseModel):
+    prompt: str
+    dtype: DType
+
+
+class DType(str, Enum):
+    INTEGER = "integer"
+    FLOAT = "float"
+    STRING = "string"
+    BOOLEAN = "boolean"
+    DATE = "date"
+    DATETIME = "datetime"
+
+
+class ForeignKeyConfig(BaseModel):
+    column: str
+    referenced_table: str
+    description: str | None = None
+
+
+def _sample_table(
+    *,
+    table_name: str,
+    table_config: TableConfig,
+    primary_keys: dict[str, str] | None,
+    sample_size: int | None,
+    context_data: pd.DataFrame | None,
+    temperature: float,
+    top_p: float,
+    batch_size: int,
+    previous_rows_size: int,
+    llm_config: LLMConfig,
+) -> pd.DataFrame:
+    assert (sample_size is None) != (context_data is None), (
+        "Exactly one of sample_size or context_data must be provided"
+    )
+    if sample_size is None:
+        sample_size = len(context_data)
+    table_rows_generator = _create_table_rows_generator(
+        table_name=table_name,
+        table_config=table_config,
+        primary_keys=primary_keys,
+        sample_size=sample_size,
+        context_data=context_data,
+        temperature=temperature,
+        top_p=top_p,
+        batch_size=batch_size,
+        previous_rows_size=previous_rows_size,
+        llm_config=llm_config,
+    )
+    table_rows_generator = tqdm(table_rows_generator, desc=f"Generating rows for table `{table_name}`".ljust(45))
+    table_df = _convert_table_rows_generator_to_df(table_rows_generator=table_rows_generator, table_config=table_config)
+    return table_df
+
+
+def _create_table_prompt(
+    *,
+    table_name: str,
+    table_description: str,
+    columns: dict[str, ColumnConfig],
+    primary_keys: dict[str, str] | None,
+    batch_size: int | None,
+    foreign_keys: list[ForeignKeyConfig] | None,
+    context_data: pd.DataFrame | None,
+    previous_rows: list[dict],
+) -> str:
+    if batch_size is not None:
+        assert foreign_keys is None
+        assert context_data is None
+    else:
+        assert foreign_keys is not None
+        assert context_data is not None
+        assert primary_keys is not None
+
+    # add description
+    prompt = f"# {table_description}\n\n"
+
+    # define table
+    prompt += f"## Table: {table_name}\n\n"
+
+    # add columns specifications
+    prompt += "## Columns Specifications:\n\n"
+    prompt += f"{json.dumps({name: config.model_dump() for name, config in columns.items()}, indent=2)}\n\n"
+
+    # define foreign keys
+    if foreign_keys is not None:
+        prompt += "## Foreign Keys:\n\n"
+        prompt += f"{json.dumps([fk.model_dump() for fk in foreign_keys], indent=2)}\n\n"
+
+    # add previous rows as context to help the LLM generate consistent data
+    if previous_rows:
+        prompt += f"\n## Previous {len(previous_rows)} Rows:\n\n"
+        prompt += json.dumps(previous_rows, indent=2)
+
+    # add context table name, primary key and data
+    if context_data is not None:
+        fk = foreign_keys[0]
+        prompt += f"## Context Table: `{fk.referenced_table}`\n\n"
+
+        prompt += f"## Context Table Primary Key: `{primary_keys[fk.referenced_table]}`\n\n"
+
+        prompt += f"## Context Table Data:\n\n"
+        prompt += f"{context_data.to_json(orient='records', indent=2)}\n\n"
+
+    # add instructions
+    prompt += "\n## Instructions:\n\n"
+    if batch_size is not None:
+        prompt += f"Generate {batch_size} rows for the `{table_name}` table.\n\n"
+    else:
+        prompt += (
+            f"Generate rows for the `{table_name}` table. "
+            f"The Foreign Key column may only contain values from Context Table Data.\n\n"
+        )
+    if previous_rows:
+        prompt += (
+            "Generate new rows that maintain consistency with the previous rows where appropriate. "
+            "Don't pay attention to the number of previous rows; there might have been more generated than provided.\n\n"
+        )
+    prompt += f"Do not use code to generate the data.\n\n"
+    prompt += f"Return the full data as a JSON string.\n"
+
+    return prompt
+
+
+def _create_table_rows_generator(
+    *,
+    table_name: str,
+    table_config: TableConfig,
+    primary_keys: dict[str, str] | None,
+    sample_size: int,
+    temperature: float,
+    top_p: float,
+    context_data: pd.DataFrame | None,
+    batch_size: int,
+    previous_rows_size: int,
+    llm_config: LLMConfig,
+) -> Generator[dict]:
+    def create_table_response_format(columns: dict[str, ColumnConfig]) -> BaseModel:
+        dtype_to_pydantic_type = {
+            DType.INTEGER: int,
+            DType.FLOAT: float,
+            DType.STRING: str,
+            DType.BOOLEAN: bool,
+            # response_format has limited support for JSON Schema features
+            # thus we represent dates and datetimes as strings
+            DType.DATE: str,
+            DType.DATETIME: str,
+        }
+        fields = {}
+        for column_name, column_config in columns.items():
+            annotation = dtype_to_pydantic_type[column_config.dtype]
+            fields[column_name] = (annotation, Field(...))
+        TableRow = create_model("TableRow", **fields)
+        TableRows = create_model("TableRows", rows=(list[TableRow], ...))
+        return TableRows
+
+    def yield_rows_from_json_chunks_stream(response: litellm.CustomStreamWrapper) -> Generator[dict]:
+        # starting with dirty buffer is to handle the `{"rows": []}` case
+        buffer = "garbage"
+        rows_json_started = False
+        in_row_json = False
+        for chunk in response:
+            delta = chunk.choices[0].delta.content
+            if delta is None:
+                continue
+            for char in delta:
+                buffer += char
+                if char == "{" and not rows_json_started:
+                    # {"rows": [{"name": "Jo\}h\{n"}]}
+                    # *                                 <- start of rows json stream
+                    rows_json_started = True
+                elif char == "{" and not in_row_json:
+                    # {"rows": [{"name": "Jo\}h\{n"}]}
+                    #           *                       <- start of single row json stream
+                    buffer = "{"
+                    in_row_json = True
+                elif char == "}":
+                    # {"rows": [{"name": "Jo\}h\{n"}]}
+                    #                        *     * *  <- any of these
+                    try:
+                        row = json.loads(buffer)
+                        yield row
+                        buffer = ""
+                        in_row_json = False
+                    except json.JSONDecodeError:
+                        continue
+
+    def batch_infinitely(data: pd.DataFrame | None) -> Generator[pd.DataFrame | None]:
+        while True:
+            if data is None:
+                yield None
+            else:
+                for i in range(0, len(data), batch_size):
+                    yield data.iloc[i : i + batch_size]
+
+    # ensure model supports response_format and json schema
+    supported_params = litellm.get_supported_openai_params(model=llm_config.model)
+    assert "response_format" in supported_params
+    assert litellm.supports_response_schema(llm_config.model), (
+        "The model does not support structured output / JSON mode."
+    )
+
+    litellm_kwargs = {
+        "response_format": create_table_response_format(columns=table_config.columns),
+        "temperature": temperature,
+        "top_p": top_p,
+        "model": llm_config.model,
+        "api_key": llm_config.api_key,
+        "stream": True,
+    }
+
+    yielded_sequences = 0
+    previous_rows = deque(maxlen=previous_rows_size)
+    for context_batch in batch_infinitely(context_data):
+        prompt_kwargs = {
+            "table_name": table_name,
+            "table_description": table_config.description,
+            "columns": table_config.columns,
+            "primary_keys": primary_keys,
+            "batch_size": batch_size if context_batch is None else None,
+            "foreign_keys": table_config.foreign_keys if context_batch is not None else None,
+            "context_data": context_batch if context_batch is not None else None,
+            "previous_rows": list(previous_rows),
+        }
+        prompt = _create_table_prompt(**prompt_kwargs)
+        messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}]
+
+        response = litellm.completion(messages=messages, **litellm_kwargs)
+        rows_stream = yield_rows_from_json_chunks_stream(response)
+
+        while True:
+            try:
+                row = next(rows_stream)
+            except StopIteration:
+                break  # move to next batch
+            previous_rows.append(row)
+            yield row
+            if context_batch is None:
+                # each subject row is considered a single sequence
+                yielded_sequences += 1
+                if yielded_sequences >= sample_size:
+                    return  # move to next table
+        if context_batch is not None:
+            # for each context_batch, full sequences are generated
+            yielded_sequences += len(context_batch)
+            if yielded_sequences >= sample_size:
+                return  # move to next table
+
+
+def _convert_table_rows_generator_to_df(
+    table_rows_generator: Generator[dict], table_config: TableConfig
+) -> pd.DataFrame:
+    def align_df_dtypes_with_mock_dtypes(df: pd.DataFrame, columns: dict[str, ColumnConfig]) -> pd.DataFrame:
+        for column_name, column_config in columns.items():
+            if column_config.dtype in [DType.DATE, DType.DATETIME]:
+                # datetime.date, datetime.datetime -> datetime64[ns] / datetime64[ns, tz]
+                df[column_name] = pd.to_datetime(df[column_name], errors="coerce")
+            elif column_config.dtype in [DType.INTEGER, DType.FLOAT]:
+                # int -> int64[pyarrow], float -> double[pyarrow]
+                df[column_name] = pd.to_numeric(df[column_name], errors="coerce", dtype_backend="pyarrow")
+            elif column_config.dtype is DType.BOOLEAN:
+                # bool -> bool
+                df[column_name] = df[column_name].astype(bool)
+            else:
+                # other -> string[pyarrow]
+                df[column_name] = df[column_name].astype("string[pyarrow]")
+        return df
+
+    df = pd.DataFrame(list(table_rows_generator))
+    df = align_df_dtypes_with_mock_dtypes(df, table_config.columns)
+    return df
+
+
+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}
+
+    if sample_size.keys() != config.root.keys():
+        raise ValueError(f"Sample size keys must match table names: {sample_size.keys()} != {config.root.keys()}")
+    return sample_size
+
+
+def sample(
+    *,
+    tables: dict[str, dict],
+    sample_size: int | dict[str, int] = 10,
+    model: str = "openai/gpt-4.1-nano",
+    api_key: str | None = None,
+    temperature: float = 1.0,
+    top_p: float = 0.95,
+) -> pd.DataFrame | dict[str, pd.DataFrame]:
+    """
+    Generate mock 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.
+            Default is 10.
+        model (str): The LiteLLM chat completion model to be used. Requires support for structured output / JSON mode.
+            Examples include:
+            - `openai/gpt-4.1-nano` (default)
+            - `openai/gpt-4.1-mini`
+            - `openai/gpt-4.1`
+            - `gemini/gemini-2.0-flash`
+            - `gemini/gemini-2.5-flash-preview-04-17`
+            See https://docs.litellm.ai/docs/providers/ for more options.
+        api_key (str | None): The API key to use for the LLM. If not provided, LiteLLM will take it from the environment variables.
+        temperature (float): The temperature to use for the LLM. Default is 1.0.
+        top_p (float): The top-p value to use for the LLM. Default is 0.95.
+
+    Returns:
+        - 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):
+    ```python
+    from mostlyai import mock
+
+    tables = {
+        "guests": {
+            "description": "Guests of an Alpine ski hotel in Austria",
+            "columns": {
+                "nationality": {"prompt": "2-letter code for the nationality", "dtype": "string"},
+                "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+                "gender": {"prompt": "gender of the guest; male or female", "dtype": "string"},
+                "age": {"prompt": "age in years; min: 18, max: 80; avg: 25", "dtype": "integer"},
+                "date_of_birth": {"prompt": "date of birth", "dtype": "date"},
+                "checkin_time": {"prompt": "the check in timestamp of the guest; may 2025", "dtype": "datetime"},
+                "is_vip": {"prompt": "is the guest a VIP", "dtype": "boolean"},
+                "price_per_night": {"prompt": "price paid per night, in EUR", "dtype": "float"},
+            },
+        }
+    }
+    df = mock.sample(tables=tables, sample_size=10, model="openai/gpt-4.1-nano")
+    ```
+
+    Example of multiple tables (with PK/FK relationships):
+    ```python
+    from mostlyai import mock
+
+    tables = {
+        "guests": {
+            "description": "Guests of an Alpine ski hotel in Austria",
+            "columns": {
+                "id": {"prompt": "the unique id of the guest", "dtype": "integer"},
+                "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+            },
+            "primary_key": "id",
+        },
+        "purchases": {
+            "description": "Purchases of a Guest during their stay",
+            "columns": {
+                "guest_id": {"prompt": "the guest id for that purchase", "dtype": "integer"},
+                "purchase_id": {"prompt": "the unique id of the purchase", "dtype": "string"},
+                "text": {"prompt": "purchase text description", "dtype": "string"},
+                "amount": {"prompt": "purchase amount in EUR", "dtype": "float"},
+            },
+            "foreign_keys": [
+                {
+                    "column": "guest_id",
+                    "referenced_table": "guests",
+                    "description": "each guest has anywhere between 1 and 10 purchases",
+                }
+            ],
+        },
+    }
+    data = mock.sample(tables=tables, sample_size=5, model="openai/gpt-4.1-nano")
+    df_guests = data["guests"]
+    df_purchases = data["purchases"]
+    ```
+    """
+
+    config = MockConfig(tables)
+
+    sample_size = _harmonize_sample_size(sample_size, config)
+    primary_keys = {table_name: table_config.primary_key for table_name, table_config in config.root.items()}
+    dfs = {}
+    for table_name, table_config in config.root.items():
+        if len(dfs) == 0:
+            # subject table
+            df = _sample_table(
+                table_name=table_name,
+                table_config=table_config,
+                primary_keys=None,
+                sample_size=sample_size[table_name],
+                context_data=None,
+                temperature=temperature,
+                top_p=top_p,
+                batch_size=20,  # generate 20 subjects at a time
+                previous_rows_size=5,
+                llm_config=LLMConfig(model=model, api_key=api_key),
+            )
+        elif len(dfs) == 1:
+            # sequence table
+            df = _sample_table(
+                table_name=table_name,
+                table_config=table_config,
+                primary_keys=primary_keys,
+                sample_size=None,
+                context_data=next(iter(dfs.values())),
+                temperature=temperature,
+                top_p=top_p,
+                batch_size=1,  # generate one sequence at a time
+                previous_rows_size=5,
+                llm_config=LLMConfig(model=model, api_key=api_key),
+            )
+        else:
+            raise RuntimeError("Only 1 or 2 table setups are supported for now")
+        dfs[table_name] = df
+
+    return dfs if len(dfs) > 1 else next(iter(dfs.values()))

mostlyai_mock-0.0.1/pyproject.toml

@@ -0,0 +1,55 @@
+[project]
+name = "mostlyai-mock"
+version = "0.0.1"
+description = "Synthetic Mock Data"
+authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
+requires-python = ">=3.10"
+readme = "README.md"
+license = "Apache-2.0"
+dependencies = [
+    "pydantic>=2.0.0,<3.0.0",
+    "numpy>=1.26.3",
+    "pandas>=2.0.0",
+    "pyarrow>=14.0.0",
+    "litellm>=1.67.0",
+]
+
+[project.urls]
+homepage = "https://github.com/mostly-ai/mostlyai-mock"
+repository = "https://github.com/mostly-ai/mostlyai-mock"
+documentation = "https://mostly-ai.github.io/mostlyai-mock/"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+    "ruff>=0.11",  # sync'ed with .pre-commit-config
+    "pre-commit>=4.0",
+    "twine>=6.1",
+    "ipykernel>=6.25",
+]
+docs = [
+    "mkdocs>=1.6",
+    "mkdocstrings[crystal, python]>=0.29",
+    "mkdocs-material>=9.0",
+    "griffe>=1.0",
+    "pymdown-extensions>=10.0",
+    "griffe-fieldz>=0.2",
+    "black>=25.0",
+]
+
+[tool.uv]
+default-groups = ["dev", "docs"]
+
+[tool.hatch.build.targets.sdist]
+include = ["mostlyai/mock"]
+
+[tool.hatch.build.targets.wheel]
+include = ["mostlyai/mock"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.ruff]
+target-version = "py310"
+line-length = 120