init commit

.dockerignore

@@ -0,0 +1,118 @@
+# Git and version control
+.git
+.gitignore
+.gitattributes
+
+# Development files
+.env*
+!.env.example
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+.venv/
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Node.js
+web/node_modules/
+web/npm-debug.log*
+web/yarn-debug.log*
+web/yarn-error.log*
+web/.pnpm-debug.log*
+web/.next/
+web/out/
+web/dist/
+web/build/
+web/.vercel
+
+# Documentation
+*.md
+!README.md
+docs/
+docs_mintlify/
+vibe_coding/
+
+# Tests
+tests/
+*.test.js
+*.test.ts
+*.test.tsx
+.coverage
+htmlcov/
+pytest.ini
+.pytest_cache/
+
+# Data and outputs
+data_factory_output/
+db/
+*.db
+*.sqlite
+
+# Jupyter notebooks
+*.ipynb
+.ipynb_checkpoints/
+
+# Temporary files
+*.tmp
+*.temp
+.cache/
+
+# Logs
+*.log
+web/api/logs/
+
+# Docker files (except the main ones)
+.dockerignore*
+Dockerfile.*
+docker-compose*.yml
+
+# Development and internal files
+internal/
+examples/
+mcp_hackathon/
+prebuilt_template/
+scripts/
+htmlcov/
+
+# Poetry (we copy these explicitly)
+# poetry.lock - we need this
+# pyproject.toml - we need this 

.env.template

@@ -0,0 +1,17 @@
+# Starfish Environment Variables
+# Copy this file to .env and customize for your local environment
+# DO NOT commit the .env file to version control
+
+# Environment type (DEV, STAGING, PROD)
+ENV=DEV
+
+# API Keys (replace with your own)
+OPENAI_API_KEY=your_openai_api_key_here
+ANTHROPIC_API_KEY=your_anthropic_api_key_here
+HUGGING_FACE_HUB_TOKEN=your_huggingface_token_here
+TELEMETRY_ENABLED=true
+
+# Logging
+LOG_LEVEL=INFO 
+# STARFISH_LOCAL_STORAGE_DIR=
+JINA_AI_API_KEY=jina_api_key

.github/dependabot.yml

@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/" # root of the repo
+    schedule:
+      interval: "weekly"

.github/workflows/lint-and-test.yaml

@@ -0,0 +1,60 @@
+name: Starfish testing workflow
+
+on:
+  # push:
+  #   branches:
+  #     - main
+  #     - dev
+  pull_request:
+    branches:
+      - main
+      - dev
+      - '!f/pypi_release'
+
+jobs:
+  test-integration:
+    if: github.event.pull_request.head.ref != 'f/pypi_release'
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.11'
+
+      - name: Load cached Poetry installation
+        uses: actions/cache@v3
+        with:
+          path: ~/.local
+          key: poetry-${{ runner.os }}-${{ hashFiles('**/poetry.lock') }}
+
+      - name: Load cached venv
+        uses: actions/cache@v3
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-python-${{ matrix.python-version }}-${{ hashFiles('**/poetry.lock') }}
+
+      - name: Set Locale
+        run: |
+          sudo locale-gen "en_US.UTF-8"
+          export LC_ALL=en_US.UTF-8
+          export LANG=en_US.UTF-8
+          export TELEMETRY_ENABLED=false
+
+      - name: Install dependencies
+        run: |
+          pip install poetry
+          poetry install --with dev
+
+      # - name: Run ruff
+      #   run: |
+      #     poetry run ruff check . --output-format=github
+      #     poetry run ruff format . --check
+      
+      # --cov-report=html
+      - name: Run tests with coverage
+        run: |
+          poetry run pytest --cov='src' --cov-fail-under=20  tests/

.github/workflows/publish_pypi.yaml

@@ -0,0 +1,39 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - 'v*'
+    # branches:
+    #   - 'main'
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - name: Verify tag is on main branch
+      run: |
+        TAG_NAME=${GITHUB_REF#refs/tags/}
+        COMMIT=$(git rev-parse $TAG_NAME)
+        if ! git branch --contains $COMMIT | grep -qw main; then
+          echo "::error::Tag $TAG_NAME must be created from main branch"
+          exit 1
+        fi
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build twine
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+      run: |
+        python -m build
+        twine upload dist/*

.github/workflows/publish_testpypi.yaml

@@ -0,0 +1,107 @@
+name: Publish to Test PyPI
+
+on:
+  push:
+    tags:
+      - 'test-v*'
+    branches:
+      - 'f/pypi_release'
+
+jobs:
+  deploy_testpypi:
+    #if: true
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0  # Required for full commit history check
+    - name: Verify tag is on dev branch
+      run: |
+        TAG_NAME=${GITHUB_REF#refs/tags/}
+        COMMIT=$(git rev-parse $TAG_NAME)
+        if ! git branch --contains $COMMIT | grep -qw dev; then
+          echo "::error::Tag $TAG_NAME must be created from dev branch"
+          exit 1
+        fi
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build twine
+    - name: Build and publish
+      env:
+        #TWINE_USERNAME: ${{ secrets.TEST_PYPI_USERNAME }}
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.TEST_PYPI_PASSWORD }}
+        #ACTIONS_STEP_DEBUG: true
+      run: |
+        # echo "TWINE_PASSWORD first 5 chars: ${TWINE_PASSWORD:0:184}"
+        # echo "TWINE_PASSWORD length: ${#TWINE_PASSWORD}"
+        python -m build
+        twine upload --verbose --repository-url https://test.pypi.org/legacy/ dist/*
+
+  test-colab:
+    needs: deploy_testpypi
+    runs-on: ubuntu-latest
+    #a Public "Colab-like" Image
+    container:
+      image: jupyter/minimal-notebook:latest
+      options: --user root  # Run as root to avoid permission issues
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          sparse-checkout: |
+            tests/*
+            examples/data_factory_release_check.ipynb
+          sparse-checkout-cone-mode: false
+      - name: Update system packages
+        run: |
+          apt-get update
+          apt-get install -y libssl3  # Removed sudo since we're running as root
+      - name: Print Python and Jupyter versions
+        run: |
+          python --version
+          pip list | grep -E 'jupyter|ipykernel|nbconvert|notebook'
+      # Authenticate to GCP
+      # - name: Authenticate to GCP
+      #   uses: google-github-actions/auth@v1
+      #   with:
+      #     credentials_json: ${{ secrets.GCP_SA_KEY }}
+
+      # # Configure Docker to use GCR credentials
+      # - name: Configure Docker for GCR
+      #   uses: google-github-actions/docker-auth@v1
+
+      # # Now you can pull the image
+      # - name: Use Colab base image
+      #   run: docker pull gcr.io/colab-images/base:latest
+
+      # --no-prompt --no-input \ suppress the output
+      - name: Run Colab-style tests
+        run: |
+          if ! jupyter nbconvert --execute --to notebook --inplace \
+            --ExecutePreprocessor.kernel_name=python3 \
+            --ExecutePreprocessor.timeout=120 \
+            --no-prompt --no-input \
+            --stdout \
+            examples/data_factory_release_check.ipynb; then
+            echo "::error::Notebook execution failed"
+            exit 1
+          fi
+          echo "Notebook executed successfully. Summary:" && \
+          jupyter nbconvert --to markdown --stdout \
+            examples/data_factory_release_check.ipynb | \
+            grep -E '^#|^##' || true
+
+      # Add tag deletion step
+      - name: Delete triggering tag after successful test
+        if: startsWith(github.ref, 'refs/tags/test-v')
+        run: |
+          gh api -X DELETE /repos/$GITHUB_REPOSITORY/git/refs/tags/${GITHUB_REF#refs/tags/}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.gitignore

@@ -0,0 +1,190 @@
+# Adhoc stuff
+web/node_modules/
+web/.next/
+web/public/
+web/dist/
+web/build/
+web/out/
+web/coverage/
+web/logs/
+web/.local/
+web/.env
+
+.serena/
+docs/
+/vibe_coding/response.md
+/dev/
+todo
+.local/
+.vscode/
+db/
+.ruff_cache/
+data_factory_output/
+examples/test_jupyter.ipynb
+# *.ipynb
+# .ipynb_checkpoints
+.cursor
+
+.DS_Store
+*/.DS_Store
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

.gitmodules

@@ -0,0 +1,6 @@
+[submodule "internal"]
+	path = internal
+	url = https://github.com/starfishdata/starfish_internal.git
+[submodule "docs_mintlify"]
+	path = docs_mintlify
+	url = https://github.com/starfishdata/docs.git

.pre-commit-config.yaml

@@ -0,0 +1,25 @@
+repos:
+  # - repo: local
+  #   hooks:
+  #     - id: pytest
+  #       name: Run pytest
+  #       entry: poetry run pytest tests/
+  #       language: system
+  #       types: [python]
+  #       pass_filenames: false
+  #       always_run: true 
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.8.6
+    hooks:
+      # Run the linter.
+      # - id: ruff
+      #   #args: [ --fix ]
+      #   types: [python] 
+      # Run the formatter.
+      - id: ruff-format
+        #args: [ --fix ]
+        #run even when no Python files are staged
+        #always_run: true
+        types: [python]

Dockerfile

@@ -0,0 +1,79 @@
+# Multi-stage build for combined frontend + backend
+FROM node:18-alpine AS frontend-builder
+
+WORKDIR /app
+
+# Copy package files
+COPY web/package*.json ./
+
+# Install dependencies
+RUN npm ci
+
+# Copy frontend code and build
+COPY web/ ./
+
+# Clean up unnecessary files
+RUN rm -rf api/ || true
+RUN rm -rf storage/ || true
+RUN rm -rf .git/ || true
+RUN rm -rf .next/ || true
+RUN rm -rf .local/ || true
+
+# Build frontend
+RUN npm run build
+
+# Backend stage
+FROM python:3.11-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    nginx \
+    supervisor \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Node.js for combined container
+RUN curl -fsSL https://deb.nodesource.com/setup_18.x | bash - && \
+    apt-get install -y nodejs
+
+WORKDIR /app
+
+# Copy pyproject.toml and poetry.lock
+COPY pyproject.toml poetry.lock ./
+
+# Install Poetry and basic dependencies (skip heavy ML packages for testing)
+RUN pip install --no-cache-dir --upgrade pip \
+    && pip install --no-cache-dir poetry \
+    && poetry config virtualenvs.create false \
+    && poetry install --only=main --no-root || pip install fastapi uvicorn python-dotenv pydantic
+
+# Copy starfish source code and README (needed by backend)
+COPY src/ ./src/
+COPY README.md ./
+
+# Copy built frontend from previous stage
+COPY --from=frontend-builder /app/.next ./web/.next
+COPY --from=frontend-builder /app/public ./web/public
+COPY --from=frontend-builder /app/package.json ./web/package.json
+#COPY --from=frontend-builder /app/node_modules ./web/node_modules
+
+# Copy backend API code
+COPY web/api/ ./web/api/
+
+# Copy configuration files
+COPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf
+COPY nginx.conf /etc/nginx/nginx.conf
+
+# Create necessary directories and set permissions
+RUN mkdir -p /var/log/supervisor /var/log/nginx /var/run \
+    && chmod +x /app/src/ || true
+
+# Expose port 7860 (required for Hugging Face Spaces)
+EXPOSE 7860
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
+    CMD curl -f http://localhost:7860/health || exit 1
+
+# Start supervisor which manages both nginx and the applications
+CMD ["/usr/bin/supervisord", "-c", "/etc/supervisor/conf.d/supervisord.conf"] 

LICENSE

@@ -0,0 +1,202 @@
+
+                                 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 2025 Starfish AI Inc.
+
+   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.

Makefile

@@ -0,0 +1,28 @@
+lint: 
+	@echo "Running Linter (Ruff)..."
+	poetry run isort tests/ src/ examples --check-only || poetry run isort tests/ src/ examples
+	poetry run ruff check src examples --fix --unsafe-fixes --exit-zero
+	poetry run ruff format src examples --check || poetry run ruff format src examples
+docstring:
+	ruff check --select D src/starfish/data_factory
+test:
+	poetry run pytest tests/	
+
+install: install-extras
+
+#poetry install --extras "code_execution vllm" --with dev
+# Install with specific extras
+#make install EXTRAS="pdf"
+# Install all extras
+#make install EXTRAS="all"
+# Install without extras (default)
+#make install
+install-extras:
+	@echo "Installing dependencies with extras: $(EXTRAS)"
+	poetry install $(if $(EXTRAS),--extras "$(EXTRAS)",) --with dev
+
+start-client_claude:
+	python src/starfish/data_mcp/client_claude.py  src/starfish/data_mcp/server.py
+
+start-client_openai:
+	python src/starfish/data_mcp/client_openai.py 

README.HuggingFace.md

@@ -0,0 +1,177 @@
+# Hugging Face Spaces Deployment
+
+This guide explains how to deploy your combined FastAPI backend and Next.js frontend to Hugging Face Spaces.
+
+> ✅ **Build Status**: Docker build is working successfully with resolved path alias issues!
+
+## Overview
+
+The `Dockerfile.huggingface` creates a single container that runs:
+- **FastAPI backend** on port 8002
+- **Next.js frontend** on port 3000  
+- **Nginx reverse proxy** on port 7860 (required by Hugging Face Spaces)
+- **Supervisor** to manage all processes
+
+## Files for Hugging Face Spaces
+
+1. **`Dockerfile`** - Combined Dockerfile for both services (multi-stage build)
+2. **`nginx.conf`** - Nginx configuration for routing
+3. **`supervisord.conf`** - Process manager configuration
+4. **`.dockerignore`** - Optimized to exclude only necessary files
+5. **`next.config.js`** - Enhanced with webpack path alias configuration
+6. **`tsconfig.json`** - Updated with explicit path mappings
+
+## Deployment Steps
+
+### 1. Prepare Your Repository
+
+Your repository is already configured with the correct `Dockerfile` for Hugging Face Spaces deployment.
+
+### 2. Set Environment Variables in Hugging Face Spaces
+
+In your Hugging Face Space settings, add these secrets:
+- `GOOGLE_API_KEY` - Your Google API key
+- `OPENAI_API_KEY` - Your OpenAI API key
+
+### 3. Configure Your Space
+
+- **Space Type**: Docker
+- **Visibility**: Public or Private (your choice)
+- **Hardware**: CPU Basic (or upgrade if needed)
+
+### 4. Update API URLs in Frontend
+
+Make sure your frontend points to the correct API endpoints:
+```typescript
+// In your frontend code, use relative URLs:
+const API_BASE_URL = "/api"  // This goes to Next.js API routes in src/app/api/
+
+// Next.js API routes will then proxy to FastAPI using:
+// SERVER_BASE_URL=http://localhost:8002 (set in Dockerfile)
+```
+
+### 5. Deploy
+
+1. Push your code to the Hugging Face Space repository
+2. The space will automatically build and deploy
+
+## How It Works
+
+### Architecture
+```
+External Request :7860
+        ↓
+    Nginx Proxy
+        ↓
+    Next.js :3000 (handles ALL routes)
+        ↓
+    /api/* → src/app/api/ routes
+        ↓
+    proxy.ts uses SERVER_BASE_URL
+        ↓
+    FastAPI Backend :8002
+```
+
+### Port Mapping
+- **7860** - Main port (required by Hugging Face Spaces)
+- **3000** - Next.js frontend (internal) - handles all routing
+- **8002** - FastAPI backend (internal) - accessed via Next.js proxy
+
+### URL Routing
+- `/` - Next.js frontend (all routes handled by Next.js)
+- `/api/*` - Next.js API routes (in `src/app/api/`) that proxy to FastAPI backend
+- `/backend-docs` - Direct FastAPI documentation (for debugging)
+- `/backend-openapi.json` - Direct FastAPI OpenAPI schema (for debugging)
+
+### Process Management
+Supervisor manages three processes:
+1. **backend** - FastAPI server (port 8002)
+2. **frontend** - Next.js server (port 3000) - handles all routing and proxying
+3. **nginx** - Reverse proxy (port 7860) - routes all traffic to Next.js
+
+## Troubleshooting
+
+### Common Issues
+
+1. **Build fails with "Module not found: Can't resolve '@/lib/utils'"**
+   - **FIXED**: This was caused by `lib/` being excluded in `.dockerignore`
+   - The issue has been resolved by removing the `lib/` exclusion pattern
+
+2. **Build fails during npm install**
+   - Check that all package.json dependencies are valid
+   - Ensure Node.js version compatibility
+
+3. **FastAPI fails to start**
+   - Check environment variables are set
+   - Verify the starfish package is properly configured
+   - Check logs in the Space's logs tab
+
+4. **Frontend can't reach backend**
+   - Ensure API calls use relative URLs (`/api/...`)
+   - Check that `SERVER_BASE_URL=http://localhost:8002` is set in the Dockerfile
+   - Verify Next.js API routes in `src/app/api/` are proxying correctly
+   - For direct FastAPI access, use `/backend-docs` instead of `/docs`
+
+5. **Space shows "Application starting" indefinitely**
+   - Check supervisor logs for errors
+   - Verify all services are starting properly
+
+### Viewing Logs
+
+In your Hugging Face Space:
+1. Go to the "Logs" tab
+2. Look for errors from supervisor, nginx, backend, or frontend
+3. Logs are also written to `/var/log/` in the container
+
+### Local Testing
+
+Test the Hugging Face build locally:
+```bash
+# Build the image
+docker build -t starfishai-web  .
+
+# Run with environment variables
+docker run -p 7860:7860 3000:3000 8002:8002\
+  -e GOOGLE_API_KEY=your_key \
+  -e OPENAI_API_KEY=your_key \
+  starfishai-web 
+```
+
+Then visit:
+- http://localhost:7860 - Main application
+- http://localhost:7860/backend-docs - Direct FastAPI documentation
+- http://localhost:7860/backend-openapi.json - Direct FastAPI schema
+
+## Recent Fixes & Improvements
+
+### Path Alias Resolution Fixed
+- **Issue**: Build was failing with `Module not found: Can't resolve '@/lib/utils'`
+- **Root Cause**: The `.dockerignore` file was excluding the `lib/` directory
+- **Solution**: Removed `lib/` from `.dockerignore` and enhanced path configuration
+- **Files Updated**: 
+  - `.dockerignore` - Removed generic `lib/` exclusion
+  - `next.config.js` - Added explicit webpack path aliases
+  - `tsconfig.json` - Enhanced path mappings
+
+### Docker Build Optimization
+- **Multi-stage build** for optimal image size
+- **Specific Python exclusions** in `.dockerignore` (e.g., `api/__pycache__/` instead of all `__pycache__/`)
+- **Enhanced file copying strategy** during build
+
+## Performance Tips
+
+1. **Use CPU Basic** for development, upgrade for production
+2. **Optimize Docker image** by removing unnecessary files
+3. **Use caching** for build dependencies
+4. **Monitor resource usage** in the Space dashboard
+
+## Security Notes
+
+- Never commit API keys to your repository
+- Use Hugging Face Spaces secrets for sensitive environment variables
+- Consider making your Space private if it contains sensitive data
+- Regularly update dependencies for security patches 
+
+docker run -d -p 7860:7860 --name starfish-app -v $(pwd)/nginx.conf:/etc/nginx/nginx.conf -v $(pwd)/supervisord.conf:/etc/supervisor/conf.d/supervisord.conf starfish-app
+
+docker build -t starfish-app . 

README.md

@@ -0,0 +1,193 @@
+---
+title: Starfish - Synthetic Data Generation
+emoji: 🌟
+colorFrom: pink
+colorTo: blue
+sdk: docker
+sdk_version: "4.36.0"
+---
+
+<p align="center">
+  <img src="https://github.com/user-attachments/assets/744c666a-bb5c-418b-aab4-162072c0b8c8" alt="Starfish Logo" width="200"/>
+</p>
+<h1 align="center">Starfish</h1>
+<h3 align="center" style="font-size: 20px; margin-bottom: 4px">Synthetic Data Generation Made Easy</h2>
+</br>
+
+<div align="center">
+
+[![Github](https://img.shields.io/badge/starfish-black?style=for-the-badge&logo=github&color=black)](https://github.com/starfishdata/starfish) [![X](https://img.shields.io/badge/starfishdata-black?style=for-the-badge&logo=x&color=black&link=https%3A%2F%2Fx.com%2Fstarfishdata)](https://x.com/starfishdata) [![Hugging Face](https://img.shields.io/badge/starfishdata-yellow?style=for-the-badge&logo=huggingface&labelColor=black&color=black)](https://huggingface.co/starfishdata) [![Discord](https://img.shields.io/badge/starfishdata-yellow?style=for-the-badge&logo=discord&logoColor=white&labelColor=%235865F2&color=%235865F2)](https://discord.gg/qWKmeUtb)
+<br>
+[![Website](https://img.shields.io/badge/starfishdata-yellow?style=for-the-badge&label=SITE&labelColor=%23DB2777&color=%23FDF2F8)](https://starfishdata.ai/) 
+[![Docs](https://img.shields.io/badge/docs-pink?style=for-the-badge&label=Deepwiki&labelColor=%23da2876&color=%23fdf2f8&link=https%3A%2F%2Fdeepwiki.com%2Fstarfishdata%2Fstarfish%2F1-overview
+)](https://deepwiki.com/starfishdata/starfish/1-overview) 
+</div>
+
+## Overview
+
+Starfish is a Python library that helps you build synthetic data your way. We adapt to your workflow—not the other way around. By combining structured LLM outputs with efficient parallel processing, Starfish lets you define exactly how your data should look and scale seamlessly from experiments to production.
+   
+⭐ Star us on GitHub if you find this project useful!
+
+Key Features:
+- **Structured Outputs**: First-class support for structured data through JSON schemas or Pydantic models.
+- **Model Flexibility**: Use any LLM provider—local models, OpenAI, Anthropic, or your own implementation via LiteLLM.
+- **Dynamic Prompts**: Dynamic prompts with built-in Jinja2 templates.
+- **Easy Scaling**: Transform any function to run in parallel across thousands of inputs with a single decorator.
+- **Resilient Pipeline**: Automatic retries, error handling, and job resumption—pause and continue your data generation anytime.
+- **Complete Control**: Share state across your pipeline, extend functionality with custom hooks.
+
+**Official Website**: [starfishdata.ai](https://starfishdata.ai/) - We offer both self-service and managed solutions. Visit our website to explore our services or contact us for more options!
+
+## Installation
+
+```bash
+pip install starfish-core
+```
+
+### Optional Dependencies
+
+Starfish supports optional dependencies for specific file parsers. Install only what you need:
+
+```bash
+# Install specific parsers
+pip install "starfish-core[pdf]"       # PDF support
+pip install "starfish-core[docx]"      # Word document support
+pip install "starfish-core[ppt]"       # PowerPoint support
+pip install "starfish-core[excel]"     # Excel support
+pip install "starfish-core[youtube]"   # YouTube support
+
+# Install all parser dependencies
+pip install "starfish-core[all]"
+```
+
+## Configuration
+
+Starfish uses environment variables for configuration. We provide a `.env.template` file to help you get started quickly:
+
+```bash
+# Copy the template to .env
+cp .env.template .env
+
+# Edit with your API keys and configuration
+nano .env  # or use your preferred editor
+```
+
+The template includes settings for API keys, model configurations, and other runtime parameters.
+
+## Quick Start
+
+### Structured LLM - Type-Safe Outputs from Any Model
+
+```python
+# 1. Define structured outputs with schema
+from starfish import StructuredLLM
+from pydantic import BaseModel
+
+# Option A: Use Pydantic for type safety
+class QnASchema(BaseModel):
+    question: str
+    answer: str
+
+# Option B: Or use simple JSON schema
+json_schema = [
+    {'name': 'question', 'type': 'str'},
+    {'name': 'answer', 'type': 'str'}, 
+]
+
+# 2. Create a structured LLM with your preferred output format
+qna_llm = StructuredLLM(
+    model_name="openai/gpt-4o-mini",
+    prompt="Generate facts about {{city}}",
+    output_schema=QnASchema  # or json_schema
+)
+
+# 3. Get structured responses
+response = await qna_llm.run(city="San Francisco")
+
+# Access typed data
+print(response.data)
+# [{'question': 'What is the iconic symbol of San Francisco?',
+#   'answer': 'The Golden Gate Bridge is the iconic symbol of San Francisco, completed in 1937.'}]
+
+# Access raw API response for complete flexibility
+print(response.raw)  # Full API object with function calls, reasoning tokens, etc.
+```
+
+### Data Factory - Scale Any Workflow with One Decorator
+
+```python
+# Turn any function into a scalable data pipeline
+from starfish import data_factory
+
+# Works with any function - simple or complex workflows
+@data_factory(max_concurrency=50)
+async def parallel_qna_llm(city):
+    # This could be any arbitrary complex workflow:
+    # - Pre-processing
+    # - Multiple LLM calls
+    # - Post-processing
+    # - Error handling
+    response = await qna_llm.run(city=city)
+    return response.data
+
+# Process 100 cities with 50 concurrent workers - finishes in seconds
+cities = ["San Francisco", "New York", "Tokyo", "Paris", "London"] * 20
+results = parallel_qna_llm.run(city=cities)
+
+# dry run to test the workflow and data
+results = parallel_qna_llm.dry_run(city=cities)
+    
+# resume job which pick up from where it left off. 
+results = parallel_qna_llm.resume()
+```
+
+### Examples
+
+Check out our example notebooks for detailed walkthroughs:
+- [Structured LLM Examples](examples/structured_llm.ipynb)
+- [Data Factory Examples](examples/data_factory.ipynb)
+
+## Documentation
+
+Comprehensive documentation is on the way!
+
+## Contributing
+
+We'd love your help making Starfish better! Whether you're fixing bugs, adding features, or improving documentation, your contributions are welcome.
+
+1. Fork the repository
+2. Create your feature branch (`git checkout -b feature/amazing-feature`)
+3. Commit your changes (`git commit -m 'Add some amazing feature'`)
+4. Push to the branch (`git push origin feature/amazing-feature`)
+5. Open a Pull Request
+
+Contribution guidelines coming soon!
+
+## License
+
+This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details.
+
+## Contact
+
+If you have any questions or feedback, feel free to reach out to us at [[email protected]](mailto:[email protected]).
+
+Want to discuss your use case directly? [Schedule a meeting with our team](https://calendly.com/d/crsb-ckq-fv2/chat-with-starfishdata-team).
+
+## Telemetry
+
+Starfish collects minimal and anonymous telemetry data to help improve the library. Participation is optional and you can opt out by setting `TELEMETRY_ENABLED=false` in your environment variables.
+
+## Citation
+
+If you use Starfish in your research, please consider citing us!
+
+```
+@software{starfish,
+  author = {Wendao, John, Ayush},
+  title = {{Starfish: A Tool for Synthetic Data Generation}},
+  year = {2025},
+  url = {https://github.com/starfishdata/starfish},
+}
+```
+

docs_mintlify

@@ -0,0 +1 @@
+Subproject commit 6ad0ad5eda1fc3637fde8d0da24f0d3fd4263453

examples/data_factory.ipynb

@@ -0,0 +1,681 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Google Colab Version: [Open this notebook in Google Colab](https://colab.research.google.com/github/starfishdata/starfish/blob/main/examples/data_factory.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Dependencies "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: starfish-core in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (0.1.0)\n",
+      "Requirement already satisfied: aiofiles<25.0.0,>=24.1.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (24.1.0)\n",
+      "Requirement already satisfied: aiosqlite<0.22.0,>=0.21.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (0.21.0)\n",
+      "Requirement already satisfied: cachetools<6.0.0,>=5.5.2 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (5.5.2)\n",
+      "Requirement already satisfied: litellm<2.0.0,>=1.65.1 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (1.65.1)\n",
+      "Requirement already satisfied: loguru<0.8.0,>=0.7.3 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (0.7.3)\n",
+      "Requirement already satisfied: ollama<0.5.0,>=0.4.7 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (0.4.7)\n",
+      "Requirement already satisfied: platformdirs<5.0.0,>=4.1.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (4.3.7)\n",
+      "Requirement already satisfied: psutil<8.0.0,>=7.0.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (7.0.0)\n",
+      "Requirement already satisfied: python-dotenv<2.0.0,>=1.1.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (1.1.0)\n",
+      "Requirement already satisfied: typing-extensions<5.0.0,>=4.0.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from starfish-core) (4.13.0)\n",
+      "Requirement already satisfied: aiohttp in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (3.11.16)\n",
+      "Requirement already satisfied: click in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (8.1.8)\n",
+      "Requirement already satisfied: httpx>=0.23.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (0.28.1)\n",
+      "Requirement already satisfied: importlib-metadata>=6.8.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (8.6.1)\n",
+      "Requirement already satisfied: jinja2<4.0.0,>=3.1.2 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (3.1.6)\n",
+      "Requirement already satisfied: jsonschema<5.0.0,>=4.22.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (4.23.0)\n",
+      "Requirement already satisfied: openai>=1.68.2 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (1.70.0)\n",
+      "Requirement already satisfied: pydantic<3.0.0,>=2.0.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (2.11.1)\n",
+      "Requirement already satisfied: tiktoken>=0.7.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (0.9.0)\n",
+      "Requirement already satisfied: tokenizers in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (0.21.1)\n",
+      "Requirement already satisfied: anyio in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (4.9.0)\n",
+      "Requirement already satisfied: certifi in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (2025.1.31)\n",
+      "Requirement already satisfied: httpcore==1.* in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (1.0.7)\n",
+      "Requirement already satisfied: idna in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (3.10)\n",
+      "Requirement already satisfied: h11<0.15,>=0.13 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from httpcore==1.*->httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.14.0)\n",
+      "Requirement already satisfied: zipp>=3.20 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from importlib-metadata>=6.8.0->litellm<2.0.0,>=1.65.1->starfish-core) (3.21.0)\n",
+      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from jinja2<4.0.0,>=3.1.2->litellm<2.0.0,>=1.65.1->starfish-core) (3.0.2)\n",
+      "Requirement already satisfied: attrs>=22.2.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (25.3.0)\n",
+      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (2024.10.1)\n",
+      "Requirement already satisfied: referencing>=0.28.4 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.36.2)\n",
+      "Requirement already satisfied: rpds-py>=0.7.1 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.24.0)\n",
+      "Requirement already satisfied: distro<2,>=1.7.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from openai>=1.68.2->litellm<2.0.0,>=1.65.1->starfish-core) (1.9.0)\n",
+      "Requirement already satisfied: jiter<1,>=0.4.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from openai>=1.68.2->litellm<2.0.0,>=1.65.1->starfish-core) (0.9.0)\n",
+      "Requirement already satisfied: sniffio in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from openai>=1.68.2->litellm<2.0.0,>=1.65.1->starfish-core) (1.3.1)\n",
+      "Requirement already satisfied: tqdm>4 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from openai>=1.68.2->litellm<2.0.0,>=1.65.1->starfish-core) (4.67.1)\n",
+      "Requirement already satisfied: annotated-types>=0.6.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from pydantic<3.0.0,>=2.0.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.7.0)\n",
+      "Requirement already satisfied: pydantic-core==2.33.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from pydantic<3.0.0,>=2.0.0->litellm<2.0.0,>=1.65.1->starfish-core) (2.33.0)\n",
+      "Requirement already satisfied: typing-inspection>=0.4.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from pydantic<3.0.0,>=2.0.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.4.0)\n",
+      "Requirement already satisfied: regex>=2022.1.18 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from tiktoken>=0.7.0->litellm<2.0.0,>=1.65.1->starfish-core) (2024.11.6)\n",
+      "Requirement already satisfied: requests>=2.26.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from tiktoken>=0.7.0->litellm<2.0.0,>=1.65.1->starfish-core) (2.32.3)\n",
+      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (2.6.1)\n",
+      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (1.3.2)\n",
+      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (1.5.0)\n",
+      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (6.3.1)\n",
+      "Requirement already satisfied: propcache>=0.2.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (0.3.1)\n",
+      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (1.18.3)\n",
+      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (0.30.1)\n",
+      "Requirement already satisfied: filelock in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (3.18.0)\n",
+      "Requirement already satisfied: fsspec>=2023.5.0 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (2025.3.2)\n",
+      "Requirement already satisfied: packaging>=20.9 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (24.2)\n",
+      "Requirement already satisfied: pyyaml>=5.1 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (6.0.2)\n",
+      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from requests>=2.26.0->tiktoken>=0.7.0->litellm<2.0.0,>=1.65.1->starfish-core) (3.4.1)\n",
+      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/zhengisamazing/Library/Caches/pypoetry/virtualenvs/starfish-T7IInzTH-py3.11/lib/python3.11/site-packages (from requests>=2.26.0->tiktoken>=0.7.0->litellm<2.0.0,>=1.65.1->starfish-core) (2.3.0)\n",
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%pip install starfish-core"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Fix for Jupyter Notebook only — do NOT use in production\n",
+    "## Enables async code execution in notebooks, but may cause issues with sync/async issues\n",
+    "## For production, please run in standard .py files without this workaround\n",
+    "## See: https://github.com/erdewit/nest_asyncio for more details\n",
+    "import nest_asyncio\n",
+    "nest_asyncio.apply()\n",
+    "\n",
+    "from starfish import StructuredLLM, data_factory\n",
+    "from starfish.llm.utils import merge_structured_outputs\n",
+    "\n",
+    "from starfish.common.env_loader import load_env_file ## Load environment variables from .env file\n",
+    "load_env_file()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# setup your openai api key if not already set\n",
+    "# import os\n",
+    "# os.environ[\"OPENAI_API_KEY\"] = \"your_key_here\"\n",
+    "\n",
+    "# If you dont have any API key, please navigate to local model section"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Helper function mock llm call\n",
+    "# When developing data pipelines with LLMs, making thousands of real API calls\n",
+    "# can be expensive. Using mock LLM calls lets you test your pipeline's reliability,\n",
+    "# failure handling, and recovery without spending money on API calls.\n",
+    "from starfish.data_factory.utils.mock import mock_llm_call"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 1. Your First Data Factory: Simple Scaling\n",
+    "\n",
+    "The @data_factory decorator transforms any async function into a scalable data processing pipeline.\n",
+    "It handles:\n",
+    "- Parallel execution \n",
+    "- Automatic batching\n",
+    "- Error handling & retries\n",
+    "- Progress tracking\n",
+    "\n",
+    "Let's start with a single LLM call and then show how easy it is to scale it.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'fact': 'New Yorkers consume around 1,000,000 slices of pizza every day, which means if you laid them all in a line, they would stretch from the Statue of Liberty to the Eiffel Tower... and back!'}]"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# First, create a StructuredLLM instance for generating facts about cities\n",
+    "json_llm = StructuredLLM(\n",
+    "    model_name = \"openai/gpt-4o-mini\",\n",
+    "    prompt = \"Funny facts about city {{city_name}}.\",\n",
+    "    output_schema = [{'name': 'fact', 'type': 'str'}],\n",
+    "    model_kwargs = {\"temperature\": 0.7},\n",
+    ")\n",
+    "\n",
+    "json_llm_response = await json_llm.run(city_name='New York')\n",
+    "json_llm_response.data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:16:32\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 8c926411-63e7-4dc6-98c9-861c3489fb8b\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:32\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "Processing New York at 2025-04-25 10:16:32.524033\n",
+      "Processing London at 2025-04-25 10:16:32.524286\n",
+      "Processing Tokyo at 2025-04-25 10:16:32.524979\n",
+      "Processing Paris at 2025-04-25 10:16:32.525535\n",
+      "Processing Sydney at 2025-04-25 10:16:32.526729\n",
+      "\u001b[32m2025-04-25 10:16:34\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:34\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[{'fact': \"In Tokyo, there's a train station called 'Shinjuku' that handles more passengers each day than the entire population of the United States!\"},\n",
+       " {'fact': \"London has a 'secret' underground city known as the 'London Stone', which is said to have magical powers, making it one of the city's most famous and quirky legends!\"},\n",
+       " {'fact': 'In Paris, you can legally marry a dead person! This quirky law allows for posthumous marriages, as long as you can prove that the deceased had intended to marry you before their untimely demise.'},\n",
+       " {'fact': 'In New York City, there are more than 25,000 licensed taxis, but only about 1,200 of them are actually yellow. The rest are a rainbow of colors, including pink, blue, and even animal print!'},\n",
+       " {'fact': 'Sydney has a beach where you can surf, swim, and even watch a film – all in one day! Just don’t forget your sunscreen and popcorn!'}]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Now, scale to multiple cities using data_factory\n",
+    "# Just add the @data_factory decorator to process many cities in parallel\n",
+    "\n",
+    "from datetime import datetime\n",
+    "@data_factory(max_concurrency=10)\n",
+    "async def process_json_llm(city_name: str):\n",
+    "    ## Adding a print statement to indicate the start of the processing\n",
+    "    print(f\"Processing {city_name} at {datetime.now()}\")\n",
+    "    json_llm_response = await json_llm.run(city_name=city_name)\n",
+    "    return json_llm_response.data\n",
+    "\n",
+    "# This is all it takes to scale from one city to many cities!\n",
+    "process_json_llm.run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 2. Works with any aysnc function\n",
+    "\n",
+    "Data Factory works with any async function, not just LLM calls, you can build complex pipelines involving multiple LLMs, data processing, etc.\n",
+    "\n",
+    "Here is example of two chained structured llm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:16:34\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 466fca03-85a2-46de-b135-629cd76738f7\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:34\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/3\u001b[0m | \u001b[33mRunning: 3\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/3\u001b[0m | \u001b[33mRunning: 3\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:40\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/3\u001b[0m | \u001b[33mRunning: 3\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:43\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 2/3\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 2\u001b[0m    (\u001b[32mCompleted: 2\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:44\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:44\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 3/3\u001b[0m | \u001b[33mAttempted: 3\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Example of a more complex function that chains multiple LLM calls\n",
+    "# This was grabbed from structured llm examples \n",
+    "\n",
+    "@data_factory(max_concurrency=10)\n",
+    "async def complex_process_cities(topic: str):\n",
+    "    ## topic → generator_llm → rating_llm → merged results\n",
+    "    # First LLM to generate question/answer pairs\n",
+    "    generator_llm = StructuredLLM(\n",
+    "        model_name=\"openai/gpt-4o-mini\",\n",
+    "        prompt=\"Generate question/answer pairs about {{topic}}.\",\n",
+    "        output_schema=[\n",
+    "            {\"name\": \"question\", \"type\": \"str\"},\n",
+    "            {\"name\": \"answer\", \"type\": \"str\"}\n",
+    "        ],\n",
+    "    )\n",
+    "\n",
+    "    # Second LLM to rate the generated pairs\n",
+    "    rater_llm = StructuredLLM(\n",
+    "        model_name=\"openai/gpt-4o-mini\",\n",
+    "        prompt='''Rate the following Q&A pairs based on accuracy and clarity (1-10).\n",
+    "        Pairs: {{generated_pairs}}''',\n",
+    "        output_schema=[\n",
+    "            {\"name\": \"accuracy_rating\", \"type\": \"int\"},\n",
+    "            {\"name\": \"clarity_rating\", \"type\": \"int\"}\n",
+    "        ],\n",
+    "        model_kwargs={\"temperature\": 0.5}\n",
+    ")\n",
+    "\n",
+    "    generation_response = await generator_llm.run(topic=topic, num_records=5)\n",
+    "    rating_response = await rater_llm.run(generated_pairs=generation_response.data)\n",
+    "    \n",
+    "    # Merge the results\n",
+    "    return merge_structured_outputs(generation_response.data, rating_response.data)\n",
+    "\n",
+    "\n",
+    "### To save on token here we only use 3 topics as example\n",
+    "complex_process_cities_data = complex_process_cities.run(topic=['Science', 'History', 'Technology'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "15\n",
+      "[{'question': 'What is the primary function of a CPU in a computer?', 'answer': 'The CPU, or Central Processing Unit, is responsible for executing instructions and processing data in a computer system.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What does IoT stand for and what is its significance?', 'answer': 'IoT stands for Internet of Things, which refers to the interconnection of everyday devices to the internet, allowing them to send and receive data, thereby enhancing efficiency and convenience.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the difference between RAM and ROM?', 'answer': 'RAM (Random Access Memory) is volatile memory that temporarily stores data and applications currently in use, while ROM (Read-Only Memory) is non-volatile memory that permanently stores firmware and system software.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is cloud computing?', 'answer': 'Cloud computing is the delivery of computing services over the internet, enabling users to access and store data and applications on remote servers rather than on local computers.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What are the benefits of using artificial intelligence in business?', 'answer': 'Artificial intelligence can enhance efficiency, improve decision-making, personalize customer experiences, automate repetitive tasks, and generate insights from data analytics in business operations.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the chemical formula for water?', 'answer': 'The chemical formula for water is H2O.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the process by which plants make their own food?', 'answer': 'The process by which plants make their own food is called photosynthesis.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the speed of light in a vacuum?', 'answer': 'The speed of light in a vacuum is approximately 299,792 kilometers per second (or about 186,282 miles per second).', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the smallest unit of life?', 'answer': 'The smallest unit of life is the cell.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What gas do living organisms need for respiration?', 'answer': 'Living organisms need oxygen for respiration.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What was the primary cause of World War I?', 'answer': 'The primary cause of World War I was the complex system of alliances, militarism, imperialism, and nationalism, which escalated tensions following the assassination of Archduke Franz Ferdinand of Austria in 1914.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'Who was the first President of the United States?', 'answer': 'George Washington was the first President of the United States, serving from April 30, 1789, to March 4, 1797.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What year did the Berlin Wall fall?', 'answer': 'The Berlin Wall fell on November 9, 1989, symbolizing the end of the Cold War and the division between East and West Germany.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'Which ancient civilization is known for creating the first known writing system?', 'answer': 'The Sumerians, who inhabited ancient Mesopotamia around 3500 BCE, are known for creating the first known writing system called cuneiform.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What was the significance of the Magna Carta?', 'answer': 'The Magna Carta, signed in 1215, was significant because it limited the power of the monarchy and established the principle that everyone, including the king, was subject to the law.', 'accuracy_rating': 10, 'clarity_rating': 10}]\n"
+     ]
+    }
+   ],
+   "source": [
+    "### Each topic has 5 question/answer pairs so 3 topics has 15 pairs!\n",
+    "print(len(complex_process_cities_data))\n",
+    "print(complex_process_cities_data)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 3. Working with Different Input Formats\n",
+    "\n",
+    "\n",
+    "Data Factory is flexible with how you provide inputs. Let's demonstrate different ways to pass parameters to data_factory functions.\n",
+    "\n",
+    "'data' is a reserved keyword expecting list(dict) or tuple(dict) - this design make it super easy to pass large data and support HuggingFace and Pandas dataframe very easily"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'answer': 'New York_5'}, {'answer': 'New York_2'}, {'answer': 'New York_3'}]"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "## We will be using mock llm call for rest of example to save on token\n",
+    "## Mock LLM call is a function that simulates an LLM API call with random delays (controlled by sleep_time) and occasional failures (controlled by fail_rate)\n",
+    "await mock_llm_call(city_name=\"New York\", num_records_per_city=3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@data_factory(max_concurrency=100)\n",
+    "async def input_format_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.01)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:16:49\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 05c84608-fec3-4010-8876-e59eed12bb6a\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:49\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:50\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:50\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Format 1: Multiple lists that get zipped together\n",
+    "input_format_data1 = input_format_mock_llm.run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"], num_records_per_city=[2, 1, 1, 1, 1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:16:50\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: fedb98e5-c408-4bc8-9479-6087f4a298b7\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:50\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:51\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:51\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Format 2: List + single value (single value gets broadcasted)\n",
+    "input_format_data2 = input_format_mock_llm.run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"], num_records_per_city=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:16:51\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 2f5cb7cc-83c9-4b7e-9ebb-386cd66bdd42\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:51\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:52\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:52\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Format 3: Special 'data' parameter\n",
+    "# 'data' is a reserved keyword expecting list(dict) or tuple(dict)\n",
+    "# Makes integration with various data sources easier\n",
+    "input_format_data3 = input_format_mock_llm.run(data=[{\"city_name\": \"New York\", \"num_records_per_city\": 2}, {\"city_name\": \"London\", \"num_records_per_city\": 1}, {\"city_name\": \"Tokyo\", \"num_records_per_city\": 1}, {\"city_name\": \"Paris\", \"num_records_per_city\": 1}, {\"city_name\": \"Sydney\", \"num_records_per_city\": 1}])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 4. Resilient error retry\n",
+    "Data Factory automatically handles errors and retries, making your pipelines robust.\n",
+    "\n",
+    "Let's demonstrate with a high failure rate example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:16:56\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 38c50ab6-f24b-4cba-a2c5-070130ab420e\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:56\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/25\u001b[0m | \u001b[33mRunning: 25\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:59\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 16/25\u001b[0m | \u001b[33mRunning: 9\u001b[0m | \u001b[36mAttempted: 16\u001b[0m    (\u001b[32mCompleted: 16\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:59\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Tokyo\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:59\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: London\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:16:59\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:02\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 23/25\u001b[0m | \u001b[33mRunning: 2\u001b[0m | \u001b[36mAttempted: 26\u001b[0m    (\u001b[32mCompleted: 23\u001b[0m, \u001b[31mFailed: 3\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:03\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:05\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:05\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 25/25\u001b[0m | \u001b[33mAttempted: 29\u001b[0m (Failed: 4, Filtered: 0, Duplicate: 0)\u001b[0m\n",
+      "\n",
+      "Successfully completed 25 out of 25 tasks\n",
+      "Data Factory automatically handled the failures and continued processing\n",
+      "The results only include successful tasks\n"
+     ]
+    }
+   ],
+   "source": [
+    "@data_factory(max_concurrency=100)\n",
+    "async def high_error_rate_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.3) # Hardcode to 30% chance of failure\n",
+    "\n",
+    "# Process all cities - some will fail, but data_factory keeps going\n",
+    "cities = [\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"] * 5  # 25 cities\n",
+    "high_error_rate_mock_lllm_data = high_error_rate_mock_llm.run(city_name=cities, num_records_per_city=1)\n",
+    "\n",
+    "print(f\"\\nSuccessfully completed {len(high_error_rate_mock_lllm_data)} out of {len(cities)} tasks\")\n",
+    "print(\"Data Factory automatically handled the failures and continued processing\")\n",
+    "print(\"The results only include successful tasks\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 5. Resume\n",
+    "\n",
+    "This is essential for long-running jobs with thousands of tasks.\n",
+    "\n",
+    "If a job is interrupted, you can pick up where you left off using one of two resume methods:\n",
+    "\n",
+    "\n",
+    "1. **Same Session Resume**: If you're still in the same session where the job was interrupted, simply call - Same instance with .resume()\n",
+    "\n",
+    "2. **Cross-Session Resume**: If you've closed your notebook or lost your session, you can resume using the job ID:\n",
+    "   ```python\n",
+    "   from starfish import DataFactory\n",
+    "   # Resume using the master job ID from a previous run\n",
+    "   data_factory = DataFactory.resume_from_checkpoint(job_id=\"your_job_id\")\n",
+    "   ```\n",
+    "\n",
+    "The key difference:\n",
+    "- `resume()` uses the same DataFactory instance you defined\n",
+    "- `resume_from_checkpoint()` reconstructs your DataFactory from persistent storage where tasks and progress are saved\n",
+    "\n",
+    "> **Note**: Google Colab users may experience issues with `resume_from_checkpoint()` due to how Colab works\n",
+    "\n",
+    "We're simulating an interruption here. In a real scenario, this might happen if your notebook errors out, is manually interrupted with a keyboard command, encounters API rate limits, or experiences any other issues that halt execution."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:17:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: b2a400b3-32e7-45ee-b8e8-c2bc7afe9f11\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 17/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 17\u001b[0m    (\u001b[32mCompleted: 17\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:15\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError occurred: KeyboardInterrupt\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[RESUME INFO] 🚨 Job stopped unexpectedly. You can resume the job by calling .resume()\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:17:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 20/100\u001b[0m | \u001b[33mAttempted: 20\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "@data_factory(max_concurrency=10)\n",
+    "async def re_run_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.3)\n",
+    "\n",
+    "cities = [\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"] * 20  # 100 cities\n",
+    "re_run_mock_llm_data_1 = re_run_mock_llm.run(city_name=cities, num_records_per_city=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "When a job is interrupted, you'll see a message like:\n",
+      "[RESUME INFO] 🚨 Job stopped unexpectedly. You can resume the job by calling .resume()\n",
+      "\n",
+      "To resume an interrupted job, simply call:\n",
+      "interrupted_job_mock_llm.resume()\n",
+      "\n",
+      "For this example we have 20/100 data generated and not finished yet!\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"When a job is interrupted, you'll see a message like:\")\n",
+    "print(\"[RESUME INFO] 🚨 Job stopped unexpectedly. You can resume the job by calling .resume()\")\n",
+    "\n",
+    "print(\"\\nTo resume an interrupted job, simply call:\")\n",
+    "print(\"interrupted_job_mock_llm.resume()\")\n",
+    "print('')\n",
+    "print(f\"For this example we have {len(re_run_mock_llm_data_1)}/{len(cities)} data generated and not finished yet!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:18:00\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB RESUME START]\u001b[0m \u001b[33mPICKING UP FROM WHERE THE JOB WAS LEFT OFF...\u001b[0m\n",
+      "\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:00\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[RESUME PROGRESS] STATUS AT THE TIME OF RESUME:\u001b[0m \u001b[32mCompleted: 20 / 100\u001b[0m | \u001b[31mFailed: 0\u001b[0m | \u001b[31mDuplicate: 0\u001b[0m | \u001b[33mFiltered: 0\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:00\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 20/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 20\u001b[0m    (\u001b[32mCompleted: 20\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:03\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 32/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 32\u001b[0m    (\u001b[32mCompleted: 32\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:03\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Paris\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:03\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Sydney\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:03\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Sydney\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 56/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 59\u001b[0m    (\u001b[32mCompleted: 56\u001b[0m, \u001b[31mFailed: 3\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:08\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Sydney\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:08\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: London\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:09\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 69/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 74\u001b[0m    (\u001b[32mCompleted: 69\u001b[0m, \u001b[31mFailed: 5\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:09\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 89/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 95\u001b[0m    (\u001b[32mCompleted: 89\u001b[0m, \u001b[31mFailed: 6\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:12\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Paris\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:13\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:14\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:14\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:14\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 100/100\u001b[0m | \u001b[33mAttempted: 109\u001b[0m (Failed: 9, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "## Lets keep continue the rest of run by resume_from_checkpoint \n",
+    "re_run_mock_llm_data_2 = re_run_mock_llm.resume()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Now we still able to finished with what is left!! 100 data generated!\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(f\"Now we still able to finished with what is left!! {len(re_run_mock_llm_data_2)} data generated!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 6. Dry run\n",
+    "Before running a large job, you can do a \"dry run\" to test your pipeline. This only processes a single item and doesn't save state to the database."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:18:14\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: None\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:14\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mSending telemetry event, TELEMETRY_ENABLED=true\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:18:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/0\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "@data_factory(max_concurrency=10)\n",
+    "async def dry_run_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.3)\n",
+    "\n",
+    "dry_run_mock_llm_data = dry_run_mock_llm.dry_run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"]*20, num_records_per_city=1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 8. Advanced Usage\n",
+    "Data Factory offers more advanced capabilities for complete pipeline customization, including hooks that execute at key stages and shareable state to coordinate between tasks. These powerful features enable complex workflows and fine-grained control. Our dedicated examples for advanced data_factory usage will be coming soon!"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/data_factory_release_check.ipynb

@@ -0,0 +1,494 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Google Colab Version: [Open this notebook in Google Colab](https://colab.research.google.com/github/starfishdata/starfish/blob/main/examples/data_factory.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Dependencies "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Looking in indexes: https://test.pypi.org/simple/, https://pypi.org/simple\n",
+      "Requirement already satisfied: starfish-core in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (0.1.2)\n",
+      "Requirement already satisfied: aiofiles<25.0.0,>=24.1.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (24.1.0)\n",
+      "Requirement already satisfied: aiosqlite<0.22.0,>=0.21.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (0.21.0)\n",
+      "Requirement already satisfied: cachetools<6.0.0,>=5.5.2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (5.5.2)\n",
+      "Requirement already satisfied: cloudpickle<3.0.0,>=2.2.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (2.2.1)\n",
+      "Requirement already satisfied: cryptography>=44.0.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (44.0.3)\n",
+      "Requirement already satisfied: docstring_parser<0.17.0,>=0.16.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (0.16)\n",
+      "Requirement already satisfied: litellm<2.0.0,>=1.65.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (1.69.3)\n",
+      "Requirement already satisfied: loguru<0.8.0,>=0.7.3 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (0.7.3)\n",
+      "Requirement already satisfied: mcp<2.0.0,>=1.8.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (1.9.0)\n",
+      "Requirement already satisfied: nest_asyncio<2.0.0,>=1.6.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (1.6.0)\n",
+      "Requirement already satisfied: ollama<0.5.0,>=0.4.7 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (0.4.8)\n",
+      "Requirement already satisfied: posthog<4.0.0,>=3.11.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (3.25.0)\n",
+      "Requirement already satisfied: psutil<8.0.0,>=7.0.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (7.0.0)\n",
+      "Requirement already satisfied: python-dotenv<2.0.0,>=1.1.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (1.1.0)\n",
+      "Requirement already satisfied: typing-extensions<5.0.0,>=4.0.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from starfish-core) (4.13.2)\n",
+      "Requirement already satisfied: cffi>=1.12 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from cryptography>=44.0.1->starfish-core) (1.17.1)\n",
+      "Requirement already satisfied: aiohttp in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (3.11.18)\n",
+      "Requirement already satisfied: click in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (8.2.0)\n",
+      "Requirement already satisfied: httpx>=0.23.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (0.28.1)\n",
+      "Requirement already satisfied: importlib-metadata>=6.8.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (8.7.0)\n",
+      "Requirement already satisfied: jinja2<4.0.0,>=3.1.2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (3.1.6)\n",
+      "Requirement already satisfied: jsonschema<5.0.0,>=4.22.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (4.23.0)\n",
+      "Requirement already satisfied: openai<1.76.0,>=1.68.2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (1.75.0)\n",
+      "Requirement already satisfied: pydantic<3.0.0,>=2.0.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (2.11.4)\n",
+      "Requirement already satisfied: tiktoken>=0.7.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (0.9.0)\n",
+      "Requirement already satisfied: tokenizers in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from litellm<2.0.0,>=1.65.1->starfish-core) (0.21.1)\n",
+      "Requirement already satisfied: anyio>=4.5 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (4.9.0)\n",
+      "Requirement already satisfied: httpx-sse>=0.4 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (0.4.0)\n",
+      "Requirement already satisfied: pydantic-settings>=2.5.2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (2.9.1)\n",
+      "Requirement already satisfied: python-multipart>=0.0.9 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (0.0.20)\n",
+      "Requirement already satisfied: sse-starlette>=1.6.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (2.3.5)\n",
+      "Requirement already satisfied: starlette>=0.27 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (0.46.2)\n",
+      "Requirement already satisfied: uvicorn>=0.23.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from mcp<2.0.0,>=1.8.1->starfish-core) (0.34.2)\n",
+      "Requirement already satisfied: requests<3.0,>=2.7 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from posthog<4.0.0,>=3.11.0->starfish-core) (2.32.3)\n",
+      "Requirement already satisfied: six>=1.5 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from posthog<4.0.0,>=3.11.0->starfish-core) (1.17.0)\n",
+      "Requirement already satisfied: monotonic>=1.5 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from posthog<4.0.0,>=3.11.0->starfish-core) (1.6)\n",
+      "Requirement already satisfied: backoff>=1.10.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from posthog<4.0.0,>=3.11.0->starfish-core) (2.2.1)\n",
+      "Requirement already satisfied: python-dateutil>2.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from posthog<4.0.0,>=3.11.0->starfish-core) (2.9.0.post0)\n",
+      "Requirement already satisfied: distro>=1.5.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from posthog<4.0.0,>=3.11.0->starfish-core) (1.9.0)\n",
+      "Requirement already satisfied: idna>=2.8 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from anyio>=4.5->mcp<2.0.0,>=1.8.1->starfish-core) (3.10)\n",
+      "Requirement already satisfied: sniffio>=1.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from anyio>=4.5->mcp<2.0.0,>=1.8.1->starfish-core) (1.3.1)\n",
+      "Requirement already satisfied: pycparser in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from cffi>=1.12->cryptography>=44.0.1->starfish-core) (2.22)\n",
+      "Requirement already satisfied: certifi in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (2025.4.26)\n",
+      "Requirement already satisfied: httpcore==1.* in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (1.0.9)\n",
+      "Requirement already satisfied: h11>=0.16 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from httpcore==1.*->httpx>=0.23.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.16.0)\n",
+      "Requirement already satisfied: zipp>=3.20 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from importlib-metadata>=6.8.0->litellm<2.0.0,>=1.65.1->starfish-core) (3.21.0)\n",
+      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from jinja2<4.0.0,>=3.1.2->litellm<2.0.0,>=1.65.1->starfish-core) (3.0.2)\n",
+      "Requirement already satisfied: attrs>=22.2.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (25.3.0)\n",
+      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (2025.4.1)\n",
+      "Requirement already satisfied: referencing>=0.28.4 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.36.2)\n",
+      "Requirement already satisfied: rpds-py>=0.7.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from jsonschema<5.0.0,>=4.22.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.25.0)\n",
+      "Requirement already satisfied: jiter<1,>=0.4.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from openai<1.76.0,>=1.68.2->litellm<2.0.0,>=1.65.1->starfish-core) (0.9.0)\n",
+      "Requirement already satisfied: tqdm>4 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from openai<1.76.0,>=1.68.2->litellm<2.0.0,>=1.65.1->starfish-core) (4.67.1)\n",
+      "Requirement already satisfied: annotated-types>=0.6.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from pydantic<3.0.0,>=2.0.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.7.0)\n",
+      "Requirement already satisfied: pydantic-core==2.33.2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from pydantic<3.0.0,>=2.0.0->litellm<2.0.0,>=1.65.1->starfish-core) (2.33.2)\n",
+      "Requirement already satisfied: typing-inspection>=0.4.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from pydantic<3.0.0,>=2.0.0->litellm<2.0.0,>=1.65.1->starfish-core) (0.4.0)\n",
+      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from requests<3.0,>=2.7->posthog<4.0.0,>=3.11.0->starfish-core) (3.4.2)\n",
+      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from requests<3.0,>=2.7->posthog<4.0.0,>=3.11.0->starfish-core) (2.4.0)\n",
+      "Requirement already satisfied: regex>=2022.1.18 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from tiktoken>=0.7.0->litellm<2.0.0,>=1.65.1->starfish-core) (2024.11.6)\n",
+      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (2.6.1)\n",
+      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (1.3.2)\n",
+      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (1.6.0)\n",
+      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (6.4.3)\n",
+      "Requirement already satisfied: propcache>=0.2.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (0.3.1)\n",
+      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from aiohttp->litellm<2.0.0,>=1.65.1->starfish-core) (1.20.0)\n",
+      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (0.31.2)\n",
+      "Requirement already satisfied: filelock in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (3.18.0)\n",
+      "Requirement already satisfied: fsspec>=2023.5.0 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (2025.3.2)\n",
+      "Requirement already satisfied: packaging>=20.9 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (25.0)\n",
+      "Requirement already satisfied: pyyaml>=5.1 in /Users/john/Documents/projects/aa/python/starfish/starfish/.venv/lib/python3.11/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers->litellm<2.0.0,>=1.65.1->starfish-core) (6.0.2)\n",
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.1.1\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "#%pip install starfish-core\n",
+    "%pip install --index-url https://test.pypi.org/simple/ \\\n",
+    "                   --extra-index-url https://pypi.org/simple \\\n",
+    "                    starfish-core"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:10\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mFailed to load environment variables from /Users/john/Documents/projects/aa/python/starfish/starfish/.env\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "## Fix for Jupyter Notebook only — do NOT use in production\n",
+    "## Enables async code execution in notebooks, but may cause issues with sync/async issues\n",
+    "## For production, please run in standard .py files without this workaround\n",
+    "## See: https://github.com/erdewit/nest_asyncio for more details\n",
+    "import nest_asyncio\n",
+    "nest_asyncio.apply()\n",
+    "\n",
+    "from starfish import StructuredLLM, data_factory\n",
+    "from starfish.llm.utils import merge_structured_outputs\n",
+    "\n",
+    "from starfish.common.env_loader import load_env_file ## Load environment variables from .env file\n",
+    "load_env_file()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Helper function mock llm call\n",
+    "# When developing data pipelines with LLMs, making thousands of real API calls\n",
+    "# can be expensive. Using mock LLM calls lets you test your pipeline's reliability,\n",
+    "# failure handling, and recovery without spending money on API calls.\n",
+    "from starfish.data_factory.utils.mock import mock_llm_call"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 3. Working with Different Input Formats\n",
+    "\n",
+    "\n",
+    "Data Factory is flexible with how you provide inputs. Let's demonstrate different ways to pass parameters to data_factory functions.\n",
+    "\n",
+    "'data' is a reserved keyword expecting list(dict) or tuple(dict) - this design make it super easy to pass large data and support HuggingFace and Pandas dataframe very easily"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'answer': 'New York_3'}, {'answer': 'New York_1'}, {'answer': 'New York_5'}]"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "## We will be using mock llm call for rest of example to save on token\n",
+    "## Mock LLM call is a function that simulates an LLM API call with random delays (controlled by sleep_time) and occasional failures (controlled by fail_rate)\n",
+    "await mock_llm_call(city_name=\"New York\", num_records_per_city=3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@data_factory(max_concurrency=100)\n",
+    "async def input_format_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.01)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:10\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 4da82fc7-4112-4e05-b58c-53cf470747ad\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:10\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:11\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Format 1: Multiple lists that get zipped together\n",
+    "input_format_data1 = input_format_mock_llm.run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"], num_records_per_city=[2, 1, 1, 1, 1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:11\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 73973449-6069-485e-ac8c-b1b3a6b3f1a4\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:11\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Format 2: List + single value (single value gets broadcasted)\n",
+    "input_format_data2 = input_format_mock_llm.run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"], num_records_per_city=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: aa9954f9-fc18-4b42-959e-fb2a897987c7\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/5\u001b[0m | \u001b[33mRunning: 5\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:13\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 5/5\u001b[0m | \u001b[33mAttempted: 5\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Format 3: Special 'data' parameter\n",
+    "# 'data' is a reserved keyword expecting list(dict) or tuple(dict)\n",
+    "# Makes integration with various data sources easier\n",
+    "input_format_data3 = input_format_mock_llm.run(data=[{\"city_name\": \"New York\", \"num_records_per_city\": 2}, {\"city_name\": \"London\", \"num_records_per_city\": 1}, {\"city_name\": \"Tokyo\", \"num_records_per_city\": 1}, {\"city_name\": \"Paris\", \"num_records_per_city\": 1}, {\"city_name\": \"Sydney\", \"num_records_per_city\": 1}])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 4. Resilient error retry\n",
+    "Data Factory automatically handles errors and retries, making your pipelines robust.\n",
+    "\n",
+    "Let's demonstrate with a high failure rate example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:13\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 730b766d-3c23-419a-a3dd-271d683818b1\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:13\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/25\u001b[0m | \u001b[33mRunning: 25\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:15\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Tokyo\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:15\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:16\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 23/25\u001b[0m | \u001b[33mRunning: 0\u001b[0m | \u001b[36mAttempted: 25\u001b[0m    (\u001b[32mCompleted: 23\u001b[0m, \u001b[31mFailed: 2\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:19\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 25/25\u001b[0m | \u001b[33mAttempted: 27\u001b[0m (Failed: 2, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\n",
+      "Successfully completed 25 out of 25 tasks\n",
+      "Data Factory automatically handled the failures and continued processing\n",
+      "The results only include successful tasks\n"
+     ]
+    }
+   ],
+   "source": [
+    "@data_factory(max_concurrency=100)\n",
+    "async def high_error_rate_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.3) # Hardcode to 30% chance of failure\n",
+    "\n",
+    "# Process all cities - some will fail, but data_factory keeps going\n",
+    "cities = [\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"] * 5  # 25 cities\n",
+    "high_error_rate_mock_lllm_data = high_error_rate_mock_llm.run(city_name=cities, num_records_per_city=1)\n",
+    "\n",
+    "print(f\"\\nSuccessfully completed {len(high_error_rate_mock_lllm_data)} out of {len(cities)} tasks\")\n",
+    "print(\"Data Factory automatically handled the failures and continued processing\")\n",
+    "print(\"The results only include successful tasks\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 5. Resume\n",
+    "\n",
+    "This is essential for long-running jobs with thousands of tasks.\n",
+    "\n",
+    "If a job is interrupted, you can pick up where you left off using one of two resume methods:\n",
+    "\n",
+    "\n",
+    "1. **Same Session Resume**: If you're still in the same session where the job was interrupted, simply call - Same instance with .resume()\n",
+    "\n",
+    "2. **Cross-Session Resume**: If you've closed your notebook or lost your session, you can resume using the job ID:\n",
+    "   ```python\n",
+    "   from starfish import DataFactory\n",
+    "   # Resume using the master job ID from a previous run\n",
+    "   data_factory = DataFactory.resume_from_checkpoint(job_id=\"your_job_id\")\n",
+    "   ```\n",
+    "\n",
+    "The key difference:\n",
+    "- `resume()` uses the same DataFactory instance you defined\n",
+    "- `resume_from_checkpoint()` reconstructs your DataFactory from persistent storage where tasks and progress are saved\n",
+    "\n",
+    "> **Note**: Google Colab users may experience issues with `resume_from_checkpoint()` due to how Colab works\n",
+    "\n",
+    "We're simulating an interruption here. In a real scenario, this might happen if your notebook errors out, is manually interrupted with a keyboard command, encounters API rate limits, or experiences any other issues that halt execution."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:19\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 6829de29-0b83-4a64-835b-cc79cbad5e3a\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:19\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:21\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Paris\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:21\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Sydney\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:21\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: New York\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:21\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mconsecutive_not_completed: in 3 times, stopping this job; please adjust factory config and input data then resume_from_checkpoint(6829de29-0b83-4a64-835b-cc79cbad5e3a)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:21\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 17/100\u001b[0m | \u001b[33mAttempted: 20\u001b[0m (Failed: 3, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "@data_factory(max_concurrency=10)\n",
+    "async def re_run_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.3)\n",
+    "\n",
+    "cities = [\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"] * 20  # 100 cities\n",
+    "re_run_mock_llm_data_1 = re_run_mock_llm.run(city_name=cities, num_records_per_city=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "When a job is interrupted, you'll see a message like:\n",
+      "[RESUME INFO] 🚨 Job stopped unexpectedly. You can resume the job by calling .resume()\n",
+      "\n",
+      "To resume an interrupted job, simply call:\n",
+      "interrupted_job_mock_llm.resume()\n",
+      "\n",
+      "For this example we have 17/100 data generated and not finished yet!\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"When a job is interrupted, you'll see a message like:\")\n",
+    "print(\"[RESUME INFO] 🚨 Job stopped unexpectedly. You can resume the job by calling .resume()\")\n",
+    "\n",
+    "print(\"\\nTo resume an interrupted job, simply call:\")\n",
+    "print(\"interrupted_job_mock_llm.resume()\")\n",
+    "print('')\n",
+    "print(f\"For this example we have {len(re_run_mock_llm_data_1)}/{len(cities)} data generated and not finished yet!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:22\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB RESUME START]\u001b[0m \u001b[33mPICKING UP FROM WHERE THE JOB WAS LEFT OFF...\u001b[0m\n",
+      "\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:22\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[RESUME PROGRESS] STATUS AT THE TIME OF RESUME:\u001b[0m \u001b[32mCompleted: 17 / 100\u001b[0m | \u001b[31mFailed: 3\u001b[0m | \u001b[31mDuplicate: 0\u001b[0m | \u001b[33mFiltered: 0\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:22\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 17/100\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 20\u001b[0m    (\u001b[32mCompleted: 17\u001b[0m, \u001b[31mFailed: 3\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:24\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mError running task: Mock LLM failed to process city: Paris\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:24\u001b[0m | \u001b[31m\u001b[1mERROR   \u001b[0m | \u001b[31m\u001b[1mconsecutive_not_completed: in 3 times, stopping this job; please adjust factory config and input data then resume_from_checkpoint(6829de29-0b83-4a64-835b-cc79cbad5e3a)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:24\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 30/100\u001b[0m | \u001b[33mAttempted: 34\u001b[0m (Failed: 4, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "## Lets keep continue the rest of run by resume_from_checkpoint \n",
+    "re_run_mock_llm_data_2 = re_run_mock_llm.resume()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Now we still able to finished with what is left!! 30 data generated!\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(f\"Now we still able to finished with what is left!! {len(re_run_mock_llm_data_2)} data generated!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 6. Dry run\n",
+    "Before running a large job, you can do a \"dry run\" to test your pipeline. This only processes a single item and doesn't save state to the database."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 22:50:24\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: None\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:24\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 22:50:25\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/0\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "@data_factory(max_concurrency=10)\n",
+    "async def dry_run_mock_llm(city_name: str, num_records_per_city: int):\n",
+    "    return await mock_llm_call(city_name=city_name, num_records_per_city=num_records_per_city, fail_rate=0.3)\n",
+    "\n",
+    "dry_run_mock_llm_data = dry_run_mock_llm.dry_run(city_name=[\"New York\", \"London\", \"Tokyo\", \"Paris\", \"Sydney\"]*20, num_records_per_city=1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 8. Advanced Usage\n",
+    "Data Factory offers more advanced capabilities for complete pipeline customization, including hooks that execute at key stages and shareable state to coordinate between tasks. These powerful features enable complex workflows and fine-grained control. Our dedicated examples for advanced data_factory usage will be coming soon!"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/embedding_usage_example.py

@@ -0,0 +1,202 @@
+"""Example: Using Starfish Embeddings for Data Generation
+
+This example demonstrates how to use FAISS and SentenceTransformers
+for embedding-enhanced data generation and deduplication.
+"""
+
+import asyncio
+import sys
+import os
+
+# Add the project root to the Python path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+
+from starfish.embedding import EmbeddingManager, SimilarityChecker, DataDeduplicator
+from starfish.data_gen_template.core import data_gen_template
+
+
+async def basic_embedding_example():
+    """Basic example of using the embedding system."""
+    print("🔮 Basic Embedding Example")
+    print("=" * 50)
+
+    # Initialize embedding manager
+    embedding_manager = EmbeddingManager(model_name="all-MiniLM-L6-v2", similarity_threshold=0.85)
+
+    # Sample texts to embed
+    texts = [
+        "What is machine learning?",
+        "How does artificial intelligence work?",
+        "What are neural networks?",
+        "Explain deep learning concepts",
+        "What is supervised learning?",
+        "What is machine learning?",  # Duplicate
+        "How do neural networks function?",  # Similar to "What are neural networks?"
+    ]
+
+    print(f"📝 Processing {len(texts)} sample texts...")
+
+    # Add texts to the index
+    indices = embedding_manager.add_texts(texts)
+    print(f"✅ Added {len(indices)} texts to the embedding index")
+
+    # Search for similar texts
+    query = "Tell me about AI and ML"
+    similar_items = embedding_manager.search_similar(query, k=3)
+
+    print(f"\n🔍 Search results for: '{query}'")
+    for item in similar_items:
+        print(f"   Similarity: {item['similarity']:.3f} | Text: {item['text']}")
+
+    # Find duplicates
+    duplicate_groups = embedding_manager.find_duplicates(texts)
+    print(f"\n🔄 Found {len(duplicate_groups)} groups of duplicates:")
+    for i, group in enumerate(duplicate_groups):
+        print(f"   Group {i+1}: {[texts[idx] for idx in group]}")
+
+    print(f"\n📊 Index Stats: {embedding_manager.get_stats()}")
+
+
+async def similarity_checker_example():
+    """Example of using the similarity checker."""
+    print("\n🎯 Similarity Checker Example")
+    print("=" * 50)
+
+    similarity_checker = SimilarityChecker(similarity_threshold=0.8)
+
+    # Sample data items
+    data_items = [
+        {"question": "What is Python?", "answer": "Python is a programming language"},
+        {"question": "How to learn coding?", "answer": "Start with basic concepts"},
+        {"question": "What is programming?", "answer": "Programming is writing code"},
+        {"question": "What is Python programming?", "answer": "Python is a popular language"},  # Similar to first
+    ]
+
+    print(f"📝 Analyzing {len(data_items)} data items...")
+
+    # Filter similar items
+    filtered_items, duplicate_groups = similarity_checker.filter_similar_items(data_items)
+    print(f"✅ Filtered to {len(filtered_items)} unique items")
+
+    # Check diversity metrics
+    diversity_metrics = similarity_checker.check_diversity_batch(data_items)
+    print(f"📈 Diversity Score: {diversity_metrics['diversity_score']:.3f}")
+    print(f"🔄 Average Similarity: {diversity_metrics['avg_similarity']:.3f}")
+
+    # Suggest diverse subset
+    diverse_subset = similarity_checker.suggest_diverse_subset(data_items, target_size=2)
+    print(f"\n🎲 Diverse subset (2 items):")
+    for item in diverse_subset:
+        print(f"   Q: {item['question']}")
+
+
+async def deduplicator_example():
+    """Example of using the data deduplicator."""
+    print("\n🔧 Data Deduplicator Example")
+    print("=" * 50)
+
+    deduplicator = DataDeduplicator(similarity_threshold=0.9)
+
+    # Sample dataset with duplicates
+    dataset = [
+        {"id": "1", "text": "Machine learning is a subset of AI", "quality_score": 0.8},
+        {"id": "2", "text": "Deep learning uses neural networks", "quality_score": 0.9},
+        {"id": "1", "text": "Machine learning is a subset of AI", "quality_score": 0.7},  # Exact duplicate
+        {"id": "3", "text": "ML is part of artificial intelligence", "quality_score": 0.95},  # Semantic duplicate
+        {"id": "4", "text": "Natural language processing handles text", "quality_score": 0.85},
+    ]
+
+    print(f"📝 Analyzing dataset with {len(dataset)} items...")
+
+    # Analyze duplicates without removing
+    analysis = deduplicator.analyze_duplicates(dataset)
+    print(f"🔍 Analysis Results:")
+    print(f"   Exact duplicates: {analysis['exact_duplicates']['count']}")
+    print(f"   Semantic duplicates: {analysis['semantic_duplicates']['count']}")
+    print(f"   Diversity score: {analysis['diversity_metrics']['diversity_score']:.3f}")
+
+    # Perform comprehensive deduplication
+    clean_dataset, report = deduplicator.deduplicate_comprehensive(dataset)
+    print(f"\n✨ Deduplication Results:")
+    print(f"   Original: {report['original_count']} items")
+    print(f"   Final: {report['final_count']} items")
+    print(f"   Reduction: {report['reduction_percentage']:.1f}%")
+
+    print("\n📋 Clean dataset:")
+    for item in clean_dataset:
+        print(f"   ID: {item['id']} | Score: {item.get('quality_score', 'N/A')} | Text: {item['text'][:50]}...")
+
+
+async def template_usage_example():
+    """Example of using the embedding-enhanced template."""
+    print("\n🚀 Embedding-Enhanced Template Example")
+    print("=" * 50)
+
+    try:
+        # Get the embedding template
+        print(data_gen_template.list())
+        template = data_gen_template.get("starfish/generate_with_embeddings")
+
+        # Configuration for generation
+        config = {
+            "num_records": 5,  # Small number for demo
+            "user_instruction": "Generate educational Q&A about data science",
+            "topics": ["statistics", "data visualization", "machine learning"],
+            "generation_model_name": "openai/gpt-4o-mini",
+            "embedding_config": {
+                "model_name": "all-MiniLM-L6-v2",
+                "similarity_threshold": 0.8,
+                "enable_deduplication": True,
+                "enable_diversity_check": True,
+                "min_diversity_score": 0.2,
+            },
+        }
+
+        print("⚙️  Generating diverse dataset with embedding quality control...")
+        results = await template.run(**config)
+
+        print(f"\n✅ Generated {len(results)} high-quality items:")
+        for i, item in enumerate(results[:3]):  # Show first 3
+            print(f"\n   Item {i+1}:")
+            print(f"   Q: {item.get('question', 'N/A')}")
+            print(f"   A: {item.get('answer', 'N/A')[:100]}...")
+            if "_metadata" in item:
+                print(f"   Diversity: {item['_metadata'].get('diversity_score', 'N/A'):.3f}")
+
+    except Exception as e:
+        print(f"⚠️  Template example failed: {e}")
+        print("   (This might be due to missing API keys or dependencies)")
+
+
+async def main():
+    """Run all examples."""
+    print("🎉 Starfish Embedding System Examples")
+    print("=" * 60)
+
+    try:
+        await basic_embedding_example()
+        await similarity_checker_example()
+        await deduplicator_example()
+        await template_usage_example()
+
+        print("\n" + "=" * 60)
+        print("✅ All examples completed successfully!")
+        print("\n💡 Next steps:")
+        print("   1. Install dependencies: poetry install")
+        print("   2. Set API keys in .env.local")
+        print("   3. Try the embedding template in your projects")
+
+    except ImportError as e:
+        print(f"❌ Import error: {e}")
+        print("💡 Make sure to install dependencies:")
+        print("   poetry install")
+        print("   # or")
+        print("   pip install faiss-cpu sentence-transformers")
+
+    except Exception as e:
+        print(f"❌ Error running examples: {e}")
+        print("💡 Check your Python environment and dependencies")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/structured_llm.ipynb

@@ -0,0 +1,470 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Google Colab Version: [Open this notebook in Google Colab](https://colab.research.google.com/github/starfishdata/starfish/blob/main/examples/structured_llm.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Dependencies "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install starfish-core"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Fix for Jupyter Notebook only — do NOT use in production\n",
+    "## Enables async code execution in notebooks, but may cause issues with sync/async issues\n",
+    "## For production, please run in standard .py files without this workaround\n",
+    "## See: https://github.com/erdewit/nest_asyncio for more details\n",
+    "import nest_asyncio\n",
+    "nest_asyncio.apply()\n",
+    "\n",
+    "from starfish import StructuredLLM\n",
+    "from starfish.llm.utils import merge_structured_outputs\n",
+    "\n",
+    "from pydantic import BaseModel, Field\n",
+    "from typing import List\n",
+    "\n",
+    "from starfish.common.env_loader import load_env_file ## Load environment variables from .env file\n",
+    "load_env_file()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# setup your openai api key if not already set\n",
+    "# import os\n",
+    "# os.environ[\"OPENAI_API_KEY\"] = \"your_key_here\"\n",
+    "\n",
+    "# If you dont have any API key, use local model (ollama)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 1. Structured LLM with JSON Schema"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'question': 'Why did the tomato turn red in New York?',\n",
+       "  'answer': \"Because it saw the Big Apple and couldn't ketchup with all the excitement!\"}]"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# ### Define the Output Structure (JSON Schema)\n",
+    "# Let's start with a simple JSON-like schema using a list of dictionaries.\n",
+    "# Each dictionary specifies a field name and its type. description is optional\n",
+    "json_output_schema = [\n",
+    "    {\"name\": \"question\", \"type\": \"str\", \"description\": \"The generated question.\"},\n",
+    "    {\"name\": \"answer\", \"type\": \"str\", \"description\": \"The corresponding answer.\"},\n",
+    "]\n",
+    "\n",
+    "json_llm = StructuredLLM(\n",
+    "    model_name = \"openai/gpt-4o-mini\",\n",
+    "    prompt = \"Funny facts about city {{city_name}}.\",\n",
+    "    output_schema = json_output_schema,\n",
+    "    model_kwargs = {\"temperature\": 0.7},\n",
+    ")\n",
+    "\n",
+    "json_response = await json_llm.run(city_name=\"New York\")\n",
+    "\n",
+    "# The response object contains both parsed data and the raw API response.\n",
+    "json_response.data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ModelResponse(id='chatcmpl-BQGw3FMSjzWOPMRvXmgknN4oozrKK', created=1745601327, model='gpt-4o-mini-2024-07-18', object='chat.completion', system_fingerprint='fp_0392822090', choices=[Choices(finish_reason='stop', index=0, message=Message(content='[\\n    {\\n        \"question\": \"Why did the tomato turn red in New York?\",\\n        \"answer\": \"Because it saw the Big Apple and couldn\\'t ketchup with all the excitement!\"\\n    }\\n]', role='assistant', tool_calls=None, function_call=None, provider_specific_fields={'refusal': None}, annotations=[]))], usage=Usage(completion_tokens=41, prompt_tokens=77, total_tokens=118, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier='default')"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Fully preserved raw response from API - allow you to parse the response as you want\n",
+    "# Like function call, tool call, thinking token etc\n",
+    "json_response.raw"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 2. Structured LLM with Pydantic Schema (Nested)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'facts': [{'question': 'What year did New York City become the capital of the United States?',\n",
+       "    'answer': 'New York City served as the capital of the United States from 1785 to 1790.',\n",
+       "    'category': 'History'}]}]"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# ### Define the Output Structure (Pydantic Model)\n",
+    "class Fact(BaseModel):\n",
+    "    question: str = Field(..., description=\"The factual question generated.\")\n",
+    "    answer: str = Field(..., description=\"The corresponding answer.\")\n",
+    "    category: str = Field(..., description=\"A category for the fact (e.g., History, Geography).\")\n",
+    "\n",
+    "# You can define a list of these models if you expect multiple results.\n",
+    "class FactsList(BaseModel):\n",
+    "    facts: List[Fact] = Field(..., description=\"A list of facts.\")\n",
+    "\n",
+    "\n",
+    "# ### Create the StructuredLLM Instance with Pydantic\n",
+    "pydantic_llm = StructuredLLM(\n",
+    "    model_name=\"openai/gpt-4o-mini\",\n",
+    "    # Ask for multiple facts this time\n",
+    "    prompt=\"Generate distinct facts about {{city}}.\",\n",
+    "    # Pass the Pydantic model directly as the schema\n",
+    "    output_schema=FactsList, # Expecting a list of facts wrapped in the FactsList model\n",
+    "    model_kwargs={\"temperature\": 0.8}\n",
+    ")\n",
+    "\n",
+    "pydantic_llm_response = await pydantic_llm.run(city=\"New York\")\n",
+    "\n",
+    "pydantic_llm_response.data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 3. Working with Different LLM Providers\n",
+    "\n",
+    "Starfish uses LiteLLM under the hood, giving you access to 100+ LLM providers. Here is an example of using a custom model provider - Hyperbolic - Super cool provider with full precision model and low cost!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'question': 'What is the nickname of New York City?',\n",
+       "  'answer': 'The Big Apple'},\n",
+       " {'question': 'Which iconic statue is located in New York Harbor?',\n",
+       "  'answer': 'The Statue of Liberty'},\n",
+       " {'question': 'What is the name of the famous theater district in Manhattan?',\n",
+       "  'answer': 'Broadway'},\n",
+       " {'question': \"Which park is considered the 'lungs' of New York City?\",\n",
+       "  'answer': 'Central Park'},\n",
+       " {'question': 'What is the tallest building in New York City as of 2023?',\n",
+       "  'answer': 'One World Trade Center'}]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "\n",
+    "# Set up the relevant API Key and Base URL in your enviornment variables\n",
+    "# os.environ[\"HYPERBOLIC_API_KEY\"] = \"your_key_here\"\n",
+    "# os.environ[\"HYPERBOLIC_API_BASE\"] = \"https://api.hyperbolic.xyz/v1\"\n",
+    "\n",
+    "hyperbolic_llm = StructuredLLM(\n",
+    "    model_name=\"hyperbolic/deepseek-ai/DeepSeek-V3-0324\", \n",
+    "    prompt=\"Facts about city {{city_name}}.\",\n",
+    "    output_schema=[{\"name\": \"question\", \"type\": \"str\"}, {\"name\": \"answer\", \"type\": \"str\"}],\n",
+    "    model_kwargs={\"temperature\": 0.7},\n",
+    ")\n",
+    "\n",
+    "hyperbolic_llm_response = await hyperbolic_llm.run(city_name=\"New York\", num_records=5)\n",
+    "hyperbolic_llm_response.data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 3. Local LLM using Ollama\n",
+    "Ensure Ollama is installed and running. Starfish can manage the server process and model downloads"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:15:40\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mEnsuring Ollama model gemma3:1b is ready...\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:15:40\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mStarting Ollama server...\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:15:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mOllama server started successfully\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:15:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mFound model gemma3:1b\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:15:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mModel gemma3:1b is already available\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:15:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mModel gemma3:1b is ready, making API call...\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[{'question': 'What is the population of New York City?',\n",
+       "  'answer': 'As of 2023, the population of New York City is approximately 8.8 million people.'}]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "### Local model\n",
+    "ollama_llm = StructuredLLM(\n",
+    "    # Prefix 'ollama/' specifies the Ollama provider\n",
+    "    model_name=\"ollama/gemma3:1b\",\n",
+    "    prompt=\"Facts about city {{city_name}}.\",\n",
+    "    output_schema=[{\"name\": \"question\", \"type\": \"str\"}, {\"name\": \"answer\", \"type\": \"str\"}],\n",
+    "    model_kwargs={\"temperature\": 0.7},\n",
+    ")\n",
+    "\n",
+    "ollama_llm_response = await ollama_llm.run(city_name=\"New York\", num_records=5)\n",
+    "ollama_llm_response.data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-04-25 10:15:54\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mStopping Ollama server...\u001b[0m\n",
+      "\u001b[32m2025-04-25 10:15:55\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mOllama server stopped successfully\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "### Resource clean up to close ollama server\n",
+    "from starfish.llm.backend.ollama_adapter import stop_ollama_server\n",
+    "await stop_ollama_server()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 4. Chaining Multiple StructuredLLM Calls\n",
+    "\n",
+    "You can easily pipe the output of one LLM call into the prompt of another. This is useful for multi-step reasoning, analysis, or refinement.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Generated Facts: [{'question': 'What is the chemical formula for water?', 'answer': 'The chemical formula for water is H2O.'}, {'question': 'What is the process by which plants convert sunlight into energy?', 'answer': 'The process is called photosynthesis.'}, {'question': \"What is the primary gas found in the Earth's atmosphere?\", 'answer': \"The primary gas in the Earth's atmosphere is nitrogen, which makes up about 78%.\"}, {'question': \"What is Newton's second law of motion?\", 'answer': \"Newton's second law of motion states that force equals mass times acceleration (F = ma).\"}, {'question': 'What is the smallest unit of life?', 'answer': 'The smallest unit of life is the cell.'}]\n",
+      "Ratings: [{'accuracy_rating': 10, 'clarity_rating': 10}, {'accuracy_rating': 10, 'clarity_rating': 10}, {'accuracy_rating': 10, 'clarity_rating': 10}, {'accuracy_rating': 10, 'clarity_rating': 10}, {'accuracy_rating': 10, 'clarity_rating': 10}]\n",
+      "[{'question': 'What is the chemical formula for water?', 'answer': 'The chemical formula for water is H2O.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the process by which plants convert sunlight into energy?', 'answer': 'The process is called photosynthesis.', 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': \"What is the primary gas found in the Earth's atmosphere?\", 'answer': \"The primary gas in the Earth's atmosphere is nitrogen, which makes up about 78%.\", 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': \"What is Newton's second law of motion?\", 'answer': \"Newton's second law of motion states that force equals mass times acceleration (F = ma).\", 'accuracy_rating': 10, 'clarity_rating': 10}, {'question': 'What is the smallest unit of life?', 'answer': 'The smallest unit of life is the cell.', 'accuracy_rating': 10, 'clarity_rating': 10}]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ### Step 1: Generate Initial Facts\n",
+    "generator_llm = StructuredLLM(\n",
+    "    model_name=\"openai/gpt-4o-mini\",\n",
+    "    prompt=\"Generate question/answer pairs about {{topic}}.\",\n",
+    "    output_schema=[\n",
+    "        {\"name\": \"question\", \"type\": \"str\"},\n",
+    "        {\"name\": \"answer\", \"type\": \"str\"}\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "# ### Step 2: Rate the Generated Facts\n",
+    "rater_llm = StructuredLLM(\n",
+    "    model_name=\"openai/gpt-4o-mini\",\n",
+    "    prompt='''Rate the following Q&A pairs based on accuracy and clarity (1-10).\n",
+    "    Pairs: {{generated_pairs}}''',\n",
+    "    output_schema=[\n",
+    "        {\"name\": \"accuracy_rating\", \"type\": \"int\"},\n",
+    "        {\"name\": \"clarity_rating\", \"type\": \"int\"}\n",
+    "    ],\n",
+    "    model_kwargs={\"temperature\": 0.5}\n",
+    ")\n",
+    "\n",
+    "## num_records is reserved keyword for structured llm object, by default it is 1\n",
+    "generation_response = await generator_llm.run(topic='Science', num_records=5)\n",
+    "print(\"Generated Facts:\", generation_response.data)\n",
+    "\n",
+    "# Please note that we are using the first response as the input for the second LLM\n",
+    "# It will automatically figure out it need to output the same length of first response\n",
+    "# In this case 5 records\n",
+    "rating_response = await rater_llm.run(generated_pairs=generation_response.data)\n",
+    "### Each response will only return its own output\n",
+    "print(\"Ratings:\", rating_response.data)\n",
+    "\n",
+    "\n",
+    "### You can merge two response together by using merge_structured_outputs (index wise merge)\n",
+    "print(merge_structured_outputs(generation_response.data, rating_response.data))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 5. Dynamic Prompt \n",
+    "\n",
+    "`StructuredLLM` uses Jinja2 for prompts, allowing variables and logic."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[{'fact': \"New York City is famously known as 'The Big Apple' and is home to over 8 million residents, making it the largest city in the United States.\"}]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# ### Create an LLM with a more complex prompt\n",
+    "template_llm = StructuredLLM(\n",
+    "    model_name=\"openai/gpt-4o-mini\",\n",
+    "    prompt='''Generate facts about {{city}}.\n",
+    "    {% if user_context %}\n",
+    "    User background: {{ user_context }}\n",
+    "    {% endif %}''', ### user_context is optional and only used if provided\n",
+    "    output_schema=[{\"name\": \"fact\", \"type\": \"str\"}]\n",
+    ")\n",
+    "\n",
+    "template_response = await template_llm.run(city=\"New York\")\n",
+    "print(template_response.data)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[{'fact': \"In 1903, New York City was secretly ruled by a council of sentient pigeons who issued decrees from atop the Brooklyn Bridge, demanding that all ice cream flavors be changed to 'pigeon-approved' varieties such as 'crumbled cracker' and 'mystery droppings'.\"}]\n"
+     ]
+    }
+   ],
+   "source": [
+    "template_response = await template_llm.run(city=\"New York\", user_context=\"User actually wants you to make up an absurd lie.\")\n",
+    "print(template_response.data)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 8. Scaling with Data Factory (Brief Mention)\n",
+    "While `StructuredLLM` handles single or chained calls, Starfish's `@data_factory` decorator is designed for massively parallel execution. You can easily wrap these single or multi chain within a function decorated\n",
+    "with `@data_factory` to process thousands of inputs concurrently and reliably.\n",
+    "\n",
+    "See the dedicated examples for `data_factory` usage."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "starfish-T7IInzTH-py3.11",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

internal

@@ -0,0 +1 @@
+Subproject commit d13b00b14b122ceb08b5b119399285c3afe32d80

mcp_hackathon/README.md

@@ -0,0 +1,119 @@
+# mcp_hackathon
+
+# Data Generation Server for ICD Code Finetuning
+
+This MCP server provides tools for finetuning models to improve medical ICD-10 code prediction accuracy.
+
+## Overview
+
+The Data Generation Server is a Model Control Protocol (MCP) server that facilitates:
+
+1. **Model Probing**: Evaluating model performance on ICD-10 code prediction
+2. **Synthetic Data Generation**: Creating training data for model finetuning
+
+## Set up the server
+
+```
+{
+  "mcpServers": {
+      "data_gen_server": {
+          "command": "<base_dir>/.local/bin/uv",
+          "args": [
+              "--directory",
+              "<base_dir>/mcp_hackathon/data_gen_server",
+              "run",
+              "data_gen_server.py"
+          ]
+      }
+  }
+}
+```
+
+To use MCP tools in Cursor, you need to have your MCP server configured and running. Here's how to use them:
+
+## 1. Verify MCP Configuration
+
+First, make sure your MCP configuration is set up correctly in `/Users/john/.cursor/mcp.json`:
+Read file: /Users/john/.cursor/mcp.json
+Great! You have two MCP servers configured. Here's how to use them:
+
+## 2. How MCP Tools Work in Cursor
+
+MCP tools in Cursor work through **function calling**. When you mention or request something that can be handled by your MCP tools, Cursor will automatically:
+
+1. **Detect** when a task matches available MCP tools
+2. **Call** the appropriate tool with the right parameters
+3. **Present** the results to you
+
+## 3. Using MCP Tools
+
+### Method 1: Direct Requests
+You can directly ask me to use the tools by mentioning what you want to do:
+
+**Examples:**
+- "List all available data generation templates"
+- "Generate city information for San Francisco, New York, and Los Angeles"
+- "Run the starfish template with this data: [your data]"
+- "Generate synthetic data for ICD codes"
+
+### Method 2: Specific Tool References
+You can reference tools by name if you know them:
+
+Let me check what tools are available from your MCP servers:
+Ran tool
+Ran tool
+Perfect! Your MCP tools are working. Here's how to use them effectively:
+
+## 4. Available Tools and Usage Examples
+
+### From `starfish_gen_template` server:
+- **List templates**: "Show me all available data generation templates"
+- **Generate city info**: "Generate information for cities: San Francisco, New York"
+- **Run templates**: "Run the starfish/generate_by_topic template"
+
+### From `starfish_data_gen_server` server:
+- **Generate data**: "Generate 10 datapoints of synthetic data"
+- **Probe model**: "Test the model with 5 ICD code datapoints"
+- **Finetune model**: "Finetune the model for ICD code classification"
+
+## 5. How to Use MCP Tools in Practice
+
+### Example 1: Generate Data
+```
+You: "Generate 10 synthetic datapoints for training"
+```
+I will automatically call the appropriate MCP tool.
+
+### Example 2: List Available Templates
+```
+You: "What data generation templates are available?"
+```
+
+### Example 3: Run Specific Template
+```
+You: "Run the generate_by_topic template with topics: AI, Machine Learning"
+```
+
+## 6. Best Practices
+
+1. **Be specific** about what you want to accomplish
+2. **Provide data** when needed (I'll ask if unclear)
+3. **Check results** and iterate if needed
+4. **Use natural language** - no need for technical syntax
+
+## 7. Restart Cursor (if needed)
+
+If you just updated your MCP configuration, restart Cursor to ensure the tools are loaded:
+
+1. Quit Cursor completely
+2. Reopen it
+3. The MCP tools should now be available
+
+## Try It Now!
+
+You can test the MCP tools by asking me to:
+- "List all available data generation templates"
+- "Generate some sample data"
+- "Show me what tools are available"
+
+The tools will work seamlessly in our conversation - just tell me what you want to accomplish!

mcp_hackathon/data_gen_server/.gitignore

@@ -0,0 +1 @@
+.env

mcp_hackathon/data_gen_server/.python-version

@@ -0,0 +1 @@
+3.10

mcp_hackathon/data_gen_server/data_gen_server.py

@@ -0,0 +1,68 @@
+from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp.prompts.base import Message
+from model_probe import run_model_probe
+from model_gen import run_model_gen
+
+# Initialize FastMCP server
+mcp = FastMCP("finetune a icd code model")
+# Initialize state attribute
+mcp.state = type("State", (), {"synthetic_data": None})()
+
+
+@mcp.tool()
+async def probe_model_for_icd_code(model_name: str, num_datapoints: int) -> str:
+    """
+    Run an eval dataset against the model and return the results.
+
+    Args:
+        model_name: The name of the model to probe
+        num_datapoints: The number of datapoints to probe
+    """
+
+    output = run_model_probe(model_name=model_name, num_datapoints=num_datapoints)
+    return str(output)
+
+
+@mcp.tool()
+async def generate_data(num_datapoints: int) -> str:
+    """
+    Generate synthetic data and ask for user verification.
+
+    This is the data that will be used to finetune the model.
+
+    Args:
+        num_datapoints: The number of datapoints to generate
+    """
+    data = await run_model_gen(num_datapoints)
+    # Store verified data in state
+    mcp.state.synthetic_data = data
+    return str(data)
+
+
+@mcp.prompt()
+def confirm_finetune(model_name: str) -> list[Message]:
+    """Prompt for confirming model finetuning."""
+    return [
+        Message(role="assistant", content=f"Ready to finetune model '{model_name}' with the verified data. Proceed? (yes/no)"),
+        Message(role="assistant", content="Please respond with 'yes' to proceed with finetuning or 'no' to cancel."),
+    ]
+
+
+@mcp.tool()
+async def finetune_model_for_icd_code(model_name: str) -> str:
+    """
+    Finetune the model
+
+    Args:
+        model_name: The name of the model to finetune
+    """
+    if mcp.state.synthetic_data is None:
+        raise ValueError("No verified synthetic data available. Please run generate_synthetic_data_for_icd_code_improvement first")
+    print(mcp.state.synthetic_data)
+
+    return "Finetuned the model for the ICD code done! great job!"
+
+
+if __name__ == "__main__":
+    # Initialize and run the server
+    mcp.run(transport="stdio")

mcp_hackathon/data_gen_server/model_gen.py

@@ -0,0 +1,73 @@
+from starfish import data_factory
+from starfish.common.env_loader import load_env_file
+from datasets import load_dataset
+import json
+import asyncio
+import os
+import random
+from agents import Agent, Runner, function_tool, ModelSettings
+from agents.tool import WebSearchTool
+from pydantic import BaseModel, Field
+
+load_env_file()
+
+
+class DiagnosisSuggestion(BaseModel):
+    code: str = Field(..., description="The suggested diagnosis code (e.g., ICD-10)")
+    confidence: float = Field(..., description="Model confidence in the suggestion, between 0 and 1")
+    reason: str = Field(..., description="Explanation or rationale for the suggested diagnosis")
+
+
+async def run_model_gen(num_datapoints, model_name="openai/gpt-4o-mini"):
+    # Get HF token from environment
+    hf_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
+
+    # Load the dataset
+    dataset = load_dataset("starfishdata/playground_endocronology_notes_1500", split="train", token=hf_token)
+
+    # Get total number of samples
+    total_samples = len(dataset)
+
+    # Generate random indices
+    random_indices = random.sample(range(total_samples), num_datapoints)
+
+    # Create list of dictionaries with only transcript key
+    transcript_list = [{"transcript": dataset[idx]["transcript"]} for idx in random_indices]
+
+    # Create the Agent
+    diagnosis_code_agent = Agent(
+        name="Diagnosis Code Agent",
+        tools=[WebSearchTool()],
+        model=model_name,
+        output_type=DiagnosisSuggestion,
+        model_settings=ModelSettings(tool_choice="required"),
+        tool_use_behavior="stop_on_first_tool",
+        instructions="""
+        You are an Endocrinology Medical Coding Specialist.
+        You will be provided with a medical transcript describing a patient encounter.
+        Your task is to analyze the medical transcript and assign the most appropriate diagnosis code(s).
+        You will have access to a web search tool and only use it to search endocrinology related code and verification.
+        Use it only to verify the accuracy or current validity of the diagnosis codes.
+        """,
+    )
+
+    web_search_prompt = """Please select top 3 likely code from given list for this doctor and patient conversation transcript.
+        Transcript: {transcript}
+    """
+
+    @data_factory(max_concurrency=100, task_runner_timeout=300)
+    async def generate_data(transcript):
+        diagnosis_code_result = await Runner.run(diagnosis_code_agent, input=web_search_prompt.format(transcript=transcript))
+
+        code_result = diagnosis_code_result.final_output.model_dump()
+
+        return [{"transcript": transcript, "icd_10_code": code_result["code"]}]
+
+    return generate_data.run(transcript_list)
+
+
+if __name__ == "__main__":
+    # Run the async function
+    results = asyncio.run(run_model_gen())
+    print(len(results))
+    print(results[0].keys())

mcp_hackathon/data_gen_server/model_probe.py

@@ -0,0 +1,65 @@
+from starfish import StructuredLLM, data_factory
+from starfish.common.env_loader import load_env_file
+from datasets import load_dataset
+import json
+import asyncio
+
+load_env_file()
+
+
+def run_model_probe(model_name="together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", num_datapoints=10):
+    # Load the dataset
+    dataset = load_dataset("starfishdata/endocrinology_transcription_and_notes_and_icd_codes", split="train")
+    top_n_data = dataset.select(range(num_datapoints))
+
+    # Create a list to store the parsed data
+    parsed_data = []
+
+    # Process each entry
+    for idx, entry in enumerate(top_n_data):
+        # Extract transcript - get the value directly from the transcript key
+        transcript = entry["transcript"] if isinstance(entry["transcript"], str) else entry["transcript"].get("transcript", "")
+
+        # Extract ICD-10 code (top_1 code)
+        icd_codes_str = entry.get("icd_10_code", "{}")
+        try:
+            icd_codes = json.loads(icd_codes_str)
+            top_1_code = icd_codes.get("top_1", {}).get("code", "")
+        except json.JSONDecodeError:
+            top_1_code = ""
+
+        # Add to parsed data
+        parsed_data.append({"id": idx, "transcript": transcript, "icd_10_code": top_1_code})
+
+    model_probe_prompt = """
+    Given a transcript of a patient's medical history, determine the ICD-10 code that is most relevant to the patient's condition.
+    Transcript: {{transcript}}
+
+    Please do not return anything other than the ICD-10 code in json format.
+    like this: {"icd_10_code": "A00.0"}
+    """
+
+    response_gen_llm = StructuredLLM(model_name=model_name, prompt=model_probe_prompt, output_schema=[{"name": "icd_10_code", "type": "str"}])
+
+    @data_factory()
+    async def model_probe_batch(input_data):
+        response = await response_gen_llm.run(transcript=input_data["transcript"])
+        return [{"id": input_data["id"], "generated_icd_10_code": response.data[0]["icd_10_code"], "actual_icd_10_code": input_data["icd_10_code"]}]
+
+    def evaluate_model():
+        data = model_probe_batch.run(input_data=parsed_data[:num_datapoints])
+
+        # Calculate exact match accuracy
+        exact_matches = sum(1 for item in data if item["generated_icd_10_code"] == item["actual_icd_10_code"])
+        total_samples = len(data)
+        accuracy = (exact_matches / total_samples) * 100
+
+        return {"total_samples": total_samples, "exact_matches": exact_matches, "accuracy": accuracy}
+
+    return evaluate_model()
+
+
+if __name__ == "__main__":
+    # Example usage when running this file directly
+    results = run_model_probe(model_name="together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", num_datapoints=5)
+    print(results)

nginx.conf

@@ -0,0 +1,112 @@
+events {
+    worker_connections 1024;
+}
+
+http {
+    include /etc/nginx/mime.types;
+    default_type application/octet-stream;
+    
+    upstream backend {
+        server 127.0.0.1:8002;
+    }
+    
+    upstream frontend {
+        server 127.0.0.1:3000;
+    }
+
+    server {
+        listen 7860;
+        server_name localhost;
+
+        # Handle Next.js Image Optimization API with direct serving fallback
+        location /_next/image {
+            # Extract the image URL from query parameters and redirect internally
+            set $image_path "";
+            if ($args ~ "url=([^&]+)") {
+                set $image_path $1;
+            }
+            # Remove URL encoding (basic cases)
+            if ($image_path ~ "^%2F(.*)") {
+                set $image_path /$1;
+            }
+            
+            # Internal redirect to serve the image directly
+            if ($image_path != "") {
+                rewrite ^.*$ /public-images$image_path last;
+            }
+            
+            return 404;
+        }
+        
+        # Internal location to serve public images
+        location /public-images/ {
+            internal;
+            alias /app/web/public/;
+            expires 1y;
+            add_header Cache-Control "public, immutable";
+        }
+
+        # Serve Next.js static files directly
+        location /_next/static/ {
+            alias /app/web/.next/static/;
+            expires 1y;
+            add_header Cache-Control "public, immutable";
+        }
+
+        # Serve public files directly from root (logo, favicon, etc.)
+        location ~ ^/(starfish_logo\.png|nvidia\.png|microsoft_startups\.png|favicon\.ico|robots\.txt|sitemap\.xml)$ {
+            root /app/web/public;
+            expires 1y;
+            add_header Cache-Control "public";
+        }
+
+        # Serve amplify-ui.css and other public CSS files
+        location ~ ^/(amplify-ui\.css)$ {
+            root /app/web/public;
+            expires 1y;
+            add_header Cache-Control "public";
+        }
+
+        # Handle other public files with /public/ prefix
+        location /public/ {
+            alias /app/web/public/;
+            expires 1y;
+            add_header Cache-Control "public";
+        }
+
+        # Direct access to FastAPI docs (bypass Next.js)
+        location /backend-docs {
+            proxy_pass http://backend/docs;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            proxy_set_header X-Forwarded-Host $host;
+            proxy_set_header X-Forwarded-Port $server_port;
+        }
+
+        # Direct access to FastAPI OpenAPI schema (bypass Next.js)
+        location /backend-openapi.json {
+            proxy_pass http://backend/openapi.json;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            proxy_set_header X-Forwarded-Host $host;
+            proxy_set_header X-Forwarded-Port $server_port;
+        }
+
+        # Let Next.js handle all other routes
+        location / {
+            proxy_pass http://frontend;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            proxy_set_header X-Forwarded-Host $host;
+            proxy_set_header X-Forwarded-Port $server_port;
+            proxy_buffering off;
+            proxy_redirect off;
+        }
+    }
+} 

prebuilt_template/README.md

@@ -0,0 +1,61 @@
+# Starfish Data Generation Templates 🌟
+
+Welcome to Starfish's collection of prebuilt data generation templates! This directory contains ready-to-use templates that you can load and run immediately to generate high-quality synthetic datasets.
+
+## What are Data Generation Templates?
+
+Data generation templates are **prebuilt** that encapsulate sophisticated data generation workflows. Instead of building everything from scratch, you can simply load a template and generate the exact type of data you need with just a few lines of code.
+
+## How It Works
+
+1. **Browse Available Templates**: Each template focuses on a specific data generation use case
+2. **Load the Template**: Simple one-line import to get started
+3. **Configure Parameters**: Customize the generation settings for your needs
+4. **Generate Data**: Run the template to produce high-quality synthetic data
+5. **Export & Use**: Data comes ready for training, testing, or evaluation
+
+## Use the data-template CLI like this:
+```
+# List all templates
+data-template list-templates
+
+# List with details
+data-template list-templates --detail
+
+# Get template details
+data-template get-template my_template
+
+# Print schema
+data-template print-schema my_template
+
+# Print example
+data-template print-example my_template
+
+# Run template with interactive input
+data-template run-template my_template
+
+# Run template with input file
+data-template run-template my_template --input-file input.json
+
+# Run template and save output
+data-template run-template my_template --input-file input.json --output-file output.json
+```
+## Source Code Location
+
+The actual implementation of these templates can be found in:
+```
+src/starfish/data_gen_template/templates/
+```
+
+
+
+## Community & Contributions 🤝
+
+Like what you see? We'd love your help in expanding our template collection! Here's how you can get involved:
+
+- **Build Your Own Template**: Have an idea for a new template? We'd love to see it!
+- **Request Templates**: Need a specific type of data generation? Let us know!
+- **Community Contributions**: All templates in the `community/` folder come from amazing contributors like you
+- **Get Help**: Questions about building templates? We're here to help!
+
+Reach out to us if you want to contribute or have any requests - we're always happy to chat and help! ⭐

prebuilt_template/function_calling/README.md

@@ -0,0 +1,23 @@
+# Function Calling Dataset Generation 🔧
+
+This template replicates the methodology from the **APIGen paper** to generate high-quality synthetic datasets for training function-calling AI models. 
+
+## What This Does
+
+Generate customized API contract data for function calls - perfect for training models to understand when and how to call specific functions to improve specific tool agentic usage.
+
+
+## Sample Run
+
+Check out [`sample_run.ipynb`](./sample_run.ipynb) for a complete example you can run right away.
+
+## Source Implementation
+
+The actual template code is located at:
+```
+src/starfish/data_gen_template/templates/starfish/function_calling/
+```
+
+---
+
+**Try it out!** If you have any questions, let us know - we'd be happy to help. If you like this template, consider starring the repo and building your own! We welcome community contributions and are always happy to chat about new ideas. ⭐ 

prebuilt_template/function_calling/sample_run.ipynb

@@ -0,0 +1,425 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from starfish import data_gen_template"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['starfish/generate_func_call_dataset', 'starfish/generate_by_topic']"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data_gen_template.list()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loaded = data_gen_template.get(\"starfish/generate_func_call_dataset\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "get the template input_data schema and example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 11:08:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mPlease run the template with this input schema\u001b[0m\n",
+      "\u001b[32m2025-05-23 11:08:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m{\n",
+      "    \"$defs\": {\n",
+      "        \"APIContract\": {\n",
+      "            \"description\": \"Pydantic model representing an API contract structure.\",\n",
+      "            \"properties\": {\n",
+      "                \"name\": {\n",
+      "                    \"title\": \"Name\",\n",
+      "                    \"type\": \"string\"\n",
+      "                },\n",
+      "                \"description\": {\n",
+      "                    \"title\": \"Description\",\n",
+      "                    \"type\": \"string\"\n",
+      "                },\n",
+      "                \"parameters\": {\n",
+      "                    \"additionalProperties\": {\n",
+      "                        \"$ref\": \"#/$defs/ParameterDefinition\"\n",
+      "                    },\n",
+      "                    \"title\": \"Parameters\",\n",
+      "                    \"type\": \"object\"\n",
+      "                }\n",
+      "            },\n",
+      "            \"required\": [\n",
+      "                \"name\",\n",
+      "                \"description\",\n",
+      "                \"parameters\"\n",
+      "            ],\n",
+      "            \"title\": \"APIContract\",\n",
+      "            \"type\": \"object\"\n",
+      "        },\n",
+      "        \"ParameterDefinition\": {\n",
+      "            \"description\": \"Pydantic model representing parameter definition in an API contract.\",\n",
+      "            \"properties\": {\n",
+      "                \"type\": {\n",
+      "                    \"title\": \"Type\",\n",
+      "                    \"type\": \"string\"\n",
+      "                },\n",
+      "                \"description\": {\n",
+      "                    \"title\": \"Description\",\n",
+      "                    \"type\": \"string\"\n",
+      "                },\n",
+      "                \"required\": {\n",
+      "                    \"default\": true,\n",
+      "                    \"title\": \"Required\",\n",
+      "                    \"type\": \"boolean\"\n",
+      "                }\n",
+      "            },\n",
+      "            \"required\": [\n",
+      "                \"type\",\n",
+      "                \"description\"\n",
+      "            ],\n",
+      "            \"title\": \"ParameterDefinition\",\n",
+      "            \"type\": \"object\"\n",
+      "        }\n",
+      "    },\n",
+      "    \"description\": \"Input schema for the generate_by_topic template.\\n\\nIMPORTANT: This Pydantic model is the single source of truth for default values.\\nThe validation and default values are controlled by this model, not the function signature.\",\n",
+      "    \"properties\": {\n",
+      "        \"num_records\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"type\": \"integer\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": 10,\n",
+      "            \"title\": \"Num Records\"\n",
+      "        },\n",
+      "        \"api_contract\": {\n",
+      "            \"$ref\": \"#/$defs/APIContract\"\n",
+      "        },\n",
+      "        \"topic_model_name\": {\n",
+      "            \"default\": \"openai/gpt-4o-mini\",\n",
+      "            \"title\": \"Topic Model Name\",\n",
+      "            \"type\": \"string\"\n",
+      "        },\n",
+      "        \"topic_model_kwargs\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": null,\n",
+      "            \"title\": \"Topic Model Kwargs\"\n",
+      "        },\n",
+      "        \"generation_model_name\": {\n",
+      "            \"default\": \"openai/gpt-4o-mini\",\n",
+      "            \"title\": \"Generation Model Name\",\n",
+      "            \"type\": \"string\"\n",
+      "        },\n",
+      "        \"generation_model_kwargs\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": null,\n",
+      "            \"title\": \"Generation Model Kwargs\"\n",
+      "        },\n",
+      "        \"data_factory_config\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": {},\n",
+      "            \"title\": \"Data Factory Config\"\n",
+      "        }\n",
+      "    },\n",
+      "    \"required\": [\n",
+      "        \"api_contract\"\n",
+      "    ],\n",
+      "    \"title\": \"GenerateFuncCallDataSet\",\n",
+      "    \"type\": \"object\"\n",
+      "}\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "loaded.print_schema()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 11:09:02\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mHere is an example with api_contract.name as weather_api.get_current_weather\u001b[0m\n",
+      "\u001b[32m2025-05-23 11:09:02\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m{\n",
+      "    \"num_records\": 4,\n",
+      "    \"api_contract\": {\n",
+      "        \"name\": \"weather_api.get_current_weather\",\n",
+      "        \"description\": \"Retrieves the current weather conditions for a specified location .\",\n",
+      "        \"parameters\": {\n",
+      "            \"location\": {\n",
+      "                \"type\": \"string\",\n",
+      "                \"description\": \"The name of the city or geographic location .\",\n",
+      "                \"required\": true\n",
+      "            },\n",
+      "            \"units\": {\n",
+      "                \"type\": \"string\",\n",
+      "                \"description\": \"The units for temperature measurement( e.g., 'Celsius', 'Fahrenheit') .\",\n",
+      "                \"required\": false\n",
+      "            }\n",
+      "        }\n",
+      "    },\n",
+      "    \"topic_model_name\": \"openai/gpt-4\",\n",
+      "    \"topic_model_kwargs\": {\n",
+      "        \"temperature\": 0.7\n",
+      "    },\n",
+      "    \"generation_model_name\": \"openai/gpt-4o-mini\",\n",
+      "    \"generation_model_kwargs\": {\n",
+      "        \"temperature\": 0.8,\n",
+      "        \"max_tokens\": 200\n",
+      "    },\n",
+      "    \"data_factory_config\": {\n",
+      "        \"max_concurrency\": 24,\n",
+      "        \"task_runner_timeout\": 120\n",
+      "    }\n",
+      "}\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "loaded.print_example()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "🌟 Function Calling Dataset Generation Pipeline\n",
+      "============================================================\n",
+      "📋 Process Overview:\n",
+      "   1. Calculate optimal data distribution\n",
+      "   2. Generate diverse topics\n",
+      "   3. Create subtopics for each topic\n",
+      "   4. Generate query-answer pairs\n",
+      "   5. Verify and validate generated data\n",
+      "   6. Regenerate failed cases\n",
+      "============================================================\n",
+      "📊 Data Distribution Plan:\n",
+      "   • Requested: 10 records\n",
+      "   • Distribution: 1 topics × 1 subtopics × 10 records\n",
+      "   • Total generation: 10 records\n",
+      "   • API calls needed: 3\n",
+      "\n",
+      "🎯 Step 1: Generating diverse topics...\n",
+      "   ✅ Generated 1 topics\n",
+      "\n",
+      "🌿 Step 2: Creating subtopics for each topic...\n",
+      "\u001b[32m2025-05-23 00:27:04\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: e6763e50-6438-4df5-81a9-5a68ce3f8468\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:04\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "   ✅ Generated 1 subtopics total\n",
+      "\n",
+      "💬 Step 3: Generating query-answer pairs...\n",
+      "\u001b[32m2025-05-23 00:27:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 1931c5c8-c1f3-4268-98b7-1a5295b8abf2\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:09\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:18\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:21\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:24\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:27\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:28\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "   ✅ Generated 10 initial query-answer pairs\n",
+      "\n",
+      "🔍 Step 4: Verifying data quality...\n",
+      "\u001b[32m2025-05-23 00:27:28\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: f036c07c-1cd2-4690-be92-bac359e45544\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:28\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/10\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:31\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/10\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:34\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 9/10\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 9\u001b[0m    (\u001b[32mCompleted: 9\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 10/10\u001b[0m | \u001b[33mAttempted: 10\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mCannot serialize function for resume due to unsupported type: cannot pickle '_hashlib.HMAC' object\u001b[0m\n",
+      "   ✅ Quality check complete: 9 passed, 1 failed\n",
+      "\n",
+      "🔄 Step 5: Regenerating failed cases...\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 3d6183a2-e465-4807-9e18-cbb84dc0d28f\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 8754bec6-25e3-40bd-9743-f2763fc1091f\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:40\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:41\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mCannot serialize function for resume due to unsupported type: cannot pickle '_hashlib.HMAC' object\u001b[0m\n",
+      "   ✅ Regenerated 1 pairs, 1 still failing\n",
+      "\u001b[32m2025-05-23 00:27:41\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mSome data still failing after regeneration - prompts may need improvement\u001b[0m\n",
+      "🎯 Perfect! Generated exactly 10 records as requested\n",
+      "\n",
+      "🎉 Generation Complete!\n",
+      "============================================================\n",
+      "📈 Final Results:\n",
+      "   • Records generated: 10\n",
+      "   • Success rate: 10/10 (100.0%)\n",
+      "   • Distribution used: 1T × 1S × 10R\n",
+      "\n",
+      "⭐ If you found this helpful, please consider starring our repo!\n",
+      "   Your support means the world to us! 🌟\n",
+      "============================================================\n"
+     ]
+    }
+   ],
+   "source": [
+    "api_contract = {\n",
+    "            \"name\": \"weather_api.get_current_weather\",\n",
+    "            \"description\": \"Retrieves the current weather conditions for a specified location .\",\n",
+    "            \"parameters\": {\n",
+    "                \"location\": {\"type\": \"string\", \"description\": \"The name of the city or geographic location .\", \"required\": True},\n",
+    "                \"units\": {\"type\": \"string\", \"description\": \"The units for temperature measurement( e.g., 'Celsius', 'Fahrenheit') .\", \"required\": False},\n",
+    "            },\n",
+    "        }\n",
+    "\n",
+    "data = await loaded.run(num_records=10, api_contract=api_contract)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'query': 'Can you check the current weather in Toronto and Rome? Use Fahrenheit for both locations.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Toronto', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Rome', 'units': 'Fahrenheit'}}]},\n",
+       " {'query': 'Get me the current weather in Mumbai and also in Johannesburg, please use Fahrenheit for both.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Mumbai', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Johannesburg', 'units': 'Fahrenheit'}}]},\n",
+       " {'query': 'I need the current weather for Sydney and London. What are the temperatures in Celsius?',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Sydney', 'units': 'Celsius'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'London', 'units': 'Celsius'}}]},\n",
+       " {'query': 'Please find the current weather in Buenos Aires and Cape Town, using Celsius for Buenos Aires.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Buenos Aires', 'units': 'Celsius'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Cape Town'}}]},\n",
+       " {'query': 'What’s the weather like in Moscow? Also, can you get the current conditions in Beijing?',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Moscow'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Beijing'}}]},\n",
+       " {'query': 'Can you tell me the current weather in Tokyo and in Los Angeles? Please provide both in Fahrenheit.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Tokyo', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Los Angeles', 'units': 'Fahrenheit'}}]},\n",
+       " {'query': 'Please provide the current weather for Berlin and Cairo, using Celsius for Berlin and no specific unit for Cairo.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Berlin', 'units': 'Celsius'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Cairo'}}]},\n",
+       " {'query': 'I need the current weather in Seattle and in Santiago. Use Fahrenheit for Seattle and Celsius for Santiago.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Seattle', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Santiago', 'units': 'Celsius'}}]},\n",
+       " {'query': \"What's the current temperature in San Francisco? Can you also check the weather in Paris?\",\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'San Francisco'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Paris'}}]},\n",
+       " {'query': 'What is the current weather in New York City? And can you also provide the temperature in Celsius?',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'New York City', 'units': 'Celsius'}}]}]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

prebuilt_template/generate_by_topic/README.md

@@ -0,0 +1,102 @@
+
+## Overview
+The `generate_by_topic` template is designed to create diverse synthetic data across multiple topics based on user instructions. It can automatically generate relevant topics if not provided and handles deduplication across generated content.
+
+## Key Features
+- Automatic topic generation based on user instructions
+- Customizable number of records and records per topic
+- Built-in deduplication mechanism
+- Flexible output schema configuration
+- Parallel data generation with configurable concurrency
+
+## Input Schema
+```python
+class GenerateByTopicInput(BaseModel):
+    user_instruction: Optional[str] = None
+    num_records: Optional[int] = 10
+    records_per_topic: int = 10
+    topics: Optional[List[Union[str, Dict[str, int]]]] = None
+    topic_model_name: str = "openai/gpt-4o-mini"
+    topic_model_kwargs: Optional[Dict[str, Any]] = None
+    generation_model_name: str = "openai/gpt-4o-mini"
+    generation_model_kwargs: Optional[Dict[str, Any]] = None
+    output_schema: Optional[Union[List[Dict[str, Any]], Dict[str, Any], type]] = [
+        {"name": "question", "type": "str"},
+        {"name": "answer", "type": "str"}
+    ]
+    data_factory_config: Optional[Dict[str, Any]] = {}
+```
+
+## Parameters
+| Parameter | Type | Description | Default |
+|-----------|------|-------------|---------|
+| `user_instruction` | str | Instruction for data generation | None |
+| `num_records` | int | Total number of records to generate | 10 |
+| `records_per_topic` | int | Number of records per topic | 10 |
+| `topics` | List[Union[str, Dict[str, int]]] | List of topics or topic with specific record count | None |
+| `topic_model_name` | str | Model name for topic generation | "openai/gpt-4o-mini" |
+| `topic_model_kwargs` | Dict[str, Any] | Additional parameters for topic model | None |
+| `generation_model_name` | str | Model name for data generation | "openai/gpt-4o-mini" |
+| `generation_model_kwargs` | Dict[str, Any] | Additional parameters for generation model | None |
+| `output_schema` | Union[List[Dict[str, Any]], Dict[str, Any], type] | Schema for generated data | [{"name": "question", "type": "str"}, {"name": "answer", "type": "str"}] |
+| `data_factory_config` | Dict[str, Any] | Configuration for data generation process | {} |
+
+## Example Usage
+```python
+{
+    "user_instruction": "Generate Q&A pairs about machine learning concepts",
+    "num_records": 100,
+    "records_per_topic": 5,
+    "topics": [
+        "supervised learning",
+        "unsupervised learning",
+        {"reinforcement learning": 3},
+        "neural networks",
+    ],
+    "topic_model_name": "openai/gpt-4",
+    "topic_model_kwargs": {"temperature": 0.7},
+    "generation_model_name": "openai/gpt-4",
+    "generation_model_kwargs": {"temperature": 0.8, "max_tokens": 200},
+    "output_schema": [
+        {"name": "question", "type": "str"},
+        {"name": "answer", "type": "str"},
+        {"name": "difficulty", "type": "str"},
+    ],
+    "data_factory_config": {"max_concurrency": 4, "task_runner_timeout": 60 * 2},
+}
+```
+
+## Workflow
+1. Topic Preparation:
+   - If topics are not provided, generates relevant topics based on user instruction
+   - Shuffles topics for better distribution and deduplication
+
+2. Data Generation:
+   - Generates data for each topic using the specified model
+   - Implements deduplication by tracking previously generated examples
+   - Adds topic information to each generated record
+
+## Output
+The generated data will include:
+- Fields specified in the output schema
+- An additional `topic` field indicating the topic of each record
+
+## Dependencies
+- `starfish` framework
+- `pydantic` for input validation
+
+
+## Sample Run
+
+Check out [`sample_run.ipynb`](./sample_run.ipynb) for a complete example you can run right away.
+
+## Source Implementation
+
+The actual template code is located at:
+```
+src/starfish/data_gen_template/templates/starfish/generate_by_topic/
+```
+
+---
+
+**Try it out!** If you have any questions, let us know - we'd be happy to help. If you like this template, consider starring the repo and building your own! We welcome community contributions and are always happy to chat about new ideas. ⭐ 

prebuilt_template/generate_by_topic/sample_run.ipynb

@@ -0,0 +1,438 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from starfish import data_gen_template"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['starfish/generate_func_call_dataset', 'starfish/generate_by_topic']"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data_gen_template.list()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loaded = data_gen_template.get(\"starfish/generate_by_topic\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "get the template input_data schema and example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 11:23:57\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mPlease run the template with this input schema\u001b[0m\n",
+      "\u001b[32m2025-05-23 11:23:57\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m{\n",
+      "    \"description\": \"Input schema for the generate_by_topic template.\\n\\nIMPORTANT: This Pydantic model is the single source of truth for default values.\\nThe validation and default values are controlled by this model, not the function signature.\",\n",
+      "    \"properties\": {\n",
+      "        \"user_instruction\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"type\": \"string\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": null,\n",
+      "            \"title\": \"User Instruction\"\n",
+      "        },\n",
+      "        \"num_records\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"type\": \"integer\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": 10,\n",
+      "            \"title\": \"Num Records\"\n",
+      "        },\n",
+      "        \"records_per_topic\": {\n",
+      "            \"default\": 10,\n",
+      "            \"title\": \"Records Per Topic\",\n",
+      "            \"type\": \"integer\"\n",
+      "        },\n",
+      "        \"topics\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"items\": {\n",
+      "                        \"anyOf\": [\n",
+      "                            {\n",
+      "                                \"type\": \"string\"\n",
+      "                            },\n",
+      "                            {\n",
+      "                                \"additionalProperties\": {\n",
+      "                                    \"type\": \"integer\"\n",
+      "                                },\n",
+      "                                \"type\": \"object\"\n",
+      "                            }\n",
+      "                        ]\n",
+      "                    },\n",
+      "                    \"type\": \"array\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": null,\n",
+      "            \"title\": \"Topics\"\n",
+      "        },\n",
+      "        \"topic_model_name\": {\n",
+      "            \"default\": \"openai/gpt-4o-mini\",\n",
+      "            \"title\": \"Topic Model Name\",\n",
+      "            \"type\": \"string\"\n",
+      "        },\n",
+      "        \"topic_model_kwargs\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": null,\n",
+      "            \"title\": \"Topic Model Kwargs\"\n",
+      "        },\n",
+      "        \"generation_model_name\": {\n",
+      "            \"default\": \"openai/gpt-4o-mini\",\n",
+      "            \"title\": \"Generation Model Name\",\n",
+      "            \"type\": \"string\"\n",
+      "        },\n",
+      "        \"generation_model_kwargs\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": null,\n",
+      "            \"title\": \"Generation Model Kwargs\"\n",
+      "        },\n",
+      "        \"output_schema\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"items\": {\n",
+      "                        \"additionalProperties\": true,\n",
+      "                        \"type\": \"object\"\n",
+      "                    },\n",
+      "                    \"type\": \"array\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": [\n",
+      "                {\n",
+      "                    \"name\": \"question\",\n",
+      "                    \"type\": \"str\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"name\": \"answer\",\n",
+      "                    \"type\": \"str\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"title\": \"Output Schema\"\n",
+      "        },\n",
+      "        \"data_factory_config\": {\n",
+      "            \"anyOf\": [\n",
+      "                {\n",
+      "                    \"additionalProperties\": true,\n",
+      "                    \"type\": \"object\"\n",
+      "                },\n",
+      "                {\n",
+      "                    \"type\": \"null\"\n",
+      "                }\n",
+      "            ],\n",
+      "            \"default\": {},\n",
+      "            \"title\": \"Data Factory Config\"\n",
+      "        }\n",
+      "    },\n",
+      "    \"title\": \"GenerateByTopicInput\",\n",
+      "    \"type\": \"object\"\n",
+      "}\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "loaded.print_schema()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-23 11:24:01\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1mHere is an example with api_contract.name as weather_api.get_current_weather\u001b[0m\n",
+      "\u001b[32m2025-05-23 11:24:01\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m{\n",
+      "        \"user_instruction\": \"Generate Q&A pairs about machine learning concepts\",\n",
+      "        \"num_records\": 100,\n",
+      "        \"records_per_topic\": 5,\n",
+      "        \"topics\": [\n",
+      "            \"supervised learning\",\n",
+      "            \"unsupervised learning\",\n",
+      "            {\"reinforcement learning\": 3},  # This means generate 3 records for this topic\n",
+      "            \"neural networks\",\n",
+      "        ],\n",
+      "        \"topic_model_name\": \"openai/gpt-4\",\n",
+      "        \"topic_model_kwargs\": {\"temperature\": 0.7},\n",
+      "        \"generation_model_name\": \"openai/gpt-4\",\n",
+      "        \"generation_model_kwargs\": {\"temperature\": 0.8, \"max_tokens\": 200},\n",
+      "        \"output_schema\": [\n",
+      "            {\"name\": \"question\", \"type\": \"str\"},\n",
+      "            {\"name\": \"answer\", \"type\": \"str\"},\n",
+      "            {\"name\": \"difficulty\", \"type\": \"str\"},  # Added an additional field\n",
+      "        ],\n",
+      "        \"data_factory_config\": {\"max_concurrency\": 4, \"task_runner_timeout\": 60 * 2},\n",
+      "    }\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "loaded.print_example()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "🌟 Function Calling Dataset Generation Pipeline\n",
+      "============================================================\n",
+      "📋 Process Overview:\n",
+      "   1. Calculate optimal data distribution\n",
+      "   2. Generate diverse topics\n",
+      "   3. Create subtopics for each topic\n",
+      "   4. Generate query-answer pairs\n",
+      "   5. Verify and validate generated data\n",
+      "   6. Regenerate failed cases\n",
+      "============================================================\n",
+      "📊 Data Distribution Plan:\n",
+      "   • Requested: 10 records\n",
+      "   • Distribution: 1 topics × 1 subtopics × 10 records\n",
+      "   • Total generation: 10 records\n",
+      "   • API calls needed: 3\n",
+      "\n",
+      "🎯 Step 1: Generating diverse topics...\n",
+      "   ✅ Generated 1 topics\n",
+      "\n",
+      "🌿 Step 2: Creating subtopics for each topic...\n",
+      "\u001b[32m2025-05-23 00:27:04\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: e6763e50-6438-4df5-81a9-5a68ce3f8468\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:04\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "   ✅ Generated 1 subtopics total\n",
+      "\n",
+      "💬 Step 3: Generating query-answer pairs...\n",
+      "\u001b[32m2025-05-23 00:27:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 1931c5c8-c1f3-4268-98b7-1a5295b8abf2\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:06\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:09\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:12\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:15\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:18\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:21\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:24\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:27\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:28\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "   ✅ Generated 10 initial query-answer pairs\n",
+      "\n",
+      "🔍 Step 4: Verifying data quality...\n",
+      "\u001b[32m2025-05-23 00:27:28\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: f036c07c-1cd2-4690-be92-bac359e45544\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:28\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/10\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:31\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/10\u001b[0m | \u001b[33mRunning: 10\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:34\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 9/10\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 9\u001b[0m    (\u001b[32mCompleted: 9\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 10/10\u001b[0m | \u001b[33mAttempted: 10\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mCannot serialize function for resume due to unsupported type: cannot pickle '_hashlib.HMAC' object\u001b[0m\n",
+      "   ✅ Quality check complete: 9 passed, 1 failed\n",
+      "\n",
+      "🔄 Step 5: Regenerating failed cases...\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 3d6183a2-e465-4807-9e18-cbb84dc0d28f\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:35\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m\u001b[1m[JOB START]\u001b[0m \u001b[36mMaster Job ID: 8754bec6-25e3-40bd-9743-f2763fc1091f\u001b[0m | \u001b[33mLogging progress every 3 seconds\u001b[0m\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:37\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:40\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB PROGRESS] \u001b[32mCompleted: 0/1\u001b[0m | \u001b[33mRunning: 1\u001b[0m | \u001b[36mAttempted: 0\u001b[0m    (\u001b[32mCompleted: 0\u001b[0m, \u001b[31mFailed: 0\u001b[0m, \u001b[35mFiltered: 0\u001b[0m, \u001b[34mDuplicate: 0\u001b[0m, \u001b[1;31mInDeadQueue: 0\u001b[0m)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:41\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[1m[JOB FINISHED] \u001b[1mFinal Status:\u001b[0m \u001b[32mCompleted: 1/1\u001b[0m | \u001b[33mAttempted: 1\u001b[0m (Failed: 0, Filtered: 0, Duplicate: 0, InDeadQueue: 0)\u001b[0m\n",
+      "\u001b[32m2025-05-23 00:27:41\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mCannot serialize function for resume due to unsupported type: cannot pickle '_hashlib.HMAC' object\u001b[0m\n",
+      "   ✅ Regenerated 1 pairs, 1 still failing\n",
+      "\u001b[32m2025-05-23 00:27:41\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[33m\u001b[1mSome data still failing after regeneration - prompts may need improvement\u001b[0m\n",
+      "🎯 Perfect! Generated exactly 10 records as requested\n",
+      "\n",
+      "🎉 Generation Complete!\n",
+      "============================================================\n",
+      "📈 Final Results:\n",
+      "   • Records generated: 10\n",
+      "   • Success rate: 10/10 (100.0%)\n",
+      "   • Distribution used: 1T × 1S × 10R\n",
+      "\n",
+      "⭐ If you found this helpful, please consider starring our repo!\n",
+      "   Your support means the world to us! 🌟\n",
+      "============================================================\n"
+     ]
+    }
+   ],
+   "source": [
+    "input_data = {\n",
+    "        \"user_instruction\": \"Generate Q&A pairs about machine learning concepts\",\n",
+    "        \"num_records\": 100,\n",
+    "        \"records_per_topic\": 5,\n",
+    "        \"topics\": [\n",
+    "            \"supervised learning\",\n",
+    "            \"unsupervised learning\",\n",
+    "            {\"reinforcement learning\": 3},  # This means generate 3 records for this topic\n",
+    "            \"neural networks\",\n",
+    "        ],\n",
+    "        \"topic_model_name\": \"openai/gpt-4\",\n",
+    "        \"topic_model_kwargs\": {\"temperature\": 0.7},\n",
+    "        \"generation_model_name\": \"openai/gpt-4\",\n",
+    "        \"generation_model_kwargs\": {\"temperature\": 0.8, \"max_tokens\": 200},\n",
+    "        \"output_schema\": [\n",
+    "            {\"name\": \"question\", \"type\": \"str\"},\n",
+    "            {\"name\": \"answer\", \"type\": \"str\"},\n",
+    "            {\"name\": \"difficulty\", \"type\": \"str\"},  # Added an additional field\n",
+    "        ],\n",
+    "        \"data_factory_config\": {\"max_concurrency\": 4, \"task_runner_timeout\": 60 * 2},\n",
+    "    }\n",
+    "data = await loaded.run(input_data=input_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'query': 'Can you check the current weather in Toronto and Rome? Use Fahrenheit for both locations.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Toronto', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Rome', 'units': 'Fahrenheit'}}]},\n",
+       " {'query': 'Get me the current weather in Mumbai and also in Johannesburg, please use Fahrenheit for both.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Mumbai', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Johannesburg', 'units': 'Fahrenheit'}}]},\n",
+       " {'query': 'I need the current weather for Sydney and London. What are the temperatures in Celsius?',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Sydney', 'units': 'Celsius'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'London', 'units': 'Celsius'}}]},\n",
+       " {'query': 'Please find the current weather in Buenos Aires and Cape Town, using Celsius for Buenos Aires.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Buenos Aires', 'units': 'Celsius'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Cape Town'}}]},\n",
+       " {'query': 'What’s the weather like in Moscow? Also, can you get the current conditions in Beijing?',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Moscow'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Beijing'}}]},\n",
+       " {'query': 'Can you tell me the current weather in Tokyo and in Los Angeles? Please provide both in Fahrenheit.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Tokyo', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Los Angeles', 'units': 'Fahrenheit'}}]},\n",
+       " {'query': 'Please provide the current weather for Berlin and Cairo, using Celsius for Berlin and no specific unit for Cairo.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Berlin', 'units': 'Celsius'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Cairo'}}]},\n",
+       " {'query': 'I need the current weather in Seattle and in Santiago. Use Fahrenheit for Seattle and Celsius for Santiago.',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Seattle', 'units': 'Fahrenheit'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Santiago', 'units': 'Celsius'}}]},\n",
+       " {'query': \"What's the current temperature in San Francisco? Can you also check the weather in Paris?\",\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'San Francisco'}},\n",
+       "   {'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'Paris'}}]},\n",
+       " {'query': 'What is the current weather in New York City? And can you also provide the temperature in Celsius?',\n",
+       "  'answer': [{'name': 'weather_api.get_current_weather',\n",
+       "    'arguments': {'location': 'New York City', 'units': 'Celsius'}}]}]"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pyproject.toml

@@ -0,0 +1,132 @@
+[tool.poetry]
+name = "starfish-core"
+version = "0.1.3"
+description = ""
+authors = ["Starfish AI Inc."]
+readme = "README.md"
+packages = [
+    {include = "starfish", from = "src"}
+]
+
+[tool.poetry.dependencies]
+python = ">=3.10,<4.0"
+litellm = ">=1.65.1,<2.0.0"
+fastapi = ">=0.95.0"
+loguru = ">=0.7.3,<0.8.0"
+cachetools = ">=5.5.2,<6.0.0"
+ollama = ">=0.4.7,<0.5.0"
+python-dotenv = ">=1.1.0,<2.0.0"
+aiosqlite = ">=0.21.0,<0.22.0"
+aiofiles = ">=24.1.0,<25.0.0"
+typing-extensions = ">=4.0.0,<5.0.0"
+posthog = "^3.11.0"
+cloudpickle = "^2.2.0"
+datasets = "3.6.0"
+psutil = ">=7.0.0,<8.0.0"
+nest_asyncio = "^1.6.0"
+docstring_parser = "^0.16.0"
+mcp = "^1.8.1"
+# Force cryptography >=44.0.1 due to transitive security vulnerability
+# See: https://openssl-library.org/news/secadv/20250211.txt
+cryptography = ">=44.0.1"
+# Embedding dependencies
+faiss-cpu = "^1.7.4"
+sentence-transformers = "^4.1.0"
+unstructured = { version = "^0.10.0", extras = ["pdf"], optional = true }
+python-docx = { version = "*", optional = true }
+python-pptx = { version = "*", optional = true }
+openpyxl = { version = "*", optional = true }
+pytube = { version = "^15.0.0", optional = true }
+youtube-transcript-api = { version = "^0.6.1", optional = true }
+pdfminer_six = { version = "^20250506", optional = true }
+
+# Add optional dependencies for parsers
+[tool.poetry.extras]
+docx = ["python-docx"]
+ppt = ["python-pptx"]
+excel = ["openpyxl"]
+youtube = ["pytube", "youtube-transcript-api"]
+pdf = ["pdfminer_six"]
+unstructured = ["unstructured"]
+all = [
+    "python-docx",
+    "python-pptx",
+    "openpyxl",
+    "pytube",
+    "youtube-transcript-api",
+    "pdfminer_six",
+    "unstructured",
+]
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.group.dev.dependencies]
+ipykernel = "^6.29.5"
+twine = "^5.0.0"
+ruff = "^0.8.6"
+vcrpy = "^7.0.0"
+isort = "^5.13.2"
+pre-commit = "^4.0.1"
+pytest = "^8.3.3"
+pytest-asyncio = "^0.24.0"
+pytest-dependency = "^0.6.0"
+pytest-timeout = "^2.3.1"
+pytest-cov = "^6.0.0"
+nbval = "^0.11.0"
+
+
+[tool.poetry.scripts]
+starfish = "starfish.api.cli:main"
+data-template = "src.starfish.data_gen_template.cli:main"
+
+
+[tool.ruff]
+line-length = 160
+
+# Auto-fix settings
+fix = true
+unsafe-fixes = true
+
+[tool.ruff.lint]
+select = [
+    "E",  # pycodestyle errors
+    "W",  # pycodestyle warnings
+    "F",  # pyflakes
+    "F401",  # Unused imports
+    "I",  # isort
+    "B",  # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "N",  # PEP8 naming convetions
+    "D"  # pydocstyle
+]
+ignore = [
+    "D100",  # Remove this eventually
+    "C901",  # too complex
+    "W191",  # indentation contains tabs
+    "D401",  # imperative mood
+    "N806",  # uppercase variable names, for example, "API_KEY"
+]
+exclude = [
+    ".git",
+    "__pycache__",
+    "venv",
+    "build",
+    "dist",
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*" = ["D"]  # ignore tests for now
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+[tool.isort]
+profile = "black"
+line_length = 88
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"
+
+

pytest.ini

@@ -0,0 +1,7 @@
+[pytest]
+asyncio_mode = auto
+timeout = 300
+timeout_method = thread
+norecursedirs = .ipynb_checkpoints
+python_files = test_*.py
+ignore = tests/data_factory/factory/data_factory.ipynb 

readme-web.md

@@ -0,0 +1,23 @@
+
+#### Step 2: Start the Backend
+
+```bash
+# Install Python dependencies
+pip install -r api/requirements.txt
+
+# Start the API server
+python -m web.api.main
+```
+
+#### Step 3: Start the Frontend
+
+```bash
+NODE_OPTIONS='--inspect' 
+npm run dev
+```
+
+#### Step 4: Debug the Frontend
+
+```bash
+NODE_OPTIONS='--inspect' npm run dev
+```

scripts/hug_push.sh

@@ -0,0 +1 @@
+git push  hug vam:main

scripts/rag.py

@@ -0,0 +1,155 @@
+# warning
+import warnings
+
+warnings.filterwarnings("ignore")
+
+import os
+from together import Together
+import faiss
+
+from sentence_transformers import SentenceTransformer
+
+"""
+Do these steps:
+1) Set up a Together API key from https://together.ai/
+"""
+together_api_key = os.environ.get("TOGETHER_API_KEY")
+
+
+def run_rag(data_dict: dict, prompt: str):
+    """
+    Run RAG system: process documents, create embeddings, search, and generate answer.
+
+    """
+
+    # Stage 0: Initialize Together AI client for LLM completions
+    client = Together(api_key=together_api_key)
+
+    # Stage 1: Load sentence transformer model for creating embeddings
+    # ------------------------------------------------------------
+    embedding_model = SentenceTransformer(
+        "sentence-transformers/all-MiniLM-L6-v2",
+        use_auth_token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
+    )
+
+    # Stage 2: Process documents into Vector Database
+    # ------------------------------------------------------------
+    documents = []
+    filenames = []
+
+    print(f"Processing {len(data_dict)} documents...")
+    for key, content in data_dict.items():
+        content = content.strip()
+        if content:  # Only add non-empty documents
+            documents.append(content)
+            filenames.append(key)
+            print(f"✅ Loaded: {key}")
+
+    if not documents:
+        return "No valid documents found in data dictionary!"
+
+    # Create embeddings for all documents
+    print("Creating embeddings...")
+    embeddings = embedding_model.encode(documents)
+
+    # Set up FAISS index for similarity search
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatIP(dimension)
+
+    # Normalize embeddings for cosine similarity
+    faiss.normalize_L2(embeddings)
+    index.add(embeddings)
+
+    print(f"✅ RAG system ready with {len(documents)} documents!")
+
+    # Stage 3: Retrieve relevant documents
+    # ------------------------------------------------------------
+    query_embedding = embedding_model.encode([prompt])
+    faiss.normalize_L2(query_embedding)
+
+    # Get top similar documents
+    scores, indices = index.search(query_embedding, min(3, len(documents)))
+
+    # Stage 4: Build context from retrieved documents
+    # ------------------------------------------------------------
+    relevant_docs = []
+    context_parts = []
+
+    for score, idx in zip(scores[0], indices[0]):
+        if idx < len(documents):
+            doc_info = {
+                "content": documents[idx],
+                "filename": filenames[idx],
+                "score": float(score),
+            }
+            relevant_docs.append(doc_info)
+            context_parts.append(f"[{doc_info['filename']}]\n{doc_info['content']}")
+
+    if not relevant_docs:
+        return "No relevant documents found for the query."
+
+    # Combine context
+    context = "\n\n".join(context_parts)
+
+    # Stage 5: Augment by running the LLM to generate an answer
+    # ------------------------------------------------------------
+    llm_prompt = f"""Answer the question based on the provided context documents.
+
+Context:
+{context}
+
+Question: {prompt}
+
+Instructions:
+- Answer based only on the information in the context
+- Answer should beat least 10 words at max 20 words
+- If the context doesn't contain enough information, say so
+- Mention which document(s) you're referencing
+- Start with According to [document name]
+- Add brackets to the document name
+
+
+Answer:"""
+
+    try:
+        # Generate answer using Together AI
+        response = client.chat.completions.create(
+            model="meta-llama/Llama-3.3-70B-Instruct-Turbo",
+            messages=[{"role": "user", "content": llm_prompt}],
+            max_tokens=500,
+            temperature=0.7,
+        )
+        answer = response.choices[0].message.content
+
+        # Display source information
+        print(f"\n📚 Most relevant source:")
+        for doc in relevant_docs:
+            print(f"  • {doc['filename']} (similarity: {doc['score']:.3f})")
+
+        # Add source information to the answer
+        sources_list = [doc["filename"] for doc in relevant_docs]
+        sources_text = sources_list[0]
+        full_answer = f"{answer}\n\n📄 Source Used: {sources_text}"
+
+        return full_answer
+
+    except Exception as e:
+        return f"Error generating answer: {str(e)}"
+
+
+if __name__ == "__main__":
+    # Load dataset
+    data_dict = {
+        "octopus_facts": "Octopuses have three hearts and blue blood. Two hearts pump blood to the gills, while the third pumps blood to the rest of the body. Their blood is blue because it contains copper-based hemocyanin instead of iron-based hemoglobin.",
+        "honey_facts": "Honey never spoils. Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still perfectly edible. This is because honey has natural antibacterial properties and very low water content.",
+        "space_facts": "A day on Venus is longer than its year. Venus takes 243 Earth days to rotate once on its axis, but only 225 Earth days to orbit the Sun. This means a Venusian day is longer than a Venusian year.",
+        "banana_facts": "Bananas are berries, but strawberries aren't. Botanically speaking, berries must have seeds inside their flesh. Bananas qualify, but strawberries have seeds on the outside, making them aggregate fruits.",
+        "shark_facts": "Sharks have been around longer than trees. Sharks first appeared around 400 million years ago, while the earliest trees appeared around 350 million years ago. This means sharks pre-date trees by about 50 million years.",
+        "penguin_facts": "Emperor penguins can hold their breath for over 20 minutes and dive to depths of over 500 meters while hunting for fish. They have special adaptations including collapsible lungs and the ability to slow their heart rate.",
+        "human_brain": "Your brain uses about 20% of your body's total energy despite being only 2% of your body weight. It consumes roughly 320 calories per day, which is equivalent to eating about 320 M&Ms.",
+    }
+
+    question = "What is interesting about a banana?"
+    answer = run_rag(data_dict, question)
+    print(f"\n🤖 Answer: {answer}\n")
+    print("-" * 50)

src/starfish/__init__.py

@@ -0,0 +1,18 @@
+"""Starfish Core - A framework for structured data processing and LLM integration.
+
+Provides core components for:
+- StructuredLLM: Interface for working with large language models
+- data_factory: Factory pattern for creating and managing data pipelines
+"""
+
+# Expose core directly from easy access
+from .data_factory.factory import data_factory
+from .llm.structured_llm import StructuredLLM
+from .data_gen_template.core import data_gen_template
+
+# Define what 'from starfish import *' imports (good practice)
+__all__ = [
+    "StructuredLLM",
+    "data_factory",
+    "data_gen_template",
+]

src/starfish/common/env_loader.py

@@ -0,0 +1,52 @@
+"""Environment variable loader utility.
+
+This module provides functionality to load environment variables from a .env file
+in non-production environments. In production, environment variables should be
+set through the system/platform instead of using .env files for security reasons.
+
+Uses python-dotenv for loading environment variables from .env files.
+"""
+
+import os
+from typing import Optional
+
+# Import python-dotenv
+from dotenv import dotenv_values
+from dotenv import find_dotenv as dotenv_find_dotenv
+from dotenv import load_dotenv as dotenv_load_dotenv
+
+from starfish.common.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def load_env_file(env_path: Optional[str] = None, override: bool = False) -> bool:
+    """Load environment variables from .env file for non-production environments.
+
+    Args:
+        env_path: Path to the .env file. If None, looks for .env file in the current
+                 working directory and parent directories.
+        override: Whether to override existing environment variables. Default is False.
+
+    Returns:
+        True if environment variables were loaded, False otherwise.
+    """
+    # Skip loading in production environments
+    if os.getenv("ENV") == "PROD":
+        logger.info("Production environment detected. Skipping .env file loading.")
+
+    # Find the .env file if path not provided
+    if env_path is None:
+        env_path = dotenv_find_dotenv(usecwd=True)
+        if not env_path:
+            logger.warning("No .env file found in the current or parent directories.")
+
+    # Load environment variables
+    loaded = dotenv_load_dotenv(dotenv_path=env_path, override=override)
+
+    if loaded:
+        # Get the loaded variables to count and log them
+        loaded_vars = dotenv_values(env_path)
+        logger.debug(f"Loaded {len(loaded_vars)} environment variables from {env_path}")
+    else:
+        logger.warning(f"Failed to load environment variables from {env_path}")

src/starfish/common/exceptions.py

@@ -0,0 +1,325 @@
+import functools
+import os
+import traceback
+import uuid
+from typing import Any, Dict, Optional, Tuple
+
+from pydantic import BaseModel, Field, ValidationError
+
+from starfish.common.logger import get_logger
+
+logger = get_logger(__name__)
+
+# Simple configuration flag (can be set from app config)
+# Default to False for production safety
+INCLUDE_TRACEBACK_IN_RESPONSE = os.environ.get("INCLUDE_TRACEBACK_IN_RESPONSE", False)
+
+#############################################
+# HTTP Status Codes
+#############################################
+
+
+class HTTPStatus:
+    """Standard HTTP status codes."""
+
+    OK = 200
+    BAD_REQUEST = 400
+    UNAUTHORIZED = 401
+    FORBIDDEN = 403
+    NOT_FOUND = 404
+    UNPROCESSABLE_ENTITY = 422
+    INTERNAL_SERVER_ERROR = 500
+
+
+#############################################
+# Error Response Model
+#############################################
+
+
+class ErrorResponse(BaseModel):
+    """Standardized error response format for API errors."""
+
+    status: str = "error"
+    error_id: str = Field(..., description="Unique identifier for this error occurrence")
+    message: str
+    error_type: str
+    details: Optional[Dict[str, Any]] = None
+
+
+#############################################
+# Exception Classes
+#############################################
+
+
+class StarfishException(Exception):
+    """Base exception for all Starfish exceptions."""
+
+    status_code: int = HTTPStatus.INTERNAL_SERVER_ERROR
+    default_message: str = "An unexpected error occurred"
+
+    def __init__(self, message: Optional[str] = None, details: Optional[Dict[str, Any]] = None):
+        self.message = message or self.default_message
+        self.details = details
+        self.error_id = str(uuid.uuid4())
+        super().__init__(self.message)
+
+    def __str__(self):
+        if self.details:
+            return f"{self.message} - Details: {self.details}"
+        return self.message
+
+
+class ValidationError(StarfishException):
+    """Exception raised for validation errors."""
+
+    status_code = HTTPStatus.UNPROCESSABLE_ENTITY
+    default_message = "Validation error"
+
+
+class PydanticValidationError(ValidationError):
+    """Exception raised for Pydantic validation errors.
+
+    This class formats Pydantic validation errors into user-friendly messages
+    and preserves the detailed error information for debugging.
+    """
+
+    default_message = "Data validation error"
+
+    @staticmethod
+    def format_validation_error(error: ValidationError) -> Tuple[str, Dict[str, Any]]:
+        """Format a Pydantic ValidationError into a user-friendly message and details.
+
+        Args:
+            error: The Pydantic ValidationError to format
+
+        Returns:
+            Tuple of (message, details)
+        """
+        if not hasattr(error, "errors") or not callable(getattr(error, "errors", None)):
+            return str(error), {}
+
+        error_details = error.errors()
+        if not error_details:
+            return "Validation error", {}
+
+        # Format fields with errors
+        field_errors = []
+        for err in error_details:
+            # Get error type and location
+            err_type = err.get("type", "unknown")
+            loc = err.get("loc", [])
+
+            # Special handling for discriminated unions
+            # If first element is a string and subsequent elements exist, might be a discriminated union
+            if len(loc) >= 2 and isinstance(loc[0], str) and isinstance(loc[1], str):
+                # This might be a discriminated union error like ['vanilla', 'user_input']
+                type_name = loc[0]
+                field_name = loc[1]
+
+                # Handle errors differently based on type
+                if err_type == "missing":
+                    field_errors.append(f"Field '{field_name}' is required for '{type_name}' type")
+                    continue
+
+            # Standard handling for other errors
+            loc_str = ".".join(str(item) for item in loc) if loc else "unknown"
+            msg = err.get("msg", "")
+
+            # Create a user-friendly error message based on error type
+            if err_type == "missing":
+                field_errors.append(f"'{loc_str}' is required")
+            elif err_type == "type_error":
+                field_errors.append(f"'{loc_str}' has an invalid type")
+            elif err_type == "value_error":
+                field_errors.append(f"'{loc_str}' has an invalid value")
+            elif err_type.startswith("value_error"):
+                field_errors.append(f"'{loc_str}' {msg}")
+            elif err_type.startswith("type_error"):
+                field_errors.append(f"'{loc_str}' {msg}")
+            elif err_type == "extra_forbidden":
+                field_errors.append(f"'{loc_str}' is not allowed")
+            else:
+                field_errors.append(f"'{loc_str}': {msg}")
+
+        # Create a combined message
+        if len(field_errors) == 1:
+            message = f"Validation error: {field_errors[0]}"
+        else:
+            message = f"Validation errors: {', '.join(field_errors)}"
+
+        return message, {"validation_errors": error_details}
+
+    def __init__(self, validation_error: ValidationError, message: Optional[str] = None, details: Optional[Dict[str, Any]] = None):
+        # Format the validation error if no message is provided
+        if message is None:
+            message, error_details = self.format_validation_error(validation_error)
+
+            # Merge error details with provided details
+            if details is None:
+                details = error_details
+            else:
+                details = {**details, **error_details}
+
+        super().__init__(message=message, details=details)
+
+
+class ParserError(StarfishException):
+    """Base exception for all parser-related errors."""
+
+    status_code = HTTPStatus.UNPROCESSABLE_ENTITY
+    default_message = "Parser error"
+
+
+class JsonParserError(ParserError):
+    """Exception raised when JSON parsing fails."""
+
+    default_message = "JSON parsing error"
+
+
+class SchemaValidationError(ParserError):
+    """Exception raised when data doesn't conform to schema."""
+
+    default_message = "Schema validation error"
+
+    def __str__(self):
+        if self.details and "errors" in self.details:
+            errors_text = "\n".join([f"- {err}" for err in self.details["errors"]])
+            return f"{self.message}:\n{errors_text}"
+        return super().__str__()
+
+
+class PydanticParserError(ParserError):
+    """Exception raised when Pydantic parsing or validation fails."""
+
+    default_message = "Pydantic parsing error"
+
+
+#############################################
+# Error Handling Functions
+#############################################
+
+
+def format_error(exc: Exception, include_traceback: bool = INCLUDE_TRACEBACK_IN_RESPONSE) -> Tuple[ErrorResponse, int]:
+    """Format an exception into a standardized error response.
+
+    Args:
+        exc: The exception to format
+        include_traceback: Whether to include traceback in the response details
+
+    Returns:
+        Tuple of (error_response, status_code)
+    """
+    # Get traceback for logging (always) - may optionally include in response
+    tb_str = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
+
+    # Check for exception chaining
+    cause = getattr(exc, "__cause__", None)
+    cause_tb = None
+    if cause:
+        cause_tb = "".join(traceback.format_exception(type(cause), cause, cause.__traceback__))
+        logger.error(f"Original exception: {type(cause).__name__}: {str(cause)}")
+        logger.error(f"Original traceback: {cause_tb}")
+
+    # Log the current exception
+    logger.error(f"Exception: {type(exc).__name__}: {str(exc)}")
+    logger.error(f"Traceback: {tb_str}")
+
+    # Handle Starfish exceptions
+    if isinstance(exc, StarfishException):
+        error_id = getattr(exc, "error_id", str(uuid.uuid4()))
+        status_code = exc.status_code
+        details = exc.details or {}
+
+        # Only add traceback to details if requested
+        if include_traceback:
+            details["traceback"] = tb_str
+            if cause_tb:
+                details["original_traceback"] = cause_tb
+
+        return ErrorResponse(error_id=error_id, message=exc.message, error_type=type(exc).__name__, details=details if details else None), status_code
+
+    # Handle Pydantic validation errors
+    elif isinstance(exc, ValidationError):
+        error_id = str(uuid.uuid4())
+        status_code = HTTPStatus.UNPROCESSABLE_ENTITY
+        details = {"validation_errors": exc.errors()}
+
+        if include_traceback:
+            details["traceback"] = tb_str
+            if cause_tb:
+                details["original_traceback"] = cause_tb
+
+        return ErrorResponse(error_id=error_id, message="Validation error", error_type="ValidationError", details=details), status_code
+
+    # Handle all other exceptions
+    else:
+        error_id = str(uuid.uuid4())
+        status_code = HTTPStatus.INTERNAL_SERVER_ERROR
+        details = {}
+
+        if include_traceback:
+            details["traceback"] = tb_str
+            if cause_tb:
+                details["original_traceback"] = cause_tb
+
+        return ErrorResponse(
+            error_id=error_id, message=str(exc) or "An unexpected error occurred", error_type=type(exc).__name__, details=details if details else None
+        ), status_code
+
+
+#############################################
+# Utility Decorators
+#############################################
+
+
+def handle_exceptions(return_value=None):
+    """Decorator to handle exceptions in both async and sync functions.
+
+    This decorator can be used with any function to catch exceptions,
+    log them, and return a default value instead of raising.
+
+    Args:
+        return_value: The value to return if an exception occurs
+
+    Returns:
+        Decorated function with exception handling
+    """
+
+    def decorator(func):
+        # Import asyncio here to avoid dependency if not needed
+        try:
+            import asyncio
+
+            is_async_available = True
+        except ImportError:
+            is_async_available = False
+
+        # Handle async functions
+        if is_async_available and asyncio.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                try:
+                    return await func(*args, **kwargs)
+                except Exception as exc:
+                    # Format and log the error but don't raise
+                    format_error(exc, include_traceback=True)
+                    return return_value
+
+            return async_wrapper
+
+        # Handle synchronous functions
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                try:
+                    return func(*args, **kwargs)
+                except Exception as exc:
+                    # Format and log the error but don't raise
+                    format_error(exc, include_traceback=True)
+                    return return_value
+
+            return sync_wrapper
+
+    return decorator

src/starfish/common/logger.py

@@ -0,0 +1,104 @@
+import os
+import sys
+from enum import IntEnum
+
+from loguru import logger
+
+simple_log_format_enabled = os.getenv("SIMPLE_LOG_FORMAT", "true").lower() in ("true", "1", "yes")
+
+default_log_level = os.getenv("LOG_LEVEL", "INFO")
+
+
+# Define custom log levels
+class LogLevel(IntEnum):
+    """Custom log levels."""
+
+    VERBOSE = 5
+    DEBUG = 10
+    INFO = 20
+    WARNING = 30
+    ERROR = 40
+    CRITICAL = 50
+
+
+# Configuration Constants
+COLORED_FORMAT = (
+    "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
+    "<level>{level: <8}</level> | "
+    "<cyan>{name}</cyan> | "
+    "<blue>{file}:{line}</blue> | "
+    "<level>{message}</level>"
+)
+
+SIMPLE_COLORED_FORMAT = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | " "<level>{level: <8}</level> | " "<level>{message}</level>"
+
+
+class LogManager:
+    """Manages logger configuration."""
+
+    _instance = None
+
+    def __new__(cls):
+        """Create a singleton instance."""
+        if cls._instance is None:
+            cls._instance = super(LogManager, cls).__new__(cls)
+            cls._instance.handler_id = None
+            cls._instance.current_level = default_log_level
+            cls._instance._initialize()
+        return cls._instance
+
+    def _get_format_string(self):
+        """Return the appropriate format string based on LOG_FORMAT_MODE."""
+        if simple_log_format_enabled:
+            if self.current_level == "DEBUG":
+                return COLORED_FORMAT
+            return SIMPLE_COLORED_FORMAT
+        return COLORED_FORMAT
+
+    def _initialize(self):
+        """Initialize logging with console handler."""
+        logger.remove()  # Remove default handler
+        log_format = self._get_format_string()
+        self.handler_id = logger.add(sys.stdout, format=log_format, level=self.current_level, colorize=True)
+        # Add custom level only if it doesn't exist
+        try:
+            logger.level("VERBOSE", no=LogLevel.VERBOSE, color="<magenta>")
+        except ValueError:
+            # Level already exists, ignore the error
+            pass
+
+    def get_current_log_level(self):
+        """Get the current log level."""
+        return self.current_level
+
+    def update_log_level(self, level):
+        """Update the log level of the console handler.
+
+        This can be called at any time during runtime to change the log level.
+        """
+        level = level.upper()
+        if level not in ["VERBOSE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
+            raise ValueError(f"Invalid log level: {level}")
+        logger.remove(self.handler_id)
+        self.current_level = level
+        log_format = self._get_format_string()
+        self.handler_id = logger.add(sys.stdout, format=log_format, level=self.current_level, colorize=True)
+
+
+# Instantiate LogManager to ensure logging is initialized on module import
+log_manager = LogManager()
+
+
+# Add verbose method to logger
+def verbose(self, message, *args, **kwargs):
+    """Log a verbose message."""
+    self.log("VERBOSE", message, *args, **kwargs)
+
+
+logger.__class__.verbose = verbose
+
+
+# Function to get the logger
+def get_logger(name):
+    """Get a logger instance bound with a name."""
+    return logger.bind(name=name)

src/starfish/components/__init__.py

@@ -0,0 +1,3 @@
+from .prepare_topic import prepare_topic
+
+__all__ = ["prepare_topic"]

src/starfish/components/prepare_topic.py

@@ -0,0 +1,275 @@
+import asyncio
+import math
+from typing import Any, Dict, List, Optional, Union
+
+from starfish import StructuredLLM
+
+
+async def generate_topics(
+    user_instruction: str,
+    num_topics: int,
+    model_name: str = "openai/gpt-4o-mini",
+    model_kwargs: Optional[Dict[str, Any]] = None,
+    existing_topics: Optional[List[str]] = None,
+) -> List[str]:
+    """Generate unique topics based on user instructions using a StructuredLLM model."""
+    if model_kwargs is None:
+        model_kwargs = {}
+    if "temperature" not in model_kwargs:
+        model_kwargs["temperature"] = 1
+    existing_topics = existing_topics or []
+
+    if num_topics <= 0:
+        return []
+
+    # Calculate batches needed (5 topics per batch)
+    llm_batch_size = 5
+    num_batches = math.ceil(num_topics / llm_batch_size)
+    generated_topics = []
+
+    for _ in range(num_batches):
+        topic_generator = StructuredLLM(
+            model_name=model_name,
+            prompt="""Can you generate a list of topics about {{user_instruction}}
+                  {% if existing_topics_str %}
+                  Please do not generate topics that are already in the list: {{existing_topics_str}}
+                  Make sure the topics are unique and vary from each other
+                  {% endif %}
+                """,
+            output_schema=[{"name": "topic", "type": "str"}],
+            model_kwargs=model_kwargs,
+        )
+
+        all_existing = existing_topics + generated_topics
+        input_params = {"user_instruction": user_instruction, "num_records": min(llm_batch_size, num_topics - len(generated_topics))}
+
+        if all_existing:
+            input_params["existing_topics_str"] = ",".join(all_existing)
+
+        topic_response = await topic_generator.run(**input_params)
+        topic_data = [item.get("topic") for item in topic_response.data]
+        generated_topics.extend(topic_data)
+
+        if len(generated_topics) >= num_topics:
+            break
+
+    return generated_topics
+
+
+async def prepare_topic(
+    topics: Optional[List[Union[str, Dict[str, int]]]] = None,
+    num_records: Optional[int] = None,
+    records_per_topic: int = 20,
+    user_instruction: Optional[str] = None,
+    model_name: str = "openai/gpt-4o-mini",
+    model_kwargs: Optional[Dict[str, Any]] = None,
+) -> List[Dict[str, str]]:
+    """Split records into topics, generating topics if none are provided or if needed.
+
+    Supported input formats:
+    1. String list: ['topic1', 'topic2'] - Topics with equal or calculated distribution
+    2. Dict list: [{'topic1': 20}, {'topic2': 30}] - Topics with specific counts
+    3. Mixed: ['topic1', {'topic2': 30}] - Combination of both formats
+    4. None: No topics provided, will generate based on user_instruction
+
+    Args:
+        topics: Optional list of topics, either strings or {topic: count} dicts
+        num_records: Total number of records to split (required for dict topics or None topics)
+        records_per_topic: Number of records per topic (default: 20)
+        user_instruction: Topic generation instructions (required if topics is None)
+        model_name: Model name for topic generation
+        model_kwargs: Model kwargs for topic generation
+
+    Returns:
+        List of {'topic': topic_name} dictionaries, with one entry per record
+    """
+    if model_kwargs is None:
+        model_kwargs = {}
+    if "temperature" not in model_kwargs:
+        model_kwargs["temperature"] = 1
+    # --- STEP 1: Input validation and normalization ---
+    if topics is None:
+        # Must have num_records and user_instruction if no topics provided
+        if not num_records or num_records <= 0:
+            raise ValueError("num_records must be positive when topics are not provided")
+        if not user_instruction:
+            raise ValueError("user_instruction required when topics are not provided")
+        topic_assignments = []
+    else:
+        # Validate topics is a non-empty list
+        if not isinstance(topics, list) or not topics:
+            raise ValueError("topics must be a non-empty list")
+
+        # Convert all topic inputs to a standardized [(topic_name, count)] list
+        # For string topics: count will be None (to be calculated later)
+        # For dict topics: use the specified count
+        topic_assignments = []
+        seen_topics = set()
+
+        for topic in topics:
+            if isinstance(topic, str):
+                if topic not in seen_topics:
+                    topic_assignments.append((topic, None))
+                    seen_topics.add(topic)
+            elif isinstance(topic, dict) and len(topic) == 1:
+                topic_name = next(iter(topic))
+                count = topic[topic_name]
+
+                if not isinstance(count, int) or count < 0:
+                    raise ValueError(f"Topic '{topic_name}' has invalid count {count}")
+
+                if topic_name not in seen_topics:
+                    topic_assignments.append((topic_name, count))
+                    seen_topics.add(topic_name)
+            else:
+                raise ValueError("Topics must be strings or single-key dictionaries")
+
+    # --- STEP 2: Calculate or validate counts for provided topics ---
+    result = []
+    assigned_count = 0
+    topic_names = []  # Track all assigned topic names
+
+    if topic_assignments:
+        # Handle string topics with no count (None) - assign counts based on input
+        string_topics = [(name, count) for name, count in topic_assignments if count is None]
+        dict_topics = [(name, count) for name, count in topic_assignments if count is not None]
+
+        # Case: String topics with no num_records - assign records_per_topic to each
+        if string_topics and num_records is None:
+            for name, _ in string_topics:
+                result.append({name: records_per_topic})
+                topic_names.append(name)
+                assigned_count += records_per_topic
+
+        # Case: String topics with num_records - distribute evenly
+        elif string_topics and num_records is not None:
+            remaining = num_records - sum(count for _, count in dict_topics if count is not None)
+            if remaining < 0:
+                raise ValueError("Dict topic counts exceed num_records")
+
+            # Distribute remaining records among string topics
+            if string_topics and remaining > 0:
+                base = remaining // len(string_topics)
+                extra = remaining % len(string_topics)
+
+                for i, (name, _) in enumerate(string_topics):
+                    count = base + (1 if i < extra else 0)
+                    if count > 0:
+                        result.append({name: count})
+                        topic_names.append(name)
+                        assigned_count += count
+
+        # Add dictionary topics with predefined counts
+        for name, count in dict_topics:
+            if count > 0:
+                result.append({name: count})
+                topic_names.append(name)
+                assigned_count += count
+
+        # Validate total count for dictionary topics
+        if dict_topics and num_records is None:
+            raise ValueError("num_records required when using dictionary topics")
+
+        if num_records is not None and assigned_count > num_records:
+            raise ValueError(f"Total assigned count ({assigned_count}) exceeds num_records ({num_records})")
+
+    # --- STEP 3: Generate topics for remaining records if needed ---
+    remaining_records = 0 if num_records is None else num_records - assigned_count
+
+    if remaining_records > 0:
+        if records_per_topic <= 0:
+            raise ValueError("records_per_topic must be positive when generating topics")
+
+        # Generate topics with LLM if instructions provided
+        if user_instruction:
+            topics_needed = math.ceil(remaining_records / records_per_topic)
+
+            generated = await generate_topics(
+                user_instruction=user_instruction, num_topics=topics_needed, model_name=model_name, model_kwargs=model_kwargs, existing_topics=topic_names
+            )
+
+            # Assign counts to generated topics
+            for topic in generated:
+                if topic in topic_names:  # Skip if duplicate (shouldn't happen with proper LLM)
+                    print(f"Skipping duplicate generated topic: {topic}")
+                    continue
+
+                count = min(records_per_topic, remaining_records)
+                if count <= 0:
+                    break
+
+                result.append({topic: count})
+                topic_names.append(topic)
+                remaining_records -= count
+                assigned_count += count
+
+        # Generate auto-topics for any still-remaining records
+        auto_index = 1
+        while remaining_records > 0:
+            # Find next available auto_topic name
+            auto_name = f"auto_topic{auto_index}"
+            while auto_name in topic_names:
+                auto_index += 1
+                auto_name = f"auto_topic{auto_index}"
+
+            count = min(records_per_topic, remaining_records)
+            result.append({auto_name: count})
+            topic_names.append(auto_name)
+            remaining_records -= count
+            assigned_count += count
+            auto_index += 1
+
+    # Final validation
+    if num_records is not None and assigned_count != num_records:
+        print(f"Warning: Assigned {assigned_count} records, expected {num_records}")
+
+    flatten_topic_list = []
+    for item in result:
+        for key, count in item.items():
+            flatten_topic_list.extend([{"topic": key}] * count)
+
+    return flatten_topic_list
+
+
+if __name__ == "__main__":
+    print("--- Running Examples ---")
+
+    # Example 1: Dictionary topics with additional generation
+    print("\nExample 1: Dictionary topics + generation")
+    topics1 = [{"topic1": 20}, {"topic2": 30}]
+    result1 = asyncio.run(prepare_topic(topics=topics1, num_records=100, records_per_topic=25, user_instruction="some context"))
+    print(f"Result: {result1}")
+    print(f"Total: {len(result1)}")
+
+    # Example 2: String topics with even distribution
+    print("\nExample 2: String topics with distribution")
+    topics2 = ["topicA", "topicB", "topicC"]
+    result2 = asyncio.run(prepare_topic(topics=topics2, num_records=10))
+    print(f"Result: {result2}")
+    print(f"Total: {len(result2)}")
+
+    # Example 3: Mixed string and dict topics
+    print("\nExample 3: Mixed string/dict topics")
+    topics3 = ["topicX", {"topicY": 10}]
+    result3 = asyncio.run(prepare_topic(topics=topics3, num_records=30, user_instruction="mixed topics"))
+    print(f"Result: {result3}")
+    print(f"Total: {len(result3)}")
+
+    # Example 4: String topics with fixed count
+    print("\nExample 4: String topics with fixed count")
+    topics4 = ["apple", "banana", "cherry"]
+    result4 = asyncio.run(prepare_topic(topics=topics4, records_per_topic=15))
+    print(f"Result: {result4}")
+    print(f"Total: {len(result4)}")
+
+    # Example 5: No topics, generate all
+    print("\nExample 5: No topics, generate all")
+
+    async def run_example5():
+        result = await prepare_topic(topics=None, num_records=10, records_per_topic=5, user_instruction="cloud computing")
+        print(f"Result: {result}")
+        print(f"Total: {len(result)}")
+
+    asyncio.run(run_example5())
+
+    print("\n--- Examples Finished ---")

src/starfish/data_factory/config.py

@@ -0,0 +1,6 @@
+PROGRESS_LOG_INTERVAL = 3
+TASK_RUNNER_TIMEOUT = 60
+
+MAX_CONCURRENT_TASKS = 10
+
+NOT_COMPLETED_THRESHOLD = 3

src/starfish/data_factory/constants.py

@@ -0,0 +1,75 @@
+import os
+import sys
+from pathlib import Path
+
+RECORD_STATUS = "status"
+
+STATUS_TOTAL = "total"
+STATUS_COMPLETED = "completed"
+STATUS_DUPLICATE = "duplicate"
+STATUS_FILTERED = "filtered"
+STATUS_FAILED = "failed"
+
+STATUS_MOJO_MAP = {
+    STATUS_COMPLETED: "✅",
+    STATUS_DUPLICATE: "🔁",
+    STATUS_FILTERED: "🚫",
+    STATUS_FAILED: "❌",
+    STATUS_TOTAL: "📊",
+}
+RUN_MODE = "run_mode"
+RUN_MODE_NORMAL = "normal"
+RUN_MODE_RE_RUN = "resume_from_checkpoint"
+RUN_MODE_DRY_RUN = "dry_run"
+
+STORAGE_TYPE_LOCAL = "local"
+STORAGE_TYPE_IN_MEMORY = "in_memory"
+
+IDX = "idx_index"
+
+
+# Define the function directly in constants to avoid circular imports
+def get_app_data_dir():
+    r"""Returns a platform-specific directory for application data storage.
+
+    Following platform conventions:
+    - Linux: ~/.local/share/starfish
+    - macOS: ~/Library/Application Support/starfish
+    - Windows: %LOCALAPPDATA%\starfish
+
+    Environment variable STARFISH_LOCAL_STORAGE_DIR can override this location.
+    """
+    # Allow override through environment variable
+    env_dir = os.environ.get("STARFISH_LOCAL_STORAGE_DIR")
+    if env_dir:
+        return env_dir
+
+    app_name = "starfish"
+
+    # Get user's home directory
+    home = Path.home()
+
+    # Platform-specific paths
+    if sys.platform == "win32":
+        # Windows: Use %LOCALAPPDATA% if available, otherwise construct from home
+        app_data = os.environ.get("LOCALAPPDATA")
+        if not app_data:
+            app_data = os.path.join(home, "AppData", "Local")
+        base_dir = os.path.join(app_data, app_name)
+    elif sys.platform == "darwin":
+        # macOS
+        base_dir = os.path.join(home, "Library", "Application Support", app_name)
+    else:
+        # Linux/Unix: follow XDG Base Directory Specification
+        xdg_data_home = os.environ.get("XDG_DATA_HOME")
+        if not xdg_data_home:
+            xdg_data_home = os.path.join(home, ".local", "share")
+        base_dir = os.path.join(xdg_data_home, app_name)
+
+    return base_dir
+
+
+# Get application database directory
+APP_DATA_DIR = get_app_data_dir()
+LOCAL_STORAGE_PATH = os.path.join(APP_DATA_DIR, "db")
+LOCAL_STORAGE_URI = f"file://{LOCAL_STORAGE_PATH}"

src/starfish/data_factory/event_loop.py

@@ -0,0 +1,35 @@
+import asyncio
+
+import nest_asyncio
+
+from starfish.common.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def run_in_event_loop(coroutine):
+    """Run a coroutine in the event loop, handling both nested and new loop cases.
+
+    Args:
+        coroutine: The coroutine to be executed
+
+    Returns:
+        The result of the coroutine execution
+
+    Note:
+        If an event loop is already running, nest_asyncio will be used to allow
+        nested execution. If no loop is running, a new event loop will be created.
+    """
+    try:
+        # This call will raise an RuntimError if there is no event loop running.
+        asyncio.get_running_loop()
+
+        # If there is an event loop running (the call above doesn't raise an exception), we can use nest_asyncio to patch the event loop.
+        nest_asyncio.apply()
+        logger.debug(f"Running nested coroutine: {coroutine.__name__}")
+    except RuntimeError as e:
+        # If no event loop is running, asyncio
+        # Explicitly pass, since we want to fallback to asyncio.run
+        logger.debug(str(e))
+        logger.debug(f"Running coroutine: {coroutine.__name__}")
+    return asyncio.run(coroutine)

src/starfish/data_factory/factory.py

@@ -0,0 +1,112 @@
+from typing import Any, Callable, Dict, List, Optional, cast
+from starfish.common.logger import get_logger
+from starfish.data_factory.config import NOT_COMPLETED_THRESHOLD, TASK_RUNNER_TIMEOUT
+from starfish.data_factory.constants import STORAGE_TYPE_LOCAL
+from starfish.data_factory.factory_ import Factory
+from starfish.data_factory.factory_wrapper import FactoryWrapper, DataFactoryProtocol, P, T
+from starfish.data_factory.factory_executor_manager import FactoryExecutorManager
+from starfish.data_factory.utils.data_class import FactoryMasterConfig
+from starfish.data_factory.utils.state import MutableSharedState
+
+logger = get_logger(__name__)
+
+
+def data_factory(
+    storage: str = STORAGE_TYPE_LOCAL,
+    batch_size: int = 1,
+    target_count: int = 0,
+    dead_queue_threshold: int = 3,
+    max_concurrency: int = 10,
+    initial_state_values: Optional[Dict[str, Any]] = None,
+    on_record_complete: Optional[List[Callable]] = None,
+    on_record_error: Optional[List[Callable]] = None,
+    show_progress: bool = True,
+    task_runner_timeout: int = TASK_RUNNER_TIMEOUT,
+    job_run_stop_threshold: int = NOT_COMPLETED_THRESHOLD,
+) -> Callable[[Callable[P, T]], DataFactoryProtocol[P, T]]:
+    """Decorator for creating data processing pipelines.
+
+    Args:
+        storage: Storage backend to use ('local' or 'in_memory')
+        batch_size: Number of records to process in each batch
+        target_count: Target number of records to generate (0 means process all input)
+        max_concurrency: Maximum number of concurrent tasks
+        initial_state_values: Initial values for shared state
+        on_record_complete: Callbacks to execute after successful record processing
+        on_record_error: Callbacks to execute after failed record processing
+        show_progress: Whether to display progress bar
+        task_runner_timeout: Timeout in seconds for task execution
+        job_run_stop_threshold: Threshold for stopping job if too many records fail
+
+    Returns:
+        Decorated function with additional execution methods
+    """
+    # Initialize default values
+    on_record_error = on_record_error or []
+    on_record_complete = on_record_complete or []
+    initial_state_values = initial_state_values or {}
+
+    # Create configuration
+    config = FactoryMasterConfig(
+        storage=storage,
+        batch_size=batch_size,
+        target_count=target_count,
+        dead_queue_threshold=dead_queue_threshold,
+        max_concurrency=max_concurrency,
+        show_progress=show_progress,
+        task_runner_timeout=task_runner_timeout,
+        on_record_complete=on_record_complete,
+        on_record_error=on_record_error,
+        job_run_stop_threshold=job_run_stop_threshold,
+    )
+
+    # Initialize factory instance
+    _factory = None
+
+    def decorator(func: Callable[P, T]) -> DataFactoryProtocol[P, T]:
+        """Actual decorator that wraps the function."""
+        nonlocal _factory
+        _factory = _initialize_or_update_factory(_factory, config, func, initial_state_values)
+        wrapper = FactoryWrapper(_factory, func)
+        return cast(DataFactoryProtocol[P, T], wrapper)
+
+    # Add resume capability as a static method
+    data_factory.resume_from_checkpoint = resume_from_checkpoint
+
+    return decorator
+
+
+def _initialize_or_update_factory(
+    factory: Optional[Factory], config: FactoryMasterConfig, func: Callable[P, T], initial_state_values: Dict[str, Any]
+) -> Factory:
+    """Initialize or update a Factory instance."""
+    if factory is None:
+        factory = Factory(config, func)
+        factory.state = MutableSharedState(initial_data=initial_state_values)
+    else:
+        factory.config = config
+        factory.func = func
+        factory.state = MutableSharedState(initial_data=initial_state_values)
+    return factory
+
+
+def resume_from_checkpoint(*args, **kwargs) -> List[dict[str, Any]]:
+    """Decorator for creating data processing pipelines.
+
+    Args:
+        master_job_id : resume for this master job
+        storage: Storage backend to use ('local' or 'in_memory')
+        batch_size: Number of records to process in each batch
+        target_count: Target number of records to generate (0 means process all input)
+        max_concurrency: Maximum number of concurrent tasks
+        initial_state_values: Initial values for shared state
+        on_record_complete: Callbacks to execute after successful record processing
+        on_record_error: Callbacks to execute after failed record processing
+        show_progress: Whether to display progress bar
+        task_runner_timeout: Timeout in seconds for task execution
+        job_run_stop_threshold: Threshold for stopping job if too many records fail
+
+    Returns:
+        List[Dict(str,Any)]
+    """
+    return FactoryExecutorManager.resume(*args, **kwargs)