mostlyai-mock 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

mostlyai/mock/__init__.py

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

mostlyai/mock/core.py

@@ -14,42 +14,29 @@
 
 from __future__ import annotations
 
-import itertools
+import asyncio
+import concurrent.futures
 import json
+import math
 from collections import deque
-from collections.abc import Generator
+from collections.abc import AsyncGenerator
 from enum import Enum
+from io import StringIO
 from typing import Any, Literal
 
+import dateutil.parser
 import litellm
 import pandas as pd
 import tenacity
 from pydantic import BaseModel, Field, RootModel, create_model, field_validator, model_validator
-from tqdm import tqdm
+from tqdm.asyncio 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.
 
-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. When enriching existing data, ensure that new values are consistent with existing values
-7. Return well-formatted JSON output that can be directly parsed
-8. 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.
-
-When enriching existing data, carefully analyze the patterns and relationships in the existing columns \
-to generate compatible and realistic values for the missing columns.
-"""
+class LLMOutputFormat(str, Enum):
+    JSON = "JSON"
+    CSV = "CSV"
 
 
 class LLMConfig(BaseModel):
@@ -162,6 +149,12 @@ class ColumnConfig(BaseModel):
             raise ValueError("At least one value must be provided when dtype is 'category'")
         return self
 
+    @model_validator(mode="after")
+    def override_values_for_boolean_dtype(self) -> ColumnConfig:
+        if self.dtype == DType.BOOLEAN:
+            self.values = [True, False]
+        return self
+
     @model_validator(mode="after")
     def harmonize_values_with_dtypes(self) -> ColumnConfig:
         if self.values:
@@ -199,18 +192,18 @@ class ForeignKeyConfig(BaseModel):
     prompt: str | None = None
 
 
-def _sample_table(
+async def _sample_table(
     *,
     name: str,
     prompt: str,
     columns: dict[str, ColumnConfig],
-    foreign_keys: list[ForeignKeyConfig] | None,
-    primary_keys: dict[str, str] | None,
+    foreign_keys: list[ForeignKeyConfig],
+    primary_keys: dict[str, str],
     data: dict[str, pd.DataFrame],
     sample_size: int,
-    batch_size: int,
     previous_rows_size: int,
     non_context_size: int | None,
+    n_workers: int,
     llm_config: LLMConfig,
 ) -> pd.DataFrame:
     table_rows_generator = _create_table_rows_generator(
@@ -221,28 +214,62 @@ def _sample_table(
         foreign_keys=foreign_keys,
         data=data,
         sample_size=sample_size,
-        batch_size=batch_size,
         previous_rows_size=previous_rows_size,
         non_context_size=non_context_size,
+        n_workers=n_workers,
         llm_config=llm_config,
     )
     table_rows_generator = tqdm(table_rows_generator, desc=f"Generating rows for table `{name}`".ljust(45))
-    table_df = _convert_table_rows_generator_to_df(table_rows_generator=table_rows_generator, columns=columns)
+    table_df = await _convert_table_rows_generator_to_df(table_rows_generator=table_rows_generator, columns=columns)
     return table_df
 
 
+def _sample_table_sync(*args, **kwargs) -> pd.DataFrame:
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        return loop.run_until_complete(_sample_table(*args, **kwargs))
+    finally:
+        loop.close()
+
+
+def _create_system_prompt(llm_output_format: LLMOutputFormat) -> str:
+    return f"""
+You are a specialized 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. When enriching existing data, ensure that new values are consistent with existing values
+7. Return well-formatted {llm_output_format.value} output that can be directly parsed
+8. 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.
+
+When enriching existing data, carefully analyze the patterns and relationships in the existing columns \
+to generate compatible and realistic values for the missing columns.
+"""
+
+
 def _create_table_prompt(
     *,
     name: str,
     prompt: str,
     columns: dict[str, ColumnConfig],
-    primary_keys: dict[str, str] | None,
+    primary_keys: dict[str, str],
     batch_size: int | None,
-    foreign_keys: list[ForeignKeyConfig] | None,
+    foreign_keys: list[ForeignKeyConfig],
     existing_data: pd.DataFrame | None,
     context_data: pd.DataFrame | None,
     non_context_data: dict[str, pd.DataFrame] | None,
     previous_rows: list[dict] | None,
+    llm_output_format: LLMOutputFormat,
 ) -> str:
     # add table prompt
     prompt = f"# {prompt}\n\n"
@@ -345,7 +372,7 @@ def _create_table_prompt(
 
     prompt += f"{verb.capitalize()} data for the Target Table `{name}`.\n\n"
     if n_rows is not None:
-        prompt += f"Number of rows to {verb}: `{n_rows}`.\n\n"
+        prompt += f"Number of data rows to {verb}: `{n_rows}`.\n\n"
 
     if has_context_table_section:
         assert foreign_keys
@@ -387,131 +414,341 @@ def _create_table_prompt(
 
     prompt += f"Do not use code to {verb} the data.\n\n"
 
-    prompt += "Return data as a JSON string."
-    prompt += " The JSON string should have 'rows' key at the top level. The value of 'rows' key should be a list of JSON objects."
-    prompt += " Each JSON object should have column names as keys and values as column values."
+    prompt += f"Return data as a {llm_output_format.value} string."
+    if llm_output_format == LLMOutputFormat.JSON:
+        prompt += " The JSON string should have 'rows' key at the top level."
+        prompt += " The value of 'rows' key should be a list of JSON objects."
+        prompt += " Each JSON object should have column names as keys and values as column values."
+    else:  # llm_output_format == LLMOutputFormat.CSV
+        prompt += " The CSV string should have a header row with column names."
+        prompt += " The CSV string should have a data row for each row to be generated."
+        prompt += " The CSV string should have a newline character at the end of each row."
+        prompt += " Each value in the CSV string should be enclosed in double quotes."
+
     if existing_data is not None:
-        prompt += (
-            f" Only include the following columns in the JSON string: {list(columns.keys() - existing_data.columns)}."
-        )
+        prompt += f" Only include the following columns in the {llm_output_format.value} string: {list(columns.keys() - existing_data.columns)}."
+
+    if llm_output_format == LLMOutputFormat.CSV and batch_size > 10:
+        prompt += " Additionally, add column called `_ROW_IDX` that is a counter from 1 to the number of rows generated so far within current batch."
+
     prompt += "\n"
     return prompt
 
 
-def _create_table_rows_generator(
+def _completion_with_retries(*args, **kwargs):
+    n_attempts = 3
+
+    def print_on_retry(_):
+        print(" * Calling LLM again... * ", end="", flush=True)
+
+    # try up to 3 times, print a message to the user on each retry
+    retryer = tenacity.AsyncRetrying(
+        stop=tenacity.stop_after_attempt(n_attempts), reraise=True, before_sleep=print_on_retry
+    )
+    return retryer(litellm.acompletion, *args, **kwargs)
+
+
+async def _yield_rows_from_json_chunks_stream(response: litellm.CustomStreamWrapper) -> AsyncGenerator[dict]:
+    def buffer_to_row(buffer: list[str]) -> dict:
+        return json.loads("".join(buffer))
+
+    # starting with dirty buffer is to handle the `{"rows": []}` case
+    buffer = list("garbage")
+    rows_json_started = False
+    in_row_json = False
+    async for chunk in response:
+        delta = chunk.choices[0].delta.content
+        if delta is None:
+            continue
+        for char in delta:
+            buffer.append(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 = list("{")
+                in_row_json = True
+            elif char == "}":
+                # {"rows": [{"name": "Jo\}h\{n"}]}
+                #                        *     * *  <- any of these
+                try:
+                    row = buffer_to_row(buffer)
+                except Exception:
+                    # in case of any error, silently drop the row
+                    continue
+                finally:
+                    buffer = list()
+                    in_row_json = False
+                yield row
+
+
+async def _yield_rows_from_csv_chunks_stream(response: litellm.CustomStreamWrapper) -> AsyncGenerator[dict]:
+    def buffer_to_row(buffer: list[str]) -> list[str]:
+        return pd.read_csv(StringIO("".join(buffer)), header=None).astype(str).iloc[0].to_list()
+
+    buffer = list()
+    header = None
+    async for chunk in response:
+        delta = chunk.choices[0].delta.content
+        if delta is None:
+            continue
+        for char in delta:
+            buffer.append(char)
+            if char == "\n":
+                try:
+                    row = buffer_to_row(buffer)
+                except Exception:
+                    # in case of any error, silently drop the row
+                    continue
+                finally:
+                    buffer = list()
+                if header is None:
+                    # column1,column2,column3\n
+                    #                        ** <- end of header row
+                    header = row
+                else:
+                    # value_1,value_2,value_3\n
+                    #                        ** <- end of data row
+                    yield dict(zip(header, row))
+    if buffer:
+        # last row might not finish with a newline, in which case the buffer would not be empty here
+        try:
+            last_row = buffer_to_row(buffer)
+            yield dict(zip(header, last_row))
+        except Exception:
+            # in case of any error, silently drop the row
+            pass
+
+
+def _create_structured_output_schema(
+    columns: dict[str, ColumnConfig], existing_data: pd.DataFrame | None
+) -> type[BaseModel]:
+    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 | None,
+            DType.FLOAT: float | None,
+            DType.STRING: str | None,
+            DType.BOOLEAN: bool | None,
+            # response_format has limited support for JSON Schema features
+            # thus we represent dates and datetimes as strings
+            DType.DATE: str | None,
+            DType.DATETIME: str | None,
+        }[column_config.dtype]
+
+    fields = {}
+    for column_name, column_config in columns.items():
+        if existing_data is not None and column_name in existing_data.columns:
+            continue  # skip columns that already exist in existing data
+        annotation = create_annotation(column_config)
+        fields[column_name] = (annotation, Field(...))
+    TableRow = create_model("TableRow", **fields)
+    TableRows = create_model("TableRows", rows=(list[TableRow], ...))
+    return TableRows
+
+
+async def _worker(
+    *,
+    name: str,
+    prompt: str,
+    columns: dict[str, ColumnConfig],
+    foreign_keys: list[ForeignKeyConfig],
+    primary_keys: dict[str, str],
+    previous_rows: deque[dict],
+    batch_queue: asyncio.Queue,
+    result_queue: asyncio.Queue,
+    retry_queue: asyncio.Queue,
+    n_workers: int,
+    llm_output_format: LLMOutputFormat,
+    llm_config: LLMConfig,
+):
+    try:
+        while True:
+            do_repeat_task = False
+
+            # get task from the batch_queue
+            batch_idx, task = await batch_queue.get()
+            if task is None:
+                # no more tasks for the worker; break the loop
+                batch_queue.task_done()
+                break
+
+            # deconstruct task
+            batch_size = task["batch_size"]
+            existing_batch = task.get("existing_batch")
+            context_batch = task.get("context_batch")
+            non_context_batch = task.get("non_context_batch")
+
+            # resolve columns to generate
+            generated_columns = set(columns.keys())
+            if existing_batch is not None:
+                generated_columns = generated_columns - set(existing_batch.columns)
+
+            # construct schema for Structured Outputs (applies to JSON LLMOutputFormat only)
+            structured_output_schema = None
+            if llm_output_format == LLMOutputFormat.JSON:
+                structured_output_schema = _create_structured_output_schema(
+                    columns=columns, existing_data=existing_batch
+                )
+
+            # construct litellm kwargs
+            litellm_kwargs = {
+                "temperature": llm_config.temperature,
+                "top_p": llm_config.top_p,
+                "model": llm_config.model,
+                "api_key": llm_config.api_key,
+                "stream": True,
+            }
+
+            # construct messages
+            system_prompt = _create_system_prompt(llm_output_format)
+            user_prompt = _create_table_prompt(
+                name=name,
+                prompt=prompt,
+                columns=columns,
+                primary_keys=primary_keys,
+                batch_size=batch_size,
+                foreign_keys=foreign_keys,
+                existing_data=existing_batch,
+                context_data=context_batch,
+                non_context_data=non_context_batch,
+                previous_rows=list(previous_rows),
+                llm_output_format=llm_output_format,
+            )
+            messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]
+
+            if generated_columns:
+                # make LLM call
+                response = await _completion_with_retries(
+                    messages=messages, response_format=structured_output_schema, **litellm_kwargs
+                )
+                yield_rows_from_chunks_stream = {
+                    LLMOutputFormat.JSON: _yield_rows_from_json_chunks_stream,
+                    LLMOutputFormat.CSV: _yield_rows_from_csv_chunks_stream,
+                }[llm_output_format]
+                rows_stream = yield_rows_from_chunks_stream(response)
+            else:
+                # skip roundtrip to LLM in case all columns are provided in existing data
+                assert existing_batch is not None
+
+                async def _yield_empty_rows(n_rows: int) -> AsyncGenerator[dict]:
+                    for _ in range(n_rows):
+                        yield {}
+
+                rows_stream = _yield_empty_rows(len(existing_batch))
+
+            # we first generate all rows in the batch, in order to run consistency checks
+            rows_generated_part = []
+            async for row_generated_part in rows_stream:
+                # remove internal columns, if exist
+                row_generated_part = {k: v for k, v in row_generated_part.items() if k in generated_columns}
+
+                if set(row_generated_part.keys()) != generated_columns:
+                    if context_batch is not None or existing_batch is not None:
+                        # in case of linked tables and data enrichment, it's critical that all rows have expected columns
+                        print(" * Malformed row, repeating batch... * ", end="", flush=True)
+                        do_repeat_task = True
+                        break
+                    else:
+                        # in case of flat tables generation, each row is independent, therefore we only skip the invalid row
+                        continue
+                rows_generated_part.append(row_generated_part)
+
+            # at least some valid rows are expected per batch, repeat the batch otherwise
+            if len(rows_generated_part) == 0:
+                print(" * No valid rows were generated, repeating batch... * ", end="", flush=True)
+                do_repeat_task = True
+
+            # in case of data enrichment, check that all rows were completed successfully
+            if existing_batch is not None and len(rows_generated_part) != len(existing_batch):
+                print(" * Some rows were not enriched successfully, repeating batch... * ", end="", flush=True)
+                do_repeat_task = True
+
+            if do_repeat_task:
+                # allow 10 retries across all workers before propagating the exception to the orchestrator
+                await retry_queue.put(1)
+                if retry_queue.qsize() < 10:
+                    # put task back to the front of the batch queue
+                    await batch_queue.put((batch_idx, task))
+                else:
+                    # inform the orchestrator that max retries were reached
+                    raise RuntimeError(
+                        "Too many malformed batches were generated. "
+                        "Consider changing the model in order to make generation more stable."
+                    )
+
+                # mark current task as done
+                batch_queue.task_done()
+                continue
+
+            # collapse existing and generated parts into coherent rows
+            rows = []
+            for row_idx, row_generated_part in enumerate(rows_generated_part):
+                row_existing_part = existing_batch.iloc[row_idx].to_dict() if existing_batch is not None else {}
+                row = {**row_generated_part, **row_existing_part}
+                # keep columns order according to user's spec
+                row = {column: row[column] for column in columns.keys()}
+                rows.append(row)
+
+            # track previous rows for improved data consistency, in case of sequential generation
+            if n_workers == 1:
+                previous_rows.extend(rows)
+
+            # put rows to the result queue and mark current task as done
+            await result_queue.put((batch_idx, rows))
+            batch_queue.task_done()
+    except Exception as e:
+        # propagate any exception through the result queue
+        await result_queue.put((batch_idx, e))
+        raise
+
+
+async def _create_table_rows_generator(
     *,
     name: str,
     prompt: str,
     columns: dict[str, ColumnConfig],
-    foreign_keys: list[ForeignKeyConfig] | None,
-    primary_keys: dict[str, str] | None,
+    foreign_keys: list[ForeignKeyConfig],
+    primary_keys: dict[str, str],
     data: dict[str, pd.DataFrame],
     sample_size: int,
-    batch_size: int,
     previous_rows_size: int,
     non_context_size: int | None,
+    n_workers: int,
     llm_config: LLMConfig,
-) -> Generator[dict]:
-    def create_table_response_format(
-        columns: dict[str, ColumnConfig], existing_data: pd.DataFrame | None
-    ) -> tuple[type[BaseModel], int]:
-        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 | None,
-                DType.FLOAT: float | None,
-                DType.STRING: str | None,
-                DType.BOOLEAN: bool | None,
-                # response_format has limited support for JSON Schema features
-                # thus we represent dates and datetimes as strings
-                DType.DATE: str | None,
-                DType.DATETIME: str | None,
-            }[column_config.dtype]
-
-        fields = {}
-        for column_name, column_config in columns.items():
-            if existing_data is not None and column_name in existing_data.columns:
-                continue  # skip columns that already exist in existing data
-            annotation = create_annotation(column_config)
-            fields[column_name] = (annotation, Field(...))
-        TableRow = create_model("TableRow", **fields)
-        TableRows = create_model("TableRows", rows=(list[TableRow], ...))
-        n_enforced_columns = len(fields)
-        return TableRows, n_enforced_columns
-
-    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]
+) -> AsyncGenerator[dict]:
+    batch_size = 20  # generate 20 root table rows at a time
 
-    def completion_with_retries(*args, **kwargs):
-        n_attempts = 3
+    def supports_structured_outputs(model: str) -> bool:
+        model = model.removeprefix("litellm_proxy/")
+        supported_params = litellm.get_supported_openai_params(model=model) or []
+        return "response_format" in supported_params and litellm.supports_response_schema(model)
 
-        def print_on_retry(_):
-            print(" * Trying again... * ", end="", flush=True)
+    llm_output_format = LLMOutputFormat.JSON if supports_structured_outputs(llm_config.model) else LLMOutputFormat.CSV
 
-        # try up to 3 times, print a message to the user on each retry
-        retryer = tenacity.Retrying(
-            stop=tenacity.stop_after_attempt(n_attempts), reraise=True, before_sleep=print_on_retry
-        )
-        return retryer(litellm.completion, *args, **kwargs)
-
-    if not llm_config.model.startswith("litellm_proxy/"):
-        # ensure model supports response_format and json schema (this check does not work with litellm_proxy)
-        supported_params = litellm.get_supported_openai_params(model=llm_config.model) or []
-        assert "response_format" in supported_params and litellm.supports_response_schema(llm_config.model), (
-            "The model does not support structured output / JSON mode."
-        )
+    previous_rows = deque(maxlen=previous_rows_size)
 
     # derive data for augmentation
     existing_data: pd.DataFrame | None = None
     if name in data:
         existing_data = data[name]
         sample_size = len(existing_data)
+        batch_size = 10  # augment 10 root table rows at a time
 
     # derive context data (if first foreign key is present) and harmonize sample size accordingly
     context_data: pd.DataFrame | None = None
+    context_batches: list[pd.DataFrame] | None = None
     if foreign_keys and foreign_keys[0].referenced_table != name:  # self-dependency is not considered as context
         context_table_name = foreign_keys[0].referenced_table
         assert context_table_name in data
         context_data = data[context_table_name]
-        batch_size = 1  # generate one sequence at a time
+        batch_size = 1  # generate 1 sequence at a time
         sample_size = len(context_data)
+        context_batches = [data.iloc[i : i + batch_size] for i in range(0, len(data), batch_size)]
 
     # derive non-context data (if more than one foreign key is present)
     non_context_data: dict[str, pd.DataFrame] = {}
@@ -524,18 +761,23 @@ def _create_table_rows_generator(
             assert non_context_table_name in data
             non_context_data[non_context_table_name] = data[non_context_table_name]
 
-    litellm_kwargs = {
-        "temperature": llm_config.temperature,
-        "top_p": llm_config.top_p,
-        "model": llm_config.model,
-        "api_key": llm_config.api_key,
-        "stream": True,
-    }
+    # calculate batch_sizes
+    n_total_batches = len(context_batches) if context_batches is not None else math.ceil(sample_size / batch_size)
+    batch_sizes = [batch_size] * n_total_batches
+    if context_batches is None:
+        # optimise the last batch size for flat tables
+        # +2 because LLM may not always count the rows correctly
+        batch_sizes[-1] = sample_size - sum(batch_sizes[:-1]) + 2
+
+    # initialize queues for async communication
+    batch_queue = asyncio.PriorityQueue()
+    result_queue = asyncio.Queue()
+    retry_queue = asyncio.Queue()
+
+    # populate batch queue
+    for batch_idx in range(n_total_batches):
+        context_batch = context_batches[batch_idx] if context_batches is not None else None
 
-    batch_idx = 0
-    yielded_sequences = 0
-    previous_rows = deque(maxlen=previous_rows_size)
-    for context_batch in batch_infinitely(context_data):
         # pick existing rows for current batch
         existing_batch: pd.DataFrame | None = None
         if existing_data is not None:
@@ -559,71 +801,94 @@ 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
-
-        response_format, n_enforced_columns = create_table_response_format(
-            columns=columns, existing_data=existing_batch
-        )
-
-        llm_prompt = _create_table_prompt(
-            name=name,
-            prompt=prompt,
-            columns=columns,
-            primary_keys=primary_keys,
-            batch_size=batch_size,
-            foreign_keys=foreign_keys,
-            existing_data=existing_batch,
-            context_data=context_batch,
-            non_context_data=non_context_batch,
-            previous_rows=list(previous_rows),
+        task = {
+            "batch_size": batch_sizes[batch_idx],
+            "existing_batch": existing_batch,
+            "context_batch": context_batch,
+            "non_context_batch": non_context_batch,
+        }
+        await batch_queue.put((batch_idx, task))
+
+    # initialize workers
+    n_workers = min(n_total_batches, n_workers)
+    workers = [
+        asyncio.create_task(
+            _worker(
+                name=name,
+                prompt=prompt,
+                columns=columns,
+                foreign_keys=foreign_keys,
+                primary_keys=primary_keys,
+                previous_rows=previous_rows,
+                batch_queue=batch_queue,
+                result_queue=result_queue,
+                retry_queue=retry_queue,
+                n_workers=n_workers,
+                llm_output_format=llm_output_format,
+                llm_config=llm_config,
+            )
         )
-        messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": llm_prompt}]
-
-        if n_enforced_columns != 0:
-            response = completion_with_retries(messages=messages, response_format=response_format, **litellm_kwargs)
-            rows_stream = yield_rows_from_json_chunks_stream(response)
-        else:
-            # skip roundtrip to LLM in case all columns are provided in existing data
-            rows_stream = itertools.repeat({})
-
-        batch_row_idx = 0
-        while True:
-            try:
-                row_generated_part = next(rows_stream)
-                row_existing_part = existing_batch.iloc[batch_row_idx].to_dict() if existing_batch is not None else {}
-                row = {**row_existing_part, **row_generated_part}
-                row = {column: row[column] for column in columns.keys()}  # keep columns order according to user's spec
-            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
-            batch_row_idx += 1
-        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
-
-        batch_idx += 1
-
-
-def _convert_table_rows_generator_to_df(
-    table_rows_generator: Generator[dict],
+        for _ in range(n_workers)
+    ]
+
+    n_completed_batches = 0
+    n_yielded_sequences = 0
+    while n_yielded_sequences < sample_size:
+        if n_completed_batches >= n_total_batches:
+            assert context_data is None, "n_total_batches is fixed for linked tables"
+            assert existing_data is None, "n_total_batches is fixed for data enrichment"
+            # LLMs may not generate exactly the number of rows requested
+            # in case of flat tables, we still accept such incomplete batches,
+            # but that means we may need to generate more batches to reach the sample size
+            # +2 because LLM may not always count the rows correctly
+            n_total_batches += 1
+            task = {
+                "batch_size": sample_size - n_yielded_sequences + 2,
+            }
+            await batch_queue.put((n_total_batches, task))
+        batch_idx, result = await result_queue.get()
+        if isinstance(result, Exception):
+            # if an exception is raised by any worker, cancel all workers and raise that exception
+            for worker in workers:
+                worker.cancel()
+            await asyncio.gather(*workers)
+            raise result
+        rows = result
+        for row_idx, row in enumerate(rows):
+            yield (batch_idx, row)
+            if context_batches is None or row_idx == len(rows) - 1:
+                # in case of flat table, each row is considered a single sequence
+                # in case of linked table, all rows are considered a single sequence
+                # NOTE: this assumes that we generate a single sequence per batch
+                n_yielded_sequences += 1
+            if n_yielded_sequences >= sample_size:
+                break
+        n_completed_batches += 1
+        result_queue.task_done()
+
+    # gracefully shutdown workers
+    await batch_queue.join()
+    for _ in workers:
+        await batch_queue.put((n_total_batches + 1, None))
+    await asyncio.gather(*workers)
+
+
+async def _convert_table_rows_generator_to_df(
+    table_rows_generator: AsyncGenerator[dict],
     columns: dict[str, ColumnConfig],
 ) -> pd.DataFrame:
     def align_df_dtypes_with_mock_dtypes(df: pd.DataFrame, columns: dict[str, ColumnConfig]) -> pd.DataFrame:
+        df = df.copy()
         for column_name, column_config in columns.items():
             if column_config.dtype in [DType.DATE, DType.DATETIME]:
-                df[column_name] = pd.to_datetime(df[column_name], errors="coerce")
+
+                def harmonize_datetime(x):
+                    try:
+                        return dateutil.parser.parse(x)
+                    except Exception:
+                        return pd.NaT
+
+                df[column_name] = pd.to_datetime(df[column_name].apply(harmonize_datetime), errors="coerce")
             elif column_config.dtype is DType.INTEGER:
                 df[column_name] = pd.to_numeric(df[column_name], errors="coerce", downcast="integer").astype(
                     "int64[pyarrow]"
@@ -631,6 +896,8 @@ def _convert_table_rows_generator_to_df(
             elif column_config.dtype is DType.FLOAT:
                 df[column_name] = pd.to_numeric(df[column_name], errors="coerce").astype("double[pyarrow]")
             elif column_config.dtype is DType.BOOLEAN:
+                df[column_name] = df[column_name].map(lambda x: True if str(x).lower() == "true" else x)
+                df[column_name] = df[column_name].map(lambda x: False if str(x).lower() == "false" else x)
                 df[column_name] = pd.to_numeric(df[column_name], errors="coerce").astype("boolean[pyarrow]")
             elif column_config.dtype is DType.CATEGORY:
                 df[column_name] = pd.Categorical(df[column_name], categories=column_config.values)
@@ -638,7 +905,13 @@ def _convert_table_rows_generator_to_df(
                 df[column_name] = df[column_name].astype("string[pyarrow]")
         return df
 
-    df = pd.DataFrame(list(table_rows_generator))
+    # consume entire generator
+    items = [{"batch_idx": batch_idx, "row": row} async for batch_idx, row in table_rows_generator]
+    # sort items by batch_idx to maintain order (relevant especially for keeping the order of existing data)
+    items = sorted(items, key=lambda x: x["batch_idx"])
+    # extract rows and convert to DataFrame
+    rows = [item["row"] for item in items]
+    df = pd.DataFrame(rows)
     df = align_df_dtypes_with_mock_dtypes(df, columns)
     return df
 
@@ -743,10 +1016,11 @@ def sample(
     api_key: str | None = None,
     temperature: float = 1.0,
     top_p: float = 0.95,
+    n_workers: int = 10,
     return_type: Literal["auto", "dict"] = "auto",
 ) -> pd.DataFrame | dict[str, pd.DataFrame]:
     """
-    Generate mock data from scratch or enrich existing data by prompting an LLM.
+    Generate synthetic data from scratch or enrich existing data with new columns.
 
     While faker and numpy are useful to create fake data, this utility is unique as it allows
     the creation of coherent, realistic multi-table tabular mock data
@@ -765,7 +1039,7 @@ def sample(
             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. Model needs to support structured output / JSON mode.
+        model (str): The LiteLLM chat completion model to be used.
             Examples include:
             - `openai/gpt-4.1-nano` (default; fast, and smart)
             - `openai/gpt-4.1-mini` (slower, but smarter)
@@ -779,6 +1053,8 @@ def sample(
         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.
+        n_workers (int): The number of concurrent workers making the LLM calls. Default is 10. The value is clamped to the range [1, 10].
+            If n_workers is 1, the generation of batches becomes sequential and certain features for better data consistency are enabled.
         return_type (Literal["auto", "dict"]): The format of the returned data. Default is "auto".
 
     Returns:
@@ -967,25 +1243,31 @@ def sample(
     sample_size: dict[str, int] = _harmonize_sample_size(sample_size, config)
     primary_keys = {table_name: table_config.primary_key for table_name, table_config in config.root.items()}
 
+    n_workers = max(min(n_workers, 10), 1)
+
     execution_plan: list[str] = _build_execution_plan(config)
 
     data: dict[str, pd.DataFrame] = existing_data or {}
 
     for table_name in execution_plan:
         table_config = config.root[table_name]
-        df = _sample_table(
-            name=table_name,
-            prompt=table_config.prompt,
-            columns=table_config.columns,
-            foreign_keys=table_config.foreign_keys,
-            primary_keys=primary_keys,
-            data=data,
-            sample_size=sample_size[table_name],
-            batch_size=20,  # generate 20 root table rows at a time
-            previous_rows_size=10,  # present 10 previously generated rows to the LLM
-            non_context_size=10,  # pick 10 rows to choose from for each non-context foreign key
-            llm_config=llm_config,
-        )
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            future = executor.submit(
+                _sample_table_sync,
+                name=table_name,
+                prompt=table_config.prompt,
+                columns=table_config.columns,
+                foreign_keys=table_config.foreign_keys,
+                primary_keys=primary_keys,
+                data=data,
+                sample_size=sample_size[table_name],
+                previous_rows_size=10,  # present 10 previously generated rows to the LLM
+                non_context_size=10,  # pick 10 rows to choose from for each non-context foreign key
+                n_workers=n_workers,
+                llm_config=llm_config,
+            )
+            df = future.result()
         data[table_name] = df
 
     return next(iter(data.values())) if len(data) == 1 and return_type == "auto" else data

mostlyai/mock/mcp_server.py

@@ -21,7 +21,9 @@ from fastmcp import FastMCP
 from mostlyai import mock
 
 SAMPLE_MOCK_TOOL_DESCRIPTION = f"""
-Generate mock data by prompting an LLM.
+Synthetic Mock Data.
+
+Use LLMs to generate any Tabular Data towards your needs. Create from scratch, expand existing datasets, or enrich tables with new columns.
 
 This tool is a proxy to the `mostlyai.mock.sample` function, but returns a dictionary of paths to the generated CSV files.
 

{mostlyai_mock-0.1.7.dist-info → mostlyai_mock-0.1.9.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: mostlyai-mock
-Version: 0.1.7
-Summary: LLM-generated Mock Data
+Version: 0.1.9
+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/
@@ -33,19 +33,20 @@ Requires-Dist: pydantic<3.0.0,>=2.0.0
 Requires-Dist: tenacity>=9.1.2
 Description-Content-Type: text/markdown
 
-# LLM-generated Mock Data 🔮
+# 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)
 
-Create data out of nothing. Prompt LLMs for Tabular Data.
+Use LLMs to generate any Tabular Data towards your needs. Create from scratch, expand existing datasets, or enrich tables with new columns. Your prompts, your rules, your data.
 
 ## Key Features
 
-* A light-weight python client for prompting LLMs for mixed-type tabular data
-* Select from a range of LLM endpoints, that provide structured output
+* A light-weight python client for prompting LLMs for mixed-type tabular data.
+* Select from a wide range of LLM endpoints and LLM models.
 * Supports single-table as well as multi-table scenarios.
 * Supports variety of data types: `string`, `categorical`, `integer`, `float`, `boolean`, `date`, and `datetime`.
 * Specify context, distributions and rules via dataset-, table- or column-level prompts.
+* Create from scratch or enrich existing datasets with new columns and/or rows.
 * Tailor the diversity and realism of your generated data via temperature and top_p.
 
 ## Getting Started

mostlyai_mock-0.1.9.dist-info/RECORD

@@ -0,0 +1,8 @@
+mostlyai/mock/__init__.py,sha256=Kh6EZ-dkw2v6SuuWM0ygkKVJmBslJNP2sTsu_CE2JXM,714
+mostlyai/mock/core.py,sha256=Os_nBMrwvB_noIwBMkA0R1z6jzvFcUByGEPoGIyoh78,54577
+mostlyai/mock/mcp_server.py,sha256=MrVUrIsAZsFzjK1suwNl1fxS1ES-wpc-YSM8cS8Fqcw,2259
+mostlyai_mock-0.1.9.dist-info/METADATA,sha256=J59K-sxrbMat0x87P1mObxjUsELCGEBYhS_k-raT_-8,14099
+mostlyai_mock-0.1.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mostlyai_mock-0.1.9.dist-info/entry_points.txt,sha256=XDbppUIAaCWW0nresVep8zb71pkzZuFA16jCBHq8CU8,61
+mostlyai_mock-0.1.9.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mostlyai_mock-0.1.9.dist-info/RECORD,,

mostlyai_mock-0.1.7.dist-info/RECORD

@@ -1,8 +0,0 @@
-mostlyai/mock/__init__.py,sha256=Cmo4Ko8-X41gSewcEpNTTvw7bpRUrtn6B5Cmnwric-Q,714
-mostlyai/mock/core.py,sha256=L-PbOTSIR1cfBeMZL8-v5k7VhxBfKAoyw230soBwQWc,42754
-mostlyai/mock/mcp_server.py,sha256=kWMIjKCwnvYfjY8B2IdP4JNs8ik_8jA6ISCDqrG9utc,2137
-mostlyai_mock-0.1.7.dist-info/METADATA,sha256=6tLpoqLx-LOI-Cr_O_xWm4LI5PBfa4nt1FkrqdNIpQA,13918
-mostlyai_mock-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mostlyai_mock-0.1.7.dist-info/entry_points.txt,sha256=XDbppUIAaCWW0nresVep8zb71pkzZuFA16jCBHq8CU8,61
-mostlyai_mock-0.1.7.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mostlyai_mock-0.1.7.dist-info/RECORD,,