mostlyai-mock 0.0.1__py3-none-any.whl → 0.0.5__py3-none-any.whl

mostlyai/mock/__init__.py

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

mostlyai/mock/core.py

@@ -18,10 +18,11 @@ import json
 from collections import deque
 from collections.abc import Generator
 from enum import Enum
+from typing import Any, Literal, Type
 
 import litellm
 import pandas as pd
-from pydantic import BaseModel, Field, RootModel, create_model, field_validator
+from pydantic import BaseModel, Field, RootModel, create_model, field_validator, model_validator
 from tqdm import tqdm
 
 SYSTEM_PROMPT = f"""
@@ -88,6 +89,31 @@ class MockConfig(RootModel[dict[str, "TableConfig"]]):
 
         return tables
 
+    @model_validator(mode="after")
+    def validate_no_circular_dependencies(self) -> MockConfig:
+        child_to_parents = {}
+        for table_name, table_config in self.root.items():
+            child_to_parents[table_name] = [fk.referenced_table for fk in table_config.foreign_keys]
+        visited = set()
+
+        def detect_cycle(table_name: str, path: list[str]) -> None:
+            if table_name in path:
+                cycle_start = path.index(table_name)
+                cycle = path[cycle_start:] + [table_name]
+                raise ValueError(f"Circular dependency detected: {' -> '.join(cycle)}")
+            if table_name in visited:
+                return
+            visited.add(table_name)
+            path.append(table_name)
+            for parent in child_to_parents[table_name]:
+                detect_cycle(parent, path)
+            path.pop()
+
+        for table_name in child_to_parents:
+            detect_cycle(table_name, [])
+
+        return self
+
 
 class TableConfig(BaseModel):
     description: str = ""
@@ -97,14 +123,59 @@ class TableConfig(BaseModel):
 
 
 class ColumnConfig(BaseModel):
-    prompt: str
+    prompt: str = ""
     dtype: DType
+    values: list[Any] = Field(default_factory=list)
+
+    @model_validator(mode="before")
+    def set_default_dtype(cls, data):
+        if isinstance(data, dict):
+            if "dtype" not in data:
+                if data.get("values"):
+                    data["dtype"] = DType.CATEGORY
+                else:
+                    data["dtype"] = DType.STRING
+        return data
+
+    @model_validator(mode="after")
+    def ensure_values_are_unique(self) -> ColumnConfig:
+        if self.values:
+            if len(self.values) != len(set(self.values)):
+                raise ValueError("Values must be unique")
+        return self
+
+    @model_validator(mode="after")
+    def ensure_values_are_provided_for_category_dtype(self) -> ColumnConfig:
+        if self.dtype == DType.CATEGORY and not self.values:
+            raise ValueError("At least one value must be provided when dtype is 'category'")
+        return self
+
+    @model_validator(mode="after")
+    def harmonize_values_with_dtypes(self) -> ColumnConfig:
+        if self.values:
+            cast_fn, convertible_to = {
+                DType.INTEGER: (int, "integers"),
+                DType.FLOAT: (float, "floats"),
+                DType.STRING: (str, "strings"),
+                DType.CATEGORY: (lambda c: c, "categories"),
+                DType.BOOLEAN: (bool, "booleans"),
+                DType.DATE: (str, "strings"),
+                DType.DATETIME: (str, "strings"),
+            }[self.dtype]
+            try:
+                self.values = [cast_fn(c) for c in self.values]
+            except ValueError:
+                raise ValueError(
+                    f"All values must be convertible to {convertible_to} when dtype is '{self.dtype.value}'"
+                )
+        return self
 
 
 class DType(str, Enum):
     INTEGER = "integer"
     FLOAT = "float"
     STRING = "string"
+    CATEGORY = "category"
     BOOLEAN = "boolean"
     DATE = "date"
     DATETIME = "datetime"
@@ -188,7 +259,7 @@ def _create_table_prompt(
     # 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)
+        prompt += f"{json.dumps(previous_rows, indent=2)}\n\n"
 
     # add context table name, primary key and data
     if context_data is not None:
@@ -206,12 +277,14 @@ def _create_table_prompt(
         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"
+            f"Generate data for the `{table_name}` table. "
+            f"The Foreign Key column may only contain values from Context Table Data. "
+            f"Pay attention to description of the Foreign Key column to understand the relationship.\n\n"
         )
     if previous_rows:
         prompt += (
             "Generate new rows that maintain consistency with the previous rows where appropriate. "
+            "Don't copy previous rows in the output. "
             "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"
@@ -234,19 +307,23 @@ def _create_table_rows_generator(
     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,
-        }
+        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 {
+                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,
+            }[column_config.dtype]
+
         fields = {}
         for column_name, column_config in columns.items():
-            annotation = dtype_to_pydantic_type[column_config.dtype]
+            annotation = create_annotation(column_config)
             fields[column_name] = (annotation, Field(...))
         TableRow = create_model("TableRow", **fields)
         TableRows = create_model("TableRows", rows=(list[TableRow], ...))
@@ -351,16 +428,14 @@ def _convert_table_rows_generator_to_df(
     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)
+            elif column_config.dtype is DType.CATEGORY:
+                df[column_name] = pd.Categorical(df[column_name], categories=column_config.values)
             else:
-                # other -> string[pyarrow]
                 df[column_name] = df[column_name].astype("string[pyarrow]")
         return df
 
@@ -378,6 +453,44 @@ def _harmonize_sample_size(sample_size: int | dict[str, int], config: MockConfig
     return sample_size
 
 
+def _build_dependency_graph(config: MockConfig) -> tuple[dict[str, list[str]], dict[str, list[str]], list[str]]:
+    child_to_parents = {}
+    parent_to_children = {}
+
+    for table_name in config.root:
+        child_to_parents[table_name] = []
+        parent_to_children[table_name] = []
+
+    for table_name, table_config in config.root.items():
+        if table_config.foreign_keys:
+            for fk in table_config.foreign_keys:
+                referenced_table = fk.referenced_table
+                child_to_parents[table_name].append(referenced_table)
+                parent_to_children[referenced_table].append(table_name)
+
+    subject_tables = [table_name for table_name, deps in child_to_parents.items() if not deps]
+    return child_to_parents, parent_to_children, subject_tables
+
+
+def _build_execution_plan(parent_to_children: dict[str, list[str]], subject_tables: list[str]) -> list[str]:
+    execution_plan = []
+    bfs_queue = list(subject_tables)
+    processed = set()
+
+    while bfs_queue:
+        table_name = bfs_queue.pop(0)
+        if table_name in processed:
+            continue
+
+        execution_plan.append(table_name)
+        processed.add(table_name)
+
+        for child in parent_to_children[table_name]:
+            if child not in bfs_queue and child not in processed:
+                bfs_queue.append(child)
+    return execution_plan
+
+
 def sample(
     *,
     tables: dict[str, dict],
@@ -404,6 +517,8 @@ def sample(
             - `openai/gpt-4.1`
             - `gemini/gemini-2.0-flash`
             - `gemini/gemini-2.5-flash-preview-04-17`
+            - `groq/llama-3.3-70b-versatile`
+            - `anthropic/claude-3-7-sonnet-latest`
             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.
@@ -423,12 +538,13 @@ def sample(
             "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"},
+                "gender": {"dtype": "category", "values": ["male", "female"]},
                 "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"},
+                "room_number": {"prompt": "room number", "dtype": "integer", "values": [101, 102, 103, 201, 202, 203, 204]}
             },
         }
     }
@@ -440,34 +556,52 @@ def sample(
     from mostlyai import mock
 
     tables = {
-        "guests": {
-            "description": "Guests of an Alpine ski hotel in Austria",
+        "customers": {
+            "description": "Customers of a hardware store",
             "columns": {
-                "id": {"prompt": "the unique id of the guest", "dtype": "integer"},
-                "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+                "customer_id": {"prompt": "the unique id of the customer", "dtype": "integer"},
+                "name": {"prompt": "first name and last name of the customer", "dtype": "string"},
             },
-            "primary_key": "id",
+            "primary_key": "customer_id",
         },
-        "purchases": {
-            "description": "Purchases of a Guest during their stay",
+        "orders": {
+            "description": "Orders of a Customer",
             "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"},
+                "customer_id": {"prompt": "the customer id for that order", "dtype": "integer"},
+                "order_id": {"prompt": "the unique id of the order", "dtype": "string"},
+                "text": {"prompt": "order text description", "dtype": "string"},
+                "amount": {"prompt": "order amount in USD", "dtype": "float"},
             },
+            "primary_key": "order_id",
             "foreign_keys": [
                 {
-                    "column": "guest_id",
-                    "referenced_table": "guests",
-                    "description": "each guest has anywhere between 1 and 10 purchases",
+                    "column": "customer_id",
+                    "referenced_table": "customers",
+                    "description": "each customer has anywhere between 1 and 3 orders",
+                }
+            ],
+        },
+        "items": {
+            "description": "Items in an Order",
+            "columns": {
+                "item_id": {"prompt": "the unique id of the item", "dtype": "string"},
+                "order_id": {"prompt": "the order id for that item", "dtype": "string"},
+                "name": {"prompt": "the name of the item", "dtype": "string"},
+                "price": {"prompt": "the price of the item in USD", "dtype": "float"},
+            },
+            "foreign_keys": [
+                {
+                    "column": "order_id",
+                    "referenced_table": "orders",
+                    "description": "each order has between 2 and 5 items",
                 }
             ],
         },
     }
-    data = mock.sample(tables=tables, sample_size=5, model="openai/gpt-4.1-nano")
-    df_guests = data["guests"]
-    df_purchases = data["purchases"]
+    data = mock.sample(tables=tables, sample_size=2, model="openai/gpt-4.1")
+    df_customers = data["customers"]
+    df_orders = data["orders"]
+    df_items = data["items"]
     ```
     """
 
@@ -475,9 +609,15 @@ def sample(
 
     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:
+
+    child_to_parents, parent_to_children, subject_tables = _build_dependency_graph(config)
+    execution_plan: list[str] = _build_execution_plan(parent_to_children, subject_tables)
+
+    results: dict[str, pd.DataFrame] = {}
+
+    for table_name in execution_plan:
+        table_config = config.root[table_name]
+        if not child_to_parents[table_name]:
             # subject table
             df = _sample_table(
                 table_name=table_name,
@@ -491,22 +631,21 @@ def sample(
                 previous_rows_size=5,
                 llm_config=LLMConfig(model=model, api_key=api_key),
             )
-        elif len(dfs) == 1:
-            # sequence table
+        else:
+            # sequencial table
+            referenced_table = table_config.foreign_keys[0].referenced_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())),
+                context_data=results[referenced_table],
                 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
+        results[table_name] = df
 
-    return dfs if len(dfs) > 1 else next(iter(dfs.values()))
+    return results if len(results) > 1 else next(iter(results.values()))

{mostlyai_mock-0.0.1.dist-info → mostlyai_mock-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-mock
-Version: 0.0.1
+Version: 0.0.5
 Summary: Synthetic Mock Data
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-mock
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-mock
@@ -8,7 +8,6 @@ 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
@@ -46,12 +45,13 @@ tables = {
         "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"},
+            "gender": {"dtype": "category", "values": ["male", "female"]},
             "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"},
+            "room_number": {"prompt": "room number", "dtype": "integer", "values": [101, 102, 103, 201, 202, 203, 204]}
         },
     }
 }
@@ -65,34 +65,53 @@ print(df)
 from mostlyai import mock
 
 tables = {
-    "guests": {
-        "description": "Guests of an Alpine ski hotel in Austria",
+    "customers": {
+        "description": "Customers of a hardware store",
         "columns": {
-            "id": {"prompt": "the unique id of the guest", "dtype": "integer"},
-            "name": {"prompt": "first name and last name of the guest", "dtype": "string"},
+            "customer_id": {"prompt": "the unique id of the customer", "dtype": "integer"},
+            "name": {"prompt": "first name and last name of the customer", "dtype": "string"},
+        },
+        "primary_key": "customer_id",
+    },
+    "orders": {
+        "description": "Orders of a Customer",
+        "columns": {
+            "customer_id": {"prompt": "the customer id for that order", "dtype": "integer"},
+            "order_id": {"prompt": "the unique id of the order", "dtype": "string"},
+            "text": {"prompt": "order text description", "dtype": "string"},
+            "amount": {"prompt": "order amount in USD", "dtype": "float"},
         },
-        "primary_key": "id",
+        "primary_key": "order_id",
+        "foreign_keys": [
+            {
+                "column": "customer_id",
+                "referenced_table": "customers",
+                "description": "each customer has anywhere between 1 and 3 orders",
+            }
+        ],
     },
-    "purchases": {
-        "description": "Purchases of a Guest during their stay",
+    "items": {
+        "description": "Items in an Order",
         "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"},
+            "item_id": {"prompt": "the unique id of the item", "dtype": "string"},
+            "order_id": {"prompt": "the order id for that item", "dtype": "string"},
+            "name": {"prompt": "the name of the item", "dtype": "string"},
+            "price": {"prompt": "the price of the item in USD", "dtype": "float"},
         },
         "foreign_keys": [
             {
-                "column": "guest_id",
-                "referenced_table": "guests",
-                "description": "each guest has anywhere between 1 and 10 purchases",
+                "column": "order_id",
+                "referenced_table": "orders",
+                "description": "each order has between 2 and 5 items",
             }
         ],
     },
 }
-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)
+data = mock.sample(tables=tables, sample_size=2, model="openai/gpt-4.1")
+df_customers = data["customers"]
+df_orders = data["orders"]
+df_items = data["items"]
+print(df_customers)
+print(df_orders)
+print(df_items)
 ```

mostlyai_mock-0.0.5.dist-info/RECORD

@@ -0,0 +1,6 @@
+mostlyai/mock/__init__.py,sha256=u9aoJD9XCxM0h1mIWBwkm_O_7VUf9d6-y8YiC4diUAM,714
+mostlyai/mock/core.py,sha256=zgUmv6yclP51NJFE5W0CtE4raXZUnKF6Fa8r8_idbpI,26351
+mostlyai_mock-0.0.5.dist-info/METADATA,sha256=Ek_7faZR-YNx7ntoiqqBZxahcwC5f2oiVVM1FseXMf8,4655
+mostlyai_mock-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mostlyai_mock-0.0.5.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mostlyai_mock-0.0.5.dist-info/RECORD,,

mostlyai_mock-0.0.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-mostlyai/mock/__init__.py,sha256=TM39__slJHAmk30gt6C8gm5D12tz8Ow7QSArHp290EI,714
-mostlyai/mock/core.py,sha256=naKteDOKCbXKtH9ldKBtPxd8-oZDgGhZPfOgtROCcJU,20626
-mostlyai_mock-0.0.1.dist-info/METADATA,sha256=TmYgQAW2nBdwfOp9QkFnEhT_UiBdSe1ABQwMJCGH4ZE,3875
-mostlyai_mock-0.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mostlyai_mock-0.0.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mostlyai_mock-0.0.1.dist-info/licenses/LICENSE_HEADER,sha256=jsSKYMHBOhEvPX1bB1Q2-EPbKZs1MAaQDI8S_AEOA2g,550
-mostlyai_mock-0.0.1.dist-info/RECORD,,

mostlyai_mock-0.0.1.dist-info/licenses/LICENSE_HEADER

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