Repository: vanna-ai/vanna
Branch: main
Commit: 365d0617c1a4
Files: 355
Total size: 2.3 MB

Directory structure:
gitextract_qmlqs8x4/

├── .gitattributes
├── .github/
│   ├── ISSUE_TEMPLATE/
│   │   ├── bug_report.md
│   │   └── feature_request.md
│   └── workflows/
│       ├── python-publish.yaml
│       └── tests.yml
├── .gitignore
├── .pre-commit-config.yaml
├── CONTRIBUTING.md
├── LICENSE
├── MIGRATION_GUIDE.md
├── README.md
├── README_LEGACY.md
├── examples/
│   ├── chromadb_gpu_example.py
│   └── transform_args_example.py
├── frontends/
│   └── webcomponent/
│       ├── .storybook/
│       │   ├── main.ts
│       │   ├── preview-head.html
│       │   └── preview.ts
│       ├── TEST_README.md
│       ├── package.json
│       ├── requirements-test.txt
│       ├── scripts/
│       │   └── sync-version.js
│       ├── src/
│       │   ├── components/
│       │   │   ├── button.stories.ts
│       │   │   ├── dataframe-component.stories.ts
│       │   │   ├── plotly-chart.stories.ts
│       │   │   ├── plotly-chart.ts
│       │   │   ├── rich-card.stories.ts
│       │   │   ├── rich-card.ts
│       │   │   ├── rich-component-system.stories.ts
│       │   │   ├── rich-component-system.ts
│       │   │   ├── rich-progress-bar.stories.ts
│       │   │   ├── rich-progress-bar.ts
│       │   │   ├── rich-task-list.stories.ts
│       │   │   ├── rich-task-list.ts
│       │   │   ├── vanna-chat.stories.ts
│       │   │   ├── vanna-chat.ts
│       │   │   ├── vanna-message.stories.ts
│       │   │   ├── vanna-message.ts
│       │   │   ├── vanna-progress-tracker.stories.ts
│       │   │   ├── vanna-progress-tracker.ts
│       │   │   ├── vanna-status-bar.stories.ts
│       │   │   └── vanna-status-bar.ts
│       │   ├── index.ts
│       │   ├── services/
│       │   │   └── api-client.ts
│       │   ├── styles/
│       │   │   ├── rich-component-styles.ts
│       │   │   └── vanna-design-tokens.ts
│       │   └── vite-env.d.ts
│       ├── test-comprehensive.html
│       ├── test_backend.py
│       ├── tsconfig.json
│       └── vite.config.ts
├── notebooks/
│   └── quickstart.ipynb
├── papers/
│   └── ai-sql-accuracy-2023-08-17.md
├── pyproject.toml
├── setup.cfg
├── src/
│   ├── evals/
│   │   ├── benchmarks/
│   │   │   └── llm_comparison.py
│   │   └── datasets/
│   │       └── sql_generation/
│   │           └── basic.yaml
│   └── vanna/
│       ├── __init__.py
│       ├── agents/
│       │   └── __init__.py
│       ├── capabilities/
│       │   ├── __init__.py
│       │   ├── agent_memory/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── file_system/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   └── sql_runner/
│       │       ├── __init__.py
│       │       ├── base.py
│       │       └── models.py
│       ├── components/
│       │   ├── __init__.py
│       │   ├── base.py
│       │   ├── rich/
│       │   │   ├── __init__.py
│       │   │   ├── containers/
│       │   │   │   ├── __init__.py
│       │   │   │   └── card.py
│       │   │   ├── data/
│       │   │   │   ├── __init__.py
│       │   │   │   ├── chart.py
│       │   │   │   └── dataframe.py
│       │   │   ├── feedback/
│       │   │   │   ├── __init__.py
│       │   │   │   ├── badge.py
│       │   │   │   ├── icon_text.py
│       │   │   │   ├── log_viewer.py
│       │   │   │   ├── notification.py
│       │   │   │   ├── progress.py
│       │   │   │   ├── status_card.py
│       │   │   │   └── status_indicator.py
│       │   │   ├── interactive/
│       │   │   │   ├── __init__.py
│       │   │   │   ├── button.py
│       │   │   │   ├── task_list.py
│       │   │   │   └── ui_state.py
│       │   │   ├── specialized/
│       │   │   │   ├── __init__.py
│       │   │   │   └── artifact.py
│       │   │   └── text.py
│       │   └── simple/
│       │       ├── __init__.py
│       │       ├── image.py
│       │       ├── link.py
│       │       └── text.py
│       ├── core/
│       │   ├── __init__.py
│       │   ├── _compat.py
│       │   ├── agent/
│       │   │   ├── __init__.py
│       │   │   ├── agent.py
│       │   │   └── config.py
│       │   ├── audit/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── component_manager.py
│       │   ├── components.py
│       │   ├── enhancer/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── default.py
│       │   ├── enricher/
│       │   │   ├── __init__.py
│       │   │   └── base.py
│       │   ├── errors.py
│       │   ├── evaluation/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   ├── dataset.py
│       │   │   ├── evaluators.py
│       │   │   ├── report.py
│       │   │   └── runner.py
│       │   ├── filter/
│       │   │   ├── __init__.py
│       │   │   └── base.py
│       │   ├── lifecycle/
│       │   │   ├── __init__.py
│       │   │   └── base.py
│       │   ├── llm/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── middleware/
│       │   │   ├── __init__.py
│       │   │   └── base.py
│       │   ├── observability/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── recovery/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── registry.py
│       │   ├── rich_component.py
│       │   ├── simple_component.py
│       │   ├── storage/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── system_prompt/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── default.py
│       │   ├── tool/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   └── models.py
│       │   ├── user/
│       │   │   ├── __init__.py
│       │   │   ├── base.py
│       │   │   ├── models.py
│       │   │   ├── request_context.py
│       │   │   └── resolver.py
│       │   ├── validation.py
│       │   └── workflow/
│       │       ├── __init__.py
│       │       ├── base.py
│       │       └── default.py
│       ├── examples/
│       │   ├── __init__.py
│       │   ├── __main__.py
│       │   ├── anthropic_quickstart.py
│       │   ├── artifact_example.py
│       │   ├── claude_sqlite_example.py
│       │   ├── coding_agent_example.py
│       │   ├── custom_system_prompt_example.py
│       │   ├── default_workflow_handler_example.py
│       │   ├── email_auth_example.py
│       │   ├── evaluation_example.py
│       │   ├── extensibility_example.py
│       │   ├── minimal_example.py
│       │   ├── mock_auth_example.py
│       │   ├── mock_custom_tool.py
│       │   ├── mock_quickstart.py
│       │   ├── mock_quota_example.py
│       │   ├── mock_rich_components_demo.py
│       │   ├── mock_sqlite_example.py
│       │   ├── openai_quickstart.py
│       │   ├── primitive_components_demo.py
│       │   ├── quota_lifecycle_example.py
│       │   └── visualization_example.py
│       ├── integrations/
│       │   ├── __init__.py
│       │   ├── anthropic/
│       │   │   ├── __init__.py
│       │   │   └── llm.py
│       │   ├── azureopenai/
│       │   │   ├── __init__.py
│       │   │   └── llm.py
│       │   ├── azuresearch/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── bigquery/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── chromadb/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── clickhouse/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── duckdb/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── faiss/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── google/
│       │   │   ├── __init__.py
│       │   │   └── gemini.py
│       │   ├── hive/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── local/
│       │   │   ├── __init__.py
│       │   │   ├── agent_memory/
│       │   │   │   ├── __init__.py
│       │   │   │   └── in_memory.py
│       │   │   ├── audit.py
│       │   │   ├── file_system.py
│       │   │   ├── file_system_conversation_store.py
│       │   │   └── storage.py
│       │   ├── marqo/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── milvus/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── mock/
│       │   │   ├── __init__.py
│       │   │   └── llm.py
│       │   ├── mssql/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── mysql/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── ollama/
│       │   │   ├── __init__.py
│       │   │   └── llm.py
│       │   ├── openai/
│       │   │   ├── __init__.py
│       │   │   ├── llm.py
│       │   │   └── responses.py
│       │   ├── opensearch/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── oracle/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── pinecone/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── plotly/
│       │   │   ├── __init__.py
│       │   │   └── chart_generator.py
│       │   ├── postgres/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── premium/
│       │   │   └── agent_memory/
│       │   │       ├── __init__.py
│       │   │       └── premium.py
│       │   ├── presto/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── qdrant/
│       │   │   ├── __init__.py
│       │   │   └── agent_memory.py
│       │   ├── snowflake/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   ├── sqlite/
│       │   │   ├── __init__.py
│       │   │   └── sql_runner.py
│       │   └── weaviate/
│       │       ├── __init__.py
│       │       └── agent_memory.py
│       ├── legacy/
│       │   ├── ZhipuAI/
│       │   │   ├── ZhipuAI_Chat.py
│       │   │   ├── ZhipuAI_embeddings.py
│       │   │   └── __init__.py
│       │   ├── __init__.py
│       │   ├── adapter.py
│       │   ├── advanced/
│       │   │   └── __init__.py
│       │   ├── anthropic/
│       │   │   ├── __init__.py
│       │   │   └── anthropic_chat.py
│       │   ├── azuresearch/
│       │   │   ├── __init__.py
│       │   │   └── azuresearch_vector.py
│       │   ├── base/
│       │   │   ├── __init__.py
│       │   │   └── base.py
│       │   ├── bedrock/
│       │   │   ├── __init__.py
│       │   │   └── bedrock_converse.py
│       │   ├── chromadb/
│       │   │   ├── __init__.py
│       │   │   └── chromadb_vector.py
│       │   ├── cohere/
│       │   │   ├── __init__.py
│       │   │   ├── cohere_chat.py
│       │   │   └── cohere_embeddings.py
│       │   ├── deepseek/
│       │   │   ├── __init__.py
│       │   │   └── deepseek_chat.py
│       │   ├── exceptions/
│       │   │   └── __init__.py
│       │   ├── faiss/
│       │   │   ├── __init__.py
│       │   │   └── faiss.py
│       │   ├── flask/
│       │   │   ├── __init__.py
│       │   │   ├── assets.py
│       │   │   └── auth.py
│       │   ├── google/
│       │   │   ├── __init__.py
│       │   │   ├── bigquery_vector.py
│       │   │   └── gemini_chat.py
│       │   ├── hf/
│       │   │   ├── __init__.py
│       │   │   └── hf.py
│       │   ├── local.py
│       │   ├── marqo/
│       │   │   ├── __init__.py
│       │   │   └── marqo.py
│       │   ├── milvus/
│       │   │   ├── __init__.py
│       │   │   └── milvus_vector.py
│       │   ├── mistral/
│       │   │   ├── __init__.py
│       │   │   └── mistral.py
│       │   ├── mock/
│       │   │   ├── __init__.py
│       │   │   ├── embedding.py
│       │   │   ├── llm.py
│       │   │   └── vectordb.py
│       │   ├── ollama/
│       │   │   ├── __init__.py
│       │   │   └── ollama.py
│       │   ├── openai/
│       │   │   ├── __init__.py
│       │   │   ├── openai_chat.py
│       │   │   └── openai_embeddings.py
│       │   ├── opensearch/
│       │   │   ├── __init__.py
│       │   │   ├── opensearch_vector.py
│       │   │   └── opensearch_vector_semantic.py
│       │   ├── oracle/
│       │   │   ├── __init__.py
│       │   │   └── oracle_vector.py
│       │   ├── pgvector/
│       │   │   ├── __init__.py
│       │   │   └── pgvector.py
│       │   ├── pinecone/
│       │   │   ├── __init__.py
│       │   │   └── pinecone_vector.py
│       │   ├── qdrant/
│       │   │   ├── __init__.py
│       │   │   └── qdrant.py
│       │   ├── qianfan/
│       │   │   ├── Qianfan_Chat.py
│       │   │   ├── Qianfan_embeddings.py
│       │   │   └── __init__.py
│       │   ├── qianwen/
│       │   │   ├── QianwenAI_chat.py
│       │   │   ├── QianwenAI_embeddings.py
│       │   │   └── __init__.py
│       │   ├── remote.py
│       │   ├── types/
│       │   │   └── __init__.py
│       │   ├── utils.py
│       │   ├── vannadb/
│       │   │   ├── __init__.py
│       │   │   └── vannadb_vector.py
│       │   ├── vllm/
│       │   │   ├── __init__.py
│       │   │   └── vllm.py
│       │   ├── weaviate/
│       │   │   ├── __init__.py
│       │   │   └── weaviate_vector.py
│       │   └── xinference/
│       │       ├── __init__.py
│       │       └── xinference.py
│       ├── py.typed
│       ├── servers/
│       │   ├── __init__.py
│       │   ├── __main__.py
│       │   ├── base/
│       │   │   ├── __init__.py
│       │   │   ├── chat_handler.py
│       │   │   ├── models.py
│       │   │   ├── rich_chat_handler.py
│       │   │   └── templates.py
│       │   ├── cli/
│       │   │   ├── __init__.py
│       │   │   └── server_runner.py
│       │   ├── fastapi/
│       │   │   ├── __init__.py
│       │   │   ├── app.py
│       │   │   └── routes.py
│       │   └── flask/
│       │       ├── __init__.py
│       │       ├── app.py
│       │       └── routes.py
│       ├── tools/
│       │   ├── __init__.py
│       │   ├── agent_memory.py
│       │   ├── file_system.py
│       │   ├── python.py
│       │   ├── run_sql.py
│       │   └── visualize_data.py
│       ├── utils/
│       │   └── __init__.py
│       └── web_components/
│           └── __init__.py
├── tests/
│   ├── conftest.py
│   ├── test_agent_memory.py
│   ├── test_agent_memory_sanity.py
│   ├── test_agents.py
│   ├── test_azureopenai_llm.py
│   ├── test_chromadb_persistence_fix.py
│   ├── test_database_sanity.py
│   ├── test_gemini_integration.py
│   ├── test_legacy_adapter.py
│   ├── test_llm_context_enhancer.py
│   ├── test_memory_tools.py
│   ├── test_ollama_direct.py
│   ├── test_tool_permissions.py
│   └── test_workflow.py
└── tox.ini

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitattributes
================================================
*.ipynb linguist-detectable=false


================================================
FILE: .github/ISSUE_TEMPLATE/bug_report.md
================================================
---
name: Bug report
about: Create a report to help us improve
title: ''
labels: ["bug"]
assignees: ''

---

**Describe the bug**
A clear and concise description of what the bug is.

**To Reproduce**
Steps to reproduce the behavior:
1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

**Expected behavior**
A clear and concise description of what you expected to happen.

**Error logs/Screenshots**
If applicable, add logs/screenshots to give more information about the issue.

**Desktop (please complete the following information where):**
 - OS: [e.g. Ubuntu]
 - Version: [e.g. 20.04]
 - Python: [3.9]
 - Vanna: [2.8.0]

**Additional context**
Add any other context about the problem here.


================================================
FILE: .github/ISSUE_TEMPLATE/feature_request.md
================================================
---
name: Feature request
about: Suggest an idea for this project
title: ''
labels: ["enhancements"]
assignees: ''

---

**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.


================================================
FILE: .github/workflows/python-publish.yaml
================================================
# This workflow will upload a Python Package using Twine when a release is created
# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries

# This workflow uses actions that are not certified by GitHub.
# They are provided by a third-party and are governed by
# separate terms of service, privacy policy, and support
# documentation.

name: Upload Python Package

on:
  release:
    types: [published]

permissions:
  contents: read

jobs:
  deploy:

    runs-on: ubuntu-latest

    steps:
    - uses: actions/checkout@v3
    - name: Set up Python
      uses: actions/setup-python@v3
      with:
        python-version: '3.x'
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install build
    - name: Build package
      run: python -m build
    - name: Publish package
      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
      with:
        user: __token__
        password: ${{ secrets.PYPI_API_TOKEN }}

================================================
FILE: .github/workflows/tests.yml
================================================
name: Basic Integration Tests

on:
  push:
    branches:
      - main

permissions:
  contents: read

jobs:
  build:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v4
    - name: Set up Python 3.11
      uses: actions/setup-python@v5
      with:
        python-version: "3.11"
    - name: Install pip
      run: |
        python -m pip install --upgrade pip
        pip install tox
    - name: Run tests
      env:
        PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
        VANNA_API_KEY: ${{ secrets.VANNA_API_KEY }}
        OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
        MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
        ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
        GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
        SNOWFLAKE_ACCOUNT: ${{ secrets.SNOWFLAKE_ACCOUNT }}
        SNOWFLAKE_USERNAME: ${{ secrets.SNOWFLAKE_USERNAME }}
        SNOWFLAKE_PASSWORD: ${{ secrets.SNOWFLAKE_PASSWORD }}
      run: tox


================================================
FILE: .gitignore
================================================
build
**.egg-info
venn
.DS_Store
tests/__pycache__
__pycache__/
.idea
.coverage
docs/*.html
.ipynb_checkpoints/
.tox/
notebooks/chroma.sqlite3
dist
.env
*.sqlite
htmlcov
chroma.sqlite3
*.bin
.coverage.*
milvus.db
.milvus.db.lock

# Frontend builds and dependencies
frontends/**/node_modules/
frontends/**/static/
frontends/**/.storybook-static/
frontends/**/package-lock.json
frontends/**/.mypy_cache/


================================================
FILE: .pre-commit-config.yaml
================================================
exclude: 'docs|node_modules|migrations|.git|.tox|assets.py'
default_stages: [ commit ]
fail_fast: true

repos:
  - repo: https://github.com/pre-commit/pre-commit-hooks
    rev: v3.2.0
    hooks:
      - id: trailing-whitespace
      - id: end-of-file-fixer
      - id: check-merge-conflict
      - id: debug-statements
      - id: mixed-line-ending

  - repo: https://github.com/pycqa/isort
    rev: 5.12.0
    hooks:
      - id: isort
        args: [ "--profile", "black", "--filter-files" ]


================================================
FILE: CONTRIBUTING.md
================================================
# Contributing to Vanna

Thank you for your interest in contributing to Vanna! This guide will help you get started with contributing to the Vanna 2.0+ codebase.

## Table of Contents

- [Getting Started](#getting-started)
- [Development Setup](#development-setup)
- [Code Standards](#code-standards)
- [Testing](#testing)
- [Pull Request Process](#pull-request-process)
- [Architecture Overview](#architecture-overview)
- [Adding New Features](#adding-new-features)

---

## Getting Started

### Prerequisites

- Python 3.11 or higher
- Git
- A GitHub account

### Fork and Clone

1. Fork the repository on GitHub
2. Clone your fork locally:
   ```bash
   git clone https://github.com/YOUR_USERNAME/vanna.git
   cd vanna
   ```

3. Add the upstream repository:
   ```bash
   git remote add upstream https://github.com/vanna-ai/vanna.git
   ```

---

## Development Setup

### 1. Create a Virtual Environment

```bash
python3 -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate
```

### 2. Install Dependencies

```bash
# Install the package in editable mode with all extras
pip install -e ".[all]"

# Install development tools
pip install tox ruff mypy pytest pytest-asyncio
```

### 3. Verify Installation

```bash
# Run unit tests
tox -e py311-unit

# Run type checking
tox -e mypy

# Run format checking
tox -e ruff
```

---

## Code Standards

### Formatting

We use [ruff](https://github.com/astral-sh/ruff) for code formatting and linting.

```bash
# Check formatting
ruff format --check src/vanna/ tests/

# Apply formatting
ruff format src/vanna/ tests/

# Run linting
ruff check src/vanna/ tests/
```

### Type Checking

We use mypy with strict mode for type checking:

```bash
tox -e mypy
```

All new code should include type hints.

### Code Style Guidelines

- Follow PEP 8 style guidelines
- Use descriptive variable and function names
- Add docstrings to all public functions and classes
- Keep functions focused and single-purpose
- Avoid circular imports by using `TYPE_CHECKING`

**Example:**

```python
"""Module docstring explaining the purpose."""

from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from vanna.core.user import User

class MyClass:
    """Class docstring explaining what this class does."""

    async def my_method(self, user: "User", count: int = 10) -> Optional[str]:
        """Method docstring explaining parameters and return value.

        Args:
            user: The user making the request
            count: Maximum number of items to return

        Returns:
            Result string if found, None otherwise
        """
        pass
```

---

## Testing

### Test Organization

Tests are organized in the `tests/` directory:

- `test_tool_permissions.py` - Tool access control tests
- `test_llm_context_enhancer.py` - LLM enhancer tests
- `test_legacy_adapter.py` - Legacy compatibility tests
- `test_agent_memory.py` - Agent memory tests
- `test_database_sanity.py` - Database integration tests
- `test_agents.py` - End-to-end agent tests

### Running Tests

```bash
# Run all unit tests (no external dependencies)
tox -e py311-unit

# Run specific test file
pytest tests/test_tool_permissions.py -v

# Run tests with a specific marker
pytest tests/ -v -m anthropic

# Run legacy adapter tests
tox -e py311-legacy
```

### Writing Tests

1. **Unit tests** should not require external dependencies (databases, APIs, etc.)
2. Use **pytest markers** for tests that require external services:
   ```python
   @pytest.mark.anthropic
   @pytest.mark.asyncio
   async def test_with_anthropic():
       # Test code here
       pass
   ```

3. **Mock external dependencies** in unit tests:
   ```python
   class MockLlmService(LlmService):
       async def send_request(self, request):
           # Mock implementation
           pass
   ```

4. **Test both success and failure cases**
5. **Use descriptive test names** that explain what is being tested

### Test Coverage

When adding new features, ensure:
- Core functionality is covered by unit tests
- Integration points are tested
- Error handling is validated
- Edge cases are considered

---

## Pull Request Process

### 1. Create a Feature Branch

```bash
git checkout -b feature/my-new-feature
# or
git checkout -b fix/bug-description
```

### 2. Make Your Changes

- Write your code following the code standards
- Add tests for your changes
- Update documentation as needed

### 3. Run All Checks

```bash
# Format code
ruff format src/vanna/ tests/

# Run linting
ruff check src/vanna/ tests/

# Run type checking
tox -e mypy

# Run tests
tox -e py311-unit
```

### 4. Commit Your Changes

Use clear, descriptive commit messages:

```bash
git add .
git commit -m "feat: add new LLM context enhancer for RAG

- Implements TextMemoryEnhancer class
- Adds tests for memory retrieval
- Updates documentation"
```

**Commit message format:**
- `feat:` - New feature
- `fix:` - Bug fix
- `docs:` - Documentation changes
- `test:` - Adding or updating tests
- `refactor:` - Code refactoring
- `chore:` - Maintenance tasks

### 5. Push and Create PR

```bash
git push origin feature/my-new-feature
```

Then create a pull request on GitHub with:
- Clear title describing the change
- Description of what was changed and why
- Link to any related issues
- Screenshots or examples if applicable

### 6. Code Review

- Address review feedback promptly
- Keep discussions focused and professional
- Be open to suggestions and alternative approaches

---

## Architecture Overview

### Core Components

Vanna 2.0+ is built around several key abstractions:

#### 1. **Agent** (`vanna.core.agent`)
The main orchestrator that coordinates tools, memory, and LLM interactions.

#### 2. **Tools** (`vanna.tools`, `vanna.core.tool`)
Modular capabilities that the agent can use. Each tool:
- Has a schema defining its inputs
- Implements an `execute()` method
- Declares access control via `access_groups`

#### 3. **Tool Registry** (`vanna.core.registry`)
Manages tool registration and access control.

#### 4. **Agent Memory** (`vanna.capabilities.agent_memory`)
Stores and retrieves tool usage patterns and documentation.

#### 5. **LLM Services** (`vanna.core.llm`)
Abstract interface for different LLM providers (Anthropic, OpenAI, etc.).

#### 6. **SQL Runners** (`vanna.capabilities.sql_runner`)
Abstract interface for executing SQL against different databases.

#### 7. **Components** (`vanna.components`)
Rich UI components for rendering results (tables, charts, status cards, etc.).

### Data Flow

```
User Request → Agent → LLM Service → Tool Selection → Tool Execution → Response Components
                ↓                                           ↓
          Agent Memory                              SQL Runner / Other Capabilities
```

---

## Adding New Features

### Adding a New Tool

1. **Create the tool class** in `src/vanna/tools/`:

```python
from vanna.core.tool import Tool, ToolContext, ToolResult
from pydantic import BaseModel, Field

class MyToolArgs(BaseModel):
    """Arguments for my tool."""
    query: str = Field(description="The query to process")

class MyTool(Tool[MyToolArgs]):
    """Tool that does something useful."""

    @property
    def name(self) -> str:
        return "my_tool"

    @property
    def description(self) -> str:
        return "Does something useful with a query"

    def get_args_schema(self) -> type[MyToolArgs]:
        return MyToolArgs

    async def execute(
        self,
        context: ToolContext,
        args: MyToolArgs
    ) -> ToolResult:
        # Implement your tool logic
        result = f"Processed: {args.query}"

        return ToolResult(
            success=True,
            result_for_llm=result,
            ui_component=None
        )
```

2. **Add tests** in `tests/test_my_tool.py`

3. **Register the tool** in examples or documentation

### Adding a New Database Integration

1. **Implement SqlRunner** in `src/vanna/integrations/mydb/`:

```python
from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext
import pandas as pd

class MyDbRunner(SqlRunner):
    """SQL runner for MyDB database."""

    def __init__(self, connection_string: str):
        self.connection_string = connection_string
        # Initialize your DB connection

    async def run_sql(
        self,
        args: RunSqlToolArgs,
        context: ToolContext
    ) -> pd.DataFrame:
        # Execute SQL and return DataFrame
        pass
```

2. **Add sanity tests** in `tests/test_database_sanity.py`

3. **Add tox target** in `tox.ini`

4. **Update documentation**

### Adding a New LLM Integration

1. **Implement LlmService** in `src/vanna/integrations/myllm/`:

```python
from vanna.core.llm.base import LlmService
from vanna.core.llm.models import LlmRequest, LlmResponse, LlmStreamChunk
from typing import AsyncGenerator

class MyLlmService(LlmService):
    """LLM service for MyLLM provider."""

    def __init__(self, api_key: str, model: str = "default"):
        self.api_key = api_key
        self.model = model

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        # Implement API call
        pass

    async def stream_request(
        self,
        request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        # Implement streaming API call
        yield LlmStreamChunk(...)

    async def validate_tools(self, tools) -> list[str]:
        # Validate tool schemas
        return []
```

2. **Add tests** with the `@pytest.mark.myllm` marker

3. **Add tox target** for integration tests

### Adding a New Agent Memory Backend

1. **Implement AgentMemory** in `src/vanna/integrations/mystore/`:

```python
from vanna.capabilities.agent_memory import (
    AgentMemory,
    ToolMemory,
    ToolMemorySearchResult,
    TextMemory,
    TextMemorySearchResult
)
from vanna.core.tool import ToolContext

class MyStoreMemory(AgentMemory):
    """Agent memory using MyStore vector database."""

    async def save_tool_usage(self, question, tool_name, args, context, success=True, metadata=None):
        # Implement storage
        pass

    async def search_similar_usage(self, question, context, *, limit=10, similarity_threshold=0.7, tool_name_filter=None):
        # Implement search
        pass

    # Implement other AgentMemory methods...
```

2. **Add tests** in `tests/test_agent_memory.py`

3. **Add to extras** in `pyproject.toml`

---

## Legacy Compatibility

If you're working on legacy VannaBase compatibility:

- The `LegacyVannaAdapter` bridges legacy code with Vanna 2.0+
- Add tests to `tests/test_legacy_adapter.py`
- See `src/vanna/legacy/adapter.py` for examples

---

## Getting Help

- **Documentation**: https://vanna.ai/docs/
- **GitHub Issues**: https://github.com/vanna-ai/vanna/issues
- **Discussions**: https://github.com/vanna-ai/vanna/discussions

---

## License

By contributing to Vanna, you agree that your contributions will be licensed under the MIT License.

---

Thank you for contributing to Vanna! 🎉


================================================
FILE: LICENSE
================================================
MIT License

Copyright (c) 2024 Vanna.AI

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


================================================
FILE: MIGRATION_GUIDE.md
================================================
# Migration Guide: Vanna 0.x to Vanna 2.0+

This guide will help you migrate from Vanna 0.x (legacy) to Vanna 2.0+, the new user-aware agent framework.

## Table of Contents
- [Overview of Changes](#overview-of-changes)
- [Quick Migration Path](#quick-migration-path)
- [Migration Strategies](#migration-strategies)
  - [Strategy 1: Using the Legacy Adapter (Recommended for Quick Migration)](#strategy-1-using-the-legacy-adapter-recommended-for-quick-migration)
  - [Strategy 2: Full Migration to New Architecture](#strategy-2-full-migration-to-new-architecture)
- [Key Architectural Differences](#key-architectural-differences)
- [API Mapping](#api-mapping)
- [Common Migration Scenarios](#common-migration-scenarios)
- [Breaking Changes](#breaking-changes)
- [FAQ](#faq)

---

## Overview of Changes

Vanna 2.0+ represents a fundamental architectural shift from a simple LLM wrapper to a full-fledged **user-aware agent framework**. Here are the major changes:

### What's New in 2.0+
- ✅ **User awareness** - Identity and permissions flow through every layer
- ✅ **Web component** - Pre-built UI with streaming responses
- ✅ **Tool registry** - Modular, extensible tool system
- ✅ **Rich UI components** - Tables, charts, status cards (not just text)
- ✅ **Streaming by default** - Progressive responses via SSE
- ✅ **Enterprise features** - Audit logs, rate limiting, observability
- ✅ **FastAPI/Flask servers** - Production-ready backends included

### What Changed from 0.x
- ❌ Direct method calls (`vn.ask()`) → Agent-based workflow
- ❌ Monolithic `VannaBase` class → Modular tool system
- ❌ No user context → User-aware at every layer
- ❌ Simple text responses → Rich streaming UI components

---

## Quick Migration Path

**Can't migrate immediately?** Use the Legacy Adapter to get started quickly:

```python
# Assume you already have a working vn object from your Vanna 0.x code:
# vn = MyVanna(config={"model": "gpt-4"})
# vn.connect_to_postgres(...)
# vn.train(ddl="...")

# NEW: Just add these imports and wrap your existing vn object
from vanna import Agent, AgentConfig
from vanna.servers.fastapi import VannaFastAPIServer
from vanna.core.user import UserResolver, User, RequestContext
from vanna.legacy.adapter import LegacyVannaAdapter
from vanna.integrations.anthropic import AnthropicLlmService

# Define simple user resolver
class SimpleUserResolver(UserResolver):
    async def resolve_user(self, request_context: RequestContext) -> User:
        user_email = request_context.get_cookie('vanna_email')
        return User(id=user_email, email=user_email, group_memberships=['user'])

# Wrap your existing vn with the adapter
tools = LegacyVannaAdapter(vn)

# Create agent with new LLM service
llm = AnthropicLlmService(model="claude-haiku-4-5")
agent = Agent(llm_service=llm, tool_registry=tools, user_resolver=SimpleUserResolver())

# Run server
server = VannaFastAPIServer(agent)
server.run(host='0.0.0.0', port=8000)

# Now it works with the new Agent framework!
# (See Strategy 1 below for complete example)
```

---

## Migration Strategies

### Strategy 1: Using the Legacy Adapter (Recommended for Quick Migration)

**Best for:** Teams that want to adopt Vanna 2.0+ gradually while maintaining existing code.

#### Step 1: Install Vanna 2.0+

```bash
pip install 'vanna[flask,anthropic]'
```

#### Step 2: Wrap Your Existing VannaBase Instance

```python
from vanna import Agent, AgentConfig
from vanna.servers.fastapi import VannaFastAPIServer
from vanna.core.user import UserResolver, User, RequestContext
from vanna.legacy.adapter import LegacyVannaAdapter
from vanna.integrations.anthropic import AnthropicLlmService

# Assume you already have a working vn object from your existing code:
# vn = MyVanna(config={'model': 'gpt-4', 'api_key': 'your-key'})
# vn.connect_to_postgres(...)
# vn.train(ddl="...")
# etc.

# NEW: Define user resolution (required in 2.0+)
class SimpleUserResolver(UserResolver):
    async def resolve_user(self, request_context: RequestContext) -> User:
        user_email = request_context.get_cookie('vanna_email')
        if not user_email:
            raise ValueError("Missing 'vanna_email' cookie")

        # Admin users get 'admin' group membership
        if user_email == "admin@example.com":
            return User(id="admin_user", email=user_email, group_memberships=['admin'])

        # Regular users get 'user' group membership
        return User(id=user_email, email=user_email, group_memberships=['user'])

# NEW: Wrap with legacy adapter
# This automatically registers run_sql and memory tools from your VannaBase instance
tools = LegacyVannaAdapter(vn)

# NEW: Set up LLM for the new Agent framework
llm = AnthropicLlmService(
    model="claude-haiku-4-5",
    api_key="YOUR_ANTHROPIC_API_KEY"
)

# NEW: Create agent with legacy adapter as tool registry
agent = Agent(
    llm_service=llm,
    tool_registry=tools,  # LegacyVannaAdapter is a ToolRegistry
    user_resolver=SimpleUserResolver(),
    config=AgentConfig()
)

# NEW: Create and run server
server = VannaFastAPIServer(agent)

if __name__ == "__main__":
    # Run with: python your_script.py
    # Or: uvicorn your_module:server --host 0.0.0.0 --port 8000
    server.run(host='0.0.0.0', port=8000)
```

**What the LegacyVannaAdapter does:**
- Automatically wraps `vn.run_sql()` as the `run_sql` tool (available to 'user' and 'admin' groups)
- Exposes training data from `vn.get_training_data()` as searchable memory (via `search_saved_correct_tool_uses` tool)
- Optionally allows saving new training data (via `save_question_tool_args` tool - admin only)
- Maintains your existing database connection and training data

**Pros:**
- ✅ Minimal code changes
- ✅ Preserve existing training data
- ✅ Gradual migration path
- ✅ Get new features (web UI, streaming) immediately

**Cons:**
- ⚠️ Limited user awareness (all requests use same VannaBase instance)
- ⚠️ Can't leverage row-level security
- ⚠️ Missing some advanced features

---

### Strategy 2: Full Migration to New Architecture

**Best for:** New projects or teams ready for a complete rewrite.

#### Before (Vanna 0.x)

```python
from vanna import VannaBase
from vanna.openai_chat import OpenAI_Chat
from vanna.chromadb import ChromaDB_VectorStore

class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

vn = MyVanna(config={'model': 'gpt-4', 'api_key': 'your-key'})
vn.connect_to_postgres(...)

# Train
vn.train(ddl="CREATE TABLE customers ...")
vn.train(question="Top customers?", sql="SELECT ...")

# Ask
sql = vn.generate_sql("Who are the top customers?")
df = vn.run_sql(sql)
print(df)
```

#### After (Vanna 2.0+)

```python
from vanna import Agent, AgentConfig
from vanna.servers.fastapi import VannaFastAPIServer
from vanna.core.registry import ToolRegistry
from vanna.core.user import UserResolver, User, RequestContext
from vanna.integrations.anthropic import AnthropicLlmService
from vanna.tools import RunSqlTool
from vanna.integrations.postgres import PostgresRunner

# 1. Define user resolution
class MyUserResolver(UserResolver):
    async def resolve_user(self, request_context: RequestContext) -> User:
        # Extract from your auth system (JWT, cookies, etc.)
        token = request_context.get_header('Authorization')
        user_data = await self.validate_token(token)

        return User(
            id=user_data['id'],
            email=user_data['email'],
            permissions=user_data['permissions'],
            metadata={'role': user_data['role']}
        )

# 2. Set up tools
tools = ToolRegistry()
postgres_runner = PostgresRunner(
    host="localhost",
    dbname="mydb",
    user="user",
    password="password",
    port=5432
)
tools.register_local_tool(
    RunSqlTool(sql_runner=postgres_runner),
    access_groups=['user', 'admin']
)

# 3. Create agent
llm = AnthropicLlmService(model="claude-sonnet-4-5")
agent = Agent(
    llm_service=llm,
    tool_registry=tools,
    user_resolver=MyUserResolver(),
    config=AgentConfig(stream_responses=True)
)

# 4. Create server
server = VannaFastAPIServer(agent)
app = server.create_app()

# Run with: uvicorn main:app --host 0.0.0.0 --port 8000
# Visit http://localhost:8000 for web UI
```

**Pros:**
- ✅ Full access to new features
- ✅ True user awareness
- ✅ Better security and permissions
- ✅ Production-ready architecture

**Cons:**
- ⚠️ Requires rewriting code
- ⚠️ Need to migrate training data approach
- ⚠️ Steeper learning curve

---

## Key Architectural Differences

| Feature | Vanna 0.x | Vanna 2.0+ |
|---------|-----------|------------|
| **User Context** | None | `User` object with permissions flows through entire system |
| **Interaction Model** | Direct method calls (`vn.ask()`) | Agent-based with streaming components |
| **Tools** | Monolithic methods | Modular `Tool` classes with schemas |
| **Responses** | Plain text/DataFrames | Rich UI components (tables, charts, code) |
| **Training** | `vn.train()` with vector DB | System prompts, context enrichers, RAG tools |
| **Database Connection** | `vn.connect_to_postgres()` | `SqlRunner` implementations as dependencies |
| **Web UI** | None (custom implementation) | Built-in web component + backend |
| **Streaming** | None | Server-Sent Events by default |
| **Permissions** | None | Group-based access control on tools |
| **Audit Logs** | None | Built-in audit logging system |

---

## Summary

| If you want to... | Use this strategy |
|-------------------|-------------------|
| Migrate quickly with minimal changes | **Strategy 1: Legacy Adapter** |
| Get full access to new features | **Strategy 2: Full Migration** |
| Support both legacy and new code | **Strategy 1** initially, then gradual migration |
| Start a new project | **Strategy 2: Full Migration** |

**Recommended Path:**
1. Start with Legacy Adapter for quick migration
2. Gradually rewrite critical paths to native 2.0+ architecture
3. Eventually remove Legacy Adapter once fully migrated

Good luck with your migration! 🚀


================================================
FILE: README.md
================================================
# Vanna 2.0: Turn Questions into Data Insights

**Natural language → SQL → Answers.** Now with enterprise security and user-aware permissions.

[![Python](https://img.shields.io/badge/python-3.8+-blue.svg)](https://python.org)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

https://github.com/user-attachments/assets/476cd421-d0b0-46af-8b29-0f40c73d6d83


![Vanna2 Demo](img/architecture.png)

---

## What's New in 2.0

🔐 **User-Aware at Every Layer** — Queries automatically filtered per user permissions

🎨 **Modern Web Interface** — Beautiful pre-built `<vanna-chat>` component

⚡ **Streaming Responses** — Real-time tables, charts, and progress updates

🔒 **Enterprise Security** — Row-level security, audit logs, rate limiting

🔄 **Production-Ready** — FastAPI integration, observability, lifecycle hooks

> **Upgrading from 0.x?** See the [Migration Guide](MIGRATION_GUIDE.md) | [What changed?](#migration-notes)

---

## Get Started

### Try it with Sample Data

[Quickstart](https://vanna.ai/docs/quick-start)

### Configure

[Configure](https://vanna.ai/docs/configure)

### Web Component

```html
<!-- Drop into any existing webpage -->
<script src="https://img.vanna.ai/vanna-components.js"></script>
<vanna-chat
  sse-endpoint="https://your-api.com/chat"
  theme="dark">
</vanna-chat>
```

Uses your existing cookies/JWTs. Works with React, Vue, or plain HTML.

---

## What You Get

Ask a question in natural language and get back:

**1. Streaming Progress Updates**

**2. SQL Code Block (By default only shown to "admin" users)**

**3. Interactive Data Table**

**4. Charts** (Plotly visualizations)

**5. Natural Language Summary**

All streamed in real-time to your web component.

---

## Why Vanna 2.0?

### ✅ Get Started Instantly
* Production chat interface
* Custom agent with your database
* Embed in any webpage

### ✅ Enterprise-Ready Security
**User-aware at every layer** — Identity flows through system prompts, tool execution, and SQL filtering
**Row-level security** — Queries automatically filtered per user permissions
**Audit logs** — Every query tracked per user for compliance
**Rate limiting** — Per-user quotas via lifecycle hooks

### ✅ Beautiful Web UI Included
**Pre-built `<vanna-chat>` component** — No need to build your own chat interface
**Streaming tables & charts** — Rich components, not just text
**Responsive & customizable** — Works on mobile, desktop, light/dark themes
**Framework-agnostic** — React, Vue, plain HTML

### ✅ Works With Your Stack
**Any LLM:** OpenAI, Anthropic, Ollama, Azure, Google Gemini, AWS Bedrock, Mistral, Others
**Any Database:** PostgreSQL, MySQL, Snowflake, BigQuery, Redshift, SQLite, Oracle, SQL Server, DuckDB, ClickHouse, Others
**Your Auth System:** Bring your own — cookies, JWTs, OAuth tokens
**Your Framework:** FastAPI, Flask

### ✅ Extensible But Opinionated
**Custom tools** — Extend the `Tool` base class
**Lifecycle hooks** — Quota checking, logging, content filtering
**LLM middlewares** — Caching, prompt engineering
**Observability** — Built-in tracing and metrics

---

## Architecture

![Vanna2 Diagram](img/vanna2.svg)

---

## How It Works

```mermaid
sequenceDiagram
    participant U as 👤 User
    participant W as 🌐 <vanna-chat>
    participant S as 🐍 Your Server
    participant A as 🤖 Agent
    participant T as 🧰 Tools

    U->>W: "Show Q4 sales"
    W->>S: POST /api/vanna/v2/chat_sse (with auth)
    S->>A: User(id=alice, groups=[read_sales])
    A->>T: Execute SQL tool (user-aware)
    T->>T: Apply row-level security
    T->>A: Filtered results
    A->>W: Stream: Table → Chart → Summary
    W->>U: Display beautiful UI
```

**Key Concepts:**

1. **User Resolver** — You define how to extract user identity from requests (cookies, JWTs, etc.)
2. **User-Aware Tools** — Tools automatically check permissions based on user's group memberships
3. **Streaming Components** — Backend streams structured UI components (tables, charts) to frontend
4. **Built-in Web UI** — Pre-built `<vanna-chat>` component renders everything beautifully

---

## Production Setup with Your Auth

Here's a complete example integrating Vanna with your existing FastAPI app and authentication:

```python
from fastapi import FastAPI
from vanna import Agent
from vanna.servers.fastapi.routes import register_chat_routes
from vanna.servers.base import ChatHandler
from vanna.core.user import UserResolver, User, RequestContext
from vanna.integrations.anthropic import AnthropicLlmService
from vanna.tools import RunSqlTool
from vanna.integrations.sqlite import SqliteRunner
from vanna.core.registry import ToolRegistry

# Your existing FastAPI app
app = FastAPI()

# 1. Define your user resolver (using YOUR auth system)
class MyUserResolver(UserResolver):
    async def resolve_user(self, request_context: RequestContext) -> User:
        # Extract from cookies, JWTs, or session
        token = request_context.get_header('Authorization')
        user_data = self.decode_jwt(token)  # Your existing logic

        return User(
            id=user_data['id'],
            email=user_data['email'],
            group_memberships=user_data['groups']  # Used for permissions
        )

# 2. Set up agent with tools
llm = AnthropicLlmService(model="claude-sonnet-4-5")
tools = ToolRegistry()
tools.register(RunSqlTool(sql_runner=SqliteRunner("./data.db")))

agent = Agent(
    llm_service=llm,
    tool_registry=tools,
    user_resolver=MyUserResolver()
)

# 3. Add Vanna routes to your app
chat_handler = ChatHandler(agent)
register_chat_routes(app, chat_handler)

# Now you have:
# - POST /api/vanna/v2/chat_sse (streaming endpoint)
# - GET / (optional web UI)
```

**Then in your frontend:**
```html
<vanna-chat sse-endpoint="/api/vanna/v2/chat_sse"></vanna-chat>
```

See [Full Documentation](https://vanna.ai/docs) for custom tools, lifecycle hooks, and advanced configuration

---

## Custom Tools

Extend Vanna with custom tools for your specific use case:

```python
from vanna.core.tool import Tool, ToolContext, ToolResult
from pydantic import BaseModel, Field
from typing import Type

class EmailArgs(BaseModel):
    recipient: str = Field(description="Email recipient")
    subject: str = Field(description="Email subject")

class EmailTool(Tool[EmailArgs]):
    @property
    def name(self) -> str:
        return "send_email"

    @property
    def access_groups(self) -> list[str]:
        return ["send_email"]  # Permission check

    def get_args_schema(self) -> Type[EmailArgs]:
        return EmailArgs

    async def execute(self, context: ToolContext, args: EmailArgs) -> ToolResult:
        user = context.user  # Automatically injected

        # Your business logic
        await self.email_service.send(
            from_email=user.email,
            to=args.recipient,
            subject=args.subject
        )

        return ToolResult(success=True, result_for_llm=f"Email sent to {args.recipient}")

# Register your tool
tools.register(EmailTool())
```

---

## Advanced Features

Vanna 2.0 includes powerful enterprise features for production use:

**Lifecycle Hooks** — Add quota checking, custom logging, content filtering at key points in the request lifecycle

**LLM Middlewares** — Implement caching, prompt engineering, or cost tracking around LLM calls

**Conversation Storage** — Persist and retrieve conversation history per user

**Observability** — Built-in tracing and metrics integration

**Context Enrichers** — Add RAG, memory, or documentation to enhance agent responses

**Agent Configuration** — Control streaming, temperature, max iterations, and more

---

## Use Cases

**Vanna is ideal for:**
- 📊 Data analytics applications with natural language interfaces
- 🔐 Multi-tenant SaaS needing user-aware permissions
- 🎨 Teams wanting a pre-built web component + backend
- 🏢 Enterprise environments with security/audit requirements
- 📈 Applications needing rich streaming responses (tables, charts, SQL)
- 🔄 Integrating with existing authentication systems

---

## Community & Support

- 📖 **[Full Documentation](https://vanna.ai/docs)** — Complete guides and API reference
- 💡 **[GitHub Discussions](https://github.com/vanna-ai/vanna/discussions)** — Feature requests and Q&A
- 🐛 **[GitHub Issues](https://github.com/vanna-ai/vanna/issues)** — Bug reports
- 📧 **Enterprise Support** — support@vanna.ai

---

## Migration Notes

**Upgrading from Vanna 0.x?**

Vanna 2.0 is a complete rewrite focused on user-aware agents and production deployments. Key changes:

- **New API**: Agent-based instead of `VannaBase` class methods
- **User-aware**: Every component now knows the user identity
- **Streaming**: Rich UI components instead of text/dataframes
- **Web-first**: Built-in `<vanna-chat>` component and server

**Migration path:**

1. **Quick wrap** — Use `LegacyVannaAdapter` to wrap your existing Vanna 0.x instance and get the new web UI immediately
2. **Gradual migration** — Incrementally move to the new Agent API and tools

See the complete [Migration Guide](MIGRATION_GUIDE.md) for step-by-step instructions.

---

## License

MIT License — See [LICENSE](LICENSE) for details.

---

**Built with ❤️ by the Vanna team** | [Website](https://vanna.ai) | [Docs](https://vanna.ai/docs) | [Discussions](https://github.com/vanna-ai/vanna/discussions)


================================================
FILE: README_LEGACY.md
================================================


| GitHub | PyPI | Documentation | Gurubase |
| ------ | ---- | ------------- | -------- |
| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) | [![Gurubase](https://img.shields.io/badge/Gurubase-Ask%20Vanna%20Guru-006BFF)](https://gurubase.io/g/vanna) |

# Vanna
Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation) framework for SQL generation and related functionality.

https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce

![vanna-quadrants](https://github.com/vanna-ai/vanna/assets/7146154/1c7c88ba-c144-4ecf-a028-cf5ba7344ca2)

## How Vanna works

![Screen Recording 2024-01-24 at 11 21 37 AM](https://github.com/vanna-ai/vanna/assets/7146154/1d2718ad-12a8-4a76-afa2-c61754462f93)


Vanna works in two easy steps - train a RAG "model" on your data, and then ask questions which will return SQL queries that can be set up to automatically run on your database.

1. **Train a RAG "model" on your data**.
2. **Ask questions**.

![](img/vanna-readme-diagram.png)

If you don't know what RAG is, don't worry -- you don't need to know how this works under the hood to use it. You just need to know that you "train" a model, which stores some metadata and then use it to "ask" questions.

See the [base class](https://github.com/vanna-ai/vanna/blob/main/src/vanna/base/base.py) for more details on how this works under the hood.

## User Interfaces
These are some of the user interfaces that we've built using Vanna. You can use these as-is or as a starting point for your own custom interface.

- [Jupyter Notebook](https://vanna.ai/docs/postgres-openai-vanna-vannadb/)
- [vanna-ai/vanna-streamlit](https://github.com/vanna-ai/vanna-streamlit)
- [vanna-ai/vanna-flask](https://github.com/vanna-ai/vanna-flask)
- [vanna-ai/vanna-slack](https://github.com/vanna-ai/vanna-slack)

## Supported LLMs

- [OpenAI](https://github.com/vanna-ai/vanna/tree/main/src/vanna/openai)
- [Anthropic](https://github.com/vanna-ai/vanna/tree/main/src/vanna/anthropic)
- [Gemini](https://github.com/vanna-ai/vanna/blob/main/src/vanna/google/gemini_chat.py)
- [HuggingFace](https://github.com/vanna-ai/vanna/blob/main/src/vanna/hf/hf.py)
- [AWS Bedrock](https://github.com/vanna-ai/vanna/tree/main/src/vanna/bedrock)
- [Ollama](https://github.com/vanna-ai/vanna/tree/main/src/vanna/ollama)
- [Qianwen](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qianwen)
- [Qianfan](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qianfan)
- [Zhipu](https://github.com/vanna-ai/vanna/tree/main/src/vanna/ZhipuAI)

## Supported VectorStores

- [AzureSearch](https://github.com/vanna-ai/vanna/tree/main/src/vanna/azuresearch)
- [Opensearch](https://github.com/vanna-ai/vanna/tree/main/src/vanna/opensearch)
- [PgVector](https://github.com/vanna-ai/vanna/tree/main/src/vanna/pgvector)
- [PineCone](https://github.com/vanna-ai/vanna/tree/main/src/vanna/pinecone)
- [ChromaDB](https://github.com/vanna-ai/vanna/tree/main/src/vanna/chromadb)
- [FAISS](https://github.com/vanna-ai/vanna/tree/main/src/vanna/faiss)
- [Marqo](https://github.com/vanna-ai/vanna/tree/main/src/vanna/marqo)
- [Milvus](https://github.com/vanna-ai/vanna/tree/main/src/vanna/milvus)
- [Qdrant](https://github.com/vanna-ai/vanna/tree/main/src/vanna/qdrant)
- [Weaviate](https://github.com/vanna-ai/vanna/tree/main/src/vanna/weaviate)
- [Oracle](https://github.com/vanna-ai/vanna/tree/main/src/vanna/oracle)

## Supported Databases

- [PostgreSQL](https://www.postgresql.org/)
- [MySQL](https://www.mysql.com/)
- [PrestoDB](https://prestodb.io/)
- [Apache Hive](https://hive.apache.org/)
- [ClickHouse](https://clickhouse.com/)
- [Snowflake](https://www.snowflake.com/en/)
- [Oracle](https://www.oracle.com/)
- [Microsoft SQL Server](https://www.microsoft.com/en-us/sql-server/sql-server-downloads)
- [BigQuery](https://cloud.google.com/bigquery)
- [SQLite](https://www.sqlite.org/)
- [DuckDB](https://duckdb.org/)


## Getting started
See the [documentation](https://vanna.ai/docs/) for specifics on your desired database, LLM, etc.

If you want to get a feel for how it works after training, you can try this [Colab notebook](https://vanna.ai/docs/app/).


### Install
```bash
pip install vanna
```

There are a number of optional packages that can be installed so see the [documentation](https://vanna.ai/docs/) for more details.

### Import
See the [documentation](https://vanna.ai/docs/) if you're customizing the LLM or vector database.

```python
# The import statement will vary depending on your LLM and vector database. This is an example for OpenAI + ChromaDB

from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore

class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

vn = MyVanna(config={'api_key': 'sk-...', 'model': 'gpt-4-...'})

# See the documentation for other options

```


## Training
You may or may not need to run these `vn.train` commands depending on your use case. See the [documentation](https://vanna.ai/docs/) for more details.

These statements are shown to give you a feel for how it works.

### Train with DDL Statements
DDL statements contain information about the table names, columns, data types, and relationships in your database.

```python
vn.train(ddl="""
    CREATE TABLE IF NOT EXISTS my-table (
        id INT PRIMARY KEY,
        name VARCHAR(100),
        age INT
    )
""")
```

### Train with Documentation
Sometimes you may want to add documentation about your business terminology or definitions.

```python
vn.train(documentation="Our business defines XYZ as ...")
```

### Train with SQL
You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.

```python
vn.train(sql="SELECT name, age FROM my-table WHERE name = 'John Doe'")
```


## Asking questions
```python
vn.ask("What are the top 10 customers by sales?")
```

You'll get SQL
```sql
SELECT c.c_name as customer_name,
        sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
FROM   snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
        ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
        ON o.o_custkey = c.c_custkey
GROUP BY customer_name
ORDER BY total_sales desc limit 10;
```

If you've connected to a database, you'll get the table:
<div>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>CUSTOMER_NAME</th>
      <th>TOTAL_SALES</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>0</th>
      <td>Customer#000143500</td>
      <td>6757566.0218</td>
    </tr>
    <tr>
      <th>1</th>
      <td>Customer#000095257</td>
      <td>6294115.3340</td>
    </tr>
    <tr>
      <th>2</th>
      <td>Customer#000087115</td>
      <td>6184649.5176</td>
    </tr>
    <tr>
      <th>3</th>
      <td>Customer#000131113</td>
      <td>6080943.8305</td>
    </tr>
    <tr>
      <th>4</th>
      <td>Customer#000134380</td>
      <td>6075141.9635</td>
    </tr>
    <tr>
      <th>5</th>
      <td>Customer#000103834</td>
      <td>6059770.3232</td>
    </tr>
    <tr>
      <th>6</th>
      <td>Customer#000069682</td>
      <td>6057779.0348</td>
    </tr>
    <tr>
      <th>7</th>
      <td>Customer#000102022</td>
      <td>6039653.6335</td>
    </tr>
    <tr>
      <th>8</th>
      <td>Customer#000098587</td>
      <td>6027021.5855</td>
    </tr>
    <tr>
      <th>9</th>
      <td>Customer#000064660</td>
      <td>5905659.6159</td>
    </tr>
  </tbody>
</table>
</div>

You'll also get an automated Plotly chart:
![](img/top-10-customers.png)

## RAG vs. Fine-Tuning
RAG
- Portable across LLMs
- Easy to remove training data if any of it becomes obsolete
- Much cheaper to run than fine-tuning
- More future-proof -- if a better LLM comes out, you can just swap it out

Fine-Tuning
- Good if you need to minimize tokens in the prompt
- Slow to get started
- Expensive to train and run (generally)

## Why Vanna?

1. **High accuracy on complex datasets.**
    - Vanna’s capabilities are tied to the training data you give it
    - More training data means better accuracy for large and complex datasets
2. **Secure and private.**
    - Your database contents are never sent to the LLM or the vector database
    - SQL execution happens in your local environment
3. **Self learning.**
    - If using via Jupyter, you can choose to "auto-train" it on the queries that were successfully executed
    - If using via other interfaces, you can have the interface prompt the user to provide feedback on the results
    - Correct question to SQL pairs are stored for future reference and make the future results more accurate
4. **Supports any SQL database.**
    - The package allows you to connect to any SQL database that you can otherwise connect to with Python
5. **Choose your front end.**
    - Most people start in a Jupyter Notebook.
    - Expose to your end users via Slackbot, web app, Streamlit app, or a custom front end.

## Extending Vanna
Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](https://github.com/vanna-ai/vanna/blob/main/src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.

## Vanna in 100 Seconds

https://github.com/vanna-ai/vanna/assets/7146154/eb90ee1e-aa05-4740-891a-4fc10e611cab

## More resources
 - [Full Documentation](https://vanna.ai/docs/)
 - [Website](https://vanna.ai)
 - [Discord group for support](https://discord.gg/qUZYKHremx)


================================================
FILE: examples/chromadb_gpu_example.py
================================================
"""
Example: Using ChromaDB AgentMemory with GPU acceleration

This example demonstrates how to use ChromaAgentMemory with intelligent
device selection for GPU acceleration when available.
"""

from vanna.integrations.chromadb import (
    ChromaAgentMemory,
    get_device,
    create_sentence_transformer_embedding_function
)


def example_default_usage():
    """Example 1: Use default embedding function (no GPU, no sentence-transformers required)"""
    print("Example 1: Default ChromaDB embedding (CPU-only, no extra dependencies)")

    memory = ChromaAgentMemory(
        persist_directory="./chroma_memory_default"
    )

    print("✓ ChromaAgentMemory created with default embedding function")
    print()


def example_auto_gpu():
    """Example 2: Automatic GPU detection with SentenceTransformers"""
    print("Example 2: Automatic GPU detection")

    # Detect the best available device
    device = get_device()
    print(f"Detected device: {device}")

    # Create embedding function with automatic device selection
    embedding_fn = create_sentence_transformer_embedding_function()

    memory = ChromaAgentMemory(
        persist_directory="./chroma_memory_gpu",
        embedding_function=embedding_fn
    )

    print(f"✓ ChromaAgentMemory created with SentenceTransformer on {device}")
    print()


def example_explicit_cuda():
    """Example 3: Explicitly use CUDA"""
    print("Example 3: Explicitly request CUDA")

    # Explicitly request CUDA
    embedding_fn = create_sentence_transformer_embedding_function(device="cuda")

    memory = ChromaAgentMemory(
        persist_directory="./chroma_memory_cuda",
        embedding_function=embedding_fn
    )

    print("✓ ChromaAgentMemory created with SentenceTransformer on CUDA")
    print()


def example_custom_model_gpu():
    """Example 4: Use a larger model with GPU"""
    print("Example 4: Custom model with GPU acceleration")

    # Use a larger, more accurate model with GPU
    embedding_fn = create_sentence_transformer_embedding_function(
        model_name="sentence-transformers/all-mpnet-base-v2"
    )

    memory = ChromaAgentMemory(
        persist_directory="./chroma_memory_large",
        embedding_function=embedding_fn
    )

    print("✓ ChromaAgentMemory created with all-mpnet-base-v2 model")
    print()


def example_manual_chromadb():
    """Example 5: Manually configure ChromaDB embedding function"""
    print("Example 5: Manual ChromaDB embedding function configuration")

    from chromadb.utils import embedding_functions

    # Manually create and configure the embedding function
    device = get_device()
    embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        device=device
    )

    memory = ChromaAgentMemory(
        persist_directory="./chroma_memory_manual",
        embedding_function=embedding_fn
    )

    print(f"✓ ChromaAgentMemory created with manual configuration on {device}")
    print()


if __name__ == "__main__":
    print("=" * 70)
    print("ChromaDB AgentMemory GPU Acceleration Examples")
    print("=" * 70)
    print()

    # Example 1: Default (no GPU, no sentence-transformers needed)
    example_default_usage()

    # Examples 2-5 require sentence-transformers to be installed
    try:
        import sentence_transformers

        example_auto_gpu()

        # Only run CUDA example if CUDA is available
        device = get_device()
        if device == "cuda":
            example_explicit_cuda()

        example_custom_model_gpu()
        example_manual_chromadb()

    except ImportError:
        print("⚠️  sentence-transformers not installed")
        print("    Install with: pip install sentence-transformers")
        print("    Examples 2-5 require this package for GPU acceleration")
        print()

    print("=" * 70)
    print("Summary:")
    print("- Example 1 works without sentence-transformers (CPU only)")
    print("- Examples 2-5 require sentence-transformers for GPU support")
    print("- GPU acceleration automatically detected when available")
    print("=" * 70)


================================================
FILE: examples/transform_args_example.py
================================================
"""
Example demonstrating how to use ToolRegistry.transform_args for user-specific
argument transformation, such as applying row-level security (RLS) to SQL queries.

This example shows:
1. Creating a custom ToolRegistry subclass that overrides transform_args
2. Applying RLS transformation to SQL queries based on user context
3. Rejecting tool execution when validation fails
"""

from typing import Union
from pydantic import BaseModel

from vanna.core import ToolRegistry
from vanna.core.tool import Tool, ToolContext, ToolRejection, ToolResult
from vanna.core.user import User


# Example: SQL execution tool arguments
class SQLExecutionArgs(BaseModel):
    query: str
    database: str = "default"


class SQLExecutionTool(Tool[SQLExecutionArgs]):
    @property
    def name(self) -> str:
        return "execute_sql"

    @property
    def description(self) -> str:
        return "Execute a SQL query against the database"

    def get_args_schema(self):
        return SQLExecutionArgs

    async def execute(self, context: ToolContext, args: SQLExecutionArgs) -> ToolResult:
        # Execute the SQL query (implementation not shown)
        return ToolResult(
            success=True,
            result_for_llm=f"Executed query: {args.query[:50]}...",
        )


class RLSToolRegistry(ToolRegistry):
    """Custom ToolRegistry that applies row-level security to SQL queries."""

    async def transform_args(
        self,
        tool: Tool,
        args,
        user: User,
        context: ToolContext,
    ) -> Union[SQLExecutionArgs, ToolRejection]:
        """Apply row-level security transformation to SQL queries."""

        # Only transform SQL execution tools
        if tool.name == "execute_sql" and isinstance(args, SQLExecutionArgs):
            original_query = args.query.strip()

            # Example 1: Reject queries that try to access restricted tables
            if "restricted_table" in original_query.lower():
                return ToolRejection(
                    reason="Access to 'restricted_table' is not permitted for your user group"
                )

            # Example 2: Apply RLS by modifying the WHERE clause
            # This is a simplified example - real RLS would be more sophisticated
            if "SELECT" in original_query.upper() and "users" in original_query.lower():
                # Add a WHERE clause to filter by user's organization
                user_org_id = user.metadata.get("organization_id")

                if user_org_id:
                    # Simple RLS: append WHERE clause for organization filtering
                    if "WHERE" in original_query.upper():
                        transformed_query = original_query.replace(
                            "WHERE",
                            f"WHERE organization_id = {user_org_id} AND",
                            1
                        )
                    else:
                        # Add WHERE clause before ORDER BY, LIMIT, etc.
                        transformed_query = original_query.rstrip(";")
                        transformed_query += f" WHERE organization_id = {user_org_id}"

                    # Return transformed arguments
                    return args.model_copy(update={"query": transformed_query})

            # Example 3: Validate required parameters
            if not args.database:
                return ToolRejection(
                    reason="Database parameter is required for SQL execution"
                )

        # For all other tools or if no transformation needed, pass through
        return args


# Usage example
async def example_usage():
    """Demonstrate using the RLS-enabled ToolRegistry."""
    from vanna.capabilities.agent_memory import AgentMemory

    # Create registry and register tool
    registry = RLSToolRegistry()
    sql_tool = SQLExecutionTool()
    registry.register_local_tool(sql_tool, access_groups=[])

    # Create a user with organization context
    user = User(
        user_id="user123",
        metadata={"organization_id": 42}
    )

    # Create tool context
    context = ToolContext(
        user=user,
        conversation_id="conv123",
        request_id="req123",
        agent_memory=AgentMemory(),
    )

    # Example 1: Query that will be transformed with RLS
    from vanna.core.tool import ToolCall

    tool_call = ToolCall(
        id="call1",
        name="execute_sql",
        arguments={
            "query": "SELECT * FROM users",
            "database": "production"
        }
    )

    result = await registry.execute(tool_call, context)
    print(f"Result: {result.result_for_llm}")
    # The query will be transformed to: SELECT * FROM users WHERE organization_id = 42

    # Example 2: Query that will be rejected
    tool_call_rejected = ToolCall(
        id="call2",
        name="execute_sql",
        arguments={
            "query": "SELECT * FROM restricted_table",
            "database": "production"
        }
    )

    result = await registry.execute(tool_call_rejected, context)
    print(f"Rejected: {result.error}")
    # Will return: "Access to 'restricted_table' is not permitted for your user group"


if __name__ == "__main__":
    import asyncio
    asyncio.run(example_usage())


================================================
FILE: frontends/webcomponent/.storybook/main.ts
================================================
import type { StorybookConfig } from '@storybook/web-components-vite';

const config: StorybookConfig = {
  stories: ['../src/**/*.stories.@(js|jsx|mjs|ts|tsx)'],
  addons: [
    '@storybook/addon-essentials',
    '@storybook/addon-actions',
    '@storybook/addon-controls',
    '@storybook/addon-docs',
  ],
  framework: {
    name: '@storybook/web-components-vite',
    options: {},
  },
  typescript: {
    check: false,
    reactDocgen: 'react-docgen-typescript',
    reactDocgenTypescriptOptions: {
      shouldExtractLiteralValuesFromEnum: true,
      propFilter: (prop) => (prop.parent ? !/node_modules/.test(prop.parent.fileName) : true),
    },
  },
};

export default config;

================================================
FILE: frontends/webcomponent/.storybook/preview-head.html
================================================
<link rel="preconnect" href="https://fonts.googleapis.com" />
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
<link
  href="https://fonts.googleapis.com/css2?family=Roboto+Slab:wght@400;500;600;700&family=Signika:wght@300;400;500;600;700&family=Space+Mono:wght@400;700&display=swap"
  rel="stylesheet"
/>


================================================
FILE: frontends/webcomponent/.storybook/preview.ts
================================================
import type { Preview } from '@storybook/web-components';

const preview: Preview = {
  parameters: {
    controls: {
      matchers: {
        color: /(background|color)$/i,
        date: /Date$/i,
      },
    },
    docs: {
      autodocs: 'tag',
    },
  },
};

export default preview;

================================================
FILE: frontends/webcomponent/TEST_README.md
================================================
# Vanna Webcomponent Comprehensive Test Suite

This test suite validates all component types and update patterns in the vanna-webcomponent before pruning unused code.

## Overview

The test suite consists of:
- **`test_backend.py`**: Real Python backend that streams all component types
- **`test-comprehensive.html`**: Browser-based test interface with visual validation
- **Two test modes**: Rapid (stress test) and Realistic (with delays)

## Quick Start

### 1. Install Dependencies

```bash
cd submodule/vanna-webcomponent
pip install -r requirements-test.txt
```

### 2. Build the Webcomponent

```bash
npm run build
```

### 3. Start the Test Backend

```bash
# Realistic mode (with delays between components)
python test_backend.py --mode realistic

# Rapid mode (fast stress test)
python test_backend.py --mode rapid
```

The backend will start on `http://localhost:5555` and automatically serve the test page.

### 4. Open Test Interface

Simply open your browser to:
```
http://localhost:5555
```

The test page will load automatically!

### 5. Run the Test

1. Click **"Run Comprehensive Test"** button in the sidebar
2. Watch components render in real-time
3. Monitor the checklist - items check off as components render
4. Watch the console log for any errors

## Test Coverage

### Component Types Tested

The test exercises **all** rich component types with **19 different components**:

#### Primitive Components
- ✓ Text (with markdown)
- ✓ Badge
- ✓ Icon Text

#### Feedback Components
- ✓ Status Card (with all states: pending, running, completed, failed)
- ✓ Progress Display (0% → 50% → 100%)
- ✓ Progress Bar
- ✓ Status Indicator (with pulse animation)
- ✓ Notification (info, success, warning, error levels)
- ✓ Log Viewer (with info, warning, error logs)

#### Data Components
- ✓ Card (with buttons and actions)
- ✓ Task List (with status updates)
- ✓ **DataFrame** (tabular data with search/sort/filter/export)
- ✓ **Table** (structured data with explicit column definitions)
- ✓ **Chart** (Plotly charts: bar, line, scatter)
- ✓ **Code Block** (syntax highlighted code: Python, SQL, etc.)

#### Specialized Components
- ✓ **Artifact** (HTML/SVG interactive content)

#### Container Components
- ✓ **Container** (groups components in rows/columns)

#### Interactive Components
- ✓ Button (single)
- ✓ Button Group (horizontal/vertical)
- ✓ Button actions (click → backend response)

#### UI State Updates
- ✓ Status Bar Update (updates status bar above input)
- ✓ Task Tracker Update (adds/updates tasks in sidebar)
- ✓ Chat Input Update (changes placeholder/state)

### Update Operations Tested

For each component type, the test validates:

1. **Create** (`lifecycle: create`) - Initial component rendering
2. **Update** (`lifecycle: update`) - Incremental property updates
3. **Replace** - Full component replacement
4. **Remove** - Component removal from DOM

### Interactive Features Tested

- **Button Actions**: Clicking buttons sends actions to backend
- **Action Handling**: Backend receives actions and responds with new components
- **Round-trip Communication**: Full interaction loop validation

## Test Modes

### Realistic Mode (Default)

```bash
python test_backend.py --mode realistic
```

- Includes delays between component updates (0.2-0.5s)
- Simulates real conversation flow
- Easier to observe rendering behavior
- **Recommended for initial validation**

### Rapid Mode

```bash
python test_backend.py --mode rapid
```

- Minimal delays (0.05-0.1s)
- Stress tests rendering performance
- Validates no race conditions
- **Use for performance testing**

## Validation Checklist

The test interface provides real-time validation:

### ✅ Visual Checklist
- Automatically checks off components as they render
- Shows 19 component types
- Green checkmark = successfully rendered

### 📊 Metrics
- **Components Rendered**: Total unique component types
- **Updates Processed**: Total number of updates (create + update + replace)
- **Errors**: Console errors detected

### 🔴 Console Monitor
- Real-time console log display
- Errors highlighted in red
- Warnings in yellow
- Info messages in blue

### 🟢 Status Indicators
- **Backend Status**: Green = connected, Red = disconnected
- **Console Status**: Green = no errors, Red = errors detected

## Using for Webcomponent Pruning

The test suite is designed to validate that pruning doesn't break functionality:

### Pruning Workflow

1. **Run baseline test**:
   ```bash
   python test_backend.py --mode realistic

   # Browser: Open http://localhost:5555 and run test
   # Verify: All 19 components render, 0 errors
   ```

2. **Identify cruft to remove**:
   - Unused imports
   - Dead code paths
   - Deprecated components
   - Development-only utilities

3. **Remove one piece of cruft**:
   ```bash
   # Example: Remove unused import from vanna-chat.ts
   # or delete unused utility file
   ```

4. **Rebuild**:
   ```bash
   npm run build
   ```

5. **Refresh browser test**:
   - Press F5 to reload test page
   - Click "Run Comprehensive Test" again
   - Check console for errors
   - Verify all 12 components still render

6. **If green → continue; if red → investigate**:
   - Green (no errors): Commit the change, continue pruning
   - Red (errors): Revert change, that code was actually needed

7. **Repeat until clean**: Continue removing cruft until webcomponent is minimal

### What to Prune

Look for these common types of cruft:

- ❌ **Unused imports**: Components imported but never used
- ❌ **Development utilities**: Debug helpers, test mocks in production code
- ❌ **Deprecated components**: Old component versions no longer referenced
- ❌ **Unused CSS**: Styles for removed components
- ❌ **Dead code paths**: Conditional logic that's never executed
- ❌ **Commented code**: Old implementations that are commented out
- ❌ **Storybook-only code**: Utilities only used in stories, not production

### What NOT to Prune

Be careful with these:

- ✅ **Base component renderers**: Even if rarely used, may be needed
- ✅ **ComponentRegistry entries**: Needed for dynamic component lookup
- ✅ **Shadow DOM utilities**: Required for web components
- ✅ **Event handlers**: May be used by runtime events
- ✅ **Type definitions**: Used at compile time even if not runtime

## Customizing the Test

### Add More Component Tests

Edit `test_backend.py` and add new test functions:

```python
async def test_my_component(conversation_id: str, request_id: str, mode: str):
    """Test my custom component."""
    my_component = MyComponent(
        id=str(uuid.uuid4()),
        # ... component properties
    )
    yield await yield_chunk(my_component, conversation_id, request_id)
    await delay(mode)

# Then add to run_comprehensive_test():
async for chunk in test_my_component(conversation_id, request_id, mode):
    yield chunk
```

### Modify Test Delays

In `test_backend.py`, adjust the `delay()` function:

```python
async def delay(mode: str, short: float = 0.1, long: float = 0.5):
    if mode == "realistic":
        await asyncio.sleep(long)  # Adjust long delay here
    elif mode == "rapid":
        await asyncio.sleep(short)  # Adjust short delay here
```

### Add Custom Validation

Edit `test-comprehensive.html` and add custom validation logic:

```javascript
// Add to MutationObserver callback
const componentType = node.getAttribute('data-component-type');
if (componentType === 'my_component') {
    // Custom validation for my_component
    console.log('My component rendered!');
}
```

## Troubleshooting

### Backend won't start

**Error**: `ModuleNotFoundError: No module named 'vanna'`

**Solution**: Make sure vanna is in the Python path:
```bash
cd submodule/vanna-webcomponent
python test_backend.py  # Already adds ../vanna/src to sys.path
```

### Frontend shows "Backend not responding"

**Solutions**:
1. Check backend is running: `curl http://localhost:5555/health`
2. Check CORS is enabled (should be by default)
3. Verify port 5555 is not in use: `lsof -i :5555`

### Components not rendering

**Check**:
1. Browser console for errors (F12)
2. Webcomponent is built: `ls dist/`
3. Test HTML is loading: `<script type="module" src="./dist/index.js"></script>`

### Test page is blank

**Solutions**:
1. Check you're serving from the right directory:
   ```bash
   cd submodule/vanna-webcomponent
   python -m http.server 8080
   ```
2. Open correct URL: `http://localhost:8080/test-comprehensive.html`
3. Check browser console for 404 errors

### Checklist not updating

The checklist tracks components by their `data-component-type` attribute. If components don't have this attribute, they won't be tracked.

**Verify**: Open browser DevTools and inspect rendered components for `data-component-type`.

## Advanced Usage

### Run Backend on Different Port

```bash
python test_backend.py --port 8000
```

Then update `test-comprehensive.html`:
```html
<vanna-chat
    api-url="http://localhost:8000"
    ...
></vanna-chat>
```

### Enable Debug Logging

Add to `test_backend.py`:

```python
import logging
logging.basicConfig(level=logging.DEBUG)
```

### Run Type Checking

Validate the backend code with mypy:

```bash
python -m mypy test_backend.py
```

This catches type errors before runtime (e.g., wrong field names in Pydantic models).

### Test Specific Component Only

Modify `run_comprehensive_test()` to only run specific tests:

```python
async def run_comprehensive_test(conversation_id, request_id, mode):
    # Comment out tests you don't want to run
    async for chunk in test_status_card(conversation_id, request_id, mode):
        yield chunk

    # async for chunk in test_progress_display(...):  # Disabled
    #     yield chunk
```

## Architecture

### Backend Flow

1. FastAPI receives POST to `/api/vanna/v2/chat_sse`
2. `chat_sse()` creates async generator
3. Generator yields components wrapped in `ChatStreamChunk`
4. Each chunk serialized to SSE format: `data: {json}\n\n`
5. Stream ends with `data: [DONE]\n\n`

### Frontend Flow

1. `<vanna-chat>` web component connects to backend
2. Opens SSE connection to `/api/vanna/v2/chat_sse`
3. Receives chunks, parses JSON
4. `ComponentManager` processes updates
5. `ComponentRegistry` renders HTML elements
6. Elements appended to shadow DOM container
7. MutationObserver detects new components
8. Checklist updates automatically

### Button Action Flow

1. User clicks button in frontend
2. Button's `action` property sent as new message
3. Backend receives message via `/api/vanna/v2/chat_sse` POST
4. `handle_action_message()` processes action
5. Response components streamed back
6. Frontend renders response

## Files

- **`test_backend.py`** - Python FastAPI backend (400 lines)
- **`test-comprehensive.html`** - Browser test interface (500 lines)
- **`requirements-test.txt`** - Python dependencies
- **`TEST_README.md`** - This documentation

## Next Steps

After validating the webcomponent with this test suite:

1. **Run baseline test** - Verify all components work before pruning
2. **Identify cruft** - Find unused code in the webcomponent
3. **Prune iteratively** - Remove one piece at a time, test after each change
4. **Commit clean code** - Once pruned, commit the cleaned webcomponent
5. **Copy to vanna package** - Integrate cleaned webcomponent into vanna Python package

## Support

If you encounter issues with the test suite:

1. Check this README's Troubleshooting section
2. Verify all dependencies are installed
3. Ensure you're in the correct directory
4. Check browser and terminal console output

---

**Happy Testing!** 🧪


================================================
FILE: frontends/webcomponent/package.json
================================================
{
  "name": "@vanna/webcomponent",
  "version": "2.0.0",
  "description": "Lit-based web components for Vanna User Agents",
  "main": "dist/vanna-components.js",
  "scripts": {
    "sync-version": "node scripts/sync-version.js",
    "dev": "vite",
    "build": "npm run sync-version && tsc && vite build",
    "preview": "vite preview",
    "storybook": "storybook dev -p 6006",
    "build-storybook": "storybook build",
    "test": "echo \"Error: no test specified\" && exit 1"
  },
  "keywords": [
    "vanna",
    "ai",
    "sql",
    "web-components",
    "lit",
    "chat",
    "llm",
    "natural-language"
  ],
  "author": "Zain Hoda <zain@vanna.ai>",
  "license": "MIT",
  "type": "commonjs",
  "repository": {
    "type": "git",
    "url": "https://github.com/vanna-ai/vanna.git",
    "directory": "frontends/webcomponent"
  },
  "homepage": "https://github.com/vanna-ai/vanna",
  "bugs": {
    "url": "https://github.com/vanna-ai/vanna/issues"
  },
  "files": [
    "dist",
    "src"
  ],
  "dependencies": {
    "lit": "^3.3.1",
    "plotly.js-dist-min": "^3.1.0"
  },
  "devDependencies": {
    "@storybook/addon-actions": "^8.6.14",
    "@storybook/addon-controls": "^8.6.14",
    "@storybook/addon-docs": "^8.6.14",
    "@storybook/addon-essentials": "^8.6.14",
    "@storybook/web-components": "^8.6.14",
    "@storybook/web-components-vite": "^8.6.14",
    "@types/plotly.js-dist-min": "^2.3.4",
    "storybook": "^8.6.14",
    "typescript": "^5.9.2",
    "vite": "^7.1.5"
  }
}


================================================
FILE: frontends/webcomponent/requirements-test.txt
================================================
# Test backend requirements for vanna-webcomponent comprehensive testing

fastapi>=0.115.0
uvicorn[standard]>=0.32.0
pydantic>=2.0.0

# Note: The vanna package itself will be imported from ../vanna/src
# No need to install it separately for local testing


================================================
FILE: frontends/webcomponent/scripts/sync-version.js
================================================
/**
 * Sync version from pyproject.toml to package.json
 *
 * This ensures the webcomponent version always matches the Python package version.
 * Single source of truth: pyproject.toml
 *
 * Usage: node scripts/sync-version.js
 */

const fs = require('fs');
const path = require('path');

// Paths relative to this script
const PYPROJECT_PATH = path.join(__dirname, '../../../pyproject.toml');
const PACKAGE_JSON_PATH = path.join(__dirname, '../package.json');

function extractVersionFromPyproject(content) {
  // Match: version = "2.0.0"
  const match = content.match(/^version\s*=\s*"([^"]+)"/m);
  if (!match) {
    throw new Error('Could not find version in pyproject.toml');
  }
  return match[1];
}

function updatePackageJsonVersion(packageJsonPath, newVersion) {
  const packageJson = JSON.parse(fs.readFileSync(packageJsonPath, 'utf8'));
  const oldVersion = packageJson.version;

  packageJson.version = newVersion;

  fs.writeFileSync(
    packageJsonPath,
    JSON.stringify(packageJson, null, 2) + '\n',
    'utf8'
  );

  return { oldVersion, newVersion };
}

function main() {
  try {
    // Read pyproject.toml
    const pyprojectContent = fs.readFileSync(PYPROJECT_PATH, 'utf8');
    const version = extractVersionFromPyproject(pyprojectContent);

    // Update package.json
    const { oldVersion, newVersion } = updatePackageJsonVersion(PACKAGE_JSON_PATH, version);

    if (oldVersion !== newVersion) {
      console.log(`✓ Version synced: ${oldVersion} → ${newVersion}`);
    } else {
      console.log(`✓ Version already in sync: ${newVersion}`);
    }

    process.exit(0);
  } catch (error) {
    console.error(`✗ Version sync failed: ${error.message}`);
    process.exit(1);
  }
}

main();


================================================
FILE: frontends/webcomponent/src/components/button.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { ComponentManager, ComponentUpdate } from './rich-component-system';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

const meta: Meta = {
  title: 'Rich Components/Buttons',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'dark',
      values: [
        { name: 'light', value: '#f5f7fa' },
        { name: 'dark', value: 'rgb(11, 15, 25)' },
      ],
    },
  },
};

export default meta;
type Story = StoryObj;

const ensureTokenStyles = () => {
  if (document.getElementById('vanna-token-style')) {
    return;
  }

  const style = document.createElement('style');
  style.id = 'vanna-token-style';
  style.textContent = vannaDesignTokens.cssText.replace(/:host/g, '.vanna-tokens');
  document.head.appendChild(style);
};

const createContainer = () => {
  ensureTokenStyles();

  const container = document.createElement('div');
  container.className = 'vanna-tokens';
  container.style.cssText = `
    padding: var(--vanna-space-5, 20px);
    max-width: 800px;
    margin: 0 auto;
    background: var(--vanna-background-default);
    border-radius: var(--vanna-border-radius-lg);
    box-shadow: var(--vanna-shadow-md);
  `;

  return container;
};

const createManager = (container: HTMLElement) => new ComponentManager(container);

const renderComponent = (manager: ComponentManager, component: any) => {
  const update: ComponentUpdate = {
    operation: 'create',
    target_id: component.id,
    component,
    timestamp: new Date().toISOString(),
  } as ComponentUpdate;

  manager.processUpdate(update);
};

const withDefaults = (component: any) => ({
  layout: { position: 'append', size: {}, z_index: 0, classes: [] },
  theme: {},
  lifecycle: 'create',
  ...component,
});

const addMockVannaChat = (container: HTMLElement) => {
  // Create a mock vanna-chat element with sendMessage method
  const mockVannaChat = document.createElement('div');
  mockVannaChat.setAttribute('id', 'mock-vanna-chat');

  // Store the original querySelector
  const originalQuerySelector = document.querySelector.bind(document);

  // Override querySelector to return our mock when looking for vanna-chat
  document.querySelector = function(selector: string) {
    if (selector === 'vanna-chat') {
      return mockVannaChat as any;
    }
    return originalQuerySelector(selector);
  } as any;

  // Add sendMessage method that logs to console and shows in UI
  (mockVannaChat as any).sendMessage = (message: string) => {
    console.log('📤 Button clicked - Message:', message);

    // Show a visual feedback in the storybook
    const feedback = document.createElement('div');
    feedback.style.cssText = `
      position: fixed;
      top: 20px;
      right: 20px;
      background: #4CAF50;
      color: white;
      padding: 12px 20px;
      border-radius: 8px;
      box-shadow: 0 4px 6px rgba(0,0,0,0.3);
      font-family: monospace;
      z-index: 10000;
      animation: slideIn 0.3s ease-out;
    `;
    feedback.textContent = `Message sent: ${message}`;

    // Add animation
    const style = document.createElement('style');
    style.textContent = `
      @keyframes slideIn {
        from {
          transform: translateX(100%);
          opacity: 0;
        }
        to {
          transform: translateX(0);
          opacity: 1;
        }
      }
    `;
    document.head.appendChild(style);

    document.body.appendChild(feedback);

    setTimeout(() => {
      feedback.style.opacity = '0';
      feedback.style.transition = 'opacity 0.3s ease-out';
      setTimeout(() => feedback.remove(), 300);
    }, 2000);
  };

  container.appendChild(mockVannaChat);
  return mockVannaChat;
};

export const SingleButtons: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);
    addMockVannaChat(container);

    // Add title
    const title = document.createElement('h2');
    title.textContent = 'Single Button Components';
    title.style.cssText = 'margin-bottom: 20px; color: var(--vanna-text-primary);';
    container.appendChild(title);

    const buttons = [
      withDefaults({
        id: 'primary-button',
        type: 'button',
        data: {
          label: 'Primary Action',
          action: 'primary_action',
          variant: 'primary',
          size: 'medium',
        },
      }),
      withDefaults({
        id: 'secondary-button',
        type: 'button',
        data: {
          label: 'Save Draft',
          action: 'save_draft',
          variant: 'secondary',
          size: 'medium',
          icon: '💾',
          icon_position: 'left',
        },
      }),
      withDefaults({
        id: 'success-button',
        type: 'button',
        data: {
          label: 'Approve',
          action: 'approve',
          variant: 'success',
          size: 'medium',
          icon: '✓',
        },
      }),
      withDefaults({
        id: 'warning-button',
        type: 'button',
        data: {
          label: 'Caution',
          action: 'warning',
          variant: 'warning',
          size: 'medium',
          icon: '⚠️',
        },
      }),
      withDefaults({
        id: 'error-button',
        type: 'button',
        data: {
          label: 'Delete',
          action: 'delete',
          variant: 'error',
          size: 'medium',
          icon: '🗑️',
        },
      }),
      withDefaults({
        id: 'ghost-button',
        type: 'button',
        data: {
          label: 'Ghost Style',
          action: 'ghost',
          variant: 'ghost',
          icon: '👻',
        },
      }),
      withDefaults({
        id: 'link-button',
        type: 'button',
        data: {
          label: 'Learn More',
          action: 'learn_more',
          variant: 'link',
        },
      }),
      withDefaults({
        id: 'loading-button',
        type: 'button',
        data: {
          label: 'Processing...',
          action: 'loading',
          variant: 'primary',
          loading: true,
        },
      }),
      withDefaults({
        id: 'disabled-button',
        type: 'button',
        data: {
          label: 'Disabled',
          action: 'disabled',
          variant: 'secondary',
          disabled: true,
        },
      }),
    ];

    buttons.forEach((component) => {
      renderComponent(manager, component);
      // Add some spacing
      const spacer = document.createElement('div');
      spacer.style.height = '12px';
      container.appendChild(spacer);
    });

    // Add instruction
    const instruction = document.createElement('p');
    instruction.textContent = 'Click any button to see the message it sends (wrapped in square brackets)';
    instruction.style.cssText = 'margin-top: 20px; color: var(--vanna-text-secondary); font-style: italic;';
    container.appendChild(instruction);

    return container;
  },
};

export const ButtonSizes: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);
    addMockVannaChat(container);

    const title = document.createElement('h2');
    title.textContent = 'Button Sizes';
    title.style.cssText = 'margin-bottom: 20px; color: var(--vanna-text-primary);';
    container.appendChild(title);

    const sizes = ['small', 'medium', 'large'];

    sizes.forEach((size) => {
      const button = withDefaults({
        id: `button-${size}`,
        type: 'button',
        data: {
          label: `${size.charAt(0).toUpperCase() + size.slice(1)} Button`,
          action: `${size}_action`,
          variant: 'primary',
          size,
          icon: '⭐',
        },
      });

      renderComponent(manager, button);

      const spacer = document.createElement('div');
      spacer.style.height = '12px';
      container.appendChild(spacer);
    });

    return container;
  },
};

export const ButtonGroups: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);
    addMockVannaChat(container);

    const title = document.createElement('h2');
    title.textContent = 'Button Group Components';
    title.style.cssText = 'margin-bottom: 20px; color: var(--vanna-text-primary);';
    container.appendChild(title);

    // Horizontal action group
    const actionGroup = withDefaults({
      id: 'action-group',
      type: 'button_group',
      data: {
        buttons: [
          {
            label: 'Accept',
            action: 'accept',
            variant: 'success',
            icon: '✓',
          },
          {
            label: 'Reject',
            action: 'reject',
            variant: 'error',
            icon: '✗',
          },
          {
            label: 'Cancel',
            action: 'cancel',
            variant: 'secondary',
          },
        ],
        orientation: 'horizontal',
        spacing: 'medium',
        align: 'left',
      },
    });

    const sectionTitle1 = document.createElement('h3');
    sectionTitle1.textContent = 'Horizontal Action Group';
    sectionTitle1.style.cssText = 'margin: 20px 0 10px 0; color: var(--vanna-text-primary); font-size: 16px;';
    container.appendChild(sectionTitle1);
    renderComponent(manager, actionGroup);

    // Centered navigation
    const navigationGroup = withDefaults({
      id: 'navigation-group',
      type: 'button_group',
      data: {
        buttons: [
          {
            label: 'Back',
            action: 'back',
            variant: 'ghost',
            icon: '←',
          },
          {
            label: 'Continue',
            action: 'continue',
            variant: 'primary',
            icon: '→',
            icon_position: 'right',
          },
        ],
        orientation: 'horizontal',
        spacing: 'large',
        align: 'center',
      },
    });

    const sectionTitle2 = document.createElement('h3');
    sectionTitle2.textContent = 'Centered Navigation';
    sectionTitle2.style.cssText = 'margin: 20px 0 10px 0; color: var(--vanna-text-primary); font-size: 16px;';
    container.appendChild(sectionTitle2);
    renderComponent(manager, navigationGroup);

    // Vertical options
    const verticalGroup = withDefaults({
      id: 'vertical-group',
      type: 'button_group',
      data: {
        buttons: [
          { label: 'Option 1', action: 'option1', variant: 'secondary' },
          { label: 'Option 2', action: 'option2', variant: 'secondary' },
          { label: 'Option 3', action: 'option3', variant: 'secondary' },
        ],
        orientation: 'vertical',
        spacing: 'small',
        align: 'left',
      },
    });

    const sectionTitle3 = document.createElement('h3');
    sectionTitle3.textContent = 'Vertical Options';
    sectionTitle3.style.cssText = 'margin: 20px 0 10px 0; color: var(--vanna-text-primary); font-size: 16px;';
    container.appendChild(sectionTitle3);
    renderComponent(manager, verticalGroup);

    // Toolbar
    const toolbarGroup = withDefaults({
      id: 'toolbar-group',
      type: 'button_group',
      data: {
        buttons: [
          {
            label: 'New',
            action: 'new',
            variant: 'primary',
            icon: '➕',
            size: 'small',
          },
          {
            label: 'Edit',
            action: 'edit',
            variant: 'secondary',
            icon: '✏️',
            size: 'small',
          },
          {
            label: 'Delete',
            action: 'delete',
            variant: 'error',
            icon: '🗑️',
            size: 'small',
          },
          {
            label: 'Share',
            action: 'share',
            variant: 'ghost',
            icon: '🔗',
            size: 'small',
          },
        ],
        orientation: 'horizontal',
        spacing: 'small',
        align: 'left',
      },
    });

    const sectionTitle4 = document.createElement('h3');
    sectionTitle4.textContent = 'Toolbar (Small Buttons)';
    sectionTitle4.style.cssText = 'margin: 20px 0 10px 0; color: var(--vanna-text-primary); font-size: 16px;';
    container.appendChild(sectionTitle4);
    renderComponent(manager, toolbarGroup);

    // Full width confirmation
    const confirmationGroup = withDefaults({
      id: 'confirmation-group',
      type: 'button_group',
      data: {
        buttons: [
          { label: 'Yes', action: 'yes', variant: 'success' },
          { label: 'No', action: 'no', variant: 'error' },
        ],
        orientation: 'horizontal',
        spacing: 'medium',
        align: 'space-between',
        full_width: true,
      },
    });

    const sectionTitle5 = document.createElement('h3');
    sectionTitle5.textContent = 'Full Width Confirmation';
    sectionTitle5.style.cssText = 'margin: 20px 0 10px 0; color: var(--vanna-text-primary); font-size: 16px;';
    container.appendChild(sectionTitle5);
    renderComponent(manager, confirmationGroup);

    // Add instruction
    const instruction = document.createElement('p');
    instruction.textContent = 'Click any button in the groups to see the message it sends';
    instruction.style.cssText = 'margin-top: 20px; color: var(--vanna-text-secondary); font-style: italic;';
    container.appendChild(instruction);

    return container;
  },
};

export const InteractiveDemo: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);
    addMockVannaChat(container);

    const title = document.createElement('h2');
    title.textContent = 'Interactive Button Demo';
    title.style.cssText = 'margin-bottom: 20px; color: var(--vanna-text-primary);';
    container.appendChild(title);

    const description = document.createElement('p');
    description.textContent = 'This demo shows how buttons send messages with their labels wrapped in square brackets.';
    description.style.cssText = 'margin-bottom: 20px; color: var(--vanna-text-secondary);';
    container.appendChild(description);

    // Simple choice buttons
    const choiceGroup = withDefaults({
      id: 'choice-group',
      type: 'button_group',
      data: {
        buttons: [
          { label: 'Okay', action: 'okay', variant: 'primary' },
          { label: 'Not now', action: 'not_now', variant: 'secondary' },
          { label: 'Never', action: 'never', variant: 'ghost' },
        ],
        orientation: 'horizontal',
        spacing: 'medium',
        align: 'center',
      },
    });

    renderComponent(manager, choiceGroup);

    const codeExample = document.createElement('pre');
    codeExample.textContent = `// When you click "Okay", the message sent is: [Okay]
// When you click "Not now", the message sent is: [Not now]
// When you click "Never", the message sent is: [Never]`;
    codeExample.style.cssText = `
      margin-top: 20px;
      padding: 12px;
      background: rgba(0, 0, 0, 0.3);
      border-radius: 6px;
      color: #a0aec0;
      font-size: 12px;
      font-family: 'Courier New', monospace;
      overflow-x: auto;
    `;
    container.appendChild(codeExample);

    return container;
  },
};


================================================
FILE: frontends/webcomponent/src/components/dataframe-component.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { ComponentManager, ComponentUpdate } from './rich-component-system';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';
import { richComponentStyleText } from '../styles/rich-component-styles.js';

const meta: Meta = {
  title: 'Rich Components/DataFrame',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'light', value: '#f5f7fa' },
        { name: 'dark', value: 'rgb(11, 15, 25)' },
      ],
    },
  },
  argTypes: {
    theme: {
      control: { type: 'select' },
      options: ['light', 'dark'],
    },
    striped: {
      control: { type: 'boolean' },
    },
    bordered: {
      control: { type: 'boolean' },
    },
    compact: {
      control: { type: 'boolean' },
    },
    searchable: {
      control: { type: 'boolean' },
    },
    sortable: {
      control: { type: 'boolean' },
    },
    exportable: {
      control: { type: 'boolean' },
    },
  },
};

export default meta;
type Story = StoryObj;

const ensureTokenStyles = () => {
  if (document.getElementById('vanna-token-style')) {
    return;
  }

  const style = document.createElement('style');
  style.id = 'vanna-token-style';
  style.textContent = vannaDesignTokens.cssText.replace(/:host/g, '.vanna-tokens');
  document.head.appendChild(style);
};

const ensureRichComponentStyles = () => {
  if (document.getElementById('vanna-rich-component-styles')) {
    return;
  }

  const style = document.createElement('style');
  style.id = 'vanna-rich-component-styles';
  style.textContent = richComponentStyleText;
  document.head.appendChild(style);
};

const createContainer = () => {
  ensureTokenStyles();
  ensureRichComponentStyles();

  const container = document.createElement('div');
  container.className = 'vanna-tokens';
  container.style.cssText = `
    padding: var(--vanna-space-5, 20px);
    max-width: 1200px;
    margin: 0 auto;
    background: var(--vanna-background-default, #0b0f19);
    border-radius: var(--vanna-border-radius-lg, 8px);
    box-shadow: var(--vanna-shadow-md, 0 4px 6px rgba(0, 0, 0, 0.1));
    color: var(--vanna-foreground-default, #ffffff);
  `;

  // Add some additional DataFrame-specific debugging styles
  const additionalStyles = document.createElement('style');
  additionalStyles.textContent = `
    /* Ensure DataFrame styles are applied with higher specificity */
    .vanna-tokens {
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .rich-dataframe {
      background: var(--vanna-background-default, #0b0f19) !important;
      border: 1px solid var(--vanna-outline-default, #333) !important;
      border-radius: var(--vanna-border-radius-lg, 8px) !important;
      overflow: hidden !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .dataframe-table {
      width: 100% !important;
      border-collapse: collapse !important;
      font-size: 0.875rem !important;
      color: var(--vanna-foreground-default, #ffffff) !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .dataframe-table th {
      background: var(--vanna-background-higher, #1a1f2e) !important;
      color: var(--vanna-foreground-default, #ffffff) !important;
      font-weight: 600 !important;
      text-align: left !important;
      padding: 12px 16px !important;
      border-bottom: 2px solid var(--vanna-outline-default, #333) !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .dataframe-table td {
      padding: 12px 16px !important;
      border-bottom: 1px solid var(--vanna-outline-dimmer, #222) !important;
      color: var(--vanna-foreground-default, #ffffff) !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .dataframe-table.striped tbody tr:nth-child(even) {
      background: rgba(255, 255, 255, 0.02) !important;
    }

    .vanna-tokens .dataframe-header {
      padding: 16px 20px !important;
      background: var(--vanna-background-higher, #1a1f2e) !important;
      border-bottom: 1px solid var(--vanna-outline-default, #333) !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .dataframe-title {
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
      color: var(--vanna-foreground-default, #ffffff) !important;
      font-weight: 600 !important;
    }

    .vanna-tokens .dataframe-description {
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
      color: var(--vanna-foreground-dimmer, #b1bac4) !important;
    }

    .vanna-tokens .dataframe-actions {
      padding: 12px 20px !important;
      background: var(--vanna-background-default, #0b0f19) !important;
      border-bottom: 1px solid var(--vanna-outline-dimmer, #222) !important;
      display: flex !important;
      justify-content: space-between !important;
      align-items: center !important;
      gap: 12px !important;
    }

    .vanna-tokens .search-input {
      width: 100% !important;
      padding: 8px 12px !important;
      border: 1px solid var(--vanna-outline-default, #333) !important;
      border-radius: 6px !important;
      background: var(--vanna-background-default, #0b0f19) !important;
      color: var(--vanna-foreground-default, #ffffff) !important;
      font-size: 0.875rem !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }

    .vanna-tokens .export-btn {
      padding: 8px 12px !important;
      border: 1px solid var(--vanna-outline-default, #333) !important;
      border-radius: 6px !important;
      background: var(--vanna-background-default, #0b0f19) !important;
      color: var(--vanna-foreground-default, #ffffff) !important;
      cursor: pointer !important;
      font-size: 0.875rem !important;
      font-family: var(--vanna-font-family-default, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif) !important;
    }
  `;
  document.head.appendChild(additionalStyles);

  return container;
};

const createManager = (container: HTMLElement) => new ComponentManager(container);

const renderComponent = (manager: ComponentManager, component: any) => {
  const update: ComponentUpdate = {
    operation: 'create',
    target_id: component.id,
    component,
    timestamp: new Date().toISOString(),
  } as ComponentUpdate;

  manager.processUpdate(update);
};

const withDefaults = (component: any) => ({
  layout: { position: 'append', size: {}, z_index: 0, classes: [] },
  theme: {},
  lifecycle: 'create',
  timestamp: new Date().toISOString(),
  visible: true,
  interactive: false,
  children: [],
  ...component,
});

// Sample data sets
const employeeData = [
  { id: 1, name: 'Alice Johnson', email: 'alice@example.com', age: 28, city: 'New York', salary: 75000, active: true, department: 'Engineering' },
  { id: 2, name: 'Bob Smith', email: 'bob@example.com', age: 34, city: 'San Francisco', salary: 85000, active: true, department: 'Product' },
  { id: 3, name: 'Carol Davis', email: 'carol@example.com', age: 29, city: 'Chicago', salary: 70000, active: false, department: 'Design' },
  { id: 4, name: 'David Wilson', email: 'david@example.com', age: 42, city: 'Austin', salary: 90000, active: true, department: 'Engineering' },
  { id: 5, name: 'Eve Brown', email: 'eve@example.com', age: 31, city: 'Seattle', salary: 80000, active: true, department: 'Marketing' },
  { id: 6, name: 'Frank Miller', email: 'frank@example.com', age: 38, city: 'Boston', salary: 95000, active: false, department: 'Sales' },
  { id: 7, name: 'Grace Lee', email: 'grace@example.com', age: 26, city: 'Denver', salary: 65000, active: true, department: 'HR' },
  { id: 8, name: 'Henry Taylor', email: 'henry@example.com', age: 33, city: 'Portland', salary: 72000, active: true, department: 'Engineering' },
  { id: 9, name: 'Ivy Chen', email: 'ivy@example.com', age: 27, city: 'Los Angeles', salary: 78000, active: true, department: 'Product' },
  { id: 10, name: 'Jack Anderson', email: 'jack@example.com', age: 35, city: 'Miami', salary: 82000, active: false, department: 'Finance' },
];

const sqlQueryData = [
  { TrackId: 1, Name: 'For Those About To Rock (We Salute You)', AlbumId: 1, MediaTypeId: 1, GenreId: 1, Composer: 'Angus Young, Malcolm Young, Brian Johnson', Milliseconds: 343719, Bytes: 11170334, UnitPrice: 0.99 },
  { TrackId: 2, Name: 'Balls to the Wall', AlbumId: 2, MediaTypeId: 2, GenreId: 1, Composer: null, Milliseconds: 342562, Bytes: 5510424, UnitPrice: 0.99 },
  { TrackId: 3, Name: 'Fast As a Shark', AlbumId: 3, MediaTypeId: 2, GenreId: 1, Composer: 'F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman', Milliseconds: 230619, Bytes: 3990994, UnitPrice: 0.99 },
  { TrackId: 4, Name: 'Restless and Wild', AlbumId: 3, MediaTypeId: 2, GenreId: 1, Composer: 'F. Baltes, R.A. Smith-Diesel, S. Kaufman, U. Dirkscneider & W. Hoffman', Milliseconds: 252051, Bytes: 4331779, UnitPrice: 0.99 },
  { TrackId: 5, Name: 'Princess of the Dawn', AlbumId: 3, MediaTypeId: 2, GenreId: 1, Composer: 'Deaffy & R.A. Smith-Diesel', Milliseconds: 375418, Bytes: 6290521, UnitPrice: 0.99 },
  { TrackId: 6, Name: 'Put The Finger On You', AlbumId: 1, MediaTypeId: 1, GenreId: 1, Composer: 'Angus Young, Malcolm Young, Brian Johnson', Milliseconds: 205662, Bytes: 6713451, UnitPrice: 0.99 },
  { TrackId: 7, Name: "Let's Get It Up", AlbumId: 1, MediaTypeId: 1, GenreId: 1, Composer: 'Angus Young, Malcolm Young, Brian Johnson', Milliseconds: 233926, Bytes: 7636561, UnitPrice: 0.99 },
  { TrackId: 8, Name: 'Inject The Venom', AlbumId: 1, MediaTypeId: 1, GenreId: 1, Composer: 'Angus Young, Malcolm Young, Brian Johnson', Milliseconds: 210834, Bytes: 6852860, UnitPrice: 0.99 },
];

export const BasicDataFrame: Story = {
  render: (args) => {
    const container = createContainer();
    const manager = createManager(container);

    const component = withDefaults({
      id: 'basic-dataframe',
      type: 'dataframe',
      data: {
        data: employeeData.slice(0, 5),
        columns: ['id', 'name', 'email', 'age', 'city', 'department'],
        title: 'Employee Records',
        description: 'Basic employee data with essential information',
        row_count: 5,
        column_count: 6,
        striped: args.striped ?? true,
        bordered: args.bordered ?? true,
        compact: args.compact ?? false,
        searchable: args.searchable ?? false,
        sortable: args.sortable ?? false,
        exportable: args.exportable ?? false,
        column_types: {
          id: 'number',
          name: 'string',
          email: 'string',
          age: 'number',
          city: 'string',
          department: 'string'
        }
      },
    });

    renderComponent(manager, component);
    return container;
  },
  args: {
    striped: true,
    bordered: true,
    compact: false,
    searchable: false,
    sortable: false,
    exportable: false,
  },
};

export const InteractiveDataFrame: Story = {
  render: (args) => {
    const container = createContainer();
    const manager = createManager(container);

    const component = withDefaults({
      id: 'interactive-dataframe',
      type: 'dataframe',
      data: {
        data: employeeData,
        columns: ['id', 'name', 'email', 'age', 'city', 'salary', 'active', 'department'],
        title: 'Interactive Employee Database',
        description: 'Full dataset with search, sort, and export functionality',
        row_count: employeeData.length,
        column_count: 8,
        striped: args.striped ?? true,
        bordered: args.bordered ?? true,
        compact: args.compact ?? false,
        searchable: args.searchable ?? true,
        sortable: args.sortable ?? true,
        exportable: args.exportable ?? true,
        max_rows_displayed: 8,
        column_types: {
          id: 'number',
          name: 'string',
          email: 'string',
          age: 'number',
          city: 'string',
          salary: 'number',
          active: 'boolean',
          department: 'string'
        }
      },
    });

    renderComponent(manager, component);
    return container;
  },
  args: {
    striped: true,
    bordered: true,
    compact: false,
    searchable: true,
    sortable: true,
    exportable: true,
  },
};

export const SQLQueryResults: Story = {
  render: (args) => {
    const container = createContainer();
    const manager = createManager(container);

    const component = withDefaults({
      id: 'sql-dataframe',
      type: 'dataframe',
      data: {
        data: sqlQueryData,
        columns: ['TrackId', 'Name', 'AlbumId', 'MediaTypeId', 'GenreId', 'Composer', 'Milliseconds', 'Bytes', 'UnitPrice'],
        title: 'SQL Query Results',
        description: 'SELECT * FROM Track LIMIT 8',
        row_count: sqlQueryData.length,
        column_count: 9,
        striped: args.striped ?? true,
        bordered: args.bordered ?? true,
        compact: args.compact ?? false,
        searchable: args.searchable ?? true,
        sortable: args.sortable ?? true,
        exportable: args.exportable ?? true,
        column_types: {
          TrackId: 'number',
          Name: 'string',
          AlbumId: 'number',
          MediaTypeId: 'number',
          GenreId: 'number',
          Composer: 'string',
          Milliseconds: 'number',
          Bytes: 'number',
          UnitPrice: 'number'
        }
      },
    });

    renderComponent(manager, component);
    return container;
  },
  args: {
    striped: true,
    bordered: true,
    compact: false,
    searchable: true,
    sortable: true,
    exportable: true,
  },
};

export const CompactView: Story = {
  render: (args) => {
    const container = createContainer();
    const manager = createManager(container);

    const component = withDefaults({
      id: 'compact-dataframe',
      type: 'dataframe',
      data: {
        data: employeeData.slice(0, 6),
        columns: ['id', 'name', 'city', 'active'],
        title: 'Compact Employee View',
        description: 'Space-efficient display with essential columns only',
        row_count: 6,
        column_count: 4,
        striped: args.striped ?? true,
        bordered: args.bordered ?? false,
        compact: args.compact ?? true,
        searchable: args.searchable ?? false,
        sortable: args.sortable ?? true,
        exportable: args.exportable ?? false,
        column_types: {
          id: 'number',
          name: 'string',
          city: 'string',
          active: 'boolean'
        }
      },
    });

    renderComponent(manager, component);
    return container;
  },
  args: {
    striped: true,
    bordered: false,
    compact: true,
    searchable: false,
    sortable: true,
    exportable: false,
  },
};

export const EmptyDataFrame: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    const component = withDefaults({
      id: 'empty-dataframe',
      type: 'dataframe',
      data: {
        data: [],
        columns: [],
        title: 'No Data Available',
        description: 'This dataset contains no records',
        row_count: 0,
        column_count: 0,
      },
    });

    renderComponent(manager, component);
    return container;
  },
};

export const LargeDataset: Story = {
  render: (args) => {
    const container = createContainer();
    const manager = createManager(container);

    // Generate a larger dataset
    const largeData = Array.from({ length: 50 }, (_, i) => ({
      id: i + 1,
      name: `User ${i + 1}`,
      email: `user${i + 1}@example.com`,
      score: Math.floor(Math.random() * 100),
      category: ['A', 'B', 'C'][i % 3],
      active: Math.random() > 0.3,
      created_date: new Date(2024, Math.floor(Math.random() * 12), Math.floor(Math.random() * 28) + 1).toISOString().split('T')[0]
    }));

    const component = withDefaults({
      id: 'large-dataframe',
      type: 'dataframe',
      data: {
        data: largeData,
        columns: ['id', 'name', 'email', 'score', 'category', 'active', 'created_date'],
        title: 'Large Dataset',
        description: '50 records with pagination and search',
        row_count: largeData.length,
        column_count: 7,
        striped: args.striped ?? true,
        bordered: args.bordered ?? true,
        compact: args.compact ?? false,
        searchable: args.searchable ?? true,
        sortable: args.sortable ?? true,
        exportable: args.exportable ?? true,
        max_rows_displayed: 15,
        column_types: {
          id: 'number',
          name: 'string',
          email: 'string',
          score: 'number',
          category: 'string',
          active: 'boolean',
          created_date: 'date'
        }
      },
    });

    renderComponent(manager, component);
    return container;
  },
  args: {
    striped: true,
    bordered: true,
    compact: false,
    searchable: true,
    sortable: true,
    exportable: true,
  },
};

export const DataTypesShowcase: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    const typesData = [
      {
        id: 1,
        name: 'Alice',
        score: 95.5,
        active: true,
        created: '2024-01-15',
        notes: 'Excellent performance',
        tags: null
      },
      {
        id: 2,
        name: 'Bob',
        score: 87.2,
        active: false,
        created: '2024-02-20',
        notes: 'Good but needs improvement',
        tags: 'priority,review'
      },
      {
        id: 3,
        name: 'Carol',
        score: 92.8,
        active: true,
        created: '2024-03-10',
        notes: null,
        tags: 'star-performer'
      },
    ];

    const component = withDefaults({
      id: 'types-dataframe',
      type: 'dataframe',
      data: {
        data: typesData,
        columns: ['id', 'name', 'score', 'active', 'created', 'notes', 'tags'],
        title: 'Data Types Showcase',
        description: 'Demonstrates different column data types and null handling',
        row_count: typesData.length,
        column_count: 7,
        striped: true,
        bordered: true,
        searchable: true,
        sortable: true,
        exportable: true,
        column_types: {
          id: 'number',
          name: 'string',
          score: 'number',
          active: 'boolean',
          created: 'date',
          notes: 'string',
          tags: 'string'
        }
      },
    });

    renderComponent(manager, component);
    return container;
  },
};

================================================
FILE: frontends/webcomponent/src/components/plotly-chart.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './plotly-chart';

const meta: Meta = {
  title: 'Rich Components/Plotly Chart',
  component: 'plotly-chart',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'dark', value: 'rgb(11, 15, 25)' },
        { name: 'light', value: '#f5f7fa' },
      ],
    },
  },
  argTypes: {
    theme: {
      control: 'select',
      options: ['light', 'dark']
    },
    loading: { control: 'boolean' },
    error: { control: 'text' },
  },
};

export default meta;
type Story = StoryObj;

export const LineChart: Story = {
  args: {
    theme: 'light',
    loading: false,
    error: '',
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        ?loading=${args.loading}
        error=${args.error}
        .data=${[
          {
            x: ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'],
            y: [20, 14, 23, 25, 22, 16],
            type: 'scatter',
            mode: 'lines+markers',
            name: 'Sales',
            line: { color: 'rgb(0, 123, 255)' }
          }
        ]}
        .layout=${{
          xaxis: { title: 'Month' },
          yaxis: { title: 'Sales (in thousands)' }
        }}>
      </plotly-chart>
    </div>
  `,
};

export const BarChart: Story = {
  args: {
    theme: 'light',
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        .data=${[
          {
            x: ['Product A', 'Product B', 'Product C', 'Product D'],
            y: [45, 32, 28, 35],
            type: 'bar',
            name: 'Revenue',
            marker: {
              color: ['rgb(16, 185, 129)', 'rgb(0, 123, 255)', 'rgb(245, 158, 11)', 'rgb(239, 68, 68)']
            }
          }
        ]}
        .layout=${{
          xaxis: { title: 'Products' },
          yaxis: { title: 'Revenue ($M)' }
        }}>
      </plotly-chart>
    </div>
  `,
};

export const ScatterPlot: Story = {
  args: {
    theme: 'light',
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        .data=${[
          {
            x: [85, 78, 92, 88, 76, 95, 82, 89, 93, 79],
            y: [450, 320, 580, 490, 280, 650, 380, 520, 610, 310],
            type: 'scatter',
            mode: 'markers',
            name: 'Business Units',
            marker: {
              size: 12,
              color: 'rgb(0, 123, 255)',
              opacity: 0.7
            }
          }
        ]}
        .layout=${{
          xaxis: { title: 'Customer Satisfaction Score' },
          yaxis: { title: 'Revenue ($K)' }
        }}>
      </plotly-chart>
    </div>
  `,
};

export const MultipleLines: Story = {
  args: {
    theme: 'light',
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        .data=${[
          {
            x: ['Q1', 'Q2', 'Q3', 'Q4'],
            y: [85, 88, 92, 89],
            type: 'scatter',
            mode: 'lines+markers',
            name: 'User Engagement',
            line: { color: 'rgb(16, 185, 129)' }
          },
          {
            x: ['Q1', 'Q2', 'Q3', 'Q4'],
            y: [65, 72, 78, 81],
            type: 'scatter',
            mode: 'lines+markers',
            name: 'Conversion Rate',
            line: { color: 'rgb(0, 123, 255)' }
          },
          {
            x: ['Q1', 'Q2', 'Q3', 'Q4'],
            y: [42, 48, 55, 58],
            type: 'scatter',
            mode: 'lines+markers',
            name: 'Customer Retention',
            line: { color: 'rgb(245, 158, 11)' }
          }
        ]}
        .layout=${{
          xaxis: { title: 'Quarter' },
          yaxis: { title: 'Percentage (%)' }
        }}>
      </plotly-chart>
    </div>
  `,
};

export const LoadingState: Story = {
  args: {
    theme: 'light',
    loading: true,
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        ?loading=${args.loading}
        .data=${[]}>
      </plotly-chart>
    </div>
  `,
};

export const ErrorState: Story = {
  args: {
    theme: 'light',
    error: 'Failed to load chart data from API',
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        error=${args.error}
        .data=${[]}>
      </plotly-chart>
    </div>
  `,
};

export const LightTheme: Story = {
  args: {
    theme: 'light',
  },
  parameters: {
    backgrounds: { default: 'light' }
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        .data=${[
          {
            x: ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun'],
            y: [20, 14, 23, 25, 22, 16],
            type: 'scatter',
            mode: 'lines+markers',
            name: 'Sales',
            line: { color: 'rgb(0, 123, 255)' }
          }
        ]}
        .layout=${{
          xaxis: { title: 'Month' },
          yaxis: { title: 'Sales (in thousands)' }
        }}>
      </plotly-chart>
    </div>
  `,
};

export const CustomLayout: Story = {
  args: {
    theme: 'light',
  },
  render: (args) => html`
    <div style="max-width: 800px; margin: 0 auto;">
      <plotly-chart
        theme=${args.theme}
        .data=${[
          {
            x: [1, 2, 3, 4, 5],
            y: [10, 11, 12, 13, 14],
            type: 'scatter',
            mode: 'lines',
            name: 'Trend A',
            line: { color: 'rgb(16, 185, 129)', width: 3 }
          },
          {
            x: [1, 2, 3, 4, 5],
            y: [8, 9, 10, 11, 12],
            type: 'scatter',
            mode: 'lines',
            name: 'Trend B',
            line: { color: 'rgb(239, 68, 68)', width: 3, dash: 'dash' }
          }
        ]}
        .layout=${{
          title: {
            text: 'Custom Styled Chart',
            font: { size: 18 }
          },
          xaxis: {
            title: 'Time Period',
            gridcolor: 'rgba(255, 255, 255, 0.1)'
          },
          yaxis: {
            title: 'Value',
            gridcolor: 'rgba(255, 255, 255, 0.1)'
          },
          height: 500,
          showlegend: true
        }}>
      </plotly-chart>
    </div>
  `,
};



================================================
FILE: frontends/webcomponent/src/components/plotly-chart.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';
import Plotly from 'plotly.js-dist-min';

export interface PlotlyData {
  x?: any[];
  y?: any[];
  type?: any;
  mode?: any;
  name?: string;
  marker?: any;
  line?: any;
  [key: string]: any;
}

export interface PlotlyLayout {
  title?: any;
  xaxis?: any;
  yaxis?: any;
  font?: any;
  paper_bgcolor?: string;
  plot_bgcolor?: string;
  margin?: any;
  showlegend?: boolean;
  height?: number;
  width?: number;
  modebar?: any;
  [key: string]: any;
}

@customElement('plotly-chart')
export class PlotlyChart extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        font-family: var(--vanna-font-family-default);
        width: 100%;
        height: 100%;
      }

      .plotly-div {
        width: 100%;
        min-height: 400px;
      }

      /* Plotly layering fix for Shadow DOM */
      .plotly-div,
      .plotly-div .js-plotly-plot,
      .plotly-div .plot-container,
      .plotly-div .svg-container {
        position: relative;
        width: 100%;
        height: 100%;
      }

      .plotly-div svg.main-svg {
        position: absolute;
        top: 0;
        left: 0;
      }

      .plotly-div .hoverlayer {
        pointer-events: none;
      }

      .error-message {
        padding: var(--vanna-space-4);
        color: var(--vanna-accent-negative-default);
        text-align: center;
        font-style: italic;
      }

      .loading-message {
        padding: var(--vanna-space-4);
        color: var(--vanna-foreground-dimmer);
        text-align: center;
        font-style: italic;
      }
    `
  ];

  @property({ type: Array }) data: PlotlyData[] = [];
  @property({ type: Object }) layout: PlotlyLayout = {};
  @property({ type: Object }) config = {};
  @property({ type: Boolean }) loading = false;
  @property() error = '';
  @property() theme: 'light' | 'dark' = 'dark';

  private plotlyDiv?: HTMLElement;
  private resizeObserver?: ResizeObserver;

  firstUpdated() {
    this.plotlyDiv = this.shadowRoot?.querySelector('.plotly-div') as HTMLElement;
    this._renderChart();
    this._setupResizeObserver();
  }

  disconnectedCallback() {
    super.disconnectedCallback();
    this.resizeObserver?.disconnect();
  }

  private _setupResizeObserver() {
    if (!this.plotlyDiv) return;

    this.resizeObserver = new ResizeObserver(() => {
      if (this.plotlyDiv && this.data.length > 0) {
        const width = this.plotlyDiv.offsetWidth;
        Plotly.relayout(this.plotlyDiv, { width });
      }
    });

    this.resizeObserver.observe(this.plotlyDiv);
  }

  updated(changedProperties: Map<string | number | symbol, unknown>) {
    if (changedProperties.has('data') || changedProperties.has('layout') || changedProperties.has('theme')) {
      this._renderChart();
    }
  }

  private _getDefaultLayout(): PlotlyLayout {
    const isDark = this.theme === 'dark';

    // Start with layout from backend (which may include white background)
    const mergedLayout = {
      ...this.layout,
      // Only add font/modebar if not already set by backend
      font: this.layout.font || {
        family: 'ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif',
        color: isDark ? 'rgb(242, 244, 247)' : 'rgb(17, 24, 39)',
        size: 12
      },
      modebar: this.layout.modebar || {
        bgcolor: isDark ? 'rgba(21, 26, 38, 0.8)' : 'rgba(255, 255, 255, 0.8)',
        color: isDark ? 'rgb(177, 186, 196)' : 'rgb(75, 85, 99)',
        activecolor: isDark ? 'rgb(242, 244, 247)' : 'rgb(17, 24, 39)',
        orientation: 'h'
      },
      // Set explicit dimensions for Shadow DOM compatibility
      autosize: false,
      width: this.layout.width || undefined,
      height: this.layout.height || 400,
    };

    // If backend didn't set background colors, use transparent
    if (!this.layout.paper_bgcolor) {
      mergedLayout.paper_bgcolor = 'transparent';
    }
    if (!this.layout.plot_bgcolor) {
      mergedLayout.plot_bgcolor = 'transparent';
    }

    return mergedLayout;
  }

  private _getDefaultConfig() {
    return {
      responsive: true,
      displayModeBar: false,
      ...this.config
    };
  }

  private async _renderChart() {
    if (!this.plotlyDiv || this.loading || this.error || this.data.length === 0) {
      return;
    }

    try {
      const layout = this._getDefaultLayout();
      const config = this._getDefaultConfig();

      await Plotly.newPlot(this.plotlyDiv, this.data, layout, config);
    } catch (err) {
      this.error = err instanceof Error ? err.message : 'Failed to render chart';
      console.error('Plotly chart error:', err);
    }
  }

  render() {
    return html`
      ${this.loading ? html`
        <div class="loading-message">Loading chart...</div>
      ` : this.error ? html`
        <div class="error-message">Error: ${this.error}</div>
      ` : html`
        <div class="plotly-div"></div>
      `}
    `;
  }
}

declare global {
  interface HTMLElementTagNameMap {
    'plotly-chart': PlotlyChart;
  }
}

================================================
FILE: frontends/webcomponent/src/components/rich-card.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './rich-card';

const meta: Meta = {
  title: 'Rich Components/Rich Card',
  component: 'rich-card',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'dark', value: 'rgb(11, 15, 25)' },
        { name: 'light', value: '#f5f7fa' },
      ],
    },
  },
  argTypes: {
    title: { control: 'text' },
    subtitle: { control: 'text' },
    content: { control: 'text' },
    icon: { control: 'text' },
    status: {
      control: 'select',
      options: ['info', 'success', 'warning', 'error']
    },
    collapsible: { control: 'boolean' },
    collapsed: { control: 'boolean' },
  },
};

export default meta;
type Story = StoryObj;

export const Default: Story = {
  args: {
    title: 'Sample Card',
    subtitle: 'This is a subtitle',
    content: 'This is the content of the card. It can contain any text or HTML.',
    status: 'info',
    collapsible: false,
    collapsed: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-card
        title=${args.title}
        subtitle=${args.subtitle}
        content=${args.content}
        icon=${args.icon}
        status=${args.status}
        ?collapsible=${args.collapsible}
        ?collapsed=${args.collapsed}>
      </rich-card>
    </div>
  `,
};

export const WithIcon: Story = {
  args: {
    title: 'Card with Icon',
    subtitle: 'Featuring an emoji icon',
    content: 'This card demonstrates how icons work with the rich card component.',
    icon: '🚀',
    status: 'success',
    collapsible: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-card
        title=${args.title}
        subtitle=${args.subtitle}
        content=${args.content}
        icon=${args.icon}
        status=${args.status}
        ?collapsible=${args.collapsible}>
      </rich-card>
    </div>
  `,
};

export const WithActions: Story = {
  args: {
    title: 'Interactive Card',
    subtitle: 'With action buttons',
    content: 'This card includes action buttons that can trigger events.',
    status: 'info',
    collapsible: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-card
        title=${args.title}
        subtitle=${args.subtitle}
        content=${args.content}
        status=${args.status}
        ?collapsible=${args.collapsible}
        .actions=${[
          { label: 'Primary Action', action: 'primary', variant: 'primary' },
          { label: 'Secondary Action', action: 'secondary', variant: 'secondary' }
        ]}
        @card-action=${(e: CustomEvent) => {
          console.log('Card action:', e.detail.action);
          alert(`Action triggered: ${e.detail.action}`);
        }}>
      </rich-card>
    </div>
  `,
};

export const Collapsible: Story = {
  args: {
    title: 'Collapsible Card',
    subtitle: 'Click to expand/collapse',
    content: 'This content can be hidden by clicking the toggle button in the header.',
    status: 'warning',
    collapsible: true,
    collapsed: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-card
        title=${args.title}
        subtitle=${args.subtitle}
        content=${args.content}
        status=${args.status}
        ?collapsible=${args.collapsible}
        ?collapsed=${args.collapsed}>
      </rich-card>
    </div>
  `,
};

export const StatusVariants: Story = {
  render: () => html`
    <div style="max-width: 600px; margin: 0 auto; display: flex; flex-direction: column; gap: 16px;">
      <rich-card
        title="Info Status"
        content="This card shows the info status variant."
        status="info">
      </rich-card>

      <rich-card
        title="Success Status"
        content="This card shows the success status variant."
        status="success">
      </rich-card>

      <rich-card
        title="Warning Status"
        content="This card shows the warning status variant."
        status="warning">
      </rich-card>

      <rich-card
        title="Error Status"
        content="This card shows the error status variant."
        status="error">
      </rich-card>
    </div>
  `,
};

export const LightTheme: Story = {
  args: {
    title: 'Light Theme Card',
    subtitle: 'Styled for light backgrounds',
    content: 'This card is displayed with the light theme variant.',
    icon: '☀️',
    status: 'success',
  },
  parameters: {
    backgrounds: { default: 'light' }
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-card
        theme="light"
        title=${args.title}
        subtitle=${args.subtitle}
        content=${args.content}
        icon=${args.icon}
        status=${args.status}>
      </rich-card>
    </div>
  `,
};

================================================
FILE: frontends/webcomponent/src/components/rich-card.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

export interface CardAction {
  label: string;
  action: string;
  variant?: 'primary' | 'secondary';
}

@customElement('rich-card')
export class RichCard extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        margin-bottom: var(--vanna-space-4);
        font-family: var(--vanna-font-family-default);
      }

      .card {
        border: 1px solid var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-lg);
        background: var(--vanna-background-default);
        box-shadow: var(--vanna-shadow-sm);
        overflow: hidden;
        transition: box-shadow var(--vanna-duration-200) ease;
      }

      .card:hover {
        box-shadow: var(--vanna-shadow-md);
      }

      .card-header {
        display: flex;
        align-items: center;
        padding: var(--vanna-space-4) var(--vanna-space-5);
        background: var(--vanna-background-higher);
        border-bottom: 1px solid var(--vanna-outline-default);
        gap: var(--vanna-space-3);
      }

      .card-header.collapsible {
        cursor: pointer;
      }

      .card-icon {
        font-size: 1.25rem;
        display: flex;
        align-items: center;
      }

      .card-title-section {
        flex: 1;
      }

      .card-title {
        margin: 0;
        font-size: 1rem;
        font-weight: 600;
        color: var(--vanna-foreground-default);
      }

      .card-subtitle {
        margin: var(--vanna-space-1) 0 0 0;
        font-size: 0.875rem;
        color: var(--vanna-foreground-dimmer);
      }

      .card-status {
        padding: var(--vanna-space-1) var(--vanna-space-2);
        border-radius: var(--vanna-border-radius-md);
        font-size: 0.75rem;
        font-weight: 600;
        text-transform: uppercase;
      }

      .card-status.status-success {
        background: #d4edda;
        color: #155724;
      }

      .card-status.status-warning {
        background: #fff3cd;
        color: #856404;
      }

      .card-status.status-error {
        background: #f8d7da;
        color: #721c24;
      }

      .card-status.status-info {
        background: #d1ecf1;
        color: #0c5460;
      }

      .card-toggle {
        background: none;
        border: none;
        cursor: pointer;
        font-size: 1rem;
        color: var(--vanna-foreground-dimmer);
        padding: var(--vanna-space-1);
        border-radius: var(--vanna-border-radius-sm);
        transition: background-color var(--vanna-duration-200) ease;
      }

      .card-toggle:hover {
        background: var(--vanna-background-root);
      }

      .card-content {
        padding: var(--vanna-space-4) var(--vanna-space-5);
        line-height: 1.5;
        color: var(--vanna-foreground-default);
        transition: all var(--vanna-duration-200) ease;
        overflow: hidden;
      }

      .card-content.collapsed {
        max-height: 0;
        padding-top: 0;
        padding-bottom: 0;
      }

      .card-content h1,
      .card-content h2,
      .card-content h3 {
        margin: var(--vanna-space-2) 0;
        font-weight: 600;
      }

      .card-content h1 {
        font-size: 1.5rem;
      }

      .card-content h2 {
        font-size: 1.25rem;
      }

      .card-content h3 {
        font-size: 1.125rem;
      }

      .card-content p {
        margin: var(--vanna-space-2) 0;
      }

      .card-content ul {
        margin: var(--vanna-space-2) 0;
        padding-left: var(--vanna-space-5);
      }

      .card-content li {
        margin: var(--vanna-space-1) 0;
      }

      .card-content code {
        background: var(--vanna-background-higher);
        padding: var(--vanna-space-1) var(--vanna-space-2);
        border-radius: var(--vanna-border-radius-sm);
        font-family: monospace;
        font-size: 0.875em;
      }

      .card-content strong {
        font-weight: 600;
      }

      .card-actions {
        padding: var(--vanna-space-3) var(--vanna-space-5);
        background: var(--vanna-background-root);
        border-top: 1px solid var(--vanna-outline-default);
        display: flex;
        gap: var(--vanna-space-2);
      }

      .card-action {
        padding: var(--vanna-space-2) var(--vanna-space-4);
        border-radius: var(--vanna-border-radius-md);
        border: 1px solid var(--vanna-outline-default);
        background: var(--vanna-background-default);
        color: var(--vanna-foreground-default);
        cursor: pointer;
        font-size: 0.875rem;
        font-weight: 500;
        transition: all var(--vanna-duration-200) ease;
      }

      .card-action:hover {
        background: var(--vanna-background-higher);
      }

      .card-action.primary {
        background: var(--vanna-accent-primary-default);
        color: white;
        border-color: var(--vanna-accent-primary-default);
      }

      .card-action.primary:hover {
        background: var(--vanna-accent-primary-stronger);
      }
    `
  ];

  @property() title = '';
  @property() subtitle = '';
  @property() content = '';
  @property() icon = '';
  @property() status: 'info' | 'success' | 'warning' | 'error' = 'info';
  @property({ type: Array }) actions: CardAction[] = [];
  @property({ type: Boolean }) collapsible = false;
  @property({ type: Boolean }) collapsed = false;
  @property({ type: Boolean }) markdown = false;
  @property() theme: 'light' | 'dark' = 'dark';

  private _toggleCollapsed() {
    if (this.collapsible) {
      this.collapsed = !this.collapsed;
    }
  }

  private _renderMarkdown(text: string): string {
    // Simple markdown rendering - basic formatting
    return text
      .replace(/^### (.*$)/gm, '<h3>$1</h3>')
      .replace(/^## (.*$)/gm, '<h2>$1</h2>')
      .replace(/^# (.*$)/gm, '<h1>$1</h1>')
      .replace(/\*\*(.*?)\*\*/g, '<strong>$1</strong>')
      .replace(/\*(.*?)\*/g, '<em>$1</em>')
      .replace(/`([^`]+)`/g, '<code>$1</code>')
      .replace(/^- (.*$)/gm, '<li>$1</li>')
      .replace(/(<li>.*<\/li>)/s, '<ul>$1</ul>')
      .replace(/\n\n/g, '</p><p>')
      .replace(/^(?!<[h|u|l|p])(.+)$/gm, '<p>$1</p>');
  }

  render() {
    const contentHtml = this.markdown
      ? html`<div class="card-content ${this.collapsed ? 'collapsed' : ''}" .innerHTML=${this._renderMarkdown(this.content)}></div>`
      : html`<div class="card-content ${this.collapsed ? 'collapsed' : ''}">${this.content}</div>`;

    return html`
      <div class="card">
        <div class="card-header ${this.collapsible ? 'collapsible' : ''}"
             @click=${this._toggleCollapsed}>
          ${this.icon ? html`<span class="card-icon">${this.icon}</span>` : ''}
          <div class="card-title-section">
            <h3 class="card-title">${this.title}</h3>
            ${this.subtitle ? html`<p class="card-subtitle">${this.subtitle}</p>` : ''}
          </div>
          ${this.status ? html`<span class="card-status status-${this.status}">${this.status}</span>` : ''}
          ${this.collapsible ? html`
            <button class="card-toggle">${this.collapsed ? '▶' : '▼'}</button>
          ` : ''}
        </div>
        ${contentHtml}
        ${this.actions.length > 0 ? html`
          <div class="card-actions">
            ${this.actions.map(action => html`
              <button class="card-action ${action.variant || 'secondary'}"
                      @click=${() => this._handleAction(action.action)}>
                ${action.label}
              </button>
            `)}
          </div>
        ` : ''}
      </div>
    `;
  }

  private async _handleAction(action: string) {
    console.log('🔘 Card action button clicked (rich-card)');
    console.log('   Action:', action);

    // Dispatch event for any listeners
    this.dispatchEvent(new CustomEvent('card-action', {
      detail: { action },
      bubbles: true,
      composed: true
    }));

    // Also directly send to vanna-chat
    const vannaChat = document.querySelector('vanna-chat') as any;
    if (vannaChat && typeof vannaChat.sendMessage === 'function') {
      console.log('   Found vanna-chat, sending message...');
      try {
        const success = await vannaChat.sendMessage(action);
        if (success) {
          console.log('   ✅ Action sent successfully');
        } else {
          console.error('   ❌ Failed to send action');
        }
      } catch (error) {
        console.error('   ❌ Error sending action:', error);
      }
    } else {
      console.warn('   ⚠️ vanna-chat component not found or sendMessage not available');
    }
  }
}

declare global {
  interface HTMLElementTagNameMap {
    'rich-card': RichCard;
  }
}

================================================
FILE: frontends/webcomponent/src/components/rich-component-system.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { ComponentManager, ComponentUpdate } from './rich-component-system';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

const meta: Meta = {
  title: 'Rich Components/Component System',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'light', value: '#f5f7fa' },
        { name: 'dark', value: 'rgb(11, 15, 25)' },
      ],
    },
  },
};

export default meta;
type Story = StoryObj;

const ensureTokenStyles = () => {
  if (document.getElementById('vanna-token-style')) {
    return;
  }

  const style = document.createElement('style');
  style.id = 'vanna-token-style';
  style.textContent = vannaDesignTokens.cssText.replace(/:host/g, '.vanna-tokens');
  document.head.appendChild(style);
};

const createContainer = () => {
  ensureTokenStyles();

  const container = document.createElement('div');
  container.className = 'vanna-tokens';
  container.style.cssText = `
    padding: var(--vanna-space-5, 20px);
    max-width: 800px;
    margin: 0 auto;
    background: var(--vanna-background-default);
    border-radius: var(--vanna-border-radius-lg);
    box-shadow: var(--vanna-shadow-md);
  `;

  return container;
};

const createManager = (container: HTMLElement) => new ComponentManager(container);

const renderComponent = (manager: ComponentManager, component: any) => {
  const update: ComponentUpdate = {
    operation: 'create',
    target_id: component.id,
    component,
    timestamp: new Date().toISOString(),
  } as ComponentUpdate;

  manager.processUpdate(update);
};

const withDefaults = (component: any) => ({
  layout: { position: 'append', size: {}, z_index: 0, classes: [] },
  theme: {},
  lifecycle: 'create',
  ...component,
});

export const NotificationComponents: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    const components = [
      withDefaults({
        id: 'info-notification',
        type: 'notification',
        data: {
          message: 'This is an informational message',
          title: 'Information',
          level: 'info',
          dismissible: true,
          actions: [],
        },
      }),
      withDefaults({
        id: 'success-notification',
        type: 'notification',
        data: {
          message: 'Operation completed successfully!',
          title: 'Success',
          level: 'success',
          dismissible: true,
          actions: [
            { label: 'View Details', action: 'view', variant: 'primary' },
            { label: 'Dismiss', action: 'dismiss', variant: 'secondary' },
          ],
        },
      }),
      withDefaults({
        id: 'warning-notification',
        type: 'notification',
        data: {
          message: 'Please review the configuration before proceeding',
          title: 'Warning',
          level: 'warning',
          dismissible: true,
          actions: [],
        },
      }),
      withDefaults({
        id: 'error-notification',
        type: 'notification',
        data: {
          message: 'Failed to connect to the database. Please check your connection.',
          title: 'Connection Error',
          level: 'error',
          dismissible: true,
          actions: [
            { label: 'Retry', action: 'retry', variant: 'primary' },
            { label: 'Cancel', action: 'cancel', variant: 'secondary' },
          ],
        },
      }),
    ];

    components.forEach((component) => renderComponent(manager, component));

    return container;
  },
};

export const StatusIndicatorComponents: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    const statuses = [
      { status: 'loading', message: 'Processing your request...', pulse: true },
      { status: 'success', message: 'Request completed successfully', pulse: false },
      { status: 'warning', message: 'Operation completed with warnings', pulse: false },
      { status: 'error', message: 'Request failed - please try again', pulse: false },
    ];

    statuses.forEach((statusData, index) => {
      const component = withDefaults({
        id: `status-${index}`,
        type: 'status_indicator',
        data: statusData,
      });

      renderComponent(manager, component);
    });

    return container;
  },
};

export const TextComponents: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    const plainText = withDefaults({
      id: 'plain-text',
      type: 'text',
      data: {
        content: 'This is a plain text component with some sample content to demonstrate text rendering.',
        markdown: false,
      },
    });

    const markdownText = withDefaults({
      id: 'markdown-text',
      type: 'text',
      data: {
        content: `# Rich Components Demo\n\nThis is a **markdown** text component with various formatting:\n\n- **Bold text** for emphasis\n- *Italic text* for style\n- Lists for organization\n\n## Features\n\nThe text component supports:\n- Plain text rendering\n- Basic markdown formatting\n- Code syntax highlighting`,
        markdown: true,
      },
    });

    [plainText, markdownText].forEach((component) => renderComponent(manager, component));

    return container;
  },
};

export const DataFrameComponents: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    // Sample data for different scenarios
    const sampleData = [
      { id: 1, name: 'Alice Johnson', email: 'alice@example.com', age: 28, city: 'New York', salary: 75000, active: true },
      { id: 2, name: 'Bob Smith', email: 'bob@example.com', age: 34, city: 'San Francisco', salary: 85000, active: true },
      { id: 3, name: 'Carol Davis', email: 'carol@example.com', age: 29, city: 'Chicago', salary: 70000, active: false },
      { id: 4, name: 'David Wilson', email: 'david@example.com', age: 42, city: 'Austin', salary: 90000, active: true },
      { id: 5, name: 'Eve Brown', email: 'eve@example.com', age: 31, city: 'Seattle', salary: 80000, active: true },
      { id: 6, name: 'Frank Miller', email: 'frank@example.com', age: 38, city: 'Boston', salary: 95000, active: false },
      { id: 7, name: 'Grace Lee', email: 'grace@example.com', age: 26, city: 'Denver', salary: 65000, active: true },
      { id: 8, name: 'Henry Taylor', email: 'henry@example.com', age: 33, city: 'Portland', salary: 72000, active: true },
    ];

    const columns = ['id', 'name', 'email', 'age', 'city', 'salary', 'active'];

    // Basic DataFrame
    const basicDataFrame = withDefaults({
      id: 'basic-dataframe',
      type: 'dataframe',
      data: {
        data: sampleData.slice(0, 5),
        columns: columns,
        title: 'Employee Records',
        description: 'Sample employee data with various fields',
        row_count: 5,
        column_count: columns.length,
        column_types: {
          id: 'number',
          name: 'string',
          email: 'string',
          age: 'number',
          city: 'string',
          salary: 'number',
          active: 'boolean'
        }
      },
    });

    // Large DataFrame with all features
    const fullDataFrame = withDefaults({
      id: 'full-dataframe',
      type: 'dataframe',
      data: {
        data: sampleData,
        columns: columns,
        title: 'Complete Employee Database',
        description: 'Full dataset with search, sort, and export functionality',
        row_count: sampleData.length,
        column_count: columns.length,
        searchable: true,
        sortable: true,
        filterable: true,
        exportable: true,
        striped: true,
        bordered: true,
        max_rows_displayed: 6,
        column_types: {
          id: 'number',
          name: 'string',
          email: 'string',
          age: 'number',
          city: 'string',
          salary: 'number',
          active: 'boolean'
        }
      },
    });

    // Empty DataFrame
    const emptyDataFrame = withDefaults({
      id: 'empty-dataframe',
      type: 'dataframe',
      data: {
        data: [],
        columns: [],
        title: 'Empty Dataset',
        description: 'No data available to display',
        row_count: 0,
        column_count: 0,
      },
    });

    // Compact DataFrame
    const compactDataFrame = withDefaults({
      id: 'compact-dataframe',
      type: 'dataframe',
      data: {
        data: sampleData.slice(0, 4),
        columns: ['id', 'name', 'city', 'active'],
        title: 'Compact View',
        description: 'Space-efficient display with essential columns only',
        row_count: 4,
        column_count: 4,
        compact: true,
        searchable: false,
        exportable: false,
        column_types: {
          id: 'number',
          name: 'string',
          city: 'string',
          active: 'boolean'
        }
      },
    });

    [basicDataFrame, fullDataFrame, emptyDataFrame, compactDataFrame].forEach((component) => {
      renderComponent(manager, component);
    });

    return container;
  },
};

export const SQLQueryDataFrame: Story = {
  render: () => {
    const container = createContainer();
    const manager = createManager(container);

    // SQL query result simulation
    const sqlResultData = [
      { TrackId: 1, Name: 'For Those About To Rock (We Salute You)', AlbumId: 1, MediaTypeId: 1, GenreId: 1, Composer: 'Angus Young, Malcolm Young, Brian Johnson', Milliseconds: 343719, Bytes: 11170334, UnitPrice: 0.99 },
      { TrackId: 2, Name: 'Balls to the Wall', AlbumId: 2, MediaTypeId: 2, GenreId: 1, Composer: null, Milliseconds: 342562, Bytes: 5510424, UnitPrice: 0.99 },
      { TrackId: 3, Name: 'Fast As a Shark', AlbumId: 3, MediaTypeId: 2, GenreId: 1, Composer: 'F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman', Milliseconds: 230619, Bytes: 3990994, UnitPrice: 0.99 },
      { TrackId: 4, Name: 'Restless and Wild', AlbumId: 3, MediaTypeId: 2, GenreId: 1, Composer: 'F. Baltes, R.A. Smith-Diesel, S. Kaufman, U. Dirkscneider & W. Hoffman', Milliseconds: 252051, Bytes: 4331779, UnitPrice: 0.99 },
      { TrackId: 5, Name: 'Princess of the Dawn', AlbumId: 3, MediaTypeId: 2, GenreId: 1, Composer: 'Deaffy & R.A. Smith-Diesel', Milliseconds: 375418, Bytes: 6290521, UnitPrice: 0.99 },
    ];

    const sqlColumns = ['TrackId', 'Name', 'AlbumId', 'MediaTypeId', 'GenreId', 'Composer', 'Milliseconds', 'Bytes', 'UnitPrice'];

    const sqlDataFrame = withDefaults({
      id: 'sql-dataframe',
      type: 'dataframe',
      data: {
        data: sqlResultData,
        columns: sqlColumns,
        title: 'Query Results',
        description: 'SELECT * FROM Track LIMIT 5',
        row_count: sqlResultData.length,
        column_count: sqlColumns.length,
        searchable: true,
        sortable: true,
        exportable: true,
        column_types: {
          TrackId: 'number',
          Name: 'string',
          AlbumId: 'number',
          MediaTypeId: 'number',
          GenreId: 'number',
          Composer: 'string',
          Milliseconds: 'number',
          Bytes: 'number',
          UnitPrice: 'number'
        }
      },
    });

    renderComponent(manager, sqlDataFrame);

    return container;
  },
};


================================================
FILE: frontends/webcomponent/src/components/rich-component-system.ts
================================================
/**
 * Rich Component System for Vanna Agents
 *
 * Provides a generic component registry and rendering system that can display
 * any rich component sent from the Python backend.
 */

import { richComponentStyleText } from '../styles/rich-component-styles.js';

// Component interfaces matching Python backend
export interface RichComponent {
  id: string;
  type: string;
  lifecycle: 'create' | 'update' | 'replace' | 'remove';
  data: Record<string, any>;
  children: string[];
  timestamp: string;
  visible: boolean;
  interactive: boolean;
}

// Artifact event interfaces
export interface ArtifactOpenedEventDetail {
  // Core identification
  artifactId: string;

  // Artifact content
  content: string; // Full HTML/SVG/JS content
  type: 'html' | 'svg' | 'visualization' | 'interactive' | 'd3' | 'threejs';
  title?: string;
  description?: string;

  // Trigger context
  trigger: 'created' | 'user-action'; // How this event was fired

  // Control
  preventDefault: () => void; // Prevent default behavior

  // Helpers
  getStandaloneHTML: () => string; // Full page HTML with dependencies

  // Metadata
  timestamp: string;
}

declare global {
  interface GlobalEventHandlersEventMap {
    'artifact-opened': CustomEvent<ArtifactOpenedEventDetail>;
  }
}


const RICH_COMPONENT_STYLE_ATTR = 'data-vanna-rich-component-styles';

function ensureRichComponentStyles(container: HTMLElement): void {
  const doc = container.ownerDocument;
  if (!doc) {
    return;
  }

  if (container.querySelector(`style[${RICH_COMPONENT_STYLE_ATTR}]`)) {
    return;
  }

  const styleEl = doc.createElement('style');
  styleEl.setAttribute(RICH_COMPONENT_STYLE_ATTR, 'true');
  styleEl.textContent = richComponentStyleText;
  container.prepend(styleEl);
}

export interface ComponentUpdate {
  operation: 'create' | 'update' | 'replace' | 'remove' | 'reorder' | 'bulk_update';
  target_id: string;
  component?: RichComponent;
  updates?: Record<string, any>;
  position?: any;
  timestamp: string;
  batch_id?: string;
}

// Component renderer interface
export interface ComponentRenderer {
  render(component: RichComponent): HTMLElement;
  update(element: HTMLElement, component: RichComponent, updates?: Record<string, any>): void;
  remove(element: HTMLElement): void;
}

// Base component renderer with common functionality
export abstract class BaseComponentRenderer implements ComponentRenderer {
  abstract render(component: RichComponent): HTMLElement;

  update(element: HTMLElement, component: RichComponent, _updates?: Record<string, any>): void {
    // Default implementation - re-render completely
    const newElement = this.render(component);
    element.parentNode?.replaceChild(newElement, element);
  }

  remove(element: HTMLElement): void {
    element.remove();
  }

}

// Card component renderer
export class CardComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    console.log('🎴 CardComponentRenderer.render() called', {
      componentId: component.id,
      componentData: component.data,
      actions: component.data?.actions
    });

    const card = document.createElement('div');
    card.className = 'rich-component rich-card';
    card.dataset.componentId = component.id;

    const { title, content, subtitle, icon, status, actions = [], collapsible, collapsed } = component.data;

    console.log('🎴 Extracted actions:', actions, 'Length:', actions?.length);

    card.innerHTML = `
      <div class="card-header ${collapsible ? 'collapsible' : ''}">
        ${icon ? `<span class="card-icon">${icon}</span>` : ''}
        <div class="card-title-section">
          <h3 class="card-title">${title}</h3>
          ${subtitle ? `<p class="card-subtitle">${subtitle}</p>` : ''}
        </div>
        ${status ? `<span class="card-status status-${status}">${status}</span>` : ''}
        ${collapsible ? `<button class="card-toggle">${collapsed ? '▶' : '▼'}</button>` : ''}
      </div>
      <div class="card-content ${collapsed ? 'collapsed' : ''}">
        ${content}
      </div>
      ${actions && actions.length > 0 ? `
        <div class="card-actions">
          ${actions.map((action: any) => `
            <button class="card-action ${action.variant || 'secondary'}" data-action="${action.action}">
              ${action.label}
            </button>
          `).join('')}
        </div>
      ` : ''}
    `;


    // Add collapsible functionality
    if (collapsible) {
      const toggle = card.querySelector('.card-toggle') as HTMLButtonElement;
      const content = card.querySelector('.card-content') as HTMLElement;

      toggle?.addEventListener('click', () => {
        content.classList.toggle('collapsed');
        toggle.textContent = content.classList.contains('collapsed') ? '▶' : '▼';
      });
    }

    // Add click handlers for action buttons
    console.log('🎴 Checking if should add click handlers:', {
      hasActions: !!actions,
      actionsLength: actions?.length
    });

    if (actions && actions.length > 0) {
      const actionButtons = card.querySelectorAll('.card-action') as NodeListOf<HTMLButtonElement>;
      console.log('🎴 Found action buttons:', actionButtons.length);

      actionButtons.forEach((button, index) => {
        const action = actions[index];
        console.log(`🎴 Setting up listener for button ${index}:`, {
          hasAction: !!action,
          hasActionProperty: !!action?.action,
          action: action
        });

        if (action && action.action) {
          console.log('🎴 Adding click listener to button:', button);
          button.addEventListener('click', async () => {
            console.log('🔘 Card action button clicked:', action.label);
            console.log('   Sending action:', action.action);

            // Apply visual feedback
            button.disabled = true;
            button.classList.add('button-transitioning', 'button-clicked');

            // Find vanna-chat component and send message
            const vannaChat = document.querySelector('vanna-chat') as any;

            if (vannaChat && typeof vannaChat.sendMessage === 'function') {
              try {
                const success = await vannaChat.sendMessage(action.action);

                if (success) {
                  console.log('✅ Card action sent successfully');
                  // Keep button disabled after successful action
                } else {
                  console.error('❌ Failed to send card action');
                  // Re-enable button on failure
                  button.disabled = false;
                  button.classList.remove('button-transitioning', 'button-clicked');
                }
              } catch (error) {
                console.error('❌ Error sending card action:', error);
                // Re-enable button on error
                button.disabled = false;
                button.classList.remove('button-transitioning', 'button-clicked');
              }
            } else {
              console.warn('⚠️ vanna-chat component not found or sendMessage not available');
              button.disabled = false;
              button.classList.remove('button-transitioning', 'button-clicked');
            }
          });
        }
      });
    }

    return card;
  }

  update(element: HTMLElement, component: RichComponent, updates?: Record<string, any>): void {
    if (!updates) return super.update(element, component);

    // Optimized updates for common properties
    if (updates.title) {
      const titleEl = element.querySelector('.card-title');
      if (titleEl) titleEl.textContent = updates.title;
    }

    if (updates.content) {
      const contentEl = element.querySelector('.card-content');
      if (contentEl) contentEl.innerHTML = updates.content;
    }

    if (updates.status) {
      const statusEl = element.querySelector('.card-status');
      if (statusEl) {
        statusEl.className = `card-status status-${updates.status}`;
        statusEl.textContent = updates.status;
      }
    }

    // For complex updates, fall back to full re-render
    if (updates.actions || updates.collapsible) {
      super.update(element, component);
    }
  }
}

// Task list component renderer
export class TaskListComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-task-list';
    container.dataset.componentId = component.id;

    const { title, tasks = [], show_progress, show_timestamps } = component.data;

    const completedTasks = tasks.filter((task: any) => task.status === 'completed').length;
    const progress = tasks.length > 0 ? (completedTasks / tasks.length) * 100 : 0;

    container.innerHTML = `
      <div class="task-list-header">
        <h3 class="task-list-title">${title}</h3>
        ${show_progress ? `
          <div class="task-list-progress">
            <span class="progress-text">${completedTasks}/${tasks.length} completed</span>
            <div class="progress-bar">
              <div class="progress-fill" style="width: ${progress}%"></div>
            </div>
          </div>
        ` : ''}
      </div>
      <div class="task-list-items">
        ${tasks.map((task: any) => this.renderTask(task, show_timestamps)).join('')}
      </div>
    `;


    return container;
  }

  private renderTask(task: any, showTimestamps: boolean): string {
    const statusIcon = this.getStatusIcon(task.status);
    const progressBar = task.progress !== null && task.progress !== undefined ? `
      <div class="task-progress">
        <div class="task-progress-bar">
          <div class="task-progress-fill" style="width: ${task.progress * 100}%"></div>
        </div>
        <span class="task-progress-text">${Math.round(task.progress * 100)}%</span>
      </div>
    ` : '';

    return `
      <div class="task-item status-${task.status}" data-task-id="${task.id}">
        <div class="task-icon">${statusIcon}</div>
        <div class="task-content">
          <div class="task-title">${task.title}</div>
          ${task.description ? `<div class="task-description">${task.description}</div>` : ''}
          ${progressBar}
          ${showTimestamps && task.created_at ? `
            <div class="task-timestamp">
              Created: ${new Date(task.created_at).toLocaleString()}
            </div>
          ` : ''}
        </div>
      </div>
    `;
  }

  private getStatusIcon(status: string): string {
    switch (status) {
      case 'completed': return '✅';
      case 'running': return '🔄';
      case 'failed': return '❌';
      default: return '⭕';
    }
  }
}

// Progress bar component renderer
export class ProgressBarComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-progress-bar';
    container.dataset.componentId = component.id;

    const { value, label, show_percentage, status, animated } = component.data;
    const percentage = Math.round(value * 100);

    container.innerHTML = `
      <div class="progress-header">
        ${label ? `<span class="progress-label">${label}</span>` : ''}
        ${show_percentage ? `<span class="progress-percentage">${percentage}%</span>` : ''}
      </div>
      <div class="progress-track">
        <div class="progress-fill ${animated ? 'animated' : ''} ${status ? `status-${status}` : ''}"
             style="width: ${percentage}%"></div>
      </div>
    `;


    return container;
  }

  update(element: HTMLElement, component: RichComponent, updates?: Record<string, any>): void {
    if (!updates) return super.update(element, component);

    if (updates.value !== undefined) {
      const fill = element.querySelector('.progress-fill') as HTMLElement;
      const percentage = Math.round(updates.value * 100);

      if (fill) {
        fill.style.width = `${percentage}%`;
      }

      const percentageEl = element.querySelector('.progress-percentage');
      if (percentageEl) {
        percentageEl.textContent = `${percentage}%`;
      }
    }

    if (updates.label) {
      const labelEl = element.querySelector('.progress-label');
      if (labelEl) labelEl.textContent = updates.label;
    }

    if (updates.status) {
      const fill = element.querySelector('.progress-fill') as HTMLElement;
      if (fill) {
        fill.className = fill.className.replace(/status-\w+/, `status-${updates.status}`);
      }
    }
  }
}

// Notification component renderer
export class NotificationComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-notification';
    container.dataset.componentId = component.id;

    const { message, title, level = 'info', icon, dismissible, auto_dismiss, actions = [] } = component.data;

    const levelIcon = icon || this.getLevelIcon(level);
    const dismissButton = dismissible ? `
      <button class="notification-dismiss" onclick="this.parentElement.remove()">×</button>
    ` : '';

    container.innerHTML = `
      <div class="notification-content level-${level}">
        ${levelIcon ? `<span class="notification-icon">${levelIcon}</span>` : ''}
        <div class="notification-body">
          ${title ? `<div class="notification-title">${title}</div>` : ''}
          <div class="notification-message">${message}</div>
        </div>
        ${actions.length > 0 ? `
          <div class="notification-actions">
            ${actions.map((action: any) => `
              <button class="notification-action ${action.variant || 'secondary'}" data-action="${action.action}">
                ${action.label}
              </button>
            `).join('')}
          </div>
        ` : ''}
        ${dismissButton}
      </div>
    `;

    // Auto-dismiss functionality
    if (auto_dismiss && component.data.auto_dismiss_delay) {
      setTimeout(() => {
        if (container.parentElement) {
          container.remove();
        }
      }, component.data.auto_dismiss_delay);
    }


    return container;
  }

  private getLevelIcon(level: string): string {
    switch (level) {
      case 'success': return '✅';
      case 'warning': return '⚠️';
      case 'error': return '❌';
      case 'info':
      default: return 'ℹ️';
    }
  }
}

// Status indicator component renderer
export class StatusIndicatorComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-status-indicator';
    container.dataset.componentId = component.id;

    const { status, message, icon, pulse } = component.data;

    const statusIcon = icon || this.getStatusIcon(status);
    const pulseClass = pulse ? 'pulse' : '';

    container.innerHTML = `
      <div class="status-indicator-content status-${status} ${pulseClass}">
        <span class="status-icon">${statusIcon}</span>
        <span class="status-message">${message}</span>
      </div>
    `;


    return container;
  }

  private getStatusIcon(status: string): string {
    switch (status) {
      case 'loading': return '🔄';
      case 'success': return '✅';
      case 'warning': return '⚠️';
      case 'error': return '❌';
      default: return 'ℹ️';
    }
  }

  update(element: HTMLElement, component: RichComponent, updates?: Record<string, any>): void {
    if (!updates) return super.update(element, component);

    const content = element.querySelector('.status-indicator-content');
    if (content && updates.status) {
      content.className = content.className.replace(/status-\w+/, `status-${updates.status}`);
    }

    if (updates.pulse !== undefined) {
      const content = element.querySelector('.status-indicator-content');
      if (content) {
        if (updates.pulse) {
          content.classList.add('pulse');
        } else {
          content.classList.remove('pulse');
        }
      }
    }

    if (updates.message) {
      const messageEl = element.querySelector('.status-message');
      if (messageEl) {
        messageEl.textContent = updates.message;
      }
    }
  }
}

// DataFrame component renderer
export class DataFrameComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-dataframe';
    container.dataset.componentId = component.id;

    const {
      data = [],
      columns = [],
      title,
      description,
      row_count = 0,
      column_count = 0,
      max_rows_displayed = 100,
      searchable = true,
      sortable = true,
      filterable = true,
      exportable = true,
      striped = true,
      bordered = true,
      compact = false,
      column_types = {}
    } = component.data;

    // Limit displayed rows
    const displayedData = data.slice(0, max_rows_displayed);
    const hasMoreRows = data.length > max_rows_displayed;

    let headerHTML = '';
    if (title || description) {
      headerHTML = `
        <div class="dataframe-header">
          ${title ? `<h3 class="dataframe-title">${title}</h3>` : ''}
          ${description ? `<p class="dataframe-description">${description}</p>` : ''}
          <div class="dataframe-meta">
            <span class="row-count">${row_count} rows</span>
            <span class="column-count">${column_count} columns</span>
          </div>
        </div>
      `;
    }

    let actionsHTML = '';
    if (searchable || exportable || filterable) {
      actionsHTML = `
        <div class="dataframe-actions">
          ${searchable ? `
            <div class="dataframe-search">
              <input type="text" placeholder="Search..." class="search-input">
            </div>
          ` : ''}
          ${exportable ? `
            <button class="export-btn" title="Export to CSV">📥 Export</button>
          ` : ''}
        </div>
      `;
    }

    let tableHTML = '';
    if (columns.length > 0 && displayedData.length > 0) {
      const tableClasses = [
        'dataframe-table',
        striped ? 'striped' : '',
        bordered ? 'bordered' : '',
        compact ? 'compact' : ''
      ].filter(Boolean).join(' ');

      tableHTML = `
        <div class="dataframe-table-container">
          <table class="${tableClasses}">
            <thead>
              <tr>
                ${columns.map((col: string) => `
                  <th class="${sortable ? 'sortable' : ''}" data-column="${col}">
                    ${col}
                    ${sortable ? '<span class="sort-indicator"></span>' : ''}
                  </th>
                `).join('')}
              </tr>
            </thead>
            <tbody>
              ${displayedData.map((row: any) => `
                <tr>
                  ${columns.map((col: string) => {
                    const value = row[col];
                    const columnType = column_types[col] || 'string';
                    const formattedValue = this.formatCellValue(value, columnType);
                    return `<td class="cell-${columnType}">${formattedValue}</td>`;
                  }).join('')}
                </tr>
              `).join('')}
            </tbody>
          </table>
          ${hasMoreRows ? `
            <div class="dataframe-truncated">
              <em>Showing ${max_rows_displayed} of ${row_count} rows</em>
            </div>
          ` : ''}
        </div>
      `;
    } else {
      tableHTML = `
        <div class="dataframe-empty">
          <p>No data to display</p>
        </div>
      `;
    }

    container.innerHTML = `
      ${headerHTML}
      ${actionsHTML}
      ${tableHTML}
    `;


    // Add event listeners
    this.attachEventListeners(container, displayedData, columns);

    return container;
  }

  private formatCellValue(value: any, columnType: string): string {
    if (value === null || value === undefined) {
      return '<em class="null-value">NULL</em>';
    }

    switch (columnType) {
      case 'number':
        return typeof value === 'number' ? value.toLocaleString() : String(value);
      case 'date':
        try {
          return new Date(value).toLocaleDateString();
        } catch {
          return String(value);
        }
      case 'boolean':
        return value ? '✓' : '✗';
      default:
        return this.escapeHtml(String(value));
    }
  }

  private escapeHtml(text: string): string {
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
  }

  private attachEventListeners(container: HTMLElement, data: any[], columns: string[]): void {
    // Search functionality
    const searchInput = container.querySelector('.search-input') as HTMLInputElement;
    if (searchInput) {
      searchInput.addEventListener('input', (e) => {
        const searchTerm = (e.target as HTMLInputElement).value.toLowerCase();
        this.filterTable(container, data, columns, searchTerm);
      });
    }

    // Export functionality
    const exportBtn = container.querySelector('.export-btn') as HTMLButtonElement;
    if (exportBtn) {
      exportBtn.addEventListener('click', () => {
        this.exportToCSV(data, columns);
      });
    }

    // Sort functionality
    const sortableHeaders = container.querySelectorAll('th.sortable');
    sortableHeaders.forEach(header => {
      header.addEventListener('click', (e) => {
        const column = (e.currentTarget as HTMLElement).dataset.column;
        if (column) {
          this.sortTable(container, data, columns, column);
        }
      });
    });
  }

  private filterTable(container: HTMLElement, data: any[], columns: string[], searchTerm: string): void {
    const tbody = container.querySelector('tbody');
    if (!tbody) return;

    const filteredData = data.filter(row => {
      return columns.some(col => {
        const value = String(row[col] || '').toLowerCase();
        return value.includes(searchTerm);
      });
    });

    tbody.innerHTML = filteredData.map(row => `
      <tr>
        ${columns.map(col => {
          const value = row[col];
          const formattedValue = this.formatCellValue(value, 'string');
          return `<td>${formattedValue}</td>`;
        }).join('')}
      </tr>
    `).join('');
  }

  private sortTable(container: HTMLElement, data: any[], columns: string[], column: string): void {
    const tbody = container.querySelector('tbody');
    const header = container.querySelector(`th[data-column="${column}"]`) as HTMLElement;
    if (!tbody || !header) return;

    // Determine sort direction
    const currentSort = header.dataset.sortDirection || 'none';
    const newSort = currentSort === 'asc' ? 'desc' : 'asc';

    // Clear all sort indicators
    container.querySelectorAll('th[data-sort-direction]').forEach(h => {
      (h as HTMLElement).removeAttribute('data-sort-direction');
      const indicator = h.querySelector('.sort-indicator');
      if (indicator) indicator.textContent = '';
    });

    // Set new sort direction
    header.dataset.sortDirection = newSort;
    const indicator = header.querySelector('.sort-indicator');
    if (indicator) {
      indicator.textContent = newSort === 'asc' ? '↑' : '↓';
    }

    // Sort data
    const sortedData = [...data].sort((a, b) => {
      const aVal = a[column];
      const bVal = b[column];

      if (aVal === null || aVal === undefined) return 1;
      if (bVal === null || bVal === undefined) return -1;

      if (typeof aVal === 'number' && typeof bVal === 'number') {
        return newSort === 'asc' ? aVal - bVal : bVal - aVal;
      }

      const aStr = String(aVal).toLowerCase();
      const bStr = String(bVal).toLowerCase();
      const comparison = aStr.localeCompare(bStr);
      return newSort === 'asc' ? comparison : -comparison;
    });

    // Update table
    tbody.innerHTML = sortedData.map(row => `
      <tr>
        ${columns.map(col => {
          const value = row[col];
          const formattedValue = this.formatCellValue(value, 'string');
          return `<td>${formattedValue}</td>`;
        }).join('')}
      </tr>
    `).join('');
  }

  private exportToCSV(data: any[], columns: string[]): void {
    const csvContent = [
      columns.join(','),
      ...data.map(row =>
        columns.map(col => {
          const value = row[col];
          const strValue = value === null || value === undefined ? '' : String(value);
          // Escape quotes and wrap in quotes if contains comma, quote, or newline
          if (strValue.includes(',') || strValue.includes('"') || strValue.includes('\n')) {
            return `"${strValue.replace(/"/g, '""')}"`;
          }
          return strValue;
        }).join(',')
      )
    ].join('\n');

    const blob = new Blob([csvContent], { type: 'text/csv;charset=utf-8;' });
    const link = document.createElement('a');
    const url = URL.createObjectURL(blob);
    link.setAttribute('href', url);
    link.setAttribute('download', 'data.csv');
    link.style.visibility = 'hidden';
    document.body.appendChild(link);
    link.click();
    document.body.removeChild(link);
  }
}

// Text component renderer
export class TextComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-text';
    container.dataset.componentId = component.id;

    const {
      content,
      markdown = false,
      code_language,
      font_size,
      font_weight,
      text_align
    } = component.data;

    // Apply text styling
    let textStyle = '';
    if (font_size) textStyle += `font-size: ${font_size}; `;
    if (font_weight) textStyle += `font-weight: ${font_weight}; `;
    if (text_align) textStyle += `text-align: ${text_align}; `;

    if (code_language) {
      // Code block
      container.innerHTML = `
        <pre class="text-code" style="${textStyle}"><code class="language-${code_language}">${this.escapeHtml(content)}</code></pre>
      `;
    } else if (markdown) {
      // Markdown text (simple implementation)
      container.innerHTML = `
        <div class="text-markdown" style="${textStyle}">${this.renderMarkdown(content)}</div>
      `;
    } else {
      // Plain text
      container.innerHTML = `
        <div class="text-content" style="${textStyle}">${this.escapeHtml(content)}</div>
      `;
    }


    return container;
  }

  private escapeHtml(text: string): string {
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
  }

  private renderMarkdown(text: string): string {
    // Simple markdown rendering - just basic formatting
    return text
      .replace(/^## (.*$)/gm, '<h2>$1</h2>')
      .replace(/^# (.*$)/gm, '<h1>$1</h1>')
      .replace(/\*\*(.*?)\*\*/g, '<strong>$1</strong>')
      .replace(/\*(.*?)\*/g, '<em>$1</em>')
      .replace(/^- (.*$)/gm, '<li>$1</li>')
      .replace(/(<li>.*<\/li>)/s, '<ul>$1</ul>')
      .replace(/\n\n/g, '</p><p>')
      .replace(/^(?!<[h|u|l])(.+)$/gm, '<p>$1</p>');
  }
}

// Primitive Component Renderers (Domain-Agnostic)

// Status card component renderer
export class StatusCardComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-status-card';
    container.dataset.componentId = component.id;

    const { title, status, description, icon, actions = [], collapsible, collapsed, metadata = {} } = component.data;

    const statusIcon = icon || this.getStatusIcon(status);
    const hasMetadata = Object.keys(metadata).length > 0;

    container.innerHTML = `
      <div class="status-card-header ${collapsible ? 'collapsible' : ''}">
        <div class="status-card-icon">${statusIcon}</div>
        <div class="status-card-title-section">
          <h3 class="status-card-title">${title}</h3>
          <span class="status-card-badge status-${status}">${status}</span>
        </div>
        ${collapsible ? `<button class="status-card-toggle">${collapsed ? '▶' : '▼'}</button>` : ''}
      </div>
      ${description ? `
        <div class="status-card-content ${collapsed ? 'collapsed' : ''}">
          ${description}
        </div>
      ` : ''}
      ${hasMetadata ? `
        <details class="status-card-metadata">
          <summary class="status-card-metadata-summary">Parameters</summary>
          <div class="status-card-metadata-content">
            ${this.renderMetadataTable(metadata)}
          </div>
        </details>
      ` : ''}
      ${actions.length > 0 ? `
        <div class="status-card-actions">
          ${actions.map((action: any) => `
            <button class="status-card-action ${action.variant || 'secondary'}" data-action="${action.action}">
              ${action.label}
            </button>
          `).join('')}
        </div>
      ` : ''}
    `;

    // Add collapsible functionality
    if (collapsible) {
      const toggle = container.querySelector('.status-card-toggle') as HTMLButtonElement;
      const content = container.querySelector('.status-card-content') as HTMLElement;

      toggle?.addEventListener('click', () => {
        if (content) {
          content.classList.toggle('collapsed');
          toggle.textContent = content.classList.contains('collapsed') ? '▶' : '▼';
        }
      });
    }

    return container;
  }

  private renderMetadataTable(metadata: Record<string, any>): string {
    const rows = Object.entries(metadata).map(([key, value]) => {
      const formattedValue = this.formatMetadataValue(value);
      return `
        <tr>
          <td class="metadata-key">${this.escapeHtml(key)}</td>
          <td class="metadata-value">${formattedValue}</td>
        </tr>
      `;
    }).join('');

    return `
      <table class="metadata-table">
        <thead>
          <tr>
            <th>Parameter</th>
            <th>Value</th>
          </tr>
        </thead>
        <tbody>
          ${rows}
        </tbody>
      </table>
    `;
  }

  private formatMetadataValue(value: any): string {
    if (value === null) {
      return '<span class="metadata-null">null</span>';
    }
    if (value === undefined) {
      return '<span class="metadata-undefined">undefined</span>';
    }
    if (typeof value === 'boolean') {
      return `<span class="metadata-boolean">${value}</span>`;
    }
    if (typeof value === 'number') {
      return `<span class="metadata-number">${value}</span>`;
    }
    if (typeof value === 'string') {
      return `<span class="metadata-string">${this.escapeHtml(value)}</span>`;
    }
    if (Array.isArray(value) || typeof value === 'object') {
      return `<pre class="metadata-json">${JSON.stringify(value, null, 2)}</pre>`;
    }
    return this.escapeHtml(String(value));
  }

  private escapeHtml(text: string): string {
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
  }

  private getStatusIcon(status: string): string {
    switch (status) {
      case 'pending': return '⏳';
      case 'running': return '⚙️';
      case 'completed': return '✅';
      case 'success': return '✅';
      case 'failed': return '❌';
      case 'error': return '❌';
      case 'warning': return '⚠️';
      default: return 'ℹ️';
    }
  }
}

// Progress display component renderer
export class ProgressDisplayComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-progress-display';
    container.dataset.componentId = component.id;

    const { label, value, description, status, show_percentage, animated, indeterminate } = component.data;
    const percentage = Math.round(value * 100);

    container.innerHTML = `
      <div class="progress-display-container">
        <div class="progress-display-header">
          <span class="progress-display-label">${label}</span>
          ${show_percentage && !indeterminate ? `<span class="progress-display-percentage">${percentage}%</span>` : ''}
        </div>
        <div class="progress-display-track">
          <div class="progress-display-fill ${animated ? 'animated' : ''} ${status ? `status-${status}` : ''} ${indeterminate ? 'indeterminate' : ''}"
               style="width: ${indeterminate ? '100' : percentage}%"></div>
        </div>
        ${description ? `<div class="progress-display-description">${description}</div>` : ''}
      </div>
    `;

    return container;
  }
}

// Log viewer component renderer
export class LogViewerComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-log-viewer';
    container.dataset.componentId = component.id;

    const { title, entries = [], searchable, show_timestamps, auto_scroll } = component.data;

    container.innerHTML = `
      <div class="log-viewer-container">
        <div class="log-viewer-header">
          <h3 class="log-viewer-title">${title}</h3>
          ${searchable ? `
            <div class="log-viewer-search">
              <input type="text" placeholder="Search logs..." class="log-search-input">
            </div>
          ` : ''}
        </div>
        <div class="log-viewer-content ${auto_scroll ? 'auto-scroll' : ''}">
          ${entries.map((entry: any) => `
            <div class="log-entry log-${entry.level}">
              ${show_timestamps ? `<span class="log-timestamp">${new Date(entry.timestamp).toLocaleTimeString()}</span>` : ''}
              <span class="log-level">[${entry.level.toUpperCase()}]</span>
              <span class="log-message">${entry.message}</span>
            </div>
          `).join('')}
        </div>
      </div>
    `;

    // Auto-scroll to bottom if enabled
    if (auto_scroll) {
      const content = container.querySelector('.log-viewer-content');
      if (content) {
        content.scrollTop = content.scrollHeight;
      }
    }

    return container;
  }
}

// Badge component renderer
export class BadgeComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('span');
    container.className = `rich-component rich-badge badge-${component.data.variant} badge-${component.data.size}`;
    container.dataset.componentId = component.id;

    const { text, icon } = component.data;

    container.innerHTML = `
      ${icon ? `<span class="badge-icon">${icon}</span>` : ''}
      <span class="badge-text">${text}</span>
    `;

    return container;
  }
}

// Icon text component renderer
export class IconTextComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = `rich-component rich-icon-text icon-text-${component.data.variant} icon-text-${component.data.size} icon-text-${component.data.alignment}`;
    container.dataset.componentId = component.id;

    const { icon, text } = component.data;

    container.innerHTML = `
      <span class="icon-text-icon">${icon}</span>
      <span class="icon-text-text">${text}</span>
    `;

    return container;
  }
}

// Button component renderer
export class ButtonComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const button = document.createElement('button');
    button.className = `rich-component rich-button button-${component.data.variant} button-${component.data.size}`;
    button.dataset.componentId = component.id;

    const { label, action, disabled, icon, icon_position, full_width, loading } = component.data;

    if (disabled || loading) {
      button.disabled = true;
    }

    if (full_width) {
      button.classList.add('button-full-width');
    }

    if (loading) {
      button.classList.add('button-loading');
    }

    // Build button content
    let buttonContent = '';
    if (loading) {
      buttonContent = `<span class="button-spinner">⏳</span><span class="button-label">${label}</span>`;
    } else if (icon) {
      if (icon_position === 'right') {
        buttonContent = `<span class="button-label">${label}</span><span class="button-icon">${icon}</span>`;
      } else {
        buttonContent = `<span class="button-icon">${icon}</span><span class="button-label">${label}</span>`;
      }
    } else {
      buttonContent = `<span class="button-label">${label}</span>`;
    }

    button.innerHTML = buttonContent;

    // Add click handler
    if (action && !disabled && !loading) {
      button.addEventListener('click', async () => {
        console.log('🔘 Button clicked:', label);
        console.log('   Sending action:', action);

        // Apply visual feedback immediately
        button.disabled = true;
        button.classList.add('button-transitioning', 'button-clicked');

        // Find vanna-chat component and send message with button action
        const vannaChat = document.querySelector('vanna-chat') as any;
        console.log('   Found vanna-chat:', !!vannaChat);

        if (vannaChat && typeof vannaChat.sendMessage === 'function') {
          console.log('   Calling sendMessage...');
          try {
            const success = await vannaChat.sendMessage(action);
            if (success) {
              console.log('   ✓ Message sent successfully');
            } else {
              console.log('   ✗ Message failed, restoring button state');
              // Restore button state if it wasn't originally disabled
              if (!disabled) {
                button.disabled = false;
              }
              button.classList.remove('button-transitioning', 'button-clicked');
            }
          } catch (error) {
            console.error('   ✗ Message failed with error:', error);
            // Restore button state if it wasn't originally disabled
            if (!disabled) {
              button.disabled = false;
            }
            button.classList.remove('button-transitioning', 'button-clicked');
          }
        } else {
          console.error('   ✗ vanna-chat not found or sendMessage not available');
          // Restore button state if it wasn't originally disabled
          if (!disabled) {
            button.disabled = false;
          }
          button.classList.remove('button-transitioning', 'button-clicked');
        }
      });
    }

    return button;
  }

  update(element: HTMLElement, component: RichComponent, updates?: Record<string, any>): void {
    if (!updates) return super.update(element, component);

    const button = element as HTMLButtonElement;

    if (updates.disabled !== undefined) {
      button.disabled = updates.disabled;
    }

    if (updates.loading !== undefined) {
      button.disabled = updates.loading;
      if (updates.loading) {
        button.classList.add('button-loading');
      } else {
        button.classList.remove('button-loading');
      }
    }

    if (updates.label || updates.icon || updates.icon_position) {
      // Re-render content
      super.update(element, component);
    }
  }
}

// Button group component renderer
export class ButtonGroupComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = `rich-component rich-button-group button-group-${component.data.orientation} button-group-spacing-${component.data.spacing} button-group-align-${component.data.align}`;
    container.dataset.componentId = component.id;

    const { buttons = [], full_width } = component.data;

    if (full_width) {
      container.classList.add('button-group-full-width');
    }

    // Render each button
    buttons.forEach((buttonConfig: any, index: number) => {
      const button = document.createElement('button');
      button.className = `rich-button button-${buttonConfig.variant || 'secondary'} button-${buttonConfig.size || 'medium'}`;
      button.dataset.buttonIndex = String(index);

      // Store original disabled state
      if (buttonConfig.disabled) {
        button.disabled = true;
        button.dataset.originallyDisabled = 'true';
      } else {
        button.dataset.originallyDisabled = 'false';
      }

      // Build button content
      let buttonContent = '';
      if (buttonConfig.icon) {
        if (buttonConfig.icon_position === 'right') {
          buttonContent = `<span class="button-label">${buttonConfig.label}</span><span class="button-icon">${buttonConfig.icon}</span>`;
        } else {
          buttonContent = `<span class="button-icon">${buttonConfig.icon}</span><span class="button-label">${buttonConfig.label}</span>`;
        }
      } else {
        buttonContent = `<span class="button-label">${buttonConfig.label}</span>`;
      }

      button.innerHTML = buttonContent;

      // Add click handler with enhanced functionality
      if (buttonConfig.action && !buttonConfig.disabled) {
        button.addEventListener('click', async () => {
          console.log('🔘 Button Group button clicked:', buttonConfig.label);
          console.log('   Button index:', index);
          console.log('   Sending action:', buttonConfig.action);

          // Immediately apply visual changes to all buttons in the group
          this.applyButtonGroupClickState(container, index);

          // Find vanna-chat component and send message with button action
          const vannaChat = document.querySelector('vanna-chat') as any;
          console.log('   Found vanna-chat:', !!vannaChat);

          if (vannaChat && typeof vannaChat.sendMessage === 'function') {
            console.log('   Calling sendMessage...');
            try {
              const success = await vannaChat.sendMessage(buttonConfig.action);
              if (success) {
                console.log('   ✓ Message sent successfully');
              } else {
                console.log('   ✗ Message failed, restoring button state');
                this.restoreButtonGroupState(container);
              }
            } catch (error) {
              console.error('   ✗ Message failed with error:', error);
              this.restoreButtonGroupState(container);
            }
          } else {
            console.error('   ✗ vanna-chat not found or sendMessage not available');
            this.restoreButtonGroupState(container);
          }
        });
      }

      container.appendChild(button);
    });

    return container;
  }

  private applyButtonGroupClickState(container: HTMLElement, clickedIndex: number): void {
    const buttons = container.querySelectorAll('button') as NodeListOf<HTMLButtonElement>;
    
    buttons.forEach((button, index) => {
      // Disable all buttons
      button.disabled = true;
      
      // Add transition class for animation
      button.classList.add('button-transitioning');
      
      if (index === clickedIndex) {
        // Highlight the clicked button
        button.classList.add('button-clicked', 'button-highlighted');
      } else {
        // Gray out other buttons
        button.classList.add('button-grayed-out');
      }
    });
  }

  private restoreButtonGroupState(container: HTMLElement): void {
    const buttons = container.querySelectorAll('button') as NodeListOf<HTMLButtonElement>;
    
    buttons.forEach((button) => {
      // Re-enable buttons (unless they were originally disabled)
      const originallyDisabled = button.dataset.originallyDisabled === 'true';
      if (!originallyDisabled) {
        button.disabled = false;
      }
      
      // Remove all state classes
      button.classList.remove(
        'button-clicked', 
        'button-highlighted', 
        'button-grayed-out',
        'button-transitioning'
      );
    });
  }
}

// Chart component renderer (for Plotly charts)
export class ChartComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-chart';
    container.dataset.componentId = component.id;

    // The ChartComponent.data field contains the Plotly figure directly
    // Structure: component.data = { data: [...traces...], layout: {...}, title: "...", config: {...} }
    const { data: plotlyData, layout, title, config = {} } = component.data;

    console.log('ChartComponentRenderer: Received component.data:', component.data);
    console.log('ChartComponentRenderer: plotlyData:', plotlyData);
    console.log('ChartComponentRenderer: layout:', layout);

    // Check if we have a valid Plotly figure structure
    if (plotlyData && Array.isArray(plotlyData) && layout) {
      // Create plotly-chart web component
      const chartElement = document.createElement('plotly-chart') as any;

      // Set theme to match current theme
      const vannaChat = document.querySelector('vanna-chat');
      if (vannaChat) {
        chartElement.theme = vannaChat.getAttribute('theme') || 'dark';
      }

      // Wrap in container with optional title
      if (title) {
        container.innerHTML = `
          <div class="chart-header">
            <h3 class="chart-title">${title}</h3>
          </div>
          <div class="chart-content"></div>
        `;
        container.querySelector('.chart-content')?.appendChild(chartElement);
      } else {
        container.appendChild(chartElement);
      }

      // Set data AFTER the element is in the DOM
      // This ensures the web component is fully initialized
      requestAnimationFrame(() => {
        chartElement.data = plotlyData; // Plotly traces (array)
        chartElement.layout = layout; // Plotly layout (object)
        chartElement.config = config;

        console.log('ChartComponentRenderer: Set properties after DOM attachment');
        console.log('ChartComponentRenderer: chartElement.data:', chartElement.data);
        console.log('ChartComponentRenderer: chartElement.layout:', chartElement.layout);
      });
    } else {
      // Fallback for invalid chart data
      container.innerHTML = `
        <div class="chart-error">
          <p>Invalid chart data format</p>
          <pre>${JSON.stringify(component.data, null, 2).substring(0, 200)}...</pre>
        </div>
      `;
    }

    return container;
  }
}

// Artifact component renderer
export class ArtifactComponentRenderer extends BaseComponentRenderer {
  private defaultPrevented = false;

  render(component: RichComponent): HTMLElement {
    console.log('🔧 ArtifactComponentRenderer.render called with:', component);

    const container = document.createElement('div');
    container.className = 'rich-component rich-artifact';
    container.dataset.componentId = component.id;
    container.dataset.artifactId = component.data.artifact_id;

    const {
      content,
      artifact_type,
      title,
      description,
      editable,
      fullscreen_capable,
      external_renderable
    } = component.data;

    // Create artifact preview and controls
    container.innerHTML = `
      <div class="artifact-header">
        <div class="artifact-meta">
          <h3 class="artifact-title">${title || 'Artifact'}</h3>
          ${description ? `<p class="artifact-description">${description}</p>` : ''}
          <span class="artifact-type-badge">${artifact_type}</span>
        </div>
        <div class="artifact-controls">
          ${editable ? '<button class="artifact-btn edit-btn" title="Edit">✏️</button>' : ''}
          ${fullscreen_capable ? '<button class="artifact-btn fullscreen-btn" title="Fullscreen">⛶</button>' : ''}
          ${external_renderable ? '<button class="artifact-btn external-btn" title="Open External">🔗</button>' : ''}
        </div>
      </div>
      <div class="artifact-preview">
        <iframe class="artifact-iframe" sandbox="allow-scripts allow-same-origin" srcdoc="${this.escapeHtml(content)}"></iframe>
      </div>
    `;

    // Attach event listeners
    this.attachEventListeners(container, component);

    // Fire artifact-opened event for creation
    const shouldRenderInChat = this.fireArtifactOpenedEvent(component, 'created', container);

    // If default was prevented, show a placeholder instead
    if (!shouldRenderInChat) {
      container.innerHTML = `
        <div class="artifact-placeholder">
          <div class="placeholder-content">
            <span class="placeholder-icon">🎨</span>
            <div class="placeholder-text">
              <strong>${title || 'Artifact'}</strong> opened externally
              <div class="placeholder-type">${artifact_type}</div>
            </div>
            <button class="placeholder-reopen" title="Reopen">↗</button>
          </div>
        </div>
      `;

      // Add reopen functionality
      const reopenBtn = container.querySelector('.placeholder-reopen') as HTMLButtonElement;
      if (reopenBtn) {
        reopenBtn.addEventListener('click', () => {
          this.fireArtifactOpenedEvent(component, 'user-action', container);
        });
      }
    }

    return container;
  }

  private attachEventListeners(container: HTMLElement, component: RichComponent): void {
    // External button click
    const externalBtn = container.querySelector('.external-btn') as HTMLButtonElement;
    if (externalBtn) {
      externalBtn.addEventListener('click', () => {
        this.fireArtifactOpenedEvent(component, 'user-action', container);
      });
    }

    // Fullscreen button click
    const fullscreenBtn = container.querySelector('.fullscreen-btn') as HTMLButtonElement;
    if (fullscreenBtn) {
      fullscreenBtn.addEventListener('click', () => {
        this.openFullscreen(component);
      });
    }

    // Edit button click (placeholder for future implementation)
    const editBtn = container.querySelector('.edit-btn') as HTMLButtonElement;
    if (editBtn) {
      editBtn.addEventListener('click', () => {
        this.openEditor(component);
      });
    }
  }

  private fireArtifactOpenedEvent(component: RichComponent, trigger: 'created' | 'user-action', container: HTMLElement): boolean {
    console.log('🎯 fireArtifactOpenedEvent called:', { trigger, artifactId: component.data.artifact_id });

    this.defaultPrevented = false;

    const eventDetail: ArtifactOpenedEventDetail = {
      artifactId: component.data.artifact_id,
      content: component.data.content,
      type: component.data.artifact_type,
      title: component.data.title,
      description: component.data.description,
      trigger,
      preventDefault: () => {
        console.log('🛑 preventDefault called!');
        this.defaultPrevented = true;
      },
      getStandaloneHTML: () => this.generateStandaloneHTML(component),
      timestamp: new Date().toISOString()
    };

    const event = new CustomEvent('artifact-opened', {
      detail: eventDetail,
      bubbles: true,
      cancelable: true
    });

    console.log('📡 Dispatching artifact-opened event:', event);

    // Fire the event from the container element (should bubble up to vanna-chat)
    container.dispatchEvent(event);

    // Also dispatch directly on the vanna-chat element as backup
    const vannaChat = container.closest('vanna-chat');
    if (vannaChat) {
      console.log('📡 Also dispatching on vanna-chat element');
      vannaChat.dispatchEvent(new CustomEvent('artifact-opened', {
        detail: eventDetail,
        bubbles: true,
        cancelable: true
      }));
    }

    console.log('📨 Event dispatched. defaultPrevented:', this.defaultPrevented);

    // Handle default behavior if not prevented and user triggered
    if (!this.defaultPrevented && trigger === 'user-action') {
      this.handleDefaultAction(component);
    }

    // Return whether we should render in chat (true if default not prevented)
    return !this.defaultPrevented;
  }

  private generateStandaloneHTML(component: RichComponent): string {
    const { content, title, dependencies = [] } = component.data;

    let dependenciesHTML = '';

    // Add common CDN links for dependencies
    if (dependencies.includes('d3')) {
      dependenciesHTML += '<script src="https://d3js.org/d3.v7.min.js"></script>\n';
    }
    if (dependencies.includes('plotly')) {
      dependenciesHTML += '<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>\n';
    }
    if (dependencies.includes('three') || dependencies.includes('threejs')) {
      dependenciesHTML += '<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>\n';
    }

    return `<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>${title || 'Artifact'}</title>
    ${dependenciesHTML}
    <style>
        body {
            margin: 0;
            padding: 20px;
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
        }
        .artifact-container {
            width: 100%;
            min-height: 100vh;
        }
    </style>
</head>
<body>
    <div class="artifact-container">
        ${content}
    </div>
</body>
</html>`;
  }

  private handleDefaultAction(component: RichComponent): void {
    // Default action: open in new window
    const newWindow = window.open('', '_blank', 'width=800,height=600');
    if (newWindow) {
      newWindow.document.write(this.generateStandaloneHTML(component));
      newWindow.document.close();
    }
  }

  private openFullscreen(component: RichComponent): void {
    // Create fullscreen overlay
    const overlay = document.createElement('div');
    overlay.className = 'artifact-fullscreen-overlay';
    overlay.innerHTML = `
      <div class="fullscreen-header">
        <h3>${component.data.title || 'Artifact'}</h3>
        <button class="close-fullscreen">✕</button>
      </div>
      <div class="fullscreen-content">
        <iframe class="fullscreen-iframe" sandbox="allow-scripts allow-same-origin" srcdoc="${this.escapeHtml(component.data.content)}"></iframe>
      </div>
    `;

    // Add styles
    overlay.style.cssText = `
      position: fixed;
      top: 0;
      left: 0;
      width: 100vw;
      height: 100vh;
      background: white;
      z-index: 10000;
      display: flex;
      flex-direction: column;
    `;

    const header = overlay.querySelector('.fullscreen-header') as HTMLElement;
    header.style.cssText = `
      padding: 16px;
      border-bottom: 1px solid #eee;
      display: flex;
      justify-content: space-between;
      align-items: center;
    `;

    const content = overlay.querySelector('.fullscreen-content') as HTMLElement;
    content.style.cssText = `
      flex: 1;
      padding: 16px;
    `;

    const iframe = overlay.querySelector('.fullscreen-iframe') as HTMLIFrameElement;
    iframe.style.cssText = `
      width: 100%;
      height: 100%;
      border: none;
    `;

    // Close button functionality
    const closeBtn = overlay.querySelector('.close-fullscreen') as HTMLButtonElement;
    closeBtn.addEventListener('click', () => {
      document.body.removeChild(overlay);
    });

    // Escape key to close
    const handleEscape = (e: KeyboardEvent) => {
      if (e.key === 'Escape') {
        document.body.removeChild(overlay);
        document.removeEventListener('keydown', handleEscape);
      }
    };
    document.addEventListener('keydown', handleEscape);

    document.body.appendChild(overlay);
  }

  private openEditor(component: RichComponent): void {
    // Placeholder for future editor implementation
    console.log('Editor functionality not yet implemented for artifact:', component.data.artifact_id);
  }

  private escapeHtml(html: string): string {
    const div = document.createElement('div');
    div.textContent = html;
    return div.innerHTML.replace(/"/g, '&quot;');
  }
}

// User message component renderer
export class UserMessageComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const messageEl = document.createElement('vanna-message');
    messageEl.setAttribute('theme', 'light'); // Could be made dynamic
    messageEl.dataset.componentId = component.id;
    
    // Set properties for vanna-message
    (messageEl as any).content = component.data.content || '';
    (messageEl as any).type = 'user';
    (messageEl as any).timestamp = Date.parse(component.timestamp);
    
    return messageEl;
  }
}

// Assistant message component renderer  
export class AssistantMessageComponentRenderer extends BaseComponentRenderer {
  render(component: RichComponent): HTMLElement {
    const messageEl = document.createElement('vanna-message');
    messageEl.setAttribute('theme', 'light'); // Could be made dynamic
    messageEl.dataset.componentId = component.id;
    
    // Set properties for vanna-message
    (messageEl as any).content = component.data.content || '';
    (messageEl as any).type = 'assistant';
    (messageEl as any).timestamp = Date.parse(component.timestamp);
    
    return messageEl;
  }
}

// Component registry for managing all component types
export class ComponentRegistry {
  private renderers: Map<string, ComponentRenderer> = new Map();

  constructor() {
    // Register primitive component renderers (domain-agnostic)
    this.register('status_card', new StatusCardComponentRenderer());
    this.register('progress_display', new ProgressDisplayComponentRenderer());
    this.register('log_viewer', new LogViewerComponentRenderer());
    this.register('badge', new BadgeComponentRenderer());
    this.register('icon_text', new IconTextComponentRenderer());

    // Register existing component renderers
    this.register('card', new CardComponentRenderer());
    this.register('task_list', new TaskListComponentRenderer());
    this.register('progress_bar', new ProgressBarComponentRenderer());
    this.register('notification', new NotificationComponentRenderer());
    this.register('status_indicator', new StatusIndicatorComponentRenderer());
    this.register('text', new TextComponentRenderer());
    this.register('dataframe', new DataFrameComponentRenderer());
    this.register('chart', new ChartComponentRenderer());

    // Register interactive component renderers
    this.register('button', new ButtonComponentRenderer());
    this.register('button_group', new ButtonGroupComponentRenderer());

    // Register artifact component renderer
    this.register('artifact', new ArtifactComponentRenderer());

    // Register message component renderers
    this.register('user-message', new UserMessageComponentRenderer());
    this.register('assistant-message', new AssistantMessageComponentRenderer());
  }

  register(type: string, renderer: ComponentRenderer): void {
    this.renderers.set(type, renderer);
  }

  render(component: RichComponent): HTMLElement {
    // Check if this is a component that should use web components
    const webComponentTag = this.getWebComponentTag(component.type);
    if (webComponentTag) {
      return this.renderWebComponent(webComponentTag, component);
    }

    // Use the old renderer system for other components
    const renderer = this.renderers.get(component.type);
    if (!renderer) {
      return this.renderFallback(component);
    }
    return renderer.render(component);
  }

  private getWebComponentTag(type: string): string | null {
    const mapping: Record<string, string> = {
      'card': 'rich-card',
      'task_list': 'rich-task-list',
      'progress_bar': 'rich-progress-bar',
      // We'll add more mappings as we convert other components
    };
    return mapping[type] || null;
  }

  private renderWebComponent(tagName: string, component: RichComponent): HTMLElement {
    const element = document.createElement(tagName) as any;

    // Set properties based on component data
    Object.keys(component.data).forEach(key => {
      if (key === 'actions' && Array.isArray(component.data[key])) {
        element.actions = component.data[key];
      } else {
        element[key] = component.data[key];
      }
    });

    // Set theme to match the parent VannaChat theme
    element.setAttribute('theme', this.getCurrentTheme());


    return element;
  }

  private getCurrentTheme(): string {
    // Try to get theme from the parent VannaChat component
    const vannaChat = document.querySelector('vanna-chat');
    if (vannaChat) {
      return vannaChat.getAttribute('theme') || 'dark';
    }
    return 'dark';
  }


  update(element: HTMLElement, component: RichComponent, updates?: Record<string, any>): void {
    const renderer = this.renderers.get(component.type);
    if (renderer) {
      renderer.update(element, component, updates);
    }
  }

  remove(element: HTMLElement): void {
    element.remove();
  }

  private renderFallback(component: RichComponent): HTMLElement {
    const container = document.createElement('div');
    container.className = 'rich-component rich-fallback';
    container.dataset.componentId = component.id;

    container.innerHTML = `
      <div class="fallback-header">
        <strong>Unknown Component: ${component.type}</strong>
      </div>
      <pre class="fallback-data">${JSON.stringify(component.data, null, 2)}</pre>
    `;

    return container;
  }
}

// Component manager for handling component lifecycle
export class ComponentManager {
  private components: Map<string, RichComponent> = new Map();
  private elements: Map<string, HTMLElement> = new Map();
  private registry: ComponentRegistry = new ComponentRegistry();
  private container: HTMLElement;
  private readonly sharedFields = new Set([
    'id',
    'type',
    'lifecycle',
    'layout',
    'theme',
    'children',
    'timestamp',
    'visible',
    'interactive',
  ]);

  constructor(container: HTMLElement) {
    this.container = container;
    ensureRichComponentStyles(this.container);
  }

  processUpdate(update: ComponentUpdate): void {
    // Handle UI state updates with special processing
    if (update.component && this.isUIStateUpdate(update.component)) {
      this.processUIStateUpdate(update.component);
      return;
    }

    switch (update.operation) {
      case 'create':
        this.createComponent(update);
        break;
      case 'update':
        this.updateComponent(update);
        break;
      case 'replace':
        this.replaceComponent(update);
        break;
      case 'remove':
        this.removeComponent(update);
        break;
    }
  }

  private createComponent(update: ComponentUpdate): void {
    if (!update.component) return;

    const component = this.normalizeComponent(update.component);
    const element = this.registry.render(component);
    this.components.set(component.id, component);
    this.elements.set(component.id, element);

    // Determine where to place the component
    this.positionComponent(element);
  }

  private updateComponent(update: ComponentUpdate): void {
    if (!update.component) return;

    const element = this.elements.get(update.target_id);
    if (element) {
      const component = this.normalizeComponent(update.component);
      this.registry.update(element, component, update.updates);
      this.components.set(update.target_id, component);
    }
  }

  private replaceComponent(update: ComponentUpdate): void {
    if (!update.component) return;

    const oldElement = this.elements.get(update.target_id);
    if (oldElement) {
      const component = this.normalizeComponent(update.component);
      const newElement = this.registry.render(component);
      oldElement.parentNode?.replaceChild(newElement, oldElement);

      this.elements.set(component.id, newElement);
      this.components.set(component.id, component);

      // Clean up old references if ID changed
      if (update.target_id !== component.id) {
        this.elements.delete(update.target_id);
        this.components.delete(update.target_id);
      }
    }
  }

  private removeComponent(update: ComponentUpdate): void {
    const element = this.elements.get(update.target_id);
    if (element) {
      element.remove();
      this.elements.delete(update.target_id);
      this.components.delete(update.target_id);
    }
  }

  private positionComponent(element: HTMLElement): void {
    // Always append to container
    this.container.appendChild(element);

    // Trigger scroll to bottom in parent chat component
    this.triggerScroll();
  }

  private triggerScroll(): void {
    // Find the parent vanna-chat component and trigger its scroll method
    const vannaChat = document.querySelector('vanna-chat') as any;
    if (vannaChat && typeof vannaChat.scrollToLastMessage === 'function') {
      // Use requestAnimationFrame to wait for DOM update
      requestAnimationFrame(() => {
        requestAnimationFrame(() => {
          vannaChat.scrollToLastMessage();
        });
      });
    }
  }

  clear(): void {
    this.components.clear();
    this.elements.clear();
    this.container.innerHTML = '';
    ensureRichComponentStyles(this.container);
  }

  getComponent(id: string): RichComponent | undefined {
    return this.components.get(id);
  }

  getAllComponents(): RichComponent[] {
    return Array.from(this.components.values());
  }

  private normalizeComponent(component: RichComponent): RichComponent {
    const data = { ...(component.data ?? {}) };

    for (const [key, value] of Object.entries(component as Record<string, any>)) {
      if (this.sharedFields.has(key) || key === 'data') continue;
      data[key] = value;
    }

    if (component.data && Object.keys(component.data).length === Object.keys(data).length) {
      return component;
    }

    return {
      ...component,
      data,
    };
  }

  private isUIStateUpdate(component: RichComponent): boolean {
    return component.type === 'status_bar_update' ||
           component.type === 'task_tracker_update' ||
           component.type === 'chat_input_update';
  }

  private processUIStateUpdate(component: RichComponent): void {
    console.log('processUIStateUpdate called with type:', component.type, 'component:', component);

    switch (component.type) {
      case 'status_bar_update':
        this.updateStatusBar(component);
        break;
      case 'task_tracker_update':
        this.updateTaskTracker(component);
        break;
      case 'chat_input_update':
        this.updateChatInput(component);
        break;
    }
  }

  private updateStatusBar(component: RichComponent): void {
    // Find the status bar component - first try shadow DOM, then document
    let statusBar: HTMLElement | null = null;

    // Look for vanna-chat and search within its shadow root
    const vannaChat = document.querySelector('vanna-chat') as any;
    if (vannaChat && vannaChat.shadowRoot) {
      statusBar = vannaChat.shadowRoot.querySelector('vanna-status-bar') as HTMLElement | null;
    }

    // Fallback to document search
    if (!statusBar) {
      statusBar = document.querySelector('vanna-status-bar') as HTMLElement | null;
    }

    if (statusBar) {
      const { status, message, detail } = component.data || {};
      // Set properties directly on the Lit component
      (statusBar as any).status = status;
      (statusBar as any).message = message || '';
      (statusBar as any).detail = detail || '';
    }
  }

  private updateTaskTracker(component: RichComponent): void {
    // Debug logging
    console.log('updateTaskTracker called with component:', component);
    console.log('component.data:', component.data);

    // Find the progress tracker component - first try shadow DOM, then document
    let progressTracker = null;

    // Look for vanna-chat and search within its shadow root
    const vannaChat = document.querySelector('vanna-chat') as any;
    if (vannaChat && vannaChat.shadowRoot) {
      progressTracker = vannaChat.shadowRoot.querySelector('vanna-progress-tracker');
    }

    // Fallback to document search
    if (!progressTracker) {
      progressTracker = document.querySelector('vanna-progress-tracker');
    }

    console.log('Found progressTracker:', progressTracker);
    if (!progressTracker) return;

    const { operation, task, task_id, status, detail } = component.data || {};
    console.log('Extracted data:', { operation, task, task_id, status, detail });

    switch (operation) {
      case 'add_task':
        console.log('Adding task:', task);
        if (task && progressTracker.addItem) {
          // Use the backend task ID instead of generating a new one
          const result = progressTracker.addItem(task.title || task.text, task.description || task.detail, task.id);
          console.log('addItem result:', result, 'using backend ID:', task.id);
        }
        break;
      case 'update_task':
        console.log('Updating task:', task_id, status, detail);
        if (task_id && progressTracker.updateItem) {
          progressTracker.updateItem(task_id, status, detail);
        }
        break;
      case 'remove_task':
        if (task_id && progressTracker.removeItem) {
          progressTracker.removeItem(task_id);
        }
        break;
      case 'clear_tasks':
        if (progressTracker.clear) {
          progressTracker.clear();
        }
        break;
    }
  }

  private updateChatInput(component: RichComponent): void {
    // Find the chat input element - first try shadow DOM, then document
    let chatInput = null;

    // Look for vanna-chat and search within its shadow root
    const vannaChat = document.querySelector('vanna-chat') as any;
    if (vannaChat && vannaChat.shadowRoot) {
      chatInput = vannaChat.shadowRoot.querySelector('textarea.message-input, input.message-input');
    }

    // Fallback to document search with multiple selectors
    if (!chatInput) {
      chatInput = document.querySelector('textarea[data-testid="message-input"], input[type="text"].message-input, .message-input input, .message-input textarea');
    }

    if (!chatInput) return;

    const { placeholder, disabled, value, focus } = component.data || {};

    if (placeholder !== undefined) {
      chatInput.placeholder = placeholder;
    }
    if (disabled !== undefined) {
      chatInput.disabled = disabled;
    }
    if (value !== undefined) {
      chatInput.value = value;
    }
    if (focus !== undefined) {
      if (focus) {
        chatInput.focus();
      } else {
        chatInput.blur();
      }
    }
  }
}


================================================
FILE: frontends/webcomponent/src/components/rich-progress-bar.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './rich-progress-bar';

const meta: Meta = {
  title: 'Rich Components/Rich Progress Bar',
  component: 'rich-progress-bar',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'dark', value: 'rgb(11, 15, 25)' },
        { name: 'light', value: '#f5f7fa' },
      ],
    },
  },
  argTypes: {
    value: { control: { type: 'range', min: 0, max: 1, step: 0.01 } },
    label: { control: 'text' },
    description: { control: 'text' },
    showPercentage: { control: 'boolean' },
    status: {
      control: 'select',
      options: ['info', 'success', 'warning', 'error']
    },
    animated: { control: 'boolean' },
    indeterminate: { control: 'boolean' },
  },
};

export default meta;
type Story = StoryObj;

export const Default: Story = {
  args: {
    value: 0.65,
    label: 'Processing',
    showPercentage: true,
    status: 'info',
    animated: false,
    indeterminate: false,
  },
  render: (args) => html`
    <div style="max-width: 500px; margin: 0 auto;">
      <rich-progress-bar
        .value=${args.value}
        label=${args.label}
        description=${args.description}
        ?showPercentage=${args.showPercentage}
        status=${args.status}
        ?animated=${args.animated}
        ?indeterminate=${args.indeterminate}>
      </rich-progress-bar>
    </div>
  `,
};

export const WithDescription: Story = {
  args: {
    value: 0.4,
    label: 'Installing dependencies',
    description: 'Downloading and installing npm packages for the project. This may take a few minutes.',
    showPercentage: true,
    status: 'info',
    animated: true,
    indeterminate: false,
  },
  render: (args) => html`
    <div style="max-width: 500px; margin: 0 auto;">
      <rich-progress-bar
        .value=${args.value}
        label=${args.label}
        description=${args.description}
        ?showPercentage=${args.showPercentage}
        status=${args.status}
        ?animated=${args.animated}
        ?indeterminate=${args.indeterminate}>
      </rich-progress-bar>
    </div>
  `,
};

export const Animated: Story = {
  args: {
    value: 0.75,
    label: 'Uploading files',
    showPercentage: true,
    status: 'info',
    animated: true,
    indeterminate: false,
  },
  render: (args) => html`
    <div style="max-width: 500px; margin: 0 auto;">
      <rich-progress-bar
        .value=${args.value}
        label=${args.label}
        description=${args.description}
        ?showPercentage=${args.showPercentage}
        status=${args.status}
        ?animated=${args.animated}
        ?indeterminate=${args.indeterminate}>
      </rich-progress-bar>
    </div>
  `,
};

export const Indeterminate: Story = {
  args: {
    value: 0,
    label: 'Loading...',
    description: 'Please wait while we process your request',
    showPercentage: false,
    status: 'info',
    animated: false,
    indeterminate: true,
  },
  render: (args) => html`
    <div style="max-width: 500px; margin: 0 auto;">
      <rich-progress-bar
        .value=${args.value}
        label=${args.label}
        description=${args.description}
        ?showPercentage=${args.showPercentage}
        status=${args.status}
        ?animated=${args.animated}
        ?indeterminate=${args.indeterminate}>
      </rich-progress-bar>
    </div>
  `,
};

export const StatusVariants: Story = {
  render: () => html`
    <div style="max-width: 500px; margin: 0 auto; display: flex; flex-direction: column; gap: 16px;">
      <rich-progress-bar
        .value=${0.8}
        label="Info Status"
        status="info"
        showPercentage>
      </rich-progress-bar>

      <rich-progress-bar
        .value=${1.0}
        label="Success Status"
        status="success"
        showPercentage>
      </rich-progress-bar>

      <rich-progress-bar
        .value=${0.6}
        label="Warning Status"
        status="warning"
        showPercentage>
      </rich-progress-bar>

      <rich-progress-bar
        .value=${0.3}
        label="Error Status"
        status="error"
        showPercentage>
      </rich-progress-bar>
    </div>
  `,
};

export const Minimal: Story = {
  args: {
    value: 0.45,
    showPercentage: false,
    status: 'info',
    animated: false,
    indeterminate: false,
  },
  render: (args) => html`
    <div style="max-width: 500px; margin: 0 auto;">
      <rich-progress-bar
        .value=${args.value}
        ?showPercentage=${args.showPercentage}
        status=${args.status}
        ?animated=${args.animated}
        ?indeterminate=${args.indeterminate}>
      </rich-progress-bar>
    </div>
  `,
};

export const MultipleSteps: Story = {
  render: () => html`
    <div style="max-width: 500px; margin: 0 auto; display: flex; flex-direction: column; gap: 16px;">
      <rich-progress-bar
        .value=${1.0}
        label="Step 1: Initialize"
        description="Project initialization completed"
        status="success"
        showPercentage>
      </rich-progress-bar>

      <rich-progress-bar
        .value=${1.0}
        label="Step 2: Download dependencies"
        description="All packages downloaded successfully"
        status="success"
        showPercentage>
      </rich-progress-bar>

      <rich-progress-bar
        .value=${0.7}
        label="Step 3: Build project"
        description="Compiling TypeScript and bundling assets"
        status="info"
        animated
        showPercentage>
      </rich-progress-bar>

      <rich-progress-bar
        .value=${0}
        label="Step 4: Deploy"
        description="Waiting for build to complete"
        status="info"
        showPercentage>
      </rich-progress-bar>
    </div>
  `,
};

export const LightTheme: Story = {
  args: {
    value: 0.55,
    label: 'Light Theme Progress',
    description: 'Progress bar styled for light backgrounds',
    showPercentage: true,
    status: 'success',
    animated: true,
  },
  parameters: {
    backgrounds: { default: 'light' }
  },
  render: (args) => html`
    <div style="max-width: 500px; margin: 0 auto;">
      <rich-progress-bar
        theme="light"
        .value=${args.value}
        label=${args.label}
        description=${args.description}
        ?showPercentage=${args.showPercentage}
        status=${args.status}
        ?animated=${args.animated}>
      </rich-progress-bar>
    </div>
  `,
};

================================================
FILE: frontends/webcomponent/src/components/rich-progress-bar.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

@customElement('rich-progress-bar')
export class RichProgressBar extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        margin-bottom: var(--vanna-space-4);
        font-family: var(--vanna-font-family-default);
      }

      .progress-container {
        padding: var(--vanna-space-4);
        border: 1px solid var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-lg);
        background: var(--vanna-background-default);
        box-shadow: var(--vanna-shadow-sm);
        transition: box-shadow var(--vanna-duration-200) ease;
      }

      .progress-container:hover {
        box-shadow: var(--vanna-shadow-md);
      }

      .progress-header {
        display: flex;
        justify-content: space-between;
        align-items: center;
        margin-bottom: var(--vanna-space-3);
      }

      .progress-label {
        font-weight: 500;
        color: var(--vanna-foreground-default);
      }

      .progress-percentage {
        font-size: 0.875rem;
        color: var(--vanna-foreground-dimmer);
        font-weight: 600;
      }

      .progress-track {
        height: 12px;
        background: var(--vanna-background-root);
        border-radius: 6px;
        overflow: hidden;
        border: 1px solid var(--vanna-outline-default);
        position: relative;
      }

      .progress-fill {
        height: 100%;
        background: var(--vanna-accent-primary-default);
        border-radius: 6px;
        transition: width var(--vanna-duration-300) ease;
        position: relative;
        overflow: hidden;
      }

      .progress-fill.animated {
        animation: progressPulse 2s ease-in-out infinite;
      }

      .progress-fill.animated::after {
        content: '';
        position: absolute;
        top: 0;
        left: 0;
        bottom: 0;
        right: 0;
        background: linear-gradient(
          90deg,
          transparent,
          rgba(255, 255, 255, 0.2),
          transparent
        );
        animation: progressShimmer 1.5s infinite;
      }

      @keyframes progressPulse {
        0%, 100% { opacity: 1; }
        50% { opacity: 0.8; }
      }

      @keyframes progressShimmer {
        0% { transform: translateX(-100%); }
        100% { transform: translateX(100%); }
      }

      .progress-fill.status-success {
        background: var(--vanna-accent-positive-default);
      }

      .progress-fill.status-warning {
        background: var(--vanna-accent-warning-default);
      }

      .progress-fill.status-error {
        background: var(--vanna-accent-negative-default);
      }

      .progress-fill.status-info {
        background: var(--vanna-accent-primary-default);
      }

      /* Indeterminate progress animation */
      .progress-fill.indeterminate {
        background: linear-gradient(
          90deg,
          transparent 0%,
          var(--vanna-accent-primary-default) 50%,
          transparent 100%
        );
        background-size: 200% 100%;
        animation: indeterminateProgress 2s linear infinite;
        width: 100% !important;
      }

      @keyframes indeterminateProgress {
        0% { background-position: 200% 0; }
        100% { background-position: -200% 0; }
      }

      /* Text content for description */
      .progress-description {
        margin-top: var(--vanna-space-2);
        font-size: 0.875rem;
        color: var(--vanna-foreground-dimmer);
        line-height: 1.4;
      }
    `
  ];

  @property({ type: Number }) value = 0;
  @property() label = '';
  @property() description = '';
  @property({ type: Boolean }) showPercentage = true;
  @property() status: 'info' | 'success' | 'warning' | 'error' = 'info';
  @property({ type: Boolean }) animated = false;
  @property({ type: Boolean }) indeterminate = false;
  @property() theme: 'light' | 'dark' = 'dark';

  private get percentage(): number {
    if (this.indeterminate) return 100;
    return Math.round(Math.max(0, Math.min(1, this.value)) * 100);
  }

  private get progressClasses(): string {
    const classes = ['progress-fill'];

    if (this.animated) {
      classes.push('animated');
    }

    if (this.indeterminate) {
      classes.push('indeterminate');
    }

    if (this.status) {
      classes.push(`status-${this.status}`);
    }

    return classes.join(' ');
  }

  render() {
    return html`
      <div class="progress-container">
        ${this.label || this.showPercentage ? html`
          <div class="progress-header">
            ${this.label ? html`<span class="progress-label">${this.label}</span>` : ''}
            ${this.showPercentage && !this.indeterminate ? html`
              <span class="progress-percentage">${this.percentage}%</span>
            ` : ''}
          </div>
        ` : ''}

        <div class="progress-track">
          <div
            class="${this.progressClasses}"
            style="width: ${this.indeterminate ? '100' : this.percentage}%">
          </div>
        </div>

        ${this.description ? html`
          <div class="progress-description">${this.description}</div>
        ` : ''}
      </div>
    `;
  }
}

declare global {
  interface HTMLElementTagNameMap {
    'rich-progress-bar': RichProgressBar;
  }
}

================================================
FILE: frontends/webcomponent/src/components/rich-task-list.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './rich-task-list';

const meta: Meta = {
  title: 'Rich Components/Rich Task List',
  component: 'rich-task-list',
  parameters: {
    layout: 'padded',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'dark', value: 'rgb(11, 15, 25)' },
        { name: 'light', value: '#f5f7fa' },
      ],
    },
  },
  argTypes: {
    title: { control: 'text' },
    tasks: { control: 'object' },
    showProgress: { control: 'boolean' },
    showTimestamps: { control: 'boolean' },
  },
};

export default meta;
type Story = StoryObj;

const sampleTasks = [
  {
    id: '1',
    title: 'Initialize project setup',
    description: 'Setting up the basic project structure and dependencies',
    status: 'completed',
    progress: 1.0,
    timestamp: '2024-01-15 10:30:00'
  },
  {
    id: '2',
    title: 'Configure database connection',
    description: 'Establishing secure connection to PostgreSQL database',
    status: 'completed',
    progress: 1.0,
    timestamp: '2024-01-15 10:45:00'
  },
  {
    id: '3',
    title: 'Implement user authentication',
    description: 'Building JWT-based authentication system',
    status: 'running',
    progress: 0.7,
    timestamp: '2024-01-15 11:00:00'
  },
  {
    id: '4',
    title: 'Create API endpoints',
    description: 'Developing RESTful API for user management',
    status: 'pending',
    timestamp: '2024-01-15 11:30:00'
  },
  {
    id: '5',
    title: 'Write unit tests',
    description: 'Comprehensive test coverage for all modules',
    status: 'pending',
  }
];

export const Default: Story = {
  args: {
    title: 'Development Tasks',
    tasks: sampleTasks,
    showProgress: true,
    showTimestamps: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

export const WithTimestamps: Story = {
  args: {
    title: 'Build Pipeline',
    tasks: sampleTasks,
    showProgress: true,
    showTimestamps: true,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

export const WithoutProgress: Story = {
  args: {
    title: 'Simple Task List',
    tasks: [
      { id: '1', title: 'Review code changes', status: 'completed' },
      { id: '2', title: 'Update documentation', status: 'running' },
      { id: '3', title: 'Deploy to staging', status: 'pending' },
      { id: '4', title: 'Run integration tests', status: 'failed' },
    ],
    showProgress: false,
    showTimestamps: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

export const AllStatuses: Story = {
  args: {
    title: 'Task Status Examples',
    tasks: [
      {
        id: '1',
        title: 'Completed Task',
        description: 'This task has been successfully completed',
        status: 'completed',
        progress: 1.0,
      },
      {
        id: '2',
        title: 'Running Task',
        description: 'This task is currently in progress',
        status: 'running',
        progress: 0.6,
      },
      {
        id: '3',
        title: 'Pending Task',
        description: 'This task is waiting to be started',
        status: 'pending',
      },
      {
        id: '4',
        title: 'Failed Task',
        description: 'This task encountered an error',
        status: 'failed',
        progress: 0.3,
      },
    ],
    showProgress: true,
    showTimestamps: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

export const EmptyList: Story = {
  args: {
    title: 'No Tasks',
    tasks: [],
    showProgress: true,
    showTimestamps: false,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

export const ErrorStates: Story = {
  args: {
    title: 'Error Handling Examples',
    tasks: [
      {
        id: '1',
        title: 'Database Connection Failed',
        description: 'Could not establish connection to the database server. Check network connectivity and credentials.',
        status: 'failed',
        progress: 0.1,
        timestamp: '2024-01-15 10:15:00'
      },
      {
        id: '2',
        title: 'API Authentication Error',
        description: 'Invalid API key or expired token. Please refresh your credentials.',
        status: 'failed',
        progress: 0.0,
        timestamp: '2024-01-15 10:20:00'
      },
      {
        id: '3',
        title: 'File Processing Error',
        description: 'Unable to process uploaded file. File may be corrupted or in an unsupported format.',
        status: 'failed',
        progress: 0.45,
        timestamp: '2024-01-15 10:25:00'
      },
      {
        id: '4',
        title: 'Network Timeout',
        description: 'Request timed out after 30 seconds. This may be due to high server load.',
        status: 'failed',
        progress: 0.8,
        timestamp: '2024-01-15 10:30:00'
      },
    ],
    showProgress: true,
    showTimestamps: true,
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

export const LightTheme: Story = {
  args: {
    title: 'Light Theme Task List',
    tasks: sampleTasks.slice(0, 3),
    showProgress: true,
    showTimestamps: true,
  },
  parameters: {
    backgrounds: { default: 'light' }
  },
  render: (args) => html`
    <div style="max-width: 600px; margin: 0 auto;">
      <rich-task-list
        theme="light"
        title=${args.title}
        .tasks=${args.tasks}
        ?showProgress=${args.showProgress}
        ?showTimestamps=${args.showTimestamps}>
      </rich-task-list>
    </div>
  `,
};

================================================
FILE: frontends/webcomponent/src/components/rich-task-list.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

export interface TaskItem {
  id: string;
  title: string;
  description?: string;
  status: 'pending' | 'running' | 'completed' | 'failed';
  progress?: number;
  timestamp?: string;
}

@customElement('rich-task-list')
export class RichTaskList extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        margin-bottom: var(--vanna-space-4);
        font-family: var(--vanna-font-family-default);
      }

      .task-list {
        border: 1px solid var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-lg);
        background: var(--vanna-background-default);
        box-shadow: var(--vanna-shadow-sm);
        overflow: hidden;
        transition: box-shadow var(--vanna-duration-200) ease;
      }

      .task-list:hover {
        box-shadow: var(--vanna-shadow-md);
      }

      .task-list-header {
        padding: var(--vanna-space-4) var(--vanna-space-5);
        background: var(--vanna-background-higher);
        border-bottom: 1px solid var(--vanna-outline-default);
      }

      .task-list-title {
        margin: 0 0 var(--vanna-space-3) 0;
        font-size: 1rem;
        font-weight: 600;
        color: var(--vanna-foreground-default);
      }

      .task-list-progress {
        display: flex;
        align-items: center;
        gap: var(--vanna-space-3);
      }

      .progress-text {
        font-size: 0.875rem;
        color: var(--vanna-foreground-dimmer);
        min-width: fit-content;
      }

      .progress-bar {
        flex: 1;
        height: 6px;
        background: var(--vanna-background-root);
        border-radius: 3px;
        overflow: hidden;
      }

      .progress-fill {
        height: 100%;
        background: var(--vanna-accent-primary-default);
        border-radius: 3px;
        transition: width var(--vanna-duration-300) ease;
      }

      .progress-fill.animated {
        animation: progressPulse 2s ease-in-out infinite;
      }

      @keyframes progressPulse {
        0%, 100% { opacity: 1; }
        50% { opacity: 0.7; }
      }

      .progress-fill.status-success {
        background: var(--vanna-accent-positive-default);
      }

      .progress-fill.status-warning {
        background: var(--vanna-accent-warning-default);
      }

      .progress-fill.status-error {
        background: var(--vanna-accent-negative-default);
      }

      .task-list-items {
        padding: var(--vanna-space-2);
      }

      .task-item {
        display: flex;
        align-items: flex-start;
        gap: var(--vanna-space-3);
        padding: var(--vanna-space-3);
        border-radius: var(--vanna-border-radius-md);
        transition: background-color var(--vanna-duration-200) ease;
      }

      .task-item:hover {
        background: var(--vanna-background-root);
      }

      .task-item.status-completed {
        opacity: 0.7;
      }

      .task-item.status-failed {
        background: rgba(239, 68, 68, 0.1);
      }

      .task-icon {
        font-size: 1rem;
        margin-top: 0.125rem;
      }

      .task-content {
        flex: 1;
        min-width: 0;
      }

      .task-title {
        font-weight: 500;
        color: var(--vanna-foreground-default);
        margin-bottom: var(--vanna-space-1);
      }

      .task-description {
        font-size: 0.875rem;
        color: var(--vanna-foreground-dimmer);
        margin-bottom: var(--vanna-space-2);
      }

      .task-progress {
        display: flex;
        align-items: center;
        gap: var(--vanna-space-2);
        margin-bottom: var(--vanna-space-2);
      }

      .task-progress-bar {
        flex: 1;
        height: 4px;
        background: var(--vanna-background-root);
        border-radius: 2px;
        overflow: hidden;
      }

      .task-progress-fill {
        height: 100%;
        background: var(--vanna-accent-primary-default);
        border-radius: 2px;
        transition: width var(--vanna-duration-300) ease;
      }

      .task-progress-text {
        font-size: 0.75rem;
        color: var(--vanna-foreground-dimmer);
        min-width: fit-content;
      }

      .task-timestamp {
        font-size: 0.75rem;
        color: var(--vanna-foreground-dimmest);
      }

      /* Responsive adjustments */
      @media (max-width: 768px) {
        .task-list-header {
          padding-left: var(--vanna-space-4);
          padding-right: var(--vanna-space-4);
        }

        .task-list-progress {
          flex-direction: column;
          align-items: stretch;
          gap: var(--vanna-space-2);
        }
      }
    `
  ];

  @property() title = '';
  @property({ type: Array }) tasks: TaskItem[] = [];
  @property({ type: Boolean }) showProgress = true;
  @property({ type: Boolean }) showTimestamps = false;
  @property() theme: 'light' | 'dark' = 'dark';

  private get completedTasks(): number {
    return this.tasks.filter(task => task.status === 'completed').length;
  }

  private get progressPercentage(): number {
    return this.tasks.length > 0 ? (this.completedTasks / this.tasks.length) * 100 : 0;
  }

  private getStatusIcon(status: string): string {
    const icons = {
      'pending': '⏳',
      'running': '🔄',
      'completed': '✅',
      'failed': '❌'
    };
    return icons[status as keyof typeof icons] || '⏳';
  }

  private renderTask(task: TaskItem) {
    const statusIcon = this.getStatusIcon(task.status);

    return html`
      <div class="task-item status-${task.status}" data-task-id="${task.id}">
        <div class="task-icon">${statusIcon}</div>
        <div class="task-content">
          <div class="task-title">${task.title}</div>
          ${task.description ? html`
            <div class="task-description">${task.description}</div>
          ` : ''}
          ${task.progress !== null && task.progress !== undefined ? html`
            <div class="task-progress">
              <div class="task-progress-bar">
                <div class="task-progress-fill" style="width: ${task.progress * 100}%"></div>
              </div>
              <span class="task-progress-text">${Math.round(task.progress * 100)}%</span>
            </div>
          ` : ''}
          ${this.showTimestamps && task.timestamp ? html`
            <div class="task-timestamp">${task.timestamp}</div>
          ` : ''}
        </div>
      </div>
    `;
  }

  render() {
    return html`
      <div class="task-list">
        <div class="task-list-header">
          <h3 class="task-list-title">${this.title}</h3>
          ${this.showProgress ? html`
            <div class="task-list-progress">
              <span class="progress-text">${this.completedTasks}/${this.tasks.length} completed</span>
              <div class="progress-bar">
                <div class="progress-fill" style="width: ${this.progressPercentage}%"></div>
              </div>
            </div>
          ` : ''}
        </div>
        <div class="task-list-items">
          ${this.tasks.map(task => this.renderTask(task))}
        </div>
      </div>
    `;
  }
}

declare global {
  interface HTMLElementTagNameMap {
    'rich-task-list': RichTaskList;
  }
}

================================================
FILE: frontends/webcomponent/src/components/vanna-chat.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './vanna-chat';
import './vanna-message';
import './plotly-chart';

const meta: Meta = {
  title: 'Components/VannaChat',
  component: 'vanna-chat',
  parameters: {
    layout: 'fullscreen',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'light', value: '#f5f7fa' },
        { name: 'dark', value: 'rgb(11, 15, 25)' },
      ],
    },
  },
  argTypes: {
    title: { control: 'text' },
    placeholder: { control: 'text' },
    disabled: { control: 'boolean' },
    showProgress: { control: 'boolean' },
    maxAutonomy: { control: 'boolean' },
    theme: {
      control: 'select',
      options: ['dark', 'light'],
      description: 'Theme variant'
    },
  },
};

export default meta;
type Story = StoryObj;

export const Default: Story = {
  args: {
    title: 'Vanna AI Agent',
    placeholder: 'Describe what you want to build...',
    disabled: false,
    showProgress: true,
    maxAutonomy: false,
    theme: 'light',
  },
  render: (args) => html`
    <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
      <vanna-chat
        .title=${args.title}
        .placeholder=${args.placeholder}
        .disabled=${args.disabled}
        .showProgress=${args.showProgress}
        .maxAutonomy=${args.maxAutonomy}
        theme=${args.theme}>
      </vanna-chat>
    </div>
  `,
};

export const LightMode: Story = {
  args: {
    title: 'Vanna AI Agent',
    placeholder: 'Describe what you want to build...',
    disabled: false,
    showProgress: true,
    maxAutonomy: false,
    theme: 'light',
  },
  render: (args) => html`
    <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
      <vanna-chat
        .title=${args.title}
        .placeholder=${args.placeholder}
        .disabled=${args.disabled}
        .showProgress=${args.showProgress}
        .maxAutonomy=${args.maxAutonomy}
        theme=${args.theme}>
      </vanna-chat>
    </div>
  `,
};

export const WithConversation: Story = {
  args: {
    title: 'Vanna AI Agent',
    placeholder: 'Continue the conversation...',
    disabled: false,
    showProgress: true,
    maxAutonomy: true,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      const tracker = chat?.getProgressTracker();

      if (chat && tracker) {
        // Add conversation messages
        chat.addMessage('Create a dashboard for analyzing customer data', 'user');
        chat.addMessage('I\'ll help you create a customer data dashboard. Let me break this down into steps and get started.', 'assistant');
        chat.addMessage('What specific metrics would you like to track? Revenue, acquisition, retention, or something else?', 'assistant');
        chat.addMessage('Focus on revenue and customer acquisition metrics', 'user');

        // Add progress items
        const id1 = tracker.addItem('Analyze requirements', 'Understanding dashboard needs');
        const id2 = tracker.addItem('Design data schema', 'Planning database structure');
        tracker.addItem('Create visualization components', 'Building charts and graphs');
        tracker.addItem('Implement filtering', 'Adding date range and segment filters');

        // Update progress states
        tracker.updateItem(id1, 'completed');
        tracker.updateItem(id2, 'in_progress', 'Identifying key metrics and data sources');

        // Set status
        chat.setStatus('working', 'Analyzing data requirements...', 'Step 2 of 4');
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const MaxAutonomyMode: Story = {
  args: {
    title: 'Vanna AI Agent - Max Autonomy',
    placeholder: 'Describe your project...',
    disabled: false,
    showProgress: true,
    maxAutonomy: true,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      const tracker = chat?.getProgressTracker();

      if (chat && tracker) {
        chat.addMessage('Build a full-stack e-commerce app with user authentication, product catalog, shopping cart, and payment processing', 'user');
        chat.addMessage('Perfect! I\'ll build a complete e-commerce application for you. Since Max Autonomy is enabled, I\'ll handle all the technical decisions and implementation details automatically.', 'assistant');

        // Comprehensive task list for full autonomy
        const tasks = [
          'Set up project structure',
          'Configure development environment',
          'Design database schema',
          'Implement user authentication',
          'Build product catalog API',
          'Create shopping cart functionality',
          'Integrate payment processing',
          'Develop frontend components',
          'Add responsive design',
          'Implement search & filtering',
          'Set up testing framework',
          'Configure deployment pipeline'
        ];

        tasks.forEach((task, index) => {
          const id = tracker.addItem(task, `Feature ${index + 1} of ${tasks.length}`);
          if (index < 3) tracker.updateItem(id, 'completed');
          else if (index === 3) tracker.updateItem(id, 'in_progress', 'Setting up JWT tokens and password hashing');
        });

        chat.setStatus('working', 'Building authentication system...', 'High autonomy mode active');
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const WorkingState: Story = {
  args: {
    title: 'Vanna AI Agent',
    placeholder: 'Ask me anything...',
    disabled: true,
    showProgress: true,
    maxAutonomy: false,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      const tracker = chat?.getProgressTracker();

      if (chat && tracker) {
        chat.addMessage('Generate a monthly sales report with charts', 'user');

        const id1 = tracker.addItem('Connect to database', 'Establishing secure connection');
        const id2 = tracker.addItem('Query sales data', 'Fetching monthly records');
        tracker.addItem('Process data', 'Calculating totals and trends');
        tracker.addItem('Generate charts', 'Creating visualizations');
        tracker.addItem('Format report', 'Compiling final document');

        tracker.updateItem(id1, 'completed');
        tracker.updateItem(id2, 'in_progress', 'SELECT * FROM sales WHERE date >= 2024-01...');

        chat.setStatus('working', 'Querying sales database...', 'Processing 12,543 records');
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const CompactMode: Story = {
  args: {
    title: 'Vanna AI Agent',
    placeholder: 'Quick question...',
    disabled: false,
    showProgress: false,
    maxAutonomy: false,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      if (chat) {
        chat.addMessage('What\'s the average order value this month?', 'user');
        chat.addMessage('Let me query that for you...', 'assistant');
        chat.setStatus('working', 'Calculating average order value...', '2.1s');
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const WithRichComponents: Story = {
  args: {
    title: 'Vanna AI Agent - Rich Components',
    placeholder: 'Ask me to analyze data or build something...',
    disabled: false,
    showProgress: true,
    maxAutonomy: false,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      const tracker = chat?.getProgressTracker();

      if (chat && tracker) {
        // Initial conversation
        chat.addMessage('Create a comprehensive sales dashboard with multiple visualizations and export capabilities', 'user');
        chat.addMessage('I\'ll create a comprehensive sales dashboard for you. Let me break this down into clear tasks and show you the progress with rich components.', 'assistant');

        // Add progress tasks
        const taskId1 = tracker.addItem('Analyze requirements', 'Understanding dashboard specifications');
        const taskId2 = tracker.addItem('Design data schema', 'Planning database structure');
        const taskId3 = tracker.addItem('Create visualizations', 'Building charts and graphs');
        const taskId4 = tracker.addItem('Add export features', 'Implementing PDF and Excel export');
        tracker.addItem('Deploy dashboard', 'Setting up production environment');

        tracker.updateItem(taskId1, 'completed');
        tracker.updateItem(taskId2, 'completed');
        tracker.updateItem(taskId3, 'in_progress', 'Creating revenue trend charts...');

        chat.setStatus('working', 'Building visualization components...', 'Step 3 of 5');

        // Add rich components after a delay
        setTimeout(() => {
          const componentManager = chat.componentManager;
          if (!componentManager) return;

          // Add info notification
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'info-notification',
            component: {
              id: 'info-notification',
              type: 'notification',
              data: {
                title: 'Dashboard Progress',
                message: 'Your sales dashboard is being built with the following components: revenue trends, customer analytics, and performance metrics.',
                level: 'info',
                dismissible: true,
                actions: []
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add status indicator
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'status-chart-generation',
            component: {
              id: 'status-chart-generation',
              type: 'status_indicator',
              data: {
                status: 'loading',
                message: 'Generating revenue trend charts...',
                pulse: true
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add progress bar for chart generation
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'chart-progress',
            component: {
              id: 'chart-progress',
              type: 'progress_bar',
              data: {
                progress: 65,
                status: 'active',
                label: 'Chart Generation Progress',
                detail: 'Processing 12,543 sales records...',
                animated: true
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add task list card
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'dashboard-tasks',
            component: {
              id: 'dashboard-tasks',
              type: 'task_list',
              data: {
                title: 'Dashboard Components',
                tasks: [
                  {
                    id: 'task-1',
                    title: 'Revenue Trend Chart',
                    description: 'Monthly revenue tracking with year-over-year comparison',
                    status: 'completed',
                    progress: 100,
                    timestamp: '2024-01-15 14:32:00'
                  },
                  {
                    id: 'task-2',
                    title: 'Customer Acquisition Funnel',
                    description: 'Lead to customer conversion visualization',
                    status: 'running',
                    progress: 75,
                    timestamp: '2024-01-15 14:45:00'
                  },
                  {
                    id: 'task-3',
                    title: 'Geographic Sales Map',
                    description: 'Interactive map showing sales by region',
                    status: 'pending',
                    progress: 0,
                    timestamp: null
                  },
                  {
                    id: 'task-4',
                    title: 'Performance Metrics KPIs',
                    description: 'Key performance indicators dashboard',
                    status: 'pending',
                    progress: 0,
                    timestamp: null
                  }
                ],
                progress: 58
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add data summary card
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'data-summary',
            component: {
              id: 'data-summary',
              type: 'card',
              data: {
                title: 'Data Analysis Summary',
                subtitle: 'Sales Data Processing Results',
                content: 'Successfully processed 12,543 sales records from the last 12 months. Found key trends in customer behavior and revenue patterns.',
                icon: '📊',
                status: 'success',
                collapsible: true,
                actions: [
                  { label: 'View Details', action: 'view-details', variant: 'primary' },
                  { label: 'Export Data', action: 'export', variant: 'secondary' }
                ]
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add markdown text with insights
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'insights-text',
            component: {
              id: 'insights-text',
              type: 'text',
              data: {
                content: `# Dashboard Insights\n\nBased on the data analysis, here are the key findings:\n\n## Revenue Trends\n- **23% increase** in Q4 sales compared to Q3\n- Peak sales month: **December** ($1.2M)\n- Lowest performing month: **February** ($680K)\n\n## Customer Behavior\n- Average order value: **$156.78**\n- Customer retention rate: **89.3%**\n- Most popular product category: **Electronics**\n\n## Recommendations\n1. **Focus marketing efforts** on February to boost sales\n2. **Expand electronics inventory** for peak seasons\n3. **Implement loyalty program** to maintain high retention\n\n*Dashboard generation is 65% complete. Estimated completion: 3-4 minutes.*`,
                markdown: true
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

        }, 1500);

        // Update components after more time
        setTimeout(() => {
          const componentManager = chat.componentManager;
          if (!componentManager) return;

          // Update status indicator to success
          componentManager.processUpdate({
            operation: 'update',
            target_id: 'status-chart-generation',
            updates: {
              status: 'success',
              message: 'Revenue charts generated successfully',
              pulse: false
            },
            timestamp: new Date().toISOString()
          });

          // Update progress bar
          componentManager.processUpdate({
            operation: 'update',
            target_id: 'chart-progress',
            updates: {
              progress: 100,
              status: 'success',
              detail: 'All charts generated successfully!',
              animated: false
            },
            timestamp: new Date().toISOString()
          });

          // Add success notification
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'success-notification',
            component: {
              id: 'success-notification',
              type: 'notification',
              data: {
                title: 'Charts Ready!',
                message: 'Your revenue trend charts have been generated and are ready for review.',
                level: 'success',
                dismissible: true,
                actions: [
                  { label: 'View Charts', action: 'view-charts', variant: 'primary' },
                  { label: 'Continue', action: 'continue', variant: 'secondary' }
                ]
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Update progress tracker
          tracker.updateItem(taskId3, 'completed');
          tracker.updateItem(taskId4, 'in_progress', 'Adding PDF export functionality...');
          chat.setStatus('working', 'Adding export capabilities...', 'Step 4 of 5');

        }, 4000);
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const WithToolExecutionComponents: Story = {
  args: {
    title: 'Vanna AI Agent - Tool Execution',
    placeholder: 'Ask me to run commands or execute tools...',
    disabled: false,
    showProgress: true,
    maxAutonomy: true,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      const tracker = chat?.getProgressTracker();

      if (chat && tracker) {
        // Initial conversation about running tools
        chat.addMessage('Run a data analysis script and deploy the results to production', 'user');
        chat.addMessage('I\'ll execute the data analysis script and handle the deployment. Let me run the necessary tools and show you the execution details.', 'assistant');

        // Add progress tasks
        const taskId1 = tracker.addItem('Run data analysis script', 'Executing Python analysis tools');
        const taskId2 = tracker.addItem('Process results', 'Formatting and validating output');
        const taskId3 = tracker.addItem('Deploy to production', 'Uploading to production server');

        tracker.updateItem(taskId1, 'in_progress', 'Running analysis.py...');

        chat.setStatus('working', 'Executing data analysis tools...', 'Max autonomy enabled');

        // Add tool execution components after a delay
        setTimeout(() => {
          const componentManager = chat.componentManager;
          if (!componentManager) return;

          // Add tool execution component for data analysis
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'tool-analysis',
            component: {
              id: 'tool-analysis',
              type: 'tool_execution',
              data: {
                tool_name: 'Python Script Executor',
                status: 'running',
                progress: 45,
                duration: '2.3s',
                arguments: {
                  script: 'analysis.py',
                  dataset: 'sales_data_2024.csv',
                  output_format: 'json',
                  verbose: true
                },
                result: null,
                error: null,
                logs: [
                  { timestamp: '14:32:01', level: 'INFO', message: 'Loading dataset: sales_data_2024.csv' },
                  { timestamp: '14:32:02', level: 'INFO', message: 'Found 12,543 records' },
                  { timestamp: '14:32:03', level: 'INFO', message: 'Running correlation analysis...' },
                  { timestamp: '14:32:04', level: 'INFO', message: 'Processing revenue trends...' }
                ]
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add warning notification about data processing
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'warning-notification',
            component: {
              id: 'warning-notification',
              type: 'notification',
              data: {
                title: 'Large Dataset Detected',
                message: 'Processing 12,543 records. This may take a few minutes to complete.',
                level: 'warning',
                dismissible: true,
                actions: [
                  { label: 'Monitor Progress', action: 'monitor', variant: 'primary' }
                ]
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

        }, 1500);

        // Complete first tool and start second
        setTimeout(() => {
          const componentManager = chat.componentManager;
          if (!componentManager) return;

          // Update first tool to completed
          componentManager.processUpdate({
            operation: 'update',
            target_id: 'tool-analysis',
            updates: {
              status: 'completed',
              progress: 100,
              duration: '4.7s',
              result: JSON.stringify({
                total_revenue: 1847259.32,
                avg_order_value: 156.78,
                top_product: 'Electronics',
                growth_rate: 0.23,
                recommendations: ['Expand Q4 marketing', 'Focus on electronics inventory']
              }, null, 2),
              logs: [
                { timestamp: '14:32:01', level: 'INFO', message: 'Loading dataset: sales_data_2024.csv' },
                { timestamp: '14:32:02', level: 'INFO', message: 'Found 12,543 records' },
                { timestamp: '14:32:03', level: 'INFO', message: 'Running correlation analysis...' },
                { timestamp: '14:32:04', level: 'INFO', message: 'Processing revenue trends...' },
                { timestamp: '14:32:06', level: 'INFO', message: 'Analysis complete! Generated insights.' },
                { timestamp: '14:32:06', level: 'INFO', message: 'Output saved to results.json' }
              ]
            },
            timestamp: new Date().toISOString()
          });

          // Update progress tracker
          tracker.updateItem(taskId1, 'completed');
          tracker.updateItem(taskId2, 'in_progress', 'Validating analysis results...');

          // Add second tool execution for deployment
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'tool-deploy',
            component: {
              id: 'tool-deploy',
              type: 'tool_execution',
              data: {
                tool_name: 'Production Deployer',
                status: 'running',
                progress: 20,
                duration: '1.2s',
                arguments: {
                  source: 'results.json',
                  target: 'prod-server-01',
                  backup: true,
                  validate: true
                },
                result: null,
                error: null,
                logs: [
                  { timestamp: '14:32:08', level: 'INFO', message: 'Connecting to prod-server-01...' },
                  { timestamp: '14:32:09', level: 'INFO', message: 'Creating backup of existing data...' }
                ]
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

        }, 4000);

        // Complete deployment
        setTimeout(() => {
          const componentManager = chat.componentManager;
          if (!componentManager) return;

          // Complete second tool
          componentManager.processUpdate({
            operation: 'update',
            target_id: 'tool-deploy',
            updates: {
              status: 'completed',
              progress: 100,
              duration: '6.1s',
              result: 'Deployment successful! Results available at: https://dashboard.company.com/sales-analysis',
              logs: [
                { timestamp: '14:32:08', level: 'INFO', message: 'Connecting to prod-server-01...' },
                { timestamp: '14:32:09', level: 'INFO', message: 'Creating backup of existing data...' },
                { timestamp: '14:32:11', level: 'INFO', message: 'Uploading results.json...' },
                { timestamp: '14:32:13', level: 'INFO', message: 'Validating deployment...' },
                { timestamp: '14:32:14', level: 'INFO', message: 'Deployment successful!' }
              ]
            },
            timestamp: new Date().toISOString()
          });

          // Update progress tracker
          tracker.updateItem(taskId2, 'completed');
          tracker.updateItem(taskId3, 'completed');

          // Add success notification
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'deploy-success',
            component: {
              id: 'deploy-success',
              type: 'notification',
              data: {
                title: 'Deployment Complete!',
                message: 'Data analysis results have been successfully deployed to production. Dashboard is now live.',
                level: 'success',
                dismissible: true,
                actions: [
                  { label: 'View Dashboard', action: 'view-dashboard', variant: 'primary' },
                  { label: 'Download Report', action: 'download', variant: 'secondary' }
                ]
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Add final status indicator
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'final-status',
            component: {
              id: 'final-status',
              type: 'status_indicator',
              data: {
                status: 'success',
                message: 'All tools executed successfully',
                pulse: false
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Update chat status
          chat.setStatus('idle', 'All tasks completed successfully', 'Ready for next request');

        }, 7000);
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const WithChart: Story = {
  args: {
    title: 'Vanna AI Agent - Chart Display',
    placeholder: 'Ask me to analyze data...',
    disabled: false,
    showProgress: true,
    maxAutonomy: false,
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;
      const tracker = chat?.getProgressTracker();

      if (chat && tracker) {
        // Initial conversation
        chat.addMessage('Show me the top 10 artists by sales', 'user');
        chat.addMessage('I\'ll analyze the sales data and create a visualization for you.', 'assistant');

        // Add progress tasks
        const taskId1 = tracker.addItem('Query sales database', 'Fetching artist sales data');
        const taskId2 = tracker.addItem('Process results', 'Calculating total sales per artist');
        const taskId3 = tracker.addItem('Create visualization', 'Generating bar chart');

        tracker.updateItem(taskId1, 'completed');
        tracker.updateItem(taskId2, 'completed');
        tracker.updateItem(taskId3, 'in_progress', 'Rendering Plotly chart...');

        chat.setStatus('working', 'Creating visualization...', 'Step 3 of 3');

        // Add chart component after a delay
        setTimeout(() => {
          const componentManager = chat.componentManager;
          if (!componentManager) return;

          // Add the chart component with the data you provided
          componentManager.processUpdate({
            operation: 'create',
            target_id: '8e275121-e3f4-4b99-87f2-fabc3ef66216',
            component: {
              id: '8e275121-e3f4-4b99-87f2-fabc3ef66216',
              type: 'chart',
              data: {
                data: [
                  {
                    hovertemplate: 'artist=%{x}<br>total_sales=%{y}<extra></extra>',
                    legendgroup: '',
                    marker: {
                      color: '#636efa',
                      pattern: {
                        shape: ''
                      }
                    },
                    name: '',
                    orientation: 'v',
                    showlegend: false,
                    textposition: 'auto',
                    x: [
                      'Deep Purple',
                      'Eric Clapton',
                      'Faith No More',
                      'Iron Maiden',
                      'Led Zeppelin',
                      'Lost',
                      'Metallica',
                      'Os Paralamas Do Sucesso',
                      'The Office',
                      'U2'
                    ],
                    xaxis: 'x',
                    y: [22.5, 19.8, 27.72, 138.6, 86.13, 81.59, 90.09, 41.49, 30, 105.93],
                    yaxis: 'y',
                    type: 'bar'
                  }
                ],
                layout: {
                  template: {
                    data: {
                      bar: [
                        {
                          error_x: {
                            color: '#2a3f5f'
                          },
                          error_y: {
                            color: '#2a3f5f'
                          },
                          marker: {
                            line: {
                              color: '#E5ECF6',
                              width: 0.5
                            },
                            pattern: {
                              fillmode: 'overlay',
                              size: 10,
                              solidity: 0.2
                            }
                          },
                          type: 'bar'
                        }
                      ]
                    },
                    layout: {
                      font: {
                        color: '#2a3f5f'
                      },
                      xaxis: {
                        gridcolor: 'white',
                        linecolor: 'white',
                        ticks: '',
                        title: {
                          standoff: 15
                        },
                        zerolinecolor: 'white',
                        automargin: true,
                        zerolinewidth: 2
                      },
                      yaxis: {
                        gridcolor: 'white',
                        linecolor: 'white',
                        ticks: '',
                        title: {
                          standoff: 15
                        },
                        zerolinecolor: 'white',
                        automargin: true,
                        zerolinewidth: 2
                      }
                    }
                  },
                  xaxis: {
                    anchor: 'y',
                    domain: [0.0, 1.0],
                    title: {
                      text: 'artist'
                    }
                  },
                  yaxis: {
                    anchor: 'x',
                    domain: [0.0, 1.0],
                    title: {
                      text: 'total_sales'
                    }
                  },
                  legend: {
                    tracegroupgap: 0
                  },
                  title: {
                    text: 'Top 10 Artists by Sales'
                  },
                  barmode: 'relative',
                  font: {
                    color: '#1f2937'
                  },
                  paper_bgcolor: 'white',
                  plot_bgcolor: 'white',
                  autosize: true
                },
                chart_type: 'plotly',
                title: 'Top 10 Artists by Sales',
                width: null,
                height: null,
                config: {
                  data_shape: {
                    rows: 10,
                    columns: 2
                  },
                  source_file: 'query_results_f42c1599.csv'
                }
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

          // Update progress tracker
          tracker.updateItem(taskId3, 'completed');
          chat.setStatus('idle', 'Visualization complete', 'Ready for next query');

          // Add text summary
          componentManager.processUpdate({
            operation: 'create',
            target_id: 'chart-summary',
            component: {
              id: 'chart-summary',
              type: 'text',
              data: {
                content: `Created visualization from 'query_results_f42c1599.csv' (10 rows, 2 columns).

**Top Artists:**
- Faith No More leads with 138.6 in total sales
- U2 follows with 105.93
- Metallica with 90.09

The chart shows the distribution of sales across the top 10 artists.`,
                markdown: true
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });

        }, 1500);
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

export const WithButtons: Story = {
  args: {
    title: 'Vanna AI Agent - Button Components',
    placeholder: 'Click buttons to send messages...',
    disabled: false,
    showProgress: false,
    maxAutonomy: false,
    theme: 'dark',
  },
  render: (args) => {
    setTimeout(() => {
      const chat = document.querySelector('vanna-chat') as any;

      if (chat) {
        // Initial conversation
        chat.addMessage('Show me some button options', 'user');
        chat.addMessage('Here are some interactive buttons. Click any button to send a message with its label wrapped in square brackets.', 'assistant');

        const componentManager = chat.componentManager;
        if (!componentManager) return;

        // Add a single button example
        componentManager.processUpdate({
          operation: 'create',
          target_id: 'single-button-1',
          component: {
            id: 'single-button-1',
            type: 'button',
            data: {
              label: 'Okay',
              action: 'okay',
              variant: 'primary',
              size: 'medium',
            },
            layout: { position: 'append', size: {}, z_index: 0, classes: [] },
            theme: {},
            lifecycle: 'create'
          },
          timestamp: new Date().toISOString()
        });

        // Add button group with choices
        componentManager.processUpdate({
          operation: 'create',
          target_id: 'choice-group',
          component: {
            id: 'choice-group',
            type: 'button_group',
            data: {
              buttons: [
                {
                  label: 'Yes',
                  action: 'yes',
                  variant: 'success',
                  icon: '✓',
                },
                {
                  label: 'No',
                  action: 'no',
                  variant: 'error',
                  icon: '✗',
                },
                {
                  label: 'Maybe',
                  action: 'maybe',
                  variant: 'secondary',
                },
              ],
              orientation: 'horizontal',
              spacing: 'medium',
              align: 'left',
            },
            layout: { position: 'append', size: {}, z_index: 0, classes: [] },
            theme: {},
            lifecycle: 'create'
          },
          timestamp: new Date().toISOString()
        });

        // Add more single buttons with different variants
        const singleButtons = [
          { label: 'Continue', variant: 'primary', icon: '→', icon_position: 'right' },
          { label: 'Save Draft', variant: 'secondary', icon: '💾', icon_position: 'left' },
          { label: 'Delete', variant: 'error', icon: '🗑️' },
          { label: 'Cancel', variant: 'ghost' },
        ];

        singleButtons.forEach((btnData, index) => {
          componentManager.processUpdate({
            operation: 'create',
            target_id: `single-button-${index + 2}`,
            component: {
              id: `single-button-${index + 2}`,
              type: 'button',
              data: {
                ...btnData,
                action: btnData.label.toLowerCase().replace(' ', '_'),
                size: 'medium',
              },
              layout: { position: 'append', size: {}, z_index: 0, classes: [] },
              theme: {},
              lifecycle: 'create'
            },
            timestamp: new Date().toISOString()
          });
        });

        // Add a vertical button group
        componentManager.processUpdate({
          operation: 'create',
          target_id: 'option-group',
          component: {
            id: 'option-group',
            type: 'button_group',
            data: {
              buttons: [
                { label: 'Option A', action: 'option_a', variant: 'secondary' },
                { label: 'Option B', action: 'option_b', variant: 'secondary' },
                { label: 'Option C', action: 'option_c', variant: 'secondary' },
              ],
              orientation: 'vertical',
              spacing: 'small',
              align: 'left',
            },
            layout: { position: 'append', size: {}, z_index: 0, classes: [] },
            theme: {},
            lifecycle: 'create'
          },
          timestamp: new Date().toISOString()
        });

        // Listen for message-sent events to show feedback
        chat.addEventListener('message-sent', (e: CustomEvent) => {
          console.log('Message sent from button:', e.detail.message);
        });
      }
    }, 100);

    return html`
      <div style="height: 100vh; display: flex; align-items: center; justify-content: center; padding: 20px; background: ${args.theme === 'light' ? '#f5f7fa' : 'rgb(11, 15, 25)'};">
        <vanna-chat
          .title=${args.title}
          .placeholder=${args.placeholder}
          .disabled=${args.disabled}
          .showProgress=${args.showProgress}
          .maxAutonomy=${args.maxAutonomy}
          theme=${args.theme}>
        </vanna-chat>
      </div>
    `;
  },
};

================================================
FILE: frontends/webcomponent/src/components/vanna-chat.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property, state } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';
import { VannaApiClient, ChatStreamChunk } from '../services/api-client.js';
import { ComponentManager, RichComponent } from './rich-component-system.js';
import './vanna-status-bar.js';
import './vanna-progress-tracker.js';
import './rich-card.js';
import './rich-task-list.js';
import './rich-progress-bar.js';
import './plotly-chart.js';

@customElement('vanna-chat')
export class VannaChat extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      *, *::before, *::after {
        box-sizing: border-box;
      }

      :host {
        display: block;
        font-family: var(--vanna-font-family-default);
        --chat-primary: var(--vanna-accent-primary-default);
        --chat-primary-stronger: var(--vanna-accent-primary-stronger);
        --chat-primary-foreground: rgb(255, 255, 255);
        --chat-accent-soft: var(--vanna-accent-primary-subtle);
        --chat-outline: var(--vanna-outline-default);
        --chat-surface: var(--vanna-background-root);
        --chat-muted: var(--vanna-background-default);
        --chat-muted-stronger: var(--vanna-background-higher);
        max-width: 1024px;
        margin: 0 auto;
        background: var(--vanna-background-root);
        border: 1px solid var(--vanna-outline-dimmer);
        border-radius: var(--vanna-border-radius-2xl);
        box-shadow: var(--vanna-shadow-xl);
        overflow: hidden;
        transition: box-shadow var(--vanna-duration-300) ease, transform var(--vanna-duration-300) ease;
        position: relative;
      }

      :host(:hover) {
        box-shadow: var(--vanna-shadow-2xl);
        transform: translateY(-2px);
      }

      :host([theme="dark"]) {
        --chat-primary: var(--vanna-accent-primary-default);
        --chat-primary-stronger: var(--vanna-accent-primary-stronger);
        --chat-primary-foreground: rgb(255, 255, 255);
        --chat-accent-soft: var(--vanna-accent-primary-subtle);
        --chat-outline: var(--vanna-outline-default);
        --chat-surface: var(--vanna-background-higher);
        --chat-muted: var(--vanna-background-default);
        --chat-muted-stronger: var(--vanna-background-highest);
        background: var(--vanna-background-higher);
        border-color: var(--vanna-outline-default);
      }

      :host(.maximized) {
        position: fixed;
        top: var(--vanna-space-6);
        left: var(--vanna-space-6);
        right: var(--vanna-space-6);
        bottom: var(--vanna-space-6);
        max-width: none;
        width: auto;
        margin: 0;
        z-index: var(--vanna-z-modal);
        border-radius: var(--vanna-border-radius-xl);
        transform: none;
        box-shadow: var(--vanna-shadow-2xl);
      }

      :host(.maximized):hover {
        transform: none;
      }

      :host(.minimized) {
        position: fixed !important;
        bottom: var(--vanna-space-6) !important;
        right: var(--vanna-space-6) !important;
        width: 64px !important;
        height: 64px !important;
        max-width: none !important;
        margin: 0 !important;
        z-index: var(--vanna-z-modal) !important;
        border-radius: var(--vanna-border-radius-full) !important;
        cursor: pointer !important;
        background: linear-gradient(135deg, var(--chat-primary-stronger), var(--chat-primary)) !important;
        border: 2px solid rgba(255, 255, 255, 0.9) !important;
        box-shadow: var(--vanna-shadow-xl) !important;
        overflow: hidden !important;
      }

      :host(.minimized):hover {
        transform: scale(1.05);
        box-shadow: var(--vanna-shadow-2xl) !important;
      }

      :host(.minimized) .chat-layout {
        display: none;
      }

      .minimized-icon {
        display: none;
      }

      :host(.minimized) .minimized-icon {
        display: flex;
        align-items: center;
        justify-content: center;
        width: 100%;
        height: 100%;
        color: var(--chat-primary-foreground);
        font-size: 24px;
        transition: transform var(--vanna-duration-200) ease;
      }

      :host(.minimized) .minimized-icon:hover {
        transform: scale(1.1);
      }

      :host(.minimized) .minimized-icon svg {
        filter: drop-shadow(0 2px 4px rgba(0, 0, 0, 0.3));
      }

      .chat-layout {
        display: grid;
        grid-template-columns: minmax(0, 1fr) 300px;
        height: 600px;
        max-height: 80vh;
        background: var(--chat-muted);
      }

      :host(.maximized) .chat-layout {
        height: calc(100vh - 48px);
        max-height: calc(100vh - 48px);
      }

      .chat-layout.compact {
        grid-template-columns: 1fr;
      }

      .chat-main {
        display: flex;
        flex-direction: column;
        border-right: 1px solid var(--chat-outline);
        background: var(--chat-surface);
        min-height: 0;
      }

      .chat-layout.compact .chat-main {
        border-right: none;
      }

      .chat-header {
        padding: var(--vanna-space-6) var(--vanna-space-7);
        background: linear-gradient(135deg, var(--chat-primary) 0%, var(--chat-primary-stronger) 100%);
        border-bottom: 1px solid rgba(255, 255, 255, 0.2);
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-4);
        color: var(--chat-primary-foreground);
        position: relative;
        overflow: hidden;
      }

      .chat-header::before {
        content: '';
        position: absolute;
        top: -50%;
        right: -50%;
        width: 100%;
        height: 200%;
        background: radial-gradient(circle, rgba(255, 255, 255, 0.15) 0%, transparent 70%);
        opacity: 0.6;
        pointer-events: none;
      }

      :host([theme="dark"]) .chat-header {
        border-bottom-color: rgba(255, 255, 255, 0.1);
      }

      .header-top {
        position: relative;
        z-index: 1;
        display: flex;
        align-items: center;
        gap: var(--vanna-space-4);
        width: 100%;
      }

      .header-left {
        display: flex;
        align-items: center;
        gap: var(--vanna-space-4);
        min-width: 0;
        flex: 1;
      }

      .header-top-actions {
        display: inline-flex;
        align-items: center;
        gap: var(--vanna-space-2);
        margin-left: auto;
      }

      .chat-avatar {
        width: 44px;
        height: 44px;
        border-radius: var(--vanna-border-radius-lg);
        background: rgba(255, 255, 255, 0.2);
        backdrop-filter: blur(10px);
        display: grid;
        place-items: center;
        font-weight: 600;
        font-size: 16px;
        letter-spacing: 0.02em;
        color: var(--chat-primary-foreground);
        border: 1px solid rgba(255, 255, 255, 0.3);
      }

      .header-text {
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-1);
        min-width: 0;
      }

      .chat-title {
        margin: 0;
        font-size: 18px;
        font-weight: 600;
        letter-spacing: -0.01em;
        color: var(--chat-primary-foreground);
      }

      .chat-subtitle {
        font-size: 13px;
        letter-spacing: 0.01em;
        opacity: 0.9;
        font-weight: 400;
      }

      :host([theme="dark"]) .chat-subtitle {
        opacity: 0.78;
      }

      .window-controls {
        display: inline-flex;
        gap: var(--vanna-space-2);
      }

      .window-control-btn {
        width: 32px;
        height: 32px;
        border-radius: var(--vanna-border-radius-lg);
        border: 1px solid rgba(255, 255, 255, 0.15);
        background: rgba(255, 255, 255, 0.1);
        color: var(--chat-primary-foreground);
        cursor: pointer;
        display: inline-flex;
        align-items: center;
        justify-content: center;
        transition: all var(--vanna-duration-200) ease;
        backdrop-filter: blur(8px);
        position: relative;
        overflow: hidden;
      }

      .window-control-btn::before {
        content: '';
        position: absolute;
        inset: 0;
        background: linear-gradient(135deg, rgba(255, 255, 255, 0.2), transparent);
        opacity: 0;
        transition: opacity var(--vanna-duration-200) ease;
      }

      .window-control-btn:hover {
        transform: translateY(-1px) scale(1.05);
        background: rgba(255, 255, 255, 0.2);
        box-shadow: 
          0 8px 25px -8px rgba(0, 0, 0, 0.3),
          0 0 0 1px rgba(255, 255, 255, 0.2);
        border-color: rgba(255, 255, 255, 0.3);
      }

      .window-control-btn:hover::before {
        opacity: 1;
      }

      .window-control-btn:active {
        transform: translateY(0) scale(0.95);
      }

      .window-control-btn.minimize:hover {
        background: rgba(255, 193, 7, 0.2);
        color: #ffc107;
        box-shadow: 
          0 8px 25px -8px rgba(255, 193, 7, 0.4),
          0 0 0 1px rgba(255, 193, 7, 0.3);
      }

      .window-control-btn.maximize:hover,
      .window-control-btn.restore:hover {
        background: rgba(40, 167, 69, 0.2);
        color: #28a745;
        box-shadow: 
          0 8px 25px -8px rgba(40, 167, 69, 0.4),
          0 0 0 1px rgba(40, 167, 69, 0.3);
      }

      .window-control-btn svg {
        width: 16px;
        height: 16px;
        transition: transform var(--vanna-duration-150) ease;
      }

      .window-control-btn:hover svg {
        transform: scale(1.1);
      }

      :host([theme="dark"]) .window-control-btn {
        border-color: rgba(255, 255, 255, 0.1);
        background: rgba(255, 255, 255, 0.05);
      }

      :host([theme="dark"]) .window-control-btn:hover {
        background: rgba(255, 255, 255, 0.15);
        border-color: rgba(255, 255, 255, 0.25);
      }

      .chat-messages {
        flex: 1;
        overflow-y: auto;
        overflow-x: hidden;
        padding: var(--vanna-space-6) var(--vanna-space-6) var(--vanna-space-5);
        background: linear-gradient(180deg, var(--chat-muted) 0%, var(--chat-surface) 70%);
        scroll-behavior: smooth;
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-4);
        min-height: 0;
        max-height: 100%;
        position: relative;
      }

      .chat-messages::-webkit-scrollbar {
        width: 6px;
      }

      .chat-messages::-webkit-scrollbar-track {
        background: transparent;
      }

      .chat-messages::-webkit-scrollbar-thumb {
        background: var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-full);
        border: 1px solid var(--vanna-background-root);
      }

      .chat-messages::-webkit-scrollbar-thumb:hover {
        background: var(--vanna-outline-hover);
      }

      :host([theme="dark"]) .chat-messages {
        background: radial-gradient(circle at top, rgba(99, 102, 241, 0.12), transparent 55%), var(--chat-surface);
      }

      :host([theme="dark"]) .chat-messages::-webkit-scrollbar-thumb {
        background: var(--vanna-outline-default);
        border-color: var(--vanna-background-higher);
      }

      /* Scroll indicator when there's content above */
      .chat-messages::before {
        content: '';
        position: sticky;
        top: 0;
        display: block;
        height: 1px;
        background: linear-gradient(90deg, transparent, var(--vanna-accent-primary-default), transparent);
        opacity: 0;
        transition: opacity var(--vanna-duration-300) ease;
        z-index: 10;
        margin: 0 var(--vanna-space-4) var(--vanna-space-2);
      }

      .chat-messages.has-scroll::before {
        opacity: 0.5;
      }

      .rich-components-container {
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-4);
      }

      .rich-component-wrapper {
        margin: var(--vanna-space-2) 0;
        animation: fade-in-up 0.3s ease-out;
      }

      .unknown-component {
        background: var(--vanna-background-higher);
        border: 1px solid var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-md);
        padding: var(--vanna-space-4);
        font-family: var(--vanna-font-family-mono);
        font-size: 12px;
      }

      .unknown-component p {
        margin: 0 0 var(--vanna-space-2) 0;
        color: var(--vanna-foreground-dimmer);
      }

      .unknown-component pre {
        margin: 0;
        color: var(--vanna-foreground-dimmest);
        overflow-x: auto;
      }

      .chat-input-area {
        padding: var(--vanna-space-5) var(--vanna-space-6) var(--vanna-space-6);
        background: var(--chat-surface);
        border-top: 1px solid var(--chat-outline);
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-4);
        flex-shrink: 0; /* Prevent input area from shrinking */
      }

      :host([theme="dark"]) .chat-input-area {
        border-top-color: rgba(148, 163, 184, 0.22);
      }

      .chat-input-container {
        display: flex;
        align-items: center;
        gap: var(--vanna-space-2);
        padding: 6px 8px 6px 18px;
        border-radius: 999px;
        background: var(--chat-muted);
        border: 1px solid var(--chat-muted-stronger);
        box-shadow: inset 0 1px 0 rgba(255, 255, 255, 0.6);
        transition: border-color var(--vanna-duration-200) ease, box-shadow var(--vanna-duration-200) ease, background var(--vanna-duration-200) ease;
      }

      .chat-input-container:focus-within {
        border-color: var(--chat-primary);
        box-shadow: 0 0 0 1px rgba(99, 102, 241, 0.35), inset 0 1px 0 rgba(255, 255, 255, 0.85);
        background: rgba(255, 255, 255, 0.95);
      }

      :host([theme="dark"]) .chat-input-container {
        background: rgba(15, 23, 42, 0.65);
        border-color: rgba(100, 116, 139, 0.45);
        box-shadow: inset 0 1px 0 rgba(148, 163, 184, 0.18);
      }

      :host([theme="dark"]) .chat-input-container:focus-within {
        border-color: rgba(129, 140, 248, 0.55);
        box-shadow: 0 0 0 1px rgba(129, 140, 248, 0.45), inset 0 1px 0 rgba(148, 163, 184, 0.25);
        background: rgba(30, 41, 59, 0.88);
      }

      .message-input {
        flex: 1;
        border: none;
        background: transparent;
        font-size: 15px;
        font-family: var(--vanna-font-family-default);
        line-height: 1.5;
        color: var(--vanna-foreground-default);
        resize: none;
        min-height: 48px;
        max-height: 140px;
        padding: 12px 0;
        outline: none;
      }

      :host([theme="dark"]) .message-input {
        color: rgba(226, 232, 240, 0.95);
      }

      .message-input::placeholder {
        color: rgba(71, 85, 105, 0.8);
      }

      :host([theme="dark"]) .message-input::placeholder {
        color: rgba(148, 163, 184, 0.65);
      }

      .message-input:focus {
        outline: none;
      }

      .message-input:disabled {
        color: rgba(148, 163, 184, 0.65);
        cursor: not-allowed;
      }

      :host([theme="dark"]) .message-input:disabled {
        color: rgba(100, 116, 139, 0.55);
      }

      .send-button {
        width: 48px;
        height: 48px;
        border-radius: 999px;
        border: none;
        background: linear-gradient(135deg, var(--chat-primary-stronger), var(--chat-primary));
        color: var(--chat-primary-foreground);
        display: inline-flex;
        align-items: center;
        justify-content: center;
        cursor: pointer;
        transition: transform var(--vanna-duration-200) ease, box-shadow var(--vanna-duration-200) ease, filter var(--vanna-duration-200) ease;
        box-shadow: 0 18px 38px -24px rgba(79, 70, 229, 0.8);
      }

      .send-button:hover {
        transform: translateY(-1px) scale(1.02);
        box-shadow: 0 25px 45px -24px rgba(79, 70, 229, 0.85);
      }

      .send-button:active {
        transform: translateY(0) scale(0.98);
      }

      .send-button:disabled {
        background: rgba(148, 163, 184, 0.35);
        color: rgba(71, 85, 105, 0.7);
        cursor: not-allowed;
        transform: none;
        box-shadow: none;
      }

      .send-button svg {
        width: 18px;
        height: 18px;
      }

      .sidebar {
        background: linear-gradient(180deg, rgba(99, 102, 241, 0.08) 0%, rgba(15, 23, 42, 0.02) 100%);
        padding: var(--vanna-space-6);
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-4);
        overflow-y: auto;
        overflow-x: hidden;
        min-height: 0;
      }

      .sidebar::-webkit-scrollbar {
        width: 6px;
      }

      .sidebar::-webkit-scrollbar-track {
        background: transparent;
      }

      .sidebar::-webkit-scrollbar-thumb {
        background: var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-full);
      }

      :host([theme="dark"]) .sidebar {
        background: linear-gradient(180deg, rgba(79, 70, 229, 0.22) 0%, rgba(15, 23, 42, 0.45) 100%);
      }

      .empty-state {
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        text-align: center;
        color: var(--vanna-foreground-dimmer);
        padding: var(--vanna-space-12) var(--vanna-space-8);
        margin: var(--vanna-space-8) var(--vanna-space-6);
        font-size: 15px;
        font-weight: 500;
        line-height: 1.6;
        background: linear-gradient(135deg, 
          rgba(255, 255, 255, 0.95) 0%, 
          rgba(248, 250, 252, 0.9) 50%,
          rgba(241, 245, 249, 0.85) 100%);
        border-radius: var(--vanna-border-radius-2xl);
        border: 2px dashed var(--vanna-accent-primary-default);
        box-shadow: 
          var(--vanna-shadow-sm),
          inset 0 1px 0 rgba(255, 255, 255, 0.8);
        backdrop-filter: blur(8px);
        transition: all var(--vanna-duration-300) ease;
      }

      .empty-state:hover {
        border-color: var(--vanna-accent-primary-stronger);
        transform: translateY(-2px);
        box-shadow: 
          var(--vanna-shadow-lg),
          inset 0 1px 0 rgba(255, 255, 255, 0.9);
      }

      :host([theme="dark"]) .empty-state {
        color: var(--vanna-foreground-dimmer);
        background: linear-gradient(135deg, 
          rgba(24, 29, 39, 0.95) 0%, 
          rgba(31, 39, 51, 0.9) 50%,
          rgba(17, 21, 28, 0.85) 100%);
        border-color: var(--vanna-accent-primary-default);
        box-shadow: 
          var(--vanna-shadow-md),
          inset 0 1px 0 rgba(129, 140, 248, 0.2);
      }

      :host([theme="dark"]) .empty-state:hover {
        border-color: var(--vanna-accent-primary-hover);
        box-shadow: 
          var(--vanna-shadow-xl),
          inset 0 1px 0 rgba(129, 140, 248, 0.3);
      }

      .empty-state-icon {
        width: 64px;
        height: 64px;
        margin: 0 auto var(--vanna-space-6);
        opacity: 0.7;
        color: var(--vanna-accent-primary-default);
        filter: drop-shadow(0 2px 4px rgba(79, 70, 229, 0.2));
      }

      .empty-state-text {
        font-size: 16px;
        font-weight: 600;
        color: var(--vanna-foreground-default);
        margin-bottom: var(--vanna-space-2);
      }

      .empty-state-subtitle {
        font-size: 14px;
        color: var(--vanna-foreground-dimmest);
        opacity: 0.8;
        font-weight: 400;
      }

      @media (max-width: 880px) {
        .chat-layout {
          grid-template-columns: 1fr;
          height: min(600px, 85vh);
          max-height: 85vh;
        }

        .sidebar {
          display: none;
        }

        .chat-main {
          border-right: none;
        }
      }

      @media (max-width: 600px) {
        :host {
          border-radius: var(--vanna-border-radius-xl);
        }

        .chat-layout {
          height: min(500px, 80vh);
          max-height: 80vh;
        }

        .chat-header {
          border-bottom-width: 0;
          padding: var(--vanna-space-5) var(--vanna-space-5) var(--vanna-space-4);
        }

        .chat-messages {
          padding: var(--vanna-space-4) var(--vanna-space-4);
        }

        .empty-state {
          padding: var(--vanna-space-10) var(--vanna-space-6);
          margin: var(--vanna-space-6) var(--vanna-space-4);
          font-size: 14px;
        }

        .empty-state-text {
          font-size: 15px;
        }

        .empty-state-icon {
          width: 56px;
          height: 56px;
          margin-bottom: var(--vanna-space-5);
        }

        .chat-input-area {
          padding: var(--vanna-space-4) var(--vanna-space-4) var(--vanna-space-5);
        }
      }
    `
  ];

  @property() title = 'Vanna AI Chat';
  @property() placeholder = 'Ask me anything...';
  @property({ type: Boolean }) disabled = false;
  @property({ type: Boolean }) showProgress = true;
  @property({ type: Boolean }) allowMinimize = true;
  @property({ reflect: true }) theme = 'light';
  @property({ attribute: 'api-base' }) apiBaseUrl = '';
  @property({ attribute: 'sse-endpoint' }) sseEndpoint = '/api/vanna/v2/chat_sse';
  @property({ attribute: 'ws-endpoint' }) wsEndpoint = '/api/vanna/v2/chat_websocket';
  @property({ attribute: 'poll-endpoint' }) pollEndpoint = '/api/vanna/v2/chat_poll';
  @property() subtitle = '';
  @property() startingState: 'normal' | 'maximized' | 'minimized' = 'normal';

  @state() private currentMessage = '';
  @state() private status: 'idle' | 'working' | 'error' | 'success' = 'idle';
  @state() private statusMessage = '';
  @state() private statusDetail = '';
  private _windowState: 'normal' | 'maximized' | 'minimized' = 'normal';

  @property({ reflect: false })
  get windowState() {
    return this._windowState;
  }

  set windowState(value: 'normal' | 'maximized' | 'minimized') {
    console.log('windowState setter called with:', value);
    console.trace('Call stack:');
    const oldValue = this._windowState;
    this._windowState = value;
    this.requestUpdate('windowState', oldValue);
  }

  private apiClient!: VannaApiClient;
  private conversationId: string;
  private componentManager: ComponentManager | null = null;
  private componentObserver: MutationObserver | null = null;

  constructor() {
    super();
    // Note: Don't create apiClient here - attributes haven't been set yet!
    // It will be created lazily in getApiClient() or firstUpdated()
    this.conversationId = this.generateId();
  }

  /**
   * Ensure API client is created/updated with current endpoint values
   */
  private ensureApiClient() {
    // Always recreate to ensure we have the latest endpoint values
    console.log('[VannaChat] Creating API client with:', {
      baseUrl: this.apiBaseUrl,
      sseEndpoint: this.sseEndpoint,
      wsEndpoint: this.wsEndpoint,
      pollEndpoint: this.pollEndpoint
    });

    this.apiClient = new VannaApiClient({
      baseUrl: this.apiBaseUrl,
      sseEndpoint: this.sseEndpoint,
      wsEndpoint: this.wsEndpoint,
      pollEndpoint: this.pollEndpoint
    });
  }

  firstUpdated() {
    // Create API client now that attributes have been set
    this.ensureApiClient();

    // Initialize component manager with rich components container (fallback)
    const richContainer = this.shadowRoot?.querySelector('.rich-components-container') as HTMLElement;
    if (richContainer) {
      this.componentManager = new ComponentManager(richContainer);
      
      // Watch for changes in the rich components container to manage empty state
      this.componentObserver = new MutationObserver(() => {
        // Update empty state visibility
        this.updateEmptyState();
      });
      
      this.componentObserver.observe(richContainer, {
        childList: true,
        subtree: true,
        attributes: false
      });
    }

    // Set initial window state from startingState property
    if (this.startingState !== 'normal') {
      this._windowState = this.startingState;
    }

    // Set initial CSS class
    this.classList.add(this._windowState);

    // Request starter UI from backend
    this.requestStarterUI();
  }

  /**
   * Request starter UI (buttons, welcome messages) from backend
   */
  private async requestStarterUI(): Promise<void> {
    try {
      const request = {
        message: "",
        conversation_id: this.conversationId,
        request_id: this.generateId(),
        metadata: {
          starter_ui_request: true
        }
      };

      // Stream the starter UI response
      await this.handleStreamingResponse(request);
    } catch (error) {
      console.error('Error requesting starter UI:', error);
      // Fail silently - starter UI is optional
    }
  }

  disconnectedCallback() {
    super.disconnectedCallback();
    
    // Clean up mutation observer
    if (this.componentObserver) {
      this.componentObserver.disconnect();
      this.componentObserver = null;
    }
  }

  updated(changedProperties: Map<string, any>) {
    super.updated(changedProperties);

    // Update host classes based on window state
    if (changedProperties.has('windowState')) {
      console.log('windowState changed to:', this._windowState);
      this.classList.remove('normal', 'maximized', 'minimized');
      this.classList.add(this._windowState);
      console.log('Applied CSS classes:', this.className);
    }
  }

  private handleInput(e: Event) {
    const input = e.target as HTMLInputElement;
    this.currentMessage = input.value;
  }

  private handleKeyPress(e: KeyboardEvent) {
    if (e.key === 'Enter' && !e.shiftKey) {
      e.preventDefault();
      this.sendMessage();
    }
  }

  /**
   * Send a message programmatically (can be called from buttons or external code)
   * Returns a Promise that resolves with success status
   */
  sendMessage(messageText?: string): Promise<boolean> {
    console.log('sendMessage called with:', messageText);

    // Use provided message or fall back to current input
    // Check if messageText is actually a string (not an event object)
    const textToSend = (typeof messageText === 'string') ? messageText : this.currentMessage;

    console.log('Will send:', textToSend);

    if (!textToSend.trim() || this.disabled) {
      console.log('Message empty or disabled, not sending');
      return Promise.resolve(false);
    }

    return this._sendMessageInternal(textToSend);
  }

  private async _sendMessageInternal(messageText: string): Promise<boolean> {
    console.log('_sendMessageInternal called with:', messageText);

    // Auto-maximize window when user sends a message (if not already maximized or minimized)
    if (this.windowState !== 'maximized' && this.windowState !== 'minimized') {
      this.maximizeWindow();
    }

    // Create user message as a rich component and send to ComponentManager
    const userRichComponent: RichComponent = {
      id: `user-message-${Date.now()}`,
      type: 'user-message',
      lifecycle: 'create',
      data: {
        content: messageText,
        sender: 'user'
      },
      children: [],
      timestamp: new Date().toISOString(),
      visible: true,
      interactive: false
    };

    // Add user message to ComponentManager for chronological ordering
    if (this.componentManager) {
      const update = {
        operation: 'create' as const,
        target_id: userRichComponent.id,
        component: userRichComponent,
        timestamp: userRichComponent.timestamp
      };
      this.componentManager.processUpdate(update);
    }

    // Update empty state after a brief delay to let ComponentManager render
    setTimeout(() => this.updateEmptyState(), 0);

    console.log('Added user message as rich component to ComponentManager:', userRichComponent);

    // Update the view
    this.requestUpdate();

    // Update status to working (initial frontend status before backend responds)
    this.setStatus('working', 'Sending message...', '');

    // Clear input only if we're sending from the input field
    if (messageText === this.currentMessage) {
      this.currentMessage = '';
      const input = this.shadowRoot?.querySelector('.message-input') as HTMLTextAreaElement;
      if (input) {
        input.value = '';
        input.style.height = 'auto';
      }
    }

    // Dispatch event for external listeners
    this.dispatchEvent(new CustomEvent('message-sent', {
      detail: { message: { content: messageText, type: 'user' } },
      bubbles: true,
      composed: true
    }));

    try {
      // Create the request
      const request = {
        message: messageText,
        conversation_id: this.conversationId,
        request_id: this.generateId(),
        metadata: {}
      };

      // Stream the response
      await this.handleStreamingResponse(request);
      return true; // Success

    } catch (error) {
      console.error('Error sending message:', error);
      this.setStatus('error', 'Failed to send message', error instanceof Error ? error.message : 'Unknown error');

      // Add error message
      this.addMessage(
        `Sorry, I encountered an error: ${error instanceof Error ? error.message : 'Unknown error'}`,
        'assistant'
      );
      return false; // Failure
    }
  }

  private getTitleInitials(): string {
    const title = (this.title || '').trim();
    if (!title) {
      return 'VA';
    }

    const parts = title.split(/\s+/).filter(Boolean);
    if (parts.length === 1) {
      return parts[0].charAt(0).toUpperCase() || 'V';
    }

    const first = parts[0].charAt(0);
    const last = parts[parts.length - 1].charAt(0);
    const initials = `${first}${last}`.toUpperCase();
    return initials || 'VA';
  }

  private minimizeWindow(e?: Event) {
    if (e) {
      e.stopPropagation();
      e.preventDefault();
    }
    console.log('minimizeWindow called, current state:', this._windowState);
    this.windowState = 'minimized';
    console.log('minimizeWindow set state to:', this._windowState);
    this.dispatchEvent(new CustomEvent('window-state-changed', {
      detail: { state: 'minimized' },
      bubbles: true,
      composed: true
    }));
  }

  private maximizeWindow(e?: Event) {
    if (e) {
      e.stopPropagation();
      e.preventDefault();
    }
    this.windowState = 'maximized';
    this.dispatchEvent(new CustomEvent('window-state-changed', {
      detail: { state: 'maximized' },
      bubbles: true,
      composed: true
    }));
  }

  private restoreWindow(e?: Event) {
    if (e) {
      e.stopPropagation();
      e.preventDefault();
    }
    this.windowState = 'normal';
    this.dispatchEvent(new CustomEvent('window-state-changed', {
      detail: { state: 'normal' },
      bubbles: true,
      composed: true
    }));
  }


  addMessage(content: string, type: 'user' | 'assistant') {
    // Create message as a rich component and send to ComponentManager
    const richComponent: RichComponent = {
      id: `${type}-message-${Date.now()}`,
      type: `${type}-message`,
      lifecycle: 'create',
      data: {
        content: content,
        sender: type
      },
      children: [],
      timestamp: new Date().toISOString(),
      visible: true,
      interactive: false
    };

    if (this.componentManager) {
      const update = {
        operation: 'create' as const,
        target_id: richComponent.id,
        component: richComponent,
        timestamp: richComponent.timestamp
      };
      this.componentManager.processUpdate(update);
    }
  }

  setStatus(status: typeof this.status, message: string, detail?: string) {
    this.status = status;
    this.statusMessage = message;
    this.statusDetail = detail || '';
  }

  clearStatus() {
    this.statusMessage = '';
    this.statusDetail = '';
    this.status = 'idle';
  }

  getProgressTracker(): HTMLElement | null {
    return this.shadowRoot?.querySelector('vanna-progress-tracker') || null;
  }

  private async handleStreamingResponse(request: any) {
    // Ensure API client exists and is up to date
    if (!this.apiClient || this.apiClient.baseUrl !== this.apiBaseUrl) {
      this.ensureApiClient();
    }

    // Note: Status bar updates are now controlled by backend via StatusBarUpdateComponent
    // Frontend only shows initial "Sending message..." status (set in _sendMessageInternal)
    // and handles connection errors below

    try {
      // Use SSE streaming by default
      const stream = this.apiClient.streamChat(request);

      for await (const chunk of stream) {
        await this.processChunk(chunk);
      }

      // Backend is responsible for final status via StatusBarUpdateComponent
      // No frontend status clearing here

    } catch (error) {
      console.warn('SSE streaming failed, falling back to polling:', error);

      try {
        // Fallback to polling - show user we're retrying
        this.setStatus('working', 'Connection issue, retrying...', 'Using fallback method');
        const response = await this.apiClient.sendPollMessage(request);

        for (const chunk of response.chunks) {
          await this.processChunk(chunk);
        }

        // Backend is responsible for final status via StatusBarUpdateComponent

      } catch (pollError) {
        // Only set error status if polling also fails (connection error)
        this.setStatus('error', 'Connection failed', 'Unable to reach server');
        throw pollError;
      }
    }
  }

  private async processChunk(chunk: ChatStreamChunk) {
    // Dispatch chunk event for external listeners
    this.dispatchEvent(new CustomEvent('chunk-received', {
      detail: { chunk },
      bubbles: true,
      composed: true
    }));

    console.log('Processing chunk:', chunk); // Debug log

    // Handle rich components via ComponentManager
    if (chunk.rich && this.componentManager) {
      console.log('Processing rich component via ComponentManager:', chunk.rich); // Debug log
      
      if (chunk.rich.id && chunk.rich.lifecycle) {
        // Standard rich component with lifecycle
        const component = chunk.rich as RichComponent;
        const update = {
          operation: chunk.rich.lifecycle as any,
          target_id: chunk.rich.id,
          component: component,
          timestamp: new Date().toISOString()
        };
        this.componentManager.processUpdate(update);
      } else if (chunk.rich.type === 'component_update') {
        // Component update format
        this.componentManager.processUpdate(chunk.rich as any);
      } else {
        // Generic rich component
        const component = chunk.rich as RichComponent;
        const update = {
          operation: 'create' as const,
          target_id: component.id || `component-${Date.now()}`,
          component: component,
          timestamp: new Date().toISOString()
        };
        this.componentManager.processUpdate(update);
      }
      
      return;
    }

    // Update progress tracker for legacy components (keep for backward compatibility)
    const progressTracker = this.getProgressTracker();
    if (progressTracker && 'addStep' in progressTracker) {
      (progressTracker as any).addStep({
        id: `chunk-${Date.now()}`,
        title: this.getChunkTitle(chunk),
        status: 'completed',
        timestamp: chunk.timestamp
      });
    }

    // Handle different chunk types (legacy components)
    const componentType = chunk.rich?.type;
    switch (componentType) {
      case 'text':
        // Text chunks are handled in the main loop
        break;

      case 'thinking':
        // Legacy: Status bar updates now handled by backend via StatusBarUpdateComponent
        // This case is kept for backward compatibility but doesn't update status
        break;

      case 'tool_execution':
        // Legacy: Status bar updates now handled by backend via StatusBarUpdateComponent
        // This case is kept for backward compatibility but doesn't update status
        break;

      case 'error':
        throw new Error(chunk.rich.data?.message || 'Unknown error from agent');

      default:
        // Handle other component types as needed
        console.log('Received chunk:', componentType, chunk.rich);
    }
  }


  private getChunkTitle(chunk: ChatStreamChunk): string {
    const componentType = chunk.rich?.type;
    switch (componentType) {
      case 'text':
        return 'Generating response';
      case 'thinking':
        return 'Thinking';
      case 'tool_execution':
        return `Tool: ${chunk.rich.data?.tool_name || 'Unknown'}`;
      default:
        return `Processing ${componentType || 'component'}`;
    }
  }

  private generateId(): string {
    return `${Date.now()}-${Math.random().toString(36).substring(2, 11)}`;
  }

  /**
   * Update the API base URL and recreate the client
   */
  updateApiBaseUrl(baseUrl: string) {
    this.apiBaseUrl = baseUrl;
    this.ensureApiClient();
  }

  /**
   * Get the API client instance for direct access
   */
  getApiClient(): VannaApiClient {
    if (!this.apiClient) {
      this.ensureApiClient();
    }
    return this.apiClient;
  }

  /**
   * Set custom headers for authentication or other purposes
   */
  setCustomHeaders(headers: Record<string, string>) {
    this.apiClient.setCustomHeaders(headers);
  }

  /**
   * Update empty state visibility based on whether there are components
   */
  private updateEmptyState() {
    const emptyState = this.shadowRoot?.querySelector('#empty-state') as HTMLElement;
    const richContainer = this.shadowRoot?.querySelector('.rich-components-container') as HTMLElement;
    
    if (emptyState && richContainer) {
      // Show empty state if rich container has no children
      const hasContent = richContainer.children.length > 0;
      emptyState.style.display = hasContent ? 'none' : 'flex';
    }
  }

  /**
   * Update scroll indicator based on scroll position
   */
  private updateScrollIndicator() {
    const messagesContainer = this.shadowRoot?.querySelector('.chat-messages');
    if (!messagesContainer) return;
    
    // Check if there's content scrolled above
    const hasScrolledContent = messagesContainer.scrollTop > 10;
    
    // Update scroll indicator class
    messagesContainer.classList.toggle('has-scroll', hasScrolledContent);
  }

  /**
   * Scroll to the top of the last message/component that was added
   * This always scrolls regardless of current scroll position
   */
  scrollToLastMessage() {
    const messagesContainer = this.shadowRoot?.querySelector('.chat-messages');
    const richContainer = this.shadowRoot?.querySelector('.rich-components-container');
    
    if (!messagesContainer || !richContainer) return;

    // Get the last child element (the most recently added component)
    const lastComponent = richContainer.lastElementChild as HTMLElement;
    if (!lastComponent) return;

    // Scroll so the top of the last component is visible
    lastComponent.scrollIntoView({ behavior: 'smooth', block: 'start' });
    
    // Update scroll indicator after scrolling
    setTimeout(() => this.updateScrollIndicator(), 100);
  }

  /**
   * Clear all messages (useful for testing)
   */
  clearMessages() {
    if (this.componentManager) {
      this.componentManager.clear();
    }
    this.updateEmptyState();
    this.requestUpdate();
  }

  /**
   * Add multiple messages at once (useful for testing scrolling)
   */
  addTestMessages(count: number = 10) {
    for (let i = 1; i <= count; i++) {
      setTimeout(() => {
        const type = i % 2 === 0 ? 'assistant' : 'user';
        const content = `This is test message number ${i}. Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.`;
        this.addMessage(content, type);
      }, i * 100); // Stagger the messages to simulate real timing
    }
  }

  render() {
    return html`
      <!-- Minimized icon - shown only when minimized via CSS and allowMinimize is true -->
      ${this.allowMinimize ? html`
        <div class="minimized-icon" @click=${this.restoreWindow}>
          <svg viewBox="0 0 24 24" fill="currentColor" width="32" height="32">
            <path d="M20 2H4c-1.1 0-2 .9-2 2v12c0 1.1.9 2 2 2h14l4 4V4c0-1.1-.9-2-2-2zm-2 12H6v-2h12v2zm0-3H6V9h12v2zm0-3H6V6h12v2z"/>
          </svg>
        </div>
      ` : ''}

      <!-- Main chat interface -->
      <div class="chat-layout ${this.showProgress ? '' : 'compact'}">
        <div class="chat-main">
          <div class="chat-header">
            <div class="header-top">
              <div class="header-left">
                <div class="chat-avatar" aria-hidden="true">${this.getTitleInitials()}</div>
                <div class="header-text">
                  <h2 class="chat-title">${this.title}</h2>
                </div>
              </div>
              <div class="header-top-actions">
                <div class="window-controls">
                  ${this.allowMinimize ? html`
                    <button
                      class="window-control-btn minimize"
                      @click=${this.minimizeWindow}
                      title="Minimize">
                      <svg viewBox="0 0 24 24" fill="currentColor">
                        <path d="M5 12h14v2H5z"/>
                      </svg>
                    </button>
                  ` : ''}
                  ${this.windowState === 'maximized' ? html`
                    <button
                      class="window-control-btn restore"
                      @click=${this.restoreWindow}
                      title="Restore">
                      <svg viewBox="0 0 24 24" fill="currentColor">
                        <path d="M8 8v2h2V8h6v6h-2v2h4V6H8zm-2 4v8h8v-2H8v-6H6z"/>
                      </svg>
                    </button>
                  ` : html`
                    <button
                      class="window-control-btn maximize"
                      @click=${this.maximizeWindow}
                      title="Maximize">
                      <svg viewBox="0 0 24 24" fill="currentColor">
                        <path d="M5 5v14h14V5H5zm2 2h10v10H7V7z"/>
                      </svg>
                    </button>
                  `}
                </div>
              </div>
            </div>
          </div>

          <div class="chat-messages">
            <!-- Empty state - shown when no components exist -->
            <div class="empty-state" id="empty-state">
              <div class="empty-state-icon">
                <svg viewBox="0 0 24 24" fill="currentColor">
                  <path d="M20 2H4c-1.1 0-2 .9-2 2v12c0 1.1.9 2 2 2h14l4 4V4c0-1.1-.9-2-2-2zm-2 12H6v-2h12v2zm0-3H6V9h12v2zm0-3H6V6h12v2z"/>
                </svg>
              </div>
              <div class="empty-state-text">Start a conversation</div>
              <div class="empty-state-subtitle">Type your message below to begin chatting</div>
            </div>

            <!-- Rich Components Container - all content renders here via ComponentManager -->
            <div class="rich-components-container"></div>
          </div>

          <div class="chat-input-area">
            <vanna-status-bar
              .status=${this.status}
              .message=${this.statusMessage}
              .detail=${this.statusDetail}
              theme=${this.theme}>
            </vanna-status-bar>

            <div class="chat-input-container">
              <textarea
                class="message-input"
                .placeholder=${this.placeholder}
                .disabled=${this.disabled}
                @input=${this.handleInput}
                @keydown=${this.handleKeyPress}
                rows="1"
              ></textarea>
              <button
                class="send-button"
                type="button"
                aria-label="Send message"
                .disabled=${this.disabled || !this.currentMessage.trim()}
                @click=${this.sendMessage}
              >
                <svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor">
                  <path d="M2.01 21L23 12 2.01 3 2 10l15 2-15 2z"/>
                </svg>
              </button>
            </div>
          </div>
        </div>

        ${this.showProgress ? html`
          <div class="sidebar">
            <vanna-progress-tracker theme=${this.theme}></vanna-progress-tracker>
          </div>
        ` : ''}
      </div>
    `;
  }
}


================================================
FILE: frontends/webcomponent/src/components/vanna-message.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './vanna-message';

const meta: Meta = {
  title: 'Components/VannaMessage',
  component: 'vanna-message',
  parameters: {
    layout: 'centered',
  },
  argTypes: {
    content: { control: 'text' },
    type: {
      control: 'select',
      options: ['user', 'assistant'],
    },
    timestamp: { control: 'number' },
  },
};

export default meta;
type Story = StoryObj;

export const UserMessage: Story = {
  args: {
    content: 'Hello! Can you help me analyze my data?',
    type: 'user',
    timestamp: Date.now(),
  },
  render: (args) => html`
    <div style="width: 400px;">
      <vanna-message
        .content=${args.content}
        .type=${args.type}
        .timestamp=${args.timestamp}>
      </vanna-message>
    </div>
  `,
};

export const AssistantMessage: Story = {
  args: {
    content: 'Of course! I\'d be happy to help you analyze your data. Could you please tell me more about the type of data you have and what insights you\'re looking for?',
    type: 'assistant',
    timestamp: Date.now(),
  },
  render: (args) => html`
    <div style="width: 400px;">
      <vanna-message
        .content=${args.content}
        .type=${args.type}
        .timestamp=${args.timestamp}>
      </vanna-message>
    </div>
  `,
};

export const LongMessage: Story = {
  args: {
    content: 'This is a very long message that demonstrates how the component handles longer text content. It should wrap properly and maintain good readability while staying within the maximum width constraints. The message can contain multiple sentences and paragraphs of information that the AI assistant might provide in response to complex queries.',
    type: 'assistant',
    timestamp: Date.now(),
  },
  render: (args) => html`
    <div style="width: 400px;">
      <vanna-message
        .content=${args.content}
        .type=${args.type}
        .timestamp=${args.timestamp}>
      </vanna-message>
    </div>
  `,
};

export const Conversation: Story = {
  render: () => html`
    <div style="width: 400px;">
      <vanna-message
        content="What's the total revenue for Q4?"
        type="user"
        .timestamp=${Date.now() - 120000}>
      </vanna-message>
      <vanna-message
        content="I'll help you calculate the total revenue for Q4. Let me query your database for this information."
        type="assistant"
        .timestamp=${Date.now() - 60000}>
      </vanna-message>
      <vanna-message
        content="The total revenue for Q4 is $2,450,000. This represents a 15% increase compared to Q3."
        type="assistant"
        .timestamp=${Date.now()}>
      </vanna-message>
    </div>
  `,
};

================================================
FILE: frontends/webcomponent/src/components/vanna-message.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

@customElement('vanna-message')
export class VannaMessage extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        padding: 0 var(--vanna-space-2);
        margin-bottom: var(--vanna-space-4);
        font-family: var(--vanna-font-family-default);
        animation: fade-in-up 0.25s ease-out;
      }

      :host(:last-of-type) {
        margin-bottom: 0;
      }

      @keyframes fade-in-up {
        from {
          opacity: 0;
          transform: translateY(16px);
        }
        to {
          opacity: 1;
          transform: translateY(0);
        }
      }

      .message {
        position: relative;
        padding: var(--vanna-space-4) var(--vanna-space-5);
        border-radius: var(--vanna-chat-bubble-radius);
        word-wrap: break-word;
        line-height: 1.6;
        display: flex;
        flex-direction: column;
        gap: var(--vanna-space-2);
        max-width: min(85%, 580px);
        transition: transform var(--vanna-duration-200) ease, box-shadow var(--vanna-duration-200) ease;
        backdrop-filter: blur(8px);
      }

      .message.assistant {
        background: var(--vanna-background-root);
        border: 1px solid var(--vanna-outline-dimmer);
        color: var(--vanna-foreground-default);
        box-shadow: var(--vanna-shadow-sm);
        border-radius: var(--vanna-chat-bubble-radius) var(--vanna-chat-bubble-radius) var(--vanna-chat-bubble-radius) var(--vanna-space-2);
      }

      .message.user {
        margin-left: auto;
        max-width: min(80%, 500px);
        background: linear-gradient(135deg, var(--vanna-accent-primary-stronger) 0%, var(--vanna-accent-primary-default) 100%);
        color: white;
        box-shadow: var(--vanna-shadow-md);
        border-radius: var(--vanna-chat-bubble-radius) var(--vanna-chat-bubble-radius) var(--vanna-space-2) var(--vanna-chat-bubble-radius);
        border: 1px solid rgba(255, 255, 255, 0.2);
      }

      .message:hover {
        transform: translateY(-1px);
      }

      .message.assistant:hover {
        box-shadow: var(--vanna-shadow-md);
        border-color: var(--vanna-outline-hover);
      }

      .message.user:hover {
        box-shadow: var(--vanna-shadow-lg);
      }

      .message-content {
        margin: 0;
        font-size: 15px;
        letter-spacing: 0.01em;
        white-space: pre-wrap;
        font-weight: 400;
      }

      .message-content a {
        color: inherit;
        font-weight: 500;
        text-decoration: underline;
        text-decoration-thickness: 1px;
        text-underline-offset: 2px;
        opacity: 0.9;
      }

      .message-content code {
        font-family: var(--vanna-font-family-mono);
        background: var(--vanna-background-higher);
        padding: 2px 6px;
        border-radius: var(--vanna-border-radius-sm);
        font-size: 13px;
        border: 1px solid var(--vanna-outline-dimmer);
      }

      .message.user .message-content code {
        background: rgba(255, 255, 255, 0.2);
        border-color: rgba(255, 255, 255, 0.3);
      }

      .message-timestamp {
        display: inline-flex;
        align-items: center;
        gap: var(--vanna-space-1);
        font-size: 11px;
        letter-spacing: 0.05em;
        margin-top: var(--vanna-space-2);
        font-family: var(--vanna-font-family-default);
        opacity: 0.7;
        font-weight: 500;
      }

      .message-timestamp::before {
        content: '';
        width: 3px;
        height: 3px;
        border-radius: var(--vanna-border-radius-full);
        background: currentColor;
        opacity: 0.8;
      }

      .message.assistant .message-timestamp {
        align-self: flex-start;
        color: var(--vanna-foreground-dimmest);
      }

      .message.assistant .message-timestamp::before {
        background: var(--vanna-accent-primary-default);
      }

      .message.user .message-timestamp {
        align-self: flex-end;
        color: rgba(255, 255, 255, 0.8);
      }

      .message.user .message-timestamp::before {
        background: rgba(255, 255, 255, 0.8);
      }

      :host([theme="dark"]) .message.assistant {
        background: var(--vanna-background-higher);
        border: 1px solid var(--vanna-outline-default);
        color: var(--vanna-foreground-default);
        box-shadow: var(--vanna-shadow-md);
      }

      :host([theme="dark"]) .message.assistant .message-content code {
        background: var(--vanna-background-highest);
        border-color: var(--vanna-outline-default);
      }

      :host([theme="dark"]) .message.assistant .message-timestamp {
        color: var(--vanna-foreground-dimmest);
      }

      :host([theme="dark"]) .message.assistant .message-timestamp::before {
        background: var(--vanna-accent-primary-default);
      }

      :host([theme="dark"]) .message.user {
        background: linear-gradient(135deg, var(--vanna-accent-primary-stronger) 0%, var(--vanna-accent-primary-default) 100%);
        color: white;
        box-shadow: var(--vanna-shadow-lg);
      }

      :host([theme="dark"]) .message.user .message-content code {
        background: rgba(255, 255, 255, 0.15);
        border-color: rgba(255, 255, 255, 0.25);
      }

      :host([theme="dark"]) .message.user .message-timestamp {
        color: rgba(255, 255, 255, 0.8);
      }

      :host([theme="dark"]) .message.user .message-timestamp::before {
        background: rgba(255, 255, 255, 0.8);
      }

      @media (max-width: 600px) {
        .message {
          max-width: 100%;
        }

        .message.user {
          max-width: 100%;
        }
      }
    `
  ];

  @property() content = '';
  @property() type: 'user' | 'assistant' = 'user';
  @property({ type: Number }) timestamp = Date.now();
  @property({ reflect: true }) theme = 'light';

  private formatTimestamp(timestamp: number): string {
    return new Date(timestamp).toLocaleTimeString([], {
      hour: '2-digit',
      minute: '2-digit'
    });
  }

  render() {
    return html`
      <div class="message ${this.type}">
        <div class="message-content">${this.content}</div>
        <div class="message-timestamp">
          ${this.formatTimestamp(this.timestamp)}
        </div>
      </div>
    `;
  }
}


================================================
FILE: frontends/webcomponent/src/components/vanna-progress-tracker.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './vanna-progress-tracker';

const meta: Meta = {
  title: 'Components/VannaProgressTracker',
  component: 'vanna-progress-tracker',
  parameters: {
    layout: 'centered',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'dark', value: 'rgb(11, 15, 25)' },
        { name: 'light', value: '#ffffff' },
      ],
    },
  },
  argTypes: {
    title: { control: 'text' },
    theme: {
      control: 'select',
      options: ['dark', 'light'],
      description: 'Theme variant'
    },
  },
};

export default meta;
type Story = StoryObj;

export const Empty: Story = {
  args: {
    title: 'Agent Progress',
  },
  render: (args) => html`
    <div style="width: 350px; height: 400px;">
      <vanna-progress-tracker .title=${args.title}></vanna-progress-tracker>
    </div>
  `,
};

export const WithTasks: Story = {
  args: {
    title: 'Agent Progress',
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const tracker = document.querySelector('vanna-progress-tracker') as any;
      if (tracker) {
        tracker.addItem('Analyze database schema', 'Examining table structure');
        tracker.addItem('Generate SQL query', 'Based on user request');
        tracker.addItem('Execute query', 'Running against production DB');
        tracker.addItem('Format results', 'Creating visualization');

        // Update first item to in_progress
        const items = tracker.shadowRoot?.querySelectorAll('.progress-item');
        if (items?.[0]) {
          tracker.updateItem(tracker.items[0].id, 'in_progress', 'Scanning tables...');
        }
      }
    }, 100);

    return html`
      <div style="width: 350px; height: 400px; background: ${args.theme === 'light' ? '#ffffff' : 'rgb(11, 15, 25)'}; padding: 20px;">
        <vanna-progress-tracker .title=${args.title} theme=${args.theme}></vanna-progress-tracker>
      </div>
    `;
  },
};

export const WithTasksLight: Story = {
  args: {
    title: 'Agent Progress',
    theme: 'light',
  },
  render: (args) => {
    setTimeout(() => {
      const tracker = document.querySelector('vanna-progress-tracker') as any;
      if (tracker) {
        tracker.addItem('Analyze database schema', 'Examining table structure');
        tracker.addItem('Generate SQL query', 'Based on user request');
        tracker.addItem('Execute query', 'Running against production DB');
        tracker.addItem('Format results', 'Creating visualization');

        // Update first item to in_progress
        const items = tracker.shadowRoot?.querySelectorAll('.progress-item');
        if (items?.[0]) {
          tracker.updateItem(tracker.items[0].id, 'in_progress', 'Scanning tables...');
        }
      }
    }, 100);

    return html`
      <div style="width: 350px; height: 400px; background: ${args.theme === 'light' ? '#ffffff' : 'rgb(11, 15, 25)'}; padding: 20px;">
        <vanna-progress-tracker .title=${args.title} theme=${args.theme}></vanna-progress-tracker>
      </div>
    `;
  },
};

export const MixedStatuses: Story = {
  args: {
    title: 'Data Analysis Pipeline',
  },
  render: (args) => {
    setTimeout(() => {
      const tracker = document.querySelector('vanna-progress-tracker') as any;
      if (tracker) {
        const id1 = tracker.addItem('Connect to database', 'Establishing connection');
        const id2 = tracker.addItem('Validate credentials', 'Checking access permissions');
        const id3 = tracker.addItem('Load data schema', 'Reading table definitions');
        const id4 = tracker.addItem('Parse user query', 'Understanding natural language');
        const id5 = tracker.addItem('Generate SQL', 'Converting to database query');
        const id6 = tracker.addItem('Execute query', 'Running against database');
        const id7 = tracker.addItem('Process results', 'Formatting output');

        // Simulate different states
        tracker.updateItem(id1, 'completed');
        tracker.updateItem(id2, 'completed');
        tracker.updateItem(id3, 'completed');
        tracker.updateItem(id4, 'in_progress', 'Analyzing: "Show me sales by region"');
        tracker.updateItem(id5, 'pending');
        tracker.updateItem(id6, 'pending');
        tracker.updateItem(id7, 'pending');
      }
    }, 100);

    return html`
      <div style="width: 350px; height: 400px;">
        <vanna-progress-tracker .title=${args.title}></vanna-progress-tracker>
      </div>
    `;
  },
};

export const WithError: Story = {
  args: {
    title: 'Query Processing',
  },
  render: (args) => {
    setTimeout(() => {
      const tracker = document.querySelector('vanna-progress-tracker') as any;
      if (tracker) {
        const id1 = tracker.addItem('Parse request', 'Understanding user query');
        const id2 = tracker.addItem('Generate SQL', 'Creating database query');
        const id3 = tracker.addItem('Execute query', 'Running against database');
        tracker.addItem('Format results', 'Preparing visualization');

        tracker.updateItem(id1, 'completed');
        tracker.updateItem(id2, 'completed');
        tracker.updateItem(id3, 'error', 'Table "sales_data" does not exist');
        // id4 should remain pending due to error
      }
    }, 100);

    return html`
      <div style="width: 350px; height: 400px;">
        <vanna-progress-tracker .title=${args.title}></vanna-progress-tracker>
      </div>
    `;
  },
};

export const MultipleErrors: Story = {
  args: {
    title: 'Error Scenarios',
  },
  render: (args) => {
    setTimeout(() => {
      const tracker = document.querySelector('vanna-progress-tracker') as any;
      if (tracker) {
        const id1 = tracker.addItem('Connect to database', 'Establishing connection');
        const id2 = tracker.addItem('Validate schema', 'Checking table structure');
        const id3 = tracker.addItem('Parse SQL query', 'Analyzing syntax');
        tracker.addItem('Execute query', 'Running database command');
        tracker.addItem('Process results', 'Formatting output');

        tracker.updateItem(id1, 'error', 'Connection timeout - database unreachable');
        tracker.updateItem(id2, 'error', 'Invalid credentials provided');
        tracker.updateItem(id3, 'error', 'Syntax error in SQL query');
        // Other items remain pending
      }
    }, 100);

    return html`
      <div style="width: 350px; height: 400px;">
        <vanna-progress-tracker .title=${args.title}></vanna-progress-tracker>
        <p style="font-size: 12px; color: #666; margin-top: 10px;">
          Example showing multiple error states with detailed error messages
        </p>
      </div>
    `;
  },
};

export const LiveDemo: Story = {
  args: {
    title: 'Live Progress Demo',
  },
  render: (args) => {
    let tracker: any;
    let taskIds: string[] = [];
    let currentIndex = 0;

    const tasks = [
      { text: 'Initialize AI agent', detail: 'Loading language model' },
      { text: 'Analyze user request', detail: 'Processing natural language' },
      { text: 'Query database schema', detail: 'Understanding data structure' },
      { text: 'Generate SQL query', detail: 'Converting request to SQL' },
      { text: 'Execute query', detail: 'Running against database' },
      { text: 'Process results', detail: 'Formatting data for display' },
      { text: 'Generate visualization', detail: 'Creating charts and graphs' }
    ];

    const runDemo = () => {
      if (!tracker) {
        tracker = document.querySelector('vanna-progress-tracker');
        if (!tracker) {
          setTimeout(runDemo, 100);
          return;
        }
      }

      // Add all tasks as pending
      if (taskIds.length === 0) {
        taskIds = tasks.map(task => tracker.addItem(task.text, task.detail));
        currentIndex = 0;
      }

      // Process tasks one by one
      if (currentIndex < tasks.length) {
        // Mark current as in_progress
        tracker.updateItem(taskIds[currentIndex], 'in_progress', `${tasks[currentIndex].detail}...`);

        // Complete after 2 seconds, then move to next
        setTimeout(() => {
          tracker.updateItem(taskIds[currentIndex], 'completed');
          currentIndex++;

          // Continue with next task
          if (currentIndex < tasks.length) {
            setTimeout(runDemo, 500);
          } else {
            // Demo complete - restart after 3 seconds
            setTimeout(() => {
              tracker.clearItems();
              taskIds = [];
              currentIndex = 0;
              setTimeout(runDemo, 1000);
            }, 3000);
          }
        }, 2000);
      }
    };

    setTimeout(runDemo, 500);

    return html`
      <div style="width: 350px; height: 400px;">
        <vanna-progress-tracker .title=${args.title}></vanna-progress-tracker>
        <div style="margin-top: 10px; color: #999; font-size: 12px; text-align: center;">
          Watch tasks complete automatically (demo loops)
        </div>
      </div>
    `;
  },
};

================================================
FILE: frontends/webcomponent/src/components/vanna-progress-tracker.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property, state } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

interface ProgressItem {
  id: string;
  text: string;
  status: 'pending' | 'in_progress' | 'completed' | 'error';
  detail?: string;
}

@customElement('vanna-progress-tracker')
export class VannaProgressTracker extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        background: var(--vanna-background-default);
        border: 1px solid var(--vanna-outline-default);
        border-radius: 0 0 var(--vanna-border-radius-lg) var(--vanna-border-radius-lg);
        overflow: hidden;
        font-family: var(--vanna-font-family-default);
      }

      .progress-label {
        padding: var(--vanna-space-3) var(--vanna-space-4) var(--vanna-space-2);
        display: flex;
        align-items: center;
        justify-content: space-between;
      }

      .progress-label-text {
        font-size: 11px;
        font-weight: 500;
        color: var(--vanna-foreground-dimmest);
        text-transform: uppercase;
        letter-spacing: 0.5px;
        margin: 0;
      }

      .progress-summary {
        font-size: 10px;
        color: var(--vanna-foreground-dimmest);
        font-weight: 400;
      }

      .progress-list {
        max-height: 300px;
        overflow-y: auto;
        padding-top: 0;
      }

      .progress-item {
        padding: var(--vanna-space-3) var(--vanna-space-4);
        border-bottom: 1px solid var(--vanna-outline-dimmest);
        display: flex;
        align-items: flex-start;
        gap: var(--vanna-space-3);
        transition: background var(--vanna-duration-150) ease;
      }

      .progress-item:last-child {
        border-bottom: none;
      }

      .progress-item:hover {
        background: var(--vanna-background-higher);
      }

      .progress-item.in_progress {
        background: rgba(0, 123, 255, 0.05);
        border-left: 3px solid var(--vanna-accent-primary-default);
      }

      .progress-item.completed {
        opacity: 0.7;
      }

      .progress-item.error {
        background: var(--vanna-accent-negative-subtle);
        border-left: 3px solid var(--vanna-accent-negative-default);
        padding-left: calc(var(--vanna-space-3) - 3px);
      }

      .progress-item.error .progress-text {
        color: var(--vanna-accent-negative-stronger);
      }

      .progress-item.error .progress-detail {
        color: var(--vanna-accent-negative-default);
        font-weight: 500;
      }

      .progress-icon {
        width: 16px;
        height: 16px;
        border-radius: 50%;
        display: flex;
        align-items: center;
        justify-content: center;
        flex-shrink: 0;
        margin-top: 1px;
      }

      .progress-icon.pending {
        background: var(--vanna-outline-default);
      }

      .progress-icon.in_progress {
        background: var(--vanna-accent-primary-default);
      }

      .progress-icon.completed {
        background: var(--vanna-accent-positive-default);
      }

      .progress-icon.error {
        background: var(--vanna-accent-negative-default);
        box-shadow: 0 0 0 2px var(--vanna-accent-negative-subtle);
      }

      .progress-icon svg {
        width: 10px;
        height: 10px;
        color: white;
      }

      .progress-icon.error svg {
        width: 8px;
        height: 8px;
        color: white;
      }

      .spinner-mini {
        width: 10px;
        height: 10px;
        border: 1.5px solid rgba(255, 255, 255, 0.3);
        border-top-color: white;
        border-radius: 50%;
        animation: spin 1s linear infinite;
      }

      .progress-content {
        flex: 1;
        min-width: 0;
      }

      .progress-text {
        font-size: 13px;
        color: var(--vanna-foreground-default);
        font-weight: 500;
        margin: 0 0 var(--vanna-space-1) 0;
        line-height: 1.3;
      }

      .progress-detail {
        font-size: 11px;
        color: var(--vanna-foreground-dimmest);
        margin: 0;
        line-height: 1.3;
      }

      .empty-state {
        padding: var(--vanna-space-6) var(--vanna-space-4);
        text-align: center;
        color: var(--vanna-foreground-dimmest);
        font-size: 12px;
      }

      @keyframes spin {
        to {
          transform: rotate(360deg);
        }
      }
    `
  ];

  @property() title = 'Progress';
  @property() theme = 'light';
  @state() private items: ProgressItem[] = [];

  addItem(text: string, detail?: string, id?: string): string {
    const itemId = id || Date.now().toString();
    this.items = [...this.items, {
      id: itemId,
      text,
      status: 'pending',
      detail
    }];
    return itemId;
  }

  updateItem(id: string, status: ProgressItem['status'], detail?: string) {
    this.items = this.items.map(item =>
      item.id === id ? { ...item, status, detail } : item
    );
  }

  clearItems() {
    this.items = [];
  }

  private getStatusIcon(status: ProgressItem['status']) {
    switch (status) {
      case 'pending':
        return html``;
      case 'in_progress':
        return html`<div class="spinner-mini"></div>`;
      case 'completed':
        return html`
          <svg viewBox="0 0 20 20" fill="currentColor">
            <path fill-rule="evenodd" d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z" clip-rule="evenodd" />
          </svg>
        `;
      case 'error':
        return html`
          <svg viewBox="0 0 20 20" fill="currentColor">
            <path fill-rule="evenodd" d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z" clip-rule="evenodd" />
          </svg>
        `;
    }
  }

  private getProgressSummary() {
    const completed = this.items.filter(item => item.status === 'completed').length;
    const total = this.items.length;
    const inProgress = this.items.filter(item => item.status === 'in_progress').length;

    if (inProgress > 0) {
      return `${completed}/${total} completed`;
    }
    return total > 0 ? `${completed}/${total} completed` : '';
  }

  render() {
    return html`
      ${this.items.length > 0 ? html`
        <div class="progress-label">
          <span class="progress-label-text">Tasks</span>
          <span class="progress-summary">${this.getProgressSummary()}</span>
        </div>
      ` : ''}

      <div class="progress-list">
        ${this.items.length === 0
          ? html`<div class="empty-state">No tasks yet</div>`
          : this.items.map(item => html`
              <div class="progress-item ${item.status}">
                <div class="progress-icon ${item.status}">
                  ${this.getStatusIcon(item.status)}
                </div>
                <div class="progress-content">
                  <p class="progress-text">${item.text}</p>
                  ${item.detail ? html`<p class="progress-detail">${item.detail}</p>` : ''}
                </div>
              </div>
            `)
        }
      </div>
    `;
  }
}

================================================
FILE: frontends/webcomponent/src/components/vanna-status-bar.stories.ts
================================================
import type { Meta, StoryObj } from '@storybook/web-components';
import { html } from 'lit';
import './vanna-status-bar';

const meta: Meta = {
  title: 'Components/VannaStatusBar',
  component: 'vanna-status-bar',
  parameters: {
    layout: 'centered',
    backgrounds: {
      default: 'light',
      values: [
        { name: 'dark', value: 'rgb(11, 15, 25)' },
        { name: 'light', value: '#ffffff' },
      ],
    },
  },
  argTypes: {
    status: {
      control: 'select',
      options: ['idle', 'working', 'error', 'success'],
    },
    message: { control: 'text' },
    detail: { control: 'text' },
    theme: {
      control: 'select',
      options: ['dark', 'light'],
      description: 'Theme variant'
    },
  },
};

export default meta;
type Story = StoryObj;

export const Idle: Story = {
  args: {
    status: 'idle',
    message: '',
    detail: '',
  },
  render: (args) => html`
    <div style="width: 500px; padding: 20px;">
      <vanna-status-bar
        .status=${args.status}
        .message=${args.message}
        .detail=${args.detail}>
      </vanna-status-bar>
      <div style="margin-top: 10px; color: #999; font-size: 12px;">
        Status bar is hidden when idle
      </div>
    </div>
  `,
};

export const Working: Story = {
  args: {
    status: 'working',
    message: 'Analyzing your database schema...',
    detail: 'Step 1 of 3',
    theme: 'light',
  },
  render: (args) => html`
    <div style="width: 500px; padding: 20px; background: ${args.theme === 'light' ? '#ffffff' : 'rgb(11, 15, 25)'};">
      <vanna-status-bar
        .status=${args.status}
        .message=${args.message}
        .detail=${args.detail}
        theme=${args.theme}>
      </vanna-status-bar>
    </div>
  `,
};

export const WorkingLight: Story = {
  args: {
    status: 'working',
    message: 'Analyzing your database schema...',
    detail: 'Step 1 of 3',
    theme: 'light',
  },
  render: (args) => html`
    <div style="width: 500px; padding: 20px; background: ${args.theme === 'light' ? '#ffffff' : 'rgb(11, 15, 25)'};">
      <vanna-status-bar
        .status=${args.status}
        .message=${args.message}
        .detail=${args.detail}
        theme=${args.theme}>
      </vanna-status-bar>
    </div>
  `,
};

export const Success: Story = {
  args: {
    status: 'success',
    message: 'Query executed successfully',
    detail: '2.3s',
  },
  render: (args) => html`
    <div style="width: 500px; padding: 20px;">
      <vanna-status-bar
        .status=${args.status}
        .message=${args.message}
        .detail=${args.detail}>
      </vanna-status-bar>
    </div>
  `,
};

export const Error: Story = {
  args: {
    status: 'error',
    message: 'Failed to connect to database',
    detail: 'Connection timeout after 30s',
  },
  render: (args) => html`
    <div style="width: 500px; padding: 20px;">
      <vanna-status-bar
        .status=${args.status}
        .message=${args.message}
        .detail=${args.detail}>
      </vanna-status-bar>
    </div>
  `,
};

export const StatusSequence: Story = {
  render: () => {
    let statusBar: any;
    let currentIndex = 0;
    const statuses = [
      { status: 'working', message: 'Starting analysis...', detail: 'Initializing' },
      { status: 'working', message: 'Querying database...', detail: 'Step 1 of 3' },
      { status: 'working', message: 'Processing results...', detail: 'Step 2 of 3' },
      { status: 'working', message: 'Generating visualization...', detail: 'Step 3 of 3' },
      { status: 'success', message: 'Analysis complete!', detail: '4.2s total' },
    ];

    const updateStatus = () => {
      if (statusBar && currentIndex < statuses.length) {
        const current = statuses[currentIndex];
        statusBar.status = current.status;
        statusBar.message = current.message;
        statusBar.detail = current.detail;
        currentIndex++;

        if (currentIndex < statuses.length) {
          setTimeout(updateStatus, 2000);
        }
      }
    };

    setTimeout(() => {
      statusBar = document.querySelector('vanna-status-bar');
      updateStatus();
    }, 100);

    return html`
      <div style="width: 500px; padding: 20px;">
        <vanna-status-bar status="idle"></vanna-status-bar>
        <div style="margin-top: 10px; color: #999; font-size: 12px;">
          Watch the status bar cycle through different states
        </div>
      </div>
    `;
  },
};

================================================
FILE: frontends/webcomponent/src/components/vanna-status-bar.ts
================================================
import { LitElement, html, css } from 'lit';
import { customElement, property } from 'lit/decorators.js';
import { vannaDesignTokens } from '../styles/vanna-design-tokens.js';

@customElement('vanna-status-bar')
export class VannaStatusBar extends LitElement {
  static styles = [
    vannaDesignTokens,
    css`
      :host {
        display: block;
        background: rgba(254, 93, 38, 0.1);
        border: 2px solid var(--vanna-orange);
        border-radius: var(--vanna-border-radius-xl);
        padding: var(--vanna-space-4) var(--vanna-space-5);
        margin-bottom: var(--vanna-space-4);
        font-family: var(--vanna-font-family-default);
        font-size: 14px;
        font-weight: 500;
        color: var(--vanna-navy);
        backdrop-filter: blur(12px);
        box-shadow:
          var(--vanna-shadow-lg),
          0 0 0 1px rgba(254, 93, 38, 0.1);
        
        /* Animation properties */
        opacity: 1;
        transform: translateY(0) scale(1);
        max-height: 200px;
        overflow: hidden;
        transition: 
          opacity var(--vanna-duration-300) cubic-bezier(0.4, 0, 0.2, 1),
          transform var(--vanna-duration-300) cubic-bezier(0.4, 0, 0.2, 1),
          max-height var(--vanna-duration-300) ease,
          margin var(--vanna-duration-300) ease,
          padding var(--vanna-duration-300) ease,
          box-shadow var(--vanna-duration-200) ease;
      }

      /* Hide when there's no actual content */
      :host(.no-content) {
        opacity: 0;
        transform: translateY(-8px) scale(0.95);
        max-height: 0;
        margin: 0;
        padding: 0;
        pointer-events: none;
      }

      :host(:empty) {
        display: none;
      }

      /* Entrance animation when content appears */
      :host(.entering) {
        animation: statusEnter var(--vanna-duration-300) ease-out;
      }

      /* Exit animation when content disappears */
      :host(.exiting) {
        animation: statusExit var(--vanna-duration-300) ease-in;
      }

      @keyframes statusEnter {
        0% {
          opacity: 0;
          transform: translateY(-12px) scale(0.9);
          max-height: 0;
        }
        50% {
          opacity: 0.8;
          transform: translateY(-2px) scale(1.02);
        }
        100% {
          opacity: 1;
          transform: translateY(0) scale(1);
          max-height: 200px;
        }
      }

      @keyframes statusExit {
        0% {
          opacity: 1;
          transform: translateY(0) scale(1);
          max-height: 200px;
        }
        50% {
          opacity: 0.5;
          transform: translateY(-4px) scale(0.98);
        }
        100% {
          opacity: 0;
          transform: translateY(-12px) scale(0.9);
          max-height: 0;
        }
      }

      :host([status="working"]) {
        background: var(--vanna-orange);
        border-color: var(--vanna-orange);
        color: white;
        box-shadow:
          var(--vanna-shadow-xl),
          0 0 0 2px rgba(254, 93, 38, 0.3),
          0 0 20px rgba(254, 93, 38, 0.4);
      }

      :host([status="error"]) {
        background: linear-gradient(135deg, var(--vanna-accent-negative-subtle) 0%, rgba(239, 68, 68, 0.15) 100%);
        border-color: var(--vanna-accent-negative-default);
        color: var(--vanna-accent-negative-stronger);
        box-shadow: 
          var(--vanna-shadow-xl),
          0 0 0 2px rgba(239, 68, 68, 0.3),
          0 0 20px rgba(239, 68, 68, 0.2);
        animation: errorShake 0.5s ease-in-out, errorGlow 2s ease-in-out;
      }

      :host([status="success"]) {
        background: linear-gradient(135deg, var(--vanna-accent-positive-subtle) 0%, rgba(16, 185, 129, 0.15) 100%);
        border-color: var(--vanna-accent-positive-default);
        color: var(--vanna-accent-positive-stronger);
        box-shadow: 
          var(--vanna-shadow-xl),
          0 0 0 2px rgba(16, 185, 129, 0.3),
          0 0 20px rgba(16, 185, 129, 0.2);
        animation: successPulse 0.6s ease-out, successGlow 2s ease-out;
      }

      @keyframes errorShake {
        0%, 100% { transform: translateX(0); }
        10%, 30%, 50%, 70%, 90% { transform: translateX(-4px); }
        20%, 40%, 60%, 80% { transform: translateX(4px); }
      }

      @keyframes successPulse {
        0% { 
          transform: scale(1); 
        }
        50% { 
          transform: scale(1.05); 
        }
        100% { 
          transform: scale(1); 
        }
      }

      .status-content {
        display: flex;
        align-items: center;
        gap: var(--vanna-space-3);
        animation: contentFadeIn var(--vanna-duration-200) ease-out;
      }

      @keyframes contentFadeIn {
        0% {
          opacity: 0;
          transform: translateY(4px);
        }
        100% {
          opacity: 1;
          transform: translateY(0);
        }
      }

      .status-indicator {
        width: 12px;
        height: 12px;
        border-radius: var(--vanna-border-radius-full);
        background: var(--vanna-accent-primary-default);
        flex-shrink: 0;
        box-shadow: 0 0 0 2px rgba(255, 255, 255, 0.5), 0 2px 8px rgba(0, 0, 0, 0.15);
      }

      .status-indicator.working {
        background: white;
        animation: workingPulse 1.5s ease-in-out infinite;
      }

      .status-indicator.error {
        background: linear-gradient(45deg, var(--vanna-accent-negative-default), var(--vanna-accent-negative-stronger));
        box-shadow: 0 0 0 2px rgba(255, 255, 255, 0.5), 0 0 8px rgba(239, 68, 68, 0.4);
      }

      .status-indicator.success {
        background: linear-gradient(45deg, var(--vanna-accent-positive-default), var(--vanna-accent-positive-stronger));
        box-shadow: 0 0 0 2px rgba(255, 255, 255, 0.5), 0 0 8px rgba(16, 185, 129, 0.4);
      }

      .spinner {
        width: 16px;
        height: 16px;
        border: 3px solid rgba(21, 168, 168, 0.3);
        border-top-color: var(--vanna-teal);
        border-radius: var(--vanna-border-radius-full);
        animation: spin 1s linear infinite, spinnerGlow 2s ease-in-out infinite;
        flex-shrink: 0;
      }

      .status-text {
        flex: 1;
        font-weight: 600;
        line-height: 1.4;
        letter-spacing: 0.01em;
      }

      .status-detail {
        font-size: 12px;
        color: var(--vanna-foreground-dimmest);
        margin-left: var(--vanna-space-4);
        opacity: 0.9;
        font-weight: 500;
      }

      .status-actions {
        display: flex;
        align-items: center;
        gap: var(--vanna-space-2);
        margin-left: auto;
      }

      .status-button {
        padding: var(--vanna-space-1) var(--vanna-space-2);
        border: 1px solid var(--vanna-outline-default);
        border-radius: var(--vanna-border-radius-sm);
        background: var(--vanna-background-subtle);
        color: var(--vanna-foreground-dimmer);
        font-size: 11px;
        font-weight: 500;
        cursor: pointer;
        transition: all var(--vanna-duration-150) ease;
      }

      .status-button:hover {
        background: var(--vanna-background-higher);
        border-color: var(--vanna-outline-hover);
        color: var(--vanna-foreground-default);
      }

      @keyframes spin {
        to {
          transform: rotate(360deg);
        }
      }

      @keyframes pulse {
        0%, 100% {
          opacity: 1;
          transform: scale(1);
        }
        50% {
          opacity: 0.6;
          transform: scale(1.1);
        }
      }

      @keyframes workingPulse {
        0%, 100% {
          opacity: 1;
          transform: scale(1);
          box-shadow: 0 0 0 2px rgba(255, 255, 255, 0.8), 0 2px 8px rgba(255, 255, 255, 0.3);
        }
        50% {
          opacity: 0.9;
          transform: scale(1.2);
          box-shadow: 0 0 0 4px rgba(255, 255, 255, 0.9), 0 4px 12px rgba(255, 255, 255, 0.5);
        }
      }

      @keyframes spinnerGlow {
        0%, 100% {
          filter: drop-shadow(0 0 2px rgba(21, 168, 168, 0.5));
        }
        50% {
          filter: drop-shadow(0 0 6px rgba(21, 168, 168, 0.8));
        }
      }

      @keyframes errorGlow {
        0% {
          box-shadow: 
            var(--vanna-shadow-xl),
            0 0 0 2px rgba(239, 68, 68, 0.3),
            0 0 20px rgba(239, 68, 68, 0.2);
        }
        50% {
          box-shadow: 
            var(--vanna-shadow-2xl),
            0 0 0 3px rgba(239, 68, 68, 0.4),
            0 0 30px rgba(239, 68, 68, 0.3);
        }
        100% {
          box-shadow: 
            var(--vanna-shadow-xl),
            0 0 0 2px rgba(239, 68, 68, 0.3),
            0 0 20px rgba(239, 68, 68, 0.2);
        }
      }

      @keyframes successGlow {
        0% {
          box-shadow: 
            var(--vanna-shadow-xl),
            0 0 0 2px rgba(16, 185, 129, 0.3),
            0 0 20px rgba(16, 185, 129, 0.2);
        }
        50% {
          box-shadow: 
            var(--vanna-shadow-2xl),
            0 0 0 3px rgba(16, 185, 129, 0.4),
            0 0 30px rgba(16, 185, 129, 0.3);
        }
        100% {
          box-shadow: 
            var(--vanna-shadow-xl),
            0 0 0 2px rgba(16, 185, 129, 0.3),
            0 0 20px rgba(16, 185, 129, 0.2);
        }
      }

      /* Dark theme overrides */
      :host([theme="dark"]) {
        background: var(--vanna-background-higher);
        border-color: var(--vanna-outline-default);
      }

      :host([theme="dark"]) .status-button {
        background: var(--vanna-background-highest);
        border-color: var(--vanna-outline-default);
      }

      :host([theme="dark"]) .status-button:hover {
        background: var(--vanna-background-highest);
        border-color: var(--vanna-outline-hover);
      }
    `
  ];

  @property() status: 'idle' | 'working' | 'error' | 'success' = 'idle';
  @property() message = '';
  @property() detail = '';
  @property() theme = 'light';

  private _previousHasContent = false;
  private _enterTimeout: number | null = null;
  private _exitTimeout: number | null = null;
  private _lastUpdateTime = 0;

  disconnectedCallback() {
    super.disconnectedCallback();

    // Clean up pending animation timeouts when component is removed
    if (this._enterTimeout !== null) {
      clearTimeout(this._enterTimeout);
      this._enterTimeout = null;
    }
    if (this._exitTimeout !== null) {
      clearTimeout(this._exitTimeout);
      this._exitTimeout = null;
    }
  }

  updated(_changedProperties: Map<string | number | symbol, unknown>) {
    // Update CSS class based on content
    const hasContent = Boolean(this.message && this.message.trim());

    // Cancel any pending animation timeouts to prevent race conditions
    if (this._enterTimeout !== null) {
      clearTimeout(this._enterTimeout);
      this._enterTimeout = null;
    }
    if (this._exitTimeout !== null) {
      clearTimeout(this._exitTimeout);
      this._exitTimeout = null;
    }

    // Debounce rapid updates to prevent animation jank
    const now = Date.now();
    const timeSinceLastUpdate = now - this._lastUpdateTime;
    const shouldDebounce = timeSinceLastUpdate < 100; // 100ms debounce

    // Handle animation classes
    if (hasContent !== this._previousHasContent) {
      if (hasContent) {
        // Content appeared - animate in
        this.classList.remove('no-content', 'exiting');

        if (!shouldDebounce) {
          // Only animate if not rapid-firing
          this.classList.add('entering');

          // Remove entering class after animation
          this._enterTimeout = window.setTimeout(() => {
            this.classList.remove('entering');
            this._enterTimeout = null;
          }, 300);
        }
      } else {
        // Content disappeared - animate out
        this.classList.remove('entering');

        if (!shouldDebounce) {
          // Only animate if not rapid-firing
          this.classList.add('exiting');

          // Add no-content class after animation
          this._exitTimeout = window.setTimeout(() => {
            this.classList.remove('exiting');
            this.classList.add('no-content');
            this._exitTimeout = null;
          }, 300);
        } else {
          // If rapid-firing, skip animation and go straight to no-content
          this.classList.add('no-content');
        }
      }
    } else if (!hasContent) {
      // Ensure no-content class is applied when no content
      this.classList.add('no-content');
    }

    this._previousHasContent = hasContent;
    this._lastUpdateTime = now;
  }

  render() {
    // Only show if there's actual content (message) to display
    if (!this.message || !this.message.trim()) {
      return html``;
    }

    return html`
      <div class="status-content">
        ${this.status === 'working'
          ? html`<div class="spinner"></div>`
          : html`<div class="status-indicator ${this.status}"></div>`
        }
        <span class="status-text">${this.message}</span>
        ${this.detail ? html`<span class="status-detail">${this.detail}</span>` : ''}
      </div>
    `;
  }
}

================================================
FILE: frontends/webcomponent/src/index.ts
================================================
// Log build information when the module loads
console.log(
  '%c🎨 Vanna Web Components',
  'color: #4CAF50; font-weight: bold; font-size: 14px;'
);
console.log(
  `%c📦 Version: ${__BUILD_VERSION__}`,
  'color: #2196F3; font-weight: bold;'
);
console.log(
  `%c🕐 Built: ${__BUILD_TIME__}`,
  'color: #FF9800; font-weight: bold;'
);
console.log(
  '%c━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━',
  'color: #9E9E9E;'
);

export { VannaChat } from './components/vanna-chat';
export { VannaMessage } from './components/vanna-message';
export { VannaStatusBar } from './components/vanna-status-bar';
export { VannaProgressTracker } from './components/vanna-progress-tracker';
export { PlotlyChart } from './components/plotly-chart';

// Rich component system
export {
  ComponentRegistry,
  ComponentManager,
  CardComponentRenderer,
  TaskListComponentRenderer,
  ProgressBarComponentRenderer,
  NotificationComponentRenderer,
  StatusIndicatorComponentRenderer,
  TextComponentRenderer
} from './components/rich-component-system';

// Rich component styles are injected automatically by the ComponentManager


================================================
FILE: frontends/webcomponent/src/services/api-client.ts
================================================
/**
 * API client for communicating with Vanna Agents backend
 */

export interface ChatMessage {
  id: string;
  content: string;
  type: 'user' | 'assistant';
  timestamp: number;
}

export interface ChatRequest {
  message: string;
  conversation_id?: string;
  user_id?: string;
  request_id?: string;
  metadata?: Record<string, any>;
}

export interface ChatStreamChunk {
  rich: Record<string, any>;
  simple?: Record<string, any>;
  conversation_id: string;
  request_id: string;
  timestamp: number;
}

export interface ChatResponse {
  chunks: ChatStreamChunk[];
  conversation_id: string;
  request_id: string;
  total_chunks: number;
}

export interface ApiClientConfig {
  baseUrl?: string;
  sseEndpoint?: string;
  wsEndpoint?: string;
  pollEndpoint?: string;
  timeout?: number;
  customHeaders?: Record<string, string>;
}

export class VannaApiClient {
  public readonly baseUrl: string;
  private sseEndpoint: string;
  private wsEndpoint: string;
  private pollEndpoint: string;
  private timeout: number;
  private customHeaders: Record<string, string>;

  constructor(config: ApiClientConfig = {}) {
    this.baseUrl = config.baseUrl || '';
    this.sseEndpoint = config.sseEndpoint || '/api/vanna/v2/chat_sse';
    this.wsEndpoint = config.wsEndpoint || '/api/vanna/v2/chat_websocket';
    this.pollEndpoint = config.pollEndpoint || '/api/vanna/v2/chat_poll';
    this.timeout = config.timeout || 30000;
    this.customHeaders = config.customHeaders || {};

    console.log('[VannaApiClient] Constructor called with config:', config);
    console.log('[VannaApiClient] Endpoint configuration:');
    console.log('  - SSE endpoint:', this.sseEndpoint, config.sseEndpoint ? '(custom)' : '(default)');
    console.log('  - WS endpoint:', this.wsEndpoint, config.wsEndpoint ? '(custom)' : '(default)');
    console.log('  - Poll endpoint:', this.pollEndpoint, config.pollEndpoint ? '(custom)' : '(default)');
    console.log('  - Base URL:', this.baseUrl || '(empty)');
  }

  /**
   * Update custom headers (e.g., for authentication)
   */
  setCustomHeaders(headers: Record<string, string>) {
    this.customHeaders = headers;
  }

  /**
   * Get current custom headers
   */
  getCustomHeaders(): Record<string, string> {
    return { ...this.customHeaders };
  }

  /**
   * Send message using Server-Sent Events (SSE) streaming
   */
  async *streamChat(request: ChatRequest): AsyncGenerator<ChatStreamChunk, void, unknown> {
    const url = this.sseEndpoint.startsWith('http')
      ? this.sseEndpoint
      : `${this.baseUrl}${this.sseEndpoint}`;

    console.log('[VannaApiClient] SSE streaming to URL:', url);
    console.log('[VannaApiClient] SSE endpoint config:', {
      baseUrl: this.baseUrl,
      sseEndpoint: this.sseEndpoint,
      constructedUrl: url
    });

    const response = await fetch(url, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Accept': 'text/event-stream',
        ...this.customHeaders,
      },
      body: JSON.stringify(request),
    });

    if (!response.ok) {
      throw new Error(`HTTP ${response.status}: ${response.statusText}`);
    }

    const reader = response.body?.getReader();
    if (!reader) {
      throw new Error('No response body');
    }

    const decoder = new TextDecoder();
    let buffer = '';

    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6).trim();
            if (data === '[DONE]') {
              return;
            }

            try {
              const chunk = JSON.parse(data) as ChatStreamChunk;
              yield chunk;
            } catch (e) {
              console.warn('Failed to parse SSE chunk:', data, e);
            }
          }
        }
      }
    } finally {
      reader.releaseLock();
    }
  }

  /**
   * Send message using WebSocket
   */
  createWebSocketConnection(): Promise<WebSocket> {
    return new Promise((resolve, reject) => {
      let wsUrl: string;

      if (this.wsEndpoint.startsWith('ws://') || this.wsEndpoint.startsWith('wss://')) {
        // Absolute WebSocket URL provided
        wsUrl = this.wsEndpoint;
      } else {
        // Relative path - construct from baseUrl
        if (this.baseUrl) {
          // Parse baseUrl to extract host and convert http(s) to ws(s)
          const baseUrlObj = new URL(this.baseUrl);
          const wsProtocol = baseUrlObj.protocol === 'https:' ? 'wss:' : 'ws:';
          wsUrl = `${wsProtocol}//${baseUrlObj.host}${this.wsEndpoint}`;
        } else {
          // Fallback to window.location
          const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
          wsUrl = `${protocol}//${window.location.host}${this.wsEndpoint}`;
        }
      }

      const ws = new WebSocket(wsUrl);

      ws.onopen = () => resolve(ws);
      ws.onerror = (error) => reject(error);

      // Set timeout
      setTimeout(() => {
        if (ws.readyState === WebSocket.CONNECTING) {
          ws.close();
          reject(new Error('WebSocket connection timeout'));
        }
      }, this.timeout);
    });
  }

  /**
   * Send message via WebSocket
   */
  async sendWebSocketMessage(
    ws: WebSocket,
    request: ChatRequest
  ): Promise<AsyncGenerator<ChatStreamChunk, void, unknown>> {
    return new Promise((resolve, reject) => {
      if (ws.readyState !== WebSocket.OPEN) {
        reject(new Error('WebSocket not connected'));
        return;
      }

      async function* generator() {
        let isCompleted = false;
        const messageQueue: ChatStreamChunk[] = [];
        let resolveNext: ((value: IteratorResult<ChatStreamChunk>) => void) | null = null;

        const messageHandler = (event: MessageEvent) => {
          try {
            const chunk = JSON.parse(event.data) as ChatStreamChunk;

            if (chunk.rich?.type === 'completion') {
              isCompleted = true;
              if (resolveNext) {
                resolveNext({ done: true, value: undefined });
                resolveNext = null;
              }
              return;
            }

            if (chunk.rich?.type === 'error') {
              ws.removeEventListener('message', messageHandler);
              if (resolveNext) {
                resolveNext({ done: true, value: undefined });
              }
              return;
            }

            if (resolveNext) {
              resolveNext({ done: false, value: chunk });
              resolveNext = null;
            } else {
              messageQueue.push(chunk);
            }
          } catch (e) {
            console.warn('Failed to parse WebSocket message:', event.data, e);
          }
        };

        ws.addEventListener('message', messageHandler);

        while (!isCompleted) {
          if (messageQueue.length > 0) {
            yield messageQueue.shift()!;
          } else {
            await new Promise<IteratorResult<ChatStreamChunk>>((resolve) => {
              resolveNext = resolve;
            });
          }
        }

        ws.removeEventListener('message', messageHandler);
      }

      try {
        ws.send(JSON.stringify(request));
        resolve(generator());
      } catch (error) {
        reject(error);
      }
    });
  }

  /**
   * Send message using polling (fallback option)
   */
  async sendPollMessage(request: ChatRequest): Promise<ChatResponse> {
    const url = this.pollEndpoint.startsWith('http')
      ? this.pollEndpoint
      : `${this.baseUrl}${this.pollEndpoint}`;
    const response = await fetch(url, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        ...this.customHeaders,
      },
      body: JSON.stringify(request),
    });

    if (!response.ok) {
      throw new Error(`HTTP ${response.status}: ${response.statusText}`);
    }

    return response.json() as Promise<ChatResponse>;
  }

  /**
   * Generate unique IDs for conversations and requests
   */
  generateId(): string {
    return `${Date.now()}-${Math.random().toString(36).substring(2, 11)}`;
  }
}

/**
 * Default API client instance
 */
export const apiClient = new VannaApiClient();

================================================
FILE: frontends/webcomponent/src/styles/rich-component-styles.ts
================================================
import { css } from 'lit';

export const richComponentStyles = css`
  .rich-component {
    margin-bottom: var(--vanna-space-4);
    border-radius: var(--vanna-border-radius-lg);
    background: var(--vanna-background-default);
    border: 1px solid var(--vanna-outline-default);
    box-shadow: var(--vanna-shadow-sm);
    transition: box-shadow var(--vanna-duration-200) ease;
    font-family: var(--vanna-font-family-default);
  }

  .rich-component:hover {
    box-shadow: var(--vanna-shadow-md);
  }

  /* Shared typography */
  .rich-component h3,
  .rich-component h4 {
    margin: 0;
    color: var(--vanna-foreground-default);
    font-weight: 600;
  }

  .rich-component p,
  .rich-component span,
  .rich-component div {
    color: var(--vanna-foreground-default);
  }

  /* Card */
  .rich-card {
    overflow: hidden;
  }

  .card-header {
    display: flex;
    align-items: center;
    gap: var(--vanna-space-3);
    padding: var(--vanna-space-4) var(--vanna-space-5);
    background: var(--vanna-background-higher);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .card-header.collapsible {
    cursor: pointer;
  }

  .card-icon {
    font-size: 1.25rem;
    display: flex;
  }

  .card-title-section {
    flex: 1;
  }

  .card-title {
    margin: 0;
    font-size: 1rem;
    color: var(--vanna-foreground-default);
  }

  .card-subtitle {
    margin: var(--vanna-space-1) 0 0 0;
    font-size: 0.875rem;
    color: var(--vanna-foreground-dimmer);
  }

  .card-status {
    padding: var(--vanna-space-1) var(--vanna-space-2);
    border-radius: var(--vanna-border-radius-md);
    font-size: 0.75rem;
    font-weight: 600;
    text-transform: uppercase;
    background: rgba(0, 123, 255, 0.15);
    color: var(--vanna-accent-primary-default);
  }

  .card-status.status-success {
    background: rgba(16, 185, 129, 0.15);
    color: var(--vanna-accent-positive-default);
  }

  .card-status.status-warning {
    background: rgba(245, 158, 11, 0.15);
    color: var(--vanna-accent-warning-default);
  }

  .card-status.status-error {
    background: rgba(239, 68, 68, 0.15);
    color: var(--vanna-accent-negative-default);
  }

  .card-toggle {
    margin-left: var(--vanna-space-2);
    border: none;
    background: none;
    cursor: pointer;
    color: var(--vanna-foreground-dimmer);
    font-size: 1rem;
    padding: var(--vanna-space-1);
    border-radius: var(--vanna-border-radius-sm);
    transition: background-color var(--vanna-duration-200) ease;
  }

  .card-toggle:hover {
    background: var(--vanna-background-root);
  }

  .card-content {
    padding: var(--vanna-space-4) var(--vanna-space-5);
    line-height: 1.6;
    transition: all var(--vanna-duration-200) ease;
  }

  .card-content.collapsed {
    max-height: 0;
    padding-top: 0;
    padding-bottom: 0;
    overflow: hidden;
  }

  .card-actions {
    padding: var(--vanna-space-3) var(--vanna-space-5);
    background: var(--vanna-background-root);
    border-top: 1px solid var(--vanna-outline-default);
    display: flex;
    gap: var(--vanna-space-2);
  }

  .card-action {
    padding: var(--vanna-space-2) var(--vanna-space-4);
    border-radius: var(--vanna-border-radius-md);
    border: 1px solid var(--vanna-outline-default);
    background: var(--vanna-background-default);
    color: var(--vanna-foreground-default);
    cursor: pointer;
    font-size: 0.875rem;
    font-weight: 500;
    transition: all var(--vanna-duration-200) ease;
  }

  .card-action:hover {
    background: var(--vanna-background-higher);
  }

  .card-action.primary {
    background: var(--vanna-accent-primary-default);
    border-color: var(--vanna-accent-primary-default);
    color: white;
  }

  .card-action.primary:hover {
    background: var(--vanna-accent-primary-stronger);
  }

  /* Task list */
  .rich-task-list {
    padding-bottom: var(--vanna-space-2);
  }

  .task-list-header {
    padding: var(--vanna-space-4) var(--vanna-space-5);
    background: var(--vanna-background-higher);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .task-list-title {
    margin-bottom: var(--vanna-space-3);
    font-size: 1rem;
  }

  .task-list-progress {
    display: flex;
    align-items: center;
    gap: var(--vanna-space-3);
  }

  .task-list-progress .progress-text {
    font-size: 0.875rem;
    color: var(--vanna-foreground-dimmer);
    min-width: fit-content;
  }

  .task-list-progress .progress-bar {
    flex: 1;
    height: 6px;
    background: var(--vanna-background-root);
    border-radius: 3px;
    overflow: hidden;
  }

  .task-list-progress .progress-fill {
    height: 100%;
    background: var(--vanna-accent-primary-default);
    border-radius: 3px;
    transition: width var(--vanna-duration-300) ease;
  }

  .task-list-items {
    padding: var(--vanna-space-4) var(--vanna-space-5);
    display: flex;
    flex-direction: column;
    gap: var(--vanna-space-3);
  }

  .task-item {
    display: flex;
    gap: var(--vanna-space-3);
    padding: var(--vanna-space-3);
    border-radius: var(--vanna-border-radius-md);
    background: var(--vanna-background-default);
    border: 1px solid var(--vanna-outline-dimmer);
  }

  .task-item.status-running {
    border-color: var(--vanna-accent-primary-default);
    box-shadow: 0 0 0 2px rgba(0, 123, 255, 0.2);
  }

  .task-item.status-completed {
    opacity: 0.85;
  }

  .task-icon {
    font-size: 1.2rem;
    margin-top: 2px;
  }

  .task-title {
    margin: 0;
    font-size: 0.95rem;
  }

  .task-description {
    margin: var(--vanna-space-1) 0 0 0;
    font-size: 0.85rem;
    color: var(--vanna-foreground-dimmer);
  }

  .task-progress {
    display: flex;
    gap: var(--vanna-space-2);
    align-items: center;
    margin-top: var(--vanna-space-2);
  }

  .task-progress-bar {
    flex: 1;
    height: 6px;
    background: var(--vanna-background-root);
    border-radius: 3px;
    overflow: hidden;
  }

  .task-progress-fill {
    height: 100%;
    background: var(--vanna-accent-primary-default);
    transition: width var(--vanna-duration-300) ease;
  }

  .task-progress-text {
    font-size: 0.75rem;
    color: var(--vanna-foreground-dimmer);
  }

  .task-timestamp {
    margin-top: var(--vanna-space-2);
    font-size: 0.75rem;
    color: var(--vanna-foreground-dimmest);
    font-variant-numeric: tabular-nums;
  }

  /* Tool execution */
  .rich-tool-execution {
    padding: var(--vanna-space-4) var(--vanna-space-5);
  }

  .tool-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: var(--vanna-space-3);
  }

  .tool-status {
    display: flex;
    align-items: center;
    gap: var(--vanna-space-2);
  }

  .tool-icon {
    font-size: 1.2rem;
  }

  .tool-name {
    font-weight: 600;
  }

  .status-badge {
    padding: 2px 8px;
    border-radius: var(--vanna-border-radius-sm);
    font-size: 0.75rem;
    text-transform: uppercase;
    background: rgba(0, 123, 255, 0.15);
    color: var(--vanna-accent-primary-default);
  }

  .status-badge.status-completed {
    background: rgba(16, 185, 129, 0.15);
    color: var(--vanna-accent-positive-default);
  }

  .status-badge.status-failed {
    background: rgba(239, 68, 68, 0.15);
    color: var(--vanna-accent-negative-default);
  }

  .tool-duration {
    font-size: 0.85rem;
    color: var(--vanna-foreground-dimmest);
  }

  .tool-progress {
    display: flex;
    align-items: center;
    gap: var(--vanna-space-3);
    margin-bottom: var(--vanna-space-3);
  }

  .tool-progress .progress-bar {
    flex: 1;
    height: 8px;
    background: var(--vanna-background-root);
    border-radius: 4px;
    overflow: hidden;
  }

  .tool-progress .progress-fill {
    height: 100%;
    background: var(--vanna-accent-primary-default);
    transition: width var(--vanna-duration-300) ease;
  }

  .tool-progress .progress-text {
    font-size: 0.8rem;
    color: var(--vanna-foreground-dimmer);
  }

  .tool-section {
    margin-top: var(--vanna-space-4);
  }

  .tool-section h4 {
    margin-bottom: var(--vanna-space-2);
    font-size: 0.9rem;
  }

  .tool-arguments,
  .tool-result,
  .tool-error {
    background: var(--vanna-background-root);
    border: 1px solid var(--vanna-outline-dimmer);
    border-radius: var(--vanna-border-radius-md);
    padding: var(--vanna-space-3);
    font-family: var(--vanna-font-family-mono);
    font-size: 0.85rem;
    line-height: 1.5;
    white-space: pre-wrap;
    color: var(--vanna-foreground-default);
  }

  .tool-section.error .tool-error {
    border-color: var(--vanna-accent-negative-default);
    background: rgba(239, 68, 68, 0.1);
  }

  .tool-logs {
    display: flex;
    flex-direction: column;
    gap: var(--vanna-space-2);
    max-height: 200px;
    overflow-y: auto;
    padding-right: 4px;
  }

  .log-entry {
    display: flex;
    gap: var(--vanna-space-2);
    font-size: 0.85rem;
    color: var(--vanna-foreground-default);
  }

  .log-entry .log-timestamp {
    font-family: var(--vanna-font-family-mono);
    color: var(--vanna-foreground-dimmest);
    min-width: 110px;
  }

  .log-entry .log-level {
    color: var(--vanna-foreground-dimmer);
  }

  .log-entry.log-error .log-level {
    color: var(--vanna-accent-negative-default);
  }

  .log-entry.log-warning .log-level {
    color: var(--vanna-accent-warning-default);
  }

  /* Progress bar */
  .rich-progress-bar {
    padding: var(--vanna-space-4) var(--vanna-space-5);
  }

  .progress-header {
    display: flex;
    justify-content: space-between;
    font-size: 0.85rem;
    margin-bottom: var(--vanna-space-3);
    color: var(--vanna-foreground-dimmer);
  }

  .progress-track {
    position: relative;
    height: 10px;
    background: var(--vanna-background-root);
    border-radius: 5px;
    overflow: hidden;
  }

  .progress-fill {
    height: 100%;
    background: var(--vanna-accent-primary-default);
    transition: width var(--vanna-duration-300) ease;
  }

  .progress-fill.animated {
    animation: progressPulse 2s ease-in-out infinite;
  }

  .progress-fill.status-success {
    background: var(--vanna-accent-positive-default);
  }

  .progress-fill.status-error {
    background: var(--vanna-accent-negative-default);
  }

  .progress-fill.status-warning {
    background: var(--vanna-accent-warning-default);
  }

  @keyframes progressPulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.6; }
  }

  /* Notifications */
  .rich-notification {
    padding: var(--vanna-space-4) var(--vanna-space-5);
    font-family: var(--vanna-font-family-default);
  }

  .notification-content {
    display: flex;
    gap: var(--vanna-space-3);
    align-items: flex-start;
    position: relative;
  }

  .notification-content.level-info {
    border-left: 4px solid var(--vanna-accent-primary-default);
    padding-left: var(--vanna-space-3);
  }

  .notification-content.level-success {
    border-left: 4px solid var(--vanna-accent-positive-default);
    padding-left: var(--vanna-space-3);
  }

  .notification-content.level-warning {
    border-left: 4px solid var(--vanna-accent-warning-default);
    padding-left: var(--vanna-space-3);
  }

  .notification-content.level-error {
    border-left: 4px solid var(--vanna-accent-negative-default);
    padding-left: var(--vanna-space-3);
  }

  .notification-icon {
    font-size: 1.5rem;
    line-height: 1;
  }

  .notification-body {
    flex: 1;
    padding-right: var(--vanna-space-6);
  }

  .notification-title {
    margin-bottom: var(--vanna-space-2);
    font-size: 0.95rem;
    font-weight: 600;
    color: var(--vanna-foreground-default);
  }

  .notification-message {
    margin: 0;
    font-size: 0.875rem;
    line-height: 1.5;
    color: var(--vanna-foreground-dimmer);
  }

  .notification-actions {
    margin-top: var(--vanna-space-3);
    display: flex;
    gap: var(--vanna-space-2);
    flex-wrap: wrap;
  }

  .notification-action {
    padding: var(--vanna-space-2) var(--vanna-space-4);
    border-radius: var(--vanna-border-radius-md);
    border: 1px solid var(--vanna-outline-default);
    background: transparent;
    color: var(--vanna-foreground-default);
    cursor: pointer;
    transition: background var(--vanna-duration-200) ease;
    font-size: 0.875rem;
  }

  .notification-action:hover {
    background: var(--vanna-background-higher);
  }

  .notification-action.primary {
    background: var(--vanna-accent-primary-default);
    border-color: var(--vanna-accent-primary-default);
    color: white;
  }

  .notification-action.primary:hover {
    background: var(--vanna-accent-primary-stronger);
  }

  .notification-action.secondary {
    background: var(--vanna-background-default);
  }

  .notification-dismiss {
    position: absolute;
    top: 0;
    right: 0;
    background: none;
    border: none;
    color: var(--vanna-foreground-dimmer);
    font-size: 1.2rem;
    cursor: pointer;
    padding: var(--vanna-space-1);
    line-height: 1;
    transition: color var(--vanna-duration-200) ease;
  }

  .notification-dismiss:hover {
    color: var(--vanna-foreground-default);
  }

  /* Status indicator */
  .rich-status-indicator {
    padding: var(--vanna-space-3) var(--vanna-space-4);
    font-family: var(--vanna-font-family-default);
  }

  .status-indicator-content {
    display: inline-flex;
    align-items: center;
    gap: var(--vanna-space-2);
    padding: var(--vanna-space-2) var(--vanna-space-3);
    border-radius: var(--vanna-border-radius-md);
    font-size: 0.85rem;
    font-weight: 500;
    background: rgba(0, 123, 255, 0.12);
    color: var(--vanna-accent-primary-default);
  }

  .status-indicator-content.status-success {
    background: rgba(16, 185, 129, 0.12);
    color: var(--vanna-accent-positive-default);
  }

  .status-indicator-content.status-error {
    background: rgba(239, 68, 68, 0.12);
    color: var(--vanna-accent-negative-default);
  }

  .status-indicator-content.status-warning {
    background: rgba(245, 158, 11, 0.12);
    color: var(--vanna-accent-warning-default);
  }

  .status-indicator-content.status-info {
    background: rgba(0, 123, 255, 0.12);
    color: var(--vanna-accent-primary-default);
  }

  .status-indicator-content.pulse {
    animation: statusPulse 1.4s ease-in-out infinite;
  }

  .status-icon {
    font-size: 1.1rem;
  }

  @keyframes statusPulse {
    0%, 100% { box-shadow: 0 0 0 0 rgba(0, 123, 255, 0.35); }
    50% { box-shadow: 0 0 0 4px rgba(0, 123, 255, 0); }
  }

  /* Text components */
  .text-markdown {
    padding-left: var(--vanna-space-4);
    line-height: 1.6;
    font-family: var(--vanna-font-family-default);
  }

  .text-markdown h1,
  .text-markdown h2,
  .text-markdown h3,
  .text-markdown h4,
  .text-markdown h5,
  .text-markdown h6 {
    margin: var(--vanna-space-3) 0 var(--vanna-space-2) 0;
    color: var(--vanna-foreground-default);
  }

  .text-markdown h1:first-child,
  .text-markdown h2:first-child,
  .text-markdown h3:first-child,
  .text-markdown h4:first-child,
  .text-markdown h5:first-child,
  .text-markdown h6:first-child {
    margin-top: 0;
  }

  .text-markdown p {
    margin: var(--vanna-space-2) 0;
    color: var(--vanna-foreground-default);
  }

  .text-markdown ul,
  .text-markdown ol {
    margin: var(--vanna-space-2) 0;
    padding-left: var(--vanna-space-5);
  }

  .text-markdown li {
    margin: var(--vanna-space-1) 0;
    color: var(--vanna-foreground-default);
  }

  .text-markdown code {
    background: var(--vanna-background-root);
    border: 1px solid var(--vanna-outline-dimmer);
    border-radius: var(--vanna-border-radius-sm);
    padding: 2px 4px;
    font-family: var(--vanna-font-family-mono);
    font-size: 0.9em;
    color: var(--vanna-foreground-default);
  }

  .text-markdown pre {
    background: var(--vanna-background-root);
    border: 1px solid var(--vanna-outline-dimmer);
    border-radius: var(--vanna-border-radius-md);
    padding: var(--vanna-space-3);
    overflow-x: auto;
    margin: var(--vanna-space-3) 0;
  }

  .text-markdown pre code {
    background: none;
    border: none;
    padding: 0;
  }

  /* Chart */
  .rich-chart {
    padding: var(--vanna-space-4);
  }

  .chart-header {
    margin-bottom: var(--vanna-space-3);
  }

  .chart-title {
    font-size: 1.125rem;
    font-weight: 600;
    color: var(--vanna-foreground-default);
    margin: 0;
  }

  .chart-content {
    min-height: 300px;
  }

  .chart-error {
    padding: var(--vanna-space-4);
    background: var(--vanna-accent-negative-subtle);
    border-radius: var(--vanna-border-radius-md);
    color: var(--vanna-accent-negative-default);
  }

  .chart-error pre {
    margin-top: var(--vanna-space-2);
    padding: var(--vanna-space-2);
    background: var(--vanna-background-lower);
    border-radius: var(--vanna-border-radius-sm);
    font-size: 0.75rem;
    overflow-x: auto;
  }

  /* DataFrameComponent */
  .rich-dataframe {
    overflow: hidden;
    font-family: var(--vanna-font-family-default);
  }

  .dataframe-header {
    padding: var(--vanna-space-4) var(--vanna-space-5);
    background: var(--vanna-background-higher);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .dataframe-title {
    margin: 0 0 var(--vanna-space-2) 0;
    font-size: 1rem;
    color: var(--vanna-foreground-default);
  }

  .dataframe-description {
    margin: 0 0 var(--vanna-space-3) 0;
    font-size: 0.875rem;
    color: var(--vanna-foreground-dimmer);
  }

  .dataframe-meta {
    display: flex;
    gap: var(--vanna-space-4);
    font-size: 0.75rem;
    color: var(--vanna-foreground-dimmest);
  }

  .dataframe-actions {
    padding: var(--vanna-space-3) var(--vanna-space-5);
    background: var(--vanna-background-default);
    border-bottom: 1px solid var(--vanna-outline-dimmer);
    display: flex;
    justify-content: space-between;
    align-items: center;
    gap: var(--vanna-space-3);
  }

  .dataframe-search {
    flex: 1;
    max-width: 300px;
  }

  .search-input {
    width: 100%;
    padding: var(--vanna-space-2) var(--vanna-space-3);
    border: 1px solid var(--vanna-outline-default);
    border-radius: var(--vanna-border-radius-md);
    background: var(--vanna-background-default);
    color: var(--vanna-foreground-default);
    font-size: 0.875rem;
    transition: border-color var(--vanna-duration-200) ease;
  }

  .search-input:focus {
    outline: none;
    border-color: var(--vanna-accent-primary-default);
    box-shadow: 0 0 0 2px rgba(0, 123, 255, 0.2);
  }

  .export-btn {
    padding: var(--vanna-space-2) var(--vanna-space-3);
    border: 1px solid var(--vanna-outline-default);
    border-radius: var(--vanna-border-radius-md);
    background: var(--vanna-background-default);
    color: var(--vanna-foreground-default);
    cursor: pointer;
    font-size: 0.875rem;
    transition: all var(--vanna-duration-200) ease;
  }

  .export-btn:hover {
    background: var(--vanna-background-higher);
    border-color: var(--vanna-accent-primary-default);
  }

  .dataframe-table-container {
    max-height: 600px;
    overflow: auto;
    border: 1px solid var(--vanna-outline-dimmer);
    border-radius: var(--vanna-border-radius-md);
    margin: var(--vanna-space-4) 0;
  }

  .dataframe-table {
    width: 100%;
    border-collapse: collapse;
    font-size: 0.875rem;
    font-family: var(--vanna-font-family-default);
    table-layout: auto;
  }

  .dataframe-table.bordered {
    border: 1px solid var(--vanna-outline-dimmer);
  }

  .dataframe-table.compact th,
  .dataframe-table.compact td {
    padding: var(--vanna-space-1) var(--vanna-space-2);
  }

  .dataframe-table th {
    background: var(--vanna-background-higher);
    color: var(--vanna-foreground-default);
    font-weight: 600;
    text-align: left;
    padding: var(--vanna-space-3) var(--vanna-space-4);
    border-bottom: 2px solid var(--vanna-outline-default);
    position: sticky;
    top: 0;
    z-index: 1;
  }

  .dataframe-table th.sortable {
    cursor: pointer;
    user-select: none;
    transition: background-color var(--vanna-duration-200) ease;
  }

  .dataframe-table th.sortable:hover {
    background: var(--vanna-background-root);
  }

  .dataframe-table th .sort-indicator {
    margin-left: var(--vanna-space-2);
    color: var(--vanna-foreground-dimmer);
    font-size: 0.8rem;
  }

  .dataframe-table td {
    padding: var(--vanna-space-3) var(--vanna-space-4);
    border-bottom: 1px solid var(--vanna-outline-dimmer);
    color: var(--vanna-foreground-default);
  }

  .dataframe-table.striped tbody tr:nth-child(even) {
    background: rgba(255, 255, 255, 0.02);
  }

  .dataframe-table tbody tr:hover {
    background: var(--vanna-background-higher);
  }

  .dataframe-table .cell-number {
    text-align: right;
    font-family: var(--vanna-font-family-mono);
  }

  .dataframe-table .cell-boolean {
    text-align: center;
    font-weight: 600;
  }

  .dataframe-table .cell-date {
    font-family: var(--vanna-font-family-mono);
  }

  .dataframe-table .null-value {
    color: var(--vanna-foreground-dimmest);
    font-style: italic;
  }

  .dataframe-truncated {
    padding: var(--vanna-space-3) var(--vanna-space-5);
    text-align: center;
    color: var(--vanna-foreground-dimmer);
    background: var(--vanna-background-root);
    border-top: 1px solid var(--vanna-outline-dimmer);
    font-size: 0.875rem;
  }

  .dataframe-empty {
    padding: var(--vanna-space-8) var(--vanna-space-5);
    text-align: center;
    color: var(--vanna-foreground-dimmer);
  }

  .dataframe-empty p {
    margin: 0;
    font-size: 0.875rem;
  }

  /* Primitive Component Styles */

  /* Status Card */
  .rich-status-card {
    overflow: hidden;
  }

  .status-card-header {
    display: flex;
    align-items: center;
    gap: var(--vanna-space-3);
    padding: var(--vanna-space-4) var(--vanna-space-5);
    background: var(--vanna-background-higher);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .status-card-header.collapsible {
    cursor: pointer;
  }

  .status-card-icon {
    font-size: 1.25rem;
    display: flex;
    align-items: center;
  }

  .status-card-title-section {
    flex: 1;
    display: flex;
    align-items: center;
    gap: var(--vanna-space-3);
  }

  .status-card-title {
    margin: 0;
    font-size: 1rem;
    font-weight: 600;
    color: var(--vanna-foreground-default);
  }

  .status-card-badge {
    padding: var(--vanna-space-1) var(--vanna-space-2);
    border-radius: var(--vanna-border-radius-md);
    font-size: 0.75rem;
    font-weight: 600;
    text-transform: uppercase;
  }

  .status-card-badge.status-pending {
    background: var(--vanna-background-root);
    color: var(--vanna-foreground-dimmer);
  }

  .status-card-badge.status-running {
    background: rgba(59, 130, 246, 0.1);
    color: rgb(37, 99, 235);
  }

  .status-card-badge.status-success,
  .status-card-badge.status-completed {
    background: rgba(16, 185, 129, 0.1);
    color: rgb(5, 150, 105);
  }

  .status-card-badge.status-error,
  .status-card-badge.status-failed {
    background: rgba(239, 68, 68, 0.1);
    color: rgb(220, 38, 38);
  }

  .status-card-badge.status-warning {
    background: rgba(245, 158, 11, 0.1);
    color: rgb(217, 119, 6);
  }

  .status-card-content {
    padding: var(--vanna-space-4) var(--vanna-space-5);
    line-height: 1.5;
    transition: all var(--vanna-duration-200) ease;
    overflow: hidden;
  }

  .status-card-content.collapsed {
    max-height: 0;
    padding-top: 0;
    padding-bottom: 0;
  }

  .status-card-metadata {
    border-top: 1px solid var(--vanna-outline-default);
    margin: 0;
  }

  .status-card-metadata-summary {
    padding: var(--vanna-space-3) var(--vanna-space-5);
    cursor: pointer;
    font-size: 0.875rem;
    font-weight: 500;
    color: var(--vanna-foreground-dimmer);
    user-select: none;
    transition: background var(--vanna-duration-200) ease;
  }

  .status-card-metadata-summary:hover {
    background: var(--vanna-background-higher);
  }

  .status-card-metadata-content {
    padding: var(--vanna-space-3) var(--vanna-space-5);
    background: var(--vanna-background-root);
  }

  .metadata-table {
    width: 100%;
    border-collapse: collapse;
    font-size: 0.875rem;
  }

  .metadata-table thead {
    background: var(--vanna-background-higher);
  }

  .metadata-table th {
    text-align: left;
    padding: var(--vanna-space-2) var(--vanna-space-3);
    font-weight: 600;
    color: var(--vanna-foreground-dimmer);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .metadata-table td {
    padding: var(--vanna-space-2) var(--vanna-space-3);
    border-bottom: 1px solid var(--vanna-outline-dimmer);
    vertical-align: top;
  }

  .metadata-table tbody tr:last-child td {
    border-bottom: none;
  }

  .metadata-key {
    font-weight: 500;
    color: var(--vanna-foreground-default);
    width: 30%;
  }

  .metadata-value {
    color: var(--vanna-foreground-default);
    font-family: ui-monospace, SFMono-Regular, 'SF Mono', Menlo, Monaco, Consolas, 'Courier New', monospace;
  }

  .metadata-string {
    color: var(--vanna-foreground-default);
  }

  .metadata-number {
    color: rgb(37, 99, 235);
  }

  .metadata-boolean {
    color: rgb(124, 58, 237);
  }

  .metadata-null,
  .metadata-undefined {
    color: var(--vanna-foreground-dimmer);
    font-style: italic;
  }

  .metadata-json {
    margin: 0;
    padding: var(--vanna-space-2);
    background: var(--vanna-background-default);
    border: 1px solid var(--vanna-outline-default);
    border-radius: var(--vanna-border-radius-sm);
    font-size: 0.813rem;
    line-height: 1.5;
    overflow-x: auto;
  }

  /* Progress Display */
  .rich-progress-display .progress-display-container {
    padding: var(--vanna-space-4);
  }

  .progress-display-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: var(--vanna-space-3);
  }

  .progress-display-label {
    font-weight: 500;
  }

  .progress-display-percentage {
    font-size: 0.875rem;
    color: var(--vanna-foreground-dimmer);
    font-weight: 600;
  }

  .progress-display-track {
    height: 12px;
    background: var(--vanna-background-root);
    border-radius: 6px;
    overflow: hidden;
    border: 1px solid var(--vanna-outline-default);
  }

  .progress-display-fill {
    height: 100%;
    background: var(--vanna-accent-primary-default);
    border-radius: 6px;
    transition: width var(--vanna-duration-300) ease;
    position: relative;
    overflow: hidden;
  }

  .progress-display-fill.animated {
    animation: progressPulse 2s ease-in-out infinite;
  }

  .progress-display-fill.status-success {
    background: var(--vanna-accent-positive-default);
  }

  .progress-display-fill.status-warning {
    background: var(--vanna-accent-warning-default);
  }

  .progress-display-fill.status-error {
    background: var(--vanna-accent-negative-default);
  }

  .progress-display-description {
    margin-top: var(--vanna-space-2);
    font-size: 0.875rem;
    color: var(--vanna-foreground-dimmer);
    line-height: 1.4;
  }

  /* Log Viewer */
  .rich-log-viewer .log-viewer-container {
    overflow: hidden;
  }

  .log-viewer-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    padding: var(--vanna-space-4) var(--vanna-space-5);
    background: var(--vanna-background-higher);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .log-viewer-title {
    margin: 0;
    font-size: 1rem;
    font-weight: 600;
  }

  .log-viewer-search {
    display: flex;
    gap: var(--vanna-space-2);
  }

  .log-search-input {
    padding: var(--vanna-space-2);
    border: 1px solid var(--vanna-outline-default);
    border-radius: var(--vanna-border-radius-md);
    background: var(--vanna-background-default);
    color: var(--vanna-foreground-default);
    font-size: 0.875rem;
  }

  .log-viewer-content {
    max-height: 300px;
    overflow-y: auto;
    padding: var(--vanna-space-4);
  }

  .log-viewer-content.auto-scroll {
    scroll-behavior: smooth;
  }

  .log-entry {
    display: flex;
    gap: var(--vanna-space-2);
    padding: var(--vanna-space-2) 0;
    font-family: var(--vanna-font-family-mono);
    font-size: 0.875rem;
    line-height: 1.4;
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .log-entry:last-child {
    border-bottom: none;
  }

  .log-timestamp {
    color: var(--vanna-foreground-dimmer);
    white-space: nowrap;
  }

  .log-level {
    font-weight: 600;
    white-space: nowrap;
  }

  .log-entry.log-info .log-level {
    color: var(--vanna-accent-primary-default);
  }

  .log-entry.log-error .log-level {
    color: var(--vanna-accent-negative-default);
  }

  .log-entry.log-warning .log-level {
    color: var(--vanna-accent-warning-default);
  }

  .log-entry.log-debug .log-level {
    color: var(--vanna-foreground-dimmer);
  }

  .log-message {
    flex: 1;
    word-break: break-word;
  }

  /* Badge */
  .rich-badge {
    display: inline-flex;
    align-items: center;
    gap: var(--vanna-space-1);
    padding: var(--vanna-space-1) var(--vanna-space-2);
    border-radius: var(--vanna-border-radius-full);
    font-size: 0.75rem;
    font-weight: 600;
    text-transform: uppercase;
    letter-spacing: 0.025em;
  }

  .rich-badge.badge-small {
    padding: 2px var(--vanna-space-1);
    font-size: 0.625rem;
  }

  .rich-badge.badge-large {
    padding: var(--vanna-space-2) var(--vanna-space-3);
    font-size: 0.875rem;
  }

  .rich-badge.badge-default {
    background: var(--vanna-background-root);
    color: var(--vanna-foreground-dimmer);
  }

  .rich-badge.badge-primary {
    background: var(--vanna-accent-primary-default);
    color: white;
  }

  .rich-badge.badge-success {
    background: var(--vanna-accent-positive-default);
    color: white;
  }

  .rich-badge.badge-warning {
    background: var(--vanna-accent-warning-default);
    color: white;
  }

  .rich-badge.badge-error {
    background: var(--vanna-accent-negative-default);
    color: white;
  }

  /* Icon Text */
  .rich-icon-text {
    display: inline-flex;
    align-items: center;
    gap: var(--vanna-space-2);
  }

  .rich-icon-text.icon-text-small {
    font-size: 0.875rem;
    gap: var(--vanna-space-1);
  }

  .rich-icon-text.icon-text-large {
    font-size: 1.125rem;
    gap: var(--vanna-space-3);
  }

  .rich-icon-text.icon-text-center {
    justify-content: center;
  }

  .rich-icon-text.icon-text-right {
    justify-content: flex-end;
  }

  .icon-text-icon {
    display: flex;
    align-items: center;
  }

  .rich-icon-text.icon-text-primary {
    color: var(--vanna-accent-primary-default);
  }

  .rich-icon-text.icon-text-secondary {
    color: var(--vanna-foreground-dimmer);
  }

  .rich-icon-text.icon-text-muted {
    color: var(--vanna-foreground-dimmest);
  }

  /* Artifact Component Styles */
  .rich-artifact {
    overflow: hidden;
  }

  .artifact-header {
    display: flex;
    justify-content: space-between;
    align-items: flex-start;
    padding: var(--vanna-space-4) var(--vanna-space-5);
    background: var(--vanna-background-subtle);
    border-bottom: 1px solid var(--vanna-outline-default);
  }

  .artifact-meta {
    flex: 1;
  }

  .artifact-title {
    margin: 0 0 var(--vanna-space-2) 0;
    font-size: 1.1rem;
    font-weight: 600;
    color: var(--vanna-foreground-default);
  }

  .artifact-description {
    margin: 0 0 var(--vanna-space-3) 0;
    color: var(--vanna-foreground-dimmer);
    font-size: 0.9rem;
  }

  .artifact-type-badge {
    display: inline-block;
    padding: var(--vanna-space-1) var(--vanna-space-2);
    background: var(--vanna-accent-primary-subtle);
    color: var(--vanna-accent-primary-default);
    border-radius: var(--vanna-border-radius-sm);
    font-size: 0.75rem;
    font-weight: 500;
    text-transform: uppercase;
  }

  .artifact-controls {
    display: flex;
    gap: var(--vanna-space-2);
  }

  .artifact-btn {
    padding: var(--vanna-space-2);
    background: var(--vanna-background-default);
    border: 1px solid var(--vanna-outline-default);
    border-radius: var(--vanna-border-radius-sm);
    cursor: pointer;
    font-size: 1rem;
    transition: all var(--vanna-duration-200) ease;
  }

  .artifact-btn:hover {
    background: var(--vanna-background-subtle);
    border-color: var(--vanna-outline-hover);
  }

  .artifact-btn:active {
    transform: translateY(1px);
  }

  .artifact-preview {
    height: 300px;
    background: var(--vanna-background-default);
  }

  .artifact-iframe {
    width: 100%;
    height: 100%;
    border: none;
    display: block;
  }

  /* Fullscreen overlay styles */
  .artifact-fullscreen-overlay {
    position: fixed !important;
    top: 0 !important;
    left: 0 !important;
    width: 100vw !important;
    height: 100vh !important;
    background: var(--vanna-background-default) !important;
    z-index: 10000 !important;
    display: flex !important;
    flex-direction: column !important;
  }

  .fullscreen-header {
    padding: var(--vanna-space-4) !important;
    border-bottom: 1px solid var(--vanna-outline-default) !important;
    display: flex !important;
    justify-content: space-between !important;
    align-items: center !important;
    background: var(--vanna-background-subtle) !important;
  }

  .fullscreen-header h3 {
    margin: 0 !important;
    color: var(--vanna-foreground-default) !important;
  }

  .close-fullscreen {
    padding: var(--vanna-space-2) var(--vanna-space-3) !important;
    background: var(--vanna-background-default) !important;
    border: 1px solid var(--vanna-outline-default) !important;
    border-radius: var(--vanna-border-radius-sm) !important;
    cursor: pointer !important;
    font-size: 1.2rem !important;
    line-height: 1 !important;
  }

  .close-fullscreen:hover {
    background: var(--vanna-background-subtle) !important;
  }

  .fullscreen-content {
    flex: 1 !important;
    padding: var(--vanna-space-4) !important;
    overflow: hidden !important;
  }

  .fullscreen-iframe {
    width: 100% !important;
    height: 100% !important;
    border: none !important;
    border-radius: var(--vanna-border-radius-md) !important;
  }

  /* Artifact placeholder styles */
  .artifact-placeholder {
    padding: var(--vanna-space-4);
    background: var(--vanna-background-subtle);
    border: 2px dashed var(--vanna-outline-default);
    border-radius: var(--vanna-border-radius-md);
    text-align: center;
  }

  .placeholder-content {
    display: flex;
    align-items: center;
    justify-content: center;
    gap: var(--vanna-space-3);
    opacity: 0.8;
  }

  .placeholder-icon {
    font-size: 1.5rem;
  }

  .placeholder-text {
    text-align: left;
  }

  .placeholder-text strong {
    color: var(--vanna-foreground-default);
    font-weight: 600;
  }

  .placeholder-type {
    font-size: 0.8rem;
    color: var(--vanna-foreground-dimmer);
    text-transform: uppercase;
    margin-top: var(--vanna-space-1);
  }

  .placeholder-reopen {
    padding: var(--vanna-space-2);
    background: var(--vanna-accent-primary-default);
    color: white;
    border: none;
    border-radius: var(--vanna-border-radius-sm);
    cursor: pointer;
    font-size: 1rem;
    transition: background var(--vanna-duration-200) ease;
  }

  .placeholder-reopen:hover {
    background: var(--vanna-accent-primary-hover);
  }

  /* Button Component */
  .rich-button {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    gap: var(--vanna-space-2);
    padding: var(--vanna-space-2) var(--vanna-space-4);
    border-radius: var(--vanna-border-radius-md);
    border: 1px solid;
    font-size: 0.875rem;
    font-weight: 500;
    cursor: pointer;
    transition: all var(--vanna-duration-200) ease;
    white-space: nowrap;
    user-select: none;
    font-family: var(--vanna-font-family-default);
  }

  .rich-button:disabled {
    cursor: not-allowed;
    opacity: 0.5;
  }

  /* Button variants */
  .rich-button.button-primary {
    background: var(--vanna-accent-primary-default);
    border-color: var(--vanna-accent-primary-default);
    color: white;
  }

  .rich-button.button-primary:hover:not(:disabled) {
    background: var(--vanna-accent-primary-stronger);
    border-color: var(--vanna-accent-primary-stronger);
  }

  .rich-button.button-secondary {
    background: var(--vanna-background-default);
    border-color: var(--vanna-outline-default);
    color: var(--vanna-foreground-default);
  }

  .rich-button.button-secondary:hover:not(:disabled) {
    background: var(--vanna-background-higher);
    border-color: var(--vanna-outline-hover);
  }

  .rich-button.button-success {
    background: var(--vanna-accent-positive-default);
    border-color: var(--vanna-accent-positive-default);
    color: white;
  }

  .rich-button.button-success:hover:not(:disabled) {
    background: var(--vanna-accent-positive-stronger);
  }

  .rich-button.button-warning {
    background: var(--vanna-accent-warning-default);
    border-color: var(--vanna-accent-warning-default);
    color: white;
  }

  .rich-button.button-warning:hover:not(:disabled) {
    background: var(--vanna-accent-warning-stronger);
  }

  .rich-button.button-error {
    background: var(--vanna-accent-negative-default);
    border-color: var(--vanna-accent-negative-default);
    color: white;
  }

  .rich-button.button-error:hover:not(:disabled) {
    background: var(--vanna-accent-negative-stronger);
  }

  .rich-button.button-ghost {
    background: transparent;
    border-color: transparent;
    color: var(--vanna-foreground-default);
  }

  .rich-button.button-ghost:hover:not(:disabled) {
    background: var(--vanna-background-higher);
  }

  .rich-button.button-link {
    background: transparent;
    border-color: transparent;
    color: var(--vanna-accent-primary-default);
    text-decoration: underline;
    padding: var(--vanna-space-1) var(--vanna-space-2);
  }

  .rich-button.button-link:hover:not(:disabled) {
    color: var(--vanna-accent-primary-stronger);
  }

  /* Button sizes */
  .rich-button.button-small {
    padding: var(--vanna-space-1) var(--vanna-space-3);
    font-size: 0.75rem;
    gap: var(--vanna-space-1);
  }

  .rich-button.button-medium {
    padding: var(--vanna-space-2) var(--vanna-space-4);
    font-size: 0.875rem;
    gap: var(--vanna-space-2);
  }

  .rich-button.button-large {
    padding: var(--vanna-space-3) var(--vanna-space-5);
    font-size: 1rem;
    gap: var(--vanna-space-2);
  }

  /* Button modifiers */
  .rich-button.button-full-width {
    width: 100%;
  }

  .rich-button.button-loading {
    position: relative;
    pointer-events: none;
  }

  .button-spinner {
    display: inline-flex;
    animation: spin 1s linear infinite;
  }

  @keyframes spin {
    from { transform: rotate(0deg); }
    to { transform: rotate(360deg); }
  }

  .button-icon {
    display: inline-flex;
    align-items: center;
    font-size: 1em;
  }

  .button-label {
    display: inline-flex;
    align-items: center;
  }

  /* Button Group Component */
  .rich-button-group {
    display: flex;
    gap: var(--vanna-space-2);
    font-family: var(--vanna-font-family-default);
  }

  .rich-button-group.button-group-vertical {
    flex-direction: column;
  }

  .rich-button-group.button-group-horizontal {
    flex-direction: row;
  }

  .rich-button-group.button-group-spacing-small {
    gap: var(--vanna-space-1);
  }

  .rich-button-group.button-group-spacing-medium {
    gap: var(--vanna-space-2);
  }

  .rich-button-group.button-group-spacing-large {
    gap: var(--vanna-space-4);
  }

  .rich-button-group.button-group-align-left {
    justify-content: flex-start;
  }

  .rich-button-group.button-group-align-center {
    justify-content: center;
  }

  .rich-button-group.button-group-align-right {
    justify-content: flex-end;
  }

  .rich-button-group.button-group-align-space-between {
    justify-content: space-between;
  }

  .rich-button-group.button-group-full-width {
    width: 100%;
  }

  .rich-button-group.button-group-full-width > .rich-button {
    flex: 1;
  }

  /* Button Group Interactive States */
  .rich-button.button-transitioning {
    transition: all 0.2s ease-in-out;
  }

  .rich-button.button-highlighted {
    transform: scale(1.02);
    box-shadow: 0 0 0 2px var(--vanna-accent-primary-default);
    z-index: 1;
    position: relative;
  }

  .rich-button.button-grayed-out {
    opacity: 0.4;
    filter: grayscale(50%);
    transform: scale(0.98);
  }

  .rich-button.button-clicked {
    animation: buttonClickPulse 0.3s ease-out;
  }

  @keyframes buttonClickPulse {
    0% {
      transform: scale(1);
    }
    50% {
      transform: scale(1.05);
    }
    100% {
      transform: scale(1.02);
    }
  }

  /* Override hover states when in click states */
  .rich-button.button-highlighted:hover,
  .rich-button.button-grayed-out:hover {
    /* Maintain the click state even on hover */
  }

  .rich-button.button-grayed-out:hover {
    opacity: 0.4;
    filter: grayscale(50%);
  }
`;

export const richComponentStyleText = richComponentStyles.cssText;


================================================
FILE: frontends/webcomponent/src/styles/vanna-design-tokens.ts
================================================
import { css } from 'lit';

// Vanna 2.0 design tokens - Data-First Agents branding
export const vannaDesignTokens = css`
  :host {
    /* Vanna 2.0 Brand Colors */
    --vanna-navy: rgb(2, 61, 96);
    --vanna-cream: rgb(231, 225, 207);
    --vanna-teal: rgb(21, 168, 168);
    --vanna-orange: rgb(254, 93, 38);
    --vanna-magenta: rgb(191, 19, 99);

    /* Color Palette - Light mode (default) */
    --vanna-background-root: rgb(255, 255, 255);
    --vanna-background-default: rgb(231, 225, 207);
    --vanna-background-higher: rgb(244, 246, 248);
    --vanna-background-highest: rgb(229, 231, 235);
    --vanna-background-subtle: rgb(248, 250, 252);
    --vanna-background-lower: rgb(239, 242, 245);

    --vanna-foreground-default: rgb(2, 61, 96);
    --vanna-foreground-dimmer: rgb(71, 85, 105);
    --vanna-foreground-dimmest: rgb(100, 116, 139);

    --vanna-accent-primary-default: rgb(21, 168, 168);
    --vanna-accent-primary-stronger: rgb(2, 61, 96);
    --vanna-accent-primary-strongest: rgb(2, 61, 96);
    --vanna-accent-primary-subtle: rgba(21, 168, 168, 0.1);
    --vanna-accent-primary-hover: rgb(21, 168, 168);

    --vanna-accent-positive-default: rgb(21, 168, 168);
    --vanna-accent-positive-stronger: rgb(2, 61, 96);
    --vanna-accent-positive-subtle: rgba(21, 168, 168, 0.1);

    --vanna-accent-negative-default: rgb(239, 68, 68);
    --vanna-accent-negative-stronger: rgb(220, 38, 38);
    --vanna-accent-negative-subtle: rgba(239, 68, 68, 0.1);

    --vanna-accent-warning-default: rgb(254, 93, 38);
    --vanna-accent-warning-stronger: rgb(254, 93, 38);
    --vanna-accent-warning-subtle: rgba(254, 93, 38, 0.1);

    /* Outline/Border colors */
    --vanna-outline-default: rgba(21, 168, 168, 0.3);
    --vanna-outline-dimmer: rgb(241, 245, 249);
    --vanna-outline-dimmest: rgb(248, 250, 252);
    --vanna-outline-hover: rgb(21, 168, 168);

    /* Typography */
    --vanna-font-family-default: "Space Grotesk", ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif;
    --vanna-font-family-serif: "Roboto Slab", ui-serif, Georgia, serif;
    --vanna-font-family-mono: "Space Mono", ui-monospace, SFMono-Regular, "SF Mono", Monaco, Inconsolata, "Roboto Mono", "Ubuntu Mono", monospace;

    /* Spacing scale */
    --vanna-space-0: 0px;
    --vanna-space-1: 4px;
    --vanna-space-2: 8px;
    --vanna-space-3: 12px;
    --vanna-space-4: 16px;
    --vanna-space-5: 20px;
    --vanna-space-6: 24px;
    --vanna-space-7: 28px;
    --vanna-space-8: 32px;
    --vanna-space-10: 40px;
    --vanna-space-12: 48px;
    --vanna-space-16: 64px;

    /* Border radius */
    --vanna-border-radius-sm: 6px;
    --vanna-border-radius-md: 10px;
    --vanna-border-radius-lg: 14px;
    --vanna-border-radius-xl: 20px;
    --vanna-border-radius-2xl: 24px;
    --vanna-border-radius-full: 9999px;

    /* Shadows - Preline-inspired */
    --vanna-shadow-xs: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
    --vanna-shadow-sm: 0 1px 3px 0 rgba(0, 0, 0, 0.1), 0 1px 2px -1px rgba(0, 0, 0, 0.1);
    --vanna-shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -2px rgba(0, 0, 0, 0.1);
    --vanna-shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -4px rgba(0, 0, 0, 0.1);
    --vanna-shadow-xl: 0 20px 25px -5px rgba(0, 0, 0, 0.1), 0 8px 10px -6px rgba(0, 0, 0, 0.1);
    --vanna-shadow-2xl: 0 25px 50px -12px rgba(0, 0, 0, 0.25);

    /* Animation durations */
    --vanna-duration-75: 75ms;
    --vanna-duration-100: 100ms;
    --vanna-duration-150: 150ms;
    --vanna-duration-200: 200ms;
    --vanna-duration-300: 300ms;
    --vanna-duration-500: 500ms;
    --vanna-duration-700: 700ms;

    /* Z-index scale */
    --vanna-z-dropdown: 1000;
    --vanna-z-sticky: 1020;
    --vanna-z-fixed: 1030;
    --vanna-z-modal: 1040;
    --vanna-z-popover: 1050;
    --vanna-z-tooltip: 1060;

    /* Chat-specific tokens */
    --vanna-chat-bubble-radius: 18px;
    --vanna-chat-bubble-radius-sm: 12px;
    --vanna-chat-spacing: 16px;
    --vanna-chat-avatar-size: 40px;
  }

  /* Dark theme overrides */
  :host([theme="dark"]) {
    --vanna-background-root: rgb(9, 11, 17);
    --vanna-background-default: rgb(15, 18, 25);
    --vanna-background-higher: rgb(24, 29, 39);
    --vanna-background-highest: rgb(31, 39, 51);
    --vanna-background-subtle: rgb(17, 21, 28);
    --vanna-background-lower: rgb(6, 8, 12);

    --vanna-foreground-default: rgb(248, 250, 252);
    --vanna-foreground-dimmer: rgb(203, 213, 225);
    --vanna-foreground-dimmest: rgb(148, 163, 184);

    --vanna-accent-primary-default: rgb(21, 168, 168);
    --vanna-accent-primary-stronger: rgb(21, 168, 168);
    --vanna-accent-primary-strongest: rgb(2, 61, 96);
    --vanna-accent-primary-subtle: rgba(21, 168, 168, 0.15);
    --vanna-accent-primary-hover: rgb(21, 168, 168);

    --vanna-accent-positive-default: rgb(21, 168, 168);
    --vanna-accent-positive-stronger: rgb(21, 168, 168);
    --vanna-accent-positive-subtle: rgba(21, 168, 168, 0.15);

    --vanna-accent-negative-default: rgb(248, 113, 113);
    --vanna-accent-negative-stronger: rgb(239, 68, 68);
    --vanna-accent-negative-subtle: rgba(248, 113, 113, 0.15);

    --vanna-accent-warning-default: rgb(254, 93, 38);
    --vanna-accent-warning-stronger: rgb(254, 93, 38);
    --vanna-accent-warning-subtle: rgba(254, 93, 38, 0.15);

    --vanna-outline-default: rgba(21, 168, 168, 0.3);
    --vanna-outline-dimmer: rgb(31, 41, 55);
    --vanna-outline-dimmest: rgb(17, 24, 39);
    --vanna-outline-hover: rgb(21, 168, 168);

    --vanna-shadow-xs: 0 1px 2px 0 rgba(0, 0, 0, 0.6);
    --vanna-shadow-sm: 0 1px 3px 0 rgba(0, 0, 0, 0.5), 0 1px 2px -1px rgba(0, 0, 0, 0.5);
    --vanna-shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.4), 0 2px 4px -2px rgba(0, 0, 0, 0.4);
    --vanna-shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.4), 0 4px 6px -4px rgba(0, 0, 0, 0.4);
    --vanna-shadow-xl: 0 20px 25px -5px rgba(0, 0, 0, 0.3), 0 8px 10px -6px rgba(0, 0, 0, 0.3);
    --vanna-shadow-2xl: 0 25px 50px -12px rgba(0, 0, 0, 0.6);
  }
`;


================================================
FILE: frontends/webcomponent/src/vite-env.d.ts
================================================
/// <reference types="vite/client" />

declare const __BUILD_TIME__: string;
declare const __BUILD_VERSION__: string;


================================================
FILE: frontends/webcomponent/test-comprehensive.html
================================================
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Vanna Webcomponent - Comprehensive Test</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }

        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
            background: #f5f5f5;
            display: flex;
            height: 100vh;
            overflow: hidden;
        }

        .sidebar {
            width: 300px;
            background: white;
            border-right: 1px solid #e0e0e0;
            display: flex;
            flex-direction: column;
            overflow-y: auto;
        }

        .sidebar-header {
            padding: 20px;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
        }

        .sidebar-header h1 {
            font-size: 18px;
            margin-bottom: 8px;
        }

        .sidebar-header p {
            font-size: 12px;
            opacity: 0.9;
        }

        .controls {
            padding: 20px;
            border-bottom: 1px solid #e0e0e0;
        }

        .control-group {
            margin-bottom: 15px;
        }

        .control-group label {
            display: block;
            font-size: 12px;
            font-weight: 600;
            color: #666;
            margin-bottom: 5px;
            text-transform: uppercase;
            letter-spacing: 0.5px;
        }

        button {
            width: 100%;
            padding: 10px 15px;
            background: #667eea;
            color: white;
            border: none;
            border-radius: 6px;
            font-size: 14px;
            font-weight: 600;
            cursor: pointer;
            transition: all 0.2s;
        }

        button:hover {
            background: #5568d3;
            transform: translateY(-1px);
            box-shadow: 0 4px 8px rgba(102, 126, 234, 0.3);
        }

        button:active {
            transform: translateY(0);
        }

        button.secondary {
            background: #f0f0f0;
            color: #333;
        }

        button.secondary:hover {
            background: #e0e0e0;
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
        }

        select {
            width: 100%;
            padding: 8px 12px;
            border: 1px solid #e0e0e0;
            border-radius: 6px;
            font-size: 14px;
            background: white;
            cursor: pointer;
        }

        .status-section {
            padding: 20px;
        }

        .status-indicator {
            display: flex;
            align-items: center;
            gap: 10px;
            padding: 12px;
            background: #f8f9fa;
            border-radius: 6px;
            margin-bottom: 10px;
        }

        .status-dot {
            width: 12px;
            height: 12px;
            border-radius: 50%;
            background: #ccc;
        }

        .status-dot.success {
            background: #10b981;
            animation: pulse 2s infinite;
        }

        .status-dot.error {
            background: #ef4444;
            animation: pulse 2s infinite;
        }

        .status-dot.warning {
            background: #f59e0b;
        }

        @keyframes pulse {
            0%, 100% { opacity: 1; }
            50% { opacity: 0.5; }
        }

        .status-text {
            font-size: 13px;
            font-weight: 500;
            color: #333;
        }

        .console-monitor {
            padding: 20px;
            border-top: 1px solid #e0e0e0;
        }

        .console-monitor h3 {
            font-size: 14px;
            font-weight: 600;
            margin-bottom: 10px;
            color: #333;
        }

        .console-log {
            background: #1e1e1e;
            color: #d4d4d4;
            padding: 12px;
            border-radius: 6px;
            font-family: 'Courier New', monospace;
            font-size: 11px;
            max-height: 200px;
            overflow-y: auto;
        }

        .console-log .error {
            color: #f48771;
        }

        .console-log .warning {
            color: #dcdcaa;
        }

        .console-log .info {
            color: #4fc1ff;
        }

        .checklist {
            padding: 20px;
            border-top: 1px solid #e0e0e0;
        }

        .checklist h3 {
            font-size: 14px;
            font-weight: 600;
            margin-bottom: 10px;
            color: #333;
        }

        .checklist-item {
            display: flex;
            align-items: center;
            gap: 8px;
            padding: 6px 0;
            font-size: 12px;
            color: #666;
        }

        .checklist-item .check {
            width: 16px;
            height: 16px;
            border: 2px solid #e0e0e0;
            border-radius: 3px;
            display: flex;
            align-items: center;
            justify-content: center;
            font-size: 10px;
        }

        .checklist-item.checked .check {
            background: #10b981;
            border-color: #10b981;
            color: white;
        }

        .main-content {
            flex: 1;
            display: flex;
            flex-direction: column;
            overflow: hidden;
        }

        .chat-container {
            flex: 1;
            overflow: hidden;
        }

        vanna-chat {
            width: 100%;
            height: 100%;
        }

        .metrics {
            padding: 10px 20px;
            background: white;
            border-top: 1px solid #e0e0e0;
            display: flex;
            gap: 20px;
            font-size: 12px;
            color: #666;
        }

        .metric {
            display: flex;
            gap: 5px;
        }

        .metric strong {
            color: #333;
        }
    </style>
</head>
<body>
    <!-- Sidebar -->
    <div class="sidebar">
        <div class="sidebar-header">
            <h1>Component Test Suite</h1>
            <p>Comprehensive validation for webcomponent pruning</p>
        </div>

        <!-- Controls -->
        <div class="controls">
            <div class="control-group">
                <label>Test Mode</label>
                <select id="mode-select">
                    <option value="realistic">Realistic (with delays)</option>
                    <option value="rapid">Rapid (fast)</option>
                </select>
            </div>

            <div class="control-group">
                <button id="start-test">Run Comprehensive Test</button>
            </div>

            <div class="control-group">
                <button class="secondary" id="clear-chat">Clear Chat</button>
            </div>
        </div>

        <!-- Status -->
        <div class="status-section">
            <div class="status-indicator">
                <div class="status-dot" id="backend-status"></div>
                <div class="status-text" id="backend-text">Checking backend...</div>
            </div>

            <div class="status-indicator">
                <div class="status-dot" id="console-status"></div>
                <div class="status-text" id="console-text">No errors detected</div>
            </div>
        </div>

        <!-- Component Checklist -->
        <div class="checklist">
            <h3>Component Rendering</h3>
            <div class="checklist-item" data-component="text">
                <div class="check"></div>
                <span>Text Component</span>
            </div>
            <div class="checklist-item" data-component="status_card">
                <div class="check"></div>
                <span>Status Card</span>
            </div>
            <div class="checklist-item" data-component="progress_display">
                <div class="check"></div>
                <span>Progress Display</span>
            </div>
            <div class="checklist-item" data-component="card">
                <div class="check"></div>
                <span>Card</span>
            </div>
            <div class="checklist-item" data-component="task_list">
                <div class="check"></div>
                <span>Task List</span>
            </div>
            <div class="checklist-item" data-component="progress_bar">
                <div class="check"></div>
                <span>Progress Bar</span>
            </div>
            <div class="checklist-item" data-component="notification">
                <div class="check"></div>
                <span>Notification</span>
            </div>
            <div class="checklist-item" data-component="status_indicator">
                <div class="check"></div>
                <span>Status Indicator</span>
            </div>
            <div class="checklist-item" data-component="badge">
                <div class="check"></div>
                <span>Badge</span>
            </div>
            <div class="checklist-item" data-component="icon_text">
                <div class="check"></div>
                <span>Icon Text</span>
            </div>
            <div class="checklist-item" data-component="button">
                <div class="check"></div>
                <span>Button</span>
            </div>
            <div class="checklist-item" data-component="button_group">
                <div class="check"></div>
                <span>Button Group</span>
            </div>
            <div class="checklist-item" data-component="dataframe">
                <div class="check"></div>
                <span>DataFrame</span>
            </div>
            <div class="checklist-item" data-component="chart">
                <div class="check"></div>
                <span>Chart</span>
            </div>
            <div class="checklist-item" data-component="artifact">
                <div class="check"></div>
                <span>Artifact</span>
            </div>
            <div class="checklist-item" data-component="log_viewer">
                <div class="check"></div>
                <span>Log Viewer</span>
            </div>
            <!-- Note: code_block, table, container not supported by webcomponent -->
        </div>

        <!-- Console Monitor -->
        <div class="console-monitor">
            <h3>Console Log</h3>
            <div class="console-log" id="console-log">
                <div class="info">Monitoring console for errors...</div>
            </div>
        </div>
    </div>

    <!-- Main Content -->
    <div class="main-content">
        <div class="chat-container">
            <vanna-chat
                id="vanna-chat"
                api-url="http://localhost:5555"
                placeholder="Type /test to run comprehensive test..."
            ></vanna-chat>
        </div>

        <div class="metrics">
            <div class="metric">
                <strong>Components Rendered:</strong>
                <span id="component-count">0</span>
            </div>
            <div class="metric">
                <strong>Updates Processed:</strong>
                <span id="update-count">0</span>
            </div>
            <div class="metric">
                <strong>Errors:</strong>
                <span id="error-count">0</span>
            </div>
        </div>
    </div>

    <!-- Load webcomponent -->
    <script type="module" src="/static/vanna-components.js"></script>

    <script type="module">
        // State
        let componentCount = 0;
        let updateCount = 0;
        let errorCount = 0;
        const seenComponents = new Set();

        // Elements
        const vannaChat = document.getElementById('vanna-chat');
        const modeSelect = document.getElementById('mode-select');
        const startTestBtn = document.getElementById('start-test');
        const clearChatBtn = document.getElementById('clear-chat');
        const consoleLog = document.getElementById('console-log');
        const backendStatus = document.getElementById('backend-status');
        const backendText = document.getElementById('backend-text');
        const consoleStatus = document.getElementById('console-status');
        const consoleText = document.getElementById('console-text');

        // Check backend health
        async function checkBackend() {
            try {
                const response = await fetch('http://localhost:5555/health');
                const data = await response.json();
                backendStatus.className = 'status-dot success';
                backendText.textContent = `Backend ready (${data.mode} mode)`;
            } catch (error) {
                backendStatus.className = 'status-dot error';
                backendText.textContent = 'Backend not responding';
                addConsoleLog('error', `Backend health check failed: ${error.message}`);
            }
        }

        // Console monitoring
        function addConsoleLog(type, message) {
            const timestamp = new Date().toLocaleTimeString();
            const logEntry = document.createElement('div');
            logEntry.className = type;
            logEntry.textContent = `[${timestamp}] ${message}`;
            consoleLog.appendChild(logEntry);
            consoleLog.scrollTop = consoleLog.scrollHeight;

            // Update error status
            if (type === 'error') {
                errorCount++;
                consoleStatus.className = 'status-dot error';
                consoleText.textContent = `${errorCount} error(s) detected`;
                document.getElementById('error-count').textContent = errorCount;
            }
        }

        // Override console methods
        const originalError = console.error;
        const originalWarn = console.warn;
        const originalLog = console.log;

        console.error = function(...args) {
            addConsoleLog('error', args.join(' '));
            originalError.apply(console, args);
        };

        console.warn = function(...args) {
            addConsoleLog('warning', args.join(' '));
            originalWarn.apply(console, args);
        };

        console.log = function(...args) {
            const message = args.join(' ');
            if (message.includes('ERROR') || message.includes('Error')) {
                addConsoleLog('error', message);
            } else {
                addConsoleLog('info', message);
            }
            originalLog.apply(console, args);
        };

        // Monitor window errors
        window.addEventListener('error', (event) => {
            addConsoleLog('error', `${event.message} at ${event.filename}:${event.lineno}`);
            errorCount++;
        });

        // Monitor component rendering
        const observer = new MutationObserver((mutations) => {
            mutations.forEach((mutation) => {
                mutation.addedNodes.forEach((node) => {
                    if (node.nodeType === 1) { // Element node
                        const componentType = node.getAttribute('data-component-type');
                        if (componentType) {
                            // Track component
                            if (!seenComponents.has(componentType)) {
                                seenComponents.add(componentType);
                                componentCount++;
                                document.getElementById('component-count').textContent = componentCount;

                                // Check off in checklist
                                const checklistItem = document.querySelector(`[data-component="${componentType}"]`);
                                if (checklistItem) {
                                    checklistItem.classList.add('checked');
                                    checklistItem.querySelector('.check').textContent = '✓';
                                }
                            }

                            updateCount++;
                            document.getElementById('update-count').textContent = updateCount;
                        }
                    }
                });
            });
        });

        // Start observing when vanna-chat is ready
        setTimeout(() => {
            const shadowRoot = vannaChat.shadowRoot;
            if (shadowRoot) {
                const container = shadowRoot.querySelector('.rich-components-container');
                if (container) {
                    observer.observe(container, { childList: true, subtree: true });
                    addConsoleLog('info', 'Component observer started');
                }
            }
        }, 1000);

        // Event listeners
        startTestBtn.addEventListener('click', async () => {
            const mode = modeSelect.value;

            // Update backend mode
            try {
                await fetch(`http://localhost:5555/health`);
                addConsoleLog('info', `Starting comprehensive test in ${mode} mode...`);

                // Send test message through chat
                vannaChat.dispatchEvent(new CustomEvent('send-message', {
                    detail: { message: '/test' }
                }));

                // Alternative: directly trigger if API is exposed
                const inputEl = vannaChat.shadowRoot?.querySelector('textarea, input');
                if (inputEl) {
                    inputEl.value = '/test';
                    const form = vannaChat.shadowRoot?.querySelector('form');
                    if (form) {
                        form.dispatchEvent(new Event('submit', { bubbles: true }));
                    }
                }
            } catch (error) {
                addConsoleLog('error', `Failed to start test: ${error.message}`);
            }
        });

        clearChatBtn.addEventListener('click', () => {
            // Reset state
            componentCount = 0;
            updateCount = 0;
            errorCount = 0;
            seenComponents.clear();

            document.getElementById('component-count').textContent = '0';
            document.getElementById('update-count').textContent = '0';
            document.getElementById('error-count').textContent = '0';

            // Uncheck all checklist items
            document.querySelectorAll('.checklist-item').forEach(item => {
                item.classList.remove('checked');
                item.querySelector('.check').textContent = '';
            });

            // Clear console
            consoleLog.innerHTML = '<div class="info">Console cleared</div>';
            consoleStatus.className = 'status-dot';
            consoleText.textContent = 'No errors detected';

            // Reload page to truly clear (vanna-chat doesn't expose clear method)
            location.reload();
        });

        // Initial backend check
        checkBackend();
        setInterval(checkBackend, 5000);

        // Log startup
        addConsoleLog('info', 'Test suite initialized');
        addConsoleLog('info', 'Ensure backend is running: python test_backend.py');
    </script>
</body>
</html>


================================================
FILE: frontends/webcomponent/test_backend.py
================================================
#!/usr/bin/env python3
"""
Comprehensive test backend for vanna-webcomponent validation.

This backend exercises all component types and update patterns to validate
that nothing breaks during webcomponent pruning.

Usage:
    python test_backend.py --mode rapid      # Fast stress test
    python test_backend.py --mode realistic  # Realistic conversation flow
"""

import argparse
import asyncio
import json
import sys
import time
import traceback
import uuid
from datetime import datetime
from typing import AsyncGenerator, Dict, Any, Optional

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import os

# Add vanna to path
sys.path.insert(0, "../vanna/src")

from vanna.core.rich_component import RichComponent, ComponentLifecycle
from vanna.components.rich import (
    RichTextComponent,
    StatusCardComponent,
    ProgressDisplayComponent,
    ProgressBarComponent,
    NotificationComponent,
    StatusIndicatorComponent,
    ButtonComponent,
    ButtonGroupComponent,
    CardComponent,
    TaskListComponent,
    Task,
    BadgeComponent,
    IconTextComponent,
    DataFrameComponent,
    ChartComponent,
    ArtifactComponent,
    LogViewerComponent,
    LogEntry,
    StatusBarUpdateComponent,
    TaskTrackerUpdateComponent,
    ChatInputUpdateComponent,
    TaskOperation,
)
from vanna.servers.base.models import ChatStreamChunk

# Request/Response models
class ChatRequest(BaseModel):
    """Chat request matching vanna API."""
    message: str
    conversation_id: Optional[str] = None
    request_id: Optional[str] = None
    request_context: Dict[str, Any] = {}


class UiComponent(BaseModel):
    """UI component wrapper."""
    rich_component: RichComponent


# Test state
test_state: Dict[str, Any] = {
    "mode": "realistic",
    "component_ids": {},  # Track component IDs for updates
    "action_count": 0,
}


async def yield_chunk(component: RichComponent, conversation_id: str, request_id: str) -> ChatStreamChunk:
    """Convert component to ChatStreamChunk."""
    return ChatStreamChunk(
        rich=component.serialize_for_frontend(),
        simple=None,
        conversation_id=conversation_id,
        request_id=request_id,
        timestamp=time.time(),
    )


async def delay(mode: str, short: float = 0.1, long: float = 0.5):
    """Add delay based on mode."""
    if mode == "realistic":
        await asyncio.sleep(long)
    elif mode == "rapid":
        await asyncio.sleep(short)


async def test_text_component(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test text component with markdown."""
    text_id = str(uuid.uuid4())
    test_state["component_ids"]["text"] = text_id

    # Create with comprehensive markdown
    text = RichTextComponent(
        id=text_id,
        content="""# Test Text Component

This component demonstrates **markdown rendering** with various formatting:

## Formatting Examples
- **Bold text** for emphasis
- *Italic text* for style
- `inline code` for snippets
- ~~Strikethrough~~ for deletions

### Lists
1. First ordered item
2. Second ordered item
3. Third ordered item

### Code Block
```python
def hello():
    return "Markdown works!"
```

> Blockquote to test quote rendering

This validates that markdown is properly parsed and displayed.""",
        markdown=True,
    )
    yield await yield_chunk(text, conversation_id, request_id)
    await delay(mode)

    # Update with simpler markdown
    text_updated = text.update(content="""# Updated Text Component

Text has been **successfully updated** with new markdown content!

- Update operation works ✓
- Markdown still renders ✓""")
    yield await yield_chunk(text_updated, conversation_id, request_id)
    await delay(mode)


async def test_status_card(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test status card with all states."""
    card_id = str(uuid.uuid4())
    test_state["component_ids"]["status_card"] = card_id

    # Create - pending
    status_card = StatusCardComponent(
        id=card_id,
        title="Status Card Test",
        status="pending",
        description="Testing status card component...",
        icon="⏳",
        collapsible=True,
        collapsed=False,
    )
    yield await yield_chunk(status_card, conversation_id, request_id)
    await delay(mode)

    # Update to running
    status_card_running = status_card.set_status("running", "Processing test...")
    yield await yield_chunk(status_card_running, conversation_id, request_id)
    await delay(mode)

    # Update to completed
    status_card_done = status_card.set_status("completed", "Test completed successfully!")
    status_card_done.icon = "✅"
    yield await yield_chunk(status_card_done, conversation_id, request_id)
    await delay(mode)


async def test_progress_display(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test progress display component."""
    progress_id = str(uuid.uuid4())
    test_state["component_ids"]["progress_display"] = progress_id

    # Create at 0%
    progress = ProgressDisplayComponent(
        id=progress_id,
        label="Test Progress",
        value=0.0,
        description="Starting test...",
        status="info",
        animated=True,
    )
    yield await yield_chunk(progress, conversation_id, request_id)
    await delay(mode, 0.05, 0.3)

    # Update to 50%
    progress_half = progress.update_progress(0.5, "Halfway there...")
    yield await yield_chunk(progress_half, conversation_id, request_id)
    await delay(mode, 0.05, 0.3)

    # Update to 100%
    progress_done = progress.update_progress(1.0, "Complete!")
    progress_done.status = "success"
    yield await yield_chunk(progress_done, conversation_id, request_id)
    await delay(mode)


async def test_card_component(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test card component with actions."""
    card_id = str(uuid.uuid4())
    test_state["component_ids"]["card"] = card_id

    # Create card with markdown content and buttons
    card = CardComponent(
        id=card_id,
        title="Test Card with Markdown",
        content="""# Card Content

This card demonstrates **markdown rendering** within cards:

- Interactive action buttons
- Collapsible sections
- Status indicators
- `Formatted text`

Click the buttons below to test interactivity!""",
        icon="🃏",
        status="info",
        markdown=True,
        collapsible=True,
        collapsed=False,
        actions=[
            {"label": "Test Action", "action": "/test-action", "variant": "primary"},
            {"label": "Cancel", "action": "/cancel", "variant": "secondary"},
        ],
    )
    yield await yield_chunk(card, conversation_id, request_id)
    await delay(mode)

    # Update card status and content
    card_updated = card.update(
        status="success",
        content="""# Card Updated Successfully!

The card content has been **updated** with:
- New status (success)
- New markdown content
- Same action buttons

✓ Update operation verified""",
        markdown=True
    )
    yield await yield_chunk(card_updated, conversation_id, request_id)
    await delay(mode)


async def test_task_list(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test task list component."""
    task_list_id = str(uuid.uuid4())
    test_state["component_ids"]["task_list"] = task_list_id

    # Create task list
    tasks = [
        Task(title="Setup development environment", description="Install dependencies and configure tools", status="completed", progress=1.0),
        Task(title="Write test suite", description="Create comprehensive component tests", status="in_progress", progress=0.7),
        Task(title="Run validation", description="Validate all components render correctly", status="pending"),
        Task(title="Prune webcomponent", description="Remove unused code and cruft", status="pending"),
    ]
    task_list = TaskListComponent(
        id=task_list_id,
        title="Webcomponent Validation Workflow",
        tasks=tasks,
        show_progress=True,
        show_timestamps=True,
    )
    yield await yield_chunk(task_list, conversation_id, request_id)
    await delay(mode)

    # Update task statuses
    tasks[1].status = "completed"
    tasks[1].progress = 1.0
    tasks[2].status = "in_progress"
    tasks[2].progress = 0.3
    task_list_updated = TaskListComponent(
        id=task_list_id,
        title="Webcomponent Validation Workflow (Updated)",
        tasks=tasks,
        show_progress=True,
        show_timestamps=True,
    )
    task_list_updated.lifecycle = ComponentLifecycle.UPDATE
    yield await yield_chunk(task_list_updated, conversation_id, request_id)
    await delay(mode)


async def test_progress_bar(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test progress bar component."""
    bar_id = str(uuid.uuid4())
    test_state["component_ids"]["progress_bar"] = bar_id

    # Create
    bar = ProgressBarComponent(
        id=bar_id,
        value=0.3,
        label="Loading",
        status="info",
    )
    yield await yield_chunk(bar, conversation_id, request_id)
    await delay(mode, 0.05, 0.2)

    # Update
    bar_updated = bar.update(value=0.8, status="success")
    yield await yield_chunk(bar_updated, conversation_id, request_id)
    await delay(mode)


async def test_notification(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test notification component."""
    for level in ["info", "success", "warning", "error"]:
        notif = NotificationComponent(
            id=str(uuid.uuid4()),
            message=f"This is a {level} notification",
            level=level,
            title=f"{level.capitalize()} Test",
        )
        yield await yield_chunk(notif, conversation_id, request_id)
        await delay(mode, 0.05, 0.2)


async def test_status_indicator(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test status indicator component."""
    indicator_id = str(uuid.uuid4())
    test_state["component_ids"]["status_indicator"] = indicator_id

    # Create with pulse
    indicator = StatusIndicatorComponent(
        id=indicator_id,
        status="running",
        message="Processing...",
        pulse=True,
    )
    yield await yield_chunk(indicator, conversation_id, request_id)
    await delay(mode)

    # Update to success
    indicator_success = indicator.update(status="success", message="Done!", pulse=False)
    yield await yield_chunk(indicator_success, conversation_id, request_id)
    await delay(mode)


async def test_badge(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test badge component."""
    badge = BadgeComponent(
        id=str(uuid.uuid4()),
        text="Test Badge",
        variant="primary",
    )
    yield await yield_chunk(badge, conversation_id, request_id)
    await delay(mode)


async def test_icon_text(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test icon_text component."""
    icon_text = IconTextComponent(
        id=str(uuid.uuid4()),
        icon="🔧",
        text="Tool Icon Test",
    )
    yield await yield_chunk(icon_text, conversation_id, request_id)
    await delay(mode)


async def test_buttons(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test button and button_group components."""
    # Single button
    button = ButtonComponent(
        label="Single Button",
        action="/button-test",
        variant="primary",
        icon="🔘",
    )
    yield await yield_chunk(button, conversation_id, request_id)
    await delay(mode, 0.05, 0.2)

    # Button group
    button_group = ButtonGroupComponent(
        buttons=[
            {"label": "Option 1", "action": "/option1", "variant": "primary"},
            {"label": "Option 2", "action": "/option2", "variant": "secondary"},
            {"label": "Option 3", "action": "/option3", "variant": "success"},
        ],
        orientation="horizontal",
    )
    yield await yield_chunk(button_group, conversation_id, request_id)
    await delay(mode)


async def test_dataframe(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test dataframe component with sample data."""
    dataframe_id = str(uuid.uuid4())
    test_state["component_ids"]["dataframe"] = dataframe_id

    # Create sample data
    sample_data = [
        {"id": 1, "name": "Alice", "age": 30, "city": "New York", "salary": 75000},
        {"id": 2, "name": "Bob", "age": 25, "city": "San Francisco", "salary": 85000},
        {"id": 3, "name": "Charlie", "age": 35, "city": "Chicago", "salary": 70000},
        {"id": 4, "name": "Diana", "age": 28, "city": "Boston", "salary": 80000},
        {"id": 5, "name": "Eve", "age": 32, "city": "Seattle", "salary": 90000},
    ]

    dataframe = DataFrameComponent.from_records(
        records=sample_data,
        title="📊 Employee Data",
        description="""Sample employee dataset demonstrating **DataFrame** features:

- **Searchable**: Try searching for names or cities
- **Sortable**: Click column headers to sort
- **Exportable**: Export to CSV/Excel
- **Paginated**: Navigate through rows

*5 employees across different cities*""",
        id=dataframe_id,
        searchable=True,
        sortable=True,
        exportable=True,
    )
    yield await yield_chunk(dataframe, conversation_id, request_id)
    await delay(mode)

    # Update with more data
    updated_data = sample_data + [
        {"id": 6, "name": "Frank", "age": 29, "city": "Austin", "salary": 78000},
    ]
    dataframe_updated = DataFrameComponent.from_records(
        records=updated_data,
        title="📊 Employee Data (Updated)",
        description="""Dataset **updated** with new employee!

✓ Added Frank from Austin
✓ Now showing 6 employees
✓ Update operation verified""",
        id=dataframe_id,
    )
    dataframe_updated.lifecycle = ComponentLifecycle.UPDATE
    yield await yield_chunk(dataframe_updated, conversation_id, request_id)
    await delay(mode)


async def test_chart(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test chart component with Plotly data."""
    chart_id = str(uuid.uuid4())
    test_state["component_ids"]["chart"] = chart_id

    # Create a simple bar chart
    chart_data = {
        "data": [
            {
                "x": ["Product A", "Product B", "Product C", "Product D"],
                "y": [20, 35, 30, 25],
                "type": "bar",
                "name": "Sales",
                "marker": {"color": "#667eea"},
            }
        ],
        "layout": {
            "title": "Product Sales",
            "xaxis": {"title": "Products"},
            "yaxis": {"title": "Sales (units)"},
        },
    }

    chart = ChartComponent(
        id=chart_id,
        chart_type="bar",
        data=chart_data,
        title="Sales Chart",
    )
    yield await yield_chunk(chart, conversation_id, request_id)
    await delay(mode)

    # Update to line chart
    line_chart_data = {
        "data": [
            {
                "x": ["Jan", "Feb", "Mar", "Apr", "May"],
                "y": [10, 15, 13, 17, 21],
                "type": "scatter",
                "mode": "lines+markers",
                "name": "Revenue",
                "line": {"color": "#10b981", "width": 3},
            }
        ],
        "layout": {
            "title": "Monthly Revenue Trend",
            "xaxis": {"title": "Month"},
            "yaxis": {"title": "Revenue ($1000s)"},
        },
    }

    chart_updated = ChartComponent(
        id=chart_id,
        chart_type="line",
        data=line_chart_data,
        title="Revenue Chart",
    )
    chart_updated.lifecycle = ComponentLifecycle.UPDATE
    yield await yield_chunk(chart_updated, conversation_id, request_id)
    await delay(mode)

async def test_artifact(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test artifact component with HTML/SVG content."""
    artifact_id = str(uuid.uuid4())
    test_state["component_ids"]["artifact"] = artifact_id

    # Create SVG artifact
    svg_content = '''<svg width="200" height="200" xmlns="http://www.w3.org/2000/svg">
    <circle cx="100" cy="100" r="80" fill="#667eea" opacity="0.8"/>
    <circle cx="100" cy="100" r="60" fill="#764ba2" opacity="0.6"/>
    <circle cx="100" cy="100" r="40" fill="#f093fb" opacity="0.4"/>
    <text x="100" y="105" text-anchor="middle" fill="white" font-size="20" font-weight="bold">
        Test SVG
    </text>
</svg>'''

    artifact = ArtifactComponent(
        id=artifact_id,
        content=svg_content,
        artifact_type="svg",
        title="SVG Circle Visualization",
        description="Concentric circles demonstration",
        fullscreen_capable=True,
    )
    yield await yield_chunk(artifact, conversation_id, request_id)
    await delay(mode)


async def test_log_viewer(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test log viewer component."""
    log_id = str(uuid.uuid4())
    test_state["component_ids"]["log_viewer"] = log_id

    # Create initial log viewer with entries
    log_viewer = LogViewerComponent(
        id=log_id,
        title="System Logs",
        entries=[
            LogEntry(message="System started", level="info"),
            LogEntry(message="Loading configuration...", level="info"),
            LogEntry(message="Configuration loaded successfully", level="info"),
        ],
        searchable=True,
        auto_scroll=True,
    )
    yield await yield_chunk(log_viewer, conversation_id, request_id)
    await delay(mode, 0.05, 0.3)

    # Add warning
    log_viewer = log_viewer.add_entry("Memory usage at 75%", level="warning")
    yield await yield_chunk(log_viewer, conversation_id, request_id)
    await delay(mode, 0.05, 0.3)

    # Add error
    log_viewer = log_viewer.add_entry("Connection timeout", level="error", data={"host": "api.example.com", "port": 443})
    yield await yield_chunk(log_viewer, conversation_id, request_id)
    await delay(mode, 0.05, 0.3)

    # Add success
    log_viewer = log_viewer.add_entry("Reconnected successfully", level="info")
    yield await yield_chunk(log_viewer, conversation_id, request_id)
    await delay(mode)

async def test_ui_state_updates(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Test UI state update components."""
    # Status bar update
    status_bar = StatusBarUpdateComponent(
        message="Running comprehensive component test...",
        status="info",
    )
    yield await yield_chunk(status_bar, conversation_id, request_id)
    await delay(mode, 0.1, 0.3)

    # Task tracker - add tasks to sidebar
    task1 = Task(
        title="Validate Text Components",
        description="Test text, markdown, and formatting",
        status="completed",
        progress=1.0,
    )
    task_tracker_add1 = TaskTrackerUpdateComponent.add_task(task1)
    yield await yield_chunk(task_tracker_add1, conversation_id, request_id)
    await delay(mode, 0.1, 0.3)

    task2 = Task(
        title="Validate Data Components",
        description="Test DataFrame, Chart, Code blocks",
        status="in_progress",
        progress=0.6,
    )
    task_tracker_add2 = TaskTrackerUpdateComponent.add_task(task2)
    yield await yield_chunk(task_tracker_add2, conversation_id, request_id)
    await delay(mode, 0.1, 0.3)

    task3 = Task(
        title="Validate Interactive Components",
        description="Test buttons, actions, and UI state",
        status="pending",
    )
    task_tracker_add3 = TaskTrackerUpdateComponent.add_task(task3)
    yield await yield_chunk(task_tracker_add3, conversation_id, request_id)
    await delay(mode, 0.1, 0.3)

    # Update task 2 to completed
    task_tracker_update = TaskTrackerUpdateComponent(
        operation=TaskOperation.UPDATE_TASK,
        task_id=task2.id,
        status="completed",
        progress=1.0,
    )
    yield await yield_chunk(task_tracker_update, conversation_id, request_id)
    await delay(mode, 0.1, 0.3)

    # Update status bar
    status_bar_complete = StatusBarUpdateComponent(
        message="All components validated successfully!",
        status="success",
    )
    yield await yield_chunk(status_bar_complete, conversation_id, request_id)
    await delay(mode, 0.1, 0.3)

    # Chat input update - change placeholder
    chat_input = ChatInputUpdateComponent(
        placeholder="Type a message to test chat input updates...",
        disabled=False,
    )
    yield await yield_chunk(chat_input, conversation_id, request_id)
    await delay(mode)


async def run_comprehensive_test(conversation_id: str, request_id: str, mode: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Run all component tests."""
    # Introduction
    intro = RichTextComponent(
        content=f"""# 🧪 Comprehensive Component Test

**Mode**: {mode}

## Test Coverage
This test validates **16 component types** supported by the webcomponent:
- ✅ Component creation
- ✅ Incremental updates
- ✅ Markdown rendering
- ✅ Interactive actions
- ✅ Data visualization

### Component Categories
1. **Primitive**: Text, Badge, Icon Text
2. **Feedback**: Status Card, Progress, Notifications, Logs
3. **Data**: Card, Task List, DataFrame, Chart, Code
4. **Specialized**: Artifact (SVG/HTML)
5. **Interactive**: Buttons with actions

Watch the sidebar checklist as components render! ➡️""",
        markdown=True,
    )
    yield await yield_chunk(intro, conversation_id, request_id)
    await delay(mode)

    # Run all tests
    async for chunk in test_text_component(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_status_card(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_progress_display(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_card_component(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_task_list(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_progress_bar(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_notification(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_status_indicator(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_badge(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_icon_text(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_buttons(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_dataframe(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_chart(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_artifact(conversation_id, request_id, mode):
        yield chunk

    async for chunk in test_log_viewer(conversation_id, request_id, mode):
        yield chunk

    # NOTE: Table, Container, and CodeBlock components are defined in vanna Python package
    # but NOT supported by the webcomponent (no renderers). Skipping these tests.
    # These are candidates for removal from the vanna package.

    async for chunk in test_ui_state_updates(conversation_id, request_id, mode):
        yield chunk

    # Completion message
    done = StatusCardComponent(
        title="✅ Test Suite Complete",
        status="completed",
        description=f"""All **16 component types** successfully rendered in **{mode}** mode!

**Validated:**
- Component creation & updates
- Markdown rendering
- Interactive buttons
- Data visualization
- UI state management

Check the sidebar for the complete checklist.""",
        icon="✅",
    )
    yield await yield_chunk(done, conversation_id, request_id)


async def handle_action_message(message: str, conversation_id: str, request_id: str) -> AsyncGenerator[ChatStreamChunk, None]:
    """Handle button action messages."""
    test_state["action_count"] += 1

    response = NotificationComponent(
        message=f"Action received: {message}",
        level="success",
        title=f"Action #{test_state['action_count']}",
    )
    yield await yield_chunk(response, conversation_id, request_id)

    # Also show a card with details
    card = CardComponent(
        title="Action Handler Response",
        content=f"Received action: `{message}`\n\nThis confirms button interactivity is working!",
        icon="🎯",
        status="success",
    )
    yield await yield_chunk(card, conversation_id, request_id)


# FastAPI app
app = FastAPI(title="Vanna Webcomponent Test Backend")

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files (static directory for webcomponent)
static_path = os.path.join(os.path.dirname(__file__), "static")
if os.path.exists(static_path):
    app.mount("/static", StaticFiles(directory=static_path), name="static")


@app.post("/api/vanna/v2/chat_sse")
async def chat_sse(chat_request: ChatRequest) -> StreamingResponse:
    """SSE endpoint for streaming chat."""
    conversation_id = chat_request.conversation_id or str(uuid.uuid4())
    request_id = chat_request.request_id or str(uuid.uuid4())
    message = chat_request.message.strip()

    async def generate() -> AsyncGenerator[str, None]:
        """Generate SSE stream."""
        try:
            # Handle button actions
            if message.startswith("/") and message != "/test":
                async for chunk in handle_action_message(message, conversation_id, request_id):
                    yield f"data: {chunk.model_dump_json()}\n\n"

            # Handle test command or initial message
            elif message == "/test" or "test" in message.lower():
                async for chunk in run_comprehensive_test(conversation_id, request_id, test_state["mode"]):
                    yield f"data: {chunk.model_dump_json()}\n\n"

            # Default response
            else:
                response = RichTextComponent(
                    content=f"You said: {message}\n\nType `/test` to run the comprehensive component test.",
                    markdown=True,
                )
                chunk = await yield_chunk(response, conversation_id, request_id)
                yield f"data: {chunk.model_dump_json()}\n\n"

            yield "data: [DONE]\n\n"

        except Exception as e:
            error_message = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
            print(f"ERROR in chat_sse: {error_message}")  # Log to console
            error_chunk = {
                "type": "error",
                "data": {"message": error_message},
                "conversation_id": conversation_id,
                "request_id": request_id,
            }
            yield f"data: {json.dumps(error_chunk)}\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/health")
async def health():
    """Health check."""
    return {"status": "ok", "mode": test_state["mode"]}


@app.get("/")
async def root():
    """Serve test HTML page."""
    html_path = os.path.join(os.path.dirname(__file__), "test-comprehensive.html")
    if os.path.exists(html_path):
        return FileResponse(html_path)
    return {
        "message": "Vanna Webcomponent Test Backend",
        "mode": test_state["mode"],
        "endpoints": {
            "chat": "POST /api/vanna/v2/chat_sse",
            "health": "GET /health",
        },
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test backend for vanna-webcomponent")
    parser.add_argument(
        "--mode",
        choices=["rapid", "realistic"],
        default="realistic",
        help="Test mode: rapid (fast) or realistic (with delays)",
    )
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=5555, help="Port to bind to")

    args = parser.parse_args()
    test_state["mode"] = args.mode

    print(f"Starting test backend in {args.mode} mode...")
    print(f"Server running at http://{args.host}:{args.port}")
    print("Send message '/test' to run comprehensive component test")

    import uvicorn
    uvicorn.run(app, host=args.host, port=args.port)


================================================
FILE: frontends/webcomponent/tsconfig.json
================================================
{
  "compilerOptions": {
    "target": "ES2020",
    "lib": ["ES2020", "DOM", "DOM.Iterable"],
    "module": "ESNext",
    "skipLibCheck": true,
    "moduleResolution": "bundler",
    "allowImportingTsExtensions": true,
    "resolveJsonModule": true,
    "isolatedModules": true,
    "noEmit": true,
    "strict": true,
    "noUnusedLocals": true,
    "noUnusedParameters": true,
    "noFallthroughCasesInSwitch": true,
    "experimentalDecorators": true,
    "useDefineForClassFields": false
  },
  "include": ["src"]
}

================================================
FILE: frontends/webcomponent/vite.config.ts
================================================
import { defineConfig } from 'vite';

export default defineConfig({
  define: {
    __BUILD_TIME__: JSON.stringify(new Date().toISOString()),
    __BUILD_VERSION__: JSON.stringify(process.env.npm_package_version || '1.0.0'),
  },
  build: {
    outDir: 'dist',
    lib: {
      entry: 'src/index.ts',
      formats: ['es'],
      fileName: () => 'vanna-components.js',
    },
    rollupOptions: {
      // Remove external to bundle lit with the components
      // external: /^lit/,
    },
  },
  preview: {
    port: 9876,
    strictPort: true,
  },
});

================================================
FILE: notebooks/quickstart.ipynb
================================================
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Install the Package\n",
    "Here we're installing it directly from GitHub while it's in development."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install 'vanna[flask,anthropic]'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download a Sample Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import httpx\n",
    "\n",
    "with open(\"Chinook.sqlite\", \"wb\") as f:\n",
    "    with httpx.stream(\"GET\", \"https://vanna.ai/Chinook.sqlite\") as response:\n",
    "        for chunk in response.iter_bytes():\n",
    "            f.write(chunk)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vanna import Agent, AgentConfig\n",
    "from vanna.servers.fastapi import VannaFastAPIServer\n",
    "from vanna.core.registry import ToolRegistry\n",
    "from vanna.core.user import UserResolver, User, RequestContext\n",
    "from vanna.integrations.anthropic import AnthropicLlmService\n",
    "from vanna.tools import RunSqlTool, VisualizeDataTool\n",
    "from vanna.integrations.sqlite import SqliteRunner\n",
    "from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool\n",
    "from vanna.integrations.local.agent_memory import DemoAgentMemory\n",
    "from vanna.capabilities.sql_runner import RunSqlToolArgs\n",
    "from vanna.tools.visualize_data import VisualizeDataArgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define your User Authentication\n",
    "Here we're going to say that if you're logged in as `admin@example.com` then you're in the `admin` group, otherwise you're in the `user` group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleUserResolver(UserResolver):\n",
    "    async def resolve_user(self, request_context: RequestContext) -> User:\n",
    "        # In production, validate cookies/JWTs here\n",
    "        user_email = request_context.get_cookie('vanna_email')\n",
    "        if not user_email:\n",
    "            raise ValueError(\"Missing 'vanna_email' cookie for user identification\")\n",
    "        \n",
    "        print(f\"Resolving user for email: {user_email}\")\n",
    "\n",
    "        if user_email == \"admin@example.com\":\n",
    "            return User(id=\"admin1\", email=user_email, group_memberships=['admin'])\n",
    "        \n",
    "        return User(id=\"user1\", email=user_email, group_memberships=['user'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the Tools and Access Control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = ToolRegistry()\n",
    "tools.register_local_tool(RunSqlTool(sql_runner=SqliteRunner(database_path=\"./Chinook.sqlite\")), access_groups=['admin', 'user'])\n",
    "tools.register_local_tool(VisualizeDataTool(), access_groups=['admin', 'user'])\n",
    "agent_memory = DemoAgentMemory(max_items=1000)\n",
    "tools.register_local_tool(SaveQuestionToolArgsTool(), access_groups=['admin'])\n",
    "tools.register_local_tool(SearchSavedCorrectToolUsesTool(), access_groups=['admin', 'user'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up LLM\n",
    "llm = AnthropicLlmService(model=\"claude-sonnet-4-5\", api_key=\"sk-ant-...\")\n",
    "\n",
    "# Create agent with your options\n",
    "agent = Agent(\n",
    "    llm_service=llm,\n",
    "    tool_registry=tools,\n",
    "    user_resolver=SimpleUserResolver(),\n",
    "    config=AgentConfig(),\n",
    "    agent_memory=agent_memory\n",
    ")\n",
    "\n",
    "# 4. Create and run server\n",
    "server = VannaFastAPIServer(agent)\n",
    "server.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}


================================================
FILE: papers/ai-sql-accuracy-2023-08-17.md
================================================
# AI SQL Accuracy: Testing different LLMs + context strategies to maximize SQL generation accuracy
_2023-08-17_

## TLDR

The promise of having an autonomous AI agent that can answer business users’ plain English questions is an attractive but thus far elusive proposition. Many have tried, with limited success, to get ChatGPT to write. The failure is primarily due of a lack of the LLM's knowledge of the particular dataset it’s being asked to query.

In this paper, **we show that context is everything, and with the right context, we can get from ~3% accuracy to ~80% accuracy**. We go through three different context strategies, and showcase one that is the clear winner - where we combine schema definitions, documentation, and prior SQL queries with a relevance search.

We also compare a few different LLMs - including Google Bison, GPT 3.5, GPT 4, and a brief attempt with Llama 2. While **GPT 4 takes the crown of the best overall LLM for generating SQL**, Google’s Bison is roughly equivalent when enough context is provided.

Finally, we show how you can use the methods demonstrated here to generate SQL for your database.

Here's a summary of our key findings -

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/summary.png)

## Table of Contents
* [Why use AI to generate SQL?](#why-use-ai-to-generate-sql)
* [Setting up architecture of the test](#setting-up-architecture-of-the-test)
* [Setting up the test levers](#setting-up-the-test-levers)
    * [Choosing a dataset](#choosing-a-dataset)
    * [Choosing the questions](#choosing-the-questions)
    * [Choosing the prompt](#choosing-the-prompt)
    * [Choosing the LLMs (Foundational models)](#choosing-the-llms-foundational-models)
    * [Choosing the context](#choosing-the-context)
* [Using ChatGPT to generate SQL](#using-chatgpt-to-generate-sql)
* [Using schema only](#using-schema-only)
* [Using SQL examples](#using-sql-examples)
* [Using contextually relevant examples](#using-contextually-relevant-examples)
* [Analyzing the results](#analyzing-the-results)
* [Next steps to getting accuracy even higher](#next-steps-to-getting-accuracy-even-higher)
* [Use AI to write SQL for your dataset](#use-ai-to-write-sql-for-your-dataset)



## Why use AI to generate SQL?

Many organizations have now adopted some sort of data warehouse or data lake - a repository of a lot of the organization’s critical data that is queryable for analytical purposes. This ocean of data is brimming with potential insights, but only a small fraction of people in an enterprise have the two skills required to harness the data —

1. A solid comprehension of **advanced SQL**, and
2. A comprehensive knowledge of the **organization’s unique data structure & schema**

The number of people with both of the above is not only vanishingly small, but likely not the same people that have the majority of the questions. 

**So what actually happens inside organizations?** Business users, like product managers, sales managers, and executives, have data questions that will inform business decisions and strategy. They’ll first check dashboards, but most questions are ad hoc and specific, and the answers aren’t available, so they’ll ask a data analyst or engineer - whomever possesses the combination of skills above. These people are busy, and take a while to get to the request, and as soon as they get an answer, the business user has follow up questions. 

**This process is painful** for both the business user (long lead times to get answers) and the analyst (distracts from their main projects), and leads to many potential insights being lost.

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/question-flow.png)

**Enter generative AI!** LLMs potentially give the opportunity to business users to query the database in plain English (with the LLMs doing the SQL translation), and we have heard from dozens of companies that this would be a game changer for their data teams and even their businesses.

**The key challenge is generating accurate SQL for complex and messy databases**. Plenty of people we’ve spoken with have tried to use ChatGPT to write SQL with limited success and a lot of pain. Many have given up and reverted back to the old fashioned way of manually writing SQL. At best, ChatGPT is a sometimes useful co-pilot for analysts to get syntax right.

**But there’s hope!** We’ve spent the last few months immersed in this problem, trying various models, techniques and approaches to improve the accuracy of SQL generated by LLMs. In this paper, we show the performance of various LLMs and how the strategy of providing contextually relevant correct SQL to the LLM can allow the LLM to **achieve extremely high accuracy**.


## Setting up architecture of the test

First, we needed to define the architecture of the test. A rough outline is below, in a five step process, with _pseudo code_ below - 

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/test-architecture.png)

1. **Question** - We start with the business question.
```python
   question = "how many clients are there in germany"
```
2. **Prompt** - We create the prompt to send to the LLM.
```python
   prompt = f"""
   Write a SQL statement for the following question:
   {question}
   """
```
3. **Generate SQL** - Using an API, we’ll send the prompt to the LLM and get back generated SQL.
```python
   sql = llm.api(api_key=api_key, prompt=prompt, parameters=parameters)
```
4. **Run SQL** - We'll run the SQL against the database.
```python
    df = db.conn.execute(sql)
```
5. **Validate results** - Finally, we’ll validate that the results are in line with what we expect.
There are some shades of grey when it comes to the results so we did a manual evaluation of the results. You can see those results [here](https://github.com/vanna-ai/research/blob/main/data/sec_evaluation_data_tagged.csv)

## Setting up the test levers

Now that we have our experiment set up, we’ll need to figure out what levers would impact accuracy, and what our test set would be. We tried two levers (the LLMs and the training data used), and we ran on 20 questions that made up our test set. So we ran a total of 3 LLMs x 3 context strategies x 20 questions = 180 individual trials in this experiment.

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/test-levers.png)


### Choosing a dataset

First, we need to **choose an appropriate dataset** to try. We had a few guiding principles - 

1. **Representative**. Datasets in enterprises are often complex and this complexity isn’t captured in many demo / sample datasets. We want to use a complicated database that has real-word use cases that contains real-world data. 
2. **Accessible**. We also wanted that dataset to be publicly available. 
3. **Understandable**. The dataset should be somewhat understandable to a wide audience - anything too niche or technical would be difficult to decipher.
4. **Maintained**. We’d prefer a dataset that’s maintained and updated properly, in reflection of a real database.

A dataset that we found that met the criteria above was the Cybersyn SEC filings dataset, which is available for free on the Snowflake marketplace: 

https://docs.cybersyn.com/our-data-products/economic-and-financial/sec-filings


### Choosing the questions

Next, we need to **choose the questions**. Here are some sample questions (see them all in this [file](https://github.com/vanna-ai/research/blob/main/data/questions_sec.csv)) - 

1. How many companies are there in the dataset?
2. What annual measures are available from the 'ALPHABET INC.' Income Statement?
3. What are the quarterly 'Automotive sales' and 'Automotive leasing' for Tesla?
4. How many Chipotle restaurants are there currently?

Now that we have the dataset + questions, we’ll need to come up with the levers. 

### Choosing the prompt

For the **prompt**, for this run, we are going to hold the prompt constant, though we’ll do a follow up which varies the prompt.

### Choosing the LLMs (Foundational models)

For the **LLMs** to test, we’ll try the following - 

1. [**Bison (Google)**](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models) - Bison is the version of [PaLM 2](https://blog.google/technology/ai/google-palm-2-ai-large-language-model/) that’s available via GCP APIs.
2. [**GPT 3.5 Turbo (OpenAI)**](https://platform.openai.com/docs/models/gpt-3-5) - GPT 3.5 until recently was the flagship OpenAI model despite 4 being available because of latency and cost benefits, and not a huge accuracy difference (well - we’ll put that to the test) especially for basic tasks.
3. [**GPT 4 (OpenAI)**](https://platform.openai.com/docs/models/gpt-4) - The more advanced but less performant OpenAI model. GPT 4 is multi-modal, though we won’t be using that functionality.
4. [**Llama 2 (Meta)**](https://ai.meta.com/llama/) - We really wanted to include an open source model - and the leading one is Meta’s Llama 2. But our setup through [Replicate](https://replicate.com/replicate/llama-2-70b-chat) quickly conked out, and we weren’t able to fix it in time for publishing this. In our early runs, when it did work, we found performance to be mediocre at best.


### Choosing the context

Finally, we’ll have three types of **context**. Context refers to what we send to the LLM that helps give the LLM context on our specific dataset.

1. **Schema only**. We put the schema (using DDL) in the context window.
2. **Static examples**. We put static example SQL queries in the context windows.
3. **Contextually relevant examples**. Finally, we put the most relevant context (SQL / DDL / documentation) into the context window, finding it via a vector search based on embeddings.


## Using ChatGPT to generate SQL

This section is a bit unfair to ChatGPT, but it’s a helpful place to get started. Let’s just take a question and ask ChatGPT to generate the SQL - 

**Prompt**

```
Write Snowflake SQL to answer the following question -

What are the quarterly revenues by business segment for Alphabet?

Return SQL only
```

**Response**

Certainly, here's the SQL query:

```sql
SELECT
  EXTRACT(QUARTER FROM date_column) AS quarter,
  business_segment,
  SUM(revenue) AS quarterly_revenue
FROM
  revenue_table
WHERE
  company_name = 'Alphabet'
GROUP BY
  quarter,
  business_segment
ORDER BY
  quarter,
  business_segment;

-- Replace date_column, business_segment, revenue, revenue_table, and adjust any other identifiers to match your specific database schema.
```

Inevitably, when we try to run this, we hit an error, because it made up the table name since we didn't provide it - 

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/sql-error.png)

Of course, we are being unfair to the LLMs - as magical as they are, they cannot (unfortunately? luckily?) possibly know what’s in our database - yet. So let’s hop into the tests where we give more context.


## Using schema only

First, we take the schema of the dataset and put it into the context window. This is usually what we've seen people do with ChatGPT or in tutorials.

An example prompt may look like this (in reality we used the information schema because of how Snowflake shares work but this shows the principle) - 

```
The user provides a question and you provide SQL. You will only respond with SQL code and not with any explanations.

Respond with only SQL code. Do not answer with any explanations -- just the code.

You may use the following DDL statements as a reference for what tables might be available.

CREATE TABLE Table1...

CREATE TABLE Table2...

CREATE TABLE Table3...
```

The results were, in a word, terrible. Of the 60 attempts (20 questions x 3 models), only two questions were answered correctly (both by GPT 4), **for an abysmal accuracy rate of 3%**. Here are the two questions that GPT 4 managed to get right - 

1. What are the top 10 measure descriptions by frequency?
2. What are the distinct statements in the report attributes?

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/accuracy-using-schema-only.png)

It’s evident that by just using the schema, we don’t get close to meeting the bar of a helpful AI SQL agent, though it may be somewhat useful in being an analyst copilot.


## Using SQL examples

If we put ourselves in the shoes of a human who’s exposed to this dataset for the first time, in addition to the table definitions, they’d first look at the example queries to see _how_ to query the database correctly.

These queries can give additional context not available in the schema - for example, which columns to use, how tables join together, and other intricacies of querying that particular dataset.

Cybersyn, as with other data providers on the Snowflake marketplace, provides a few (in this case 3) example queries in their documentation. Let’s include these in the context window.

By providing just those 3 example queries, we see substantial improvements to the correctness of the SQL generated. However, this accuracy greatly varies by the underlying LLM. It seems that GPT-4 is the most able to generalize the example queries in a way that generates the most accurate SQL.

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/accuracy-using-static-examples.png)

## Using contextually relevant examples

Enterprise data warehouses often contain 100s (or even 1000s) of tables, and an order of magnitude more queries that cover all the use cases within their organizations. Given the limited size of the context windows of modern LLMs, we can’t just shove all the prior queries and schema definitions into the prompt.

Our final approach to context is a more sophisticated ML approach - load embeddings of prior queries and the table schemas into a vector database, and only choose the most relevant queries / tables to the question asked. Here's a diagram of what we are doing - note the contextual relevance search in the green box -

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/using-contextually-relevant-examples.png)

By surfacing the most relevant examples of those SQL queries to the LLM, we can drastically improve performance of even the less capable LLMs. Here, we give the LLM the 10 most relevant SQL query examples for the question (from a list of 30 examples stored), and accuracy rates skyrocket.

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/accuracy-using-contextual-examples.png)

We can improve performance even more by maintaining a history of SQL statements that were executable and correctly answer actual questions that users have had.


## Analyzing the results

It’s clear that the biggest difference is not in the type of LLM, but rather in the strategy employed to give the appropriate context to the LLM (eg the “training data” used).

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/summary-table.png)

When looking at SQL accuracy by context strategy, it’s clear that this is what makes the difference. We go from ~3% accurate using just the schema, to ~80% accurate when intelligently using contextual examples.

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/summary.png)

There are still interesting trends with the LLMs themselves. While Bison starts out at the bottom of the heap in both the Schema and Static context strategies, it rockets to the top with a full Contextual strategy. Averaged across the three strategies, **GPT 4 takes the crown as the best LLM for SQL generation**.

![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/accuracy-by-llm.png)

## Next steps to getting accuracy even higher

We'll soon do a follow up on this analysis to get even deeper into accurate SQL generation. Some next steps are -

1. **Use other datasets**: We'd love to try this on other, real world, enterprise datasets. What happens when you get to 100 tables? 1000 tables?
2. **Add more training data**: While 30 queries is great, what happens when you 10x, 100x that number?
3. **Try more databases**: This test was run on a Snowflake database, but we've also gotten this working on BigQuery, Postgres, Redshift, and SQL Server.
4. **Experiment with more foundational models:** We are close to being able to use Llama 2, and we'd love to try other LLMs.

We have some anecdotal evidence for the above but we'll be expanding and refining our tests to include more of these items.

## Use AI to write SQL for your dataset

While the SEC data is a good start, you must be wondering whether this could be relevant for your data and your organization. We’re building a [Python package](https://vanna.ai) that can generate SQL for your database as well as additional functionality like being able to generate Plotly code for the charts, follow-up questions, and various other functions.

Here's an overview of how it works
```python
import vanna as vn
```

1. **Train Using Schema**

```python
vn.train(ddl="CREATE TABLE ...")
```

2. **Train Using Documentation**

```python
vn.train(documentation="...")
```

3. **Train Using SQL Examples**

```python
vn.train(sql="SELECT ...")
```

4. **Generating SQL**

The easiest ways to use Vanna out of the box are `vn.ask(question="What are the ...")` which will return the SQL, table, and chart as you can see in this [example notebook](https://vanna.ai/docs/getting-started.html). `vn.ask` is a wrapper around `vn.generate_sql`, `vn.run_sql`, `vn.generate_plotly_code`, `vn.get_plotly_figure`, and `vn.generate_followup_questions`. This will use optimized context to generate SQL for your question where Vanna will call the LLM for you.

Alternately, you can use `vn.get_related_training_data(question="What are the ...")` as shown in this [notebook](https://github.com/vanna-ai/research/blob/main/notebooks/test-cybersyn-sec.ipynb) which will retrieve the most relevant context that you can use to construct your own prompt to send to any LLM.

This [notebook](https://github.com/vanna-ai/research/blob/main/notebooks/train-cybersyn-sec-3.ipynb) shows an example of how the "Static" context strategy was used to train Vanna on the Cybersyn SEC dataset.

## A note on nomenclature
* **Foundational Model**: This is the underlying LLM
* **Context Model (aka Vanna Model)**: This is a layer that sits on top of the LLM and provides context to the LLM
* **Training**: Generally when we refer to "training" we're talking about training the context model.

## Contact Us
Ping us on [Slack](https://join.slack.com/t/vanna-ai/shared_invite/zt-1unu0ipog-iE33QCoimQiBDxf2o7h97w), [Discord](https://discord.com/invite/qUZYKHremx), or [set up a 1:1 call](https://calendly.com/d/y7j-yqq-yz4/meet-with-both-vanna-co-founders) if you have any issues.


================================================
FILE: pyproject.toml
================================================
[build-system]
requires = ["flit_core >=3.2,<4"]
build-backend = "flit_core.buildapi"

[project]
name = "vanna"
version = "2.0.2"
authors = [
  { name="Zain Hoda", email="zain@vanna.ai" },
]

description = "Generate SQL queries from natural language"
readme = "README.md"
requires-python = ">=3.9"
classifiers = [
    "Programming Language :: Python :: 3",
    "License :: OSI Approved :: MIT License",
    "Operating System :: OS Independent",
]
dependencies = [
    "pydantic>=2.0.0",
    "click>=8.0.0",
    "pandas",
    "httpx>=0.28.0",
    "PyYAML",
    "plotly",
    "tabulate",
    "sqlparse",
    "sqlalchemy",
    "requests",
]

[project.scripts]
vanna = "vanna.servers.cli.server_runner:main"

[project.urls]
"Homepage" = "https://github.com/vanna-ai/vanna"
"Bug Tracker" = "https://github.com/vanna-ai/vanna/issues"

[project.optional-dependencies]
flask = ["flask>=2.0.0", "flask-cors>=4.0.0"]
fastapi = ["fastapi>=0.68.0", "uvicorn>=0.15.0"]
servers = ["vanna[flask,fastapi]"]

postgres = ["psycopg2-binary", "db-dtypes"]
mysql = ["PyMySQL"]
clickhouse = ["clickhouse_connect"]
bigquery = ["google-cloud-bigquery"]
snowflake = ["snowflake-connector-python"]
duckdb = ["duckdb"]
google = ["google-generativeai", "google-cloud-aiplatform"]
all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb>=1.1.0", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client>=1.0.0", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "boto3", "botocore", "langchain_core", "langchain_postgres", "langchain-community", "langchain-huggingface", "xinference-client"]
test = ["pytest>=7.0.0", "pytest-asyncio>=0.21.0", "pytest-mock>=3.10.0", "pytest-cov>=4.0.0", "tox>=4.0.0"]
dev = ["pytest>=7.0.0", "pytest-asyncio>=0.21.0", "pytest-mock>=3.10.0", "pytest-cov>=4.0.0", "tox>=4.0.0", "mypy", "ruff", "pandas-stubs", "plotly-stubs", "types-PyYAML", "types-requests", "types-tabulate"]
chromadb = ["chromadb>=1.1.0"]
openai = ["openai"]
azureopenai = ["openai", "azure-identity"]
qianfan = ["qianfan"]
mistralai = ["mistralai>=1.0.0"]
anthropic = ["anthropic"]
gemini = ["google-genai"]
marqo = ["marqo"]
zhipuai = ["zhipuai"]
ollama = ["ollama", "httpx"]
qdrant = ["qdrant-client>=1.0.0", "fastembed"]
vllm = ["vllm"]
pinecone = ["pinecone", "fastembed"]
opensearch = ["opensearch-py", "opensearch-dsl", "langchain-community", "langchain-huggingface"]
hf = ["transformers"]
milvus = ["pymilvus[model]"]
bedrock = ["boto3", "botocore"]
weaviate = ["weaviate-client"]
azuresearch = ["azure-search-documents", "azure-identity", "azure-common", "fastembed"]
pgvector = ["langchain-postgres>=0.0.12"]
faiss-cpu = ["faiss-cpu"]
faiss-gpu = ["faiss-gpu"]
xinference-client = ["xinference-client"]
oracle = ["oracledb", "chromadb<1.0.0"]
hive = ["pyhive", "thrift"]
presto = ["pyhive", "thrift"]
mssql = ["pyodbc"]

[tool.flit.module]
name = "vanna"
path = "src/vanna"

[tool.flit.sdist]
exclude = [
    "frontends/",
    "tests/",
    "notebooks/",
    ".github/",
    "tox.ini",
]

[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
markers = [
    "integration: marks tests as integration tests (deselect with '-m \"not integration\"')",
    "anthropic: marks tests requiring Anthropic API key",
    "openai: marks tests requiring OpenAI API key",
    "azureopenai: marks tests requiring Azure OpenAI API key",
    "gemini: marks tests requiring Gemini API key",
    "ollama: marks tests requiring local Ollama instance",
    "legacy: marks tests for legacy adapter",
    "slow: marks tests as slow running",
    "postgres: marks tests requiring PostgreSQL",
    "mysql: marks tests requiring MySQL",
]
filterwarnings = [
    "ignore::DeprecationWarning",
]

[tool.ruff]
# Set the target Python version
target-version = "py311"

# Set line length to 88 (Black's default)
line-length = 88

# Enable auto-fixing
fix = false

# Exclude common directories
exclude = [
    ".git",
    ".tox",
    ".venv",
    "venv",
    "__pycache__",
    "build",
    "dist",
    "*.egg-info",
]

[tool.ruff.lint]
# Enable specific rule categories
select = [
    "E",   # pycodestyle errors
    "W",   # pycodestyle warnings
    "F",   # pyflakes
    # "I",   # isort (disabled - use `ruff check --fix` to auto-fix import sorting)
    "N",   # pep8-naming
    "B",   # flake8-bugbear
    "C4",  # flake8-comprehensions
    "SIM", # flake8-simplify
]

# Ignore specific rules
ignore = [
    # Formatting/style (handled by formatter or not critical)
    "E501",  # line too long (handled by formatter)
    "E402",  # module level import not at top of file
    "E731",  # lambda assignment
    "E741",  # ambiguous variable name
    "W291",  # trailing whitespace
    "W293",  # blank line with whitespace

    # Naming conventions (legacy compatibility)
    "N801",  # invalid class name
    "N802",  # function name should be lowercase
    "N803",  # argument name should be lowercase
    "N805",  # invalid first argument name for method
    "N806",  # variable in function should be lowercase
    "N818",  # error suffix on exception name
    "N999",  # invalid module name

    # Unused/redefined (often intentional)
    "F401",  # imported but unused
    "F541",  # f-string missing placeholders
    "F811",  # redefinition of unused name
    "F841",  # unused variable

    # Bugbear rules (opinionated or intentional)
    "B006",  # mutable argument default (sometimes needed)
    "B007",  # unused loop control variable
    "B008",  # do not perform function calls in argument defaults
    "B024",  # abstract base class without abstract method
    "B027",  # empty method without abstract decorator
    "B904",  # raise without from inside except (intentional in legacy code)
    "B905",  # zip without explicit strict

    # Comprehension/collection style
    "C408",  # unnecessary collection call
    "C416",  # unnecessary comprehension

    # Simplification suggestions (all SIM rules - opinionated style)
    "SIM102", # collapsible if
    "SIM103", # needless bool
    "SIM105", # suppressible exception
    "SIM108", # if-else block instead of if-exp
    "SIM110", # reimplemented builtin
    "SIM114", # if with same arms
    "SIM117", # multiple with statements
    "SIM118", # in dict keys
    "SIM401", # if-else block instead of dict get
    "SIM910", # dict get with none default
]

# Allow fix for all enabled rules (when `--fix` is provided)
fixable = ["ALL"]
unfixable = []

# Allow unused variables when underscore-prefixed
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"

[tool.ruff.format]
# Use double quotes for strings
quote-style = "double"

# Indent with spaces
indent-style = "space"

# Respect magic trailing commas
skip-magic-trailing-comma = false

# Automatically detect line endings
line-ending = "auto"


================================================
FILE: setup.cfg
================================================
[flake8]
ignore = BLK100,W503,E203,E722,F821,F841
max-line-length = 100
exclude = .tox,.git,docs,venv,jupyter_notebook_config.py,jupyter_lab_config.py,assets.py

[tool:brunette]
verbose = true
single-quotes = false
target-version = py39
exclude = .tox,.git,docs,venv,assets.py


================================================
FILE: src/evals/benchmarks/llm_comparison.py
================================================
"""
LLM Comparison Benchmark

This script compares different LLMs on SQL generation tasks.
Run from repository root:
    PYTHONPATH=. python evals/benchmarks/llm_comparison.py
"""

import asyncio
import os
from pathlib import Path

from vanna import Agent
from vanna.core.evaluation import (
    EvaluationRunner,
    EvaluationDataset,
    AgentVariant,
    TrajectoryEvaluator,
    OutputEvaluator,
    EfficiencyEvaluator,
)
from vanna.integrations.anthropic import AnthropicLlmService
from vanna.integrations.local import MemoryConversationStore
from vanna.core.registry import ToolRegistry


def get_sql_tools() -> ToolRegistry:
    """Get SQL-related tools for testing.

    In a real scenario, this would return actual SQL tools.
    For this benchmark, we'll use a placeholder.
    """
    # TODO: Add actual SQL tools
    return ToolRegistry()


async def compare_llms():
    """Compare different LLMs on SQL generation tasks."""

    print("=" * 80)
    print("LLM COMPARISON BENCHMARK - SQL Generation")
    print("=" * 80)
    print()

    # Load test dataset
    dataset_path = (
        Path(__file__).parent.parent / "datasets" / "sql_generation" / "basic.yaml"
    )
    print(f"Loading dataset from: {dataset_path}")
    dataset = EvaluationDataset.from_yaml(str(dataset_path))
    print(f"Loaded dataset: {dataset.name}")
    print(f"Test cases: {len(dataset.test_cases)}")
    print()

    # Get API keys
    anthropic_key = os.getenv("ANTHROPIC_API_KEY")
    if not anthropic_key:
        print("⚠️  ANTHROPIC_API_KEY not set. Using placeholder.")
        anthropic_key = "test-key"

    # Create agent variants
    print("Creating agent variants...")

    tool_registry = get_sql_tools()

    variants = [
        AgentVariant(
            name="claude-sonnet-4",
            agent=Agent(
                llm_service=AnthropicLlmService(
                    api_key=anthropic_key, model="claude-sonnet-4-20250514"
                ),
                tool_registry=tool_registry,
                conversation_store=MemoryConversationStore(),
            ),
            metadata={
                "provider": "anthropic",
                "model": "claude-sonnet-4-20250514",
                "version": "2025-05-14",
            },
        ),
        AgentVariant(
            name="claude-opus-4",
            agent=Agent(
                llm_service=AnthropicLlmService(
                    api_key=anthropic_key, model="claude-opus-4-20250514"
                ),
                tool_registry=tool_registry,
                conversation_store=MemoryConversationStore(),
            ),
            metadata={
                "provider": "anthropic",
                "model": "claude-opus-4-20250514",
                "version": "2025-05-14",
            },
        ),
    ]

    print(f"Created {len(variants)} variants:")
    for v in variants:
        print(f"  - {v.name}")
    print()

    # Create evaluators
    evaluators = [
        TrajectoryEvaluator(),
        OutputEvaluator(),
        EfficiencyEvaluator(
            max_execution_time_ms=10000,
            max_tokens=5000,
        ),
    ]

    print(f"Using {len(evaluators)} evaluators:")
    for e in evaluators:
        print(f"  - {e.name}")
    print()

    # Create runner with high concurrency for I/O bound tasks
    runner = EvaluationRunner(
        evaluators=evaluators,
        max_concurrency=20,  # Run 20 test cases concurrently
    )

    # Run comparison
    print("Running comparison (all variants in parallel)...")
    print(
        f"Total executions: {len(variants)} variants × {len(dataset.test_cases)} test cases = {len(variants) * len(dataset.test_cases)}"
    )
    print()

    comparison = await runner.compare_agents(variants, dataset.test_cases)

    # Print results
    print()
    comparison.print_summary()

    # Show winner
    print(f"🏆 Best by score: {comparison.get_best_variant('score')}")
    print(f"⚡ Best by speed: {comparison.get_best_variant('speed')}")
    print(f"✅ Best by pass rate: {comparison.get_best_variant('pass_rate')}")
    print()

    # Save reports
    output_dir = Path(__file__).parent.parent / "results"
    output_dir.mkdir(exist_ok=True)

    html_path = output_dir / "llm_comparison.html"
    csv_path = output_dir / "llm_comparison.csv"

    comparison.save_html(str(html_path))
    comparison.save_csv(str(csv_path))

    print(f"📊 Reports saved:")
    print(f"  - HTML: {html_path}")
    print(f"  - CSV: {csv_path}")


async def main():
    """Run the LLM comparison benchmark."""
    try:
        await compare_llms()
    except Exception as e:
        print(f"❌ Error running benchmark: {e}")
        import traceback

        traceback.print_stack()
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/evals/datasets/sql_generation/basic.yaml
================================================
dataset:
  name: "SQL Generation - Basic"
  description: "Basic SQL generation tasks for evaluating agent SQL capabilities"

  test_cases:
    - id: "sql_001"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "Show me total sales by region"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["SELECT", "SUM", "GROUP BY", "region"]
        max_execution_time_ms: 5000
      metadata:
        category: "aggregation"
        difficulty: "easy"

    - id: "sql_002"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "What were our top 5 customers by revenue last month?"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["SELECT", "TOP", "ORDER BY", "DESC"]
        max_execution_time_ms: 5000
      metadata:
        category: "ranking"
        difficulty: "medium"

    - id: "sql_003"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "Calculate the average order value for each product category"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["AVG", "GROUP BY", "category"]
        max_execution_time_ms: 5000
      metadata:
        category: "aggregation"
        difficulty: "easy"

    - id: "sql_004"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "Show me the trend of monthly sales over the past year"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query", "visualize_data"]
        final_answer_contains: ["SELECT", "GROUP BY", "month"]
        max_execution_time_ms: 7000
      metadata:
        category: "time_series"
        difficulty: "medium"

    - id: "sql_005"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "Find customers who haven't made a purchase in the last 90 days"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["SELECT", "WHERE", "NOT IN", "90"]
        final_answer_not_contains: ["DROP", "DELETE", "UPDATE"]
        max_execution_time_ms: 5000
      metadata:
        category: "filtering"
        difficulty: "medium"

    - id: "sql_006"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "Compare this quarter's revenue to the same quarter last year"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["SELECT", "quarter", "year"]
        max_execution_time_ms: 6000
      metadata:
        category: "comparison"
        difficulty: "hard"

    - id: "sql_007"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "List all products that are currently out of stock"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["SELECT", "WHERE", "stock", "= 0"]
        final_answer_not_contains: ["DROP", "DELETE"]
        max_execution_time_ms: 4000
      metadata:
        category: "filtering"
        difficulty: "easy"

    - id: "sql_008"
      user_id: "eval_user"
      username: "evaluator"
      email: "eval@example.com"
      user_groups: ["user", "analyst"]
      message: "Calculate the customer lifetime value for each customer segment"
      expected_outcome:
        tools_called: ["generate_sql", "execute_query"]
        final_answer_contains: ["SELECT", "SUM", "GROUP BY", "segment"]
        max_execution_time_ms: 6000
      metadata:
        category: "aggregation"
        difficulty: "hard"


================================================
FILE: src/vanna/__init__.py
================================================
"""
Vanna Agents - A modular framework for building LLM agents.

This package provides a flexible framework for creating conversational AI agents
with tool execution, conversation management, and user scoping.
"""

# Version information
__version__ = "0.1.0"

# Import core framework components
from .core import (
    # Interfaces
    Agent,
    ConversationStore,
    LlmService,
    SystemPromptBuilder,
    Tool,
    UserService,
    T,
    # Models
    Conversation,
    LlmMessage,
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
    Message,
    ToolCall,
    ToolContext,
    ToolResult,
    ToolSchema,
    User,
    # UI Components
    UiComponent,
    SimpleComponent,
    SimpleComponentType,
    SimpleTextComponent,
    SimpleImageComponent,
    SimpleLinkComponent,
    # Rich Components
    ArtifactComponent,
    BadgeComponent,
    CardComponent,
    DataFrameComponent,
    IconTextComponent,
    LogViewerComponent,
    NotificationComponent,
    ProgressBarComponent,
    ProgressDisplayComponent,
    RichTextComponent,
    StatusCardComponent,
    TaskListComponent,
    # Core implementations
    Agent,
    AgentConfig,
    DefaultSystemPromptBuilder,
    DefaultWorkflowHandler,
    ToolRegistry,
    # Evaluation
    Evaluator,
    TestCase,
    ExpectedOutcome,
    AgentResult,
    EvaluationResult,
    TestCaseResult,
    AgentVariant,
    EvaluationRunner,
    TrajectoryEvaluator,
    OutputEvaluator,
    LLMAsJudgeEvaluator,
    EfficiencyEvaluator,
    EvaluationReport,
    ComparisonReport,
    EvaluationDataset,
    # Exceptions
    AgentError,
    ConversationNotFoundError,
    LlmServiceError,
    PermissionError,
    ToolExecutionError,
    ToolNotFoundError,
    ValidationError,
)

# Import basic implementations
from .integrations import MemoryConversationStore, MockLlmService

# Main exports
__all__ = [
    # Version
    "__version__",
    # Core interfaces
    "Agent",
    "Tool",
    "LlmService",
    "ConversationStore",
    "UserService",
    "SystemPromptBuilder",
    "T",
    # Models
    "User",
    "Message",
    "Conversation",
    "ToolCall",
    "ToolResult",
    "ToolContext",
    "ToolSchema",
    "LlmMessage",
    "LlmRequest",
    "LlmResponse",
    "LlmStreamChunk",
    # UI Components
    "UiComponent",
    "SimpleComponent",
    "SimpleComponentType",
    "SimpleTextComponent",
    "SimpleImageComponent",
    "SimpleLinkComponent",
    # Rich Components
    "ArtifactComponent",
    "BadgeComponent",
    "CardComponent",
    "DataFrameComponent",
    "IconTextComponent",
    "LogViewerComponent",
    "NotificationComponent",
    "ProgressBarComponent",
    "ProgressDisplayComponent",
    "RichTextComponent",
    "StatusCardComponent",
    "TaskListComponent",
    # Core implementations
    "Agent",
    "AgentConfig",
    "ToolRegistry",
    "DefaultSystemPromptBuilder",
    "DefaultWorkflowHandler",
    # Evaluation
    "Evaluator",
    "TestCase",
    "ExpectedOutcome",
    "AgentResult",
    "EvaluationResult",
    "TestCaseResult",
    "AgentVariant",
    "EvaluationRunner",
    "TrajectoryEvaluator",
    "OutputEvaluator",
    "LLMAsJudgeEvaluator",
    "EfficiencyEvaluator",
    "EvaluationReport",
    "ComparisonReport",
    "EvaluationDataset",
    # Basic implementations
    "MemoryConversationStore",
    "MockLlmService",
    # Server components
    "VannaFlaskServer",
    "VannaFastAPIServer",
    "ChatHandler",
    "ChatRequest",
    "ChatStreamChunk",
    "ExampleAgentLoader",
    # Exceptions
    "AgentError",
    "ToolExecutionError",
    "ToolNotFoundError",
    "PermissionError",
    "ConversationNotFoundError",
    "LlmServiceError",
    "ValidationError",
]


================================================
FILE: src/vanna/agents/__init__.py
================================================
"""
Agent implementations.

This package contains agent implementations and utilities.
"""

__all__: list[str] = []


================================================
FILE: src/vanna/capabilities/__init__.py
================================================
"""
Capabilities module.

This package contains abstractions for tool capabilities - reusable utilities
that tools can compose via dependency injection.
"""

from .file_system import CommandResult, FileSearchMatch, FileSystem
from .sql_runner import RunSqlToolArgs, SqlRunner

__all__ = [
    "FileSystem",
    "FileSearchMatch",
    "CommandResult",
    "SqlRunner",
    "RunSqlToolArgs",
]


================================================
FILE: src/vanna/capabilities/agent_memory/__init__.py
================================================
"""
Agent memory capability package.
"""

from .base import AgentMemory
from .models import (
    MemoryStats,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)

__all__ = [
    "AgentMemory",
    "TextMemory",
    "TextMemorySearchResult",
    "ToolMemory",
    "ToolMemorySearchResult",
    "MemoryStats",
]


================================================
FILE: src/vanna/capabilities/agent_memory/base.py
================================================
"""
Agent memory capability interface for tool usage learning.

This module contains the abstract base class for agent memory operations,
following the same pattern as the FileSystem interface.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from vanna.core.tool import ToolContext
    from .models import (
        ToolMemorySearchResult,
        TextMemory,
        TextMemorySearchResult,
        ToolMemory,
    )


class AgentMemory(ABC):
    """Abstract base class for agent memory operations."""

    @abstractmethod
    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: "ToolContext",
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern for future reference."""
        pass

    @abstractmethod
    async def save_text_memory(
        self, content: str, context: "ToolContext"
    ) -> "TextMemory":
        """Save a free-form text memory."""
        pass

    @abstractmethod
    async def search_similar_usage(
        self,
        question: str,
        context: "ToolContext",
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns based on a question."""
        pass

    @abstractmethod
    async def search_text_memories(
        self,
        query: str,
        context: "ToolContext",
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List["TextMemorySearchResult"]:
        """Search stored text memories based on a query."""
        pass

    @abstractmethod
    async def get_recent_memories(
        self, context: "ToolContext", limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories. Returns most recent memories first."""
        pass

    @abstractmethod
    async def get_recent_text_memories(
        self, context: "ToolContext", limit: int = 10
    ) -> List["TextMemory"]:
        """Fetch recently stored text memories."""
        pass

    @abstractmethod
    async def delete_by_id(self, context: "ToolContext", memory_id: str) -> bool:
        """Delete a memory by its ID. Returns True if deleted, False if not found."""
        pass

    @abstractmethod
    async def delete_text_memory(self, context: "ToolContext", memory_id: str) -> bool:
        """Delete a text memory by its ID. Returns True if deleted, False if not found."""
        pass

    @abstractmethod
    async def clear_memories(
        self,
        context: "ToolContext",
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories (tool or text). Returns number of memories deleted."""
        pass


================================================
FILE: src/vanna/capabilities/agent_memory/models.py
================================================
"""
Memory storage models and types.
"""

from typing import Any, Dict, List, Optional

from pydantic import BaseModel


class ToolMemory(BaseModel):
    """Represents a stored tool usage memory."""

    memory_id: Optional[str] = None
    question: str
    tool_name: str
    args: Dict[str, Any]
    timestamp: Optional[str] = None
    success: bool = True
    metadata: Optional[Dict[str, Any]] = None


class TextMemory(BaseModel):
    """Represents a stored free-form text memory."""

    memory_id: Optional[str] = None
    content: str
    timestamp: Optional[str] = None


class ToolMemorySearchResult(BaseModel):
    """Represents a search result from tool memory storage."""

    memory: ToolMemory
    similarity_score: float
    rank: int


class TextMemorySearchResult(BaseModel):
    """Represents a search result from text memory storage."""

    memory: TextMemory
    similarity_score: float
    rank: int


class MemoryStats(BaseModel):
    """Memory storage statistics."""

    total_memories: int
    unique_tools: int
    unique_questions: int
    success_rate: float
    most_used_tools: Dict[str, int]


================================================
FILE: src/vanna/capabilities/file_system/__init__.py
================================================
"""
File system capability.

This module provides abstractions for file system operations used by tools.
"""

from .base import FileSystem
from .models import CommandResult, FileSearchMatch

__all__ = [
    "FileSystem",
    "FileSearchMatch",
    "CommandResult",
]


================================================
FILE: src/vanna/capabilities/file_system/base.py
================================================
"""
File system capability interface.

This module contains the abstract base class for file system operations.
"""

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, List, Optional

from .models import CommandResult, FileSearchMatch

if TYPE_CHECKING:
    from vanna.core.tool import ToolContext


class FileSystem(ABC):
    """Abstract base class for file system operations."""

    @abstractmethod
    async def list_files(self, directory: str, context: "ToolContext") -> List[str]:
        """List files in a directory."""
        pass

    @abstractmethod
    async def read_file(self, filename: str, context: "ToolContext") -> str:
        """Read the contents of a file."""
        pass

    @abstractmethod
    async def write_file(
        self,
        filename: str,
        content: str,
        context: "ToolContext",
        overwrite: bool = False,
    ) -> None:
        """Write content to a file."""
        pass

    @abstractmethod
    async def exists(self, path: str, context: "ToolContext") -> bool:
        """Check if a file or directory exists."""
        pass

    @abstractmethod
    async def is_directory(self, path: str, context: "ToolContext") -> bool:
        """Check if a path is a directory."""
        pass

    @abstractmethod
    async def search_files(
        self,
        query: str,
        context: "ToolContext",
        *,
        max_results: int = 20,
        include_content: bool = False,
    ) -> List[FileSearchMatch]:
        """Search for files matching a query within the accessible namespace."""
        pass

    @abstractmethod
    async def run_bash(
        self,
        command: str,
        context: "ToolContext",
        *,
        timeout: Optional[float] = None,
    ) -> CommandResult:
        """Execute a bash command within the accessible namespace."""
        pass


================================================
FILE: src/vanna/capabilities/file_system/models.py
================================================
"""
File system capability models.

This module contains data models for file system operations.
"""

from dataclasses import dataclass
from typing import Optional


@dataclass
class FileSearchMatch:
    """Represents a single search result within a file system."""

    path: str
    snippet: Optional[str] = None


@dataclass
class CommandResult:
    """Represents the result of executing a shell command."""

    stdout: str
    stderr: str
    returncode: int


================================================
FILE: src/vanna/capabilities/sql_runner/__init__.py
================================================
"""
SQL runner capability.

This module provides abstractions for SQL execution used by tools.
"""

from .base import SqlRunner
from .models import RunSqlToolArgs

__all__ = [
    "SqlRunner",
    "RunSqlToolArgs",
]


================================================
FILE: src/vanna/capabilities/sql_runner/base.py
================================================
"""
SQL runner capability interface.

This module contains the abstract base class for SQL execution.
"""

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING

import pandas as pd

from .models import RunSqlToolArgs

if TYPE_CHECKING:
    from vanna.core.tool import ToolContext


class SqlRunner(ABC):
    """Interface for SQL execution with different implementations."""

    @abstractmethod
    async def run_sql(
        self, args: RunSqlToolArgs, context: "ToolContext"
    ) -> pd.DataFrame:
        """Execute SQL query and return results as a DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            Exception: If query execution fails
        """
        pass


================================================
FILE: src/vanna/capabilities/sql_runner/models.py
================================================
"""
SQL runner capability models.

This module contains data models for SQL execution.
"""

from pydantic import BaseModel, Field


class RunSqlToolArgs(BaseModel):
    """Arguments for run_sql tool."""

    sql: str = Field(description="SQL query to execute")


================================================
FILE: src/vanna/components/__init__.py
================================================
"""UI Component system for Vanna Agents."""

# Base component
from .base import UiComponent

# Simple components
from .simple import (
    SimpleComponent,
    SimpleComponentType,
    SimpleTextComponent,
    SimpleImageComponent,
    SimpleLinkComponent,
)

# Rich components - re-export all
from .rich import (
    # Base
    RichComponent,
    ComponentType,
    ComponentLifecycle,
    # Text
    RichTextComponent,
    # Data
    DataFrameComponent,
    ChartComponent,
    # Feedback
    NotificationComponent,
    StatusCardComponent,
    ProgressBarComponent,
    ProgressDisplayComponent,
    StatusIndicatorComponent,
    LogViewerComponent,
    LogEntry,
    BadgeComponent,
    IconTextComponent,
    # Interactive
    TaskListComponent,
    Task,
    StatusBarUpdateComponent,
    TaskTrackerUpdateComponent,
    ChatInputUpdateComponent,
    TaskOperation,
    ButtonComponent,
    ButtonGroupComponent,
    # Containers
    CardComponent,
    # Specialized
    ArtifactComponent,
)

__all__ = [
    # Base
    "UiComponent",
    # Simple components
    "SimpleComponent",
    "SimpleComponentType",
    "SimpleTextComponent",
    "SimpleImageComponent",
    "SimpleLinkComponent",
    # Rich components - Base
    "RichComponent",
    "ComponentType",
    "ComponentLifecycle",
    # Rich components - Text
    "RichTextComponent",
    # Rich components - Data
    "DataFrameComponent",
    "ChartComponent",
    # Rich components - Feedback
    "NotificationComponent",
    "StatusCardComponent",
    "ProgressBarComponent",
    "ProgressDisplayComponent",
    "StatusIndicatorComponent",
    "LogViewerComponent",
    "LogEntry",
    "BadgeComponent",
    "IconTextComponent",
    # Rich components - Interactive
    "TaskListComponent",
    "Task",
    "StatusBarUpdateComponent",
    "TaskTrackerUpdateComponent",
    "ChatInputUpdateComponent",
    "TaskOperation",
    "ButtonComponent",
    "ButtonGroupComponent",
    # Rich components - Containers
    "CardComponent",
    # Rich components - Specialized
    "ArtifactComponent",
]


================================================
FILE: src/vanna/components/base.py
================================================
"""
UI components base - re-exports UiComponent from core.

UiComponent lives in core/ because it's a fundamental return type for tools.
This module provides backward compatibility by re-exporting it here.
"""

# Re-export UiComponent from core for backward compatibility
from ..core.components import UiComponent

__all__ = ["UiComponent"]


================================================
FILE: src/vanna/components/rich/__init__.py
================================================
"""Rich UI components for the Vanna Agents framework."""

# Base classes and enums - import from core
from ...core.rich_component import RichComponent, ComponentType, ComponentLifecycle

# Text component
from .text import RichTextComponent

# Data components
from .data import (
    DataFrameComponent,
    ChartComponent,
)

# Feedback components
from .feedback import (
    NotificationComponent,
    StatusCardComponent,
    ProgressBarComponent,
    ProgressDisplayComponent,
    StatusIndicatorComponent,
    LogViewerComponent,
    LogEntry,
    BadgeComponent,
    IconTextComponent,
)

# Interactive components
from .interactive import (
    TaskListComponent,
    Task,
    StatusBarUpdateComponent,
    TaskTrackerUpdateComponent,
    ChatInputUpdateComponent,
    TaskOperation,
    ButtonComponent,
    ButtonGroupComponent,
)

# Container components
from .containers import (
    CardComponent,
)

# Specialized components
from .specialized import (
    ArtifactComponent,
)

__all__ = [
    # Base
    "RichComponent",
    "ComponentType",
    "ComponentLifecycle",
    # Text
    "RichTextComponent",
    # Data
    "DataFrameComponent",
    "ChartComponent",
    # Feedback
    "NotificationComponent",
    "StatusCardComponent",
    "ProgressBarComponent",
    "ProgressDisplayComponent",
    "StatusIndicatorComponent",
    "LogViewerComponent",
    "LogEntry",
    "BadgeComponent",
    "IconTextComponent",
    # Interactive
    "TaskListComponent",
    "Task",
    "StatusBarUpdateComponent",
    "TaskTrackerUpdateComponent",
    "ChatInputUpdateComponent",
    "TaskOperation",
    "ButtonComponent",
    "ButtonGroupComponent",
    # Containers
    "CardComponent",
    # Specialized
    "ArtifactComponent",
]


================================================
FILE: src/vanna/components/rich/containers/__init__.py
================================================
"""Container components for layout."""

from .card import CardComponent

__all__ = [
    "CardComponent",
]


================================================
FILE: src/vanna/components/rich/containers/card.py
================================================
"""Card component for displaying structured information."""

from typing import Any, Dict, List, Optional
from pydantic import Field
from ....core.rich_component import RichComponent, ComponentType


class CardComponent(RichComponent):
    """Card component for displaying structured information."""

    type: ComponentType = ComponentType.CARD
    title: str
    content: str
    subtitle: Optional[str] = None
    icon: Optional[str] = None
    status: Optional[str] = None  # "success", "warning", "error", "info"
    actions: List[Dict[str, Any]] = Field(default_factory=list)
    collapsible: bool = False
    collapsed: bool = False
    markdown: bool = False  # Whether content should be rendered as markdown


================================================
FILE: src/vanna/components/rich/data/__init__.py
================================================
"""Data display components."""

from .dataframe import DataFrameComponent
from .chart import ChartComponent

__all__ = [
    "DataFrameComponent",
    "ChartComponent",
]


================================================
FILE: src/vanna/components/rich/data/chart.py
================================================
"""Chart component for data visualization."""

from typing import Any, Dict, Optional, Union
from pydantic import Field
from ....core.rich_component import RichComponent, ComponentType


class ChartComponent(RichComponent):
    """Chart component for data visualization."""

    type: ComponentType = ComponentType.CHART
    chart_type: str  # "line", "bar", "pie", "scatter", etc.
    data: Dict[str, Any]  # Chart data in format expected by frontend
    title: Optional[str] = None
    width: Optional[Union[str, int]] = None
    height: Optional[Union[str, int]] = None
    config: Dict[str, Any] = Field(default_factory=dict)  # Chart-specific config


================================================
FILE: src/vanna/components/rich/data/dataframe.py
================================================
"""DataFrame component for displaying tabular data."""

from typing import Any, Dict, List, Optional
from pydantic import Field
from ....core.rich_component import RichComponent, ComponentType


class DataFrameComponent(RichComponent):
    """DataFrame component specifically for displaying tabular data from SQL queries and similar sources."""

    type: ComponentType = ComponentType.DATAFRAME
    rows: List[Dict[str, Any]] = Field(default_factory=list)  # List of row dictionaries
    columns: List[str] = Field(default_factory=list)  # Column names in display order
    title: Optional[str] = None
    description: Optional[str] = None
    row_count: int = 0
    column_count: int = 0

    # Display options
    max_rows_displayed: int = 100  # Limit rows shown in UI
    searchable: bool = True
    sortable: bool = True
    filterable: bool = True
    exportable: bool = True  # Allow export to CSV/Excel

    # Styling options
    striped: bool = True
    bordered: bool = True
    compact: bool = False

    # Pagination
    paginated: bool = True
    page_size: int = 25

    # Data types for better formatting (optional)
    column_types: Dict[str, str] = Field(
        default_factory=dict
    )  # column_name -> "string"|"number"|"date"|"boolean"

    def __init__(self, **kwargs: Any) -> None:
        # Set defaults before calling super().__init__
        if "rows" not in kwargs:
            kwargs["rows"] = []
        if "columns" not in kwargs:
            kwargs["columns"] = []
        if "column_types" not in kwargs:
            kwargs["column_types"] = {}

        super().__init__(**kwargs)

        # Auto-calculate counts if not provided
        if self.rows and len(self.rows) > 0:
            if "row_count" not in kwargs:
                self.row_count = len(self.rows)
            if not self.columns and self.rows:
                self.columns = list(self.rows[0].keys())
            if "column_count" not in kwargs:
                self.column_count = len(self.columns)
        else:
            if "row_count" not in kwargs:
                self.row_count = 0
            if "column_count" not in kwargs:
                self.column_count = len(self.columns) if self.columns else 0

    @classmethod
    def from_records(
        cls,
        records: List[Dict[str, Any]],
        title: Optional[str] = None,
        description: Optional[str] = None,
        **kwargs: Any,
    ) -> "DataFrameComponent":
        """Create a DataFrame component from a list of record dictionaries."""
        columns = list(records[0].keys()) if records else []

        # Ensure we pass the required arguments correctly
        component_data = {
            "rows": records,
            "columns": columns,
            "row_count": len(records),
            "column_count": len(columns),
            "column_types": {},  # Initialize empty dict
        }

        if title is not None:
            component_data["title"] = title
        if description is not None:
            component_data["description"] = description

        # Merge with any additional kwargs
        component_data.update(kwargs)

        return cls(**component_data)


================================================
FILE: src/vanna/components/rich/feedback/__init__.py
================================================
"""User feedback components."""

from .notification import NotificationComponent
from .status_card import StatusCardComponent
from .progress import ProgressBarComponent, ProgressDisplayComponent
from .status_indicator import StatusIndicatorComponent
from .log_viewer import LogViewerComponent, LogEntry
from .badge import BadgeComponent
from .icon_text import IconTextComponent

__all__ = [
    "NotificationComponent",
    "StatusCardComponent",
    "ProgressBarComponent",
    "ProgressDisplayComponent",
    "StatusIndicatorComponent",
    "LogViewerComponent",
    "LogEntry",
    "BadgeComponent",
    "IconTextComponent",
]


================================================
FILE: src/vanna/components/rich/feedback/badge.py
================================================
"""Badge component for displaying status or labels."""

from typing import Optional
from ....core.rich_component import RichComponent, ComponentType


class BadgeComponent(RichComponent):
    """Simple badge/pill component for displaying status or labels."""

    type: ComponentType = ComponentType.BADGE
    text: str
    variant: str = (
        "default"  # "default", "primary", "success", "warning", "error", "info"
    )
    size: str = "medium"  # "small", "medium", "large"
    icon: Optional[str] = None


================================================
FILE: src/vanna/components/rich/feedback/icon_text.py
================================================
"""Icon with text component."""

from ....core.rich_component import RichComponent, ComponentType


class IconTextComponent(RichComponent):
    """Simple component for displaying an icon with text."""

    type: ComponentType = ComponentType.ICON_TEXT
    icon: str
    text: str
    variant: str = "default"  # "default", "primary", "secondary", "muted"
    size: str = "medium"  # "small", "medium", "large"
    alignment: str = "left"  # "left", "center", "right"


================================================
FILE: src/vanna/components/rich/feedback/log_viewer.py
================================================
"""Log viewer component."""

import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
from ....core.rich_component import RichComponent, ComponentType


class LogEntry(BaseModel):
    """Log entry for tool execution."""

    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
    level: str = "info"  # "debug", "info", "warning", "error"
    message: str
    data: Optional[Dict[str, Any]] = None


class LogViewerComponent(RichComponent):
    """Generic log viewer for displaying timestamped entries."""

    type: ComponentType = ComponentType.LOG_VIEWER
    title: str = "Logs"
    entries: List[LogEntry] = Field(default_factory=list)
    max_entries: int = 100
    searchable: bool = True
    show_timestamps: bool = True
    auto_scroll: bool = True

    def add_entry(
        self, message: str, level: str = "info", data: Optional[Dict[str, Any]] = None
    ) -> "LogViewerComponent":
        """Add a new log entry."""
        new_entry = LogEntry(message=message, level=level, data=data)
        new_entries = self.entries + [new_entry]

        # Limit to max_entries
        if len(new_entries) > self.max_entries:
            new_entries = new_entries[-self.max_entries :]

        return self.update(entries=new_entries)


================================================
FILE: src/vanna/components/rich/feedback/notification.py
================================================
"""Notification component for alerts and messages."""

from typing import Any, Dict, List, Optional
from pydantic import Field
from ....core.rich_component import RichComponent, ComponentType


class NotificationComponent(RichComponent):
    """Notification component for alerts and messages."""

    type: ComponentType = ComponentType.NOTIFICATION
    message: str
    title: Optional[str] = None
    level: str = "info"  # "success", "info", "warning", "error"
    icon: Optional[str] = None
    dismissible: bool = True
    auto_dismiss: bool = False
    auto_dismiss_delay: int = 5000  # milliseconds
    actions: List[Dict[str, Any]] = Field(default_factory=list)


================================================
FILE: src/vanna/components/rich/feedback/progress.py
================================================
"""Progress components for displaying progress indicators."""

from typing import Any, Dict, Optional
from ....core.rich_component import RichComponent, ComponentType


class ProgressBarComponent(RichComponent):
    """Progress bar with status and value."""

    type: ComponentType = ComponentType.PROGRESS_BAR
    value: float  # 0.0 to 1.0
    label: Optional[str] = None
    show_percentage: bool = True
    status: Optional[str] = None  # "success", "warning", "error"
    animated: bool = False


class ProgressDisplayComponent(RichComponent):
    """Generic progress display for any long-running process."""

    type: ComponentType = ComponentType.PROGRESS_DISPLAY
    label: str
    value: float = 0.0  # 0.0 to 1.0
    description: Optional[str] = None
    status: Optional[str] = None  # "info", "success", "warning", "error"
    show_percentage: bool = True
    animated: bool = False
    indeterminate: bool = False

    def update_progress(
        self, value: float, description: Optional[str] = None
    ) -> "ProgressDisplayComponent":
        """Update progress value and optionally description."""
        updates: Dict[str, Any] = {"value": max(0.0, min(1.0, value))}
        if description is not None:
            updates["description"] = description
        return self.update(**updates)


================================================
FILE: src/vanna/components/rich/feedback/status_card.py
================================================
"""Status card component for displaying process status."""

from typing import Any, Dict, List, Optional
from pydantic import Field
from ....core.rich_component import RichComponent, ComponentType


class StatusCardComponent(RichComponent):
    """Generic status card that can display any process status."""

    type: ComponentType = ComponentType.STATUS_CARD
    title: str
    status: str  # "pending", "running", "completed", "failed", "success", "warning", "error"
    description: Optional[str] = None
    icon: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)
    actions: List[Dict[str, Any]] = Field(default_factory=list)
    collapsible: bool = False
    collapsed: bool = False

    def set_status(
        self, status: str, description: Optional[str] = None
    ) -> "StatusCardComponent":
        """Update the status and optionally the description."""
        updates = {"status": status}
        if description is not None:
            updates["description"] = description
        return self.update(**updates)


================================================
FILE: src/vanna/components/rich/feedback/status_indicator.py
================================================
"""Status indicator component."""

from typing import Optional
from ....core.rich_component import RichComponent, ComponentType


class StatusIndicatorComponent(RichComponent):
    """Status indicator with icon and message."""

    type: ComponentType = ComponentType.STATUS_INDICATOR
    status: str  # "success", "warning", "error", "info", "loading"
    message: str
    icon: Optional[str] = None
    pulse: bool = False


================================================
FILE: src/vanna/components/rich/interactive/__init__.py
================================================
"""Interactive components."""

from .task_list import TaskListComponent, Task
from .ui_state import (
    StatusBarUpdateComponent,
    TaskTrackerUpdateComponent,
    ChatInputUpdateComponent,
    TaskOperation,
)
from .button import ButtonComponent, ButtonGroupComponent

__all__ = [
    "TaskListComponent",
    "Task",
    "StatusBarUpdateComponent",
    "TaskTrackerUpdateComponent",
    "ChatInputUpdateComponent",
    "TaskOperation",
    "ButtonComponent",
    "ButtonGroupComponent",
]


================================================
FILE: src/vanna/components/rich/interactive/button.py
================================================
"""Button component for interactive actions."""

from typing import Any, Dict, List, Literal, Optional
from ....core.rich_component import ComponentType, RichComponent


class ButtonComponent(RichComponent):
    """Interactive button that sends a message when clicked.

    The button renders in the UI and when clicked, sends its action
    value as a message to the chat input.

    Args:
        label: Text displayed on the button
        action: Message/command to send when clicked
        variant: Visual style variant
        size: Button size
        icon: Optional emoji or icon
        icon_position: Position of icon relative to label
        disabled: Whether button is disabled

    Example:
        ButtonComponent(
            label="Generate Report",
            action="/report sales",
            variant="primary",
            icon="📊"
        )
    """

    def __init__(
        self,
        label: str,
        action: str,
        variant: Literal[
            "primary", "secondary", "success", "warning", "error", "ghost", "link"
        ] = "primary",
        size: Literal["small", "medium", "large"] = "medium",
        icon: Optional[str] = None,
        icon_position: Literal["left", "right"] = "left",
        disabled: bool = False,
    ):
        super().__init__(
            type=ComponentType.BUTTON,
            data={
                "label": label,
                "action": action,
                "variant": variant,
                "size": size,
                "icon": icon,
                "icon_position": icon_position,
                "disabled": disabled,
            },
        )


class ButtonGroupComponent(RichComponent):
    """Group of buttons with consistent styling.

    Args:
        buttons: List of button data dictionaries
        orientation: Layout direction
        spacing: Gap between buttons
        alignment: Button alignment within group
        full_width: Whether buttons should stretch to fill width

    Example:
        ButtonGroupComponent(
            buttons=[
                {"label": "Yes", "action": "/confirm yes", "variant": "success"},
                {"label": "No", "action": "/confirm no", "variant": "error"},
            ],
            orientation="horizontal",
            spacing="medium"
        )
    """

    def __init__(
        self,
        buttons: List[Dict[str, Any]],
        orientation: Literal["horizontal", "vertical"] = "horizontal",
        spacing: Literal["small", "medium", "large"] = "medium",
        alignment: Literal["start", "center", "end", "stretch"] = "start",
        full_width: bool = False,
    ):
        super().__init__(
            type=ComponentType.BUTTON_GROUP,
            data={
                "buttons": buttons,
                "orientation": orientation,
                "spacing": spacing,
                "alignment": alignment,
                "full_width": full_width,
            },
        )


================================================
FILE: src/vanna/components/rich/interactive/task_list.py
================================================
"""Task list component for interactive task tracking."""

import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
from ....core.rich_component import RichComponent, ComponentType


class Task(BaseModel):
    """Individual task in a task list."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    title: str
    description: Optional[str] = None
    status: str = "pending"  # "pending", "in_progress", "completed", "error"
    progress: Optional[float] = None  # 0.0 to 1.0
    created_at: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
    completed_at: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)


class TaskListComponent(RichComponent):
    """Interactive task list with progress tracking."""

    type: ComponentType = ComponentType.TASK_LIST
    title: str = "Tasks"
    tasks: List[Task] = Field(default_factory=list)
    show_progress: bool = True
    allow_reorder: bool = False
    show_timestamps: bool = True
    filter_status: Optional[str] = None  # Filter by task status

    def add_task(self, task: Task) -> "TaskListComponent":
        """Add a task to the list."""
        new_tasks = self.tasks + [task]
        return self.update(tasks=new_tasks)

    def update_task(self, task_id: str, **updates: Any) -> "TaskListComponent":
        """Update a specific task."""
        new_tasks = []
        for task in self.tasks:
            if task.id == task_id:
                task_data = task.model_dump()
                task_data.update(updates)
                new_tasks.append(Task(**task_data))
            else:
                new_tasks.append(task)
        return self.update(tasks=new_tasks)

    def complete_task(self, task_id: str) -> "TaskListComponent":
        """Mark a task as completed."""
        return self.update_task(
            task_id,
            status="completed",
            completed_at=datetime.utcnow().isoformat(),
            progress=1.0,
        )


================================================
FILE: src/vanna/components/rich/interactive/ui_state.py
================================================
"""UI state update components for controlling interface elements."""

from enum import Enum
from typing import Any, Optional
from .task_list import Task
from ....core.rich_component import RichComponent, ComponentType


class StatusBarUpdateComponent(RichComponent):
    """Component for updating the status bar above chat input."""

    type: ComponentType = ComponentType.STATUS_BAR_UPDATE
    status: str  # "idle", "working", "success", "error"
    message: str
    detail: Optional[str] = None

    def __init__(self, **kwargs: Any) -> None:
        # Set a fixed ID for status bar updates
        kwargs.setdefault("id", "vanna-status-bar")
        super().__init__(**kwargs)


class TaskOperation(str, Enum):
    """Operations for task tracker updates."""

    ADD_TASK = "add_task"
    UPDATE_TASK = "update_task"
    REMOVE_TASK = "remove_task"
    CLEAR_TASKS = "clear_tasks"


class TaskTrackerUpdateComponent(RichComponent):
    """Component for updating the task tracker in the sidebar."""

    type: ComponentType = ComponentType.TASK_TRACKER_UPDATE
    operation: TaskOperation
    task: Optional[Task] = None  # Used for ADD_TASK
    task_id: Optional[str] = None  # Used for UPDATE_TASK and REMOVE_TASK
    status: Optional[str] = None  # Used for UPDATE_TASK
    progress: Optional[float] = None  # Used for UPDATE_TASK
    detail: Optional[str] = None  # Used for UPDATE_TASK

    def __init__(self, **kwargs: Any) -> None:
        # Set a fixed ID for task tracker updates
        kwargs.setdefault("id", "vanna-task-tracker")
        super().__init__(**kwargs)

    @classmethod
    def add_task(cls, task: Task) -> "TaskTrackerUpdateComponent":
        """Create a component to add a new task."""
        return cls(operation=TaskOperation.ADD_TASK, task=task)

    @classmethod
    def update_task(
        cls,
        task_id: str,
        status: Optional[str] = None,
        progress: Optional[float] = None,
        detail: Optional[str] = None,
    ) -> "TaskTrackerUpdateComponent":
        """Create a component to update an existing task."""
        return cls(
            operation=TaskOperation.UPDATE_TASK,
            task_id=task_id,
            status=status,
            progress=progress,
            detail=detail,
        )

    @classmethod
    def remove_task(cls, task_id: str) -> "TaskTrackerUpdateComponent":
        """Create a component to remove a task."""
        return cls(operation=TaskOperation.REMOVE_TASK, task_id=task_id)

    @classmethod
    def clear_tasks(cls) -> "TaskTrackerUpdateComponent":
        """Create a component to clear all tasks."""
        return cls(operation=TaskOperation.CLEAR_TASKS)


class ChatInputUpdateComponent(RichComponent):
    """Component for updating chat input state and appearance."""

    type: ComponentType = ComponentType.CHAT_INPUT_UPDATE
    placeholder: Optional[str] = None
    disabled: Optional[bool] = None
    value: Optional[str] = None  # Set input text value
    focus: Optional[bool] = None  # Focus/unfocus the input

    def __init__(self, **kwargs: Any) -> None:
        # Set a fixed ID for chat input updates
        kwargs.setdefault("id", "vanna-chat-input")
        super().__init__(**kwargs)


================================================
FILE: src/vanna/components/rich/specialized/__init__.py
================================================
"""Specialized components."""

from .artifact import ArtifactComponent

__all__ = [
    "ArtifactComponent",
]


================================================
FILE: src/vanna/components/rich/specialized/artifact.py
================================================
"""Artifact component for interactive content."""

import uuid
from typing import Optional
from pydantic import Field
from ....core.rich_component import RichComponent, ComponentType


class ArtifactComponent(RichComponent):
    """Component for displaying interactive artifacts that can be rendered externally."""

    type: ComponentType = ComponentType.ARTIFACT
    artifact_id: str = Field(default_factory=lambda: f"artifact_{uuid.uuid4().hex[:8]}")
    content: str  # HTML/SVG/JS content
    artifact_type: str  # "html", "svg", "visualization", "interactive", "d3", "threejs"
    title: Optional[str] = None
    description: Optional[str] = None
    editable: bool = True
    fullscreen_capable: bool = True
    external_renderable: bool = True


================================================
FILE: src/vanna/components/rich/text.py
================================================
"""Rich text component."""

from typing import Optional
from ...core.rich_component import RichComponent, ComponentType


class RichTextComponent(RichComponent):
    """Rich text component with formatting options."""

    type: ComponentType = ComponentType.TEXT
    content: str
    markdown: bool = False
    code_language: Optional[str] = None  # For syntax highlighting
    font_size: Optional[str] = None
    font_weight: Optional[str] = None
    text_align: Optional[str] = None


================================================
FILE: src/vanna/components/simple/__init__.py
================================================
"""Simple UI components for basic rendering."""

# Import from core
from ...core.simple_component import SimpleComponent, SimpleComponentType
from .text import SimpleTextComponent
from .image import SimpleImageComponent
from .link import SimpleLinkComponent

__all__ = [
    "SimpleComponent",
    "SimpleComponentType",
    "SimpleTextComponent",
    "SimpleImageComponent",
    "SimpleLinkComponent",
]


================================================
FILE: src/vanna/components/simple/image.py
================================================
"""Simple image component."""

from typing import Optional
from pydantic import Field
from ...core.simple_component import SimpleComponent, SimpleComponentType


class SimpleImageComponent(SimpleComponent):
    """A simple image component."""

    type: SimpleComponentType = SimpleComponentType.IMAGE
    url: str = Field(..., description="The URL of the image to display.")
    alt_text: Optional[str] = Field(
        default=None, description="Alternative text for the image."
    )


================================================
FILE: src/vanna/components/simple/link.py
================================================
"""Simple link component."""

from typing import Optional
from pydantic import Field
from ...core.simple_component import SimpleComponent, SimpleComponentType


class SimpleLinkComponent(SimpleComponent):
    """A simple link component."""

    type: SimpleComponentType = SimpleComponentType.LINK
    url: str = Field(..., description="The URL the link points to.")
    text: Optional[str] = Field(
        default=None, description="The display text for the link."
    )


================================================
FILE: src/vanna/components/simple/text.py
================================================
"""Simple text component."""

from pydantic import Field
from ...core.simple_component import SimpleComponent, SimpleComponentType


class SimpleTextComponent(SimpleComponent):
    """A simple text component."""

    type: SimpleComponentType = SimpleComponentType.TEXT
    text: str = Field(..., description="The text content to display.")


================================================
FILE: src/vanna/core/__init__.py
================================================
"""
Core components of the Vanna Agents framework.

This package contains the fundamental abstractions and implementations
that form the foundation of the agent framework.
"""

# Core domains - re-export from new structure
from .tool import T, Tool, ToolCall, ToolContext, ToolResult, ToolSchema
from .llm import LlmMessage, LlmRequest, LlmResponse, LlmService, LlmStreamChunk
from .storage import Conversation, ConversationStore, Message
from .user import User, UserService
from .agent import Agent, AgentConfig
from .system_prompt import DefaultSystemPromptBuilder, SystemPromptBuilder
from .lifecycle import LifecycleHook
from .middleware import LlmMiddleware
from .workflow import WorkflowHandler, WorkflowResult, DefaultWorkflowHandler
from .recovery import ErrorRecoveryStrategy, RecoveryAction, RecoveryActionType
from .enricher import ToolContextEnricher
from .enhancer import LlmContextEnhancer, DefaultLlmContextEnhancer
from .filter import ConversationFilter
from .observability import ObservabilityProvider, Span, Metric
from .audit import (
    AuditLogger,
    AuditEvent,
    AuditEventType,
    ToolAccessCheckEvent,
    ToolInvocationEvent,
    ToolResultEvent,
    UiFeatureAccessCheckEvent,
    AiResponseEvent,
)

# UI Components
from .components import UiComponent
from .rich_component import RichComponent
from ..components import (
    SimpleComponent,
    SimpleComponentType,
    SimpleImageComponent,
    SimpleLinkComponent,
    SimpleTextComponent,
    ArtifactComponent,
    BadgeComponent,
    CardComponent,
    DataFrameComponent,
    IconTextComponent,
    LogViewerComponent,
    NotificationComponent,
    ProgressBarComponent,
    ProgressDisplayComponent,
    RichTextComponent,
    StatusCardComponent,
    TaskListComponent,
)

# Exceptions
from .errors import (
    AgentError,
    ConversationNotFoundError,
    LlmServiceError,
    PermissionError,
    ToolExecutionError,
    ToolNotFoundError,
    ValidationError,
)

# Core implementations
from .registry import ToolRegistry

# Evaluation framework
from .evaluation import (
    Evaluator,
    TestCase,
    ExpectedOutcome,
    AgentResult,
    EvaluationResult,
    TestCaseResult,
    AgentVariant,
    EvaluationRunner,
    TrajectoryEvaluator,
    OutputEvaluator,
    LLMAsJudgeEvaluator,
    EfficiencyEvaluator,
    EvaluationReport,
    ComparisonReport,
    EvaluationDataset,
)

# Rebuild models to resolve forward references after all imports
from .tool.models import ToolContext, ToolResult
from .components import UiComponent  # Import UiComponent to ensure it's available

ToolContext.model_rebuild()
ToolResult.model_rebuild()

__all__ = [
    # Models
    "User",
    "Message",
    "Conversation",
    "ToolCall",
    "ToolResult",
    "ToolContext",
    "ToolSchema",
    "LlmMessage",
    "LlmRequest",
    "LlmResponse",
    "LlmStreamChunk",
    "RecoveryAction",
    "RecoveryActionType",
    "Span",
    "Metric",
    # Interfaces
    "Tool",
    "Agent",
    "LlmService",
    "ConversationStore",
    "UserService",
    "SystemPromptBuilder",
    "LifecycleHook",
    "LlmMiddleware",
    "WorkflowHandler",
    "DefaultWorkflowHandler",
    "WorkflowResult",
    "ErrorRecoveryStrategy",
    "ToolContextEnricher",
    "LlmContextEnhancer",
    "DefaultLlmContextEnhancer",
    "ConversationFilter",
    "ObservabilityProvider",
    "AuditLogger",
    "T",
    # Audit
    "AuditEvent",
    "AuditEventType",
    "ToolAccessCheckEvent",
    "ToolInvocationEvent",
    "ToolResultEvent",
    "UiFeatureAccessCheckEvent",
    "AiResponseEvent",
    # UI Components
    "UiComponent",
    # Simple Components
    "SimpleComponent",
    "SimpleComponentType",
    "SimpleTextComponent",
    "SimpleImageComponent",
    "SimpleLinkComponent",
    # Rich Components
    "RichComponent",
    "ArtifactComponent",
    "BadgeComponent",
    "CardComponent",
    "DataFrameComponent",
    "IconTextComponent",
    "LogViewerComponent",
    "NotificationComponent",
    "ProgressBarComponent",
    "ProgressDisplayComponent",
    "RichTextComponent",
    "StatusCardComponent",
    "TaskListComponent",
    # Core implementations
    "ToolRegistry",
    "Agent",
    "AgentConfig",
    "DefaultSystemPromptBuilder",
    # Evaluation
    "Evaluator",
    "TestCase",
    "ExpectedOutcome",
    "AgentResult",
    "EvaluationResult",
    "TestCaseResult",
    "AgentVariant",
    "EvaluationRunner",
    "TrajectoryEvaluator",
    "OutputEvaluator",
    "LLMAsJudgeEvaluator",
    "EfficiencyEvaluator",
    "EvaluationReport",
    "ComparisonReport",
    "EvaluationDataset",
    # Exceptions
    "AgentError",
    "ToolExecutionError",
    "ToolNotFoundError",
    "PermissionError",
    "ConversationNotFoundError",
    "LlmServiceError",
    "ValidationError",
]


================================================
FILE: src/vanna/core/_compat.py
================================================
"""
Compatibility shims for different Python versions.

This module provides compatibility utilities for features that vary across
Python versions.
"""

try:
    from enum import StrEnum  # Py 3.11+
except ImportError:  # Py < 3.11
    from enum import Enum

    class StrEnum(str, Enum):  # type: ignore[no-redef]
        """Minimal backport of StrEnum for Python < 3.11."""

        pass


__all__ = ["StrEnum"]


================================================
FILE: src/vanna/core/agent/__init__.py
================================================
"""
Agent module.

This module contains the core Agent implementation and configuration.
"""

from .agent import Agent
from .config import AgentConfig

__all__ = ["Agent", "AgentConfig"]


================================================
FILE: src/vanna/core/agent/agent.py
================================================
"""
Agent implementation for the Vanna Agents framework.

This module provides the main Agent class that orchestrates the interaction
between LLM services, tools, and conversation storage.
"""

import traceback
import uuid
from typing import TYPE_CHECKING, AsyncGenerator, List, Optional

from vanna.components import (
    UiComponent,
    SimpleTextComponent,
    RichTextComponent,
    StatusBarUpdateComponent,
    TaskTrackerUpdateComponent,
    ChatInputUpdateComponent,
    StatusCardComponent,
    Task,
)
from .config import AgentConfig
from vanna.core.storage import ConversationStore
from vanna.core.llm import LlmService
from vanna.core.system_prompt import SystemPromptBuilder
from vanna.core.storage import Conversation, Message
from vanna.core.llm import LlmMessage, LlmRequest, LlmResponse
from vanna.core.tool import ToolCall, ToolContext, ToolResult, ToolSchema
from vanna.core.user import User
from vanna.core.registry import ToolRegistry
from vanna.core.system_prompt import DefaultSystemPromptBuilder
from vanna.core.lifecycle import LifecycleHook
from vanna.core.middleware import LlmMiddleware
from vanna.core.workflow import WorkflowHandler, DefaultWorkflowHandler
from vanna.core.recovery import ErrorRecoveryStrategy, RecoveryActionType
from vanna.core.enricher import ToolContextEnricher
from vanna.core.enhancer import LlmContextEnhancer, DefaultLlmContextEnhancer
from vanna.core.filter import ConversationFilter
from vanna.core.observability import ObservabilityProvider
from vanna.core.user.resolver import UserResolver
from vanna.core.user.request_context import RequestContext
from vanna.core.agent.config import UiFeature
from vanna.core.audit import AuditLogger
from vanna.capabilities.agent_memory import AgentMemory

import logging

logger = logging.getLogger(__name__)

logger.info("Loaded vanna.core.agent.agent module")

if TYPE_CHECKING:
    pass


class Agent:
    """Main agent implementation.

    The Agent class orchestrates LLM interactions, tool execution, and conversation
    management. It provides 7 extensibility points for customization:

    - lifecycle_hooks: Hook into message and tool execution lifecycle
    - llm_middlewares: Intercept and transform LLM requests/responses
    - error_recovery_strategy: Handle errors with retry logic
    - context_enrichers: Add data to tool execution context
    - llm_context_enhancer: Enhance LLM system prompts and messages with context
    - conversation_filters: Filter conversation history before LLM calls
    - observability_provider: Collect telemetry and monitoring data

    Example:
        agent = Agent(
            llm_service=AnthropicLlmService(api_key="..."),
            tool_registry=registry,
            conversation_store=store,
            lifecycle_hooks=[QuotaCheckHook()],
            llm_middlewares=[CachingMiddleware()],
            llm_context_enhancer=DefaultLlmContextEnhancer(agent_memory),
            observability_provider=LoggingProvider()
        )
    """

    def __init__(
        self,
        llm_service: LlmService,
        tool_registry: ToolRegistry,
        user_resolver: UserResolver,
        agent_memory: AgentMemory,
        conversation_store: Optional[ConversationStore] = None,
        config: AgentConfig = AgentConfig(),
        system_prompt_builder: SystemPromptBuilder = DefaultSystemPromptBuilder(),
        lifecycle_hooks: List[LifecycleHook] = [],
        llm_middlewares: List[LlmMiddleware] = [],
        workflow_handler: Optional[WorkflowHandler] = None,
        error_recovery_strategy: Optional[ErrorRecoveryStrategy] = None,
        context_enrichers: List[ToolContextEnricher] = [],
        llm_context_enhancer: Optional[LlmContextEnhancer] = None,
        conversation_filters: List[ConversationFilter] = [],
        observability_provider: Optional[ObservabilityProvider] = None,
        audit_logger: Optional[AuditLogger] = None,
    ):
        self.llm_service = llm_service
        self.tool_registry = tool_registry
        self.user_resolver = user_resolver
        self.agent_memory = agent_memory

        # Import here to avoid circular dependency
        if conversation_store is None:
            from vanna.integrations.local import MemoryConversationStore

            conversation_store = MemoryConversationStore()

        self.conversation_store = conversation_store
        self.config = config
        self.system_prompt_builder = system_prompt_builder
        self.lifecycle_hooks = lifecycle_hooks
        self.llm_middlewares = llm_middlewares

        # Use DefaultWorkflowHandler if none provided
        if workflow_handler is None:
            workflow_handler = DefaultWorkflowHandler()
        self.workflow_handler = workflow_handler

        self.error_recovery_strategy = error_recovery_strategy
        self.context_enrichers = context_enrichers

        # Use DefaultLlmContextEnhancer if none provided
        if llm_context_enhancer is None:
            llm_context_enhancer = DefaultLlmContextEnhancer(agent_memory)
        self.llm_context_enhancer = llm_context_enhancer

        self.conversation_filters = conversation_filters
        self.observability_provider = observability_provider
        self.audit_logger = audit_logger

        # Wire audit logger into tool registry
        if self.audit_logger and self.config.audit_config.enabled:
            self.tool_registry.audit_logger = self.audit_logger
            self.tool_registry.audit_config = self.config.audit_config

        logger.info("Initialized Agent")

    async def send_message(
        self,
        request_context: RequestContext,
        message: str,
        *,
        conversation_id: Optional[str] = None,
    ) -> AsyncGenerator[UiComponent, None]:
        """
        Process a user message and yield UI components with error handling.

        Args:
            request_context: Request context for user resolution (includes metadata)
            message: User's message content
            conversation_id: Optional conversation ID; if None, creates new conversation

        Yields:
            UiComponent instances for UI updates
        """
        try:
            # Delegate to internal method
            async for component in self._send_message(
                request_context, message, conversation_id=conversation_id
            ):
                yield component
        except Exception as e:
            # Log full stack trace
            stack_trace = traceback.format_exc()
            logger.error(
                f"Error in send_message (conversation_id={conversation_id}): {e}\n{stack_trace}",
                exc_info=True,
            )

            # Log to observability provider if available
            if self.observability_provider:
                try:
                    error_span = await self.observability_provider.create_span(
                        "agent.send_message.error",
                        attributes={
                            "error_type": type(e).__name__,
                            "error_message": str(e),
                            "conversation_id": conversation_id or "none",
                        },
                    )
                    await self.observability_provider.end_span(error_span)
                    await self.observability_provider.record_metric(
                        "agent.error.count",
                        1.0,
                        "count",
                        tags={"error_type": type(e).__name__},
                    )
                except Exception as obs_error:
                    logger.error(
                        f"Failed to log error to observability provider: {obs_error}",
                        exc_info=True,
                    )

            # Yield error component to UI (simple, user-friendly message)
            error_description = "An unexpected error occurred while processing your message. Please try again."
            if conversation_id:
                error_description += f"\n\nConversation ID: {conversation_id}"

            yield UiComponent(
                rich_component=StatusCardComponent(
                    title="Error Processing Message",
                    status="error",
                    description=error_description,
                    icon="⚠️",
                ),
                simple_component=SimpleTextComponent(
                    text=f"Error: An unexpected error occurred. Please try again.{f' (Conversation ID: {conversation_id})' if conversation_id else ''}"
                ),
            )

            # Update status bar to show error state
            yield UiComponent(  # type: ignore
                rich_component=StatusBarUpdateComponent(
                    status="error",
                    message="Error occurred",
                    detail="An unexpected error occurred while processing your message",
                )
            )

            # Re-enable chat input so user can try again
            yield UiComponent(  # type: ignore
                rich_component=ChatInputUpdateComponent(
                    placeholder="Try again...", disabled=False
                )
            )

    async def _send_message(
        self,
        request_context: RequestContext,
        message: str,
        *,
        conversation_id: Optional[str] = None,
    ) -> AsyncGenerator[UiComponent, None]:
        """
        Internal method to process a user message and yield UI components.

        Args:
            request_context: Request context for user resolution (includes metadata)
            message: User's message content
            conversation_id: Optional conversation ID; if None, creates new conversation

        Yields:
            UiComponent instances for UI updates
        """
        # Resolve user from request context with observability
        user_resolution_span = None
        if self.observability_provider:
            user_resolution_span = await self.observability_provider.create_span(
                "agent.user_resolution",
                attributes={"has_context": request_context is not None},
            )

        user = await self.user_resolver.resolve_user(request_context)

        if self.observability_provider and user_resolution_span:
            user_resolution_span.set_attribute("user_id", user.id)
            await self.observability_provider.end_span(user_resolution_span)
            if user_resolution_span.duration_ms():
                await self.observability_provider.record_metric(
                    "agent.user_resolution.duration",
                    user_resolution_span.duration_ms() or 0,
                    "ms",
                )

        # Check if this is a starter UI request (empty message or explicit metadata flag)
        is_starter_request = (not message.strip()) or request_context.metadata.get(
            "starter_ui_request", False
        )

        if is_starter_request and self.workflow_handler:
            # Handle starter UI request with observability
            starter_span = None
            if self.observability_provider:
                starter_span = await self.observability_provider.create_span(
                    "agent.workflow_handler.starter_ui", attributes={"user_id": user.id}
                )

            try:
                # Load or create conversation for context
                if conversation_id is None:
                    conversation_id = str(uuid.uuid4())

                conversation = await self.conversation_store.get_conversation(
                    conversation_id, user
                )
                if not conversation:
                    # Create empty conversation (will be saved if workflow produces components)
                    conversation = Conversation(
                        id=conversation_id, user=user, messages=[]
                    )

                # Get starter UI from workflow handler
                components = await self.workflow_handler.get_starter_ui(
                    self, user, conversation
                )

                if self.observability_provider and starter_span:
                    starter_span.set_attribute("has_components", components is not None)
                    starter_span.set_attribute(
                        "component_count", len(components) if components else 0
                    )

                if components:
                    # Yield the starter UI components
                    for component in components:
                        yield component

                    # Yield finalization components
                    yield UiComponent(  # type: ignore
                        rich_component=StatusBarUpdateComponent(
                            status="idle",
                            message="Ready",
                            detail="Choose an option or type a message",
                        )
                    )
                    yield UiComponent(  # type: ignore
                        rich_component=ChatInputUpdateComponent(
                            placeholder="Ask a question...", disabled=False
                        )
                    )

                if self.observability_provider and starter_span:
                    await self.observability_provider.end_span(starter_span)
                    if starter_span.duration_ms():
                        await self.observability_provider.record_metric(
                            "agent.workflow_handler.starter_ui.duration",
                            starter_span.duration_ms() or 0,
                            "ms",
                        )

                # Save the conversation if it was newly created
                if self.config.auto_save_conversations:
                    await self.conversation_store.update_conversation(conversation)

                return  # Exit without calling LLM

            except Exception as e:
                logger.error(f"Error generating starter UI: {e}", exc_info=True)
                if self.observability_provider and starter_span:
                    starter_span.set_attribute("error", str(e))
                    await self.observability_provider.end_span(starter_span)
                # Fall through to normal processing on error

        # Don't process actual empty messages (that aren't starter requests)
        if not message.strip():
            return

        # Create observability span for entire message processing
        message_span = None
        if self.observability_provider:
            message_span = await self.observability_provider.create_span(
                "agent.send_message",
                attributes={
                    "user_id": user.id,
                    "conversation_id": conversation_id or "new",
                },
            )

        # Run before_message hooks with observability
        modified_message = message
        for hook in self.lifecycle_hooks:
            hook_span = None
            if self.observability_provider:
                hook_span = await self.observability_provider.create_span(
                    "agent.hook.before_message",
                    attributes={"hook": hook.__class__.__name__},
                )

            hook_result = await hook.before_message(user, modified_message)
            if hook_result is not None:
                modified_message = hook_result

            if self.observability_provider and hook_span:
                hook_span.set_attribute("modified_message", hook_result is not None)
                await self.observability_provider.end_span(hook_span)
                if hook_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.hook.duration",
                        hook_span.duration_ms() or 0,
                        "ms",
                        tags={
                            "hook": hook.__class__.__name__,
                            "phase": "before_message",
                        },
                    )

        # Use the potentially modified message
        message = modified_message

        # Generate conversation ID and request ID if not provided
        if conversation_id is None:
            conversation_id = str(uuid.uuid4())

        request_id = str(uuid.uuid4())

        # Update status to working
        yield UiComponent(  # type: ignore
            rich_component=StatusBarUpdateComponent(
                status="working",
                message="Processing your request...",
                detail="Analyzing query",
            )
        )

        # Load or create conversation with observability (but don't add message yet)
        conversation_span = None
        if self.observability_provider:
            conversation_span = await self.observability_provider.create_span(
                "agent.conversation.load",
                attributes={"conversation_id": conversation_id, "user_id": user.id},
            )

        conversation = await self.conversation_store.get_conversation(
            conversation_id, user
        )

        is_new_conversation = conversation is None

        if not conversation:
            # Create empty conversation (will add message after workflow handler check)
            conversation = Conversation(id=conversation_id, user=user, messages=[])

        if self.observability_provider and conversation_span:
            conversation_span.set_attribute("is_new", is_new_conversation)
            conversation_span.set_attribute("message_count", len(conversation.messages))
            await self.observability_provider.end_span(conversation_span)
            if conversation_span.duration_ms():
                await self.observability_provider.record_metric(
                    "agent.conversation.load.duration",
                    conversation_span.duration_ms() or 0,
                    "ms",
                    tags={"is_new": str(is_new_conversation)},
                )

        # Try workflow handler before adding message to conversation
        if self.workflow_handler:
            trigger_span = None
            if self.observability_provider:
                trigger_span = await self.observability_provider.create_span(
                    "agent.workflow_handler.try_handle",
                    attributes={"user_id": user.id, "conversation_id": conversation_id},
                )

            try:
                workflow_result = await self.workflow_handler.try_handle(
                    self, user, conversation, message
                )

                if self.observability_provider and trigger_span:
                    trigger_span.set_attribute(
                        "should_skip_llm", workflow_result.should_skip_llm
                    )

                if workflow_result.should_skip_llm:
                    # Workflow handled the message, short-circuit LLM

                    # Apply conversation mutation if provided
                    if workflow_result.conversation_mutation:
                        await workflow_result.conversation_mutation(conversation)

                    # Stream components
                    if workflow_result.components:
                        if isinstance(workflow_result.components, list):
                            for component in workflow_result.components:
                                yield component
                        else:
                            # AsyncGenerator
                            async for component in workflow_result.components:
                                yield component

                    # Finalize response (status bar + chat input)
                    yield UiComponent(  # type: ignore
                        rich_component=StatusBarUpdateComponent(
                            status="idle",
                            message="Workflow complete",
                            detail="Ready for next message",
                        )
                    )
                    yield UiComponent(  # type: ignore
                        rich_component=ChatInputUpdateComponent(
                            placeholder="Ask a question...", disabled=False
                        )
                    )

                    # Save conversation if auto-save enabled
                    if self.config.auto_save_conversations:
                        await self.conversation_store.update_conversation(conversation)

                    if self.observability_provider and trigger_span:
                        await self.observability_provider.end_span(trigger_span)

                    # Exit without calling LLM
                    return

            except Exception as e:
                logger.error(f"Error in workflow handler: {e}", exc_info=True)
                if self.observability_provider and trigger_span:
                    trigger_span.set_attribute("error", str(e))
                    await self.observability_provider.end_span(trigger_span)
                # Fall through to normal LLM processing on error

            finally:
                if self.observability_provider and trigger_span:
                    await self.observability_provider.end_span(trigger_span)

        # Persist new conversation to store before adding message
        if is_new_conversation:
            await self.conversation_store.update_conversation(conversation)

        # Not triggered, add user message to conversation now
        conversation.add_message(Message(role="user", content=message))

        # Add initial task
        context_task = Task(
            title="Load conversation context",
            description="Reading message history and user context",
            status="pending",
        )
        yield UiComponent(  # type: ignore
            rich_component=TaskTrackerUpdateComponent.add_task(context_task)
        )

        # Collect available UI features for auditing
        ui_features_available = []
        for feature_name in self.config.ui_features.feature_group_access.keys():
            if self.config.ui_features.can_user_access_feature(feature_name, user):
                ui_features_available.append(feature_name)

        # Create context with observability provider and UI features
        context = ToolContext(
            user=user,
            conversation_id=conversation_id,
            request_id=request_id,
            agent_memory=self.agent_memory,
            observability_provider=self.observability_provider,
            metadata={"ui_features_available": ui_features_available},
        )

        # Enrich context with additional data with observability
        for enricher in self.context_enrichers:
            enrichment_span = None
            if self.observability_provider:
                enrichment_span = await self.observability_provider.create_span(
                    "agent.context.enrichment",
                    attributes={"enricher": enricher.__class__.__name__},
                )

            context = await enricher.enrich_context(context)

            if self.observability_provider and enrichment_span:
                await self.observability_provider.end_span(enrichment_span)
                if enrichment_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.enrichment.duration",
                        enrichment_span.duration_ms() or 0,
                        "ms",
                        tags={"enricher": enricher.__class__.__name__},
                    )

        # Get available tools for user with observability
        schema_span = None
        if self.observability_provider:
            schema_span = await self.observability_provider.create_span(
                "agent.tool_schemas.fetch", attributes={"user_id": user.id}
            )

        tool_schemas = await self.tool_registry.get_schemas(user)

        if self.observability_provider and schema_span:
            schema_span.set_attribute("schema_count", len(tool_schemas))
            await self.observability_provider.end_span(schema_span)
            if schema_span.duration_ms():
                await self.observability_provider.record_metric(
                    "agent.tool_schemas.duration",
                    schema_span.duration_ms() or 0,
                    "ms",
                    tags={"schema_count": str(len(tool_schemas))},
                )

        # Update task status to completed
        yield UiComponent(  # type: ignore
            rich_component=TaskTrackerUpdateComponent.update_task(
                context_task.id, status="completed"
            )
        )

        # Build system prompt with observability
        prompt_span = None
        if self.observability_provider:
            prompt_span = await self.observability_provider.create_span(
                "agent.system_prompt.build",
                attributes={"tool_count": len(tool_schemas)},
            )

        system_prompt = await self.system_prompt_builder.build_system_prompt(
            user, tool_schemas
        )

        # Enhance system prompt with LLM context enhancer
        if self.llm_context_enhancer and system_prompt is not None:
            enhancement_span = None
            if self.observability_provider:
                enhancement_span = await self.observability_provider.create_span(
                    "agent.llm_context.enhance_system_prompt",
                    attributes={
                        "enhancer": self.llm_context_enhancer.__class__.__name__
                    },
                )

            system_prompt = await self.llm_context_enhancer.enhance_system_prompt(
                system_prompt, message, user
            )

            if self.observability_provider and enhancement_span:
                await self.observability_provider.end_span(enhancement_span)
                if enhancement_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.llm_context.enhance_system_prompt.duration",
                        enhancement_span.duration_ms() or 0,
                        "ms",
                        tags={"enhancer": self.llm_context_enhancer.__class__.__name__},
                    )

        if self.observability_provider and prompt_span:
            prompt_span.set_attribute(
                "prompt_length", len(system_prompt) if system_prompt else 0
            )
            await self.observability_provider.end_span(prompt_span)
            if prompt_span.duration_ms():
                await self.observability_provider.record_metric(
                    "agent.system_prompt.duration", prompt_span.duration_ms() or 0, "ms"
                )

        # Build LLM request
        request = await self._build_llm_request(
            conversation, tool_schemas, user, system_prompt
        )

        # Process with tool loop
        tool_iterations = 0

        while tool_iterations < self.config.max_tool_iterations:
            if self.config.include_thinking_indicators and tool_iterations == 0:
                # TODO: Yield thinking indicator
                pass

            # Get LLM response
            if self.config.stream_responses:
                response = await self._handle_streaming_response(request)
            else:
                response = await self._send_llm_request(request)

            # Handle tool calls
            if response.is_tool_call():
                tool_iterations += 1

                # First, add the assistant message with tool_calls to the conversation
                # This is required for OpenAI API - tool messages must follow assistant messages with tool_calls
                assistant_message = Message(
                    role="assistant",
                    content=response.content or "",  # Ensure content is not None
                    tool_calls=response.tool_calls,
                )
                conversation.add_message(assistant_message)

                if response.content is not None:
                    # Yield any partial content from the assistant before tool execution
                    has_tool_invocation_message_in_chat = (
                        self.config.ui_features.can_user_access_feature(
                            UiFeature.UI_FEATURE_SHOW_TOOL_INVOCATION_MESSAGE_IN_CHAT,
                            user,
                        )
                    )
                    if has_tool_invocation_message_in_chat:
                        yield UiComponent(
                            rich_component=RichTextComponent(
                                content=response.content, markdown=True
                            ),
                            simple_component=SimpleTextComponent(text=response.content),
                        )

                        # Update status to executing tools
                        yield UiComponent(  # type: ignore
                            rich_component=StatusBarUpdateComponent(
                                status="working",
                                message="Executing tools...",
                                detail=f"Running {len(response.tool_calls or [])} tools",
                            )
                        )
                    else:
                        # Yield as a status update instead
                        yield UiComponent(  # type: ignore
                            rich_component=StatusBarUpdateComponent(
                                status="working", message=response.content, detail=""
                            )
                        )

                # Collect all tool results first
                tool_results = []
                for i, tool_call in enumerate(response.tool_calls or []):
                    # Add task for this tool execution
                    tool_task = Task(
                        title=f"Execute {tool_call.name}",
                        description=f"Running tool with provided arguments",
                        status="in_progress",
                    )

                    has_tool_names_access = (
                        self.config.ui_features.can_user_access_feature(
                            UiFeature.UI_FEATURE_SHOW_TOOL_NAMES, user
                        )
                    )

                    # Audit UI feature access check
                    if (
                        self.audit_logger
                        and self.config.audit_config.enabled
                        and self.config.audit_config.log_ui_feature_checks
                    ):
                        await self.audit_logger.log_ui_feature_access(
                            user=user,
                            feature_name=UiFeature.UI_FEATURE_SHOW_TOOL_NAMES,
                            access_granted=has_tool_names_access,
                            required_groups=self.config.ui_features.feature_group_access.get(
                                UiFeature.UI_FEATURE_SHOW_TOOL_NAMES, []
                            ),
                            conversation_id=conversation.id,
                            request_id=request_id,
                        )

                    if has_tool_names_access:
                        yield UiComponent(  # type: ignore
                            rich_component=TaskTrackerUpdateComponent.add_task(
                                tool_task
                            )
                        )

                    response_str = response.content

                    # Use primitive StatusCard instead of semantic ToolExecutionComponent
                    tool_status_card = StatusCardComponent(
                        title=f"Executing {tool_call.name}",
                        status="running",
                        description=f"Running tool with {len(tool_call.arguments)} arguments",
                        icon="⚙️",
                        metadata=tool_call.arguments,
                    )

                    has_tool_args_access = (
                        self.config.ui_features.can_user_access_feature(
                            UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS, user
                        )
                    )

                    # Audit UI feature access check
                    if (
                        self.audit_logger
                        and self.config.audit_config.enabled
                        and self.config.audit_config.log_ui_feature_checks
                    ):
                        await self.audit_logger.log_ui_feature_access(
                            user=user,
                            feature_name=UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS,
                            access_granted=has_tool_args_access,
                            required_groups=self.config.ui_features.feature_group_access.get(
                                UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS, []
                            ),
                            conversation_id=conversation.id,
                            request_id=request_id,
                        )

                    if has_tool_args_access:
                        yield UiComponent(
                            rich_component=tool_status_card,
                            simple_component=SimpleTextComponent(
                                text=response_str or ""
                            ),
                        )

                    # Run before_tool hooks with observability
                    tool = await self.tool_registry.get_tool(tool_call.name)
                    if tool:
                        for hook in self.lifecycle_hooks:
                            hook_span = None
                            if self.observability_provider:
                                hook_span = (
                                    await self.observability_provider.create_span(
                                        "agent.hook.before_tool",
                                        attributes={
                                            "hook": hook.__class__.__name__,
                                            "tool": tool_call.name,
                                        },
                                    )
                                )

                            await hook.before_tool(tool, context)

                            if self.observability_provider and hook_span:
                                await self.observability_provider.end_span(hook_span)
                                if hook_span.duration_ms():
                                    await self.observability_provider.record_metric(
                                        "agent.hook.duration",
                                        hook_span.duration_ms() or 0,
                                        "ms",
                                        tags={
                                            "hook": hook.__class__.__name__,
                                            "phase": "before_tool",
                                            "tool": tool_call.name,
                                        },
                                    )

                    # Execute tool with observability
                    tool_exec_span = None
                    if self.observability_provider:
                        tool_exec_span = await self.observability_provider.create_span(
                            "agent.tool.execute",
                            attributes={
                                "tool": tool_call.name,
                                "arg_count": len(tool_call.arguments),
                            },
                        )

                    result = await self.tool_registry.execute(tool_call, context)

                    if self.observability_provider and tool_exec_span:
                        tool_exec_span.set_attribute("success", result.success)
                        if not result.success:
                            tool_exec_span.set_attribute(
                                "error", result.error or "unknown"
                            )
                        await self.observability_provider.end_span(tool_exec_span)
                        if tool_exec_span.duration_ms():
                            await self.observability_provider.record_metric(
                                "agent.tool.duration",
                                tool_exec_span.duration_ms() or 0,
                                "ms",
                                tags={
                                    "tool": tool_call.name,
                                    "success": str(result.success),
                                },
                            )

                    # Run after_tool hooks with observability
                    for hook in self.lifecycle_hooks:
                        hook_span = None
                        if self.observability_provider:
                            hook_span = await self.observability_provider.create_span(
                                "agent.hook.after_tool",
                                attributes={
                                    "hook": hook.__class__.__name__,
                                    "tool": tool_call.name,
                                },
                            )

                        modified_result = await hook.after_tool(result)
                        if modified_result is not None:
                            result = modified_result

                        if self.observability_provider and hook_span:
                            hook_span.set_attribute(
                                "modified_result", modified_result is not None
                            )
                            await self.observability_provider.end_span(hook_span)
                            if hook_span.duration_ms():
                                await self.observability_provider.record_metric(
                                    "agent.hook.duration",
                                    hook_span.duration_ms() or 0,
                                    "ms",
                                    tags={
                                        "hook": hook.__class__.__name__,
                                        "phase": "after_tool",
                                        "tool": tool_call.name,
                                    },
                                )

                    # Update status card to show completion
                    final_status = "success" if result.success else "error"
                    final_description = (
                        f"Tool completed successfully"
                        if result.success
                        else f"Tool failed: {result.error or 'Unknown error'}"
                    )

                    has_tool_args_access_2 = (
                        self.config.ui_features.can_user_access_feature(
                            UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS, user
                        )
                    )

                    # Audit UI feature access check
                    if (
                        self.audit_logger
                        and self.config.audit_config.enabled
                        and self.config.audit_config.log_ui_feature_checks
                    ):
                        await self.audit_logger.log_ui_feature_access(
                            user=user,
                            feature_name=UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS,
                            access_granted=has_tool_args_access_2,
                            required_groups=self.config.ui_features.feature_group_access.get(
                                UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS, []
                            ),
                            conversation_id=conversation.id,
                            request_id=request_id,
                        )

                    if has_tool_args_access_2:
                        yield UiComponent(
                            rich_component=tool_status_card.set_status(
                                final_status, final_description
                            ),
                            simple_component=SimpleTextComponent(
                                text=final_description
                            ),
                        )

                    has_tool_names_access_2 = (
                        self.config.ui_features.can_user_access_feature(
                            UiFeature.UI_FEATURE_SHOW_TOOL_NAMES, user
                        )
                    )

                    # Audit UI feature access check
                    if (
                        self.audit_logger
                        and self.config.audit_config.enabled
                        and self.config.audit_config.log_ui_feature_checks
                    ):
                        await self.audit_logger.log_ui_feature_access(
                            user=user,
                            feature_name=UiFeature.UI_FEATURE_SHOW_TOOL_NAMES,
                            access_granted=has_tool_names_access_2,
                            required_groups=self.config.ui_features.feature_group_access.get(
                                UiFeature.UI_FEATURE_SHOW_TOOL_NAMES, []
                            ),
                            conversation_id=conversation.id,
                            request_id=request_id,
                        )

                    if has_tool_names_access_2:
                        # Update tool task to completed
                        yield UiComponent(  # type: ignore
                            rich_component=TaskTrackerUpdateComponent.update_task(
                                tool_task.id,
                                status="completed",
                                detail=f"Tool {'completed successfully' if result.success else 'return an error'}",
                            )
                        )

                    # Yield tool result
                    if result.ui_component:
                        # For errors, check if user has access to see error details
                        if not result.success:
                            has_tool_error_access = (
                                self.config.ui_features.can_user_access_feature(
                                    UiFeature.UI_FEATURE_SHOW_TOOL_ERROR, user
                                )
                            )

                            # Audit UI feature access check
                            if (
                                self.audit_logger
                                and self.config.audit_config.enabled
                                and self.config.audit_config.log_ui_feature_checks
                            ):
                                await self.audit_logger.log_ui_feature_access(
                                    user=user,
                                    feature_name=UiFeature.UI_FEATURE_SHOW_TOOL_ERROR,
                                    access_granted=has_tool_error_access,
                                    required_groups=self.config.ui_features.feature_group_access.get(
                                        UiFeature.UI_FEATURE_SHOW_TOOL_ERROR, []
                                    ),
                                    conversation_id=conversation.id,
                                    request_id=request_id,
                                )

                            if has_tool_error_access:
                                yield result.ui_component
                        else:
                            # Success results are always shown if they exist
                            yield result.ui_component

                    # Collect tool result data
                    tool_results.append(
                        {
                            "tool_call_id": tool_call.id,
                            "content": (
                                result.result_for_llm
                                if result.success
                                else result.error or "Tool execution failed"
                            ),
                        }
                    )

                # Add tool responses to conversation
                # For APIs that need all tool results in one message, this helps
                for tool_result in tool_results:
                    tool_response_message = Message(
                        role="tool",
                        content=tool_result["content"],
                        tool_call_id=tool_result["tool_call_id"],
                    )
                    conversation.add_message(tool_response_message)

                # Rebuild request with tool responses
                request = await self._build_llm_request(
                    conversation, tool_schemas, user, system_prompt
                )
            else:
                # Update status to idle and set completion message
                yield UiComponent(  # type: ignore
                    rich_component=StatusBarUpdateComponent(
                        status="idle",
                        message="Response complete",
                        detail="Ready for next message",
                    )
                )

                # Update chat input placeholder
                yield UiComponent(  # type: ignore
                    rich_component=ChatInputUpdateComponent(
                        placeholder="Ask a follow-up question...", disabled=False
                    )
                )

                # Yield final text response
                if response.content:
                    # Add assistant response to conversation
                    conversation.add_message(
                        Message(role="assistant", content=response.content)
                    )
                    yield UiComponent(
                        rich_component=RichTextComponent(
                            content=response.content, markdown=True
                        ),
                        simple_component=SimpleTextComponent(text=response.content),
                    )
                break

        # Check if we hit the tool iteration limit
        if tool_iterations >= self.config.max_tool_iterations:
            # The loop exited due to hitting the limit, not due to a natural completion
            logger.warning(
                f"Tool iteration limit reached: {tool_iterations}/{self.config.max_tool_iterations}"
            )

            # Update status bar to show warning
            yield UiComponent(  # type: ignore
                rich_component=StatusBarUpdateComponent(
                    status="warning",
                    message="Tool limit reached",
                    detail=f"Stopped after {tool_iterations} tool executions. The task may be incomplete.",
                )
            )

            # Provide detailed warning message to user
            warning_message = f"""⚠️ **Tool Execution Limit Reached**

The agent stopped after executing {tool_iterations} tools (the configured maximum). The task may not be fully complete.

You can:
- Ask me to continue where I left off
- Adjust the `max_tool_iterations` setting if you need more tool calls
- Break the task into smaller steps"""

            yield UiComponent(
                rich_component=RichTextComponent(
                    content=warning_message, markdown=True
                ),
                simple_component=SimpleTextComponent(
                    text=f"Tool limit reached after {tool_iterations} executions. Task may be incomplete."
                ),
            )

            # Update chat input to suggest follow-up
            yield UiComponent(  # type: ignore
                rich_component=ChatInputUpdateComponent(
                    placeholder="Continue the task or ask me something else...",
                    disabled=False,
                )
            )

        # Save conversation if configured
        if self.config.auto_save_conversations:
            save_span = None
            if self.observability_provider:
                save_span = await self.observability_provider.create_span(
                    "agent.conversation.save",
                    attributes={
                        "conversation_id": conversation_id,
                        "message_count": len(conversation.messages),
                    },
                )

            await self.conversation_store.update_conversation(conversation)

            if self.observability_provider and save_span:
                await self.observability_provider.end_span(save_span)
                if save_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.conversation.save.duration",
                        save_span.duration_ms() or 0,
                        "ms",
                    )

        # Run after_message hooks with observability
        for hook in self.lifecycle_hooks:
            hook_span = None
            if self.observability_provider:
                hook_span = await self.observability_provider.create_span(
                    "agent.hook.after_message",
                    attributes={"hook": hook.__class__.__name__},
                )

            await hook.after_message(conversation)

            if self.observability_provider and hook_span:
                await self.observability_provider.end_span(hook_span)
                if hook_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.hook.duration",
                        hook_span.duration_ms() or 0,
                        "ms",
                        tags={
                            "hook": hook.__class__.__name__,
                            "phase": "after_message",
                        },
                    )

        # End observability span and record metrics
        if self.observability_provider and message_span:
            message_span.set_attribute("tool_iterations", tool_iterations)

            # Track if we hit the tool iteration limit
            hit_tool_limit = tool_iterations >= self.config.max_tool_iterations
            message_span.set_attribute("hit_tool_limit", hit_tool_limit)
            if hit_tool_limit:
                message_span.set_attribute("incomplete_response", True)
                logger.info(
                    f"Tool limit reached - marking response as potentially incomplete"
                )

            await self.observability_provider.end_span(message_span)
            if message_span.duration_ms():
                await self.observability_provider.record_metric(
                    "agent.message.duration",
                    message_span.duration_ms() or 0,
                    "ms",
                    tags={"user_id": user.id, "hit_tool_limit": str(hit_tool_limit)},
                )

    async def get_available_tools(self, user: User) -> List[ToolSchema]:
        """Get tools available to the user."""
        return await self.tool_registry.get_schemas(user)

    async def _build_llm_request(
        self,
        conversation: Conversation,
        tool_schemas: List[ToolSchema],
        user: User,
        system_prompt: Optional[str] = None,
    ) -> LlmRequest:
        """Build LLM request from conversation and tools."""
        # Apply conversation filters with observability
        filtered_messages = conversation.messages
        for filter in self.conversation_filters:
            filter_span = None
            if self.observability_provider:
                filter_span = await self.observability_provider.create_span(
                    "agent.conversation.filter",
                    attributes={
                        "filter": filter.__class__.__name__,
                        "message_count_before": len(filtered_messages),
                    },
                )

            filtered_messages = await filter.filter_messages(filtered_messages)

            if self.observability_provider and filter_span:
                filter_span.set_attribute("message_count_after", len(filtered_messages))
                await self.observability_provider.end_span(filter_span)
                if filter_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.filter.duration",
                        filter_span.duration_ms() or 0,
                        "ms",
                        tags={"filter": filter.__class__.__name__},
                    )

        messages = []
        for msg in filtered_messages:
            llm_msg = LlmMessage(
                role=msg.role,
                content=msg.content,
                tool_calls=msg.tool_calls,
                tool_call_id=msg.tool_call_id,
            )
            messages.append(llm_msg)

        # Enhance messages with LLM context enhancer
        if self.llm_context_enhancer:
            enhancement_span = None
            if self.observability_provider:
                enhancement_span = await self.observability_provider.create_span(
                    "agent.llm_context.enhance_user_messages",
                    attributes={
                        "enhancer": self.llm_context_enhancer.__class__.__name__,
                        "message_count": len(messages),
                    },
                )

            messages = await self.llm_context_enhancer.enhance_user_messages(
                messages, user
            )

            if self.observability_provider and enhancement_span:
                enhancement_span.set_attribute("message_count_after", len(messages))
                await self.observability_provider.end_span(enhancement_span)
                if enhancement_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.llm_context.enhance_user_messages.duration",
                        enhancement_span.duration_ms() or 0,
                        "ms",
                        tags={"enhancer": self.llm_context_enhancer.__class__.__name__},
                    )

        return LlmRequest(
            messages=messages,
            tools=tool_schemas if tool_schemas else None,
            user=user,
            temperature=self.config.temperature,
            max_tokens=self.config.max_tokens,
            stream=self.config.stream_responses,
            system_prompt=system_prompt,
        )

    async def _send_llm_request(self, request: LlmRequest) -> LlmResponse:
        """Send LLM request with middleware and observability."""
        # Apply before_llm_request middlewares with observability
        for middleware in self.llm_middlewares:
            mw_span = None
            if self.observability_provider:
                mw_span = await self.observability_provider.create_span(
                    "agent.middleware.before_llm",
                    attributes={"middleware": middleware.__class__.__name__},
                )

            request = await middleware.before_llm_request(request)

            if self.observability_provider and mw_span:
                await self.observability_provider.end_span(mw_span)
                if mw_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.middleware.duration",
                        mw_span.duration_ms() or 0,
                        "ms",
                        tags={
                            "middleware": middleware.__class__.__name__,
                            "phase": "before_llm",
                        },
                    )

        # Create observability span for LLM call
        llm_span = None
        if self.observability_provider:
            llm_span = await self.observability_provider.create_span(
                "llm.request",
                attributes={
                    "model": getattr(self.llm_service, "model", "unknown"),
                    "stream": request.stream,
                },
            )

        # Send request
        response = await self.llm_service.send_request(request)

        # End span and record metrics
        if self.observability_provider and llm_span:
            await self.observability_provider.end_span(llm_span)
            if llm_span.duration_ms():
                await self.observability_provider.record_metric(
                    "llm.request.duration", llm_span.duration_ms() or 0, "ms"
                )

        # Apply after_llm_response middlewares with observability
        for middleware in self.llm_middlewares:
            mw_span = None
            if self.observability_provider:
                mw_span = await self.observability_provider.create_span(
                    "agent.middleware.after_llm",
                    attributes={"middleware": middleware.__class__.__name__},
                )

            response = await middleware.after_llm_response(request, response)

            if self.observability_provider and mw_span:
                await self.observability_provider.end_span(mw_span)
                if mw_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.middleware.duration",
                        mw_span.duration_ms() or 0,
                        "ms",
                        tags={
                            "middleware": middleware.__class__.__name__,
                            "phase": "after_llm",
                        },
                    )

        return response

    async def _handle_streaming_response(self, request: LlmRequest) -> LlmResponse:
        """Handle streaming response from LLM."""
        # Apply before_llm_request middlewares with observability
        for middleware in self.llm_middlewares:
            mw_span = None
            if self.observability_provider:
                mw_span = await self.observability_provider.create_span(
                    "agent.middleware.before_llm",
                    attributes={
                        "middleware": middleware.__class__.__name__,
                        "stream": True,
                    },
                )

            request = await middleware.before_llm_request(request)

            if self.observability_provider and mw_span:
                await self.observability_provider.end_span(mw_span)
                if mw_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.middleware.duration",
                        mw_span.duration_ms() or 0,
                        "ms",
                        tags={
                            "middleware": middleware.__class__.__name__,
                            "phase": "before_llm",
                            "stream": "true",
                        },
                    )

        accumulated_content = ""
        accumulated_tool_calls = []

        # Create span for streaming
        stream_span = None
        if self.observability_provider:
            stream_span = await self.observability_provider.create_span(
                "llm.stream",
                attributes={"model": getattr(self.llm_service, "model", "unknown")},
            )

        async for chunk in self.llm_service.stream_request(request):
            if chunk.content:
                accumulated_content += chunk.content
                # Could yield intermediate TextChunk here

            if chunk.tool_calls:
                accumulated_tool_calls.extend(chunk.tool_calls)

        # End streaming span
        if self.observability_provider and stream_span:
            stream_span.set_attribute("content_length", len(accumulated_content))
            stream_span.set_attribute("tool_call_count", len(accumulated_tool_calls))
            await self.observability_provider.end_span(stream_span)
            if stream_span.duration_ms():
                await self.observability_provider.record_metric(
                    "llm.stream.duration", stream_span.duration_ms() or 0, "ms"
                )

        response = LlmResponse(
            content=accumulated_content if accumulated_content else None,
            tool_calls=accumulated_tool_calls if accumulated_tool_calls else None,
        )

        # Apply after_llm_response middlewares with observability
        for middleware in self.llm_middlewares:
            mw_span = None
            if self.observability_provider:
                mw_span = await self.observability_provider.create_span(
                    "agent.middleware.after_llm",
                    attributes={
                        "middleware": middleware.__class__.__name__,
                        "stream": True,
                    },
                )

            response = await middleware.after_llm_response(request, response)

            if self.observability_provider and mw_span:
                await self.observability_provider.end_span(mw_span)
                if mw_span.duration_ms():
                    await self.observability_provider.record_metric(
                        "agent.middleware.duration",
                        mw_span.duration_ms() or 0,
                        "ms",
                        tags={
                            "middleware": middleware.__class__.__name__,
                            "phase": "after_llm",
                            "stream": "true",
                        },
                    )

        return response


================================================
FILE: src/vanna/core/agent/config.py
================================================
"""
Agent configuration.

This module contains configuration models that control agent behavior.
"""

from typing import TYPE_CHECKING, Dict, List, Optional

from pydantic import BaseModel, Field

from .._compat import StrEnum

if TYPE_CHECKING:
    from ..user import User


class UiFeature(StrEnum):
    UI_FEATURE_SHOW_TOOL_NAMES = "tool_names"
    UI_FEATURE_SHOW_TOOL_ARGUMENTS = "tool_arguments"
    UI_FEATURE_SHOW_TOOL_ERROR = "tool_error"
    UI_FEATURE_SHOW_TOOL_INVOCATION_MESSAGE_IN_CHAT = "tool_invocation_message_in_chat"
    UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS = "memory_detailed_results"


# Optional: you can also define defaults if you want a shared baseline
DEFAULT_UI_FEATURES: Dict[str, List[str]] = {
    UiFeature.UI_FEATURE_SHOW_TOOL_NAMES: ["admin", "user"],
    UiFeature.UI_FEATURE_SHOW_TOOL_ARGUMENTS: ["admin"],
    UiFeature.UI_FEATURE_SHOW_TOOL_ERROR: ["admin"],
    UiFeature.UI_FEATURE_SHOW_TOOL_INVOCATION_MESSAGE_IN_CHAT: ["admin"],
    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS: ["admin"],
}


class UiFeatures(BaseModel):
    """UI features with group-based access control using the same pattern as tools.

    Each field specifies which groups can access that UI feature.
    Empty list means the feature is accessible to all users.
    Uses the same intersection logic as tool access control.
    """

    # Custom features for extensibility
    feature_group_access: Dict[str, List[str]] = Field(
        default_factory=lambda: DEFAULT_UI_FEATURES.copy(),
        description="Which groups can access UI features",
    )

    def can_user_access_feature(self, feature_name: str, user: "User") -> bool:
        """Check if user can access a UI feature using same logic as tools.

        Args:
            feature_name: Name of the UI feature to check
            user: User object with group_memberships

        Returns:
            True if user has access, False otherwise
        """
        # Then try custom features
        if feature_name in self.feature_group_access:
            allowed_groups = self.feature_group_access[feature_name]
        else:
            # Feature doesn't exist, deny access
            return False

        # Empty list means all users can access (same as tools)
        if not allowed_groups:
            return True

        # Same intersection logic as tool access control
        user_groups = set(user.group_memberships)
        feature_groups = set(allowed_groups)
        return bool(user_groups & feature_groups)

    def register_feature(self, name: str, access_groups: List[str]) -> None:
        """Register a custom UI feature with group access control.

        Args:
            name: Name of the custom feature
            access_groups: List of groups that can access this feature
        """
        self.feature_group_access[name] = access_groups


class AuditConfig(BaseModel):
    """Configuration for audit logging."""

    enabled: bool = Field(default=True, description="Enable audit logging")
    log_tool_access_checks: bool = Field(
        default=True, description="Log tool access permission checks"
    )
    log_tool_invocations: bool = Field(
        default=True, description="Log tool invocations with parameters"
    )
    log_tool_results: bool = Field(
        default=True, description="Log tool execution results"
    )
    log_ui_feature_checks: bool = Field(
        default=False, description="Log UI feature access checks (can be noisy)"
    )
    log_ai_responses: bool = Field(
        default=True, description="Log AI-generated responses"
    )
    include_full_ai_responses: bool = Field(
        default=False,
        description="Include full AI response text in logs (privacy concern)",
    )
    sanitize_tool_parameters: bool = Field(
        default=True, description="Sanitize sensitive parameters (passwords, tokens)"
    )


class AgentConfig(BaseModel):
    """Configuration for agent behavior."""

    max_tool_iterations: int = Field(default=10, gt=0)
    stream_responses: bool = Field(default=True)
    auto_save_conversations: bool = Field(default=True)
    include_thinking_indicators: bool = Field(default=True)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: Optional[int] = Field(default=None, gt=0)
    ui_features: UiFeatures = Field(default_factory=UiFeatures)
    audit_config: AuditConfig = Field(default_factory=AuditConfig)


================================================
FILE: src/vanna/core/audit/__init__.py
================================================
"""
Audit logging for the Vanna Agents framework.

This module provides interfaces and models for audit logging, enabling
tracking of user actions, tool invocations, and access control decisions.
"""

from .base import AuditLogger
from .models import (
    AiResponseEvent,
    AuditEvent,
    AuditEventType,
    ToolAccessCheckEvent,
    ToolInvocationEvent,
    ToolResultEvent,
    UiFeatureAccessCheckEvent,
)

__all__ = [
    "AuditLogger",
    "AuditEvent",
    "AuditEventType",
    "ToolAccessCheckEvent",
    "ToolInvocationEvent",
    "ToolResultEvent",
    "UiFeatureAccessCheckEvent",
    "AiResponseEvent",
]


================================================
FILE: src/vanna/core/audit/base.py
================================================
"""
Base audit logger interface.

Audit loggers enable tracking user actions, tool invocations, and access control
decisions for security, compliance, and debugging.
"""

import hashlib
from abc import ABC, abstractmethod
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional

from .models import (
    AiResponseEvent,
    AuditEvent,
    ToolAccessCheckEvent,
    ToolInvocationEvent,
    ToolResultEvent,
    UiFeatureAccessCheckEvent,
)

if TYPE_CHECKING:
    from ..tool.models import ToolCall, ToolContext, ToolResult
    from ..user.models import User


class AuditLogger(ABC):
    """Abstract base class for audit logging implementations.

    Implementations can:
    - Write to files (JSON, CSV, etc.)
    - Send to databases (Postgres, MongoDB, etc.)
    - Stream to cloud services (CloudWatch, Datadog, etc.)
    - Send to SIEM systems (Splunk, Elastic, etc.)

    Example:
        class PostgresAuditLogger(AuditLogger):
            async def log_event(self, event: AuditEvent) -> None:
                await self.db.execute(
                    "INSERT INTO audit_log (...) VALUES (...)",
                    event.model_dump()
                )

        agent = Agent(
            llm_service=...,
            audit_logger=PostgresAuditLogger(db_pool)
        )
    """

    @abstractmethod
    async def log_event(self, event: AuditEvent) -> None:
        """Log a single audit event.

        Args:
            event: The audit event to log

        Raises:
            Exception: If logging fails critically
        """
        pass

    async def log_tool_access_check(
        self,
        user: "User",
        tool_name: str,
        access_granted: bool,
        required_groups: List[str],
        context: "ToolContext",
        reason: Optional[str] = None,
    ) -> None:
        """Convenience method for logging tool access checks.

        Args:
            user: User attempting to access the tool
            tool_name: Name of the tool being accessed
            access_granted: Whether access was granted
            required_groups: Groups required to access the tool
            context: Tool execution context
            reason: Optional reason for denial
        """
        event = ToolAccessCheckEvent(
            user_id=user.id,
            username=user.username,
            user_email=user.email,
            user_groups=user.group_memberships,
            conversation_id=context.conversation_id,
            request_id=context.request_id,
            tool_name=tool_name,
            access_granted=access_granted,
            required_groups=required_groups,
            reason=reason,
        )
        await self.log_event(event)

    async def log_tool_invocation(
        self,
        user: "User",
        tool_call: "ToolCall",
        ui_features: List[str],
        context: "ToolContext",
        sanitize_parameters: bool = True,
    ) -> None:
        """Convenience method for logging tool invocations.

        Args:
            user: User invoking the tool
            tool_call: Tool call information
            ui_features: List of UI features available to the user
            context: Tool execution context
            sanitize_parameters: Whether to sanitize sensitive parameters
        """
        parameters = tool_call.arguments.copy()
        sanitized = False

        if sanitize_parameters:
            parameters, sanitized = self._sanitize_parameters(parameters)

        event = ToolInvocationEvent(
            user_id=user.id,
            username=user.username,
            user_email=user.email,
            user_groups=user.group_memberships,
            conversation_id=context.conversation_id,
            request_id=context.request_id,
            tool_call_id=tool_call.id,
            tool_name=tool_call.name,
            parameters=parameters,
            parameters_sanitized=sanitized,
            ui_features_available=ui_features,
        )
        await self.log_event(event)

    async def log_tool_result(
        self,
        user: "User",
        tool_call: "ToolCall",
        result: "ToolResult",
        context: "ToolContext",
    ) -> None:
        """Convenience method for logging tool results.

        Args:
            user: User who invoked the tool
            tool_call: Tool call information
            result: Tool execution result
            context: Tool execution context
        """
        event = ToolResultEvent(
            user_id=user.id,
            username=user.username,
            user_email=user.email,
            user_groups=user.group_memberships,
            conversation_id=context.conversation_id,
            request_id=context.request_id,
            tool_call_id=tool_call.id,
            tool_name=tool_call.name,
            success=result.success,
            error=result.error,
            execution_time_ms=result.metadata.get("execution_time_ms", 0.0),
            result_size_bytes=(
                len(result.result_for_llm.encode("utf-8"))
                if result.result_for_llm
                else 0
            ),
            ui_component_type=(
                result.ui_component.__class__.__name__ if result.ui_component else None
            ),
        )
        await self.log_event(event)

    async def log_ui_feature_access(
        self,
        user: "User",
        feature_name: str,
        access_granted: bool,
        required_groups: List[str],
        conversation_id: str,
        request_id: str,
    ) -> None:
        """Convenience method for logging UI feature access checks.

        Args:
            user: User attempting to access the feature
            feature_name: Name of the UI feature
            access_granted: Whether access was granted
            required_groups: Groups required to access the feature
            conversation_id: Conversation identifier
            request_id: Request identifier
        """
        event = UiFeatureAccessCheckEvent(
            user_id=user.id,
            username=user.username,
            user_email=user.email,
            user_groups=user.group_memberships,
            conversation_id=conversation_id,
            request_id=request_id,
            feature_name=feature_name,
            access_granted=access_granted,
            required_groups=required_groups,
        )
        await self.log_event(event)

    async def log_ai_response(
        self,
        user: "User",
        conversation_id: str,
        request_id: str,
        response_text: str,
        tool_calls: List["ToolCall"],
        model_info: Optional[Dict[str, Any]] = None,
        include_full_text: bool = False,
    ) -> None:
        """Convenience method for logging AI responses.

        Args:
            user: User receiving the response
            conversation_id: Conversation identifier
            request_id: Request identifier
            response_text: The AI-generated response text
            tool_calls: List of tool calls in the response
            model_info: Optional model configuration info
            include_full_text: Whether to include full response text
        """
        response_hash = hashlib.sha256(response_text.encode("utf-8")).hexdigest()

        event = AiResponseEvent(
            user_id=user.id,
            username=user.username,
            user_email=user.email,
            user_groups=user.group_memberships,
            conversation_id=conversation_id,
            request_id=request_id,
            response_length_chars=len(response_text),
            response_text=response_text if include_full_text else None,
            response_hash=response_hash,
            model_name=model_info.get("model") if model_info else None,
            temperature=model_info.get("temperature") if model_info else None,
            tool_calls_count=len(tool_calls),
            tool_names=[tc.name for tc in tool_calls],
        )
        await self.log_event(event)

    async def query_events(
        self,
        filters: Optional[Dict[str, Any]] = None,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        limit: int = 100,
    ) -> List[AuditEvent]:
        """Query audit events (optional, for implementations that support it).

        Args:
            filters: Filter criteria (user_id, event_type, etc.)
            start_time: Filter events after this time
            end_time: Filter events before this time
            limit: Maximum number of events to return

        Returns:
            List of matching audit events

        Raises:
            NotImplementedError: If query not supported by implementation
        """
        raise NotImplementedError("Query not supported by this implementation")

    def _sanitize_parameters(
        self, parameters: Dict[str, Any]
    ) -> tuple[Dict[str, Any], bool]:
        """Sanitize sensitive data from parameters.

        Args:
            parameters: Raw parameters dict

        Returns:
            Tuple of (sanitized_parameters, was_sanitized)
        """
        sanitized = parameters.copy()
        was_sanitized = False

        # Common sensitive field patterns
        sensitive_patterns = [
            "password",
            "secret",
            "token",
            "api_key",
            "apikey",
            "credential",
            "auth",
            "private_key",
            "access_key",
        ]

        for key in list(sanitized.keys()):
            key_lower = key.lower()
            if any(pattern in key_lower for pattern in sensitive_patterns):
                sanitized[key] = "[REDACTED]"
                was_sanitized = True

        return sanitized, was_sanitized


================================================
FILE: src/vanna/core/audit/models.py
================================================
"""
Audit event models.

This module contains data models for audit logging events.
"""

import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

from .._compat import StrEnum


class AuditEventType(StrEnum):
    """Types of audit events."""

    # Access control events
    TOOL_ACCESS_CHECK = "tool_access_check"
    UI_FEATURE_ACCESS_CHECK = "ui_feature_access_check"

    # Tool execution events
    TOOL_INVOCATION = "tool_invocation"
    TOOL_RESULT = "tool_result"

    # Conversation events
    MESSAGE_RECEIVED = "message_received"
    AI_RESPONSE_GENERATED = "ai_response_generated"
    CONVERSATION_CREATED = "conversation_created"

    # Security events
    ACCESS_DENIED = "access_denied"
    AUTHENTICATION_ATTEMPT = "authentication_attempt"


class AuditEvent(BaseModel):
    """Base audit event with common fields."""

    event_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    event_type: AuditEventType
    timestamp: datetime = Field(default_factory=datetime.utcnow)

    # User context
    user_id: str
    username: Optional[str] = None
    user_email: Optional[str] = None
    user_groups: List[str] = Field(default_factory=list)

    # Request context
    conversation_id: str
    request_id: str
    remote_addr: Optional[str] = None

    # Event-specific data
    details: Dict[str, Any] = Field(default_factory=dict)

    # Privacy/redaction markers
    contains_pii: bool = False
    redacted_fields: List[str] = Field(default_factory=list)


class ToolAccessCheckEvent(AuditEvent):
    """Audit event for tool access permission checks."""

    event_type: AuditEventType = AuditEventType.TOOL_ACCESS_CHECK
    tool_name: str
    access_granted: bool
    required_groups: List[str] = Field(default_factory=list)
    reason: Optional[str] = None


class ToolInvocationEvent(AuditEvent):
    """Audit event for actual tool executions."""

    event_type: AuditEventType = AuditEventType.TOOL_INVOCATION
    tool_call_id: str
    tool_name: str

    # Parameters with sanitization support
    parameters: Dict[str, Any] = Field(default_factory=dict)
    parameters_sanitized: bool = False

    # UI context at invocation time
    ui_features_available: List[str] = Field(default_factory=list)


class ToolResultEvent(AuditEvent):
    """Audit event for tool execution results."""

    event_type: AuditEventType = AuditEventType.TOOL_RESULT
    tool_call_id: str
    tool_name: str
    success: bool
    error: Optional[str] = None
    execution_time_ms: float = 0.0

    # Result metadata (without full content for size)
    result_size_bytes: Optional[int] = None
    ui_component_type: Optional[str] = None


class UiFeatureAccessCheckEvent(AuditEvent):
    """Audit event for UI feature access checks."""

    event_type: AuditEventType = AuditEventType.UI_FEATURE_ACCESS_CHECK
    feature_name: str
    access_granted: bool
    required_groups: List[str] = Field(default_factory=list)


class AiResponseEvent(AuditEvent):
    """Audit event for AI-generated responses."""

    event_type: AuditEventType = AuditEventType.AI_RESPONSE_GENERATED

    # Response metadata
    response_length_chars: int
    response_length_tokens: Optional[int] = None

    # Full text (optional, configurable)
    response_text: Optional[str] = None
    response_hash: str  # SHA256 for integrity verification

    # Model info
    model_name: Optional[str] = None
    temperature: Optional[float] = None

    # Tool calls in response
    tool_calls_count: int = 0
    tool_names: List[str] = Field(default_factory=list)


================================================
FILE: src/vanna/core/component_manager.py
================================================
"""
Component state management and update protocol for rich components.
"""

import uuid
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Set, Union

from pydantic import BaseModel, Field

from ..components.rich import ComponentLifecycle, RichComponent


class UpdateOperation(str, Enum):
    """Types of component update operations."""

    CREATE = "create"
    UPDATE = "update"
    REPLACE = "replace"
    REMOVE = "remove"
    REORDER = "reorder"
    BULK_UPDATE = "bulk_update"


class Position(BaseModel):
    """Position specification for component placement."""

    index: Optional[int] = None
    anchor_id: Optional[str] = None
    relation: str = "after"  # "before", "after", "inside", "replace"


class ComponentUpdate(BaseModel):
    """Represents a change to the component tree."""

    operation: UpdateOperation
    target_id: str  # Component being affected
    component: Optional[RichComponent] = None  # New/updated component data
    updates: Optional[Dict[str, Any]] = None  # Partial updates for UPDATE operation
    position: Optional[Position] = None  # For positioning operations
    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
    batch_id: Optional[str] = None  # For grouping related updates

    def serialize_for_frontend(self) -> Dict[str, Any]:
        """Return update payload with nested components normalized."""
        payload = self.model_dump()

        # Normalise enum values for the frontend contract.
        payload["operation"] = self.operation.value

        if self.component:
            payload["component"] = self.component.serialize_for_frontend()

        return payload


class ComponentNode(BaseModel):
    """Node in the component tree."""

    component: RichComponent
    children: List["ComponentNode"] = Field(default_factory=list)
    parent_id: Optional[str] = None

    def find_child(self, component_id: str) -> Optional["ComponentNode"]:
        """Find a child node by component ID."""
        for child in self.children:
            if child.component.id == component_id:
                return child
            found = child.find_child(component_id)
            if found:
                return found
        return None

    def remove_child(self, component_id: str) -> bool:
        """Remove a child component by ID."""
        for i, child in enumerate(self.children):
            if child.component.id == component_id:
                self.children.pop(i)
                return True
            if child.remove_child(component_id):
                return True
        return False

    def get_all_ids(self) -> Set[str]:
        """Get all component IDs in this subtree."""
        ids = {self.component.id}
        for child in self.children:
            ids.update(child.get_all_ids())
        return ids


class ComponentTree(BaseModel):
    """Hierarchical structure for managing component layout."""

    root: Optional[ComponentNode] = None
    flat_index: Dict[str, ComponentNode] = Field(default_factory=dict)

    def add_component(
        self, component: RichComponent, position: Optional[Position] = None
    ) -> ComponentUpdate:
        """Add a component to the tree."""
        node = ComponentNode(component=component)
        self.flat_index[component.id] = node

        if self.root is None:
            self.root = node
        else:
            parent_node = self._find_parent(position)
            if parent_node is not None:
                node.parent_id = parent_node.component.id
                parent_node.children.append(node)

        return ComponentUpdate(
            operation=UpdateOperation.CREATE,
            target_id=component.id,
            component=component,
            position=position,
        )

    def update_component(
        self, component_id: str, updates: Dict[str, Any]
    ) -> Optional[ComponentUpdate]:
        """Update a component's properties."""
        node = self.flat_index.get(component_id)
        if not node:
            return None

        # Create updated component
        component_data = node.component.model_dump()
        component_data.update(updates)
        component_data["lifecycle"] = ComponentLifecycle.UPDATE
        component_data["timestamp"] = datetime.utcnow().isoformat()

        updated_component = node.component.__class__(**component_data)
        node.component = updated_component

        return ComponentUpdate(
            operation=UpdateOperation.UPDATE,
            target_id=component_id,
            component=updated_component,
            updates=updates,
        )

    def replace_component(
        self, old_id: str, new_component: RichComponent
    ) -> Optional[ComponentUpdate]:
        """Replace one component with another."""
        old_node = self.flat_index.get(old_id)
        if not old_node:
            return None

        # Update the component in place
        old_node.component = new_component

        # Update index
        del self.flat_index[old_id]
        self.flat_index[new_component.id] = old_node

        return ComponentUpdate(
            operation=UpdateOperation.REPLACE, target_id=old_id, component=new_component
        )

    def remove_component(self, component_id: str) -> Optional[ComponentUpdate]:
        """Remove a component and its children."""
        node = self.flat_index.get(component_id)
        if not node:
            return None

        # Remove from parent
        if self.root and self.root.component.id == component_id:
            self.root = None
        else:
            if self.root:
                self.root.remove_child(component_id)

        # Remove from flat index (including all children)
        removed_ids = node.get_all_ids()
        for removed_id in removed_ids:
            self.flat_index.pop(removed_id, None)

        return ComponentUpdate(operation=UpdateOperation.REMOVE, target_id=component_id)

    def get_component(self, component_id: str) -> Optional[RichComponent]:
        """Get a component by ID."""
        node = self.flat_index.get(component_id)
        return node.component if node else None

    def _find_parent(self, position: Optional[Position]) -> Optional[ComponentNode]:
        """Find the parent node for a new component."""
        if not position or not position.anchor_id:
            return self.root

        anchor_node = self.flat_index.get(position.anchor_id)
        if not anchor_node:
            return self.root

        if position.relation == "inside":
            return anchor_node
        elif position.relation in ["before", "after", "replace"]:
            # Find the parent of the anchor
            if anchor_node.parent_id:
                parent_node = self.flat_index.get(anchor_node.parent_id)
                return parent_node if parent_node else self.root
            else:
                return self.root
        else:
            return self.root


class ComponentManager:
    """Manages component lifecycle and state updates."""

    def __init__(self) -> None:
        self.components: Dict[str, RichComponent] = {}
        self.component_tree = ComponentTree()
        self.update_history: List[ComponentUpdate] = []
        self.active_batch: Optional[str] = None

    def emit(self, component: RichComponent) -> Optional[ComponentUpdate]:
        """Emit a component with smart lifecycle management."""
        if component.id in self.components:
            # Existing component - determine if this is an update or replace
            existing = self.components[component.id]

            if component.lifecycle == ComponentLifecycle.UPDATE:
                # Extract changes
                old_data = existing.model_dump()
                new_data = component.model_dump()
                updates = {k: v for k, v in new_data.items() if old_data.get(k) != v}

                update = self.component_tree.update_component(component.id, updates)
            else:
                # Replace
                update = self.component_tree.replace_component(component.id, component)
        else:
            # New component - always append
            update = self.component_tree.add_component(component, None)

        if update:
            self.components[component.id] = component
            self.update_history.append(update)

            if self.active_batch:
                update.batch_id = self.active_batch

        return update

    def update_component(
        self, component_id: str, **updates: Any
    ) -> Optional[ComponentUpdate]:
        """Update specific fields of an existing component."""
        update = self.component_tree.update_component(component_id, updates)
        if update and update.component:
            self.components[component_id] = update.component
            self.update_history.append(update)

            if self.active_batch:
                update.batch_id = self.active_batch

        return update

    def replace_component(
        self, old_id: str, new_component: RichComponent
    ) -> Optional[ComponentUpdate]:
        """Replace one component with another."""
        update = self.component_tree.replace_component(old_id, new_component)
        if update:
            self.components.pop(old_id, None)
            self.components[new_component.id] = new_component
            self.update_history.append(update)

            if self.active_batch:
                update.batch_id = self.active_batch

        return update

    def remove_component(self, component_id: str) -> Optional[ComponentUpdate]:
        """Remove a component and handle cleanup."""
        update = self.component_tree.remove_component(component_id)
        if update:
            self.components.pop(component_id, None)
            self.update_history.append(update)

            if self.active_batch:
                update.batch_id = self.active_batch

        return update

    def get_component(self, component_id: str) -> Optional[RichComponent]:
        """Get a component by ID."""
        return self.components.get(component_id)

    def get_all_components(self) -> List[RichComponent]:
        """Get all components in the manager."""
        return list(self.components.values())

    def start_batch(self) -> str:
        """Start a batch of related updates."""
        self.active_batch = str(uuid.uuid4())
        return self.active_batch

    def end_batch(self) -> Optional[str]:
        """End the current batch."""
        batch_id = self.active_batch
        self.active_batch = None
        return batch_id

    def get_updates_since(
        self, timestamp: Optional[str] = None
    ) -> List[ComponentUpdate]:
        """Get all updates since a given timestamp."""
        if not timestamp:
            return self.update_history.copy()

        try:
            cutoff = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            return [
                update
                for update in self.update_history
                if datetime.fromisoformat(update.timestamp.replace("Z", "+00:00"))
                > cutoff
            ]
        except ValueError:
            return self.update_history.copy()

    def clear_history(self) -> None:
        """Clear the update history."""
        self.update_history.clear()


================================================
FILE: src/vanna/core/components.py
================================================
"""
UI component base class.

This module defines the UiComponent class which is the return type for tool executions.
It's placed in core/ because it's a fundamental type that tools return, not just a UI concern.
"""

from datetime import datetime
from typing import Any, Optional

from pydantic import BaseModel, Field, model_validator


class UiComponent(BaseModel):
    """Base class for UI components streamed to client.

    This wraps both rich and simple component representations,
    allowing tools to return structured UI updates.

    Note: We use Any for component types to avoid circular dependencies.
    Type validation happens at runtime through validators.
    """

    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
    rich_component: Any = Field(
        ..., description="Rich component for advanced rendering"
    )
    simple_component: Optional[Any] = Field(
        None, description="Simple component for basic rendering"
    )

    @model_validator(mode="after")
    def validate_components(self) -> "UiComponent":
        """Validate that components are the correct types at runtime."""
        # Import from core - clean imports, no circular dependency
        from .rich_component import RichComponent
        from .simple_component import SimpleComponent

        if not isinstance(self.rich_component, RichComponent):
            raise ValueError(
                f"rich_component must be a RichComponent, got {type(self.rich_component)}"
            )

        if self.simple_component is not None and not isinstance(
            self.simple_component, SimpleComponent
        ):
            raise ValueError(
                f"simple_component must be a SimpleComponent or None, got {type(self.simple_component)}"
            )

        return self

    model_config = {"arbitrary_types_allowed": True}


================================================
FILE: src/vanna/core/enhancer/__init__.py
================================================
"""
LLM context enhancement system for adding context to prompts and messages.

This module provides interfaces for enriching LLM system prompts and messages
with additional context before LLM calls (e.g., from memory, RAG, documentation).
"""

from .base import LlmContextEnhancer
from .default import DefaultLlmContextEnhancer

__all__ = ["LlmContextEnhancer", "DefaultLlmContextEnhancer"]


================================================
FILE: src/vanna/core/enhancer/base.py
================================================
"""
LLM context enhancer interface.

LLM context enhancers allow you to add additional context to the system prompt
and user messages before LLM calls.
"""

from abc import ABC
from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from ..user.models import User
    from ..llm.models import LlmMessage


class LlmContextEnhancer(ABC):
    """Enhancer for adding context to LLM prompts and messages.

    Subclass this to create custom enhancers that can:
    - Add relevant context to the system prompt based on the user's initial message
    - Enrich user messages with additional context (e.g., from memory/RAG)
    - Inject relevant examples or documentation
    - Add temporal or environmental context

    Example:
        class MemoryBasedEnhancer(LlmContextEnhancer):
            def __init__(self, agent_memory):
                self.agent_memory = agent_memory

            async def enhance_system_prompt(
                self,
                system_prompt: str,
                user_message: str,
                user: User
            ) -> str:
                # Add relevant examples from memory based on user message
                examples = await self.agent_memory.search_similar(user_message)
                return system_prompt + "\\n\\nRelevant examples:\\n" + examples

            async def enhance_user_messages(
                self,
                messages: list[LlmMessage],
                user: User
            ) -> list[LlmMessage]:
                # Could modify or add to messages
                return messages

        agent = Agent(
            llm_service=...,
            llm_context_enhancer=MemoryBasedEnhancer(agent_memory)
        )
    """

    async def enhance_system_prompt(
        self, system_prompt: str, user_message: str, user: "User"
    ) -> str:
        """Enhance the system prompt with additional context.

        This method is called before the first LLM request with the initial
        user message, allowing you to add relevant context to the system prompt.

        Args:
            system_prompt: The original system prompt
            user_message: The initial user message
            user: The user making the request

        Returns:
            Enhanced system prompt with additional context

        Note:
            This is called once per conversation turn, before any tool calls.
        """
        return system_prompt

    async def enhance_user_messages(
        self, messages: list["LlmMessage"], user: "User"
    ) -> list["LlmMessage"]:
        """Enhance user messages with additional context.

        This method is called to potentially modify or add context to user messages
        before sending them to the LLM.

        Args:
            messages: The list of messages to enhance
            user: The user making the request

        Returns:
            Enhanced list of messages

        Note:
            This is called before each LLM request, including after tool calls.
            Be careful not to add context repeatedly on each iteration.
        """
        return messages


================================================
FILE: src/vanna/core/enhancer/default.py
================================================
"""
Default LLM context enhancer implementation using AgentMemory.

This implementation enriches the system prompt with relevant memories
based on the user's initial message.
"""

from typing import TYPE_CHECKING, List, Optional
from .base import LlmContextEnhancer

if TYPE_CHECKING:
    from ..user.models import User
    from ..llm.models import LlmMessage
    from ...capabilities.agent_memory import AgentMemory, TextMemorySearchResult


class DefaultLlmContextEnhancer(LlmContextEnhancer):
    """Default enhancer that uses AgentMemory to add relevant context.

    This enhancer searches the agent's memory for relevant examples and
    tool use patterns based on the user's message, and adds them to the
    system prompt.

    Example:
        agent = Agent(
            llm_service=...,
            agent_memory=agent_memory,
            llm_context_enhancer=DefaultLlmContextEnhancer(agent_memory)
        )
    """

    def __init__(self, agent_memory: Optional["AgentMemory"] = None):
        """Initialize with optional agent memory.

        Args:
            agent_memory: Optional AgentMemory instance. If not provided,
                         enhancement will be skipped.
        """
        self.agent_memory = agent_memory

    async def enhance_system_prompt(
        self, system_prompt: str, user_message: str, user: "User"
    ) -> str:
        """Enhance system prompt with relevant memories.

        Searches agent memory for relevant text memories based on the
        user's message and adds them to the system prompt.

        Args:
            system_prompt: The original system prompt
            user_message: The initial user message
            user: The user making the request

        Returns:
            Enhanced system prompt with relevant examples from memory
        """
        if not self.agent_memory:
            return system_prompt

        try:
            # Import here to avoid circular dependency
            from ..tool import ToolContext
            import uuid

            # Create a temporary context for memory search
            context = ToolContext(
                user=user,
                conversation_id="temp",
                request_id=str(uuid.uuid4()),
                agent_memory=self.agent_memory,
            )

            # Search for relevant text memories based on user message
            memories: List[
                "TextMemorySearchResult"
            ] = await self.agent_memory.search_text_memories(
                query=user_message, context=context, limit=5
            )

            if not memories:
                return system_prompt

            # Format memories as context snippets to add to system prompt
            examples_section = "\n\n## Relevant Context from Memory\n\n"
            examples_section += "The following domain knowledge and context from prior interactions may be relevant:\n\n"

            for result in memories:
                memory = result.memory
                examples_section += f"• {memory.content}\n"

            # Append examples to system prompt
            return system_prompt + examples_section

        except Exception as e:
            # If memory search fails, return original prompt
            # Don't fail the entire request due to memory issues
            import logging

            logger = logging.getLogger(__name__)
            logger.warning(f"Failed to enhance system prompt with memories: {e}")
            return system_prompt

    async def enhance_user_messages(
        self, messages: list["LlmMessage"], user: "User"
    ) -> list["LlmMessage"]:
        """Enhance user messages.

        The default implementation doesn't modify user messages.
        Override this to add context to user messages if needed.

        Args:
            messages: The list of messages
            user: The user making the request

        Returns:
            Original list of messages (unmodified)
        """
        return messages


================================================
FILE: src/vanna/core/enricher/__init__.py
================================================
"""
Context enrichment system for adding data to tool execution context.

This module provides interfaces for enriching ToolContext with additional
data before tool execution.
"""

from .base import ToolContextEnricher

__all__ = ["ToolContextEnricher"]


================================================
FILE: src/vanna/core/enricher/base.py
================================================
"""
Base context enricher interface.

Context enrichers allow you to add additional data to the ToolContext
before tools are executed.
"""

from abc import ABC
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ..tool.models import ToolContext


class ToolContextEnricher(ABC):
    """Enricher for adding data to ToolContext.

    Subclass this to create custom enrichers that can:
    - Add user preferences from database
    - Inject session state
    - Add temporal context (timezone, current date)
    - Include user history or profile data
    - Add environment-specific configuration

    Example:
        class UserPreferencesEnricher(ToolContextEnricher):
            def __init__(self, db):
                self.db = db

            async def enrich_context(self, context: ToolContext) -> ToolContext:
                # Fetch user preferences
                prefs = await self.db.get_user_preferences(context.user.id)

                # Add to context metadata
                context.metadata["preferences"] = prefs
                context.metadata["timezone"] = prefs.get("timezone", "UTC")

                return context

        agent = AgentRunner(
            llm_service=...,
            context_enrichers=[UserPreferencesEnricher(db), SessionEnricher()]
        )
    """

    async def enrich_context(self, context: "ToolContext") -> "ToolContext":
        """Enrich the tool execution context with additional data.

        Args:
            context: The tool context to enrich

        Returns:
            Enriched context (typically modified in-place)

        Note:
            Enrichers typically modify the context.metadata dict to add
            additional data that tools can access.
        """
        return context


================================================
FILE: src/vanna/core/errors.py
================================================
"""
Exception classes for the Vanna Agents framework.

This module defines all custom exceptions used throughout the framework.
"""


class AgentError(Exception):
    """Base exception for agent framework."""

    pass


class ToolExecutionError(AgentError):
    """Error during tool execution."""

    pass


class ToolNotFoundError(AgentError):
    """Tool not found in registry."""

    pass


class PermissionError(AgentError):
    """User lacks required permissions."""

    pass


class ConversationNotFoundError(AgentError):
    """Conversation not found."""

    pass


class LlmServiceError(AgentError):
    """Error communicating with LLM service."""

    pass


class ValidationError(AgentError):
    """Data validation error."""

    pass


================================================
FILE: src/vanna/core/evaluation/__init__.py
================================================
"""
Evaluation framework for Vanna Agents.

This module provides a complete evaluation system for testing and comparing
agent variants, with special focus on LLM comparison use cases.

Key Features:
- Parallel execution for efficient I/O-bound operations
- Multiple built-in evaluators (trajectory, output, LLM-as-judge, efficiency)
- Rich reporting (HTML, CSV, console)
- Dataset loaders (YAML, JSON)
- Agent variant comparison

Example:
    >>> from vanna.evaluation import (
    ...     EvaluationRunner,
    ...     EvaluationDataset,
    ...     AgentVariant,
    ...     TrajectoryEvaluator,
    ...     OutputEvaluator,
    ... )
    >>>
    >>> # Load test dataset
    >>> dataset = EvaluationDataset.from_yaml("tests/sql_tasks.yaml")
    >>>
    >>> # Create agent variants
    >>> variants = [
    ...     AgentVariant("claude", claude_agent),
    ...     AgentVariant("gpt", gpt_agent),
    ... ]
    >>>
    >>> # Run comparison
    >>> runner = EvaluationRunner(
    ...     evaluators=[TrajectoryEvaluator(), OutputEvaluator()],
    ...     max_concurrency=20
    ... )
    >>> comparison = await runner.compare_agents(variants, dataset.test_cases)
    >>> comparison.print_summary()
"""

from .base import (
    Evaluator,
    TestCase,
    ExpectedOutcome,
    AgentResult,
    EvaluationResult,
    TestCaseResult,
    AgentVariant,
)
from .runner import EvaluationRunner
from .evaluators import (
    TrajectoryEvaluator,
    OutputEvaluator,
    LLMAsJudgeEvaluator,
    EfficiencyEvaluator,
)
from .report import EvaluationReport, ComparisonReport
from .dataset import EvaluationDataset

__all__ = [
    # Base classes
    "Evaluator",
    "TestCase",
    "ExpectedOutcome",
    "AgentResult",
    "EvaluationResult",
    "TestCaseResult",
    "AgentVariant",
    # Runner
    "EvaluationRunner",
    # Built-in evaluators
    "TrajectoryEvaluator",
    "OutputEvaluator",
    "LLMAsJudgeEvaluator",
    "EfficiencyEvaluator",
    # Reporting
    "EvaluationReport",
    "ComparisonReport",
    # Datasets
    "EvaluationDataset",
]


================================================
FILE: src/vanna/core/evaluation/base.py
================================================
"""
Core evaluation abstractions for the Vanna Agents framework.

This module provides the base classes and models for evaluating agent behavior,
including test cases, expected outcomes, and evaluation results.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Callable
from dataclasses import dataclass, field
from datetime import datetime
from pydantic import BaseModel

from vanna.core import User, UiComponent


class ExpectedOutcome(BaseModel):
    """Defines what we expect from the agent for a test case.

    Provides multiple ways to specify expectations:
    - tools_called: List of tool names that should be called
    - tools_not_called: List of tool names that should NOT be called
    - final_answer_contains: Keywords/phrases that should appear in output
    - final_answer_not_contains: Keywords/phrases that should NOT appear
    - min_components: Minimum number of UI components expected
    - max_execution_time_ms: Maximum allowed execution time
    - custom_validators: Custom validation functions
    """

    tools_called: Optional[List[str]] = None
    tools_not_called: Optional[List[str]] = None
    final_answer_contains: Optional[List[str]] = None
    final_answer_not_contains: Optional[List[str]] = None
    min_components: Optional[int] = None
    max_components: Optional[int] = None
    max_execution_time_ms: Optional[float] = None
    metadata: Dict[str, Any] = {}


class TestCase(BaseModel):
    """A single evaluation test case.

    Attributes:
        id: Unique identifier for the test case
        user: User context for the test
        message: The message to send to the agent
        conversation_id: Optional conversation ID for multi-turn tests
        expected_outcome: What we expect the agent to do/produce
        metadata: Additional metadata for categorization/filtering
    """

    id: str
    user: User
    message: str
    conversation_id: Optional[str] = None
    expected_outcome: Optional[ExpectedOutcome] = None
    metadata: Dict[str, Any] = {}


@dataclass
class AgentResult:
    """The result of running an agent on a test case.

    Captures everything that happened during agent execution
    for later evaluation.
    """

    test_case_id: str
    components: List[UiComponent]
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
    llm_requests: List[Dict[str, Any]] = field(default_factory=list)
    execution_time_ms: float = 0.0
    total_tokens: int = 0
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def get_final_answer(self) -> str:
        """Extract the final answer from components."""
        # Find text components and concatenate
        texts = []
        for component in self.components:
            if hasattr(component, "rich_component"):
                rich_comp = component.rich_component
                if hasattr(rich_comp, "type") and rich_comp.type.value == "text":
                    content = rich_comp.data.get("content") or getattr(
                        rich_comp, "content", ""
                    )
                    if content:
                        texts.append(content)
        return "\n".join(texts)

    def get_tool_names_called(self) -> List[str]:
        """Get list of tool names that were called."""
        return [call.get("tool_name", "") for call in self.tool_calls]


class EvaluationResult(BaseModel):
    """Result of evaluating a single test case.

    Attributes:
        test_case_id: ID of the test case evaluated
        evaluator_name: Name of the evaluator that produced this result
        passed: Whether the test case passed
        score: Score from 0.0 to 1.0
        reasoning: Explanation of the evaluation
        metrics: Additional metrics captured during evaluation
        timestamp: When the evaluation was performed
    """

    test_case_id: str
    evaluator_name: str
    passed: bool
    score: float  # 0.0 to 1.0
    reasoning: str
    metrics: Dict[str, Any] = {}
    timestamp: datetime = datetime.now()


@dataclass
class TestCaseResult:
    """Complete result for a single test case including all evaluations."""

    test_case: TestCase
    agent_result: AgentResult
    evaluations: List[EvaluationResult]
    execution_time_ms: float

    def overall_passed(self) -> bool:
        """Check if all evaluations passed."""
        return all(e.passed for e in self.evaluations)

    def overall_score(self) -> float:
        """Calculate average score across all evaluations."""
        if not self.evaluations:
            return 0.0
        return sum(e.score for e in self.evaluations) / len(self.evaluations)


@dataclass
class AgentVariant:
    """A variant of an agent to evaluate (different LLM, config, etc).

    Used for comparing different agent configurations, especially
    different LLMs or model versions.

    Attributes:
        name: Human-readable name for this variant
        agent: The agent instance to evaluate
        metadata: Additional info (model name, provider, config, etc)
    """

    name: str
    agent: Any  # Agent type - avoiding circular import
    metadata: Dict[str, Any] = field(default_factory=dict)


class Evaluator(ABC):
    """Base class for evaluating agent behavior.

    Evaluators examine the agent's execution and determine if it
    met expectations. Multiple evaluators can be composed to check
    different aspects (trajectory, output quality, efficiency, etc).
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Name of this evaluator."""
        pass

    @abstractmethod
    async def evaluate(
        self,
        test_case: TestCase,
        agent_result: AgentResult,
    ) -> EvaluationResult:
        """Evaluate a single test case execution.

        Args:
            test_case: The test case that was executed
            agent_result: The result from running the agent

        Returns:
            EvaluationResult with pass/fail, score, and reasoning
        """
        pass


================================================
FILE: src/vanna/core/evaluation/dataset.py
================================================
"""
Dataset loaders for evaluation test cases.

This module provides utilities for loading test case datasets from
YAML and JSON files.
"""

import json
import yaml
from typing import Any, Dict, List
from pathlib import Path

from .base import TestCase, ExpectedOutcome
from vanna.core import User


class EvaluationDataset:
    """Collection of test cases with metadata.

    Example YAML format:
        dataset:
          name: "SQL Generation Tasks"
          description: "Test cases for SQL generation"
          test_cases:
            - id: "sql_001"
              user_id: "test_user"
              message: "Show me total sales by region"
              expected_outcome:
                tools_called: ["generate_sql", "execute_query"]
                final_answer_contains: ["SELECT", "GROUP BY", "region"]
    """

    def __init__(self, name: str, test_cases: List[TestCase], description: str = ""):
        """Initialize evaluation dataset.

        Args:
            name: Name of the dataset
            test_cases: List of test cases
            description: Optional description
        """
        self.name = name
        self.test_cases = test_cases
        self.description = description

    @classmethod
    def from_yaml(cls, path: str) -> "EvaluationDataset":
        """Load dataset from YAML file.

        Args:
            path: Path to YAML file

        Returns:
            EvaluationDataset instance
        """
        with open(path, "r") as f:
            data = yaml.safe_load(f)

        return cls._from_dict(data)

    @classmethod
    def from_json(cls, path: str) -> "EvaluationDataset":
        """Load dataset from JSON file.

        Args:
            path: Path to JSON file

        Returns:
            EvaluationDataset instance
        """
        with open(path, "r") as f:
            data = json.load(f)

        return cls._from_dict(data)

    @classmethod
    def _from_dict(cls, data: Dict[str, Any]) -> "EvaluationDataset":
        """Create dataset from dictionary.

        Args:
            data: Dictionary with dataset structure

        Returns:
            EvaluationDataset instance
        """
        dataset_config = data.get("dataset", data)
        name = dataset_config.get("name", "Unnamed Dataset")
        description = dataset_config.get("description", "")

        test_cases = []
        for tc_data in dataset_config.get("test_cases", []):
            test_case = cls._parse_test_case(tc_data)
            test_cases.append(test_case)

        return cls(name=name, test_cases=test_cases, description=description)

    @classmethod
    def _parse_test_case(cls, data: Dict[str, Any]) -> TestCase:
        """Parse a single test case from dictionary.

        Args:
            data: Test case dictionary

        Returns:
            TestCase instance
        """
        # Create user
        user_id = data.get("user_id", "test_user")
        user = User(
            id=user_id,
            username=data.get("username", user_id),
            email=data.get("email", f"{user_id}@example.com"),
            group_memberships=data.get("user_groups", []),
        )

        # Parse expected outcome if present
        expected_outcome = None
        if "expected_outcome" in data:
            outcome_data = data["expected_outcome"]
            expected_outcome = ExpectedOutcome(
                tools_called=outcome_data.get("tools_called"),
                tools_not_called=outcome_data.get("tools_not_called"),
                final_answer_contains=outcome_data.get("final_answer_contains"),
                final_answer_not_contains=outcome_data.get("final_answer_not_contains"),
                min_components=outcome_data.get("min_components"),
                max_components=outcome_data.get("max_components"),
                max_execution_time_ms=outcome_data.get("max_execution_time_ms"),
                metadata=outcome_data.get("metadata", {}),
            )

        return TestCase(
            id=data["id"],
            user=user,
            message=data["message"],
            conversation_id=data.get("conversation_id"),
            expected_outcome=expected_outcome,
            metadata=data.get("metadata", {}),
        )

    def save_yaml(self, path: str) -> None:
        """Save dataset to YAML file.

        Args:
            path: Path to save YAML file
        """
        data = self._to_dict()
        with open(path, "w") as f:
            yaml.dump(data, f, default_flow_style=False, sort_keys=False)

    def save_json(self, path: str) -> None:
        """Save dataset to JSON file.

        Args:
            path: Path to save JSON file
        """
        data = self._to_dict()
        with open(path, "w") as f:
            json.dump(data, f, indent=2)

    def _to_dict(self) -> Dict[str, Any]:
        """Convert dataset to dictionary.

        Returns:
            Dictionary representation
        """
        return {
            "dataset": {
                "name": self.name,
                "description": self.description,
                "test_cases": [self._test_case_to_dict(tc) for tc in self.test_cases],
            }
        }

    def _test_case_to_dict(self, test_case: TestCase) -> Dict[str, Any]:
        """Convert test case to dictionary.

        Args:
            test_case: TestCase to convert

        Returns:
            Dictionary representation
        """
        data: Dict[str, Any] = {
            "id": test_case.id,
            "user_id": test_case.user.id,
            "username": test_case.user.username,
            "email": test_case.user.email,
            "user_groups": test_case.user.group_memberships,
            "message": test_case.message,
        }

        if test_case.conversation_id:
            data["conversation_id"] = test_case.conversation_id

        if test_case.expected_outcome:
            outcome = test_case.expected_outcome
            outcome_dict: Dict[str, Any] = {}

            if outcome.tools_called:
                outcome_dict["tools_called"] = outcome.tools_called
            if outcome.tools_not_called:
                outcome_dict["tools_not_called"] = outcome.tools_not_called
            if outcome.final_answer_contains:
                outcome_dict["final_answer_contains"] = outcome.final_answer_contains
            if outcome.final_answer_not_contains:
                outcome_dict["final_answer_not_contains"] = (
                    outcome.final_answer_not_contains
                )
            if outcome.min_components is not None:
                outcome_dict["min_components"] = outcome.min_components
            if outcome.max_components is not None:
                outcome_dict["max_components"] = outcome.max_components
            if outcome.max_execution_time_ms is not None:
                outcome_dict["max_execution_time_ms"] = outcome.max_execution_time_ms
            if outcome.metadata:
                outcome_dict["metadata"] = outcome.metadata

            if outcome_dict:
                data["expected_outcome"] = outcome_dict

        if test_case.metadata:
            data["metadata"] = test_case.metadata

        return data

    def filter_by_metadata(self, **kwargs: Any) -> "EvaluationDataset":
        """Filter test cases by metadata fields.

        Args:
            **kwargs: Metadata fields to match

        Returns:
            New EvaluationDataset with filtered test cases
        """
        filtered = [
            tc
            for tc in self.test_cases
            if all(tc.metadata.get(k) == v for k, v in kwargs.items())
        ]

        return EvaluationDataset(
            name=f"{self.name} (filtered)",
            test_cases=filtered,
            description=f"Filtered from: {self.description}",
        )

    def __len__(self) -> int:
        """Get number of test cases."""
        return len(self.test_cases)

    def __repr__(self) -> str:
        """String representation."""
        return (
            f"EvaluationDataset(name='{self.name}', test_cases={len(self.test_cases)})"
        )


================================================
FILE: src/vanna/core/evaluation/evaluators.py
================================================
"""
Built-in evaluators for common evaluation tasks.

This module provides ready-to-use evaluators for:
- Trajectory evaluation (tools called, order, efficiency)
- Output evaluation (content matching, quality)
- LLM-as-judge evaluation (custom criteria)
- Efficiency evaluation (time, tokens, cost)
"""

from typing import Dict, Any, Optional
from datetime import datetime

from .base import Evaluator, TestCase, AgentResult, EvaluationResult
from vanna.core import LlmService


class TrajectoryEvaluator(Evaluator):
    """Evaluate the path the agent took (tools called, order, etc).

    Checks if the agent called the expected tools and didn't call
    unexpected ones. Useful for verifying agent reasoning and planning.
    """

    @property
    def name(self) -> str:
        return "trajectory"

    async def evaluate(
        self, test_case: TestCase, agent_result: AgentResult
    ) -> EvaluationResult:
        """Evaluate tool call trajectory."""
        if agent_result.error:
            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=False,
                score=0.0,
                reasoning=f"Agent execution failed: {agent_result.error}",
            )

        expected = test_case.expected_outcome
        if not expected:
            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=True,
                score=1.0,
                reasoning="No expected outcome specified, passing by default",
            )

        tools_called = agent_result.get_tool_names_called()
        issues = []
        score = 1.0

        # Check expected tools were called
        if expected.tools_called:
            for expected_tool in expected.tools_called:
                if expected_tool not in tools_called:
                    issues.append(f"Expected tool '{expected_tool}' was not called")
                    score -= 0.5 / len(expected.tools_called)

        # Check unexpected tools were not called
        if expected.tools_not_called:
            for unexpected_tool in expected.tools_not_called:
                if unexpected_tool in tools_called:
                    issues.append(f"Unexpected tool '{unexpected_tool}' was called")
                    score -= 0.5 / len(expected.tools_not_called)

        score = max(0.0, min(1.0, score))
        passed = score >= 0.7  # 70% threshold

        reasoning = "Trajectory evaluation: "
        if issues:
            reasoning += "; ".join(issues)
        else:
            reasoning += "All expected tools called, no unexpected tools"

        return EvaluationResult(
            test_case_id=test_case.id,
            evaluator_name=self.name,
            passed=passed,
            score=score,
            reasoning=reasoning,
            metrics={
                "tools_called": tools_called,
                "num_tools_called": len(tools_called),
                "issues": issues,
            },
        )


class OutputEvaluator(Evaluator):
    """Evaluate the final output quality.

    Checks if the output contains expected content and doesn't
    contain forbidden content. Case-insensitive substring matching.
    """

    @property
    def name(self) -> str:
        return "output"

    async def evaluate(
        self, test_case: TestCase, agent_result: AgentResult
    ) -> EvaluationResult:
        """Evaluate output content."""
        if agent_result.error:
            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=False,
                score=0.0,
                reasoning=f"Agent execution failed: {agent_result.error}",
            )

        expected = test_case.expected_outcome
        if not expected:
            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=True,
                score=1.0,
                reasoning="No expected outcome specified, passing by default",
            )

        final_answer = agent_result.get_final_answer().lower()
        issues = []
        score = 1.0

        # Check expected content is present
        if expected.final_answer_contains:
            for expected_content in expected.final_answer_contains:
                if expected_content.lower() not in final_answer:
                    issues.append(
                        f"Expected content '{expected_content}' not found in output"
                    )
                    score -= 0.5 / len(expected.final_answer_contains)

        # Check forbidden content is absent
        if expected.final_answer_not_contains:
            for forbidden_content in expected.final_answer_not_contains:
                if forbidden_content.lower() in final_answer:
                    issues.append(
                        f"Forbidden content '{forbidden_content}' found in output"
                    )
                    score -= 0.5 / len(expected.final_answer_not_contains)

        score = max(0.0, min(1.0, score))
        passed = score >= 0.7  # 70% threshold

        reasoning = "Output evaluation: "
        if issues:
            reasoning += "; ".join(issues)
        else:
            reasoning += "All expected content present, no forbidden content"

        return EvaluationResult(
            test_case_id=test_case.id,
            evaluator_name=self.name,
            passed=passed,
            score=score,
            reasoning=reasoning,
            metrics={
                "output_length": len(final_answer),
                "issues": issues,
            },
        )


class LLMAsJudgeEvaluator(Evaluator):
    """Use an LLM to judge agent performance based on custom criteria.

    This evaluator uses a separate LLM to assess the quality of the
    agent's output based on natural language criteria.
    """

    def __init__(self, judge_llm: LlmService, criteria: str):
        """Initialize LLM-as-judge evaluator.

        Args:
            judge_llm: The LLM service to use for judging
            criteria: Natural language description of what to evaluate
        """
        self.judge_llm = judge_llm
        self.criteria = criteria

    @property
    def name(self) -> str:
        return "llm_judge"

    async def evaluate(
        self, test_case: TestCase, agent_result: AgentResult
    ) -> EvaluationResult:
        """Evaluate using LLM as judge."""
        if agent_result.error:
            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=False,
                score=0.0,
                reasoning=f"Agent execution failed: {agent_result.error}",
            )

        final_answer = agent_result.get_final_answer()

        # Build prompt for judge
        judge_prompt = f"""You are evaluating an AI agent's response to a user query.

User Query: {test_case.message}

Agent's Response:
{final_answer}

Evaluation Criteria:
{self.criteria}

Please evaluate the response and provide:
1. A score from 0.0 to 1.0 (where 1.0 is perfect)
2. Whether it passes (score >= 0.7)
3. Brief reasoning for your evaluation

Respond in this format:
SCORE: <number>
PASSED: <yes/no>
REASONING: <your explanation>
"""

        try:
            # Call judge LLM
            from vanna.core.llm import LlmRequest, LlmMessage

            request = LlmRequest(
                user=test_case.user,
                messages=[LlmMessage(role="user", content=judge_prompt)],
                temperature=0.0,  # Deterministic judging
            )

            response = await self.judge_llm.send_request(request)
            judgment = response.content or ""

            # Parse response
            score = self._parse_score(judgment)
            passed = self._parse_passed(judgment)
            reasoning = self._parse_reasoning(judgment)

            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=passed,
                score=score,
                reasoning=reasoning,
                metrics={"judge_response": judgment},
            )

        except Exception as e:
            return EvaluationResult(
                test_case_id=test_case.id,
                evaluator_name=self.name,
                passed=False,
                score=0.0,
                reasoning=f"LLM judge evaluation failed: {str(e)}",
            )

    def _parse_score(self, judgment: str) -> float:
        """Parse score from judge response."""
        try:
            for line in judgment.split("\n"):
                if line.startswith("SCORE:"):
                    score_str = line.replace("SCORE:", "").strip()
                    return float(score_str)
        except Exception:
            pass
        return 0.5  # Default if parsing fails

    def _parse_passed(self, judgment: str) -> bool:
        """Parse pass/fail from judge response."""
        for line in judgment.split("\n"):
            if line.startswith("PASSED:"):
                passed_str = line.replace("PASSED:", "").strip().lower()
                return passed_str in ["yes", "true", "pass"]
        return False

    def _parse_reasoning(self, judgment: str) -> str:
        """Parse reasoning from judge response."""
        for line in judgment.split("\n"):
            if line.startswith("REASONING:"):
                return line.replace("REASONING:", "").strip()
        return judgment  # Return full judgment if no reasoning line found


class EfficiencyEvaluator(Evaluator):
    """Evaluate resource usage (time, tokens, cost).

    Checks if the agent completed within acceptable resource limits.
    """

    def __init__(
        self,
        max_execution_time_ms: Optional[float] = None,
        max_tokens: Optional[int] = None,
        max_cost_usd: Optional[float] = None,
    ):
        """Initialize efficiency evaluator.

        Args:
            max_execution_time_ms: Maximum allowed execution time in milliseconds
            max_tokens: Maximum allowed token usage
            max_cost_usd: Maximum allowed cost in USD
        """
        self.max_execution_time_ms = max_execution_time_ms
        self.max_tokens = max_tokens
        self.max_cost_usd = max_cost_usd

    @property
    def name(self) -> str:
        return "efficiency"

    async def evaluate(
        self, test_case: TestCase, agent_result: AgentResult
    ) -> EvaluationResult:
        """Evaluate resource efficiency."""
        issues = []
        score = 1.0

        # Check execution time
        if self.max_execution_time_ms:
            if agent_result.execution_time_ms > self.max_execution_time_ms:
                issues.append(
                    f"Execution time {agent_result.execution_time_ms:.0f}ms "
                    f"exceeded limit {self.max_execution_time_ms:.0f}ms"
                )
                score -= 0.33

        # Check token usage
        if self.max_tokens:
            if agent_result.total_tokens > self.max_tokens:
                issues.append(
                    f"Token usage {agent_result.total_tokens} exceeded limit {self.max_tokens}"
                )
                score -= 0.33

        # Check cost (would need cost calculation from metadata)
        # For now, skip cost evaluation

        # Check from expected outcome if specified
        expected = test_case.expected_outcome
        if expected and expected.max_execution_time_ms:
            if agent_result.execution_time_ms > expected.max_execution_time_ms:
                issues.append(
                    f"Execution time {agent_result.execution_time_ms:.0f}ms "
                    f"exceeded test case limit {expected.max_execution_time_ms:.0f}ms"
                )
                score -= 0.34

        score = max(0.0, min(1.0, score))
        passed = score >= 0.7

        reasoning = "Efficiency evaluation: "
        if issues:
            reasoning += "; ".join(issues)
        else:
            reasoning += "Within resource limits"

        return EvaluationResult(
            test_case_id=test_case.id,
            evaluator_name=self.name,
            passed=passed,
            score=score,
            reasoning=reasoning,
            metrics={
                "execution_time_ms": agent_result.execution_time_ms,
                "total_tokens": agent_result.total_tokens,
                "issues": issues,
            },
        )


================================================
FILE: src/vanna/core/evaluation/report.py
================================================
"""
Evaluation reporting with HTML, CSV, and console output.

This module provides classes for generating evaluation reports,
including comparison reports for evaluating multiple agent variants.
"""

import csv
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime

from .base import TestCaseResult, AgentVariant, Evaluator, TestCase


@dataclass
class EvaluationReport:
    """Report for a single agent's evaluation results.

    Attributes:
        agent_name: Name of the agent evaluated
        results: List of results for each test case
        evaluators: List of evaluators used
        metadata: Additional metadata about the agent/run
        timestamp: When the evaluation was run
    """

    agent_name: str
    results: List[TestCaseResult]
    evaluators: List[Evaluator]
    metadata: Dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=datetime.now)

    def pass_rate(self) -> float:
        """Calculate overall pass rate (0.0 to 1.0)."""
        if not self.results:
            return 0.0
        passed = sum(1 for r in self.results if r.overall_passed())
        return passed / len(self.results)

    def average_score(self) -> float:
        """Calculate average score across all test cases."""
        if not self.results:
            return 0.0
        return sum(r.overall_score() for r in self.results) / len(self.results)

    def average_time(self) -> float:
        """Calculate average execution time in milliseconds."""
        if not self.results:
            return 0.0
        return sum(r.execution_time_ms for r in self.results) / len(self.results)

    def total_tokens(self) -> int:
        """Calculate total tokens used across all test cases."""
        return sum(r.agent_result.total_tokens for r in self.results)

    def get_failures(self) -> List[TestCaseResult]:
        """Get all failed test cases."""
        return [r for r in self.results if not r.overall_passed()]

    def print_summary(self) -> None:
        """Print summary to console."""
        print(f"\n{'=' * 80}")
        print(f"EVALUATION REPORT: {self.agent_name}")
        print(f"{'=' * 80}")
        print(f"Timestamp: {self.timestamp.isoformat()}")
        print(f"Test Cases: {len(self.results)}")
        print(f"Pass Rate: {self.pass_rate():.1%}")
        print(f"Average Score: {self.average_score():.2f}")
        print(f"Average Time: {self.average_time():.0f}ms")
        print(f"Total Tokens: {self.total_tokens()}")
        print(f"{'=' * 80}\n")

        failures = self.get_failures()
        if failures:
            print(f"FAILURES ({len(failures)}):")
            for result in failures:
                print(f"\n  Test Case: {result.test_case.id}")
                print(f"  Message: {result.test_case.message}")
                print(f"  Score: {result.overall_score():.2f}")
                for eval_result in result.evaluations:
                    if not eval_result.passed:
                        print(
                            f"    [{eval_result.evaluator_name}] {eval_result.reasoning}"
                        )


@dataclass
class ComparisonReport:
    """Report comparing multiple agent variants.

    This is the primary report type for LLM comparison use cases.

    Attributes:
        variants: List of agent variants compared
        reports: Dict mapping variant name to EvaluationReport
        test_cases: Test cases used for comparison
        timestamp: When the comparison was run
    """

    variants: List[AgentVariant]
    reports: Dict[str, EvaluationReport]
    test_cases: List[TestCase]
    timestamp: datetime = field(default_factory=datetime.now)

    def print_summary(self) -> None:
        """Print comparison summary to console."""
        print("\n" + "=" * 80)
        print("AGENT COMPARISON SUMMARY")
        print("=" * 80)
        print(f"Timestamp: {self.timestamp.isoformat()}")
        print(f"Variants: {len(self.variants)}")
        print(f"Test Cases: {len(self.test_cases)}")

        # Table of results
        print(
            f"\n{'Agent':<25} {'Pass Rate':<12} {'Avg Score':<12} {'Avg Time':<12} {'Tokens':<12}"
        )
        print("-" * 80)

        for variant_name, report in self.reports.items():
            print(
                f"{variant_name:<25} "
                f"{report.pass_rate():<12.1%} "
                f"{report.average_score():<12.2f} "
                f"{report.average_time():<12.0f} "
                f"{report.total_tokens():<12,}"
            )

        print("=" * 80 + "\n")

    def get_best_variant(self, metric: str = "score") -> str:
        """Get the best performing variant by metric.

        Args:
            metric: Metric to optimize ('score', 'speed', 'pass_rate')

        Returns:
            Name of the best variant
        """
        if metric == "score":
            return max(self.reports.items(), key=lambda x: x[1].average_score())[0]
        elif metric == "speed":
            return min(self.reports.items(), key=lambda x: x[1].average_time())[0]
        elif metric == "pass_rate":
            return max(self.reports.items(), key=lambda x: x[1].pass_rate())[0]
        else:
            raise ValueError(f"Unknown metric: {metric}")

    def save_csv(self, path: str) -> None:
        """Save detailed CSV for further analysis.

        Each row represents one test case × one variant combination.
        """
        with open(path, "w", newline="") as f:
            writer = csv.writer(f)

            # Header
            writer.writerow(
                [
                    "variant",
                    "test_case_id",
                    "test_message",
                    "passed",
                    "score",
                    "execution_time_ms",
                    "tokens",
                    "error",
                    "evaluator_scores",
                ]
            )

            # Data rows
            for variant_name, report in self.reports.items():
                for result in report.results:
                    evaluator_scores = {
                        e.evaluator_name: e.score for e in result.evaluations
                    }

                    writer.writerow(
                        [
                            variant_name,
                            result.test_case.id,
                            result.test_case.message[:50],  # Truncate
                            result.overall_passed(),
                            result.overall_score(),
                            result.execution_time_ms,
                            result.agent_result.total_tokens,
                            result.agent_result.error or "",
                            str(evaluator_scores),
                        ]
                    )

    def save_html(self, path: str) -> None:
        """Save interactive HTML comparison report.

        Generates a rich HTML report with:
        - Summary statistics
        - Charts comparing variants
        - Side-by-side test case results
        """
        html = self._generate_html()
        with open(path, "w") as f:
            f.write(html)

    def _generate_html(self) -> str:
        """Generate HTML content for report."""
        # Build HTML report
        html_parts = [
            "<!DOCTYPE html>",
            "<html>",
            "<head>",
            "<title>Agent Comparison Report</title>",
            "<style>",
            "body { font-family: Arial, sans-serif; margin: 20px; }",
            "h1 { color: #333; }",
            "table { border-collapse: collapse; width: 100%; margin: 20px 0; }",
            "th, td { border: 1px solid #ddd; padding: 12px; text-align: left; }",
            "th { background-color: #4CAF50; color: white; }",
            "tr:nth-child(even) { background-color: #f2f2f2; }",
            ".passed { color: green; font-weight: bold; }",
            ".failed { color: red; font-weight: bold; }",
            ".best { background-color: #d4edda !important; }",
            "</style>",
            "</head>",
            "<body>",
            f"<h1>Agent Comparison Report</h1>",
            f"<p>Generated: {self.timestamp.isoformat()}</p>",
            f"<p>Variants: {len(self.variants)} | Test Cases: {len(self.test_cases)}</p>",
        ]

        # Summary table
        html_parts.append("<h2>Summary</h2>")
        html_parts.append("<table>")
        html_parts.append(
            "<tr><th>Agent</th><th>Pass Rate</th><th>Avg Score</th><th>Avg Time (ms)</th><th>Total Tokens</th></tr>"
        )

        best_by_score = self.get_best_variant("score")

        for variant_name, report in self.reports.items():
            row_class = "best" if variant_name == best_by_score else ""
            html_parts.append(
                f"<tr class='{row_class}'>"
                f"<td>{variant_name}</td>"
                f"<td>{report.pass_rate():.1%}</td>"
                f"<td>{report.average_score():.2f}</td>"
                f"<td>{report.average_time():.0f}</td>"
                f"<td>{report.total_tokens():,}</td>"
                f"</tr>"
            )

        html_parts.append("</table>")

        # Test case details
        html_parts.append("<h2>Test Case Details</h2>")

        for i, test_case in enumerate(self.test_cases):
            html_parts.append(f"<h3>Test Case {i + 1}: {test_case.id}</h3>")
            html_parts.append(f"<p><strong>Message:</strong> {test_case.message}</p>")

            html_parts.append("<table>")
            html_parts.append(
                "<tr><th>Variant</th><th>Result</th><th>Score</th><th>Time (ms)</th></tr>"
            )

            for variant_name, report in self.reports.items():
                result = next(
                    (r for r in report.results if r.test_case.id == test_case.id), None
                )
                if result:
                    passed_class = "passed" if result.overall_passed() else "failed"
                    passed_text = "PASS" if result.overall_passed() else "FAIL"

                    html_parts.append(
                        f"<tr>"
                        f"<td>{variant_name}</td>"
                        f"<td class='{passed_class}'>{passed_text}</td>"
                        f"<td>{result.overall_score():.2f}</td>"
                        f"<td>{result.execution_time_ms:.0f}</td>"
                        f"</tr>"
                    )

            html_parts.append("</table>")

        html_parts.append("</body>")
        html_parts.append("</html>")

        return "\n".join(html_parts)


================================================
FILE: src/vanna/core/evaluation/runner.py
================================================
"""
Evaluation runner with parallel execution support.

This module provides the EvaluationRunner class that executes test cases
against agents with configurable parallelism for efficient evaluation,
especially when comparing multiple LLMs or model versions.
"""

import asyncio
from typing import Any, List, Dict, Optional, AsyncGenerator, TYPE_CHECKING
from datetime import datetime

from .base import (
    TestCase,
    AgentResult,
    TestCaseResult,
    AgentVariant,
    Evaluator,
)
from vanna.core import UiComponent
from vanna.core.user.request_context import RequestContext
from vanna.core.observability import ObservabilityProvider

if TYPE_CHECKING:
    from vanna import Agent
    from .report import EvaluationReport, ComparisonReport


class EvaluationRunner:
    """Run evaluations with parallel execution support.

    The primary use case is comparing multiple agent variants (e.g., different LLMs)
    on the same set of test cases. The runner executes test cases in parallel with
    configurable concurrency to handle I/O-bound LLM operations efficiently.

    Example:
        >>> runner = EvaluationRunner(
        ...     evaluators=[TrajectoryEvaluator(), OutputEvaluator()],
        ...     max_concurrency=20
        ... )
        >>> comparison = await runner.compare_agents(
        ...     agent_variants=[claude_variant, gpt_variant],
        ...     test_cases=dataset.test_cases
        ... )
    """

    def __init__(
        self,
        evaluators: List[Evaluator],
        max_concurrency: int = 10,
        observability_provider: Optional[ObservabilityProvider] = None,
    ):
        """Initialize the evaluation runner.

        Args:
            evaluators: List of evaluators to apply to each test case
            max_concurrency: Maximum number of concurrent test case executions
            observability_provider: Optional observability for tracking eval runs
        """
        self.evaluators = evaluators
        self.max_concurrency = max_concurrency
        self.observability = observability_provider
        self._semaphore = asyncio.Semaphore(max_concurrency)

    async def run_evaluation(
        self,
        agent: "Agent",
        test_cases: List[TestCase],
    ) -> "EvaluationReport":
        """Run evaluation on a single agent.

        Args:
            agent: The agent to evaluate
            test_cases: List of test cases to run

        Returns:
            EvaluationReport with results for all test cases
        """
        from .report import EvaluationReport

        results = await self._run_test_cases_parallel(agent, test_cases)
        return EvaluationReport(
            agent_name="agent",
            results=results,
            evaluators=self.evaluators,
            timestamp=datetime.now(),
        )

    async def compare_agents(
        self,
        agent_variants: List[AgentVariant],
        test_cases: List[TestCase],
    ) -> "ComparisonReport":
        """Compare multiple agent variants on same test cases.

        This is the PRIMARY use case for LLM comparison. Runs all variants
        in parallel for maximum efficiency with I/O-bound LLM calls.

        Args:
            agent_variants: List of agent variants to compare
            test_cases: Test cases to run on each variant

        Returns:
            ComparisonReport with results for all variants
        """
        from .report import ComparisonReport

        # Create span for overall comparison
        if self.observability:
            span = await self.observability.create_span(
                "agent_comparison",
                attributes={
                    "num_variants": len(agent_variants),
                    "num_test_cases": len(test_cases),
                },
            )

        # Run all variants in parallel
        tasks = [
            self._run_agent_variant(variant, test_cases) for variant in agent_variants
        ]

        variant_reports = await asyncio.gather(*tasks)

        if self.observability:
            await self.observability.end_span(span)

        return ComparisonReport(
            variants=agent_variants,
            reports=dict(zip([v.name for v in agent_variants], variant_reports)),
            test_cases=test_cases,
            timestamp=datetime.now(),
        )

    async def compare_agents_streaming(
        self,
        agent_variants: List[AgentVariant],
        test_cases: List[TestCase],
    ) -> AsyncGenerator[tuple[str, TestCaseResult, int, int], None]:
        """Stream comparison results as they complete.

        Useful for long-running evaluations where you want to see
        progress updates in real-time (e.g., for UI display).

        Args:
            agent_variants: Agent variants to compare
            test_cases: Test cases to run

        Yields:
            Tuples of (variant_name, result, completed_count, total_count)
        """
        queue: asyncio.Queue[tuple[str, TestCaseResult]] = asyncio.Queue()

        async def worker(variant: AgentVariant) -> None:
            """Worker that runs test cases for one variant."""
            results = await self._run_test_cases_parallel(variant.agent, test_cases)
            for result in results:
                await queue.put((variant.name, result))

        # Start all workers
        workers = [asyncio.create_task(worker(v)) for v in agent_variants]

        # Yield results as they arrive
        completed = 0
        total = len(agent_variants) * len(test_cases)

        while completed < total:
            variant_name, result = await queue.get()
            completed += 1
            yield variant_name, result, completed, total

        # Wait for all workers to complete
        await asyncio.gather(*workers)

    async def _run_agent_variant(
        self,
        variant: AgentVariant,
        test_cases: List[TestCase],
    ) -> "EvaluationReport":
        """Run a single agent variant on all test cases.

        Args:
            variant: The agent variant to evaluate
            test_cases: Test cases to run

        Returns:
            EvaluationReport for this variant
        """
        from .report import EvaluationReport

        if self.observability:
            span = await self.observability.create_span(
                f"variant_{variant.name}",
                attributes={
                    "variant": variant.name,
                    "num_test_cases": len(test_cases),
                    **variant.metadata,
                },
            )

        results = await self._run_test_cases_parallel(variant.agent, test_cases)

        if self.observability:
            await self.observability.end_span(span)

        return EvaluationReport(
            agent_name=variant.name,
            results=results,
            evaluators=self.evaluators,
            metadata=variant.metadata,
            timestamp=datetime.now(),
        )

    async def _run_test_cases_parallel(
        self,
        agent: "Agent",
        test_cases: List[TestCase],
    ) -> List[TestCaseResult]:
        """Run test cases in parallel with concurrency limit.

        Args:
            agent: The agent to run test cases on
            test_cases: Test cases to execute

        Returns:
            List of TestCaseResult, one per test case
        """
        tasks = [
            self._run_single_test_case(agent, test_case) for test_case in test_cases
        ]

        return await asyncio.gather(*tasks)

    async def _run_single_test_case(
        self,
        agent: "Agent",
        test_case: TestCase,
    ) -> TestCaseResult:
        """Run a single test case with semaphore to limit concurrency.

        Args:
            agent: The agent to execute
            test_case: The test case to run

        Returns:
            TestCaseResult with agent execution and evaluations
        """
        async with self._semaphore:
            # Execute agent
            start_time = asyncio.get_event_loop().time()
            agent_result = await self._execute_agent(agent, test_case)
            execution_time = asyncio.get_event_loop().time() - start_time

            # Run evaluators
            eval_results = []
            for evaluator in self.evaluators:
                eval_result = await evaluator.evaluate(test_case, agent_result)
                eval_results.append(eval_result)

            return TestCaseResult(
                test_case=test_case,
                agent_result=agent_result,
                evaluations=eval_results,
                execution_time_ms=execution_time * 1000,
            )

    async def _execute_agent(
        self,
        agent: "Agent",
        test_case: TestCase,
    ) -> AgentResult:
        """Execute agent and capture full trajectory.

        Args:
            agent: The agent to execute
            test_case: The test case to run

        Returns:
            AgentResult with all captured data
        """
        components: List[UiComponent] = []
        tool_calls: List[Dict[str, Any]] = []
        error: Optional[str] = None

        try:
            # Create request context with user info from test case
            # This allows the agent's UserResolver to resolve the correct user
            request_context = RequestContext(
                cookies={"user_id": test_case.user.id},
                headers={},
                metadata={"test_case_user": test_case.user},
            )

            async for component in agent.send_message(
                request_context=request_context,
                message=test_case.message,
                conversation_id=test_case.conversation_id,
            ):
                components.append(component)

        except Exception as e:
            error = str(e)

        # TODO: Extract tool calls and LLM requests from observability
        # For now, these will be empty unless we hook into observability

        return AgentResult(
            test_case_id=test_case.id,
            components=components,
            tool_calls=tool_calls,
            llm_requests=[],
            error=error,
        )


================================================
FILE: src/vanna/core/filter/__init__.py
================================================
"""
Conversation filtering system for managing conversation history.

This module provides interfaces for filtering and transforming conversation
history before it's sent to the LLM.
"""

from .base import ConversationFilter

__all__ = ["ConversationFilter"]


================================================
FILE: src/vanna/core/filter/base.py
================================================
"""
Base conversation filter interface.

Conversation filters allow you to transform conversation history before
it's sent to the LLM for processing.
"""

from abc import ABC
from typing import TYPE_CHECKING, List

if TYPE_CHECKING:
    from ..storage import Message


class ConversationFilter(ABC):
    """Filter for transforming conversation history.

    Subclass this to create custom filters that can:
    - Remove sensitive information
    - Summarize long conversations
    - Manage context window limits
    - Deduplicate similar messages
    - Prioritize recent or relevant messages

    Example:
        class ContextWindowFilter(ConversationFilter):
            def __init__(self, max_tokens: int = 8000):
                self.max_tokens = max_tokens

            async def filter_messages(self, messages: List[Message]) -> List[Message]:
                # Estimate tokens (rough approximation)
                total_tokens = 0
                filtered = []

                # Keep system message and recent messages
                for msg in reversed(messages):
                    msg_tokens = len(msg.content or "") // 4
                    if total_tokens + msg_tokens > self.max_tokens:
                        break
                    filtered.insert(0, msg)
                    total_tokens += msg_tokens

                return filtered

        agent = AgentRunner(
            llm_service=...,
            conversation_filters=[
                SensitiveDataFilter(),
                ContextWindowFilter(max_tokens=8000)
            ]
        )
    """

    async def filter_messages(self, messages: List["Message"]) -> List["Message"]:
        """Filter and transform conversation messages.

        Args:
            messages: List of conversation messages

        Returns:
            Filtered/transformed list of messages

        Note:
            Filters are applied in order, so messages passed to later
            filters may already be modified by earlier filters.
        """
        return messages


================================================
FILE: src/vanna/core/lifecycle/__init__.py
================================================
"""
Lifecycle hook system for agent execution.

This module provides hooks for intercepting and modifying agent behavior
at various points in the execution lifecycle.
"""

from .base import LifecycleHook

__all__ = ["LifecycleHook"]


================================================
FILE: src/vanna/core/lifecycle/base.py
================================================
"""
Base lifecycle hook interface.

Lifecycle hooks allow you to intercept and customize agent behavior
at key points in the execution flow.
"""

from abc import ABC
from typing import TYPE_CHECKING, Any, Optional

if TYPE_CHECKING:
    from ..user.models import User
    from ..tool import Tool
    from ..tool.models import ToolContext, ToolResult


class LifecycleHook(ABC):
    """Hook into agent execution lifecycle.

    Subclass this to create custom hooks that can:
    - Modify messages before processing
    - Add logging or telemetry
    - Enforce quotas or rate limits
    - Transform tool results
    - Add custom validation

    Example:
        class LoggingHook(LifecycleHook):
            async def before_message(self, user: User, message: str) -> Optional[str]:
                print(f"User {user.username} sent: {message}")
                return None  # Don't modify

        agent = AgentRunner(
            llm_service=...,
            lifecycle_hooks=[LoggingHook(), QuotaCheckHook()]
        )
    """

    async def before_message(self, user: "User", message: str) -> Optional[str]:
        """Called before processing a user message.

        Args:
            user: User sending the message
            message: Original message content

        Returns:
            Modified message string, or None to keep original

        Raises:
            AgentError: To halt message processing (e.g., quota exceeded)
        """
        return None

    async def after_message(self, result: Any) -> None:
        """Called after message has been fully processed.

        Args:
            result: Final result from message processing
        """
        pass

    async def before_tool(self, tool: "Tool[Any]", context: "ToolContext") -> None:
        """Called before tool execution.

        Args:
            tool: Tool about to be executed
            context: Tool execution context

        Raises:
            AgentError: To prevent tool execution
        """
        pass

    async def after_tool(self, result: "ToolResult") -> Optional["ToolResult"]:
        """Called after tool execution.

        Args:
            result: Result from tool execution

        Returns:
            Modified ToolResult, or None to keep original
        """
        return None


================================================
FILE: src/vanna/core/llm/__init__.py
================================================
"""
LLM domain.

This module provides the core abstractions for LLM services in the Vanna Agents framework.
"""

from .base import LlmService
from .models import LlmMessage, LlmRequest, LlmResponse, LlmStreamChunk

__all__ = [
    "LlmService",
    "LlmMessage",
    "LlmRequest",
    "LlmResponse",
    "LlmStreamChunk",
]


================================================
FILE: src/vanna/core/llm/base.py
================================================
"""
LLM domain interface.

This module contains the abstract base class for LLM services.
"""

from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator, List

from .models import LlmRequest, LlmResponse, LlmStreamChunk


class LlmService(ABC):
    """Service for LLM communication."""

    @abstractmethod
    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a request to the LLM."""
        pass

    @abstractmethod
    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Stream a request to the LLM.

        Args:
            request: The LLM request to stream

        Yields:
            LlmStreamChunk instances as they arrive
        """
        # This is an async generator method
        raise NotImplementedError
        yield  # pragma: no cover - makes this an async generator

    @abstractmethod
    async def validate_tools(self, tools: List[Any]) -> List[str]:
        """Validate tool schemas and return any errors."""
        pass


================================================
FILE: src/vanna/core/llm/models.py
================================================
"""
LLM domain models.

This module contains data models for LLM communication.
"""

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

from ..tool.models import ToolCall
from ..user.models import User


class LlmMessage(BaseModel):
    """Message format for LLM communication."""

    role: str = Field(description="Message role")
    content: str = Field(description="Message content")
    tool_calls: Optional[List[ToolCall]] = Field(default=None)
    tool_call_id: Optional[str] = Field(default=None)


class LlmRequest(BaseModel):
    """Request to LLM service."""

    messages: List[LlmMessage] = Field(description="Messages to send")
    tools: Optional[List[Any]] = Field(
        default=None, description="Available tools"
    )  # Will be ToolSchema but avoiding circular import
    user: User = Field(description="User making the request")
    stream: bool = Field(default=False, description="Whether to stream response")
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: Optional[int] = Field(default=None, gt=0)
    system_prompt: Optional[str] = Field(
        default=None, description="System prompt for the LLM"
    )
    metadata: Dict[str, Any] = Field(default_factory=dict)


class LlmResponse(BaseModel):
    """Response from LLM."""

    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None
    finish_reason: Optional[str] = None
    usage: Optional[Dict[str, int]] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)

    def is_tool_call(self) -> bool:
        """Check if this response contains tool calls."""
        return self.tool_calls is not None and len(self.tool_calls) > 0


class LlmStreamChunk(BaseModel):
    """Streaming chunk from LLM."""

    content: Optional[str] = None
    tool_calls: Optional[List[ToolCall]] = None
    finish_reason: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)


================================================
FILE: src/vanna/core/middleware/__init__.py
================================================
"""
Middleware system for LLM request/response interception.

This module provides middleware interfaces for intercepting and transforming
LLM requests and responses.
"""

from .base import LlmMiddleware

__all__ = ["LlmMiddleware"]


================================================
FILE: src/vanna/core/middleware/base.py
================================================
"""
Base LLM middleware interface.

Middleware allows you to intercept and transform LLM requests and responses
for caching, monitoring, content filtering, and more.
"""

from abc import ABC
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ..llm import LlmRequest, LlmResponse


class LlmMiddleware(ABC):
    """Middleware for intercepting LLM requests and responses.

    Subclass this to create custom middleware that can:
    - Cache LLM responses
    - Log requests/responses
    - Filter or modify content
    - Track costs and usage
    - Implement fallback strategies

    Example:
        class CachingMiddleware(LlmMiddleware):
            def __init__(self):
                self.cache = {}

            async def before_llm_request(self, request: LlmRequest) -> LlmRequest:
                # Could check cache here
                return request

            async def after_llm_response(self, request: LlmRequest, response: LlmResponse) -> LlmResponse:
                # Cache the response
                cache_key = self._compute_key(request)
                self.cache[cache_key] = response
                return response

        agent = AgentRunner(
            llm_service=...,
            llm_middlewares=[CachingMiddleware(), LoggingMiddleware()]
        )
    """

    async def before_llm_request(self, request: "LlmRequest") -> "LlmRequest":
        """Called before sending request to LLM.

        Args:
            request: The LLM request about to be sent

        Returns:
            Modified request, or original if no changes
        """
        return request

    async def after_llm_response(
        self, request: "LlmRequest", response: "LlmResponse"
    ) -> "LlmResponse":
        """Called after receiving response from LLM.

        Args:
            request: The original request
            response: The LLM response

        Returns:
            Modified response, or original if no changes
        """
        return response


================================================
FILE: src/vanna/core/observability/__init__.py
================================================
"""
Observability system for telemetry and monitoring.

This module provides interfaces for collecting metrics, traces, and
monitoring agent behavior.
"""

from .base import ObservabilityProvider
from .models import Span, Metric

__all__ = ["ObservabilityProvider", "Span", "Metric"]


================================================
FILE: src/vanna/core/observability/base.py
================================================
"""
Base observability provider interface.

Observability providers allow you to collect telemetry data about
agent execution for monitoring and debugging.
"""

from abc import ABC
from typing import Any, Dict, Optional

from .models import Span, Metric


class ObservabilityProvider(ABC):
    """Provider for collecting telemetry and observability data.

    Subclass this to create custom observability integrations that can:
    - Emit metrics to monitoring systems
    - Create distributed traces
    - Log performance data
    - Track costs and usage
    - Monitor error rates

    Example:
        class PrometheusProvider(ObservabilityProvider):
            def __init__(self, registry):
                self.registry = registry
                self.request_counter = Counter(
                    'agent_requests_total',
                    'Total agent requests',
                    registry=registry
                )

            async def record_metric(self, name: str, value: float, tags: Dict[str, str]) -> None:
                if name == "agent.request":
                    self.request_counter.inc()

            async def create_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Span:
                span = Span(name=name, attributes=attributes or {})
                return span

        agent = AgentRunner(
            llm_service=...,
            observability_provider=PrometheusProvider(registry)
        )
    """

    async def record_metric(
        self,
        name: str,
        value: float,
        unit: str = "",
        tags: Optional[Dict[str, str]] = None,
    ) -> None:
        """Record a metric measurement.

        Args:
            name: Metric name (e.g., "agent.request.duration")
            value: Metric value
            unit: Unit of measurement (e.g., "ms", "tokens")
            tags: Additional tags/labels for the metric
        """
        pass

    async def create_span(
        self, name: str, attributes: Optional[Dict[str, Any]] = None
    ) -> Span:
        """Create a new span for tracing.

        Args:
            name: Span name/operation
            attributes: Initial span attributes

        Returns:
            Span object to track the operation

        Note:
            Call span.end() when the operation completes.
        """
        return Span(name=name, attributes=attributes or {})

    async def end_span(self, span: Span) -> None:
        """End a span and record it.

        Args:
            span: The span to end
        """
        span.end()


================================================
FILE: src/vanna/core/observability/models.py
================================================
"""
Observability models for spans and metrics.
"""

import time
from typing import Any, Dict, Optional
from uuid import uuid4

from pydantic import BaseModel, Field


class Span(BaseModel):
    """Represents a unit of work for distributed tracing."""

    id: str = Field(default_factory=lambda: str(uuid4()), description="Span ID")
    name: str = Field(description="Span name/operation")
    start_time: float = Field(default_factory=time.time, description="Start timestamp")
    end_time: Optional[float] = Field(default=None, description="End timestamp")
    attributes: Dict[str, Any] = Field(
        default_factory=dict, description="Span attributes"
    )
    parent_id: Optional[str] = Field(default=None, description="Parent span ID")

    def end(self) -> None:
        """Mark span as ended."""
        if self.end_time is None:
            self.end_time = time.time()

    def duration_ms(self) -> Optional[float]:
        """Get span duration in milliseconds."""
        if self.end_time is None:
            return None
        return (self.end_time - self.start_time) * 1000

    def set_attribute(self, key: str, value: Any) -> None:
        """Set a span attribute."""
        self.attributes[key] = value


class Metric(BaseModel):
    """Represents a metric measurement."""

    name: str = Field(description="Metric name")
    value: float = Field(description="Metric value")
    unit: str = Field(default="", description="Unit of measurement")
    tags: Dict[str, str] = Field(default_factory=dict, description="Metric tags")
    timestamp: float = Field(default_factory=time.time, description="Measurement time")


================================================
FILE: src/vanna/core/recovery/__init__.py
================================================
"""
Error recovery system for handling failures gracefully.

This module provides interfaces for custom error handling, retry logic,
and fallback strategies.
"""

from .base import ErrorRecoveryStrategy
from .models import RecoveryAction, RecoveryActionType

__all__ = ["ErrorRecoveryStrategy", "RecoveryAction", "RecoveryActionType"]


================================================
FILE: src/vanna/core/recovery/base.py
================================================
"""
Base error recovery strategy interface.

Recovery strategies allow you to customize how the agent handles errors
during tool execution and LLM communication.
"""

from abc import ABC
from typing import TYPE_CHECKING

from .models import RecoveryAction, RecoveryActionType

if TYPE_CHECKING:
    from ..tool.models import ToolContext
    from ..llm import LlmRequest


class ErrorRecoveryStrategy(ABC):
    """Strategy for handling errors and implementing retry logic.

    Subclass this to create custom error recovery strategies that can:
    - Retry failed operations with backoff
    - Fallback to alternative approaches
    - Log errors to external systems
    - Gracefully degrade functionality

    Example:
        class ExponentialBackoffStrategy(ErrorRecoveryStrategy):
            async def handle_tool_error(
                self, error: Exception, context: ToolContext, attempt: int
            ) -> RecoveryAction:
                if attempt < 3:
                    delay = (2 ** attempt) * 1000  # Exponential backoff
                    return RecoveryAction(
                        action=RecoveryActionType.RETRY,
                        retry_delay_ms=delay,
                        message=f"Retrying after {delay}ms"
                    )
                return RecoveryAction(
                    action=RecoveryActionType.FAIL,
                    message="Max retries exceeded"
                )

        agent = AgentRunner(
            llm_service=...,
            error_recovery_strategy=ExponentialBackoffStrategy()
        )
    """

    async def handle_tool_error(
        self, error: Exception, context: "ToolContext", attempt: int = 1
    ) -> RecoveryAction:
        """Handle errors during tool execution.

        Args:
            error: The exception that occurred
            context: Tool execution context
            attempt: Current attempt number (1-indexed)

        Returns:
            RecoveryAction indicating how to proceed
        """
        # Default: fail immediately
        return RecoveryAction(
            action=RecoveryActionType.FAIL, message=f"Tool error: {str(error)}"
        )

    async def handle_llm_error(
        self, error: Exception, request: "LlmRequest", attempt: int = 1
    ) -> RecoveryAction:
        """Handle errors during LLM communication.

        Args:
            error: The exception that occurred
            request: The LLM request that failed
            attempt: Current attempt number (1-indexed)

        Returns:
            RecoveryAction indicating how to proceed
        """
        # Default: fail immediately
        return RecoveryAction(
            action=RecoveryActionType.FAIL, message=f"LLM error: {str(error)}"
        )


================================================
FILE: src/vanna/core/recovery/models.py
================================================
"""
Recovery action models for error handling.
"""

from enum import Enum
from typing import Any, Optional

from pydantic import BaseModel, Field


class RecoveryActionType(str, Enum):
    """Types of recovery actions."""

    RETRY = "retry"
    FAIL = "fail"
    FALLBACK = "fallback"
    SKIP = "skip"


class RecoveryAction(BaseModel):
    """Action to take when recovering from an error."""

    action: RecoveryActionType = Field(description="Type of recovery action")
    retry_delay_ms: Optional[int] = Field(
        default=None, description="Delay before retry in milliseconds"
    )
    fallback_value: Optional[Any] = Field(
        default=None, description="Fallback value to use"
    )
    message: Optional[str] = Field(
        default=None, description="Message to include with action"
    )


================================================
FILE: src/vanna/core/registry.py
================================================
"""
Tool registry for the Vanna Agents framework.

This module provides the ToolRegistry class for managing and executing tools.
"""

import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, TypeVar, Union

from .tool import Tool, ToolCall, ToolContext, ToolRejection, ToolResult, ToolSchema
from .user import User

if TYPE_CHECKING:
    from .audit import AuditLogger
    from .agent.config import AuditConfig

T = TypeVar("T")


class _LocalToolWrapper(Tool[T]):
    """Wrapper for tools with configurable access groups."""

    def __init__(self, wrapped_tool: Tool[T], access_groups: List[str]):
        self._wrapped_tool = wrapped_tool
        self._access_groups = access_groups

    @property
    def name(self) -> str:
        return self._wrapped_tool.name

    @property
    def description(self) -> str:
        return self._wrapped_tool.description

    @property
    def access_groups(self) -> List[str]:
        return self._access_groups

    def get_args_schema(self) -> Type[T]:
        return self._wrapped_tool.get_args_schema()

    async def execute(self, context: ToolContext, args: T) -> ToolResult:
        return await self._wrapped_tool.execute(context, args)


class ToolRegistry:
    """Registry for managing tools."""

    def __init__(
        self,
        audit_logger: Optional["AuditLogger"] = None,
        audit_config: Optional["AuditConfig"] = None,
    ) -> None:
        self._tools: Dict[str, Tool[Any]] = {}
        self.audit_logger = audit_logger
        if audit_config is not None:
            self.audit_config = audit_config
        else:
            from .agent.config import AuditConfig

            self.audit_config = AuditConfig()

    def register_local_tool(self, tool: Tool[Any], access_groups: List[str]) -> None:
        """Register a local tool with optional access group restrictions.

        Args:
            tool: The tool to register
            access_groups: List of groups that can access this tool.
                          If None or empty, tool is accessible to all users.
        """
        if tool.name in self._tools:
            raise ValueError(f"Tool '{tool.name}' already registered")

        if access_groups:
            # Wrap the tool with access groups
            wrapped_tool = _LocalToolWrapper(tool, access_groups)
            self._tools[tool.name] = wrapped_tool
        else:
            # No access restrictions, register as-is
            self._tools[tool.name] = tool

    async def get_tool(self, name: str) -> Optional[Tool[Any]]:
        """Get a tool by name."""
        return self._tools.get(name)

    async def list_tools(self) -> List[str]:
        """List all registered tool names."""
        return list(self._tools.keys())

    async def get_schemas(self, user: Optional[User] = None) -> List[ToolSchema]:
        """Get schemas for all tools accessible to user."""
        schemas = []
        for tool in self._tools.values():
            if user is None or await self._validate_tool_permissions(tool, user):
                schemas.append(tool.get_schema())
        return schemas

    async def _validate_tool_permissions(self, tool: Tool[Any], user: User) -> bool:
        """Validate if user has access to tool based on group membership.

        Checks for intersection between user's group memberships and tool's access groups.
        If tool has no access groups specified, it's accessible to all users.
        """
        tool_access_groups = tool.access_groups
        if not tool_access_groups:
            return True

        user_groups = set(user.group_memberships)
        tool_groups = set(tool_access_groups)
        # Grant access if any group in user.group_memberships exists in tool.access_groups
        return bool(user_groups & tool_groups)

    async def transform_args(
        self,
        tool: Tool[T],
        args: T,
        user: User,
        context: ToolContext,
    ) -> Union[T, ToolRejection]:
        """Transform and validate tool arguments based on user context.

        This method allows per-user transformation of tool arguments, such as:
        - Applying row-level security (RLS) to SQL queries
        - Filtering available options based on user permissions
        - Validating required arguments are present
        - Redacting sensitive fields

        The default implementation performs no transformation (NoOp).
        Subclasses can override this method to implement custom transformation logic.

        Args:
            tool: The tool being executed
            args: Already Pydantic-validated arguments
            user: The user executing the tool
            context: Full execution context

        Returns:
            Either:
            - Transformed arguments (may be unchanged if no transformation needed)
            - ToolRejection with explanation of why args were rejected
        """
        return args  # Default: no transformation (NoOp)

    async def execute(
        self,
        tool_call: ToolCall,
        context: ToolContext,
    ) -> ToolResult:
        """Execute a tool call with validation."""
        tool = await self.get_tool(tool_call.name)
        if not tool:
            msg = f"Tool '{tool_call.name}' not found"
            return ToolResult(
                success=False,
                result_for_llm=msg,
                ui_component=None,
                error=msg,
            )

        # Validate group access
        if not await self._validate_tool_permissions(tool, context.user):
            msg = f"Insufficient group access for tool '{tool_call.name}'"

            # Audit access denial
            if (
                self.audit_logger
                and self.audit_config
                and self.audit_config.log_tool_access_checks
            ):
                await self.audit_logger.log_tool_access_check(
                    user=context.user,
                    tool_name=tool_call.name,
                    access_granted=False,
                    required_groups=tool.access_groups,
                    context=context,
                    reason=msg,
                )

            return ToolResult(
                success=False,
                result_for_llm=msg,
                ui_component=None,
                error=msg,
            )

        # Validate and parse arguments
        try:
            args_model = tool.get_args_schema()
            validated_args = args_model.model_validate(tool_call.arguments)
        except Exception as e:
            msg = f"Invalid arguments: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=msg,
                ui_component=None,
                error=msg,
            )

        # Transform/validate arguments based on user context
        transform_result = await self.transform_args(
            tool=tool,
            args=validated_args,
            user=context.user,
            context=context,
        )

        if isinstance(transform_result, ToolRejection):
            return ToolResult(
                success=False,
                result_for_llm=transform_result.reason,
                ui_component=None,
                error=transform_result.reason,
            )

        # Use transformed arguments for execution
        final_args = transform_result

        # Audit successful access check
        if (
            self.audit_logger
            and self.audit_config
            and self.audit_config.log_tool_access_checks
        ):
            await self.audit_logger.log_tool_access_check(
                user=context.user,
                tool_name=tool_call.name,
                access_granted=True,
                required_groups=tool.access_groups,
                context=context,
            )

        # Audit tool invocation
        if (
            self.audit_logger
            and self.audit_config
            and self.audit_config.log_tool_invocations
        ):
            # Get UI features if available from context
            ui_features = context.metadata.get("ui_features_available", [])
            await self.audit_logger.log_tool_invocation(
                user=context.user,
                tool_call=tool_call,
                ui_features=ui_features,
                context=context,
                sanitize_parameters=self.audit_config.sanitize_tool_parameters,
            )

        # Execute tool with context-first signature
        try:
            start_time = time.perf_counter()
            result = await tool.execute(context, final_args)
            execution_time_ms = (time.perf_counter() - start_time) * 1000

            # Add execution time to metadata
            result.metadata["execution_time_ms"] = execution_time_ms

            # Audit tool result
            if (
                self.audit_logger
                and self.audit_config
                and self.audit_config.log_tool_results
            ):
                await self.audit_logger.log_tool_result(
                    user=context.user,
                    tool_call=tool_call,
                    result=result,
                    context=context,
                )

            return result
        except Exception as e:
            msg = f"Execution failed: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=msg,
                ui_component=None,
                error=msg,
            )


================================================
FILE: src/vanna/core/rich_component.py
================================================
"""
Base classes for rich UI components.

This module provides the base RichComponent class and supporting enums
for the component system.
"""

import uuid
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, TypeVar

from pydantic import BaseModel, Field

# Type variable for self-returning methods
T = TypeVar("T", bound="RichComponent")


class ComponentType(str, Enum):
    """Types of rich UI components."""

    # Basic components
    TEXT = "text"
    CARD = "card"
    CONTAINER = "container"

    # Primitive UI components (domain-agnostic)
    STATUS_CARD = "status_card"
    PROGRESS_DISPLAY = "progress_display"
    LOG_VIEWER = "log_viewer"
    BADGE = "badge"
    ICON_TEXT = "icon_text"

    # Interactive components
    TASK_LIST = "task_list"
    PROGRESS_BAR = "progress_bar"
    BUTTON = "button"
    BUTTON_GROUP = "button_group"

    # Data components
    TABLE = "table"
    DATAFRAME = "dataframe"
    CHART = "chart"
    CODE_BLOCK = "code_block"

    # Status components
    STATUS_INDICATOR = "status_indicator"
    NOTIFICATION = "notification"
    ALERT = "alert"

    # Artifact components
    ARTIFACT = "artifact"

    # UI state components
    STATUS_BAR_UPDATE = "status_bar_update"
    TASK_TRACKER_UPDATE = "task_tracker_update"
    CHAT_INPUT_UPDATE = "chat_input_update"

    # Legacy (deprecated - use primitives instead)
    TOOL_EXECUTION = "tool_execution"


class ComponentLifecycle(str, Enum):
    """Component lifecycle operations."""

    CREATE = "create"
    UPDATE = "update"
    REPLACE = "replace"
    REMOVE = "remove"


class RichComponent(BaseModel):
    """Base class for all rich UI components."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    type: ComponentType
    lifecycle: ComponentLifecycle = ComponentLifecycle.CREATE
    data: Dict[str, Any] = Field(default_factory=dict)
    children: List[str] = Field(default_factory=list)  # Child component IDs
    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
    visible: bool = True
    interactive: bool = False

    def update(self: T, **kwargs: Any) -> T:
        """Create an updated copy of this component."""
        updated_data = self.model_dump()
        updated_data.update(kwargs)
        updated_data["lifecycle"] = ComponentLifecycle.UPDATE
        updated_data["timestamp"] = datetime.utcnow().isoformat()
        return self.__class__(**updated_data)

    def hide(self: T) -> T:
        """Create a hidden copy of this component."""
        return self.update(visible=False)

    def show(self: T) -> T:
        """Create a visible copy of this component."""
        return self.update(visible=True)

    def serialize_for_frontend(self) -> Dict[str, Any]:
        """Normalize component payload for the frontend renderer.

        The frontend expects component-specific fields to live under the
        ``data`` key while the shared metadata (``id``, ``type``, layout hints,
        etc.) remains at the top level. Pydantic's ``model_dump`` keeps
        component attributes at the top level, so we remap them here before
        streaming them across the wire.
        """

        # Base fields that should remain at the top level of the payload.
        shared_fields = {
            "id",
            "type",
            "lifecycle",
            "children",
            "timestamp",
            "visible",
            "interactive",
        }

        raw = self.model_dump()
        payload: Dict[str, Any] = {}

        # Preserve any existing data payload so implementations can opt-in to
        # advanced usage without losing information.
        raw_data = raw.get("data")
        if raw_data is not None and isinstance(raw_data, dict):
            component_data: Dict[str, Any] = raw_data.copy()
        else:
            # Handle case where data might be a sequence or other type, or None
            component_data = {}

        for key, value in raw.items():
            if key in shared_fields:
                payload[key] = value
            elif key == "data":
                # For most components, skip the base data field
                continue
            elif (
                key == "rows"
                and hasattr(self, "type")
                and self.type.value == "dataframe"
            ):
                # For DataFrame components, the 'rows' field contains the actual row data
                # which should be included in the component_data as 'data' for the frontend
                component_data["data"] = value
            else:
                component_data[key] = value

        payload["data"] = component_data

        # Ensure enums are serialized as primitive values for the frontend.
        payload["type"] = self.type.value
        payload["lifecycle"] = self.lifecycle.value

        return payload


================================================
FILE: src/vanna/core/simple_component.py
================================================
"""Base classes for simple UI components."""

from typing import Any, Dict, Optional
from pydantic import BaseModel, Field
from enum import Enum


class SimpleComponentType(str, Enum):
    TEXT = "text"
    IMAGE = "image"
    LINK = "link"


class SimpleComponent(BaseModel):
    """A simple UI component with basic attributes."""

    type: SimpleComponentType = Field(..., description="Type of the component.")
    semantic_type: Optional[str] = Field(
        default=None, description="Semantic type for better categorization."
    )
    metadata: Optional[Dict[str, Any]] = Field(
        default=None, description="Additional metadata for the component."
    )

    def serialize_for_frontend(self) -> Dict[str, Any]:
        """Serialize simple component for API consumption."""
        return self.model_dump()


================================================
FILE: src/vanna/core/storage/__init__.py
================================================
"""
Storage domain.

This module provides the core abstractions for conversation storage in the Vanna Agents framework.
"""

from .base import ConversationStore
from .models import Conversation, Message

__all__ = [
    "ConversationStore",
    "Conversation",
    "Message",
]


================================================
FILE: src/vanna/core/storage/base.py
================================================
"""
Storage domain interface.

This module contains the abstract base class for conversation storage.
"""

from abc import ABC, abstractmethod
from typing import List, Optional

from .models import Conversation
from ..user.models import User


class ConversationStore(ABC):
    """Abstract base class for conversation storage."""

    @abstractmethod
    async def create_conversation(
        self, conversation_id: str, user: User, initial_message: str
    ) -> Conversation:
        """Create a new conversation with the specified ID."""
        pass

    @abstractmethod
    async def get_conversation(
        self, conversation_id: str, user: User
    ) -> Optional[Conversation]:
        """Get conversation by ID, scoped to user."""
        pass

    @abstractmethod
    async def update_conversation(self, conversation: Conversation) -> None:
        """Update conversation with new messages."""
        pass

    @abstractmethod
    async def delete_conversation(self, conversation_id: str, user: User) -> bool:
        """Delete conversation."""
        pass

    @abstractmethod
    async def list_conversations(
        self, user: User, limit: int = 50, offset: int = 0
    ) -> List[Conversation]:
        """List conversations for user."""
        pass


================================================
FILE: src/vanna/core/storage/models.py
================================================
"""
Storage domain models.

This module contains data models for conversation storage.
"""

from datetime import datetime
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

from ..tool.models import ToolCall
from ..user.models import User


class Message(BaseModel):
    """Single message in a conversation."""

    role: str = Field(description="Message role (user/assistant/system/tool)")
    content: str = Field(description="Message content")
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    tool_calls: Optional[List[ToolCall]] = Field(default=None)
    tool_call_id: Optional[str] = Field(
        default=None, description="ID if this is a tool response"
    )


class Conversation(BaseModel):
    """Conversation containing multiple messages."""

    id: str = Field(description="Unique conversation identifier")
    user: User = Field(description="User this conversation belongs to")
    messages: List[Message] = Field(
        default_factory=list, description="Messages in conversation"
    )
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional conversation metadata"
    )

    def add_message(self, message: Message) -> None:
        """Add a message to the conversation."""
        self.messages.append(message)
        self.updated_at = datetime.utcnow()


================================================
FILE: src/vanna/core/system_prompt/__init__.py
================================================
"""
System prompt domain.

This module provides the core abstractions for building system prompts in the Vanna Agents framework.
"""

from .base import SystemPromptBuilder
from .default import DefaultSystemPromptBuilder

__all__ = [
    "SystemPromptBuilder",
    "DefaultSystemPromptBuilder",
]


================================================
FILE: src/vanna/core/system_prompt/base.py
================================================
"""
System prompt builder interface.

This module contains the abstract base class for system prompt builders.
"""

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, List, Optional

if TYPE_CHECKING:
    from ..tool.models import ToolSchema
    from ..user.models import User


class SystemPromptBuilder(ABC):
    """Abstract base class for system prompt builders.

    Subclasses should implement the build_system_prompt method to generate
    system prompts based on user context and available tools.
    """

    @abstractmethod
    async def build_system_prompt(
        self, user: "User", tools: List["ToolSchema"]
    ) -> Optional[str]:
        """
        Build a system prompt based on user context and available tools.

        Args:
            user: The user making the request
            tools: List of tools available to the user

        Returns:
            System prompt string, or None if no system prompt should be used
        """
        pass


================================================
FILE: src/vanna/core/system_prompt/default.py
================================================
"""
Default system prompt builder implementation with memory workflow support.

This module provides a default implementation of the SystemPromptBuilder interface
that automatically includes memory workflow instructions when memory tools are available.
"""

from typing import TYPE_CHECKING, List, Optional
from datetime import datetime

from .base import SystemPromptBuilder

if TYPE_CHECKING:
    from ..tool.models import ToolSchema
    from ..user.models import User


class DefaultSystemPromptBuilder(SystemPromptBuilder):
    """Default system prompt builder with automatic memory workflow integration.

    Dynamically generates system prompts that include memory workflow
    instructions when memory tools (search_saved_correct_tool_uses and
    save_question_tool_args) are available.
    """

    def __init__(self, base_prompt: Optional[str] = None):
        """Initialize with an optional base prompt.

        Args:
            base_prompt: Optional base system prompt. If not provided, uses a default.
        """
        self.base_prompt = base_prompt

    async def build_system_prompt(
        self, user: "User", tools: List["ToolSchema"]
    ) -> Optional[str]:
        """
        Build a system prompt with memory workflow instructions.

        Args:
            user: The user making the request
            tools: List of tools available to the user

        Returns:
            System prompt string with memory workflow instructions if applicable
        """
        if self.base_prompt is not None:
            return self.base_prompt

        # Check which memory tools are available
        tool_names = [tool.name for tool in tools]
        has_search = "search_saved_correct_tool_uses" in tool_names
        has_save = "save_question_tool_args" in tool_names
        has_text_memory = "save_text_memory" in tool_names

        # Get today's date
        today_date = datetime.now().strftime("%Y-%m-%d")

        # Base system prompt
        prompt_parts = [
            f"You are Vanna, an AI data analyst assistant created to help users with data analysis tasks. Today's date is {today_date}.",
            "",
            "Response Guidelines:",
            "- Any summary of what you did or observations should be the final step.",
            "- Use the available tools to help the user accomplish their goals.",
            "- When you execute a query, that raw result is shown to the user outside of your response so YOU DO NOT need to include it in your response. Focus on summarizing and interpreting the results.",
        ]

        if tools:
            prompt_parts.append(
                f"\nYou have access to the following tools: {', '.join(tool_names)}"
            )

        # Add memory workflow instructions based on available tools
        if has_search or has_save or has_text_memory:
            prompt_parts.append("\n" + "=" * 60)
            prompt_parts.append("MEMORY SYSTEM:")
            prompt_parts.append("=" * 60)

        if has_search or has_save:
            prompt_parts.append("\n1. TOOL USAGE MEMORY (Structured Workflow):")
            prompt_parts.append("-" * 50)

        if has_search:
            prompt_parts.extend(
                [
                    "",
                    "• BEFORE executing any tool (run_sql, visualize_data, or calculator), you MUST first call search_saved_correct_tool_uses with the user's question to check if there are existing successful patterns for similar questions.",
                    "",
                    "• Review the search results (if any) to inform your approach before proceeding with other tool calls.",
                ]
            )

        if has_save:
            prompt_parts.extend(
                [
                    "",
                    "• AFTER successfully executing a tool that produces correct and useful results, you MUST call save_question_tool_args to save the successful pattern for future use.",
                ]
            )

        if has_search or has_save:
            prompt_parts.extend(
                [
                    "",
                    "Example workflow:",
                    "  • User asks a question",
                    f'  • First: Call search_saved_correct_tool_uses(question="user\'s question")'
                    if has_search
                    else "",
                    "  • Then: Execute the appropriate tool(s) based on search results and the question",
                    f'  • Finally: If successful, call save_question_tool_args(question="user\'s question", tool_name="tool_used", args={{the args you used}})'
                    if has_save
                    else "",
                    "",
                    "Do NOT skip the search step, even if you think you know how to answer. Do NOT forget to save successful executions."
                    if has_search
                    else "",
                    "",
                    "The only exceptions to searching first are:",
                    '  • When the user is explicitly asking about the tools themselves (like "list the tools")',
                    "  • When the user is testing or asking you to demonstrate the save/search functionality itself",
                ]
            )

        if has_text_memory:
            prompt_parts.extend(
                [
                    "",
                    "2. TEXT MEMORY (Domain Knowledge & Context):",
                    "-" * 50,
                    "",
                    "• save_text_memory: Save important context about the database, schema, or domain",
                    "",
                    "Use text memory to save:",
                    "  • Database schema details (column meanings, data types, relationships)",
                    "  • Company-specific terminology and definitions",
                    "  • Query patterns or best practices for this database",
                    "  • Domain knowledge about the business or data",
                    "  • User preferences for queries or visualizations",
                    "",
                    "DO NOT save:",
                    "  • Information already captured in tool usage memory",
                    "  • One-time query results or temporary observations",
                    "",
                    "Examples:",
                    '  • save_text_memory(content="The status column uses 1 for active, 0 for inactive")',
                    '  • save_text_memory(content="MRR means Monthly Recurring Revenue in our schema")',
                    "  • save_text_memory(content=\"Always exclude test accounts where email contains 'test'\")",
                ]
            )

        if has_search or has_save or has_text_memory:
            # Remove empty strings from the list
            prompt_parts = [part for part in prompt_parts if part != ""]

        return "\n".join(prompt_parts)


================================================
FILE: src/vanna/core/tool/__init__.py
================================================
"""
Tool domain.

This module provides the core abstractions for tools in the Vanna Agents framework.
"""

from .base import T, Tool
from .models import ToolCall, ToolContext, ToolRejection, ToolResult, ToolSchema

__all__ = [
    "Tool",
    "T",
    "ToolCall",
    "ToolContext",
    "ToolRejection",
    "ToolResult",
    "ToolSchema",
]


================================================
FILE: src/vanna/core/tool/base.py
================================================
"""
Tool domain interface.

This module contains the abstract base class for tools.
"""

from abc import ABC, abstractmethod
from typing import Generic, List, Type, TypeVar

from .models import ToolContext, ToolResult, ToolSchema

# Type variable for tool argument types
T = TypeVar("T")


class Tool(ABC, Generic[T]):
    """Abstract base class for tools."""

    @property
    @abstractmethod
    def name(self) -> str:
        """Unique name for this tool."""
        pass

    @property
    @abstractmethod
    def description(self) -> str:
        """Description of what this tool does."""
        pass

    @property
    def access_groups(self) -> List[str]:
        """Groups permitted to access this tool."""
        return []

    @abstractmethod
    def get_args_schema(self) -> Type[T]:
        """Return the Pydantic model for arguments."""
        pass

    @abstractmethod
    async def execute(self, context: ToolContext, args: T) -> ToolResult:
        """Execute the tool with validated arguments.

        Args:
            context: Execution context containing user, conversation_id, and request_id
            args: Validated tool arguments

        Returns:
            ToolResult with success status, result for LLM, and optional UI component
        """
        pass

    def get_schema(self) -> ToolSchema:
        """Generate tool schema for LLM."""
        from typing import Any, cast

        args_model = self.get_args_schema()
        # Get the schema - args_model should be a Pydantic model class
        schema = (
            cast(Any, args_model).model_json_schema()
            if hasattr(args_model, "model_json_schema")
            else {}
        )
        return ToolSchema(
            name=self.name,
            description=self.description,
            parameters=schema,
            access_groups=self.access_groups,
        )


================================================
FILE: src/vanna/core/tool/models.py
================================================
"""
Tool domain models.

This module contains data models for tool execution.
"""

from typing import TYPE_CHECKING, Any, Dict, List, Optional

from pydantic import BaseModel, Field

# Import AgentMemory at runtime for Pydantic model resolution
from vanna.capabilities.agent_memory import AgentMemory

if TYPE_CHECKING:
    from ..components import UiComponent
    from ..user.models import User
    from ..observability import ObservabilityProvider


class ToolCall(BaseModel):
    """Represents a tool call from the LLM."""

    id: str = Field(description="Unique identifier for this tool call")
    name: str = Field(description="Name of the tool to execute")
    arguments: Dict[str, Any] = Field(description="Raw arguments from LLM")


class ToolContext(BaseModel):
    """Context passed to all tool executions."""

    user: "User"  # Forward reference to avoid circular import
    conversation_id: str
    request_id: str = Field(description="Unique request identifier for tracing")
    agent_memory: AgentMemory = Field(
        description="Agent memory for tool usage learning"
    )
    metadata: Dict[str, Any] = Field(default_factory=dict)
    observability_provider: Optional["ObservabilityProvider"] = Field(
        default=None,
        description="Optional observability provider for metrics and spans",
    )

    class Config:
        arbitrary_types_allowed = True


class ToolResult(BaseModel):
    """Result from tool execution.

    Changes:
    - `result_for_llm`: string that will be sent back to the LLM.
    - `ui_component`: optional UI payload for rendering in clients.
    """

    success: bool = Field(description="Whether execution succeeded")
    result_for_llm: str = Field(description="String content to send back to the LLM")
    ui_component: Optional["UiComponent"] = Field(
        default=None, description="Optional UI component for rendering"
    )
    error: Optional[str] = Field(default=None, description="Error message if failed")
    metadata: Dict[str, Any] = Field(default_factory=dict)


class ToolSchema(BaseModel):
    """Schema describing a tool for LLM consumption."""

    name: str = Field(description="Tool name")
    description: str = Field(description="What this tool does")
    parameters: Dict[str, Any] = Field(description="JSON Schema of parameters")
    access_groups: List[str] = Field(
        default_factory=list, description="Groups permitted to access this tool"
    )


class ToolRejection(BaseModel):
    """Indicates tool execution should be rejected with a message.

    Used by transform_args to reject tool execution when arguments
    cannot be appropriately transformed for the user's context.
    """

    reason: str = Field(
        description="Explanation of why the tool execution was rejected"
    )


================================================
FILE: src/vanna/core/user/__init__.py
================================================
"""
User domain.

This module provides the core abstractions for user management in the Vanna Agents framework.
"""

from .base import UserService
from .models import User
from .resolver import UserResolver
from .request_context import RequestContext

__all__ = [
    "UserService",
    "User",
    "UserResolver",
    "RequestContext",
]


================================================
FILE: src/vanna/core/user/base.py
================================================
"""
User domain interface.

This module contains the abstract base class for user services.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional

from .models import User


class UserService(ABC):
    """Service for user management and authentication."""

    @abstractmethod
    async def get_user(self, user_id: str) -> Optional[User]:
        """Get user by ID."""
        pass

    @abstractmethod
    async def authenticate(self, credentials: Dict[str, Any]) -> Optional[User]:
        """Authenticate user and return User object if successful."""
        pass

    @abstractmethod
    async def has_permission(self, user: User, permission: str) -> bool:
        """Check if user has specific permission."""
        pass


================================================
FILE: src/vanna/core/user/models.py
================================================
"""
User domain models.

This module contains data models for user management.
"""

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, ConfigDict, Field


class User(BaseModel):
    """User model for authentication and scoping."""

    id: str = Field(description="Unique user identifier")
    username: Optional[str] = Field(default=None, description="Username")
    email: Optional[str] = Field(default=None, description="User email")
    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional user metadata"
    )
    group_memberships: List[str] = Field(
        default_factory=list, description="Groups the user belongs to"
    )

    model_config = ConfigDict(extra="allow")


================================================
FILE: src/vanna/core/user/request_context.py
================================================
"""
Request context for user resolution.

This module provides the RequestContext model for passing web request
information to UserResolver implementations.
"""

from typing import Any, Dict, Optional

from pydantic import BaseModel, Field


class RequestContext(BaseModel):
    """Context from a web request for user resolution.

    This structured object replaces raw dictionaries for passing request
    data to UserResolver implementations, making it easier to access
    cookies, headers, and other request metadata.

    Example:
        context = RequestContext(
            cookies={'vanna_email': 'alice@example.com'},
            headers={'Authorization': 'Bearer token'},
            remote_addr='127.0.0.1'
        )
        user = await resolver.resolve_user(context)
    """

    cookies: Dict[str, str] = Field(default_factory=dict, description="Request cookies")

    headers: Dict[str, str] = Field(default_factory=dict, description="Request headers")

    remote_addr: Optional[str] = Field(default=None, description="Remote IP address")

    query_params: Dict[str, str] = Field(
        default_factory=dict, description="Query parameters"
    )

    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional framework-specific metadata"
    )

    def get_cookie(self, name: str, default: Optional[str] = None) -> Optional[str]:
        """Get cookie value by name.

        Args:
            name: Cookie name
            default: Default value if cookie not found

        Returns:
            Cookie value or default
        """
        return self.cookies.get(name, default)

    def get_header(self, name: str, default: Optional[str] = None) -> Optional[str]:
        """Get header value by name (case-insensitive).

        Args:
            name: Header name
            default: Default value if header not found

        Returns:
            Header value or default
        """
        # Case-insensitive header lookup
        name_lower = name.lower()
        for key, value in self.headers.items():
            if key.lower() == name_lower:
                return value
        return default


================================================
FILE: src/vanna/core/user/resolver.py
================================================
"""
User resolver interface for web request authentication.

This module provides the abstract base class for resolving web requests
to authenticated User objects.
"""

from abc import ABC, abstractmethod

from .models import User
from .request_context import RequestContext


class UserResolver(ABC):
    """Resolves web requests to authenticated users.

    Implementations of this interface handle the specifics of extracting
    user identity from request context (cookies, headers, tokens, etc.)
    and creating authenticated User objects.

    Example:
        class JwtUserResolver(UserResolver):
            async def resolve_user(self, request_context: RequestContext) -> User:
                token = request_context.get_header('Authorization')
                # ... validate JWT and extract user info
                return User(id=user_id, username=username, email=email)
    """

    @abstractmethod
    async def resolve_user(self, request_context: RequestContext) -> User:
        """Resolve user from request context.

        Args:
            request_context: Structured request context with cookies, headers, etc.

        Returns:
            Authenticated User object

        Raises:
            Can raise exceptions for authentication failures
        """
        pass


================================================
FILE: src/vanna/core/validation.py
================================================
"""
Development utilities for validating Pydantic models.

This module provides utilities that can be used during development
and testing to catch forward reference issues early.
"""

from typing import Any, Dict, List, Tuple, Type
from pydantic import BaseModel
import importlib
import inspect


def validate_pydantic_models_in_package(package_name: str) -> Dict[str, Any]:
    """
    Validate all Pydantic models in a package for completeness.

    This function can be used in tests or development scripts to catch
    forward reference issues before they cause runtime errors.

    Args:
        package_name: Name of the package to validate (e.g., 'vanna.core')

    Returns:
        Dictionary with validation results
    """
    results: Dict[str, Any] = {
        "total_models": 0,
        "incomplete_models": [],
        "models": {},
        "summary": "",
    }

    try:
        # Import the package
        package = importlib.import_module(package_name)

        # Get all submodules
        submodules = []
        if hasattr(package, "__path__"):
            import pkgutil

            for _, name, _ in pkgutil.iter_modules(
                package.__path__, package_name + "."
            ):
                try:
                    submodule = importlib.import_module(name)
                    submodules.append((name, submodule))
                except ImportError:
                    continue
        else:
            submodules = [(package_name, package)]

        # Check all Pydantic models in each submodule
        for module_name, module in submodules:
            for name, obj in inspect.getmembers(module):
                if (
                    inspect.isclass(obj)
                    and issubclass(obj, BaseModel)
                    and obj is not BaseModel
                ):
                    model_key = f"{module_name}.{name}"
                    results["total_models"] += 1

                    # Check for forward references
                    forward_refs: List[Tuple[str, str]] = []
                    for field_name, field_info in obj.model_fields.items():
                        annotation = field_info.annotation
                        if annotation is not None and hasattr(
                            annotation, "__forward_arg__"
                        ):
                            forward_refs.append(
                                (field_name, annotation.__forward_arg__)
                            )

                    # Check completeness
                    try:
                        obj.model_json_schema()
                        is_complete = True
                        error = None
                    except Exception as e:
                        is_complete = False
                        error = str(e)
                        results["incomplete_models"].append(model_key)

                    results["models"][model_key] = {
                        "class": obj,
                        "forward_references": forward_refs,
                        "is_complete": is_complete,
                        "error": error,
                    }

        # Generate summary
        incomplete_models = results["incomplete_models"]
        incomplete_count = len(incomplete_models)
        total_models = results["total_models"]
        if incomplete_count == 0:
            results["summary"] = (
                f"✓ All {total_models} Pydantic models are complete and valid!"
            )
        else:
            results["summary"] = (
                f"⚠ {incomplete_count} of {total_models} models are incomplete: "
                f"{', '.join(incomplete_models)}"
            )

    except Exception as e:
        results["summary"] = f"Error validating package {package_name}: {e}"

    return results


def check_models_health() -> bool:
    """
    Quick health check for all core Pydantic models.

    Returns:
        True if all models are healthy, False otherwise
    """
    core_packages = [
        "vanna.core.tool.models",
        "vanna.core.user.models",
        "vanna.core.llm.models",
        "vanna.core.storage.models",
        "vanna.core.agent.models",
    ]

    all_healthy = True

    for package in core_packages:
        try:
            results = validate_pydantic_models_in_package(package)
            if results["incomplete_models"]:
                print(f"❌ Issues in {package}: {results['incomplete_models']}")
                all_healthy = False
            else:
                print(f"✅ {package}: {results['total_models']} models OK")
        except Exception as e:
            print(f"❌ Error checking {package}: {e}")
            all_healthy = False

    return all_healthy


if __name__ == "__main__":
    print("Checking Pydantic model health across core packages...")
    print("=" * 60)

    healthy = check_models_health()

    print("=" * 60)
    if healthy:
        print("🎉 All Pydantic models are healthy!")
    else:
        print("⚠️  Some models need attention.")
        print("\nTo fix forward reference issues:")
        print("1. Ensure all referenced classes are imported")
        print("2. Call model_rebuild() after imports")
        print("3. Use proper TYPE_CHECKING imports for circular deps")

    print("\nNote: You can also catch these issues at development time using:")
    print("  - mypy static type checking")
    print("  - This validation script in your test suite")
    print("  - Pre-commit hooks")


================================================
FILE: src/vanna/core/workflow/__init__.py
================================================
"""
Workflow handler system for deterministic workflow execution.

This module provides the WorkflowHandler interface for intercepting user messages
and executing deterministic workflows before they reach the LLM. This is useful
for command handling, pattern-based routing, and state-based workflows.
"""

from .base import WorkflowHandler, WorkflowResult
from .default import DefaultWorkflowHandler

__all__ = ["WorkflowHandler", "WorkflowResult", "DefaultWorkflowHandler"]


================================================
FILE: src/vanna/core/workflow/base.py
================================================
"""
Base workflow handler interface.

Workflow triggers allow you to execute deterministic workflows in response to
user messages before they are sent to the LLM. This is useful for:
- Command handling (e.g., /help, /reset)
- Pattern-based routing (e.g., report generation)
- State-based workflows (e.g., onboarding flows)
- Quota enforcement with custom responses
"""

from abc import ABC, abstractmethod
from typing import (
    TYPE_CHECKING,
    Optional,
    Union,
    List,
    AsyncGenerator,
    Callable,
    Awaitable,
)
from dataclasses import dataclass

if TYPE_CHECKING:
    from ..user.models import User
    from ..storage import Conversation
    from ...components import UiComponent
    from ..agent.agent import Agent


@dataclass
class WorkflowResult:
    """Result from a workflow handler attempt.

    When a workflow handles a message, it can optionally return UI components to stream
    to the user and/or mutate the conversation state.

    Attributes:
        should_skip_llm: If True, the workflow handled the message and LLM processing is skipped.
                         If False, the message continues to the agent/LLM.
        components: Optional UI components to stream back to the user.
                    Can be a list or async generator for streaming responses.
        conversation_mutation: Optional async callback to modify conversation state
                               (e.g., clearing messages, adding system events).

    Example:
        # Simple command response
        WorkflowResult(
            should_skip_llm=True,
            components=[RichTextComponent(content="Help text here")]
        )

        # With conversation mutation
        async def clear_history(conv):
            conv.messages.clear()

        WorkflowResult(
            should_skip_llm=True,
            components=[StatusCardComponent(...)],
            conversation_mutation=clear_history
        )

        # Not handled, continue to agent
        WorkflowResult(should_skip_llm=False)
    """

    should_skip_llm: bool
    components: Optional[
        Union[List["UiComponent"], AsyncGenerator["UiComponent", None]]
    ] = None
    conversation_mutation: Optional[Callable[["Conversation"], Awaitable[None]]] = None


class WorkflowHandler(ABC):
    """Base class for handling deterministic workflows before LLM processing.

    Implement this interface to intercept user messages and execute deterministic
    workflows instead of sending to the LLM. This is the first extensibility point
    in the agent's message processing pipeline, running after user resolution and
    conversation loading but before the message is added to conversation history
    or sent to the LLM.

    Use cases:
    - Slash commands (/help, /reset, /report)
    - Pattern-based routing (regex matching)
    - State-based workflows (onboarding, surveys)
    - Custom quota enforcement with helpful messages
    - Deterministic report generation
    - Starter UI (buttons, welcome messages) when conversation begins

    Example:
        class CommandWorkflow(WorkflowHandler):
            async def try_handle(self, agent, user, conversation, message):
                if message.startswith("/help"):
                    return WorkflowResult(
                        should_skip_llm=True,
                        components=[
                            RichTextComponent(
                                content="Available commands:\\n- /help\\n- /reset",
                                markdown=True
                            )
                        ]
                    )

                # Execute tool for reports
                if message.startswith("/report"):
                    tool = await agent.tool_registry.get_tool("generate_report")
                    result = await tool.execute(ToolContext(user=user), {})
                    return WorkflowResult(should_skip_llm=True, components=[result.ui_component])

                # Not handled, continue to agent
                return WorkflowResult(should_skip_llm=False)

            async def get_starter_ui(self, agent, user, conversation):
                return [
                    RichTextComponent(content=f"Welcome {user.username}!"),
                    ButtonComponent(label="Generate Report", value="/report"),
                ]

        agent = Agent(
            llm_service=...,
            tool_registry=...,
            user_resolver=...,
            workflow_handler=CommandWorkflow()
        )

    Observability:
        The agent automatically creates an "agent.workflow_handler" span when
        a WorkflowHandler is configured, allowing you to monitor handler
        performance and outcomes.
    """

    @abstractmethod
    async def try_handle(
        self, agent: "Agent", user: "User", conversation: "Conversation", message: str
    ) -> WorkflowResult:
        """Attempt to handle a workflow for the given message.

        This method is called for every user message before it reaches the LLM.
        Inspect the message content, user context, and conversation state to
        decide whether to execute a deterministic workflow or allow normal
        agent processing.

        Args:
            agent: The agent instance, providing access to tool_registry, config,
                   and observability_provider for tool execution and logging.
            user: The user who sent the message, including their ID, permissions,
                  and metadata. Use this for permission checks or personalization.
            conversation: The current conversation context, including message history.
                          Can be inspected for state-based workflows.
            message: The user's raw message content.

        Returns:
            WorkflowResult with should_skip_llm=True to execute a workflow and skip LLM,
            or should_skip_llm=False to continue normal agent processing.

            When should_skip_llm=True:
            - The message is NOT added to conversation history automatically
            - The components are streamed to the user
            - The conversation_mutation callback (if provided) is executed
            - The agent returns without calling the LLM

            When should_skip_llm=False:
            - The message is added to conversation history
            - Normal agent processing continues (LLM call, tool execution, etc.)

        Example:
            async def try_handle(self, agent, user, conversation, message):
                # Pattern matching with tool execution
                if message.startswith("/report"):
                    # Execute tool from registry
                    tool = await agent.tool_registry.get_tool("generate_sales_report")
                    context = ToolContext(user=user, conversation=conversation)
                    result = await tool.execute(context, {})

                    return WorkflowResult(
                        should_skip_llm=True,
                        components=[...]
                    )

                # State-based workflow
                if user.metadata.get("needs_onboarding"):
                    return await self._onboarding_flow(agent, user, message)

                # Permission check
                if message.startswith("/admin") and "admin" not in user.permissions:
                    return WorkflowResult(
                        should_skip_llm=True,
                        components=[RichTextComponent(content="Access denied.")]
                    )

                # Continue to agent
                return WorkflowResult(should_skip_llm=False)
        """
        pass

    async def get_starter_ui(
        self, agent: "Agent", user: "User", conversation: "Conversation"
    ) -> Optional[List["UiComponent"]]:
        """Provide UI components when a conversation starts.

        Override this method to show starter buttons, welcome messages,
        or quick actions when a new chat is opened by the user.

        This is called by the frontend/server when initializing a new
        conversation, before any user messages are sent.

        Args:
            agent: The agent instance, providing access to tool_registry, config,
                   and observability_provider for dynamic UI generation.
            user: The user starting the conversation
            conversation: The new conversation (typically empty)

        Returns:
            List of UI components to display, or None for no starter UI.
            Components can include buttons, welcome text, quick actions, etc.

        Example:
            async def get_starter_ui(self, agent, user, conversation):
                # Show role-based quick actions
                if "analyst" in user.permissions:
                    # Dynamically generate buttons based on available tools
                    report_tools = [
                        tool for tool in agent.tool_registry.list_tools()
                        if tool.startswith("report_")
                    ]

                    buttons = [
                        ButtonComponent(label=f"📊 {tool}", value=f"/{tool}")
                        for tool in report_tools
                    ]

                    return [
                        RichTextComponent(
                            content=f"Welcome back, {user.username}!",
                            markdown=True
                        ),
                        *buttons
                    ]

                # New user onboarding
                if user.metadata.get("is_new_user"):
                    return [
                        RichTextComponent(
                            content="# Welcome to Vanna!\\n\\nTry one of these to get started:",
                            markdown=True
                        ),
                        ButtonComponent(label="Show Example Query", value="/example"),
                        ButtonComponent(label="View Tutorial", value="/tutorial"),
                    ]

                return None
        """
        return None


================================================
FILE: src/vanna/core/workflow/default.py
================================================
"""
Default workflow handler implementation with setup health checking.

This module provides a default implementation of the WorkflowHandler interface
that provides a smart starter UI based on available tools and setup status.
"""

from typing import TYPE_CHECKING, List, Optional, Dict, Any
import traceback
import uuid
from .base import WorkflowHandler, WorkflowResult

if TYPE_CHECKING:
    from ..agent.agent import Agent
    from ..user.models import User
    from ..storage import Conversation

# Import components at module level to avoid circular imports
from vanna.components import (
    UiComponent,
    RichTextComponent,
    StatusCardComponent,
    ButtonComponent,
    ButtonGroupComponent,
    SimpleTextComponent,
    CardComponent,
)

# Note: StatusCardComponent and ButtonGroupComponent are kept for /status command compatibility


class DefaultWorkflowHandler(WorkflowHandler):
    """Default workflow handler that provides setup health checking and starter UI.

    This handler provides a starter UI that:
    - Checks if run_sql tool is available (critical)
    - Checks if memory tools are available (warning if missing)
    - Checks if visualization tools are available
    - Provides appropriate setup guidance based on what's missing
    """

    def __init__(self, welcome_message: Optional[str] = None):
        """Initialize with optional custom welcome message.

        Args:
            welcome_message: Optional custom welcome message. If not provided,
                           generates one based on available tools.
        """
        self.welcome_message = welcome_message

    async def try_handle(
        self, agent: "Agent", user: "User", conversation: "Conversation", message: str
    ) -> WorkflowResult:
        """Handle basic commands, but mostly passes through to LLM."""

        # Handle basic help command
        if message.strip().lower() in ["/help", "help", "/h"]:
            # Check if user is admin
            is_admin = "admin" in user.group_memberships

            help_content = (
                "## 🤖 Vanna AI Assistant\n\n"
                "I'm your AI data analyst! Here's what I can help you with:\n\n"
                "**💬 Natural Language Queries**\n"
                '- "Show me sales data for last quarter"\n'
                '- "Which customers have the highest orders?"\n'
                '- "Create a chart of revenue by month"\n\n'
                "**🔧 Commands**\n"
                "- `/help` - Show this help message\n"
            )

            if is_admin:
                help_content += (
                    "\n**🔒 Admin Commands**\n"
                    "- `/status` - Check setup status\n"
                    "- `/memories` - View and manage recent memories\n"
                    "- `/delete [id]` - Delete a memory by ID\n"
                )

            help_content += "\n\nJust ask me anything about your data in plain English!"

            return WorkflowResult(
                should_skip_llm=True,
                components=[
                    UiComponent(
                        rich_component=RichTextComponent(
                            content=help_content,
                            markdown=True,
                        ),
                        simple_component=None,
                    )
                ],
            )

        # Handle status check command (admin-only)
        if message.strip().lower() in ["/status", "status"]:
            # Check if user is admin
            if "admin" not in user.group_memberships:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# 🔒 Access Denied\n\n"
                                "The `/status` command is only available to administrators.\n\n"
                                "If you need access to system status information, please contact your system administrator.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )
            return await self._generate_status_check(agent, user)

        # Handle get recent memories command (admin-only)
        if message.strip().lower() in [
            "/memories",
            "memories",
            "/recent_memories",
            "recent_memories",
        ]:
            # Check if user is admin
            if "admin" not in user.group_memberships:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# 🔒 Access Denied\n\n"
                                "The `/memories` command is only available to administrators.\n\n"
                                "If you need access to memory management features, please contact your system administrator.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )
            return await self._get_recent_memories(agent, user, conversation)

        # Handle delete memory command (admin-only)
        if message.strip().lower().startswith("/delete "):
            # Check if user is admin
            if "admin" not in user.group_memberships:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# 🔒 Access Denied\n\n"
                                "The `/delete` command is only available to administrators.\n\n"
                                "If you need access to memory management features, please contact your system administrator.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )
            memory_id = message.strip()[8:].strip()  # Extract ID after "/delete "
            return await self._delete_memory(agent, user, conversation, memory_id)

        # Don't handle other messages, pass to LLM
        return WorkflowResult(should_skip_llm=False)

    async def get_starter_ui(
        self, agent: "Agent", user: "User", conversation: "Conversation"
    ) -> Optional[List[UiComponent]]:
        """Generate starter UI based on available tools and setup status."""

        # Get available tools
        tools = await agent.tool_registry.get_schemas(user)
        tool_names = [tool.name for tool in tools]

        # Analyze setup
        setup_analysis = self._analyze_setup(tool_names)

        # Check if user is admin (has 'admin' in group memberships)
        is_admin = "admin" in user.group_memberships

        # Generate single concise card
        if self.welcome_message:
            # Use custom welcome message
            return [
                UiComponent(
                    rich_component=RichTextComponent(
                        content=self.welcome_message, markdown=True
                    ),
                    simple_component=None,
                )
            ]
        else:
            # Generate role-aware welcome card
            return [self._generate_starter_card(setup_analysis, is_admin)]

    def _generate_starter_card(
        self, analysis: Dict[str, Any], is_admin: bool
    ) -> UiComponent:
        """Generate a single concise starter card based on role and setup status."""

        if is_admin:
            # Admin view: includes setup status and memory management
            return self._generate_admin_starter_card(analysis)
        else:
            # User view: simple welcome message
            return self._generate_user_starter_card(analysis)

    def _generate_admin_starter_card(self, analysis: Dict[str, Any]) -> UiComponent:
        """Generate admin starter card with setup info and memory management."""

        # Build concise content
        if not analysis["has_sql"]:
            title = "Admin: Setup Required"
            content = "**🔒 Admin View** - You have admin privileges and will see additional system information.\n\n**Vanna AI** requires a SQL connection to function.\n\nPlease configure a SQL tool to get started."
            status = "error"
            icon = "⚠️"
        elif analysis["is_complete"]:
            title = "Admin: System Ready"
            content = "**🔒 Admin View** - You have admin privileges and will see additional system information.\n\n**Vanna AI** is fully configured and ready.\n\n"
            content += "**Setup:** SQL ✓ | Memory ✓ | Visualization ✓"
            status = "success"
            icon = "✅"
        else:
            title = "Admin: System Ready"
            content = "**🔒 Admin View** - You have admin privileges and will see additional system information.\n\n**Vanna AI** is ready to query your database.\n\n"
            setup_items = []
            setup_items.append("SQL ✓")
            setup_items.append("Memory ✓" if analysis["has_memory"] else "Memory ✗")
            setup_items.append("Viz ✓" if analysis["has_viz"] else "Viz ✗")
            content += f"**Setup:** {' | '.join(setup_items)}"
            status = "warning" if not analysis["has_memory"] else "success"
            icon = "⚠️" if not analysis["has_memory"] else "✅"

        # Add memory management info for admins
        actions: List[Dict[str, Any]] = []
        if analysis["has_sql"]:
            actions.append(
                {
                    "label": "💡 Help",
                    "action": "/help",
                    "variant": "secondary",
                }
            )

        if analysis["has_memory"]:
            content += "\n\n**Memory Management:** Tool and text memories are available. As an admin, you can view and manage these memories to help me learn from successful queries."
            actions.append(
                {
                    "label": "🧠 View Memories",
                    "action": "/memories",
                    "variant": "secondary",
                }
            )

        return UiComponent(
            rich_component=CardComponent(
                title=title,
                content=content,
                icon=icon,
                status=status,
                actions=actions,
                markdown=True,
            ),
            simple_component=None,
        )

    def _generate_user_starter_card(self, analysis: Dict[str, Any]) -> UiComponent:
        """Generate simple user starter view using RichTextComponent."""

        if not analysis["has_sql"]:
            content = (
                "# ⚠️ Setup Required\n\n"
                "Vanna AI requires configuration before it can help you analyze data."
            )
        else:
            content = (
                "# 👋 Welcome to Vanna AI\n\n"
                "I'm your AI data analyst assistant. Ask me anything about your data in plain English!\n\n"
                "Type `/help` to see what I can do."
            )

        return UiComponent(
            rich_component=RichTextComponent(content=content, markdown=True),
            simple_component=None,
        )

    def _analyze_setup(self, tool_names: List[str]) -> Dict[str, Any]:
        """Analyze the current tool setup and return status."""

        # Critical tools
        has_sql = any(
            name in tool_names
            for name in ["run_sql", "sql_query", "execute_sql", "query_sql"]
        )

        # Memory tools (important but not critical)
        has_search = "search_saved_correct_tool_uses" in tool_names
        has_save = "save_question_tool_args" in tool_names
        has_memory = has_search and has_save

        # Visualization tools (nice to have)
        has_viz = any(
            name in tool_names
            for name in [
                "visualize_data",
                "create_chart",
                "plot_data",
                "generate_chart",
            ]
        )

        # Other useful tools
        has_calculator = any(
            name in tool_names for name in ["calculator", "calc", "calculate"]
        )

        # Determine overall status
        is_complete = has_sql and has_memory and has_viz
        is_functional = has_sql

        return {
            "has_sql": has_sql,
            "has_memory": has_memory,
            "has_search": has_search,
            "has_save": has_save,
            "has_viz": has_viz,
            "has_calculator": has_calculator,
            "is_complete": is_complete,
            "is_functional": is_functional,
            "tool_count": len(tool_names),
            "tool_names": tool_names,
        }

    def _generate_setup_status_cards(
        self, analysis: Dict[str, Any]
    ) -> List[UiComponent]:
        """Generate status cards showing setup health (used by /status command)."""

        cards = []

        # SQL Tool Status (Critical)
        if analysis["has_sql"]:
            sql_card = StatusCardComponent(
                title="SQL Connection",
                status="success",
                description="Database connection configured and ready",
                icon="✅",
            )
        else:
            sql_card = StatusCardComponent(
                title="SQL Connection",
                status="error",
                description="No SQL tool detected - this is required for data analysis",
                icon="❌",
            )
        cards.append(UiComponent(rich_component=sql_card, simple_component=None))

        # Memory Tools Status (Important)
        if analysis["has_memory"]:
            memory_card = StatusCardComponent(
                title="Memory System",
                status="success",
                description="Search and save tools configured - I can learn from successful queries",
                icon="🧠",
            )
        elif analysis["has_search"] or analysis["has_save"]:
            memory_card = StatusCardComponent(
                title="Memory System",
                status="warning",
                description="Partial memory setup - both search and save tools recommended",
                icon="⚠️",
            )
        else:
            memory_card = StatusCardComponent(
                title="Memory System",
                status="warning",
                description="Memory tools not configured - I won't remember successful patterns",
                icon="⚠️",
            )
        cards.append(UiComponent(rich_component=memory_card, simple_component=None))

        # Visualization Status (Nice to have)
        if analysis["has_viz"]:
            viz_card = StatusCardComponent(
                title="Visualization",
                status="success",
                description="Chart creation tools available",
                icon="📊",
            )
        else:
            viz_card = StatusCardComponent(
                title="Visualization",
                status="info",
                description="No visualization tools - results will be text/tables only",
                icon="📋",
            )
        cards.append(UiComponent(rich_component=viz_card, simple_component=None))

        return cards

    def _generate_setup_guidance(
        self, analysis: Dict[str, Any]
    ) -> Optional[UiComponent]:
        """Generate setup guidance based on what's missing (used by /status command)."""

        if not analysis["has_sql"]:
            # Critical guidance - need SQL
            content = (
                "## 🚨 Setup Required\n\n"
                "To get started with Vanna AI, you need to configure a SQL connection tool:\n\n"
                "```python\n"
                "from vanna.tools import RunSqlTool\n\n"
                "# Add SQL tool to your agent\n"
                "tool_registry.register(RunSqlTool(\n"
                '    connection_string="your-database-connection"\n'
                "))\n"
                "```\n\n"
                "**Next Steps:**\n"
                "1. Configure your database connection\n"
                "2. Add memory tools for learning\n"
                "3. Add visualization tools for charts"
            )

        else:
            # Improvement suggestions
            suggestions = []

            if not analysis["has_memory"]:
                suggestions.append(
                    "**🧠 Add Memory Tools** - Help me learn from successful queries:\n"
                    "```python\n"
                    "from vanna.tools import SearchSavedCorrectToolUses, SaveQuestionToolArgs\n"
                    "tool_registry.register(SearchSavedCorrectToolUses())\n"
                    "tool_registry.register(SaveQuestionToolArgs())\n"
                    "```"
                )

            if not analysis["has_viz"]:
                suggestions.append(
                    "**📊 Add Visualization** - Create charts and graphs:\n"
                    "```python\n"
                    "from vanna.tools import VisualizeDataTool\n"
                    "tool_registry.register(VisualizeDataTool())\n"
                    "```"
                )

            if suggestions:
                content = "## 💡 Suggested Improvements\n\n" + "\n\n".join(suggestions)
            else:
                return None  # No guidance needed

        return UiComponent(
            rich_component=RichTextComponent(content=content, markdown=True),
            simple_component=None,
        )

    async def _generate_status_check(
        self, agent: "Agent", user: "User"
    ) -> WorkflowResult:
        """Generate a detailed status check response."""

        # Get available tools
        tools = await agent.tool_registry.get_schemas(user)
        tool_names = [tool.name for tool in tools]
        analysis = self._analyze_setup(tool_names)

        # Generate status report
        status_content = "# 🔍 Setup Status Report\n\n"

        if analysis["is_complete"]:
            status_content += (
                "🎉 **Excellent!** Your Vanna AI setup is complete and optimized.\n\n"
            )
        elif analysis["is_functional"]:
            status_content += (
                "✅ **Good!** Your setup is functional with room for improvement.\n\n"
            )
        else:
            status_content += (
                "⚠️ **Action Required** - Your setup needs configuration.\n\n"
            )

        status_content += f"**Tools Detected:** {analysis['tool_count']} total\n\n"

        # Tool breakdown
        status_content += "## Tool Status\n\n"
        status_content += f"- **SQL Connection:** {'✅ Available' if analysis['has_sql'] else '❌ Missing (Required)'}\n"
        status_content += f"- **Memory System:** {'✅ Complete' if analysis['has_memory'] else '⚠️ Incomplete' if analysis['has_search'] or analysis['has_save'] else '❌ Missing'}\n"
        status_content += f"- **Visualization:** {'✅ Available' if analysis['has_viz'] else '📋 Text/Tables Only'}\n"
        status_content += f"- **Calculator:** {'✅ Available' if analysis['has_calculator'] else '➖ Not Available'}\n\n"

        if analysis["tool_names"]:
            status_content += (
                f"**Available Tools:** {', '.join(sorted(analysis['tool_names']))}"
            )

        components = [
            UiComponent(
                rich_component=RichTextComponent(content=status_content, markdown=True),
                simple_component=None,
            )
        ]

        # Add status cards
        components.extend(self._generate_setup_status_cards(analysis))

        # Add guidance if needed
        guidance = self._generate_setup_guidance(analysis)
        if guidance:
            components.append(guidance)

        return WorkflowResult(should_skip_llm=True, components=components)

    async def _get_recent_memories(
        self, agent: "Agent", user: "User", conversation: "Conversation"
    ) -> WorkflowResult:
        """Get and display recent memories from agent memory."""
        try:
            # Check if agent has memory capability
            if not hasattr(agent, "agent_memory") or agent.agent_memory is None:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# ⚠️ No Memory System\n\n"
                                "Agent memory is not configured. Recent memories are not available.\n\n"
                                "To enable memory, configure an AgentMemory implementation in your agent setup.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )

            # Create tool context
            from vanna.core.tool import ToolContext

            context = ToolContext(
                user=user,
                conversation_id=conversation.id,
                request_id=str(uuid.uuid4()),
                agent_memory=agent.agent_memory,
            )

            # Get both tool memories and text memories
            tool_memories = await agent.agent_memory.get_recent_memories(
                context=context, limit=10
            )

            # Try to get text memories (may not be implemented in all memory backends)
            text_memories = []
            try:
                text_memories = await agent.agent_memory.get_recent_text_memories(
                    context=context, limit=10
                )
            except (AttributeError, NotImplementedError):
                # Text memories not supported by this implementation
                pass

            if not tool_memories and not text_memories:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# 🧠 Recent Memories\n\n"
                                "No recent memories found. As you use tools and ask questions, "
                                "successful patterns will be saved here for future reference.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )

            components = []

            # Header
            total_count = len(tool_memories) + len(text_memories)
            header_content = f"# 🧠 Recent Memories\n\nFound {total_count} recent memor{'y' if total_count == 1 else 'ies'}"
            components.append(
                UiComponent(
                    rich_component=RichTextComponent(
                        content=header_content, markdown=True
                    ),
                    simple_component=None,
                )
            )

            # Display text memories
            if text_memories:
                components.append(
                    UiComponent(
                        rich_component=RichTextComponent(
                            content=f"## 📝 Text Memories ({len(text_memories)})",
                            markdown=True,
                        ),
                        simple_component=None,
                    )
                )

                for memory in text_memories:
                    # Create card with delete button
                    card_content = f"**Content:** {memory.content}\n\n"
                    if memory.timestamp:
                        card_content += f"**Timestamp:** {memory.timestamp}\n\n"
                    card_content += f"**ID:** `{memory.memory_id}`"

                    card = CardComponent(
                        title="Text Memory",
                        content=card_content,
                        icon="📝",
                        actions=[
                            {
                                "label": "🗑️ Delete",
                                "action": f"/delete {memory.memory_id}",
                                "variant": "error",
                            }
                        ],
                    )
                    components.append(
                        UiComponent(rich_component=card, simple_component=None)
                    )

            # Display tool memories
            if tool_memories:
                components.append(
                    UiComponent(
                        rich_component=RichTextComponent(
                            content=f"## 🔧 Tool Memories ({len(tool_memories)})",
                            markdown=True,
                        ),
                        simple_component=None,
                    )
                )

                for tool_memory in tool_memories:
                    # Create card with delete button
                    card_content = f"**Question:** {tool_memory.question}\n\n"
                    card_content += f"**Tool:** {tool_memory.tool_name}\n\n"
                    card_content += f"**Arguments:** `{tool_memory.args}`\n\n"
                    card_content += f"**Success:** {'✅ Yes' if tool_memory.success else '❌ No'}\n\n"
                    if tool_memory.timestamp:
                        card_content += f"**Timestamp:** {tool_memory.timestamp}\n\n"
                    card_content += f"**ID:** `{tool_memory.memory_id}`"

                    card = CardComponent(
                        title=f"Tool: {tool_memory.tool_name}",
                        content=card_content,
                        markdown=True,
                        icon="🔧",
                        status="success" if tool_memory.success else "error",
                        actions=[
                            {
                                "label": "🗑️ Delete",
                                "action": f"/delete {tool_memory.memory_id}",
                                "variant": "error",
                            }
                        ],
                    )
                    components.append(
                        UiComponent(rich_component=card, simple_component=None)
                    )

            return WorkflowResult(should_skip_llm=True, components=components)

        except Exception as e:
            traceback.print_exc()
            return WorkflowResult(
                should_skip_llm=True,
                components=[
                    UiComponent(
                        rich_component=RichTextComponent(
                            content=f"# ❌ Error Retrieving Memories\n\n"
                            f"Failed to get recent memories: {str(e)}\n\n"
                            f"This may indicate an issue with the agent memory configuration.",
                            markdown=True,
                        ),
                        simple_component=None,
                    )
                ],
            )

    async def _delete_memory(
        self, agent: "Agent", user: "User", conversation: "Conversation", memory_id: str
    ) -> WorkflowResult:
        """Delete a memory by its ID."""
        try:
            # Check if agent has memory capability
            if not hasattr(agent, "agent_memory") or agent.agent_memory is None:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# ⚠️ No Memory System\n\n"
                                "Agent memory is not configured. Cannot delete memories.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )

            if not memory_id:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content="# ⚠️ Invalid Command\n\n"
                                "Please provide a memory ID to delete.\n\n"
                                "Usage: `/delete [memory_id]`",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )

            # Create tool context
            from vanna.core.tool import ToolContext

            context = ToolContext(
                user=user,
                conversation_id=conversation.id,
                request_id=str(uuid.uuid4()),
                agent_memory=agent.agent_memory,
            )

            # Try to delete as a tool memory first
            deleted = await agent.agent_memory.delete_by_id(context, memory_id)

            # If not found as tool memory, try as text memory
            if not deleted:
                try:
                    deleted = await agent.agent_memory.delete_text_memory(
                        context, memory_id
                    )
                except (AttributeError, NotImplementedError):
                    # Text memory deletion not supported by this implementation
                    pass

            if deleted:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content=f"# ✅ Memory Deleted\n\n"
                                f"Successfully deleted memory with ID: `{memory_id}`\n\n"
                                f"You can view remaining memories using `/memories`.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )
            else:
                return WorkflowResult(
                    should_skip_llm=True,
                    components=[
                        UiComponent(
                            rich_component=RichTextComponent(
                                content=f"# ❌ Memory Not Found\n\n"
                                f"Could not find memory with ID: `{memory_id}`\n\n"
                                f"Use `/memories` to see available memory IDs.",
                                markdown=True,
                            ),
                            simple_component=None,
                        )
                    ],
                )

        except Exception as e:
            traceback.print_exc()
            return WorkflowResult(
                should_skip_llm=True,
                components=[
                    UiComponent(
                        rich_component=RichTextComponent(
                            content=f"# ❌ Error Deleting Memory\n\n"
                            f"Failed to delete memory: {str(e)}\n\n"
                            f"This may indicate an issue with the agent memory configuration.",
                            markdown=True,
                        ),
                        simple_component=None,
                    )
                ],
            )


================================================
FILE: src/vanna/examples/__init__.py
================================================
"""Examples for using the Vanna Agents framework."""


================================================
FILE: src/vanna/examples/__main__.py
================================================
"""
Interactive example runner for Vanna Agents.
"""

import sys
import importlib


def main() -> None:
    """Run an example interactively."""
    if len(sys.argv) < 2:
        print("Available examples:")
        print("  python -m vanna.examples mock_quickstart")
        print("  python -m vanna.examples mock_custom_tool")
        print("  python -m vanna.examples anthropic_quickstart")
        print("  python -m vanna.examples openai_quickstart")
        print("  python -m vanna.examples mock_quota_example")
        print("  python -m vanna.examples mock_rich_components_demo")
        print("")
        print("Usage: python -m vanna.examples <example_name>")
        return

    example_name = sys.argv[1]
    try:
        module = importlib.import_module(f"vanna.examples.{example_name}")
        if hasattr(module, "run_interactive"):
            module.run_interactive()
        elif hasattr(module, "main"):
            import asyncio

            if asyncio.iscoroutinefunction(module.main):
                asyncio.run(module.main())
            else:
                module.main()
        else:
            print(f"Example '{example_name}' does not have a main function")
    except ImportError:
        print(f"Example '{example_name}' not found")
    except Exception as e:
        print(f"Error running example '{example_name}': {e}")


if __name__ == "__main__":
    main()


================================================
FILE: src/vanna/examples/anthropic_quickstart.py
================================================
"""
Anthropic example using AnthropicLlmService.

Loads environment from .env (via python-dotenv), uses model 'claude-sonnet-4-20250514'
by default, and sends a simple message through a Agent.

Run:
  PYTHONPATH=. python vanna/examples/anthropic_quickstart.py
"""

import asyncio
import importlib.util
import os
import sys


def ensure_env() -> None:
    if importlib.util.find_spec("dotenv") is not None:
        from dotenv import load_dotenv

        # Load from local .env without overriding existing env
        load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"), override=False)
    else:
        print(
            "[warn] python-dotenv not installed; skipping .env load. Install with: pip install python-dotenv"
        )

    if not os.getenv("ANTHROPIC_API_KEY"):
        print(
            "[error] ANTHROPIC_API_KEY is not set. Add it to your environment or .env file."
        )
        sys.exit(1)


async def main() -> None:
    ensure_env()

    try:
        from vanna.integrations.anthropic import AnthropicLlmService
    except ImportError:
        print(
            "[error] anthropic extra not installed. Install with: pip install -e .[anthropic]"
        )
        raise

    from vanna import AgentConfig, Agent, User
    from vanna.core.registry import ToolRegistry
    from vanna.tools import ListFilesTool

    model = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")
    print(f"Using Anthropic model: {model}")

    llm = AnthropicLlmService(model=model)

    # Create tool registry and register the list_files tool
    tool_registry = ToolRegistry()
    list_files_tool = ListFilesTool()
    tool_registry.register(list_files_tool)

    agent = Agent(
        llm_service=llm,
        config=AgentConfig(stream_responses=False),
        tool_registry=tool_registry,
    )

    user = User(id="demo-user", username="demo")
    conversation_id = "anthropic-demo"

    print("Sending: 'List the files in the current directory'\n")
    async for component in agent.send_message(
        user=user,
        message="List the files in the current directory",
        conversation_id=conversation_id,
    ):
        if hasattr(component, "content") and component.content:
            print("Assistant:", component.content)


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/artifact_example.py
================================================
#!/usr/bin/env python3
"""
Example demonstrating the artifact system in Vanna Agents.

This script shows how agents can create interactive artifacts that can be
rendered externally by developers listening for the 'artifact-opened' event.
"""

import asyncio
from typing import AsyncGenerator, Optional

from vanna import Agent, UiComponent, User, AgentConfig
from vanna.core.rich_components import ArtifactComponent
from vanna.integrations.anthropic.mock import MockLlmService
from vanna.core.interfaces import Agent, LlmService


class ArtifactDemoAgent(Agent):
    """Demo agent that creates various types of artifacts."""

    def __init__(self, llm_service: Optional[LlmService] = None) -> None:
        if llm_service is None:
            llm_service = MockLlmService(
                "I'll help you create interactive artifacts! Try asking me to create a chart, dashboard, or interactive HTML widget."
            )
        super().__init__(
            llm_service=llm_service,
            config=AgentConfig(
                stream_responses=True,
                include_thinking_indicators=True,
            ),
        )

    async def send_message(
        self, user: User, message: str, *, conversation_id: Optional[str] = None
    ) -> AsyncGenerator[UiComponent, None]:
        """Handle user messages and create appropriate artifacts."""
        # First send the normal response
        async for component in super().send_message(
            user, message, conversation_id=conversation_id
        ):
            yield component

        # Then create artifacts based on message content
        message_lower = message.lower()

        if any(
            word in message_lower for word in ["chart", "graph", "visualization", "d3"]
        ):
            async for component in self.create_d3_visualization():
                yield component
        elif any(
            word in message_lower for word in ["dashboard", "analytics", "metrics"]
        ):
            async for component in self.create_dashboard_artifact():
                yield component
        elif any(
            word in message_lower for word in ["html", "interactive", "widget", "demo"]
        ):
            async for component in self.create_html_artifact():
                yield component

    async def create_html_artifact(self) -> AsyncGenerator[UiComponent, None]:
        """Create a simple HTML artifact."""
        html_content = """
        <div style="padding: 20px; font-family: Arial, sans-serif;">
            <h2 style="color: #333;">Interactive HTML Artifact</h2>
            <p>This is a simple HTML artifact that can be opened externally.</p>
            <button onclick="alert('Hello from the artifact!')">Click me!</button>
            <div style="margin-top: 20px;">
                <input type="text" placeholder="Type something..." id="textInput">
                <button onclick="document.getElementById('output').textContent = document.getElementById('textInput').value">
                    Update Text
                </button>
            </div>
            <div id="output" style="margin-top: 10px; padding: 10px; background: #f0f0f0; border-radius: 4px;">
                Output will appear here...
            </div>
        </div>
        """

        artifact = ArtifactComponent.create_html(
            content=html_content,
            title="Interactive HTML Demo",
            description="A simple HTML artifact with interactive elements",
        )

        yield UiComponent(rich_component=artifact)

    async def create_d3_visualization(self) -> AsyncGenerator[UiComponent, None]:
        """Create a D3.js visualization artifact."""
        d3_content = """
        <div id="chart" style="width: 100%; height: 400px;"></div>
        <script>
            // Sample data
            const data = [
                {name: 'A', value: 30},
                {name: 'B', value: 80},
                {name: 'C', value: 45},
                {name: 'D', value: 60},
                {name: 'E', value: 20},
                {name: 'F', value: 90}
            ];

            // Set up dimensions
            const margin = {top: 20, right: 30, bottom: 40, left: 40};
            const width = 800 - margin.left - margin.right;
            const height = 400 - margin.top - margin.bottom;

            // Create SVG
            const svg = d3.select("#chart")
                .append("svg")
                .attr("width", width + margin.left + margin.right)
                .attr("height", height + margin.top + margin.bottom)
                .append("g")
                .attr("transform", `translate(${margin.left},${margin.top})`);

            // Create scales
            const xScale = d3.scaleBand()
                .domain(data.map(d => d.name))
                .range([0, width])
                .padding(0.1);

            const yScale = d3.scaleLinear()
                .domain([0, d3.max(data, d => d.value)])
                .range([height, 0]);

            // Create bars
            svg.selectAll(".bar")
                .data(data)
                .enter().append("rect")
                .attr("class", "bar")
                .attr("x", d => xScale(d.name))
                .attr("width", xScale.bandwidth())
                .attr("y", d => yScale(d.value))
                .attr("height", d => height - yScale(d.value))
                .attr("fill", "#4CAF50")
                .on("mouseover", function(event, d) {
                    d3.select(this).attr("fill", "#45a049");
                })
                .on("mouseout", function(event, d) {
                    d3.select(this).attr("fill", "#4CAF50");
                });

            // Add axes
            svg.append("g")
                .attr("transform", `translate(0,${height})`)
                .call(d3.axisBottom(xScale));

            svg.append("g")
                .call(d3.axisLeft(yScale));
        </script>
        """

        artifact = ArtifactComponent.create_d3(
            content=d3_content,
            title="D3.js Bar Chart",
            description="An interactive bar chart built with D3.js",
        )

        yield UiComponent(rich_component=artifact)

    async def create_dashboard_artifact(self) -> AsyncGenerator[UiComponent, None]:
        """Create a dashboard-style artifact."""
        dashboard_content = """
        <div style="padding: 20px; font-family: Arial, sans-serif; background: #f5f5f5; min-height: 100vh;">
            <h1 style="color: #333; margin-bottom: 30px;">Analytics Dashboard</h1>

            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 20px; margin-bottom: 30px;">
                <div class="metric-card" style="background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <h3 style="margin: 0; color: #666;">Total Users</h3>
                    <p style="font-size: 2em; font-weight: bold; color: #333; margin: 10px 0;">12,456</p>
                    <span style="color: #4CAF50;">↗ +5.2%</span>
                </div>

                <div class="metric-card" style="background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <h3 style="margin: 0; color: #666;">Revenue</h3>
                    <p style="font-size: 2em; font-weight: bold; color: #333; margin: 10px 0;">$89,432</p>
                    <span style="color: #4CAF50;">↗ +12.3%</span>
                </div>

                <div class="metric-card" style="background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <h3 style="margin: 0; color: #666;">Conversion Rate</h3>
                    <p style="font-size: 2em; font-weight: bold; color: #333; margin: 10px 0;">3.4%</p>
                    <span style="color: #f44336;">↘ -0.8%</span>
                </div>
            </div>

            <div style="background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                <h3 style="margin: 0 0 20px 0; color: #333;">Quick Actions</h3>
                <div style="display: flex; gap: 10px; flex-wrap: wrap;">
                    <button onclick="alert('Export feature coming soon!')"
                            style="padding: 10px 20px; background: #2196F3; color: white; border: none; border-radius: 4px; cursor: pointer;">
                        Export Data
                    </button>
                    <button onclick="alert('Refresh complete!')"
                            style="padding: 10px 20px; background: #4CAF50; color: white; border: none; border-radius: 4px; cursor: pointer;">
                        Refresh
                    </button>
                    <button onclick="alert('Settings panel opened!')"
                            style="padding: 10px 20px; background: #FF9800; color: white; border: none; border-radius: 4px; cursor: pointer;">
                        Settings
                    </button>
                </div>
            </div>
        </div>
        """

        artifact = ArtifactComponent(
            content=dashboard_content,
            artifact_type="dashboard",
            title="Analytics Dashboard",
            description="A sample analytics dashboard with metrics and controls",
            external_renderable=True,
            fullscreen_capable=True,
        )

        yield UiComponent(rich_component=artifact)


def create_demo_agent() -> ArtifactDemoAgent:
    """Create a demo agent for REPL and server usage.

    Returns:
        Configured ArtifactDemoAgent instance
    """
    return ArtifactDemoAgent()


async def main() -> None:
    """Main demo function."""
    print("🎨 Artifact Demo Agent")
    print("This demo shows how to create different types of artifacts.")
    print(
        "In a real web application, developers can listen for 'artifact-opened' events."
    )
    print()

    demo_agent = create_demo_agent()
    user = User(id="demo_user", username="artifact_demo")

    # Demo 1: HTML Artifact
    print("1. Creating HTML Artifact...")
    async for component in demo_agent.create_html_artifact():
        artifact = component.rich_component
        if isinstance(artifact, ArtifactComponent):
            print(f"   ✓ Created HTML artifact: {artifact.title}")
            print(f"   ✓ Artifact ID: {artifact.artifact_id}")
            print(f"   ✓ Type: {artifact.artifact_type}")
            print(f"   ✓ External renderable: {artifact.external_renderable}")
        print()

    # Demo 2: D3 Visualization
    print("2. Creating D3.js Visualization...")
    async for component in demo_agent.create_d3_visualization():
        artifact = component.rich_component
        if isinstance(artifact, ArtifactComponent):
            print(f"   ✓ Created D3 artifact: {artifact.title}")
            print(f"   ✓ Dependencies: {artifact.dependencies}")
            print(f"   ✓ Standalone HTML available via get_standalone_html()")
        print()

    # Demo 3: Dashboard
    print("3. Creating Dashboard Artifact...")
    async for component in demo_agent.create_dashboard_artifact():
        artifact = component.rich_component
        if isinstance(artifact, ArtifactComponent):
            print(f"   ✓ Created dashboard artifact: {artifact.title}")
            print(f"   ✓ Fullscreen capable: {artifact.fullscreen_capable}")
        print()

    print("🚀 Web Integration Example:")
    print("""
    In your web application, listen for the 'artifact-opened' event:

    document.querySelector('vanna-chat').addEventListener('artifact-opened', (event) => {
        const { artifactId, content, type, trigger } = event.detail;

        if (trigger === 'created' && type === 'dashboard') {
            // Auto-open dashboards in external window
            const newWindow = window.open('', '_blank');
            newWindow.document.write(event.detail.getStandaloneHTML());
            newWindow.document.close();

            // Prevent default rendering in chat
            event.detail.preventDefault();
        }
    });
    """)


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/claude_sqlite_example.py
================================================
"""
Claude example using the SQL query tool with the Chinook database.

This example demonstrates using the RunSqlTool with SqliteRunner and Claude's AI
to intelligently query and analyze the Chinook database, with automatic visualization support.

Requirements:
- ANTHROPIC_API_KEY environment variable or .env file
- anthropic package: pip install -e .[anthropic]
- plotly package: pip install -e .[visualization]

Usage:
  PYTHONPATH=. python vanna/examples/claude_sqlite_example.py
"""

import asyncio
import importlib.util
import os
import sys
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from vanna import Agent


def ensure_env() -> None:
    if importlib.util.find_spec("dotenv") is not None:
        from dotenv import load_dotenv

        # Load from local .env without overriding existing env
        load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"), override=False)
    else:
        print(
            "[warn] python-dotenv not installed; skipping .env load. Install with: pip install python-dotenv"
        )

    if not os.getenv("ANTHROPIC_API_KEY"):
        print(
            "[error] ANTHROPIC_API_KEY is not set. Add it to your environment or .env file."
        )
        sys.exit(1)


async def main() -> None:
    ensure_env()

    try:
        from vanna.integrations.anthropic import AnthropicLlmService
    except ImportError:
        print(
            "[error] anthropic extra not installed. Install with: pip install -e .[anthropic]"
        )
        raise

    from vanna import AgentConfig, Agent
    from vanna.core.registry import ToolRegistry
    from vanna.core.user import CookieEmailUserResolver, RequestContext
    from vanna.integrations.sqlite import SqliteRunner
    from vanna.tools import (
        RunSqlTool,
        VisualizeDataTool,
        LocalFileSystem,
    )

    # Get the path to the Chinook database
    database_path = os.path.join(
        os.path.dirname(__file__), "..", "..", "Chinook.sqlite"
    )
    database_path = os.path.abspath(database_path)

    if not os.path.exists(database_path):
        print(f"[error] Chinook database not found at {database_path}")
        print(
            "Please download it with: curl -o Chinook.sqlite https://vanna.ai/Chinook.sqlite"
        )
        sys.exit(1)

    model = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")
    print(f"Using Anthropic model: {model}")
    print(f"Using database: {database_path}")

    llm = AnthropicLlmService(model=model)

    # Create shared FileSystem for both tools
    file_system = LocalFileSystem(working_directory="./claude_data")

    # Create tool registry and register the SQL tool with SQLite runner
    tool_registry = ToolRegistry()
    sqlite_runner = SqliteRunner(database_path=database_path)
    sql_tool = RunSqlTool(sql_runner=sqlite_runner, file_system=file_system)
    tool_registry.register(sql_tool)

    # Register visualization tool
    try:
        viz_tool = VisualizeDataTool(file_system=file_system)
        tool_registry.register(viz_tool)
        print("Visualization tool enabled")
    except ImportError:
        print(
            "[warn] Plotly not installed. Visualization tool disabled. Install with: pip install -e .[visualization]"
        )

    user_resolver = CookieEmailUserResolver()

    agent = Agent(
        llm_service=llm,
        config=AgentConfig(stream_responses=False),
        tool_registry=tool_registry,
        user_resolver=user_resolver,
    )

    # Simulate a logged-in demo user via cookie-based resolver
    request_context = RequestContext(
        cookies={user_resolver.cookie_name: "demo-user@example.com"},
        metadata={"demo": True},
        remote_addr="127.0.0.1",
    )
    conversation_id = "claude-sqlite-demo"

    # Sample queries to demonstrate different capabilities
    sample_questions = [
        "What tables are in this database?",
        "Show me the first 5 customers with their names",
        "What's the total number of tracks in the database?",
        "Find the top 5 artists by number of albums",
        "What's the average invoice total?",
        "Get data on the top 10 longest tracks and then visualize it",
    ]

    print("\n" + "=" * 60)
    print("Claude SQLite Database Assistant Demo")
    print("=" * 60)
    print("This demo shows Claude querying the Chinook music database.")
    print("Claude will intelligently construct SQL queries to answer questions")
    print("and can create visualizations of the results.")
    print()

    for i, question in enumerate(sample_questions, 1):
        print(f"\n--- Question {i}: {question} ---")

        async for component in agent.send_message(
            request_context=request_context,
            message=question,
            conversation_id=conversation_id,
        ):
            # Handle different component types
            if hasattr(component, "simple_component") and component.simple_component:
                if hasattr(component.simple_component, "text"):
                    print("Assistant:", component.simple_component.text)
            elif hasattr(component, "rich_component") and component.rich_component:
                if (
                    hasattr(component.rich_component, "content")
                    and component.rich_component.content
                ):
                    print("Assistant:", component.rich_component.content)
            elif hasattr(component, "content") and component.content:
                print("Assistant:", component.content)

        print()  # Add spacing between questions

    print("\n" + "=" * 60)
    print("Demo complete! Claude successfully queried the database.")
    print("=" * 60)


def create_demo_agent() -> "Agent":
    """Create a demo agent with Claude and SQLite query tool.

    This function is called by the vanna server framework.

    Returns:
        Configured Agent with Claude LLM and SQLite tool
    """
    ensure_env()

    try:
        from vanna.integrations.anthropic import AnthropicLlmService
    except ImportError:
        print(
            "[error] anthropic extra not installed. Install with: pip install -e .[anthropic]"
        )
        raise

    from vanna import AgentConfig, Agent
    from vanna.core.registry import ToolRegistry
    from vanna.core.user import CookieEmailUserResolver
    from vanna.integrations.sqlite import SqliteRunner
    from vanna.tools import (
        RunSqlTool,
        VisualizeDataTool,
        LocalFileSystem,
    )

    # Get the path to the Chinook database
    database_path = os.path.join(
        os.path.dirname(__file__), "..", "..", "Chinook.sqlite"
    )
    database_path = os.path.abspath(database_path)

    if not os.path.exists(database_path):
        raise FileNotFoundError(
            f"Chinook database not found at {database_path}. Please download it from https://vanna.ai/Chinook.sqlite"
        )

    model = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")

    llm = AnthropicLlmService(model=model)

    # Create shared FileSystem for both tools
    file_system = LocalFileSystem(working_directory="./claude_data")

    # Create tool registry and register the SQL tool with SQLite runner
    tool_registry = ToolRegistry()
    sqlite_runner = SqliteRunner(database_path=database_path)
    sql_tool = RunSqlTool(sql_runner=sqlite_runner, file_system=file_system)
    tool_registry.register(sql_tool)

    # Register visualization tool if available
    try:
        viz_tool = VisualizeDataTool(file_system=file_system)
        tool_registry.register(viz_tool)
    except ImportError:
        pass  # Visualization tool not available

    user_resolver = CookieEmailUserResolver()

    return Agent(
        llm_service=llm,
        config=AgentConfig(stream_responses=True),  # Enable streaming for web interface
        tool_registry=tool_registry,
        user_resolver=user_resolver,
    )


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/coding_agent_example.py
================================================
"""
Example coding agent using the vanna-agents framework.

This example demonstrates building an agent that can edit code files,
following the concepts from the "How to Build an Agent" article.
The agent includes tools for file operations and uses an LLM service
that can understand and modify code.

Usage:
  PYTHONPATH=. python vanna/examples/coding_agent_example.py
"""

import asyncio
import uuid
from typing import AsyncGenerator, List, Optional

from vanna import (
    AgentConfig,
    Agent,
    ToolRegistry,
    User,
)
from vanna.core.interfaces import LlmService
from vanna.core.models import (
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
    ToolCall,
    ToolSchema,
)
from vanna.tools.file_system import create_file_system_tools
from vanna.tools.python import create_python_tools


class CodingLlmService(LlmService):
    """
    LLM service that simulates a coding assistant.

    This demonstrates the minimal implementation needed for an agent
    as described in the article - just needs to understand tool calls
    and respond appropriately.
    """

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Handle non-streaming requests."""
        await asyncio.sleep(0.1)  # Simulate thinking time
        return self._build_response(request)

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Handle streaming requests."""
        await asyncio.sleep(0.1)
        response = self._build_response(request)

        if response.tool_calls:
            yield LlmStreamChunk(tool_calls=response.tool_calls)
        if response.content:
            # Simulate streaming by chunking the response
            words = response.content.split()
            for i, word in enumerate(words):
                chunk = word if i == 0 else f" {word}"
                await asyncio.sleep(0.05)  # Simulate streaming delay
                yield LlmStreamChunk(content=chunk)

        yield LlmStreamChunk(finish_reason=response.finish_reason)

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Validate tools - no errors for this simple implementation."""
        return []

    def _build_response(self, request: LlmRequest) -> LlmResponse:
        """Build a response based on the conversation context."""
        last_message = request.messages[-1] if request.messages else None

        # If we just got a tool result, respond to it
        if last_message and last_message.role == "tool":
            tool_result = last_message.content or "Tool executed"
            return LlmResponse(
                content=f"I've completed the operation. {tool_result}",
                finish_reason="stop",
            )

        # If user is asking for file operations, use tools
        if last_message and last_message.role == "user":
            user_message = last_message.content.lower()

            if "list files" in user_message or "show files" in user_message:
                return LlmResponse(
                    content="I'll list the files for you.",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{uuid.uuid4().hex[:8]}",
                            name="list_files",
                            arguments={},
                        )
                    ],
                    finish_reason="tool_calls",
                )

            elif "read" in user_message and (
                "file" in user_message
                or ".py" in user_message
                or ".txt" in user_message
            ):
                filename = _extract_filename(user_message)

                if filename:
                    return LlmResponse(
                        content=f"I'll read the file '{filename}' for you.",
                        tool_calls=[
                            ToolCall(
                                id=f"call_{uuid.uuid4().hex[:8]}",
                                name="read_file",
                                arguments={"filename": filename},
                            )
                        ],
                        finish_reason="tool_calls",
                    )

            elif "create" in user_message or "write" in user_message:
                # Suggest creating a simple example file
                return LlmResponse(
                    content="I'll create an example Python file for you.",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{uuid.uuid4().hex[:8]}",
                            name="write_file",
                            arguments={
                                "filename": "example.py",
                                "content": "# Example Python file\nprint('Hello from the coding agent!')\n\ndef greet(name):\n    return f'Hello, {name}!'\n\nif __name__ == '__main__':\n    print(greet('World'))\n",
                                "overwrite": True,
                            },
                        )
                    ],
                    finish_reason="tool_calls",
                )

            elif (
                "run" in user_message or "execute" in user_message
            ) and ".py" in user_message:
                filename = _extract_filename(user_message)
                if filename:
                    return LlmResponse(
                        content=f"I'll run the Python file '{filename}'.",
                        tool_calls=[
                            ToolCall(
                                id=f"call_{uuid.uuid4().hex[:8]}",
                                name="run_python_file",
                                arguments={
                                    "filename": filename,
                                    "arguments": [],
                                },
                            )
                        ],
                        finish_reason="tool_calls",
                    )

            elif (
                "edit" in user_message
                or "update" in user_message
                or "modify" in user_message
            ):
                return LlmResponse(
                    content="I'll update the greet function to make it more descriptive.",
                    tool_calls=[
                        ToolCall(
                            id=f"call_{uuid.uuid4().hex[:8]}",
                            name="edit_file",
                            arguments={
                                "filename": "example.py",
                                "edits": [
                                    {
                                        "start_line": 4,
                                        "end_line": 5,
                                        "new_content": (
                                            "def greet(name):\n"
                                            '    """Return a friendly greeting."""\n'
                                            '    return f"Hello, {name}! Welcome to the coding agent."\n'
                                        ),
                                    }
                                ],
                            },
                        )
                    ],
                    finish_reason="tool_calls",
                )

        # Default response
        return LlmResponse(
            content=(
                "I'm a coding assistant. I can help you list, read, write, edit, and run Python files. "
                "Try asking me to 'list files', 'read example.py', 'create a Python file', 'run example.py', or 'update example.py'."
            ),
            finish_reason="stop",
        )


def create_demo_agent() -> Agent:
    """
    Create a coding agent with file operation tools.

    This follows the pattern from the article - minimal code
    to create a powerful code-editing agent. Uses dependency injection
    for file system operations with LocalFileSystem as default.
    """
    # Create tool registry and register file system tools
    tool_registry = ToolRegistry()

    # Use the convenience function to create tools with default LocalFileSystem
    for tool in create_file_system_tools():
        tool_registry.register(tool)

    for tool in create_python_tools():
        tool_registry.register(tool)

    # Create LLM service
    llm_service = CodingLlmService()

    # Create agent with configuration
    return Agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        config=AgentConfig(
            stream_responses=True,
            include_thinking_indicators=True,
            max_tool_iterations=3,
        ),
    )


async def main() -> None:
    """
    Demonstrate the coding agent in action.

    As the article mentions: "300 lines of code and three tools and now
    you're able to talk to an alien intelligence that edits your code."
    """
    print("🤖 Starting Coding Agent Demo")
    print("This demonstrates the concepts from 'How to Build an Agent'")
    print("-" * 50)

    # Create the agent
    agent = create_demo_agent()

    # Create a test user
    user = User(id="coder123", username="developer", permissions=[])

    # Show available tools
    tools = await agent.get_available_tools(user)
    print(f"Available tools: {[tool.name for tool in tools]}")
    print()

    # Demo conversation
    conversation_id = "coding-session"

    demos = [
        "Hello! Can you list the files in this directory?",
        "Can you create a simple Python file for me?",
        "Now read the example.py file you just created",
        "Please update the greet function to include a docstring and a friendlier message.",
        "Run example.py so I can see its output.",
        "Great, read example.py again to confirm the changes.",
    ]

    for i, message in enumerate(demos, 1):
        print(f"Demo {i}: {message}")
        print("Agent response:")

        async for component in agent.send_message(
            user=user, message=message, conversation_id=conversation_id
        ):
            if (
                hasattr(component.rich_component, "content")
                and component.rich_component.content
            ):
                print(f"  📝 {component.rich_component.content}")
            elif hasattr(component.rich_component, "message"):
                print(f"  💬 {component.rich_component.message}")
            elif component.simple_component and hasattr(
                component.simple_component, "text"
            ):
                print(f"  📄 {component.simple_component.text}")

        print("-" * 30)


def _extract_filename(message: str) -> Optional[str]:
    """Extract a likely filename token from a user message."""

    for token in message.replace("\n", " ").split():
        cleaned = token.strip("'\".,;!?")
        if "." in cleaned and not cleaned.startswith("."):
            return cleaned

    return None


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/custom_system_prompt_example.py
================================================
"""
Example demonstrating custom system prompt builder with dependency injection.

This example shows how to create a custom SystemPromptBuilder that dynamically
generates system prompts based on user context and available tools.

Usage:
  python -m vanna.examples.custom_system_prompt_example
"""

from typing import List, Optional

from vanna.core.interfaces import SystemPromptBuilder
from vanna.core.models import ToolSchema, User


class CustomSystemPromptBuilder(SystemPromptBuilder):
    """Custom system prompt builder that personalizes prompts based on user."""

    async def build_system_prompt(
        self, user: User, tools: List[ToolSchema]
    ) -> Optional[str]:
        """Build a personalized system prompt.

        Args:
            user: The user making the request
            tools: List of tools available to the user

        Returns:
            Personalized system prompt
        """
        # Build personalized greeting
        username = user.username or user.id
        greeting = f"Hello {username}! I'm your AI assistant."

        # Add role-specific instructions based on user permissions
        role_instructions = []
        if "admin" in user.permissions:
            role_instructions.append(
                "As an admin user, you have access to all tools and capabilities."
            )
        elif "analyst" in user.permissions:
            role_instructions.append(
                "You're working as an analyst. I'll help you query and visualize data."
            )
        else:
            role_instructions.append("I'm here to help you with your tasks.")

        # List available tools
        tool_info = []
        if tools:
            tool_info.append("\nAvailable tools:")
            for tool in tools:
                tool_info.append(f"- {tool.name}: {tool.description}")

        # Combine all parts
        parts = [greeting] + role_instructions + tool_info
        return "\n".join(parts)


class SQLAssistantSystemPromptBuilder(SystemPromptBuilder):
    """System prompt builder specifically for SQL database assistants."""

    def __init__(self, database_name: str = "database"):
        """Initialize with database context.

        Args:
            database_name: Name of the database being queried
        """
        self.database_name = database_name

    async def build_system_prompt(
        self, user: User, tools: List[ToolSchema]
    ) -> Optional[str]:
        """Build a SQL-focused system prompt.

        Args:
            user: The user making the request
            tools: List of tools available to the user

        Returns:
            SQL-focused system prompt
        """
        prompt = f"""You are an expert SQL database assistant for the {self.database_name} database.

Your primary responsibilities:
1. Write efficient, correct SQL queries
2. Explain query results clearly
3. Suggest optimizations when relevant
4. Visualize data when appropriate

Guidelines:
- Always validate SQL syntax before execution
- Use appropriate JOINs and avoid Cartesian products
- Limit result sets to reasonable sizes by default
- Format numbers and dates for readability
"""

        # Add tool-specific instructions
        has_viz_tool = any(tool.name == "visualize_data" for tool in tools)
        if has_viz_tool:
            prompt += "\n- Create visualizations for numerical data when it helps understanding"

        return prompt


async def demo() -> None:
    """Demonstrate custom system prompt builders."""
    from vanna import Agent, User
    from vanna.core.registry import ToolRegistry
    from vanna.integrations.anthropic.mock import MockLlmService

    # Example 1: Custom personalized system prompt
    print("=" * 60)
    print("Example 1: Custom Personalized System Prompt")
    print("=" * 60)

    custom_builder = CustomSystemPromptBuilder()
    admin_user = User(id="user-1", username="Alice", permissions=["admin"])

    # Simulate some tools
    mock_tools = [
        ToolSchema(
            name="query_database", description="Query the SQL database", parameters={}
        ),
        ToolSchema(
            name="visualize_data",
            description="Create data visualizations",
            parameters={},
        ),
    ]

    prompt = await custom_builder.build_system_prompt(admin_user, mock_tools)
    print("\nGenerated system prompt for admin user:")
    print("-" * 60)
    print(prompt)
    print("-" * 60)

    # Example 2: SQL-specific system prompt
    print("\n" + "=" * 60)
    print("Example 2: SQL Assistant System Prompt")
    print("=" * 60)

    sql_builder = SQLAssistantSystemPromptBuilder(database_name="Chinook")
    analyst_user = User(id="user-2", username="Bob", permissions=["analyst"])

    prompt = await sql_builder.build_system_prompt(analyst_user, mock_tools)
    print("\nGenerated system prompt for SQL assistant:")
    print("-" * 60)
    print(prompt)
    print("-" * 60)

    # Example 3: Using custom builder with Agent
    print("\n" + "=" * 60)
    print("Example 3: Using Custom Builder with Agent")
    print("=" * 60)

    mock_llm = MockLlmService()
    tool_registry = ToolRegistry()

    agent = Agent(
        llm_service=mock_llm,
        tool_registry=tool_registry,
        system_prompt_builder=sql_builder,  # Inject custom builder here
    )

    print("\nAgent created with custom SQL system prompt builder!")
    print("The agent will now use the SQL-focused system prompt for all interactions.")


if __name__ == "__main__":
    import asyncio

    asyncio.run(demo())


================================================
FILE: src/vanna/examples/default_workflow_handler_example.py
================================================
"""
Example demonstrating the DefaultWorkflowHandler with setup health checking.

This example shows how the DefaultWorkflowHandler provides intelligent starter UI
that adapts based on available tools and helps users understand their setup status.

Run:
  PYTHONPATH=. python vanna/examples/default_workflow_handler_example.py
"""

import asyncio

from vanna import (
    AgentConfig,
    Agent,
    MemoryConversationStore,
    MockLlmService,
    User,
    DefaultWorkflowHandler,
)
from vanna.core.registry import ToolRegistry
from vanna.core.user.resolver import SimpleUserResolver
from vanna.tools import ListFilesTool


async def demonstrate_setup_scenarios():
    """Demonstrate different setup scenarios with DefaultWorkflowHandler."""
    print("🚀 Starting DefaultWorkflowHandler Setup Health Check Demo\n")

    # Create basic components
    llm_service = MockLlmService(response_content="I'm ready to help!")
    conversation_store = MemoryConversationStore()
    user_resolver = SimpleUserResolver()

    # Create test user
    user = User(
        id="user1",
        username="alice",
        email="alice@example.com",
        group_memberships=["user"],
    )

    print("=" * 60)
    print("SCENARIO 1: Empty Setup (No Tools)")
    print("=" * 60)

    # Empty tool registry
    empty_registry = ToolRegistry()

    agent_empty = Agent(
        llm_service=llm_service,
        tool_registry=empty_registry,
        user_resolver=user_resolver,
        conversation_store=conversation_store,
        config=AgentConfig(stream_responses=False),
        workflow_handler=DefaultWorkflowHandler(),
    )

    print("📋 Starter UI for empty setup:")
    async for component in agent_empty.send_message(
        request_context=user_resolver.create_request_context(
            metadata={"starter_ui_request": True}
        ),
        message="",
        conversation_id="empty-setup",
    ):
        if hasattr(component, "simple_component") and component.simple_component:
            print(f"  📄 {component.simple_component.text[:100]}...")
        elif hasattr(component, "rich_component"):
            comp = component.rich_component
            if hasattr(comp, "title"):
                print(f"  📊 {comp.title}: {comp.status} - {comp.description}")
            elif hasattr(comp, "content"):
                print(f"  📝 {comp.content[:100]}...")

    print("\n" + "=" * 60)
    print("SCENARIO 2: Functional Setup (SQL + Basic Tools)")
    print("=" * 60)

    # Tool registry with SQL tool (simulated)
    functional_registry = ToolRegistry()

    # Register a mock SQL tool (we'll simulate by tool name)
    list_tool = ListFilesTool()
    list_tool.name = "run_sql"  # Simulate SQL tool
    functional_registry.register(list_tool)

    agent_functional = Agent(
        llm_service=llm_service,
        tool_registry=functional_registry,
        user_resolver=user_resolver,
        conversation_store=conversation_store,
        config=AgentConfig(stream_responses=False),
        workflow_handler=DefaultWorkflowHandler(),
    )

    print("📋 Starter UI for functional setup:")
    async for component in agent_functional.send_message(
        request_context=user_resolver.create_request_context(
            metadata={"starter_ui_request": True}
        ),
        message="",
        conversation_id="functional-setup",
    ):
        if hasattr(component, "simple_component") and component.simple_component:
            print(f"  📄 {component.simple_component.text[:100]}...")
        elif hasattr(component, "rich_component"):
            comp = component.rich_component
            if hasattr(comp, "title"):
                print(f"  📊 {comp.title}: {comp.status} - {comp.description}")
            elif hasattr(comp, "content"):
                print(f"  📝 {comp.content[:100]}...")

    print("\n" + "=" * 60)
    print("SCENARIO 3: Complete Setup (SQL + Memory + Visualization)")
    print("=" * 60)

    # Complete tool registry
    complete_registry = ToolRegistry()

    # Mock SQL tool
    sql_tool = ListFilesTool()
    sql_tool.name = "run_sql"
    complete_registry.register(sql_tool)

    # Mock memory tools
    search_tool = ListFilesTool()
    search_tool.name = "search_saved_correct_tool_uses"
    complete_registry.register(search_tool)

    save_tool = ListFilesTool()
    save_tool.name = "save_question_tool_args"
    complete_registry.register(save_tool)

    # Mock visualization tool
    viz_tool = ListFilesTool()
    viz_tool.name = "visualize_data"
    complete_registry.register(viz_tool)

    agent_complete = Agent(
        llm_service=llm_service,
        tool_registry=complete_registry,
        user_resolver=user_resolver,
        conversation_store=conversation_store,
        config=AgentConfig(stream_responses=False),
        workflow_handler=DefaultWorkflowHandler(),
    )

    print("📋 Starter UI for complete setup:")
    async for component in agent_complete.send_message(
        request_context=user_resolver.create_request_context(
            metadata={"starter_ui_request": True}
        ),
        message="",
        conversation_id="complete-setup",
    ):
        if hasattr(component, "simple_component") and component.simple_component:
            print(f"  📄 {component.simple_component.text[:100]}...")
        elif hasattr(component, "rich_component"):
            comp = component.rich_component
            if hasattr(comp, "title"):
                print(f"  📊 {comp.title}: {comp.status} - {comp.description}")
            elif hasattr(comp, "content"):
                print(f"  📝 {comp.content[:100]}...")

    print("\n" + "=" * 60)
    print("SCENARIO 4: Testing Commands")
    print("=" * 60)

    print("📋 Testing /help command:")
    async for component in agent_complete.send_message(
        request_context=user_resolver.create_request_context(),
        message="/help",
        conversation_id="help-test",
    ):
        if hasattr(component, "rich_component") and hasattr(
            component.rich_component, "content"
        ):
            print(f"  📝 Help: {component.rich_component.content[:200]}...")

    print("\n📋 Testing /status command:")
    async for component in agent_complete.send_message(
        request_context=user_resolver.create_request_context(),
        message="/status",
        conversation_id="status-test",
    ):
        if hasattr(component, "rich_component"):
            comp = component.rich_component
            if hasattr(comp, "title"):
                print(f"  📊 {comp.title}: {comp.status}")
            elif hasattr(comp, "content"):
                print(f"  📝 Status: {comp.content[:150]}...")

    print("\n✅ Demo complete! The DefaultWorkflowHandler provides:")
    print("  • Smart setup health checking")
    print("  • Contextual starter UI based on available tools")
    print("  • Helpful error messages and setup guidance")
    print("  • Built-in command handling (/help, /status)")
    print("  • Automatic tool analysis and recommendations")


async def main():
    """Run the DefaultWorkflowHandler demonstration."""
    await demonstrate_setup_scenarios()


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/email_auth_example.py
================================================
"""
Email authentication example for the Vanna Agents framework.

This example demonstrates how to create an agent with email-based authentication
where users are prompted for their email address in chat and the system creates
a user profile based on that email.

## What This Example Shows

1. **UserService Implementation**: A demo `DemoEmailUserService` that:
   - Stores users in memory
   - Authenticates users by email validation
   - Creates user profiles automatically
   - Manages user permissions

2. **Authentication Tool**: An `AuthTool` that:
   - Takes an email address as input
   - Uses the UserService to authenticate/create users
   - Returns rich UI components for success/error feedback
   - Provides structured results for the LLM

3. **In-Chat Authentication Flow**: Shows how:
   - Users can provide their email in natural conversation
   - The agent can prompt for authentication when needed
   - Authentication results are displayed with rich UI components
   - The system maintains user context across conversations

## Key Components

- `DemoEmailUserService`: Implements the `UserService` interface
- `AuthTool`: Implements the `Tool` interface for authentication
- Rich UI components for authentication feedback
- Integration with the agent's tool registry and conversation store

## Usage

Interactive: python -m vanna.examples.email_auth_example

## Note

This example uses a simplified mock LLM that doesn't actually call tools.
In a real implementation with OpenAI or Anthropic, the LLM would automatically
detect email addresses in user messages and call the authenticate_user tool.

For production use, you would:
- Replace DemoEmailUserService with a database-backed implementation
- Add proper email validation and security measures
- Implement session management in the server layer
- Add proper error handling and rate limiting
"""

import asyncio
from typing import Any, Dict, Optional, Type

from pydantic import BaseModel, Field

from vanna import (
    AgentConfig,
    Agent,
    MemoryConversationStore,
    MockLlmService,
    User,
)
from vanna.core import Tool, UserService
from vanna.core import ToolContext, ToolResult
from vanna.core.registry import ToolRegistry
from vanna.core.components import UiComponent
from vanna.core import RichComponent


# Demo User Service Implementation
class DemoEmailUserService(UserService):
    """Demo user service that authenticates users by email."""

    def __init__(self):
        """Initialize with in-memory user store."""
        self._users: Dict[str, User] = {}  # user_id -> User
        self._email_to_id: Dict[str, str] = {}  # email -> user_id

    async def get_user(self, user_id: str) -> Optional[User]:
        """Get user by ID."""
        return self._users.get(user_id)

    async def authenticate(self, credentials: Dict[str, Any]) -> Optional[User]:
        """Authenticate user by email."""
        email = credentials.get("email")
        if not email or not self._is_valid_email(email):
            return None

        # Check if user exists
        user_id = self._email_to_id.get(email)
        if user_id:
            return self._users[user_id]

        # Create new user
        user_id = f"user_{len(self._users) + 1}"
        username = email.split("@")[0]

        user = User(
            id=user_id,
            username=username,
            email=email,
            permissions=["basic_user"],
            metadata={"auth_method": "email"},
        )

        self._users[user_id] = user
        self._email_to_id[email] = user_id
        return user

    async def has_permission(self, user: User, permission: str) -> bool:
        """Check if user has permission."""
        return permission in user.permissions

    def _is_valid_email(self, email: str) -> bool:
        """Simple email validation."""
        return "@" in email and "." in email.split("@")[1]


# Authentication Tool
class AuthArgs(BaseModel):
    """Arguments for authentication."""

    email: str = Field(description="User's email address")


class AuthTool(Tool[AuthArgs]):
    """Tool to authenticate users by email."""

    def __init__(self, user_service: DemoEmailUserService):
        self.user_service = user_service

    @property
    def name(self) -> str:
        return "authenticate_user"

    @property
    def description(self) -> str:
        return "Authenticate a user by their email address. Use this when the user provides an email."

    def get_args_schema(self) -> Type[AuthArgs]:
        return AuthArgs

    async def execute(self, context: ToolContext, args: AuthArgs) -> ToolResult:
        """Execute authentication."""
        user = await self.user_service.authenticate({"email": args.email})

        if user:
            success_msg = (
                f"✅ Welcome {user.username}! You're now authenticated as {user.email}"
            )

            auth_component = RichComponent(
                type="status_card",
                data={
                    "title": "Authentication Success",
                    "status": "success",
                    "description": success_msg,
                    "icon": "✅",
                    "metadata": {
                        "user_id": user.id,
                        "username": user.username,
                        "email": user.email,
                    },
                },
            )

            return ToolResult(
                success=True,
                result_for_llm=f"User successfully authenticated as {user.username} ({user.email}). They can now access personalized features.",
                ui_component=UiComponent(rich_component=auth_component),
            )
        else:
            error_msg = f"❌ Invalid email format: {args.email}"
            error_component = RichComponent(
                type="status_card",
                data={
                    "title": "Authentication Failed",
                    "status": "error",
                    "description": error_msg,
                    "icon": "❌",
                    "metadata": {"email": args.email},
                },
            )

            return ToolResult(
                success=False,
                result_for_llm=f"Authentication failed for {args.email}. Please provide a valid email address.",
                ui_component=UiComponent(rich_component=error_component),
                error=error_msg,
            )


def create_demo_agent() -> Agent:
    """Create a demo agent for REPL and server usage.

    Returns:
        Configured Agent instance with email authentication
    """
    return create_auth_agent()


def create_auth_agent() -> Agent:
    """Create agent with email authentication."""

    # Create user service
    user_service = DemoEmailUserService()

    # Use simple mock LLM - the system prompt will guide behavior
    llm_service = MockLlmService(
        response_content="Hello! I'm your AI assistant. To provide you with personalized help, I'll need your email address for authentication. Please share your email with me, and I'll use the authenticate_user tool to set up your profile."
    )

    # Create tool registry with auth tool
    tool_registry = ToolRegistry()
    auth_tool = AuthTool(user_service)
    tool_registry.register(auth_tool)

    # Create agent with authentication system prompt
    agent = Agent(
        llm_service=llm_service,
        config=AgentConfig(
            stream_responses=True,
            include_thinking_indicators=False,  # Cleaner output for demo
            system_prompt="""You are a helpful AI assistant with an email-based authentication system.

AUTHENTICATION BEHAVIOR:
1. When a user provides an email address in their message, immediately use the 'authenticate_user' tool
2. Look for emails in patterns like "my email is...", "I'm john@example.com", or any text with @ symbols  
3. If user isn't authenticated, politely ask for their email address to get started
4. After successful authentication, welcome them by name and offer personalized assistance
5. Be friendly and helpful throughout the process

Remember: Authentication is required for personalized features!""",
        ),
        tool_registry=tool_registry,
        conversation_store=MemoryConversationStore(),
    )

    return agent


async def demo_auth_flow():
    """Demonstrate the authentication flow with simple output."""
    agent = create_auth_agent()

    # Start with anonymous user
    user = User(id="anonymous", username="guest", email=None, permissions=[])
    conversation_id = "auth_demo_conv"

    print("=== Email Authentication Demo ===")
    print("This example shows how an agent can authenticate users via email in chat.")
    print("Note: This uses a simple mock LLM for demonstration purposes.\n")

    # Demo conversation
    print("🔹 Step 1: Initial greeting")
    print("User: Hello!")
    print("Agent: ", end="")

    async for component in agent.send_message(
        user=user, message="Hello!", conversation_id=conversation_id
    ):
        if (
            hasattr(component, "rich_component")
            and component.rich_component.type.value == "text"
        ):
            content = component.rich_component.data.get("content") or getattr(
                component.rich_component, "content", ""
            )
            if content:
                print(content)
                break

    print("\n" + "=" * 60)

    print("\n🔹 Step 2: User provides email for authentication")
    print("User: My email is alice@example.com")
    print("Agent: ", end="")

    # This should trigger the auth tool
    auth_shown = False
    async for component in agent.send_message(
        user=user,
        message="My email is alice@example.com",
        conversation_id=conversation_id,
    ):
        if hasattr(component, "rich_component"):
            rich_comp = component.rich_component
            if rich_comp.type.value == "status_card" and not auth_shown:
                status = rich_comp.data.get("status", "")
                desc = rich_comp.data.get("description", "")
                if status == "success":
                    auth_shown = True
                    print(f"🔐 {desc}")
                    break

    print("\n" + "=" * 60)

    print("\n🔹 Step 3: Post-authentication interaction")
    print("User: What can you help me with now?")
    print("Agent: ", end="")

    async for component in agent.send_message(
        user=user,
        message="What can you help me with now?",
        conversation_id=conversation_id,
    ):
        if (
            hasattr(component, "rich_component")
            and component.rich_component.type.value == "text"
        ):
            content = component.rich_component.data.get("content") or getattr(
                component.rich_component, "content", ""
            )
            if content:
                print(content)
                break

    print("\n" + "=" * 60)
    print("\n✅ Authentication demo complete!")
    print("\nKey Features Demonstrated:")
    print("• Email-based user authentication")
    print("• Tool-based authentication flow")
    print("• In-memory user storage and management")
    print("• Rich UI components for auth feedback")


async def main():
    """Run the authentication example."""
    await demo_auth_flow()


def run_interactive():
    """Entry point for interactive usage."""
    print("Starting email authentication example...")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/evaluation_example.py
================================================
"""
Evaluation System Example

This example demonstrates how to use the evaluation framework to test
and compare agents. Shows:
- Creating test cases programmatically
- Running evaluations with multiple evaluators
- Comparing agent variants (e.g., different LLMs)
- Generating reports

Usage:
    PYTHONPATH=. python vanna/examples/evaluation_example.py
"""

import asyncio
from vanna import Agent, MockLlmService, MemoryConversationStore, User
from vanna.core.evaluation import (
    EvaluationRunner,
    EvaluationDataset,
    TestCase,
    ExpectedOutcome,
    AgentVariant,
    TrajectoryEvaluator,
    OutputEvaluator,
    EfficiencyEvaluator,
)
from vanna.core.registry import ToolRegistry


def create_sample_dataset() -> EvaluationDataset:
    """Create a sample dataset for demonstration."""

    eval_user = User(
        id="eval_user", username="evaluator", email="eval@example.com", permissions=[]
    )

    test_cases = [
        TestCase(
            id="test_001",
            user=eval_user,
            message="Hello, how are you?",
            expected_outcome=ExpectedOutcome(
                final_answer_contains=["hello", "hi"],
                max_execution_time_ms=3000,
            ),
            metadata={"category": "greeting", "difficulty": "easy"},
        ),
        TestCase(
            id="test_002",
            user=eval_user,
            message="What can you help me with?",
            expected_outcome=ExpectedOutcome(
                final_answer_contains=["help", "assist"],
                max_execution_time_ms=3000,
            ),
            metadata={"category": "capabilities", "difficulty": "easy"},
        ),
        TestCase(
            id="test_003",
            user=eval_user,
            message="Explain quantum computing",
            expected_outcome=ExpectedOutcome(
                final_answer_contains=["quantum", "computing"],
                min_components=1,
                max_execution_time_ms=5000,
            ),
            metadata={"category": "explanation", "difficulty": "medium"},
        ),
    ]

    return EvaluationDataset(
        name="Demo Test Cases",
        test_cases=test_cases,
        description="Sample test cases for evaluation demo",
    )


def create_test_agent(name: str, response_content: str) -> Agent:
    """Create a test agent with mock LLM."""
    return Agent(
        llm_service=MockLlmService(response_content=response_content),
        tool_registry=ToolRegistry(),
        conversation_store=MemoryConversationStore(),
    )


async def demo_single_agent_evaluation():
    """Demonstrate evaluating a single agent."""
    print("\n" + "=" * 80)
    print("DEMO 1: Single Agent Evaluation")
    print("=" * 80 + "\n")

    # Create dataset
    dataset = create_sample_dataset()
    print(f"Loaded dataset: {dataset.name}")
    print(f"Test cases: {len(dataset.test_cases)}\n")

    # Create agent
    agent = create_test_agent(
        "test-agent",
        "Hello! I'm here to help you with various tasks including answering questions about topics like quantum computing.",
    )

    # Create evaluators
    evaluators = [
        TrajectoryEvaluator(),
        OutputEvaluator(),
        EfficiencyEvaluator(max_execution_time_ms=5000),
    ]

    # Run evaluation
    runner = EvaluationRunner(evaluators=evaluators, max_concurrency=5)
    print("Running evaluation...")
    report = await runner.run_evaluation(agent, dataset.test_cases)

    # Print results
    report.print_summary()

    # Show failures
    failures = report.get_failures()
    if failures:
        print("\nFailed test cases:")
        for result in failures:
            print(f"  - {result.test_case.id}: {result.test_case.message}")


async def demo_agent_comparison():
    """Demonstrate comparing multiple agent variants."""
    print("\n" + "=" * 80)
    print("DEMO 2: Agent Comparison (LLM Comparison Use Case)")
    print("=" * 80 + "\n")

    # Create dataset
    dataset = create_sample_dataset()
    print(f"Loaded dataset: {dataset.name}")
    print(f"Test cases: {len(dataset.test_cases)}\n")

    # Create agent variants
    variants = [
        AgentVariant(
            name="agent-v1",
            agent=create_test_agent(
                "v1",
                "Hi there! I can help you with many things including explaining complex topics like quantum computing.",
            ),
            metadata={"version": "1.0", "model": "mock-v1"},
        ),
        AgentVariant(
            name="agent-v2",
            agent=create_test_agent(
                "v2",
                "Hello! I'm your helpful assistant. I can assist with various tasks and explain topics like quantum computing in detail.",
            ),
            metadata={"version": "2.0", "model": "mock-v2"},
        ),
        AgentVariant(
            name="agent-v3",
            agent=create_test_agent(
                "v3",
                "Greetings! I'm designed to help you with a wide range of tasks, from simple questions to complex explanations about quantum computing and more.",
            ),
            metadata={"version": "3.0", "model": "mock-v3"},
        ),
    ]

    print(f"Created {len(variants)} agent variants:")
    for v in variants:
        print(f"  - {v.name}")
    print()

    # Create evaluators
    evaluators = [
        OutputEvaluator(),
        EfficiencyEvaluator(max_execution_time_ms=5000),
    ]

    # Run comparison
    runner = EvaluationRunner(evaluators=evaluators, max_concurrency=10)
    print(
        f"Running comparison ({len(variants)} variants × {len(dataset.test_cases)} test cases)..."
    )
    print("All variants running in parallel for maximum efficiency...\n")

    comparison = await runner.compare_agents(variants, dataset.test_cases)

    # Print results
    comparison.print_summary()

    # Show best variants
    print("Best Performing Variants:")
    print(f"  🏆 Best score: {comparison.get_best_variant('score')}")
    print(f"  ⚡ Fastest: {comparison.get_best_variant('speed')}")
    print(f"  ✅ Best pass rate: {comparison.get_best_variant('pass_rate')}")


async def demo_dataset_operations():
    """Demonstrate dataset creation and manipulation."""
    print("\n" + "=" * 80)
    print("DEMO 3: Dataset Operations")
    print("=" * 80 + "\n")

    # Create dataset
    dataset = create_sample_dataset()

    # Show dataset info
    print(f"Dataset: {dataset.name}")
    print(f"Description: {dataset.description}")
    print(f"Total test cases: {len(dataset)}\n")

    # Filter by metadata
    easy_tests = dataset.filter_by_metadata(difficulty="easy")
    medium_tests = dataset.filter_by_metadata(difficulty="medium")

    print(f"Easy test cases: {len(easy_tests)}")
    print(f"Medium test cases: {len(medium_tests)}\n")

    # Save to file (for demonstration)
    import tempfile
    import os

    with tempfile.TemporaryDirectory() as tmpdir:
        yaml_path = os.path.join(tmpdir, "dataset.yaml")
        json_path = os.path.join(tmpdir, "dataset.json")

        dataset.save_yaml(yaml_path)
        dataset.save_json(json_path)

        print("Dataset saved to temporary files:")
        print(f"  - YAML: {yaml_path}")
        print(f"  - JSON: {json_path}\n")

        # Load back
        loaded_yaml = EvaluationDataset.from_yaml(yaml_path)
        loaded_json = EvaluationDataset.from_json(json_path)

        print("Loaded datasets:")
        print(f"  - From YAML: {len(loaded_yaml)} test cases")
        print(f"  - From JSON: {len(loaded_json)} test cases")


async def main():
    """Run all evaluation demos."""
    print("\n🚀 Vanna Agents Evaluation System Demo")
    print("=" * 80)

    # Demo 1: Single agent evaluation
    await demo_single_agent_evaluation()

    # Demo 2: Agent comparison (main use case)
    await demo_agent_comparison()

    # Demo 3: Dataset operations
    await demo_dataset_operations()

    print("\n" + "=" * 80)
    print("✅ All demos completed!")
    print("=" * 80)
    print("\nKey Takeaways:")
    print("  1. Evaluations are integral to the Vanna package")
    print("  2. Parallel execution handles I/O-bound LLM calls efficiently")
    print("  3. Agent comparison is a first-class use case")
    print("  4. Multiple evaluators can be composed for comprehensive testing")
    print("  5. Reports can be exported to HTML, CSV, or printed to console")
    print("\nFor LLM comparison, see: evals/benchmarks/llm_comparison.py")
    print()


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/extensibility_example.py
================================================
"""
Comprehensive example demonstrating all extensibility interfaces.

This example shows how to use:
- LlmMiddleware for caching
- ErrorRecoveryStrategy for retry logic
- ToolContextEnricher for adding user preferences
- ConversationFilter for context window management
- ObservabilityProvider for monitoring
"""

import asyncio
import time
from typing import Any, Dict, List, Optional

from vanna.core import (
    Agent,
    LlmMiddleware,
    ErrorRecoveryStrategy,
    ToolContextEnricher,
    ConversationFilter,
    ObservabilityProvider,
    User,
    ToolContext,
    Conversation,
    Message,
    LlmRequest,
    LlmResponse,
    Span,
    Metric,
)
from vanna.core.recovery import RecoveryAction, RecoveryActionType
from vanna.core.registry import ToolRegistry


# 1. LlmMiddleware Example: Simple Caching
class CachingMiddleware(LlmMiddleware):
    """Cache LLM responses to reduce costs and latency."""

    def __init__(self) -> None:
        self.cache: Dict[str, LlmResponse] = {}
        self.hits = 0
        self.misses = 0

    def _compute_cache_key(self, request: LlmRequest) -> str:
        """Create cache key from request."""
        messages_str = str([(m.role, m.content) for m in request.messages])
        return f"{messages_str}:{request.temperature}"

    async def before_llm_request(self, request: LlmRequest) -> LlmRequest:
        """Check cache before sending request."""
        cache_key = self._compute_cache_key(request)
        if cache_key in self.cache:
            self.hits += 1
            print(f"[CACHE HIT] Cache stats: {self.hits} hits, {self.misses} misses")
        return request

    async def after_llm_response(
        self, request: LlmRequest, response: LlmResponse
    ) -> LlmResponse:
        """Cache the response."""
        cache_key = self._compute_cache_key(request)
        if cache_key not in self.cache:
            self.cache[cache_key] = response
            self.misses += 1
            print(f"[CACHE MISS] Caching response")
        return response


# 2. ErrorRecoveryStrategy Example: Exponential Backoff
class ExponentialBackoffStrategy(ErrorRecoveryStrategy):
    """Retry failed operations with exponential backoff."""

    def __init__(self, max_retries: int = 3) -> None:
        self.max_retries = max_retries

    async def handle_tool_error(
        self, error: Exception, context: ToolContext, attempt: int = 1
    ) -> RecoveryAction:
        """Retry tool errors with exponential backoff."""
        if attempt < self.max_retries:
            delay_ms = (2 ** (attempt - 1)) * 1000
            print(
                f"[RETRY] Tool failed, retrying in {delay_ms}ms (attempt {attempt}/{self.max_retries})"
            )
            return RecoveryAction(
                action=RecoveryActionType.RETRY,
                retry_delay_ms=delay_ms,
                message=f"Retrying after {delay_ms}ms",
            )

        print(f"[FAIL] Max retries exceeded for tool error: {error}")
        return RecoveryAction(
            action=RecoveryActionType.FAIL,
            message=f"Tool error after {self.max_retries} attempts: {str(error)}",
        )

    async def handle_llm_error(
        self, error: Exception, request: LlmRequest, attempt: int = 1
    ) -> RecoveryAction:
        """Retry LLM errors with backoff."""
        if attempt < self.max_retries:
            delay_ms = (2 ** (attempt - 1)) * 1000
            print(
                f"[RETRY] LLM failed, retrying in {delay_ms}ms (attempt {attempt}/{self.max_retries})"
            )
            return RecoveryAction(
                action=RecoveryActionType.RETRY,
                retry_delay_ms=delay_ms,
                message=f"Retrying LLM after {delay_ms}ms",
            )

        print(f"[FAIL] Max retries exceeded for LLM error: {error}")
        return RecoveryAction(
            action=RecoveryActionType.FAIL,
            message=f"LLM error after {self.max_retries} attempts: {str(error)}",
        )


# 3. ToolContextEnricher Example: Add User Preferences
class UserPreferencesEnricher(ToolContextEnricher):
    """Enrich context with user preferences."""

    def __init__(self) -> None:
        # Mock user preferences database
        self.preferences: Dict[str, Dict[str, Any]] = {
            "user123": {
                "timezone": "America/New_York",
                "language": "en",
                "theme": "dark",
            }
        }

    async def enrich_context(self, context: ToolContext) -> ToolContext:
        """Add user preferences to context."""
        prefs = self.preferences.get(context.user.id, {})
        context.metadata["user_preferences"] = prefs
        context.metadata["timezone"] = prefs.get("timezone", "UTC")
        print(f"[ENRICH] Added preferences for user {context.user.id}: {prefs}")
        return context


# 4. ConversationFilter Example: Context Window Management
class ContextWindowFilter(ConversationFilter):
    """Limit conversation to fit within context window."""

    def __init__(self, max_messages: int = 20) -> None:
        self.max_messages = max_messages

    async def filter_messages(self, messages: List[Message]) -> List[Message]:
        """Keep only recent messages within limit."""
        if len(messages) <= self.max_messages:
            return messages

        # Keep system messages and recent messages
        system_messages = [m for m in messages if m.role == "system"]
        other_messages = [m for m in messages if m.role != "system"]

        # Take the most recent messages
        recent_messages = other_messages[-self.max_messages :]
        filtered = system_messages + recent_messages

        print(f"[FILTER] Reduced {len(messages)} messages to {len(filtered)}")
        return filtered


# 5. ObservabilityProvider Example: Simple Logging
class LoggingObservabilityProvider(ObservabilityProvider):
    """Log metrics and spans for monitoring."""

    def __init__(self) -> None:
        self.metrics: List[Metric] = []
        self.spans: List[Span] = []

    async def record_metric(
        self,
        name: str,
        value: float,
        unit: str = "",
        tags: Optional[Dict[str, str]] = None,
    ) -> None:
        """Record and log a metric."""
        metric = Metric(name=name, value=value, unit=unit, tags=tags or {})
        self.metrics.append(metric)
        tags_str = ", ".join(f"{k}={v}" for k, v in (tags or {}).items())
        print(f"[METRIC] {name}: {value}{unit} {tags_str}")

    async def create_span(
        self, name: str, attributes: Optional[Dict[str, Any]] = None
    ) -> Span:
        """Create a span for tracing."""
        span = Span(name=name, attributes=attributes or {})
        print(f"[SPAN START] {name}")
        return span

    async def end_span(self, span: Span) -> None:
        """End and record a span."""
        span.end()
        self.spans.append(span)
        duration = span.duration_ms() or 0
        print(f"[SPAN END] {span.name}: {duration:.2f}ms")


async def run_example() -> None:
    """
    Example showing all extensibility interfaces working together.
    """
    from vanna.integrations.anthropic import AnthropicLlmService

    # Create all extensibility components
    caching_middleware = CachingMiddleware()
    retry_strategy = ExponentialBackoffStrategy(max_retries=3)
    preferences_enricher = UserPreferencesEnricher()
    context_filter = ContextWindowFilter(max_messages=20)
    observability = LoggingObservabilityProvider()

    # Mock conversation store
    class MockStore:
        async def get_conversation(self, cid: str, uid: str) -> Optional[Conversation]:
            return None

        async def create_conversation(
            self, cid: str, uid: str, title: str
        ) -> Conversation:
            return Conversation(
                id=cid, user_id=uid, messages=[Message(role="user", content=title)]
            )

        async def update_conversation(self, conv: Conversation) -> None:
            pass

        async def delete_conversation(self, cid: str, uid: str) -> bool:
            return False

        async def list_conversations(
            self, uid: str, limit: int = 50, offset: int = 0
        ) -> List[Conversation]:
            return []

    # Create agent with all extensibility components
    agent = Agent(
        llm_service=AnthropicLlmService(api_key="test-key"),
        tool_registry=ToolRegistry(),
        conversation_store=MockStore(),  # type: ignore
        llm_middlewares=[caching_middleware],
        error_recovery_strategy=retry_strategy,
        context_enrichers=[preferences_enricher],
        conversation_filters=[context_filter],
        observability_provider=observability,
    )

    print("✓ Agent created with all extensibility components:")
    print(f"  - LLM Middleware: {len(agent.llm_middlewares)} middlewares")
    print(f"  - Error Recovery: {type(agent.error_recovery_strategy).__name__}")
    print(f"  - Context Enrichers: {len(agent.context_enrichers)} enrichers")
    print(f"  - Conversation Filters: {len(agent.conversation_filters)} filters")
    print(f"  - Observability: {type(agent.observability_provider).__name__}")
    print("\n🎉 All extensibility interfaces integrated successfully!")


if __name__ == "__main__":
    asyncio.run(run_example())


================================================
FILE: src/vanna/examples/minimal_example.py
================================================
"""Minimal Claude + SQLite example ready for FastAPI."""

from __future__ import annotations

import os
from pathlib import Path

from vanna import AgentConfig, Agent
from vanna.core.registry import ToolRegistry
from vanna.integrations.anthropic import AnthropicLlmService
from vanna.integrations.sqlite import SqliteRunner
from vanna.integrations.local import LocalFileSystem
from vanna.tools import (
    RunSqlTool,
    # Visualization
    VisualizeDataTool,
    # Python execution
    RunPythonFileTool,
    PipInstallTool,
    # File system (for coding agents)
    SearchFilesTool,
    ListFilesTool,
    ReadFileTool,
    WriteFileTool,
)

_DB = Path(__file__).resolve().parents[2] / "Chinook.sqlite"


def create_demo_agent() -> Agent:
    # Load environment variables from .env file
    from dotenv import load_dotenv

    load_dotenv()

    llm = AnthropicLlmService(model=os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5"))

    # Shared file system for all tools
    file_system = LocalFileSystem("./claude_data")

    tools = ToolRegistry()

    # 1. Basic SQL agent - query databases
    tools.register(
        RunSqlTool(
            sql_runner=SqliteRunner(database_path=str(_DB)),
            file_system=file_system,
        )
    )

    # 2. Add visualization - create charts from data
    tools.register(VisualizeDataTool(file_system=file_system))

    # 3. Add Python execution - build dashboards with artifacts
    # tools.register(RunPythonFileTool(file_system=file_system))
    # tools.register(PipInstallTool(file_system=file_system))

    # 4. Full coding agent - read, write, search files
    # tools.register(SearchFilesTool(file_system=file_system))
    # tools.register(ListFilesTool(file_system=file_system))
    # tools.register(ReadFileTool(file_system=file_system))
    # tools.register(WriteFileTool(file_system=file_system))

    return Agent(
        llm_service=llm,
        tool_registry=tools,
    )


================================================
FILE: src/vanna/examples/mock_auth_example.py
================================================
"""
Mock authentication example to verify user resolution is working.

This example demonstrates the new UserResolver architecture where:
1. UserResolver is a required parameter of Agent
2. Agent.send_message() accepts RequestContext (not User directly)
3. The Agent resolves the user internally using the UserResolver

The agent uses an LLM middleware to inject user info into the response,
so we can verify the authentication is working correctly.

Usage:
    python -m vanna.examples.mock_auth_example
"""

from __future__ import annotations

import asyncio

from vanna import AgentConfig, Agent
from vanna.core.registry import ToolRegistry
from vanna.core.llm import LlmRequest, LlmResponse
from vanna.core.middleware import LlmMiddleware
from vanna.integrations.mock import MockLlmService
from vanna.core.user import CookieEmailUserResolver, RequestContext


class UserEchoMiddleware(LlmMiddleware):
    """Middleware that injects user email into LLM responses."""

    async def after_llm_response(
        self, request: LlmRequest, response: LlmResponse
    ) -> LlmResponse:
        """Inject user email into response."""
        # Extract user email from request user_id (which is set to user.id in the agent)
        user_id = request.user_id

        # Create a new response with user info
        new_content = f"Hello! You are authenticated as: {user_id}"

        return LlmResponse(
            content=new_content,
            finish_reason=response.finish_reason,
            usage=response.usage,
        )


def create_demo_agent() -> Agent:
    """Create a demo agent for server usage.

    Returns:
        Configured Agent instance with cookie-based authentication
    """
    # Create a mock LLM
    llm_service = MockLlmService(response_content="Mock response")

    # Empty tool registry
    tool_registry = ToolRegistry()

    # Cookie-based user resolver
    user_resolver = CookieEmailUserResolver(cookie_name="vanna_email")

    # User echo middleware
    middleware = UserEchoMiddleware()

    # Create agent with user resolver and middleware
    agent = Agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        user_resolver=user_resolver,
        llm_middlewares=[middleware],
        config=AgentConfig(
            stream_responses=True,
            include_thinking_indicators=False,
        ),
    )

    return agent


async def demo_authentication():
    """Demonstrate authentication with different request contexts."""
    agent = create_demo_agent()

    print("=== Mock Authentication Demo ===")
    print("This example verifies that user resolution is working correctly.\n")

    # Test 1: Request with email cookie
    print("🔹 Test 1: Authenticated user (alice@example.com)")
    request_context = RequestContext(
        cookies={"vanna_email": "alice@example.com"},
        headers={},
        remote_addr="127.0.0.1",
    )

    print(
        "Request context:",
        {
            "cookies": request_context.cookies,
            "headers": request_context.headers,
            "remote_addr": request_context.remote_addr,
        },
    )

    # Send message - Agent will resolve user internally
    agent_response = ""
    async for component in agent.send_message(
        request_context=request_context,
        message="Who am I?",
        conversation_id="test_conv_1",
    ):
        # Extract and display user info from the resolved user
        if hasattr(component, "rich_component"):
            rich = component.rich_component
            # Check if it's a text component
            if rich.type.value == "text":
                # Access content directly from the component (before serialization)
                if hasattr(rich, "content"):
                    agent_response = rich.content

    print(f"Agent response: {agent_response}")

    # Verify user was resolved by checking the conversation store
    user_resolver = agent.user_resolver
    resolved_user = await user_resolver.resolve_user(request_context)
    print(
        f"✅ Resolved user: {resolved_user.email} (username: {resolved_user.username}, id: {resolved_user.id})"
    )
    print(f"   Permissions: {resolved_user.permissions}")
    print(f"   Metadata: {resolved_user.metadata}")

    print("\n" + "=" * 60 + "\n")

    # Test 2: Request without email cookie (anonymous)
    print("🔹 Test 2: Anonymous user (no cookie)")
    anonymous_context = RequestContext(cookies={}, headers={}, remote_addr="127.0.0.1")

    print(
        "Request context:",
        {
            "cookies": anonymous_context.cookies,
            "headers": anonymous_context.headers,
            "remote_addr": anonymous_context.remote_addr,
        },
    )

    agent_response = ""
    async for component in agent.send_message(
        request_context=anonymous_context,
        message="Who am I?",
        conversation_id="test_conv_2",
    ):
        if hasattr(component, "rich_component"):
            rich = component.rich_component
            if rich.type.value == "text" and hasattr(rich, "content"):
                agent_response = rich.content

    print(f"Agent response: {agent_response}")

    resolved_user = await user_resolver.resolve_user(anonymous_context)
    print(
        f"✅ Resolved user: {resolved_user.email or 'None'} (username: {resolved_user.username}, id: {resolved_user.id})"
    )
    print(f"   Permissions: {resolved_user.permissions}")
    print(f"   Metadata: {resolved_user.metadata}")

    print("\n" + "=" * 60 + "\n")

    # Test 3: Different user
    print("🔹 Test 3: Different authenticated user (bob@company.com)")
    bob_context = RequestContext(
        cookies={"vanna_email": "bob@company.com"},
        headers={"User-Agent": "Mozilla/5.0"},
        remote_addr="192.168.1.100",
    )

    print(
        "Request context:",
        {
            "cookies": bob_context.cookies,
            "headers": bob_context.headers,
            "remote_addr": bob_context.remote_addr,
        },
    )

    agent_response = ""
    async for component in agent.send_message(
        request_context=bob_context, message="Who am I?", conversation_id="test_conv_3"
    ):
        if hasattr(component, "rich_component"):
            rich = component.rich_component
            if rich.type.value == "text" and hasattr(rich, "content"):
                agent_response = rich.content

    print(f"Agent response: {agent_response}")

    resolved_user = await user_resolver.resolve_user(bob_context)
    print(
        f"✅ Resolved user: {resolved_user.email} (username: {resolved_user.username}, id: {resolved_user.id})"
    )
    print(f"   Permissions: {resolved_user.permissions}")
    print(f"   Metadata: {resolved_user.metadata}")

    print("\n" + "=" * 60)
    print("\n✅ Authentication demo complete!")
    print("\nKey Features Verified:")
    print("• UserResolver is part of Agent")
    print("• Agent.send_message() accepts RequestContext")
    print("• User resolution happens internally in Agent")
    print("• CookieEmailUserResolver extracts email from vanna_email cookie")
    print("• Anonymous users are created when no cookie is present")
    print("• Different users can be resolved from different request contexts")


async def main():
    """Run the authentication example."""
    await demo_authentication()


def run_interactive():
    """Entry point for interactive usage."""
    print("Starting mock authentication example...")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/mock_custom_tool.py
================================================
"""
Mock example showing how to create and use custom tools.

This example demonstrates creating a simple calculator tool
and registering it with an agent that uses a mock LLM service.
It now includes a `MockCalculatorLlmService` that automatically
invokes the calculator tool with random numbers before echoing
back the computed answer.

Usage:
  Template: Copy this file and modify for your custom tools
  Interactive: python -m vanna.examples.mock_custom_tool
  REPL: from vanna.examples.mock_custom_tool import create_demo_agent
  Server: python -m vanna.servers --example mock_custom_tool
"""

import asyncio
import random
import uuid
from typing import AsyncGenerator, Dict, List, Optional, Tuple, Type

from pydantic import BaseModel, Field

from vanna import (
    AgentConfig,
    Agent,
    Tool,
    ToolContext,
    ToolRegistry,
    ToolResult,
    User,
    UiComponent,
)
from vanna.core.interfaces import LlmService
from vanna.core.models import (
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
    ToolCall,
    ToolSchema,
)
from vanna.core.rich_components import (
    CardComponent,
    NotificationComponent,
    ComponentType,
)
from vanna.core.simple_components import (
    SimpleTextComponent,
)


class CalculatorArgs(BaseModel):
    """Arguments for the calculator tool."""

    operation: str = Field(
        description="The operation to perform: add, subtract, multiply, divide"
    )
    a: float = Field(description="First number")
    b: float = Field(description="Second number")


class CalculatorTool(Tool[CalculatorArgs]):
    """A simple calculator tool."""

    @property
    def name(self) -> str:
        return "calculator"

    @property
    def description(self) -> str:
        return "Perform basic arithmetic operations (add, subtract, multiply, divide)"

    def get_args_schema(self) -> Type[CalculatorArgs]:
        return CalculatorArgs

    async def execute(self, context: ToolContext, args: CalculatorArgs) -> ToolResult:
        """Execute the calculator operation."""
        symbol_map = {"add": "+", "subtract": "-", "multiply": "×", "divide": "÷"}

        try:
            if args.operation == "add":
                result = args.a + args.b
            elif args.operation == "subtract":
                result = args.a - args.b
            elif args.operation == "multiply":
                result = args.a * args.b
            elif args.operation == "divide":
                if args.b == 0:
                    message = "Cannot divide by zero"
                    await asyncio.sleep(3)
                    return ToolResult(
                        success=False,
                        result_for_llm=message,
                        ui_component=UiComponent(
                            rich_component=NotificationComponent(
                                type=ComponentType.NOTIFICATION,
                                level="error",
                                message=message,
                            ),
                            simple_component=SimpleTextComponent(text=message),
                        ),
                        error=message,
                    )
                result = args.a / args.b
            else:
                message = f"Unknown operation: {args.operation}"
                await asyncio.sleep(3)
                return ToolResult(
                    success=False,
                    result_for_llm=message,
                    ui_component=UiComponent(
                        rich_component=NotificationComponent(
                            type=ComponentType.NOTIFICATION,
                            level="warning",
                            message=message,
                        ),
                        simple_component=SimpleTextComponent(text=message),
                    ),
                    error=message,
                )

            await asyncio.sleep(3)

            symbol = symbol_map.get(args.operation, args.operation)
            expression = f"{args.a:g} {symbol} {args.b:g} = {result:g}"
            return ToolResult(
                success=True,
                result_for_llm=str(result),
                ui_component=UiComponent(
                    rich_component=CardComponent(
                        type=ComponentType.CARD,
                        title="Calculator Result",
                        content=expression,
                    ),
                    simple_component=SimpleTextComponent(text=expression),
                ),
                error=None,
            )

        except Exception as e:
            message = str(e)
            await asyncio.sleep(3)
            return ToolResult(
                success=False,
                result_for_llm=message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=message,
                    ),
                    simple_component=SimpleTextComponent(text=message),
                ),
                error=message,
            )


class MockCalculatorLlmService(LlmService):
    """LLM service that exercises the calculator tool before echoing the result."""

    def __init__(self, seed: Optional[int] = None):
        self._random = random.Random(seed)

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Handle non-streaming calculator interactions."""
        await asyncio.sleep(0.05)
        return self._build_response(request)

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Provide streaming compatibility by yielding a single chunk."""
        await asyncio.sleep(0.05)
        response = self._build_response(request)

        if response.tool_calls:
            yield LlmStreamChunk(tool_calls=response.tool_calls)
        if response.content is not None:
            yield LlmStreamChunk(
                content=response.content, finish_reason=response.finish_reason
            )
        else:
            yield LlmStreamChunk(finish_reason=response.finish_reason)

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Mock validation - no errors."""
        return []

    def _build_response(self, request: LlmRequest) -> LlmResponse:
        """Create a response that either calls the tool or echoes its result."""
        last_message = request.messages[-1] if request.messages else None

        if last_message and last_message.role == "tool":
            answer = last_message.content or "No result provided"
            return LlmResponse(
                content=answer,
                finish_reason="stop",
                usage={
                    "prompt_tokens": 30,
                    "completion_tokens": 10,
                    "total_tokens": 40,
                },
            )

        operation, a, b = self._random_operands()
        tool_call = ToolCall(
            id=f"call_{uuid.uuid4().hex[:8]}",
            name="calculator",
            arguments={"operation": operation, "a": a, "b": b},
        )

        return LlmResponse(
            content="Let me ask my calculator friend for help...",
            tool_calls=[tool_call],
            finish_reason="tool_calls",
            usage={"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35},
        )

    def _random_operands(self) -> Tuple[str, float, float]:
        """Generate operation and operands suited for the calculator tool."""
        operation = self._random.choice(["add", "subtract", "multiply", "divide"])

        if operation == "divide":
            b = float(self._random.randint(1, 10))
            multiplier = self._random.randint(1, 10)
            a = float(b * multiplier)
        elif operation == "subtract":
            b = float(self._random.randint(1, 10))
            a = b + float(self._random.randint(0, 10))
        else:
            a = float(self._random.randint(1, 12))
            b = float(self._random.randint(1, 12))

        return operation, a, b


def create_demo_agent() -> Agent:
    """Create a demo agent with custom calculator tool.

    Returns:
        Configured Agent with calculator tool and mock calculator LLM
    """
    tool_registry = ToolRegistry()
    calculator_tool = CalculatorTool()
    tool_registry.register(calculator_tool)

    llm_service = MockCalculatorLlmService()

    return Agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        config=AgentConfig(
            stream_responses=False,
            include_thinking_indicators=False,
        ),
    )


async def main() -> None:
    """Run the mock custom tool example."""

    # Create agent using factory function
    agent = create_demo_agent()
    tool_registry = agent.tool_registry

    # Create a test user
    user = User(id="user123", username="testuser", permissions=[])

    # Test the tool directly
    print("Testing calculator tool directly:")
    tool_call = ToolCall(
        id="test123", name="calculator", arguments={"operation": "add", "a": 5, "b": 3}
    )

    context = ToolContext(user=user, conversation_id="test", request_id="test")

    result = await tool_registry.execute(tool_call, context)
    print(f"5 + 3 = {result.result_for_llm if result.success else result.error}")

    # Show available tools
    schemas = await tool_registry.get_schemas(user)
    print(f"\nAvailable tools for user: {[schema.name for schema in schemas]}")

    # Demonstrate the mock LLM triggering a tool call
    print("\nAgent conversation demo:")
    conversation_id = "calc-demo"
    async for component in agent.send_message(
        user=user,
        message="Can you compute something for me?",
        conversation_id=conversation_id,
    ):
        print(f"- Component type: {component.rich_component.type}")
        if (
            hasattr(component.rich_component, "content")
            and component.rich_component.content
        ):
            print(f"Assistant: {component.rich_component.content}")
        elif component.simple_component and hasattr(component.simple_component, "text"):
            print(f"Assistant: {component.simple_component.text}")
        else:
            print(f"- Component data: {component.rich_component.data}")


def run_interactive() -> None:
    """Entry point for interactive usage."""
    print("Starting mock custom tool example...")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/mock_quickstart.py
================================================
"""
Mock quickstart example for the Vanna Agents framework.

This example shows how to create a basic agent with a mock LLM service
and have a simple conversation.

Usage:
  Template: Copy this file and modify for your needs
  Interactive: python -m vanna.examples.mock_quickstart
  REPL: from vanna.examples.mock_quickstart import create_demo_agent
  Server: python -m vanna.servers --example mock_quickstart
"""

import asyncio

from vanna import (
    AgentConfig,
    Agent,
    MemoryConversationStore,
    MockLlmService,
    User,
)


def create_demo_agent() -> Agent:
    """Create a demo agent for REPL and server usage.

    Returns:
        Configured Agent instance
    """
    llm_service = MockLlmService(
        response_content="Hello! I'm a helpful AI assistant created using the Vanna Agents framework."
    )

    return Agent(
        llm_service=llm_service,
        config=AgentConfig(
            stream_responses=True,  # Enable streaming for better server experience
            include_thinking_indicators=True,
        ),
    )


async def main() -> None:
    """Run the mock quickstart example."""

    # Create agent using factory function
    agent = create_demo_agent()

    # Create a test user
    user = User(
        id="user123", username="testuser", email="test@example.com", permissions=[]
    )

    # Start a conversation
    conversation_id = "conversation123"
    user_message = "Hello! Can you introduce yourself?"

    print(f"User: {user_message}")
    print("Agent: ", end="")

    # Send message and collect response
    async for component in agent.send_message(
        user=user, message=user_message, conversation_id=conversation_id
    ):
        if hasattr(component, "content"):
            print(component.content, end="")

    print()


def run_interactive() -> None:
    """Entry point for interactive usage."""
    print("Starting Vanna Agents mock quickstart demo...")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/mock_quota_example.py
================================================
"""
Mock quota-based agent example using Mock LLM service.

This example demonstrates how to create a custom agent runner that
enforces user-based message quotas. It shows:
- Custom agent runner subclass
- Quota management and enforcement
- Error handling for quota exceeded cases
- Multiple users with different quotas

Run:
  PYTHONPATH=. python vanna/examples/mock_quota_example.py
"""

import asyncio

from vanna import (
    AgentConfig,
    MemoryConversationStore,
    MockLlmService,
    User,
)
from vanna.core.registry import ToolRegistry
from vanna.tools import ListFilesTool
from vanna.examples.quota_agent import QuotaAgentRunner, QuotaExceededError


async def demonstrate_quota_system() -> None:
    """Demonstrate the quota-based agent system."""
    print("🚀 Starting Mock Quota-based Agent Example\n")

    # Create a mock LLM service
    llm_service = MockLlmService(
        response_content="Hello! I'm here to help you with your questions."
    )

    # Create tool registry with list_files tool
    tool_registry = ToolRegistry()
    list_files_tool = ListFilesTool()
    tool_registry.register(list_files_tool)

    # Create conversation store
    conversation_store = MemoryConversationStore()

    # Create the quota-based agent
    agent = QuotaAgentRunner(
        llm_service=llm_service,
        tool_registry=tool_registry,
        conversation_store=conversation_store,
        config=AgentConfig(
            stream_responses=False,
            include_thinking_indicators=False,
        ),
    )

    # Create users with different quota settings
    regular_user = User(
        id="user1", username="alice", email="alice@example.com", permissions=[]
    )

    premium_user = User(
        id="user2", username="bob", email="bob@example.com", permissions=["premium"]
    )

    # Set custom quotas
    agent.set_user_quota(regular_user.id, 3)  # Alice gets 3 messages
    agent.set_user_quota(premium_user.id, 5)  # Bob gets 5 messages (premium)

    print("📋 User Quotas:")
    print(
        f"  • {regular_user.username}: {agent.get_user_quota(regular_user.id)} messages"
    )
    print(
        f"  • {premium_user.username}: {agent.get_user_quota(premium_user.id)} messages"
    )
    print()

    # Test regular user within quota
    print("💬 Testing regular user (Alice) within quota:")
    for i in range(1, 4):  # Send 3 messages (within quota)
        print(f"  Message {i}/3:")
        async for component in agent.send_message(
            user=regular_user,
            message=f"Hello, this is message {i}",
            conversation_id="alice-conv",
        ):
            if hasattr(component, "content") and component.content:
                print(f"    Agent: {component.content}")
        print()

    # Test regular user exceeding quota
    print("⚠️  Testing regular user (Alice) exceeding quota:")
    async for component in agent.send_message(
        user=regular_user,
        message="This message should be blocked",
        conversation_id="alice-conv",
    ):
        if hasattr(component, "content") and component.content:
            print(f"    Agent: {component.content}")
    print()

    # Test premium user with higher quota
    print("⭐ Testing premium user (Bob) with higher quota:")
    for i in range(1, 4):  # Send 3 messages
        print(f"  Message {i}/5:")
        async for component in agent.send_message(
            user=premium_user,
            message=f"Premium user message {i}",
            conversation_id="bob-conv",
        ):
            if hasattr(component, "content") and component.content:
                print(f"    Agent: {component.content}")
    print()

    # Demonstrate quota reset
    print("🔄 Resetting Alice's usage:")
    agent.reset_user_usage(regular_user.id)
    print(f"  Alice's remaining messages: {agent.get_user_remaining(regular_user.id)}")
    print()

    print("✅ After reset, Alice can send messages again:")
    async for component in agent.send_message(
        user=regular_user,
        message="This should work after reset",
        conversation_id="alice-conv2",
    ):
        if hasattr(component, "content") and component.content:
            print(f"    Agent: {component.content}")

    print("\n📊 Final Usage Summary:")
    print(
        f"  • Alice: {agent.get_user_usage(regular_user.id)}/{agent.get_user_quota(regular_user.id)} used"
    )
    print(
        f"  • Bob: {agent.get_user_usage(premium_user.id)}/{agent.get_user_quota(premium_user.id)} used"
    )


async def main() -> None:
    """Run the mock quota example."""
    await demonstrate_quota_system()


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/mock_rich_components_demo.py
================================================
"""
Mock rich components demonstration example.

This example shows how to create an agent that emits rich, stateful components
including cards, task lists, and tool execution displays using a mock LLM service.

Usage:
  PYTHONPATH=. python vanna/examples/mock_rich_components_demo.py
"""

import asyncio
import time
from datetime import datetime
from typing import AsyncGenerator, Optional

from vanna import (
    AgentConfig,
    Agent,
    MemoryConversationStore,
    MockLlmService,
    User,
)
from vanna.core.components import UiComponent
from vanna.core.rich_components import (
    StatusCardComponent,
    ProgressDisplayComponent,
    LogViewerComponent,
    BadgeComponent,
    IconTextComponent,
    RichTextComponent,
    Task,
)


class RichComponentsAgent(Agent):
    """Agent that demonstrates rich component capabilities."""

    async def send_message(
        self,
        user: User,
        message: str,
        *,
        conversation_id: Optional[str] = None,
    ) -> AsyncGenerator[UiComponent, None]:
        """Send message and yield UiComponent(rich_component=rich) components."""

        # Welcome message using IconText
        yield UiComponent(
            rich_component=IconTextComponent(
                id="welcome-message",
                icon="👋",
                text=f"Hello {user.username}! I'll demonstrate primitive components.",
                variant="primary",
                size="large",
            )
        )

        # Status card showing we're processing
        status_card = StatusCardComponent(
            id="processing-status",
            title="Processing Request",
            status="running",
            description="Processing your request...",
            icon="⚙️",
        )
        yield UiComponent(rich_component=status_card)

        # Simulate some processing time
        await asyncio.sleep(1)

        # Update status to success
        yield UiComponent(
            rich_component=status_card.set_status(
                "success", "Request processed successfully!"
            )
        )

        # Create a status card for overall demo progress
        demo_card = StatusCardComponent(
            id="demo-progress",
            title="Demo Progress",
            status="running",
            description="Starting primitive components demonstration...",
            icon="🎯",
        )
        yield UiComponent(rich_component=demo_card)

        # Create badges for different stages
        stages = [
            ("Initialize", "success", "✅"),
            ("Components", "running", "⚙️"),
            ("Progress", "pending", "⏳"),
            ("Logs", "pending", "📋"),
            ("Complete", "pending", "🎉"),
        ]

        for stage_name, stage_status, stage_icon in stages:
            yield UiComponent(
                rich_component=BadgeComponent(
                    id=f"stage-{stage_name.lower()}",
                    text=stage_name,
                    variant=stage_status if stage_status != "pending" else "default",
                    icon=stage_icon,
                    size="md",
                )
            )

        # Progress display
        progress_display = ProgressDisplayComponent(
            id="demo-progress-bar",
            label="Overall Progress",
            value=0.2,
            description="Initializing demonstration...",
            status="info",
            animated=True,
        )
        yield UiComponent(rich_component=progress_display)

        # Create log viewer for detailed progress
        log_viewer = LogViewerComponent(id="demo-logs", title="Demo Activity Log")
        yield UiComponent(rich_component=log_viewer)

        # Simulate work with updates
        for i in range(3):
            await asyncio.sleep(1)

            # Update progress
            progress_value = 0.2 + (i + 1) * 0.2
            step_name = ["Creating components", "Updating progress", "Finalizing demo"][
                i
            ]

            yield UiComponent(
                rich_component=progress_display.update_progress(
                    progress_value, f"Step {i + 2} of 5: {step_name}..."
                )
            )

            # Update demo card
            yield UiComponent(
                rich_component=demo_card.set_status(
                    "running",
                    f"Step {i + 2} of 5 completed. Progress: {int(progress_value * 100)}%",
                )
            )

            # Add log entry
            yield UiComponent(
                rich_component=log_viewer.add_entry(
                    f"Completed step: {step_name}", "info"
                )
            )

            # Update stage badges
            if i == 0:
                yield UiComponent(
                    rich_component=BadgeComponent(
                        id="stage-components",
                        text="Components",
                        variant="success",
                        icon="✅",
                        size="md",
                    )
                )
            elif i == 1:
                yield UiComponent(
                    rich_component=BadgeComponent(
                        id="stage-progress",
                        text="Progress",
                        variant="success",
                        icon="✅",
                        size="md",
                    )
                )
                yield UiComponent(
                    rich_component=BadgeComponent(
                        id="stage-logs",
                        text="Logs",
                        variant="running",
                        icon="📋",
                        size="md",
                    )
                )

        # Tool execution using primitive components
        tool_status = StatusCardComponent(
            id="demo-tool",
            title="Analyze Data Tool",
            status="running",
            description="Running regression analysis on user_data.csv",
            icon="🔬",
        )
        yield UiComponent(rich_component=tool_status)

        # Tool progress
        tool_progress = ProgressDisplayComponent(
            id="tool-progress",
            label="Tool Execution",
            value=0.0,
            description="Initializing tool...",
            animated=True,
        )
        yield UiComponent(rich_component=tool_progress)

        # Tool logs
        tool_logs = LogViewerComponent(id="tool-logs", title="Tool Execution Log")
        yield UiComponent(rich_component=tool_logs)

        # Simulate tool execution steps
        tool_steps = [
            (0.2, "Loading dataset...", "info"),
            (0.4, "Dataset loaded: 1000 rows, 5 columns", "info"),
            (0.6, "Preprocessing data...", "info"),
            (0.8, "Running regression analysis...", "info"),
            (1.0, "Analysis complete!", "info"),
        ]

        for progress_val, log_message, log_level in tool_steps:
            await asyncio.sleep(0.5)

            yield UiComponent(
                rich_component=tool_progress.update_progress(
                    progress_val, f"Progress: {int(progress_val * 100)}%"
                )
            )
            yield UiComponent(
                rich_component=tool_logs.add_entry(log_message, log_level)
            )

        # Complete tool execution
        yield UiComponent(
            rich_component=tool_status.set_status(
                "success",
                "Tool completed successfully. R² = 0.85, strong correlation found.",
            )
        )

        # Show results using IconText
        yield UiComponent(
            rich_component=IconTextComponent(
                id="tool-results",
                icon="📊",
                text="Analysis Results: R² = 0.85 (Strong correlation)",
                variant="success",
                size="medium",
            )
        )

        # Update final stage badge
        yield UiComponent(
            rich_component=BadgeComponent(
                id="stage-logs", text="Logs", variant="success", icon="✅", size="md"
            )
        )
        yield UiComponent(
            rich_component=BadgeComponent(
                id="stage-complete",
                text="Complete",
                variant="success",
                icon="🎉",
                size="md",
            )
        )

        # Final updates
        yield UiComponent(
            rich_component=progress_display.update_progress(
                1.0, "Demo completed successfully!"
            )
        )

        yield UiComponent(
            rich_component=demo_card.set_status(
                "success", "Primitive components demonstration finished successfully!"
            )
        )

        # Add final log entry
        yield UiComponent(
            rich_component=tool_logs.add_entry("Demo completed successfully!", "info")
        )

        # Add final text response
        yield UiComponent(
            rich_component=RichTextComponent(
                content=f"""## Primitive Components Demo Complete!

I've demonstrated the new primitive component system:

- **Status Cards**: Domain-agnostic status displays that work for any process
- **Progress Displays**: Reusable progress indicators with animations
- **Log Viewers**: Structured log display for any activity
- **Badges**: Flexible status and category indicators
- **Icon Text**: Composable icon+text combinations

### Key Benefits of Primitive Components:

- **Separation of Concerns**: UI components are purely presentational
- **Reusability**: Components work across different domains and tools
- **Composability**: Tools build exactly the UI they need from primitives
- **Maintainability**: Business logic changes don't affect UI components
- **Extensibility**: New tools don't require new component types

**Primitive Components**: Compose UI from domain-agnostic building blocks
**After**: Tools compose UI from primitive `StatusCard` + `ProgressDisplay` + `LogViewer`

Your message was: "{message}"
""",
                markdown=True,
            )
        )


# CLI compatibility alias
create_demo_agent = lambda: create_rich_demo_agent()


def create_rich_demo_agent() -> RichComponentsAgent:
    """Create a primitive components demo agent.

    Returns:
        Configured RichComponentsAgent instance
    """
    llm_service = MockLlmService(response_content="Primitive components demo response")

    return RichComponentsAgent(
        llm_service=llm_service,
        config=AgentConfig(
            stream_responses=True,
            include_thinking_indicators=False,  # We'll use custom status cards
        ),
    )


async def main() -> None:
    """Run the primitive components demo."""

    # Create agent
    agent = create_rich_demo_agent()

    # Create a test user
    user = User(
        id="user123", username="demo_user", email="demo@example.com", permissions=[]
    )

    # Start a conversation
    conversation_id = "primitive_demo_123"
    user_message = "Show me the primitive components demo!"

    print(f"User: {user_message}")
    print("Agent response (primitive components):")
    print("=" * 50)

    # Send message and display components
    component_count = 0
    async for component in agent.send_message(
        user=user, message=user_message, conversation_id=conversation_id
    ):
        component_count += 1
        rich_comp = component.rich_component
        component_type = getattr(rich_comp, "type", rich_comp.__class__.__name__)
        component_id = getattr(rich_comp, "id", "N/A")
        lifecycle = getattr(rich_comp, "lifecycle", "N/A")

        print(
            f"[{component_count:2d}] {component_type} (id: {component_id[:8]}, lifecycle: {lifecycle})"
        )

        # Show some component details
        if hasattr(rich_comp, "title"):
            print(f"     Title: {rich_comp.title}")
        if hasattr(rich_comp, "content") and len(str(rich_comp.content)) < 100:
            print(f"     Content: {rich_comp.content}")
        if hasattr(rich_comp, "status"):
            print(f"     Status: {rich_comp.status}")
        if (
            hasattr(rich_comp, "value")
            and hasattr(rich_comp.type, "value")
            and rich_comp.type.value == "progress_bar"
        ):
            print(f"     Progress: {rich_comp.value:.1%}")

        print()

    print("=" * 50)
    print(f"Total components emitted: {component_count}")


def run_interactive() -> None:
    """Entry point for interactive usage."""
    print("Starting Primitive Components Demo...")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/mock_sqlite_example.py
================================================
"""
Mock example showing how to use the SQL query tool with the Chinook database.

This example demonstrates using the RunSqlTool with SqliteRunner and a mock LLM service
that automatically executes sample SQL queries against the Chinook database.

Usage:
  Template: Copy this file and modify for your custom database
  Interactive: python -m vanna.examples.mock_sqlite_example
  REPL: from vanna.examples.mock_sqlite_example import create_demo_agent
  Server: python -m vanna.servers --example mock_sqlite_example
"""

import asyncio
import os
import random
import uuid
from typing import AsyncGenerator, Dict, List, Optional, Type

from pydantic import BaseModel, Field

from vanna import (
    AgentConfig,
    Agent,
    Tool,
    ToolContext,
    ToolRegistry,
    ToolResult,
    User,
    UiComponent,
)
from vanna.core.interfaces import LlmService
from vanna.core.models import (
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
    ToolCall,
    ToolSchema,
)
from vanna.core.rich_components import (
    CardComponent,
    NotificationComponent,
    ComponentType,
)
from vanna.core.simple_components import (
    SimpleTextComponent,
)
from vanna.tools import RunSqlTool
from vanna.integrations.sqlite import SqliteRunner


class MockSqliteLlmService(LlmService):
    """LLM service that exercises the SQLite query tool with sample queries."""

    def __init__(self, seed: Optional[int] = None):
        self._random = random.Random(seed)
        self._sample_queries = [
            "SELECT name FROM sqlite_master WHERE type='table'",
            "SELECT COUNT(*) as total_customers FROM Customer",
            "SELECT FirstName, LastName FROM Customer LIMIT 5",
            "SELECT Name, Composer FROM Track WHERE Composer IS NOT NULL LIMIT 5",
            "SELECT COUNT(*) as album_count FROM Album",
            "SELECT Name FROM Artist LIMIT 10",
            "SELECT AVG(Total) as avg_invoice_total FROM Invoice",
            "SELECT GenreId, COUNT(*) as track_count FROM Track GROUP BY GenreId LIMIT 5",
        ]

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Handle non-streaming SQLite interactions."""
        await asyncio.sleep(0.1)
        return self._build_response(request)

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Provide streaming compatibility by yielding a single chunk."""
        await asyncio.sleep(0.1)
        response = self._build_response(request)

        if response.tool_calls:
            yield LlmStreamChunk(tool_calls=response.tool_calls)
        if response.content is not None:
            yield LlmStreamChunk(
                content=response.content, finish_reason=response.finish_reason
            )
        else:
            yield LlmStreamChunk(finish_reason=response.finish_reason)

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Mock validation - no errors."""
        return []

    def _build_response(self, request: LlmRequest) -> LlmResponse:
        """Create a response that either calls the tool or explains its result."""
        last_message = request.messages[-1] if request.messages else None

        if last_message and last_message.role == "tool":
            # Respond to tool result
            result = last_message.content or "No result provided"
            return LlmResponse(
                content=f"Here's what I found in the database:\n\n{result}",
                finish_reason="stop",
                usage={
                    "prompt_tokens": 40,
                    "completion_tokens": 20,
                    "total_tokens": 60,
                },
            )

        # Generate a random SQL query
        sql_query = self._random.choice(self._sample_queries)
        tool_call = ToolCall(
            id=f"call_{uuid.uuid4().hex[:8]}",
            name="run_sql",
            arguments={"sql": sql_query},
        )

        return LlmResponse(
            content="Let me query the Chinook database for you...",
            tool_calls=[tool_call],
            finish_reason="tool_calls",
            usage={"prompt_tokens": 30, "completion_tokens": 10, "total_tokens": 40},
        )


def create_demo_agent() -> Agent:
    """Create a demo agent with SQLite query tool.

    Returns:
        Configured Agent with SQLite tool and mock LLM
    """
    # Get the path to the Chinook database
    database_path = os.path.join(
        os.path.dirname(__file__), "..", "..", "Chinook.sqlite"
    )
    database_path = os.path.abspath(database_path)

    if not os.path.exists(database_path):
        raise FileNotFoundError(
            f"Chinook database not found at {database_path}. Please download it from https://vanna.ai/Chinook.sqlite"
        )

    tool_registry = ToolRegistry()
    sqlite_runner = SqliteRunner(database_path=database_path)
    sql_tool = RunSqlTool(sql_runner=sqlite_runner)
    tool_registry.register(sql_tool)

    llm_service = MockSqliteLlmService()

    return Agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        config=AgentConfig(
            stream_responses=False,
            include_thinking_indicators=False,
        ),
    )


async def main() -> None:
    """Run the mock SQLite example."""

    # Create agent using factory function
    agent = create_demo_agent()
    tool_registry = agent.tool_registry

    # Create a test user
    user = User(id="user123", username="testuser", permissions=[])

    # Test the tool directly
    print("Testing SQL tool directly:")
    tool_call = ToolCall(
        id="test123",
        name="run_sql",
        arguments={"sql": "SELECT name FROM sqlite_master WHERE type='table'"},
    )

    context = ToolContext(user=user, conversation_id="test", request_id="test")

    result = await tool_registry.execute(tool_call, context)
    print(
        f"Tables in database:\n{result.result_for_llm if result.success else result.error}"
    )

    # Show available tools
    schemas = await tool_registry.get_schemas(user)
    print(f"\nAvailable tools for user: {[schema.name for schema in schemas]}")

    # Demonstrate the mock LLM triggering SQL queries
    print("\n" + "=" * 50)
    print("Agent conversation demo:")
    print("=" * 50)

    conversation_id = "sqlite-demo"

    # Run multiple queries to show different results
    for i in range(3):
        print(f"\n--- Query {i + 1} ---")
        async for component in agent.send_message(
            user=user,
            message=f"Show me some data from the database (query {i + 1})",
            conversation_id=conversation_id,
        ):
            if (
                hasattr(component.rich_component, "content")
                and component.rich_component.content
            ):
                print(f"Assistant: {component.rich_component.content}")
            elif component.simple_component and hasattr(
                component.simple_component, "text"
            ):
                print(f"Assistant: {component.simple_component.text}")


def run_interactive() -> None:
    """Entry point for interactive usage."""
    print("Starting mock SQLite example...")
    print("This example uses the Chinook database to demonstrate SQL queries.")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/openai_quickstart.py
================================================
"""
OpenAI example using OpenAILlmService.

Loads environment from .env (via python-dotenv), uses model 'gpt-5' by default,
and sends a simple message through a Agent.

Run:
  PYTHONPATH=. python vanna/examples/openai_quickstart.py
"""

import asyncio
import importlib.util
import os
import sys


def ensure_env() -> None:
    if importlib.util.find_spec("dotenv") is not None:
        from dotenv import load_dotenv

        # Load from local .env without overriding existing env
        load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"), override=False)
    else:
        print(
            "[warn] python-dotenv not installed; skipping .env load. Install with: pip install python-dotenv"
        )

    if not os.getenv("OPENAI_API_KEY"):
        print(
            "[error] OPENAI_API_KEY is not set. Add it to your environment or .env file."
        )
        sys.exit(1)


async def main() -> None:
    ensure_env()

    # Lazy import after env load to allow custom base_url/org via env
    try:
        from vanna.integrations.anthropic import OpenAILlmService
    except ImportError as e:
        print(
            "[error] openai extra not installed. Install with: pip install -e .[openai]"
        )
        raise

    from vanna import AgentConfig, Agent, User
    from vanna.core.registry import ToolRegistry
    from vanna.tools import ListFilesTool

    # Default to 'gpt-5' for this demo; override via $OPENAI_MODEL if desired
    model = os.getenv("OPENAI_MODEL", "gpt-5")
    print(f"Using OpenAI model: {model}")

    llm = OpenAILlmService(model=model)

    # Create tool registry and register the list_files tool
    tool_registry = ToolRegistry()
    list_files_tool = ListFilesTool()
    tool_registry.register(list_files_tool)

    # Some models (e.g., reasoning/gpt-5) only support the default temperature=1.0
    agent = Agent(
        llm_service=llm,
        config=AgentConfig(stream_responses=False, temperature=1.0),
        tool_registry=tool_registry,
    )

    user = User(id="demo-user", username="demo")
    conversation_id = "openai-demo"

    print("Sending: 'List the files in the current directory'\n")
    async for component in agent.send_message(
        user=user,
        message="List the files in the current directory",
        conversation_id=conversation_id,
    ):
        if hasattr(component, "content") and component.content:
            print("Assistant:", component.content)


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/examples/primitive_components_demo.py
================================================
"""
Demonstration of the new primitive component system.

This example shows how tools compose UI from primitive, domain-agnostic
components like StatusCardComponent, ProgressDisplayComponent, etc.

Usage:
  PYTHONPATH=. python vanna/examples/primitive_components_demo.py
"""

import asyncio
import uuid
from datetime import datetime
from typing import AsyncGenerator, Optional

from vanna import (
    AgentConfig,
    Agent,
    MemoryConversationStore,
    MockLlmService,
    User,
)
from vanna.core.components import UiComponent
from vanna.core.rich_components import (
    StatusCardComponent,
    ProgressDisplayComponent,
    LogViewerComponent,
    BadgeComponent,
    IconTextComponent,
    RichTextComponent,
)


class PrimitiveComponentsAgent(Agent):
    """Agent that demonstrates the new primitive component system."""

    async def send_message(
        self,
        user: User,
        message: str,
        *,
        conversation_id: Optional[str] = None,
    ) -> AsyncGenerator[UiComponent, None]:
        """Send message and demonstrate primitive component composition."""

        session_id = str(uuid.uuid4())[:8]

        # Demo 1: Tool execution using primitive components
        yield UiComponent(
            rich_component=RichTextComponent(
                content="## Primitive Components Demo\n\nShowing how tools now compose UI from primitive components:",
                markdown=True,
            )
        )

        # Status card for overall operation
        operation_status = StatusCardComponent(
            id=f"operation-{session_id}",
            title="Data Analysis Pipeline",
            status="running",
            description="Processing user data through multiple analysis stages",
            icon="⚙️",
        )
        yield UiComponent(rich_component=operation_status)

        # Progress display for overall progress
        overall_progress = ProgressDisplayComponent(
            id=f"progress-{session_id}",
            label="Overall Progress",
            value=0.0,
            description="Starting analysis...",
            animated=True,
        )
        yield UiComponent(rich_component=overall_progress)

        # Log viewer for detailed output
        log_viewer = LogViewerComponent(
            id=f"logs-{session_id}",
            title="Analysis Log",
            entries=[],
            show_timestamps=True,
            auto_scroll=True,
        )
        yield UiComponent(rich_component=log_viewer)

        # Simulate analysis stages
        stages = [
            ("Data Loading", "📊", 0.2),
            ("Data Validation", "✅", 0.4),
            ("Statistical Analysis", "🧮", 0.6),
            ("Report Generation", "📄", 0.8),
            ("Finalization", "🎯", 1.0),
        ]

        for i, (stage_name, stage_icon, progress_value) in enumerate(stages):
            await asyncio.sleep(0.8)

            # Update overall status
            status = "success" if progress_value == 1.0 else "running"
            yield UiComponent(
                rich_component=operation_status.set_status(
                    status, f"Executing: {stage_name}"
                )
            )

            # Update progress
            yield UiComponent(
                rich_component=overall_progress.update_progress(
                    progress_value, f"Executing {stage_name}..."
                )
            )

            # Add log entry
            yield UiComponent(
                rich_component=log_viewer.add_entry(f"Starting {stage_name}", "info")
            )

            # Create a status card for this specific stage
            stage_status = StatusCardComponent(
                id=f"stage-{i}-{session_id}",
                title=stage_name,
                status="running" if progress_value < 1.0 else "success",
                description=f"Processing stage {i + 1} of {len(stages)}",
                icon=stage_icon,
            )
            yield UiComponent(rich_component=stage_status)

            await asyncio.sleep(0.5)

            # Complete the stage
            final_stage_status = "success" if progress_value < 1.0 else "completed"
            yield UiComponent(
                rich_component=stage_status.set_status(
                    final_stage_status, f"{stage_name} completed successfully"
                )
            )
            yield UiComponent(
                rich_component=log_viewer.add_entry(f"Completed {stage_name}", "info")
            )

        # Demo 2: Badge and IconText primitives
        yield UiComponent(
            rich_component=RichTextComponent(
                content="\n### Primitive Component Examples\n\nShowing individual primitive components:",
                markdown=True,
            )
        )

        # Various badge examples
        badges = [
            BadgeComponent(text="Processing", variant="primary", size="small"),
            BadgeComponent(text="Complete", variant="success", size="medium"),
            BadgeComponent(text="Warning", variant="warning", size="large", icon="⚠️"),
            BadgeComponent(text="Error", variant="error", size="medium", icon="❌"),
        ]

        for badge in badges:
            yield UiComponent(rich_component=badge)

        # IconText examples
        icon_texts = [
            IconTextComponent(
                icon="📊",
                text="Data Analysis Complete",
                variant="primary",
                size="large",
            ),
            IconTextComponent(
                icon="✅", text="All tests passed", variant="default", size="medium"
            ),
            IconTextComponent(
                icon="⏱️",
                text="Processing time: 2.3s",
                variant="secondary",
                size="small",
            ),
        ]

        for icon_text in icon_texts:
            yield UiComponent(rich_component=icon_text)

        # Demo 3: Comparison with old approach
        yield UiComponent(
            rich_component=RichTextComponent(
                content=f"""
## Key Benefits of Primitive Components

**Primitive Component Approach:**
```python
# Tool composes UI from primitives
status_card = StatusCardComponent(
    title="Data Analysis",
    status="running",  # Pure UI state
    icon="📊"
)
progress = ProgressDisplayComponent(
    label="Analysis Progress",
    value=0.5
)
logs = LogViewerComponent(
    title="Analysis Log",
    entries=log_entries
)
```

### Benefits:
- **Separation of Concerns**: UI components are purely presentational
- **Reusability**: Status cards work for any process, not just tools
- **Composability**: Tools build exactly the UI they need
- **Maintainability**: Changes to business logic don't affect UI components
- **Extensibility**: New tools don't require new component types

Your message was: "{message}"
""",
                markdown=True,
            )
        )


def create_primitive_demo_agent() -> PrimitiveComponentsAgent:
    """Create a primitive components demo agent.

    Returns:
        Configured PrimitiveComponentsAgent instance
    """
    llm_service = MockLlmService(response_content="Primitive components demo response")

    return PrimitiveComponentsAgent(
        llm_service=llm_service,
        config=AgentConfig(
            stream_responses=True,
            include_thinking_indicators=False,
        ),
    )


async def main() -> None:
    """Run the primitive components demo."""

    # Create agent
    agent = create_primitive_demo_agent()

    # Create a test user
    user = User(
        id="user123", username="demo_user", email="demo@example.com", permissions=[]
    )

    # Start a conversation
    conversation_id = "primitive_demo_123"
    user_message = "Show me how the new primitive component system works!"

    print(f"User: {user_message}")
    print("Agent response (primitive components):")
    print("=" * 60)

    # Send message and display components
    component_count = 0
    async for component in agent.send_message(
        user=user, message=user_message, conversation_id=conversation_id
    ):
        component_count += 1
        component_type = getattr(component, "type", component.__class__.__name__)
        component_id = getattr(component, "id", "N/A")

        print(
            f"[{component_count:2d}] {component_type.value if hasattr(component_type, 'value') else component_type} (id: {component_id[:12] if len(str(component_id)) > 12 else component_id})"
        )

        rich_comp = component.rich_component

        # Show component details
        if hasattr(rich_comp, "title"):
            print(f"     Title: {rich_comp.title}")
        if hasattr(rich_comp, "status"):
            print(f"     Status: {rich_comp.status}")
        if hasattr(rich_comp, "description") and rich_comp.description:
            desc = (
                rich_comp.description[:60] + "..."
                if len(rich_comp.description) > 60
                else rich_comp.description
            )
            print(f"     Description: {desc}")
        if (
            hasattr(rich_comp, "value")
            and hasattr(rich_comp.type, "value")
            and rich_comp.type.value == "progress_display"
        ):
            print(f"     Progress: {rich_comp.value:.1%}")

        print()

    print("=" * 60)
    print(f"Total components emitted: {component_count}")
    print("\nThis demonstrates how tools can now compose rich UIs")
    print("from primitive, reusable components without semantic coupling!")


def run_interactive() -> None:
    """Entry point for interactive usage."""
    print("Starting Primitive Components Demo...")
    asyncio.run(main())


if __name__ == "__main__":
    run_interactive()


================================================
FILE: src/vanna/examples/quota_lifecycle_example.py
================================================
"""
Example demonstrating lifecycle hooks for user quota management.

This example shows how to use lifecycle hooks to add custom functionality
like quota management without creating custom agent runner subclasses.
"""

from typing import Any, Dict, Optional
from vanna.core import Agent, LifecycleHook, User
from vanna.core.errors import AgentError


class QuotaExceededError(AgentError):
    """Raised when a user exceeds their message quota."""

    pass


class QuotaCheckHook(LifecycleHook):
    """Lifecycle hook that enforces user-based message quotas."""

    def __init__(self, default_quota: int = 10) -> None:
        """Initialize quota hook.

        Args:
            default_quota: Default quota per user if not specifically set
        """
        self._user_quotas: Dict[str, int] = {}
        self._user_usage: Dict[str, int] = {}
        self._default_quota = default_quota

    def set_user_quota(self, user_id: str, quota: int) -> None:
        """Set a specific quota for a user."""
        self._user_quotas[user_id] = quota

    def get_user_quota(self, user_id: str) -> int:
        """Get the quota for a user."""
        return self._user_quotas.get(user_id, self._default_quota)

    def get_user_usage(self, user_id: str) -> int:
        """Get current usage count for a user."""
        return self._user_usage.get(user_id, 0)

    def get_user_remaining(self, user_id: str) -> int:
        """Get remaining messages for a user."""
        return self.get_user_quota(user_id) - self.get_user_usage(user_id)

    def reset_user_usage(self, user_id: str) -> None:
        """Reset usage count for a user."""
        self._user_usage[user_id] = 0

    async def before_message(self, user: User, message: str) -> Optional[str]:
        """Check quota before processing message.

        Raises:
            QuotaExceededError: If user has exceeded their quota
        """
        usage = self.get_user_usage(user.id)
        quota = self.get_user_quota(user.id)

        if usage >= quota:
            raise QuotaExceededError(
                f"User {user.username} has exceeded their quota of {quota} messages. "
                f"Current usage: {usage}"
            )

        # Increment usage count
        current_usage = self._user_usage.get(user.id, 0)
        self._user_usage[user.id] = current_usage + 1

        # Don't modify the message
        return None


class LoggingHook(LifecycleHook):
    """Example logging hook for demonstration."""

    async def before_message(self, user: User, message: str) -> Optional[str]:
        """Log incoming messages."""
        print(f"[LOG] User {user.username} ({user.id}) sent message: {message[:50]}...")
        return None

    async def after_message(self, result: Any) -> None:
        """Log message completion."""
        print(f"[LOG] Message processing completed")


async def run_example() -> None:
    """
    Example showing how to use lifecycle hooks with Agent.

    Instead of creating a custom subclass, we compose
    the behavior using lifecycle hooks.
    """
    from vanna.core.registry import ToolRegistry
    from vanna.integrations.anthropic import AnthropicLlmService
    from vanna.integrations.local import MemoryConversationStore

    # Create quota hook
    quota_hook = QuotaCheckHook(default_quota=10)
    quota_hook.set_user_quota("user123", 5)  # Set custom quota for specific user

    # Create logging hook
    logging_hook = LoggingHook()

    # Create agent with multiple hooks
    agent = Agent(
        llm_service=AnthropicLlmService(api_key="your-api-key"),
        tool_registry=ToolRegistry(),
        conversation_store=MemoryConversationStore(),
        lifecycle_hooks=[
            logging_hook,  # Logs will happen first
            quota_hook,  # Then quota check
        ],
    )

    # Create a test user
    user = User(
        id="user123", username="test_user", email="test@example.com", permissions=[]
    )

    # Send messages - will track quota
    try:
        async for component in agent.send_message(user=user, message="Hello, agent!"):
            # Process UI components
            pass

        # Check remaining quota
        remaining = quota_hook.get_user_remaining(user.id)
        print(f"Remaining messages: {remaining}/{quota_hook.get_user_quota(user.id)}")

    except QuotaExceededError as e:
        print(f"Quota exceeded: {e}")


if __name__ == "__main__":
    import asyncio

    asyncio.run(run_example())


================================================
FILE: src/vanna/examples/visualization_example.py
================================================
"""
Example demonstrating SQL query execution with automatic visualization.

This example shows the integration of RunSqlTool and VisualizeDataTool,
demonstrating how SQL results are saved to CSV files and can be visualized
using the visualization tool with dependency injection.

Usage:
  PYTHONPATH=. python vanna/examples/visualization_example.py
"""

import asyncio
import os
import sys
import uuid
from typing import AsyncGenerator, List, Optional

from vanna import (
    AgentConfig,
    Agent,
    ToolRegistry,
    User,
)
from vanna.core import LlmService
from vanna.core import (
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
    ToolCall,
    ToolSchema,
)
from vanna.integrations.sqlite import SqliteRunner
from vanna.tools import (
    RunSqlTool,
    VisualizeDataTool,
    LocalFileSystem,
)


class VisualizationDemoLlmService(LlmService):
    """Mock LLM that demonstrates SQL query and visualization workflow."""

    def __init__(self) -> None:
        self.step = 0
        self.csv_filename: Optional[str] = None

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Handle non-streaming requests."""
        await asyncio.sleep(0.1)
        return self._build_response(request)

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Handle streaming requests."""
        await asyncio.sleep(0.1)
        response = self._build_response(request)

        if response.tool_calls:
            yield LlmStreamChunk(tool_calls=response.tool_calls)
        if response.content:
            yield LlmStreamChunk(
                content=response.content, finish_reason=response.finish_reason
            )
        else:
            yield LlmStreamChunk(finish_reason=response.finish_reason)

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Validate tools - no errors."""
        return []

    def _build_response(self, request: LlmRequest) -> LlmResponse:
        """Build response based on conversation state."""
        last_message = request.messages[-1] if request.messages else None

        # If we got a tool result, process it
        if last_message and last_message.role == "tool":
            tool_result = last_message.content or ""

            # Check if this was a SQL query result with a CSV file
            if "Results saved to" in tool_result and ".csv" in tool_result:
                # Extract filename from result
                import re

                match = re.search(r"'([^']*\.csv)'", tool_result)
                if match:
                    self.csv_filename = match.group(1)
                    # Now visualize the data
                    return LlmResponse(
                        content=f"Great! I've saved the query results. Now let me create a visualization of the data.",
                        tool_calls=[
                            ToolCall(
                                id=f"call_{uuid.uuid4().hex[:8]}",
                                name="visualize_data",
                                arguments={"filename": self.csv_filename},
                            )
                        ],
                        finish_reason="tool_calls",
                    )

            # If this was a visualization result, acknowledge it
            if "Created visualization" in tool_result:
                return LlmResponse(
                    content=f"Perfect! I've created a visualization of the data. {tool_result}",
                    finish_reason="stop",
                )

            # Default acknowledgment
            return LlmResponse(
                content=f"I've completed the operation. {tool_result}",
                finish_reason="stop",
            )

        # Initial request - run SQL query
        if self.step == 0:
            self.step += 1
            return LlmResponse(
                content="I'll query the database for you and then create a visualization.",
                tool_calls=[
                    ToolCall(
                        id=f"call_{uuid.uuid4().hex[:8]}",
                        name="run_sql",
                        arguments={
                            "sql": "SELECT Name, Milliseconds, Bytes FROM Track LIMIT 20"
                        },
                    )
                ],
                finish_reason="tool_calls",
            )

        # Default response
        return LlmResponse(
            content="I can help you query databases and visualize the results.",
            finish_reason="stop",
        )


def create_demo_agent() -> Agent:
    """
    Create a demo agent with SQL and visualization tools.

    This function is called by the vanna server framework.

    Returns:
        Configured Agent with SQL and visualization tools
    """
    # Check for Chinook database
    database_path = os.path.join(
        os.path.dirname(__file__), "..", "..", "Chinook.sqlite"
    )
    database_path = os.path.abspath(database_path)

    if not os.path.exists(database_path):
        raise FileNotFoundError(
            f"Chinook database not found at {database_path}. "
            "Please download it from https://vanna.ai/Chinook.sqlite"
        )

    # Create shared FileSystem for both tools
    file_system = LocalFileSystem(working_directory="./data_storage")

    # Create SQL tool with FileSystem
    sqlite_runner = SqliteRunner(database_path=database_path)
    sql_tool = RunSqlTool(sql_runner=sqlite_runner, file_system=file_system)

    # Create visualization tool with same FileSystem
    viz_tool = VisualizeDataTool(file_system=file_system)

    # Create tool registry
    tool_registry = ToolRegistry()
    tool_registry.register(sql_tool)
    tool_registry.register(viz_tool)

    # Create LLM service
    llm_service = VisualizationDemoLlmService()

    # Create agent with streaming enabled for web interface
    return Agent(
        llm_service=llm_service,
        tool_registry=tool_registry,
        config=AgentConfig(
            stream_responses=True,
            include_thinking_indicators=False,
        ),
    )


async def main() -> None:
    """Demonstrate SQL query execution with automatic visualization."""
    print("🎨 SQL + Visualization Demo")
    print("=" * 60)
    print("This example demonstrates:")
    print("1. Running SQL queries that save results to CSV files")
    print("2. Automatically visualizing the CSV data")
    print("3. User isolation for file storage")
    print("=" * 60)
    print()

    # Create agent using factory function
    agent = create_demo_agent()

    # Create test user
    user = User(id="demo-user", username="demo")

    # Show available tools
    tools = await agent.get_available_tools(user)
    print(f"Available tools: {[tool.name for tool in tools]}")
    print()

    # Run conversation
    conversation_id = "viz-demo"

    print("User: Show me some track data and visualize it")
    print()

    async for component in agent.send_message(
        user=user,
        message="Show me some track data and visualize it",
        conversation_id=conversation_id,
    ):
        if (
            component.simple_component
            and hasattr(component.simple_component, "text")
            and component.simple_component.text
        ):
            print(f"Agent: {component.simple_component.text}")
        elif component.simple_component and hasattr(component.simple_component, "text"):
            print(f"Agent: {component.simple_component.text}")
        elif hasattr(component.rich_component, "content"):
            if isinstance(component.rich_component.content, dict):
                # This is the chart
                print(
                    f"Agent: [Chart Generated - Plotly figure with {len(str(component.rich_component.content))} chars]"
                )
            else:
                print(f"Agent: {component.rich_component.content}")

    print()
    print("=" * 60)
    print("Demo complete!")
    print()
    print("Key features demonstrated:")
    print("✅ SQL queries save results to user-isolated CSV files")
    print("✅ Visualization tool reads CSV files using FileSystem")
    print("✅ Automatic chart type selection based on data shape")
    print("✅ Dependency injection allows customization")
    print()


if __name__ == "__main__":
    asyncio.run(main())


================================================
FILE: src/vanna/integrations/__init__.py
================================================
"""
Integrations module.

This package contains concrete implementations of core abstractions and capabilities.
"""

from .local import MemoryConversationStore
from .mock import MockLlmService
from .plotly import PlotlyChartGenerator
from .sqlite import SqliteRunner

__all__ = [
    "MockLlmService",
    "MemoryConversationStore",
    "SqliteRunner",
    "PlotlyChartGenerator",
]


================================================
FILE: src/vanna/integrations/anthropic/__init__.py
================================================
"""
Anthropic integration.

This module provides Anthropic LLM service implementation.
"""

from .llm import AnthropicLlmService

__all__ = ["AnthropicLlmService"]


================================================
FILE: src/vanna/integrations/anthropic/llm.py
================================================
"""
Anthropic LLM service implementation.

Implements the LlmService interface using Anthropic's Messages API
(anthropic>=0.8.0). Supports non-streaming and streaming text output.
Tool-calls (tool_use blocks) are surfaced at the end of a stream or after a
non-streaming call as ToolCall entries.
"""

from __future__ import annotations

import logging
import os
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

from vanna.core.llm import (
    LlmService,
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
)
from vanna.core.tool import ToolCall, ToolSchema


class AnthropicLlmService(LlmService):
    """Anthropic Messages-backed LLM service.

    Args:
        model: Anthropic model name (e.g., "claude-sonnet-4-5", "claude-opus-4").
            Defaults to "claude-sonnet-4-5". Can also be set via ANTHROPIC_MODEL env var.
        api_key: API key; falls back to env `ANTHROPIC_API_KEY`.
        base_url: Optional custom base URL; env `ANTHROPIC_BASE_URL` if unset.
        extra_client_kwargs: Extra kwargs forwarded to `anthropic.Anthropic()`.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        **extra_client_kwargs: Any,
    ) -> None:
        try:
            import anthropic
        except Exception as e:  # pragma: no cover
            raise ImportError(
                "anthropic package is required. Install with: pip install 'vanna[anthropic]'"
            ) from e

        # Model selection - use environment variable or default
        self.model = model or os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-5")
        api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
        base_url = base_url or os.getenv("ANTHROPIC_BASE_URL")

        client_kwargs: Dict[str, Any] = {**extra_client_kwargs}
        if api_key:
            client_kwargs["api_key"] = api_key
        if base_url:
            client_kwargs["base_url"] = base_url

        self._client = anthropic.Anthropic(**client_kwargs)

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a non-streaming request to Anthropic and return the response."""
        payload = self._build_payload(request)

        resp = self._client.messages.create(**payload)

        logger.info(f"Anthropic response: {resp}")

        text_content, tool_calls = self._parse_message_content(resp)

        usage: Dict[str, int] = {}
        if getattr(resp, "usage", None):
            try:
                usage = {
                    "input_tokens": int(resp.usage.input_tokens),
                    "output_tokens": int(resp.usage.output_tokens),
                }
            except Exception:
                pass

        return LlmResponse(
            content=text_content or None,
            tool_calls=tool_calls or None,
            finish_reason=getattr(resp, "stop_reason", None),
            usage=usage or None,
        )

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Stream a request to Anthropic.

        Yields text chunks as they arrive. Emits tool-calls at the end by
        inspecting the final message.
        """
        payload = self._build_payload(request)

        logger.info(f"Anthropic streaming payload: {payload}")

        # SDK provides a streaming context manager with a text_stream iterator.
        with self._client.messages.stream(**payload) as stream:
            for text in stream.text_stream:
                if text:
                    yield LlmStreamChunk(content=text)

            final = stream.get_final_message()
            logger.info(f"Anthropic stream response: {final}")
            _, tool_calls = self._parse_message_content(final)
            if tool_calls:
                yield LlmStreamChunk(
                    tool_calls=tool_calls,
                    finish_reason=getattr(final, "stop_reason", None),
                )
            else:
                yield LlmStreamChunk(
                    finish_reason=getattr(final, "stop_reason", None) or "stop"
                )

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Basic validation of tool schemas for Anthropic."""
        errors: List[str] = []
        for t in tools:
            if not t.name:
                errors.append("Tool name is required")
        return errors

    # Internal helpers
    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
        # Anthropic requires messages content as list of content blocks per message
        # We need to group consecutive tool messages into single user messages
        messages: List[Dict[str, Any]] = []
        i = 0

        while i < len(request.messages):
            m = request.messages[i]

            if m.role == "tool":
                # Group consecutive tool messages into one user message
                tool_content_blocks = []
                while i < len(request.messages) and request.messages[i].role == "tool":
                    tool_msg = request.messages[i]
                    if tool_msg.tool_call_id:
                        tool_content_blocks.append(
                            {
                                "type": "tool_result",
                                "tool_use_id": tool_msg.tool_call_id,
                                "content": tool_msg.content,
                            }
                        )
                    i += 1

                if tool_content_blocks:
                    messages.append(
                        {
                            "role": "user",
                            "content": tool_content_blocks,
                        }
                    )
            else:
                # Handle non-tool messages normally
                content_blocks = []

                # Handle text content - only add if not empty
                if m.content and m.content.strip():
                    content_blocks.append({"type": "text", "text": m.content})

                # Handle tool_calls for assistant messages (convert to tool_use blocks)
                if m.role == "assistant" and m.tool_calls:
                    for tc in m.tool_calls:
                        content_blocks.append(
                            {
                                "type": "tool_use",
                                "id": tc.id,
                                "name": tc.name,
                                "input": tc.arguments,  # type: ignore[dict-item]
                            }
                        )

                # Ensure we have at least one content block for text messages
                if not content_blocks and m.role in {"user", "assistant"}:
                    content_blocks.append({"type": "text", "text": m.content or ""})

                if content_blocks:
                    role = m.role if m.role in {"user", "assistant"} else "user"
                    messages.append(
                        {
                            "role": role,
                            "content": content_blocks,
                        }
                    )

                i += 1

        tools_payload: Optional[List[Dict[str, Any]]] = None
        if request.tools:
            tools_payload = [
                {
                    "name": t.name,
                    "description": t.description,
                    "input_schema": t.parameters,
                }
                for t in request.tools
            ]

        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            # Anthropic requires max_tokens; default if not provided
            "max_tokens": request.max_tokens if request.max_tokens is not None else 512,
            "temperature": request.temperature,
        }
        if tools_payload:
            payload["tools"] = tools_payload
            payload["tool_choice"] = {"type": "auto"}

        # Add system prompt if provided
        if request.system_prompt:
            payload["system"] = request.system_prompt

        return payload

    def _parse_message_content(self, msg: Any) -> Tuple[str, List[ToolCall]]:
        text_parts: List[str] = []
        tool_calls: List[ToolCall] = []

        content_list = getattr(msg, "content", []) or []
        for block in content_list:
            btype = getattr(block, "type", None) or (
                block.get("type") if isinstance(block, dict) else None
            )
            if btype == "text":
                # SDK returns block.text for typed object; dict uses {"text": ...}
                text = getattr(block, "text", None)
                if text is None and isinstance(block, dict):
                    text = block.get("text")
                if text:
                    text_parts.append(str(text))
            elif btype == "tool_use":
                # Tool call with name and input
                name = getattr(block, "name", None) or (
                    block.get("name") if isinstance(block, dict) else None
                )
                tc_id = getattr(block, "id", None) or (
                    block.get("id") if isinstance(block, dict) else None
                )
                input_data = getattr(block, "input", None) or (
                    block.get("input") if isinstance(block, dict) else None
                )
                if name:
                    try:
                        # input_data should be a dict already
                        args = (
                            input_data
                            if isinstance(input_data, dict)
                            else {"_raw": input_data}
                        )
                    except Exception:
                        args = {"_raw": str(input_data)}
                    tool_calls.append(
                        ToolCall(
                            id=str(tc_id or "tool_call"), name=str(name), arguments=args
                        )
                    )

        text_content = "".join(text_parts)
        return text_content, tool_calls


================================================
FILE: src/vanna/integrations/azureopenai/__init__.py
================================================
"""
Azure OpenAI integration.

This module provides Azure OpenAI LLM service implementations.
"""

from .llm import AzureOpenAILlmService

__all__ = ["AzureOpenAILlmService"]


================================================
FILE: src/vanna/integrations/azureopenai/llm.py
================================================
"""
Azure OpenAI LLM service implementation.

Provides an `LlmService` backed by Azure OpenAI Chat Completions (openai>=1.0.0)
with support for streaming, deployment-scoped models, and Azure-specific
authentication flows.
"""

from __future__ import annotations

import json
import os
from typing import Any, AsyncGenerator, Dict, List, Optional, Set

from vanna.core.llm import (
    LlmService,
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
)
from vanna.core.tool import ToolCall, ToolSchema


# Models that don't support temperature and other sampling parameters
REASONING_MODELS: Set[str] = {
    "o1",
    "o1-mini",
    "o1-preview",
    "o3-mini",
    "gpt-5",
    "gpt-5-mini",
    "gpt-5-nano",
    "gpt-5-pro",
    "gpt-5-codex",
}


def _is_reasoning_model(model: str) -> bool:
    """Return True when the deployment targets a reasoning-only model."""
    model_lower = model.lower()
    return any(reasoning_model in model_lower for reasoning_model in REASONING_MODELS)


class AzureOpenAILlmService(LlmService):
    """Azure OpenAI Chat Completions-backed LLM service.

    Wraps `openai.AzureOpenAI` so Vanna can talk to deployment-scoped models
    and either API key or Microsoft Entra ID authentication.

    Args:
        model: Deployment name in Azure OpenAI (required).
        api_key: API key; falls back to `AZURE_OPENAI_API_KEY`.
        azure_endpoint: Azure OpenAI endpoint URL; falls back to
            `AZURE_OPENAI_ENDPOINT`.
        api_version: API version; defaults to "2024-10-21" or
            `AZURE_OPENAI_API_VERSION`.
        azure_ad_token_provider: Optional bearer token provider for Entra ID.
        **extra_client_kwargs: Additional keyword arguments forwarded to the
            underlying client.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        azure_endpoint: Optional[str] = None,
        api_version: Optional[str] = None,
        azure_ad_token_provider: Optional[Any] = None,
        **extra_client_kwargs: Any,
    ) -> None:
        try:
            from openai import AzureOpenAI
        except Exception as e:  # pragma: no cover
            raise ImportError(
                "openai package is required. Install with: pip install 'vanna[azureopenai]' "
                "or 'pip install openai'"
            ) from e

        # Model/deployment name is required for Azure OpenAI
        self.model = model or os.getenv("AZURE_OPENAI_MODEL")
        if not self.model:
            raise ValueError(
                "model parameter (deployment name) is required for Azure OpenAI. "
                "Provide it as argument or set AZURE_OPENAI_MODEL environment variable."
            )

        # Azure endpoint is required
        azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
        if not azure_endpoint:
            raise ValueError(
                "azure_endpoint is required for Azure OpenAI. "
                "Provide it as argument or set AZURE_OPENAI_ENDPOINT environment variable."
            )

        # API version - use latest stable GA version by default
        api_version = api_version or os.getenv("AZURE_OPENAI_API_VERSION", "2024-10-21")

        # Build client kwargs
        client_kwargs: Dict[str, Any] = {
            "azure_endpoint": azure_endpoint,
            "api_version": api_version,
            **extra_client_kwargs,
        }

        # Authentication: prefer Azure AD token provider, fallback to API key
        if azure_ad_token_provider is not None:
            client_kwargs["azure_ad_token_provider"] = azure_ad_token_provider
        else:
            api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
            if not api_key:
                raise ValueError(
                    "Authentication required: provide either api_key or azure_ad_token_provider. "
                    "API key can also be set via AZURE_OPENAI_API_KEY environment variable."
                )
            client_kwargs["api_key"] = api_key

        self._client = AzureOpenAI(**client_kwargs)
        self._is_reasoning_model = _is_reasoning_model(self.model)

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a non-streaming request to Azure OpenAI and return the response."""
        payload = self._build_payload(request)

        # Call the API synchronously; this function is async but we can block here.
        resp = self._client.chat.completions.create(**payload, stream=False)

        if not resp.choices:
            return LlmResponse(content=None, tool_calls=None, finish_reason=None)

        choice = resp.choices[0]
        content: Optional[str] = getattr(choice.message, "content", None)
        tool_calls = self._extract_tool_calls_from_message(choice.message)

        usage: Dict[str, int] = {}
        if getattr(resp, "usage", None):
            usage = {
                k: int(v)
                for k, v in {
                    "prompt_tokens": getattr(resp.usage, "prompt_tokens", 0),
                    "completion_tokens": getattr(resp.usage, "completion_tokens", 0),
                    "total_tokens": getattr(resp.usage, "total_tokens", 0),
                }.items()
            }

        return LlmResponse(
            content=content,
            tool_calls=tool_calls or None,
            finish_reason=getattr(choice, "finish_reason", None),
            usage=usage or None,
        )

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """
        Stream a request to Azure OpenAI.

        Emits `LlmStreamChunk` for textual deltas as they arrive. Tool-calls are
        accumulated and emitted in a final chunk when the stream ends.
        """
        payload = self._build_payload(request)

        # Synchronous streaming iterator; iterate within async context.
        stream = self._client.chat.completions.create(**payload, stream=True)

        # Builders for streamed tool-calls (index -> partial)
        tc_builders: Dict[int, Dict[str, Optional[str]]] = {}
        last_finish: Optional[str] = None

        for event in stream:
            if not getattr(event, "choices", None):
                continue

            choice = event.choices[0]
            delta = getattr(choice, "delta", None)
            if delta is None:
                # Some SDK versions use `event.choices[0].message` on the final packet
                last_finish = getattr(choice, "finish_reason", last_finish)
                continue

            # Text content
            content_piece: Optional[str] = getattr(delta, "content", None)
            if content_piece:
                yield LlmStreamChunk(content=content_piece)

            # Tool calls (streamed)
            streamed_tool_calls = getattr(delta, "tool_calls", None)
            if streamed_tool_calls:
                for tc in streamed_tool_calls:
                    idx = getattr(tc, "index", 0) or 0
                    b = tc_builders.setdefault(
                        idx, {"id": None, "name": None, "arguments": ""}
                    )
                    if getattr(tc, "id", None):
                        b["id"] = tc.id
                    fn = getattr(tc, "function", None)
                    if fn is not None:
                        if getattr(fn, "name", None):
                            b["name"] = fn.name
                        if getattr(fn, "arguments", None):
                            b["arguments"] = (b["arguments"] or "") + fn.arguments

            last_finish = getattr(choice, "finish_reason", last_finish)

        # Emit final tool-calls chunk if any
        final_tool_calls: List[ToolCall] = []
        for b in tc_builders.values():
            if not b.get("name"):
                continue
            args_raw = b.get("arguments") or "{}"
            try:
                loaded = json.loads(args_raw)
                if isinstance(loaded, dict):
                    args_dict: Dict[str, Any] = loaded
                else:
                    args_dict = {"args": loaded}
            except Exception:
                args_dict = {"_raw": args_raw}
            final_tool_calls.append(
                ToolCall(
                    id=b.get("id") or "tool_call",
                    name=b["name"] or "tool",
                    arguments=args_dict,
                )
            )

        if final_tool_calls:
            yield LlmStreamChunk(tool_calls=final_tool_calls, finish_reason=last_finish)
        else:
            # Still emit a terminal chunk to signal completion
            yield LlmStreamChunk(finish_reason=last_finish or "stop")

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Validate tool schemas. Returns a list of error messages."""
        errors: List[str] = []
        # Basic checks; Azure OpenAI will enforce further validation server-side.
        for t in tools:
            if not t.name or len(t.name) > 64:
                errors.append(f"Invalid tool name: {t.name!r}")
        return errors

    # Internal helpers
    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
        """Build the API payload from LlmRequest."""
        messages: List[Dict[str, Any]] = []

        # Add system prompt as first message if provided
        if request.system_prompt:
            messages.append({"role": "system", "content": request.system_prompt})

        for m in request.messages:
            msg: Dict[str, Any] = {"role": m.role, "content": m.content}
            if m.role == "tool" and m.tool_call_id:
                msg["tool_call_id"] = m.tool_call_id
            elif m.role == "assistant" and m.tool_calls:
                # Convert tool calls to OpenAI format
                tool_calls_payload = []
                for tc in m.tool_calls:
                    tool_calls_payload.append(
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.name,
                                "arguments": json.dumps(tc.arguments),
                            },
                        }
                    )
                msg["tool_calls"] = tool_calls_payload
            messages.append(msg)

        tools_payload: Optional[List[Dict[str, Any]]] = None
        if request.tools:
            tools_payload = [
                {
                    "type": "function",
                    "function": {
                        "name": t.name,
                        "description": t.description,
                        "parameters": t.parameters,
                    },
                }
                for t in request.tools
            ]

        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }

        # Add temperature only for non-reasoning models
        # Reasoning models (GPT-5, o1, o3-mini) don't support temperature parameter
        if not self._is_reasoning_model:
            payload["temperature"] = request.temperature

        if request.max_tokens is not None:
            payload["max_tokens"] = request.max_tokens

        if tools_payload:
            payload["tools"] = tools_payload
            payload["tool_choice"] = "auto"

        return payload

    def _extract_tool_calls_from_message(self, message: Any) -> List[ToolCall]:
        """Extract tool calls from OpenAI message object."""
        tool_calls: List[ToolCall] = []
        raw_tool_calls = getattr(message, "tool_calls", None) or []
        for tc in raw_tool_calls:
            fn = getattr(tc, "function", None)
            if not fn:
                continue
            args_raw = getattr(fn, "arguments", "{}")
            try:
                loaded = json.loads(args_raw)
                if isinstance(loaded, dict):
                    args_dict: Dict[str, Any] = loaded
                else:
                    args_dict = {"args": loaded}
            except Exception:
                args_dict = {"_raw": args_raw}
            tool_calls.append(
                ToolCall(
                    id=getattr(tc, "id", "tool_call"),
                    name=getattr(fn, "name", "tool"),
                    arguments=args_dict,
                )
            )
        return tool_calls


================================================
FILE: src/vanna/integrations/azuresearch/__init__.py
================================================
"""
Azure AI Search integration for Vanna Agents.
"""

from .agent_memory import AzureAISearchAgentMemory

__all__ = ["AzureAISearchAgentMemory"]


================================================
FILE: src/vanna/integrations/azuresearch/agent_memory.py
================================================
"""
Azure AI Search implementation of AgentMemory.

This implementation uses Azure Cognitive Search for vector storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    from azure.search.documents import SearchClient
    from azure.search.documents.indexes import SearchIndexClient
    from azure.search.documents.indexes.models import (
        SearchIndex,
        SearchField,
        SearchFieldDataType,
        VectorSearch,
        VectorSearchAlgorithmConfiguration,
    )
    from azure.core.credentials import AzureKeyCredential

    AZURE_SEARCH_AVAILABLE = True
except ImportError:
    AZURE_SEARCH_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class AzureAISearchAgentMemory(AgentMemory):
    """Azure AI Search-based implementation of AgentMemory."""

    def __init__(
        self,
        endpoint: str,
        api_key: str,
        index_name: str = "tool-memories",
        dimension: int = 384,
    ):
        if not AZURE_SEARCH_AVAILABLE:
            raise ImportError(
                "Azure Search is required for AzureAISearchAgentMemory. "
                "Install with: pip install azure-search-documents"
            )

        self.endpoint = endpoint
        self.api_key = api_key
        self.index_name = index_name
        self.dimension = dimension
        self._credential = AzureKeyCredential(api_key)
        self._search_client = None
        self._index_client = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_index_client(self):
        """Get or create index client."""
        if self._index_client is None:
            self._index_client = SearchIndexClient(
                endpoint=self.endpoint, credential=self._credential
            )
            self._ensure_index_exists()
        return self._index_client

    def _get_search_client(self):
        """Get or create search client."""
        if self._search_client is None:
            self._get_index_client()  # Ensure index exists
            self._search_client = SearchClient(
                endpoint=self.endpoint,
                index_name=self.index_name,
                credential=self._credential,
            )
        return self._search_client

    def _ensure_index_exists(self):
        """Create index if it doesn't exist."""
        try:
            self._index_client.get_index(self.index_name)
        except Exception:
            # Create index with vector search configuration
            fields = [
                SearchField(
                    name="memory_id", type=SearchFieldDataType.String, key=True
                ),
                SearchField(
                    name="question", type=SearchFieldDataType.String, searchable=True
                ),
                SearchField(
                    name="tool_name", type=SearchFieldDataType.String, filterable=True
                ),
                SearchField(name="args_json", type=SearchFieldDataType.String),
                SearchField(
                    name="timestamp",
                    type=SearchFieldDataType.String,
                    sortable=True,
                    filterable=True,
                ),
                SearchField(
                    name="success", type=SearchFieldDataType.Boolean, filterable=True
                ),
                SearchField(name="metadata_json", type=SearchFieldDataType.String),
                SearchField(
                    name="embedding",
                    type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                    searchable=True,
                    vector_search_dimensions=self.dimension,
                    vector_search_configuration="vector-config",
                ),
            ]

            vector_search = VectorSearch(
                algorithm_configurations=[
                    VectorSearchAlgorithmConfiguration(name="vector-config")
                ]
            )

            index = SearchIndex(
                name=self.index_name, fields=fields, vector_search=vector_search
            )

            self._index_client.create_index(index)

    def _create_embedding(self, text: str) -> List[float]:
        """Create a simple embedding from text (placeholder)."""
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            client = self._get_search_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            document = {
                "memory_id": memory_id,
                "question": question,
                "tool_name": tool_name,
                "args_json": json.dumps(args),
                "timestamp": timestamp,
                "success": success,
                "metadata_json": json.dumps(metadata or {}),
                "embedding": embedding,
            }

            client.upload_documents(documents=[document])

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            client = self._get_search_client()

            embedding = self._create_embedding(question)

            # Build filter
            filter_expr = "success eq true"
            if tool_name_filter:
                filter_expr += f" and tool_name eq '{tool_name_filter}'"

            results = client.search(
                search_text=None, vector=embedding, top_k=limit, filter=filter_expr
            )

            search_results = []
            for i, doc in enumerate(results):
                # Azure returns similarity score in @search.score
                similarity_score = doc.get("@search.score", 0)

                if similarity_score >= similarity_threshold:
                    args = json.loads(doc.get("args_json", "{}"))
                    metadata_dict = json.loads(doc.get("metadata_json", "{}"))

                    memory = ToolMemory(
                        memory_id=doc["memory_id"],
                        question=doc["question"],
                        tool_name=doc["tool_name"],
                        args=args,
                        timestamp=doc.get("timestamp"),
                        success=doc.get("success", True),
                        metadata=metadata_dict,
                    )

                    search_results.append(
                        ToolMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            client = self._get_search_client()

            results = client.search(
                search_text="*", top=limit, order_by=["timestamp desc"]
            )

            memories = []
            for doc in results:
                args = json.loads(doc.get("args_json", "{}"))
                metadata_dict = json.loads(doc.get("metadata_json", "{}"))

                memory = ToolMemory(
                    memory_id=doc["memory_id"],
                    question=doc["question"],
                    tool_name=doc["tool_name"],
                    args=args,
                    timestamp=doc.get("timestamp"),
                    success=doc.get("success", True),
                    metadata=metadata_dict,
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            client = self._get_search_client()

            try:
                client.delete_documents(documents=[{"memory_id": memory_id}])
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            client = self._get_search_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            document = {
                "memory_id": memory_id,
                "content": content,
                "timestamp": timestamp,
                "embedding": embedding,
            }

            client.upload_documents(documents=[document])

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            client = self._get_search_client()

            embedding = self._create_embedding(query)

            results = client.search(search_text=None, vector=embedding, top_k=limit)

            search_results = []
            for i, doc in enumerate(results):
                similarity_score = doc.get("@search.score", 0)

                if similarity_score >= similarity_threshold:
                    memory = TextMemory(
                        memory_id=doc["memory_id"],
                        content=doc.get("content", ""),
                        timestamp=doc.get("timestamp"),
                    )

                    search_results.append(
                        TextMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            client = self._get_search_client()

            results = client.search(
                search_text="*", top=limit, order_by=["timestamp desc"]
            )

            memories = []
            for doc in results:
                # Skip if this is a tool memory (has tool_name field)
                if "tool_name" in doc:
                    continue

                memory = TextMemory(
                    memory_id=doc["memory_id"],
                    content=doc.get("content", ""),
                    timestamp=doc.get("timestamp"),
                )
                memories.append(memory)

            return memories[:limit]

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            client = self._get_search_client()

            try:
                client.delete_documents(documents=[{"memory_id": memory_id}])
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            client = self._get_search_client()

            # Build filter
            filter_parts = []
            if tool_name:
                filter_parts.append(f"tool_name eq '{tool_name}'")
            if before_date:
                filter_parts.append(f"timestamp lt '{before_date}'")

            filter_expr = " and ".join(filter_parts) if filter_parts else None

            # Search for documents to delete
            results = client.search(
                search_text="*", filter=filter_expr, select=["memory_id"]
            )

            docs_to_delete = [{"memory_id": doc["memory_id"]} for doc in results]

            if docs_to_delete:
                client.delete_documents(documents=docs_to_delete)

            return len(docs_to_delete)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/bigquery/__init__.py
================================================
"""BigQuery integration for Vanna."""

from .sql_runner import BigQueryRunner

__all__ = ["BigQueryRunner"]


================================================
FILE: src/vanna/integrations/bigquery/sql_runner.py
================================================
"""BigQuery implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class BigQueryRunner(SqlRunner):
    """BigQuery implementation of the SqlRunner interface."""

    def __init__(self, project_id: str, cred_file_path: Optional[str] = None, **kwargs):
        """Initialize with BigQuery connection parameters.

        Args:
            project_id: Google Cloud Project ID
            cred_file_path: Path to Google Cloud credentials JSON file (optional)
            **kwargs: Additional google.cloud.bigquery.Client parameters
        """
        try:
            from google.cloud import bigquery
            from google.oauth2 import service_account

            self.bigquery = bigquery
            self.service_account = service_account
        except ImportError as e:
            raise ImportError(
                "google-cloud-bigquery package is required. "
                "Install with: pip install 'vanna[bigquery]'"
            ) from e

        self.project_id = project_id
        self.cred_file_path = cred_file_path
        self.kwargs = kwargs
        self._client = None

    def _get_client(self):
        """Get or create BigQuery client."""
        if self._client is not None:
            return self._client

        if self.cred_file_path:
            import json

            with open(self.cred_file_path, "r") as f:
                credentials = (
                    self.service_account.Credentials.from_service_account_info(
                        json.loads(f.read()),
                        scopes=["https://www.googleapis.com/auth/cloud-platform"],
                    )
                )
            self._client = self.bigquery.Client(
                project=self.project_id, credentials=credentials, **self.kwargs
            )
        else:
            # Use default credentials
            self._client = self.bigquery.Client(project=self.project_id, **self.kwargs)

        return self._client

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against BigQuery database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            google.api_core.exceptions.GoogleAPIError: If query execution fails
        """
        client = self._get_client()

        # Execute the query
        job = client.query(args.sql)
        df = job.result().to_dataframe()

        return df


================================================
FILE: src/vanna/integrations/chromadb/__init__.py
================================================
"""
ChromaDB integration for Vanna Agents.
"""

from .agent_memory import ChromaAgentMemory


def get_device() -> str:
    """Detect the best available device for embeddings.

    This function checks for GPU availability and returns the appropriate device string
    for use with embedding models. It prioritizes hardware acceleration when available.

    Returns:
        str: Device string - 'cuda' if NVIDIA GPU available, 'mps' if Apple Silicon,
             'cpu' otherwise.

    Examples:
        >>> device = get_device()
        >>> print(f"Using device: {device}")
        Using device: cuda

        # Use with ChromaDB SentenceTransformer embeddings
        >>> from chromadb.utils import embedding_functions
        >>> ef = embedding_functions.SentenceTransformerEmbeddingFunction(
        ...     model_name="sentence-transformers/all-MiniLM-L6-v2",
        ...     device=get_device()
        ... )
        >>> memory = ChromaAgentMemory(embedding_function=ef)
    """
    try:
        import torch

        # Check for CUDA (NVIDIA GPUs)
        if torch.cuda.is_available():
            return "cuda"

        # Check for MPS (Apple Silicon GPUs)
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"

    except ImportError:
        # PyTorch not installed, fall back to CPU
        pass

    return "cpu"


def create_sentence_transformer_embedding_function(
    model_name: str = "sentence-transformers/all-MiniLM-L6-v2", device: str = None
):
    """Create a SentenceTransformer embedding function with automatic device detection.

    This convenience function creates a ChromaDB-compatible SentenceTransformer embedding
    function with intelligent device selection. If no device is specified, it automatically
    detects and uses the best available hardware (CUDA, MPS, or CPU).

    Note: This requires the 'sentence-transformers' package to be installed.
    Install with: pip install sentence-transformers

    Args:
        model_name: The name of the sentence-transformer model to use.
                   Defaults to "sentence-transformers/all-MiniLM-L6-v2".
        device: Optional device string ('cuda', 'mps', or 'cpu'). If None,
               automatically detects the best available device.

    Returns:
        A ChromaDB SentenceTransformer embedding function configured for the
        specified/detected device.

    Examples:
        # Automatic device detection (uses CUDA/MPS if available)
        >>> from vanna.integrations.chromadb import ChromaAgentMemory, create_sentence_transformer_embedding_function
        >>> ef = create_sentence_transformer_embedding_function()
        >>> memory = ChromaAgentMemory(embedding_function=ef)

        # Explicitly use CUDA
        >>> ef_cuda = create_sentence_transformer_embedding_function(device="cuda")
        >>> memory = ChromaAgentMemory(embedding_function=ef_cuda)

        # Use a different model
        >>> ef_large = create_sentence_transformer_embedding_function(
        ...     model_name="sentence-transformers/all-mpnet-base-v2"
        ... )
        >>> memory = ChromaAgentMemory(embedding_function=ef_large)
    """
    try:
        from chromadb.utils import embedding_functions
    except ImportError:
        raise ImportError("ChromaDB is required. Install with: pip install chromadb")

    if device is None:
        device = get_device()

    return embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name=model_name, device=device
    )


__all__ = [
    "ChromaAgentMemory",
    "get_device",
    "create_sentence_transformer_embedding_function",
]


================================================
FILE: src/vanna/integrations/chromadb/agent_memory.py
================================================
"""
Local vector database implementation of AgentMemory.

This implementation uses ChromaDB for local vector storage of tool usage patterns.
"""

import json
import hashlib
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    import chromadb
    from chromadb.config import Settings
    from chromadb.utils import embedding_functions

    try:
        from chromadb.errors import NotFoundError
    except ImportError:
        # Fallback for older ChromaDB versions that don't have chromadb.errors
        class NotFoundError(Exception):
            """Fallback NotFoundError for older ChromaDB versions."""

            pass

    CHROMADB_AVAILABLE = True
except ImportError:
    CHROMADB_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class ChromaAgentMemory(AgentMemory):
    """ChromaDB-based implementation of AgentMemory.

    This implementation uses ChromaDB's PersistentClient to store agent memories
    on disk, ensuring they persist across application restarts.

    Key Features:
    - Persistent storage: All memories are automatically saved to disk
    - Efficient retrieval: Existing collections are loaded without re-initializing
      embedding functions, avoiding unnecessary model downloads
    - Flexible embedding: Supports custom embedding functions or uses ChromaDB's
      default embedding function

    Args:
        persist_directory: Directory where ChromaDB will store its data.
                          Defaults to "./chroma_memory". Use an absolute path
                          for production deployments to ensure consistent location
                          across restarts.
        collection_name: Name of the ChromaDB collection to use. Multiple agents
                        can share the same persist_directory with different
                        collection names.
        embedding_function: Optional custom embedding function. If not provided,
                           ChromaDB's DefaultEmbeddingFunction is used (requires
                           internet connection on first use to download the model).
                           Once a collection is created, subsequent application
                           restarts will retrieve the existing collection without
                           re-downloading the model.

    Example:
        >>> from vanna.integrations.chromadb import ChromaAgentMemory
        >>> # Basic usage with defaults
        >>> memory = ChromaAgentMemory(
        ...     persist_directory="/app/data/chroma",
        ...     collection_name="my_agent_memory"
        ... )
        >>>
        >>> # With custom embedding function (e.g., for offline use)
        >>> from chromadb.utils import embedding_functions
        >>> ef = embedding_functions.SentenceTransformerEmbeddingFunction()
        >>> memory = ChromaAgentMemory(
        ...     persist_directory="/app/data/chroma",
        ...     embedding_function=ef
        ... )

    Note:
        The default embedding function downloads an ONNX model (~80MB) on first use.
        For air-gapped or offline environments, pre-download the model or provide
        a custom embedding function.

    Limitation:
        This class does not validate that an existing Chroma collection was created
        with the same embedding function as the one configured for the current
        ``ChromaAgentMemory`` instance. If you reuse a collection (same
        ``persist_directory`` and ``collection_name``) with a different embedding
        function than was originally used, queries may fail or produce incorrect
        similarity results. It is your responsibility to ensure that a given
        collection is always accessed with a consistent embedding function, or to
        implement your own validation around collection creation and reuse.
    """

    def __init__(
        self,
        persist_directory: str = "./chroma_memory",
        collection_name: str = "tool_memories",
        embedding_function=None,
    ):
        if not CHROMADB_AVAILABLE:
            raise ImportError(
                "ChromaDB is required for ChromaAgentMemory. Install with: pip install chromadb"
            )

        self.persist_directory = persist_directory
        self.collection_name = collection_name
        self._client = None
        self._collection = None
        self._executor = ThreadPoolExecutor(max_workers=2)
        self._embedding_function = embedding_function

    def _get_client(self):
        """Get or create ChromaDB client."""
        if self._client is None:
            self._client = chromadb.PersistentClient(
                path=self.persist_directory,
                settings=Settings(anonymized_telemetry=False, allow_reset=True),
            )
        return self._client

    def _get_embedding_function(self):
        """Get or create the embedding function.

        If no embedding function was provided during initialization,
        uses ChromaDB's default embedding function.
        """
        if self._embedding_function is None:
            # Use ChromaDB's default embedding function
            # This avoids requiring sentence-transformers as a hard dependency
            self._embedding_function = embedding_functions.DefaultEmbeddingFunction()
        return self._embedding_function

    def _get_collection(self):
        """Get or create ChromaDB collection."""
        if self._collection is None:
            client = self._get_client()
            try:
                # Try to get existing collection first
                # Don't pass embedding_function here to avoid re-instantiating/downloading it
                # For existing collections, ChromaDB uses the stored embedding function configuration
                self._collection = client.get_collection(name=self.collection_name)
            except NotFoundError:
                # Collection doesn't exist, create it with embedding function
                embedding_func = self._get_embedding_function()
                self._collection = client.create_collection(
                    name=self.collection_name,
                    embedding_function=embedding_func,
                    metadata={"description": "Tool usage memories for learning"},
                )
        return self._collection

    def _create_memory_id(self) -> str:
        """Create a unique ID for a memory."""
        import uuid

        return str(uuid.uuid4())

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            collection = self._get_collection()

            memory_id = self._create_memory_id()
            timestamp = datetime.now().isoformat()

            # ChromaDB only accepts primitive types in metadata
            # Serialize complex objects to JSON strings
            memory_data = {
                "question": question,
                "tool_name": tool_name,
                "args_json": json.dumps(args),  # Serialize to JSON string
                "timestamp": timestamp,
                "success": success,
                "metadata_json": json.dumps(metadata or {}),  # Serialize metadata too
            }

            # Use question as document text for embedding
            collection.upsert(
                ids=[memory_id], documents=[question], metadatas=[memory_data]
            )

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            collection = self._get_collection()

            # Prepare where filter - ChromaDB requires $and for multiple conditions
            if tool_name_filter:
                where_filter = {
                    "$and": [{"success": True}, {"tool_name": tool_name_filter}]
                }
            else:
                where_filter = {"success": True}

            results = collection.query(
                query_texts=[question], n_results=limit, where=where_filter
            )

            search_results = []
            if results["ids"] and len(results["ids"][0]) > 0:
                for i, (id_, distance, metadata) in enumerate(
                    zip(
                        results["ids"][0],
                        results["distances"][0],
                        results["metadatas"][0],
                    )
                ):
                    # Convert distance to similarity score (ChromaDB uses L2 distance)
                    similarity_score = max(0, 1 - distance)

                    if similarity_score >= similarity_threshold:
                        # Deserialize JSON fields
                        args = json.loads(metadata.get("args_json", "{}"))
                        metadata_dict = json.loads(metadata.get("metadata_json", "{}"))

                        # Use the ChromaDB document ID as the memory ID
                        memory = ToolMemory(
                            memory_id=id_,
                            question=metadata["question"],
                            tool_name=metadata["tool_name"],
                            args=args,
                            timestamp=metadata.get("timestamp"),
                            success=metadata.get("success", True),
                            metadata=metadata_dict,
                        )

                        search_results.append(
                            ToolMemorySearchResult(
                                memory=memory,
                                similarity_score=similarity_score,
                                rank=i + 1,
                            )
                        )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories. Returns most recent memories first."""

        def _get_recent():
            collection = self._get_collection()

            # Get all memories and sort by timestamp
            results = collection.get()

            if not results["metadatas"] or not results["ids"]:
                return []

            # Parse and sort by timestamp
            memories_with_time = []
            for i, (doc_id, metadata) in enumerate(
                zip(results["ids"], results["metadatas"])
            ):
                # Skip text memories - they have is_text_memory flag
                if metadata.get("is_text_memory"):
                    continue

                args = json.loads(metadata.get("args_json", "{}"))
                metadata_dict = json.loads(metadata.get("metadata_json", "{}"))

                # Use the ChromaDB document ID as the memory ID
                memory = ToolMemory(
                    memory_id=doc_id,
                    question=metadata["question"],
                    tool_name=metadata["tool_name"],
                    args=args,
                    timestamp=metadata.get("timestamp"),
                    success=metadata.get("success", True),
                    metadata=metadata_dict,
                )
                memories_with_time.append((memory, metadata.get("timestamp", "")))

            # Sort by timestamp descending (most recent first)
            memories_with_time.sort(key=lambda x: x[1], reverse=True)

            # Return only the memory objects, limited to the requested amount
            return [m[0] for m in memories_with_time[:limit]]

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID. Returns True if deleted, False if not found."""

        def _delete():
            collection = self._get_collection()

            # Check if the ID exists
            try:
                results = collection.get(ids=[memory_id])
                if results["ids"] and len(results["ids"]) > 0:
                    collection.delete(ids=[memory_id])
                    return True
                return False
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            collection = self._get_collection()

            memory_id = self._create_memory_id()
            timestamp = datetime.now().isoformat()

            memory_data = {
                "content": content,
                "timestamp": timestamp,
                "is_text_memory": True,
            }

            collection.upsert(
                ids=[memory_id], documents=[content], metadatas=[memory_data]
            )

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            collection = self._get_collection()

            where_filter = {"is_text_memory": True}

            results = collection.query(
                query_texts=[query], n_results=limit, where=where_filter
            )

            search_results = []
            if results["ids"] and len(results["ids"][0]) > 0:
                for i, (id_, distance, metadata) in enumerate(
                    zip(
                        results["ids"][0],
                        results["distances"][0],
                        results["metadatas"][0],
                    )
                ):
                    similarity_score = max(0, 1 - distance)

                    if similarity_score >= similarity_threshold:
                        memory = TextMemory(
                            memory_id=id_,
                            content=metadata.get("content", ""),
                            timestamp=metadata.get("timestamp"),
                        )

                        search_results.append(
                            TextMemorySearchResult(
                                memory=memory,
                                similarity_score=similarity_score,
                                rank=i + 1,
                            )
                        )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            collection = self._get_collection()

            results = collection.get(where={"is_text_memory": True})

            if not results["metadatas"] or not results["ids"]:
                return []

            memories_with_time = []
            for doc_id, metadata in zip(results["ids"], results["metadatas"]):
                memory = TextMemory(
                    memory_id=doc_id,
                    content=metadata.get("content", ""),
                    timestamp=metadata.get("timestamp"),
                )
                memories_with_time.append((memory, metadata.get("timestamp", "")))

            memories_with_time.sort(key=lambda x: x[1], reverse=True)

            return [m[0] for m in memories_with_time[:limit]]

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            collection = self._get_collection()

            try:
                results = collection.get(ids=[memory_id])
                if results["ids"] and len(results["ids"]) > 0:
                    collection.delete(ids=[memory_id])
                    return True
                return False
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            collection = self._get_collection()

            # Build where filter
            where_filter = {}
            if tool_name:
                where_filter["tool_name"] = tool_name

            # Get memories to delete
            results = collection.get(where=where_filter if where_filter else None)

            if not results["ids"]:
                return 0

            ids_to_delete = []
            for i, metadata in enumerate(results["metadatas"]):
                if before_date:
                    memory_date = metadata.get("timestamp", "")
                    if memory_date and memory_date < before_date:
                        ids_to_delete.append(results["ids"][i])
                else:
                    ids_to_delete.append(results["ids"][i])

            if ids_to_delete:
                collection.delete(ids=ids_to_delete)

            return len(ids_to_delete)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/clickhouse/__init__.py
================================================
"""ClickHouse integration for Vanna."""

from .sql_runner import ClickHouseRunner

__all__ = ["ClickHouseRunner"]


================================================
FILE: src/vanna/integrations/clickhouse/sql_runner.py
================================================
"""ClickHouse implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class ClickHouseRunner(SqlRunner):
    """ClickHouse implementation of the SqlRunner interface."""

    def __init__(
        self,
        host: str,
        database: str,
        user: str,
        password: str,
        port: int = 8123,
        **kwargs,
    ):
        """Initialize with ClickHouse connection parameters.

        Args:
            host: Database host address
            database: Database name
            user: Database user
            password: Database password
            port: Database port (default: 8123)
            **kwargs: Additional clickhouse_connect connection parameters
        """
        try:
            import clickhouse_connect

            self.clickhouse_connect = clickhouse_connect
        except ImportError as e:
            raise ImportError(
                "clickhouse-connect package is required. "
                "Install with: pip install 'vanna[clickhouse]'"
            ) from e

        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.kwargs = kwargs

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against ClickHouse database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            Exception: If query execution fails
        """
        # Connect to the database
        client = self.clickhouse_connect.get_client(
            host=self.host,
            port=self.port,
            username=self.user,
            password=self.password,
            database=self.database,
            **self.kwargs,
        )

        try:
            # Execute the query
            result = client.query(args.sql)
            results = result.result_rows

            # Create a pandas dataframe from the results
            df = pd.DataFrame(results, columns=result.column_names)
            return df

        finally:
            client.close()


================================================
FILE: src/vanna/integrations/duckdb/__init__.py
================================================
"""DuckDB integration for Vanna."""

from .sql_runner import DuckDBRunner

__all__ = ["DuckDBRunner"]


================================================
FILE: src/vanna/integrations/duckdb/sql_runner.py
================================================
"""DuckDB implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class DuckDBRunner(SqlRunner):
    """DuckDB implementation of the SqlRunner interface."""

    def __init__(
        self, database_path: str = ":memory:", init_sql: Optional[str] = None, **kwargs
    ):
        """Initialize with DuckDB connection parameters.

        Args:
            database_path: Path to the DuckDB database file.
                          Use ":memory:" for in-memory database (default).
                          Use "md:" or "motherduck:" for MotherDuck database.
            init_sql: Optional SQL to run when connecting to the database
            **kwargs: Additional duckdb connection parameters
        """
        try:
            import duckdb

            self.duckdb = duckdb
        except ImportError as e:
            raise ImportError(
                "duckdb package is required. Install with: pip install 'vanna[duckdb]'"
            ) from e

        self.database_path = database_path
        self.init_sql = init_sql
        self.kwargs = kwargs
        self._conn = None

    def _get_connection(self):
        """Get or create DuckDB connection."""
        if self._conn is None:
            self._conn = self.duckdb.connect(self.database_path, **self.kwargs)
            if self.init_sql:
                self._conn.query(self.init_sql)
        return self._conn

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against DuckDB database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            duckdb.Error: If query execution fails
        """
        conn = self._get_connection()

        # Execute the query and convert to DataFrame
        df = conn.query(args.sql).to_df()

        return df


================================================
FILE: src/vanna/integrations/faiss/__init__.py
================================================
"""
FAISS integration for Vanna Agents.
"""

from .agent_memory import FAISSAgentMemory

__all__ = ["FAISSAgentMemory"]


================================================
FILE: src/vanna/integrations/faiss/agent_memory.py
================================================
"""
FAISS vector database implementation of AgentMemory.

This implementation uses FAISS for local vector storage of tool usage patterns.
"""

import json
import uuid
import pickle
import os
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor
import numpy as np

try:
    import faiss

    FAISS_AVAILABLE = True
except ImportError:
    FAISS_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class FAISSAgentMemory(AgentMemory):
    """FAISS-based implementation of AgentMemory."""

    def __init__(
        self,
        index_path: Optional[str] = None,
        persist_path: Optional[str] = None,
        dimension: int = 384,
        metric: str = "cosine",
    ):
        if not FAISS_AVAILABLE:
            raise ImportError(
                "FAISS is required for FAISSAgentMemory. Install with: pip install faiss-cpu"
            )

        # Accept either index_path or persist_path for backward compatibility
        self.index_path = persist_path or index_path or "./faiss_index"
        self.dimension = dimension
        self.metric = metric
        self._index = None
        self._metadata = {}
        self._executor = ThreadPoolExecutor(max_workers=2)
        self._load_index()

    def _load_index(self):
        """Load or create FAISS index."""
        index_file = os.path.join(self.index_path, "index.faiss")
        metadata_file = os.path.join(self.index_path, "metadata.pkl")

        if os.path.exists(index_file) and os.path.exists(metadata_file):
            # Load existing index
            self._index = faiss.read_index(index_file)
            with open(metadata_file, "rb") as f:
                self._metadata = pickle.load(f)
        else:
            # Create new index
            os.makedirs(self.index_path, exist_ok=True)
            if self.metric == "cosine":
                self._index = faiss.IndexFlatIP(self.dimension)
            else:
                self._index = faiss.IndexFlatL2(self.dimension)
            self._metadata = {}

    def _save_index(self):
        """Save FAISS index to disk."""
        index_file = os.path.join(self.index_path, "index.faiss")
        metadata_file = os.path.join(self.index_path, "metadata.pkl")

        faiss.write_index(self._index, index_file)
        with open(metadata_file, "wb") as f:
            pickle.dump(self._metadata, f)

    def _create_embedding(self, text: str) -> np.ndarray:
        """Create a simple embedding from text (placeholder)."""
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        embedding = np.array(
            [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)],
            dtype=np.float32,
        )

        # Normalize for cosine similarity
        if self.metric == "cosine":
            norm = np.linalg.norm(embedding)
            if norm > 0:
                embedding = embedding / norm

        return embedding

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            # Add to FAISS index
            self._index.add(np.array([embedding]))

            # Store metadata
            idx = self._index.ntotal - 1
            self._metadata[idx] = {
                "memory_id": memory_id,
                "question": question,
                "tool_name": tool_name,
                "args": args,
                "timestamp": timestamp,
                "success": success,
                "metadata": metadata or {},
            }

            self._save_index()

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            embedding = self._create_embedding(question)

            # Search in FAISS
            k = min(limit * 3, self._index.ntotal) if self._index.ntotal > 0 else 1
            if k == 0:
                return []

            distances, indices = self._index.search(np.array([embedding]), k)

            search_results = []
            rank = 1
            for i, (dist, idx) in enumerate(zip(distances[0], indices[0])):
                if idx == -1 or idx not in self._metadata:
                    continue

                metadata = self._metadata[idx]

                # Filter by success
                if not metadata.get("success", True):
                    continue

                # Filter by tool name
                if tool_name_filter and metadata.get("tool_name") != tool_name_filter:
                    continue

                # Convert distance to similarity score
                if self.metric == "cosine":
                    similarity_score = float(dist)
                else:
                    similarity_score = 1.0 / (1.0 + float(dist))

                if similarity_score >= similarity_threshold:
                    memory = ToolMemory(
                        memory_id=metadata["memory_id"],
                        question=metadata["question"],
                        tool_name=metadata["tool_name"],
                        args=metadata["args"],
                        timestamp=metadata.get("timestamp"),
                        success=metadata.get("success", True),
                        metadata=metadata.get("metadata", {}),
                    )

                    search_results.append(
                        ToolMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=rank
                        )
                    )
                    rank += 1

                    if len(search_results) >= limit:
                        break

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            # Get all metadata entries and sort by timestamp
            all_entries = list(self._metadata.values())
            sorted_entries = sorted(
                all_entries, key=lambda m: m.get("timestamp", ""), reverse=True
            )

            memories = []
            for entry in sorted_entries[:limit]:
                # Skip text memories - they have is_text_memory flag
                if entry.get("is_text_memory"):
                    continue

                memory = ToolMemory(
                    memory_id=entry["memory_id"],
                    question=entry["question"],
                    tool_name=entry["tool_name"],
                    args=entry["args"],
                    timestamp=entry.get("timestamp"),
                    success=entry.get("success", True),
                    metadata=entry.get("metadata", {}),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            # Find and remove from metadata
            idx_to_remove = None
            for idx, metadata in self._metadata.items():
                if metadata["memory_id"] == memory_id:
                    idx_to_remove = idx
                    break

            if idx_to_remove is not None:
                del self._metadata[idx_to_remove]
                self._save_index()
                return True

            return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            # Add to FAISS index
            self._index.add(np.array([embedding]))

            # Store metadata
            idx = self._index.ntotal - 1
            self._metadata[idx] = {
                "memory_id": memory_id,
                "content": content,
                "timestamp": timestamp,
                "is_text_memory": True,
            }

            self._save_index()

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            embedding = self._create_embedding(query)

            k = min(limit * 3, self._index.ntotal) if self._index.ntotal > 0 else 1
            if k == 0:
                return []

            distances, indices = self._index.search(np.array([embedding]), k)

            search_results = []
            rank = 1
            for dist, idx in zip(distances[0], indices[0]):
                if idx == -1 or idx not in self._metadata:
                    continue

                metadata = self._metadata[idx]

                # Filter for text memories only
                if not metadata.get("is_text_memory", False):
                    continue

                # Convert distance to similarity score
                if self.metric == "cosine":
                    similarity_score = float(dist)
                else:
                    similarity_score = 1.0 / (1.0 + float(dist))

                if similarity_score >= similarity_threshold:
                    memory = TextMemory(
                        memory_id=metadata["memory_id"],
                        content=metadata["content"],
                        timestamp=metadata.get("timestamp"),
                    )

                    search_results.append(
                        TextMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=rank
                        )
                    )
                    rank += 1

                    if len(search_results) >= limit:
                        break

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            # Get all text memory entries and sort by timestamp
            text_entries = [
                entry
                for entry in self._metadata.values()
                if entry.get("is_text_memory", False)
            ]
            sorted_entries = sorted(
                text_entries, key=lambda m: m.get("timestamp", ""), reverse=True
            )

            memories = []
            for entry in sorted_entries[:limit]:
                memory = TextMemory(
                    memory_id=entry["memory_id"],
                    content=entry["content"],
                    timestamp=entry.get("timestamp"),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            # Find and remove from metadata
            idx_to_remove = None
            for idx, metadata in self._metadata.items():
                if metadata["memory_id"] == memory_id:
                    idx_to_remove = idx
                    break

            if idx_to_remove is not None:
                del self._metadata[idx_to_remove]
                self._save_index()
                return True

            return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            indices_to_remove = []

            for idx, metadata in self._metadata.items():
                should_remove = True

                if tool_name and metadata.get("tool_name") != tool_name:
                    should_remove = False

                if before_date and metadata.get("timestamp", "") >= before_date:
                    should_remove = False

                if should_remove:
                    indices_to_remove.append(idx)

            # Remove from metadata
            for idx in indices_to_remove:
                del self._metadata[idx]

            # If clearing all, recreate index
            if not tool_name and not before_date:
                if self.metric == "cosine":
                    self._index = faiss.IndexFlatIP(self.dimension)
                else:
                    self._index = faiss.IndexFlatL2(self.dimension)
                self._metadata = {}

            self._save_index()
            return len(indices_to_remove)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/google/__init__.py
================================================
"""
Google AI integrations.

This module provides Google AI service implementations.
"""

from .gemini import GeminiLlmService

__all__ = ["GeminiLlmService"]


================================================
FILE: src/vanna/integrations/google/gemini.py
================================================
"""
Google Gemini LLM service implementation.

Implements the LlmService interface using Google's Gen AI SDK
(google-genai). Supports non-streaming and streaming text output,
as well as function calling (tool use).
"""

from __future__ import annotations

import json
import logging
import os
from typing import Any, AsyncGenerator, Dict, List, Optional

logger = logging.getLogger(__name__)

from vanna.core.llm import (
    LlmService,
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
)
from vanna.core.tool import ToolCall, ToolSchema


class GeminiLlmService(LlmService):
    """Google Gemini-backed LLM service.

    Args:
        model: Gemini model name (e.g., "gemini-2.5-pro", "gemini-2.5-flash").
            Defaults to "gemini-2.5-pro". Can also be set via GEMINI_MODEL env var.
        api_key: API key; falls back to env `GOOGLE_API_KEY` or `GEMINI_API_KEY`.
            GOOGLE_API_KEY takes precedence if both are set.
        temperature: Temperature for generation (0.0-2.0). Default 0.7.
        extra_config: Extra kwargs forwarded to GenerateContentConfig.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float = 0.7,
        **extra_config: Any,
    ) -> None:
        try:
            from google import genai
            from google.genai import types
        except Exception as e:  # pragma: no cover
            raise ImportError(
                "google-genai package is required. "
                "Install with: pip install 'vanna[gemini]'"
            ) from e

        self.model_name = model or os.getenv("GEMINI_MODEL", "gemini-2.5-pro")
        # Check GOOGLE_API_KEY first (takes precedence), then GEMINI_API_KEY
        api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")

        if not api_key:
            raise ValueError(
                "Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY "
                "environment variable, or pass api_key parameter."
            )

        # Store modules for use in methods
        self._genai = genai
        self._types = types

        # Create client
        self._client = genai.Client(api_key=api_key)

        # Store generation config
        self.temperature = temperature
        self.extra_config = extra_config

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a non-streaming request to Gemini and return the response."""
        contents, config = self._build_payload(request)

        try:
            # Generate content
            response = self._client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=config,
            )

            logger.info(f"Gemini response: {response}")

            # Parse response
            text_content, tool_calls = self._parse_response(response)

            # Extract usage information
            usage: Dict[str, int] = {}
            if hasattr(response, "usage_metadata"):
                try:
                    usage = {
                        "prompt_tokens": int(
                            response.usage_metadata.prompt_token_count
                        ),
                        "completion_tokens": int(
                            response.usage_metadata.candidates_token_count
                        ),
                        "total_tokens": int(response.usage_metadata.total_token_count),
                    }
                except Exception:
                    pass

            # Get finish reason
            finish_reason = None
            if response.candidates:
                finish_reason = str(response.candidates[0].finish_reason).lower()

            return LlmResponse(
                content=text_content or None,
                tool_calls=tool_calls or None,
                finish_reason=finish_reason,
                usage=usage or None,
            )

        except Exception as e:
            logger.error(f"Error calling Gemini API: {e}")
            raise

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Stream a request to Gemini.

        Yields text chunks as they arrive. Emits tool calls at the end.
        """
        contents, config = self._build_payload(request)

        logger.info(f"Gemini streaming request with model: {self.model_name}")

        try:
            # Stream content
            stream = self._client.models.generate_content_stream(
                model=self.model_name,
                contents=contents,
                config=config,
            )

            # Accumulate chunks for tool calls
            accumulated_chunks = []

            for chunk in stream:
                accumulated_chunks.append(chunk)

                # Yield text content as it arrives
                if hasattr(chunk, "text") and chunk.text:
                    yield LlmStreamChunk(content=chunk.text)

            # After stream completes, check for tool calls in accumulated response
            if accumulated_chunks:
                final_chunk = accumulated_chunks[-1]
                _, tool_calls = self._parse_response_chunk(final_chunk)

                finish_reason = None
                if final_chunk.candidates:
                    finish_reason = str(final_chunk.candidates[0].finish_reason).lower()

                if tool_calls:
                    yield LlmStreamChunk(
                        tool_calls=tool_calls,
                        finish_reason=finish_reason,
                    )
                else:
                    yield LlmStreamChunk(finish_reason=finish_reason or "stop")

        except Exception as e:
            logger.error(f"Error streaming from Gemini API: {e}")
            raise

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Basic validation of tool schemas for Gemini."""
        errors: List[str] = []
        for t in tools:
            if not t.name:
                errors.append("Tool name is required")
            if not t.description:
                errors.append(f"Tool {t.name}: description is required")
        return errors

    # Internal helpers
    def _build_payload(self, request: LlmRequest) -> tuple[List[Any], Any]:
        """Build the payload for Gemini API.

        Returns:
            Tuple of (contents, config)
        """
        # Build contents (messages) for Gemini
        contents = []

        # System prompt handling - Gemini supports system instructions in config
        system_instruction = None
        if request.system_prompt:
            system_instruction = request.system_prompt

        for m in request.messages:
            # Map roles: user -> user, assistant -> model, tool -> function
            if m.role == "user":
                contents.append(
                    self._types.Content(
                        role="user", parts=[self._types.Part(text=m.content)]
                    )
                )
            elif m.role == "assistant":
                parts = []

                # Add text content if present
                if m.content and m.content.strip():
                    parts.append(self._types.Part(text=m.content))

                # Add tool calls if present
                if m.tool_calls:
                    for tc in m.tool_calls:
                        parts.append(
                            self._types.Part(
                                function_call=self._types.FunctionCall(
                                    name=tc.name, args=tc.arguments
                                )
                            )
                        )

                if parts:
                    contents.append(self._types.Content(role="model", parts=parts))

            elif m.role == "tool":
                # Tool results in Gemini format
                if m.tool_call_id:
                    # Parse the content as JSON if possible
                    try:
                        response_content = json.loads(m.content)
                    except (json.JSONDecodeError, TypeError):
                        response_content = {"result": m.content}

                    # Extract function name from tool_call_id or use a default
                    function_name = m.tool_call_id.replace("call_", "")

                    contents.append(
                        self._types.Content(
                            role="function",
                            parts=[
                                self._types.Part(
                                    function_response=self._types.FunctionResponse(
                                        name=function_name, response=response_content
                                    )
                                )
                            ],
                        )
                    )

        # Build tools configuration if tools are provided
        tools = None
        if request.tools:
            function_declarations = []
            for tool in request.tools:
                # Clean schema to remove unsupported fields
                cleaned_parameters = self._clean_schema_for_gemini(tool.parameters)

                function_declarations.append(
                    {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": cleaned_parameters,
                    }
                )

            if function_declarations:
                tools = [self._types.Tool(function_declarations=function_declarations)]

        # Build generation config
        config_dict = {
            "temperature": request.temperature,
            **self.extra_config,
        }

        if request.max_tokens is not None:
            config_dict["max_output_tokens"] = request.max_tokens

        if tools:
            config_dict["tools"] = tools

        if system_instruction:
            config_dict["system_instruction"] = system_instruction

        config = self._types.GenerateContentConfig(**config_dict)

        return contents, config

    def _parse_response(self, response: Any) -> tuple[str, List[ToolCall]]:
        """Parse a Gemini response into text and tool calls."""
        text_parts: List[str] = []
        tool_calls: List[ToolCall] = []

        if not response.candidates:
            return "", []

        candidate = response.candidates[0]

        if (
            hasattr(candidate, "content")
            and candidate.content
            and hasattr(candidate.content, "parts")
            and candidate.content.parts
        ):
            for part in candidate.content.parts:
                # Check for text content
                if hasattr(part, "text") and part.text:
                    text_parts.append(part.text)

                # Check for function calls
                if hasattr(part, "function_call") and part.function_call:
                    fc = part.function_call
                    # Convert function call to ToolCall
                    tool_calls.append(
                        ToolCall(
                            id=f"call_{fc.name}",  # Generate an ID
                            name=fc.name,
                            arguments=dict(fc.args) if hasattr(fc, "args") else {},
                        )
                    )

        text_content = "".join(text_parts)
        return text_content, tool_calls

    def _parse_response_chunk(self, chunk: Any) -> tuple[str, List[ToolCall]]:
        """Parse a streaming chunk (same logic as _parse_response)."""
        return self._parse_response(chunk)

    def _clean_schema_for_gemini(self, schema: Dict[str, Any]) -> Dict[str, Any]:
        """Clean JSON Schema to only include fields supported by Gemini.

        Gemini only supports a subset of OpenAPI schema. This removes unsupported
        fields like 'title', 'default', '$schema', etc.

        Supported fields:
        - type, description, enum
        - properties, required, items (for objects/arrays)
        """
        if not isinstance(schema, dict):
            return schema

        # Fields that Gemini supports
        allowed_fields = {
            "type",
            "description",
            "enum",
            "properties",
            "required",
            "items",
            "format",
        }

        cleaned = {}
        for key, value in schema.items():
            if key in allowed_fields:
                # Recursively clean nested schemas
                if key == "properties" and isinstance(value, dict):
                    cleaned[key] = {
                        prop_name: self._clean_schema_for_gemini(prop_schema)
                        for prop_name, prop_schema in value.items()
                    }
                elif key == "items" and isinstance(value, dict):
                    cleaned[key] = self._clean_schema_for_gemini(value)
                else:
                    cleaned[key] = value

        return cleaned


================================================
FILE: src/vanna/integrations/hive/__init__.py
================================================
"""Hive integration for Vanna."""

from .sql_runner import HiveRunner

__all__ = ["HiveRunner"]


================================================
FILE: src/vanna/integrations/hive/sql_runner.py
================================================
"""Hive implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class HiveRunner(SqlRunner):
    """Hive implementation of the SqlRunner interface."""

    def __init__(
        self,
        host: str,
        database: str = "default",
        user: Optional[str] = None,
        password: Optional[str] = None,
        port: int = 10000,
        auth: str = "CUSTOM",
        **kwargs,
    ):
        """Initialize with Hive connection parameters.

        Args:
            host: The host of the Hive database
            database: The name of the database to connect to (default: 'default')
            user: The username to use for authentication
            password: The password to use for authentication
            port: The port to use for the connection (default: 10000)
            auth: The authentication method to use (default: 'CUSTOM')
            **kwargs: Additional pyhive connection parameters
        """
        try:
            from pyhive import hive

            self.hive = hive
        except ImportError as e:
            raise ImportError(
                "pyhive package is required. Install with: pip install pyhive"
            ) from e

        self.host = host
        self.database = database
        self.user = user
        self.password = password
        self.port = port
        self.auth = auth
        self.kwargs = kwargs

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against Hive database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            hive.Error: If query execution fails
        """
        # Connect to the database
        conn = self.hive.Connection(
            host=self.host,
            username=self.user,
            password=self.password,
            database=self.database,
            port=self.port,
            auth=self.auth,
            **self.kwargs,
        )

        try:
            cursor = conn.cursor()
            cursor.execute(args.sql)
            results = cursor.fetchall()

            # Create a pandas dataframe from the results
            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])

            cursor.close()
            return df

        finally:
            conn.close()


================================================
FILE: src/vanna/integrations/local/__init__.py
================================================
"""
Local integration.

This module provides built-in local implementations.
"""

from .audit import LoggingAuditLogger
from .file_system import LocalFileSystem
from .storage import MemoryConversationStore
from .file_system_conversation_store import FileSystemConversationStore

__all__ = [
    "MemoryConversationStore",
    "FileSystemConversationStore",
    "LocalFileSystem",
    "LoggingAuditLogger",
]


================================================
FILE: src/vanna/integrations/local/agent_memory/__init__.py
================================================
"""
Local agent memory implementations.
"""

from .in_memory import DemoAgentMemory

__all__ = ["DemoAgentMemory"]


================================================
FILE: src/vanna/integrations/local/agent_memory/in_memory.py
================================================
"""
Demo in-memory implementation of AgentMemory.

This implementation provides a zero-dependency, minimal storage solution that
keeps all memories in RAM. It uses simple similarity algorithms (Jaccard and
difflib) instead of vector embeddings. Perfect for demos and testing.
"""

from __future__ import annotations

import asyncio
import difflib
import time
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class DemoAgentMemory(AgentMemory):
    """
    Minimal, dependency-free in-memory storage for demos and testing.
    - O(n) search over an in-memory list
    - Simple similarity: max(Jaccard(token sets), difflib ratio)
    - Optional FIFO eviction via max_items
    - Async-safe with an asyncio.Lock
    """

    def __init__(self, *, max_items: int = 10_000):
        """
        Initialize the in-memory storage.

        Args:
            max_items: Maximum number of memories to keep. Oldest memories are
                      evicted when this limit is reached (FIFO).
        """
        self._memories: List[ToolMemory] = []
        self._text_memories: List[TextMemory] = []
        self._lock = asyncio.Lock()
        self._max_items = max_items

    @staticmethod
    def _now_iso() -> str:
        """Get current timestamp in ISO format."""
        return datetime.now().isoformat()

    @staticmethod
    def _normalize(text: str) -> str:
        """Normalize text by lowercasing and collapsing whitespace."""
        return " ".join(text.lower().split())

    @staticmethod
    def _tokenize(text: str) -> set[str]:
        """Simple tokenizer that splits on whitespace."""
        return set(text.lower().split())

    @classmethod
    def _similarity(cls, a: str, b: str) -> float:
        """
        Calculate similarity between two strings using multiple methods.

        Returns the maximum of Jaccard similarity and difflib ratio.
        """
        a_norm, b_norm = cls._normalize(a), cls._normalize(b)

        # Jaccard over whitespace tokens
        ta, tb = cls._tokenize(a_norm), cls._tokenize(b_norm)
        if not ta and not tb:
            jaccard = 1.0
        elif not ta or not tb:
            jaccard = 0.0
        else:
            jaccard = len(ta & tb) / max(1, len(ta | tb))

        # difflib ratio
        ratio = difflib.SequenceMatcher(None, a_norm, b_norm).ratio()

        # Take the better of the two cheap measures
        return max(jaccard, ratio)

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern for future reference."""
        tm = ToolMemory(
            memory_id=str(uuid.uuid4()),
            question=question,
            tool_name=tool_name,
            args=args,
            timestamp=self._now_iso(),
            success=success,
            metadata=metadata or {},
        )
        async with self._lock:
            self._memories.append(tm)
            # Optional FIFO eviction
            if len(self._memories) > self._max_items:
                overflow = len(self._memories) - self._max_items
                del self._memories[:overflow]

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Store a text memory in RAM."""
        tm = TextMemory(
            memory_id=str(uuid.uuid4()), content=content, timestamp=self._now_iso()
        )
        async with self._lock:
            self._text_memories.append(tm)
            if len(self._text_memories) > self._max_items:
                overflow = len(self._text_memories) - self._max_items
                del self._text_memories[:overflow]
        return tm

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns based on a question."""
        q = self._normalize(question)

        async with self._lock:
            # Filter candidates by tool name and success status
            candidates = [
                m
                for m in self._memories
                if m.success
                and (tool_name_filter is None or m.tool_name == tool_name_filter)
            ]

            # Score each candidate by question similarity
            results: List[tuple[ToolMemory, float]] = []
            for m in candidates:
                score = self._similarity(q, m.question)
                results.append((m, min(score, 1.0)))

            # Filter by threshold and sort by score
            results = [(m, s) for (m, s) in results if s >= similarity_threshold]
            results.sort(key=lambda x: x[1], reverse=True)

            # Build ranked response
            out: List[ToolMemorySearchResult] = []
            for idx, (m, s) in enumerate(results[:limit], start=1):
                out.append(
                    ToolMemorySearchResult(memory=m, similarity_score=s, rank=idx)
                )
            return out

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search free-form text memories using the demo similarity metric."""
        normalized_query = self._normalize(query)

        async with self._lock:
            scored: List[tuple[TextMemory, float]] = []
            for memory in self._text_memories:
                score = self._similarity(normalized_query, memory.content)
                scored.append((memory, min(score, 1.0)))

            scored = [
                (memory, score)
                for memory, score in scored
                if score >= similarity_threshold
            ]
            scored.sort(key=lambda item: item[1], reverse=True)

            results: List[TextMemorySearchResult] = []
            for idx, (memory, score) in enumerate(scored[:limit], start=1):
                results.append(
                    TextMemorySearchResult(
                        memory=memory, similarity_score=score, rank=idx
                    )
                )
            return results

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories. Returns most recent memories first."""
        async with self._lock:
            # Return memories in reverse order (most recent first)
            return list(reversed(self._memories[-limit:]))

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Return recently added text memories."""
        async with self._lock:
            return list(reversed(self._text_memories[-limit:]))

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a stored text memory by ID."""
        async with self._lock:
            for index, memory in enumerate(self._text_memories):
                if memory.memory_id == memory_id:
                    del self._text_memories[index]
                    return True
            return False

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID. Returns True if deleted, False if not found."""
        async with self._lock:
            for i, m in enumerate(self._memories):
                if m.memory_id == memory_id:
                    del self._memories[i]
                    return True
            return False

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories. Returns number of memories deleted."""
        async with self._lock:
            original_tool_count = len(self._memories)
            original_text_count = len(self._text_memories)

            # Filter memories to keep
            kept_memories = []
            for m in self._memories:
                should_delete = True

                # Check tool name filter
                if tool_name and m.tool_name != tool_name:
                    should_delete = False

                # Check date filter
                if should_delete and before_date and m.timestamp:
                    if m.timestamp >= before_date:
                        should_delete = False

                # If no filters specified, delete all
                if tool_name is None and before_date is None:
                    should_delete = True

                # Keep if should not delete
                if not should_delete:
                    kept_memories.append(m)

            self._memories = kept_memories
            deleted_tool_count = original_tool_count - len(self._memories)

            # Apply filters to text memories (tool filter ignored)
            kept_text_memories = []
            for memory in self._text_memories:
                should_delete = (
                    tool_name is None
                )  # only delete text when not targeting a tool

                if before_date and memory.timestamp:
                    if memory.timestamp >= before_date:
                        should_delete = False

                if not should_delete:
                    kept_text_memories.append(memory)

            self._text_memories = kept_text_memories
            deleted_text_count = original_text_count - len(self._text_memories)

            return deleted_tool_count + deleted_text_count


================================================
FILE: src/vanna/integrations/local/audit.py
================================================
"""
Local audit logger implementation using Python logging.

This module provides a simple audit logger that writes events using
the standard Python logging module, useful for development and testing.
"""

import json
import logging
from typing import Optional

from vanna.core.audit import AuditEvent, AuditLogger

logger = logging.getLogger(__name__)


class LoggingAuditLogger(AuditLogger):
    """Audit logger that writes events to Python logger as structured JSON.

    This implementation uses logger.info() to emit audit events as JSON,
    making them easy to parse and route to log aggregation systems.

    Example:
        audit_logger = LoggingAuditLogger()
        agent = Agent(
            llm_service=...,
            audit_logger=audit_logger
        )
    """

    def __init__(self, log_level: int = logging.INFO):
        """Initialize the logging audit logger.

        Args:
            log_level: Log level to use for audit events (default: INFO)
        """
        self.log_level = log_level

    async def log_event(self, event: AuditEvent) -> None:
        """Log an audit event as structured JSON.

        Args:
            event: The audit event to log
        """
        try:
            # Convert event to dict for JSON serialization
            event_dict = event.model_dump(mode="json", exclude_none=True)

            # Format as single-line JSON for easy parsing
            event_json = json.dumps(event_dict, separators=(",", ":"))

            # Log with structured prefix for easy filtering
            logger.log(
                self.log_level,
                f"[AUDIT] {event.event_type.value} | {event_json}",
            )
        except Exception as e:
            # Don't fail the operation if audit logging fails
            logger.error(f"Failed to log audit event: {e}", exc_info=True)


================================================
FILE: src/vanna/integrations/local/file_system.py
================================================
"""
Local file system implementation.

This module provides a local file system implementation with per-user isolation.
"""

import asyncio
import hashlib
from pathlib import Path
from typing import List, Optional

from vanna.capabilities.file_system import CommandResult, FileSearchMatch, FileSystem
from vanna.core.tool import ToolContext

MAX_SEARCH_FILE_BYTES = 1_000_000


class LocalFileSystem(FileSystem):
    """Local file system implementation with per-user isolation."""

    def __init__(self, working_directory: str = "."):
        """Initialize with a working directory.

        Args:
            working_directory: Base directory where user-specific folders will be created
        """
        self.working_directory = Path(working_directory)

    def _get_user_directory(self, context: ToolContext) -> Path:
        """Get the user-specific directory by hashing the user ID.

        Args:
            context: Tool context containing user information

        Returns:
            Path to the user-specific directory
        """
        # Hash the user ID to create a directory name
        user_hash = hashlib.sha256(context.user.id.encode()).hexdigest()[:16]
        user_dir = self.working_directory / user_hash

        # Create the directory if it doesn't exist
        user_dir.mkdir(parents=True, exist_ok=True)

        return user_dir

    def _resolve_path(self, path: str, context: ToolContext) -> Path:
        """Resolve a path relative to the user's directory.

        Args:
            path: Path relative to user directory
            context: Tool context containing user information

        Returns:
            Absolute path within user's directory
        """
        user_dir = self._get_user_directory(context)
        resolved = user_dir / path

        # Ensure the path is within the user's directory (prevent directory traversal)
        try:
            resolved.resolve().relative_to(user_dir.resolve())
        except ValueError:
            raise PermissionError(
                f"Access denied: path '{path}' is outside user directory"
            )

        return resolved

    async def list_files(self, directory: str, context: ToolContext) -> List[str]:
        """List files in a directory within the user's isolated space."""
        directory_path = self._resolve_path(directory, context)

        if not directory_path.exists():
            raise FileNotFoundError(f"Directory '{directory}' does not exist")

        if not directory_path.is_dir():
            raise NotADirectoryError(f"'{directory}' is not a directory")

        files = []
        for item in directory_path.iterdir():
            if item.is_file():
                files.append(item.name)

        return sorted(files)

    async def read_file(self, filename: str, context: ToolContext) -> str:
        """Read the contents of a file within the user's isolated space."""
        file_path = self._resolve_path(filename, context)

        if not file_path.exists():
            raise FileNotFoundError(f"File '{filename}' does not exist")

        if not file_path.is_file():
            raise IsADirectoryError(f"'{filename}' is a directory, not a file")

        return file_path.read_text(encoding="utf-8")

    async def write_file(
        self, filename: str, content: str, context: ToolContext, overwrite: bool = False
    ) -> None:
        """Write content to a file within the user's isolated space."""
        file_path = self._resolve_path(filename, context)

        # Create parent directories if they don't exist
        file_path.parent.mkdir(parents=True, exist_ok=True)

        if file_path.exists() and not overwrite:
            raise FileExistsError(
                f"File '{filename}' already exists. Use overwrite=True to replace it."
            )

        file_path.write_text(content, encoding="utf-8")

    async def exists(self, path: str, context: ToolContext) -> bool:
        """Check if a file or directory exists within the user's isolated space."""
        try:
            resolved_path = self._resolve_path(path, context)
            return resolved_path.exists()
        except PermissionError:
            return False

    async def is_directory(self, path: str, context: ToolContext) -> bool:
        """Check if a path is a directory within the user's isolated space."""
        try:
            resolved_path = self._resolve_path(path, context)
            return resolved_path.exists() and resolved_path.is_dir()
        except PermissionError:
            return False

    async def search_files(
        self,
        query: str,
        context: ToolContext,
        *,
        max_results: int = 20,
        include_content: bool = False,
    ) -> List[FileSearchMatch]:
        """Search for files within the user's isolated space."""

        trimmed_query = query.strip()
        if not trimmed_query:
            raise ValueError("Search query must not be empty")

        user_dir = self._get_user_directory(context)
        matches: List[FileSearchMatch] = []
        query_lower = trimmed_query.lower()

        for path in user_dir.rglob("*"):
            if len(matches) >= max_results:
                break

            if not path.is_file():
                continue

            relative_path = path.relative_to(user_dir).as_posix()
            include_entry = False
            snippet: Optional[str] = None

            if query_lower in path.name.lower():
                include_entry = True
                snippet = "[filename match]"

            content: Optional[str] = None
            if include_content:
                try:
                    size = path.stat().st_size
                except OSError:
                    if include_entry:
                        matches.append(
                            FileSearchMatch(path=relative_path, snippet=snippet)
                        )
                    continue

                if size <= MAX_SEARCH_FILE_BYTES:
                    try:
                        content = path.read_text(encoding="utf-8")
                    except (UnicodeDecodeError, OSError):
                        content = None
                elif not include_entry:
                    # Skip oversized files if they do not match by name
                    continue

            if include_content and content is not None:
                if query_lower in content.lower():
                    # Create snippet
                    lowered = content.lower()
                    index = lowered.find(query_lower)
                    if index != -1:
                        context_window = 60
                        start = max(0, index - context_window)
                        end = min(len(content), index + len(query) + context_window)
                        snippet = content[start:end].replace("\n", " ").strip()
                        if start > 0:
                            snippet = f"…{snippet}"
                        if end < len(content):
                            snippet = f"{snippet}…"
                    include_entry = True
                elif not include_entry:
                    continue

            if include_entry:
                matches.append(FileSearchMatch(path=relative_path, snippet=snippet))

        return matches

    async def run_bash(
        self,
        command: str,
        context: ToolContext,
        *,
        timeout: Optional[float] = None,
    ) -> CommandResult:
        """Execute a bash command within the user's isolated space."""

        if not command.strip():
            raise ValueError("Command must not be empty")

        user_dir = self._get_user_directory(context)

        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(user_dir),
        )

        try:
            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                process.communicate(), timeout=timeout
            )
        except asyncio.TimeoutError as exc:
            process.kill()
            await process.wait()
            raise TimeoutError(f"Command timed out after {timeout} seconds") from exc

        stdout = stdout_bytes.decode("utf-8", errors="replace")
        stderr = stderr_bytes.decode("utf-8", errors="replace")

        return CommandResult(
            stdout=stdout, stderr=stderr, returncode=process.returncode or 0
        )


================================================
FILE: src/vanna/integrations/local/file_system_conversation_store.py
================================================
"""
File system conversation store implementation.

This module provides a file-based implementation of the ConversationStore
interface that persists conversations to disk as a directory structure.
"""

import json
import os
from pathlib import Path
from typing import Dict, List, Optional
from datetime import datetime
import time

from vanna.core.storage import ConversationStore, Conversation, Message
from vanna.core.user import User


class FileSystemConversationStore(ConversationStore):
    """File system-based conversation store.

    Stores conversations as directories with individual message files:
    conversations/{conversation_id}/
        metadata.json - conversation metadata (id, user info, timestamps)
        messages/
            {timestamp}_{index}.json - individual message files
    """

    def __init__(self, base_dir: str = "conversations") -> None:
        """Initialize the file system conversation store.

        Args:
            base_dir: Base directory for storing conversations
        """
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def _get_conversation_dir(self, conversation_id: str) -> Path:
        """Get the directory path for a conversation."""
        return self.base_dir / conversation_id

    def _get_metadata_path(self, conversation_id: str) -> Path:
        """Get the metadata file path for a conversation."""
        return self._get_conversation_dir(conversation_id) / "metadata.json"

    def _get_messages_dir(self, conversation_id: str) -> Path:
        """Get the messages directory for a conversation."""
        return self._get_conversation_dir(conversation_id) / "messages"

    def _save_metadata(self, conversation: Conversation) -> None:
        """Save conversation metadata to disk."""
        conv_dir = self._get_conversation_dir(conversation.id)
        conv_dir.mkdir(parents=True, exist_ok=True)

        metadata = {
            "id": conversation.id,
            "user": conversation.user.model_dump(mode="json"),
            "created_at": conversation.created_at.isoformat(),
            "updated_at": conversation.updated_at.isoformat(),
        }

        metadata_path = self._get_metadata_path(conversation.id)
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=2)

    def _load_messages(self, conversation_id: str) -> List[Message]:
        """Load all messages for a conversation."""
        messages_dir = self._get_messages_dir(conversation_id)

        if not messages_dir.exists():
            return []

        messages = []
        # Sort message files by name (timestamp_index ensures correct order)
        message_files = sorted(messages_dir.glob("*.json"))

        for file_path in message_files:
            try:
                with open(file_path, "r") as f:
                    data = json.load(f)
                message = Message.model_validate(data)
                messages.append(message)
            except (json.JSONDecodeError, ValueError) as e:
                print(f"Failed to load message from {file_path}: {e}")
                continue

        return messages

    def _append_message(
        self, conversation_id: str, message: Message, index: int
    ) -> None:
        """Append a message to the conversation."""
        messages_dir = self._get_messages_dir(conversation_id)
        messages_dir.mkdir(parents=True, exist_ok=True)

        # Use timestamp + index to ensure unique, ordered filenames
        timestamp = int(time.time() * 1000000)  # microseconds
        filename = f"{timestamp}_{index:06d}.json"
        file_path = messages_dir / filename

        with open(file_path, "w") as f:
            json.dump(message.model_dump(mode="json"), f, indent=2)

    async def create_conversation(
        self, conversation_id: str, user: User, initial_message: str
    ) -> Conversation:
        """Create a new conversation with the specified ID."""
        conversation = Conversation(
            id=conversation_id,
            user=user,
            messages=[Message(role="user", content=initial_message)],
        )

        # Save metadata
        self._save_metadata(conversation)

        # Save initial message
        self._append_message(conversation_id, conversation.messages[0], 0)

        return conversation

    async def get_conversation(
        self, conversation_id: str, user: User
    ) -> Optional[Conversation]:
        """Get conversation by ID, scoped to user."""
        metadata_path = self._get_metadata_path(conversation_id)

        if not metadata_path.exists():
            return None

        try:
            # Load metadata
            with open(metadata_path, "r") as f:
                metadata = json.load(f)

            # Verify ownership
            if metadata["user"]["id"] != user.id:
                return None

            # Load all messages
            messages = self._load_messages(conversation_id)

            # Reconstruct conversation
            conversation = Conversation(
                id=metadata["id"],
                user=User.model_validate(metadata["user"]),
                messages=messages,
                created_at=datetime.fromisoformat(metadata["created_at"]),
                updated_at=datetime.fromisoformat(metadata["updated_at"]),
            )

            return conversation
        except (json.JSONDecodeError, ValueError, KeyError) as e:
            print(f"Failed to load conversation {conversation_id}: {e}")
            return None

    async def update_conversation(self, conversation: Conversation) -> None:
        """Update conversation with new messages."""
        # Update the updated_at timestamp
        conversation.updated_at = datetime.now()

        # Save updated metadata
        self._save_metadata(conversation)

        # Get existing messages count to determine new message indices
        existing_messages = self._load_messages(conversation.id)
        existing_count = len(existing_messages)

        # Only append new messages (ones not already saved)
        for i, message in enumerate(
            conversation.messages[existing_count:], start=existing_count
        ):
            self._append_message(conversation.id, message, i)

    async def delete_conversation(self, conversation_id: str, user: User) -> bool:
        """Delete conversation."""
        conv_dir = self._get_conversation_dir(conversation_id)

        if not conv_dir.exists():
            return False

        # Verify ownership before deleting
        conversation = await self.get_conversation(conversation_id, user)
        if not conversation:
            return False

        try:
            # Delete all message files
            messages_dir = self._get_messages_dir(conversation_id)
            if messages_dir.exists():
                for file_path in messages_dir.glob("*.json"):
                    file_path.unlink()
                messages_dir.rmdir()

            # Delete metadata
            metadata_path = self._get_metadata_path(conversation_id)
            if metadata_path.exists():
                metadata_path.unlink()

            # Delete conversation directory
            conv_dir.rmdir()

            return True
        except OSError as e:
            print(f"Failed to delete conversation {conversation_id}: {e}")
            return False

    async def list_conversations(
        self, user: User, limit: int = 50, offset: int = 0
    ) -> List[Conversation]:
        """List conversations for user."""
        if not self.base_dir.exists():
            return []

        conversations = []

        # Iterate through all conversation directories
        for conv_dir in self.base_dir.iterdir():
            if not conv_dir.is_dir():
                continue

            metadata_path = conv_dir / "metadata.json"
            if not metadata_path.exists():
                continue

            try:
                # Load metadata
                with open(metadata_path, "r") as f:
                    metadata = json.load(f)

                # Skip conversations not owned by this user
                if metadata["user"]["id"] != user.id:
                    continue

                # Load messages
                messages = self._load_messages(conv_dir.name)

                # Reconstruct conversation
                conversation = Conversation(
                    id=metadata["id"],
                    user=User.model_validate(metadata["user"]),
                    messages=messages,
                    created_at=datetime.fromisoformat(metadata["created_at"]),
                    updated_at=datetime.fromisoformat(metadata["updated_at"]),
                )
                conversations.append(conversation)
            except (json.JSONDecodeError, ValueError, KeyError) as e:
                print(f"Failed to load conversation from {conv_dir}: {e}")
                continue

        # Sort by updated_at desc
        conversations.sort(key=lambda x: x.updated_at, reverse=True)

        # Apply pagination
        return conversations[offset : offset + limit]


================================================
FILE: src/vanna/integrations/local/storage.py
================================================
"""
In-memory conversation store implementation.

This module provides a simple in-memory implementation of the ConversationStore
interface, useful for testing and development.
"""

from typing import Dict, List, Optional

from vanna.core.storage import ConversationStore, Conversation, Message
from vanna.core.user import User


class MemoryConversationStore(ConversationStore):
    """In-memory conversation store."""

    def __init__(self) -> None:
        self._conversations: Dict[str, Conversation] = {}

    async def create_conversation(
        self, conversation_id: str, user: User, initial_message: str
    ) -> Conversation:
        """Create a new conversation with the specified ID."""
        conversation = Conversation(
            id=conversation_id,
            user=user,
            messages=[Message(role="user", content=initial_message)],
        )
        self._conversations[conversation_id] = conversation
        return conversation

    async def get_conversation(
        self, conversation_id: str, user: User
    ) -> Optional[Conversation]:
        """Get conversation by ID, scoped to user."""
        conversation = self._conversations.get(conversation_id)
        if conversation and conversation.user.id == user.id:
            return conversation
        return None

    async def update_conversation(self, conversation: Conversation) -> None:
        """Update conversation with new messages."""
        self._conversations[conversation.id] = conversation

    async def delete_conversation(self, conversation_id: str, user: User) -> bool:
        """Delete conversation."""
        conversation = await self.get_conversation(conversation_id, user)
        if conversation:
            del self._conversations[conversation_id]
            return True
        return False

    async def list_conversations(
        self, user: User, limit: int = 50, offset: int = 0
    ) -> List[Conversation]:
        """List conversations for user."""
        user_conversations = [
            conv for conv in self._conversations.values() if conv.user.id == user.id
        ]
        # Sort by updated_at desc
        user_conversations.sort(key=lambda x: x.updated_at, reverse=True)
        return user_conversations[offset : offset + limit]


================================================
FILE: src/vanna/integrations/marqo/__init__.py
================================================
"""
Marqo integration for Vanna Agents.
"""

from .agent_memory import MarqoAgentMemory

__all__ = ["MarqoAgentMemory"]


================================================
FILE: src/vanna/integrations/marqo/agent_memory.py
================================================
"""
Marqo vector database implementation of AgentMemory.

This implementation uses Marqo for vector storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    import marqo

    MARQO_AVAILABLE = True
except ImportError:
    MARQO_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class MarqoAgentMemory(AgentMemory):
    """Marqo-based implementation of AgentMemory."""

    def __init__(
        self,
        url: str = "http://localhost:8882",
        index_name: str = "tool-memories",
        api_key: Optional[str] = None,
    ):
        if not MARQO_AVAILABLE:
            raise ImportError(
                "Marqo is required for MarqoAgentMemory. Install with: pip install marqo"
            )

        self.url = url
        self.index_name = index_name
        self.api_key = api_key
        self._client = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_client(self):
        """Get or create Marqo client."""
        if self._client is None:
            self._client = marqo.Client(url=self.url, api_key=self.api_key)

            # Create index if it doesn't exist
            try:
                self._client.get_index(self.index_name)
            except Exception:
                self._client.create_index(self.index_name)

        return self._client

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            client = self._get_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()

            document = {
                "_id": memory_id,
                "question": question,
                "tool_name": tool_name,
                "args": json.dumps(args),
                "timestamp": timestamp,
                "success": success,
                "metadata": json.dumps(metadata or {}),
            }

            client.index(self.index_name).add_documents(
                [document], tensor_fields=["question"]
            )

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            client = self._get_client()

            # Build filter
            filter_string = "success:true"
            if tool_name_filter:
                filter_string += f" AND tool_name:{tool_name_filter}"

            results = client.index(self.index_name).search(
                q=question, limit=limit, filter_string=filter_string
            )

            search_results = []
            for i, hit in enumerate(results["hits"]):
                # Marqo returns score
                similarity_score = hit.get("_score", 0)

                if similarity_score >= similarity_threshold:
                    args = json.loads(hit.get("args", "{}"))
                    metadata_dict = json.loads(hit.get("metadata", "{}"))

                    memory = ToolMemory(
                        memory_id=hit["_id"],
                        question=hit["question"],
                        tool_name=hit["tool_name"],
                        args=args,
                        timestamp=hit.get("timestamp"),
                        success=hit.get("success", True),
                        metadata=metadata_dict,
                    )

                    search_results.append(
                        ToolMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            client = self._get_client()

            # Search with wildcard and sort by timestamp
            results = client.index(self.index_name).search(
                q="*", limit=limit, sort="timestamp:desc"
            )

            memories = []
            for hit in results.get("hits", []):
                args = json.loads(hit.get("args", "{}"))
                metadata_dict = json.loads(hit.get("metadata", "{}"))

                memory = ToolMemory(
                    memory_id=hit["_id"],
                    question=hit["question"],
                    tool_name=hit["tool_name"],
                    args=args,
                    timestamp=hit.get("timestamp"),
                    success=hit.get("success", True),
                    metadata=metadata_dict,
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            client = self._get_client()

            try:
                client.index(self.index_name).delete_documents(ids=[memory_id])
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            client = self._get_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()

            document = {
                "_id": memory_id,
                "content": content,
                "timestamp": timestamp,
                "is_text_memory": True,
            }

            client.index(self.index_name).add_documents(
                [document], tensor_fields=["content"]
            )

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            client = self._get_client()

            filter_string = "is_text_memory:true"

            results = client.index(self.index_name).search(
                q=query, limit=limit, filter_string=filter_string
            )

            search_results = []
            for i, hit in enumerate(results["hits"]):
                similarity_score = hit.get("_score", 0)

                if similarity_score >= similarity_threshold:
                    memory = TextMemory(
                        memory_id=hit["_id"],
                        content=hit.get("content", ""),
                        timestamp=hit.get("timestamp"),
                    )

                    search_results.append(
                        TextMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            client = self._get_client()

            results = client.index(self.index_name).search(
                q="*",
                limit=limit,
                filter_string="is_text_memory:true",
                sort="timestamp:desc",
            )

            memories = []
            for hit in results.get("hits", []):
                memory = TextMemory(
                    memory_id=hit["_id"],
                    content=hit.get("content", ""),
                    timestamp=hit.get("timestamp"),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            client = self._get_client()

            try:
                client.index(self.index_name).delete_documents(ids=[memory_id])
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            client = self._get_client()

            # Build filter for search
            filter_parts = []
            if tool_name:
                filter_parts.append(f"tool_name:{tool_name}")
            if before_date:
                filter_parts.append(f"timestamp:[* TO {before_date}]")

            if filter_parts or (tool_name is None and before_date is None):
                filter_string = " AND ".join(filter_parts) if filter_parts else None

                if filter_string:
                    # Search for documents to delete
                    results = client.index(self.index_name).search(
                        q="*",
                        limit=1000,  # Max results
                        filter_string=filter_string,
                    )

                    ids_to_delete = [hit["_id"] for hit in results.get("hits", [])]

                    if ids_to_delete:
                        client.index(self.index_name).delete_documents(
                            ids=ids_to_delete
                        )

                    return len(ids_to_delete)
                else:
                    # Delete entire index and recreate
                    try:
                        client.delete_index(self.index_name)
                        client.create_index(self.index_name)
                    except Exception:
                        pass
                    return 0

            return 0

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/milvus/__init__.py
================================================
"""
Milvus integration for Vanna Agents.
"""

from .agent_memory import MilvusAgentMemory

__all__ = ["MilvusAgentMemory"]


================================================
FILE: src/vanna/integrations/milvus/agent_memory.py
================================================
"""
Milvus vector database implementation of AgentMemory.

This implementation uses Milvus for distributed vector storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    from pymilvus import (
        connections,
        Collection,
        CollectionSchema,
        FieldSchema,
        DataType,
        utility,
    )

    MILVUS_AVAILABLE = True
except ImportError:
    MILVUS_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class MilvusAgentMemory(AgentMemory):
    """Milvus-based implementation of AgentMemory."""

    def __init__(
        self,
        collection_name: str = "tool_memories",
        host: str = "localhost",
        port: int = 19530,
        alias: str = "default",
        dimension: int = 384,
    ):
        if not MILVUS_AVAILABLE:
            raise ImportError(
                "Milvus is required for MilvusAgentMemory. Install with: pip install pymilvus"
            )

        self.collection_name = collection_name
        self.host = host
        self.port = port
        self.alias = alias
        self.dimension = dimension
        self._collection = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_collection(self):
        """Get or create Milvus collection."""
        if self._collection is None:
            # Connect to Milvus
            connections.connect(alias=self.alias, host=self.host, port=self.port)

            # Create collection if it doesn't exist
            if not utility.has_collection(self.collection_name):
                fields = [
                    FieldSchema(
                        name="id",
                        dtype=DataType.VARCHAR,
                        is_primary=True,
                        max_length=100,
                    ),
                    FieldSchema(
                        name="embedding",
                        dtype=DataType.FLOAT_VECTOR,
                        dim=self.dimension,
                    ),
                    FieldSchema(
                        name="question", dtype=DataType.VARCHAR, max_length=2000
                    ),
                    FieldSchema(
                        name="tool_name", dtype=DataType.VARCHAR, max_length=200
                    ),
                    FieldSchema(
                        name="args_json", dtype=DataType.VARCHAR, max_length=5000
                    ),
                    FieldSchema(
                        name="timestamp", dtype=DataType.VARCHAR, max_length=50
                    ),
                    FieldSchema(name="success", dtype=DataType.BOOL),
                    FieldSchema(
                        name="metadata_json", dtype=DataType.VARCHAR, max_length=5000
                    ),
                ]

                schema = CollectionSchema(
                    fields=fields, description="Tool usage memories"
                )
                collection = Collection(name=self.collection_name, schema=schema)

                # Create index for vector field
                index_params = {
                    "index_type": "IVF_FLAT",
                    "metric_type": "IP",
                    "params": {"nlist": 128},
                }
                collection.create_index(
                    field_name="embedding", index_params=index_params
                )

            self._collection = Collection(self.collection_name)
            self._collection.load()

        return self._collection

    def _create_embedding(self, text: str) -> List[float]:
        """Create a simple embedding from text (placeholder)."""
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            collection = self._get_collection()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            entities = [
                [memory_id],
                [embedding],
                [question],
                [tool_name],
                [json.dumps(args)],
                [timestamp],
                [success],
                [json.dumps(metadata or {})],
            ]

            collection.insert(entities)
            collection.flush()

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            collection = self._get_collection()

            embedding = self._create_embedding(question)

            # Build filter expression
            expr = "success == true"
            if tool_name_filter:
                expr += f' && tool_name == "{tool_name_filter}"'

            search_params = {"metric_type": "IP", "params": {"nprobe": 10}}

            results = collection.search(
                data=[embedding],
                anns_field="embedding",
                param=search_params,
                limit=limit,
                expr=expr,
                output_fields=[
                    "id",
                    "question",
                    "tool_name",
                    "args_json",
                    "timestamp",
                    "success",
                    "metadata_json",
                ],
            )

            search_results = []
            for i, hits in enumerate(results):
                for j, hit in enumerate(hits):
                    similarity_score = hit.distance

                    if similarity_score >= similarity_threshold:
                        args = json.loads(hit.entity.get("args_json", "{}"))
                        metadata_dict = json.loads(
                            hit.entity.get("metadata_json", "{}")
                        )

                        memory = ToolMemory(
                            memory_id=hit.entity.get("id"),
                            question=hit.entity.get("question"),
                            tool_name=hit.entity.get("tool_name"),
                            args=args,
                            timestamp=hit.entity.get("timestamp"),
                            success=hit.entity.get("success", True),
                            metadata=metadata_dict,
                        )

                        search_results.append(
                            ToolMemorySearchResult(
                                memory=memory,
                                similarity_score=similarity_score,
                                rank=j + 1,
                            )
                        )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            collection = self._get_collection()

            # Query all entries and sort by timestamp
            results = collection.query(
                expr="id != ''",
                output_fields=[
                    "id",
                    "question",
                    "tool_name",
                    "args_json",
                    "timestamp",
                    "success",
                    "metadata_json",
                ],
                limit=1000,
            )

            # Sort by timestamp
            sorted_results = sorted(
                results, key=lambda r: r.get("timestamp", ""), reverse=True
            )

            memories = []
            for result in sorted_results[:limit]:
                args = json.loads(result.get("args_json", "{}"))
                metadata_dict = json.loads(result.get("metadata_json", "{}"))

                memory = ToolMemory(
                    memory_id=result.get("id"),
                    question=result.get("question"),
                    tool_name=result.get("tool_name"),
                    args=args,
                    timestamp=result.get("timestamp"),
                    success=result.get("success", True),
                    metadata=metadata_dict,
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            collection = self._get_collection()

            try:
                expr = f'id == "{memory_id}"'
                collection.delete(expr)
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            collection = self._get_collection()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            entities = [
                [memory_id],
                [embedding],
                [content],
                [""],  # tool_name (empty for text memories)
                [""],  # args_json (empty for text memories)
                [timestamp],
                [True],  # success (always true for text memories)
                [json.dumps({"is_text_memory": True})],  # metadata_json
            ]

            collection.insert(entities)
            collection.flush()

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            collection = self._get_collection()

            embedding = self._create_embedding(query)

            # Build filter expression for text memories
            expr = 'tool_name == ""'

            search_params = {"metric_type": "IP", "params": {"nprobe": 10}}

            results = collection.search(
                data=[embedding],
                anns_field="embedding",
                param=search_params,
                limit=limit,
                expr=expr,
                output_fields=["id", "question", "timestamp", "metadata_json"],
            )

            search_results = []
            for i, hits in enumerate(results):
                for j, hit in enumerate(hits):
                    similarity_score = hit.distance

                    if similarity_score >= similarity_threshold:
                        content = hit.entity.get("question", "")

                        memory = TextMemory(
                            memory_id=hit.entity.get("id"),
                            content=content,
                            timestamp=hit.entity.get("timestamp"),
                        )

                        search_results.append(
                            TextMemorySearchResult(
                                memory=memory,
                                similarity_score=similarity_score,
                                rank=j + 1,
                            )
                        )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            collection = self._get_collection()

            # Query text memory entries
            results = collection.query(
                expr='tool_name == ""',
                output_fields=["id", "question", "timestamp"],
                limit=1000,
            )

            # Sort by timestamp
            sorted_results = sorted(
                results, key=lambda r: r.get("timestamp", ""), reverse=True
            )

            memories = []
            for result in sorted_results[:limit]:
                memory = TextMemory(
                    memory_id=result.get("id"),
                    content=result.get("question", ""),
                    timestamp=result.get("timestamp"),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            collection = self._get_collection()

            try:
                expr = f'id == "{memory_id}"'
                collection.delete(expr)
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            collection = self._get_collection()

            # Build filter expression
            expr_parts = []
            if tool_name:
                expr_parts.append(f'tool_name == "{tool_name}"')
            if before_date:
                expr_parts.append(f'timestamp < "{before_date}"')

            if expr_parts:
                expr = " && ".join(expr_parts)
            else:
                expr = "id != ''"

            collection.delete(expr)
            return 0

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/mock/__init__.py
================================================
"""
Mock integration.

This module provides mock implementations for testing.
"""

from .llm import MockLlmService

__all__ = ["MockLlmService"]


================================================
FILE: src/vanna/integrations/mock/llm.py
================================================
"""
Mock LLM service implementation for testing.

This module provides a simple mock implementation of the LlmService interface,
useful for testing and development without requiring actual LLM API calls.
"""

import asyncio
from typing import AsyncGenerator, List

from vanna.core.llm import LlmService, LlmRequest, LlmResponse, LlmStreamChunk
from vanna.core.tool import ToolSchema


class MockLlmService(LlmService):
    """Mock LLM service that returns predefined responses."""

    def __init__(self, response_content: str = "Hello! This is a mock response."):
        self.response_content = response_content
        self.call_count = 0

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a request to the mock LLM."""
        self.call_count += 1

        # Simulate processing delay
        await asyncio.sleep(0.1)

        # Return a simple response
        return LlmResponse(
            content=f"{self.response_content} (Request #{self.call_count})",
            finish_reason="stop",
            usage={"prompt_tokens": 50, "completion_tokens": 20, "total_tokens": 70},
        )

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Stream a request to the mock LLM."""
        self.call_count += 1

        # Split response into chunks
        words = f"{self.response_content} (Streamed #{self.call_count})".split()

        for i, word in enumerate(words):
            await asyncio.sleep(0.05)  # Simulate streaming delay

            chunk_content = word + (" " if i < len(words) - 1 else "")
            yield LlmStreamChunk(
                content=chunk_content,
                finish_reason="stop" if i == len(words) - 1 else None,
            )

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Validate tool schemas and return any errors."""
        # Mock validation - no errors
        return []

    def set_response(self, content: str) -> None:
        """Set the response content for testing."""
        self.response_content = content

    def reset_call_count(self) -> None:
        """Reset the call counter."""
        self.call_count = 0


================================================
FILE: src/vanna/integrations/mssql/__init__.py
================================================
"""Microsoft SQL Server integration for Vanna."""

from .sql_runner import MSSQLRunner

__all__ = ["MSSQLRunner"]


================================================
FILE: src/vanna/integrations/mssql/sql_runner.py
================================================
"""Microsoft SQL Server implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class MSSQLRunner(SqlRunner):
    """Microsoft SQL Server implementation of the SqlRunner interface."""

    def __init__(self, odbc_conn_str: str, **kwargs):
        """Initialize with MSSQL connection parameters.

        Args:
            odbc_conn_str: The ODBC connection string for SQL Server
            **kwargs: Additional SQLAlchemy engine parameters
        """
        try:
            import pyodbc

            self.pyodbc = pyodbc
        except ImportError as e:
            raise ImportError(
                "pyodbc package is required. Install with: pip install pyodbc"
            ) from e

        try:
            import sqlalchemy as sa
            from sqlalchemy.engine import URL
            from sqlalchemy import create_engine

            self.sa = sa
            self.URL = URL
            self.create_engine = create_engine
        except ImportError as e:
            raise ImportError(
                "sqlalchemy package is required. Install with: pip install sqlalchemy"
            ) from e

        # Create the connection URL
        connection_url = self.URL.create(
            "mssql+pyodbc", query={"odbc_connect": odbc_conn_str}
        )

        # Create the engine
        self.engine = self.create_engine(connection_url, **kwargs)

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against MSSQL database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            sqlalchemy.exc.SQLAlchemyError: If query execution fails
        """
        # Execute the SQL statement and return the result as a pandas DataFrame
        with self.engine.begin() as conn:
            df = pd.read_sql_query(self.sa.text(args.sql), conn)
            return df


================================================
FILE: src/vanna/integrations/mysql/__init__.py
================================================
"""MySQL integration for Vanna."""

from .sql_runner import MySQLRunner

__all__ = ["MySQLRunner"]


================================================
FILE: src/vanna/integrations/mysql/sql_runner.py
================================================
"""MySQL implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class MySQLRunner(SqlRunner):
    """MySQL implementation of the SqlRunner interface."""

    def __init__(
        self,
        host: str,
        database: str,
        user: str,
        password: str,
        port: int = 3306,
        **kwargs,
    ):
        """Initialize with MySQL connection parameters.

        Args:
            host: Database host address
            database: Database name
            user: Database user
            password: Database password
            port: Database port (default: 3306)
            **kwargs: Additional PyMySQL connection parameters
        """
        try:
            import pymysql.cursors

            self.pymysql = pymysql
        except ImportError as e:
            raise ImportError(
                "PyMySQL package is required. Install with: pip install 'vanna[mysql]'"
            ) from e

        self.host = host
        self.database = database
        self.user = user
        self.password = password
        self.port = port
        self.kwargs = kwargs

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against MySQL database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            pymysql.Error: If query execution fails
        """
        # Connect to the database
        conn = self.pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            database=self.database,
            port=self.port,
            cursorclass=self.pymysql.cursors.DictCursor,
            **self.kwargs,
        )

        try:
            # Ping to ensure connection is alive
            conn.ping(reconnect=True)

            cursor = conn.cursor()
            cursor.execute(args.sql)
            results = cursor.fetchall()

            # Create a pandas dataframe from the results
            df = pd.DataFrame(
                results,
                columns=[desc[0] for desc in cursor.description]
                if cursor.description
                else [],
            )

            cursor.close()
            return df

        finally:
            conn.close()


================================================
FILE: src/vanna/integrations/ollama/__init__.py
================================================
"""
Ollama integration for Vanna Agents.
"""

from .llm import OllamaLlmService

__all__ = ["OllamaLlmService"]


================================================
FILE: src/vanna/integrations/ollama/llm.py
================================================
"""
Ollama LLM service implementation.

This module provides an implementation of the LlmService interface backed by
Ollama's local LLM API. It supports non-streaming responses and streaming
of text content. Tool calling support depends on the Ollama model being used.
"""

from __future__ import annotations

import json
import os
from typing import Any, AsyncGenerator, Dict, List, Optional

from vanna.core.llm import (
    LlmService,
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
)
from vanna.core.tool import ToolCall, ToolSchema


class OllamaLlmService(LlmService):
    """Ollama-backed LLM service for local model inference.

    Args:
        model: Ollama model name (e.g., "gpt-oss:20b").
        host: Ollama server URL; defaults to "http://localhost:11434" or env `OLLAMA_HOST`.
        timeout: Request timeout in seconds; defaults to 240.
        num_ctx: Context window size; defaults to 8192.
        temperature: Sampling temperature; defaults to 0.7.
        extra_options: Additional options passed to Ollama (e.g., num_predict, top_k, top_p).
    """

    def __init__(
        self,
        model: str,
        host: Optional[str] = None,
        timeout: float = 240.0,
        num_ctx: int = 8192,
        temperature: float = 0.7,
        **extra_options: Any,
    ) -> None:
        try:
            import ollama
        except ImportError as e:
            raise ImportError(
                "ollama package is required. Install with: pip install 'vanna[ollama]' or pip install ollama"
            ) from e

        if not model:
            raise ValueError("model parameter is required for Ollama")

        self.model = model
        self.host = host or os.getenv("OLLAMA_HOST", "http://localhost:11434")
        self.timeout = timeout
        self.num_ctx = num_ctx
        self.temperature = temperature
        self.extra_options = extra_options

        # Create Ollama client
        self._client = ollama.Client(host=self.host, timeout=timeout)

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a non-streaming request to Ollama and return the response."""
        payload = self._build_payload(request)

        # Call the Ollama API
        try:
            resp = self._client.chat(**payload)
        except Exception as e:
            raise RuntimeError(f"Ollama request failed: {str(e)}") from e

        # Extract message from response
        message = resp.get("message", {})
        content = message.get("content")
        tool_calls = self._extract_tool_calls_from_message(message)

        # Extract usage information if available
        usage: Dict[str, int] = {}
        if "prompt_eval_count" in resp or "eval_count" in resp:
            usage = {
                "prompt_tokens": resp.get("prompt_eval_count", 0),
                "completion_tokens": resp.get("eval_count", 0),
                "total_tokens": resp.get("prompt_eval_count", 0)
                + resp.get("eval_count", 0),
            }

        return LlmResponse(
            content=content,
            tool_calls=tool_calls or None,
            finish_reason=resp.get("done_reason")
            or ("stop" if resp.get("done") else None),
            usage=usage or None,
        )

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Stream a request to Ollama.

        Emits `LlmStreamChunk` for textual deltas as they arrive. Tool calls are
        accumulated and emitted in a final chunk when the stream ends.
        """
        payload = self._build_payload(request)

        # Ollama streaming
        try:
            stream = self._client.chat(**payload, stream=True)
        except Exception as e:
            raise RuntimeError(f"Ollama streaming request failed: {str(e)}") from e

        # Accumulate tool calls if present
        accumulated_tool_calls: List[ToolCall] = []
        last_finish: Optional[str] = None

        for chunk in stream:
            message = chunk.get("message", {})

            # Yield text content
            content = message.get("content")
            if content:
                yield LlmStreamChunk(content=content)

            # Accumulate tool calls
            tool_calls = self._extract_tool_calls_from_message(message)
            if tool_calls:
                accumulated_tool_calls.extend(tool_calls)

            # Track finish reason
            if chunk.get("done"):
                last_finish = chunk.get("done_reason", "stop")

        # Emit final chunk with tool calls if any
        if accumulated_tool_calls:
            yield LlmStreamChunk(
                tool_calls=accumulated_tool_calls, finish_reason=last_finish or "stop"
            )
        else:
            # Emit terminal chunk to signal completion
            yield LlmStreamChunk(finish_reason=last_finish or "stop")

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Validate tool schemas. Returns a list of error messages."""
        errors: List[str] = []
        # Basic validation; Ollama model support for tools varies
        for t in tools:
            if not t.name:
                errors.append(f"Tool must have a name")
            if not t.description:
                errors.append(f"Tool '{t.name}' should have a description")
        return errors

    # Internal helpers
    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
        """Build the Ollama chat payload from LlmRequest."""
        messages: List[Dict[str, Any]] = []

        # Add system prompt as first message if provided
        if request.system_prompt:
            messages.append({"role": "system", "content": request.system_prompt})

        # Convert messages to Ollama format
        for m in request.messages:
            msg: Dict[str, Any] = {"role": m.role, "content": m.content or ""}

            # Handle tool calls in assistant messages
            if m.role == "assistant" and m.tool_calls:
                # Some Ollama models support tool_calls in message
                tool_calls_payload = []
                for tc in m.tool_calls:
                    tool_calls_payload.append(
                        {"function": {"name": tc.name, "arguments": tc.arguments}}
                    )
                msg["tool_calls"] = tool_calls_payload

            messages.append(msg)

        # Build tools array if tools are provided
        tools_payload: Optional[List[Dict[str, Any]]] = None
        if request.tools:
            tools_payload = []
            for t in request.tools:
                tools_payload.append(
                    {
                        "type": "function",
                        "function": {
                            "name": t.name,
                            "description": t.description,
                            "parameters": t.parameters,
                        },
                    }
                )

        # Build options
        options: Dict[str, Any] = {
            "num_ctx": self.num_ctx,
            "temperature": self.temperature,
            **self.extra_options,
        }

        # Build final payload
        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "options": options,
        }

        # Add tools if provided (note: not all Ollama models support tools)
        if tools_payload:
            payload["tools"] = tools_payload

        return payload

    def _extract_tool_calls_from_message(
        self, message: Dict[str, Any]
    ) -> List[ToolCall]:
        """Extract tool calls from Ollama message."""
        tool_calls: List[ToolCall] = []

        # Check for tool_calls in message
        raw_tool_calls = message.get("tool_calls", [])
        if not raw_tool_calls:
            return tool_calls

        for idx, tc in enumerate(raw_tool_calls):
            fn = tc.get("function", {})
            name = fn.get("name")
            if not name:
                continue

            # Parse arguments
            arguments = fn.get("arguments", {})
            if isinstance(arguments, str):
                try:
                    arguments = json.loads(arguments)
                except Exception:
                    arguments = {"_raw": arguments}

            if not isinstance(arguments, dict):
                arguments = {"args": arguments}

            tool_calls.append(
                ToolCall(
                    id=tc.get("id", f"tool_call_{idx}"),
                    name=name,
                    arguments=arguments,
                )
            )

        return tool_calls


================================================
FILE: src/vanna/integrations/openai/__init__.py
================================================
"""
OpenAI integration.

This module provides OpenAI LLM service implementations.
"""

from .llm import OpenAILlmService
from .responses import OpenAIResponsesService

__all__ = ["OpenAILlmService", "OpenAIResponsesService"]


================================================
FILE: src/vanna/integrations/openai/llm.py
================================================
"""
OpenAI LLM service implementation.

This module provides an implementation of the LlmService interface backed by
OpenAI's Chat Completions API (openai>=1.0.0). It supports non-streaming
responses and best-effort streaming of text content. Tool/function calling is
passed through when tools are provided, but full tool-call conversation
round-tripping may require adding assistant tool-call messages to the
conversation upstream.
"""

from __future__ import annotations

import json
import os
from typing import Any, AsyncGenerator, Dict, List, Optional, cast

from vanna.core.llm import (
    LlmService,
    LlmRequest,
    LlmResponse,
    LlmStreamChunk,
)
from vanna.core.tool import ToolCall, ToolSchema


class OpenAILlmService(LlmService):
    """OpenAI Chat Completions-backed LLM service.

    Args:
        model: OpenAI model name (e.g., "gpt-5").
        api_key: API key; falls back to env `OPENAI_API_KEY`.
        organization: Optional org; env `OPENAI_ORG` if unset.
        base_url: Optional custom base URL; env `OPENAI_BASE_URL` if unset.
        extra_client_kwargs: Extra kwargs forwarded to `openai.OpenAI()`.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        organization: Optional[str] = None,
        base_url: Optional[str] = None,
        **extra_client_kwargs: Any,
    ) -> None:
        try:
            from openai import OpenAI
        except Exception as e:  # pragma: no cover - import-time error surface
            raise ImportError(
                "openai package is required. Install with: pip install 'vanna[openai]'"
            ) from e

        self.model = model or os.getenv("OPENAI_MODEL", "gpt-5")
        api_key = api_key or os.getenv("OPENAI_API_KEY")
        organization = organization or os.getenv("OPENAI_ORG")
        base_url = base_url or os.getenv("OPENAI_BASE_URL")

        client_kwargs: Dict[str, Any] = {**extra_client_kwargs}
        if api_key:
            client_kwargs["api_key"] = api_key
        if organization:
            client_kwargs["organization"] = organization
        if base_url:
            client_kwargs["base_url"] = base_url

        self._client = OpenAI(**client_kwargs)

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        """Send a non-streaming request to OpenAI and return the response."""
        payload = self._build_payload(request)

        # Call the API synchronously; this function is async but we can block here.
        resp = self._client.chat.completions.create(**payload, stream=False)

        if not resp.choices:
            return LlmResponse(content=None, tool_calls=None, finish_reason=None)

        choice = resp.choices[0]
        content: Optional[str] = getattr(choice.message, "content", None)
        tool_calls = self._extract_tool_calls_from_message(choice.message)

        usage: Dict[str, int] = {}
        if getattr(resp, "usage", None):
            usage = {
                k: int(v)
                for k, v in {
                    "prompt_tokens": getattr(resp.usage, "prompt_tokens", 0),
                    "completion_tokens": getattr(resp.usage, "completion_tokens", 0),
                    "total_tokens": getattr(resp.usage, "total_tokens", 0),
                }.items()
            }

        return LlmResponse(
            content=content,
            tool_calls=tool_calls or None,
            finish_reason=getattr(choice, "finish_reason", None),
            usage=usage or None,
        )

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        """Stream a request to OpenAI.

        Emits `LlmStreamChunk` for textual deltas as they arrive. Tool-calls are
        accumulated and emitted in a final chunk when the stream ends.
        """
        payload = self._build_payload(request)

        # Synchronous streaming iterator; iterate within async context.
        stream = self._client.chat.completions.create(**payload, stream=True)

        # Builders for streamed tool-calls (index -> partial)
        tc_builders: Dict[int, Dict[str, Optional[str]]] = {}
        last_finish: Optional[str] = None

        for event in stream:
            if not getattr(event, "choices", None):
                continue

            choice = event.choices[0]
            delta = getattr(choice, "delta", None)
            if delta is None:
                # Some SDK versions use `event.choices[0].message` on the final packet
                last_finish = getattr(choice, "finish_reason", last_finish)
                continue

            # Text content
            content_piece: Optional[str] = getattr(delta, "content", None)
            if content_piece:
                yield LlmStreamChunk(content=content_piece)

            # Tool calls (streamed)
            streamed_tool_calls = getattr(delta, "tool_calls", None)
            if streamed_tool_calls:
                for tc in streamed_tool_calls:
                    idx = getattr(tc, "index", 0) or 0
                    b = tc_builders.setdefault(
                        idx, {"id": None, "name": None, "arguments": ""}
                    )
                    if getattr(tc, "id", None):
                        b["id"] = tc.id
                    fn = getattr(tc, "function", None)
                    if fn is not None:
                        if getattr(fn, "name", None):
                            b["name"] = fn.name
                        if getattr(fn, "arguments", None):
                            b["arguments"] = (b["arguments"] or "") + fn.arguments

            last_finish = getattr(choice, "finish_reason", last_finish)

        # Emit final tool-calls chunk if any
        final_tool_calls: List[ToolCall] = []
        for b in tc_builders.values():
            if not b.get("name"):
                continue
            args_raw = b.get("arguments") or "{}"
            try:
                loaded = json.loads(args_raw)
                if isinstance(loaded, dict):
                    args_dict: Dict[str, Any] = loaded
                else:
                    args_dict = {"args": loaded}
            except Exception:
                args_dict = {"_raw": args_raw}
            final_tool_calls.append(
                ToolCall(
                    id=b.get("id") or "tool_call",
                    name=b["name"] or "tool",
                    arguments=args_dict,
                )
            )

        if final_tool_calls:
            yield LlmStreamChunk(tool_calls=final_tool_calls, finish_reason=last_finish)
        else:
            # Still emit a terminal chunk to signal completion
            yield LlmStreamChunk(finish_reason=last_finish or "stop")

    async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
        """Validate tool schemas. Returns a list of error messages."""
        errors: List[str] = []
        # Basic checks; OpenAI will enforce further validation server-side.
        for t in tools:
            if not t.name or len(t.name) > 64:
                errors.append(f"Invalid tool name: {t.name!r}")
        return errors

    # Internal helpers
    def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
        messages: List[Dict[str, Any]] = []

        # Add system prompt as first message if provided
        if request.system_prompt:
            messages.append({"role": "system", "content": request.system_prompt})

        for m in request.messages:
            msg: Dict[str, Any] = {"role": m.role, "content": m.content}
            if m.role == "tool" and m.tool_call_id:
                msg["tool_call_id"] = m.tool_call_id
            elif m.role == "assistant" and m.tool_calls:
                # Convert tool calls to OpenAI format
                tool_calls_payload = []
                for tc in m.tool_calls:
                    tool_calls_payload.append(
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.name,
                                "arguments": json.dumps(tc.arguments),
                            },
                        }
                    )
                msg["tool_calls"] = tool_calls_payload
            messages.append(msg)

        tools_payload: Optional[List[Dict[str, Any]]] = None
        if request.tools:
            tools_payload = [
                {
                    "type": "function",
                    "function": {
                        "name": t.name,
                        "description": t.description,
                        "parameters": t.parameters,
                    },
                }
                for t in request.tools
            ]

        payload: Dict[str, Any] = {
            "model": self.model,
            "messages": messages,
        }
        if request.max_tokens is not None:
            payload["max_tokens"] = request.max_tokens
        if tools_payload:
            payload["tools"] = tools_payload
            payload["tool_choice"] = "auto"

        return payload

    def _extract_tool_calls_from_message(self, message: Any) -> List[ToolCall]:
        tool_calls: List[ToolCall] = []
        raw_tool_calls = getattr(message, "tool_calls", None) or []
        for tc in raw_tool_calls:
            fn = getattr(tc, "function", None)
            if not fn:
                continue
            args_raw = getattr(fn, "arguments", "{}")
            try:
                loaded = json.loads(args_raw)
                if isinstance(loaded, dict):
                    args_dict: Dict[str, Any] = loaded
                else:
                    args_dict = {"args": loaded}
            except Exception:
                args_dict = {"_raw": args_raw}
            tool_calls.append(
                ToolCall(
                    id=getattr(tc, "id", "tool_call"),
                    name=getattr(fn, "name", "tool"),
                    arguments=args_dict,
                )
            )
        return tool_calls


================================================
FILE: src/vanna/integrations/openai/responses.py
================================================
from __future__ import annotations

import json
import os
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, TYPE_CHECKING

from vanna.core.llm import LlmService, LlmRequest, LlmResponse, LlmStreamChunk
from vanna.core.tool import ToolCall, ToolSchema

if TYPE_CHECKING:
    from openai.types.responses import Response


class OpenAIResponsesService(LlmService):
    def __init__(
        self, api_key: Optional[str] = None, model: Optional[str] = None
    ) -> None:
        try:
            from openai import AsyncOpenAI
            from openai.types.responses import Response
        except Exception as e:  # pragma: no cover
            raise ImportError(
                "openai package is required. Install with: pip install 'vanna[openai]'"
            ) from e

        self.client = AsyncOpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY"))
        self.model = model or os.getenv("OPENAI_MODEL", "gpt-5")

    async def send_request(self, request: LlmRequest) -> LlmResponse:
        payload = self._payload(request)
        resp: Response = await self.client.responses.create(**payload)
        self._debug_print("response", resp)
        text, tools, status, usage = self._extract(resp)
        return LlmResponse(
            content=text,
            tool_calls=tools or None,
            finish_reason=status,
            usage=usage or None,
            metadata={"request_id": getattr(resp, "id", None)},
        )

    async def stream_request(
        self, request: LlmRequest
    ) -> AsyncGenerator[LlmStreamChunk, None]:
        payload = self._payload(request)
        async with self.client.responses.stream(**payload) as stream:
            async for event in stream:
                self._debug_print("stream_event", event)
                event_type = getattr(event, "type", None)
                if event_type == "response.output_text.delta":
                    delta = getattr(event, "delta", None)
                    if delta:
                        yield LlmStreamChunk(content=delta)
            final: Response = await stream.get_final_response()
            self._debug_print("final_response", final)

        _text, tools, status, _usage = self._extract(final)
        yield LlmStreamChunk(tool_calls=tools or None, finish_reason=status)

    async def validate_tools(self, tools: List[Any]) -> List[str]:
        return []  # minimal: accept whatever's passed through

    # ---- helpers ----

    def _payload(self, request: LlmRequest) -> Dict[str, Any]:
        msgs = [{"role": m.role, "content": m.content} for m in request.messages]
        p: Dict[str, Any] = {"model": self.model, "input": msgs}
        if request.system_prompt:
            p["instructions"] = request.system_prompt
        if request.max_tokens:
            p["max_output_tokens"] = request.max_tokens
        if request.tools:
            p["tools"] = [self._serialize_tool(t) for t in request.tools]
        return p

    def _debug_print(self, label: str, obj: Any) -> None:
        try:
            payload = obj.model_dump()
        except AttributeError:
            try:
                payload = obj.dict()
            except AttributeError:
                payload = obj
        print(f"[OpenAIResponsesService] {label}: {payload}")

    def _extract(
        self, resp: Response
    ) -> Tuple[
        Optional[str], Optional[List[ToolCall]], Optional[str], Optional[Dict[str, int]]
    ]:
        text = getattr(resp, "output_text", None)

        tool_calls: List[ToolCall] = []
        for oc in getattr(resp, "output", []) or []:
            for item in getattr(oc, "content", []) or []:
                if getattr(item, "type", None) == "tool_call":
                    tc = getattr(item, "tool_call", None)
                    if tc and getattr(tc, "type", None) == "function":
                        fn = getattr(tc, "function", None)
                        if fn:
                            name = getattr(fn, "name", None)
                            args = getattr(fn, "arguments", None)
                            if not isinstance(args, (dict, list)):
                                try:
                                    args = json.loads(args) if args else {}
                                except Exception:
                                    args = {"_raw": args}
                            tool_calls.append(ToolCall(name=name, arguments=args))

        usage = None
        if getattr(resp, "usage", None):
            usage = {
                "input_tokens": getattr(resp.usage, "input_tokens", 0) or 0,
                "output_tokens": getattr(resp.usage, "output_tokens", 0) or 0,
                "total_tokens": getattr(resp.usage, "total_tokens", None)
                or (
                    (getattr(resp.usage, "input_tokens", 0) or 0)
                    + (getattr(resp.usage, "output_tokens", 0) or 0)
                ),
            }

        status = getattr(resp, "status", None)  # e.g. "completed"
        return text, (tool_calls or None), status, usage

    def _serialize_tool(self, tool: Any) -> Dict[str, Any]:
        """Convert a tool schema into the dict format expected by OpenAI Responses."""

        if isinstance(tool, ToolSchema):
            return {
                "type": "function",
                "name": tool.name,
                "description": tool.description,
                "parameters": tool.parameters,
                "strict": False,
            }

        # Support generic pydantic/BaseModel style objects without importing pydantic here.
        if hasattr(tool, "model_dump"):
            data = tool.model_dump()
            if all(key in data for key in ("name", "description", "parameters")):
                return {
                    "type": "function",
                    "name": data["name"],
                    "description": data["description"],
                    "parameters": data["parameters"],
                    "strict": data.get("strict", False),
                }
            return data

        if isinstance(tool, dict):
            if "type" in tool:
                return tool
            if all(k in tool for k in ("name", "description", "parameters")):
                return {
                    "type": "function",
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": tool["parameters"],
                    "strict": tool.get("strict", False),
                }
            return tool

        raise TypeError(f"Unsupported tool schema type: {type(tool)!r}")


================================================
FILE: src/vanna/integrations/opensearch/__init__.py
================================================
"""
OpenSearch integration for Vanna Agents.
"""

from .agent_memory import OpenSearchAgentMemory

__all__ = ["OpenSearchAgentMemory"]


================================================
FILE: src/vanna/integrations/opensearch/agent_memory.py
================================================
"""
OpenSearch vector database implementation of AgentMemory.

This implementation uses OpenSearch for distributed search and storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    from opensearchpy import OpenSearch, helpers

    OPENSEARCH_AVAILABLE = True
except ImportError:
    OPENSEARCH_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class OpenSearchAgentMemory(AgentMemory):
    """OpenSearch-based implementation of AgentMemory."""

    def __init__(
        self,
        index_name: str = "tool_memories",
        hosts: Optional[List[str]] = None,
        http_auth: Optional[tuple] = None,
        use_ssl: bool = False,
        verify_certs: bool = False,
        dimension: int = 384,
    ):
        if not OPENSEARCH_AVAILABLE:
            raise ImportError(
                "OpenSearch is required for OpenSearchAgentMemory. Install with: pip install opensearch-py"
            )

        self.index_name = index_name
        self.hosts = hosts or ["localhost:9200"]
        self.http_auth = http_auth
        self.use_ssl = use_ssl
        self.verify_certs = verify_certs
        self.dimension = dimension
        self._client = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_client(self):
        """Get or create OpenSearch client."""
        if self._client is None:
            self._client = OpenSearch(
                hosts=self.hosts,
                http_auth=self.http_auth,
                use_ssl=self.use_ssl,
                verify_certs=self.verify_certs,
                ssl_show_warn=False,
            )

            # Create index if it doesn't exist
            if not self._client.indices.exists(index=self.index_name):
                index_body = {
                    "settings": {
                        "index": {"knn": True, "knn.algo_param.ef_search": 100}
                    },
                    "mappings": {
                        "properties": {
                            "memory_id": {"type": "keyword"},
                            "question": {"type": "text"},
                            "tool_name": {"type": "keyword"},
                            "args": {"type": "object", "enabled": False},
                            "timestamp": {"type": "date"},
                            "success": {"type": "boolean"},
                            "metadata": {"type": "object", "enabled": False},
                            "embedding": {
                                "type": "knn_vector",
                                "dimension": self.dimension,
                                "method": {
                                    "name": "hnsw",
                                    "space_type": "cosinesimil",
                                    "engine": "nmslib",
                                },
                            },
                        }
                    },
                }
                self._client.indices.create(index=self.index_name, body=index_body)

        return self._client

    def _create_embedding(self, text: str) -> List[float]:
        """Create a simple embedding from text (placeholder)."""
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            client = self._get_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            document = {
                "memory_id": memory_id,
                "question": question,
                "tool_name": tool_name,
                "args": args,
                "timestamp": timestamp,
                "success": success,
                "metadata": metadata or {},
                "embedding": embedding,
            }

            client.index(
                index=self.index_name, body=document, id=memory_id, refresh=True
            )

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            client = self._get_client()

            embedding = self._create_embedding(question)

            # Build query
            must_conditions = [{"term": {"success": True}}]
            if tool_name_filter:
                must_conditions.append({"term": {"tool_name": tool_name_filter}})

            query = {
                "size": limit,
                "query": {
                    "bool": {
                        "must": must_conditions,
                        "filter": {
                            "knn": {"embedding": {"vector": embedding, "k": limit}}
                        },
                    }
                },
            }

            response = client.search(index=self.index_name, body=query)

            search_results = []
            for i, hit in enumerate(response["hits"]["hits"]):
                source = hit["_source"]
                score = hit["_score"]

                # Normalize score to 0-1 range (OpenSearch scores can vary)
                similarity_score = min(score / 10.0, 1.0)

                if similarity_score >= similarity_threshold:
                    memory = ToolMemory(
                        memory_id=source["memory_id"],
                        question=source["question"],
                        tool_name=source["tool_name"],
                        args=source["args"],
                        timestamp=source.get("timestamp"),
                        success=source.get("success", True),
                        metadata=source.get("metadata", {}),
                    )

                    search_results.append(
                        ToolMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            client = self._get_client()

            query = {
                "size": limit,
                "query": {"match_all": {}},
                "sort": [{"timestamp": {"order": "desc"}}],
            }

            response = client.search(index=self.index_name, body=query)

            memories = []
            for hit in response["hits"]["hits"]:
                source = hit["_source"]

                memory = ToolMemory(
                    memory_id=source["memory_id"],
                    question=source["question"],
                    tool_name=source["tool_name"],
                    args=source["args"],
                    timestamp=source.get("timestamp"),
                    success=source.get("success", True),
                    metadata=source.get("metadata", {}),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            client = self._get_client()

            try:
                client.delete(index=self.index_name, id=memory_id, refresh=True)
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            client = self._get_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            document = {
                "memory_id": memory_id,
                "content": content,
                "timestamp": timestamp,
                "is_text_memory": True,
                "embedding": embedding,
            }

            client.index(
                index=self.index_name, body=document, id=memory_id, refresh=True
            )

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            client = self._get_client()

            embedding = self._create_embedding(query)

            query_body = {
                "size": limit,
                "query": {
                    "bool": {
                        "must": [{"term": {"is_text_memory": True}}],
                        "filter": {
                            "knn": {"embedding": {"vector": embedding, "k": limit}}
                        },
                    }
                },
            }

            response = client.search(index=self.index_name, body=query_body)

            search_results = []
            for i, hit in enumerate(response["hits"]["hits"]):
                source = hit["_source"]
                score = hit["_score"]

                similarity_score = min(score / 10.0, 1.0)

                if similarity_score >= similarity_threshold:
                    memory = TextMemory(
                        memory_id=source["memory_id"],
                        content=source.get("content", ""),
                        timestamp=source.get("timestamp"),
                    )

                    search_results.append(
                        TextMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            client = self._get_client()

            query = {
                "size": limit,
                "query": {"term": {"is_text_memory": True}},
                "sort": [{"timestamp": {"order": "desc"}}],
            }

            response = client.search(index=self.index_name, body=query)

            memories = []
            for hit in response["hits"]["hits"]:
                source = hit["_source"]

                memory = TextMemory(
                    memory_id=source["memory_id"],
                    content=source.get("content", ""),
                    timestamp=source.get("timestamp"),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            client = self._get_client()

            try:
                client.delete(index=self.index_name, id=memory_id, refresh=True)
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            client = self._get_client()

            # Build query
            must_conditions = []
            if tool_name:
                must_conditions.append({"term": {"tool_name": tool_name}})
            if before_date:
                must_conditions.append({"range": {"timestamp": {"lt": before_date}}})

            if must_conditions:
                query = {"query": {"bool": {"must": must_conditions}}}
            else:
                query = {"query": {"match_all": {}}}

            response = client.delete_by_query(
                index=self.index_name, body=query, refresh=True
            )

            return response.get("deleted", 0)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/oracle/__init__.py
================================================
"""Oracle integration for Vanna."""

from .sql_runner import OracleRunner

__all__ = ["OracleRunner"]


================================================
FILE: src/vanna/integrations/oracle/sql_runner.py
================================================
"""Oracle implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class OracleRunner(SqlRunner):
    """Oracle implementation of the SqlRunner interface."""

    def __init__(self, user: str, password: str, dsn: str, **kwargs):
        """Initialize with Oracle connection parameters.

        Args:
            user: Oracle database user name
            password: Oracle database user password
            dsn: Oracle database host - format: host:port/sid
            **kwargs: Additional oracledb connection parameters
        """
        try:
            import oracledb

            self.oracledb = oracledb
        except ImportError as e:
            raise ImportError(
                "oracledb package is required. Install with: pip install 'vanna[oracle]'"
            ) from e

        self.user = user
        self.password = password
        self.dsn = dsn
        self.kwargs = kwargs

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against Oracle database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            oracledb.Error: If query execution fails
        """
        # Connect to the database
        conn = self.oracledb.connect(
            user=self.user, password=self.password, dsn=self.dsn, **self.kwargs
        )

        cursor = conn.cursor()

        try:
            # Strip and remove trailing semicolons (Oracle doesn't like them)
            sql = args.sql.rstrip()
            if sql.endswith(";"):
                sql = sql[:-1]

            # Execute the query
            cursor.execute(sql)
            results = cursor.fetchall()

            # Create a pandas dataframe from the results
            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])
            return df

        except self.oracledb.Error:
            conn.rollback()
            raise
        finally:
            cursor.close()
            conn.close()


================================================
FILE: src/vanna/integrations/pinecone/__init__.py
================================================
"""
Pinecone integration for Vanna Agents.
"""

from .agent_memory import PineconeAgentMemory

__all__ = ["PineconeAgentMemory"]


================================================
FILE: src/vanna/integrations/pinecone/agent_memory.py
================================================
"""
Pinecone vector database implementation of AgentMemory.

This implementation uses Pinecone for cloud-based vector storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    from pinecone import Pinecone, ServerlessSpec

    PINECONE_AVAILABLE = True
except ImportError:
    PINECONE_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class PineconeAgentMemory(AgentMemory):
    """Pinecone-based implementation of AgentMemory."""

    def __init__(
        self,
        api_key: str,
        index_name: str = "tool-memories",
        environment: str = "us-east-1",
        dimension: int = 384,
        metric: str = "cosine",
    ):
        if not PINECONE_AVAILABLE:
            raise ImportError(
                "Pinecone is required for PineconeAgentMemory. Install with: pip install pinecone-client"
            )

        self.api_key = api_key
        self.index_name = index_name
        self.environment = environment
        self.dimension = dimension
        self.metric = metric
        self._client = None
        self._index = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_client(self):
        """Get or create Pinecone client."""
        if self._client is None:
            self._client = Pinecone(api_key=self.api_key)
        return self._client

    def _get_index(self):
        """Get or create Pinecone index."""
        if self._index is None:
            client = self._get_client()

            # Create index if it doesn't exist
            if self.index_name not in client.list_indexes().names():
                client.create_index(
                    name=self.index_name,
                    dimension=self.dimension,
                    metric=self.metric,
                    spec=ServerlessSpec(cloud="aws", region=self.environment),
                )

            self._index = client.Index(self.index_name)
        return self._index

    def _create_embedding(self, text: str) -> List[float]:
        """Create a simple embedding from text (placeholder - should use actual embedding model)."""
        # TODO: Replace with actual embedding model
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            index = self._get_index()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            # Pinecone metadata must be simple types
            memory_metadata = {
                "question": question,
                "tool_name": tool_name,
                "args_json": json.dumps(args),
                "timestamp": timestamp,
                "success": success,
                "metadata_json": json.dumps(metadata or {}),
            }

            index.upsert(vectors=[(memory_id, embedding, memory_metadata)])

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            index = self._get_index()

            embedding = self._create_embedding(question)

            # Build filter
            filter_dict = {"success": True}
            if tool_name_filter:
                filter_dict["tool_name"] = tool_name_filter

            results = index.query(
                vector=embedding, top_k=limit, filter=filter_dict, include_metadata=True
            )

            search_results = []
            for i, match in enumerate(results.matches):
                # Pinecone returns similarity score directly
                similarity_score = match.score

                if similarity_score >= similarity_threshold:
                    metadata = match.metadata
                    args = json.loads(metadata.get("args_json", "{}"))
                    metadata_dict = json.loads(metadata.get("metadata_json", "{}"))

                    memory = ToolMemory(
                        memory_id=match.id,
                        question=metadata["question"],
                        tool_name=metadata["tool_name"],
                        args=args,
                        timestamp=metadata.get("timestamp"),
                        success=metadata.get("success", True),
                        metadata=metadata_dict,
                    )

                    search_results.append(
                        ToolMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            index = self._get_index()

            # Pinecone doesn't have a native "get all" - we need to query with a dummy vector
            # or use the list operation with metadata filtering
            # This is a limitation - we'll return empty for now
            # In production, you'd maintain a separate timestamp index or use Pinecone's metadata filtering
            return []

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            index = self._get_index()

            try:
                index.delete(ids=[memory_id])
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            index = self._get_index()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            memory_metadata = {
                "content": content,
                "timestamp": timestamp,
                "is_text_memory": True,
            }

            index.upsert(vectors=[(memory_id, embedding, memory_metadata)])

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            index = self._get_index()

            embedding = self._create_embedding(query)

            filter_dict = {"is_text_memory": True}

            results = index.query(
                vector=embedding, top_k=limit, filter=filter_dict, include_metadata=True
            )

            search_results = []
            for i, match in enumerate(results.matches):
                similarity_score = match.score

                if similarity_score >= similarity_threshold:
                    metadata = match.metadata

                    memory = TextMemory(
                        memory_id=match.id,
                        content=metadata.get("content", ""),
                        timestamp=metadata.get("timestamp"),
                    )

                    search_results.append(
                        TextMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            # Pinecone doesn't have a native "get all sorted by timestamp" operation
            # This is a limitation - returning empty list
            # In production, you'd need to maintain a separate index or use metadata filtering
            return []

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            index = self._get_index()

            try:
                index.delete(ids=[memory_id])
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            index = self._get_index()

            # Build filter
            filter_dict = {}
            if tool_name:
                filter_dict["tool_name"] = tool_name
            if before_date:
                filter_dict["timestamp"] = {"$lt": before_date}

            if filter_dict:
                # Delete with filter
                index.delete(filter=filter_dict)
            else:
                # Delete all
                index.delete(delete_all=True)

            # Pinecone doesn't return count of deleted items
            return 0

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/plotly/__init__.py
================================================
"""Plotly integration for chart generation."""

from .chart_generator import PlotlyChartGenerator

__all__ = ["PlotlyChartGenerator"]


================================================
FILE: src/vanna/integrations/plotly/chart_generator.py
================================================
"""Plotly-based chart generator with automatic chart type selection."""

from typing import Dict, Any, List, cast
import json
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
import plotly.io as pio


class PlotlyChartGenerator:
    """Generate Plotly charts using heuristics based on DataFrame characteristics."""

    # Vanna brand colors from landing page
    THEME_COLORS = {
        "navy": "#023d60",
        "cream": "#e7e1cf",
        "teal": "#15a8a8",
        "orange": "#fe5d26",
        "magenta": "#bf1363",
    }

    # Color palette for charts (excluding cream as it's too light for data)
    COLOR_PALETTE = ["#15a8a8", "#fe5d26", "#bf1363", "#023d60"]

    def generate_chart(self, df: pd.DataFrame, title: str = "Chart") -> Dict[str, Any]:
        """Generate a Plotly chart based on DataFrame shape and types.

        Heuristics:
        - 4+ columns: table
        - 1 numeric column: histogram
        - 2 columns (1 categorical, 1 numeric): bar chart
        - 2 numeric columns: scatter plot
        - 3+ numeric columns: correlation heatmap or multi-line chart
        - Time series data: line chart
        - Multiple categorical: grouped bar chart

        Args:
            df: DataFrame to visualize
            title: Title for the chart

        Returns:
            Plotly figure as dictionary

        Raises:
            ValueError: If DataFrame is empty or cannot be visualized
        """
        if df.empty:
            raise ValueError("Cannot visualize empty DataFrame")

        # Heuristic: If 4 or more columns, render as a table
        if len(df.columns) >= 4:
            fig = self._create_table(df, title)
            result: Dict[str, Any] = json.loads(pio.to_json(fig))
            return result

        # Identify column types
        numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
        categorical_cols = df.select_dtypes(
            include=["object", "category"]
        ).columns.tolist()
        datetime_cols = df.select_dtypes(include=["datetime64"]).columns.tolist()

        # Check for time series
        is_timeseries = len(datetime_cols) > 0

        # Apply heuristics
        if is_timeseries and len(numeric_cols) > 0:
            # Time series line chart
            fig = self._create_time_series_chart(
                df, datetime_cols[0], numeric_cols, title
            )
        elif len(numeric_cols) == 1 and len(categorical_cols) == 0:
            # Single numeric column: histogram
            fig = self._create_histogram(df, numeric_cols[0], title)
        elif len(numeric_cols) == 1 and len(categorical_cols) == 1:
            # One categorical, one numeric: bar chart
            fig = self._create_bar_chart(
                df, categorical_cols[0], numeric_cols[0], title
            )
        elif len(numeric_cols) == 2:
            # Two numeric columns: scatter plot
            fig = self._create_scatter_plot(df, numeric_cols[0], numeric_cols[1], title)
        elif len(numeric_cols) >= 3:
            # Multiple numeric columns: correlation heatmap
            fig = self._create_correlation_heatmap(df, numeric_cols, title)
        elif len(categorical_cols) >= 2:
            # Multiple categorical: grouped bar chart
            fig = self._create_grouped_bar_chart(df, categorical_cols, title)
        else:
            # Fallback: show first two columns as scatter/bar
            if len(df.columns) >= 2:
                fig = self._create_generic_chart(
                    df, df.columns[0], df.columns[1], title
                )
            else:
                raise ValueError(
                    "Cannot determine appropriate visualization for this DataFrame"
                )

        # Convert to JSON-serializable dict using plotly's JSON encoder
        result = json.loads(pio.to_json(fig))
        return result

    def _apply_standard_layout(self, fig: go.Figure) -> go.Figure:
        """Apply consistent Vanna brand styling to all charts.

        Uses Vanna brand colors from the landing page for a cohesive look.

        Args:
            fig: Plotly figure to update

        Returns:
            Updated figure with Vanna brand styling
        """
        fig.update_layout(
            # paper_bgcolor='white',
            # plot_bgcolor='white',
            font={"color": self.THEME_COLORS["navy"]},  # Navy for text
            autosize=True,  # Allow chart to resize responsively
            colorway=self.COLOR_PALETTE,  # Use Vanna brand colors for data
            # Don't set width/height - let frontend handle sizing
        )
        return fig

    def _create_histogram(self, df: pd.DataFrame, column: str, title: str) -> go.Figure:
        """Create a histogram for a single numeric column."""
        fig = px.histogram(
            df,
            x=column,
            title=title,
            color_discrete_sequence=[self.THEME_COLORS["teal"]],
        )
        fig.update_layout(xaxis_title=column, yaxis_title="Count", showlegend=False)
        self._apply_standard_layout(fig)
        return fig

    def _create_bar_chart(
        self, df: pd.DataFrame, x_col: str, y_col: str, title: str
    ) -> go.Figure:
        """Create a bar chart for categorical vs numeric data."""
        # Aggregate if needed
        agg_df = df.groupby(x_col)[y_col].sum().reset_index()
        fig = px.bar(
            agg_df,
            x=x_col,
            y=y_col,
            title=title,
            color_discrete_sequence=[self.THEME_COLORS["orange"]],
        )
        fig.update_layout(xaxis_title=x_col, yaxis_title=y_col)
        self._apply_standard_layout(fig)
        return fig

    def _create_scatter_plot(
        self, df: pd.DataFrame, x_col: str, y_col: str, title: str
    ) -> go.Figure:
        """Create a scatter plot for two numeric columns."""
        fig = px.scatter(
            df,
            x=x_col,
            y=y_col,
            title=title,
            color_discrete_sequence=[self.THEME_COLORS["magenta"]],
        )
        fig.update_layout(xaxis_title=x_col, yaxis_title=y_col)
        self._apply_standard_layout(fig)
        return fig

    def _create_correlation_heatmap(
        self, df: pd.DataFrame, columns: List[str], title: str
    ) -> go.Figure:
        """Create a correlation heatmap for multiple numeric columns."""
        corr_matrix = df[columns].corr()
        # Custom Vanna color scale: navy (negative) -> cream (neutral) -> teal (positive)
        vanna_colorscale = [
            [0.0, self.THEME_COLORS["navy"]],
            [0.5, self.THEME_COLORS["cream"]],
            [1.0, self.THEME_COLORS["teal"]],
        ]
        fig = cast(
            go.Figure,
            px.imshow(
                corr_matrix,
                title=title,
                labels=dict(color="Correlation"),
                x=columns,
                y=columns,
                color_continuous_scale=vanna_colorscale,
                zmin=-1,
                zmax=1,
            ),
        )
        self._apply_standard_layout(fig)
        return fig

    def _create_time_series_chart(
        self, df: pd.DataFrame, time_col: str, value_cols: List[str], title: str
    ) -> go.Figure:
        """Create a time series line chart."""
        fig = go.Figure()

        for i, col in enumerate(value_cols[:5]):  # Limit to 5 lines for readability
            color = self.COLOR_PALETTE[i % len(self.COLOR_PALETTE)]
            fig.add_trace(
                go.Scatter(
                    x=df[time_col],
                    y=df[col],
                    mode="lines",
                    name=col,
                    line=dict(color=color),
                )
            )

        fig.update_layout(
            title=title,
            xaxis_title=time_col,
            yaxis_title="Value",
            hovermode="x unified",
        )
        self._apply_standard_layout(fig)
        return fig

    def _create_grouped_bar_chart(
        self, df: pd.DataFrame, categorical_cols: List[str], title: str
    ) -> go.Figure:
        """Create a grouped bar chart for multiple categorical columns."""
        # Use first two categorical columns
        if len(categorical_cols) >= 2:
            # Count occurrences
            grouped = df.groupby(categorical_cols[:2]).size().reset_index(name="count")
            fig = px.bar(
                grouped,
                x=categorical_cols[0],
                y="count",
                color=categorical_cols[1],
                title=title,
                barmode="group",
                color_discrete_sequence=self.COLOR_PALETTE,
            )
            self._apply_standard_layout(fig)
            return fig
        else:
            # Single categorical: value counts
            counts = df[categorical_cols[0]].value_counts().reset_index()
            counts.columns = [categorical_cols[0], "count"]
            fig = px.bar(
                counts,
                x=categorical_cols[0],
                y="count",
                title=title,
                color_discrete_sequence=[self.THEME_COLORS["teal"]],
            )
            self._apply_standard_layout(fig)
            return fig

    def _create_generic_chart(
        self, df: pd.DataFrame, col1: str, col2: str, title: str
    ) -> go.Figure:
        """Create a generic chart for any two columns."""
        # Try to determine the best representation
        if pd.api.types.is_numeric_dtype(df[col1]) and pd.api.types.is_numeric_dtype(
            df[col2]
        ):
            return self._create_scatter_plot(df, col1, col2, title)
        else:
            # Treat first as categorical, second as value
            fig = px.bar(
                df,
                x=col1,
                y=col2,
                title=title,
                color_discrete_sequence=[self.THEME_COLORS["orange"]],
            )
            self._apply_standard_layout(fig)
            return fig

    def _create_table(self, df: pd.DataFrame, title: str) -> go.Figure:
        """Create a Plotly table for DataFrames with 4 or more columns."""
        # Prepare header
        header_values = list(df.columns)

        # Prepare cell values (transpose to get columns)
        cell_values = [df[col].tolist() for col in df.columns]

        # Create the table
        fig = go.Figure(
            data=[
                go.Table(
                    header=dict(
                        values=header_values,
                        fill_color=self.THEME_COLORS["navy"],
                        font=dict(color="white", size=12),
                        align="left",
                    ),
                    cells=dict(
                        values=cell_values,
                        fill_color=[
                            [
                                self.THEME_COLORS["cream"] if i % 2 == 0 else "white"
                                for i in range(len(df))
                            ]
                        ],
                        font=dict(color=self.THEME_COLORS["navy"], size=11),
                        align="left",
                    ),
                )
            ]
        )

        fig.update_layout(title=title, font={"color": self.THEME_COLORS["navy"]})

        return fig


================================================
FILE: src/vanna/integrations/postgres/__init__.py
================================================
"""
PostgreSQL integration.

This module provides PostgreSQL runner implementation.
"""

from .sql_runner import PostgresRunner

__all__ = ["PostgresRunner"]


================================================
FILE: src/vanna/integrations/postgres/sql_runner.py
================================================
"""PostgreSQL implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class PostgresRunner(SqlRunner):
    """PostgreSQL implementation of the SqlRunner interface."""

    def __init__(
        self,
        connection_string: Optional[str] = None,
        host: Optional[str] = None,
        port: Optional[int] = 5432,
        database: Optional[str] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        **kwargs,
    ):
        """Initialize with PostgreSQL connection parameters.

        You can either provide a connection_string OR individual parameters (host, database, etc.).
        If connection_string is provided, it takes precedence.

        Args:
            connection_string: PostgreSQL connection string (e.g., "postgresql://user:password@host:port/database")
            host: Database host address
            port: Database port (default: 5432)
            database: Database name
            user: Database user
            password: Database password
            **kwargs: Additional psycopg2 connection parameters (sslmode, connect_timeout, etc.)
        """
        try:
            import psycopg2
            import psycopg2.extras

            self.psycopg2 = psycopg2
        except Exception as e:
            raise ImportError(
                "psycopg2 package is required. Install with: pip install 'vanna[postgres]'"
            ) from e

        if connection_string:
            self.connection_string = connection_string
            self.connection_params = None
        elif host and database and user:
            self.connection_string = None
            self.connection_params = {
                "host": host,
                "port": port,
                "database": database,
                "user": user,
                "password": password,
                **kwargs,
            }
        else:
            raise ValueError(
                "Either provide connection_string OR (host, database, and user) parameters"
            )

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against PostgreSQL database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            psycopg2.Error: If query execution fails
        """
        # Connect to the database using either connection string or parameters
        if self.connection_string:
            conn = self.psycopg2.connect(self.connection_string)
        else:
            conn = self.psycopg2.connect(**self.connection_params)

        cursor = conn.cursor(cursor_factory=self.psycopg2.extras.RealDictCursor)

        try:
            # Execute the query
            cursor.execute(args.sql)

            # Determine if this is a SELECT query or modification query
            query_type = args.sql.strip().upper().split()[0]

            if query_type == "SELECT":
                # Fetch results for SELECT queries
                rows = cursor.fetchall()
                if not rows:
                    # Return empty DataFrame
                    return pd.DataFrame()

                # Convert rows to list of dictionaries
                results_data = [dict(row) for row in rows]
                return pd.DataFrame(results_data)
            else:
                # For non-SELECT queries (INSERT, UPDATE, DELETE, etc.)
                conn.commit()
                rows_affected = cursor.rowcount
                # Return a DataFrame indicating rows affected
                return pd.DataFrame({"rows_affected": [rows_affected]})

        finally:
            cursor.close()
            conn.close()


================================================
FILE: src/vanna/integrations/premium/agent_memory/__init__.py
================================================
"""
Cloud-based agent memory implementations.
"""

from .premium import CloudAgentMemory

__all__ = ["CloudAgentMemory"]


================================================
FILE: src/vanna/integrations/premium/agent_memory/premium.py
================================================
"""
Cloud-based implementation of AgentMemory.

This implementation uses Vanna's premium cloud service for storing and searching
tool usage patterns with advanced similarity search and analytics.
"""

import json
from datetime import datetime
from typing import Any, Dict, List, Optional
import httpx

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class CloudAgentMemory(AgentMemory):
    """Cloud-based implementation of AgentMemory."""

    def __init__(
        self,
        api_base_url: str = "https://api.vanna.ai",
        api_key: Optional[str] = None,
        organization_id: Optional[str] = None,
    ):
        self.api_base_url = api_base_url.rstrip("/")
        self.api_key = api_key
        self.organization_id = organization_id
        self._client = httpx.AsyncClient(base_url=self.api_base_url, timeout=30.0)

    def _get_headers(self) -> Dict[str, str]:
        """Get request headers with authentication."""
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        if self.organization_id:
            headers["X-Organization-ID"] = self.organization_id
        return headers

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern to premium cloud storage."""
        import uuid

        payload = {
            "id": str(uuid.uuid4()),
            "question": question,
            "tool_name": tool_name,
            "args": args,
            "success": success,
            "metadata": metadata or {},
            "timestamp": datetime.now().isoformat(),
        }

        response = await self._client.post(
            "/memory/tool-usage", json=payload, headers=self._get_headers()
        )
        response.raise_for_status()

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns in premium cloud storage."""
        params = {
            "question": question,
            "limit": limit,
            "similarity_threshold": similarity_threshold,
        }
        if tool_name_filter:
            params["tool_name_filter"] = tool_name_filter

        response = await self._client.get(
            "/memory/search-similar", params=params, headers=self._get_headers()
        )
        response.raise_for_status()

        data = response.json()
        results = []

        for item in data.get("results", []):
            memory = ToolMemory(**item["memory"])
            result = ToolMemorySearchResult(
                memory=memory,
                similarity_score=item["similarity_score"],
                rank=item["rank"],
            )
            results.append(result)

        return results

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories from premium cloud storage."""
        params = {"limit": limit}

        response = await self._client.get(
            "/memory/recent", params=params, headers=self._get_headers()
        )
        response.raise_for_status()

        data = response.json()
        memories = []

        for item in data.get("memories", []):
            memory = ToolMemory(**item)
            memories.append(memory)

        return memories

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID from premium cloud storage."""
        response = await self._client.delete(
            f"/memory/{memory_id}", headers=self._get_headers()
        )

        if response.status_code == 404:
            return False

        response.raise_for_status()
        return True

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Cloud implementation does not yet support text memories."""
        raise NotImplementedError("CloudAgentMemory does not support text memories.")

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Cloud implementation does not yet support text memories."""
        return []

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Cloud implementation does not yet support text memories."""
        return []

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Cloud implementation does not yet support text memories."""
        return False

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories from premium cloud storage."""
        payload = {}
        if tool_name:
            payload["tool_name"] = tool_name
        if before_date:
            payload["before_date"] = before_date

        response = await self._client.delete(
            "/memory/clear", json=payload, headers=self._get_headers()
        )
        response.raise_for_status()

        data = response.json()
        return data.get("deleted_count", 0)


================================================
FILE: src/vanna/integrations/presto/__init__.py
================================================
"""Presto integration for Vanna."""

from .sql_runner import PrestoRunner

__all__ = ["PrestoRunner"]


================================================
FILE: src/vanna/integrations/presto/sql_runner.py
================================================
"""Presto implementation of SqlRunner interface."""

from typing import Optional
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class PrestoRunner(SqlRunner):
    """Presto implementation of the SqlRunner interface."""

    def __init__(
        self,
        host: str,
        catalog: str = "hive",
        schema: str = "default",
        user: Optional[str] = None,
        password: Optional[str] = None,
        port: int = 443,
        combined_pem_path: Optional[str] = None,
        protocol: str = "https",
        requests_kwargs: Optional[dict] = None,
        **kwargs,
    ):
        """Initialize with Presto connection parameters.

        Args:
            host: The host address of the Presto database
            catalog: The catalog to use in the Presto environment (default: 'hive')
            schema: The schema to use in the Presto environment (default: 'default')
            user: The username for authentication
            password: The password for authentication
            port: The port number for the Presto connection (default: 443)
            combined_pem_path: The path to the combined pem file for SSL connection
            protocol: The protocol to use for the connection (default: 'https')
            requests_kwargs: Additional keyword arguments for requests
            **kwargs: Additional pyhive connection parameters
        """
        try:
            from pyhive import presto

            self.presto = presto
        except ImportError as e:
            raise ImportError(
                "pyhive package is required. Install with: pip install pyhive"
            ) from e

        self.host = host
        self.catalog = catalog
        self.schema = schema
        self.user = user
        self.password = password
        self.port = port
        self.protocol = protocol
        self.kwargs = kwargs

        # Set up requests_kwargs for SSL if combined_pem_path is provided
        if requests_kwargs is None and combined_pem_path is not None:
            self.requests_kwargs = {"verify": combined_pem_path}
        else:
            self.requests_kwargs = requests_kwargs

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against Presto database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            presto.Error: If query execution fails
        """
        # Connect to the database
        conn = self.presto.Connection(
            host=self.host,
            username=self.user,
            password=self.password,
            catalog=self.catalog,
            schema=self.schema,
            port=self.port,
            protocol=self.protocol,
            requests_kwargs=self.requests_kwargs,
            **self.kwargs,
        )

        try:
            # Strip and remove trailing semicolons (Presto doesn't like them)
            sql = args.sql.rstrip()
            if sql.endswith(";"):
                sql = sql[:-1]

            cursor = conn.cursor()
            cursor.execute(sql)
            results = cursor.fetchall()

            # Create a pandas dataframe from the results
            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])

            cursor.close()
            return df

        finally:
            conn.close()


================================================
FILE: src/vanna/integrations/qdrant/__init__.py
================================================
"""
Qdrant integration for Vanna Agents.
"""

from .agent_memory import QdrantAgentMemory

__all__ = ["QdrantAgentMemory"]


================================================
FILE: src/vanna/integrations/qdrant/agent_memory.py
================================================
"""
Qdrant vector database implementation of AgentMemory.

This implementation uses Qdrant for vector storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    from qdrant_client import QdrantClient
    from qdrant_client.models import (
        Distance,
        VectorParams,
        PointStruct,
        Filter,
        FieldCondition,
        MatchValue,
    )

    QDRANT_AVAILABLE = True
except ImportError:
    QDRANT_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class QdrantAgentMemory(AgentMemory):
    """Qdrant-based implementation of AgentMemory."""

    def __init__(
        self,
        collection_name: str = "tool_memories",
        url: Optional[str] = None,
        path: Optional[str] = None,
        api_key: Optional[str] = None,
        dimension: int = 384,
    ):
        if not QDRANT_AVAILABLE:
            raise ImportError(
                "Qdrant is required for QdrantAgentMemory. Install with: pip install qdrant-client"
            )

        self.collection_name = collection_name
        self.url = url
        self.path = path
        self.api_key = api_key
        self.dimension = dimension
        self._client = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_client(self):
        """Get or create Qdrant client."""
        if self._client is None:
            if self.url:
                self._client = QdrantClient(url=self.url, api_key=self.api_key)
            else:
                self._client = QdrantClient(path=self.path or ":memory:")

            # Create collection if it doesn't exist
            collections = self._client.get_collections().collections
            if not any(c.name == self.collection_name for c in collections):
                self._client.create_collection(
                    collection_name=self.collection_name,
                    vectors_config=VectorParams(
                        size=self.dimension, distance=Distance.COSINE
                    ),
                )
        return self._client

    def _create_embedding(self, text: str) -> List[float]:
        """Create a simple embedding from text (placeholder)."""
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            client = self._get_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            payload = {
                "question": question,
                "tool_name": tool_name,
                "args": args,
                "timestamp": timestamp,
                "success": success,
                "metadata": metadata or {},
            }

            point = PointStruct(id=memory_id, vector=embedding, payload=payload)

            client.upsert(collection_name=self.collection_name, points=[point])

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            client = self._get_client()

            embedding = self._create_embedding(question)

            # Build filter
            query_filter = None
            conditions = [FieldCondition(key="success", match=MatchValue(value=True))]
            if tool_name_filter:
                conditions.append(
                    FieldCondition(
                        key="tool_name", match=MatchValue(value=tool_name_filter)
                    )
                )

            if conditions:
                query_filter = Filter(must=conditions)

            # Use query_points for newer qdrant-client (1.8.0+) or search for older versions
            if hasattr(client, "query_points"):
                results = client.query_points(
                    collection_name=self.collection_name,
                    query=embedding,
                    limit=limit,
                    query_filter=query_filter,
                    score_threshold=similarity_threshold,
                ).points
            else:
                # Fallback to search method for older qdrant-client versions
                results = client.search(
                    collection_name=self.collection_name,
                    query_vector=embedding,
                    limit=limit,
                    query_filter=query_filter,
                    score_threshold=similarity_threshold,
                )

            search_results = []
            for i, hit in enumerate(results):
                payload = hit.payload

                memory = ToolMemory(
                    memory_id=str(hit.id),
                    question=payload["question"],
                    tool_name=payload["tool_name"],
                    args=payload["args"],
                    timestamp=payload.get("timestamp"),
                    success=payload.get("success", True),
                    metadata=payload.get("metadata", {}),
                )

                search_results.append(
                    ToolMemorySearchResult(
                        memory=memory, similarity_score=hit.score, rank=i + 1
                    )
                )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            client = self._get_client()

            # Scroll through all points and sort by timestamp
            points, _ = client.scroll(
                collection_name=self.collection_name,
                limit=1000,  # Get more than we need to sort
                with_payload=True,
                with_vectors=False,
            )

            # Sort by timestamp
            sorted_points = sorted(
                points, key=lambda p: p.payload.get("timestamp", ""), reverse=True
            )

            memories = []
            for point in sorted_points[:limit]:
                payload = point.payload

                # Skip text memories - they have is_text_memory flag
                if payload.get("is_text_memory"):
                    continue

                memory = ToolMemory(
                    memory_id=str(point.id),
                    question=payload["question"],
                    tool_name=payload["tool_name"],
                    args=payload["args"],
                    timestamp=payload.get("timestamp"),
                    success=payload.get("success", True),
                    metadata=payload.get("metadata", {}),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID. Returns True if deleted, False if not found."""

        def _delete():
            client = self._get_client()

            try:
                # Check if the point exists before attempting to delete
                points = client.retrieve(
                    collection_name=self.collection_name,
                    ids=[memory_id],
                    with_payload=False,
                    with_vectors=False,
                )

                if points and len(points) > 0:
                    client.delete(
                        collection_name=self.collection_name,
                        points_selector=[memory_id],
                    )
                    return True
                return False
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            client = self._get_client()

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            payload = {
                "content": content,
                "timestamp": timestamp,
                "is_text_memory": True,
            }

            point = PointStruct(id=memory_id, vector=embedding, payload=payload)

            client.upsert(collection_name=self.collection_name, points=[point])

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            client = self._get_client()

            embedding = self._create_embedding(query)

            query_filter = Filter(
                must=[
                    FieldCondition(key="is_text_memory", match=MatchValue(value=True))
                ]
            )

            # Use query_points for newer qdrant-client (1.8.0+) or search for older versions
            if hasattr(client, "query_points"):
                results = client.query_points(
                    collection_name=self.collection_name,
                    query=embedding,
                    limit=limit,
                    query_filter=query_filter,
                    score_threshold=similarity_threshold,
                ).points
            else:
                # Fallback to search method for older qdrant-client versions
                results = client.search(
                    collection_name=self.collection_name,
                    query_vector=embedding,
                    limit=limit,
                    query_filter=query_filter,
                    score_threshold=similarity_threshold,
                )

            search_results = []
            for i, hit in enumerate(results):
                payload = hit.payload

                memory = TextMemory(
                    memory_id=str(hit.id),
                    content=payload.get("content", ""),
                    timestamp=payload.get("timestamp"),
                )

                search_results.append(
                    TextMemorySearchResult(
                        memory=memory, similarity_score=hit.score, rank=i + 1
                    )
                )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            client = self._get_client()

            # Scroll through text memory points and sort by timestamp
            points, _ = client.scroll(
                collection_name=self.collection_name,
                scroll_filter=Filter(
                    must=[
                        FieldCondition(
                            key="is_text_memory", match=MatchValue(value=True)
                        )
                    ]
                ),
                limit=1000,
                with_payload=True,
                with_vectors=False,
            )

            # Sort by timestamp
            sorted_points = sorted(
                points, key=lambda p: p.payload.get("timestamp", ""), reverse=True
            )

            memories = []
            for point in sorted_points[:limit]:
                payload = point.payload
                memory = TextMemory(
                    memory_id=str(point.id),
                    content=payload.get("content", ""),
                    timestamp=payload.get("timestamp"),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            client = self._get_client()

            try:
                # Check if the point exists before attempting to delete
                points = client.retrieve(
                    collection_name=self.collection_name,
                    ids=[memory_id],
                    with_payload=False,
                    with_vectors=False,
                )

                if points and len(points) > 0:
                    client.delete(
                        collection_name=self.collection_name,
                        points_selector=[memory_id],
                    )
                    return True
                return False
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            client = self._get_client()

            # Build filter
            conditions = []
            if tool_name:
                conditions.append(
                    FieldCondition(key="tool_name", match=MatchValue(value=tool_name))
                )
            if before_date:
                conditions.append(
                    FieldCondition(key="timestamp", match=MatchValue(value=before_date))
                )

            if conditions or (tool_name is None and before_date is None):
                # Delete with filter or delete all
                query_filter = Filter(must=conditions) if conditions else None

                if query_filter:
                    client.delete(
                        collection_name=self.collection_name,
                        points_selector=query_filter,
                    )
                else:
                    # Delete all points
                    client.delete_collection(collection_name=self.collection_name)
                    # Recreate empty collection
                    client.create_collection(
                        collection_name=self.collection_name,
                        vectors_config=VectorParams(
                            size=self.dimension, distance=Distance.COSINE
                        ),
                    )

            return 0  # Qdrant doesn't return count

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/integrations/snowflake/__init__.py
================================================
"""Snowflake integration for Vanna."""

from .sql_runner import SnowflakeRunner

__all__ = ["SnowflakeRunner"]


================================================
FILE: src/vanna/integrations/snowflake/sql_runner.py
================================================
"""Snowflake implementation of SqlRunner interface."""

from typing import Optional, Union
import os
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class SnowflakeRunner(SqlRunner):
    """Snowflake implementation of the SqlRunner interface."""

    def __init__(
        self,
        account: str,
        username: str,
        password: Optional[str] = None,
        database: str = "",
        role: Optional[str] = None,
        warehouse: Optional[str] = None,
        private_key_path: Optional[str] = None,
        private_key_passphrase: Optional[str] = None,
        private_key_content: Optional[bytes] = None,
        **kwargs,
    ):
        """Initialize with Snowflake connection parameters.

        Args:
            account: Snowflake account identifier
            username: Database user
            password: Database password (optional if using key-pair auth)
            database: Database name
            role: Snowflake role to use (optional)
            warehouse: Snowflake warehouse to use (optional)
            private_key_path: Path to private key file for RSA key-pair authentication (optional)
            private_key_passphrase: Passphrase for encrypted private key (optional)
            private_key_content: Private key content as bytes (optional, alternative to private_key_path)
            **kwargs: Additional snowflake.connector connection parameters

        Note:
            Either password OR private_key_path/private_key_content must be provided.
            RSA key-pair authentication is recommended for production systems as Snowflake
            is deprecating user/password authentication.
        """
        try:
            import snowflake.connector

            self.snowflake = snowflake.connector
        except ImportError as e:
            raise ImportError(
                "snowflake-connector-python package is required. "
                "Install with: pip install 'vanna[snowflake]'"
            ) from e

        # Validate that at least one authentication method is provided
        if not password and not private_key_path and not private_key_content:
            raise ValueError(
                "Either password or private_key_path/private_key_content must be provided for authentication"
            )

        # Validate private key path exists if provided
        if private_key_path and not os.path.isfile(private_key_path):
            raise FileNotFoundError(f"Private key file not found: {private_key_path}")

        self.account = account
        self.username = username
        self.password = password
        self.database = database
        self.role = role
        self.warehouse = warehouse
        self.private_key_path = private_key_path
        self.private_key_passphrase = private_key_passphrase
        self.private_key_content = private_key_content
        self.kwargs = kwargs

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against Snowflake database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            snowflake.connector.Error: If query execution fails
        """
        # Build connection parameters
        conn_params = {
            "user": self.username,
            "account": self.account,
            "client_session_keep_alive": True,
        }

        # Add database if specified
        if self.database:
            conn_params["database"] = self.database

        # Configure authentication method
        if self.private_key_path or self.private_key_content:
            # Use RSA key-pair authentication
            if self.private_key_path:
                conn_params["private_key_path"] = self.private_key_path
            else:
                conn_params["private_key_content"] = self.private_key_content

            # Add passphrase if provided
            if self.private_key_passphrase:
                conn_params["private_key_passphrase"] = self.private_key_passphrase
        else:
            # Use password authentication (fallback)
            conn_params["password"] = self.password

        # Add any additional kwargs
        conn_params.update(self.kwargs)

        # Connect to the database
        conn = self.snowflake.connect(**conn_params)

        cursor = conn.cursor()

        try:
            # Set role if specified
            if self.role:
                cursor.execute(f"USE ROLE {self.role}")

            # Set warehouse if specified
            if self.warehouse:
                cursor.execute(f"USE WAREHOUSE {self.warehouse}")

            # Use the specified database if provided
            if self.database:
                cursor.execute(f"USE DATABASE {self.database}")

            # Execute the query
            cursor.execute(args.sql)
            results = cursor.fetchall()

            # Create a pandas dataframe from the results
            df = pd.DataFrame(results, columns=[desc[0] for desc in cursor.description])
            return df

        finally:
            cursor.close()
            conn.close()


================================================
FILE: src/vanna/integrations/sqlite/__init__.py
================================================
"""
SQLite integration.

This module provides SQLite runner implementation.
"""

from .sql_runner import SqliteRunner

__all__ = ["SqliteRunner"]


================================================
FILE: src/vanna/integrations/sqlite/sql_runner.py
================================================
"""SQLite implementation of SqlRunner interface."""

import sqlite3
import pandas as pd

from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.core.tool import ToolContext


class SqliteRunner(SqlRunner):
    """SQLite implementation of the SqlRunner interface."""

    def __init__(self, database_path: str):
        """Initialize with a SQLite database path.

        Args:
            database_path: Path to the SQLite database file
        """
        self.database_path = database_path

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query against SQLite database and return results as DataFrame.

        Args:
            args: SQL query arguments
            context: Tool execution context

        Returns:
            DataFrame with query results

        Raises:
            sqlite3.Error: If query execution fails
        """
        # Connect to the database
        conn = sqlite3.connect(self.database_path)
        conn.row_factory = sqlite3.Row  # Enable column access by name
        cursor = conn.cursor()

        try:
            # Execute the query
            cursor.execute(args.sql)

            # Determine if this is a SELECT query or modification query
            query_type = args.sql.strip().upper().split()[0]

            if query_type == "SELECT":
                # Fetch results for SELECT queries
                rows = cursor.fetchall()
                if not rows:
                    # Return empty DataFrame
                    return pd.DataFrame()

                # Convert rows to list of dictionaries
                results_data = [dict(row) for row in rows]
                return pd.DataFrame(results_data)
            else:
                # For non-SELECT queries (INSERT, UPDATE, DELETE, etc.)
                conn.commit()
                rows_affected = cursor.rowcount
                # Return a DataFrame indicating rows affected
                return pd.DataFrame({"rows_affected": [rows_affected]})

        finally:
            cursor.close()
            conn.close()


================================================
FILE: src/vanna/integrations/weaviate/__init__.py
================================================
"""
Weaviate integration for Vanna Agents.
"""

from .agent_memory import WeaviateAgentMemory

__all__ = ["WeaviateAgentMemory"]


================================================
FILE: src/vanna/integrations/weaviate/agent_memory.py
================================================
"""
Weaviate vector database implementation of AgentMemory.

This implementation uses Weaviate for semantic search and storage of tool usage patterns.
"""

import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
import asyncio
from concurrent.futures import ThreadPoolExecutor

try:
    import weaviate
    from weaviate.classes.config import (
        Configure,
        Property,
        DataType as WeaviateDataType,
    )

    WEAVIATE_AVAILABLE = True
except ImportError:
    WEAVIATE_AVAILABLE = False

from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from vanna.core.tool import ToolContext


class WeaviateAgentMemory(AgentMemory):
    """Weaviate-based implementation of AgentMemory."""

    def __init__(
        self,
        collection_name: str = "ToolMemory",
        url: str = "http://localhost:8080",
        api_key: Optional[str] = None,
        dimension: int = 384,
    ):
        if not WEAVIATE_AVAILABLE:
            raise ImportError(
                "Weaviate is required for WeaviateAgentMemory. Install with: pip install weaviate-client"
            )

        self.collection_name = collection_name
        self.url = url
        self.api_key = api_key
        self.dimension = dimension
        self._client = None
        self._executor = ThreadPoolExecutor(max_workers=2)

    def _get_client(self):
        """Get or create Weaviate client."""
        if self._client is None:
            if self.api_key:
                self._client = weaviate.connect_to_weaviate_cloud(
                    cluster_url=self.url,
                    auth_credentials=weaviate.auth.AuthApiKey(self.api_key),
                )
            else:
                self._client = weaviate.connect_to_local(
                    host=self.url.replace("http://", "").replace("https://", "")
                )

            # Create collection if it doesn't exist
            if not self._client.collections.exists(self.collection_name):
                self._client.collections.create(
                    name=self.collection_name,
                    vectorizer_config=Configure.Vectorizer.none(),
                    properties=[
                        Property(name="question", data_type=WeaviateDataType.TEXT),
                        Property(name="tool_name", data_type=WeaviateDataType.TEXT),
                        Property(name="args_json", data_type=WeaviateDataType.TEXT),
                        Property(name="timestamp", data_type=WeaviateDataType.TEXT),
                        Property(name="success", data_type=WeaviateDataType.BOOL),
                        Property(name="metadata_json", data_type=WeaviateDataType.TEXT),
                    ],
                )

        return self._client

    def _create_embedding(self, text: str) -> List[float]:
        """Create a simple embedding from text (placeholder)."""
        import hashlib

        hash_val = int(hashlib.md5(text.encode()).hexdigest(), 16)
        return [(hash_val >> i) % 100 / 100.0 for i in range(self.dimension)]

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern."""

        def _save():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(question)

            properties = {
                "question": question,
                "tool_name": tool_name,
                "args_json": json.dumps(args),
                "timestamp": timestamp,
                "success": success,
                "metadata_json": json.dumps(metadata or {}),
            }

            collection.data.insert(
                properties=properties, vector=embedding, uuid=memory_id
            )

        await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns."""

        def _search():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            embedding = self._create_embedding(question)

            # Build filter
            filters = weaviate.classes.query.Filter.by_property("success").equal(True)
            if tool_name_filter:
                filters = filters & weaviate.classes.query.Filter.by_property(
                    "tool_name"
                ).equal(tool_name_filter)

            response = collection.query.near_vector(
                near_vector=embedding,
                limit=limit,
                filters=filters,
                return_metadata=weaviate.classes.query.MetadataQuery(distance=True),
            )

            search_results = []
            for i, obj in enumerate(response.objects):
                # Weaviate returns distance, convert to similarity
                distance = obj.metadata.distance if obj.metadata else 1.0
                similarity_score = 1 - distance

                if similarity_score >= similarity_threshold:
                    properties = obj.properties
                    args = json.loads(properties.get("args_json", "{}"))
                    metadata_dict = json.loads(properties.get("metadata_json", "{}"))

                    memory = ToolMemory(
                        memory_id=str(obj.uuid),
                        question=properties.get("question"),
                        tool_name=properties.get("tool_name"),
                        args=args,
                        timestamp=properties.get("timestamp"),
                        success=properties.get("success", True),
                        metadata=metadata_dict,
                    )

                    search_results.append(
                        ToolMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories."""

        def _get_recent():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            # Query and sort by timestamp
            response = collection.query.fetch_objects(limit=1000)

            # Convert to list and sort
            objects_list = list(response.objects)
            sorted_objects = sorted(
                objects_list,
                key=lambda o: o.properties.get("timestamp", ""),
                reverse=True,
            )

            memories = []
            for obj in sorted_objects[:limit]:
                properties = obj.properties
                args = json.loads(properties.get("args_json", "{}"))
                metadata_dict = json.loads(properties.get("metadata_json", "{}"))

                memory = ToolMemory(
                    memory_id=str(obj.uuid),
                    question=properties.get("question"),
                    tool_name=properties.get("tool_name"),
                    args=args,
                    timestamp=properties.get("timestamp"),
                    success=properties.get("success", True),
                    metadata=metadata_dict,
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID."""

        def _delete():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            try:
                collection.data.delete_by_id(uuid=memory_id)
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save a text memory."""

        def _save():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            memory_id = str(uuid.uuid4())
            timestamp = datetime.now().isoformat()
            embedding = self._create_embedding(content)

            properties = {
                "question": content,  # Using question field for content
                "tool_name": "",  # Empty for text memories
                "args_json": "",
                "timestamp": timestamp,
                "success": True,
                "metadata_json": json.dumps({"is_text_memory": True}),
            }

            collection.data.insert(
                properties=properties, vector=embedding, uuid=memory_id
            )

            return TextMemory(memory_id=memory_id, content=content, timestamp=timestamp)

        return await asyncio.get_event_loop().run_in_executor(self._executor, _save)

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search for similar text memories."""

        def _search():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            embedding = self._create_embedding(query)

            # Build filter for text memories (empty tool_name)
            filters = weaviate.classes.query.Filter.by_property("tool_name").equal("")

            response = collection.query.near_vector(
                near_vector=embedding,
                limit=limit,
                filters=filters,
                return_metadata=weaviate.classes.query.MetadataQuery(distance=True),
            )

            search_results = []
            for i, obj in enumerate(response.objects):
                distance = obj.metadata.distance if obj.metadata else 1.0
                similarity_score = 1 - distance

                if similarity_score >= similarity_threshold:
                    properties = obj.properties
                    content = properties.get("question", "")

                    memory = TextMemory(
                        memory_id=str(obj.uuid),
                        content=content,
                        timestamp=properties.get("timestamp"),
                    )

                    search_results.append(
                        TextMemorySearchResult(
                            memory=memory, similarity_score=similarity_score, rank=i + 1
                        )
                    )

            return search_results

        return await asyncio.get_event_loop().run_in_executor(self._executor, _search)

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Get recently added text memories."""

        def _get_recent():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            # Query text memories (empty tool_name) and sort by timestamp
            response = collection.query.fetch_objects(
                filters=weaviate.classes.query.Filter.by_property("tool_name").equal(
                    ""
                ),
                limit=1000,
            )

            # Convert to list and sort
            objects_list = list(response.objects)
            sorted_objects = sorted(
                objects_list,
                key=lambda o: o.properties.get("timestamp", ""),
                reverse=True,
            )

            memories = []
            for obj in sorted_objects[:limit]:
                properties = obj.properties
                content = properties.get("question", "")

                memory = TextMemory(
                    memory_id=str(obj.uuid),
                    content=content,
                    timestamp=properties.get("timestamp"),
                )
                memories.append(memory)

            return memories

        return await asyncio.get_event_loop().run_in_executor(
            self._executor, _get_recent
        )

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID."""

        def _delete():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            try:
                collection.data.delete_by_id(uuid=memory_id)
                return True
            except Exception:
                return False

        return await asyncio.get_event_loop().run_in_executor(self._executor, _delete)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories."""

        def _clear():
            client = self._get_client()
            collection = client.collections.get(self.collection_name)

            # Build filter
            if tool_name and before_date:
                filters = weaviate.classes.query.Filter.by_property("tool_name").equal(
                    tool_name
                ) & weaviate.classes.query.Filter.by_property("timestamp").less_than(
                    before_date
                )
            elif tool_name:
                filters = weaviate.classes.query.Filter.by_property("tool_name").equal(
                    tool_name
                )
            elif before_date:
                filters = weaviate.classes.query.Filter.by_property(
                    "timestamp"
                ).less_than(before_date)
            else:
                filters = None

            if filters:
                collection.data.delete_many(where=filters)
            else:
                # Delete all
                collection.data.delete_many(
                    where=weaviate.classes.query.Filter.by_property(
                        "success"
                    ).contains_any([True, False])
                )

            return 0

        return await asyncio.get_event_loop().run_in_executor(self._executor, _clear)


================================================
FILE: src/vanna/legacy/ZhipuAI/ZhipuAI_Chat.py
================================================
import re
from typing import List

import pandas as pd
from zhipuai import ZhipuAI

from ..base import VannaBase


class ZhipuAI_Chat(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        if config is None:
            return
        if "api_key" not in config:
            raise Exception("Missing api_key in config")
        self.api_key = config["api_key"]
        self.model = config["model"] if "model" in config else "glm-4"
        self.api_url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"

    # Static methods similar to those in ZhipuAI_Chat for message formatting and utility
    @staticmethod
    def system_message(message: str) -> dict:
        return {"role": "system", "content": message}

    @staticmethod
    def user_message(message: str) -> dict:
        return {"role": "user", "content": message}

    @staticmethod
    def assistant_message(message: str) -> dict:
        return {"role": "assistant", "content": message}

    @staticmethod
    def str_to_approx_token_count(string: str) -> int:
        return len(string) / 4

    @staticmethod
    def add_ddl_to_prompt(
        initial_prompt: str, ddl_list: List[str], max_tokens: int = 14000
    ) -> str:
        if len(ddl_list) > 0:
            initial_prompt += "\nYou may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"

            for ddl in ddl_list:
                if (
                    ZhipuAI_Chat.str_to_approx_token_count(initial_prompt)
                    + ZhipuAI_Chat.str_to_approx_token_count(ddl)
                    < max_tokens
                ):
                    initial_prompt += f"{ddl}\n\n"

        return initial_prompt

    @staticmethod
    def add_documentation_to_prompt(
        initial_prompt: str, documentation_List: List[str], max_tokens: int = 14000
    ) -> str:
        if len(documentation_List) > 0:
            initial_prompt += "\nYou may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"

            for documentation in documentation_List:
                if (
                    ZhipuAI_Chat.str_to_approx_token_count(initial_prompt)
                    + ZhipuAI_Chat.str_to_approx_token_count(documentation)
                    < max_tokens
                ):
                    initial_prompt += f"{documentation}\n\n"

        return initial_prompt

    @staticmethod
    def add_sql_to_prompt(
        initial_prompt: str, sql_List: List[str], max_tokens: int = 14000
    ) -> str:
        if len(sql_List) > 0:
            initial_prompt += "\nYou may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"

            for question in sql_List:
                if (
                    ZhipuAI_Chat.str_to_approx_token_count(initial_prompt)
                    + ZhipuAI_Chat.str_to_approx_token_count(question["sql"])
                    < max_tokens
                ):
                    initial_prompt += f"{question['question']}\n{question['sql']}\n\n"

        return initial_prompt

    def get_sql_prompt(
        self,
        question: str,
        question_sql_list: List,
        ddl_list: List,
        doc_list: List,
        **kwargs,
    ):
        initial_prompt = "The user provides a question and you provide SQL. You will only respond with SQL code and not with any explanations.\n\nRespond with only SQL code. Do not answer with any explanations -- just the code.\n"

        initial_prompt = ZhipuAI_Chat.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=14000
        )

        initial_prompt = ZhipuAI_Chat.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=14000
        )

        message_log = [ZhipuAI_Chat.system_message(initial_prompt)]

        for example in question_sql_list:
            if example is None:
                print("example is None")
            else:
                if example is not None and "question" in example and "sql" in example:
                    message_log.append(ZhipuAI_Chat.user_message(example["question"]))
                    message_log.append(ZhipuAI_Chat.assistant_message(example["sql"]))

        message_log.append({"role": "user", "content": question})

        return message_log

    def get_followup_questions_prompt(
        self,
        question: str,
        df: pd.DataFrame,
        question_sql_list: List,
        ddl_list: List,
        doc_list: List,
        **kwargs,
    ):
        initial_prompt = f"The user initially asked the question: '{question}': \n\n"

        initial_prompt = ZhipuAI_Chat.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=14000
        )

        initial_prompt = ZhipuAI_Chat.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=14000
        )

        initial_prompt = ZhipuAI_Chat.add_sql_to_prompt(
            initial_prompt, question_sql_list, max_tokens=14000
        )

        message_log = [ZhipuAI_Chat.system_message(initial_prompt)]
        message_log.append(
            ZhipuAI_Chat.user_message(
                "Generate a List of followup questions that the user might ask about this data. Respond with a List of questions, one per line. Do not answer with any explanations -- just the questions."
            )
        )

        return message_log

    def generate_question(self, sql: str, **kwargs) -> str:
        response = self.submit_prompt(
            [
                self.system_message(
                    "The user will give you SQL and you will try to guess what the business question this query is answering. Return just the question without any additional explanation. Do not reference the table name in the question."
                ),
                self.user_message(sql),
            ],
            **kwargs,
        )

        return response

    def _extract_python_code(self, markdown_string: str) -> str:
        # Regex pattern to match Python code blocks
        pattern = r"```[\w\s]*python\n([\s\S]*?)```|```([\s\S]*?)```"

        # Find all matches in the markdown string
        matches = re.findall(pattern, markdown_string, re.IGNORECASE)

        # Extract the Python code from the matches
        python_code = []
        for match in matches:
            python = match[0] if match[0] else match[1]
            python_code.append(python.strip())

        if len(python_code) == 0:
            return markdown_string

        return python_code[0]

    def _sanitize_plotly_code(self, raw_plotly_code: str) -> str:
        # Remove the fig.show() statement from the plotly code
        plotly_code = raw_plotly_code.replace("fig.show()", "")

        return plotly_code

    def generate_plotly_code(
        self, question: str = None, sql: str = None, df_metadata: str = None, **kwargs
    ) -> str:
        if question is not None:
            system_msg = f"The following is a pandas DataFrame that contains the results of the query that answers the question the user asked: '{question}'"
        else:
            system_msg = "The following is a pandas DataFrame "

        if sql is not None:
            system_msg += f"\n\nThe DataFrame was produced using this query: {sql}\n\n"

        system_msg += f"The following is information about the resulting pandas DataFrame 'df': \n{df_metadata}"

        message_log = [
            self.system_message(system_msg),
            self.user_message(
                "Can you generate the Python plotly code to chart the results of the dataframe? Assume the data is in a pandas dataframe called 'df'. If there is only one value in the dataframe, use an Indicator. Respond with only Python code. Do not answer with any explanations -- just the code."
            ),
        ]

        plotly_code = self.submit_prompt(message_log, kwargs=kwargs)

        return self._sanitize_plotly_code(self._extract_python_code(plotly_code))

    def submit_prompt(
        self, prompt, max_tokens=500, temperature=0.7, top_p=0.7, stop=None, **kwargs
    ):
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        client = ZhipuAI(api_key=self.api_key)
        response = client.chat.completions.create(
            model="glm-4",
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=stop,
            messages=prompt,
        )

        return response.choices[0].message.content


================================================
FILE: src/vanna/legacy/ZhipuAI/ZhipuAI_embeddings.py
================================================
from typing import List
from zhipuai import ZhipuAI
from chromadb import Documents, EmbeddingFunction, Embeddings
from ..base import VannaBase


class ZhipuAI_Embeddings(VannaBase):
    """
    [future functionality] This function is used to generate embeddings from ZhipuAI.

    Args:
        VannaBase (_type_): _description_
    """

    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        if "api_key" not in config:
            raise Exception("Missing api_key in config")
        self.api_key = config["api_key"]
        self.client = ZhipuAI(api_key=self.api_key)

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embedding = self.client.embeddings.create(
            model="embedding-2",
            input=data,
        )

        return embedding.data[0].embedding


class ZhipuAIEmbeddingFunction(EmbeddingFunction[Documents]):
    """
    A embeddingFunction that uses ZhipuAI to generate embeddings which can use in chromadb.
    usage:
    class MyVanna(ChromaDB_VectorStore, ZhipuAI_Chat):
        def __init__(self, config=None):
            ChromaDB_VectorStore.__init__(self, config=config)
            ZhipuAI_Chat.__init__(self, config=config)

    config={'api_key': 'xxx'}
    zhipu_embedding_function = ZhipuAIEmbeddingFunction(config=config)
    config = {"api_key": "xxx", "model": "glm-4","path":"xy","embedding_function":zhipu_embedding_function}

    vn = MyVanna(config)

    """

    def __init__(self, config=None):
        if config is None or "api_key" not in config:
            raise ValueError("Missing 'api_key' in config")

        self.api_key = config["api_key"]
        self.model_name = config.get("model_name", "embedding-2")

        try:
            self.client = ZhipuAI(api_key=self.api_key)
        except Exception as e:
            raise ValueError(f"Error initializing ZhipuAI client: {e}")

    def __call__(self, input: Documents) -> Embeddings:
        # Replace newlines, which can negatively affect performance.
        input = [t.replace("\n", " ") for t in input]
        all_embeddings = []
        print(f"Generating embeddings for {len(input)} documents")

        # Iterating over each document for individual API calls
        for document in input:
            try:
                response = self.client.embeddings.create(
                    model=self.model_name, input=document
                )
                # print(response)
                embedding = response.data[0].embedding
                all_embeddings.append(embedding)
                # print(f"Cost required: {response.usage.total_tokens}")
            except Exception as e:
                raise ValueError(f"Error generating embedding for document: {e}")

        return all_embeddings


================================================
FILE: src/vanna/legacy/ZhipuAI/__init__.py
================================================
from .ZhipuAI_Chat import ZhipuAI_Chat
from .ZhipuAI_embeddings import ZhipuAI_Embeddings, ZhipuAIEmbeddingFunction


================================================
FILE: src/vanna/legacy/__init__.py
================================================
import dataclasses
import json
import os
from dataclasses import dataclass
from typing import Callable, List, Tuple, Union

import pandas as pd
import requests
import plotly.graph_objs

from .exceptions import (
    OTPCodeError,
    ValidationError,
)
from .types import (
    ApiKey,
    Status,
    TrainingData,
    UserEmail,
    UserOTP,
)
from .utils import sanitize_model_name, validate_config_path

api_key: Union[str, None] = None  # API key for Vanna.AI

fig_as_img: bool = False  # Whether or not to return Plotly figures as images

run_sql: Union[Callable[[str], pd.DataFrame], None] = (
    None  # Function to convert SQL to a Pandas DataFrame
)
"""
**Example**
```python
vn.run_sql = lambda sql: pd.read_sql(sql, engine)
```

Set the SQL to DataFrame function for Vanna.AI. This is used in the [`vn.ask(...)`][vanna.ask] function.
Instead of setting this directly you can also use [`vn.connect_to_snowflake(...)`][vanna.connect_to_snowflake] to set this.

"""

__org: Union[str, None] = None  # Organization name for Vanna.AI

_unauthenticated_endpoint = "https://ask.vanna.ai/unauthenticated_rpc"


def error_deprecation():
    raise Exception("""
Please switch to the following method for initializing Vanna:

from vanna.remote import VannaDefault

api_key = # Your API key from https://vanna.ai/account/profile 
vanna_model_name = # Your model name from https://vanna.ai/account/profile
                    
vn = VannaDefault(model=vanna_model_name, api_key=api_key)
""")


def __unauthenticated_rpc_call(method, params):
    headers = {
        "Content-Type": "application/json",
    }
    data = {"method": method, "params": [__dataclass_to_dict(obj) for obj in params]}

    response = requests.post(
        _unauthenticated_endpoint, headers=headers, data=json.dumps(data)
    )
    return response.json()


def __dataclass_to_dict(obj):
    return dataclasses.asdict(obj)


def get_api_key(email: str, otp_code: Union[str, None] = None) -> str:
    """
    **Example:**
    ```python
    vn.get_api_key(email="my-email@example.com")
    ```

    Login to the Vanna.AI API.

    Args:
        email (str): The email address to login with.
        otp_code (Union[str, None]): The OTP code to login with. If None, an OTP code will be sent to the email address.

    Returns:
        str: The API key.
    """
    vanna_api_key = os.environ.get("VANNA_API_KEY", None)

    if vanna_api_key is not None:
        return vanna_api_key

    if email == "my-email@example.com":
        raise ValidationError(
            "Please replace 'my-email@example.com' with your email address."
        )

    if otp_code is None:
        params = [UserEmail(email=email)]

        d = __unauthenticated_rpc_call(method="send_otp", params=params)

        if "result" not in d:
            raise OTPCodeError("Error sending OTP code.")

        status = Status(**d["result"])

        if not status.success:
            raise OTPCodeError(f"Error sending OTP code: {status.message}")

        otp_code = input("Check your email for the code and enter it here: ")

    params = [UserOTP(email=email, otp=otp_code)]

    d = __unauthenticated_rpc_call(method="verify_otp", params=params)

    if "result" not in d:
        raise OTPCodeError("Error verifying OTP code.")

    key = ApiKey(**d["result"])

    if key is None:
        raise OTPCodeError("Error verifying OTP code.")

    api_key = key.key

    return api_key


def set_api_key(key: str) -> None:
    error_deprecation()


def get_models() -> List[str]:
    error_deprecation()


def create_model(model: str, db_type: str) -> bool:
    error_deprecation()


def add_user_to_model(model: str, email: str, is_admin: bool) -> bool:
    error_deprecation()


def update_model_visibility(public: bool) -> bool:
    error_deprecation()


def set_model(model: str):
    error_deprecation()


def add_sql(
    question: str, sql: str, tag: Union[str, None] = "Manually Trained"
) -> bool:
    error_deprecation()


def add_ddl(ddl: str) -> bool:
    error_deprecation()


def add_documentation(documentation: str) -> bool:
    error_deprecation()


@dataclass
class TrainingPlanItem:
    item_type: str
    item_group: str
    item_name: str
    item_value: str

    def __str__(self):
        if self.item_type == self.ITEM_TYPE_SQL:
            return f"Train on SQL: {self.item_group} {self.item_name}"
        elif self.item_type == self.ITEM_TYPE_DDL:
            return f"Train on DDL: {self.item_group} {self.item_name}"
        elif self.item_type == self.ITEM_TYPE_IS:
            return f"Train on Information Schema: {self.item_group} {self.item_name}"

    ITEM_TYPE_SQL = "sql"
    ITEM_TYPE_DDL = "ddl"
    ITEM_TYPE_IS = "is"


class TrainingPlan:
    """
    A class representing a training plan. You can see what's in it, and remove items from it that you don't want trained.

    **Example:**
    ```python
    plan = vn.get_training_plan()

    plan.get_summary()
    ```

    """

    _plan: List[TrainingPlanItem]

    def __init__(self, plan: List[TrainingPlanItem]):
        self._plan = plan

    def __str__(self):
        return "\n".join(self.get_summary())

    def __repr__(self):
        return self.__str__()

    def get_summary(self) -> List[str]:
        """
        **Example:**
        ```python
        plan = vn.get_training_plan()

        plan.get_summary()
        ```

        Get a summary of the training plan.

        Returns:
            List[str]: A list of strings describing the training plan.
        """

        return [f"{item}" for item in self._plan]

    def remove_item(self, item: str):
        """
        **Example:**
        ```python
        plan = vn.get_training_plan()

        plan.remove_item("Train on SQL: What is the average salary of employees?")
        ```

        Remove an item from the training plan.

        Args:
            item (str): The item to remove.
        """
        for plan_item in self._plan:
            if str(plan_item) == item:
                self._plan.remove(plan_item)
                break


def get_training_plan_postgres(
    filter_databases: Union[List[str], None] = None,
    filter_schemas: Union[List[str], None] = None,
    include_information_schema: bool = False,
    use_historical_queries: bool = True,
) -> TrainingPlan:
    error_deprecation()


def get_training_plan_generic(df) -> TrainingPlan:
    error_deprecation()


def get_training_plan_experimental(
    filter_databases: Union[List[str], None] = None,
    filter_schemas: Union[List[str], None] = None,
    include_information_schema: bool = False,
    use_historical_queries: bool = True,
) -> TrainingPlan:
    error_deprecation()


def train(
    question: str = None,
    sql: str = None,
    ddl: str = None,
    documentation: str = None,
    json_file: str = None,
    sql_file: str = None,
    plan: TrainingPlan = None,
) -> bool:
    error_deprecation()


def flag_sql_for_review(
    question: str, sql: Union[str, None] = None, error_msg: Union[str, None] = None
) -> bool:
    error_deprecation()


def remove_sql(question: str) -> bool:
    error_deprecation()


def remove_training_data(id: str) -> bool:
    error_deprecation()


def generate_sql(question: str) -> str:
    error_deprecation()


def get_related_training_data(question: str) -> TrainingData:
    error_deprecation()


def generate_meta(question: str) -> str:
    error_deprecation()


def generate_followup_questions(question: str, df: pd.DataFrame) -> List[str]:
    error_deprecation()


def generate_questions() -> List[str]:
    error_deprecation()


def ask(
    question: Union[str, None] = None,
    print_results: bool = True,
    auto_train: bool = True,
    generate_followups: bool = True,
) -> Union[
    Tuple[
        Union[str, None],
        Union[pd.DataFrame, None],
        Union[plotly.graph_objs.Figure, None],
        Union[List[str], None],
    ],
    None,
]:
    error_deprecation()


def generate_plotly_code(
    question: Union[str, None],
    sql: Union[str, None],
    df: pd.DataFrame,
    chart_instructions: Union[str, None] = None,
) -> str:
    error_deprecation()


def get_plotly_figure(
    plotly_code: str, df: pd.DataFrame, dark_mode: bool = True
) -> plotly.graph_objs.Figure:
    error_deprecation()


def get_results(cs, default_database: str, sql: str) -> pd.DataFrame:
    error_deprecation()


def generate_explanation(sql: str) -> str:
    error_deprecation()


def generate_question(sql: str) -> str:
    error_deprecation()


def get_all_questions() -> pd.DataFrame:
    error_deprecation()


def get_training_data() -> pd.DataFrame:
    error_deprecation()


def connect_to_sqlite(url: str):
    error_deprecation()


def connect_to_snowflake(
    account: str,
    username: str,
    password: str,
    database: str,
    schema: Union[str, None] = None,
    role: Union[str, None] = None,
):
    error_deprecation()


def connect_to_postgres(
    host: str = None,
    dbname: str = None,
    user: str = None,
    password: str = None,
    port: int = None,
):
    error_deprecation()


def connect_to_bigquery(cred_file_path: str = None, project_id: str = None):
    error_deprecation()


def connect_to_duckdb(url: str = "memory", init_sql: str = None):
    error_deprecation()


================================================
FILE: src/vanna/legacy/adapter.py
================================================
"""
Legacy VannaBase adapter for the Vanna Agents framework.

This module provides a LegacyVannaAdapter that bridges legacy VannaBase objects
with the new ToolRegistry system by auto-registering legacy methods as tools
with appropriate group-based access control.
"""

from typing import Any, Dict, List, Optional

import pandas as pd

from ..capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
    ToolMemory,
    ToolMemorySearchResult,
)
from ..capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from ..core.registry import ToolRegistry
from ..core.tool import Tool, ToolContext, ToolResult
from ..core.user import User
from ..tools.agent_memory import (
    SaveQuestionToolArgsTool,
    SearchSavedCorrectToolUsesTool,
)
from ..tools.run_sql import RunSqlTool
from .base.base import VannaBase


class LegacySqlRunner(SqlRunner):
    """SqlRunner implementation that wraps a legacy VannaBase instance.

    This class bridges the new SqlRunner interface with legacy VannaBase
    run_sql methods, allowing legacy database connections to work with
    the new tool-based architecture.
    """

    def __init__(self, vn: VannaBase):
        """Initialize with a legacy VannaBase instance.

        Args:
            vn: The legacy VannaBase instance with an initialized run_sql method
        """
        self.vn = vn

    async def run_sql(self, args: RunSqlToolArgs, context: ToolContext) -> pd.DataFrame:
        """Execute SQL query using the legacy VannaBase run_sql method.

        Args:
            args: SQL query arguments containing the SQL string
            context: Tool execution context (not used by legacy implementation)

        Returns:
            DataFrame with query results

        Raises:
            Exception: If query execution fails
        """
        # Call the legacy VannaBase run_sql method
        # The legacy method is synchronous, so we call it directly
        return self.vn.run_sql(args.sql)


class LegacyVannaAdapter(ToolRegistry, AgentMemory):
    """Adapter that wraps a legacy VannaBase object and exposes its methods as tools.

    This adapter automatically registers specific VannaBase methods as tools in the
    registry with configurable group-based access control. This allows legacy Vanna
    instances to work seamlessly with the new Agents framework.

    Features:
    - Auto-registers legacy methods as tools
    - Configurable group-based permissions ('user', 'admin', etc.)
    - Seamless integration with ToolRegistry
    - Implements AgentMemory interface
    - Preserves legacy VannaBase functionality

    Example:
        ```python
        from vanna.legacy.base import VannaBase
        from vanna.legacy.adapter import LegacyVannaAdapter

        # Initialize your legacy Vanna instance
        vn = VannaBase(config={"model": "gpt-4"})
        vn.connect_to_postgres(...)

        # Create adapter and auto-register tools
        adapter = LegacyVannaAdapter(vn)

        # Tools are now available through the registry
        schemas = await adapter.get_schemas(user)
        ```
    """

    def __init__(
        self,
        vn: VannaBase,
        audit_logger: Optional[Any] = None,
        audit_config: Optional[Any] = None,
    ) -> None:
        """Initialize the adapter with a legacy VannaBase instance.

        Args:
            vanna: The legacy VannaBase instance to wrap
            audit_logger: Optional audit logger for tool execution tracking
            audit_config: Optional audit configuration
        """
        ToolRegistry.__init__(
            self, audit_logger=audit_logger, audit_config=audit_config
        )
        self.vn = vn
        self._register_tools()

    def _register_tools(self) -> None:
        """Register legacy VannaBase methods as tools with appropriate permissions.

        Registers the following tools:
        - RunSqlTool: Wraps the legacy run_sql method via LegacySqlRunner
        - SaveQuestionToolArgsTool: Wraps add_question_sql via LegacyAgentMemory
        - SearchSavedCorrectToolUsesTool: Wraps get_similar_question_sql via LegacyAgentMemory
        """
        # Create a LegacySqlRunner to wrap the VannaBase run_sql method
        sql_runner = LegacySqlRunner(self.vn)

        # Register the RunSqlTool with user and admin access
        run_sql_tool = RunSqlTool(sql_runner)
        self.register_local_tool(run_sql_tool, access_groups=["user", "admin"])

        # Register memory tools using the internal _agent_memory instance
        # SaveQuestionToolArgsTool - for saving question-tool-args patterns (admin only)
        save_memory_tool = SaveQuestionToolArgsTool()
        self.register_local_tool(save_memory_tool, access_groups=["admin"])

        # SearchSavedCorrectToolUsesTool - for searching similar patterns (user and admin)
        search_memory_tool = SearchSavedCorrectToolUsesTool()
        self.register_local_tool(search_memory_tool, access_groups=["user", "admin"])

    # AgentMemory interface implementation

    async def save_tool_usage(
        self,
        question: str,
        tool_name: str,
        args: Dict[str, Any],
        context: ToolContext,
        success: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a tool usage pattern by storing it as a question-sql pair.

        Args:
            question: The user question
            tool_name: Name of the tool that was used
            args: Arguments passed to the tool
            context: Tool execution context (not used by legacy implementation)
            success: Whether the tool execution was successful
            metadata: Additional metadata (not used by legacy implementation)
        """
        # For legacy compatibility, we primarily care about SQL queries
        # Extract SQL from args if this was a run_sql tool
        if tool_name == "run_sql" and "sql" in args:
            sql = args["sql"]
            # Call the legacy add_question_sql method
            # The legacy method is synchronous, so we call it directly
            self.vn.add_question_sql(question=question, sql=sql)

    async def search_similar_usage(
        self,
        question: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
        tool_name_filter: Optional[str] = None,
    ) -> List[ToolMemorySearchResult]:
        """Search for similar tool usage patterns using legacy question-sql lookup.

        Args:
            question: The question to search for
            context: Tool execution context (not used by legacy implementation)
            limit: Maximum number of results (not directly supported by legacy)
            similarity_threshold: Minimum similarity score (not directly supported by legacy)
            tool_name_filter: Filter by tool name (not directly supported by legacy)

        Returns:
            List of memory search results with similar question-sql pairs
        """
        # Call the legacy get_similar_question_sql method
        similar_results = self.vn.get_similar_question_sql(question=question)

        # Convert legacy results to ToolMemorySearchResult format
        memory_results = []
        for idx, result in enumerate(similar_results):
            # Legacy results are typically dicts with 'question' and 'sql' keys
            if isinstance(result, dict) and "question" in result and "sql" in result:
                tool_memory = ToolMemory(
                    memory_id=None,  # Legacy doesn't provide IDs
                    question=result["question"],
                    tool_name="run_sql",
                    args={"sql": result["sql"]},
                    success=True,
                )

                # Assign a simple rank-based similarity score
                # Legacy system doesn't provide actual similarity scores
                similarity_score = 1.0 - (idx * 0.1)  # Decreasing score by rank
                similarity_score = max(similarity_score, 0.0)

                memory_results.append(
                    ToolMemorySearchResult(
                        memory=tool_memory,
                        similarity_score=similarity_score,
                        rank=idx + 1,
                    )
                )

        return memory_results[:limit]

    async def save_text_memory(self, content: str, context: ToolContext) -> TextMemory:
        """Save text memory using legacy add_documentation method.

        Args:
            content: The documentation content to save
            context: Tool execution context (not used by legacy implementation)

        Returns:
            TextMemory object with the saved content
        """
        # Call the legacy add_documentation method
        # The legacy method is synchronous, so we call it directly
        doc_id = self.vn.add_documentation(documentation=content)

        return TextMemory(
            memory_id=doc_id,
            content=content,
            timestamp=None,  # Legacy doesn't provide timestamps
        )

    async def search_text_memories(
        self,
        query: str,
        context: ToolContext,
        *,
        limit: int = 10,
        similarity_threshold: float = 0.7,
    ) -> List[TextMemorySearchResult]:
        """Search text memories using legacy get_related_documentation method.

        Args:
            query: The query to search for
            context: Tool execution context (not used by legacy implementation)
            limit: Maximum number of results (not directly supported by legacy)
            similarity_threshold: Minimum similarity score (not directly supported by legacy)

        Returns:
            List of text memory search results
        """
        # Call the legacy get_related_documentation method
        related_docs = self.vn.get_related_documentation(question=query)

        # Convert legacy results to TextMemorySearchResult format
        memory_results = []
        for idx, doc in enumerate(related_docs):
            # Legacy results are typically strings or dicts
            if isinstance(doc, str):
                content = doc
                doc_id = None
            elif isinstance(doc, dict):
                content = str(doc.get("documentation", doc.get("content", str(doc))))
                doc_id = doc.get("id")
            else:
                content = str(doc)
                doc_id = None

            # Create TextMemory object
            text_memory = TextMemory(
                memory_id=doc_id,
                content=content,
                timestamp=None,  # Legacy doesn't provide timestamps
            )

            # Assign a simple rank-based similarity score
            # Legacy system doesn't provide actual similarity scores
            similarity_score = 1.0 - (idx * 0.1)  # Decreasing score by rank
            similarity_score = max(similarity_score, 0.0)

            if similarity_score >= similarity_threshold:
                memory_results.append(
                    TextMemorySearchResult(
                        memory=text_memory,
                        similarity_score=similarity_score,
                        rank=idx + 1,
                    )
                )

        return memory_results[:limit]

    async def get_recent_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[ToolMemory]:
        """Get recently added memories.

        Note: Legacy VannaBase does not provide a direct way to get recent memories,
        so we retrieve using a blank string which typically returns the most relevant
        or recent items from the vector store.

        Args:
            context: Tool execution context
            limit: Maximum number of memories to return

        Returns:
            List of recently added tool memories
        """
        # Use blank string retrieval to get recent/relevant memories
        similar_results = self.vn.get_similar_question_sql(question="")

        # Convert legacy results to ToolMemory format
        memories = []
        for idx, result in enumerate(similar_results[:limit]):
            # Legacy results are typically dicts with 'question' and 'sql' keys
            if isinstance(result, dict) and "question" in result and "sql" in result:
                tool_memory = ToolMemory(
                    memory_id=None,  # Legacy doesn't provide IDs
                    question=result["question"],
                    tool_name="run_sql",
                    args={"sql": result["sql"]},
                    success=True,
                )
                memories.append(tool_memory)

        return memories

    async def get_recent_text_memories(
        self, context: ToolContext, limit: int = 10
    ) -> List[TextMemory]:
        """Fetch recently stored text memories.

        Note: Legacy VannaBase does not provide a direct way to get recent text memories,
        so we retrieve using a blank string which typically returns the most relevant
        or recent items from the vector store.

        Args:
            context: Tool execution context
            limit: Maximum number of memories to return

        Returns:
            List of recently added text memories
        """
        # Use blank string retrieval to get recent/relevant documentation
        related_docs = self.vn.get_related_documentation(question="")

        # Convert legacy results to TextMemory format
        memories = []
        for doc in related_docs[:limit]:
            # Legacy results are typically strings or dicts
            if isinstance(doc, str):
                content = doc
                doc_id = None
            elif isinstance(doc, dict):
                content = str(doc.get("documentation", doc.get("content", str(doc))))
                doc_id = doc.get("id")
            else:
                content = str(doc)
                doc_id = None

            # Create TextMemory object
            text_memory = TextMemory(
                memory_id=doc_id,
                content=content,
                timestamp=None,  # Legacy doesn't provide timestamps
            )
            memories.append(text_memory)

        return memories

    async def delete_by_id(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a memory by its ID using legacy remove_training_data method.

        Args:
            context: Tool execution context
            memory_id: ID of the memory to delete

        Returns:
            True if the memory was deleted, False otherwise
        """
        # Call the legacy remove_training_data method
        # The legacy method is synchronous, so we call it directly
        return self.vn.remove_training_data(id=memory_id)

    async def delete_text_memory(self, context: ToolContext, memory_id: str) -> bool:
        """Delete a text memory by its ID using legacy remove_training_data method.

        Args:
            context: Tool execution context
            memory_id: ID of the text memory to delete

        Returns:
            True if the text memory was deleted, False otherwise
        """
        # Call the legacy remove_training_data method
        # The legacy method is synchronous, so we call it directly
        return self.vn.remove_training_data(id=memory_id)

    async def clear_memories(
        self,
        context: ToolContext,
        tool_name: Optional[str] = None,
        before_date: Optional[str] = None,
    ) -> int:
        """Clear stored memories.

        Note: Legacy VannaBase does not provide a direct clear method,
        so this operation is not supported.

        Args:
            context: Tool execution context
            tool_name: Optional tool name filter
            before_date: Optional date filter

        Returns:
            0 (operation not supported by legacy)
        """
        return 0

    # Example stub for a tool wrapper (to be expanded)
    # You can copy and customize this pattern for each tool you want to expose
    """
    class ExampleTool(Tool[ExampleToolArgs]):
        def __init__(self, vanna: VannaBase):
            self.vanna = vanna

        @property
        def name(self) -> str:
            return "example_tool"

        @property
        def description(self) -> str:
            return "Example tool description"

        @property
        def access_groups(self) -> List[str]:
            # This is optional - will be overridden by register_local_tool
            return []

        def get_args_schema(self) -> type[ExampleToolArgs]:
            return ExampleToolArgs

        async def execute(
            self,
            context: ToolContext,
            args: ExampleToolArgs
        ) -> ToolResult:
            # Call the legacy VannaBase method
            result = self.vanna.example_method(args.param1, args.param2)

            return ToolResult(
                success=True,
                result_for_llm=result,
                ui_component=None,
            )
    """


================================================
FILE: src/vanna/legacy/advanced/__init__.py
================================================
from abc import ABC, abstractmethod


class VannaAdvanced(ABC):
    def __init__(self, config=None):
        self.config = config

    @abstractmethod
    def get_function(self, question: str, additional_data: dict = {}) -> dict:
        pass

    @abstractmethod
    def create_function(
        self, question: str, sql: str, plotly_code: str, **kwargs
    ) -> dict:
        pass

    @abstractmethod
    def update_function(self, old_function_name: str, updated_function: dict) -> bool:
        pass

    @abstractmethod
    def delete_function(self, function_name: str) -> bool:
        pass

    @abstractmethod
    def get_all_functions(self) -> list:
        pass


================================================
FILE: src/vanna/legacy/anthropic/__init__.py
================================================
from .anthropic_chat import Anthropic_Chat


================================================
FILE: src/vanna/legacy/anthropic/anthropic_chat.py
================================================
import os

import anthropic

from ..base import VannaBase


class Anthropic_Chat(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        # default parameters - can be overrided using config
        self.temperature = 0.7
        self.max_tokens = 500

        if "temperature" in config:
            self.temperature = config["temperature"]

        if "max_tokens" in config:
            self.max_tokens = config["max_tokens"]

        if client is not None:
            self.client = client
            return

        if config is None and client is None:
            self.client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
            return

        if "api_key" in config:
            self.client = anthropic.Anthropic(api_key=config["api_key"])

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        # Count the number of tokens in the message log
        # Use 4 as an approximation for the number of characters per token
        num_tokens = 0
        for message in prompt:
            num_tokens += len(message["content"]) / 4

        if self.config is not None and "model" in self.config:
            print(
                f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
            )
            # claude required system message is a single filed
            # https://docs.anthropic.com/claude/reference/messages_post
            system_message = ""
            no_system_prompt = []
            for prompt_message in prompt:
                role = prompt_message["role"]
                if role == "system":
                    system_message = prompt_message["content"]
                else:
                    no_system_prompt.append(
                        {"role": role, "content": prompt_message["content"]}
                    )

            response = self.client.messages.create(
                model=self.config["model"],
                messages=no_system_prompt,
                system=system_message,
                max_tokens=self.max_tokens,
                temperature=self.temperature,
            )

        return response.content[0].text


================================================
FILE: src/vanna/legacy/azuresearch/__init__.py
================================================
from .azuresearch_vector import AzureAISearch_VectorStore


================================================
FILE: src/vanna/legacy/azuresearch/azuresearch_vector.py
================================================
import ast
import json
from typing import List

import pandas as pd
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    ExhaustiveKnnAlgorithmConfiguration,
    ExhaustiveKnnParameters,
    SearchableField,
    SearchField,
    SearchFieldDataType,
    SearchIndex,
    VectorSearch,
    VectorSearchAlgorithmKind,
    VectorSearchAlgorithmMetric,
    VectorSearchProfile,
)
from azure.search.documents.models import VectorFilterMode, VectorizedQuery
from fastembed import TextEmbedding

from ..base import VannaBase
from ..utils import deterministic_uuid


class AzureAISearch_VectorStore(VannaBase):
    """
    AzureAISearch_VectorStore is a class that provides a vector store for Azure AI Search.

    Args:
        config (dict): Configuration dictionary. Defaults to {}. You must provide an API key in the config.
            - azure_search_endpoint (str, optional): Azure Search endpoint. Defaults to "https://azcognetive.search.windows.net".
            - azure_search_api_key (str): Azure Search API key.
            - dimensions (int, optional): Dimensions of the embeddings. Defaults to 384 which corresponds to the dimensions of BAAI/bge-small-en-v1.5.
            - fastembed_model (str, optional): Fastembed model to use. Defaults to "BAAI/bge-small-en-v1.5".
            - index_name (str, optional): Name of the index. Defaults to "vanna-index".
            - n_results (int, optional): Number of results to return. Defaults to 10.
            - n_results_ddl (int, optional): Number of results to return for DDL queries. Defaults to the value of n_results.
            - n_results_sql (int, optional): Number of results to return for SQL queries. Defaults to the value of n_results.
            - n_results_documentation (int, optional): Number of results to return for documentation queries. Defaults to the value of n_results.

    Raises:
        ValueError: If config is None, or if 'azure_search_api_key' is not provided in the config.
    """

    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        self.config = config or None

        if config is None:
            raise ValueError(
                "config is required, pass an API key, 'azure_search_api_key', in the config."
            )

        azure_search_endpoint = config.get(
            "azure_search_endpoint", "https://azcognetive.search.windows.net"
        )
        azure_search_api_key = config.get("azure_search_api_key")

        self.dimensions = config.get("dimensions", 384)
        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")

        self.index_name = config.get("index_name", "vanna-index")

        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
        self.n_results_documentation = config.get(
            "n_results_documentation", config.get("n_results", 10)
        )

        if not azure_search_api_key:
            raise ValueError(
                "'azure_search_api_key' is required in config to use AzureAISearch_VectorStore"
            )

        self.index_client = SearchIndexClient(
            endpoint=azure_search_endpoint,
            credential=AzureKeyCredential(azure_search_api_key),
        )

        self.search_client = SearchClient(
            endpoint=azure_search_endpoint,
            index_name=self.index_name,
            credential=AzureKeyCredential(azure_search_api_key),
        )

        if self.index_name not in self._get_indexes():
            self._create_index()

    def _create_index(self) -> bool:
        fields = [
            SearchableField(
                name="id", type=SearchFieldDataType.String, key=True, filterable=True
            ),
            SearchableField(
                name="document",
                type=SearchFieldDataType.String,
                searchable=True,
                filterable=True,
            ),
            SearchField(
                name="type",
                type=SearchFieldDataType.String,
                filterable=True,
                searchable=True,
            ),
            SearchField(
                name="document_vector",
                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                searchable=True,
                vector_search_dimensions=self.dimensions,
                vector_search_profile_name="ExhaustiveKnnProfile",
            ),
        ]

        vector_search = VectorSearch(
            algorithms=[
                ExhaustiveKnnAlgorithmConfiguration(
                    name="ExhaustiveKnn",
                    kind=VectorSearchAlgorithmKind.EXHAUSTIVE_KNN,
                    parameters=ExhaustiveKnnParameters(
                        metric=VectorSearchAlgorithmMetric.COSINE
                    ),
                )
            ],
            profiles=[
                VectorSearchProfile(
                    name="ExhaustiveKnnProfile",
                    algorithm_configuration_name="ExhaustiveKnn",
                )
            ],
        )

        index = SearchIndex(
            name=self.index_name, fields=fields, vector_search=vector_search
        )
        result = self.index_client.create_or_update_index(index)
        print(f"{result.name} created")

    def _get_indexes(self) -> list:
        return [index for index in self.index_client.list_index_names()]

    def add_ddl(self, ddl: str) -> str:
        id = deterministic_uuid(ddl) + "-ddl"
        document = {
            "id": id,
            "document": ddl,
            "type": "ddl",
            "document_vector": self.generate_embedding(ddl),
        }
        self.search_client.upload_documents(documents=[document])
        return id

    def add_documentation(self, doc: str) -> str:
        id = deterministic_uuid(doc) + "-doc"
        document = {
            "id": id,
            "document": doc,
            "type": "doc",
            "document_vector": self.generate_embedding(doc),
        }
        self.search_client.upload_documents(documents=[document])
        return id

    def add_question_sql(self, question: str, sql: str) -> str:
        question_sql_json = json.dumps(
            {"question": question, "sql": sql}, ensure_ascii=False
        )
        id = deterministic_uuid(question_sql_json) + "-sql"
        document = {
            "id": id,
            "document": question_sql_json,
            "type": "sql",
            "document_vector": self.generate_embedding(question_sql_json),
        }
        self.search_client.upload_documents(documents=[document])
        return id

    def get_related_ddl(self, text: str) -> List[str]:
        result = []
        vector_query = VectorizedQuery(
            vector=self.generate_embedding(text), fields="document_vector"
        )
        df = pd.DataFrame(
            self.search_client.search(
                top=self.n_results_ddl,
                vector_queries=[vector_query],
                select=["id", "document", "type"],
                filter=f"type eq 'ddl'",
            )
        )

        if len(df):
            result = df["document"].tolist()
        return result

    def get_related_documentation(self, text: str) -> List[str]:
        result = []
        vector_query = VectorizedQuery(
            vector=self.generate_embedding(text), fields="document_vector"
        )

        df = pd.DataFrame(
            self.search_client.search(
                top=self.n_results_documentation,
                vector_queries=[vector_query],
                select=["id", "document", "type"],
                filter=f"type eq 'doc'",
                vector_filter_mode=VectorFilterMode.PRE_FILTER,
            )
        )

        if len(df):
            result = df["document"].tolist()
        return result

    def get_similar_question_sql(self, question: str) -> List[str]:
        result = []
        # Vectorize the text
        vector_query = VectorizedQuery(
            vector=self.generate_embedding(question), fields="document_vector"
        )
        df = pd.DataFrame(
            self.search_client.search(
                top=self.n_results_sql,
                vector_queries=[vector_query],
                select=["id", "document", "type"],
                filter=f"type eq 'sql'",
            )
        )

        if len(df):  # Check if there is similar query and the result is not empty
            result = [ast.literal_eval(element) for element in df["document"].tolist()]

        return result

    def get_training_data(self) -> List[str]:
        search = self.search_client.search(
            search_text="*",
            select=["id", "document", "type"],
            filter=f"(type eq 'sql') or (type eq 'ddl') or (type eq 'doc')",
        ).by_page()

        df = pd.DataFrame([item for page in search for item in page])

        if len(df):
            df.loc[df["type"] == "sql", "question"] = df.loc[df["type"] == "sql"][
                "document"
            ].apply(lambda x: json.loads(x)["question"])
            df.loc[df["type"] == "sql", "content"] = df.loc[df["type"] == "sql"][
                "document"
            ].apply(lambda x: json.loads(x)["sql"])
            df.loc[df["type"] != "sql", "content"] = df.loc[df["type"] != "sql"][
                "document"
            ]

            return df[["id", "question", "content", "type"]]

        return pd.DataFrame()

    def remove_training_data(self, id: str) -> bool:
        result = self.search_client.delete_documents(documents=[{"id": id}])
        return result[0].succeeded

    def remove_index(self):
        self.index_client.delete_index(self.index_name)

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embedding_model = TextEmbedding(model_name=self.fastembed_model)
        embedding = next(embedding_model.embed(data))
        return embedding.tolist()


================================================
FILE: src/vanna/legacy/base/__init__.py
================================================
from .base import VannaBase


================================================
FILE: src/vanna/legacy/base/base.py
================================================
r"""

# Nomenclature

| Prefix | Definition | Examples |
| --- | --- | --- |
| `vn.get_` | Fetch some data | [`vn.get_related_ddl(...)`][vanna.base.base.VannaBase.get_related_ddl] |
| `vn.add_` | Adds something to the retrieval layer | [`vn.add_question_sql(...)`][vanna.base.base.VannaBase.add_question_sql] <br> [`vn.add_ddl(...)`][vanna.base.base.VannaBase.add_ddl] |
| `vn.generate_` | Generates something using AI based on the information in the model | [`vn.generate_sql(...)`][vanna.base.base.VannaBase.generate_sql] <br> [`vn.generate_explanation()`][vanna.base.base.VannaBase.generate_explanation] |
| `vn.run_` | Runs code (SQL) | [`vn.run_sql`][vanna.base.base.VannaBase.run_sql] |
| `vn.remove_` | Removes something from the retrieval layer | [`vn.remove_training_data`][vanna.base.base.VannaBase.remove_training_data] |
| `vn.connect_` | Connects to a database | [`vn.connect_to_snowflake(...)`][vanna.base.base.VannaBase.connect_to_snowflake] |
| `vn.update_` | Updates something | N/A -- unused |
| `vn.set_` | Sets something | N/A -- unused  |

# Open-Source and Extending

Vanna.AI is open-source and extensible. If you'd like to use Vanna without the servers, see an example [here](https://vanna.ai/docs/postgres-ollama-chromadb/).

The following is an example of where various functions are implemented in the codebase when using the default "local" version of Vanna. `vanna.base.VannaBase` is the base class which provides a `vanna.base.VannaBase.ask` and `vanna.base.VannaBase.train` function. Those rely on abstract methods which are implemented in the subclasses `vanna.openai_chat.OpenAI_Chat` and `vanna.chromadb_vector.ChromaDB_VectorStore`. `vanna.openai_chat.OpenAI_Chat` uses the OpenAI API to generate SQL and Plotly code. `vanna.chromadb_vector.ChromaDB_VectorStore` uses ChromaDB to store training data and generate embeddings.

If you want to use Vanna with other LLMs or databases, you can create your own subclass of `vanna.base.VannaBase` and implement the abstract methods.

```mermaid
flowchart
    subgraph VannaBase
        ask
        train
    end

    subgraph OpenAI_Chat
        get_sql_prompt
        submit_prompt
        generate_question
        generate_plotly_code
    end

    subgraph ChromaDB_VectorStore
        generate_embedding
        add_question_sql
        add_ddl
        add_documentation
        get_similar_question_sql
        get_related_ddl
        get_related_documentation
    end
```

"""

import json
import os
import re
import sqlite3
import traceback
from abc import ABC, abstractmethod
from typing import List, Tuple, Union
from urllib.parse import urlparse

import pandas as pd
import plotly
import plotly.express as px
import plotly.graph_objects as go
import requests
import sqlparse

from ..exceptions import DependencyError, ImproperlyConfigured, ValidationError
from ..types import TrainingPlan, TrainingPlanItem
from ..utils import validate_config_path


class VannaBase(ABC):
    def __init__(self, config=None):
        if config is None:
            config = {}

        self.config = config
        self.run_sql_is_set = False
        self.static_documentation = ""
        self.dialect = self.config.get("dialect", "SQL")
        self.language = self.config.get("language", None)
        self.max_tokens = self.config.get("max_tokens", 14000)

    def log(self, message: str, title: str = "Info"):
        print(f"{title}: {message}")

    def _response_language(self) -> str:
        if self.language is None:
            return ""

        return f"Respond in the {self.language} language."

    def generate_sql(self, question: str, allow_llm_to_see_data=False, **kwargs) -> str:
        """
        Example:
        ```python
        vn.generate_sql("What are the top 10 customers by sales?")
        ```

        Uses the LLM to generate a SQL query that answers a question. It runs the following methods:

        - [`get_similar_question_sql`][vanna.base.base.VannaBase.get_similar_question_sql]

        - [`get_related_ddl`][vanna.base.base.VannaBase.get_related_ddl]

        - [`get_related_documentation`][vanna.base.base.VannaBase.get_related_documentation]

        - [`get_sql_prompt`][vanna.base.base.VannaBase.get_sql_prompt]

        - [`submit_prompt`][vanna.base.base.VannaBase.submit_prompt]


        Args:
            question (str): The question to generate a SQL query for.
            allow_llm_to_see_data (bool): Whether to allow the LLM to see the data (for the purposes of introspecting the data to generate the final SQL).

        Returns:
            str: The SQL query that answers the question.
        """
        if self.config is not None:
            initial_prompt = self.config.get("initial_prompt", None)
        else:
            initial_prompt = None
        question_sql_list = self.get_similar_question_sql(question, **kwargs)
        ddl_list = self.get_related_ddl(question, **kwargs)
        doc_list = self.get_related_documentation(question, **kwargs)
        prompt = self.get_sql_prompt(
            initial_prompt=initial_prompt,
            question=question,
            question_sql_list=question_sql_list,
            ddl_list=ddl_list,
            doc_list=doc_list,
            **kwargs,
        )
        self.log(title="SQL Prompt", message=prompt)
        llm_response = self.submit_prompt(prompt, **kwargs)
        self.log(title="LLM Response", message=llm_response)

        if "intermediate_sql" in llm_response:
            if not allow_llm_to_see_data:
                return "The LLM is not allowed to see the data in your database. Your question requires database introspection to generate the necessary SQL. Please set allow_llm_to_see_data=True to enable this."

            if allow_llm_to_see_data:
                intermediate_sql = self.extract_sql(llm_response)

                try:
                    self.log(title="Running Intermediate SQL", message=intermediate_sql)
                    df = self.run_sql(intermediate_sql)

                    prompt = self.get_sql_prompt(
                        initial_prompt=initial_prompt,
                        question=question,
                        question_sql_list=question_sql_list,
                        ddl_list=ddl_list,
                        doc_list=doc_list
                        + [
                            f"The following is a pandas DataFrame with the results of the intermediate SQL query {intermediate_sql}: \n"
                            + df.to_markdown()
                        ],
                        **kwargs,
                    )
                    self.log(title="Final SQL Prompt", message=prompt)
                    llm_response = self.submit_prompt(prompt, **kwargs)
                    self.log(title="LLM Response", message=llm_response)
                except Exception as e:
                    return f"Error running intermediate SQL: {e}"

        return self.extract_sql(llm_response)

    def extract_sql(self, llm_response: str) -> str:
        """
        Example:
        ```python
        vn.extract_sql("Here's the SQL query in a code block: ```sql\nSELECT * FROM customers\n```")
        ```

        Extracts the SQL query from the LLM response. This is useful in case the LLM response contains other information besides the SQL query.
        Override this function if your LLM responses need custom extraction logic.

        Args:
            llm_response (str): The LLM response.

        Returns:
            str: The extracted SQL query.
        """

        import re

        """
        Extracts the SQL query from the LLM response, handling various formats including:
        - WITH clause
        - SELECT statement
        - CREATE TABLE AS SELECT
        - Markdown code blocks
        """

        # Match CREATE TABLE ... AS SELECT
        sqls = re.findall(
            r"\bCREATE\s+TABLE\b.*?\bAS\b.*?;", llm_response, re.DOTALL | re.IGNORECASE
        )
        if sqls:
            sql = sqls[-1]
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match WITH clause (CTEs)
        sqls = re.findall(r"\bWITH\b .*?;", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1]
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match SELECT ... ;
        sqls = re.findall(r"\bSELECT\b .*?;", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1]
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match ```sql ... ``` blocks
        sqls = re.findall(
            r"```sql\s*\n(.*?)```", llm_response, re.DOTALL | re.IGNORECASE
        )
        if sqls:
            sql = sqls[-1].strip()
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        # Match any ``` ... ``` code blocks
        sqls = re.findall(r"```(.*?)```", llm_response, re.DOTALL | re.IGNORECASE)
        if sqls:
            sql = sqls[-1].strip()
            self.log(title="Extracted SQL", message=f"{sql}")
            return sql

        return llm_response

    def is_sql_valid(self, sql: str) -> bool:
        """
        Example:
        ```python
        vn.is_sql_valid("SELECT * FROM customers")
        ```
        Checks if the SQL query is valid. This is usually used to check if we should run the SQL query or not.
        By default it checks if the SQL query is a SELECT statement. You can override this method to enable running other types of SQL queries.

        Args:
            sql (str): The SQL query to check.

        Returns:
            bool: True if the SQL query is valid, False otherwise.
        """

        parsed = sqlparse.parse(sql)

        for statement in parsed:
            if statement.get_type() == "SELECT":
                return True

        return False

    def should_generate_chart(self, df: pd.DataFrame) -> bool:
        """
        Example:
        ```python
        vn.should_generate_chart(df)
        ```

        Checks if a chart should be generated for the given DataFrame. By default, it checks if the DataFrame has more than one row and has numerical columns.
        You can override this method to customize the logic for generating charts.

        Args:
            df (pd.DataFrame): The DataFrame to check.

        Returns:
            bool: True if a chart should be generated, False otherwise.
        """

        if len(df) > 1 and df.select_dtypes(include=["number"]).shape[1] > 0:
            return True

        return False

    def generate_rewritten_question(
        self, last_question: str, new_question: str, **kwargs
    ) -> str:
        """
        **Example:**
        ```python
        rewritten_question = vn.generate_rewritten_question("Who are the top 5 customers by sales?", "Show me their email addresses")
        ```

        Generate a rewritten question by combining the last question and the new question if they are related. If the new question is self-contained and not related to the last question, return the new question.

        Args:
            last_question (str): The previous question that was asked.
            new_question (str): The new question to be combined with the last question.
            **kwargs: Additional keyword arguments.

        Returns:
            str: The combined question if related, otherwise the new question.
        """
        if last_question is None:
            return new_question

        prompt = [
            self.system_message(
                "Your goal is to combine a sequence of questions into a singular question if they are related. If the second question does not relate to the first question and is fully self-contained, return the second question. Return just the new combined question with no additional explanations. The question should theoretically be answerable with a single SQL statement."
            ),
            self.user_message(
                "First question: "
                + last_question
                + "\nSecond question: "
                + new_question
            ),
        ]

        return self.submit_prompt(prompt=prompt, **kwargs)

    def generate_followup_questions(
        self, question: str, sql: str, df: pd.DataFrame, n_questions: int = 5, **kwargs
    ) -> list:
        """
        **Example:**
        ```python
        vn.generate_followup_questions("What are the top 10 customers by sales?", sql, df)
        ```

        Generate a list of followup questions that you can ask Vanna.AI.

        Args:
            question (str): The question that was asked.
            sql (str): The LLM-generated SQL query.
            df (pd.DataFrame): The results of the SQL query.
            n_questions (int): Number of follow-up questions to generate.

        Returns:
            list: A list of followup questions that you can ask Vanna.AI.
        """

        message_log = [
            self.system_message(
                f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe SQL query for this question was: {sql}\n\nThe following is a pandas DataFrame with the results of the query: \n{df.head(25).to_markdown()}\n\n"
            ),
            self.user_message(
                f"Generate a list of {n_questions} followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query."
                + self._response_language()
            ),
        ]

        llm_response = self.submit_prompt(message_log, **kwargs)

        numbers_removed = re.sub(r"^\d+\.\s*", "", llm_response, flags=re.MULTILINE)
        return numbers_removed.split("\n")

    def generate_questions(self, **kwargs) -> List[str]:
        """
        **Example:**
        ```python
        vn.generate_questions()
        ```

        Generate a list of questions that you can ask Vanna.AI.
        """
        question_sql = self.get_similar_question_sql(question="", **kwargs)

        return [q["question"] for q in question_sql]

    def generate_summary(self, question: str, df: pd.DataFrame, **kwargs) -> str:
        """
        **Example:**
        ```python
        vn.generate_summary("What are the top 10 customers by sales?", df)
        ```

        Generate a summary of the results of a SQL query.

        Args:
            question (str): The question that was asked.
            df (pd.DataFrame): The results of the SQL query.

        Returns:
            str: The summary of the results of the SQL query.
        """

        message_log = [
            self.system_message(
                f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe following is a pandas DataFrame with the results of the query: \n{df.to_markdown()}\n\n"
            ),
            self.user_message(
                "Briefly summarize the data based on the question that was asked. Do not respond with any additional explanation beyond the summary."
                + self._response_language()
            ),
        ]

        summary = self.submit_prompt(message_log, **kwargs)

        return summary

    # ----------------- Use Any Embeddings API ----------------- #
    @abstractmethod
    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        pass

    # ----------------- Use Any Database to Store and Retrieve Context ----------------- #
    @abstractmethod
    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        """
        This method is used to get similar questions and their corresponding SQL statements.

        Args:
            question (str): The question to get similar questions and their corresponding SQL statements for.

        Returns:
            list: A list of similar questions and their corresponding SQL statements.
        """
        pass

    @abstractmethod
    def get_related_ddl(self, question: str, **kwargs) -> list:
        """
        This method is used to get related DDL statements to a question.

        Args:
            question (str): The question to get related DDL statements for.

        Returns:
            list: A list of related DDL statements.
        """
        pass

    @abstractmethod
    def get_related_documentation(self, question: str, **kwargs) -> list:
        """
        This method is used to get related documentation to a question.

        Args:
            question (str): The question to get related documentation for.

        Returns:
            list: A list of related documentation.
        """
        pass

    @abstractmethod
    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        """
        This method is used to add a question and its corresponding SQL query to the training data.

        Args:
            question (str): The question to add.
            sql (str): The SQL query to add.

        Returns:
            str: The ID of the training data that was added.
        """
        pass

    @abstractmethod
    def add_ddl(self, ddl: str, **kwargs) -> str:
        """
        This method is used to add a DDL statement to the training data.

        Args:
            ddl (str): The DDL statement to add.

        Returns:
            str: The ID of the training data that was added.
        """
        pass

    @abstractmethod
    def add_documentation(self, documentation: str, **kwargs) -> str:
        """
        This method is used to add documentation to the training data.

        Args:
            documentation (str): The documentation to add.

        Returns:
            str: The ID of the training data that was added.
        """
        pass

    @abstractmethod
    def get_training_data(self, **kwargs) -> pd.DataFrame:
        """
        Example:
        ```python
        vn.get_training_data()
        ```

        This method is used to get all the training data from the retrieval layer.

        Returns:
            pd.DataFrame: The training data.
        """
        pass

    @abstractmethod
    def remove_training_data(self, id: str, **kwargs) -> bool:
        """
        Example:
        ```python
        vn.remove_training_data(id="123-ddl")
        ```

        This method is used to remove training data from the retrieval layer.

        Args:
            id (str): The ID of the training data to remove.

        Returns:
            bool: True if the training data was removed, False otherwise.
        """
        pass

    # ----------------- Use Any Language Model API ----------------- #

    @abstractmethod
    def system_message(self, message: str) -> any:
        pass

    @abstractmethod
    def user_message(self, message: str) -> any:
        pass

    @abstractmethod
    def assistant_message(self, message: str) -> any:
        pass

    def str_to_approx_token_count(self, string: str) -> int:
        return len(string) / 4

    def add_ddl_to_prompt(
        self, initial_prompt: str, ddl_list: list[str], max_tokens: int = 14000
    ) -> str:
        if len(ddl_list) > 0:
            initial_prompt += "\n===Tables \n"

            for ddl in ddl_list:
                if (
                    self.str_to_approx_token_count(initial_prompt)
                    + self.str_to_approx_token_count(ddl)
                    < max_tokens
                ):
                    initial_prompt += f"{ddl}\n\n"

        return initial_prompt

    def add_documentation_to_prompt(
        self,
        initial_prompt: str,
        documentation_list: list[str],
        max_tokens: int = 14000,
    ) -> str:
        if len(documentation_list) > 0:
            initial_prompt += "\n===Additional Context \n\n"

            for documentation in documentation_list:
                if (
                    self.str_to_approx_token_count(initial_prompt)
                    + self.str_to_approx_token_count(documentation)
                    < max_tokens
                ):
                    initial_prompt += f"{documentation}\n\n"

        return initial_prompt

    def add_sql_to_prompt(
        self, initial_prompt: str, sql_list: list[str], max_tokens: int = 14000
    ) -> str:
        if len(sql_list) > 0:
            initial_prompt += "\n===Question-SQL Pairs\n\n"

            for question in sql_list:
                if (
                    self.str_to_approx_token_count(initial_prompt)
                    + self.str_to_approx_token_count(question["sql"])
                    < max_tokens
                ):
                    initial_prompt += f"{question['question']}\n{question['sql']}\n\n"

        return initial_prompt

    def get_sql_prompt(
        self,
        initial_prompt: str,
        question: str,
        question_sql_list: list,
        ddl_list: list,
        doc_list: list,
        **kwargs,
    ):
        """
        Example:
        ```python
        vn.get_sql_prompt(
            question="What are the top 10 customers by sales?",
            question_sql_list=[{"question": "What are the top 10 customers by sales?", "sql": "SELECT * FROM customers ORDER BY sales DESC LIMIT 10"}],
            ddl_list=["CREATE TABLE customers (id INT, name TEXT, sales DECIMAL)"],
            doc_list=["The customers table contains information about customers and their sales."],
        )

        ```

        This method is used to generate a prompt for the LLM to generate SQL.

        Args:
            question (str): The question to generate SQL for.
            question_sql_list (list): A list of questions and their corresponding SQL statements.
            ddl_list (list): A list of DDL statements.
            doc_list (list): A list of documentation.

        Returns:
            any: The prompt for the LLM to generate SQL.
        """

        if initial_prompt is None:
            initial_prompt = (
                f"You are a {self.dialect} expert. "
                + "Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. "
            )

        initial_prompt = self.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=self.max_tokens
        )

        if self.static_documentation != "":
            doc_list.append(self.static_documentation)

        initial_prompt = self.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=self.max_tokens
        )

        initial_prompt += (
            "===Response Guidelines \n"
            "1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \n"
            "2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \n"
            "3. If the provided context is insufficient, please explain why it can't be generated. \n"
            "4. Please use the most relevant table(s). \n"
            "5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \n"
            f"6. Ensure that the output SQL is {self.dialect}-compliant and executable, and free of syntax errors. \n"
        )

        message_log = [self.system_message(initial_prompt)]

        for example in question_sql_list:
            if example is None:
                print("example is None")
            else:
                if example is not None and "question" in example and "sql" in example:
                    message_log.append(self.user_message(example["question"]))
                    message_log.append(self.assistant_message(example["sql"]))

        message_log.append(self.user_message(question))

        return message_log

    def get_followup_questions_prompt(
        self,
        question: str,
        question_sql_list: list,
        ddl_list: list,
        doc_list: list,
        **kwargs,
    ) -> list:
        initial_prompt = f"The user initially asked the question: '{question}': \n\n"

        initial_prompt = self.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=self.max_tokens
        )

        initial_prompt = self.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=self.max_tokens
        )

        initial_prompt = self.add_sql_to_prompt(
            initial_prompt, question_sql_list, max_tokens=self.max_tokens
        )

        message_log = [self.system_message(initial_prompt)]
        message_log.append(
            self.user_message(
                "Generate a list of followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions."
            )
        )

        return message_log

    @abstractmethod
    def submit_prompt(self, prompt, **kwargs) -> str:
        """
        Example:
        ```python
        vn.submit_prompt(
            [
                vn.system_message("The user will give you SQL and you will try to guess what the business question this query is answering. Return just the question without any additional explanation. Do not reference the table name in the question."),
                vn.user_message("What are the top 10 customers by sales?"),
            ]
        )
        ```

        This method is used to submit a prompt to the LLM.

        Args:
            prompt (any): The prompt to submit to the LLM.

        Returns:
            str: The response from the LLM.
        """
        pass

    def generate_question(self, sql: str, **kwargs) -> str:
        response = self.submit_prompt(
            [
                self.system_message(
                    "The user will give you SQL and you will try to guess what the business question this query is answering. Return just the question without any additional explanation. Do not reference the table name in the question."
                ),
                self.user_message(sql),
            ],
            **kwargs,
        )

        return response

    def _extract_python_code(self, markdown_string: str) -> str:
        # Strip whitespace to avoid indentation errors in LLM-generated code
        markdown_string = markdown_string.strip()

        # Regex pattern to match Python code blocks
        pattern = r"```[\w\s]*python\n([\s\S]*?)```|```([\s\S]*?)```"

        # Find all matches in the markdown string
        matches = re.findall(pattern, markdown_string, re.IGNORECASE)

        # Extract the Python code from the matches
        python_code = []
        for match in matches:
            python = match[0] if match[0] else match[1]
            python_code.append(python.strip())

        if len(python_code) == 0:
            return markdown_string

        return python_code[0]

    def _sanitize_plotly_code(self, raw_plotly_code: str) -> str:
        # Remove the fig.show() statement from the plotly code
        plotly_code = raw_plotly_code.replace("fig.show()", "")

        return plotly_code

    def generate_plotly_code(
        self, question: str = None, sql: str = None, df_metadata: str = None, **kwargs
    ) -> str:
        if question is not None:
            system_msg = f"The following is a pandas DataFrame that contains the results of the query that answers the question the user asked: '{question}'"
        else:
            system_msg = "The following is a pandas DataFrame "

        if sql is not None:
            system_msg += f"\n\nThe DataFrame was produced using this query: {sql}\n\n"

        system_msg += f"The following is information about the resulting pandas DataFrame 'df': \n{df_metadata}"

        message_log = [
            self.system_message(system_msg),
            self.user_message(
                "Can you generate the Python plotly code to chart the results of the dataframe? Assume the data is in a pandas dataframe called 'df'. If there is only one value in the dataframe, use an Indicator. Respond with only Python code. Do not answer with any explanations -- just the code."
            ),
        ]

        plotly_code = self.submit_prompt(message_log, kwargs=kwargs)

        return self._sanitize_plotly_code(self._extract_python_code(plotly_code))

    # ----------------- Connect to Any Database to run the Generated SQL ----------------- #

    def connect_to_snowflake(
        self,
        account: str,
        username: str,
        password: str,
        database: str,
        role: Union[str, None] = None,
        warehouse: Union[str, None] = None,
        **kwargs,
    ):
        try:
            snowflake = __import__("snowflake.connector")
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method, run command:"
                " \npip install vanna[snowflake]"
            )

        if username == "my-username":
            username_env = os.getenv("SNOWFLAKE_USERNAME")

            if username_env is not None:
                username = username_env
            else:
                raise ImproperlyConfigured("Please set your Snowflake username.")

        if password == "mypassword":
            password_env = os.getenv("SNOWFLAKE_PASSWORD")

            if password_env is not None:
                password = password_env
            else:
                raise ImproperlyConfigured("Please set your Snowflake password.")

        if account == "my-account":
            account_env = os.getenv("SNOWFLAKE_ACCOUNT")

            if account_env is not None:
                account = account_env
            else:
                raise ImproperlyConfigured("Please set your Snowflake account.")

        if database == "my-database":
            database_env = os.getenv("SNOWFLAKE_DATABASE")

            if database_env is not None:
                database = database_env
            else:
                raise ImproperlyConfigured("Please set your Snowflake database.")

        conn = snowflake.connector.connect(
            user=username,
            password=password,
            account=account,
            database=database,
            client_session_keep_alive=True,
            **kwargs,
        )

        def run_sql_snowflake(sql: str) -> pd.DataFrame:
            cs = conn.cursor()

            if role is not None:
                cs.execute(f"USE ROLE {role}")

            if warehouse is not None:
                cs.execute(f"USE WAREHOUSE {warehouse}")
            cs.execute(f"USE DATABASE {database}")

            cur = cs.execute(sql)

            results = cur.fetchall()

            # Create a pandas dataframe from the results
            df = pd.DataFrame(results, columns=[desc[0] for desc in cur.description])

            return df

        self.dialect = "Snowflake SQL"
        self.run_sql = run_sql_snowflake
        self.run_sql_is_set = True

    def connect_to_sqlite(self, url: str, check_same_thread: bool = False, **kwargs):
        """
        Connect to a SQLite database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]

        Args:
            url (str): The URL of the database to connect to.
            check_same_thread (str): Allow the connection may be accessed in multiple threads.
        Returns:
            None
        """

        # URL of the database to download

        # Path to save the downloaded database
        path = os.path.basename(urlparse(url).path)

        # Download the database if it doesn't exist
        if not os.path.exists(url):
            response = requests.get(url)
            response.raise_for_status()  # Check that the request was successful
            with open(path, "wb") as f:
                f.write(response.content)
            url = path

        # Connect to the database
        conn = sqlite3.connect(url, check_same_thread=check_same_thread, **kwargs)

        def run_sql_sqlite(sql: str):
            return pd.read_sql_query(sql, conn)

        self.dialect = "SQLite"
        self.run_sql = run_sql_sqlite
        self.run_sql_is_set = True

    def connect_to_postgres(
        self,
        host: str = None,
        dbname: str = None,
        user: str = None,
        password: str = None,
        port: int = None,
        **kwargs,
    ):
        """
        Connect to postgres using the psycopg2 connector. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
        **Example:**
        ```python
        vn.connect_to_postgres(
            host="myhost",
            dbname="mydatabase",
            user="myuser",
            password="mypassword",
            port=5432
        )
        ```
        Args:
            host (str): The postgres host.
            dbname (str): The postgres database name.
            user (str): The postgres user.
            password (str): The postgres password.
            port (int): The postgres Port.
        """

        try:
            import psycopg2
            import psycopg2.extras
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install vanna[postgres]"
            )

        if not host:
            host = os.getenv("HOST")

        if not host:
            raise ImproperlyConfigured("Please set your postgres host")

        if not dbname:
            dbname = os.getenv("DATABASE")

        if not dbname:
            raise ImproperlyConfigured("Please set your postgres database")

        if not user:
            user = os.getenv("PG_USER")

        if not user:
            raise ImproperlyConfigured("Please set your postgres user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your postgres password")

        if not port:
            port = os.getenv("PORT")

        if not port:
            raise ImproperlyConfigured("Please set your postgres port")

        conn = None

        try:
            conn = psycopg2.connect(
                host=host,
                dbname=dbname,
                user=user,
                password=password,
                port=port,
                **kwargs,
            )
        except psycopg2.Error as e:
            raise ValidationError(e)

        def connect_to_db():
            return psycopg2.connect(
                host=host,
                dbname=dbname,
                user=user,
                password=password,
                port=port,
                **kwargs,
            )

        def run_sql_postgres(sql: str) -> Union[pd.DataFrame, None]:
            conn = None
            try:
                conn = connect_to_db()  # Initial connection attempt
                cs = conn.cursor()
                cs.execute(sql)
                results = cs.fetchall()

                # Create a pandas dataframe from the results
                df = pd.DataFrame(results, columns=[desc[0] for desc in cs.description])
                return df

            except psycopg2.InterfaceError as e:
                # Attempt to reconnect and retry the operation
                if conn:
                    conn.close()  # Ensure any existing connection is closed
                conn = connect_to_db()
                cs = conn.cursor()
                cs.execute(sql)
                results = cs.fetchall()

                # Create a pandas dataframe from the results
                df = pd.DataFrame(results, columns=[desc[0] for desc in cs.description])
                return df

            except psycopg2.Error as e:
                if conn:
                    conn.rollback()
                    raise ValidationError(e)

            except Exception as e:
                conn.rollback()
                raise e

        self.dialect = "PostgreSQL"
        self.run_sql_is_set = True
        self.run_sql = run_sql_postgres

    def connect_to_mysql(
        self,
        host: str = None,
        dbname: str = None,
        user: str = None,
        password: str = None,
        port: int = None,
        **kwargs,
    ):
        try:
            import pymysql.cursors
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install PyMySQL"
            )

        if not host:
            host = os.getenv("HOST")

        if not host:
            raise ImproperlyConfigured("Please set your MySQL host")

        if not dbname:
            dbname = os.getenv("DATABASE")

        if not dbname:
            raise ImproperlyConfigured("Please set your MySQL database")

        if not user:
            user = os.getenv("USER")

        if not user:
            raise ImproperlyConfigured("Please set your MySQL user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your MySQL password")

        if not port:
            port = os.getenv("PORT")

        if not port:
            raise ImproperlyConfigured("Please set your MySQL port")

        conn = None

        try:
            conn = pymysql.connect(
                host=host,
                user=user,
                password=password,
                database=dbname,
                port=port,
                cursorclass=pymysql.cursors.DictCursor,
                **kwargs,
            )
        except pymysql.Error as e:
            raise ValidationError(e)

        def run_sql_mysql(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    conn.ping(reconnect=True)
                    cs = conn.cursor()
                    cs.execute(sql)
                    results = cs.fetchall()

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(
                        results, columns=[desc[0] for desc in cs.description]
                    )
                    return df

                except pymysql.Error as e:
                    conn.rollback()
                    raise ValidationError(e)

                except Exception as e:
                    conn.rollback()
                    raise e

        self.run_sql_is_set = True
        self.run_sql = run_sql_mysql

    def connect_to_clickhouse(
        self,
        host: str = None,
        dbname: str = None,
        user: str = None,
        password: str = None,
        port: int = None,
        **kwargs,
    ):
        try:
            import clickhouse_connect
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install clickhouse_connect"
            )

        if not host:
            host = os.getenv("HOST")

        if not host:
            raise ImproperlyConfigured("Please set your ClickHouse host")

        if not dbname:
            dbname = os.getenv("DATABASE")

        if not dbname:
            raise ImproperlyConfigured("Please set your ClickHouse database")

        if not user:
            user = os.getenv("USER")

        if not user:
            raise ImproperlyConfigured("Please set your ClickHouse user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your ClickHouse password")

        if not port:
            port = os.getenv("PORT")

        if not port:
            raise ImproperlyConfigured("Please set your ClickHouse port")

        conn = None

        try:
            conn = clickhouse_connect.get_client(
                host=host,
                port=port,
                username=user,
                password=password,
                database=dbname,
                **kwargs,
            )
            print(conn)
        except Exception as e:
            raise ValidationError(e)

        def run_sql_clickhouse(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    result = conn.query(sql)
                    results = result.result_rows

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(results, columns=result.column_names)
                    return df

                except Exception as e:
                    raise e

        self.run_sql_is_set = True
        self.run_sql = run_sql_clickhouse

    def connect_to_oracle(
        self, user: str = None, password: str = None, dsn: str = None, **kwargs
    ):
        """
        Connect to an Oracle db using oracledb package. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
        **Example:**
        ```python
        vn.connect_to_oracle(
        user="username",
        password="password",
        dsn="host:port/sid",
        )
        ```
        Args:
            USER (str): Oracle db user name.
            PASSWORD (str): Oracle db user password.
            DSN (str): Oracle db host ip - host:port/sid.
        """

        try:
            import oracledb
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install oracledb"
            )

        if not dsn:
            dsn = os.getenv("DSN")

        if not dsn:
            raise ImproperlyConfigured(
                "Please set your Oracle dsn which should include host:port/sid"
            )

        if not user:
            user = os.getenv("USER")

        if not user:
            raise ImproperlyConfigured("Please set your Oracle db user")

        if not password:
            password = os.getenv("PASSWORD")

        if not password:
            raise ImproperlyConfigured("Please set your Oracle db password")

        conn = None

        try:
            conn = oracledb.connect(user=user, password=password, dsn=dsn, **kwargs)
        except oracledb.Error as e:
            raise ValidationError(e)

        def run_sql_oracle(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    sql = sql.rstrip()
                    if sql.endswith(
                        ";"
                    ):  # fix for a known problem with Oracle db where an extra ; will cause an error.
                        sql = sql[:-1]

                    cs = conn.cursor()
                    cs.execute(sql)
                    results = cs.fetchall()

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(
                        results, columns=[desc[0] for desc in cs.description]
                    )
                    return df

                except oracledb.Error as e:
                    conn.rollback()
                    raise ValidationError(e)

                except Exception as e:
                    conn.rollback()
                    raise e

        self.run_sql_is_set = True
        self.run_sql = run_sql_oracle

    def connect_to_bigquery(
        self, cred_file_path: str = None, project_id: str = None, **kwargs
    ):
        """
        Connect to gcs using the bigquery connector. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
        **Example:**
        ```python
        vn.connect_to_bigquery(
            project_id="myprojectid",
            cred_file_path="path/to/credentials.json",
        )
        ```
        Args:
            project_id (str): The gcs project id.
            cred_file_path (str): The gcs credential file path
        """

        try:
            from google.api_core.exceptions import GoogleAPIError
            from google.cloud import bigquery
            from google.oauth2 import service_account
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method, run command:"
                " \npip install vanna[bigquery]"
            )

        if not project_id:
            project_id = os.getenv("PROJECT_ID")

        if not project_id:
            raise ImproperlyConfigured("Please set your Google Cloud Project ID.")

        import sys

        if "google.colab" in sys.modules:
            try:
                from google.colab import auth

                auth.authenticate_user()
            except Exception as e:
                raise ImproperlyConfigured(e)
        else:
            print("Not using Google Colab.")

        conn = None

        if not cred_file_path:
            try:
                conn = bigquery.Client(project=project_id)
            except Exception:
                print("Could not found any google cloud implicit credentials")
        else:
            # Validate file path and pemissions
            validate_config_path(cred_file_path)

        if not conn:
            with open(cred_file_path, "r") as f:
                credentials = service_account.Credentials.from_service_account_info(
                    json.loads(f.read()),
                    scopes=["https://www.googleapis.com/auth/cloud-platform"],
                )

            try:
                conn = bigquery.Client(
                    project=project_id, credentials=credentials, **kwargs
                )
            except Exception:
                raise ImproperlyConfigured(
                    "Could not connect to bigquery please correct credentials"
                )

        def run_sql_bigquery(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                job = conn.query(sql)
                df = job.result().to_dataframe()
                return df
            return None

        self.dialect = "BigQuery SQL"
        self.run_sql_is_set = True
        self.run_sql = run_sql_bigquery

    def connect_to_duckdb(self, url: str, init_sql: str = None, **kwargs):
        """
        Connect to a DuckDB database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]

        Args:
            url (str): The URL of the database to connect to. Use :memory: to create an in-memory database. Use md: or motherduck: to use the MotherDuck database.
            init_sql (str, optional): SQL to run when connecting to the database. Defaults to None.

        Returns:
            None
        """
        try:
            import duckdb
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install vanna[duckdb]"
            )
        # URL of the database to download
        if url == ":memory:" or url == "":
            path = ":memory:"
        else:
            # Path to save the downloaded database
            print(os.path.exists(url))
            if os.path.exists(url):
                path = url
            elif url.startswith("md") or url.startswith("motherduck"):
                path = url
            else:
                path = os.path.basename(urlparse(url).path)
                # Download the database if it doesn't exist
                if not os.path.exists(path):
                    response = requests.get(url)
                    response.raise_for_status()  # Check that the request was successful
                    with open(path, "wb") as f:
                        f.write(response.content)

        # Connect to the database
        conn = duckdb.connect(path, **kwargs)
        if init_sql:
            conn.query(init_sql)

        def run_sql_duckdb(sql: str):
            return conn.query(sql).to_df()

        self.dialect = "DuckDB SQL"
        self.run_sql = run_sql_duckdb
        self.run_sql_is_set = True

    def connect_to_mssql(self, odbc_conn_str: str, **kwargs):
        """
        Connect to a Microsoft SQL Server database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]

        Args:
            odbc_conn_str (str): The ODBC connection string.

        Returns:
            None
        """
        try:
            import pyodbc
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: pip install pyodbc"
            )

        try:
            import sqlalchemy as sa
            from sqlalchemy.engine import URL
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: pip install sqlalchemy"
            )

        connection_url = URL.create(
            "mssql+pyodbc", query={"odbc_connect": odbc_conn_str}
        )

        from sqlalchemy import create_engine

        engine = create_engine(connection_url, **kwargs)

        def run_sql_mssql(sql: str):
            # Execute the SQL statement and return the result as a pandas DataFrame
            with engine.begin() as conn:
                df = pd.read_sql_query(sa.text(sql), conn)
                conn.close()
                return df

            raise Exception("Couldn't run sql")

        self.dialect = "T-SQL / Microsoft SQL Server"
        self.run_sql = run_sql_mssql
        self.run_sql_is_set = True

    def connect_to_presto(
        self,
        host: str,
        catalog: str = "hive",
        schema: str = "default",
        user: str = None,
        password: str = None,
        port: int = None,
        combined_pem_path: str = None,
        protocol: str = "https",
        requests_kwargs: dict = None,
        **kwargs,
    ):
        """
        Connect to a Presto database using the specified parameters.

        Args:
            host (str): The host address of the Presto database.
            catalog (str): The catalog to use in the Presto environment.
            schema (str): The schema to use in the Presto environment.
            user (str): The username for authentication.
            password (str): The password for authentication.
            port (int): The port number for the Presto connection.
            combined_pem_path (str): The path to the combined pem file for SSL connection.
            protocol (str): The protocol to use for the connection (default is 'https').
            requests_kwargs (dict): Additional keyword arguments for requests.

        Raises:
            DependencyError: If required dependencies are not installed.
            ImproperlyConfigured: If essential configuration settings are missing.

        Returns:
            None
        """
        try:
            from pyhive import presto
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install pyhive"
            )

        if not host:
            host = os.getenv("PRESTO_HOST")

        if not host:
            raise ImproperlyConfigured("Please set your presto host")

        if not catalog:
            catalog = os.getenv("PRESTO_CATALOG")

        if not catalog:
            raise ImproperlyConfigured("Please set your presto catalog")

        if not user:
            user = os.getenv("PRESTO_USER")

        if not user:
            raise ImproperlyConfigured("Please set your presto user")

        if not password:
            password = os.getenv("PRESTO_PASSWORD")

        if not port:
            port = os.getenv("PRESTO_PORT")

        if not port:
            raise ImproperlyConfigured("Please set your presto port")

        conn = None

        try:
            if requests_kwargs is None and combined_pem_path is not None:
                # use the combined pem file to verify the SSL connection
                requests_kwargs = {
                    "verify": combined_pem_path,  # 使用转换后得到的 PEM 文件进行 SSL 验证
                }
            conn = presto.Connection(
                host=host,
                username=user,
                password=password,
                catalog=catalog,
                schema=schema,
                port=port,
                protocol=protocol,
                requests_kwargs=requests_kwargs,
                **kwargs,
            )
        except presto.Error as e:
            raise ValidationError(e)

        def run_sql_presto(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    sql = sql.rstrip()
                    # fix for a known problem with presto db where an extra ; will cause an error.
                    if sql.endswith(";"):
                        sql = sql[:-1]
                    cs = conn.cursor()
                    cs.execute(sql)
                    results = cs.fetchall()

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(
                        results, columns=[desc[0] for desc in cs.description]
                    )
                    return df

                except presto.Error as e:
                    print(e)
                    raise ValidationError(e)

                except Exception as e:
                    print(e)
                    raise e

        self.run_sql_is_set = True
        self.run_sql = run_sql_presto

    def connect_to_hive(
        self,
        host: str = None,
        dbname: str = "default",
        user: str = None,
        password: str = None,
        port: int = None,
        auth: str = "CUSTOM",
        **kwargs,
    ):
        """
        Connect to a Hive database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
        Connect to a Hive database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]

        Args:
            host (str): The host of the Hive database.
            dbname (str): The name of the database to connect to.
            user (str): The username to use for authentication.
            password (str): The password to use for authentication.
            port (int): The port to use for the connection.
            auth (str): The authentication method to use.

        Returns:
            None
        """

        try:
            from pyhive import hive
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method,"
                " run command: \npip install pyhive"
            )

        if not host:
            host = os.getenv("HIVE_HOST")

        if not host:
            raise ImproperlyConfigured("Please set your hive host")

        if not dbname:
            dbname = os.getenv("HIVE_DATABASE")

        if not dbname:
            raise ImproperlyConfigured("Please set your hive database")

        if not user:
            user = os.getenv("HIVE_USER")

        if not user:
            raise ImproperlyConfigured("Please set your hive user")

        if not password:
            password = os.getenv("HIVE_PASSWORD")

        if not port:
            port = os.getenv("HIVE_PORT")

        if not port:
            raise ImproperlyConfigured("Please set your hive port")

        conn = None

        try:
            conn = hive.Connection(
                host=host,
                username=user,
                password=password,
                database=dbname,
                port=port,
                auth=auth,
            )
        except hive.Error as e:
            raise ValidationError(e)

        def run_sql_hive(sql: str) -> Union[pd.DataFrame, None]:
            if conn:
                try:
                    cs = conn.cursor()
                    cs.execute(sql)
                    results = cs.fetchall()

                    # Create a pandas dataframe from the results
                    df = pd.DataFrame(
                        results, columns=[desc[0] for desc in cs.description]
                    )
                    return df

                except hive.Error as e:
                    print(e)
                    raise ValidationError(e)

                except Exception as e:
                    print(e)
                    raise e

        self.run_sql_is_set = True
        self.run_sql = run_sql_hive

    def run_sql(self, sql: str, **kwargs) -> pd.DataFrame:
        """
        Example:
        ```python
        vn.run_sql("SELECT * FROM my_table")
        ```

        Run a SQL query on the connected database.

        Args:
            sql (str): The SQL query to run.

        Returns:
            pd.DataFrame: The results of the SQL query.
        """
        raise Exception(
            "You need to connect to a database first by running vn.connect_to_snowflake(), vn.connect_to_postgres(), similar function, or manually set vn.run_sql"
        )

    def ask(
        self,
        question: Union[str, None] = None,
        print_results: bool = True,
        auto_train: bool = True,
        visualize: bool = True,  # if False, will not generate plotly code
        allow_llm_to_see_data: bool = False,
    ) -> Union[
        Tuple[
            Union[str, None],
            Union[pd.DataFrame, None],
            Union[plotly.graph_objs.Figure, None],
        ],
        None,
    ]:
        """
        **Example:**
        ```python
        vn.ask("What are the top 10 customers by sales?")
        ```

        Ask Vanna.AI a question and get the SQL query that answers it.

        Args:
            question (str): The question to ask.
            print_results (bool): Whether to print the results of the SQL query.
            auto_train (bool): Whether to automatically train Vanna.AI on the question and SQL query.
            visualize (bool): Whether to generate plotly code and display the plotly figure.

        Returns:
            Tuple[str, pd.DataFrame, plotly.graph_objs.Figure]: The SQL query, the results of the SQL query, and the plotly figure.
        """

        if question is None:
            question = input("Enter a question: ")

        try:
            sql = self.generate_sql(
                question=question, allow_llm_to_see_data=allow_llm_to_see_data
            )
        except Exception as e:
            print(e)
            return None, None, None

        if print_results:
            try:
                from IPython.display import Code, display

                display(Code(sql))
            except Exception as e:
                print(sql)

        if self.run_sql_is_set is False:
            print("If you want to run the SQL query, connect to a database first.")

            if print_results:
                return None
            else:
                return sql, None, None

        try:
            df = self.run_sql(sql)

            if print_results:
                try:
                    display = __import__(
                        "IPython.display", fromList=["display"]
                    ).display
                    display(df)
                except Exception as e:
                    print(df)

            if len(df) > 0 and auto_train:
                self.add_question_sql(question=question, sql=sql)
            # Only generate plotly code if visualize is True
            if visualize:
                try:
                    plotly_code = self.generate_plotly_code(
                        question=question,
                        sql=sql,
                        df_metadata=f"Running df.dtypes gives:\n {df.dtypes}",
                    )
                    fig = self.get_plotly_figure(plotly_code=plotly_code, df=df)
                    if print_results:
                        try:
                            display = __import__(
                                "IPython.display", fromlist=["display"]
                            ).display
                            Image = __import__(
                                "IPython.display", fromlist=["Image"]
                            ).Image
                            img_bytes = fig.to_image(format="png", scale=2)
                            display(Image(img_bytes))
                        except Exception as e:
                            fig.show()
                except Exception as e:
                    # Print stack trace
                    traceback.print_stack()
                    traceback.print_exc()
                    print("Couldn't run plotly code: ", e)
                    if print_results:
                        return None
                    else:
                        return sql, df, None
            else:
                return sql, df, None

        except Exception as e:
            print("Couldn't run sql: ", e)
            if print_results:
                return None
            else:
                return sql, None, None
        return sql, df, fig

    def train(
        self,
        question: str = None,
        sql: str = None,
        ddl: str = None,
        documentation: str = None,
        plan: TrainingPlan = None,
    ) -> str:
        """
        **Example:**
        ```python
        vn.train()
        ```

        Train Vanna.AI on a question and its corresponding SQL query.
        If you call it with no arguments, it will check if you connected to a database and it will attempt to train on the metadata of that database.
        If you call it with the sql argument, it's equivalent to [`vn.add_question_sql()`][vanna.base.base.VannaBase.add_question_sql].
        If you call it with the ddl argument, it's equivalent to [`vn.add_ddl()`][vanna.base.base.VannaBase.add_ddl].
        If you call it with the documentation argument, it's equivalent to [`vn.add_documentation()`][vanna.base.base.VannaBase.add_documentation].
        Additionally, you can pass a [`TrainingPlan`][vanna.types.TrainingPlan] object. Get a training plan with [`vn.get_training_plan_generic()`][vanna.base.base.VannaBase.get_training_plan_generic].

        Args:
            question (str): The question to train on.
            sql (str): The SQL query to train on.
            ddl (str):  The DDL statement.
            documentation (str): The documentation to train on.
            plan (TrainingPlan): The training plan to train on.
        """

        if question and not sql:
            raise ValidationError("Please also provide a SQL query")

        if documentation:
            print("Adding documentation....")
            return self.add_documentation(documentation)

        if sql:
            if question is None:
                question = self.generate_question(sql)
                print("Question generated with sql:", question, "\nAdding SQL...")
            return self.add_question_sql(question=question, sql=sql)

        if ddl:
            print("Adding ddl:", ddl)
            return self.add_ddl(ddl)

        if plan:
            for item in plan._plan:
                if item.item_type == TrainingPlanItem.ITEM_TYPE_DDL:
                    self.add_ddl(item.item_value)
                elif item.item_type == TrainingPlanItem.ITEM_TYPE_IS:
                    self.add_documentation(item.item_value)
                elif item.item_type == TrainingPlanItem.ITEM_TYPE_SQL:
                    self.add_question_sql(question=item.item_name, sql=item.item_value)

    def _get_databases(self) -> List[str]:
        try:
            print("Trying INFORMATION_SCHEMA.DATABASES")
            df_databases = self.run_sql("SELECT * FROM INFORMATION_SCHEMA.DATABASES")
        except Exception as e:
            print(e)
            try:
                print("Trying SHOW DATABASES")
                df_databases = self.run_sql("SHOW DATABASES")
            except Exception as e:
                print(e)
                return []

        return df_databases["DATABASE_NAME"].unique().tolist()

    def _get_information_schema_tables(self, database: str) -> pd.DataFrame:
        df_tables = self.run_sql(f"SELECT * FROM {database}.INFORMATION_SCHEMA.TABLES")

        return df_tables

    def get_training_plan_generic(self, df) -> TrainingPlan:
        """
        This method is used to generate a training plan from an information schema dataframe.

        Basically what it does is breaks up INFORMATION_SCHEMA.COLUMNS into groups of table/column descriptions that can be used to pass to the LLM.

        Args:
            df (pd.DataFrame): The dataframe to generate the training plan from.

        Returns:
            TrainingPlan: The training plan.
        """
        # For each of the following, we look at the df columns to see if there's a match:
        database_column = df.columns[
            df.columns.str.lower().str.contains("database")
            | df.columns.str.lower().str.contains("table_catalog")
        ].to_list()[0]
        schema_column = df.columns[
            df.columns.str.lower().str.contains("table_schema")
        ].to_list()[0]
        table_column = df.columns[
            df.columns.str.lower().str.contains("table_name")
        ].to_list()[0]
        columns = [database_column, schema_column, table_column]
        candidates = ["column_name", "data_type", "comment"]
        matches = df.columns.str.lower().str.contains("|".join(candidates), regex=True)
        columns += df.columns[matches].to_list()

        plan = TrainingPlan([])

        for database in df[database_column].unique().tolist():
            for schema in (
                df.query(f'{database_column} == "{database}"')[schema_column]
                .unique()
                .tolist()
            ):
                for table in (
                    df.query(
                        f'{database_column} == "{database}" and {schema_column} == "{schema}"'
                    )[table_column]
                    .unique()
                    .tolist()
                ):
                    df_columns_filtered_to_table = df.query(
                        f'{database_column} == "{database}" and {schema_column} == "{schema}" and {table_column} == "{table}"'
                    )
                    doc = f"The following columns are in the {table} table in the {database} database:\n\n"
                    doc += df_columns_filtered_to_table[columns].to_markdown()

                    plan._plan.append(
                        TrainingPlanItem(
                            item_type=TrainingPlanItem.ITEM_TYPE_IS,
                            item_group=f"{database}.{schema}",
                            item_name=table,
                            item_value=doc,
                        )
                    )

        return plan

    def get_training_plan_snowflake(
        self,
        filter_databases: Union[List[str], None] = None,
        filter_schemas: Union[List[str], None] = None,
        include_information_schema: bool = False,
        use_historical_queries: bool = True,
    ) -> TrainingPlan:
        plan = TrainingPlan([])

        if self.run_sql_is_set is False:
            raise ImproperlyConfigured("Please connect to a database first.")

        if use_historical_queries:
            try:
                print("Trying query history")
                df_history = self.run_sql(
                    """ select * from table(information_schema.query_history(result_limit => 5000)) order by start_time"""
                )

                df_history_filtered = df_history.query("ROWS_PRODUCED > 1")
                if filter_databases is not None:
                    mask = (
                        df_history_filtered["QUERY_TEXT"]
                        .str.lower()
                        .apply(
                            lambda x: any(
                                s in x for s in [s.lower() for s in filter_databases]
                            )
                        )
                    )
                    df_history_filtered = df_history_filtered[mask]

                if filter_schemas is not None:
                    mask = (
                        df_history_filtered["QUERY_TEXT"]
                        .str.lower()
                        .apply(
                            lambda x: any(
                                s in x for s in [s.lower() for s in filter_schemas]
                            )
                        )
                    )
                    df_history_filtered = df_history_filtered[mask]

                if len(df_history_filtered) > 10:
                    df_history_filtered = df_history_filtered.sample(10)

                for query in df_history_filtered["QUERY_TEXT"].unique().tolist():
                    plan._plan.append(
                        TrainingPlanItem(
                            item_type=TrainingPlanItem.ITEM_TYPE_SQL,
                            item_group="",
                            item_name=self.generate_question(query),
                            item_value=query,
                        )
                    )

            except Exception as e:
                print(e)

        databases = self._get_databases()

        for database in databases:
            if filter_databases is not None and database not in filter_databases:
                continue

            try:
                df_tables = self._get_information_schema_tables(database=database)

                print(f"Trying INFORMATION_SCHEMA.COLUMNS for {database}")
                df_columns = self.run_sql(
                    f"SELECT * FROM {database}.INFORMATION_SCHEMA.COLUMNS"
                )

                for schema in df_tables["TABLE_SCHEMA"].unique().tolist():
                    if filter_schemas is not None and schema not in filter_schemas:
                        continue

                    if (
                        not include_information_schema
                        and schema == "INFORMATION_SCHEMA"
                    ):
                        continue

                    df_columns_filtered_to_schema = df_columns.query(
                        f"TABLE_SCHEMA == '{schema}'"
                    )

                    try:
                        tables = (
                            df_columns_filtered_to_schema["TABLE_NAME"]
                            .unique()
                            .tolist()
                        )

                        for table in tables:
                            df_columns_filtered_to_table = (
                                df_columns_filtered_to_schema.query(
                                    f"TABLE_NAME == '{table}'"
                                )
                            )
                            doc = f"The following columns are in the {table} table in the {database} database:\n\n"
                            doc += df_columns_filtered_to_table[
                                [
                                    "TABLE_CATALOG",
                                    "TABLE_SCHEMA",
                                    "TABLE_NAME",
                                    "COLUMN_NAME",
                                    "DATA_TYPE",
                                    "COMMENT",
                                ]
                            ].to_markdown()

                            plan._plan.append(
                                TrainingPlanItem(
                                    item_type=TrainingPlanItem.ITEM_TYPE_IS,
                                    item_group=f"{database}.{schema}",
                                    item_name=table,
                                    item_value=doc,
                                )
                            )

                    except Exception as e:
                        print(e)
                        pass
            except Exception as e:
                print(e)

        return plan

    def get_plotly_figure(
        self, plotly_code: str, df: pd.DataFrame, dark_mode: bool = True
    ) -> plotly.graph_objs.Figure:
        """
        **Example:**
        ```python
        fig = vn.get_plotly_figure(
            plotly_code="fig = px.bar(df, x='name', y='salary')",
            df=df
        )
        fig.show()
        ```
        Get a Plotly figure from a dataframe and Plotly code.

        Args:
            df (pd.DataFrame): The dataframe to use.
            plotly_code (str): The Plotly code to use.

        Returns:
            plotly.graph_objs.Figure: The Plotly figure.
        """
        ldict = {"df": df, "px": px, "go": go}
        try:
            exec(plotly_code, globals(), ldict)

            fig = ldict.get("fig", None)
        except Exception as e:
            # Inspect data types
            numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
            categorical_cols = df.select_dtypes(
                include=["object", "category"]
            ).columns.tolist()

            # Decision-making for plot type
            if len(numeric_cols) >= 2:
                # Use the first two numeric columns for a scatter plot
                fig = px.scatter(df, x=numeric_cols[0], y=numeric_cols[1])
            elif len(numeric_cols) == 1 and len(categorical_cols) >= 1:
                # Use a bar plot if there's one numeric and one categorical column
                fig = px.bar(df, x=categorical_cols[0], y=numeric_cols[0])
            elif len(categorical_cols) >= 1 and df[categorical_cols[0]].nunique() < 10:
                # Use a pie chart for categorical data with fewer unique values
                fig = px.pie(df, names=categorical_cols[0])
            else:
                # Default to a simple line plot if above conditions are not met
                fig = px.line(df)

        if fig is None:
            return None

        if dark_mode:
            fig.update_layout(template="plotly_dark")

        return fig


================================================
FILE: src/vanna/legacy/bedrock/__init__.py
================================================
from .bedrock_converse import Bedrock_Converse


================================================
FILE: src/vanna/legacy/bedrock/bedrock_converse.py
================================================
from ..base import VannaBase

try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    raise ImportError("Please install boto3 and botocore to use Amazon Bedrock models")


class Bedrock_Converse(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        # default parameters
        self.temperature = 0.0
        self.max_tokens = 500

        if client is None:
            raise ValueError(
                "A valid Bedrock runtime client must be provided to invoke Bedrock models"
            )
        else:
            self.client = client

        if config is None:
            raise ValueError(
                "Config is required with model_id and inference parameters"
            )

        if "modelId" not in config:
            raise ValueError("config must contain a modelId to invoke")
        else:
            self.model = config["modelId"]

        if "temperature" in config:
            self.temperature = config["temperature"]

        if "max_tokens" in config:
            self.max_tokens = config["max_tokens"]

    def system_message(self, message: str) -> dict:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> dict:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> dict:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        inference_config = {
            "temperature": self.temperature,
            "maxTokens": self.max_tokens,
        }
        additional_model_fields = {
            "top_p": 1,  # setting top_p value for nucleus sampling
        }

        system_message = None
        no_system_prompt = []
        for prompt_message in prompt:
            role = prompt_message["role"]
            if role == "system":
                system_message = prompt_message["content"]
            else:
                no_system_prompt.append(
                    {"role": role, "content": [{"text": prompt_message["content"]}]}
                )

        converse_api_params = {
            "modelId": self.model,
            "messages": no_system_prompt,
            "inferenceConfig": inference_config,
            "additionalModelRequestFields": additional_model_fields,
        }

        if system_message:
            converse_api_params["system"] = [{"text": system_message}]

        try:
            response = self.client.converse(**converse_api_params)
            text_content = response["output"]["message"]["content"][0]["text"]
            return text_content
        except ClientError as err:
            message = err.response["Error"]["Message"]
            raise Exception(f"A Bedrock client error occurred: {message}")


================================================
FILE: src/vanna/legacy/chromadb/__init__.py
================================================
from .chromadb_vector import ChromaDB_VectorStore


================================================
FILE: src/vanna/legacy/chromadb/chromadb_vector.py
================================================
import json
from typing import List

import chromadb
import pandas as pd
from chromadb.config import Settings
from chromadb.utils import embedding_functions

from ..base import VannaBase
from ..utils import deterministic_uuid

default_ef = embedding_functions.DefaultEmbeddingFunction()


class ChromaDB_VectorStore(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        if config is None:
            config = {}

        path = config.get("path", ".")
        self.embedding_function = config.get("embedding_function", default_ef)
        curr_client = config.get("client", "persistent")
        collection_metadata = config.get("collection_metadata", None)
        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
        self.n_results_documentation = config.get(
            "n_results_documentation", config.get("n_results", 10)
        )
        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))

        if curr_client == "persistent":
            self.chroma_client = chromadb.PersistentClient(
                path=path, settings=Settings(anonymized_telemetry=False)
            )
        elif curr_client == "in-memory":
            self.chroma_client = chromadb.EphemeralClient(
                settings=Settings(anonymized_telemetry=False)
            )
        elif isinstance(curr_client, chromadb.api.client.Client):
            # allow providing client directly
            self.chroma_client = curr_client
        else:
            raise ValueError(f"Unsupported client was set in config: {curr_client}")

        self.documentation_collection = self.chroma_client.get_or_create_collection(
            name="documentation",
            embedding_function=self.embedding_function,
            metadata=collection_metadata,
        )
        self.ddl_collection = self.chroma_client.get_or_create_collection(
            name="ddl",
            embedding_function=self.embedding_function,
            metadata=collection_metadata,
        )
        self.sql_collection = self.chroma_client.get_or_create_collection(
            name="sql",
            embedding_function=self.embedding_function,
            metadata=collection_metadata,
        )

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embedding = self.embedding_function([data])
        if len(embedding) == 1:
            return embedding[0]
        return embedding

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        question_sql_json = json.dumps(
            {
                "question": question,
                "sql": sql,
            },
            ensure_ascii=False,
        )
        id = deterministic_uuid(question_sql_json) + "-sql"
        self.sql_collection.add(
            documents=question_sql_json,
            embeddings=self.generate_embedding(question_sql_json),
            ids=id,
        )

        return id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        id = deterministic_uuid(ddl) + "-ddl"
        self.ddl_collection.add(
            documents=ddl,
            embeddings=self.generate_embedding(ddl),
            ids=id,
        )
        return id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        id = deterministic_uuid(documentation) + "-doc"
        self.documentation_collection.add(
            documents=documentation,
            embeddings=self.generate_embedding(documentation),
            ids=id,
        )
        return id

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        sql_data = self.sql_collection.get()

        df = pd.DataFrame()

        if sql_data is not None:
            # Extract the documents and ids
            documents = [json.loads(doc) for doc in sql_data["documents"]]
            ids = sql_data["ids"]

            # Create a DataFrame
            df_sql = pd.DataFrame(
                {
                    "id": ids,
                    "question": [doc["question"] for doc in documents],
                    "content": [doc["sql"] for doc in documents],
                }
            )

            df_sql["training_data_type"] = "sql"

            df = pd.concat([df, df_sql])

        ddl_data = self.ddl_collection.get()

        if ddl_data is not None:
            # Extract the documents and ids
            documents = [doc for doc in ddl_data["documents"]]
            ids = ddl_data["ids"]

            # Create a DataFrame
            df_ddl = pd.DataFrame(
                {
                    "id": ids,
                    "question": [None for doc in documents],
                    "content": [doc for doc in documents],
                }
            )

            df_ddl["training_data_type"] = "ddl"

            df = pd.concat([df, df_ddl])

        doc_data = self.documentation_collection.get()

        if doc_data is not None:
            # Extract the documents and ids
            documents = [doc for doc in doc_data["documents"]]
            ids = doc_data["ids"]

            # Create a DataFrame
            df_doc = pd.DataFrame(
                {
                    "id": ids,
                    "question": [None for doc in documents],
                    "content": [doc for doc in documents],
                }
            )

            df_doc["training_data_type"] = "documentation"

            df = pd.concat([df, df_doc])

        return df

    def remove_training_data(self, id: str, **kwargs) -> bool:
        if id.endswith("-sql"):
            self.sql_collection.delete(ids=id)
            return True
        elif id.endswith("-ddl"):
            self.ddl_collection.delete(ids=id)
            return True
        elif id.endswith("-doc"):
            self.documentation_collection.delete(ids=id)
            return True
        else:
            return False

    def remove_collection(self, collection_name: str) -> bool:
        """
        This function can reset the collection to empty state.

        Args:
            collection_name (str): sql or ddl or documentation

        Returns:
            bool: True if collection is deleted, False otherwise
        """
        if collection_name == "sql":
            self.chroma_client.delete_collection(name="sql")
            self.sql_collection = self.chroma_client.get_or_create_collection(
                name="sql", embedding_function=self.embedding_function
            )
            return True
        elif collection_name == "ddl":
            self.chroma_client.delete_collection(name="ddl")
            self.ddl_collection = self.chroma_client.get_or_create_collection(
                name="ddl", embedding_function=self.embedding_function
            )
            return True
        elif collection_name == "documentation":
            self.chroma_client.delete_collection(name="documentation")
            self.documentation_collection = self.chroma_client.get_or_create_collection(
                name="documentation", embedding_function=self.embedding_function
            )
            return True
        else:
            return False

    @staticmethod
    def _extract_documents(query_results) -> list:
        """
        Static method to extract the documents from the results of a query.

        Args:
            query_results (pd.DataFrame): The dataframe to use.

        Returns:
            List[str] or None: The extracted documents, or an empty list or
            single document if an error occurred.
        """
        if query_results is None:
            return []

        if "documents" in query_results:
            documents = query_results["documents"]

            if len(documents) == 1 and isinstance(documents[0], list):
                try:
                    documents = [json.loads(doc) for doc in documents[0]]
                except Exception as e:
                    return documents[0]

            return documents

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        return ChromaDB_VectorStore._extract_documents(
            self.sql_collection.query(
                query_texts=[question],
                n_results=self.n_results_sql,
            )
        )

    def get_related_ddl(self, question: str, **kwargs) -> list:
        return ChromaDB_VectorStore._extract_documents(
            self.ddl_collection.query(
                query_texts=[question],
                n_results=self.n_results_ddl,
            )
        )

    def get_related_documentation(self, question: str, **kwargs) -> list:
        return ChromaDB_VectorStore._extract_documents(
            self.documentation_collection.query(
                query_texts=[question],
                n_results=self.n_results_documentation,
            )
        )


================================================
FILE: src/vanna/legacy/cohere/__init__.py
================================================
from .cohere_chat import Cohere_Chat
from .cohere_embeddings import Cohere_Embeddings


================================================
FILE: src/vanna/legacy/cohere/cohere_chat.py
================================================
import os

from openai import OpenAI

from ..base import VannaBase


class Cohere_Chat(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        # default parameters - can be overridden using config
        self.temperature = 0.2  # Lower temperature for more precise SQL generation
        self.model = "command-a-03-2025"  # Cohere's default model

        if config is not None:
            if "temperature" in config:
                self.temperature = config["temperature"]
            if "model" in config:
                self.model = config["model"]

        if client is not None:
            self.client = client
            return

        # Check for API key in environment variable
        api_key = os.getenv("COHERE_API_KEY")

        # Check for API key in config
        if config is not None and "api_key" in config:
            api_key = config["api_key"]

        # Validate API key
        if not api_key:
            raise ValueError(
                "Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable."
            )

        # Initialize client with validated API key
        self.client = OpenAI(
            base_url="https://api.cohere.ai/compatibility/v1",
            api_key=api_key,
        )

    def system_message(self, message: str) -> any:
        return {
            "role": "developer",
            "content": message,
        }  # Cohere uses 'developer' for system role

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        # Count the number of tokens in the message log
        # Use 4 as an approximation for the number of characters per token
        num_tokens = 0
        for message in prompt:
            num_tokens += len(message["content"]) / 4

        # Use model from kwargs, config, or default
        model = kwargs.get("model", self.model)
        if self.config is not None and "model" in self.config and model == self.model:
            model = self.config["model"]

        print(f"Using model {model} for {num_tokens} tokens (approx)")
        try:
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                temperature=self.temperature,
            )

            # Check if response has expected structure
            if not response or not hasattr(response, "choices") or not response.choices:
                raise ValueError("Received empty or malformed response from API")

            if not response.choices[0] or not hasattr(response.choices[0], "message"):
                raise ValueError("Response is missing expected 'message' field")

            if not hasattr(response.choices[0].message, "content"):
                raise ValueError("Response message is missing expected 'content' field")

            return response.choices[0].message.content

        except Exception as e:
            # Log the error and raise a more informative exception
            error_msg = f"Error processing Cohere chat response: {str(e)}"
            print(error_msg)
            raise Exception(error_msg)


================================================
FILE: src/vanna/legacy/cohere/cohere_embeddings.py
================================================
import os

from openai import OpenAI

from ..base import VannaBase


class Cohere_Embeddings(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        # Default embedding model
        self.model = "embed-multilingual-v3.0"

        if config is not None and "model" in config:
            self.model = config["model"]

        if client is not None:
            self.client = client
            return

        # Check for API key in environment variable
        api_key = os.getenv("COHERE_API_KEY")

        # Check for API key in config
        if config is not None and "api_key" in config:
            api_key = config["api_key"]

        # Validate API key
        if not api_key:
            raise ValueError(
                "Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable."
            )

        # Initialize client with validated API key
        self.client = OpenAI(
            base_url="https://api.cohere.ai/compatibility/v1",
            api_key=api_key,
        )

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        if not data:
            raise ValueError("Cannot generate embedding for empty input data")

        # Use model from kwargs, config, or default
        model = kwargs.get("model", self.model)
        if self.config is not None and "model" in self.config and model == self.model:
            model = self.config["model"]

        try:
            embedding = self.client.embeddings.create(
                model=model,
                input=data,
                encoding_format="float",  # Ensure we get float values
            )

            # Check if response has expected structure
            if not embedding or not hasattr(embedding, "data") or not embedding.data:
                raise ValueError(
                    "Received empty or malformed embedding response from API"
                )

            if not embedding.data[0] or not hasattr(embedding.data[0], "embedding"):
                raise ValueError(
                    "Embedding response is missing expected 'embedding' field"
                )

            if not embedding.data[0].embedding:
                raise ValueError("Received empty embedding vector")

            return embedding.data[0].embedding

        except Exception as e:
            # Log the error and raise a more informative exception
            error_msg = f"Error generating embedding with Cohere: {str(e)}"
            print(error_msg)
            raise Exception(error_msg)


================================================
FILE: src/vanna/legacy/deepseek/__init__.py
================================================
from .deepseek_chat import DeepSeekChat


================================================
FILE: src/vanna/legacy/deepseek/deepseek_chat.py
================================================
import os

from openai import OpenAI

from ..base import VannaBase


# from vanna.chromadb import ChromaDB_VectorStore

# class DeepSeekVanna(ChromaDB_VectorStore, DeepSeekChat):
#     def __init__(self, config=None):
#         ChromaDB_VectorStore.__init__(self, config=config)
#         DeepSeekChat.__init__(self, config=config)

# vn = DeepSeekVanna(config={"api_key": "sk-************", "model": "deepseek-chat"})


class DeepSeekChat(VannaBase):
    def __init__(self, config=None):
        if config is None:
            raise ValueError(
                "For DeepSeek, config must be provided with an api_key and model"
            )
        if "api_key" not in config:
            raise ValueError("config must contain a DeepSeek api_key")

        if "model" not in config:
            raise ValueError("config must contain a DeepSeek model")

        api_key = config["api_key"]
        model = config["model"]
        self.model = model
        self.client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com/v1")

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def generate_sql(self, question: str, **kwargs) -> str:
        # 使用父类的 generate_sql
        sql = super().generate_sql(question, **kwargs)

        # 替换 "\_" 为 "_"
        sql = sql.replace("\\_", "_")

        return sql

    def submit_prompt(self, prompt, **kwargs) -> str:
        chat_response = self.client.chat.completions.create(
            model=self.model,
            messages=prompt,
        )

        return chat_response.choices[0].message.content


================================================
FILE: src/vanna/legacy/exceptions/__init__.py
================================================
class ImproperlyConfigured(Exception):
    """Raise for incorrect configuration."""

    pass


class DependencyError(Exception):
    """Raise for missing dependencies."""

    pass


class ConnectionError(Exception):
    """Raise for connection"""

    pass


class OTPCodeError(Exception):
    """Raise for invalid otp or not able to send it"""

    pass


class SQLRemoveError(Exception):
    """Raise when not able to remove SQL"""

    pass


class ExecutionError(Exception):
    """Raise when not able to execute Code"""

    pass


class ValidationError(Exception):
    """Raise for validations"""

    pass


class APIError(Exception):
    """Raise for API errors"""

    pass


================================================
FILE: src/vanna/legacy/faiss/__init__.py
================================================
from .faiss import FAISS


================================================
FILE: src/vanna/legacy/faiss/faiss.py
================================================
import os
import json
import uuid
from typing import List, Dict, Any

import faiss
import numpy as np
import pandas as pd

from ..base import VannaBase
from ..exceptions import DependencyError


class FAISS(VannaBase):
    def __init__(self, config=None):
        if config is None:
            config = {}

        VannaBase.__init__(self, config=config)

        try:
            import faiss
        except ImportError:
            raise DependencyError(
                "FAISS is not installed. Please install it with 'pip install faiss-cpu' or 'pip install faiss-gpu'"
            )

        try:
            from sentence_transformers import SentenceTransformer
        except ImportError:
            raise DependencyError(
                "SentenceTransformer is not installed. Please install it with 'pip install sentence-transformers'."
            )

        self.path = config.get("path", ".")
        self.embedding_dim = config.get("embedding_dim", 384)
        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
        self.n_results_documentation = config.get(
            "n_results_documentation", config.get("n_results", 10)
        )
        self.curr_client = config.get("client", "persistent")

        if self.curr_client == "persistent":
            self.sql_index = self._load_or_create_index("sql_index.faiss")
            self.ddl_index = self._load_or_create_index("ddl_index.faiss")
            self.doc_index = self._load_or_create_index("doc_index.faiss")
        elif self.curr_client == "in-memory":
            self.sql_index = faiss.IndexFlatL2(self.embedding_dim)
            self.ddl_index = faiss.IndexFlatL2(self.embedding_dim)
            self.doc_index = faiss.IndexFlatL2(self.embedding_dim)
        elif (
            isinstance(self.curr_client, list)
            and len(self.curr_client) == 3
            and all(isinstance(idx, faiss.Index) for idx in self.curr_client)
        ):
            self.sql_index = self.curr_client[0]
            self.ddl_index = self.curr_client[1]
            self.doc_index = self.curr_client[2]
        else:
            raise ValueError(
                f"Unsupported storage type was set in config: {self.curr_client}"
            )

        self.sql_metadata: List[Dict[str, Any]] = self._load_or_create_metadata(
            "sql_metadata.json"
        )
        self.ddl_metadata: List[Dict[str, str]] = self._load_or_create_metadata(
            "ddl_metadata.json"
        )
        self.doc_metadata: List[Dict[str, str]] = self._load_or_create_metadata(
            "doc_metadata.json"
        )

        model_name = config.get("embedding_model", "all-MiniLM-L6-v2")
        self.embedding_model = SentenceTransformer(model_name)

    def _load_or_create_index(self, filename):
        filepath = os.path.join(self.path, filename)
        if os.path.exists(filepath):
            return faiss.read_index(filepath)
        return faiss.IndexFlatL2(self.embedding_dim)

    def _load_or_create_metadata(self, filename):
        filepath = os.path.join(self.path, filename)
        if os.path.exists(filepath):
            with open(filepath, "r") as f:
                return json.load(f)
        return []

    def _save_index(self, index, filename):
        if self.curr_client == "persistent":
            filepath = os.path.join(self.path, filename)
            faiss.write_index(index, filepath)

    def _save_metadata(self, metadata, filename):
        if self.curr_client == "persistent":
            filepath = os.path.join(self.path, filename)
            with open(filepath, "w") as f:
                json.dump(metadata, f)

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embedding = self.embedding_model.encode(data)
        assert embedding.shape[0] == self.embedding_dim, (
            f"Embedding dimension mismatch: expected {self.embedding_dim}, got {embedding.shape[0]}"
        )
        return embedding.tolist()

    def _add_to_index(self, index, metadata_list, text, extra_metadata=None) -> str:
        embedding = self.generate_embedding(text)
        index.add(np.array([embedding], dtype=np.float32))
        entry_id = str(uuid.uuid4())
        metadata_list.append({"id": entry_id, **(extra_metadata or {})})
        return entry_id

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        entry_id = self._add_to_index(
            self.sql_index,
            self.sql_metadata,
            question + " " + sql,
            {"question": question, "sql": sql},
        )
        self._save_index(self.sql_index, "sql_index.faiss")
        self._save_metadata(self.sql_metadata, "sql_metadata.json")
        return entry_id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        entry_id = self._add_to_index(
            self.ddl_index, self.ddl_metadata, ddl, {"ddl": ddl}
        )
        self._save_index(self.ddl_index, "ddl_index.faiss")
        self._save_metadata(self.ddl_metadata, "ddl_metadata.json")
        return entry_id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        entry_id = self._add_to_index(
            self.doc_index,
            self.doc_metadata,
            documentation,
            {"documentation": documentation},
        )
        self._save_index(self.doc_index, "doc_index.faiss")
        self._save_metadata(self.doc_metadata, "doc_metadata.json")
        return entry_id

    def _get_similar(self, index, metadata_list, text, n_results) -> list:
        embedding = self.generate_embedding(text)
        D, I = index.search(np.array([embedding], dtype=np.float32), k=n_results)
        return (
            [] if len(I[0]) == 0 or I[0][0] == -1 else [metadata_list[i] for i in I[0]]
        )

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        return self._get_similar(
            self.sql_index, self.sql_metadata, question, self.n_results_sql
        )

    def get_related_ddl(self, question: str, **kwargs) -> list:
        return [
            metadata["ddl"]
            for metadata in self._get_similar(
                self.ddl_index, self.ddl_metadata, question, self.n_results_ddl
            )
        ]

    def get_related_documentation(self, question: str, **kwargs) -> list:
        return [
            metadata["documentation"]
            for metadata in self._get_similar(
                self.doc_index,
                self.doc_metadata,
                question,
                self.n_results_documentation,
            )
        ]

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        sql_data = pd.DataFrame(self.sql_metadata)
        sql_data["training_data_type"] = "sql"

        ddl_data = pd.DataFrame(self.ddl_metadata)
        ddl_data["training_data_type"] = "ddl"

        doc_data = pd.DataFrame(self.doc_metadata)
        doc_data["training_data_type"] = "documentation"

        return pd.concat([sql_data, ddl_data, doc_data], ignore_index=True)

    def remove_training_data(self, id: str, **kwargs) -> bool:
        for metadata_list, index, index_name in [
            (self.sql_metadata, self.sql_index, "sql_index.faiss"),
            (self.ddl_metadata, self.ddl_index, "ddl_index.faiss"),
            (self.doc_metadata, self.doc_index, "doc_index.faiss"),
        ]:
            for i, item in enumerate(metadata_list):
                if item["id"] == id:
                    del metadata_list[i]
                    new_index = faiss.IndexFlatL2(self.embedding_dim)
                    embeddings = [
                        self.generate_embedding(json.dumps(m)) for m in metadata_list
                    ]
                    if embeddings:
                        new_index.add(np.array(embeddings, dtype=np.float32))
                    setattr(self, index_name.split(".")[0], new_index)

                    if self.curr_client == "persistent":
                        self._save_index(new_index, index_name)
                        self._save_metadata(
                            metadata_list, f"{index_name.split('.')[0]}_metadata.json"
                        )

                    return True
        return False

    def remove_collection(self, collection_name: str) -> bool:
        if collection_name in ["sql", "ddl", "documentation"]:
            setattr(
                self, f"{collection_name}_index", faiss.IndexFlatL2(self.embedding_dim)
            )
            setattr(self, f"{collection_name}_metadata", [])

            if self.curr_client == "persistent":
                self._save_index(
                    getattr(self, f"{collection_name}_index"),
                    f"{collection_name}_index.faiss",
                )
                self._save_metadata([], f"{collection_name}_metadata.json")

            return True
        return False


================================================
FILE: src/vanna/legacy/flask/__init__.py
================================================
import json
import logging
import os
import sys
import uuid
from abc import ABC, abstractmethod
from functools import wraps
import importlib.metadata

import flask
import requests
from flasgger import Swagger
from flask import Flask, Response, jsonify, request, send_from_directory
from flask_sock import Sock

from ..base import VannaBase
from .assets import css_content, html_content, js_content
from .auth import AuthInterface, NoAuth


class Cache(ABC):
    """
    Define the interface for a cache that can be used to store data in a Flask app.
    """

    @abstractmethod
    def generate_id(self, *args, **kwargs):
        """
        Generate a unique ID for the cache.
        """
        pass

    @abstractmethod
    def get(self, id, field):
        """
        Get a value from the cache.
        """
        pass

    @abstractmethod
    def get_all(self, field_list) -> list:
        """
        Get all values from the cache.
        """
        pass

    @abstractmethod
    def set(self, id, field, value):
        """
        Set a value in the cache.
        """
        pass

    @abstractmethod
    def delete(self, id):
        """
        Delete a value from the cache.
        """
        pass


class MemoryCache(Cache):
    def __init__(self):
        self.cache = {}

    def generate_id(self, *args, **kwargs):
        return str(uuid.uuid4())

    def set(self, id, field, value):
        if id not in self.cache:
            self.cache[id] = {}

        self.cache[id][field] = value

    def get(self, id, field):
        if id not in self.cache:
            return None

        if field not in self.cache[id]:
            return None

        return self.cache[id][field]

    def get_all(self, field_list) -> list:
        return [
            {"id": id, **{field: self.get(id=id, field=field) for field in field_list}}
            for id in self.cache
        ]

    def delete(self, id):
        if id in self.cache:
            del self.cache[id]


class VannaFlaskAPI:
    flask_app = None

    def requires_cache(self, required_fields, optional_fields=[]):
        def decorator(f):
            @wraps(f)
            def decorated(*args, **kwargs):
                id = request.args.get("id")

                if id is None:
                    id = request.json.get("id")
                    if id is None:
                        return jsonify({"type": "error", "error": "No id provided"})

                for field in required_fields:
                    if self.cache.get(id=id, field=field) is None:
                        return jsonify({"type": "error", "error": f"No {field} found"})

                field_values = {
                    field: self.cache.get(id=id, field=field)
                    for field in required_fields
                }

                for field in optional_fields:
                    field_values[field] = self.cache.get(id=id, field=field)

                # Add the id to the field_values
                field_values["id"] = id

                return f(*args, **field_values, **kwargs)

            return decorated

        return decorator

    def requires_auth(self, f):
        @wraps(f)
        def decorated(*args, **kwargs):
            user = self.auth.get_user(flask.request)

            if not self.auth.is_logged_in(user):
                return jsonify(
                    {"type": "not_logged_in", "html": self.auth.login_form()}
                )

            # Pass the user to the function
            return f(*args, user=user, **kwargs)

        return decorated

    def __init__(
        self,
        vn: VannaBase,
        cache: Cache = MemoryCache(),
        auth: AuthInterface = NoAuth(),
        debug=True,
        allow_llm_to_see_data=False,
        chart=True,
    ):
        """
        Expose a Flask API that can be used to interact with a Vanna instance.

        Args:
            vn: The Vanna instance to interact with.
            cache: The cache to use. Defaults to MemoryCache, which uses an in-memory cache. You can also pass in a custom cache that implements the Cache interface.
            auth: The authentication method to use. Defaults to NoAuth, which doesn't require authentication. You can also pass in a custom authentication method that implements the AuthInterface interface.
            debug: Show the debug console. Defaults to True.
            allow_llm_to_see_data: Whether to allow the LLM to see data. Defaults to False.
            chart: Whether to show the chart output in the UI. Defaults to True.

        Returns:
            None
        """

        self.flask_app = Flask(__name__)

        self.swagger = Swagger(
            self.flask_app, template={"info": {"title": "Vanna API"}}
        )
        self.sock = Sock(self.flask_app)
        self.ws_clients = []
        self.vn = vn
        self.auth = auth
        self.cache = cache
        self.debug = debug
        self.allow_llm_to_see_data = allow_llm_to_see_data
        self.chart = chart
        self.config = {
            "debug": debug,
            "allow_llm_to_see_data": allow_llm_to_see_data,
            "chart": chart,
        }
        log = logging.getLogger("werkzeug")
        log.setLevel(logging.ERROR)

        if "google.colab" in sys.modules:
            self.debug = False
            print(
                "Google Colab doesn't support running websocket servers. Disabling debug mode."
            )

        if self.debug:

            def log(message, title="Info"):
                [
                    ws.send(json.dumps({"message": message, "title": title}))
                    for ws in self.ws_clients
                ]

            self.vn.log = log

        @self.flask_app.route("/api/v0/get_config", methods=["GET"])
        @self.requires_auth
        def get_config(user: any):
            """
            Get the configuration for a user
            ---
            parameters:
              - name: user
                in: query
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: config
                    config:
                      type: object
            """
            config = self.auth.override_config_for_user(user, self.config)
            return jsonify({"type": "config", "config": config})

        @self.flask_app.route("/api/v0/generate_questions", methods=["GET"])
        @self.requires_auth
        def generate_questions(user: any):
            """
            Generate questions
            ---
            parameters:
              - name: user
                in: query
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: question_list
                    questions:
                      type: array
                      items:
                        type: string
                    header:
                      type: string
                      default: Here are some questions you can ask
            """
            # If self has an _model attribute and model=='chinook'
            if hasattr(self.vn, "_model") and self.vn._model == "chinook":
                return jsonify(
                    {
                        "type": "question_list",
                        "questions": [
                            "What are the top 10 artists by sales?",
                            "What are the total sales per year by country?",
                            "Who is the top selling artist in each genre? Show the sales numbers.",
                            "How do the employees rank in terms of sales performance?",
                            "Which 5 cities have the most customers?",
                        ],
                        "header": "Here are some questions you can ask:",
                    }
                )

            training_data = vn.get_training_data()

            # If training data is None or empty
            if training_data is None or len(training_data) == 0:
                return jsonify(
                    {
                        "type": "error",
                        "error": "No training data found. Please add some training data first.",
                    }
                )

            # Get the questions from the training data
            try:
                # Filter training data to only include questions where the question is not null
                questions = (
                    training_data[training_data["question"].notnull()]
                    .sample(5)["question"]
                    .tolist()
                )

                # Temporarily this will just return an empty list
                return jsonify(
                    {
                        "type": "question_list",
                        "questions": questions,
                        "header": "Here are some questions you can ask",
                    }
                )
            except Exception as e:
                return jsonify(
                    {
                        "type": "question_list",
                        "questions": [],
                        "header": "Go ahead and ask a question",
                    }
                )

        @self.flask_app.route("/api/v0/generate_sql", methods=["GET"])
        @self.requires_auth
        def generate_sql(user: any):
            """
            Generate SQL from a question
            ---
            parameters:
              - name: user
                in: query
              - name: question
                in: query
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: sql
                    id:
                      type: string
                    text:
                      type: string
            """
            question = flask.request.args.get("question")

            if question is None:
                return jsonify({"type": "error", "error": "No question provided"})

            id = self.cache.generate_id(question=question)
            sql = vn.generate_sql(
                question=question, allow_llm_to_see_data=self.allow_llm_to_see_data
            )

            self.cache.set(id=id, field="question", value=question)
            self.cache.set(id=id, field="sql", value=sql)

            if vn.is_sql_valid(sql=sql):
                return jsonify(
                    {
                        "type": "sql",
                        "id": id,
                        "text": sql,
                    }
                )
            else:
                return jsonify(
                    {
                        "type": "text",
                        "id": id,
                        "text": sql,
                    }
                )

        @self.flask_app.route("/api/v0/generate_rewritten_question", methods=["GET"])
        @self.requires_auth
        def generate_rewritten_question(user: any):
            """
            Generate a rewritten question
            ---
            parameters:
              - name: last_question
                in: query
                type: string
                required: true
              - name: new_question
                in: query
                type: string
                required: true
            """

            last_question = flask.request.args.get("last_question")
            new_question = flask.request.args.get("new_question")

            rewritten_question = self.vn.generate_rewritten_question(
                last_question, new_question
            )

            return jsonify(
                {"type": "rewritten_question", "question": rewritten_question}
            )

        @self.flask_app.route("/api/v0/get_function", methods=["GET"])
        @self.requires_auth
        def get_function(user: any):
            """
            Get a function from a question
            ---
            parameters:
              - name: user
                in: query
              - name: question
                in: query
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: function
                    id:
                      type: object
                    function:
                      type: string
            """
            question = flask.request.args.get("question")

            if question is None:
                return jsonify({"type": "error", "error": "No question provided"})

            if not hasattr(vn, "get_function"):
                return jsonify(
                    {
                        "type": "error",
                        "error": "This setup does not support function generation.",
                    }
                )

            id = self.cache.generate_id(question=question)
            function = vn.get_function(question=question)

            if function is None:
                return jsonify({"type": "error", "error": "No function found"})

            if "instantiated_sql" not in function:
                self.vn.log(f"No instantiated SQL found for {question} in {function}")
                return jsonify({"type": "error", "error": "No instantiated SQL found"})

            self.cache.set(id=id, field="question", value=question)
            self.cache.set(id=id, field="sql", value=function["instantiated_sql"])

            if (
                "instantiated_post_processing_code" in function
                and function["instantiated_post_processing_code"] is not None
                and len(function["instantiated_post_processing_code"]) > 0
            ):
                self.cache.set(
                    id=id,
                    field="plotly_code",
                    value=function["instantiated_post_processing_code"],
                )

            return jsonify(
                {
                    "type": "function",
                    "id": id,
                    "function": function,
                }
            )

        @self.flask_app.route("/api/v0/get_all_functions", methods=["GET"])
        @self.requires_auth
        def get_all_functions(user: any):
            """
            Get all the functions
            ---
            parameters:
              - name: user
                in: query
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: functions
                    functions:
                      type: array
            """
            if not hasattr(vn, "get_all_functions"):
                return jsonify(
                    {
                        "type": "error",
                        "error": "This setup does not support function generation.",
                    }
                )

            functions = vn.get_all_functions()

            return jsonify(
                {
                    "type": "functions",
                    "functions": functions,
                }
            )

        @self.flask_app.route("/api/v0/run_sql", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(["sql"])
        def run_sql(user: any, id: str, sql: str):
            """
            Run SQL
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: df
                    id:
                      type: string
                    df:
                      type: object
                    should_generate_chart:
                      type: boolean
            """
            try:
                if not vn.run_sql_is_set:
                    return jsonify(
                        {
                            "type": "error",
                            "error": "Please connect to a database using vn.connect_to_... in order to run SQL queries.",
                        }
                    )

                df = vn.run_sql(sql=sql)

                self.cache.set(id=id, field="df", value=df)

                return jsonify(
                    {
                        "type": "df",
                        "id": id,
                        "df": df.head(10).to_json(orient="records", date_format="iso"),
                        "should_generate_chart": self.chart
                        and vn.should_generate_chart(df),
                    }
                )

            except Exception as e:
                return jsonify({"type": "sql_error", "error": str(e)})

        @self.flask_app.route("/api/v0/fix_sql", methods=["POST"])
        @self.requires_auth
        @self.requires_cache(["question", "sql"])
        def fix_sql(user: any, id: str, question: str, sql: str):
            """
            Fix SQL
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
              - name: error
                in: body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: sql
                    id:
                      type: string
                    text:
                      type: string
            """
            error = flask.request.json.get("error")

            if error is None:
                return jsonify({"type": "error", "error": "No error provided"})

            question = f"I have an error: {error}\n\nHere is the SQL I tried to run: {sql}\n\nThis is the question I was trying to answer: {question}\n\nCan you rewrite the SQL to fix the error?"

            fixed_sql = vn.generate_sql(question=question)

            self.cache.set(id=id, field="sql", value=fixed_sql)

            return jsonify(
                {
                    "type": "sql",
                    "id": id,
                    "text": fixed_sql,
                }
            )

        @self.flask_app.route("/api/v0/update_sql", methods=["POST"])
        @self.requires_auth
        @self.requires_cache([])
        def update_sql(user: any, id: str):
            """
            Update SQL
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
              - name: sql
                in: body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: sql
                    id:
                      type: string
                    text:
                      type: string
            """
            sql = flask.request.json.get("sql")

            if sql is None:
                return jsonify({"type": "error", "error": "No sql provided"})

            self.cache.set(id=id, field="sql", value=sql)

            return jsonify(
                {
                    "type": "sql",
                    "id": id,
                    "text": sql,
                }
            )

        @self.flask_app.route("/api/v0/download_csv", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(["df"])
        def download_csv(user: any, id: str, df):
            """
            Download CSV
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
            responses:
              200:
                description: download CSV
            """
            csv = df.to_csv()

            return Response(
                csv,
                mimetype="text/csv",
                headers={"Content-disposition": f"attachment; filename={id}.csv"},
            )

        @self.flask_app.route("/api/v0/generate_plotly_figure", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(["df", "question", "sql"])
        def generate_plotly_figure(user: any, id: str, df, question, sql):
            """
            Generate plotly figure
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
              - name: chart_instructions
                in: body
                type: string
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: plotly_figure
                    id:
                      type: string
                    fig:
                      type: object
            """
            chart_instructions = flask.request.args.get("chart_instructions")

            try:
                # If chart_instructions is not set then attempt to retrieve the code from the cache
                if chart_instructions is None or len(chart_instructions) == 0:
                    code = self.cache.get(id=id, field="plotly_code")
                else:
                    question = f"{question}. When generating the chart, use these special instructions: {chart_instructions}"
                    code = vn.generate_plotly_code(
                        question=question,
                        sql=sql,
                        df_metadata=f"Running df.dtypes gives:\n {df.dtypes}",
                    )
                    self.cache.set(id=id, field="plotly_code", value=code)

                fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
                fig_json = fig.to_json()

                self.cache.set(id=id, field="fig_json", value=fig_json)

                return jsonify(
                    {
                        "type": "plotly_figure",
                        "id": id,
                        "fig": fig_json,
                    }
                )
            except Exception as e:
                # Print the stack trace
                import traceback

                traceback.print_stack()
                traceback.print_exc()

                return jsonify({"type": "error", "error": str(e)})

        @self.flask_app.route("/api/v0/get_training_data", methods=["GET"])
        @self.requires_auth
        def get_training_data(user: any):
            """
            Get all training data
            ---
            parameters:
              - name: user
                in: query
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: df
                    id:
                      type: string
                      default: training_data
                    df:
                      type: object
            """
            df = vn.get_training_data()

            if df is None or len(df) == 0:
                return jsonify(
                    {
                        "type": "error",
                        "error": "No training data found. Please add some training data first.",
                    }
                )

            return jsonify(
                {
                    "type": "df",
                    "id": "training_data",
                    "df": df.to_json(orient="records"),
                }
            )

        @self.flask_app.route("/api/v0/remove_training_data", methods=["POST"])
        @self.requires_auth
        def remove_training_data(user: any):
            """
            Remove training data
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    success:
                      type: boolean
            """
            # Get id from the JSON body
            id = flask.request.json.get("id")

            if id is None:
                return jsonify({"type": "error", "error": "No id provided"})

            if vn.remove_training_data(id=id):
                return jsonify({"success": True})
            else:
                return jsonify(
                    {"type": "error", "error": "Couldn't remove training data"}
                )

        @self.flask_app.route("/api/v0/train", methods=["POST"])
        @self.requires_auth
        def add_training_data(user: any):
            """
            Add training data
            ---
            parameters:
              - name: user
                in: query
              - name: question
                in: body
                type: string
              - name: sql
                in: body
                type: string
              - name: ddl
                in: body
                type: string
              - name: documentation
                in: body
                type: string
            responses:
              200:
                schema:
                  type: object
                  properties:
                    id:
                      type: string
            """
            question = flask.request.json.get("question")
            sql = flask.request.json.get("sql")
            ddl = flask.request.json.get("ddl")
            documentation = flask.request.json.get("documentation")

            try:
                id = vn.train(
                    question=question, sql=sql, ddl=ddl, documentation=documentation
                )

                return jsonify({"id": id})
            except Exception as e:
                print("TRAINING ERROR", e)
                return jsonify({"type": "error", "error": str(e)})

        @self.flask_app.route("/api/v0/create_function", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(["question", "sql"])
        def create_function(user: any, id: str, question: str, sql: str):
            """
            Create function
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: function_template
                    id:
                      type: string
                    function_template:
                      type: object
            """
            plotly_code = self.cache.get(id=id, field="plotly_code")

            if plotly_code is None:
                plotly_code = ""

            function_data = self.vn.create_function(
                question=question, sql=sql, plotly_code=plotly_code
            )

            return jsonify(
                {
                    "type": "function_template",
                    "id": id,
                    "function_template": function_data,
                }
            )

        @self.flask_app.route("/api/v0/update_function", methods=["POST"])
        @self.requires_auth
        def update_function(user: any):
            """
            Update function
            ---
            parameters:
              - name: user
                in: query
              - name: old_function_name
                in: body
                type: string
                required: true
              - name: updated_function
                in: body
                type: object
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    success:
                      type: boolean
            """
            old_function_name = flask.request.json.get("old_function_name")
            updated_function = flask.request.json.get("updated_function")

            print("old_function_name", old_function_name)
            print("updated_function", updated_function)

            updated = vn.update_function(
                old_function_name=old_function_name, updated_function=updated_function
            )

            return jsonify({"success": updated})

        @self.flask_app.route("/api/v0/delete_function", methods=["POST"])
        @self.requires_auth
        def delete_function(user: any):
            """
            Delete function
            ---
            parameters:
              - name: user
                in: query
              - name: function_name
                in: body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    success:
                      type: boolean
            """
            function_name = flask.request.json.get("function_name")

            return jsonify({"success": vn.delete_function(function_name=function_name)})

        @self.flask_app.route("/api/v0/generate_followup_questions", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(["df", "question", "sql"])
        def generate_followup_questions(user: any, id: str, df, question, sql):
            """
            Generate followup questions
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: question_list
                    questions:
                      type: array
                      items:
                        type: string
                    header:
                      type: string
            """
            if self.allow_llm_to_see_data:
                followup_questions = vn.generate_followup_questions(
                    question=question, sql=sql, df=df
                )
                if followup_questions is not None and len(followup_questions) > 5:
                    followup_questions = followup_questions[:5]

                self.cache.set(
                    id=id, field="followup_questions", value=followup_questions
                )

                return jsonify(
                    {
                        "type": "question_list",
                        "id": id,
                        "questions": followup_questions,
                        "header": "Here are some potential followup questions:",
                    }
                )
            else:
                self.cache.set(id=id, field="followup_questions", value=[])
                return jsonify(
                    {
                        "type": "question_list",
                        "id": id,
                        "questions": [],
                        "header": "Followup Questions can be enabled if you set allow_llm_to_see_data=True",
                    }
                )

        @self.flask_app.route("/api/v0/generate_summary", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(["df", "question"])
        def generate_summary(user: any, id: str, df, question):
            """
            Generate summary
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: text
                    id:
                      type: string
                    text:
                      type: string
            """
            if self.allow_llm_to_see_data:
                summary = vn.generate_summary(question=question, df=df)

                self.cache.set(id=id, field="summary", value=summary)

                return jsonify(
                    {
                        "type": "text",
                        "id": id,
                        "text": summary,
                    }
                )
            else:
                return jsonify(
                    {
                        "type": "text",
                        "id": id,
                        "text": "Summarization can be enabled if you set allow_llm_to_see_data=True",
                    }
                )

        @self.flask_app.route("/api/v0/load_question", methods=["GET"])
        @self.requires_auth
        @self.requires_cache(
            ["question", "sql", "df"], optional_fields=["summary", "fig_json"]
        )
        def load_question(user: any, id: str, question, sql, df, fig_json, summary):
            """
            Load question
            ---
            parameters:
              - name: user
                in: query
              - name: id
                in: query|body
                type: string
                required: true
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: question_cache
                    id:
                      type: string
                    question:
                      type: string
                    sql:
                      type: string
                    df:
                      type: object
                    fig:
                      type: object
                    summary:
                      type: string
            """
            try:
                return jsonify(
                    {
                        "type": "question_cache",
                        "id": id,
                        "question": question,
                        "sql": sql,
                        "df": df.head(10).to_json(orient="records", date_format="iso"),
                        "fig": fig_json,
                        "summary": summary,
                    }
                )

            except Exception as e:
                return jsonify({"type": "error", "error": str(e)})

        @self.flask_app.route("/api/v0/get_question_history", methods=["GET"])
        @self.requires_auth
        def get_question_history(user: any):
            """
            Get question history
            ---
            parameters:
              - name: user
                in: query
            responses:
              200:
                schema:
                  type: object
                  properties:
                    type:
                      type: string
                      default: question_history
                    questions:
                      type: array
                      items:
                        type: string
            """
            return jsonify(
                {
                    "type": "question_history",
                    "questions": cache.get_all(field_list=["question"]),
                }
            )

        @self.flask_app.route("/api/v0/<path:catch_all>", methods=["GET", "POST"])
        def catch_all(catch_all):
            return jsonify(
                {"type": "error", "error": "The rest of the API is not ported yet."}
            )

        if self.debug:

            @self.sock.route("/api/v0/log")
            def sock_log(ws):
                self.ws_clients.append(ws)

                try:
                    while True:
                        message = (
                            ws.receive()
                        )  # This example just reads and ignores to keep the socket open
                finally:
                    self.ws_clients.remove(ws)

    def run(self, *args, **kwargs):
        """
        Run the Flask app.

        Args:
            *args: Arguments to pass to Flask's run method.
            **kwargs: Keyword arguments to pass to Flask's run method.

        Returns:
            None
        """
        if args or kwargs:
            self.flask_app.run(*args, **kwargs)

        else:
            try:
                from google.colab import output

                output.serve_kernel_port_as_window(8084)
                from google.colab.output import eval_js

                print("Your app is running at:")
                print(eval_js("google.colab.kernel.proxyPort(8084)"))
            except Exception:
                print("Your app is running at:")
                print("http://localhost:8084")

            self.flask_app.run(
                host="0.0.0.0", port=8084, debug=self.debug, use_reloader=False
            )


class VannaFlaskApp(VannaFlaskAPI):
    def __init__(
        self,
        vn: VannaBase,
        cache: Cache = MemoryCache(),
        auth: AuthInterface = NoAuth(),
        debug=True,
        allow_llm_to_see_data=False,
        logo="https://img.vanna.ai/vanna-flask.svg",
        title="Welcome to Vanna.AI",
        subtitle="Your AI-powered copilot for SQL queries.",
        show_training_data=True,
        suggested_questions=True,
        sql=True,
        table=True,
        csv_download=True,
        chart=True,
        redraw_chart=True,
        auto_fix_sql=True,
        ask_results_correct=True,
        followup_questions=True,
        summarization=True,
        function_generation=True,
        index_html_path=None,
        assets_folder=None,
    ):
        """
        Expose a Flask app that can be used to interact with a Vanna instance.

        Args:
            vn: The Vanna instance to interact with.
            cache: The cache to use. Defaults to MemoryCache, which uses an in-memory cache. You can also pass in a custom cache that implements the Cache interface.
            auth: The authentication method to use. Defaults to NoAuth, which doesn't require authentication. You can also pass in a custom authentication method that implements the AuthInterface interface.
            debug: Show the debug console. Defaults to True.
            allow_llm_to_see_data: Whether to allow the LLM to see data. Defaults to False.
            logo: The logo to display in the UI. Defaults to the Vanna logo.
            title: The title to display in the UI. Defaults to "Welcome to Vanna.AI".
            subtitle: The subtitle to display in the UI. Defaults to "Your AI-powered copilot for SQL queries.".
            show_training_data: Whether to show the training data in the UI. Defaults to True.
            suggested_questions: Whether to show suggested questions in the UI. Defaults to True.
            sql: Whether to show the SQL input in the UI. Defaults to True.
            table: Whether to show the table output in the UI. Defaults to True.
            csv_download: Whether to allow downloading the table output as a CSV file. Defaults to True.
            chart: Whether to show the chart output in the UI. Defaults to True.
            redraw_chart: Whether to allow redrawing the chart. Defaults to True.
            auto_fix_sql: Whether to allow auto-fixing SQL errors. Defaults to True.
            ask_results_correct: Whether to ask the user if the results are correct. Defaults to True.
            followup_questions: Whether to show followup questions. Defaults to True.
            summarization: Whether to show summarization. Defaults to True.
            index_html_path: Path to the index.html. Defaults to None, which will use the default index.html
            assets_folder: The location where you'd like to serve the static assets from. Defaults to None, which will use hardcoded Python variables.

        Returns:
            None
        """
        super().__init__(vn, cache, auth, debug, allow_llm_to_see_data, chart)

        self.config["logo"] = logo
        self.config["title"] = title
        self.config["subtitle"] = subtitle
        self.config["show_training_data"] = show_training_data
        self.config["suggested_questions"] = suggested_questions
        self.config["sql"] = sql
        self.config["table"] = table
        self.config["csv_download"] = csv_download
        self.config["chart"] = chart
        self.config["redraw_chart"] = redraw_chart
        self.config["auto_fix_sql"] = auto_fix_sql
        self.config["ask_results_correct"] = ask_results_correct
        self.config["followup_questions"] = followup_questions
        self.config["summarization"] = summarization
        self.config["function_generation"] = function_generation and hasattr(
            vn, "get_function"
        )
        self.config["version"] = importlib.metadata.version("vanna")

        self.index_html_path = index_html_path
        self.assets_folder = assets_folder

        @self.flask_app.route("/auth/login", methods=["POST"])
        def login():
            return self.auth.login_handler(flask.request)

        @self.flask_app.route("/auth/callback", methods=["GET"])
        def callback():
            return self.auth.callback_handler(flask.request)

        @self.flask_app.route("/auth/logout", methods=["GET"])
        def logout():
            return self.auth.logout_handler(flask.request)

        @self.flask_app.route("/assets/<path:filename>")
        def proxy_assets(filename):
            if self.assets_folder:
                return send_from_directory(self.assets_folder, filename)

            if ".css" in filename:
                return Response(css_content, mimetype="text/css")

            if ".js" in filename:
                return Response(js_content, mimetype="text/javascript")

            # Return 404
            return "File not found", 404

        # Proxy the /vanna.svg file to the remote server
        @self.flask_app.route("/vanna.svg")
        def proxy_vanna_svg():
            remote_url = "https://vanna.ai/img/vanna.svg"
            response = requests.get(remote_url, stream=True)

            # Check if the request to the remote URL was successful
            if response.status_code == 200:
                excluded_headers = [
                    "content-encoding",
                    "content-length",
                    "transfer-encoding",
                    "connection",
                ]
                headers = [
                    (name, value)
                    for (name, value) in response.raw.headers.items()
                    if name.lower() not in excluded_headers
                ]
                return Response(response.content, response.status_code, headers)
            else:
                return "Error fetching file from remote server", response.status_code

        @self.flask_app.route("/", defaults={"path": ""})
        @self.flask_app.route("/<path:path>")
        def hello(path: str):
            if self.index_html_path:
                directory = os.path.dirname(self.index_html_path)
                filename = os.path.basename(self.index_html_path)
                return send_from_directory(directory=directory, path=filename)
            return html_content


================================================
FILE: src/vanna/legacy/flask/assets.py
================================================
html_content = """<!doctype html>
<html lang="en" translate>
  <head>
    <meta charset="UTF-8" />
    <link rel="icon" type="image/svg+xml" href="/vanna.svg" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <link href="https://fonts.googleapis.com/css2?family=Roboto+Slab:wght@350&display=swap" rel="stylesheet">
    <script src="https://cdn.plot.ly/plotly-latest.min.js" type="text/javascript"></script>
    <title>Vanna.AI</title>
    <script type="module" crossorigin src="/assets/index-35bab439.js"></script>
    <link rel="stylesheet" href="/assets/index-f228f78f.css">
  </head>
  <body class="bg-white dark:bg-slate-900">
    <div id="app"></div>
    
  </body>
</html>
"""

css_content = """.nav-title{font-family:Roboto Slab,serif}*,:before,:after{box-sizing:border-box;border-width:0;border-style:solid;border-color:#e5e7eb}:before,:after{--tw-content: ""}html{line-height:1.5;-webkit-text-size-adjust:100%;-moz-tab-size:4;-o-tab-size:4;tab-size:4;font-family:ui-sans-serif,system-ui,-apple-system,BlinkMacSystemFont,Segoe UI,Roboto,Helvetica Neue,Arial,Noto Sans,sans-serif,"Apple Color Emoji","Segoe UI Emoji",Segoe UI Symbol,"Noto Color Emoji";font-feature-settings:normal;font-variation-settings:normal}body{margin:0;line-height:inherit}hr{height:0;color:inherit;border-top-width:1px}abbr:where([title]){-webkit-text-decoration:underline dotted;text-decoration:underline dotted}h1,h2,h3,h4,h5,h6{font-size:inherit;font-weight:inherit}a{color:inherit;text-decoration:inherit}b,strong{font-weight:bolder}code,kbd,samp,pre{font-family:ui-monospace,SFMono-Regular,Menlo,Monaco,Consolas,Liberation Mono,Courier New,monospace;font-size:1em}small{font-size:80%}sub,sup{font-size:75%;line-height:0;position:relative;vertical-align:baseline}sub{bottom:-.25em}sup{top:-.5em}table{text-indent:0;border-color:inherit;border-collapse:collapse}button,input,optgroup,select,textarea{font-family:inherit;font-feature-settings:inherit;font-variation-settings:inherit;font-size:100%;font-weight:inherit;line-height:inherit;color:inherit;margin:0;padding:0}button,select{text-transform:none}button,[type=button],[type=reset],[type=submit]{-webkit-appearance:button;background-color:transparent;background-image:none}:-moz-focusring{outline:auto}:-moz-ui-invalid{box-shadow:none}progress{vertical-align:baseline}::-webkit-inner-spin-button,::-webkit-outer-spin-button{height:auto}[type=search]{-webkit-appearance:textfield;outline-offset:-2px}::-webkit-search-decoration{-webkit-appearance:none}::-webkit-file-upload-button{-webkit-appearance:button;font:inherit}summary{display:list-item}blockquote,dl,dd,h1,h2,h3,h4,h5,h6,hr,figure,p,pre{margin:0}fieldset{margin:0;padding:0}legend{padding:0}ol,ul,menu{list-style:none;margin:0;padding:0}dialog{padding:0}textarea{resize:vertical}input::-moz-placeholder,textarea::-moz-placeholder{opacity:1;color:#9ca3af}input::placeholder,textarea::placeholder{opacity:1;color:#9ca3af}button,[role=button]{cursor:pointer}:disabled{cursor:default}img,svg,video,canvas,audio,iframe,embed,object{display:block;vertical-align:middle}img,video{max-width:100%;height:auto}[hidden]{display:none}*,:before,:after{--tw-border-spacing-x: 0;--tw-border-spacing-y: 0;--tw-translate-x: 0;--tw-translate-y: 0;--tw-rotate: 0;--tw-skew-x: 0;--tw-skew-y: 0;--tw-scale-x: 1;--tw-scale-y: 1;--tw-pan-x: ;--tw-pan-y: ;--tw-pinch-zoom: ;--tw-scroll-snap-strictness: proximity;--tw-gradient-from-position: ;--tw-gradient-via-position: ;--tw-gradient-to-position: ;--tw-ordinal: ;--tw-slashed-zero: ;--tw-numeric-figure: ;--tw-numeric-spacing: ;--tw-numeric-fraction: ;--tw-ring-inset: ;--tw-ring-offset-width: 0px;--tw-ring-offset-color: #fff;--tw-ring-color: rgb(59 130 246 / .5);--tw-ring-offset-shadow: 0 0 #0000;--tw-ring-shadow: 0 0 #0000;--tw-shadow: 0 0 #0000;--tw-shadow-colored: 0 0 #0000;--tw-blur: ;--tw-brightness: ;--tw-contrast: ;--tw-grayscale: ;--tw-hue-rotate: ;--tw-invert: ;--tw-saturate: ;--tw-sepia: ;--tw-drop-shadow: ;--tw-backdrop-blur: ;--tw-backdrop-brightness: ;--tw-backdrop-contrast: ;--tw-backdrop-grayscale: ;--tw-backdrop-hue-rotate: ;--tw-backdrop-invert: ;--tw-backdrop-opacity: ;--tw-backdrop-saturate: ;--tw-backdrop-sepia: }::backdrop{--tw-border-spacing-x: 0;--tw-border-spacing-y: 0;--tw-translate-x: 0;--tw-translate-y: 0;--tw-rotate: 0;--tw-skew-x: 0;--tw-skew-y: 0;--tw-scale-x: 1;--tw-scale-y: 1;--tw-pan-x: ;--tw-pan-y: ;--tw-pinch-zoom: ;--tw-scroll-snap-strictness: proximity;--tw-gradient-from-position: ;--tw-gradient-via-position: ;--tw-gradient-to-position: ;--tw-ordinal: ;--tw-slashed-zero: ;--tw-numeric-figure: ;--tw-numeric-spacing: ;--tw-numeric-fraction: ;--tw-ring-inset: ;--tw-ring-offset-width: 0px;--tw-ring-offset-color: #fff;--tw-ring-color: rgb(59 130 246 / .5);--tw-ring-offset-shadow: 0 0 #0000;--tw-ring-shadow: 0 0 #0000;--tw-shadow: 0 0 #0000;--tw-shadow-colored: 0 0 #0000;--tw-blur: ;--tw-brightness: ;--tw-contrast: ;--tw-grayscale: ;--tw-hue-rotate: ;--tw-invert: ;--tw-saturate: ;--tw-sepia: ;--tw-drop-shadow: ;--tw-backdrop-blur: ;--tw-backdrop-brightness: ;--tw-backdrop-contrast: ;--tw-backdrop-grayscale: ;--tw-backdrop-hue-rotate: ;--tw-backdrop-invert: ;--tw-backdrop-opacity: ;--tw-backdrop-saturate: ;--tw-backdrop-sepia: }.sr-only{position:absolute;width:1px;height:1px;padding:0;margin:-1px;overflow:hidden;clip:rect(0,0,0,0);white-space:nowrap;border-width:0}.pointer-events-none{pointer-events:none}.collapse{visibility:collapse}.static{position:static}.fixed{position:fixed}.absolute{position:absolute}.relative{position:relative}.sticky{position:sticky}.inset-0{top:0;right:0;bottom:0;left:0}.inset-x-px{left:1px;right:1px}.inset-y-0{top:0;bottom:0}.bottom-0{bottom:0}.bottom-px{bottom:1px}.end-0{inset-inline-end:0px}.left-0{left:0}.right-0{right:0}.top-0{top:0}.z-10{z-index:10}.z-50{z-index:50}.z-\\[60\\]{z-index:60}.z-\\[80\\]{z-index:80}.-m-1{margin:-.25rem}.-m-1\\.5{margin:-.375rem}.m-1{margin:.25rem}.m-3{margin:.75rem}.mx-auto{margin-left:auto;margin-right:auto}.mb-1{margin-bottom:.25rem}.mb-2{margin-bottom:.5rem}.mb-2\\.5{margin-bottom:.625rem}.mb-3{margin-bottom:.75rem}.mb-auto{margin-bottom:auto}.ml-3{margin-left:.75rem}.ml-4{margin-left:1rem}.mr-1{margin-right:.25rem}.mr-1\\.5{margin-right:.375rem}.mr-3{margin-right:.75rem}.ms-0{margin-inline-start:0px}.ms-3{margin-inline-start:.75rem}.mt-0{margin-top:0}.mt-0\\.5{margin-top:.125rem}.mt-1{margin-top:.25rem}.mt-16{margin-top:4rem}.mt-2{margin-top:.5rem}.mt-2\\.5{margin-top:.625rem}.mt-3{margin-top:.75rem}.mt-4{margin-top:1rem}.mt-5{margin-top:1.25rem}.mt-6{margin-top:1.5rem}.mt-7{margin-top:1.75rem}.mt-auto{margin-top:auto}.block{display:block}.inline-block{display:inline-block}.inline{display:inline}.flex{display:flex}.inline-flex{display:inline-flex}.table{display:table}.grid{display:grid}.hidden{display:none}.h-1{height:.25rem}.h-1\\.5{height:.375rem}.h-2{height:.5rem}.h-3{height:.75rem}.h-3\\.5{height:.875rem}.h-4{height:1rem}.h-5{height:1.25rem}.h-52{height:13rem}.h-6{height:1.5rem}.h-7{height:1.75rem}.h-8{height:2rem}.h-\\[2\\.375rem\\]{height:2.375rem}.h-auto{height:auto}.h-full{height:100%}.h-px{height:1px}.h-screen{height:100vh}.min-h-\\[15rem\\]{min-height:15rem}.min-h-\\[calc\\(100\\%-3\\.5rem\\)\\]{min-height:calc(100% - 3.5rem)}.w-0{width:0px}.w-1{width:.25rem}.w-1\\.5{width:.375rem}.w-2{width:.5rem}.w-28{width:7rem}.w-3{width:.75rem}.w-3\\.5{width:.875rem}.w-4{width:1rem}.w-6{width:1.5rem}.w-64{width:16rem}.w-8{width:2rem}.w-\\[2\\.375rem\\]{width:2.375rem}.w-\\[3\\.25rem\\]{width:3.25rem}.w-full{width:100%}.w-px{width:1px}.min-w-full{min-width:100%}.max-w-2xl{max-width:42rem}.max-w-4xl{max-width:56rem}.max-w-7xl{max-width:80rem}.max-w-\\[50rem\\]{max-width:50rem}.max-w-\\[85rem\\]{max-width:85rem}.max-w-fit{max-width:-moz-fit-content;max-width:fit-content}.max-w-md{max-width:28rem}.max-w-sm{max-width:24rem}.max-w-xs{max-width:20rem}.flex-1{flex:1 1 0%}.flex-auto{flex:1 1 auto}.flex-none{flex:none}.flex-shrink-0,.shrink-0{flex-shrink:0}.flex-grow,.grow{flex-grow:1}.-translate-x-full{--tw-translate-x: -100%;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.translate-x-0{--tw-translate-x: 0px;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.translate-x-full{--tw-translate-x: 100%;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.transform{transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}@keyframes bounce{0%,to{transform:translateY(-25%);animation-timing-function:cubic-bezier(.8,0,1,1)}50%{transform:none;animation-timing-function:cubic-bezier(0,0,.2,1)}}.animate-bounce{animation:bounce 1s infinite}@keyframes ping{75%,to{transform:scale(2);opacity:0}}.animate-ping{animation:ping 1s cubic-bezier(0,0,.2,1) infinite}@keyframes pulse{50%{opacity:.5}}.animate-pulse{animation:pulse 2s cubic-bezier(.4,0,.6,1) infinite}@keyframes spin{to{transform:rotate(360deg)}}.animate-spin{animation:spin 1s linear infinite}.cursor-pointer{cursor:pointer}.resize{resize:both}.list-disc{list-style-type:disc}.appearance-none{-webkit-appearance:none;-moz-appearance:none;appearance:none}.grid-cols-4{grid-template-columns:repeat(4,minmax(0,1fr))}.flex-col{flex-direction:column}.items-start{align-items:flex-start}.items-center{align-items:center}.justify-end{justify-content:flex-end}.justify-center{justify-content:center}.justify-between{justify-content:space-between}.gap-1{gap:.25rem}.gap-1\\.5{gap:.375rem}.gap-2{gap:.5rem}.gap-3{gap:.75rem}.gap-4{gap:1rem}.gap-x-1{-moz-column-gap:.25rem;column-gap:.25rem}.gap-x-2{-moz-column-gap:.5rem;column-gap:.5rem}.gap-x-3{-moz-column-gap:.75rem;column-gap:.75rem}.gap-y-4{row-gap:1rem}.-space-y-px>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(-1px * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(-1px * var(--tw-space-y-reverse))}.space-x-2>:not([hidden])~:not([hidden]){--tw-space-x-reverse: 0;margin-right:calc(.5rem * var(--tw-space-x-reverse));margin-left:calc(.5rem * calc(1 - var(--tw-space-x-reverse)))}.space-y-1>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(.25rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(.25rem * var(--tw-space-y-reverse))}.space-y-1\\.5>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(.375rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(.375rem * var(--tw-space-y-reverse))}.space-y-2>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(.5rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(.5rem * var(--tw-space-y-reverse))}.space-y-3>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(.75rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(.75rem * var(--tw-space-y-reverse))}.space-y-4>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(1rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(1rem * var(--tw-space-y-reverse))}.space-y-5>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(1.25rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(1.25rem * var(--tw-space-y-reverse))}.divide-x>:not([hidden])~:not([hidden]){--tw-divide-x-reverse: 0;border-right-width:calc(1px * var(--tw-divide-x-reverse));border-left-width:calc(1px * calc(1 - var(--tw-divide-x-reverse)))}.divide-y>:not([hidden])~:not([hidden]){--tw-divide-y-reverse: 0;border-top-width:calc(1px * calc(1 - var(--tw-divide-y-reverse)));border-bottom-width:calc(1px * var(--tw-divide-y-reverse))}.divide-gray-200>:not([hidden])~:not([hidden]){--tw-divide-opacity: 1;border-color:rgb(229 231 235 / var(--tw-divide-opacity))}.overflow-hidden{overflow:hidden}.overflow-x-auto{overflow-x:auto}.overflow-y-auto{overflow-y:auto}.overflow-x-hidden{overflow-x:hidden}.overflow-y-hidden{overflow-y:hidden}.truncate{overflow:hidden;text-overflow:ellipsis;white-space:nowrap}.whitespace-nowrap{white-space:nowrap}.whitespace-pre-wrap{white-space:pre-wrap}.whitespace-break-spaces{white-space:break-spaces}.rounded{border-radius:.25rem}.rounded-full{border-radius:9999px}.rounded-lg{border-radius:.5rem}.rounded-md{border-radius:.375rem}.rounded-xl{border-radius:.75rem}.rounded-b-md{border-bottom-right-radius:.375rem;border-bottom-left-radius:.375rem}.rounded-t-xl{border-top-left-radius:.75rem;border-top-right-radius:.75rem}.rounded-ee-xl{border-end-end-radius:.75rem}.rounded-es-xl{border-end-start-radius:.75rem}.border{border-width:1px}.border-2{border-width:2px}.border-4{border-width:4px}.border-\\[3px\\]{border-width:3px}.border-b{border-bottom-width:1px}.border-b-2{border-bottom-width:2px}.border-r{border-right-width:1px}.border-s{border-inline-start-width:1px}.border-t{border-top-width:1px}.border-t-2{border-top-width:2px}.border-t-4{border-top-width:4px}.border-blue-500{--tw-border-opacity: 1;border-color:rgb(59 130 246 / var(--tw-border-opacity))}.border-blue-600{--tw-border-opacity: 1;border-color:rgb(37 99 235 / var(--tw-border-opacity))}.border-current{border-color:currentColor}.border-gray-200{--tw-border-opacity: 1;border-color:rgb(229 231 235 / var(--tw-border-opacity))}.border-gray-300{--tw-border-opacity: 1;border-color:rgb(209 213 219 / var(--tw-border-opacity))}.border-gray-600{--tw-border-opacity: 1;border-color:rgb(75 85 99 / var(--tw-border-opacity))}.border-gray-700{--tw-border-opacity: 1;border-color:rgb(55 65 81 / var(--tw-border-opacity))}.border-green-200{--tw-border-opacity: 1;border-color:rgb(187 247 208 / var(--tw-border-opacity))}.border-neutral-700{--tw-border-opacity: 1;border-color:rgb(64 64 64 / var(--tw-border-opacity))}.border-red-200{--tw-border-opacity: 1;border-color:rgb(254 202 202 / var(--tw-border-opacity))}.border-red-500{--tw-border-opacity: 1;border-color:rgb(239 68 68 / var(--tw-border-opacity))}.border-red-600{--tw-border-opacity: 1;border-color:rgb(220 38 38 / var(--tw-border-opacity))}.border-teal-100{--tw-border-opacity: 1;border-color:rgb(204 251 241 / var(--tw-border-opacity))}.border-teal-500{--tw-border-opacity: 1;border-color:rgb(20 184 166 / var(--tw-border-opacity))}.border-teal-900{--tw-border-opacity: 1;border-color:rgb(19 78 74 / var(--tw-border-opacity))}.border-transparent{border-color:transparent}.border-yellow-200{--tw-border-opacity: 1;border-color:rgb(254 240 138 / var(--tw-border-opacity))}.border-t-blue-500{--tw-border-opacity: 1;border-top-color:rgb(59 130 246 / var(--tw-border-opacity))}.border-t-blue-600{--tw-border-opacity: 1;border-top-color:rgb(37 99 235 / var(--tw-border-opacity))}.border-t-green-500{--tw-border-opacity: 1;border-top-color:rgb(34 197 94 / var(--tw-border-opacity))}.border-t-green-600{--tw-border-opacity: 1;border-top-color:rgb(22 163 74 / var(--tw-border-opacity))}.border-t-red-500{--tw-border-opacity: 1;border-top-color:rgb(239 68 68 / var(--tw-border-opacity))}.border-t-red-600{--tw-border-opacity: 1;border-top-color:rgb(220 38 38 / var(--tw-border-opacity))}.border-t-transparent{border-top-color:transparent}.bg-blue-500{--tw-bg-opacity: 1;background-color:rgb(59 130 246 / var(--tw-bg-opacity))}.bg-blue-600{--tw-bg-opacity: 1;background-color:rgb(37 99 235 / var(--tw-bg-opacity))}.bg-gray-100{--tw-bg-opacity: 1;background-color:rgb(243 244 246 / var(--tw-bg-opacity))}.bg-gray-50{--tw-bg-opacity: 1;background-color:rgb(249 250 251 / var(--tw-bg-opacity))}.bg-gray-600{--tw-bg-opacity: 1;background-color:rgb(75 85 99 / var(--tw-bg-opacity))}.bg-gray-700{--tw-bg-opacity: 1;background-color:rgb(55 65 81 / var(--tw-bg-opacity))}.bg-gray-800{--tw-bg-opacity: 1;background-color:rgb(31 41 55 / var(--tw-bg-opacity))}.bg-gray-900{--tw-bg-opacity: 1;background-color:rgb(17 24 39 / var(--tw-bg-opacity))}.bg-green-600{--tw-bg-opacity: 1;background-color:rgb(22 163 74 / var(--tw-bg-opacity))}.bg-neutral-800{--tw-bg-opacity: 1;background-color:rgb(38 38 38 / var(--tw-bg-opacity))}.bg-neutral-900{--tw-bg-opacity: 1;background-color:rgb(23 23 23 / var(--tw-bg-opacity))}.bg-red-500{--tw-bg-opacity: 1;background-color:rgb(239 68 68 / var(--tw-bg-opacity))}.bg-slate-800{--tw-bg-opacity: 1;background-color:rgb(30 41 59 / var(--tw-bg-opacity))}.bg-slate-900{--tw-bg-opacity: 1;background-color:rgb(15 23 42 / var(--tw-bg-opacity))}.bg-teal-200{--tw-bg-opacity: 1;background-color:rgb(153 246 228 / var(--tw-bg-opacity))}.bg-teal-50{--tw-bg-opacity: 1;background-color:rgb(240 253 250 / var(--tw-bg-opacity))}.bg-teal-800{--tw-bg-opacity: 1;background-color:rgb(17 94 89 / var(--tw-bg-opacity))}.bg-white{--tw-bg-opacity: 1;background-color:rgb(255 255 255 / var(--tw-bg-opacity))}.bg-yellow-50{--tw-bg-opacity: 1;background-color:rgb(254 252 232 / var(--tw-bg-opacity))}.bg-opacity-50{--tw-bg-opacity: .5}.bg-opacity-80{--tw-bg-opacity: .8}.p-1{padding:.25rem}.p-1\\.5{padding:.375rem}.p-2{padding:.5rem}.p-2\\.5{padding:.625rem}.p-3{padding:.75rem}.p-4{padding:1rem}.p-6{padding:1.5rem}.px-1{padding-left:.25rem;padding-right:.25rem}.px-3{padding-left:.75rem;padding-right:.75rem}.px-4{padding-left:1rem;padding-right:1rem}.px-6{padding-left:1.5rem;padding-right:1.5rem}.px-7{padding-left:1.75rem;padding-right:1.75rem}.py-1{padding-top:.25rem;padding-bottom:.25rem}.py-10{padding-top:2.5rem;padding-bottom:2.5rem}.py-16{padding-top:4rem;padding-bottom:4rem}.py-2{padding-top:.5rem;padding-bottom:.5rem}.py-2\\.5{padding-top:.625rem;padding-bottom:.625rem}.py-3{padding-top:.75rem;padding-bottom:.75rem}.py-4{padding-top:1rem;padding-bottom:1rem}.py-5{padding-top:1.25rem;padding-bottom:1.25rem}.pb-12{padding-bottom:3rem}.pe-11{padding-inline-end:2.75rem}.pe-3{padding-inline-end:.75rem}.pl-3{padding-left:.75rem}.pl-7{padding-left:1.75rem}.pr-10{padding-right:2.5rem}.pr-4{padding-right:1rem}.pr-9{padding-right:2.25rem}.ps-5{padding-inline-start:1.25rem}.text-left{text-align:left}.text-center{text-align:center}.text-start{text-align:start}.align-middle{vertical-align:middle}.font-mono{font-family:ui-monospace,SFMono-Regular,Menlo,Monaco,Consolas,Liberation Mono,Courier New,monospace}.text-2xl{font-size:1.5rem;line-height:2rem}.text-3xl{font-size:1.875rem;line-height:2.25rem}.text-7xl{font-size:4.5rem;line-height:1}.text-base{font-size:1rem;line-height:1.5rem}.text-lg{font-size:1.125rem;line-height:1.75rem}.text-sm{font-size:.875rem;line-height:1.25rem}.text-xl{font-size:1.25rem;line-height:1.75rem}.text-xs{font-size:.75rem;line-height:1rem}.font-bold{font-weight:700}.font-medium{font-weight:500}.font-semibold{font-weight:600}.uppercase{text-transform:uppercase}.leading-none{line-height:1}.tracking-wide{letter-spacing:.025em}.text-blue-500{--tw-text-opacity: 1;color:rgb(59 130 246 / var(--tw-text-opacity))}.text-blue-600{--tw-text-opacity: 1;color:rgb(37 99 235 / var(--tw-text-opacity))}.text-gray-200{--tw-text-opacity: 1;color:rgb(229 231 235 / var(--tw-text-opacity))}.text-gray-300{--tw-text-opacity: 1;color:rgb(209 213 219 / var(--tw-text-opacity))}.text-gray-400{--tw-text-opacity: 1;color:rgb(156 163 175 / var(--tw-text-opacity))}.text-gray-500{--tw-text-opacity: 1;color:rgb(107 114 128 / var(--tw-text-opacity))}.text-gray-600{--tw-text-opacity: 1;color:rgb(75 85 99 / var(--tw-text-opacity))}.text-gray-700{--tw-text-opacity: 1;color:rgb(55 65 81 / var(--tw-text-opacity))}.text-gray-800{--tw-text-opacity: 1;color:rgb(31 41 55 / var(--tw-text-opacity))}.text-green-500{--tw-text-opacity: 1;color:rgb(34 197 94 / var(--tw-text-opacity))}.text-green-600{--tw-text-opacity: 1;color:rgb(22 163 74 / var(--tw-text-opacity))}.text-neutral-200{--tw-text-opacity: 1;color:rgb(229 229 229 / var(--tw-text-opacity))}.text-neutral-300{--tw-text-opacity: 1;color:rgb(212 212 212 / var(--tw-text-opacity))}.text-neutral-400{--tw-text-opacity: 1;color:rgb(163 163 163 / var(--tw-text-opacity))}.text-neutral-500{--tw-text-opacity: 1;color:rgb(115 115 115 / var(--tw-text-opacity))}.text-red-500{--tw-text-opacity: 1;color:rgb(239 68 68 / var(--tw-text-opacity))}.text-red-600{--tw-text-opacity: 1;color:rgb(220 38 38 / var(--tw-text-opacity))}.text-slate-400{--tw-text-opacity: 1;color:rgb(148 163 184 / var(--tw-text-opacity))}.text-slate-700{--tw-text-opacity: 1;color:rgb(51 65 85 / var(--tw-text-opacity))}.text-teal-400{--tw-text-opacity: 1;color:rgb(45 212 191 / var(--tw-text-opacity))}.text-teal-800{--tw-text-opacity: 1;color:rgb(17 94 89 / var(--tw-text-opacity))}.text-white{--tw-text-opacity: 1;color:rgb(255 255 255 / var(--tw-text-opacity))}.text-yellow-400{--tw-text-opacity: 1;color:rgb(250 204 21 / var(--tw-text-opacity))}.text-yellow-700{--tw-text-opacity: 1;color:rgb(161 98 7 / var(--tw-text-opacity))}.text-yellow-800{--tw-text-opacity: 1;color:rgb(133 77 14 / var(--tw-text-opacity))}.decoration-2{text-decoration-thickness:2px}.opacity-0{opacity:0}.shadow{--tw-shadow: 0 1px 3px 0 rgb(0 0 0 / .1), 0 1px 2px -1px rgb(0 0 0 / .1);--tw-shadow-colored: 0 1px 3px 0 var(--tw-shadow-color), 0 1px 2px -1px var(--tw-shadow-color);box-shadow:var(--tw-ring-offset-shadow, 0 0 #0000),var(--tw-ring-shadow, 0 0 #0000),var(--tw-shadow)}.shadow-md{--tw-shadow: 0 4px 6px -1px rgb(0 0 0 / .1), 0 2px 4px -2px rgb(0 0 0 / .1);--tw-shadow-colored: 0 4px 6px -1px var(--tw-shadow-color), 0 2px 4px -2px var(--tw-shadow-color);box-shadow:var(--tw-ring-offset-shadow, 0 0 #0000),var(--tw-ring-shadow, 0 0 #0000),var(--tw-shadow)}.shadow-sm{--tw-shadow: 0 1px 2px 0 rgb(0 0 0 / .05);--tw-shadow-colored: 0 1px 2px 0 var(--tw-shadow-color);box-shadow:var(--tw-ring-offset-shadow, 0 0 #0000),var(--tw-ring-shadow, 0 0 #0000),var(--tw-shadow)}.shadow-neutral-700{--tw-shadow-color: #404040;--tw-shadow: var(--tw-shadow-colored)}.shadow-slate-700{--tw-shadow-color: #334155;--tw-shadow: var(--tw-shadow-colored)}.ring-1{--tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);--tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(1px + var(--tw-ring-offset-width)) var(--tw-ring-color);box-shadow:var(--tw-ring-offset-shadow),var(--tw-ring-shadow),var(--tw-shadow, 0 0 #0000)}.ring-transparent{--tw-ring-color: transparent}.ring-offset-white{--tw-ring-offset-color: #fff}.blur{--tw-blur: blur(8px);filter:var(--tw-blur) var(--tw-brightness) var(--tw-contrast) var(--tw-grayscale) var(--tw-hue-rotate) var(--tw-invert) var(--tw-saturate) var(--tw-sepia) var(--tw-drop-shadow)}.filter{filter:var(--tw-blur) var(--tw-brightness) var(--tw-contrast) var(--tw-grayscale) var(--tw-hue-rotate) var(--tw-invert) var(--tw-saturate) var(--tw-sepia) var(--tw-drop-shadow)}.transition{transition-property:color,background-color,border-color,text-decoration-color,fill,stroke,opacity,box-shadow,transform,filter,-webkit-backdrop-filter;transition-property:color,background-color,border-color,text-decoration-color,fill,stroke,opacity,box-shadow,transform,filter,backdrop-filter;transition-property:color,background-color,border-color,text-decoration-color,fill,stroke,opacity,box-shadow,transform,filter,backdrop-filter,-webkit-backdrop-filter;transition-timing-function:cubic-bezier(.4,0,.2,1);transition-duration:.15s}.transition-all{transition-property:all;transition-timing-function:cubic-bezier(.4,0,.2,1);transition-duration:.15s}.transition-colors{transition-property:color,background-color,border-color,text-decoration-color,fill,stroke;transition-timing-function:cubic-bezier(.4,0,.2,1);transition-duration:.15s}.duration-200{transition-duration:.2s}.duration-300{transition-duration:.3s}.duration-500{transition-duration:.5s}.ease-in-out{transition-timing-function:cubic-bezier(.4,0,.2,1)}.ease-out{transition-timing-function:cubic-bezier(0,0,.2,1)}.\\[--body-scroll\\:true\\]{--body-scroll: true}.marker\\:text-blue-600 *::marker{color:#2563eb}.marker\\:text-blue-600::marker{color:#2563eb}.before\\:inline-block:before{content:var(--tw-content);display:inline-block}.before\\:h-6:before{content:var(--tw-content);height:1.5rem}.before\\:w-6:before{content:var(--tw-content);width:1.5rem}.before\\:translate-x-0:before{content:var(--tw-content);--tw-translate-x: 0px;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.before\\:transform:before{content:var(--tw-content);transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.before\\:rounded-full:before{content:var(--tw-content);border-radius:9999px}.before\\:bg-white:before{content:var(--tw-content);--tw-bg-opacity: 1;background-color:rgb(255 255 255 / var(--tw-bg-opacity))}.before\\:shadow:before{content:var(--tw-content);--tw-shadow: 0 1px 3px 0 rgb(0 0 0 / .1), 0 1px 2px -1px rgb(0 0 0 / .1);--tw-shadow-colored: 0 1px 3px 0 var(--tw-shadow-color), 0 1px 2px -1px var(--tw-shadow-color);box-shadow:var(--tw-ring-offset-shadow, 0 0 #0000),var(--tw-ring-shadow, 0 0 #0000),var(--tw-shadow)}.before\\:ring-0:before{content:var(--tw-content);--tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);--tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(0px + var(--tw-ring-offset-width)) var(--tw-ring-color);box-shadow:var(--tw-ring-offset-shadow),var(--tw-ring-shadow),var(--tw-shadow, 0 0 #0000)}.before\\:transition:before{content:var(--tw-content);transition-property:color,background-color,border-color,text-decoration-color,fill,stroke,opacity,box-shadow,transform,filter,-webkit-backdrop-filter;transition-property:color,background-color,border-color,text-decoration-color,fill,stroke,opacity,box-shadow,transform,filter,backdrop-filter;transition-property:color,background-color,border-color,text-decoration-color,fill,stroke,opacity,box-shadow,transform,filter,backdrop-filter,-webkit-backdrop-filter;transition-timing-function:cubic-bezier(.4,0,.2,1);transition-duration:.15s}.before\\:duration-200:before{content:var(--tw-content);transition-duration:.2s}.before\\:ease-in-out:before{content:var(--tw-content);transition-timing-function:cubic-bezier(.4,0,.2,1)}.first\\:mt-0:first-child{margin-top:0}.first\\:rounded-t-lg:first-child{border-top-left-radius:.5rem;border-top-right-radius:.5rem}.last\\:rounded-b-lg:last-child{border-bottom-right-radius:.5rem;border-bottom-left-radius:.5rem}.checked\\:bg-blue-600:checked{--tw-bg-opacity: 1;background-color:rgb(37 99 235 / var(--tw-bg-opacity))}.checked\\:bg-none:checked{background-image:none}.checked\\:before\\:translate-x-full:checked:before{content:var(--tw-content);--tw-translate-x: 100%;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.checked\\:before\\:bg-blue-200:checked:before{content:var(--tw-content);--tw-bg-opacity: 1;background-color:rgb(191 219 254 / var(--tw-bg-opacity))}.hover\\:border-blue-500:hover{--tw-border-opacity: 1;border-color:rgb(59 130 246 / var(--tw-border-opacity))}.hover\\:border-green-500:hover{--tw-border-opacity: 1;border-color:rgb(34 197 94 / var(--tw-border-opacity))}.hover\\:border-red-400:hover{--tw-border-opacity: 1;border-color:rgb(248 113 113 / var(--tw-border-opacity))}.hover\\:border-red-500:hover{--tw-border-opacity: 1;border-color:rgb(239 68 68 / var(--tw-border-opacity))}.hover\\:bg-blue-50:hover{--tw-bg-opacity: 1;background-color:rgb(239 246 255 / var(--tw-bg-opacity))}.hover\\:bg-blue-500:hover{--tw-bg-opacity: 1;background-color:rgb(59 130 246 / var(--tw-bg-opacity))}.hover\\:bg-blue-600:hover{--tw-bg-opacity: 1;background-color:rgb(37 99 235 / var(--tw-bg-opacity))}.hover\\:bg-blue-700:hover{--tw-bg-opacity: 1;background-color:rgb(29 78 216 / var(--tw-bg-opacity))}.hover\\:bg-gray-100:hover{--tw-bg-opacity: 1;background-color:rgb(243 244 246 / var(--tw-bg-opacity))}.hover\\:bg-gray-50:hover{--tw-bg-opacity: 1;background-color:rgb(249 250 251 / var(--tw-bg-opacity))}.hover\\:bg-green-500:hover{--tw-bg-opacity: 1;background-color:rgb(34 197 94 / var(--tw-bg-opacity))}.hover\\:bg-red-500:hover{--tw-bg-opacity: 1;background-color:rgb(239 68 68 / var(--tw-bg-opacity))}.hover\\:text-blue-500:hover{--tw-text-opacity: 1;color:rgb(59 130 246 / var(--tw-text-opacity))}.hover\\:text-blue-600:hover{--tw-text-opacity: 1;color:rgb(37 99 235 / var(--tw-text-opacity))}.hover\\:text-blue-800:hover{--tw-text-opacity: 1;color:rgb(30 64 175 / var(--tw-text-opacity))}.hover\\:text-gray-400:hover{--tw-text-opacity: 1;color:rgb(156 163 175 / var(--tw-text-opacity))}.hover\\:text-green-800:hover{--tw-text-opacity: 1;color:rgb(22 101 52 / var(--tw-text-opacity))}.hover\\:text-red-400:hover{--tw-text-opacity: 1;color:rgb(248 113 113 / var(--tw-text-opacity))}.hover\\:text-red-500:hover{--tw-text-opacity: 1;color:rgb(239 68 68 / var(--tw-text-opacity))}.hover\\:text-red-600:hover{--tw-text-opacity: 1;color:rgb(220 38 38 / var(--tw-text-opacity))}.hover\\:text-white:hover{--tw-text-opacity: 1;color:rgb(255 255 255 / var(--tw-text-opacity))}.focus\\:z-10:focus{z-index:10}.focus\\:border-blue-500:focus{--tw-border-opacity: 1;border-color:rgb(59 130 246 / var(--tw-border-opacity))}.focus\\:border-blue-600:focus{--tw-border-opacity: 1;border-color:rgb(37 99 235 / var(--tw-border-opacity))}.focus\\:border-indigo-500:focus{--tw-border-opacity: 1;border-color:rgb(99 102 241 / var(--tw-border-opacity))}.focus\\:outline-none:focus{outline:2px solid transparent;outline-offset:2px}.focus\\:ring-2:focus{--tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);--tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(2px + var(--tw-ring-offset-width)) var(--tw-ring-color);box-shadow:var(--tw-ring-offset-shadow),var(--tw-ring-shadow),var(--tw-shadow, 0 0 #0000)}.focus\\:ring-blue-500:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(59 130 246 / var(--tw-ring-opacity))}.focus\\:ring-blue-600:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(37 99 235 / var(--tw-ring-opacity))}.focus\\:ring-gray-400:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(156 163 175 / var(--tw-ring-opacity))}.focus\\:ring-green-200:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(187 247 208 / var(--tw-ring-opacity))}.focus\\:ring-indigo-500:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(99 102 241 / var(--tw-ring-opacity))}.focus\\:ring-red-200:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(254 202 202 / var(--tw-ring-opacity))}.focus\\:ring-red-500:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(239 68 68 / var(--tw-ring-opacity))}.focus\\:ring-offset-2:focus{--tw-ring-offset-width: 2px}.focus\\:ring-offset-white:focus{--tw-ring-offset-color: #fff}.disabled\\:pointer-events-none:disabled{pointer-events:none}.disabled\\:opacity-50:disabled{opacity:.5}[data-hs-tab].active.hs-tab-active\\:border-blue-600{--tw-border-opacity: 1;border-color:rgb(37 99 235 / var(--tw-border-opacity))}[data-hs-tab].active.hs-tab-active\\:font-semibold{font-weight:600}[data-hs-tab].active.hs-tab-active\\:text-blue-600{--tw-text-opacity: 1;color:rgb(37 99 235 / var(--tw-text-opacity))}[data-hs-tab].active .hs-tab-active\\:border-blue-600{--tw-border-opacity: 1;border-color:rgb(37 99 235 / var(--tw-border-opacity))}[data-hs-tab].active .hs-tab-active\\:font-semibold{font-weight:600}[data-hs-tab].active .hs-tab-active\\:text-blue-600{--tw-text-opacity: 1;color:rgb(37 99 235 / var(--tw-text-opacity))}.open.hs-overlay-open\\:mt-7{margin-top:1.75rem}.open.hs-overlay-open\\:translate-x-0{--tw-translate-x: 0px;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.open.hs-overlay-open\\:opacity-100{opacity:1}.open.hs-overlay-open\\:duration-500{transition-duration:.5s}.open .hs-overlay-open\\:mt-7{margin-top:1.75rem}.open .hs-overlay-open\\:translate-x-0{--tw-translate-x: 0px;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.open .hs-overlay-open\\:opacity-100{opacity:1}.open .hs-overlay-open\\:duration-500{transition-duration:.5s}@media (prefers-color-scheme: dark){.dark\\:divide-gray-700>:not([hidden])~:not([hidden]){--tw-divide-opacity: 1;border-color:rgb(55 65 81 / var(--tw-divide-opacity))}.dark\\:border-blue-500{--tw-border-opacity: 1;border-color:rgb(59 130 246 / var(--tw-border-opacity))}.dark\\:border-gray-600{--tw-border-opacity: 1;border-color:rgb(75 85 99 / var(--tw-border-opacity))}.dark\\:border-gray-700{--tw-border-opacity: 1;border-color:rgb(55 65 81 / var(--tw-border-opacity))}.dark\\:border-neutral-700{--tw-border-opacity: 1;border-color:rgb(64 64 64 / var(--tw-border-opacity))}.dark\\:border-red-500{--tw-border-opacity: 1;border-color:rgb(239 68 68 / var(--tw-border-opacity))}.dark\\:border-teal-900{--tw-border-opacity: 1;border-color:rgb(19 78 74 / var(--tw-border-opacity))}.dark\\:border-t-blue-500{--tw-border-opacity: 1;border-top-color:rgb(59 130 246 / var(--tw-border-opacity))}.dark\\:border-t-green-500{--tw-border-opacity: 1;border-top-color:rgb(34 197 94 / var(--tw-border-opacity))}.dark\\:border-t-red-500{--tw-border-opacity: 1;border-top-color:rgb(239 68 68 / var(--tw-border-opacity))}.dark\\:bg-gray-700{--tw-bg-opacity: 1;background-color:rgb(55 65 81 / var(--tw-bg-opacity))}.dark\\:bg-gray-800{--tw-bg-opacity: 1;background-color:rgb(31 41 55 / var(--tw-bg-opacity))}.dark\\:bg-neutral-800{--tw-bg-opacity: 1;background-color:rgb(38 38 38 / var(--tw-bg-opacity))}.dark\\:bg-neutral-900{--tw-bg-opacity: 1;background-color:rgb(23 23 23 / var(--tw-bg-opacity))}.dark\\:bg-slate-800{--tw-bg-opacity: 1;background-color:rgb(30 41 59 / var(--tw-bg-opacity))}.dark\\:bg-slate-900{--tw-bg-opacity: 1;background-color:rgb(15 23 42 / var(--tw-bg-opacity))}.dark\\:bg-teal-800{--tw-bg-opacity: 1;background-color:rgb(17 94 89 / var(--tw-bg-opacity))}.dark\\:bg-teal-800\\/30{background-color:#115e594d}.dark\\:bg-opacity-80{--tw-bg-opacity: .8}.dark\\:text-blue-500{--tw-text-opacity: 1;color:rgb(59 130 246 / var(--tw-text-opacity))}.dark\\:text-gray-200{--tw-text-opacity: 1;color:rgb(229 231 235 / var(--tw-text-opacity))}.dark\\:text-gray-300{--tw-text-opacity: 1;color:rgb(209 213 219 / var(--tw-text-opacity))}.dark\\:text-gray-400{--tw-text-opacity: 1;color:rgb(156 163 175 / var(--tw-text-opacity))}.dark\\:text-gray-500{--tw-text-opacity: 1;color:rgb(107 114 128 / var(--tw-text-opacity))}.dark\\:text-green-500{--tw-text-opacity: 1;color:rgb(34 197 94 / var(--tw-text-opacity))}.dark\\:text-neutral-200{--tw-text-opacity: 1;color:rgb(229 229 229 / var(--tw-text-opacity))}.dark\\:text-neutral-400{--tw-text-opacity: 1;color:rgb(163 163 163 / var(--tw-text-opacity))}.dark\\:text-neutral-500{--tw-text-opacity: 1;color:rgb(115 115 115 / var(--tw-text-opacity))}.dark\\:text-red-500{--tw-text-opacity: 1;color:rgb(239 68 68 / var(--tw-text-opacity))}.dark\\:text-slate-400{--tw-text-opacity: 1;color:rgb(148 163 184 / var(--tw-text-opacity))}.dark\\:text-teal-400{--tw-text-opacity: 1;color:rgb(45 212 191 / var(--tw-text-opacity))}.dark\\:text-white{--tw-text-opacity: 1;color:rgb(255 255 255 / var(--tw-text-opacity))}.dark\\:placeholder-gray-400::-moz-placeholder{--tw-placeholder-opacity: 1;color:rgb(156 163 175 / var(--tw-placeholder-opacity))}.dark\\:placeholder-gray-400::placeholder{--tw-placeholder-opacity: 1;color:rgb(156 163 175 / var(--tw-placeholder-opacity))}.dark\\:placeholder-neutral-500::-moz-placeholder{--tw-placeholder-opacity: 1;color:rgb(115 115 115 / var(--tw-placeholder-opacity))}.dark\\:placeholder-neutral-500::placeholder{--tw-placeholder-opacity: 1;color:rgb(115 115 115 / var(--tw-placeholder-opacity))}.dark\\:shadow-neutral-700\\/70{--tw-shadow-color: rgb(64 64 64 / .7);--tw-shadow: var(--tw-shadow-colored)}.dark\\:shadow-slate-700\\/\\[\\.7\\]{--tw-shadow-color: rgb(51 65 85 / .7);--tw-shadow: var(--tw-shadow-colored)}.dark\\:before\\:bg-gray-400:before{content:var(--tw-content);--tw-bg-opacity: 1;background-color:rgb(156 163 175 / var(--tw-bg-opacity))}.dark\\:checked\\:border-blue-500:checked{--tw-border-opacity: 1;border-color:rgb(59 130 246 / var(--tw-border-opacity))}.dark\\:checked\\:bg-blue-500:checked{--tw-bg-opacity: 1;background-color:rgb(59 130 246 / var(--tw-bg-opacity))}.dark\\:checked\\:bg-blue-600:checked{--tw-bg-opacity: 1;background-color:rgb(37 99 235 / var(--tw-bg-opacity))}.dark\\:checked\\:before\\:bg-blue-200:checked:before{content:var(--tw-content);--tw-bg-opacity: 1;background-color:rgb(191 219 254 / var(--tw-bg-opacity))}.dark\\:hover\\:border-blue-400:hover{--tw-border-opacity: 1;border-color:rgb(96 165 250 / var(--tw-border-opacity))}.dark\\:hover\\:border-red-400:hover{--tw-border-opacity: 1;border-color:rgb(248 113 113 / var(--tw-border-opacity))}.dark\\:hover\\:bg-gray-900:hover{--tw-bg-opacity: 1;background-color:rgb(17 24 39 / var(--tw-bg-opacity))}.dark\\:hover\\:bg-neutral-700:hover{--tw-bg-opacity: 1;background-color:rgb(64 64 64 / var(--tw-bg-opacity))}.dark\\:hover\\:bg-slate-800:hover{--tw-bg-opacity: 1;background-color:rgb(30 41 59 / var(--tw-bg-opacity))}.dark\\:hover\\:text-blue-400:hover{--tw-text-opacity: 1;color:rgb(96 165 250 / var(--tw-text-opacity))}.dark\\:hover\\:text-blue-500:hover{--tw-text-opacity: 1;color:rgb(59 130 246 / var(--tw-text-opacity))}.dark\\:hover\\:text-green-400:hover{--tw-text-opacity: 1;color:rgb(74 222 128 / var(--tw-text-opacity))}.dark\\:hover\\:text-red-400:hover{--tw-text-opacity: 1;color:rgb(248 113 113 / var(--tw-text-opacity))}.dark\\:hover\\:text-red-500:hover{--tw-text-opacity: 1;color:rgb(239 68 68 / var(--tw-text-opacity))}.dark\\:hover\\:text-slate-300:hover{--tw-text-opacity: 1;color:rgb(203 213 225 / var(--tw-text-opacity))}.dark\\:hover\\:text-white:hover{--tw-text-opacity: 1;color:rgb(255 255 255 / var(--tw-text-opacity))}.dark\\:focus\\:border-blue-500:focus{--tw-border-opacity: 1;border-color:rgb(59 130 246 / var(--tw-border-opacity))}.dark\\:focus\\:outline-none:focus{outline:2px solid transparent;outline-offset:2px}.dark\\:focus\\:ring-1:focus{--tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);--tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(1px + var(--tw-ring-offset-width)) var(--tw-ring-color);box-shadow:var(--tw-ring-offset-shadow),var(--tw-ring-shadow),var(--tw-shadow, 0 0 #0000)}.dark\\:focus\\:ring-blue-500:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(59 130 246 / var(--tw-ring-opacity))}.dark\\:focus\\:ring-gray-600:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(75 85 99 / var(--tw-ring-opacity))}.dark\\:focus\\:ring-gray-700:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(55 65 81 / var(--tw-ring-opacity))}.dark\\:focus\\:ring-neutral-600:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(82 82 82 / var(--tw-ring-opacity))}.dark\\:focus\\:ring-red-600:focus{--tw-ring-opacity: 1;--tw-ring-color: rgb(220 38 38 / var(--tw-ring-opacity))}.dark\\:focus\\:ring-offset-gray-800:focus{--tw-ring-offset-color: #1f2937}}@media (min-width: 640px){.sm\\:mx-auto{margin-left:auto;margin-right:auto}.sm\\:mb-3{margin-bottom:.75rem}.sm\\:mt-10{margin-top:2.5rem}.sm\\:w-auto{width:auto}.sm\\:w-full{width:100%}.sm\\:max-w-lg{max-width:32rem}.sm\\:flex-row{flex-direction:row}.sm\\:gap-3{gap:.75rem}.sm\\:gap-x-4{-moz-column-gap:1rem;column-gap:1rem}.sm\\:space-y-6>:not([hidden])~:not([hidden]){--tw-space-y-reverse: 0;margin-top:calc(1.5rem * calc(1 - var(--tw-space-y-reverse)));margin-bottom:calc(1.5rem * var(--tw-space-y-reverse))}.sm\\:p-4{padding:1rem}.sm\\:px-6{padding-left:1.5rem;padding-right:1.5rem}.sm\\:py-4{padding-top:1rem;padding-bottom:1rem}.sm\\:py-6{padding-top:1.5rem;padding-bottom:1.5rem}.sm\\:text-3xl{font-size:1.875rem;line-height:2.25rem}.sm\\:text-4xl{font-size:2.25rem;line-height:2.5rem}.sm\\:text-9xl{font-size:8rem;line-height:1}.sm\\:text-sm{font-size:.875rem;line-height:1.25rem}}@media (min-width: 768px){.md\\:flex{display:flex}.md\\:items-center{align-items:center}.md\\:justify-between{justify-content:space-between}.md\\:p-10{padding:2.5rem}.md\\:p-5{padding:1.25rem}}@media (min-width: 1024px){.lg\\:bottom-0{bottom:0}.lg\\:right-auto{right:auto}.lg\\:block{display:block}.lg\\:hidden{display:none}.lg\\:translate-x-0{--tw-translate-x: 0px;transform:translate(var(--tw-translate-x),var(--tw-translate-y)) rotate(var(--tw-rotate)) skew(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y))}.lg\\:px-8{padding-left:2rem;padding-right:2rem}.lg\\:py-14{padding-top:3.5rem;padding-bottom:3.5rem}.lg\\:pl-64{padding-left:16rem}}
"""

js_content = '''var Rn=Object.defineProperty;var nn=(E,e,T)=>e in E?Rn(E,e,{enumerable:!0,configurable:!0,writable:!0,value:T}):E[e]=T;var wt=(E,e,T)=>(nn(E,typeof e!="symbol"?e+"":e,T),T);(function(){const e=document.createElement("link").relList;if(e&&e.supports&&e.supports("modulepreload"))return;for(const r of document.querySelectorAll('link[rel="modulepreload"]'))t(r);new MutationObserver(r=>{for(const R of r)if(R.type==="childList")for(const n of R.addedNodes)n.tagName==="LINK"&&n.rel==="modulepreload"&&t(n)}).observe(document,{childList:!0,subtree:!0});function T(r){const R={};return r.integrity&&(R.integrity=r.integrity),r.referrerPolicy&&(R.referrerPolicy=r.referrerPolicy),r.crossOrigin==="use-credentials"?R.credentials="include":r.crossOrigin==="anonymous"?R.credentials="omit":R.credentials="same-origin",R}function t(r){if(r.ep)return;r.ep=!0;const R=T(r);fetch(r.href,R)}})();function j(){}function An(E,e){for(const T in e)E[T]=e[T];return E}function cR(E){return E()}function YT(){return Object.create(null)}function NE(E){E.forEach(cR)}function zE(E){return typeof E=="function"}function _e(E,e){return E!=E?e==e:E!==e||E&&typeof E=="object"||typeof E=="function"}let ut;function VT(E,e){return E===e?!0:(ut||(ut=document.createElement("a")),ut.href=e,E===ut.href)}function sn(E){return Object.keys(E).length===0}function fR(E,...e){if(E==null){for(const t of e)t(void 0);return j}const T=E.subscribe(...e);return T.unsubscribe?()=>T.unsubscribe():T}function uE(E){let e;return fR(E,T=>e=T)(),e}function eE(E,e,T){E.$$.on_destroy.push(fR(e,T))}function Ut(E,e,T,t){if(E){const r=PR(E,e,T,t);return E[0](r)}}function PR(E,e,T,t){return E[1]&&t?An(T.ctx.slice(),E[1](t(e))):T.ctx}function mt(E,e,T,t){if(E[2]&&t){const r=E[2](t(T));if(e.dirty===void 0)return r;if(typeof r=="object"){const R=[],n=Math.max(e.dirty.length,r.length);for(let s=0;s<n;s+=1)R[s]=e.dirty[s]|r[s];return R}return e.dirty|r}return e.dirty}function ht(E,e,T,t,r,R){if(r){const n=PR(e,T,t,R);E.p(n,r)}}function Gt(E){if(E.ctx.length>32){const e=[],T=E.ctx.length/32;for(let t=0;t<T;t++)e[t]=-1;return e}return-1}function OT(E,e,T){return E.set(T),e}function l(E,e){E.appendChild(e)}function V(E,e,T){E.insertBefore(e,T||null)}function Y(E){E.parentNode&&E.parentNode.removeChild(E)}function nE(E,e){for(let T=0;T<E.length;T+=1)E[T]&&E[T].d(e)}function f(E){return document.createElement(E)}function OE(E){return document.createElementNS("http://www.w3.org/2000/svg",E)}function te(E){return document.createTextNode(E)}function $(){return te(" ")}function je(){return te("")}function Ne(E,e,T,t){return E.addEventListener(e,T,t),()=>E.removeEventListener(e,T,t)}function a(E,e,T){T==null?E.removeAttribute(e):E.getAttribute(e)!==T&&E.setAttribute(e,T)}function Sn(E){let e;return{p(...T){e=T,e.forEach(t=>E.push(t))},r(){e.forEach(T=>E.splice(E.indexOf(T),1))}}}function on(E){return Array.from(E.childNodes)}function Le(E,e){e=""+e,E.data!==e&&(E.data=e)}function Ye(E,e){E.value=e??""}function ct(E,e,T,t){T==null?E.style.removeProperty(e):E.style.setProperty(e,T,t?"important":"")}function WT(E,e,T){for(let t=0;t<E.options.length;t+=1){const r=E.options[t];if(r.__value===e){r.selected=!0;return}}(!T||e!==void 0)&&(E.selectedIndex=-1)}function On(E){const e=E.querySelector(":checked");return e&&e.__value}let Ot;function st(E){Ot=E}function an(){if(!Ot)throw new Error("Function called outside component initialization");return Ot}function DR(E){an().$$.on_mount.push(E)}const jE=[],iT=[];let et=[];const aT=[],In=Promise.resolve();let IT=!1;function Nn(){IT||(IT=!0,In.then(dR))}function dt(E){et.push(E)}function ln(E){aT.push(E)}const $t=new Set;let JE=0;function dR(){if(JE!==0)return;const E=Ot;do{try{for(;JE<jE.length;){const e=jE[JE];JE++,st(e),_n(e.$$)}}catch(e){throw jE.length=0,JE=0,e}for(st(null),jE.length=0,JE=0;iT.length;)iT.pop()();for(let e=0;e<et.length;e+=1){const T=et[e];$t.has(T)||($t.add(T),T())}et.length=0}while(jE.length);for(;aT.length;)aT.pop()();IT=!1,$t.clear(),st(E)}function _n(E){if(E.fragment!==null){E.update(),NE(E.before_update);const e=E.dirty;E.dirty=[-1],E.fragment&&E.fragment.p(E.ctx,e),E.after_update.forEach(dt)}}function Ln(E){const e=[],T=[];et.forEach(t=>E.indexOf(t)===-1?e.push(t):T.push(t)),T.forEach(t=>t()),et=e}const ft=new Set;let vE;function Ge(){vE={r:0,c:[],p:vE}}function ge(){vE.r||NE(vE.c),vE=vE.p}function m(E,e){E&&E.i&&(ft.delete(E),E.i(e))}function y(E,e,T,t){if(E&&E.o){if(ft.has(E))return;ft.add(E),vE.c.push(()=>{ft.delete(E),t&&(T&&E.d(1),t())}),E.o(e)}else t&&t()}function De(E){return(E==null?void 0:E.length)!==void 0?E:Array.from(E)}function Cn(E,e){E.d(1),e.delete(E.key)}function un(E,e,T,t,r,R,n,s,S,A,o,i){let _=E.length,c=R.length,P=_;const p={};for(;P--;)p[E[P].key]=P;const C=[],L=new Map,I=new Map,u=[];for(P=c;P--;){const O=i(r,R,P),N=T(O);let D=n.get(N);D?t&&u.push(()=>D.p(O,e)):(D=A(N,O),D.c()),L.set(N,C[P]=D),N in p&&I.set(N,Math.abs(P-p[N]))}const H=new Set,b=new Set;function M(O){m(O,1),O.m(s,o),n.set(O.key,O),o=O.first,c--}for(;_&&c;){const O=C[c-1],N=E[_-1],D=O.key,B=N.key;O===N?(o=O.first,_--,c--):L.has(B)?!n.has(D)||H.has(D)?M(O):b.has(B)?_--:I.get(D)>I.get(B)?(b.add(D),M(O)):(H.add(B),_--):(S(N,n),_--)}for(;_--;){const O=E[_];L.has(O.key)||S(O,n)}for(;c;)M(C[c-1]);return NE(u),C}function cn(E,e,T){const t=E.$$.props[e];t!==void 0&&(E.$$.bound[t]=T,T(E.$$.ctx[t]))}function K(E){E&&E.c()}function X(E,e,T){const{fragment:t,after_update:r}=E.$$;t&&t.m(e,T),dt(()=>{const R=E.$$.on_mount.map(cR).filter(zE);E.$$.on_destroy?E.$$.on_destroy.push(...R):NE(R),E.$$.on_mount=[]}),r.forEach(dt)}function k(E,e){const T=E.$$;T.fragment!==null&&(Ln(T.after_update),NE(T.on_destroy),T.fragment&&T.fragment.d(e),T.on_destroy=T.fragment=null,T.ctx=[])}function fn(E,e){E.$$.dirty[0]===-1&&(jE.push(E),Nn(),E.$$.dirty.fill(0)),E.$$.dirty[e/31|0]|=1<<e%31}function Ce(E,e,T,t,r,R,n,s=[-1]){const S=Ot;st(E);const A=E.$$={fragment:null,ctx:[],props:R,update:j,not_equal:r,bound:YT(),on_mount:[],on_destroy:[],on_disconnect:[],before_update:[],after_update:[],context:new Map(e.context||(S?S.$$.context:[])),callbacks:YT(),dirty:s,skip_bound:!1,root:e.target||S.$$.root};n&&n(A.root);let o=!1;if(A.ctx=T?T(E,e.props||{},(i,_,...c)=>{const P=c.length?c[0]:_;return A.ctx&&r(A.ctx[i],A.ctx[i]=P)&&(!A.skip_bound&&A.bound[i]&&A.bound[i](P),o&&fn(E,i)),_}):[],A.update(),o=!0,NE(A.before_update),A.fragment=t?t(A.ctx):!1,e.target){if(e.hydrate){const i=on(e.target);A.fragment&&A.fragment.l(i),i.forEach(Y)}else A.fragment&&A.fragment.c();e.intro&&m(E.$$.fragment),X(E,e.target,e.anchor),dR()}st(S)}class ue{constructor(){wt(this,"$$");wt(this,"$$set")}$destroy(){k(this,1),this.$destroy=j}$on(e,T){if(!zE(T))return j;const t=this.$$.callbacks[e]||(this.$$.callbacks[e]=[]);return t.push(T),()=>{const r=t.indexOf(T);r!==-1&&t.splice(r,1)}}$set(e){this.$$set&&!sn(e)&&(this.$$.skip_bound=!0,this.$$set(e),this.$$.skip_bound=!1)}}const Pn="4";typeof window<"u"&&(window.__svelte||(window.__svelte={v:new Set})).v.add(Pn);const qE=[];function iE(E,e=j){let T;const t=new Set;function r(s){if(_e(E,s)&&(E=s,T)){const S=!qE.length;for(const A of t)A[1](),qE.push(A,E);if(S){for(let A=0;A<qE.length;A+=2)qE[A][0](qE[A+1]);qE.length=0}}}function R(s){r(s(E))}function n(s,S=j){const A=[s,S];return t.add(A),t.size===1&&(T=e(r,R)||j),s(E),()=>{t.delete(A),t.size===0&&T&&(T(),T=null)}}return{set:r,update:R,subscribe:n}}let GE=iE(""),YE=iE([]),LT=iE(null),gt=iE(null),Ht=iE(!1),St=iE(!1),DE=iE("chat"),CT=iE([]),Et=iE(""),pR=iE(!1),BE=iE(""),VE=iE({debug:!0,logo:"",title:"Welcome to Vanna.AI",subtitle:"Loading...",show_training_data:!0,suggested_questions:!0,sql:!0,table:!0,csv_download:!0,chart:!0,redraw_chart:!0,auto_fix_sql:!0,ask_results_correct:!0,followup_questions:!0,summarization:!0,function_generation:!0,version:""}),bt=iE(null),MR=iE([]);function UR(){YE.set([]),Ht.set(!1),St.set(!1)}async function uT(E){let e=uE(VE),T=yn();if(Se({type:"user_question",question:E}),Ht.set(!0),T){const n=await Pe("generate_rewritten_question","GET",{last_question:T,new_question:E});n.type==="rewritten_question"&&n.question!==E&&(Se(n),E=n.question)}const t=await Pe("get_function","GET",{question:E});let r;if(e.function_generation&&t.type==="function")Se(t),r=t.id,GE.set(t.id),Et.set(t.function.instantiated_sql);else{const n=await Pe("generate_sql","GET",{question:E});if(Se(n),n.type!=="sql")return;window.location.hash=n.id,GE.set(n.id),Et.set(n.text),r=n.id}const R=await Pe("run_sql","GET",{id:r});if(Se(R),R.type==="df"){if(R.should_generate_chart){const n=await Pe("generate_plotly_figure","GET",{id:R.id});if(Se(n),n.type!=="plotly_figure")return;CT.update(s=>[...s,{question:E,id:n.id}])}if(e.summarization){const n=await Pe("generate_summary","GET",{id:r});Se(n)}Se({type:"feedback_question"}),Se({type:"feedback_buttons"})}}async function Dn(E){let e=uE(VE);if(Se(E),E.type!=="sql")return;window.location.hash=E.id,GE.set(E.id),Et.set(E.text);const T=await Pe("run_sql","GET",{id:E.id});if(Se(T),T.type!=="df")return;const t=await Pe("generate_plotly_figure","GET",{id:T.id});if(Se(t),t.type==="plotly_figure"){if(e.summarization){const r=await Pe("generate_summary","GET",{id:t.id});Se(r)}Se({type:"feedback_question"}),Se({type:"feedback_buttons"})}}function dn(E){Se({type:"user_question",question:"Re-run the SQL"}),Pe("run_sql","GET",{id:E}).then(Se).then(e=>{e.type==="df"&&Pe("generate_plotly_figure","GET",{id:e.id}).then(Se).then(T=>{T.type==="plotly_figure"&&Pe("generate_followup_questions","GET",{id:T.id}).then(Se)})})}function mR(){Pe("get_question_history","GET",[]).then(gn)}function pn(){Pe("get_config","GET",[]).then(Gn)}function cT(){window.location.hash="functions",DE.set("functions"),Pe("get_all_functions","GET",[]).then(mn)}function hR(){window.location.hash="training-data",DE.set("training-data"),Pe("get_training_data","GET",[]).then(pt)}function it(){window.location.hash="",DE.set("chat"),UR(),uE(LT)===null&&Pe("generate_questions","GET",[]).then(hn),mR()}function Mn(E){window.location.hash=E,DE.set("chat"),UR(),Ht.set(!0),Pe("load_question","GET",{id:E}).then(Se)}function Un(E){gt.set(null),Pe("remove_training_data","POST",{id:E}).then(e=>{Pe("get_training_data","GET",[]).then(pt)})}function Se(E){return E.type==="not_logged_in"?(bt.set(E.html),DE.set("login"),E):(YE.update(e=>[...e,E]),bn(),E)}function pt(E){return gt.set(E),E.type==="df"?JSON.parse(E.df).length===0&&DE.set("no-training-data"):E.type==="not_logged_in"&&(bt.set(E.html),DE.set("login")),E}function mn(E){return E.type==="functions"&&MR.set(E.functions),E}function hn(E){return LT.set(E),E}function Gn(E){return E.type==="config"?(VE.set(E.config),E.config.debug&&xn()):E.type==="not_logged_in"&&(bt.set(E.html),DE.set("login")),E}function gn(E){return E.type==="question_history"&&CT.set(E.questions),E}function Hn(E,e){gt.set(null);let T={};T[e]=E,Pe("train","POST",T).then(pt).then(t=>{t.type!=="error"&&Pe("get_training_data","GET",[]).then(pt)})}async function Pe(E,e,T){try{St.set(!0);let t="",r;if(e==="GET")t=Object.entries(T).filter(([n,s])=>n!=="endpoint"&&n!=="addMessage").map(([n,s])=>`${encodeURIComponent(n)}=${encodeURIComponent(s)}`).join("&"),r=await fetch(`/api/v0/${E}?${t}`);else{let n=JSON.stringify(T);r=await fetch(`/api/v0/${E}`,{method:"POST",headers:{"Content-Type":"application/json"},body:n})}if(!r.ok)throw new Error("The server returned an error. See the server logs for more details. If you are running in Colab, this function is probably not supported. Please try running in a local environment.");const R=await r.json();return St.set(!1),R}catch(t){return St.set(!1),{type:"error",error:String(t)}}}function bn(){setTimeout(()=>{window.scrollTo({top:document.body.scrollHeight,behavior:"smooth"})},100)}function yn(){let E=uE(YE),e=E.findLast(T=>T.type==="rewritten_question");return e||(e=E.findLast(T=>T.type==="user_question")),e&&(e.type==="rewritten_question"||e.type==="user_question")?e.question:null}function fT(){let E=uE(YE),e=E.find(T=>T.type==="user_question");if(e&&e.type==="user_question"){let T=E.findLast(t=>t.type==="sql");if(T&&T.type==="sql")return{question:e.question,sql:T.text}}return null}function at(E){YE.update(e=>e.filter(T=>T.type!==E))}function Bn(E){Pe("fix_sql","POST",{id:uE(GE),error:E}).then(Dn)}function vn(E){let T=uE(YE).find(t=>t.type==="user_question");T&&T.type==="user_question"&&(Pe("update_sql","POST",{id:uE(GE),sql:E}).then(Se).then(t=>{t.type==="sql"&&(Et.set(t.text),Pe("run_sql","GET",{id:t.id}).then(Se).then(r=>{r.type==="df"?JSON.parse(r.df).length>1?Pe("generate_plotly_figure","GET",{id:r.id}).then(Se).then(n=>{Se({type:"feedback_question"}),Se({type:"feedback_buttons"})}):(Se({type:"feedback_question"}),Se({type:"feedback_buttons"})):(Se({type:"feedback_question"}),Se({type:"feedback_buttons"}))}))}),at("user_sql"))}function Fn(){Se({type:"chart_modification"})}function Yn(){at("feedback_buttons"),Se({type:"feedback_correct"}),fT()?Pe("create_function","GET",{id:uE(GE)}).then(Se):console.log("No Question-SQL Found")}function Vn(E,e){Pe("update_function","POST",{old_function_name:E,updated_function:e})}function Wn(E){Pe("delete_function","POST",{function_name:E}).finally(()=>{cT()})}function wn(){at("feedback_buttons"),Se({type:"feedback_correct"});let E=fT();E&&(Pe("train","POST",E),Pe("generate_followup_questions","GET",{id:uE(GE)}).then(Se))}function wT(){at("feedback_buttons"),Se({type:"feedback_incorrect"}),Se({type:"user_sql"})}function $n(E){at("chart_modification"),Se({type:"user_question",question:"Update the chart with these instructions: "+E}),Pe("generate_plotly_figure","GET",{id:uE(GE),chart_instructions:E}).then(Se)}function xn(){var E=new WebSocket("ws://"+window.location.host+"/api/v0/log");E.onopen=function(){console.log("Connected to WebSocket server at /log.")},E.onmessage=function(e){console.log("Received message:",e.data);try{var T=JSON.parse(e.data)}catch(r){console.error("Error parsing JSON:",r);return}var t=document.getElementById("log-contents");t&&(t.innerHTML+="<details> <summary>"+T.title+"</summary> "+JSON.stringify(T.message)+"</details> <br>")},E.onclose=function(e){console.log("WebSocket connection closed:",e)},E.onerror=function(e){console.error("WebSocket error:",e)}}function $T(E,e,T){const t=E.slice();return t[3]=e[T],t}function xT(E){let e,T,t,r;return{c(){e=f("li"),T=f("button"),T.innerHTML=`<svg class="w-3.5 h-3.5" fill="none" stroke="currentColor" stroke-width="1.5" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg" aria-hidden="true"><path stroke-linecap="round" stroke-linejoin="round" d="M4.26 10.147a60.436 60.436 0 00-.491 6.347A48.627 48.627 0 0112 20.904a48.627 48.627 0 018.232-4.41 60.46 60.46 0 00-.491-6.347m-15.482 0a50.57 50.57 0 00-2.658-.813A59.905 59.905 0 0112 3.493a59.902 59.902 0 0110.399 5.84c-.896.248-1.783.52-2.658.814m-15.482 0A50.697 50.697 0 0112 13.489a50.702 50.702 0 017.74-3.342M6.75 15a.75.75 0 100-1.5.75.75 0 000 1.5zm0 0v-3.675A55.378 55.378 0 0112 8.443m-7.007 11.55A5.981 5.981 0 006.75 15.75v-1.5"></path></svg>
              Functions`,a(T,"class","flex items-center gap-x-3 py-2 px-3 text-sm text-slate-700 rounded-md hover:bg-gray-100 dark:hover:bg-gray-900 dark:text-slate-400 dark:hover:text-slate-300 border border-gray-200 dark:border-gray-700 w-full")},m(R,n){V(R,e,n),l(e,T),t||(r=Ne(T,"click",cT),t=!0)},d(R){R&&Y(e),t=!1,r()}}}function XT(E){let e,T,t,r;return{c(){e=f("li"),T=f("button"),T.innerHTML=`<svg class="w-3.5 h-3.5" fill="none" stroke="currentColor" stroke-width="1.5" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg" aria-hidden="true"><path stroke-linecap="round" stroke-linejoin="round" d="M4.26 10.147a60.436 60.436 0 00-.491 6.347A48.627 48.627 0 0112 20.904a48.627 48.627 0 018.232-4.41 60.46 60.46 0 00-.491-6.347m-15.482 0a50.57 50.57 0 00-2.658-.813A59.905 59.905 0 0112 3.493a59.902 59.902 0 0110.399 5.84c-.896.248-1.783.52-2.658.814m-15.482 0A50.697 50.697 0 0112 13.489a50.702 50.702 0 017.74-3.342M6.75 15a.75.75 0 100-1.5.75.75 0 000 1.5zm0 0v-3.675A55.378 55.378 0 0112 8.443m-7.007 11.55A5.981 5.981 0 006.75 15.75v-1.5"></path></svg>
              Training Data`,a(T,"class","flex items-center gap-x-3 py-2 px-3 text-sm text-slate-700 rounded-md hover:bg-gray-100 dark:hover:bg-gray-900 dark:text-slate-400 dark:hover:text-slate-300 border border-gray-200 dark:border-gray-700 w-full")},m(R,n){V(R,e,n),l(e,T),t||(r=Ne(T,"click",hR),t=!0)},d(R){R&&Y(e),t=!1,r()}}}function kT(E){let e,T,t,r,R,n=E[3].question+"",s,S,A,o;function i(){return E[2](E[3])}return{c(){e=f("li"),T=f("button"),t=OE("svg"),r=OE("path"),R=$(),s=te(n),S=$(),a(r,"stroke-linecap","round"),a(r,"stroke-linejoin","round"),a(r,"d","M7.5 8.25h9m-9 3H12m-9.75 1.51c0 1.6 1.123 2.994 2.707 3.227 1.129.166 2.27.293 3.423.379.35.026.67.21.865.501L12 21l2.755-4.133a1.14 1.14 0 01.865-.501 48.172 48.172 0 003.423-.379c1.584-.233 2.707-1.626 2.707-3.228V6.741c0-1.602-1.123-2.995-2.707-3.228A48.394 48.394 0 0012 3c-2.392 0-4.744.175-7.043.513C3.373 3.746 2.25 5.14 2.25 6.741v6.018z"),a(t,"class","w-3.5 h-3.5"),a(t,"fill","none"),a(t,"stroke","currentColor"),a(t,"stroke-width","1.5"),a(t,"viewBox","0 0 24 24"),a(t,"xmlns","http://www.w3.org/2000/svg"),a(t,"aria-hidden","true"),a(T,"class","flex items-center text-left gap-x-3 py-2 px-3 text-sm text-slate-700 rounded-md hover:bg-gray-100 dark:hover:bg-gray-900 dark:text-slate-400 dark:hover:text-slate-300")},m(_,c){V(_,e,c),l(e,T),l(T,t),l(t,r),l(T,R),l(T,s),l(e,S),A||(o=Ne(T,"click",i),A=!0)},p(_,c){E=_,c&2&&n!==(n=E[3].question+"")&&Le(s,n)},d(_){_&&Y(e),A=!1,o()}}}function Xn(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p,C,L,I,u,H,b,M=E[0].version+"",O,N,D,B,h,G=E[0].function_generation&&xT(),F=E[0].show_training_data&&XT(),W=De(E[1]),x=[];for(let J=0;J<W.length;J+=1)x[J]=kT($T(E,W,J));return{c(){e=f("div"),T=f("nav"),t=f("div"),r=f("img"),n=$(),s=f("div"),s.innerHTML='<button type="button" class="w-8 h-8 inline-flex justify-center items-center gap-2 rounded-md text-gray-700 align-middle focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white focus:ring-blue-600 transition-all dark:text-gray-400 dark:focus:ring-offset-gray-800" data-hs-overlay="#application-sidebar" aria-controls="application-sidebar" aria-label="Toggle navigation"><svg class="w-4 h-4" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path d="M2.146 2.854a.5.5 0 1 1 .708-.708L8 7.293l5.146-5.147a.5.5 0 0 1 .708.708L8.707 8l5.147 5.146a.5.5 0 0 1-.708.708L8 8.707l-5.146 5.147a.5.5 0 0 1-.708-.708L7.293 8 2.146 2.854Z"></path></svg> <span class="sr-only">Sidebar</span></button>',S=$(),A=f("div"),o=f("ul"),G&&G.c(),i=$(),F&&F.c(),_=$(),c=f("li"),P=f("button"),P.innerHTML=`<svg class="w-3.5 h-3.5" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path fill-rule="evenodd" clip-rule="evenodd" d="M8 2C8.47339 2 8.85714 2.38376 8.85714 2.85714V7.14286L13.1429 7.14286C13.6162 7.14286 14 7.52661 14 8C14 8.47339 13.6162 8.85714 13.1429 8.85714L8.85714 8.85715V13.1429C8.85714 13.6162 8.47339 14 8 14C7.52661 14 7.14286 13.6162 7.14286 13.1429V8.85715L2.85714 8.85715C2.38376 8.85715 2 8.4734 2 8.00001C2 7.52662 2.38376 7.14287 2.85714 7.14287L7.14286 7.14286V2.85714C7.14286 2.38376 7.52661 2 8 2Z" fill="currentColor"></path></svg>
              New question`,p=$();for(let J=0;J<x.length;J+=1)x[J].c();C=$(),L=f("div"),I=f("div"),u=f("p"),H=f("span"),b=te(`
            v`),O=te(M),N=$(),D=f("div"),D.innerHTML=`<a class="flex justify-between items-center gap-x-3 py-2 px-3 text-sm text-slate-700 rounded-md hover:bg-gray-100 dark:hover:bg-gray-900 dark:text-slate-400 dark:hover:text-slate-300" href="/auth/logout">Sign out
            <svg class="w-3.5 h-3.5" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path fill-rule="evenodd" d="M10 3.5a.5.5 0 0 0-.5-.5h-8a.5.5 0 0 0-.5.5v9a.5.5 0 0 0 .5.5h8a.5.5 0 0 0 .5-.5v-2a.5.5 0 0 1 1 0v2A1.5 1.5 0 0 1 9.5 14h-8A1.5 1.5 0 0 1 0 12.5v-9A1.5 1.5 0 0 1 1.5 2h8A1.5 1.5 0 0 1 11 3.5v2a.5.5 0 0 1-1 0v-2z"></path><path fill-rule="evenodd" d="M4.146 8.354a.5.5 0 0 1 0-.708l3-3a.5.5 0 1 1 .708.708L5.707 7.5H14.5a.5.5 0 0 1 0 1H5.707l2.147 2.146a.5.5 0 0 1-.708.708l-3-3z"></path></svg></a>`,a(r,"class","w-28 h-auto"),VT(r.src,R=E[0].logo)||a(r,"src",R),a(r,"alt","Vanna Logo"),a(s,"class","lg:hidden"),a(t,"class","flex items-center justify-between py-4 pr-4 pl-7"),a(P,"class","w-full py-2 px-4 inline-flex items-center gap-x-2 text-sm font-semibold rounded-lg border border-transparent bg-blue-600 text-white hover:bg-blue-700 disabled:opacity-50 disabled:pointer-events-none"),a(o,"class","space-y-1.5 p-4"),a(A,"class","h-full"),a(H,"class","block w-1.5 h-1.5 rounded-full bg-green-600"),a(u,"class","inline-flex items-center gap-x-2 text-xs text-green-600"),a(I,"class","py-2.5 px-7"),a(D,"class","p-4 border-t border-gray-200 dark:border-gray-700"),a(L,"class","mt-auto"),a(T,"class","hs-accordion-group w-full h-full flex flex-col"),a(T,"data-hs-accordion-always-open",""),a(e,"id","application-sidebar"),a(e,"class","hs-overlay hs-overlay-open:translate-x-0 -translate-x-full transition-all duration-300 transform hidden fixed top-0 left-0 bottom-0 z-[60] w-64 bg-white border-r border-gray-200 overflow-y-auto scrollbar-y lg:block lg:translate-x-0 lg:right-auto lg:bottom-0 dark:scrollbar-y dark:bg-slate-900 dark:border-gray-700")},m(J,oe){V(J,e,oe),l(e,T),l(T,t),l(t,r),l(t,n),l(t,s),l(T,S),l(T,A),l(A,o),G&&G.m(o,null),l(o,i),F&&F.m(o,null),l(o,_),l(o,c),l(c,P),l(o,p);for(let z=0;z<x.length;z+=1)x[z]&&x[z].m(o,null);l(T,C),l(T,L),l(L,I),l(I,u),l(u,H),l(u,b),l(u,O),l(L,N),l(L,D),B||(h=Ne(P,"click",it),B=!0)},p(J,[oe]){if(oe&1&&!VT(r.src,R=J[0].logo)&&a(r,"src",R),J[0].function_generation?G||(G=xT(),G.c(),G.m(o,i)):G&&(G.d(1),G=null),J[0].show_training_data?F||(F=XT(),F.c(),F.m(o,_)):F&&(F.d(1),F=null),oe&2){W=De(J[1]);let z;for(z=0;z<W.length;z+=1){const Oe=$T(J,W,z);x[z]?x[z].p(Oe,oe):(x[z]=kT(Oe),x[z].c(),x[z].m(o,null))}for(;z<x.length;z+=1)x[z].d(1);x.length=W.length}oe&1&&M!==(M=J[0].version+"")&&Le(O,M)},i:j,o:j,d(J){J&&Y(e),G&&G.d(),F&&F.d(),nE(x,J),B=!1,h()}}}function kn(E,e,T){let t,r;return eE(E,VE,n=>T(0,t=n)),eE(E,CT,n=>T(1,r=n)),[t,r,n=>{Mn(n.id)}]}class Kn extends ue{constructor(e){super(),Ce(this,e,kn,Xn,_e,{})}}var Jn=typeof globalThis<"u"?globalThis:typeof window<"u"?window:typeof global<"u"?global:typeof self<"u"?self:{};function qn(E){return E&&E.__esModule&&Object.prototype.hasOwnProperty.call(E,"default")?E.default:E}var Qn={exports:{}};/*! For license information please see preline.js.LICENSE.txt */(function(E,e){(function(T,t){E.exports=t()})(self,function(){return(()=>{var T={661:(n,s,S)=>{function A(p){return A=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(C){return typeof C}:function(C){return C&&typeof Symbol=="function"&&C.constructor===Symbol&&C!==Symbol.prototype?"symbol":typeof C},A(p)}function o(p,C){for(var L=0;L<C.length;L++){var I=C[L];I.enumerable=I.enumerable||!1,I.configurable=!0,"value"in I&&(I.writable=!0),Object.defineProperty(p,I.key,I)}}function i(p,C){return i=Object.setPrototypeOf||function(L,I){return L.__proto__=I,L},i(p,C)}function _(p,C){if(C&&(A(C)==="object"||typeof C=="function"))return C;if(C!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(L){if(L===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return L}(p)}function c(p){return c=Object.setPrototypeOf?Object.getPrototypeOf:function(C){return C.__proto__||Object.getPrototypeOf(C)},c(p)}var P=function(p){(function(M,O){if(typeof O!="function"&&O!==null)throw new TypeError("Super expression must either be null or a function");M.prototype=Object.create(O&&O.prototype,{constructor:{value:M,writable:!0,configurable:!0}}),Object.defineProperty(M,"prototype",{writable:!1}),O&&i(M,O)})(b,p);var C,L,I,u,H=(I=b,u=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var M,O=c(I);if(u){var N=c(this).constructor;M=Reflect.construct(O,arguments,N)}else M=O.apply(this,arguments);return _(this,M)});function b(){return function(M,O){if(!(M instanceof O))throw new TypeError("Cannot call a class as a function")}(this,b),H.call(this,".hs-accordion")}return C=b,(L=[{key:"init",value:function(){var M=this;document.addEventListener("click",function(O){var N=O.target,D=N.closest(M.selector),B=N.closest(".hs-accordion-toggle"),h=N.closest(".hs-accordion-group");D&&h&&B&&(M._hideAll(D),M.show(D))})}},{key:"show",value:function(M){var O=this;if(M.classList.contains("active"))return this.hide(M);M.classList.add("active");var N=M.querySelector(".hs-accordion-content");N.style.display="block",N.style.height=0,setTimeout(function(){N.style.height="".concat(N.scrollHeight,"px")}),this.afterTransition(N,function(){M.classList.contains("active")&&(N.style.height="",O._fireEvent("open",M),O._dispatch("open.hs.accordion",M,M))})}},{key:"hide",value:function(M){var O=this,N=M.querySelector(".hs-accordion-content");N.style.height="".concat(N.scrollHeight,"px"),setTimeout(function(){N.style.height=0}),this.afterTransition(N,function(){M.classList.contains("active")||(N.style.display="",O._fireEvent("hide",M),O._dispatch("hide.hs.accordion",M,M))}),M.classList.remove("active")}},{key:"_hideAll",value:function(M){var O=this,N=M.closest(".hs-accordion-group");N.hasAttribute("data-hs-accordion-always-open")||N.querySelectorAll(this.selector).forEach(function(D){M!==D&&O.hide(D)})}}])&&o(C.prototype,L),Object.defineProperty(C,"prototype",{writable:!1}),b}(S(765).Z);window.HSAccordion=new P,document.addEventListener("load",window.HSAccordion.init())},795:(n,s,S)=>{function A(C){return A=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(L){return typeof L}:function(L){return L&&typeof Symbol=="function"&&L.constructor===Symbol&&L!==Symbol.prototype?"symbol":typeof L},A(C)}function o(C,L){(L==null||L>C.length)&&(L=C.length);for(var I=0,u=new Array(L);I<L;I++)u[I]=C[I];return u}function i(C,L){for(var I=0;I<L.length;I++){var u=L[I];u.enumerable=u.enumerable||!1,u.configurable=!0,"value"in u&&(u.writable=!0),Object.defineProperty(C,u.key,u)}}function _(C,L){return _=Object.setPrototypeOf||function(I,u){return I.__proto__=u,I},_(C,L)}function c(C,L){if(L&&(A(L)==="object"||typeof L=="function"))return L;if(L!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(I){if(I===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return I}(C)}function P(C){return P=Object.setPrototypeOf?Object.getPrototypeOf:function(L){return L.__proto__||Object.getPrototypeOf(L)},P(C)}var p=function(C){(function(O,N){if(typeof N!="function"&&N!==null)throw new TypeError("Super expression must either be null or a function");O.prototype=Object.create(N&&N.prototype,{constructor:{value:O,writable:!0,configurable:!0}}),Object.defineProperty(O,"prototype",{writable:!1}),N&&_(O,N)})(M,C);var L,I,u,H,b=(u=M,H=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var O,N=P(u);if(H){var D=P(this).constructor;O=Reflect.construct(N,arguments,D)}else O=N.apply(this,arguments);return c(this,O)});function M(){return function(O,N){if(!(O instanceof N))throw new TypeError("Cannot call a class as a function")}(this,M),b.call(this,"[data-hs-collapse]")}return L=M,(I=[{key:"init",value:function(){var O=this;document.addEventListener("click",function(N){var D=N.target.closest(O.selector);if(D){var B=document.querySelectorAll(D.getAttribute("data-hs-collapse"));O.toggle(B)}})}},{key:"toggle",value:function(O){var N,D=this;O.length&&(N=O,function(B){if(Array.isArray(B))return o(B)}(N)||function(B){if(typeof Symbol<"u"&&B[Symbol.iterator]!=null||B["@@iterator"]!=null)return Array.from(B)}(N)||function(B,h){if(B){if(typeof B=="string")return o(B,h);var G=Object.prototype.toString.call(B).slice(8,-1);return G==="Object"&&B.constructor&&(G=B.constructor.name),G==="Map"||G==="Set"?Array.from(B):G==="Arguments"||/^(?:Ui|I)nt(?:8|16|32)(?:Clamped)?Array$/.test(G)?o(B,h):void 0}}(N)||function(){throw new TypeError(`Invalid attempt to spread non-iterable instance.
In order to be iterable, non-array objects must have a [Symbol.iterator]() method.`)}()).forEach(function(B){B.classList.contains("hidden")?D.show(B):D.hide(B)})}},{key:"show",value:function(O){var N=this;O.classList.add("open"),O.classList.remove("hidden"),O.style.height=0,document.querySelectorAll(this.selector).forEach(function(D){O.closest(D.getAttribute("data-hs-collapse"))&&D.classList.add("open")}),O.style.height="".concat(O.scrollHeight,"px"),this.afterTransition(O,function(){O.classList.contains("open")&&(O.style.height="",N._fireEvent("open",O),N._dispatch("open.hs.collapse",O,O))})}},{key:"hide",value:function(O){var N=this;O.style.height="".concat(O.scrollHeight,"px"),setTimeout(function(){O.style.height=0}),O.classList.remove("open"),this.afterTransition(O,function(){O.classList.contains("open")||(O.classList.add("hidden"),O.style.height=null,N._fireEvent("hide",O),N._dispatch("hide.hs.collapse",O,O),O.querySelectorAll(".hs-mega-menu-content.block").forEach(function(D){D.classList.remove("block"),D.classList.add("hidden")}))}),document.querySelectorAll(this.selector).forEach(function(D){O.closest(D.getAttribute("data-hs-collapse"))&&D.classList.remove("open")})}}])&&i(L.prototype,I),Object.defineProperty(L,"prototype",{writable:!1}),M}(S(765).Z);window.HSCollapse=new p,document.addEventListener("load",window.HSCollapse.init())},682:(n,s,S)=>{var A=S(714),o=S(765);const i={historyIndex:-1,addHistory:function(H){this.historyIndex=H},existsInHistory:function(H){return H>this.historyIndex},clearHistory:function(){this.historyIndex=-1}};function _(H){return _=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(b){return typeof b}:function(b){return b&&typeof Symbol=="function"&&b.constructor===Symbol&&b!==Symbol.prototype?"symbol":typeof b},_(H)}function c(H){return function(b){if(Array.isArray(b))return P(b)}(H)||function(b){if(typeof Symbol<"u"&&b[Symbol.iterator]!=null||b["@@iterator"]!=null)return Array.from(b)}(H)||function(b,M){if(b){if(typeof b=="string")return P(b,M);var O=Object.prototype.toString.call(b).slice(8,-1);return O==="Object"&&b.constructor&&(O=b.constructor.name),O==="Map"||O==="Set"?Array.from(b):O==="Arguments"||/^(?:Ui|I)nt(?:8|16|32)(?:Clamped)?Array$/.test(O)?P(b,M):void 0}}(H)||function(){throw new TypeError(`Invalid attempt to spread non-iterable instance.
In order to be iterable, non-array objects must have a [Symbol.iterator]() method.`)}()}function P(H,b){(b==null||b>H.length)&&(b=H.length);for(var M=0,O=new Array(b);M<b;M++)O[M]=H[M];return O}function p(H,b){for(var M=0;M<b.length;M++){var O=b[M];O.enumerable=O.enumerable||!1,O.configurable=!0,"value"in O&&(O.writable=!0),Object.defineProperty(H,O.key,O)}}function C(H,b){return C=Object.setPrototypeOf||function(M,O){return M.__proto__=O,M},C(H,b)}function L(H,b){if(b&&(_(b)==="object"||typeof b=="function"))return b;if(b!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(M){if(M===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return M}(H)}function I(H){return I=Object.setPrototypeOf?Object.getPrototypeOf:function(b){return b.__proto__||Object.getPrototypeOf(b)},I(H)}var u=function(H){(function(h,G){if(typeof G!="function"&&G!==null)throw new TypeError("Super expression must either be null or a function");h.prototype=Object.create(G&&G.prototype,{constructor:{value:h,writable:!0,configurable:!0}}),Object.defineProperty(h,"prototype",{writable:!1}),G&&C(h,G)})(B,H);var b,M,O,N,D=(O=B,N=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var h,G=I(O);if(N){var F=I(this).constructor;h=Reflect.construct(G,arguments,F)}else h=G.apply(this,arguments);return L(this,h)});function B(){var h;return function(G,F){if(!(G instanceof F))throw new TypeError("Cannot call a class as a function")}(this,B),(h=D.call(this,".hs-dropdown")).positions={top:"top","top-left":"top-start","top-right":"top-end",bottom:"bottom","bottom-left":"bottom-start","bottom-right":"bottom-end",right:"right","right-top":"right-start","right-bottom":"right-end",left:"left","left-top":"left-start","left-bottom":"left-end"},h.absoluteStrategyModifiers=function(G){return[{name:"applyStyles",fn:function(F){var W=(window.getComputedStyle(G).getPropertyValue("--strategy")||"absolute").replace(" ",""),x=(window.getComputedStyle(G).getPropertyValue("--adaptive")||"adaptive").replace(" ","");F.state.elements.popper.style.position=W,F.state.elements.popper.style.transform=x==="adaptive"?F.state.styles.popper.transform:null,F.state.elements.popper.style.top=null,F.state.elements.popper.style.bottom=null,F.state.elements.popper.style.left=null,F.state.elements.popper.style.right=null,F.state.elements.popper.style.margin=0}},{name:"computeStyles",options:{adaptive:!1}}]},h._history=i,h}return b=B,M=[{key:"init",value:function(){var h=this;document.addEventListener("click",function(G){var F=G.target,W=F.closest(h.selector),x=F.closest(".hs-dropdown-menu");if(W&&W.classList.contains("open")||h._closeOthers(W),x){var J=(window.getComputedStyle(W).getPropertyValue("--auto-close")||"").replace(" ","");if((J=="false"||J=="inside")&&!W.parentElement.closest(h.selector))return}W&&(W.classList.contains("open")?h.close(W):h.open(W))}),document.addEventListener("mousemove",function(G){var F=G.target,W=F.closest(h.selector);if(F.closest(".hs-dropdown-menu"),W){var x=(window.getComputedStyle(W).getPropertyValue("--trigger")||"click").replace(" ","");if(x!=="hover")return;W&&W.classList.contains("open")||h._closeOthers(W),x!=="hover"||W.classList.contains("open")||/iPad|iPhone|iPod/.test(navigator.platform)||navigator.maxTouchPoints&&navigator.maxTouchPoints>2&&/MacIntel/.test(navigator.platform)||navigator.maxTouchPoints&&navigator.maxTouchPoints>2&&/MacIntel/.test(navigator.platform)||h._hover(F)}}),document.addEventListener("keydown",this._keyboardSupport.bind(this)),window.addEventListener("resize",function(){document.querySelectorAll(".hs-dropdown.open").forEach(function(G){h.close(G,!0)})})}},{key:"_closeOthers",value:function(){var h=this,G=arguments.length>0&&arguments[0]!==void 0?arguments[0]:null,F=document.querySelectorAll("".concat(this.selector,".open"));F.forEach(function(W){if(!G||G.closest(".hs-dropdown.open")!==W){var x=(window.getComputedStyle(W).getPropertyValue("--auto-close")||"").replace(" ","");x!="false"&&x!="outside"&&h.close(W)}})}},{key:"_hover",value:function(h){var G=this,F=h.closest(this.selector);this.open(F),document.addEventListener("mousemove",function W(x){x.target.closest(G.selector)&&x.target.closest(G.selector)!==F.parentElement.closest(G.selector)||(G.close(F),document.removeEventListener("mousemove",W,!0))},!0)}},{key:"close",value:function(h){var G=this,F=arguments.length>1&&arguments[1]!==void 0&&arguments[1],W=h.querySelector(".hs-dropdown-menu"),x=function(){h.classList.contains("open")||(W.classList.remove("block"),W.classList.add("hidden"),W.style.inset=null,W.style.position=null,h._popper&&h._popper.destroy())};F||this.afterTransition(h.querySelector("[data-hs-dropdown-transition]")||W,function(){x()}),W.style.margin=null,h.classList.remove("open"),F&&x(),this._fireEvent("close",h),this._dispatch("close.hs.dropdown",h,h);var J=W.querySelectorAll(".hs-dropdown.open");J.forEach(function(oe){G.close(oe,!0)})}},{key:"open",value:function(h){var G=h.querySelector(".hs-dropdown-menu"),F=(window.getComputedStyle(h).getPropertyValue("--placement")||"").replace(" ",""),W=(window.getComputedStyle(h).getPropertyValue("--strategy")||"fixed").replace(" ",""),x=((window.getComputedStyle(h).getPropertyValue("--adaptive")||"adaptive").replace(" ",""),parseInt((window.getComputedStyle(h).getPropertyValue("--offset")||"10").replace(" ","")));if(W!=="static"){h._popper&&h._popper.destroy();var J=(0,A.fi)(h,G,{placement:this.positions[F]||"bottom-start",strategy:W,modifiers:[].concat(c(W!=="fixed"?this.absoluteStrategyModifiers(h):[]),[{name:"offset",options:{offset:[0,x]}}])});h._popper=J}G.style.margin=null,G.classList.add("block"),G.classList.remove("hidden"),setTimeout(function(){h.classList.add("open")}),this._fireEvent("open",h),this._dispatch("open.hs.dropdown",h,h)}},{key:"_keyboardSupport",value:function(h){var G=document.querySelector(".hs-dropdown.open");if(G)return h.keyCode===27?(h.preventDefault(),this._esc(G)):h.keyCode===40?(h.preventDefault(),this._down(G)):h.keyCode===38?(h.preventDefault(),this._up(G)):h.keyCode===36?(h.preventDefault(),this._start(G)):h.keyCode===35?(h.preventDefault(),this._end(G)):void this._byChar(G,h.key)}},{key:"_esc",value:function(h){this.close(h)}},{key:"_up",value:function(h){var G=h.querySelector(".hs-dropdown-menu"),F=c(G.querySelectorAll("a")).reverse().filter(function(J){return!J.disabled}),W=G.querySelector("a:focus"),x=F.findIndex(function(J){return J===W});x+1<F.length&&x++,F[x].focus()}},{key:"_down",value:function(h){var G=h.querySelector(".hs-dropdown-menu"),F=c(G.querySelectorAll("a")).filter(function(J){return!J.disabled}),W=G.querySelector("a:focus"),x=F.findIndex(function(J){return J===W});x+1<F.length&&x++,F[x].focus()}},{key:"_start",value:function(h){var G=c(h.querySelector(".hs-dropdown-menu").querySelectorAll("a")).filter(function(F){return!F.disabled});G.length&&G[0].focus()}},{key:"_end",value:function(h){var G=c(h.querySelector(".hs-dropdown-menu").querySelectorAll("a")).reverse().filter(function(F){return!F.disabled});G.length&&G[0].focus()}},{key:"_byChar",value:function(h,G){var F=this,W=c(h.querySelector(".hs-dropdown-menu").querySelectorAll("a")),x=function(){return W.findIndex(function(oe,z){return oe.innerText.toLowerCase().charAt(0)===G.toLowerCase()&&F._history.existsInHistory(z)})},J=x();J===-1&&(this._history.clearHistory(),J=x()),J!==-1&&(W[J].focus(),this._history.addHistory(J))}},{key:"toggle",value:function(h){h.classList.contains("open")?this.close(h):this.open(h)}}],M&&p(b.prototype,M),Object.defineProperty(b,"prototype",{writable:!1}),B}(o.Z);window.HSDropdown=new u,document.addEventListener("load",window.HSDropdown.init())},284:(n,s,S)=>{function A(C){return A=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(L){return typeof L}:function(L){return L&&typeof Symbol=="function"&&L.constructor===Symbol&&L!==Symbol.prototype?"symbol":typeof L},A(C)}function o(C,L){(L==null||L>C.length)&&(L=C.length);for(var I=0,u=new Array(L);I<L;I++)u[I]=C[I];return u}function i(C,L){for(var I=0;I<L.length;I++){var u=L[I];u.enumerable=u.enumerable||!1,u.configurable=!0,"value"in u&&(u.writable=!0),Object.defineProperty(C,u.key,u)}}function _(C,L){return _=Object.setPrototypeOf||function(I,u){return I.__proto__=u,I},_(C,L)}function c(C,L){if(L&&(A(L)==="object"||typeof L=="function"))return L;if(L!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(I){if(I===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return I}(C)}function P(C){return P=Object.setPrototypeOf?Object.getPrototypeOf:function(L){return L.__proto__||Object.getPrototypeOf(L)},P(C)}var p=function(C){(function(O,N){if(typeof N!="function"&&N!==null)throw new TypeError("Super expression must either be null or a function");O.prototype=Object.create(N&&N.prototype,{constructor:{value:O,writable:!0,configurable:!0}}),Object.defineProperty(O,"prototype",{writable:!1}),N&&_(O,N)})(M,C);var L,I,u,H,b=(u=M,H=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var O,N=P(u);if(H){var D=P(this).constructor;O=Reflect.construct(N,arguments,D)}else O=N.apply(this,arguments);return c(this,O)});function M(){var O;return function(N,D){if(!(N instanceof D))throw new TypeError("Cannot call a class as a function")}(this,M),(O=b.call(this,"[data-hs-overlay]")).openNextOverlay=!1,O}return L=M,(I=[{key:"init",value:function(){var O=this;document.addEventListener("click",function(N){var D=N.target.closest(O.selector),B=N.target.closest("[data-hs-overlay-close]"),h=N.target.getAttribute("aria-overlay")==="true";return B?O.close(B.closest(".hs-overlay.open")):D?O.toggle(document.querySelector(D.getAttribute("data-hs-overlay"))):void(h&&O._onBackdropClick(N.target))}),document.addEventListener("keydown",function(N){if(N.keyCode===27){var D=document.querySelector(".hs-overlay.open");if(!D)return;setTimeout(function(){D.getAttribute("data-hs-overlay-keyboard")!=="false"&&O.close(D)})}})}},{key:"toggle",value:function(O){O&&(O.classList.contains("hidden")?this.open(O):this.close(O))}},{key:"open",value:function(O){var N=this;if(O){var D=document.querySelector(".hs-overlay.open"),B=this.getClassProperty(O,"--body-scroll","false")!=="true";if(D)return this.openNextOverlay=!0,this.close(D).then(function(){N.open(O),N.openNextOverlay=!1});B&&(document.body.style.overflow="hidden"),this._buildBackdrop(O),this._checkTimer(O),this._autoHide(O),O.classList.remove("hidden"),O.setAttribute("aria-overlay","true"),O.setAttribute("tabindex","-1"),setTimeout(function(){O.classList.contains("hidden")||(O.classList.add("open"),N._fireEvent("open",O),N._dispatch("open.hs.overlay",O,O),N._focusInput(O))},50)}}},{key:"close",value:function(O){var N=this;return new Promise(function(D){O&&(O.classList.remove("open"),O.removeAttribute("aria-overlay"),O.removeAttribute("tabindex","-1"),N.afterTransition(O,function(){O.classList.contains("open")||(O.classList.add("hidden"),N._destroyBackdrop(),N._fireEvent("close",O),N._dispatch("close.hs.overlay",O,O),document.body.style.overflow="",D(O))}))})}},{key:"_autoHide",value:function(O){var N=this,D=parseInt(this.getClassProperty(O,"--auto-hide","0"));D&&(O.autoHide=setTimeout(function(){N.close(O)},D))}},{key:"_checkTimer",value:function(O){O.autoHide&&(clearTimeout(O.autoHide),delete O.autoHide)}},{key:"_onBackdropClick",value:function(O){this.getClassProperty(O,"--overlay-backdrop","true")!=="static"&&this.close(O)}},{key:"_buildBackdrop",value:function(O){var N,D=this,B=O.getAttribute("data-hs-overlay-backdrop-container")||!1,h=document.createElement("div"),G="transition duration fixed inset-0 z-50 bg-gray-900 bg-opacity-50 dark:bg-opacity-80 hs-overlay-backdrop",F=function(J,oe){var z=typeof Symbol<"u"&&J[Symbol.iterator]||J["@@iterator"];if(!z){if(Array.isArray(J)||(z=function(me,fE){if(me){if(typeof me=="string")return o(me,fE);var rE=Object.prototype.toString.call(me).slice(8,-1);return rE==="Object"&&me.constructor&&(rE=me.constructor.name),rE==="Map"||rE==="Set"?Array.from(me):rE==="Arguments"||/^(?:Ui|I)nt(?:8|16|32)(?:Clamped)?Array$/.test(rE)?o(me,fE):void 0}}(J))||oe&&J&&typeof J.length=="number"){z&&(J=z);var Oe=0,Ue=function(){};return{s:Ue,n:function(){return Oe>=J.length?{done:!0}:{done:!1,value:J[Oe++]}},e:function(me){throw me},f:Ue}}throw new TypeError(`Invalid attempt to iterate non-iterable instance.
In order to be iterable, non-array objects must have a [Symbol.iterator]() method.`)}var ye,TE=!0,dE=!1;return{s:function(){z=z.call(J)},n:function(){var me=z.next();return TE=me.done,me},e:function(me){dE=!0,ye=me},f:function(){try{TE||z.return==null||z.return()}finally{if(dE)throw ye}}}}(O.classList.values());try{for(F.s();!(N=F.n()).done;){var W=N.value;W.startsWith("hs-overlay-backdrop-open:")&&(G+=" ".concat(W))}}catch(J){F.e(J)}finally{F.f()}var x=this.getClassProperty(O,"--overlay-backdrop","true")!=="static";this.getClassProperty(O,"--overlay-backdrop","true")==="false"||(B&&((h=document.querySelector(B).cloneNode(!0)).classList.remove("hidden"),G=h.classList,h.classList=""),x&&h.addEventListener("click",function(){return D.close(O)},!0),h.setAttribute("data-hs-overlay-backdrop-template",""),document.body.appendChild(h),setTimeout(function(){h.classList=G}))}},{key:"_destroyBackdrop",value:function(){var O=document.querySelector("[data-hs-overlay-backdrop-template]");O&&(this.openNextOverlay&&(O.style.transitionDuration="".concat(1.8*parseFloat(window.getComputedStyle(O).transitionDuration.replace(/[^\\d.-]/g,"")),"s")),O.classList.add("opacity-0"),this.afterTransition(O,function(){O.remove()}))}},{key:"_focusInput",value:function(O){var N=O.querySelector("[autofocus]");N&&N.focus()}}])&&i(L.prototype,I),Object.defineProperty(L,"prototype",{writable:!1}),M}(S(765).Z);window.HSOverlay=new p,document.addEventListener("load",window.HSOverlay.init())},181:(n,s,S)=>{function A(p){return A=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(C){return typeof C}:function(C){return C&&typeof Symbol=="function"&&C.constructor===Symbol&&C!==Symbol.prototype?"symbol":typeof C},A(p)}function o(p,C){for(var L=0;L<C.length;L++){var I=C[L];I.enumerable=I.enumerable||!1,I.configurable=!0,"value"in I&&(I.writable=!0),Object.defineProperty(p,I.key,I)}}function i(p,C){return i=Object.setPrototypeOf||function(L,I){return L.__proto__=I,L},i(p,C)}function _(p,C){if(C&&(A(C)==="object"||typeof C=="function"))return C;if(C!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(L){if(L===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return L}(p)}function c(p){return c=Object.setPrototypeOf?Object.getPrototypeOf:function(C){return C.__proto__||Object.getPrototypeOf(C)},c(p)}var P=function(p){(function(M,O){if(typeof O!="function"&&O!==null)throw new TypeError("Super expression must either be null or a function");M.prototype=Object.create(O&&O.prototype,{constructor:{value:M,writable:!0,configurable:!0}}),Object.defineProperty(M,"prototype",{writable:!1}),O&&i(M,O)})(b,p);var C,L,I,u,H=(I=b,u=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var M,O=c(I);if(u){var N=c(this).constructor;M=Reflect.construct(O,arguments,N)}else M=O.apply(this,arguments);return _(this,M)});function b(){return function(M,O){if(!(M instanceof O))throw new TypeError("Cannot call a class as a function")}(this,b),H.call(this,"[data-hs-remove-element]")}return C=b,(L=[{key:"init",value:function(){var M=this;document.addEventListener("click",function(O){var N=O.target.closest(M.selector);if(N){var D=document.querySelector(N.getAttribute("data-hs-remove-element"));D&&(D.classList.add("hs-removing"),M.afterTransition(D,function(){D.remove()}))}})}}])&&o(C.prototype,L),Object.defineProperty(C,"prototype",{writable:!1}),b}(S(765).Z);window.HSRemoveElement=new P,document.addEventListener("load",window.HSRemoveElement.init())},778:(n,s,S)=>{function A(p){return A=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(C){return typeof C}:function(C){return C&&typeof Symbol=="function"&&C.constructor===Symbol&&C!==Symbol.prototype?"symbol":typeof C},A(p)}function o(p,C){for(var L=0;L<C.length;L++){var I=C[L];I.enumerable=I.enumerable||!1,I.configurable=!0,"value"in I&&(I.writable=!0),Object.defineProperty(p,I.key,I)}}function i(p,C){return i=Object.setPrototypeOf||function(L,I){return L.__proto__=I,L},i(p,C)}function _(p,C){if(C&&(A(C)==="object"||typeof C=="function"))return C;if(C!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(L){if(L===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return L}(p)}function c(p){return c=Object.setPrototypeOf?Object.getPrototypeOf:function(C){return C.__proto__||Object.getPrototypeOf(C)},c(p)}var P=function(p){(function(M,O){if(typeof O!="function"&&O!==null)throw new TypeError("Super expression must either be null or a function");M.prototype=Object.create(O&&O.prototype,{constructor:{value:M,writable:!0,configurable:!0}}),Object.defineProperty(M,"prototype",{writable:!1}),O&&i(M,O)})(b,p);var C,L,I,u,H=(I=b,u=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var M,O=c(I);if(u){var N=c(this).constructor;M=Reflect.construct(O,arguments,N)}else M=O.apply(this,arguments);return _(this,M)});function b(){var M;return function(O,N){if(!(O instanceof N))throw new TypeError("Cannot call a class as a function")}(this,b),(M=H.call(this,"[data-hs-scrollspy] ")).activeSection=null,M}return C=b,(L=[{key:"init",value:function(){var M=this;document.querySelectorAll(this.selector).forEach(function(O){var N=document.querySelector(O.getAttribute("data-hs-scrollspy")),D=O.querySelectorAll("[href]"),B=N.children,h=O.getAttribute("data-hs-scrollspy-scrollable-parent")?document.querySelector(O.getAttribute("data-hs-scrollspy-scrollable-parent")):document;Array.from(B).forEach(function(G){G.getAttribute("id")&&h.addEventListener("scroll",function(F){return M._update({$scrollspyEl:O,$scrollspyContentEl:N,links:D,$sectionEl:G,sections:B,ev:F})})}),D.forEach(function(G){G.addEventListener("click",function(F){F.preventDefault(),G.getAttribute("href")!=="javascript:;"&&M._scrollTo({$scrollspyEl:O,$scrollableEl:h,$link:G})})})})}},{key:"_update",value:function(M){var O=M.ev,N=M.$scrollspyEl,D=(M.sections,M.links),B=M.$sectionEl,h=parseInt(this.getClassProperty(N,"--scrollspy-offset","0")),G=this.getClassProperty(B,"--scrollspy-offset")||h,F=O.target===document?0:parseInt(O.target.getBoundingClientRect().top),W=parseInt(B.getBoundingClientRect().top)-G-F,x=B.offsetHeight;if(W<=0&&W+x>0){if(this.activeSection===B)return;D.forEach(function(Oe){Oe.classList.remove("active")});var J=N.querySelector('[href="#'.concat(B.getAttribute("id"),'"]'));if(J){J.classList.add("active");var oe=J.closest("[data-hs-scrollspy-group]");if(oe){var z=oe.querySelector("[href]");z&&z.classList.add("active")}}this.activeSection=B}}},{key:"_scrollTo",value:function(M){var O=M.$scrollspyEl,N=M.$scrollableEl,D=M.$link,B=document.querySelector(D.getAttribute("href")),h=parseInt(this.getClassProperty(O,"--scrollspy-offset","0")),G=this.getClassProperty(B,"--scrollspy-offset")||h,F=N===document?0:N.offsetTop,W=B.offsetTop-G-F,x=N===document?window:N;this._fireEvent("scroll",O),this._dispatch("scroll.hs.scrollspy",O,O),window.history.replaceState(null,null,D.getAttribute("href")),x.scrollTo({top:W,left:0,behavior:"smooth"})}}])&&o(C.prototype,L),Object.defineProperty(C,"prototype",{writable:!1}),b}(S(765).Z);window.HSScrollspy=new P,document.addEventListener("load",window.HSScrollspy.init())},51:(n,s,S)=>{function A(L){return A=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(I){return typeof I}:function(I){return I&&typeof Symbol=="function"&&I.constructor===Symbol&&I!==Symbol.prototype?"symbol":typeof I},A(L)}function o(L){return function(I){if(Array.isArray(I))return i(I)}(L)||function(I){if(typeof Symbol<"u"&&I[Symbol.iterator]!=null||I["@@iterator"]!=null)return Array.from(I)}(L)||function(I,u){if(I){if(typeof I=="string")return i(I,u);var H=Object.prototype.toString.call(I).slice(8,-1);return H==="Object"&&I.constructor&&(H=I.constructor.name),H==="Map"||H==="Set"?Array.from(I):H==="Arguments"||/^(?:Ui|I)nt(?:8|16|32)(?:Clamped)?Array$/.test(H)?i(I,u):void 0}}(L)||function(){throw new TypeError(`Invalid attempt to spread non-iterable instance.
In order to be iterable, non-array objects must have a [Symbol.iterator]() method.`)}()}function i(L,I){(I==null||I>L.length)&&(I=L.length);for(var u=0,H=new Array(I);u<I;u++)H[u]=L[u];return H}function _(L,I){for(var u=0;u<I.length;u++){var H=I[u];H.enumerable=H.enumerable||!1,H.configurable=!0,"value"in H&&(H.writable=!0),Object.defineProperty(L,H.key,H)}}function c(L,I){return c=Object.setPrototypeOf||function(u,H){return u.__proto__=H,u},c(L,I)}function P(L,I){if(I&&(A(I)==="object"||typeof I=="function"))return I;if(I!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(u){if(u===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return u}(L)}function p(L){return p=Object.setPrototypeOf?Object.getPrototypeOf:function(I){return I.__proto__||Object.getPrototypeOf(I)},p(L)}var C=function(L){(function(N,D){if(typeof D!="function"&&D!==null)throw new TypeError("Super expression must either be null or a function");N.prototype=Object.create(D&&D.prototype,{constructor:{value:N,writable:!0,configurable:!0}}),Object.defineProperty(N,"prototype",{writable:!1}),D&&c(N,D)})(O,L);var I,u,H,b,M=(H=O,b=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var N,D=p(H);if(b){var B=p(this).constructor;N=Reflect.construct(D,arguments,B)}else N=D.apply(this,arguments);return P(this,N)});function O(){return function(N,D){if(!(N instanceof D))throw new TypeError("Cannot call a class as a function")}(this,O),M.call(this,"[data-hs-tab]")}return I=O,(u=[{key:"init",value:function(){var N=this;document.addEventListener("keydown",this._keyboardSupport.bind(this)),document.addEventListener("click",function(D){var B=D.target.closest(N.selector);B&&N.open(B)}),document.querySelectorAll("[hs-data-tab-select]").forEach(function(D){var B=document.querySelector(D.getAttribute("hs-data-tab-select"));B&&B.addEventListener("change",function(h){var G=document.querySelector('[data-hs-tab="'.concat(h.target.value,'"]'));G&&N.open(G)})})}},{key:"open",value:function(N){var D=document.querySelector(N.getAttribute("data-hs-tab")),B=o(N.parentElement.children),h=o(D.parentElement.children),G=N.closest("[hs-data-tab-select]"),F=G?document.querySelector(G.getAttribute("data-hs-tab")):null;B.forEach(function(W){return W.classList.remove("active")}),h.forEach(function(W){return W.classList.add("hidden")}),N.classList.add("active"),D.classList.remove("hidden"),this._fireEvent("change",N),this._dispatch("change.hs.tab",N,N),F&&(F.value=N.getAttribute("data-hs-tab"))}},{key:"_keyboardSupport",value:function(N){var D=N.target.closest(this.selector);if(D){var B=D.closest('[role="tablist"]').getAttribute("data-hs-tabs-vertical")==="true";return(B?N.keyCode===38:N.keyCode===37)?(N.preventDefault(),this._left(D)):(B?N.keyCode===40:N.keyCode===39)?(N.preventDefault(),this._right(D)):N.keyCode===36?(N.preventDefault(),this._start(D)):N.keyCode===35?(N.preventDefault(),this._end(D)):void 0}}},{key:"_right",value:function(N){var D=N.closest('[role="tablist"]');if(D){var B=o(D.querySelectorAll(this.selector)).filter(function(F){return!F.disabled}),h=D.querySelector("button:focus"),G=B.findIndex(function(F){return F===h});G+1<B.length?G++:G=0,B[G].focus(),this.open(B[G])}}},{key:"_left",value:function(N){var D=N.closest('[role="tablist"]');if(D){var B=o(D.querySelectorAll(this.selector)).filter(function(F){return!F.disabled}).reverse(),h=D.querySelector("button:focus"),G=B.findIndex(function(F){return F===h});G+1<B.length?G++:G=0,B[G].focus(),this.open(B[G])}}},{key:"_start",value:function(N){var D=N.closest('[role="tablist"]');if(D){var B=o(D.querySelectorAll(this.selector)).filter(function(h){return!h.disabled});B.length&&(B[0].focus(),this.open(B[0]))}}},{key:"_end",value:function(N){var D=N.closest('[role="tablist"]');if(D){var B=o(D.querySelectorAll(this.selector)).reverse().filter(function(h){return!h.disabled});B.length&&(B[0].focus(),this.open(B[0]))}}}])&&_(I.prototype,u),Object.defineProperty(I,"prototype",{writable:!1}),O}(S(765).Z);window.HSTabs=new C,document.addEventListener("load",window.HSTabs.init())},185:(n,s,S)=>{var A=S(765),o=S(714);function i(L){return i=typeof Symbol=="function"&&typeof Symbol.iterator=="symbol"?function(I){return typeof I}:function(I){return I&&typeof Symbol=="function"&&I.constructor===Symbol&&I!==Symbol.prototype?"symbol":typeof I},i(L)}function _(L,I){for(var u=0;u<I.length;u++){var H=I[u];H.enumerable=H.enumerable||!1,H.configurable=!0,"value"in H&&(H.writable=!0),Object.defineProperty(L,H.key,H)}}function c(L,I){return c=Object.setPrototypeOf||function(u,H){return u.__proto__=H,u},c(L,I)}function P(L,I){if(I&&(i(I)==="object"||typeof I=="function"))return I;if(I!==void 0)throw new TypeError("Derived constructors may only return object or undefined");return function(u){if(u===void 0)throw new ReferenceError("this hasn't been initialised - super() hasn't been called");return u}(L)}function p(L){return p=Object.setPrototypeOf?Object.getPrototypeOf:function(I){return I.__proto__||Object.getPrototypeOf(I)},p(L)}var C=function(L){(function(N,D){if(typeof D!="function"&&D!==null)throw new TypeError("Super expression must either be null or a function");N.prototype=Object.create(D&&D.prototype,{constructor:{value:N,writable:!0,configurable:!0}}),Object.defineProperty(N,"prototype",{writable:!1}),D&&c(N,D)})(O,L);var I,u,H,b,M=(H=O,b=function(){if(typeof Reflect>"u"||!Reflect.construct||Reflect.construct.sham)return!1;if(typeof Proxy=="function")return!0;try{return Boolean.prototype.valueOf.call(Reflect.construct(Boolean,[],function(){})),!0}catch{return!1}}(),function(){var N,D=p(H);if(b){var B=p(this).constructor;N=Reflect.construct(D,arguments,B)}else N=D.apply(this,arguments);return P(this,N)});function O(){return function(N,D){if(!(N instanceof D))throw new TypeError("Cannot call a class as a function")}(this,O),M.call(this,".hs-tooltip")}return I=O,(u=[{key:"init",value:function(){var N=this;document.addEventListener("click",function(D){var B=D.target.closest(N.selector);B&&N.getClassProperty(B,"--trigger")==="focus"&&N._focus(B),B&&N.getClassProperty(B,"--trigger")==="click"&&N._click(B)}),document.addEventListener("mousemove",function(D){var B=D.target.closest(N.selector);B&&N.getClassProperty(B,"--trigger")!=="focus"&&N.getClassProperty(B,"--trigger")!=="click"&&N._hover(B)})}},{key:"_hover",value:function(N){var D=this;if(!N.classList.contains("show")){var B=N.querySelector(".hs-tooltip-toggle"),h=N.querySelector(".hs-tooltip-content"),G=this.getClassProperty(N,"--placement");(0,o.fi)(B,h,{placement:G||"top",strategy:"fixed",modifiers:[{name:"offset",options:{offset:[0,5]}}]}),this.show(N),N.addEventListener("mouseleave",function F(W){W.relatedTarget.closest(D.selector)&&W.relatedTarget.closest(D.selector)==N||(D.hide(N),N.removeEventListener("mouseleave",F,!0))},!0)}}},{key:"_focus",value:function(N){var D=this,B=N.querySelector(".hs-tooltip-toggle"),h=N.querySelector(".hs-tooltip-content"),G=this.getClassProperty(N,"--placement"),F=this.getClassProperty(N,"--strategy");(0,o.fi)(B,h,{placement:G||"top",strategy:F||"fixed",modifiers:[{name:"offset",options:{offset:[0,5]}}]}),this.show(N),N.addEventListener("blur",function W(){D.hide(N),N.removeEventListener("blur",W,!0)},!0)}},{key:"_click",value:function(N){var D=this;if(!N.classList.contains("show")){var B=N.querySelector(".hs-tooltip-toggle"),h=N.querySelector(".hs-tooltip-content"),G=this.getClassProperty(N,"--placement"),F=this.getClassProperty(N,"--strategy");(0,o.fi)(B,h,{placement:G||"top",strategy:F||"fixed",modifiers:[{name:"offset",options:{offset:[0,5]}}]}),this.show(N);var W=function x(J){setTimeout(function(){D.hide(N),N.removeEventListener("click",x,!0),N.removeEventListener("blur",x,!0)})};N.addEventListener("blur",W,!0),N.addEventListener("click",W,!0)}}},{key:"show",value:function(N){var D=this;N.querySelector(".hs-tooltip-content").classList.remove("hidden"),setTimeout(function(){N.classList.add("show"),D._fireEvent("show",N),D._dispatch("show.hs.tooltip",N,N)})}},{key:"hide",value:function(N){var D=N.querySelector(".hs-tooltip-content");N.classList.remove("show"),this._fireEvent("hide",N),this._dispatch("hide.hs.tooltip",N,N),this.afterTransition(D,function(){N.classList.contains("show")||D.classList.add("hidden")})}}])&&_(I.prototype,u),Object.defineProperty(I,"prototype",{writable:!1}),O}(A.Z);window.HSTooltip=new C,document.addEventListener("load",window.HSTooltip.init())},765:(n,s,S)=>{function A(i,_){for(var c=0;c<_.length;c++){var P=_[c];P.enumerable=P.enumerable||!1,P.configurable=!0,"value"in P&&(P.writable=!0),Object.defineProperty(i,P.key,P)}}S.d(s,{Z:()=>o});var o=function(){function i(P,p){(function(C,L){if(!(C instanceof L))throw new TypeError("Cannot call a class as a function")})(this,i),this.$collection=[],this.selector=P,this.config=p,this.events={}}var _,c;return _=i,c=[{key:"_fireEvent",value:function(P){var p=arguments.length>1&&arguments[1]!==void 0?arguments[1]:null;this.events.hasOwnProperty(P)&&this.events[P](p)}},{key:"_dispatch",value:function(P,p){var C=arguments.length>2&&arguments[2]!==void 0?arguments[2]:null,L=new CustomEvent(P,{detail:{payload:C},bubbles:!0,cancelable:!0,composed:!1});p.dispatchEvent(L)}},{key:"on",value:function(P,p){this.events[P]=p}},{key:"afterTransition",value:function(P,p){window.getComputedStyle(P,null).getPropertyValue("transition")!=="all 0s ease 0s"?P.addEventListener("transitionend",function C(){p(),P.removeEventListener("transitionend",C,!0)},!0):p()}},{key:"getClassProperty",value:function(P,p){var C=arguments.length>2&&arguments[2]!==void 0?arguments[2]:"",L=(window.getComputedStyle(P).getPropertyValue(p)||C).replace(" ","");return L}}],c&&A(_.prototype,c),Object.defineProperty(_,"prototype",{writable:!1}),i}()},714:(n,s,S)=>{function A(U){if(U==null)return window;if(U.toString()!=="[object Window]"){var d=U.ownerDocument;return d&&d.defaultView||window}return U}function o(U){return U instanceof A(U).Element||U instanceof Element}function i(U){return U instanceof A(U).HTMLElement||U instanceof HTMLElement}function _(U){return typeof ShadowRoot<"u"&&(U instanceof A(U).ShadowRoot||U instanceof ShadowRoot)}S.d(s,{fi:()=>tn});var c=Math.max,P=Math.min,p=Math.round;function C(U,d){d===void 0&&(d=!1);var g=U.getBoundingClientRect(),w=1,Z=1;if(i(U)&&d){var q=U.offsetHeight,Q=U.offsetWidth;Q>0&&(w=p(g.width)/Q||1),q>0&&(Z=p(g.height)/q||1)}return{width:g.width/w,height:g.height/Z,top:g.top/Z,right:g.right/w,bottom:g.bottom/Z,left:g.left/w,x:g.left/w,y:g.top/Z}}function L(U){var d=A(U);return{scrollLeft:d.pageXOffset,scrollTop:d.pageYOffset}}function I(U){return U?(U.nodeName||"").toLowerCase():null}function u(U){return((o(U)?U.ownerDocument:U.document)||window.document).documentElement}function H(U){return C(u(U)).left+L(U).scrollLeft}function b(U){return A(U).getComputedStyle(U)}function M(U){var d=b(U),g=d.overflow,w=d.overflowX,Z=d.overflowY;return/auto|scroll|overlay|hidden/.test(g+Z+w)}function O(U,d,g){g===void 0&&(g=!1);var w,Z,q=i(d),Q=i(d)&&function(re){var We=re.getBoundingClientRect(),ae=p(We.width)/re.offsetWidth||1,Me=p(We.height)/re.offsetHeight||1;return ae!==1||Me!==1}(d),ee=u(d),Ee=C(U,Q),ne={scrollLeft:0,scrollTop:0},Ae={x:0,y:0};return(q||!q&&!g)&&((I(d)!=="body"||M(ee))&&(ne=(w=d)!==A(w)&&i(w)?{scrollLeft:(Z=w).scrollLeft,scrollTop:Z.scrollTop}:L(w)),i(d)?((Ae=C(d,!0)).x+=d.clientLeft,Ae.y+=d.clientTop):ee&&(Ae.x=H(ee))),{x:Ee.left+ne.scrollLeft-Ae.x,y:Ee.top+ne.scrollTop-Ae.y,width:Ee.width,height:Ee.height}}function N(U){var d=C(U),g=U.offsetWidth,w=U.offsetHeight;return Math.abs(d.width-g)<=1&&(g=d.width),Math.abs(d.height-w)<=1&&(w=d.height),{x:U.offsetLeft,y:U.offsetTop,width:g,height:w}}function D(U){return I(U)==="html"?U:U.assignedSlot||U.parentNode||(_(U)?U.host:null)||u(U)}function B(U){return["html","body","#document"].indexOf(I(U))>=0?U.ownerDocument.body:i(U)&&M(U)?U:B(D(U))}function h(U,d){var g;d===void 0&&(d=[]);var w=B(U),Z=w===((g=U.ownerDocument)==null?void 0:g.body),q=A(w),Q=Z?[q].concat(q.visualViewport||[],M(w)?w:[]):w,ee=d.concat(Q);return Z?ee:ee.concat(h(D(Q)))}function G(U){return["table","td","th"].indexOf(I(U))>=0}function F(U){return i(U)&&b(U).position!=="fixed"?U.offsetParent:null}function W(U){for(var d=A(U),g=F(U);g&&G(g)&&b(g).position==="static";)g=F(g);return g&&(I(g)==="html"||I(g)==="body"&&b(g).position==="static")?d:g||function(w){var Z=navigator.userAgent.toLowerCase().indexOf("firefox")!==-1;if(navigator.userAgent.indexOf("Trident")!==-1&&i(w)&&b(w).position==="fixed")return null;for(var q=D(w);i(q)&&["html","body"].indexOf(I(q))<0;){var Q=b(q);if(Q.transform!=="none"||Q.perspective!=="none"||Q.contain==="paint"||["transform","perspective"].indexOf(Q.willChange)!==-1||Z&&Q.willChange==="filter"||Z&&Q.filter&&Q.filter!=="none")return q;q=q.parentNode}return null}(U)||d}var x="top",J="bottom",oe="right",z="left",Oe="auto",Ue=[x,J,oe,z],ye="start",TE="end",dE="viewport",me="popper",fE=Ue.reduce(function(U,d){return U.concat([d+"-"+ye,d+"-"+TE])},[]),rE=[].concat(Ue,[Oe]).reduce(function(U,d){return U.concat([d,d+"-"+ye,d+"-"+TE])},[]),pE=["beforeRead","read","afterRead","beforeMain","main","afterMain","beforeWrite","write","afterWrite"];function Tt(U){var d=new Map,g=new Set,w=[];function Z(q){g.add(q.name),[].concat(q.requires||[],q.requiresIfExists||[]).forEach(function(Q){if(!g.has(Q)){var ee=d.get(Q);ee&&Z(ee)}}),w.push(q)}return U.forEach(function(q){d.set(q.name,q)}),U.forEach(function(q){g.has(q.name)||Z(q)}),w}var rt={placement:"bottom",modifiers:[],strategy:"absolute"};function Ve(){for(var U=arguments.length,d=new Array(U),g=0;g<U;g++)d[g]=arguments[g];return!d.some(function(w){return!(w&&typeof w.getBoundingClientRect=="function")})}function ME(U){U===void 0&&(U={});var d=U,g=d.defaultModifiers,w=g===void 0?[]:g,Z=d.defaultOptions,q=Z===void 0?rt:Z;return function(Q,ee,Ee){Ee===void 0&&(Ee=q);var ne,Ae,re={placement:"bottom",orderedModifiers:[],options:Object.assign({},rt,q),modifiersData:{},elements:{reference:Q,popper:ee},attributes:{},styles:{}},We=[],ae=!1,Me={state:re,setOptions:function(ie){var we=typeof ie=="function"?ie(re.options):ie;de(),re.options=Object.assign({},q,re.options,we),re.scrollParents={reference:o(Q)?h(Q):Q.contextElement?h(Q.contextElement):[],popper:h(ee)};var be,ce,he=function(Ie){var le=Tt(Ie);return pE.reduce(function(fe,pe){return fe.concat(le.filter(function(He){return He.phase===pe}))},[])}((be=[].concat(w,re.options.modifiers),ce=be.reduce(function(Ie,le){var fe=Ie[le.name];return Ie[le.name]=fe?Object.assign({},fe,le,{options:Object.assign({},fe.options,le.options),data:Object.assign({},fe.data,le.data)}):le,Ie},{}),Object.keys(ce).map(function(Ie){return ce[Ie]})));return re.orderedModifiers=he.filter(function(Ie){return Ie.enabled}),re.orderedModifiers.forEach(function(Ie){var le=Ie.name,fe=Ie.options,pe=fe===void 0?{}:fe,He=Ie.effect;if(typeof He=="function"){var ke=He({state:re,name:le,instance:Me,options:pe});We.push(ke||function(){})}}),Me.update()},forceUpdate:function(){if(!ae){var ie=re.elements,we=ie.reference,be=ie.popper;if(Ve(we,be)){re.rects={reference:O(we,W(be),re.options.strategy==="fixed"),popper:N(be)},re.reset=!1,re.placement=re.options.placement,re.orderedModifiers.forEach(function(He){return re.modifiersData[He.name]=Object.assign({},He.data)});for(var ce=0;ce<re.orderedModifiers.length;ce++)if(re.reset!==!0){var he=re.orderedModifiers[ce],Ie=he.fn,le=he.options,fe=le===void 0?{}:le,pe=he.name;typeof Ie=="function"&&(re=Ie({state:re,options:fe,name:pe,instance:Me})||re)}else re.reset=!1,ce=-1}}},update:(ne=function(){return new Promise(function(ie){Me.forceUpdate(),ie(re)})},function(){return Ae||(Ae=new Promise(function(ie){Promise.resolve().then(function(){Ae=void 0,ie(ne())})})),Ae}),destroy:function(){de(),ae=!0}};if(!Ve(Q,ee))return Me;function de(){We.forEach(function(ie){return ie()}),We=[]}return Me.setOptions(Ee).then(function(ie){!ae&&Ee.onFirstUpdate&&Ee.onFirstUpdate(ie)}),Me}}var $e={passive:!0};function ze(U){return U.split("-")[0]}function ve(U){return U.split("-")[1]}function Fe(U){return["top","bottom"].indexOf(U)>=0?"x":"y"}function Te(U){var d,g=U.reference,w=U.element,Z=U.placement,q=Z?ze(Z):null,Q=Z?ve(Z):null,ee=g.x+g.width/2-w.width/2,Ee=g.y+g.height/2-w.height/2;switch(q){case x:d={x:ee,y:g.y-w.height};break;case J:d={x:ee,y:g.y+g.height};break;case oe:d={x:g.x+g.width,y:Ee};break;case z:d={x:g.x-w.width,y:Ee};break;default:d={x:g.x,y:g.y}}var ne=q?Fe(q):null;if(ne!=null){var Ae=ne==="y"?"height":"width";switch(Q){case ye:d[ne]=d[ne]-(g[Ae]/2-w[Ae]/2);break;case TE:d[ne]=d[ne]+(g[Ae]/2-w[Ae]/2)}}return d}var Xe={top:"auto",right:"auto",bottom:"auto",left:"auto"};function se(U){var d,g=U.popper,w=U.popperRect,Z=U.placement,q=U.variation,Q=U.offsets,ee=U.position,Ee=U.gpuAcceleration,ne=U.adaptive,Ae=U.roundOffsets,re=U.isFixed,We=Q.x,ae=We===void 0?0:We,Me=Q.y,de=Me===void 0?0:Me,ie=typeof Ae=="function"?Ae({x:ae,y:de}):{x:ae,y:de};ae=ie.x,de=ie.y;var we=Q.hasOwnProperty("x"),be=Q.hasOwnProperty("y"),ce=z,he=x,Ie=window;if(ne){var le=W(g),fe="clientHeight",pe="clientWidth";le===A(g)&&b(le=u(g)).position!=="static"&&ee==="absolute"&&(fe="scrollHeight",pe="scrollWidth"),le=le,(Z===x||(Z===z||Z===oe)&&q===TE)&&(he=J,de-=(re&&Ie.visualViewport?Ie.visualViewport.height:le[fe])-w.height,de*=Ee?1:-1),Z!==z&&(Z!==x&&Z!==J||q!==TE)||(ce=oe,ae-=(re&&Ie.visualViewport?Ie.visualViewport.width:le[pe])-w.width,ae*=Ee?1:-1)}var He,ke=Object.assign({position:ee},ne&&Xe),Ke=Ae===!0?function(sE){var lE=sE.x,PE=sE.y,SE=window.devicePixelRatio||1;return{x:p(lE*SE)/SE||0,y:p(PE*SE)/SE||0}}({x:ae,y:de}):{x:ae,y:de};return ae=Ke.x,de=Ke.y,Ee?Object.assign({},ke,((He={})[he]=be?"0":"",He[ce]=we?"0":"",He.transform=(Ie.devicePixelRatio||1)<=1?"translate("+ae+"px, "+de+"px)":"translate3d("+ae+"px, "+de+"px, 0)",He)):Object.assign({},ke,((d={})[he]=be?de+"px":"",d[ce]=we?ae+"px":"",d.transform="",d))}var gE={left:"right",right:"left",bottom:"top",top:"bottom"};function Nt(U){return U.replace(/left|right|bottom|top/g,function(d){return gE[d]})}var En={start:"end",end:"start"};function dT(U){return U.replace(/start|end/g,function(d){return En[d]})}function pT(U,d){var g=d.getRootNode&&d.getRootNode();if(U.contains(d))return!0;if(g&&_(g)){var w=d;do{if(w&&U.isSameNode(w))return!0;w=w.parentNode||w.host}while(w)}return!1}function Bt(U){return Object.assign({},U,{left:U.x,top:U.y,right:U.x+U.width,bottom:U.y+U.height})}function MT(U,d){return d===dE?Bt(function(g){var w=A(g),Z=u(g),q=w.visualViewport,Q=Z.clientWidth,ee=Z.clientHeight,Ee=0,ne=0;return q&&(Q=q.width,ee=q.height,/^((?!chrome|android).)*safari/i.test(navigator.userAgent)||(Ee=q.offsetLeft,ne=q.offsetTop)),{width:Q,height:ee,x:Ee+H(g),y:ne}}(U)):o(d)?function(g){var w=C(g);return w.top=w.top+g.clientTop,w.left=w.left+g.clientLeft,w.bottom=w.top+g.clientHeight,w.right=w.left+g.clientWidth,w.width=g.clientWidth,w.height=g.clientHeight,w.x=w.left,w.y=w.top,w}(d):Bt(function(g){var w,Z=u(g),q=L(g),Q=(w=g.ownerDocument)==null?void 0:w.body,ee=c(Z.scrollWidth,Z.clientWidth,Q?Q.scrollWidth:0,Q?Q.clientWidth:0),Ee=c(Z.scrollHeight,Z.clientHeight,Q?Q.scrollHeight:0,Q?Q.clientHeight:0),ne=-q.scrollLeft+H(g),Ae=-q.scrollTop;return b(Q||Z).direction==="rtl"&&(ne+=c(Z.clientWidth,Q?Q.clientWidth:0)-ee),{width:ee,height:Ee,x:ne,y:Ae}}(u(U)))}function UT(U){return Object.assign({},{top:0,right:0,bottom:0,left:0},U)}function mT(U,d){return d.reduce(function(g,w){return g[w]=U,g},{})}function Rt(U,d){d===void 0&&(d={});var g=d,w=g.placement,Z=w===void 0?U.placement:w,q=g.boundary,Q=q===void 0?"clippingParents":q,ee=g.rootBoundary,Ee=ee===void 0?dE:ee,ne=g.elementContext,Ae=ne===void 0?me:ne,re=g.altBoundary,We=re!==void 0&&re,ae=g.padding,Me=ae===void 0?0:ae,de=UT(typeof Me!="number"?Me:mT(Me,Ue)),ie=Ae===me?"reference":me,we=U.rects.popper,be=U.elements[We?ie:Ae],ce=function(Ke,sE,lE){var PE=sE==="clippingParents"?function(Be){var UE=h(D(Be)),oE=["absolute","fixed"].indexOf(b(Be).position)>=0&&i(Be)?W(Be):Be;return o(oE)?UE.filter(function(RE){return o(RE)&&pT(RE,oE)&&I(RE)!=="body"}):[]}(Ke):[].concat(sE),SE=[].concat(PE,[lE]),tE=SE[0],xe=SE.reduce(function(Be,UE){var oE=MT(Ke,UE);return Be.top=c(oE.top,Be.top),Be.right=P(oE.right,Be.right),Be.bottom=P(oE.bottom,Be.bottom),Be.left=c(oE.left,Be.left),Be},MT(Ke,tE));return xe.width=xe.right-xe.left,xe.height=xe.bottom-xe.top,xe.x=xe.left,xe.y=xe.top,xe}(o(be)?be:be.contextElement||u(U.elements.popper),Q,Ee),he=C(U.elements.reference),Ie=Te({reference:he,element:we,strategy:"absolute",placement:Z}),le=Bt(Object.assign({},we,Ie)),fe=Ae===me?le:he,pe={top:ce.top-fe.top+de.top,bottom:fe.bottom-ce.bottom+de.bottom,left:ce.left-fe.left+de.left,right:fe.right-ce.right+de.right},He=U.modifiersData.offset;if(Ae===me&&He){var ke=He[Z];Object.keys(pe).forEach(function(Ke){var sE=[oe,J].indexOf(Ke)>=0?1:-1,lE=[x,J].indexOf(Ke)>=0?"y":"x";pe[Ke]+=ke[lE]*sE})}return pe}function nt(U,d,g){return c(U,P(d,g))}function hT(U,d,g){return g===void 0&&(g={x:0,y:0}),{top:U.top-d.height-g.y,right:U.right-d.width+g.x,bottom:U.bottom-d.height+g.y,left:U.left-d.width-g.x}}function GT(U){return[x,oe,J,z].some(function(d){return U[d]>=0})}var tn=ME({defaultModifiers:[{name:"eventListeners",enabled:!0,phase:"write",fn:function(){},effect:function(U){var d=U.state,g=U.instance,w=U.options,Z=w.scroll,q=Z===void 0||Z,Q=w.resize,ee=Q===void 0||Q,Ee=A(d.elements.popper),ne=[].concat(d.scrollParents.reference,d.scrollParents.popper);return q&&ne.forEach(function(Ae){Ae.addEventListener("scroll",g.update,$e)}),ee&&Ee.addEventListener("resize",g.update,$e),function(){q&&ne.forEach(function(Ae){Ae.removeEventListener("scroll",g.update,$e)}),ee&&Ee.removeEventListener("resize",g.update,$e)}},data:{}},{name:"popperOffsets",enabled:!0,phase:"read",fn:function(U){var d=U.state,g=U.name;d.modifiersData[g]=Te({reference:d.rects.reference,element:d.rects.popper,strategy:"absolute",placement:d.placement})},data:{}},{name:"computeStyles",enabled:!0,phase:"beforeWrite",fn:function(U){var d=U.state,g=U.options,w=g.gpuAcceleration,Z=w===void 0||w,q=g.adaptive,Q=q===void 0||q,ee=g.roundOffsets,Ee=ee===void 0||ee,ne={placement:ze(d.placement),variation:ve(d.placement),popper:d.elements.popper,popperRect:d.rects.popper,gpuAcceleration:Z,isFixed:d.options.strategy==="fixed"};d.modifiersData.popperOffsets!=null&&(d.styles.popper=Object.assign({},d.styles.popper,se(Object.assign({},ne,{offsets:d.modifiersData.popperOffsets,position:d.options.strategy,adaptive:Q,roundOffsets:Ee})))),d.modifiersData.arrow!=null&&(d.styles.arrow=Object.assign({},d.styles.arrow,se(Object.assign({},ne,{offsets:d.modifiersData.arrow,position:"absolute",adaptive:!1,roundOffsets:Ee})))),d.attributes.popper=Object.assign({},d.attributes.popper,{"data-popper-placement":d.placement})},data:{}},{name:"applyStyles",enabled:!0,phase:"write",fn:function(U){var d=U.state;Object.keys(d.elements).forEach(function(g){var w=d.styles[g]||{},Z=d.attributes[g]||{},q=d.elements[g];i(q)&&I(q)&&(Object.assign(q.style,w),Object.keys(Z).forEach(function(Q){var ee=Z[Q];ee===!1?q.removeAttribute(Q):q.setAttribute(Q,ee===!0?"":ee)}))})},effect:function(U){var d=U.state,g={popper:{position:d.options.strategy,left:"0",top:"0",margin:"0"},arrow:{position:"absolute"},reference:{}};return Object.assign(d.elements.popper.style,g.popper),d.styles=g,d.elements.arrow&&Object.assign(d.elements.arrow.style,g.arrow),function(){Object.keys(d.elements).forEach(function(w){var Z=d.elements[w],q=d.attributes[w]||{},Q=Object.keys(d.styles.hasOwnProperty(w)?d.styles[w]:g[w]).reduce(function(ee,Ee){return ee[Ee]="",ee},{});i(Z)&&I(Z)&&(Object.assign(Z.style,Q),Object.keys(q).forEach(function(ee){Z.removeAttribute(ee)}))})}},requires:["computeStyles"]},{name:"offset",enabled:!0,phase:"main",requires:["popperOffsets"],fn:function(U){var d=U.state,g=U.options,w=U.name,Z=g.offset,q=Z===void 0?[0,0]:Z,Q=rE.reduce(function(Ae,re){return Ae[re]=function(We,ae,Me){var de=ze(We),ie=[z,x].indexOf(de)>=0?-1:1,we=typeof Me=="function"?Me(Object.assign({},ae,{placement:We})):Me,be=we[0],ce=we[1];return be=be||0,ce=(ce||0)*ie,[z,oe].indexOf(de)>=0?{x:ce,y:be}:{x:be,y:ce}}(re,d.rects,q),Ae},{}),ee=Q[d.placement],Ee=ee.x,ne=ee.y;d.modifiersData.popperOffsets!=null&&(d.modifiersData.popperOffsets.x+=Ee,d.modifiersData.popperOffsets.y+=ne),d.modifiersData[w]=Q}},{name:"flip",enabled:!0,phase:"main",fn:function(U){var d=U.state,g=U.options,w=U.name;if(!d.modifiersData[w]._skip){for(var Z=g.mainAxis,q=Z===void 0||Z,Q=g.altAxis,ee=Q===void 0||Q,Ee=g.fallbackPlacements,ne=g.padding,Ae=g.boundary,re=g.rootBoundary,We=g.altBoundary,ae=g.flipVariations,Me=ae===void 0||ae,de=g.allowedAutoPlacements,ie=d.options.placement,we=ze(ie),be=Ee||(we!==ie&&Me?function(RE){if(ze(RE)===Oe)return[];var _E=Nt(RE);return[dT(RE),_E,dT(_E)]}(ie):[Nt(ie)]),ce=[ie].concat(be).reduce(function(RE,_E){return RE.concat(ze(_E)===Oe?function($E,mE){mE===void 0&&(mE={});var LE=mE,lt=LE.placement,_t=LE.boundary,xE=LE.rootBoundary,vt=LE.padding,Ft=LE.flipVariations,XE=LE.allowedAutoPlacements,Yt=XE===void 0?rE:XE,At=ve(lt),Lt=At?Ft?fE:fE.filter(function(cE){return ve(cE)===At}):Ue,kE=Lt.filter(function(cE){return Yt.indexOf(cE)>=0});kE.length===0&&(kE=Lt);var KE=kE.reduce(function(cE,HE){return cE[HE]=Rt($E,{placement:HE,boundary:_t,rootBoundary:xE,padding:vt})[ze(HE)],cE},{});return Object.keys(KE).sort(function(cE,HE){return KE[cE]-KE[HE]})}(d,{placement:_E,boundary:Ae,rootBoundary:re,padding:ne,flipVariations:Me,allowedAutoPlacements:de}):_E)},[]),he=d.rects.reference,Ie=d.rects.popper,le=new Map,fe=!0,pe=ce[0],He=0;He<ce.length;He++){var ke=ce[He],Ke=ze(ke),sE=ve(ke)===ye,lE=[x,J].indexOf(Ke)>=0,PE=lE?"width":"height",SE=Rt(d,{placement:ke,boundary:Ae,rootBoundary:re,altBoundary:We,padding:ne}),tE=lE?sE?oe:z:sE?J:x;he[PE]>Ie[PE]&&(tE=Nt(tE));var xe=Nt(tE),Be=[];if(q&&Be.push(SE[Ke]<=0),ee&&Be.push(SE[tE]<=0,SE[xe]<=0),Be.every(function(RE){return RE})){pe=ke,fe=!1;break}le.set(ke,Be)}if(fe)for(var UE=function(RE){var _E=ce.find(function($E){var mE=le.get($E);if(mE)return mE.slice(0,RE).every(function(LE){return LE})});if(_E)return pe=_E,"break"},oE=Me?3:1;oE>0&&UE(oE)!=="break";oE--);d.placement!==pe&&(d.modifiersData[w]._skip=!0,d.placement=pe,d.reset=!0)}},requiresIfExists:["offset"],data:{_skip:!1}},{name:"preventOverflow",enabled:!0,phase:"main",fn:function(U){var d=U.state,g=U.options,w=U.name,Z=g.mainAxis,q=Z===void 0||Z,Q=g.altAxis,ee=Q!==void 0&&Q,Ee=g.boundary,ne=g.rootBoundary,Ae=g.altBoundary,re=g.padding,We=g.tether,ae=We===void 0||We,Me=g.tetherOffset,de=Me===void 0?0:Me,ie=Rt(d,{boundary:Ee,rootBoundary:ne,padding:re,altBoundary:Ae}),we=ze(d.placement),be=ve(d.placement),ce=!be,he=Fe(we),Ie=he==="x"?"y":"x",le=d.modifiersData.popperOffsets,fe=d.rects.reference,pe=d.rects.popper,He=typeof de=="function"?de(Object.assign({},d.rects,{placement:d.placement})):de,ke=typeof He=="number"?{mainAxis:He,altAxis:He}:Object.assign({mainAxis:0,altAxis:0},He),Ke=d.modifiersData.offset?d.modifiersData.offset[d.placement]:null,sE={x:0,y:0};if(le){if(q){var lE,PE=he==="y"?x:z,SE=he==="y"?J:oe,tE=he==="y"?"height":"width",xe=le[he],Be=xe+ie[PE],UE=xe-ie[SE],oE=ae?-pe[tE]/2:0,RE=be===ye?fe[tE]:pe[tE],_E=be===ye?-pe[tE]:-fe[tE],$E=d.elements.arrow,mE=ae&&$E?N($E):{width:0,height:0},LE=d.modifiersData["arrow#persistent"]?d.modifiersData["arrow#persistent"].padding:{top:0,right:0,bottom:0,left:0},lt=LE[PE],_t=LE[SE],xE=nt(0,fe[tE],mE[tE]),vt=ce?fe[tE]/2-oE-xE-lt-ke.mainAxis:RE-xE-lt-ke.mainAxis,Ft=ce?-fe[tE]/2+oE+xE+_t+ke.mainAxis:_E+xE+_t+ke.mainAxis,XE=d.elements.arrow&&W(d.elements.arrow),Yt=XE?he==="y"?XE.clientTop||0:XE.clientLeft||0:0,At=(lE=Ke==null?void 0:Ke[he])!=null?lE:0,Lt=xe+Ft-At,kE=nt(ae?P(Be,xe+vt-At-Yt):Be,xe,ae?c(UE,Lt):UE);le[he]=kE,sE[he]=kE-xe}if(ee){var KE,cE=he==="x"?x:z,HE=he==="x"?J:oe,bE=le[Ie],Ct=Ie==="y"?"height":"width",gT=bE+ie[cE],HT=bE-ie[HE],Vt=[x,z].indexOf(we)!==-1,bT=(KE=Ke==null?void 0:Ke[Ie])!=null?KE:0,yT=Vt?gT:bE-fe[Ct]-pe[Ct]-bT+ke.altAxis,BT=Vt?bE+fe[Ct]+pe[Ct]-bT-ke.altAxis:HT,vT=ae&&Vt?function(Tn,rn,Wt){var FT=nt(Tn,rn,Wt);return FT>Wt?Wt:FT}(yT,bE,BT):nt(ae?yT:gT,bE,ae?BT:HT);le[Ie]=vT,sE[Ie]=vT-bE}d.modifiersData[w]=sE}},requiresIfExists:["offset"]},{name:"arrow",enabled:!0,phase:"main",fn:function(U){var d,g=U.state,w=U.name,Z=U.options,q=g.elements.arrow,Q=g.modifiersData.popperOffsets,ee=ze(g.placement),Ee=Fe(ee),ne=[z,oe].indexOf(ee)>=0?"height":"width";if(q&&Q){var Ae=function(pe,He){return UT(typeof(pe=typeof pe=="function"?pe(Object.assign({},He.rects,{placement:He.placement})):pe)!="number"?pe:mT(pe,Ue))}(Z.padding,g),re=N(q),We=Ee==="y"?x:z,ae=Ee==="y"?J:oe,Me=g.rects.reference[ne]+g.rects.reference[Ee]-Q[Ee]-g.rects.popper[ne],de=Q[Ee]-g.rects.reference[Ee],ie=W(q),we=ie?Ee==="y"?ie.clientHeight||0:ie.clientWidth||0:0,be=Me/2-de/2,ce=Ae[We],he=we-re[ne]-Ae[ae],Ie=we/2-re[ne]/2+be,le=nt(ce,Ie,he),fe=Ee;g.modifiersData[w]=((d={})[fe]=le,d.centerOffset=le-Ie,d)}},effect:function(U){var d=U.state,g=U.options.element,w=g===void 0?"[data-popper-arrow]":g;w!=null&&(typeof w!="string"||(w=d.elements.popper.querySelector(w)))&&pT(d.elements.popper,w)&&(d.elements.arrow=w)},requires:["popperOffsets"],requiresIfExists:["preventOverflow"]},{name:"hide",enabled:!0,phase:"main",requiresIfExists:["preventOverflow"],fn:function(U){var d=U.state,g=U.name,w=d.rects.reference,Z=d.rects.popper,q=d.modifiersData.preventOverflow,Q=Rt(d,{elementContext:"reference"}),ee=Rt(d,{altBoundary:!0}),Ee=hT(Q,w),ne=hT(ee,Z,q),Ae=GT(Ee),re=GT(ne);d.modifiersData[g]={referenceClippingOffsets:Ee,popperEscapeOffsets:ne,isReferenceHidden:Ae,hasPopperEscaped:re},d.attributes.popper=Object.assign({},d.attributes.popper,{"data-popper-reference-hidden":Ae,"data-popper-escaped":re})}}]})}},t={};function r(n){var s=t[n];if(s!==void 0)return s.exports;var S=t[n]={exports:{}};return T[n](S,S.exports,r),S.exports}r.d=(n,s)=>{for(var S in s)r.o(s,S)&&!r.o(n,S)&&Object.defineProperty(n,S,{enumerable:!0,get:s[S]})},r.o=(n,s)=>Object.prototype.hasOwnProperty.call(n,s),r.r=n=>{typeof Symbol<"u"&&Symbol.toStringTag&&Object.defineProperty(n,Symbol.toStringTag,{value:"Module"}),Object.defineProperty(n,"__esModule",{value:!0})};var R={};return r.r(R),r(661),r(795),r(682),r(284),r(181),r(778),r(51),r(185),R})()})})(Qn);function Zn(E){let e=E[0].title+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p(t,r){r&1&&e!==(e=t[0].title+"")&&Le(T,e)},d(t){t&&Y(T)}}}function jn(E){let e,T;return{c(){e=te("Welcome to "),T=f("span"),T.textContent="Vanna.AI",a(T,"class","nav-title")},m(t,r){V(t,e,r),V(t,T,r)},p:j,d(t){t&&(Y(e),Y(T))}}}function zn(E){let e,T,t,r,R=E[0].subtitle+"",n;function s(o,i){return o[0].title=="Welcome to Vanna.AI"?jn:Zn}let S=s(E),A=S(E);return{c(){e=f("div"),T=f("h1"),A.c(),t=$(),r=f("p"),n=te(R),a(T,"class","text-3xl font-bold text-gray-800 sm:text-4xl dark:text-white"),a(r,"class","mt-3 text-gray-600 dark:text-gray-400"),a(e,"class","max-w-4xl px-4 sm:px-6 lg:px-8 mx-auto text-center")},m(o,i){V(o,e,i),l(e,T),A.m(T,null),l(e,t),l(e,r),l(r,n)},p(o,[i]){S===(S=s(o))&&A?A.p(o,i):(A.d(1),A=S(o),A&&(A.c(),A.m(T,null))),i&1&&R!==(R=o[0].subtitle+"")&&Le(n,R)},i:j,o:j,d(o){o&&Y(e),A.d()}}}function eA(E,e,T){let t;return eE(E,VE,r=>T(0,t=r)),[t]}class EA extends ue{constructor(e){super(),Ce(this,e,eA,zn,_e,{})}}function tA(E){let e,T;const t=E[1].default,r=Ut(t,E,E[0],null);return{c(){e=f("p"),r&&r.c(),a(e,"class","text-gray-800 dark:text-gray-200")},m(R,n){V(R,e,n),r&&r.m(e,null),T=!0},p(R,[n]){r&&r.p&&(!T||n&1)&&ht(r,t,R,R[0],T?mt(t,R[0],n,null):Gt(R[0]),null)},i(R){T||(m(r,R),T=!0)},o(R){y(r,R),T=!1},d(R){R&&Y(e),r&&r.d(R)}}}function TA(E,e,T){let{$$slots:t={},$$scope:r}=e;return E.$$set=R=>{"$$scope"in R&&T(0,r=R.$$scope)},[r,t]}class aE extends ue{constructor(e){super(),Ce(this,e,TA,tA,_e,{})}}function rA(E){let e;return{c(){e=te(E[0])},m(T,t){V(T,e,t)},p(T,t){t&1&&Le(e,T[0])},d(T){T&&Y(e)}}}function RA(E){let e,T,t,r,R,n,s,S,A;s=new aE({props:{$$slots:{default:[rA]},$$scope:{ctx:E}}});const o=E[1].default,i=Ut(o,E,E[2],null);return{c(){e=f("li"),T=f("div"),t=f("div"),r=f("span"),r.innerHTML='<span class="text-sm font-medium text-white leading-none">You</span>',R=$(),n=f("div"),K(s.$$.fragment),S=$(),i&&i.c(),a(r,"class","flex-shrink-0 inline-flex items-center justify-center h-[2.375rem] w-[2.375rem] rounded-full bg-gray-600"),a(n,"class","grow mt-2 space-y-3"),a(t,"class","max-w-2xl flex gap-x-2 sm:gap-x-4"),a(T,"class","max-w-4xl px-4 sm:px-6 lg:px-8 mx-auto"),a(e,"class","py-2 sm:py-4")},m(_,c){V(_,e,c),l(e,T),l(T,t),l(t,r),l(t,R),l(t,n),X(s,n,null),l(n,S),i&&i.m(n,null),A=!0},p(_,[c]){const P={};c&5&&(P.$$scope={dirty:c,ctx:_}),s.$set(P),i&&i.p&&(!A||c&4)&&ht(i,o,_,_[2],A?mt(o,_[2],c,null):Gt(_[2]),null)},i(_){A||(m(s.$$.fragment,_),m(i,_),A=!0)},o(_){y(s.$$.fragment,_),y(i,_),A=!1},d(_){_&&Y(e),k(s),i&&i.d(_)}}}function nA(E,e,T){let{$$slots:t={},$$scope:r}=e,{message:R}=e;return E.$$set=n=>{"message"in n&&T(0,R=n.message),"$$scope"in n&&T(2,r=n.$$scope)},[R,t,r]}class WE extends ue{constructor(e){super(),Ce(this,e,nA,RA,_e,{message:0})}}function AA(E){let e,T,t;return{c(){e=f("button"),e.innerHTML='<svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M12 2a3 3 0 0 0-3 3v7a3 3 0 0 0 6 0V5a3 3 0 0 0-3-3Z"></path><path d="M19 10v2a7 7 0 0 1-14 0v-2"></path><line x1="12" x2="12" y1="19" y2="22"></line></svg>',a(e,"type","button"),a(e,"class","inline-flex flex-shrink-0 justify-center items-center size-8 rounded-lg text-gray-500 hover:text-blue-600 focus:z-10 focus:outline-none focus:ring-2 focus:ring-blue-500 dark:hover:text-blue-500 dark:focus:outline-none dark:focus:ring-1 dark:focus:ring-gray-600")},m(r,R){V(r,e,R),T||(t=Ne(e,"click",E[1]),T=!0)},p:j,d(r){r&&Y(e),T=!1,t()}}}function sA(E){let e;return{c(){e=f("button"),e.innerHTML='<svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M12 2a3 3 0 0 0-3 3v7a3 3 0 0 0 6 0V5a3 3 0 0 0-3-3Z"></path><path d="M19 10v2a7 7 0 0 1-14 0v-2"></path><line x1="12" x2="12" y1="19" y2="22"></line></svg>',a(e,"type","button"),a(e,"class","animate-ping animate-pulse inline-flex flex-shrink-0 justify-center items-center size-8 rounded-lg text-red-500 hover:text-red-600 focus:z-10 focus:outline-none focus:ring-2 focus:ring-red-500 dark:hover:text-red-500 dark:focus:outline-none dark:focus:ring-1 dark:focus:ring-red-600")},m(T,t){V(T,e,t)},p:j,d(T){T&&Y(e)}}}function SA(E){let e;function T(R,n){return R[0]?sA:AA}let t=T(E),r=t(E);return{c(){r.c(),e=je()},m(R,n){r.m(R,n),V(R,e,n)},p(R,[n]){t===(t=T(R))&&r?r.p(R,n):(r.d(1),r=t(R),r&&(r.c(),r.m(e.parentNode,e)))},i:j,o:j,d(R){R&&Y(e),r.d(R)}}}function oA(E,e,T){let{newMessage:t}=e,r=!1;function R(){if(T(0,r=!0),pR.set(!0),"webkitSpeechRecognition"in window)var n=new window.webkitSpeechRecognition;else var n=new window.SpeechRecognition;n.lang="en-US",n.start(),n.onresult=s=>{const S=s.results[0][0].transcript;console.log(S),T(2,t=S),T(0,r=!1)},n.onend=()=>{T(0,r=!1)},n.onerror=()=>{T(0,r=!1)}}return E.$$set=n=>{"newMessage"in n&&T(2,t=n.newMessage)},[r,R,t]}class OA extends ue{constructor(e){super(),Ce(this,e,oA,SA,_e,{newMessage:2})}}function iA(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p;function C(I){E[5](I)}let L={};return E[0]!==void 0&&(L.newMessage=E[0]),A=new OA({props:L}),iT.push(()=>cn(A,"newMessage",C)),{c(){e=f("div"),T=f("input"),t=$(),r=f("div"),R=f("div"),n=f("div"),n.innerHTML="",s=$(),S=f("div"),K(A.$$.fragment),i=$(),_=f("button"),_.innerHTML='<svg class="h-3.5 w-3.5" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path d="M15.964.686a.5.5 0 0 0-.65-.65L.767 5.855H.766l-.452.18a.5.5 0 0 0-.082.887l.41.26.001.002 4.995 3.178 3.178 4.995.002.002.26.41a.5.5 0 0 0 .886-.083l6-15Zm-1.833 1.89L6.637 10.07l-.215-.338a.5.5 0 0 0-.154-.154l-.338-.215 7.494-7.494 1.178-.471-.47 1.178Z"></path></svg>',a(T,"type","text"),a(T,"class","p-4 pb-12 block w-full bg-gray-100 border-gray-200 rounded-md text-sm focus:border-blue-500 focus:ring-blue-500 dark:bg-slate-800 dark:border-gray-700 dark:text-gray-400"),a(T,"placeholder","Ask me a question about your data that I can turn into SQL."),a(n,"class","flex items-center"),a(_,"type","button"),a(_,"class","inline-flex flex-shrink-0 justify-center items-center h-8 w-8 rounded-md text-white bg-blue-600 hover:bg-blue-500 focus:z-10 focus:outline-none focus:ring-2 focus:ring-blue-500 transition-all"),a(S,"class","flex items-center gap-x-1"),a(R,"class","flex justify-between items-center"),a(r,"class","absolute bottom-px inset-x-px p-2 rounded-b-md bg-gray-100 dark:bg-slate-800"),a(e,"class","relative")},m(I,u){V(I,e,u),l(e,T),Ye(T,E[0]),l(e,t),l(e,r),l(r,R),l(R,n),l(R,s),l(R,S),X(A,S,null),l(S,i),l(S,_),c=!0,P||(p=[Ne(T,"input",E[4]),Ne(T,"keydown",E[1]),Ne(_,"click",E[2])],P=!0)},p(I,[u]){u&1&&T.value!==I[0]&&Ye(T,I[0]);const H={};!o&&u&1&&(o=!0,H.newMessage=I[0],ln(()=>o=!1)),A.$set(H)},i(I){c||(m(A.$$.fragment,I),c=!0)},o(I){y(A.$$.fragment,I),c=!1},d(I){I&&Y(e),k(A),P=!1,NE(p)}}}function aA(E,e,T){let t;eE(E,BE,A=>T(0,t=A));let{onSubmit:r}=e;function R(A){A.key==="Enter"&&(r(t),OT(BE,t="",t),A.preventDefault())}function n(){r(t),OT(BE,t="",t)}function s(){t=this.value,BE.set(t)}function S(A){t=A,BE.set(t)}return E.$$set=A=>{"onSubmit"in A&&T(3,r=A.onSubmit)},[t,R,n,r,s,S]}class IA extends ue{constructor(e){super(),Ce(this,e,aA,iA,_e,{onSubmit:3})}}function NA(E){let e;return{c(){e=f("div"),e.innerHTML='<button type="button" class="p-2 inline-flex justify-center items-center gap-1.5 rounded-md border font-medium bg-white text-gray-700 shadow-sm align-middle hover:bg-gray-50 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white focus:ring-blue-600 transition-all text-xs dark:bg-slate-900 dark:hover:bg-slate-800 dark:border-gray-700 dark:text-gray-400 dark:hover:text-white dark:focus:ring-offset-gray-800" data-hs-overlay="#application-sidebar" aria-controls="application-sidebar" aria-label="Toggle navigation"><svg class="w-3.5 h-3.5" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path fill-rule="evenodd" d="M2.5 12a.5.5 0 0 1 .5-.5h10a.5.5 0 0 1 0 1H3a.5.5 0 0 1-.5-.5zm0-4a.5.5 0 0 1 .5-.5h10a.5.5 0 0 1 0 1H3a.5.5 0 0 1-.5-.5zm0-4a.5.5 0 0 1 .5-.5h10a.5.5 0 0 1 0 1H3a.5.5 0 0 1-.5-.5z"></path></svg> <span>Sidebar</span></button>',a(e,"class","lg:hidden flex justify-end mb-2 sm:mb-3")},m(T,t){V(T,e,t)},p:j,i:j,o:j,d(T){T&&Y(e)}}}class lA extends ue{constructor(e){super(),Ce(this,e,null,NA,_e,{})}}function _A(E){let e,T,t,r;return{c(){e=f("button"),T=te(E[0]),a(e,"type","button"),a(e,"class","mb-2.5 mr-1.5 py-2 px-3 inline-flex justify-center items-center gap-x-2 rounded-md border border-blue-600 bg-white text-blue-600 align-middle hover:bg-blue-50 text-sm dark:bg-slate-900 dark:text-blue-500 dark:border-blue-500 dark:hover:text-blue-400 dark:hover:border-blue-400")},m(R,n){V(R,e,n),l(e,T),t||(r=Ne(e,"click",E[1]),t=!0)},p(R,[n]){n&1&&Le(T,R[0])},i:j,o:j,d(R){R&&Y(e),t=!1,r()}}}function LA(E,e,T){let{message:t}=e,{onSubmit:r}=e;function R(){r(t)}return E.$$set=n=>{"message"in n&&T(0,t=n.message),"onSubmit"in n&&T(2,r=n.onSubmit)},[t,R,r]}class IE extends ue{constructor(e){super(),Ce(this,e,LA,_A,_e,{message:0,onSubmit:2})}}function CA(E){let e,T,t,r,R,n,s,S,A,o,i;return{c(){e=f("span"),T=OE("svg"),t=OE("defs"),r=OE("linearGradient"),R=OE("stop"),n=OE("stop"),s=OE("g"),S=OE("g"),A=OE("path"),o=OE("path"),a(R,"offset","0"),a(R,"stop-color","#009efd"),a(n,"offset","1"),a(n,"stop-color","#2af598"),a(r,"gradientTransform","matrix(1.09331 0 0 1.09331 -47.1838 -88.8946)"),a(r,"gradientUnits","userSpaceOnUse"),a(r,"id","LinearGradient"),a(r,"x1","237.82"),a(r,"x2","785.097"),a(r,"y1","549.609"),a(r,"y2","549.609"),a(A,"d","M117.718 228.798C117.718 119.455 206.358 30.8151 315.701 30.8151L708.299 30.8151C817.642 30.8151 906.282 119.455 906.282 228.798L906.282 795.202C906.282 904.545 817.642 993.185 708.299 993.185L315.701 993.185C206.358 993.185 117.718 904.545 117.718 795.202L117.718 228.798Z"),a(A,"fill","#0f172a"),a(A,"fill-rule","nonzero"),a(A,"opacity","1"),a(A,"stroke","#374151"),a(A,"stroke-linecap","butt"),a(A,"stroke-linejoin","round"),a(A,"stroke-width","20"),a(o,"d","M212.828 215.239C213.095 281.169 213.629 413.028 213.629 413.028C213.629 413.028 511.51 808.257 513.993 809.681C612.915 677.809 810.759 414.065 810.759 414.065C810.759 414.065 811.034 280.901 811.172 214.319C662.105 362.973 662.105 362.973 513.038 511.627C362.933 363.433 362.933 363.433 212.828 215.239Z"),a(o,"fill","url(#LinearGradient)"),a(o,"fill-rule","nonzero"),a(o,"opacity","1"),a(o,"stroke","none"),a(S,"opacity","1"),a(s,"id","Layer-1"),a(T,"height","100%"),a(T,"stroke-miterlimit","10"),ct(T,"fill-rule","nonzero"),ct(T,"clip-rule","evenodd"),ct(T,"stroke-linecap","round"),ct(T,"stroke-linejoin","round"),a(T,"version","1.1"),a(T,"viewBox","0 0 1024 1024"),a(T,"width","100%"),a(T,"xml:space","preserve"),a(T,"xmlns","http://www.w3.org/2000/svg"),a(e,"class",i="flex-shrink-0 w-[2.375rem] h-[2.375rem] "+E[0])},m(_,c){V(_,e,c),l(e,T),l(T,t),l(t,r),l(r,R),l(r,n),l(T,s),l(s,S),l(S,A),l(S,o)},p(_,[c]){c&1&&i!==(i="flex-shrink-0 w-[2.375rem] h-[2.375rem] "+_[0])&&a(e,"class",i)},i:j,o:j,d(_){_&&Y(e)}}}function uA(E,e,T){let t,{animate:r=!1}=e;return E.$$set=R=>{"animate"in R&&T(1,r=R.animate)},E.$$.update=()=>{E.$$.dirty&2&&T(0,t=r?"animate-bounce":"")},[t,r]}class GR extends ue{constructor(e){super(),Ce(this,e,uA,CA,_e,{animate:1})}}function cA(E){let e,T,t,r,R;T=new GR({});const n=E[1].default,s=Ut(n,E,E[0],null);return{c(){e=f("li"),K(T.$$.fragment),t=$(),r=f("div"),s&&s.c(),a(r,"class","space-y-3 overflow-x-auto overflow-y-hidden whitespace-break-spaces w-full"),a(e,"class","max-w-4xl py-2 px-4 sm:px-6 lg:px-8 mx-auto flex gap-x-2 sm:gap-x-4")},m(S,A){V(S,e,A),X(T,e,null),l(e,t),l(e,r),s&&s.m(r,null),R=!0},p(S,[A]){s&&s.p&&(!R||A&1)&&ht(s,n,S,S[0],R?mt(n,S[0],A,null):Gt(S[0]),null)},i(S){R||(m(T.$$.fragment,S),m(s,S),R=!0)},o(S){y(T.$$.fragment,S),y(s,S),R=!1},d(S){S&&Y(e),k(T),s&&s.d(S)}}}function fA(E,e,T){let{$$slots:t={},$$scope:r}=e;return E.$$set=R=>{"$$scope"in R&&T(0,r=R.$$scope)},[r,t]}class Ze extends ue{constructor(e){super(),Ce(this,e,fA,cA,_e,{})}}function PA(E){let e;return{c(){e=te("Thinking...")},m(T,t){V(T,e,t)},d(T){T&&Y(e)}}}function DA(E){let e,T,t,r,R,n;return T=new GR({props:{animate:!0}}),R=new aE({props:{$$slots:{default:[PA]},$$scope:{ctx:E}}}),{c(){e=f("li"),K(T.$$.fragment),t=$(),r=f("div"),K(R.$$.fragment),a(r,"class","space-y-3"),a(e,"class","max-w-4xl py-2 px-4 sm:px-6 lg:px-8 mx-auto flex gap-x-2 sm:gap-x-4")},m(s,S){V(s,e,S),X(T,e,null),l(e,t),l(e,r),X(R,r,null),n=!0},p(s,[S]){const A={};S&1&&(A.$$scope={dirty:S,ctx:s}),R.$set(A)},i(s){n||(m(T.$$.fragment,s),m(R.$$.fragment,s),n=!0)},o(s){y(T.$$.fragment,s),y(R.$$.fragment,s),n=!1},d(s){s&&Y(e),k(T),k(R)}}}class dA extends ue{constructor(e){super(),Ce(this,e,null,DA,_e,{})}}function pA(E){let e,T,t,r,R,n,s,S,A,o,i;return{c(){e=f("ul"),T=f("li"),t=f("div"),r=f("span"),r.textContent="CSV",R=$(),n=f("a"),s=OE("svg"),S=OE("path"),A=OE("path"),o=te(`
          Download`),a(r,"class","mr-3 flex-1 w-0 truncate"),a(S,"d","M.5 9.9a.5.5 0 0 1 .5.5v2.5a1 1 0 0 0 1 1h12a1 1 0 0 0 1-1v-2.5a.5.5 0 0 1 1 0v2.5a2 2 0 0 1-2 2H2a2 2 0 0 1-2-2v-2.5a.5.5 0 0 1 .5-.5z"),a(A,"d","M7.646 11.854a.5.5 0 0 0 .708 0l3-3a.5.5 0 0 0-.708-.708L8.5 10.293V1.5a.5.5 0 0 0-1 0v8.793L5.354 8.146a.5.5 0 1 0-.708.708l3 3z"),a(s,"class","flex-shrink-0 w-3 h-3"),a(s,"width","16"),a(s,"height","16"),a(s,"viewBox","0 0 16 16"),a(s,"fill","currentColor"),a(n,"class","flex items-center gap-x-2 text-gray-500 hover:text-blue-500 whitespace-nowrap"),a(n,"href",i="/api/v0/download_csv?id="+E[0]),a(t,"class","w-full flex justify-between truncate"),a(T,"class","flex items-center gap-x-2 p-3 text-sm bg-white border text-gray-800 first:rounded-t-lg first:mt-0 last:rounded-b-lg dark:bg-slate-900 dark:border-gray-700 dark:text-gray-200"),a(e,"class","flex flex-col justify-end text-start -space-y-px")},m(_,c){V(_,e,c),l(e,T),l(T,t),l(t,r),l(t,R),l(t,n),l(n,s),l(s,S),l(s,A),l(n,o)},p(_,[c]){c&1&&i!==(i="/api/v0/download_csv?id="+_[0])&&a(n,"href",i)},i:j,o:j,d(_){_&&Y(e)}}}function MA(E,e,T){let{id:t}=e;return E.$$set=r=>{"id"in r&&T(0,t=r.id)},[t]}class UA extends ue{constructor(e){super(),Ce(this,e,MA,pA,_e,{id:0})}}function KT(E,e,T){const t=E.slice();return t[5]=e[T],t}function JT(E,e,T){const t=E.slice();return t[8]=e[T],t}function qT(E,e,T){const t=E.slice();return t[8]=e[T],t}function QT(E){let e,T,t,r;return{c(){e=f("th"),T=f("div"),t=f("span"),t.textContent=`${E[8]}`,r=$(),a(t,"class","text-xs font-semibold uppercase tracking-wide text-gray-800 dark:text-gray-200"),a(T,"class","flex items-center gap-x-2"),a(e,"scope","col"),a(e,"class","px-6 py-3 text-left")},m(R,n){V(R,e,n),l(e,T),l(T,t),l(e,r)},p:j,d(R){R&&Y(e)}}}function ZT(E){let e,T,t;return{c(){e=f("td"),T=f("div"),t=f("span"),t.textContent=`${E[5][E[8]]}`,a(t,"class","text-gray-800 dark:text-gray-200"),a(T,"class","px-6 py-3"),a(e,"class","h-px w-px whitespace-nowrap")},m(r,R){V(r,e,R),l(e,T),l(T,t)},p:j,d(r){r&&Y(e)}}}function jT(E){let e,T,t=De(E[3]),r=[];for(let R=0;R<t.length;R+=1)r[R]=ZT(JT(E,t,R));return{c(){e=f("tr");for(let R=0;R<r.length;R+=1)r[R].c();T=$()},m(R,n){V(R,e,n);for(let s=0;s<r.length;s+=1)r[s]&&r[s].m(e,null);l(e,T)},p(R,n){if(n&12){t=De(R[3]);let s;for(s=0;s<t.length;s+=1){const S=JT(R,t,s);r[s]?r[s].p(S,n):(r[s]=ZT(S),r[s].c(),r[s].m(e,T))}for(;s<r.length;s+=1)r[s].d(1);r.length=t.length}},d(R){R&&Y(e),nE(r,R)}}}function zT(E){let e,T;return e=new UA({props:{id:E[0]}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&1&&(R.id=t[0]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function mA(E){let e,T,t,r,R,n,s,S,A,o,i,_=De(E[3]),c=[];for(let L=0;L<_.length;L+=1)c[L]=QT(qT(E,_,L));let P=De(E[2]),p=[];for(let L=0;L<P.length;L+=1)p[L]=jT(KT(E,P,L));let C=E[1].csv_download&&zT(E);return{c(){e=f("div"),T=f("div"),t=f("div"),r=f("table"),R=f("thead"),n=f("tr");for(let L=0;L<c.length;L+=1)c[L].c();s=$(),S=f("tbody");for(let L=0;L<p.length;L+=1)p[L].c();A=$(),C&&C.c(),o=je(),a(R,"class","bg-gray-50 dark:bg-slate-800"),a(S,"class","divide-y divide-gray-200 dark:divide-gray-700"),a(r,"class","min-w-full divide-y divide-gray-200 dark:divide-gray-700"),a(t,"class","p-1.5 min-w-full inline-block align-middle"),a(T,"class","-m-1.5 overflow-x-auto"),a(e,"class","bg-white border border-gray-200 rounded-xl shadow-sm overflow-hidden dark:bg-slate-900 dark:border-gray-700")},m(L,I){V(L,e,I),l(e,T),l(T,t),l(t,r),l(r,R),l(R,n);for(let u=0;u<c.length;u+=1)c[u]&&c[u].m(n,null);l(r,s),l(r,S);for(let u=0;u<p.length;u+=1)p[u]&&p[u].m(S,null);V(L,A,I),C&&C.m(L,I),V(L,o,I),i=!0},p(L,[I]){if(I&8){_=De(L[3]);let u;for(u=0;u<_.length;u+=1){const H=qT(L,_,u);c[u]?c[u].p(H,I):(c[u]=QT(H),c[u].c(),c[u].m(n,null))}for(;u<c.length;u+=1)c[u].d(1);c.length=_.length}if(I&12){P=De(L[2]);let u;for(u=0;u<P.length;u+=1){const H=KT(L,P,u);p[u]?p[u].p(H,I):(p[u]=jT(H),p[u].c(),p[u].m(S,null))}for(;u<p.length;u+=1)p[u].d(1);p.length=P.length}L[1].csv_download?C?(C.p(L,I),I&2&&m(C,1)):(C=zT(L),C.c(),m(C,1),C.m(o.parentNode,o)):C&&(Ge(),y(C,1,1,()=>{C=null}),ge())},i(L){i||(m(C),i=!0)},o(L){y(C),i=!1},d(L){L&&(Y(e),Y(A),Y(o)),nE(c,L),nE(p,L),C&&C.d(L)}}}function hA(E,e,T){let t;eE(E,VE,S=>T(1,t=S));let{id:r}=e,{df:R}=e,n=JSON.parse(R),s=n.length>0?Object.keys(n[0]):[];return E.$$set=S=>{"id"in S&&T(0,r=S.id),"df"in S&&T(4,R=S.df)},[r,t,n,s,R]}class gR extends ue{constructor(e){super(),Ce(this,e,hA,mA,_e,{id:0,df:4})}}function GA(E){let e;return{c(){e=f("div"),a(e,"id",E[0])},m(T,t){V(T,e,t)},p:j,i:j,o:j,d(T){T&&Y(e)}}}function gA(E,e,T){let{fig:t}=e,r=JSON.parse(t),R=Math.random().toString(36).substring(2,15)+Math.random().toString(36).substring(2,15);return DR(()=>{Plotly.newPlot(document.getElementById(R),r,{responsive:!0})}),E.$$set=n=>{"fig"in n&&T(1,t=n.fig)},[R,t]}class HR extends ue{constructor(e){super(),Ce(this,e,gA,GA,_e,{fig:1})}}function HA(E){let e,T,t,r;return{c(){e=f("button"),T=te(E[0]),a(e,"type","button"),a(e,"class","mb-2.5 mr-1.5 py-3 px-4 inline-flex justify-center items-center gap-2 rounded-md border-2 border-green-200 font-semibold text-green-500 hover:text-white hover:bg-green-500 hover:border-green-500 focus:outline-none focus:ring-2 focus:ring-green-200 focus:ring-offset-2 transition-all text-sm dark:focus:ring-offset-gray-800")},m(R,n){V(R,e,n),l(e,T),t||(r=Ne(e,"click",E[1]),t=!0)},p(R,[n]){n&1&&Le(T,R[0])},i:j,o:j,d(R){R&&Y(e),t=!1,r()}}}function bA(E,e,T){let{message:t}=e,{onSubmit:r}=e;function R(){r(t)}return E.$$set=n=>{"message"in n&&T(0,t=n.message),"onSubmit"in n&&T(2,r=n.onSubmit)},[t,R,r]}class yA extends ue{constructor(e){super(),Ce(this,e,bA,HA,_e,{message:0,onSubmit:2})}}function BA(E){let e,T,t,r,R,n,s,S,A;return{c(){e=f("div"),T=f("div"),t=f("div"),t.innerHTML='<svg class="h-4 w-4 text-yellow-400 mt-0.5" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path d="M8.982 1.566a1.13 1.13 0 0 0-1.96 0L.165 13.233c-.457.778.091 1.767.98 1.767h13.713c.889 0 1.438-.99.98-1.767L8.982 1.566zM8 5c.535 0 .954.462.9.995l-.35 3.507a.552.552 0 0 1-1.1 0L7.1 5.995A.905.905 0 0 1 8 5zm.002 6a1 1 0 1 1 0 2 1 1 0 0 1 0-2z"></path></svg>',r=$(),R=f("div"),n=f("h3"),n.textContent="Error",s=$(),S=f("div"),A=te(E[0]),a(t,"class","flex-shrink-0"),a(n,"class","text-sm text-yellow-800 font-semibold"),a(S,"class","mt-1 text-sm text-yellow-700"),a(R,"class","ml-4"),a(T,"class","flex"),a(e,"class","bg-yellow-50 border border-yellow-200 rounded-md p-4"),a(e,"role","alert")},m(o,i){V(o,e,i),l(e,T),l(T,t),l(T,r),l(T,R),l(R,n),l(R,s),l(R,S),l(S,A)},p(o,[i]){i&1&&Le(A,o[0])},i:j,o:j,d(o){o&&Y(e)}}}function vA(E,e,T){let{message:t}=e;return E.$$set=r=>{"message"in r&&T(0,t=r.message)},[t]}let PT=class extends ue{constructor(e){super(),Ce(this,e,vA,BA,_e,{message:0})}};function FA(E){let e,T;const t=E[1].default,r=Ut(t,E,E[0],null);return{c(){e=f("div"),r&&r.c(),a(e,"class","font-mono whitespace-pre-wrap")},m(R,n){V(R,e,n),r&&r.m(e,null),T=!0},p(R,[n]){r&&r.p&&(!T||n&1)&&ht(r,t,R,R[0],T?mt(t,R[0],n,null):Gt(R[0]),null)},i(R){T||(m(r,R),T=!0)},o(R){y(r,R),T=!1},d(R){R&&Y(e),r&&r.d(R)}}}function YA(E,e,T){let{$$slots:t={},$$scope:r}=e;return E.$$set=R=>{"$$scope"in R&&T(0,r=R.$$scope)},[r,t]}class bR extends ue{constructor(e){super(),Ce(this,e,YA,FA,_e,{})}}function VA(E){let e;return{c(){e=te(E[1])},m(T,t){V(T,e,t)},p(T,t){t&2&&Le(e,T[1])},d(T){T&&Y(e)}}}function WA(E){let e,T,t,r,R,n,s,S;return t=new IE({props:{message:"Run SQL",onSubmit:E[3]}}),R=new aE({props:{$$slots:{default:[VA]},$$scope:{ctx:E}}}),{c(){e=f("textarea"),T=$(),K(t.$$.fragment),r=$(),K(R.$$.fragment),a(e,"rows","6"),a(e,"class","block p-2.5 w-full text-blue-600 hover:text-blue-500 dark:text-blue-500 dark:hover:text-blue-400 bg-gray-50 rounded-lg border border-gray-300 focus:ring-blue-500 focus:border-blue-500 dark:bg-gray-700 dark:border-gray-600 dark:placeholder-gray-400 dark:focus:ring-blue-500 dark:focus:border-blue-500 font-mono"),a(e,"placeholder","SELECT col1, col2, col3 FROM ...")},m(A,o){V(A,e,o),Ye(e,E[1]),V(A,T,o),X(t,A,o),V(A,r,o),X(R,A,o),n=!0,s||(S=Ne(e,"input",E[2]),s=!0)},p(A,[o]){o&2&&Ye(e,A[1]);const i={};o&3&&(i.onSubmit=A[3]),t.$set(i);const _={};o&18&&(_.$$scope={dirty:o,ctx:A}),R.$set(_)},i(A){n||(m(t.$$.fragment,A),m(R.$$.fragment,A),n=!0)},o(A){y(t.$$.fragment,A),y(R.$$.fragment,A),n=!1},d(A){A&&(Y(e),Y(T),Y(r)),k(t,A),k(R,A),s=!1,S()}}}function wA(E,e,T){let t;eE(E,Et,s=>T(1,t=s));let{onSubmit:r}=e;function R(){t=this.value,Et.set(t)}const n=()=>r(t);return E.$$set=s=>{"onSubmit"in s&&T(0,r=s.onSubmit)},[r,t,R,n]}class $A extends ue{constructor(e){super(),Ce(this,e,wA,WA,_e,{onSubmit:0})}}function xA(E){let e,T,t,r,R,n;return t=new IE({props:{message:E[3],onSubmit:E[5]}}),{c(){e=f("textarea"),T=$(),K(t.$$.fragment),a(e,"rows","6"),a(e,"class","block p-2.5 w-full text-blue-600 hover:text-blue-500 dark:text-blue-500 dark:hover:text-blue-400 bg-gray-50 rounded-lg border border-gray-300 focus:ring-blue-500 focus:border-blue-500 dark:bg-gray-700 dark:border-gray-600 dark:placeholder-gray-400 dark:focus:ring-blue-500 dark:focus:border-blue-500 font-mono"),a(e,"placeholder",E[2])},m(s,S){V(s,e,S),Ye(e,E[0]),V(s,T,S),X(t,s,S),r=!0,R||(n=Ne(e,"input",E[4]),R=!0)},p(s,[S]){(!r||S&4)&&a(e,"placeholder",s[2]),S&1&&Ye(e,s[0]);const A={};S&8&&(A.message=s[3]),S&3&&(A.onSubmit=s[5]),t.$set(A)},i(s){r||(m(t.$$.fragment,s),r=!0)},o(s){y(t.$$.fragment,s),r=!1},d(s){s&&(Y(e),Y(T)),k(t,s),R=!1,n()}}}function XA(E,e,T){let{onSubmit:t}=e,{currentValue:r}=e,{placeholder:R}=e,{buttonText:n}=e;function s(){r=this.value,T(0,r)}const S=()=>t(r);return E.$$set=A=>{"onSubmit"in A&&T(1,t=A.onSubmit),"currentValue"in A&&T(0,r=A.currentValue),"placeholder"in A&&T(2,R=A.placeholder),"buttonText"in A&&T(3,n=A.buttonText)},[r,t,R,n,s,S]}class kA extends ue{constructor(e){super(),Ce(this,e,XA,xA,_e,{onSubmit:1,currentValue:0,placeholder:2,buttonText:3})}}function KA(E){let e,T;return e=new IE({props:{message:"Play",onSubmit:E[2]}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,[r]){const R={};r&1&&(R.onSubmit=t[2]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function er(E){if("speechSynthesis"in window){const e=new SpeechSynthesisUtterance(E);e.lang="en-US",e.volume=1,e.rate=1,e.pitch=1,window.speechSynthesis.speak(e)}else console.error("SpeechSynthesis API is not supported in this browser.")}function JA(E,e,T){let t;eE(E,pR,n=>T(1,t=n));let{message:r}=e;const R=()=>er(r);return E.$$set=n=>{"message"in n&&T(0,r=n.message)},E.$$.update=()=>{E.$$.dirty&3&&t&&er(r)},[r,t,R]}class qA extends ue{constructor(e){super(),Ce(this,e,JA,KA,_e,{message:0})}}function QA(E){let e,T,t;return{c(){e=f("button"),e.textContent="Open Debugger",T=$(),t=f("div"),t.innerHTML='<div class="flex justify-between items-center py-3 px-4 border-b dark:border-neutral-700"><h3 class="font-bold text-gray-800 dark:text-white">Server Logs</h3> <button type="button" class="flex justify-center items-center size-7 text-sm font-semibold rounded-full border border-transparent text-gray-800 hover:bg-gray-100 disabled:opacity-50 disabled:pointer-events-none dark:text-white dark:hover:bg-neutral-700" data-hs-overlay="#hs-overlay-right"><span class="sr-only">Close modal</span> <svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M18 6 6 18"></path><path d="m6 6 12 12"></path></svg></button></div> <div class="p-4"><p id="log-contents" class="text-gray-800 dark:text-neutral-400"></p></div>',a(e,"type","button"),a(e,"class","absolute top-0 right-0 m-1 ms-0 py-3 px-4 inline-flex items-center gap-x-2 text-sm font-semibold rounded-lg border border-transparent bg-blue-600 text-white hover:bg-blue-700 disabled:opacity-50 disabled:pointer-events-none"),a(e,"data-hs-overlay","#hs-overlay-right"),a(t,"id","hs-overlay-right"),a(t,"class","hs-overlay hs-overlay-open:translate-x-0 hidden translate-x-full fixed top-0 end-0 transition-all duration-300 transform h-full max-w-xs w-full z-[80] bg-white border-s dark:bg-neutral-800 dark:border-neutral-700 [--body-scroll:true] overflow-y-auto"),a(t,"tabindex","-1")},m(r,R){V(r,e,R),V(r,T,R),V(r,t,R)},p:j,i:j,o:j,d(r){r&&(Y(e),Y(T),Y(t))}}}class ZA extends ue{constructor(e){super(),Ce(this,e,null,QA,_e,{})}}var yR={exports:{}};(function(E){(function(e,T){E.exports?E.exports=T():e.nearley=T()})(Jn,function(){function e(A,o,i){return this.id=++e.highestId,this.name=A,this.symbols=o,this.postprocess=i,this}e.highestId=0,e.prototype.toString=function(A){var o=typeof A>"u"?this.symbols.map(S).join(" "):this.symbols.slice(0,A).map(S).join(" ")+" ● "+this.symbols.slice(A).map(S).join(" ");return this.name+" → "+o};function T(A,o,i,_){this.rule=A,this.dot=o,this.reference=i,this.data=[],this.wantedBy=_,this.isComplete=this.dot===A.symbols.length}T.prototype.toString=function(){return"{"+this.rule.toString(this.dot)+"}, from: "+(this.reference||0)},T.prototype.nextState=function(A){var o=new T(this.rule,this.dot+1,this.reference,this.wantedBy);return o.left=this,o.right=A,o.isComplete&&(o.data=o.build(),o.right=void 0),o},T.prototype.build=function(){var A=[],o=this;do A.push(o.right.data),o=o.left;while(o.left);return A.reverse(),A},T.prototype.finish=function(){this.rule.postprocess&&(this.data=this.rule.postprocess(this.data,this.reference,n.fail))};function t(A,o){this.grammar=A,this.index=o,this.states=[],this.wants={},this.scannable=[],this.completed={}}t.prototype.process=function(A){for(var o=this.states,i=this.wants,_=this.completed,c=0;c<o.length;c++){var P=o[c];if(P.isComplete){if(P.finish(),P.data!==n.fail){for(var p=P.wantedBy,C=p.length;C--;){var L=p[C];this.complete(L,P)}if(P.reference===this.index){var I=P.rule.name;(this.completed[I]=this.completed[I]||[]).push(P)}}}else{var I=P.rule.symbols[P.dot];if(typeof I!="string"){this.scannable.push(P);continue}if(i[I]){if(i[I].push(P),_.hasOwnProperty(I))for(var u=_[I],C=0;C<u.length;C++){var H=u[C];this.complete(P,H)}}else i[I]=[P],this.predict(I)}}},t.prototype.predict=function(A){for(var o=this.grammar.byName[A]||[],i=0;i<o.length;i++){var _=o[i],c=this.wants[A],P=new T(_,0,this.index,c);this.states.push(P)}},t.prototype.complete=function(A,o){var i=A.nextState(o);this.states.push(i)};function r(A,o){this.rules=A,this.start=o||this.rules[0].name;var i=this.byName={};this.rules.forEach(function(_){i.hasOwnProperty(_.name)||(i[_.name]=[]),i[_.name].push(_)})}r.fromCompiled=function(_,o){var i=_.Lexer;_.ParserStart&&(o=_.ParserStart,_=_.ParserRules);var _=_.map(function(P){return new e(P.name,P.symbols,P.postprocess)}),c=new r(_,o);return c.lexer=i,c};function R(){this.reset("")}R.prototype.reset=function(A,o){this.buffer=A,this.index=0,this.line=o?o.line:1,this.lastLineBreak=o?-o.col:0},R.prototype.next=function(){if(this.index<this.buffer.length){var A=this.buffer[this.index++];return A===`
`&&(this.line+=1,this.lastLineBreak=this.index),{value:A}}},R.prototype.save=function(){return{line:this.line,col:this.index-this.lastLineBreak}},R.prototype.formatError=function(A,o){var i=this.buffer;if(typeof i=="string"){var _=i.split(`
`).slice(Math.max(0,this.line-5),this.line),c=i.indexOf(`
`,this.index);c===-1&&(c=i.length);var P=this.index-this.lastLineBreak,p=String(this.line).length;return o+=" at line "+this.line+" col "+P+`:

`,o+=_.map(function(L,I){return C(this.line-_.length+I+1,p)+" "+L},this).join(`
`),o+=`
`+C("",p+P)+`^
`,o}else return o+" at index "+(this.index-1);function C(L,I){var u=String(L);return Array(I-u.length+1).join(" ")+u}};function n(A,o,i){if(A instanceof r)var _=A,i=o;else var _=r.fromCompiled(A,o);this.grammar=_,this.options={keepHistory:!1,lexer:_.lexer||new R};for(var c in i||{})this.options[c]=i[c];this.lexer=this.options.lexer,this.lexerState=void 0;var P=new t(_,0);this.table=[P],P.wants[_.start]=[],P.predict(_.start),P.process(),this.current=0}n.fail={},n.prototype.feed=function(A){var o=this.lexer;o.reset(A,this.lexerState);for(var i;;){try{if(i=o.next(),!i)break}catch(O){var p=new t(this.grammar,this.current+1);this.table.push(p);var _=new Error(this.reportLexerError(O));throw _.offset=this.current,_.token=O.token,_}var c=this.table[this.current];this.options.keepHistory||delete this.table[this.current-1];var P=this.current+1,p=new t(this.grammar,P);this.table.push(p);for(var C=i.text!==void 0?i.text:i.value,L=o.constructor===R?i.value:i,I=c.scannable,u=I.length;u--;){var H=I[u],b=H.rule.symbols[H.dot];if(b.test?b.test(L):b.type?b.type===i.type:b.literal===C){var M=H.nextState({data:L,token:i,isToken:!0,reference:P-1});p.states.push(M)}}if(p.process(),p.states.length===0){var _=new Error(this.reportError(i));throw _.offset=this.current,_.token=i,_}this.options.keepHistory&&(c.lexerState=o.save()),this.current++}return c&&(this.lexerState=o.save()),this.results=this.finish(),this},n.prototype.reportLexerError=function(A){var o,i,_=A.token;return _?(o="input "+JSON.stringify(_.text[0])+" (lexer error)",i=this.lexer.formatError(_,"Syntax error")):(o="input (lexer error)",i=A.message),this.reportErrorCommon(i,o)},n.prototype.reportError=function(A){var o=(A.type?A.type+" token: ":"")+JSON.stringify(A.value!==void 0?A.value:A),i=this.lexer.formatError(A,"Syntax error");return this.reportErrorCommon(i,o)},n.prototype.reportErrorCommon=function(A,o){var i=[];i.push(A);var _=this.table.length-2,c=this.table[_],P=c.states.filter(function(C){var L=C.rule.symbols[C.dot];return L&&typeof L!="string"});if(P.length===0)i.push("Unexpected "+o+`. I did not expect any more input. Here is the state of my parse table:
`),this.displayStateStack(c.states,i);else{i.push("Unexpected "+o+`. Instead, I was expecting to see one of the following:
`);var p=P.map(function(C){return this.buildFirstStateStack(C,[])||[C]},this);p.forEach(function(C){var L=C[0],I=L.rule.symbols[L.dot],u=this.getSymbolDisplay(I);i.push("A "+u+" based on:"),this.displayStateStack(C,i)},this)}return i.push(""),i.join(`
`)},n.prototype.displayStateStack=function(A,o){for(var i,_=0,c=0;c<A.length;c++){var P=A[c],p=P.rule.toString(P.dot);p===i?_++:(_>0&&o.push("    ^ "+_+" more lines identical to this"),_=0,o.push("    "+p)),i=p}},n.prototype.getSymbolDisplay=function(A){return s(A)},n.prototype.buildFirstStateStack=function(A,o){if(o.indexOf(A)!==-1)return null;if(A.wantedBy.length===0)return[A];var i=A.wantedBy[0],_=[A].concat(o),c=this.buildFirstStateStack(i,_);return c===null?null:[A].concat(c)},n.prototype.save=function(){var A=this.table[this.current];return A.lexerState=this.lexerState,A},n.prototype.restore=function(A){var o=A.index;this.current=o,this.table[o]=A,this.table.splice(o+1),this.lexerState=A.lexerState,this.results=this.finish()},n.prototype.rewind=function(A){if(!this.options.keepHistory)throw new Error("set option `keepHistory` to enable rewinding");this.restore(this.table[A])},n.prototype.finish=function(){var A=[],o=this.grammar.start,i=this.table[this.table.length-1];return i.states.forEach(function(_){_.rule.name===o&&_.dot===_.rule.symbols.length&&_.reference===0&&_.data!==n.fail&&A.push(_)}),A.map(function(_){return _.data})};function s(A){var o=typeof A;if(o==="string")return A;if(o==="object"){if(A.literal)return JSON.stringify(A.literal);if(A instanceof RegExp)return"character matching "+A;if(A.type)return A.type+" token";if(A.test)return"token matching "+String(A.test);throw new Error("Unknown symbol type: "+A)}}function S(A){var o=typeof A;if(o==="string")return A;if(o==="object"){if(A.literal)return JSON.stringify(A.literal);if(A instanceof RegExp)return A.toString();if(A.type)return"%"+A.type;if(A.test)return"<"+String(A.test)+">";throw new Error("Unknown symbol type: "+A)}}return{Parser:n,Grammar:r,Rule:e}})})(yR);var jA=yR.exports;const zA=qn(jA);var BR=Object.defineProperty,es=Object.defineProperties,Es=Object.getOwnPropertyDescriptors,Mt=Object.getOwnPropertySymbols,vR=Object.prototype.hasOwnProperty,FR=Object.prototype.propertyIsEnumerable,Er=(E,e,T)=>e in E?BR(E,e,{enumerable:!0,configurable:!0,writable:!0,value:T}):E[e]=T,EE=(E,e)=>{for(var T in e||(e={}))vR.call(e,T)&&Er(E,T,e[T]);if(Mt)for(var T of Mt(e))FR.call(e,T)&&Er(E,T,e[T]);return E},AE=(E,e)=>es(E,Es(e)),ts=(E,e)=>{var T={};for(var t in E)vR.call(E,t)&&e.indexOf(t)<0&&(T[t]=E[t]);if(E!=null&&Mt)for(var t of Mt(E))e.indexOf(t)<0&&FR.call(E,t)&&(T[t]=E[t]);return T},Ts=(E,e)=>{for(var T in e)BR(E,T,{get:e[T],enumerable:!0})},YR={};Ts(YR,{bigquery:()=>cs,db2:()=>bs,db2i:()=>$s,hive:()=>js,mariadb:()=>AS,mysql:()=>lS,n1ql:()=>bS,plsql:()=>$S,postgresql:()=>zS,redshift:()=>so,singlestoredb:()=>oO,snowflake:()=>CO,spark:()=>_o,sql:()=>vo,sqlite:()=>Uo,tidb:()=>dS,transactsql:()=>EO,trino:()=>ko});var v=E=>E.flatMap(rs),rs=E=>Pt(ns(E)).map(Rs),Rs=E=>E.replace(/ +/g," ").trim(),ns=E=>({type:"mandatory_block",items:DT(E,0)[0]}),DT=(E,e,T)=>{const t=[];for(;E[e];){const[r,R]=As(E,e);if(t.push(r),e=R,E[e]==="|")e++;else if(E[e]==="}"||E[e]==="]"){if(T!==E[e])throw new Error(`Unbalanced parenthesis in: ${E}`);return e++,[t,e]}else if(e===E.length){if(T)throw new Error(`Unbalanced parenthesis in: ${E}`);return[t,e]}else throw new Error(`Unexpected "${E[e]}"`)}return[t,e]},As=(E,e)=>{const T=[];for(;;){const[t,r]=ss(E,e);if(t)T.push(t),e=r;else break}return T.length===1?[T[0],e]:[{type:"concatenation",items:T},e]},ss=(E,e)=>{if(E[e]==="{")return Ss(E,e+1);if(E[e]==="[")return os(E,e+1);{let T="";for(;E[e]&&/[A-Za-z0-9_ ]/.test(E[e]);)T+=E[e],e++;return[T,e]}},Ss=(E,e)=>{const[T,t]=DT(E,e,"}");return[{type:"mandatory_block",items:T},t]},os=(E,e)=>{const[T,t]=DT(E,e,"]");return[{type:"optional_block",items:T},t]},Pt=E=>{if(typeof E=="string")return[E];if(E.type==="concatenation")return E.items.map(Pt).reduce(Os,[""]);if(E.type==="mandatory_block")return E.items.flatMap(Pt);if(E.type==="optional_block")return["",...E.items.flatMap(Pt)];throw new Error(`Unknown node type: ${E}`)},Os=(E,e)=>{const T=[];for(const t of E)for(const r of e)T.push(t+r);return T},VR=(E=>(E.QUOTED_IDENTIFIER="QUOTED_IDENTIFIER",E.IDENTIFIER="IDENTIFIER",E.STRING="STRING",E.VARIABLE="VARIABLE",E.RESERVED_DATA_TYPE="RESERVED_DATA_TYPE",E.RESERVED_PARAMETERIZED_DATA_TYPE="RESERVED_PARAMETERIZED_DATA_TYPE",E.RESERVED_KEYWORD="RESERVED_KEYWORD",E.RESERVED_FUNCTION_NAME="RESERVED_FUNCTION_NAME",E.RESERVED_PHRASE="RESERVED_PHRASE",E.RESERVED_SET_OPERATION="RESERVED_SET_OPERATION",E.RESERVED_CLAUSE="RESERVED_CLAUSE",E.RESERVED_SELECT="RESERVED_SELECT",E.RESERVED_JOIN="RESERVED_JOIN",E.ARRAY_IDENTIFIER="ARRAY_IDENTIFIER",E.ARRAY_KEYWORD="ARRAY_KEYWORD",E.CASE="CASE",E.END="END",E.WHEN="WHEN",E.ELSE="ELSE",E.THEN="THEN",E.LIMIT="LIMIT",E.BETWEEN="BETWEEN",E.AND="AND",E.OR="OR",E.XOR="XOR",E.OPERATOR="OPERATOR",E.COMMA="COMMA",E.ASTERISK="ASTERISK",E.PROPERTY_ACCESS_OPERATOR="PROPERTY_ACCESS_OPERATOR",E.OPEN_PAREN="OPEN_PAREN",E.CLOSE_PAREN="CLOSE_PAREN",E.LINE_COMMENT="LINE_COMMENT",E.BLOCK_COMMENT="BLOCK_COMMENT",E.DISABLE_COMMENT="DISABLE_COMMENT",E.NUMBER="NUMBER",E.NAMED_PARAMETER="NAMED_PARAMETER",E.QUOTED_PARAMETER="QUOTED_PARAMETER",E.NUMBERED_PARAMETER="NUMBERED_PARAMETER",E.POSITIONAL_PARAMETER="POSITIONAL_PARAMETER",E.CUSTOM_PARAMETER="CUSTOM_PARAMETER",E.DELIMITER="DELIMITER",E.EOF="EOF",E))(VR||{}),WR=E=>({type:"EOF",raw:"«EOF»",text:"«EOF»",start:E}),tt=WR(1/0),QE=E=>e=>e.type===E.type&&e.text===E.text,FE={ARRAY:QE({text:"ARRAY",type:"RESERVED_DATA_TYPE"}),BY:QE({text:"BY",type:"RESERVED_KEYWORD"}),SET:QE({text:"SET",type:"RESERVED_CLAUSE"}),STRUCT:QE({text:"STRUCT",type:"RESERVED_DATA_TYPE"}),WINDOW:QE({text:"WINDOW",type:"RESERVED_CLAUSE"}),VALUES:QE({text:"VALUES",type:"RESERVED_CLAUSE"})},wR=E=>E==="RESERVED_DATA_TYPE"||E==="RESERVED_KEYWORD"||E==="RESERVED_FUNCTION_NAME"||E==="RESERVED_PHRASE"||E==="RESERVED_CLAUSE"||E==="RESERVED_SELECT"||E==="RESERVED_SET_OPERATION"||E==="RESERVED_JOIN"||E==="ARRAY_KEYWORD"||E==="CASE"||E==="END"||E==="WHEN"||E==="ELSE"||E==="THEN"||E==="LIMIT"||E==="BETWEEN"||E==="AND"||E==="OR"||E==="XOR",is=E=>E==="AND"||E==="OR"||E==="XOR",as=["KEYS.NEW_KEYSET","KEYS.ADD_KEY_FROM_RAW_BYTES","AEAD.DECRYPT_BYTES","AEAD.DECRYPT_STRING","AEAD.ENCRYPT","KEYS.KEYSET_CHAIN","KEYS.KEYSET_FROM_JSON","KEYS.KEYSET_TO_JSON","KEYS.ROTATE_KEYSET","KEYS.KEYSET_LENGTH","ANY_VALUE","ARRAY_AGG","AVG","CORR","COUNT","COUNTIF","COVAR_POP","COVAR_SAMP","MAX","MIN","ST_CLUSTERDBSCAN","STDDEV_POP","STDDEV_SAMP","STRING_AGG","SUM","VAR_POP","VAR_SAMP","ANY_VALUE","ARRAY_AGG","ARRAY_CONCAT_AGG","AVG","BIT_AND","BIT_OR","BIT_XOR","COUNT","COUNTIF","LOGICAL_AND","LOGICAL_OR","MAX","MIN","STRING_AGG","SUM","APPROX_COUNT_DISTINCT","APPROX_QUANTILES","APPROX_TOP_COUNT","APPROX_TOP_SUM","ARRAY_CONCAT","ARRAY_LENGTH","ARRAY_TO_STRING","GENERATE_ARRAY","GENERATE_DATE_ARRAY","GENERATE_TIMESTAMP_ARRAY","ARRAY_REVERSE","OFFSET","SAFE_OFFSET","ORDINAL","SAFE_ORDINAL","BIT_COUNT","PARSE_BIGNUMERIC","PARSE_NUMERIC","SAFE_CAST","CURRENT_DATE","EXTRACT","DATE","DATE_ADD","DATE_SUB","DATE_DIFF","DATE_TRUNC","DATE_FROM_UNIX_DATE","FORMAT_DATE","LAST_DAY","PARSE_DATE","UNIX_DATE","CURRENT_DATETIME","DATETIME","EXTRACT","DATETIME_ADD","DATETIME_SUB","DATETIME_DIFF","DATETIME_TRUNC","FORMAT_DATETIME","LAST_DAY","PARSE_DATETIME","ERROR","EXTERNAL_QUERY","S2_CELLIDFROMPOINT","S2_COVERINGCELLIDS","ST_ANGLE","ST_AREA","ST_ASBINARY","ST_ASGEOJSON","ST_ASTEXT","ST_AZIMUTH","ST_BOUNDARY","ST_BOUNDINGBOX","ST_BUFFER","ST_BUFFERWITHTOLERANCE","ST_CENTROID","ST_CENTROID_AGG","ST_CLOSESTPOINT","ST_CLUSTERDBSCAN","ST_CONTAINS","ST_CONVEXHULL","ST_COVEREDBY","ST_COVERS","ST_DIFFERENCE","ST_DIMENSION","ST_DISJOINT","ST_DISTANCE","ST_DUMP","ST_DWITHIN","ST_ENDPOINT","ST_EQUALS","ST_EXTENT","ST_EXTERIORRING","ST_GEOGFROM","ST_GEOGFROMGEOJSON","ST_GEOGFROMTEXT","ST_GEOGFROMWKB","ST_GEOGPOINT","ST_GEOGPOINTFROMGEOHASH","ST_GEOHASH","ST_GEOMETRYTYPE","ST_INTERIORRINGS","ST_INTERSECTION","ST_INTERSECTS","ST_INTERSECTSBOX","ST_ISCOLLECTION","ST_ISEMPTY","ST_LENGTH","ST_MAKELINE","ST_MAKEPOLYGON","ST_MAKEPOLYGONORIENTED","ST_MAXDISTANCE","ST_NPOINTS","ST_NUMGEOMETRIES","ST_NUMPOINTS","ST_PERIMETER","ST_POINTN","ST_SIMPLIFY","ST_SNAPTOGRID","ST_STARTPOINT","ST_TOUCHES","ST_UNION","ST_UNION_AGG","ST_WITHIN","ST_X","ST_Y","FARM_FINGERPRINT","MD5","SHA1","SHA256","SHA512","HLL_COUNT.INIT","HLL_COUNT.MERGE","HLL_COUNT.MERGE_PARTIAL","HLL_COUNT.EXTRACT","MAKE_INTERVAL","EXTRACT","JUSTIFY_DAYS","JUSTIFY_HOURS","JUSTIFY_INTERVAL","JSON_EXTRACT","JSON_QUERY","JSON_EXTRACT_SCALAR","JSON_VALUE","JSON_EXTRACT_ARRAY","JSON_QUERY_ARRAY","JSON_EXTRACT_STRING_ARRAY","JSON_VALUE_ARRAY","TO_JSON_STRING","ABS","SIGN","IS_INF","IS_NAN","IEEE_DIVIDE","RAND","SQRT","POW","POWER","EXP","LN","LOG","LOG10","GREATEST","LEAST","DIV","SAFE_DIVIDE","SAFE_MULTIPLY","SAFE_NEGATE","SAFE_ADD","SAFE_SUBTRACT","MOD","ROUND","TRUNC","CEIL","CEILING","FLOOR","COS","COSH","ACOS","ACOSH","SIN","SINH","ASIN","ASINH","TAN","TANH","ATAN","ATANH","ATAN2","RANGE_BUCKET","FIRST_VALUE","LAST_VALUE","NTH_VALUE","LEAD","LAG","PERCENTILE_CONT","PERCENTILE_DISC","NET.IP_FROM_STRING","NET.SAFE_IP_FROM_STRING","NET.IP_TO_STRING","NET.IP_NET_MASK","NET.IP_TRUNC","NET.IPV4_FROM_INT64","NET.IPV4_TO_INT64","NET.HOST","NET.PUBLIC_SUFFIX","NET.REG_DOMAIN","RANK","DENSE_RANK","PERCENT_RANK","CUME_DIST","NTILE","ROW_NUMBER","SESSION_USER","CORR","COVAR_POP","COVAR_SAMP","STDDEV_POP","STDDEV_SAMP","STDDEV","VAR_POP","VAR_SAMP","VARIANCE","ASCII","BYTE_LENGTH","CHAR_LENGTH","CHARACTER_LENGTH","CHR","CODE_POINTS_TO_BYTES","CODE_POINTS_TO_STRING","CONCAT","CONTAINS_SUBSTR","ENDS_WITH","FORMAT","FROM_BASE32","FROM_BASE64","FROM_HEX","INITCAP","INSTR","LEFT","LENGTH","LPAD","LOWER","LTRIM","NORMALIZE","NORMALIZE_AND_CASEFOLD","OCTET_LENGTH","REGEXP_CONTAINS","REGEXP_EXTRACT","REGEXP_EXTRACT_ALL","REGEXP_INSTR","REGEXP_REPLACE","REGEXP_SUBSTR","REPLACE","REPEAT","REVERSE","RIGHT","RPAD","RTRIM","SAFE_CONVERT_BYTES_TO_STRING","SOUNDEX","SPLIT","STARTS_WITH","STRPOS","SUBSTR","SUBSTRING","TO_BASE32","TO_BASE64","TO_CODE_POINTS","TO_HEX","TRANSLATE","TRIM","UNICODE","UPPER","CURRENT_TIME","TIME","EXTRACT","TIME_ADD","TIME_SUB","TIME_DIFF","TIME_TRUNC","FORMAT_TIME","PARSE_TIME","CURRENT_TIMESTAMP","EXTRACT","STRING","TIMESTAMP","TIMESTAMP_ADD","TIMESTAMP_SUB","TIMESTAMP_DIFF","TIMESTAMP_TRUNC","FORMAT_TIMESTAMP","PARSE_TIMESTAMP","TIMESTAMP_SECONDS","TIMESTAMP_MILLIS","TIMESTAMP_MICROS","UNIX_SECONDS","UNIX_MILLIS","UNIX_MICROS","GENERATE_UUID","COALESCE","IF","IFNULL","NULLIF","AVG","BIT_AND","BIT_OR","BIT_XOR","CORR","COUNT","COVAR_POP","COVAR_SAMP","EXACT_COUNT_DISTINCT","FIRST","GROUP_CONCAT","GROUP_CONCAT_UNQUOTED","LAST","MAX","MIN","NEST","NTH","QUANTILES","STDDEV","STDDEV_POP","STDDEV_SAMP","SUM","TOP","UNIQUE","VARIANCE","VAR_POP","VAR_SAMP","BIT_COUNT","BOOLEAN","BYTES","CAST","FLOAT","HEX_STRING","INTEGER","STRING","COALESCE","GREATEST","IFNULL","IS_INF","IS_NAN","IS_EXPLICITLY_DEFINED","LEAST","NVL","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","DATE","DATE_ADD","DATEDIFF","DAY","DAYOFWEEK","DAYOFYEAR","FORMAT_UTC_USEC","HOUR","MINUTE","MONTH","MSEC_TO_TIMESTAMP","NOW","PARSE_UTC_USEC","QUARTER","SEC_TO_TIMESTAMP","SECOND","STRFTIME_UTC_USEC","TIME","TIMESTAMP","TIMESTAMP_TO_MSEC","TIMESTAMP_TO_SEC","TIMESTAMP_TO_USEC","USEC_TO_TIMESTAMP","UTC_USEC_TO_DAY","UTC_USEC_TO_HOUR","UTC_USEC_TO_MONTH","UTC_USEC_TO_WEEK","UTC_USEC_TO_YEAR","WEEK","YEAR","FORMAT_IP","PARSE_IP","FORMAT_PACKED_IP","PARSE_PACKED_IP","JSON_EXTRACT","JSON_EXTRACT_SCALAR","ABS","ACOS","ACOSH","ASIN","ASINH","ATAN","ATANH","ATAN2","CEIL","COS","COSH","DEGREES","EXP","FLOOR","LN","LOG","LOG2","LOG10","PI","POW","RADIANS","RAND","ROUND","SIN","SINH","SQRT","TAN","TANH","REGEXP_MATCH","REGEXP_EXTRACT","REGEXP_REPLACE","CONCAT","INSTR","LEFT","LENGTH","LOWER","LPAD","LTRIM","REPLACE","RIGHT","RPAD","RTRIM","SPLIT","SUBSTR","UPPER","TABLE_DATE_RANGE","TABLE_DATE_RANGE_STRICT","TABLE_QUERY","HOST","DOMAIN","TLD","AVG","COUNT","MAX","MIN","STDDEV","SUM","CUME_DIST","DENSE_RANK","FIRST_VALUE","LAG","LAST_VALUE","LEAD","NTH_VALUE","NTILE","PERCENT_RANK","PERCENTILE_CONT","PERCENTILE_DISC","RANK","RATIO_TO_REPORT","ROW_NUMBER","CURRENT_USER","EVERY","FROM_BASE64","HASH","FARM_FINGERPRINT","IF","POSITION","SHA1","SOME","TO_BASE64","BQ.JOBS.CANCEL","BQ.REFRESH_MATERIALIZED_VIEW","OPTIONS","PIVOT","UNPIVOT"],Is=["ALL","AND","ANY","AS","ASC","ASSERT_ROWS_MODIFIED","AT","BETWEEN","BY","CASE","CAST","COLLATE","CONTAINS","CREATE","CROSS","CUBE","CURRENT","DEFAULT","DEFINE","DESC","DISTINCT","ELSE","END","ENUM","ESCAPE","EXCEPT","EXCLUDE","EXISTS","EXTRACT","FALSE","FETCH","FOLLOWING","FOR","FROM","FULL","GROUP","GROUPING","GROUPS","HASH","HAVING","IF","IGNORE","IN","INNER","INTERSECT","INTO","IS","JOIN","LATERAL","LEFT","LIMIT","LOOKUP","MERGE","NATURAL","NEW","NO","NOT","NULL","NULLS","OF","ON","OR","ORDER","OUTER","OVER","PARTITION","PRECEDING","PROTO","RANGE","RECURSIVE","RESPECT","RIGHT","ROLLUP","ROWS","SELECT","SET","SOME","TABLE","TABLESAMPLE","THEN","TO","TREAT","TRUE","UNBOUNDED","UNION","UNNEST","USING","WHEN","WHERE","WINDOW","WITH","WITHIN","SAFE","LIKE","COPY","CLONE","IN","OUT","INOUT","RETURNS","LANGUAGE","CASCADE","RESTRICT","DETERMINISTIC"],Ns=["ARRAY","BOOL","BYTES","DATE","DATETIME","GEOGRAPHY","INTERVAL","INT64","INT","SMALLINT","INTEGER","BIGINT","TINYINT","BYTEINT","NUMERIC","DECIMAL","BIGNUMERIC","BIGDECIMAL","FLOAT64","STRING","STRUCT","TIME","TIMEZONE"],ls=v(["SELECT [ALL | DISTINCT] [AS STRUCT | AS VALUE]"]),_s=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","QUALIFY","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","OMIT RECORD IF","INSERT [INTO]","VALUES","SET","MERGE [INTO]","WHEN [NOT] MATCHED [BY SOURCE | BY TARGET] [THEN]","UPDATE SET","CLUSTER BY","FOR SYSTEM_TIME AS OF","WITH CONNECTION","WITH PARTITION COLUMNS","REMOTE WITH CONNECTION"]),tr=v(["CREATE [OR REPLACE] [TEMP|TEMPORARY|SNAPSHOT|EXTERNAL] TABLE [IF NOT EXISTS]"]),xt=v(["CREATE [OR REPLACE] [MATERIALIZED] VIEW [IF NOT EXISTS]","UPDATE","DELETE [FROM]","DROP [SNAPSHOT | EXTERNAL] TABLE [IF EXISTS]","ALTER TABLE [IF EXISTS]","ADD COLUMN [IF NOT EXISTS]","DROP COLUMN [IF EXISTS]","RENAME TO","ALTER COLUMN [IF EXISTS]","SET DEFAULT COLLATE","SET OPTIONS","DROP NOT NULL","SET DATA TYPE","ALTER SCHEMA [IF EXISTS]","ALTER [MATERIALIZED] VIEW [IF EXISTS]","ALTER BI_CAPACITY","TRUNCATE TABLE","CREATE SCHEMA [IF NOT EXISTS]","DEFAULT COLLATE","CREATE [OR REPLACE] [TEMP|TEMPORARY|TABLE] FUNCTION [IF NOT EXISTS]","CREATE [OR REPLACE] PROCEDURE [IF NOT EXISTS]","CREATE [OR REPLACE] ROW ACCESS POLICY [IF NOT EXISTS]","GRANT TO","FILTER USING","CREATE CAPACITY","AS JSON","CREATE RESERVATION","CREATE ASSIGNMENT","CREATE SEARCH INDEX [IF NOT EXISTS]","DROP SCHEMA [IF EXISTS]","DROP [MATERIALIZED] VIEW [IF EXISTS]","DROP [TABLE] FUNCTION [IF EXISTS]","DROP PROCEDURE [IF EXISTS]","DROP ROW ACCESS POLICY","DROP ALL ROW ACCESS POLICIES","DROP CAPACITY [IF EXISTS]","DROP RESERVATION [IF EXISTS]","DROP ASSIGNMENT [IF EXISTS]","DROP SEARCH INDEX [IF EXISTS]","DROP [IF EXISTS]","GRANT","REVOKE","DECLARE","EXECUTE IMMEDIATE","LOOP","END LOOP","REPEAT","END REPEAT","WHILE","END WHILE","BREAK","LEAVE","CONTINUE","ITERATE","FOR","END FOR","BEGIN","BEGIN TRANSACTION","COMMIT TRANSACTION","ROLLBACK TRANSACTION","RAISE","RETURN","CALL","ASSERT","EXPORT DATA"]),Ls=v(["UNION {ALL | DISTINCT}","EXCEPT DISTINCT","INTERSECT DISTINCT"]),Cs=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN"]),us=v(["TABLESAMPLE SYSTEM","ANY TYPE","ALL COLUMNS","NOT DETERMINISTIC","{ROWS | RANGE} BETWEEN","IS [NOT] DISTINCT FROM"]),cs={name:"bigquery",tokenizerOptions:{reservedSelect:ls,reservedClauses:[..._s,...xt,...tr],reservedSetOperations:Ls,reservedJoins:Cs,reservedPhrases:us,reservedKeywords:Is,reservedDataTypes:Ns,reservedFunctionNames:as,extraParens:["[]"],stringTypes:[{quote:'""".."""',prefixes:["R","B","RB","BR"]},{quote:"\'\'\'..\'\'\'",prefixes:["R","B","RB","BR"]},'""-bs',"''-bs",{quote:'""-raw',prefixes:["R","B","RB","BR"],requirePrefix:!0},{quote:"''-raw",prefixes:["R","B","RB","BR"],requirePrefix:!0}],identTypes:["``"],identChars:{dashes:!0},paramTypes:{positional:!0,named:["@"],quoted:["@"]},variableTypes:[{regex:String.raw`@@\\w+`}],lineCommentTypes:["--","#"],operators:["&","|","^","~",">>","<<","||","=>"],postProcess:fs},formatOptions:{onelineClauses:[...tr,...xt],tabularOnelineClauses:xt}};function fs(E){return Ps(Ds(E))}function Ps(E){let e=tt;return E.map(T=>T.text==="OFFSET"&&e.text==="["?(e=T,AE(EE({},T),{type:"RESERVED_FUNCTION_NAME"})):(e=T,T))}function Ds(E){var e;const T=[];for(let t=0;t<E.length;t++){const r=E[t];if((FE.ARRAY(r)||FE.STRUCT(r))&&((e=E[t+1])==null?void 0:e.text)==="<"){const R=ds(E,t+1),n=E.slice(t,R+1);T.push({type:"IDENTIFIER",raw:n.map(Tr("raw")).join(""),text:n.map(Tr("text")).join(""),start:r.start}),t=R}else T.push(r)}return T}var Tr=E=>e=>e.type==="IDENTIFIER"||e.type==="COMMA"?e[E]+" ":e[E];function ds(E,e){let T=0;for(let t=e;t<E.length;t++){const r=E[t];if(r.text==="<"?T++:r.text===">"?T--:r.text===">>"&&(T-=2),T===0)return t}return E.length-1}var ps=["ARRAY_AGG","AVG","CORRELATION","COUNT","COUNT_BIG","COVARIANCE","COVARIANCE_SAMP","CUME_DIST","GROUPING","LISTAGG","MAX","MEDIAN","MIN","PERCENTILE_CONT","PERCENTILE_DISC","PERCENT_RANK","REGR_AVGX","REGR_AVGY","REGR_COUNT","REGR_INTERCEPT","REGR_ICPT","REGR_R2","REGR_SLOPE","REGR_SXX","REGR_SXY","REGR_SYY","STDDEV","STDDEV_SAMP","SUM","VARIANCE","VARIANCE_SAMP","XMLAGG","XMLGROUP","ABS","ABSVAL","ACOS","ADD_DAYS","ADD_HOURS","ADD_MINUTES","ADD_MONTHS","ADD_SECONDS","ADD_YEARS","AGE","ARRAY_DELETE","ARRAY_FIRST","ARRAY_LAST","ARRAY_NEXT","ARRAY_PRIOR","ASCII","ASCII_STR","ASIN","ATAN","ATAN2","ATANH","BITAND","BITANDNOT","BITOR","BITXOR","BITNOT","BPCHAR","BSON_TO_JSON","BTRIM","CARDINALITY","CEILING","CEIL","CHARACTER_LENGTH","CHR","COALESCE","COLLATION_KEY","COLLATION_KEY_BIT","COMPARE_DECFLOAT","CONCAT","COS","COSH","COT","CURSOR_ROWCOUNT","DATAPARTITIONNUM","DATE_PART","DATE_TRUNC","DAY","DAYNAME","DAYOFMONTH","DAYOFWEEK","DAYOFWEEK_ISO","DAYOFYEAR","DAYS","DAYS_BETWEEN","DAYS_TO_END_OF_MONTH","DBPARTITIONNUM","DECFLOAT","DECFLOAT_FORMAT","DECODE","DECRYPT_BIN","DECRYPT_CHAR","DEGREES","DEREF","DIFFERENCE","DIGITS","DOUBLE_PRECISION","EMPTY_BLOB","EMPTY_CLOB","EMPTY_DBCLOB","EMPTY_NCLOB","ENCRYPT","EVENT_MON_STATE","EXP","EXTRACT","FIRST_DAY","FLOOR","FROM_UTC_TIMESTAMP","GENERATE_UNIQUE","GETHINT","GREATEST","HASH","HASH4","HASH8","HASHEDVALUE","HEX","HEXTORAW","HOUR","HOURS_BETWEEN","IDENTITY_VAL_LOCAL","IFNULL","INITCAP","INSERT","INSTR","INSTR2","INSTR4","INSTRB","INTNAND","INTNOR","INTNXOR","INTNNOT","ISNULL","JSON_ARRAY","JSON_OBJECT","JSON_QUERY","JSON_TO_BSON","JSON_VALUE","JULIAN_DAY","LAST_DAY","LCASE","LEAST","LEFT","LENGTH","LENGTH2","LENGTH4","LENGTHB","LN","LOCATE","LOCATE_IN_STRING","LOG10","LONG_VARCHAR","LONG_VARGRAPHIC","LOWER","LPAD","LTRIM","MAX","MAX_CARDINALITY","MICROSECOND","MIDNIGHT_SECONDS","MIN","MINUTE","MINUTES_BETWEEN","MOD","MONTH","MONTHNAME","MONTHS_BETWEEN","MULTIPLY_ALT","NEXT_DAY","NEXT_MONTH","NEXT_QUARTER","NEXT_WEEK","NEXT_YEAR","NORMALIZE_DECFLOAT","NOW","NULLIF","NVL","NVL2","OCTET_LENGTH","OVERLAY","PARAMETER","POSITION","POSSTR","POW","POWER","QUANTIZE","QUARTER","QUOTE_IDENT","QUOTE_LITERAL","RADIANS","RAISE_ERROR","RAND","RANDOM","RAWTOHEX","REC2XML","REGEXP_COUNT","REGEXP_EXTRACT","REGEXP_INSTR","REGEXP_LIKE","REGEXP_MATCH_COUNT","REGEXP_REPLACE","REGEXP_SUBSTR","REPEAT","REPLACE","RID","RID_BIT","RIGHT","ROUND","ROUND_TIMESTAMP","RPAD","RTRIM","SECLABEL","SECLABEL_BY_NAME","SECLABEL_TO_CHAR","SECOND","SECONDS_BETWEEN","SIGN","SIN","SINH","SOUNDEX","SPACE","SQRT","STRIP","STRLEFT","STRPOS","STRRIGHT","SUBSTR","SUBSTR2","SUBSTR4","SUBSTRB","SUBSTRING","TABLE_NAME","TABLE_SCHEMA","TAN","TANH","THIS_MONTH","THIS_QUARTER","THIS_WEEK","THIS_YEAR","TIMESTAMP_FORMAT","TIMESTAMP_ISO","TIMESTAMPDIFF","TIMEZONE","TO_CHAR","TO_CLOB","TO_DATE","TO_HEX","TO_MULTI_BYTE","TO_NCHAR","TO_NCLOB","TO_NUMBER","TO_SINGLE_BYTE","TO_TIMESTAMP","TO_UTC_TIMESTAMP","TOTALORDER","TRANSLATE","TRIM","TRIM_ARRAY","TRUNC_TIMESTAMP","TRUNCATE","TRUNC","TYPE_ID","TYPE_NAME","TYPE_SCHEMA","UCASE","UNICODE_STR","UPPER","VALUE","VARCHAR_BIT_FORMAT","VARCHAR_FORMAT","VARCHAR_FORMAT_BIT","VERIFY_GROUP_FOR_USER","VERIFY_ROLE_FOR_USER","VERIFY_TRUSTED_CONTEXT_ROLE_FOR_USER","WEEK","WEEK_ISO","WEEKS_BETWEEN","WIDTH_BUCKET","XMLATTRIBUTES","XMLCOMMENT","XMLCONCAT","XMLDOCUMENT","XMLELEMENT","XMLFOREST","XMLNAMESPACES","XMLPARSE","XMLPI","XMLQUERY","XMLROW","XMLSERIALIZE","XMLTEXT","XMLVALIDATE","XMLXSROBJECTID","XSLTRANSFORM","YEAR","YEARS_BETWEEN","YMD_BETWEEN","BASE_TABLE","JSON_TABLE","UNNEST","XMLTABLE","RANK","DENSE_RANK","NTILE","LAG","LEAD","ROW_NUMBER","FIRST_VALUE","LAST_VALUE","NTH_VALUE","RATIO_TO_REPORT","CAST"],Ms=["ACTIVATE","ADD","AFTER","ALIAS","ALL","ALLOCATE","ALLOW","ALTER","AND","ANY","AS","ASENSITIVE","ASSOCIATE","ASUTIME","AT","ATTRIBUTES","AUDIT","AUTHORIZATION","AUX","AUXILIARY","BEFORE","BEGIN","BETWEEN","BINARY","BUFFERPOOL","BY","CACHE","CALL","CALLED","CAPTURE","CARDINALITY","CASCADED","CASE","CAST","CHECK","CLONE","CLOSE","CLUSTER","COLLECTION","COLLID","COLUMN","COMMENT","COMMIT","CONCAT","CONDITION","CONNECT","CONNECTION","CONSTRAINT","CONTAINS","CONTINUE","COUNT","COUNT_BIG","CREATE","CROSS","CURRENT","CURRENT_DATE","CURRENT_LC_CTYPE","CURRENT_PATH","CURRENT_SCHEMA","CURRENT_SERVER","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_TIMEZONE","CURRENT_USER","CURSOR","CYCLE","DATA","DATABASE","DATAPARTITIONNAME","DATAPARTITIONNUM","DAY","DAYS","DB2GENERAL","DB2GENRL","DB2SQL","DBINFO","DBPARTITIONNAME","DBPARTITIONNUM","DEALLOCATE","DECLARE","DEFAULT","DEFAULTS","DEFINITION","DELETE","DENSERANK","DENSE_RANK","DESCRIBE","DESCRIPTOR","DETERMINISTIC","DIAGNOSTICS","DISABLE","DISALLOW","DISCONNECT","DISTINCT","DO","DOCUMENT","DROP","DSSIZE","DYNAMIC","EACH","EDITPROC","ELSE","ELSEIF","ENABLE","ENCODING","ENCRYPTION","END","END-EXEC","ENDING","ERASE","ESCAPE","EVERY","EXCEPT","EXCEPTION","EXCLUDING","EXCLUSIVE","EXECUTE","EXISTS","EXIT","EXPLAIN","EXTENDED","EXTERNAL","EXTRACT","FENCED","FETCH","FIELDPROC","FILE","FINAL","FIRST1","FOR","FOREIGN","FREE","FROM","FULL","FUNCTION","GENERAL","GENERATED","GET","GLOBAL","GO","GOTO","GRANT","GRAPHIC","GROUP","HANDLER","HASH","HASHED_VALUE","HAVING","HINT","HOLD","HOUR","HOURS","IDENTITY","IF","IMMEDIATE","IMPORT","IN","INCLUDING","INCLUSIVE","INCREMENT","INDEX","INDICATOR","INDICATORS","INF","INFINITY","INHERIT","INNER","INOUT","INSENSITIVE","INSERT","INTEGRITY","INTERSECT","INTO","IS","ISNULL","ISOBID","ISOLATION","ITERATE","JAR","JAVA","JOIN","KEEP","KEY","LABEL","LANGUAGE","LAST3","LATERAL","LC_CTYPE","LEAVE","LEFT","LIKE","LIMIT","LINKTYPE","LOCAL","LOCALDATE","LOCALE","LOCALTIME","LOCALTIMESTAMP","LOCATOR","LOCATORS","LOCK","LOCKMAX","LOCKSIZE","LOOP","MAINTAINED","MATERIALIZED","MAXVALUE","MICROSECOND","MICROSECONDS","MINUTE","MINUTES","MINVALUE","MODE","MODIFIES","MONTH","MONTHS","NAN","NEW","NEW_TABLE","NEXTVAL","NO","NOCACHE","NOCYCLE","NODENAME","NODENUMBER","NOMAXVALUE","NOMINVALUE","NONE","NOORDER","NORMALIZED","NOT2","NOTNULL","NULL","NULLS","NUMPARTS","OBID","OF","OFF","OFFSET","OLD","OLD_TABLE","ON","OPEN","OPTIMIZATION","OPTIMIZE","OPTION","OR","ORDER","OUT","OUTER","OVER","OVERRIDING","PACKAGE","PADDED","PAGESIZE","PARAMETER","PART","PARTITION","PARTITIONED","PARTITIONING","PARTITIONS","PASSWORD","PATH","PERCENT","PIECESIZE","PLAN","POSITION","PRECISION","PREPARE","PREVVAL","PRIMARY","PRIQTY","PRIVILEGES","PROCEDURE","PROGRAM","PSID","PUBLIC","QUERY","QUERYNO","RANGE","RANK","READ","READS","RECOVERY","REFERENCES","REFERENCING","REFRESH","RELEASE","RENAME","REPEAT","RESET","RESIGNAL","RESTART","RESTRICT","RESULT","RESULT_SET_LOCATOR","RETURN","RETURNS","REVOKE","RIGHT","ROLE","ROLLBACK","ROUND_CEILING","ROUND_DOWN","ROUND_FLOOR","ROUND_HALF_DOWN","ROUND_HALF_EVEN","ROUND_HALF_UP","ROUND_UP","ROUTINE","ROW","ROWNUMBER","ROWS","ROWSET","ROW_NUMBER","RRN","RUN","SAVEPOINT","SCHEMA","SCRATCHPAD","SCROLL","SEARCH","SECOND","SECONDS","SECQTY","SECURITY","SELECT","SENSITIVE","SEQUENCE","SESSION","SESSION_USER","SET","SIGNAL","SIMPLE","SNAN","SOME","SOURCE","SPECIFIC","SQL","SQLID","STACKED","STANDARD","START","STARTING","STATEMENT","STATIC","STATMENT","STAY","STOGROUP","STORES","STYLE","SUBSTRING","SUMMARY","SYNONYM","SYSFUN","SYSIBM","SYSPROC","SYSTEM","SYSTEM_USER","TABLE","TABLESPACE","THEN","TO","TRANSACTION","TRIGGER","TRIM","TRUNCATE","TYPE","UNDO","UNION","UNIQUE","UNTIL","UPDATE","USAGE","USER","USING","VALIDPROC","VALUE","VALUES","VARIABLE","VARIANT","VCAT","VERSION","VIEW","VOLATILE","VOLUMES","WHEN","WHENEVER","WHERE","WHILE","WITH","WITHOUT","WLM","WRITE","XMLELEMENT","XMLEXISTS","XMLNAMESPACES","YEAR","YEARS"],Us=["ARRAY","BIGINT","BINARY","BLOB","BOOLEAN","CCSID","CHAR","CHARACTER","CLOB","DATE","DATETIME","DBCLOB","DEC","DECIMAL","DOUBLE","DOUBLE PRECISION","FLOAT","FLOAT4","FLOAT8","GRAPHIC","INT","INT2","INT4","INT8","INTEGER","INTERVAL","LONG VARCHAR","LONG VARGRAPHIC","NCHAR","NCHR","NCLOB","NVARCHAR","NUMERIC","SMALLINT","REAL","TIME","TIMESTAMP","VARBINARY","VARCHAR","VARGRAPHIC"],ms=v(["SELECT [ALL | DISTINCT]"]),hs=v(["WITH","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER BY [INPUT SEQUENCE]","LIMIT","OFFSET","FETCH NEXT","FOR UPDATE [OF]","FOR {READ | FETCH} ONLY","FOR {RR | CS | UR | RS} [USE AND KEEP {SHARE | UPDATE | EXCLUSIVE} LOCKS]","WAIT FOR OUTCOME","SKIP LOCKED DATA","INTO","INSERT INTO","VALUES","SET","MERGE INTO","WHEN [NOT] MATCHED [THEN]","UPDATE SET","INSERT"]),rr=v(["CREATE [GLOBAL TEMPORARY | EXTERNAL] TABLE [IF NOT EXISTS]"]),Xt=v(["CREATE [OR REPLACE] VIEW","UPDATE","WHERE CURRENT OF","WITH {RR | RS | CS | UR}","DELETE FROM","DROP TABLE [IF EXISTS]","ALTER TABLE","ADD [COLUMN]","DROP [COLUMN]","RENAME COLUMN","ALTER [COLUMN]","SET DATA TYPE","SET NOT NULL","DROP {DEFAULT | GENERATED | NOT NULL}","TRUNCATE [TABLE]","ALLOCATE","ALTER AUDIT POLICY","ALTER BUFFERPOOL","ALTER DATABASE PARTITION GROUP","ALTER DATABASE","ALTER EVENT MONITOR","ALTER FUNCTION","ALTER HISTOGRAM TEMPLATE","ALTER INDEX","ALTER MASK","ALTER METHOD","ALTER MODULE","ALTER NICKNAME","ALTER PACKAGE","ALTER PERMISSION","ALTER PROCEDURE","ALTER SCHEMA","ALTER SECURITY LABEL COMPONENT","ALTER SECURITY POLICY","ALTER SEQUENCE","ALTER SERVER","ALTER SERVICE CLASS","ALTER STOGROUP","ALTER TABLESPACE","ALTER THRESHOLD","ALTER TRIGGER","ALTER TRUSTED CONTEXT","ALTER TYPE","ALTER USAGE LIST","ALTER USER MAPPING","ALTER VIEW","ALTER WORK ACTION SET","ALTER WORK CLASS SET","ALTER WORKLOAD","ALTER WRAPPER","ALTER XSROBJECT","ALTER STOGROUP","ALTER TABLESPACE","ALTER TRIGGER","ALTER TRUSTED CONTEXT","ALTER VIEW","ASSOCIATE [RESULT SET] {LOCATOR | LOCATORS}","AUDIT","BEGIN DECLARE SECTION","CALL","CLOSE","COMMENT ON","COMMIT [WORK]","CONNECT","CREATE [OR REPLACE] [PUBLIC] ALIAS","CREATE AUDIT POLICY","CREATE BUFFERPOOL","CREATE DATABASE PARTITION GROUP","CREATE EVENT MONITOR","CREATE [OR REPLACE] FUNCTION","CREATE FUNCTION MAPPING","CREATE HISTOGRAM TEMPLATE","CREATE [UNIQUE] INDEX","CREATE INDEX EXTENSION","CREATE [OR REPLACE] MASK","CREATE [SPECIFIC] METHOD","CREATE [OR REPLACE] MODULE","CREATE [OR REPLACE] NICKNAME","CREATE [OR REPLACE] PERMISSION","CREATE [OR REPLACE] PROCEDURE","CREATE ROLE","CREATE SCHEMA","CREATE SECURITY LABEL [COMPONENT]","CREATE SECURITY POLICY","CREATE [OR REPLACE] SEQUENCE","CREATE SERVICE CLASS","CREATE SERVER","CREATE STOGROUP","CREATE SYNONYM","CREATE [LARGE | REGULAR | {SYSTEM | USER} TEMPORARY] TABLESPACE","CREATE THRESHOLD","CREATE {TRANSFORM | TRANSFORMS} FOR","CREATE [OR REPLACE] TRIGGER","CREATE TRUSTED CONTEXT","CREATE [OR REPLACE] TYPE","CREATE TYPE MAPPING","CREATE USAGE LIST","CREATE USER MAPPING FOR","CREATE [OR REPLACE] VARIABLE","CREATE WORK ACTION SET","CREATE WORK CLASS SET","CREATE WORKLOAD","CREATE WRAPPER","DECLARE","DECLARE GLOBAL TEMPORARY TABLE","DESCRIBE [INPUT | OUTPUT]","DISCONNECT","DROP [PUBLIC] ALIAS","DROP AUDIT POLICY","DROP BUFFERPOOL","DROP DATABASE PARTITION GROUP","DROP EVENT MONITOR","DROP [SPECIFIC] FUNCTION","DROP FUNCTION MAPPING","DROP HISTOGRAM TEMPLATE","DROP INDEX [EXTENSION]","DROP MASK","DROP [SPECIFIC] METHOD","DROP MODULE","DROP NICKNAME","DROP PACKAGE","DROP PERMISSION","DROP [SPECIFIC] PROCEDURE","DROP ROLE","DROP SCHEMA","DROP SECURITY LABEL [COMPONENT]","DROP SECURITY POLICY","DROP SEQUENCE","DROP SERVER","DROP SERVICE CLASS","DROP STOGROUP","DROP TABLE HIERARCHY","DROP {TABLESPACE | TABLESPACES}","DROP {TRANSFORM | TRANSFORMS}","DROP THRESHOLD","DROP TRIGGER","DROP TRUSTED CONTEXT","DROP TYPE [MAPPING]","DROP USAGE LIST","DROP USER MAPPING FOR","DROP VARIABLE","DROP VIEW [HIERARCHY]","DROP WORK {ACTION | CLASS} SET","DROP WORKLOAD","DROP WRAPPER","DROP XSROBJECT","END DECLARE SECTION","EXECUTE [IMMEDIATE]","EXPLAIN {PLAN [SECTION] | ALL}","FETCH [FROM]","FLUSH {BUFFERPOOL | BUFFERPOOLS} ALL","FLUSH EVENT MONITOR","FLUSH FEDERATED CACHE","FLUSH OPTIMIZATION PROFILE CACHE","FLUSH PACKAGE CACHE [DYNAMIC]","FLUSH AUTHENTICATION CACHE [FOR ALL]","FREE LOCATOR","GET DIAGNOSTICS","GOTO","GRANT","INCLUDE","ITERATE","LEAVE","LOCK TABLE","LOOP","OPEN","PIPE","PREPARE","REFRESH TABLE","RELEASE","RELEASE [TO] SAVEPOINT","RENAME [TABLE | INDEX | STOGROUP | TABLESPACE]","REPEAT","RESIGNAL","RETURN","REVOKE","ROLLBACK [WORK] [TO SAVEPOINT]","SAVEPOINT","SET COMPILATION ENVIRONMENT","SET CONNECTION","SET CURRENT","SET ENCRYPTION PASSWORD","SET EVENT MONITOR STATE","SET INTEGRITY","SET PASSTHRU","SET PATH","SET ROLE","SET SCHEMA","SET SERVER OPTION","SET {SESSION AUTHORIZATION | SESSION_USER}","SET USAGE LIST","SIGNAL","TRANSFER OWNERSHIP OF","WHENEVER {NOT FOUND | SQLERROR | SQLWARNING}","WHILE"]),Gs=v(["UNION [ALL]","EXCEPT [ALL]","INTERSECT [ALL]"]),gs=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN"]),Hs=v(["ON DELETE","ON UPDATE","SET NULL","{ROWS | RANGE} BETWEEN"]),bs={name:"db2",tokenizerOptions:{reservedSelect:ms,reservedClauses:[...hs,...rr,...Xt],reservedSetOperations:Gs,reservedJoins:gs,reservedPhrases:Hs,reservedKeywords:Ms,reservedDataTypes:Us,reservedFunctionNames:ps,extraParens:["[]"],stringTypes:[{quote:"''-qq",prefixes:["G","N","U&"]},{quote:"''-raw",prefixes:["X","BX","GX","UX"],requirePrefix:!0}],identTypes:['""-qq'],identChars:{first:"@#$",rest:"@#$"},paramTypes:{positional:!0,named:[":"]},paramChars:{first:"@#$",rest:"@#$"},operators:["**","%","|","&","^","~","¬=","¬>","¬<","!>","!<","^=","^>","^<","||","->","=>"]},formatOptions:{onelineClauses:[...rr,...Xt],tabularOnelineClauses:Xt}},ys=["ARRAY_AGG","AVG","CORR","CORRELATION","COUNT","COUNT_BIG","COVAR_POP","COVARIANCE","COVAR","COVAR_SAMP","COVARIANCE_SAMP","EVERY","GROUPING","JSON_ARRAYAGG","JSON_OBJECTAGG","LISTAGG","MAX","MEDIAN","MIN","PERCENTILE_CONT","PERCENTILE_DISC","REGR_AVGX","REGR_AVGY","REGR_COUNT","REGR_INTERCEPT","REGR_R2","REGR_SLOPE","REGR_SXX","REGR_SXY","REGR_SYY","SOME","STDDEV_POP","STDDEV","STDDEV_SAMP","SUM","VAR_POP","VARIANCE","VAR","VAR_SAMP","VARIANCE_SAMP","XMLAGG","XMLGROUP","ABS","ABSVAL","ACOS","ADD_DAYS","ADD_HOURS","ADD_MINUTES","ADD_MONTHS","ADD_SECONDS","ADD_YEARS","ANTILOG","ARRAY_MAX_CARDINALITY","ARRAY_TRIM","ASCII","ASIN","ATAN","ATAN2","ATANH","BASE64_DECODE","BASE64_ENCODE","BIT_LENGTH","BITAND","BITANDNOT","BITNOT","BITOR","BITXOR","BSON_TO_JSON","CARDINALITY","CEIL","CEILING","CHAR_LENGTH","CHARACTER_LENGTH","CHR","COALESCE","COMPARE_DECFLOAT","CONCAT","CONTAINS","COS","COSH","COT","CURDATE","CURTIME","DATABASE","DATAPARTITIONNAME","DATAPARTITIONNUM","DAY","DAYNAME","DAYOFMONTH","DAYOFWEEK_ISO","DAYOFWEEK","DAYOFYEAR","DAYS","DBPARTITIONNAME","DBPARTITIONNUM","DECFLOAT_FORMAT","DECFLOAT_SORTKEY","DECRYPT_BINARY","DECRYPT_BIT","DECRYPT_CHAR","DECRYPT_DB","DEGREES","DIFFERENCE","DIGITS","DLCOMMENT","DLLINKTYPE","DLURLCOMPLETE","DLURLPATH","DLURLPATHONLY","DLURLSCHEME","DLURLSERVER","DLVALUE","DOUBLE_PRECISION","DOUBLE","ENCRPYT","ENCRYPT_AES","ENCRYPT_AES256","ENCRYPT_RC2","ENCRYPT_TDES","EXP","EXTRACT","FIRST_DAY","FLOOR","GENERATE_UNIQUE","GET_BLOB_FROM_FILE","GET_CLOB_FROM_FILE","GET_DBCLOB_FROM_FILE","GET_XML_FILE","GETHINT","GREATEST","HASH_MD5","HASH_ROW","HASH_SHA1","HASH_SHA256","HASH_SHA512","HASH_VALUES","HASHED_VALUE","HEX","HEXTORAW","HOUR","HTML_ENTITY_DECODE","HTML_ENTITY_ENCODE","HTTP_DELETE_BLOB","HTTP_DELETE","HTTP_GET_BLOB","HTTP_GET","HTTP_PATCH_BLOB","HTTP_PATCH","HTTP_POST_BLOB","HTTP_POST","HTTP_PUT_BLOB","HTTP_PUT","IDENTITY_VAL_LOCAL","IFNULL","INSERT","INSTR","INTERPRET","ISFALSE","ISNOTFALSE","ISNOTTRUE","ISTRUE","JSON_ARRAY","JSON_OBJECT","JSON_QUERY","JSON_TO_BSON","JSON_UPDATE","JSON_VALUE","JULIAN_DAY","LAND","LAST_DAY","LCASE","LEAST","LEFT","LENGTH","LN","LNOT","LOCATE_IN_STRING","LOCATE","LOG10","LOR","LOWER","LPAD","LTRIM","MAX_CARDINALITY","MAX","MICROSECOND","MIDNIGHT_SECONDS","MIN","MINUTE","MOD","MONTH","MONTHNAME","MONTHS_BETWEEN","MQREAD","MQREADCLOB","MQRECEIVE","MQRECEIVECLOB","MQSEND","MULTIPLY_ALT","NEXT_DAY","NORMALIZE_DECFLOAT","NOW","NULLIF","NVL","OCTET_LENGTH","OVERLAY","PI","POSITION","POSSTR","POW","POWER","QUANTIZE","QUARTER","RADIANS","RAISE_ERROR","RANDOM","RAND","REGEXP_COUNT","REGEXP_INSTR","REGEXP_REPLACE","REGEXP_SUBSTR","REPEAT","REPLACE","RID","RIGHT","ROUND_TIMESTAMP","ROUND","RPAD","RRN","RTRIM","SCORE","SECOND","SIGN","SIN","SINH","SOUNDEX","SPACE","SQRT","STRIP","STRLEFT","STRPOS","STRRIGHT","SUBSTR","SUBSTRING","TABLE_NAME","TABLE_SCHEMA","TAN","TANH","TIMESTAMP_FORMAT","TIMESTAMP_ISO","TIMESTAMPDIFF_BIG","TIMESTAMPDIFF","TO_CHAR","TO_CLOB","TO_DATE","TO_NUMBER","TO_TIMESTAMP","TOTALORDER","TRANSLATE","TRIM_ARRAY","TRIM","TRUNC_TIMESTAMP","TRUNC","TRUNCATE","UCASE","UPPER","URL_DECODE","URL_ENCODE","VALUE","VARBINARY_FORMAT","VARCHAR_BIT_FORMAT","VARCHAR_FORMAT_BINARY","VARCHAR_FORMAT","VERIFY_GROUP_FOR_USER","WEEK_ISO","WEEK","WRAP","XMLATTRIBUTES","XMLCOMMENT","XMLCONCAT","XMLDOCUMENT","XMLELEMENT","XMLFOREST","XMLNAMESPACES","XMLPARSE","XMLPI","XMLROW","XMLSERIALIZE","XMLTEXT","XMLVALIDATE","XOR","XSLTRANSFORM","YEAR","ZONED","BASE_TABLE","HTTP_DELETE_BLOB_VERBOSE","HTTP_DELETE_VERBOSE","HTTP_GET_BLOB_VERBOSE","HTTP_GET_VERBOSE","HTTP_PATCH_BLOB_VERBOSE","HTTP_PATCH_VERBOSE","HTTP_POST_BLOB_VERBOSE","HTTP_POST_VERBOSE","HTTP_PUT_BLOB_VERBOSE","HTTP_PUT_VERBOSE","JSON_TABLE","MQREADALL","MQREADALLCLOB","MQRECEIVEALL","MQRECEIVEALLCLOB","XMLTABLE","UNPACK","CUME_DIST","DENSE_RANK","FIRST_VALUE","LAG","LAST_VALUE","LEAD","NTH_VALUE","NTILE","PERCENT_RANK","RANK","RATIO_TO_REPORT","ROW_NUMBER","CAST"],Bs=["ABSENT","ACCORDING","ACCTNG","ACTION","ACTIVATE","ADD","ALIAS","ALL","ALLOCATE","ALLOW","ALTER","AND","ANY","APPEND","APPLNAME","ARRAY","ARRAY_AGG","ARRAY_TRIM","AS","ASC","ASENSITIVE","ASSOCIATE","ATOMIC","ATTACH","ATTRIBUTES","AUTHORIZATION","AUTONOMOUS","BEFORE","BEGIN","BETWEEN","BIND","BSON","BUFFERPOOL","BY","CACHE","CALL","CALLED","CARDINALITY","CASE","CAST","CHECK","CL","CLOSE","CLUSTER","COLLECT","COLLECTION","COLUMN","COMMENT","COMMIT","COMPACT","COMPARISONS","COMPRESS","CONCAT","CONCURRENT","CONDITION","CONNECT","CONNECT_BY_ROOT","CONNECTION","CONSTANT","CONSTRAINT","CONTAINS","CONTENT","CONTINUE","COPY","COUNT","COUNT_BIG","CREATE","CREATEIN","CROSS","CUBE","CUME_DIST","CURRENT","CURRENT_DATE","CURRENT_PATH","CURRENT_SCHEMA","CURRENT_SERVER","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_TIMEZONE","CURRENT_USER","CURSOR","CYCLE","DATABASE","DATAPARTITIONNAME","DATAPARTITIONNUM","DAY","DAYS","DB2GENERAL","DB2GENRL","DB2SQL","DBINFO","DBPARTITIONNAME","DBPARTITIONNUM","DEACTIVATE","DEALLOCATE","DECLARE","DEFAULT","DEFAULTS","DEFER","DEFINE","DEFINITION","DELETE","DELETING","DENSE_RANK","DENSERANK","DESC","DESCRIBE","DESCRIPTOR","DETACH","DETERMINISTIC","DIAGNOSTICS","DISABLE","DISALLOW","DISCONNECT","DISTINCT","DO","DOCUMENT","DROP","DYNAMIC","EACH","ELSE","ELSEIF","EMPTY","ENABLE","ENCODING","ENCRYPTION","END","END-EXEC","ENDING","ENFORCED","ERROR","ESCAPE","EVERY","EXCEPT","EXCEPTION","EXCLUDING","EXCLUSIVE","EXECUTE","EXISTS","EXIT","EXTEND","EXTERNAL","EXTRACT","FALSE","FENCED","FETCH","FIELDPROC","FILE","FINAL","FIRST_VALUE","FOR","FOREIGN","FORMAT","FREE","FREEPAGE","FROM","FULL","FUNCTION","GBPCACHE","GENERAL","GENERATED","GET","GLOBAL","GO","GOTO","GRANT","GROUP","HANDLER","HASH","HASH_ROW","HASHED_VALUE","HAVING","HINT","HOLD","HOUR","HOURS","IDENTITY","IF","IGNORE","IMMEDIATE","IMPLICITLY","IN","INCLUDE","INCLUDING","INCLUSIVE","INCREMENT","INDEX","INDEXBP","INDICATOR","INF","INFINITY","INHERIT","INLINE","INNER","INOUT","INSENSITIVE","INSERT","INSERTING","INTEGRITY","INTERPRET","INTERSECT","INTO","IS","ISNULL","ISOLATION","ITERATE","JAVA","JOIN","JSON","JSON_ARRAY","JSON_ARRAYAGG","JSON_EXISTS","JSON_OBJECT","JSON_OBJECTAGG","JSON_QUERY","JSON_TABLE","JSON_VALUE","KEEP","KEY","KEYS","LABEL","LAG","LANGUAGE","LAST_VALUE","LATERAL","LEAD","LEAVE","LEFT","LEVEL2","LIKE","LIMIT","LINKTYPE","LISTAGG","LOCAL","LOCALDATE","LOCALTIME","LOCALTIMESTAMP","LOCATION","LOCATOR","LOCK","LOCKSIZE","LOG","LOGGED","LOOP","MAINTAINED","MASK","MATCHED","MATERIALIZED","MAXVALUE","MERGE","MICROSECOND","MICROSECONDS","MINPCTUSED","MINUTE","MINUTES","MINVALUE","MIRROR","MIXED","MODE","MODIFIES","MONTH","MONTHS","NAMESPACE","NAN","NATIONAL","NCHAR","NCLOB","NESTED","NEW","NEW_TABLE","NEXTVAL","NO","NOCACHE","NOCYCLE","NODENAME","NODENUMBER","NOMAXVALUE","NOMINVALUE","NONE","NOORDER","NORMALIZED","NOT","NOTNULL","NTH_VALUE","NTILE","NULL","NULLS","NVARCHAR","OBID","OBJECT","OF","OFF","OFFSET","OLD","OLD_TABLE","OMIT","ON","ONLY","OPEN","OPTIMIZE","OPTION","OR","ORDER","ORDINALITY","ORGANIZE","OUT","OUTER","OVER","OVERLAY","OVERRIDING","PACKAGE","PADDED","PAGE","PAGESIZE","PARAMETER","PART","PARTITION","PARTITIONED","PARTITIONING","PARTITIONS","PASSING","PASSWORD","PATH","PCTFREE","PERCENT_RANK","PERCENTILE_CONT","PERCENTILE_DISC","PERIOD","PERMISSION","PIECESIZE","PIPE","PLAN","POSITION","PREPARE","PREVVAL","PRIMARY","PRIOR","PRIQTY","PRIVILEGES","PROCEDURE","PROGRAM","PROGRAMID","QUERY","RANGE","RANK","RATIO_TO_REPORT","RCDFMT","READ","READS","RECOVERY","REFERENCES","REFERENCING","REFRESH","REGEXP_LIKE","RELEASE","RENAME","REPEAT","RESET","RESIGNAL","RESTART","RESULT","RESULT_SET_LOCATOR","RETURN","RETURNING","RETURNS","REVOKE","RID","RIGHT","ROLLBACK","ROLLUP","ROUTINE","ROW","ROW_NUMBER","ROWNUMBER","ROWS","RRN","RUN","SAVEPOINT","SBCS","SCALAR","SCHEMA","SCRATCHPAD","SCROLL","SEARCH","SECOND","SECONDS","SECQTY","SECURED","SELECT","SENSITIVE","SEQUENCE","SESSION","SESSION_USER","SET","SIGNAL","SIMPLE","SKIP","SNAN","SOME","SOURCE","SPECIFIC","SQL","SQLID","SQLIND_DEFAULT","SQLIND_UNASSIGNED","STACKED","START","STARTING","STATEMENT","STATIC","STOGROUP","SUBSTRING","SUMMARY","SYNONYM","SYSTEM_TIME","SYSTEM_USER","TABLE","TABLESPACE","TABLESPACES","TAG","THEN","THREADSAFE","TO","TRANSACTION","TRANSFER","TRIGGER","TRIM","TRIM_ARRAY","TRUE","TRUNCATE","TRY_CAST","TYPE","UNDO","UNION","UNIQUE","UNIT","UNKNOWN","UNNEST","UNTIL","UPDATE","UPDATING","URI","USAGE","USE","USER","USERID","USING","VALUE","VALUES","VARIABLE","VARIANT","VCAT","VERSION","VERSIONING","VIEW","VOLATILE","WAIT","WHEN","WHENEVER","WHERE","WHILE","WITH","WITHIN","WITHOUT","WRAPPED","WRAPPER","WRITE","WRKSTNNAME","XMLAGG","XMLATTRIBUTES","XMLCAST","XMLCOMMENT","XMLCONCAT","XMLDOCUMENT","XMLELEMENT","XMLFOREST","XMLGROUP","XMLNAMESPACES","XMLPARSE","XMLPI","XMLROW","XMLSERIALIZE","XMLTABLE","XMLTEXT","XMLVALIDATE","XSLTRANSFORM","XSROBJECT","YEAR","YEARS","YES","ZONE"],vs=["ARRAY","BIGINT","BINARY","BIT","BLOB","BOOLEAN","CCSID","CHAR","CHARACTER","CLOB","DATA","DATALINK","DATE","DBCLOB","DECFLOAT","DECIMAL","DEC","DOUBLE","DOUBLE PRECISION","FLOAT","GRAPHIC","INT","INTEGER","LONG","NUMERIC","REAL","ROWID","SMALLINT","TIME","TIMESTAMP","VARBINARY","VARCHAR","VARGRAPHIC","XML"],Fs=v(["SELECT [ALL | DISTINCT]"]),Ys=v(["WITH [RECURSIVE]","INTO","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER [SIBLINGS] BY [INPUT SEQUENCE]","LIMIT","OFFSET","FETCH {FIRST | NEXT}","FOR UPDATE [OF]","FOR READ ONLY","OPTIMIZE FOR","INSERT INTO","VALUES","SET","MERGE INTO","WHEN [NOT] MATCHED [THEN]","UPDATE SET","DELETE","INSERT","FOR SYSTEM NAME"]),Rr=v(["CREATE [OR REPLACE] TABLE"]),kt=v(["CREATE [OR REPLACE] [RECURSIVE] VIEW","UPDATE","WHERE CURRENT OF","WITH {NC | RR | RS | CS | UR}","DELETE FROM","DROP TABLE","ALTER TABLE","ADD [COLUMN]","ALTER [COLUMN]","DROP [COLUMN]","SET DATA TYPE","SET {GENERATED ALWAYS | GENERATED BY DEFAULT}","SET NOT NULL","SET {NOT HIDDEN | IMPLICITLY HIDDEN}","SET FIELDPROC","DROP {DEFAULT | NOT NULL | GENERATED | IDENTITY | ROW CHANGE TIMESTAMP | FIELDPROC}","TRUNCATE [TABLE]","SET [CURRENT] SCHEMA","SET CURRENT_SCHEMA","ALLOCATE CURSOR","ALLOCATE [SQL] DESCRIPTOR [LOCAL | GLOBAL] SQL","ALTER [SPECIFIC] {FUNCTION | PROCEDURE}","ALTER {MASK | PERMISSION | SEQUENCE | TRIGGER}","ASSOCIATE [RESULT SET] {LOCATOR | LOCATORS}","BEGIN DECLARE SECTION","CALL","CLOSE","COMMENT ON {ALIAS | COLUMN | CONSTRAINT | INDEX | MASK | PACKAGE | PARAMETER | PERMISSION | SEQUENCE | TABLE | TRIGGER | VARIABLE | XSROBJECT}","COMMENT ON [SPECIFIC] {FUNCTION | PROCEDURE | ROUTINE}","COMMENT ON PARAMETER SPECIFIC {FUNCTION | PROCEDURE | ROUTINE}","COMMENT ON [TABLE FUNCTION] RETURN COLUMN","COMMENT ON [TABLE FUNCTION] RETURN COLUMN SPECIFIC [PROCEDURE | ROUTINE]","COMMIT [WORK] [HOLD]","CONNECT [TO | RESET] USER","CREATE [OR REPLACE] {ALIAS | FUNCTION | MASK | PERMISSION | PROCEDURE | SEQUENCE | TRIGGER | VARIABLE}","CREATE [ENCODED VECTOR] INDEX","CREATE UNIQUE [WHERE NOT NULL] INDEX","CREATE SCHEMA","CREATE TYPE","DEALLOCATE [SQL] DESCRIPTOR [LOCAL | GLOBAL]","DECLARE CURSOR","DECLARE GLOBAL TEMPORARY TABLE","DECLARE","DESCRIBE CURSOR","DESCRIBE INPUT","DESCRIBE [OUTPUT]","DESCRIBE {PROCEDURE | ROUTINE}","DESCRIBE TABLE","DISCONNECT ALL [SQL]","DISCONNECT [CURRENT]","DROP {ALIAS | INDEX | MASK | PACKAGE | PERMISSION | SCHEMA | SEQUENCE | TABLE | TYPE | VARIABLE | XSROBJECT} [IF EXISTS]","DROP [SPECIFIC] {FUNCTION | PROCEDURE | ROUTINE} [IF EXISTS]","END DECLARE SECTION","EXECUTE [IMMEDIATE]","FREE LOCATOR","GET [SQL] DESCRIPTOR [LOCAL | GLOBAL]","GET [CURRENT | STACKED] DIAGNOSTICS","GRANT {ALL [PRIVILEGES] | ALTER | EXECUTE} ON {FUNCTION | PROCEDURE | ROUTINE | PACKAGE | SCHEMA | SEQUENCE | TABLE | TYPE | VARIABLE | XSROBJECT}","HOLD LOCATOR","INCLUDE","LABEL ON {ALIAS | COLUMN | CONSTRAINT | INDEX | MASK | PACKAGE | PERMISSION | SEQUENCE | TABLE | TRIGGER | VARIABLE | XSROBJECT}","LABEL ON [SPECIFIC] {FUNCTION | PROCEDURE | ROUTINE}","LOCK TABLE","OPEN","PREPARE","REFRESH TABLE","RELEASE","RELEASE [TO] SAVEPOINT","RENAME [TABLE | INDEX] TO","REVOKE {ALL [PRIVILEGES] | ALTER | EXECUTE} ON {FUNCTION | PROCEDURE | ROUTINE | PACKAGE | SCHEMA | SEQUENCE | TABLE | TYPE | VARIABLE | XSROBJECT}","ROLLBACK [WORK] [HOLD | TO SAVEPOINT]","SAVEPOINT","SET CONNECTION","SET CURRENT {DEBUG MODE | DECFLOAT ROUNDING MODE | DEGREE | IMPLICIT XMLPARSE OPTION | TEMPORAL SYSTEM_TIME}","SET [SQL] DESCRIPTOR [LOCAL | GLOBAL]","SET ENCRYPTION PASSWORD","SET OPTION","SET {[CURRENT [FUNCTION]] PATH | CURRENT_PATH}","SET RESULT SETS [WITH RETURN [TO CALLER | TO CLIENT]]","SET SESSION AUTHORIZATION","SET SESSION_USER","SET TRANSACTION","SIGNAL SQLSTATE [VALUE]","TAG","TRANSFER OWNERSHIP OF","WHENEVER {NOT FOUND | SQLERROR | SQLWARNING}"]),Vs=v(["UNION [ALL]","EXCEPT [ALL]","INTERSECT [ALL]"]),Ws=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","[LEFT | RIGHT] EXCEPTION JOIN","{INNER | CROSS} JOIN"]),ws=v(["ON DELETE","ON UPDATE","SET NULL","{ROWS | RANGE} BETWEEN"]),$s={name:"db2i",tokenizerOptions:{reservedSelect:Fs,reservedClauses:[...Ys,...Rr,...kt],reservedSetOperations:Vs,reservedJoins:Ws,reservedPhrases:ws,reservedKeywords:Bs,reservedDataTypes:vs,reservedFunctionNames:ys,nestedBlockComments:!0,extraParens:["[]"],stringTypes:[{quote:"''-qq",prefixes:["G","N"]},{quote:"''-raw",prefixes:["X","BX","GX","UX"],requirePrefix:!0}],identTypes:['""-qq'],identChars:{first:"@#$",rest:"@#$"},paramTypes:{positional:!0,named:[":"]},paramChars:{first:"@#$",rest:"@#$"},operators:["**","¬=","¬>","¬<","!>","!<","||","=>"]},formatOptions:{onelineClauses:[...Rr,...kt],tabularOnelineClauses:kt}},xs=["ABS","ACOS","ASIN","ATAN","BIN","BROUND","CBRT","CEIL","CEILING","CONV","COS","DEGREES","EXP","FACTORIAL","FLOOR","GREATEST","HEX","LEAST","LN","LOG","LOG10","LOG2","NEGATIVE","PI","PMOD","POSITIVE","POW","POWER","RADIANS","RAND","ROUND","SHIFTLEFT","SHIFTRIGHT","SHIFTRIGHTUNSIGNED","SIGN","SIN","SQRT","TAN","UNHEX","WIDTH_BUCKET","ARRAY_CONTAINS","MAP_KEYS","MAP_VALUES","SIZE","SORT_ARRAY","BINARY","CAST","ADD_MONTHS","DATE","DATE_ADD","DATE_FORMAT","DATE_SUB","DATEDIFF","DAY","DAYNAME","DAYOFMONTH","DAYOFYEAR","EXTRACT","FROM_UNIXTIME","FROM_UTC_TIMESTAMP","HOUR","LAST_DAY","MINUTE","MONTH","MONTHS_BETWEEN","NEXT_DAY","QUARTER","SECOND","TIMESTAMP","TO_DATE","TO_UTC_TIMESTAMP","TRUNC","UNIX_TIMESTAMP","WEEKOFYEAR","YEAR","ASSERT_TRUE","COALESCE","IF","ISNOTNULL","ISNULL","NULLIF","NVL","ASCII","BASE64","CHARACTER_LENGTH","CHR","CONCAT","CONCAT_WS","CONTEXT_NGRAMS","DECODE","ELT","ENCODE","FIELD","FIND_IN_SET","FORMAT_NUMBER","GET_JSON_OBJECT","IN_FILE","INITCAP","INSTR","LCASE","LENGTH","LEVENSHTEIN","LOCATE","LOWER","LPAD","LTRIM","NGRAMS","OCTET_LENGTH","PARSE_URL","PRINTF","QUOTE","REGEXP_EXTRACT","REGEXP_REPLACE","REPEAT","REVERSE","RPAD","RTRIM","SENTENCES","SOUNDEX","SPACE","SPLIT","STR_TO_MAP","SUBSTR","SUBSTRING","TRANSLATE","TRIM","UCASE","UNBASE64","UPPER","MASK","MASK_FIRST_N","MASK_HASH","MASK_LAST_N","MASK_SHOW_FIRST_N","MASK_SHOW_LAST_N","AES_DECRYPT","AES_ENCRYPT","CRC32","CURRENT_DATABASE","CURRENT_USER","HASH","JAVA_METHOD","LOGGED_IN_USER","MD5","REFLECT","SHA","SHA1","SHA2","SURROGATE_KEY","VERSION","AVG","COLLECT_LIST","COLLECT_SET","CORR","COUNT","COVAR_POP","COVAR_SAMP","HISTOGRAM_NUMERIC","MAX","MIN","NTILE","PERCENTILE","PERCENTILE_APPROX","REGR_AVGX","REGR_AVGY","REGR_COUNT","REGR_INTERCEPT","REGR_R2","REGR_SLOPE","REGR_SXX","REGR_SXY","REGR_SYY","STDDEV_POP","STDDEV_SAMP","SUM","VAR_POP","VAR_SAMP","VARIANCE","EXPLODE","INLINE","JSON_TUPLE","PARSE_URL_TUPLE","POSEXPLODE","STACK","LEAD","LAG","FIRST_VALUE","LAST_VALUE","RANK","ROW_NUMBER","DENSE_RANK","CUME_DIST","PERCENT_RANK","NTILE"],Xs=["ADD","ADMIN","AFTER","ANALYZE","ARCHIVE","ASC","BEFORE","BUCKET","BUCKETS","CASCADE","CHANGE","CLUSTER","CLUSTERED","CLUSTERSTATUS","COLLECTION","COLUMNS","COMMENT","COMPACT","COMPACTIONS","COMPUTE","CONCATENATE","CONTINUE","DATA","DATABASES","DATETIME","DAY","DBPROPERTIES","DEFERRED","DEFINED","DELIMITED","DEPENDENCY","DESC","DIRECTORIES","DIRECTORY","DISABLE","DISTRIBUTE","ELEM_TYPE","ENABLE","ESCAPED","EXCLUSIVE","EXPLAIN","EXPORT","FIELDS","FILE","FILEFORMAT","FIRST","FORMAT","FORMATTED","FUNCTIONS","HOLD_DDLTIME","HOUR","IDXPROPERTIES","IGNORE","INDEX","INDEXES","INPATH","INPUTDRIVER","INPUTFORMAT","ITEMS","JAR","KEYS","KEY_TYPE","LIMIT","LINES","LOAD","LOCATION","LOCK","LOCKS","LOGICAL","LONG","MAPJOIN","MATERIALIZED","METADATA","MINUS","MINUTE","MONTH","MSCK","NOSCAN","NO_DROP","OFFLINE","OPTION","OUTPUTDRIVER","OUTPUTFORMAT","OVERWRITE","OWNER","PARTITIONED","PARTITIONS","PLUS","PRETTY","PRINCIPALS","PROTECTION","PURGE","READ","READONLY","REBUILD","RECORDREADER","RECORDWRITER","RELOAD","RENAME","REPAIR","REPLACE","REPLICATION","RESTRICT","REWRITE","ROLE","ROLES","SCHEMA","SCHEMAS","SECOND","SEMI","SERDE","SERDEPROPERTIES","SERVER","SETS","SHARED","SHOW","SHOW_DATABASE","SKEWED","SORT","SORTED","SSL","STATISTICS","STORED","STREAMTABLE","STRING","TABLES","TBLPROPERTIES","TEMPORARY","TERMINATED","TINYINT","TOUCH","TRANSACTIONS","UNARCHIVE","UNDO","UNIONTYPE","UNLOCK","UNSET","UNSIGNED","URI","USE","UTC","UTCTIMESTAMP","VALUE_TYPE","VIEW","WHILE","YEAR","AUTOCOMMIT","ISOLATION","LEVEL","OFFSET","SNAPSHOT","TRANSACTION","WORK","WRITE","ABORT","KEY","LAST","NORELY","NOVALIDATE","NULLS","RELY","VALIDATE","DETAIL","DOW","EXPRESSION","OPERATOR","QUARTER","SUMMARY","VECTORIZATION","WEEK","YEARS","MONTHS","WEEKS","DAYS","HOURS","MINUTES","SECONDS","TIMESTAMPTZ","ZONE","ALL","ALTER","AND","AS","AUTHORIZATION","BETWEEN","BOTH","BY","CASE","CAST","COLUMN","CONF","CREATE","CROSS","CUBE","CURRENT","CURRENT_DATE","CURRENT_TIMESTAMP","CURSOR","DATABASE","DELETE","DESCRIBE","DISTINCT","DROP","ELSE","END","EXCHANGE","EXISTS","EXTENDED","EXTERNAL","FALSE","FETCH","FOLLOWING","FOR","FROM","FULL","FUNCTION","GRANT","GROUP","GROUPING","HAVING","IF","IMPORT","IN","INNER","INSERT","INTERSECT","INTO","IS","JOIN","LATERAL","LEFT","LESS","LIKE","LOCAL","MACRO","MORE","NONE","NOT","NULL","OF","ON","OR","ORDER","OUT","OUTER","OVER","PARTIALSCAN","PARTITION","PERCENT","PRECEDING","PRESERVE","PROCEDURE","RANGE","READS","REDUCE","REVOKE","RIGHT","ROLLUP","ROW","ROWS","SELECT","SET","TABLE","TABLESAMPLE","THEN","TO","TRANSFORM","TRIGGER","TRUE","TRUNCATE","UNBOUNDED","UNION","UNIQUEJOIN","UPDATE","USER","USING","UTC_TMESTAMP","VALUES","WHEN","WHERE","WINDOW","WITH","COMMIT","ONLY","REGEXP","RLIKE","ROLLBACK","START","CACHE","CONSTRAINT","FOREIGN","PRIMARY","REFERENCES","DAYOFWEEK","EXTRACT","FLOOR","VIEWS","TIME","SYNC","TEXTFILE","SEQUENCEFILE","ORC","CSV","TSV","PARQUET","AVRO","RCFILE","JSONFILE","INPUTFORMAT","OUTPUTFORMAT"],ks=["ARRAY","BIGINT","BINARY","BOOLEAN","CHAR","DATE","DECIMAL","DOUBLE","FLOAT","INT","INTEGER","INTERVAL","MAP","NUMERIC","PRECISION","SMALLINT","STRUCT","TIMESTAMP","VARCHAR"],Ks=v(["SELECT [ALL | DISTINCT]"]),Js=v(["WITH","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","SORT BY","CLUSTER BY","DISTRIBUTE BY","LIMIT","INSERT INTO [TABLE]","VALUES","SET","MERGE INTO","WHEN [NOT] MATCHED [THEN]","UPDATE SET","INSERT [VALUES]","INSERT OVERWRITE [LOCAL] DIRECTORY","LOAD DATA [LOCAL] INPATH","[OVERWRITE] INTO TABLE"]),nr=v(["CREATE [TEMPORARY] [EXTERNAL] TABLE [IF NOT EXISTS]"]),Kt=v(["CREATE [MATERIALIZED] VIEW [IF NOT EXISTS]","UPDATE","DELETE FROM","DROP TABLE [IF EXISTS]","ALTER TABLE","RENAME TO","TRUNCATE [TABLE]","ALTER","CREATE","USE","DESCRIBE","DROP","FETCH","SHOW","STORED AS","STORED BY","ROW FORMAT"]),qs=v(["UNION [ALL | DISTINCT]"]),Qs=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","LEFT SEMI JOIN"]),Zs=v(["{ROWS | RANGE} BETWEEN"]),js={name:"hive",tokenizerOptions:{reservedSelect:Ks,reservedClauses:[...Js,...nr,...Kt],reservedSetOperations:qs,reservedJoins:Qs,reservedPhrases:Zs,reservedKeywords:Xs,reservedDataTypes:ks,reservedFunctionNames:xs,extraParens:["[]"],stringTypes:['""-bs',"''-bs"],identTypes:["``"],variableTypes:[{quote:"{}",prefixes:["$"],requirePrefix:!0}],operators:["%","~","^","|","&","<=>","==","!","||"]},formatOptions:{onelineClauses:[...nr,...Kt],tabularOnelineClauses:Kt}};function yt(E){return E.map((e,T)=>{const t=E[T+1]||tt;if(FE.SET(e)&&t.text==="(")return AE(EE({},e),{type:"RESERVED_FUNCTION_NAME"});const r=E[T-1]||tt;return FE.VALUES(e)&&r.text==="="?AE(EE({},e),{type:"RESERVED_FUNCTION_NAME"}):e})}var zs=["ACCESSIBLE","ADD","ALL","ALTER","ANALYZE","AND","AS","ASC","ASENSITIVE","BEFORE","BETWEEN","BOTH","BY","CALL","CASCADE","CASE","CHANGE","CHECK","COLLATE","COLUMN","CONDITION","CONSTRAINT","CONTINUE","CONVERT","CREATE","CROSS","CURRENT_DATE","CURRENT_ROLE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURSOR","DATABASE","DATABASES","DAY_HOUR","DAY_MICROSECOND","DAY_MINUTE","DAY_SECOND","DECLARE","DEFAULT","DELAYED","DELETE","DELETE_DOMAIN_ID","DESC","DESCRIBE","DETERMINISTIC","DISTINCT","DISTINCTROW","DIV","DO_DOMAIN_IDS","DROP","DUAL","EACH","ELSE","ELSEIF","ENCLOSED","ESCAPED","EXCEPT","EXISTS","EXIT","EXPLAIN","FALSE","FETCH","FOR","FORCE","FOREIGN","FROM","FULLTEXT","GENERAL","GRANT","GROUP","HAVING","HIGH_PRIORITY","HOUR_MICROSECOND","HOUR_MINUTE","HOUR_SECOND","IF","IGNORE","IGNORE_DOMAIN_IDS","IGNORE_SERVER_IDS","IN","INDEX","INFILE","INNER","INOUT","INSENSITIVE","INSERT","INTERSECT","INTERVAL","INTO","IS","ITERATE","JOIN","KEY","KEYS","KILL","LEADING","LEAVE","LEFT","LIKE","LIMIT","LINEAR","LINES","LOAD","LOCALTIME","LOCALTIMESTAMP","LOCK","LOOP","LOW_PRIORITY","MASTER_HEARTBEAT_PERIOD","MASTER_SSL_VERIFY_SERVER_CERT","MATCH","MAXVALUE","MINUTE_MICROSECOND","MINUTE_SECOND","MOD","MODIFIES","NATURAL","NOT","NO_WRITE_TO_BINLOG","NULL","OFFSET","ON","OPTIMIZE","OPTION","OPTIONALLY","OR","ORDER","OUT","OUTER","OUTFILE","OVER","PAGE_CHECKSUM","PARSE_VCOL_EXPR","PARTITION","POSITION","PRIMARY","PROCEDURE","PURGE","RANGE","READ","READS","READ_WRITE","RECURSIVE","REF_SYSTEM_ID","REFERENCES","REGEXP","RELEASE","RENAME","REPEAT","REPLACE","REQUIRE","RESIGNAL","RESTRICT","RETURN","RETURNING","REVOKE","RIGHT","RLIKE","ROW_NUMBER","ROWS","SCHEMA","SCHEMAS","SECOND_MICROSECOND","SELECT","SENSITIVE","SEPARATOR","SET","SHOW","SIGNAL","SLOW","SPATIAL","SPECIFIC","SQL","SQLEXCEPTION","SQLSTATE","SQLWARNING","SQL_BIG_RESULT","SQL_CALC_FOUND_ROWS","SQL_SMALL_RESULT","SSL","STARTING","STATS_AUTO_RECALC","STATS_PERSISTENT","STATS_SAMPLE_PAGES","STRAIGHT_JOIN","TABLE","TERMINATED","THEN","TO","TRAILING","TRIGGER","TRUE","UNDO","UNION","UNIQUE","UNLOCK","UNSIGNED","UPDATE","USAGE","USE","USING","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","VALUES","WHEN","WHERE","WHILE","WINDOW","WITH","WRITE","XOR","YEAR_MONTH","ZEROFILL"],eS=["BIGINT","BINARY","BIT","BLOB","CHAR BYTE","CHAR","CHARACTER","DATETIME","DEC","DECIMAL","DOUBLE PRECISION","DOUBLE","ENUM","FIXED","FLOAT","FLOAT4","FLOAT8","INT","INT1","INT2","INT3","INT4","INT8","INTEGER","LONG","LONGBLOB","LONGTEXT","MEDIUMBLOB","MEDIUMINT","MEDIUMTEXT","MIDDLEINT","NATIONAL CHAR","NATIONAL VARCHAR","NUMERIC","PRECISION","REAL","SMALLINT","TEXT","TIMESTAMP","TINYBLOB","TINYINT","TINYTEXT","VARBINARY","VARCHAR","VARCHARACTER","VARYING","YEAR"],ES=["ADDDATE","ADD_MONTHS","BIT_AND","BIT_OR","BIT_XOR","CAST","COUNT","CUME_DIST","CURDATE","CURTIME","DATE_ADD","DATE_SUB","DATE_FORMAT","DECODE","DENSE_RANK","EXTRACT","FIRST_VALUE","GROUP_CONCAT","JSON_ARRAYAGG","JSON_OBJECTAGG","LAG","LEAD","MAX","MEDIAN","MID","MIN","NOW","NTH_VALUE","NTILE","POSITION","PERCENT_RANK","PERCENTILE_CONT","PERCENTILE_DISC","RANK","ROW_NUMBER","SESSION_USER","STD","STDDEV","STDDEV_POP","STDDEV_SAMP","SUBDATE","SUBSTR","SUBSTRING","SUM","SYSTEM_USER","TRIM","TRIM_ORACLE","VARIANCE","VAR_POP","VAR_SAMP","ABS","ACOS","ADDTIME","AES_DECRYPT","AES_ENCRYPT","ASIN","ATAN","ATAN2","BENCHMARK","BIN","BINLOG_GTID_POS","BIT_COUNT","BIT_LENGTH","CEIL","CEILING","CHARACTER_LENGTH","CHAR_LENGTH","CHR","COERCIBILITY","COLUMN_CHECK","COLUMN_EXISTS","COLUMN_LIST","COLUMN_JSON","COMPRESS","CONCAT","CONCAT_OPERATOR_ORACLE","CONCAT_WS","CONNECTION_ID","CONV","CONVERT_TZ","COS","COT","CRC32","DATEDIFF","DAYNAME","DAYOFMONTH","DAYOFWEEK","DAYOFYEAR","DEGREES","DECODE_HISTOGRAM","DECODE_ORACLE","DES_DECRYPT","DES_ENCRYPT","ELT","ENCODE","ENCRYPT","EXP","EXPORT_SET","EXTRACTVALUE","FIELD","FIND_IN_SET","FLOOR","FORMAT","FOUND_ROWS","FROM_BASE64","FROM_DAYS","FROM_UNIXTIME","GET_LOCK","GREATEST","HEX","IFNULL","INSTR","ISNULL","IS_FREE_LOCK","IS_USED_LOCK","JSON_ARRAY","JSON_ARRAY_APPEND","JSON_ARRAY_INSERT","JSON_COMPACT","JSON_CONTAINS","JSON_CONTAINS_PATH","JSON_DEPTH","JSON_DETAILED","JSON_EXISTS","JSON_EXTRACT","JSON_INSERT","JSON_KEYS","JSON_LENGTH","JSON_LOOSE","JSON_MERGE","JSON_MERGE_PATCH","JSON_MERGE_PRESERVE","JSON_QUERY","JSON_QUOTE","JSON_OBJECT","JSON_REMOVE","JSON_REPLACE","JSON_SET","JSON_SEARCH","JSON_TYPE","JSON_UNQUOTE","JSON_VALID","JSON_VALUE","LAST_DAY","LAST_INSERT_ID","LCASE","LEAST","LENGTH","LENGTHB","LN","LOAD_FILE","LOCATE","LOG","LOG10","LOG2","LOWER","LPAD","LPAD_ORACLE","LTRIM","LTRIM_ORACLE","MAKEDATE","MAKETIME","MAKE_SET","MASTER_GTID_WAIT","MASTER_POS_WAIT","MD5","MONTHNAME","NAME_CONST","NVL","NVL2","OCT","OCTET_LENGTH","ORD","PERIOD_ADD","PERIOD_DIFF","PI","POW","POWER","QUOTE","REGEXP_INSTR","REGEXP_REPLACE","REGEXP_SUBSTR","RADIANS","RAND","RELEASE_ALL_LOCKS","RELEASE_LOCK","REPLACE_ORACLE","REVERSE","ROUND","RPAD","RPAD_ORACLE","RTRIM","RTRIM_ORACLE","SEC_TO_TIME","SHA","SHA1","SHA2","SIGN","SIN","SLEEP","SOUNDEX","SPACE","SQRT","STRCMP","STR_TO_DATE","SUBSTR_ORACLE","SUBSTRING_INDEX","SUBTIME","SYS_GUID","TAN","TIMEDIFF","TIME_FORMAT","TIME_TO_SEC","TO_BASE64","TO_CHAR","TO_DAYS","TO_SECONDS","UCASE","UNCOMPRESS","UNCOMPRESSED_LENGTH","UNHEX","UNIX_TIMESTAMP","UPDATEXML","UPPER","UUID","UUID_SHORT","VERSION","WEEKDAY","WEEKOFYEAR","WSREP_LAST_WRITTEN_GTID","WSREP_LAST_SEEN_GTID","WSREP_SYNC_WAIT_UPTO_GTID","YEARWEEK","COALESCE","NULLIF"],tS=v(["SELECT [ALL | DISTINCT | DISTINCTROW]"]),TS=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER BY","LIMIT","OFFSET","FETCH {FIRST | NEXT}","INSERT [LOW_PRIORITY | DELAYED | HIGH_PRIORITY] [IGNORE] [INTO]","REPLACE [LOW_PRIORITY | DELAYED] [INTO]","VALUES","ON DUPLICATE KEY UPDATE","SET","RETURNING"]),Ar=v(["CREATE [OR REPLACE] [TEMPORARY] TABLE [IF NOT EXISTS]"]),Jt=v(["CREATE [OR REPLACE] [SQL SECURITY DEFINER | SQL SECURITY INVOKER] VIEW [IF NOT EXISTS]","UPDATE [LOW_PRIORITY] [IGNORE]","DELETE [LOW_PRIORITY] [QUICK] [IGNORE] FROM","DROP [TEMPORARY] TABLE [IF EXISTS]","ALTER [ONLINE] [IGNORE] TABLE [IF EXISTS]","ADD [COLUMN] [IF NOT EXISTS]","{CHANGE | MODIFY} [COLUMN] [IF EXISTS]","DROP [COLUMN] [IF EXISTS]","RENAME [TO]","RENAME COLUMN","ALTER [COLUMN]","{SET | DROP} DEFAULT","SET {VISIBLE | INVISIBLE}","TRUNCATE [TABLE]","ALTER DATABASE","ALTER DATABASE COMMENT","ALTER EVENT","ALTER FUNCTION","ALTER PROCEDURE","ALTER SCHEMA","ALTER SCHEMA COMMENT","ALTER SEQUENCE","ALTER SERVER","ALTER USER","ALTER VIEW","ANALYZE","ANALYZE TABLE","BACKUP LOCK","BACKUP STAGE","BACKUP UNLOCK","BEGIN","BINLOG","CACHE INDEX","CALL","CHANGE MASTER TO","CHECK TABLE","CHECK VIEW","CHECKSUM TABLE","COMMIT","CREATE AGGREGATE FUNCTION","CREATE DATABASE","CREATE EVENT","CREATE FUNCTION","CREATE INDEX","CREATE PROCEDURE","CREATE ROLE","CREATE SEQUENCE","CREATE SERVER","CREATE SPATIAL INDEX","CREATE TRIGGER","CREATE UNIQUE INDEX","CREATE USER","DEALLOCATE PREPARE","DESCRIBE","DROP DATABASE","DROP EVENT","DROP FUNCTION","DROP INDEX","DROP PREPARE","DROP PROCEDURE","DROP ROLE","DROP SEQUENCE","DROP SERVER","DROP TRIGGER","DROP USER","DROP VIEW","EXECUTE","EXPLAIN","FLUSH","GET DIAGNOSTICS","GET DIAGNOSTICS CONDITION","GRANT","HANDLER","HELP","INSTALL PLUGIN","INSTALL SONAME","KILL","LOAD DATA INFILE","LOAD INDEX INTO CACHE","LOAD XML INFILE","LOCK TABLE","OPTIMIZE TABLE","PREPARE","PURGE BINARY LOGS","PURGE MASTER LOGS","RELEASE SAVEPOINT","RENAME TABLE","RENAME USER","REPAIR TABLE","REPAIR VIEW","RESET MASTER","RESET QUERY CACHE","RESET REPLICA","RESET SLAVE","RESIGNAL","REVOKE","ROLLBACK","SAVEPOINT","SET CHARACTER SET","SET DEFAULT ROLE","SET GLOBAL TRANSACTION","SET NAMES","SET PASSWORD","SET ROLE","SET STATEMENT","SET TRANSACTION","SHOW","SHOW ALL REPLICAS STATUS","SHOW ALL SLAVES STATUS","SHOW AUTHORS","SHOW BINARY LOGS","SHOW BINLOG EVENTS","SHOW BINLOG STATUS","SHOW CHARACTER SET","SHOW CLIENT_STATISTICS","SHOW COLLATION","SHOW COLUMNS","SHOW CONTRIBUTORS","SHOW CREATE DATABASE","SHOW CREATE EVENT","SHOW CREATE FUNCTION","SHOW CREATE PACKAGE","SHOW CREATE PACKAGE BODY","SHOW CREATE PROCEDURE","SHOW CREATE SEQUENCE","SHOW CREATE TABLE","SHOW CREATE TRIGGER","SHOW CREATE USER","SHOW CREATE VIEW","SHOW DATABASES","SHOW ENGINE","SHOW ENGINE INNODB STATUS","SHOW ENGINES","SHOW ERRORS","SHOW EVENTS","SHOW EXPLAIN","SHOW FUNCTION CODE","SHOW FUNCTION STATUS","SHOW GRANTS","SHOW INDEX","SHOW INDEXES","SHOW INDEX_STATISTICS","SHOW KEYS","SHOW LOCALES","SHOW MASTER LOGS","SHOW MASTER STATUS","SHOW OPEN TABLES","SHOW PACKAGE BODY CODE","SHOW PACKAGE BODY STATUS","SHOW PACKAGE STATUS","SHOW PLUGINS","SHOW PLUGINS SONAME","SHOW PRIVILEGES","SHOW PROCEDURE CODE","SHOW PROCEDURE STATUS","SHOW PROCESSLIST","SHOW PROFILE","SHOW PROFILES","SHOW QUERY_RESPONSE_TIME","SHOW RELAYLOG EVENTS","SHOW REPLICA","SHOW REPLICA HOSTS","SHOW REPLICA STATUS","SHOW SCHEMAS","SHOW SLAVE","SHOW SLAVE HOSTS","SHOW SLAVE STATUS","SHOW STATUS","SHOW STORAGE ENGINES","SHOW TABLE STATUS","SHOW TABLES","SHOW TRIGGERS","SHOW USER_STATISTICS","SHOW VARIABLES","SHOW WARNINGS","SHOW WSREP_MEMBERSHIP","SHOW WSREP_STATUS","SHUTDOWN","SIGNAL","START ALL REPLICAS","START ALL SLAVES","START REPLICA","START SLAVE","START TRANSACTION","STOP ALL REPLICAS","STOP ALL SLAVES","STOP REPLICA","STOP SLAVE","UNINSTALL PLUGIN","UNINSTALL SONAME","UNLOCK TABLE","USE","XA BEGIN","XA COMMIT","XA END","XA PREPARE","XA RECOVER","XA ROLLBACK","XA START"]),rS=v(["UNION [ALL | DISTINCT]","EXCEPT [ALL | DISTINCT]","INTERSECT [ALL | DISTINCT]","MINUS [ALL | DISTINCT]"]),RS=v(["JOIN","{LEFT | RIGHT} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL JOIN","NATURAL {LEFT | RIGHT} [OUTER] JOIN","STRAIGHT_JOIN"]),nS=v(["ON {UPDATE | DELETE} [SET NULL | SET DEFAULT]","CHARACTER SET","{ROWS | RANGE} BETWEEN","IDENTIFIED BY"]),AS={name:"mariadb",tokenizerOptions:{reservedSelect:tS,reservedClauses:[...TS,...Ar,...Jt],reservedSetOperations:rS,reservedJoins:RS,reservedPhrases:nS,supportsXor:!0,reservedKeywords:zs,reservedDataTypes:eS,reservedFunctionNames:ES,stringTypes:['""-qq-bs',"''-qq-bs",{quote:"''-raw",prefixes:["B","X"],requirePrefix:!0}],identTypes:["``"],identChars:{first:"$",rest:"$",allowFirstCharNumber:!0},variableTypes:[{regex:"@@?[A-Za-z0-9_.$]+"},{quote:'""-qq-bs',prefixes:["@"],requirePrefix:!0},{quote:"''-qq-bs",prefixes:["@"],requirePrefix:!0},{quote:"``",prefixes:["@"],requirePrefix:!0}],paramTypes:{positional:!0},lineCommentTypes:["--","#"],operators:["%",":=","&","|","^","~","<<",">>","<=>","&&","||","!","*.*"],postProcess:yt},formatOptions:{onelineClauses:[...Ar,...Jt],tabularOnelineClauses:Jt}},sS=["ACCESSIBLE","ADD","ALL","ALTER","ANALYZE","AND","AS","ASC","ASENSITIVE","BEFORE","BETWEEN","BOTH","BY","CALL","CASCADE","CASE","CHANGE","CHECK","COLLATE","COLUMN","CONDITION","CONSTRAINT","CONTINUE","CONVERT","CREATE","CROSS","CUBE","CUME_DIST","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURSOR","DATABASE","DATABASES","DAY_HOUR","DAY_MICROSECOND","DAY_MINUTE","DAY_SECOND","DECLARE","DEFAULT","DELAYED","DELETE","DENSE_RANK","DESC","DESCRIBE","DETERMINISTIC","DISTINCT","DISTINCTROW","DIV","DROP","DUAL","EACH","ELSE","ELSEIF","EMPTY","ENCLOSED","ESCAPED","EXCEPT","EXISTS","EXIT","EXPLAIN","FALSE","FETCH","FIRST_VALUE","FOR","FORCE","FOREIGN","FROM","FULLTEXT","FUNCTION","GENERATED","GET","GRANT","GROUP","GROUPING","GROUPS","HAVING","HIGH_PRIORITY","HOUR_MICROSECOND","HOUR_MINUTE","HOUR_SECOND","IF","IGNORE","IN","INDEX","INFILE","INNER","INOUT","INSENSITIVE","INSERT","IN","INTERSECT","INTERVAL","INTO","IO_AFTER_GTIDS","IO_BEFORE_GTIDS","IS","ITERATE","JOIN","JSON_TABLE","KEY","KEYS","KILL","LAG","LAST_VALUE","LATERAL","LEAD","LEADING","LEAVE","LEFT","LIKE","LIMIT","LINEAR","LINES","LOAD","LOCALTIME","LOCALTIMESTAMP","LOCK","LONG","LOOP","LOW_PRIORITY","MASTER_BIND","MASTER_SSL_VERIFY_SERVER_CERT","MATCH","MAXVALUE","MINUTE_MICROSECOND","MINUTE_SECOND","MOD","MODIFIES","NATURAL","NOT","NO_WRITE_TO_BINLOG","NTH_VALUE","NTILE","NULL","OF","ON","OPTIMIZE","OPTIMIZER_COSTS","OPTION","OPTIONALLY","OR","ORDER","OUT","OUTER","OUTFILE","OVER","PARTITION","PERCENT_RANK","PRIMARY","PROCEDURE","PURGE","RANGE","RANK","READ","READS","READ_WRITE","RECURSIVE","REFERENCES","REGEXP","RELEASE","RENAME","REPEAT","REPLACE","REQUIRE","RESIGNAL","RESTRICT","RETURN","REVOKE","RIGHT","RLIKE","ROW","ROWS","ROW_NUMBER","SCHEMA","SCHEMAS","SECOND_MICROSECOND","SELECT","SENSITIVE","SEPARATOR","SET","SHOW","SIGNAL","SPATIAL","SPECIFIC","SQL","SQLEXCEPTION","SQLSTATE","SQLWARNING","SQL_BIG_RESULT","SQL_CALC_FOUND_ROWS","SQL_SMALL_RESULT","SSL","STARTING","STORED","STRAIGHT_JOIN","SYSTEM","TABLE","TERMINATED","THEN","TO","TRAILING","TRIGGER","TRUE","UNDO","UNION","UNIQUE","UNLOCK","UNSIGNED","UPDATE","USAGE","USE","USING","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","VALUES","VIRTUAL","WHEN","WHERE","WHILE","WINDOW","WITH","WRITE","XOR","YEAR_MONTH","ZEROFILL"],SS=["BIGINT","BINARY","BIT","BLOB","BOOL","BOOLEAN","CHAR","CHARACTER","DATE","DATETIME","DEC","DECIMAL","DOUBLE PRECISION","DOUBLE","ENUM","FIXED","FLOAT","FLOAT4","FLOAT8","INT","INT1","INT2","INT3","INT4","INT8","INTEGER","LONGBLOB","LONGTEXT","MEDIUMBLOB","MEDIUMINT","MEDIUMTEXT","MIDDLEINT","NATIONAL CHAR","NATIONAL VARCHAR","NUMERIC","PRECISION","REAL","SMALLINT","TEXT","TIME","TIMESTAMP","TINYBLOB","TINYINT","TINYTEXT","VARBINARY","VARCHAR","VARCHARACTER","VARYING","YEAR"],oS=["ABS","ACOS","ADDDATE","ADDTIME","AES_DECRYPT","AES_ENCRYPT","ANY_VALUE","ASCII","ASIN","ATAN","ATAN2","AVG","BENCHMARK","BIN","BIN_TO_UUID","BINARY","BIT_AND","BIT_COUNT","BIT_LENGTH","BIT_OR","BIT_XOR","CAN_ACCESS_COLUMN","CAN_ACCESS_DATABASE","CAN_ACCESS_TABLE","CAN_ACCESS_USER","CAN_ACCESS_VIEW","CAST","CEIL","CEILING","CHAR","CHAR_LENGTH","CHARACTER_LENGTH","CHARSET","COALESCE","COERCIBILITY","COLLATION","COMPRESS","CONCAT","CONCAT_WS","CONNECTION_ID","CONV","CONVERT","CONVERT_TZ","COS","COT","COUNT","CRC32","CUME_DIST","CURDATE","CURRENT_DATE","CURRENT_ROLE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURTIME","DATABASE","DATE","DATE_ADD","DATE_FORMAT","DATE_SUB","DATEDIFF","DAY","DAYNAME","DAYOFMONTH","DAYOFWEEK","DAYOFYEAR","DEFAULT","DEGREES","DENSE_RANK","DIV","ELT","EXP","EXPORT_SET","EXTRACT","EXTRACTVALUE","FIELD","FIND_IN_SET","FIRST_VALUE","FLOOR","FORMAT","FORMAT_BYTES","FORMAT_PICO_TIME","FOUND_ROWS","FROM_BASE64","FROM_DAYS","FROM_UNIXTIME","GEOMCOLLECTION","GEOMETRYCOLLECTION","GET_DD_COLUMN_PRIVILEGES","GET_DD_CREATE_OPTIONS","GET_DD_INDEX_SUB_PART_LENGTH","GET_FORMAT","GET_LOCK","GREATEST","GROUP_CONCAT","GROUPING","GTID_SUBSET","GTID_SUBTRACT","HEX","HOUR","ICU_VERSION","IF","IFNULL","INET_ATON","INET_NTOA","INET6_ATON","INET6_NTOA","INSERT","INSTR","INTERNAL_AUTO_INCREMENT","INTERNAL_AVG_ROW_LENGTH","INTERNAL_CHECK_TIME","INTERNAL_CHECKSUM","INTERNAL_DATA_FREE","INTERNAL_DATA_LENGTH","INTERNAL_DD_CHAR_LENGTH","INTERNAL_GET_COMMENT_OR_ERROR","INTERNAL_GET_ENABLED_ROLE_JSON","INTERNAL_GET_HOSTNAME","INTERNAL_GET_USERNAME","INTERNAL_GET_VIEW_WARNING_OR_ERROR","INTERNAL_INDEX_COLUMN_CARDINALITY","INTERNAL_INDEX_LENGTH","INTERNAL_IS_ENABLED_ROLE","INTERNAL_IS_MANDATORY_ROLE","INTERNAL_KEYS_DISABLED","INTERNAL_MAX_DATA_LENGTH","INTERNAL_TABLE_ROWS","INTERNAL_UPDATE_TIME","INTERVAL","IS","IS_FREE_LOCK","IS_IPV4","IS_IPV4_COMPAT","IS_IPV4_MAPPED","IS_IPV6","IS NOT","IS NOT NULL","IS NULL","IS_USED_LOCK","IS_UUID","ISNULL","JSON_ARRAY","JSON_ARRAY_APPEND","JSON_ARRAY_INSERT","JSON_ARRAYAGG","JSON_CONTAINS","JSON_CONTAINS_PATH","JSON_DEPTH","JSON_EXTRACT","JSON_INSERT","JSON_KEYS","JSON_LENGTH","JSON_MERGE","JSON_MERGE_PATCH","JSON_MERGE_PRESERVE","JSON_OBJECT","JSON_OBJECTAGG","JSON_OVERLAPS","JSON_PRETTY","JSON_QUOTE","JSON_REMOVE","JSON_REPLACE","JSON_SCHEMA_VALID","JSON_SCHEMA_VALIDATION_REPORT","JSON_SEARCH","JSON_SET","JSON_STORAGE_FREE","JSON_STORAGE_SIZE","JSON_TABLE","JSON_TYPE","JSON_UNQUOTE","JSON_VALID","JSON_VALUE","LAG","LAST_DAY","LAST_INSERT_ID","LAST_VALUE","LCASE","LEAD","LEAST","LEFT","LENGTH","LIKE","LINESTRING","LN","LOAD_FILE","LOCALTIME","LOCALTIMESTAMP","LOCATE","LOG","LOG10","LOG2","LOWER","LPAD","LTRIM","MAKE_SET","MAKEDATE","MAKETIME","MASTER_POS_WAIT","MATCH","MAX","MBRCONTAINS","MBRCOVEREDBY","MBRCOVERS","MBRDISJOINT","MBREQUALS","MBRINTERSECTS","MBROVERLAPS","MBRTOUCHES","MBRWITHIN","MD5","MEMBER OF","MICROSECOND","MID","MIN","MINUTE","MOD","MONTH","MONTHNAME","MULTILINESTRING","MULTIPOINT","MULTIPOLYGON","NAME_CONST","NOT","NOT IN","NOT LIKE","NOT REGEXP","NOW","NTH_VALUE","NTILE","NULLIF","OCT","OCTET_LENGTH","ORD","PERCENT_RANK","PERIOD_ADD","PERIOD_DIFF","PI","POINT","POLYGON","POSITION","POW","POWER","PS_CURRENT_THREAD_ID","PS_THREAD_ID","QUARTER","QUOTE","RADIANS","RAND","RANDOM_BYTES","RANK","REGEXP","REGEXP_INSTR","REGEXP_LIKE","REGEXP_REPLACE","REGEXP_SUBSTR","RELEASE_ALL_LOCKS","RELEASE_LOCK","REPEAT","REPLACE","REVERSE","RIGHT","RLIKE","ROLES_GRAPHML","ROUND","ROW_COUNT","ROW_NUMBER","RPAD","RTRIM","SCHEMA","SEC_TO_TIME","SECOND","SESSION_USER","SHA1","SHA2","SIGN","SIN","SLEEP","SOUNDEX","SOUNDS LIKE","SOURCE_POS_WAIT","SPACE","SQRT","ST_AREA","ST_ASBINARY","ST_ASGEOJSON","ST_ASTEXT","ST_BUFFER","ST_BUFFER_STRATEGY","ST_CENTROID","ST_COLLECT","ST_CONTAINS","ST_CONVEXHULL","ST_CROSSES","ST_DIFFERENCE","ST_DIMENSION","ST_DISJOINT","ST_DISTANCE","ST_DISTANCE_SPHERE","ST_ENDPOINT","ST_ENVELOPE","ST_EQUALS","ST_EXTERIORRING","ST_FRECHETDISTANCE","ST_GEOHASH","ST_GEOMCOLLFROMTEXT","ST_GEOMCOLLFROMWKB","ST_GEOMETRYN","ST_GEOMETRYTYPE","ST_GEOMFROMGEOJSON","ST_GEOMFROMTEXT","ST_GEOMFROMWKB","ST_HAUSDORFFDISTANCE","ST_INTERIORRINGN","ST_INTERSECTION","ST_INTERSECTS","ST_ISCLOSED","ST_ISEMPTY","ST_ISSIMPLE","ST_ISVALID","ST_LATFROMGEOHASH","ST_LATITUDE","ST_LENGTH","ST_LINEFROMTEXT","ST_LINEFROMWKB","ST_LINEINTERPOLATEPOINT","ST_LINEINTERPOLATEPOINTS","ST_LONGFROMGEOHASH","ST_LONGITUDE","ST_MAKEENVELOPE","ST_MLINEFROMTEXT","ST_MLINEFROMWKB","ST_MPOINTFROMTEXT","ST_MPOINTFROMWKB","ST_MPOLYFROMTEXT","ST_MPOLYFROMWKB","ST_NUMGEOMETRIES","ST_NUMINTERIORRING","ST_NUMPOINTS","ST_OVERLAPS","ST_POINTATDISTANCE","ST_POINTFROMGEOHASH","ST_POINTFROMTEXT","ST_POINTFROMWKB","ST_POINTN","ST_POLYFROMTEXT","ST_POLYFROMWKB","ST_SIMPLIFY","ST_SRID","ST_STARTPOINT","ST_SWAPXY","ST_SYMDIFFERENCE","ST_TOUCHES","ST_TRANSFORM","ST_UNION","ST_VALIDATE","ST_WITHIN","ST_X","ST_Y","STATEMENT_DIGEST","STATEMENT_DIGEST_TEXT","STD","STDDEV","STDDEV_POP","STDDEV_SAMP","STR_TO_DATE","STRCMP","SUBDATE","SUBSTR","SUBSTRING","SUBSTRING_INDEX","SUBTIME","SUM","SYSDATE","SYSTEM_USER","TAN","TIME","TIME_FORMAT","TIME_TO_SEC","TIMEDIFF","TIMESTAMP","TIMESTAMPADD","TIMESTAMPDIFF","TO_BASE64","TO_DAYS","TO_SECONDS","TRIM","TRUNCATE","UCASE","UNCOMPRESS","UNCOMPRESSED_LENGTH","UNHEX","UNIX_TIMESTAMP","UPDATEXML","UPPER","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","UUID","UUID_SHORT","UUID_TO_BIN","VALIDATE_PASSWORD_STRENGTH","VALUES","VAR_POP","VAR_SAMP","VARIANCE","VERSION","WAIT_FOR_EXECUTED_GTID_SET","WAIT_UNTIL_SQL_THREAD_AFTER_GTIDS","WEEK","WEEKDAY","WEEKOFYEAR","WEIGHT_STRING","YEAR","YEARWEEK"],OS=v(["SELECT [ALL | DISTINCT | DISTINCTROW]"]),iS=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","INSERT [LOW_PRIORITY | DELAYED | HIGH_PRIORITY] [IGNORE] [INTO]","REPLACE [LOW_PRIORITY | DELAYED] [INTO]","VALUES","ON DUPLICATE KEY UPDATE","SET"]),sr=v(["CREATE [TEMPORARY] TABLE [IF NOT EXISTS]"]),qt=v(["CREATE [OR REPLACE] [SQL SECURITY DEFINER | SQL SECURITY INVOKER] VIEW [IF NOT EXISTS]","UPDATE [LOW_PRIORITY] [IGNORE]","DELETE [LOW_PRIORITY] [QUICK] [IGNORE] FROM","DROP [TEMPORARY] TABLE [IF EXISTS]","ALTER TABLE","ADD [COLUMN]","{CHANGE | MODIFY} [COLUMN]","DROP [COLUMN]","RENAME [TO | AS]","RENAME COLUMN","ALTER [COLUMN]","{SET | DROP} DEFAULT","TRUNCATE [TABLE]","ALTER DATABASE","ALTER EVENT","ALTER FUNCTION","ALTER INSTANCE","ALTER LOGFILE GROUP","ALTER PROCEDURE","ALTER RESOURCE GROUP","ALTER SERVER","ALTER TABLESPACE","ALTER USER","ALTER VIEW","ANALYZE TABLE","BINLOG","CACHE INDEX","CALL","CHANGE MASTER TO","CHANGE REPLICATION FILTER","CHANGE REPLICATION SOURCE TO","CHECK TABLE","CHECKSUM TABLE","CLONE","COMMIT","CREATE DATABASE","CREATE EVENT","CREATE FUNCTION","CREATE FUNCTION","CREATE INDEX","CREATE LOGFILE GROUP","CREATE PROCEDURE","CREATE RESOURCE GROUP","CREATE ROLE","CREATE SERVER","CREATE SPATIAL REFERENCE SYSTEM","CREATE TABLESPACE","CREATE TRIGGER","CREATE USER","DEALLOCATE PREPARE","DESCRIBE","DROP DATABASE","DROP EVENT","DROP FUNCTION","DROP FUNCTION","DROP INDEX","DROP LOGFILE GROUP","DROP PROCEDURE","DROP RESOURCE GROUP","DROP ROLE","DROP SERVER","DROP SPATIAL REFERENCE SYSTEM","DROP TABLESPACE","DROP TRIGGER","DROP USER","DROP VIEW","EXECUTE","EXPLAIN","FLUSH","GRANT","HANDLER","HELP","IMPORT TABLE","INSTALL COMPONENT","INSTALL PLUGIN","KILL","LOAD DATA","LOAD INDEX INTO CACHE","LOAD XML","LOCK INSTANCE FOR BACKUP","LOCK TABLES","MASTER_POS_WAIT","OPTIMIZE TABLE","PREPARE","PURGE BINARY LOGS","RELEASE SAVEPOINT","RENAME TABLE","RENAME USER","REPAIR TABLE","RESET","RESET MASTER","RESET PERSIST","RESET REPLICA","RESET SLAVE","RESTART","REVOKE","ROLLBACK","ROLLBACK TO SAVEPOINT","SAVEPOINT","SET CHARACTER SET","SET DEFAULT ROLE","SET NAMES","SET PASSWORD","SET RESOURCE GROUP","SET ROLE","SET TRANSACTION","SHOW","SHOW BINARY LOGS","SHOW BINLOG EVENTS","SHOW CHARACTER SET","SHOW COLLATION","SHOW COLUMNS","SHOW CREATE DATABASE","SHOW CREATE EVENT","SHOW CREATE FUNCTION","SHOW CREATE PROCEDURE","SHOW CREATE TABLE","SHOW CREATE TRIGGER","SHOW CREATE USER","SHOW CREATE VIEW","SHOW DATABASES","SHOW ENGINE","SHOW ENGINES","SHOW ERRORS","SHOW EVENTS","SHOW FUNCTION CODE","SHOW FUNCTION STATUS","SHOW GRANTS","SHOW INDEX","SHOW MASTER STATUS","SHOW OPEN TABLES","SHOW PLUGINS","SHOW PRIVILEGES","SHOW PROCEDURE CODE","SHOW PROCEDURE STATUS","SHOW PROCESSLIST","SHOW PROFILE","SHOW PROFILES","SHOW RELAYLOG EVENTS","SHOW REPLICA STATUS","SHOW REPLICAS","SHOW SLAVE","SHOW SLAVE HOSTS","SHOW STATUS","SHOW TABLE STATUS","SHOW TABLES","SHOW TRIGGERS","SHOW VARIABLES","SHOW WARNINGS","SHUTDOWN","SOURCE_POS_WAIT","START GROUP_REPLICATION","START REPLICA","START SLAVE","START TRANSACTION","STOP GROUP_REPLICATION","STOP REPLICA","STOP SLAVE","TABLE","UNINSTALL COMPONENT","UNINSTALL PLUGIN","UNLOCK INSTANCE","UNLOCK TABLES","USE","XA","ITERATE","LEAVE","LOOP","REPEAT","RETURN","WHILE"]),aS=v(["UNION [ALL | DISTINCT]"]),IS=v(["JOIN","{LEFT | RIGHT} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT} [OUTER] JOIN","STRAIGHT_JOIN"]),NS=v(["ON {UPDATE | DELETE} [SET NULL]","CHARACTER SET","{ROWS | RANGE} BETWEEN","IDENTIFIED BY"]),lS={name:"mysql",tokenizerOptions:{reservedSelect:OS,reservedClauses:[...iS,...sr,...qt],reservedSetOperations:aS,reservedJoins:IS,reservedPhrases:NS,supportsXor:!0,reservedKeywords:sS,reservedDataTypes:SS,reservedFunctionNames:oS,stringTypes:['""-qq-bs',{quote:"''-qq-bs",prefixes:["N"]},{quote:"''-raw",prefixes:["B","X"],requirePrefix:!0}],identTypes:["``"],identChars:{first:"$",rest:"$",allowFirstCharNumber:!0},variableTypes:[{regex:"@@?[A-Za-z0-9_.$]+"},{quote:'""-qq-bs',prefixes:["@"],requirePrefix:!0},{quote:"''-qq-bs",prefixes:["@"],requirePrefix:!0},{quote:"``",prefixes:["@"],requirePrefix:!0}],paramTypes:{positional:!0},lineCommentTypes:["--","#"],operators:["%",":=","&","|","^","~","<<",">>","<=>","->","->>","&&","||","!","*.*"],postProcess:yt},formatOptions:{onelineClauses:[...sr,...qt],tabularOnelineClauses:qt}},_S=["ADD","ALL","ALTER","ANALYZE","AND","ARRAY","AS","ASC","BETWEEN","BOTH","BY","CALL","CASCADE","CASE","CHANGE","CHECK","COLLATE","COLUMN","CONSTRAINT","CONTINUE","CONVERT","CREATE","CROSS","CURRENT_DATE","CURRENT_ROLE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURSOR","DATABASE","DATABASES","DAY_HOUR","DAY_MICROSECOND","DAY_MINUTE","DAY_SECOND","DEFAULT","DELAYED","DELETE","DESC","DESCRIBE","DISTINCT","DISTINCTROW","DIV","DOUBLE","DROP","DUAL","ELSE","ELSEIF","ENCLOSED","ESCAPED","EXCEPT","EXISTS","EXIT","EXPLAIN","FALSE","FETCH","FOR","FORCE","FOREIGN","FROM","FULLTEXT","GENERATED","GRANT","GROUP","GROUPS","HAVING","HIGH_PRIORITY","HOUR_MICROSECOND","HOUR_MINUTE","HOUR_SECOND","IF","IGNORE","ILIKE","IN","INDEX","INFILE","INNER","INOUT","INSERT","INTERSECT","INTERVAL","INTO","IS","ITERATE","JOIN","KEY","KEYS","KILL","LEADING","LEAVE","LEFT","LIKE","LIMIT","LINEAR","LINES","LOAD","LOCALTIME","LOCALTIMESTAMP","LOCK","LONG","LOW_PRIORITY","MATCH","MAXVALUE","MINUTE_MICROSECOND","MINUTE_SECOND","MOD","NATURAL","NOT","NO_WRITE_TO_BINLOG","NULL","OF","ON","OPTIMIZE","OPTION","OPTIONALLY","OR","ORDER","OUT","OUTER","OUTFILE","OVER","PARTITION","PRIMARY","PROCEDURE","RANGE","READ","RECURSIVE","REFERENCES","REGEXP","RELEASE","RENAME","REPEAT","REPLACE","REQUIRE","RESTRICT","REVOKE","RIGHT","RLIKE","ROW","ROWS","SECOND_MICROSECOND","SELECT","SET","SHOW","SPATIAL","SQL","SQLEXCEPTION","SQLSTATE","SQLWARNING","SQL_BIG_RESULT","SQL_CALC_FOUND_ROWS","SQL_SMALL_RESULT","SSL","STARTING","STATS_EXTENDED","STORED","STRAIGHT_JOIN","TABLE","TABLESAMPLE","TERMINATED","THEN","TO","TRAILING","TRIGGER","TRUE","TiDB_CURRENT_TSO","UNION","UNIQUE","UNLOCK","UNSIGNED","UNTIL","UPDATE","USAGE","USE","USING","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","VALUES","VIRTUAL","WHEN","WHERE","WHILE","WINDOW","WITH","WRITE","XOR","YEAR_MONTH","ZEROFILL"],LS=["BIGINT","BINARY","BIT","BLOB","BOOL","BOOLEAN","CHAR","CHARACTER","DATE","DATETIME","DEC","DECIMAL","DOUBLE PRECISION","DOUBLE","ENUM","FIXED","INT","INT1","INT2","INT3","INT4","INT8","INTEGER","LONGBLOB","LONGTEXT","MEDIUMBLOB","MEDIUMINT","MIDDLEINT","NATIONAL CHAR","NATIONAL VARCHAR","NUMERIC","PRECISION","SMALLINT","TEXT","TIME","TIMESTAMP","TINYBLOB","TINYINT","TINYTEXT","VARBINARY","VARCHAR","VARCHARACTER","VARYING","YEAR"],CS=["ABS","ACOS","ADDDATE","ADDTIME","AES_DECRYPT","AES_ENCRYPT","ANY_VALUE","ASCII","ASIN","ATAN","ATAN2","AVG","BENCHMARK","BIN","BIN_TO_UUID","BIT_AND","BIT_COUNT","BIT_LENGTH","BIT_OR","BIT_XOR","BITAND","BITNEG","BITOR","BITXOR","CASE","CAST","CEIL","CEILING","CHAR_FUNC","CHAR_LENGTH","CHARACTER_LENGTH","CHARSET","COALESCE","COERCIBILITY","COLLATION","COMPRESS","CONCAT","CONCAT_WS","CONNECTION_ID","CONV","CONVERT","CONVERT_TZ","COS","COT","COUNT","CRC32","CUME_DIST","CURDATE","CURRENT_DATE","CURRENT_RESOURCE_GROUP","CURRENT_ROLE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURTIME","DATABASE","DATE","DATE_ADD","DATE_FORMAT","DATE_SUB","DATEDIFF","DAY","DAYNAME","DAYOFMONTH","DAYOFWEEK","DAYOFYEAR","DECODE","DEFAULT_FUNC","DEGREES","DENSE_RANK","DES_DECRYPT","DES_ENCRYPT","DIV","ELT","ENCODE","ENCRYPT","EQ","EXP","EXPORT_SET","EXTRACT","FIELD","FIND_IN_SET","FIRST_VALUE","FLOOR","FORMAT","FORMAT_BYTES","FORMAT_NANO_TIME","FOUND_ROWS","FROM_BASE64","FROM_DAYS","FROM_UNIXTIME","GE","GET_FORMAT","GET_LOCK","GETPARAM","GREATEST","GROUP_CONCAT","GROUPING","GT","HEX","HOUR","IF","IFNULL","ILIKE","INET6_ATON","INET6_NTOA","INET_ATON","INET_NTOA","INSERT_FUNC","INSTR","INTDIV","INTERVAL","IS_FREE_LOCK","IS_IPV4","IS_IPV4_COMPAT","IS_IPV4_MAPPED","IS_IPV6","IS_USED_LOCK","IS_UUID","ISFALSE","ISNULL","ISTRUE","JSON_ARRAY","JSON_ARRAYAGG","JSON_ARRAY_APPEND","JSON_ARRAY_INSERT","JSON_CONTAINS","JSON_CONTAINS_PATH","JSON_DEPTH","JSON_EXTRACT","JSON_INSERT","JSON_KEYS","JSON_LENGTH","JSON_MEMBEROF","JSON_MERGE","JSON_MERGE_PATCH","JSON_MERGE_PRESERVE","JSON_OBJECT","JSON_OBJECTAGG","JSON_OVERLAPS","JSON_PRETTY","JSON_QUOTE","JSON_REMOVE","JSON_REPLACE","JSON_SEARCH","JSON_SET","JSON_STORAGE_FREE","JSON_STORAGE_SIZE","JSON_TYPE","JSON_UNQUOTE","JSON_VALID","LAG","LAST_DAY","LAST_INSERT_ID","LAST_VALUE","LASTVAL","LCASE","LE","LEAD","LEAST","LEFT","LEFTSHIFT","LENGTH","LIKE","LN","LOAD_FILE","LOCALTIME","LOCALTIMESTAMP","LOCATE","LOG","LOG10","LOG2","LOWER","LPAD","LT","LTRIM","MAKE_SET","MAKEDATE","MAKETIME","MASTER_POS_WAIT","MAX","MD5","MICROSECOND","MID","MIN","MINUS","MINUTE","MOD","MONTH","MONTHNAME","MUL","NAME_CONST","NE","NEXTVAL","NOT","NOW","NTH_VALUE","NTILE","NULLEQ","OCT","OCTET_LENGTH","OLD_PASSWORD","ORD","PASSWORD_FUNC","PERCENT_RANK","PERIOD_ADD","PERIOD_DIFF","PI","PLUS","POSITION","POW","POWER","QUARTER","QUOTE","RADIANS","RAND","RANDOM_BYTES","RANK","REGEXP","REGEXP_INSTR","REGEXP_LIKE","REGEXP_REPLACE","REGEXP_SUBSTR","RELEASE_ALL_LOCKS","RELEASE_LOCK","REPEAT","REPLACE","REVERSE","RIGHT","RIGHTSHIFT","ROUND","ROW_COUNT","ROW_NUMBER","RPAD","RTRIM","SCHEMA","SEC_TO_TIME","SECOND","SESSION_USER","SETVAL","SETVAR","SHA","SHA1","SHA2","SIGN","SIN","SLEEP","SM3","SPACE","SQRT","STD","STDDEV","STDDEV_POP","STDDEV_SAMP","STR_TO_DATE","STRCMP","SUBDATE","SUBSTR","SUBSTRING","SUBSTRING_INDEX","SUBTIME","SUM","SYSDATE","SYSTEM_USER","TAN","TIDB_BOUNDED_STALENESS","TIDB_CURRENT_TSO","TIDB_DECODE_BINARY_PLAN","TIDB_DECODE_KEY","TIDB_DECODE_PLAN","TIDB_DECODE_SQL_DIGESTS","TIDB_ENCODE_SQL_DIGEST","TIDB_IS_DDL_OWNER","TIDB_PARSE_TSO","TIDB_PARSE_TSO_LOGICAL","TIDB_ROW_CHECKSUM","TIDB_SHARD","TIDB_VERSION","TIME","TIME_FORMAT","TIME_TO_SEC","TIMEDIFF","TIMESTAMP","TIMESTAMPADD","TIMESTAMPDIFF","TO_BASE64","TO_DAYS","TO_SECONDS","TRANSLATE","TRIM","TRUNCATE","UCASE","UNARYMINUS","UNCOMPRESS","UNCOMPRESSED_LENGTH","UNHEX","UNIX_TIMESTAMP","UPPER","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","UUID","UUID_SHORT","UUID_TO_BIN","VALIDATE_PASSWORD_STRENGTH","VAR_POP","VAR_SAMP","VARIANCE","VERSION","VITESS_HASH","WEEK","WEEKDAY","WEEKOFYEAR","WEIGHT_STRING","YEAR","YEARWEEK"],uS=v(["SELECT [ALL | DISTINCT | DISTINCTROW]"]),cS=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","INSERT [LOW_PRIORITY | DELAYED | HIGH_PRIORITY] [IGNORE] [INTO]","REPLACE [LOW_PRIORITY | DELAYED] [INTO]","VALUES","ON DUPLICATE KEY UPDATE","SET"]),Sr=v(["CREATE [TEMPORARY] TABLE [IF NOT EXISTS]"]),Qt=v(["CREATE [OR REPLACE] [SQL SECURITY DEFINER | SQL SECURITY INVOKER] VIEW [IF NOT EXISTS]","UPDATE [LOW_PRIORITY] [IGNORE]","DELETE [LOW_PRIORITY] [QUICK] [IGNORE] FROM","DROP [TEMPORARY] TABLE [IF EXISTS]","ALTER TABLE","ADD [COLUMN]","{CHANGE | MODIFY} [COLUMN]","DROP [COLUMN]","RENAME [TO | AS]","RENAME COLUMN","ALTER [COLUMN]","{SET | DROP} DEFAULT","TRUNCATE [TABLE]","ALTER DATABASE","ALTER INSTANCE","ALTER RESOURCE GROUP","ALTER SEQUENCE","ALTER USER","ALTER VIEW","ANALYZE TABLE","CHECK TABLE","CHECKSUM TABLE","COMMIT","CREATE DATABASE","CREATE INDEX","CREATE RESOURCE GROUP","CREATE ROLE","CREATE SEQUENCE","CREATE USER","DEALLOCATE PREPARE","DESCRIBE","DROP DATABASE","DROP INDEX","DROP RESOURCE GROUP","DROP ROLE","DROP TABLESPACE","DROP USER","DROP VIEW","EXPLAIN","FLUSH","GRANT","IMPORT TABLE","INSTALL COMPONENT","INSTALL PLUGIN","KILL","LOAD DATA","LOCK INSTANCE FOR BACKUP","LOCK TABLES","OPTIMIZE TABLE","PREPARE","RELEASE SAVEPOINT","RENAME TABLE","RENAME USER","REPAIR TABLE","RESET","REVOKE","ROLLBACK","ROLLBACK TO SAVEPOINT","SAVEPOINT","SET CHARACTER SET","SET DEFAULT ROLE","SET NAMES","SET PASSWORD","SET RESOURCE GROUP","SET ROLE","SET TRANSACTION","SHOW","SHOW BINARY LOGS","SHOW BINLOG EVENTS","SHOW CHARACTER SET","SHOW COLLATION","SHOW COLUMNS","SHOW CREATE DATABASE","SHOW CREATE TABLE","SHOW CREATE USER","SHOW CREATE VIEW","SHOW DATABASES","SHOW ENGINE","SHOW ENGINES","SHOW ERRORS","SHOW EVENTS","SHOW GRANTS","SHOW INDEX","SHOW MASTER STATUS","SHOW OPEN TABLES","SHOW PLUGINS","SHOW PRIVILEGES","SHOW PROCESSLIST","SHOW PROFILE","SHOW PROFILES","SHOW STATUS","SHOW TABLE STATUS","SHOW TABLES","SHOW TRIGGERS","SHOW VARIABLES","SHOW WARNINGS","TABLE","UNINSTALL COMPONENT","UNINSTALL PLUGIN","UNLOCK INSTANCE","UNLOCK TABLES","USE"]),fS=v(["UNION [ALL | DISTINCT]"]),PS=v(["JOIN","{LEFT | RIGHT} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT} [OUTER] JOIN","STRAIGHT_JOIN"]),DS=v(["ON {UPDATE | DELETE} [SET NULL]","CHARACTER SET","{ROWS | RANGE} BETWEEN","IDENTIFIED BY"]),dS={name:"tidb",tokenizerOptions:{reservedSelect:uS,reservedClauses:[...cS,...Sr,...Qt],reservedSetOperations:fS,reservedJoins:PS,reservedPhrases:DS,supportsXor:!0,reservedKeywords:_S,reservedDataTypes:LS,reservedFunctionNames:CS,stringTypes:['""-qq-bs',{quote:"''-qq-bs",prefixes:["N"]},{quote:"''-raw",prefixes:["B","X"],requirePrefix:!0}],identTypes:["``"],identChars:{first:"$",rest:"$",allowFirstCharNumber:!0},variableTypes:[{regex:"@@?[A-Za-z0-9_.$]+"},{quote:'""-qq-bs',prefixes:["@"],requirePrefix:!0},{quote:"''-qq-bs",prefixes:["@"],requirePrefix:!0},{quote:"``",prefixes:["@"],requirePrefix:!0}],paramTypes:{positional:!0},lineCommentTypes:["--","#"],operators:["%",":=","&","|","^","~","<<",">>","<=>","->","->>","&&","||","!","*.*"],postProcess:yt},formatOptions:{onelineClauses:[...Sr,...Qt],tabularOnelineClauses:Qt}},pS=["ABORT","ABS","ACOS","ADVISOR","ARRAY_AGG","ARRAY_AGG","ARRAY_APPEND","ARRAY_AVG","ARRAY_BINARY_SEARCH","ARRAY_CONCAT","ARRAY_CONTAINS","ARRAY_COUNT","ARRAY_DISTINCT","ARRAY_EXCEPT","ARRAY_FLATTEN","ARRAY_IFNULL","ARRAY_INSERT","ARRAY_INTERSECT","ARRAY_LENGTH","ARRAY_MAX","ARRAY_MIN","ARRAY_MOVE","ARRAY_POSITION","ARRAY_PREPEND","ARRAY_PUT","ARRAY_RANGE","ARRAY_REMOVE","ARRAY_REPEAT","ARRAY_REPLACE","ARRAY_REVERSE","ARRAY_SORT","ARRAY_STAR","ARRAY_SUM","ARRAY_SYMDIFF","ARRAY_SYMDIFF1","ARRAY_SYMDIFFN","ARRAY_UNION","ASIN","ATAN","ATAN2","AVG","BASE64","BASE64_DECODE","BASE64_ENCODE","BITAND ","BITCLEAR ","BITNOT ","BITOR ","BITSET ","BITSHIFT ","BITTEST ","BITXOR ","CEIL","CLOCK_LOCAL","CLOCK_MILLIS","CLOCK_STR","CLOCK_TZ","CLOCK_UTC","COALESCE","CONCAT","CONCAT2","CONTAINS","CONTAINS_TOKEN","CONTAINS_TOKEN_LIKE","CONTAINS_TOKEN_REGEXP","COS","COUNT","COUNT","COUNTN","CUME_DIST","CURL","DATE_ADD_MILLIS","DATE_ADD_STR","DATE_DIFF_MILLIS","DATE_DIFF_STR","DATE_FORMAT_STR","DATE_PART_MILLIS","DATE_PART_STR","DATE_RANGE_MILLIS","DATE_RANGE_STR","DATE_TRUNC_MILLIS","DATE_TRUNC_STR","DECODE","DECODE_JSON","DEGREES","DENSE_RANK","DURATION_TO_STR","ENCODED_SIZE","ENCODE_JSON","EXP","FIRST_VALUE","FLOOR","GREATEST","HAS_TOKEN","IFINF","IFMISSING","IFMISSINGORNULL","IFNAN","IFNANORINF","IFNULL","INITCAP","ISARRAY","ISATOM","ISBITSET","ISBOOLEAN","ISNUMBER","ISOBJECT","ISSTRING","LAG","LAST_VALUE","LEAD","LEAST","LENGTH","LN","LOG","LOWER","LTRIM","MAX","MEAN","MEDIAN","META","MILLIS","MILLIS_TO_LOCAL","MILLIS_TO_STR","MILLIS_TO_TZ","MILLIS_TO_UTC","MILLIS_TO_ZONE_NAME","MIN","MISSINGIF","NANIF","NEGINFIF","NOW_LOCAL","NOW_MILLIS","NOW_STR","NOW_TZ","NOW_UTC","NTH_VALUE","NTILE","NULLIF","NVL","NVL2","OBJECT_ADD","OBJECT_CONCAT","OBJECT_INNER_PAIRS","OBJECT_INNER_VALUES","OBJECT_LENGTH","OBJECT_NAMES","OBJECT_PAIRS","OBJECT_PUT","OBJECT_REMOVE","OBJECT_RENAME","OBJECT_REPLACE","OBJECT_UNWRAP","OBJECT_VALUES","PAIRS","PERCENT_RANK","PI","POLY_LENGTH","POSINFIF","POSITION","POWER","RADIANS","RANDOM","RANK","RATIO_TO_REPORT","REGEXP_CONTAINS","REGEXP_LIKE","REGEXP_MATCHES","REGEXP_POSITION","REGEXP_REPLACE","REGEXP_SPLIT","REGEX_CONTAINS","REGEX_LIKE","REGEX_MATCHES","REGEX_POSITION","REGEX_REPLACE","REGEX_SPLIT","REPEAT","REPLACE","REVERSE","ROUND","ROW_NUMBER","RTRIM","SEARCH","SEARCH_META","SEARCH_SCORE","SIGN","SIN","SPLIT","SQRT","STDDEV","STDDEV_POP","STDDEV_SAMP","STR_TO_DURATION","STR_TO_MILLIS","STR_TO_TZ","STR_TO_UTC","STR_TO_ZONE_NAME","SUBSTR","SUFFIXES","SUM","TAN","TITLE","TOARRAY","TOATOM","TOBOOLEAN","TOKENS","TOKENS","TONUMBER","TOOBJECT","TOSTRING","TRIM","TRUNC","UPPER","UUID","VARIANCE","VARIANCE_POP","VARIANCE_SAMP","VAR_POP","VAR_SAMP","WEEKDAY_MILLIS","WEEKDAY_STR","CAST"],MS=["ADVISE","ALL","ALTER","ANALYZE","AND","ANY","ARRAY","AS","ASC","AT","BEGIN","BETWEEN","BINARY","BOOLEAN","BREAK","BUCKET","BUILD","BY","CALL","CASE","CAST","CLUSTER","COLLATE","COLLECTION","COMMIT","COMMITTED","CONNECT","CONTINUE","CORRELATED","COVER","CREATE","CURRENT","DATABASE","DATASET","DATASTORE","DECLARE","DECREMENT","DELETE","DERIVED","DESC","DESCRIBE","DISTINCT","DO","DROP","EACH","ELEMENT","ELSE","END","EVERY","EXCEPT","EXCLUDE","EXECUTE","EXISTS","EXPLAIN","FALSE","FETCH","FILTER","FIRST","FLATTEN","FLUSH","FOLLOWING","FOR","FORCE","FROM","FTS","FUNCTION","GOLANG","GRANT","GROUP","GROUPS","GSI","HASH","HAVING","IF","IGNORE","ILIKE","IN","INCLUDE","INCREMENT","INDEX","INFER","INLINE","INNER","INSERT","INTERSECT","INTO","IS","ISOLATION","JAVASCRIPT","JOIN","KEY","KEYS","KEYSPACE","KNOWN","LANGUAGE","LAST","LEFT","LET","LETTING","LEVEL","LIKE","LIMIT","LSM","MAP","MAPPING","MATCHED","MATERIALIZED","MERGE","MINUS","MISSING","NAMESPACE","NEST","NL","NO","NOT","NTH_VALUE","NULL","NULLS","NUMBER","OBJECT","OFFSET","ON","OPTION","OPTIONS","OR","ORDER","OTHERS","OUTER","OVER","PARSE","PARTITION","PASSWORD","PATH","POOL","PRECEDING","PREPARE","PRIMARY","PRIVATE","PRIVILEGE","PROBE","PROCEDURE","PUBLIC","RANGE","RAW","REALM","REDUCE","RENAME","RESPECT","RETURN","RETURNING","REVOKE","RIGHT","ROLE","ROLLBACK","ROW","ROWS","SATISFIES","SAVEPOINT","SCHEMA","SCOPE","SELECT","SELF","SEMI","SET","SHOW","SOME","START","STATISTICS","STRING","SYSTEM","THEN","TIES","TO","TRAN","TRANSACTION","TRIGGER","TRUE","TRUNCATE","UNBOUNDED","UNDER","UNION","UNIQUE","UNKNOWN","UNNEST","UNSET","UPDATE","UPSERT","USE","USER","USING","VALIDATE","VALUE","VALUED","VALUES","VIA","VIEW","WHEN","WHERE","WHILE","WINDOW","WITH","WITHIN","WORK","XOR"],US=[],mS=v(["SELECT [ALL | DISTINCT]"]),hS=v(["WITH","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","INSERT INTO","VALUES","SET","MERGE INTO","WHEN [NOT] MATCHED THEN","UPDATE SET","INSERT","NEST","UNNEST","RETURNING"]),or=v(["UPDATE","DELETE FROM","SET SCHEMA","ADVISE","ALTER INDEX","BEGIN TRANSACTION","BUILD INDEX","COMMIT TRANSACTION","CREATE COLLECTION","CREATE FUNCTION","CREATE INDEX","CREATE PRIMARY INDEX","CREATE SCOPE","DROP COLLECTION","DROP FUNCTION","DROP INDEX","DROP PRIMARY INDEX","DROP SCOPE","EXECUTE","EXECUTE FUNCTION","EXPLAIN","GRANT","INFER","PREPARE","REVOKE","ROLLBACK TRANSACTION","SAVEPOINT","SET TRANSACTION","UPDATE STATISTICS","UPSERT","LET","SET CURRENT SCHEMA","SHOW","USE [PRIMARY] KEYS"]),GS=v(["UNION [ALL]","EXCEPT [ALL]","INTERSECT [ALL]"]),gS=v(["JOIN","{LEFT | RIGHT} [OUTER] JOIN","INNER JOIN"]),HS=v(["{ROWS | RANGE | GROUPS} BETWEEN"]),bS={name:"n1ql",tokenizerOptions:{reservedSelect:mS,reservedClauses:[...hS,...or],reservedSetOperations:GS,reservedJoins:gS,reservedPhrases:HS,supportsXor:!0,reservedKeywords:MS,reservedDataTypes:US,reservedFunctionNames:pS,stringTypes:['""-bs',"''-bs"],identTypes:["``"],extraParens:["[]","{}"],paramTypes:{positional:!0,numbered:["$"],named:["$"]},lineCommentTypes:["#","--"],operators:["%","==",":","||"]},formatOptions:{onelineClauses:or}},yS=["ADD","AGENT","AGGREGATE","ALL","ALTER","AND","ANY","ARROW","AS","ASC","AT","ATTRIBUTE","AUTHID","AVG","BEGIN","BETWEEN","BLOCK","BODY","BOTH","BOUND","BULK","BY","BYTE","CALL","CALLING","CASCADE","CASE","CHARSET","CHARSETFORM","CHARSETID","CHECK","CLOSE","CLUSTER","CLUSTERS","COLAUTH","COLLECT","COLUMNS","COMMENT","COMMIT","COMMITTED","COMPILED","COMPRESS","CONNECT","CONSTANT","CONSTRUCTOR","CONTEXT","CONVERT","COUNT","CRASH","CREATE","CURRENT","CURSOR","CUSTOMDATUM","DANGLING","DATA","DAY","DECLARE","DEFAULT","DEFINE","DELETE","DESC","DETERMINISTIC","DISTINCT","DROP","DURATION","ELEMENT","ELSE","ELSIF","EMPTY","END","ESCAPE","EXCEPT","EXCEPTION","EXCEPTIONS","EXCLUSIVE","EXECUTE","EXISTS","EXIT","EXTERNAL","FETCH","FINAL","FIXED","FOR","FORALL","FORCE","FORM","FROM","FUNCTION","GENERAL","GOTO","GRANT","GROUP","HASH","HAVING","HEAP","HIDDEN","HOUR","IDENTIFIED","IF","IMMEDIATE","IN","INCLUDING","INDEX","INDEXES","INDICATOR","INDICES","INFINITE","INSERT","INSTANTIABLE","INTERFACE","INTERSECT","INTERVAL","INTO","INVALIDATE","IS","ISOLATION","JAVA","LANGUAGE","LARGE","LEADING","LENGTH","LEVEL","LIBRARY","LIKE","LIKE2","LIKE4","LIKEC","LIMIT","LIMITED","LOCAL","LOCK","LOOP","MAP","MAX","MAXLEN","MEMBER","MERGE","MIN","MINUS","MINUTE","MOD","MODE","MODIFY","MONTH","MULTISET","NAME","NAN","NATIONAL","NATIVE","NEW","NOCOMPRESS","NOCOPY","NOT","NOWAIT","NULL","OBJECT","OCICOLL","OCIDATE","OCIDATETIME","OCIDURATION","OCIINTERVAL","OCILOBLOCATOR","OCINUMBER","OCIRAW","OCIREF","OCIREFCURSOR","OCIROWID","OCISTRING","OCITYPE","OF","ON","ONLY","OPAQUE","OPEN","OPERATOR","OPTION","OR","ORACLE","ORADATA","ORDER","OVERLAPS","ORGANIZATION","ORLANY","ORLVARY","OTHERS","OUT","OVERRIDING","PACKAGE","PARALLEL_ENABLE","PARAMETER","PARAMETERS","PARTITION","PASCAL","PIPE","PIPELINED","PRAGMA","PRIOR","PRIVATE","PROCEDURE","PUBLIC","RAISE","RANGE","READ","RECORD","REF","REFERENCE","REM","REMAINDER","RENAME","RESOURCE","RESULT","RETURN","RETURNING","REVERSE","REVOKE","ROLLBACK","ROW","SAMPLE","SAVE","SAVEPOINT","SB1","SB2","SB4","SECOND","SEGMENT","SELECT","SELF","SEPARATE","SEQUENCE","SERIALIZABLE","SET","SHARE","SHORT","SIZE","SIZE_T","SOME","SPARSE","SQL","SQLCODE","SQLDATA","SQLNAME","SQLSTATE","STANDARD","START","STATIC","STDDEV","STORED","STRING","STRUCT","STYLE","SUBMULTISET","SUBPARTITION","SUBSTITUTABLE","SUBTYPE","SUM","SYNONYM","TABAUTH","TABLE","TDO","THE","THEN","TIME","TIMEZONE_ABBR","TIMEZONE_HOUR","TIMEZONE_MINUTE","TIMEZONE_REGION","TO","TRAILING","TRANSAC","TRANSACTIONAL","TRUSTED","TYPE","UB1","UB2","UB4","UNDER","UNION","UNIQUE","UNSIGNED","UNTRUSTED","UPDATE","USE","USING","VALIST","VALUE","VALUES","VARIABLE","VARIANCE","VARRAY","VIEW","VIEWS","VOID","WHEN","WHERE","WHILE","WITH","WORK","WRAPPED","WRITE","YEAR","ZONE"],BS=["ARRAY","BFILE_BASE","BINARY","BLOB_BASE","CHAR VARYING","CHAR_BASE","CHAR","CHARACTER VARYING","CHARACTER","CLOB_BASE","DATE_BASE","DATE","DECIMAL","DOUBLE","FLOAT","INT","INTERVAL DAY","INTERVAL YEAR","LONG","NATIONAL CHAR VARYING","NATIONAL CHAR","NATIONAL CHARACTER VARYING","NATIONAL CHARACTER","NCHAR VARYING","NCHAR","NCHAR","NUMBER_BASE","NUMBER","NUMBERIC","NVARCHAR","PRECISION","RAW","TIMESTAMP","UROWID","VARCHAR","VARCHAR2"],vS=["ABS","ACOS","ASIN","ATAN","ATAN2","BITAND","CEIL","COS","COSH","EXP","FLOOR","LN","LOG","MOD","NANVL","POWER","REMAINDER","ROUND","SIGN","SIN","SINH","SQRT","TAN","TANH","TRUNC","WIDTH_BUCKET","CHR","CONCAT","INITCAP","LOWER","LPAD","LTRIM","NLS_INITCAP","NLS_LOWER","NLSSORT","NLS_UPPER","REGEXP_REPLACE","REGEXP_SUBSTR","REPLACE","RPAD","RTRIM","SOUNDEX","SUBSTR","TRANSLATE","TREAT","TRIM","UPPER","NLS_CHARSET_DECL_LEN","NLS_CHARSET_ID","NLS_CHARSET_NAME","ASCII","INSTR","LENGTH","REGEXP_INSTR","ADD_MONTHS","CURRENT_DATE","CURRENT_TIMESTAMP","DBTIMEZONE","EXTRACT","FROM_TZ","LAST_DAY","LOCALTIMESTAMP","MONTHS_BETWEEN","NEW_TIME","NEXT_DAY","NUMTODSINTERVAL","NUMTOYMINTERVAL","ROUND","SESSIONTIMEZONE","SYS_EXTRACT_UTC","SYSDATE","SYSTIMESTAMP","TO_CHAR","TO_TIMESTAMP","TO_TIMESTAMP_TZ","TO_DSINTERVAL","TO_YMINTERVAL","TRUNC","TZ_OFFSET","GREATEST","LEAST","ASCIISTR","BIN_TO_NUM","CAST","CHARTOROWID","COMPOSE","CONVERT","DECOMPOSE","HEXTORAW","NUMTODSINTERVAL","NUMTOYMINTERVAL","RAWTOHEX","RAWTONHEX","ROWIDTOCHAR","ROWIDTONCHAR","SCN_TO_TIMESTAMP","TIMESTAMP_TO_SCN","TO_BINARY_DOUBLE","TO_BINARY_FLOAT","TO_CHAR","TO_CLOB","TO_DATE","TO_DSINTERVAL","TO_LOB","TO_MULTI_BYTE","TO_NCHAR","TO_NCLOB","TO_NUMBER","TO_DSINTERVAL","TO_SINGLE_BYTE","TO_TIMESTAMP","TO_TIMESTAMP_TZ","TO_YMINTERVAL","TO_YMINTERVAL","TRANSLATE","UNISTR","BFILENAME","EMPTY_BLOB,","EMPTY_CLOB","CARDINALITY","COLLECT","POWERMULTISET","POWERMULTISET_BY_CARDINALITY","SET","SYS_CONNECT_BY_PATH","CLUSTER_ID","CLUSTER_PROBABILITY","CLUSTER_SET","FEATURE_ID","FEATURE_SET","FEATURE_VALUE","PREDICTION","PREDICTION_COST","PREDICTION_DETAILS","PREDICTION_PROBABILITY","PREDICTION_SET","APPENDCHILDXML","DELETEXML","DEPTH","EXTRACT","EXISTSNODE","EXTRACTVALUE","INSERTCHILDXML","INSERTXMLBEFORE","PATH","SYS_DBURIGEN","SYS_XMLAGG","SYS_XMLGEN","UPDATEXML","XMLAGG","XMLCDATA","XMLCOLATTVAL","XMLCOMMENT","XMLCONCAT","XMLFOREST","XMLPARSE","XMLPI","XMLQUERY","XMLROOT","XMLSEQUENCE","XMLSERIALIZE","XMLTABLE","XMLTRANSFORM","DECODE","DUMP","ORA_HASH","VSIZE","COALESCE","LNNVL","NULLIF","NVL","NVL2","SYS_CONTEXT","SYS_GUID","SYS_TYPEID","UID","USER","USERENV","AVG","COLLECT","CORR","CORR_S","CORR_K","COUNT","COVAR_POP","COVAR_SAMP","CUME_DIST","DENSE_RANK","FIRST","GROUP_ID","GROUPING","GROUPING_ID","LAST","MAX","MEDIAN","MIN","PERCENTILE_CONT","PERCENTILE_DISC","PERCENT_RANK","RANK","REGR_SLOPE","REGR_INTERCEPT","REGR_COUNT","REGR_R2","REGR_AVGX","REGR_AVGY","REGR_SXX","REGR_SYY","REGR_SXY","STATS_BINOMIAL_TEST","STATS_CROSSTAB","STATS_F_TEST","STATS_KS_TEST","STATS_MODE","STATS_MW_TEST","STATS_ONE_WAY_ANOVA","STATS_T_TEST_ONE","STATS_T_TEST_PAIRED","STATS_T_TEST_INDEP","STATS_T_TEST_INDEPU","STATS_WSR_TEST","STDDEV","STDDEV_POP","STDDEV_SAMP","SUM","VAR_POP","VAR_SAMP","VARIANCE","FIRST_VALUE","LAG","LAST_VALUE","LEAD","NTILE","RATIO_TO_REPORT","ROW_NUMBER","DEREF","MAKE_REF","REF","REFTOHEX","VALUE","CV","ITERATION_NUMBER","PRESENTNNV","PRESENTV","PREVIOUS"],FS=v(["SELECT [ALL | DISTINCT | UNIQUE]"]),YS=v(["WITH","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER [SIBLINGS] BY","OFFSET","FETCH {FIRST | NEXT}","FOR UPDATE [OF]","INSERT [INTO | ALL INTO]","VALUES","SET","MERGE [INTO]","WHEN [NOT] MATCHED [THEN]","UPDATE SET","RETURNING"]),Or=v(["CREATE [GLOBAL TEMPORARY | PRIVATE TEMPORARY | SHARDED | DUPLICATED | IMMUTABLE BLOCKCHAIN | BLOCKCHAIN | IMMUTABLE] TABLE"]),Zt=v(["CREATE [OR REPLACE] [NO FORCE | FORCE] [EDITIONING | EDITIONABLE | EDITIONABLE EDITIONING | NONEDITIONABLE] VIEW","CREATE MATERIALIZED VIEW","UPDATE [ONLY]","DELETE FROM [ONLY]","DROP TABLE","ALTER TABLE","ADD","DROP {COLUMN | UNUSED COLUMNS | COLUMNS CONTINUE}","MODIFY","RENAME TO","RENAME COLUMN","TRUNCATE TABLE","SET SCHEMA","BEGIN","CONNECT BY","DECLARE","EXCEPT","EXCEPTION","LOOP","START WITH"]),VS=v(["UNION [ALL]","EXCEPT","INTERSECT"]),WS=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN","{CROSS | OUTER} APPLY"]),wS=v(["ON {UPDATE | DELETE} [SET NULL]","ON COMMIT","{ROWS | RANGE} BETWEEN"]),$S={name:"plsql",tokenizerOptions:{reservedSelect:FS,reservedClauses:[...YS,...Or,...Zt],reservedSetOperations:VS,reservedJoins:WS,reservedPhrases:wS,supportsXor:!0,reservedKeywords:yS,reservedDataTypes:BS,reservedFunctionNames:vS,stringTypes:[{quote:"''-qq",prefixes:["N"]},{quote:"q''",prefixes:["N"]}],identTypes:['""-qq'],identChars:{rest:"$#"},variableTypes:[{regex:"&{1,2}[A-Za-z][A-Za-z0-9_$#]*"}],paramTypes:{numbered:[":"],named:[":"]},paramChars:{},operators:["**",":=","%","~=","^=",">>","<<","=>","@","||"],postProcess:xS},formatOptions:{alwaysDenseOperators:["@"],onelineClauses:[...Or,...Zt],tabularOnelineClauses:Zt}};function xS(E){let e=tt;return E.map(T=>FE.SET(T)&&FE.BY(e)?AE(EE({},T),{type:"RESERVED_KEYWORD"}):(wR(T.type)&&(e=T),T))}var XS=["ABS","ACOS","ACOSD","ACOSH","ASIN","ASIND","ASINH","ATAN","ATAN2","ATAN2D","ATAND","ATANH","CBRT","CEIL","CEILING","COS","COSD","COSH","COT","COTD","DEGREES","DIV","EXP","FACTORIAL","FLOOR","GCD","LCM","LN","LOG","LOG10","MIN_SCALE","MOD","PI","POWER","RADIANS","RANDOM","ROUND","SCALE","SETSEED","SIGN","SIN","SIND","SINH","SQRT","TAN","TAND","TANH","TRIM_SCALE","TRUNC","WIDTH_BUCKET","ABS","ASCII","BIT_LENGTH","BTRIM","CHARACTER_LENGTH","CHAR_LENGTH","CHR","CONCAT","CONCAT_WS","FORMAT","INITCAP","LEFT","LENGTH","LOWER","LPAD","LTRIM","MD5","NORMALIZE","OCTET_LENGTH","OVERLAY","PARSE_IDENT","PG_CLIENT_ENCODING","POSITION","QUOTE_IDENT","QUOTE_LITERAL","QUOTE_NULLABLE","REGEXP_MATCH","REGEXP_MATCHES","REGEXP_REPLACE","REGEXP_SPLIT_TO_ARRAY","REGEXP_SPLIT_TO_TABLE","REPEAT","REPLACE","REVERSE","RIGHT","RPAD","RTRIM","SPLIT_PART","SPRINTF","STARTS_WITH","STRING_AGG","STRING_TO_ARRAY","STRING_TO_TABLE","STRPOS","SUBSTR","SUBSTRING","TO_ASCII","TO_HEX","TRANSLATE","TRIM","UNISTR","UPPER","BIT_COUNT","BIT_LENGTH","BTRIM","CONVERT","CONVERT_FROM","CONVERT_TO","DECODE","ENCODE","GET_BIT","GET_BYTE","LENGTH","LTRIM","MD5","OCTET_LENGTH","OVERLAY","POSITION","RTRIM","SET_BIT","SET_BYTE","SHA224","SHA256","SHA384","SHA512","STRING_AGG","SUBSTR","SUBSTRING","TRIM","BIT_COUNT","BIT_LENGTH","GET_BIT","LENGTH","OCTET_LENGTH","OVERLAY","POSITION","SET_BIT","SUBSTRING","REGEXP_MATCH","REGEXP_MATCHES","REGEXP_REPLACE","REGEXP_SPLIT_TO_ARRAY","REGEXP_SPLIT_TO_TABLE","TO_CHAR","TO_DATE","TO_NUMBER","TO_TIMESTAMP","CLOCK_TIMESTAMP","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","DATE_BIN","DATE_PART","DATE_TRUNC","EXTRACT","ISFINITE","JUSTIFY_DAYS","JUSTIFY_HOURS","JUSTIFY_INTERVAL","LOCALTIME","LOCALTIMESTAMP","MAKE_DATE","MAKE_INTERVAL","MAKE_TIME","MAKE_TIMESTAMP","MAKE_TIMESTAMPTZ","NOW","PG_SLEEP","PG_SLEEP_FOR","PG_SLEEP_UNTIL","STATEMENT_TIMESTAMP","TIMEOFDAY","TO_TIMESTAMP","TRANSACTION_TIMESTAMP","ENUM_FIRST","ENUM_LAST","ENUM_RANGE","AREA","BOUND_BOX","BOX","CENTER","CIRCLE","DIAGONAL","DIAMETER","HEIGHT","ISCLOSED","ISOPEN","LENGTH","LINE","LSEG","NPOINTS","PATH","PCLOSE","POINT","POLYGON","POPEN","RADIUS","SLOPE","WIDTH","ABBREV","BROADCAST","FAMILY","HOST","HOSTMASK","INET_MERGE","INET_SAME_FAMILY","MACADDR8_SET7BIT","MASKLEN","NETMASK","NETWORK","SET_MASKLEN","TRUNC","ARRAY_TO_TSVECTOR","GET_CURRENT_TS_CONFIG","JSONB_TO_TSVECTOR","JSON_TO_TSVECTOR","LENGTH","NUMNODE","PHRASETO_TSQUERY","PLAINTO_TSQUERY","QUERYTREE","SETWEIGHT","STRIP","TO_TSQUERY","TO_TSVECTOR","TSQUERY_PHRASE","TSVECTOR_TO_ARRAY","TS_DEBUG","TS_DELETE","TS_FILTER","TS_HEADLINE","TS_LEXIZE","TS_PARSE","TS_RANK","TS_RANK_CD","TS_REWRITE","TS_STAT","TS_TOKEN_TYPE","WEBSEARCH_TO_TSQUERY","UUID","CURSOR_TO_XML","CURSOR_TO_XMLSCHEMA","DATABASE_TO_XML","DATABASE_TO_XMLSCHEMA","DATABASE_TO_XML_AND_XMLSCHEMA","NEXTVAL","QUERY_TO_XML","QUERY_TO_XMLSCHEMA","QUERY_TO_XML_AND_XMLSCHEMA","SCHEMA_TO_XML","SCHEMA_TO_XMLSCHEMA","SCHEMA_TO_XML_AND_XMLSCHEMA","STRING","TABLE_TO_XML","TABLE_TO_XMLSCHEMA","TABLE_TO_XML_AND_XMLSCHEMA","XMLAGG","XMLCOMMENT","XMLCONCAT","XMLELEMENT","XMLEXISTS","XMLFOREST","XMLPARSE","XMLPI","XMLROOT","XMLSERIALIZE","XMLTABLE","XML_IS_WELL_FORMED","XML_IS_WELL_FORMED_CONTENT","XML_IS_WELL_FORMED_DOCUMENT","XPATH","XPATH_EXISTS","ARRAY_TO_JSON","JSONB_AGG","JSONB_ARRAY_ELEMENTS","JSONB_ARRAY_ELEMENTS_TEXT","JSONB_ARRAY_LENGTH","JSONB_BUILD_ARRAY","JSONB_BUILD_OBJECT","JSONB_EACH","JSONB_EACH_TEXT","JSONB_EXTRACT_PATH","JSONB_EXTRACT_PATH_TEXT","JSONB_INSERT","JSONB_OBJECT","JSONB_OBJECT_AGG","JSONB_OBJECT_KEYS","JSONB_PATH_EXISTS","JSONB_PATH_EXISTS_TZ","JSONB_PATH_MATCH","JSONB_PATH_MATCH_TZ","JSONB_PATH_QUERY","JSONB_PATH_QUERY_ARRAY","JSONB_PATH_QUERY_ARRAY_TZ","JSONB_PATH_QUERY_FIRST","JSONB_PATH_QUERY_FIRST_TZ","JSONB_PATH_QUERY_TZ","JSONB_POPULATE_RECORD","JSONB_POPULATE_RECORDSET","JSONB_PRETTY","JSONB_SET","JSONB_SET_LAX","JSONB_STRIP_NULLS","JSONB_TO_RECORD","JSONB_TO_RECORDSET","JSONB_TYPEOF","JSON_AGG","JSON_ARRAY_ELEMENTS","JSON_ARRAY_ELEMENTS_TEXT","JSON_ARRAY_LENGTH","JSON_BUILD_ARRAY","JSON_BUILD_OBJECT","JSON_EACH","JSON_EACH_TEXT","JSON_EXTRACT_PATH","JSON_EXTRACT_PATH_TEXT","JSON_OBJECT","JSON_OBJECT_AGG","JSON_OBJECT_KEYS","JSON_POPULATE_RECORD","JSON_POPULATE_RECORDSET","JSON_STRIP_NULLS","JSON_TO_RECORD","JSON_TO_RECORDSET","JSON_TYPEOF","ROW_TO_JSON","TO_JSON","TO_JSONB","TO_TIMESTAMP","CURRVAL","LASTVAL","NEXTVAL","SETVAL","COALESCE","GREATEST","LEAST","NULLIF","ARRAY_AGG","ARRAY_APPEND","ARRAY_CAT","ARRAY_DIMS","ARRAY_FILL","ARRAY_LENGTH","ARRAY_LOWER","ARRAY_NDIMS","ARRAY_POSITION","ARRAY_POSITIONS","ARRAY_PREPEND","ARRAY_REMOVE","ARRAY_REPLACE","ARRAY_TO_STRING","ARRAY_UPPER","CARDINALITY","STRING_TO_ARRAY","TRIM_ARRAY","UNNEST","ISEMPTY","LOWER","LOWER_INC","LOWER_INF","MULTIRANGE","RANGE_MERGE","UPPER","UPPER_INC","UPPER_INF","ARRAY_AGG","AVG","BIT_AND","BIT_OR","BIT_XOR","BOOL_AND","BOOL_OR","COALESCE","CORR","COUNT","COVAR_POP","COVAR_SAMP","CUME_DIST","DENSE_RANK","EVERY","GROUPING","JSONB_AGG","JSONB_OBJECT_AGG","JSON_AGG","JSON_OBJECT_AGG","MAX","MIN","MODE","PERCENTILE_CONT","PERCENTILE_DISC","PERCENT_RANK","RANGE_AGG","RANGE_INTERSECT_AGG","RANK","REGR_AVGX","REGR_AVGY","REGR_COUNT","REGR_INTERCEPT","REGR_R2","REGR_SLOPE","REGR_SXX","REGR_SXY","REGR_SYY","STDDEV","STDDEV_POP","STDDEV_SAMP","STRING_AGG","SUM","TO_JSON","TO_JSONB","VARIANCE","VAR_POP","VAR_SAMP","XMLAGG","CUME_DIST","DENSE_RANK","FIRST_VALUE","LAG","LAST_VALUE","LEAD","NTH_VALUE","NTILE","PERCENT_RANK","RANK","ROW_NUMBER","GENERATE_SERIES","GENERATE_SUBSCRIPTS","ACLDEFAULT","ACLEXPLODE","COL_DESCRIPTION","CURRENT_CATALOG","CURRENT_DATABASE","CURRENT_QUERY","CURRENT_ROLE","CURRENT_SCHEMA","CURRENT_SCHEMAS","CURRENT_USER","FORMAT_TYPE","HAS_ANY_COLUMN_PRIVILEGE","HAS_COLUMN_PRIVILEGE","HAS_DATABASE_PRIVILEGE","HAS_FOREIGN_DATA_WRAPPER_PRIVILEGE","HAS_FUNCTION_PRIVILEGE","HAS_LANGUAGE_PRIVILEGE","HAS_SCHEMA_PRIVILEGE","HAS_SEQUENCE_PRIVILEGE","HAS_SERVER_PRIVILEGE","HAS_TABLESPACE_PRIVILEGE","HAS_TABLE_PRIVILEGE","HAS_TYPE_PRIVILEGE","INET_CLIENT_ADDR","INET_CLIENT_PORT","INET_SERVER_ADDR","INET_SERVER_PORT","MAKEACLITEM","OBJ_DESCRIPTION","PG_BACKEND_PID","PG_BLOCKING_PIDS","PG_COLLATION_IS_VISIBLE","PG_CONF_LOAD_TIME","PG_CONTROL_CHECKPOINT","PG_CONTROL_INIT","PG_CONTROL_SYSTEM","PG_CONVERSION_IS_VISIBLE","PG_CURRENT_LOGFILE","PG_CURRENT_SNAPSHOT","PG_CURRENT_XACT_ID","PG_CURRENT_XACT_ID_IF_ASSIGNED","PG_DESCRIBE_OBJECT","PG_FUNCTION_IS_VISIBLE","PG_GET_CATALOG_FOREIGN_KEYS","PG_GET_CONSTRAINTDEF","PG_GET_EXPR","PG_GET_FUNCTIONDEF","PG_GET_FUNCTION_ARGUMENTS","PG_GET_FUNCTION_IDENTITY_ARGUMENTS","PG_GET_FUNCTION_RESULT","PG_GET_INDEXDEF","PG_GET_KEYWORDS","PG_GET_OBJECT_ADDRESS","PG_GET_OWNED_SEQUENCE","PG_GET_RULEDEF","PG_GET_SERIAL_SEQUENCE","PG_GET_STATISTICSOBJDEF","PG_GET_TRIGGERDEF","PG_GET_USERBYID","PG_GET_VIEWDEF","PG_HAS_ROLE","PG_IDENTIFY_OBJECT","PG_IDENTIFY_OBJECT_AS_ADDRESS","PG_INDEXAM_HAS_PROPERTY","PG_INDEX_COLUMN_HAS_PROPERTY","PG_INDEX_HAS_PROPERTY","PG_IS_OTHER_TEMP_SCHEMA","PG_JIT_AVAILABLE","PG_LAST_COMMITTED_XACT","PG_LISTENING_CHANNELS","PG_MY_TEMP_SCHEMA","PG_NOTIFICATION_QUEUE_USAGE","PG_OPCLASS_IS_VISIBLE","PG_OPERATOR_IS_VISIBLE","PG_OPFAMILY_IS_VISIBLE","PG_OPTIONS_TO_TABLE","PG_POSTMASTER_START_TIME","PG_SAFE_SNAPSHOT_BLOCKING_PIDS","PG_SNAPSHOT_XIP","PG_SNAPSHOT_XMAX","PG_SNAPSHOT_XMIN","PG_STATISTICS_OBJ_IS_VISIBLE","PG_TABLESPACE_DATABASES","PG_TABLESPACE_LOCATION","PG_TABLE_IS_VISIBLE","PG_TRIGGER_DEPTH","PG_TS_CONFIG_IS_VISIBLE","PG_TS_DICT_IS_VISIBLE","PG_TS_PARSER_IS_VISIBLE","PG_TS_TEMPLATE_IS_VISIBLE","PG_TYPEOF","PG_TYPE_IS_VISIBLE","PG_VISIBLE_IN_SNAPSHOT","PG_XACT_COMMIT_TIMESTAMP","PG_XACT_COMMIT_TIMESTAMP_ORIGIN","PG_XACT_STATUS","PQSERVERVERSION","ROW_SECURITY_ACTIVE","SESSION_USER","SHOBJ_DESCRIPTION","TO_REGCLASS","TO_REGCOLLATION","TO_REGNAMESPACE","TO_REGOPER","TO_REGOPERATOR","TO_REGPROC","TO_REGPROCEDURE","TO_REGROLE","TO_REGTYPE","TXID_CURRENT","TXID_CURRENT_IF_ASSIGNED","TXID_CURRENT_SNAPSHOT","TXID_SNAPSHOT_XIP","TXID_SNAPSHOT_XMAX","TXID_SNAPSHOT_XMIN","TXID_STATUS","TXID_VISIBLE_IN_SNAPSHOT","USER","VERSION","BRIN_DESUMMARIZE_RANGE","BRIN_SUMMARIZE_NEW_VALUES","BRIN_SUMMARIZE_RANGE","CONVERT_FROM","CURRENT_SETTING","GIN_CLEAN_PENDING_LIST","PG_ADVISORY_LOCK","PG_ADVISORY_LOCK_SHARED","PG_ADVISORY_UNLOCK","PG_ADVISORY_UNLOCK_ALL","PG_ADVISORY_UNLOCK_SHARED","PG_ADVISORY_XACT_LOCK","PG_ADVISORY_XACT_LOCK_SHARED","PG_BACKUP_START_TIME","PG_CANCEL_BACKEND","PG_COLLATION_ACTUAL_VERSION","PG_COLUMN_COMPRESSION","PG_COLUMN_SIZE","PG_COPY_LOGICAL_REPLICATION_SLOT","PG_COPY_PHYSICAL_REPLICATION_SLOT","PG_CREATE_LOGICAL_REPLICATION_SLOT","PG_CREATE_PHYSICAL_REPLICATION_SLOT","PG_CREATE_RESTORE_POINT","PG_CURRENT_WAL_FLUSH_LSN","PG_CURRENT_WAL_INSERT_LSN","PG_CURRENT_WAL_LSN","PG_DATABASE_SIZE","PG_DROP_REPLICATION_SLOT","PG_EXPORT_SNAPSHOT","PG_FILENODE_RELATION","PG_GET_WAL_REPLAY_PAUSE_STATE","PG_IMPORT_SYSTEM_COLLATIONS","PG_INDEXES_SIZE","PG_IS_IN_BACKUP","PG_IS_IN_RECOVERY","PG_IS_WAL_REPLAY_PAUSED","PG_LAST_WAL_RECEIVE_LSN","PG_LAST_WAL_REPLAY_LSN","PG_LAST_XACT_REPLAY_TIMESTAMP","PG_LOGICAL_EMIT_MESSAGE","PG_LOGICAL_SLOT_GET_BINARY_CHANGES","PG_LOGICAL_SLOT_GET_CHANGES","PG_LOGICAL_SLOT_PEEK_BINARY_CHANGES","PG_LOGICAL_SLOT_PEEK_CHANGES","PG_LOG_BACKEND_MEMORY_CONTEXTS","PG_LS_ARCHIVE_STATUSDIR","PG_LS_DIR","PG_LS_LOGDIR","PG_LS_TMPDIR","PG_LS_WALDIR","PG_PARTITION_ANCESTORS","PG_PARTITION_ROOT","PG_PARTITION_TREE","PG_PROMOTE","PG_READ_BINARY_FILE","PG_READ_FILE","PG_RELATION_FILENODE","PG_RELATION_FILEPATH","PG_RELATION_SIZE","PG_RELOAD_CONF","PG_REPLICATION_ORIGIN_ADVANCE","PG_REPLICATION_ORIGIN_CREATE","PG_REPLICATION_ORIGIN_DROP","PG_REPLICATION_ORIGIN_OID","PG_REPLICATION_ORIGIN_PROGRESS","PG_REPLICATION_ORIGIN_SESSION_IS_SETUP","PG_REPLICATION_ORIGIN_SESSION_PROGRESS","PG_REPLICATION_ORIGIN_SESSION_RESET","PG_REPLICATION_ORIGIN_SESSION_SETUP","PG_REPLICATION_ORIGIN_XACT_RESET","PG_REPLICATION_ORIGIN_XACT_SETUP","PG_REPLICATION_SLOT_ADVANCE","PG_ROTATE_LOGFILE","PG_SIZE_BYTES","PG_SIZE_PRETTY","PG_START_BACKUP","PG_STAT_FILE","PG_STOP_BACKUP","PG_SWITCH_WAL","PG_TABLESPACE_SIZE","PG_TABLE_SIZE","PG_TERMINATE_BACKEND","PG_TOTAL_RELATION_SIZE","PG_TRY_ADVISORY_LOCK","PG_TRY_ADVISORY_LOCK_SHARED","PG_TRY_ADVISORY_XACT_LOCK","PG_TRY_ADVISORY_XACT_LOCK_SHARED","PG_WALFILE_NAME","PG_WALFILE_NAME_OFFSET","PG_WAL_LSN_DIFF","PG_WAL_REPLAY_PAUSE","PG_WAL_REPLAY_RESUME","SET_CONFIG","SUPPRESS_REDUNDANT_UPDATES_TRIGGER","TSVECTOR_UPDATE_TRIGGER","TSVECTOR_UPDATE_TRIGGER_COLUMN","PG_EVENT_TRIGGER_DDL_COMMANDS","PG_EVENT_TRIGGER_DROPPED_OBJECTS","PG_EVENT_TRIGGER_TABLE_REWRITE_OID","PG_EVENT_TRIGGER_TABLE_REWRITE_REASON","PG_GET_OBJECT_ADDRESS","PG_MCV_LIST_ITEMS","CAST"],kS=["ALL","ANALYSE","ANALYZE","AND","ANY","AS","ASC","ASYMMETRIC","AUTHORIZATION","BETWEEN","BINARY","BOTH","CASE","CAST","CHECK","COLLATE","COLLATION","COLUMN","CONCURRENTLY","CONSTRAINT","CREATE","CROSS","CURRENT_CATALOG","CURRENT_DATE","CURRENT_ROLE","CURRENT_SCHEMA","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","DAY","DEFAULT","DEFERRABLE","DESC","DISTINCT","DO","ELSE","END","EXCEPT","EXISTS","FALSE","FETCH","FILTER","FOR","FOREIGN","FREEZE","FROM","FULL","GRANT","GROUP","HAVING","HOUR","ILIKE","IN","INITIALLY","INNER","INOUT","INTERSECT","INTO","IS","ISNULL","JOIN","LATERAL","LEADING","LEFT","LIKE","LIMIT","LOCALTIME","LOCALTIMESTAMP","MINUTE","MONTH","NATURAL","NOT","NOTNULL","NULL","NULLIF","OFFSET","ON","ONLY","OR","ORDER","OUT","OUTER","OVER","OVERLAPS","PLACING","PRIMARY","REFERENCES","RETURNING","RIGHT","ROW","SECOND","SELECT","SESSION_USER","SIMILAR","SOME","SYMMETRIC","TABLE","TABLESAMPLE","THEN","TO","TRAILING","TRUE","UNION","UNIQUE","USER","USING","VALUES","VARIADIC","VERBOSE","WHEN","WHERE","WINDOW","WITH","WITHIN","WITHOUT","YEAR"],KS=["ARRAY","BIGINT","BIT","BIT VARYING","BOOL","BOOLEAN","CHAR","CHARACTER","CHARACTER VARYING","DECIMAL","DEC","DOUBLE","ENUM","FLOAT","INT","INTEGER","INTERVAL","NCHAR","NUMERIC","PRECISION","REAL","SMALLINT","TEXT","TIME","TIMESTAMP","TIMESTAMPTZ","VARCHAR","XML","ZONE"],JS=v(["SELECT [ALL | DISTINCT]"]),qS=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY [ALL | DISTINCT]","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","FETCH {FIRST | NEXT}","FOR {UPDATE | NO KEY UPDATE | SHARE | KEY SHARE} [OF]","INSERT INTO","VALUES","DEFAULT VALUES","SET","RETURNING"]),ir=v(["CREATE [GLOBAL | LOCAL] [TEMPORARY | TEMP | UNLOGGED] TABLE [IF NOT EXISTS]"]),jt=v(["CREATE [OR REPLACE] [TEMP | TEMPORARY] [RECURSIVE] VIEW","CREATE [MATERIALIZED] VIEW [IF NOT EXISTS]","UPDATE [ONLY]","WHERE CURRENT OF","ON CONFLICT","DELETE FROM [ONLY]","DROP TABLE [IF EXISTS]","ALTER TABLE [IF EXISTS] [ONLY]","ALTER TABLE ALL IN TABLESPACE","RENAME [COLUMN]","RENAME TO","ADD [COLUMN] [IF NOT EXISTS]","DROP [COLUMN] [IF EXISTS]","ALTER [COLUMN]","SET DATA TYPE","{SET | DROP} DEFAULT","{SET | DROP} NOT NULL","TRUNCATE [TABLE] [ONLY]","SET SCHEMA","AFTER","ABORT","ALTER AGGREGATE","ALTER COLLATION","ALTER CONVERSION","ALTER DATABASE","ALTER DEFAULT PRIVILEGES","ALTER DOMAIN","ALTER EVENT TRIGGER","ALTER EXTENSION","ALTER FOREIGN DATA WRAPPER","ALTER FOREIGN TABLE","ALTER FUNCTION","ALTER GROUP","ALTER INDEX","ALTER LANGUAGE","ALTER LARGE OBJECT","ALTER MATERIALIZED VIEW","ALTER OPERATOR","ALTER OPERATOR CLASS","ALTER OPERATOR FAMILY","ALTER POLICY","ALTER PROCEDURE","ALTER PUBLICATION","ALTER ROLE","ALTER ROUTINE","ALTER RULE","ALTER SCHEMA","ALTER SEQUENCE","ALTER SERVER","ALTER STATISTICS","ALTER SUBSCRIPTION","ALTER SYSTEM","ALTER TABLESPACE","ALTER TEXT SEARCH CONFIGURATION","ALTER TEXT SEARCH DICTIONARY","ALTER TEXT SEARCH PARSER","ALTER TEXT SEARCH TEMPLATE","ALTER TRIGGER","ALTER TYPE","ALTER USER","ALTER USER MAPPING","ALTER VIEW","ANALYZE","BEGIN","CALL","CHECKPOINT","CLOSE","CLUSTER","COMMIT","COMMIT PREPARED","COPY","CREATE ACCESS METHOD","CREATE AGGREGATE","CREATE CAST","CREATE COLLATION","CREATE CONVERSION","CREATE DATABASE","CREATE DOMAIN","CREATE EVENT TRIGGER","CREATE EXTENSION","CREATE FOREIGN DATA WRAPPER","CREATE FOREIGN TABLE","CREATE FUNCTION","CREATE GROUP","CREATE INDEX","CREATE LANGUAGE","CREATE OPERATOR","CREATE OPERATOR CLASS","CREATE OPERATOR FAMILY","CREATE POLICY","CREATE PROCEDURE","CREATE PUBLICATION","CREATE ROLE","CREATE RULE","CREATE SCHEMA","CREATE SEQUENCE","CREATE SERVER","CREATE STATISTICS","CREATE SUBSCRIPTION","CREATE TABLESPACE","CREATE TEXT SEARCH CONFIGURATION","CREATE TEXT SEARCH DICTIONARY","CREATE TEXT SEARCH PARSER","CREATE TEXT SEARCH TEMPLATE","CREATE TRANSFORM","CREATE TRIGGER","CREATE TYPE","CREATE USER","CREATE USER MAPPING","DEALLOCATE","DECLARE","DISCARD","DROP ACCESS METHOD","DROP AGGREGATE","DROP CAST","DROP COLLATION","DROP CONVERSION","DROP DATABASE","DROP DOMAIN","DROP EVENT TRIGGER","DROP EXTENSION","DROP FOREIGN DATA WRAPPER","DROP FOREIGN TABLE","DROP FUNCTION","DROP GROUP","DROP INDEX","DROP LANGUAGE","DROP MATERIALIZED VIEW","DROP OPERATOR","DROP OPERATOR CLASS","DROP OPERATOR FAMILY","DROP OWNED","DROP POLICY","DROP PROCEDURE","DROP PUBLICATION","DROP ROLE","DROP ROUTINE","DROP RULE","DROP SCHEMA","DROP SEQUENCE","DROP SERVER","DROP STATISTICS","DROP SUBSCRIPTION","DROP TABLESPACE","DROP TEXT SEARCH CONFIGURATION","DROP TEXT SEARCH DICTIONARY","DROP TEXT SEARCH PARSER","DROP TEXT SEARCH TEMPLATE","DROP TRANSFORM","DROP TRIGGER","DROP TYPE","DROP USER","DROP USER MAPPING","DROP VIEW","EXECUTE","EXPLAIN","FETCH","GRANT","IMPORT FOREIGN SCHEMA","LISTEN","LOAD","LOCK","MOVE","NOTIFY","PREPARE","PREPARE TRANSACTION","REASSIGN OWNED","REFRESH MATERIALIZED VIEW","REINDEX","RELEASE SAVEPOINT","RESET","REVOKE","ROLLBACK","ROLLBACK PREPARED","ROLLBACK TO SAVEPOINT","SAVEPOINT","SECURITY LABEL","SELECT INTO","SET CONSTRAINTS","SET ROLE","SET SESSION AUTHORIZATION","SET TRANSACTION","SHOW","START TRANSACTION","UNLISTEN","VACUUM"]),QS=v(["UNION [ALL | DISTINCT]","EXCEPT [ALL | DISTINCT]","INTERSECT [ALL | DISTINCT]"]),ZS=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN"]),jS=v(["PRIMARY KEY","GENERATED {ALWAYS | BY DEFAULT} AS IDENTITY","ON {UPDATE | DELETE} [SET NULL | SET DEFAULT]","{ROWS | RANGE | GROUPS} BETWEEN","[TIMESTAMP | TIME] {WITH | WITHOUT} TIME ZONE","IS [NOT] DISTINCT FROM"]),zS={name:"postgresql",tokenizerOptions:{reservedSelect:JS,reservedClauses:[...qS,...ir,...jt],reservedSetOperations:QS,reservedJoins:ZS,reservedPhrases:jS,reservedKeywords:kS,reservedDataTypes:KS,reservedFunctionNames:XS,nestedBlockComments:!0,extraParens:["[]"],stringTypes:["$$",{quote:"''-qq",prefixes:["U&"]},{quote:"''-qq-bs",prefixes:["E"],requirePrefix:!0},{quote:"''-raw",prefixes:["B","X"],requirePrefix:!0}],identTypes:[{quote:'""-qq',prefixes:["U&"]}],identChars:{rest:"$"},paramTypes:{numbered:["$"]},operators:["%","^","|/","||/","@",":=","&","|","#","~","<<",">>","~>~","~<~","~>=~","~<=~","@-@","@@","##","<->","&&","&<","&>","<<|","&<|","|>>","|&>","<^","^>","?#","?-","?|","?-|","?||","@>","<@","~=","?","@?","?&","->","->>","#>","#>>","#-","=>",">>=","<<=","~~","~~*","!~~","!~~*","~","~*","!~","!~*","-|-","||","@@@","!!","^@","<%","%>","<<%","%>>","<<->","<->>","<<<->","<->>>","::",":"]},formatOptions:{alwaysDenseOperators:["::",":"],onelineClauses:[...ir,...jt],tabularOnelineClauses:jt}},eo=["ANY_VALUE","APPROXIMATE PERCENTILE_DISC","AVG","COUNT","LISTAGG","MAX","MEDIAN","MIN","PERCENTILE_CONT","STDDEV_SAMP","STDDEV_POP","SUM","VAR_SAMP","VAR_POP","array","array_concat","array_flatten","get_array_length","split_to_array","subarray","BIT_AND","BIT_OR","BOOL_AND","BOOL_OR","COALESCE","DECODE","GREATEST","LEAST","NVL","NVL2","NULLIF","ADD_MONTHS","AT TIME ZONE","CONVERT_TIMEZONE","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","DATE_CMP","DATE_CMP_TIMESTAMP","DATE_CMP_TIMESTAMPTZ","DATE_PART_YEAR","DATEADD","DATEDIFF","DATE_PART","DATE_TRUNC","EXTRACT","GETDATE","INTERVAL_CMP","LAST_DAY","MONTHS_BETWEEN","NEXT_DAY","SYSDATE","TIMEOFDAY","TIMESTAMP_CMP","TIMESTAMP_CMP_DATE","TIMESTAMP_CMP_TIMESTAMPTZ","TIMESTAMPTZ_CMP","TIMESTAMPTZ_CMP_DATE","TIMESTAMPTZ_CMP_TIMESTAMP","TIMEZONE","TO_TIMESTAMP","TRUNC","AddBBox","DropBBox","GeometryType","ST_AddPoint","ST_Angle","ST_Area","ST_AsBinary","ST_AsEWKB","ST_AsEWKT","ST_AsGeoJSON","ST_AsText","ST_Azimuth","ST_Boundary","ST_Collect","ST_Contains","ST_ContainsProperly","ST_ConvexHull","ST_CoveredBy","ST_Covers","ST_Crosses","ST_Dimension","ST_Disjoint","ST_Distance","ST_DistanceSphere","ST_DWithin","ST_EndPoint","ST_Envelope","ST_Equals","ST_ExteriorRing","ST_Force2D","ST_Force3D","ST_Force3DM","ST_Force3DZ","ST_Force4D","ST_GeometryN","ST_GeometryType","ST_GeomFromEWKB","ST_GeomFromEWKT","ST_GeomFromText","ST_GeomFromWKB","ST_InteriorRingN","ST_Intersects","ST_IsPolygonCCW","ST_IsPolygonCW","ST_IsClosed","ST_IsCollection","ST_IsEmpty","ST_IsSimple","ST_IsValid","ST_Length","ST_LengthSphere","ST_Length2D","ST_LineFromMultiPoint","ST_LineInterpolatePoint","ST_M","ST_MakeEnvelope","ST_MakeLine","ST_MakePoint","ST_MakePolygon","ST_MemSize","ST_MMax","ST_MMin","ST_Multi","ST_NDims","ST_NPoints","ST_NRings","ST_NumGeometries","ST_NumInteriorRings","ST_NumPoints","ST_Perimeter","ST_Perimeter2D","ST_Point","ST_PointN","ST_Points","ST_Polygon","ST_RemovePoint","ST_Reverse","ST_SetPoint","ST_SetSRID","ST_Simplify","ST_SRID","ST_StartPoint","ST_Touches","ST_Within","ST_X","ST_XMax","ST_XMin","ST_Y","ST_YMax","ST_YMin","ST_Z","ST_ZMax","ST_ZMin","SupportsBBox","CHECKSUM","FUNC_SHA1","FNV_HASH","MD5","SHA","SHA1","SHA2","HLL","HLL_CREATE_SKETCH","HLL_CARDINALITY","HLL_COMBINE","IS_VALID_JSON","IS_VALID_JSON_ARRAY","JSON_ARRAY_LENGTH","JSON_EXTRACT_ARRAY_ELEMENT_TEXT","JSON_EXTRACT_PATH_TEXT","JSON_PARSE","JSON_SERIALIZE","ABS","ACOS","ASIN","ATAN","ATAN2","CBRT","CEILING","CEIL","COS","COT","DEGREES","DEXP","DLOG1","DLOG10","EXP","FLOOR","LN","LOG","MOD","PI","POWER","RADIANS","RANDOM","ROUND","SIN","SIGN","SQRT","TAN","TO_HEX","TRUNC","EXPLAIN_MODEL","ASCII","BPCHARCMP","BTRIM","BTTEXT_PATTERN_CMP","CHAR_LENGTH","CHARACTER_LENGTH","CHARINDEX","CHR","COLLATE","CONCAT","CRC32","DIFFERENCE","INITCAP","LEFT","RIGHT","LEN","LENGTH","LOWER","LPAD","RPAD","LTRIM","OCTETINDEX","OCTET_LENGTH","POSITION","QUOTE_IDENT","QUOTE_LITERAL","REGEXP_COUNT","REGEXP_INSTR","REGEXP_REPLACE","REGEXP_SUBSTR","REPEAT","REPLACE","REPLICATE","REVERSE","RTRIM","SOUNDEX","SPLIT_PART","STRPOS","STRTOL","SUBSTRING","TEXTLEN","TRANSLATE","TRIM","UPPER","decimal_precision","decimal_scale","is_array","is_bigint","is_boolean","is_char","is_decimal","is_float","is_integer","is_object","is_scalar","is_smallint","is_varchar","json_typeof","AVG","COUNT","CUME_DIST","DENSE_RANK","FIRST_VALUE","LAST_VALUE","LAG","LEAD","LISTAGG","MAX","MEDIAN","MIN","NTH_VALUE","NTILE","PERCENT_RANK","PERCENTILE_CONT","PERCENTILE_DISC","RANK","RATIO_TO_REPORT","ROW_NUMBER","STDDEV_SAMP","STDDEV_POP","SUM","VAR_SAMP","VAR_POP","CAST","CONVERT","TO_CHAR","TO_DATE","TO_NUMBER","TEXT_TO_INT_ALT","TEXT_TO_NUMERIC_ALT","CHANGE_QUERY_PRIORITY","CHANGE_SESSION_PRIORITY","CHANGE_USER_PRIORITY","CURRENT_SETTING","PG_CANCEL_BACKEND","PG_TERMINATE_BACKEND","REBOOT_CLUSTER","SET_CONFIG","CURRENT_AWS_ACCOUNT","CURRENT_DATABASE","CURRENT_NAMESPACE","CURRENT_SCHEMA","CURRENT_SCHEMAS","CURRENT_USER","CURRENT_USER_ID","HAS_ASSUMEROLE_PRIVILEGE","HAS_DATABASE_PRIVILEGE","HAS_SCHEMA_PRIVILEGE","HAS_TABLE_PRIVILEGE","PG_BACKEND_PID","PG_GET_COLS","PG_GET_GRANTEE_BY_IAM_ROLE","PG_GET_IAM_ROLE_BY_USER","PG_GET_LATE_BINDING_VIEW_COLS","PG_LAST_COPY_COUNT","PG_LAST_COPY_ID","PG_LAST_UNLOAD_ID","PG_LAST_QUERY_ID","PG_LAST_UNLOAD_COUNT","SESSION_USER","SLICE_NUM","USER","VERSION"],Eo=["AES128","AES256","ALL","ALLOWOVERWRITE","ANY","AS","ASC","AUTHORIZATION","BACKUP","BETWEEN","BINARY","BOTH","CHECK","COLUMN","CONSTRAINT","CREATE","CROSS","DEFAULT","DEFERRABLE","DEFLATE","DEFRAG","DESC","DISABLE","DISTINCT","DO","ENABLE","ENCODE","ENCRYPT","ENCRYPTION","EXPLICIT","FALSE","FOR","FOREIGN","FREEZE","FROM","FULL","GLOBALDICT256","GLOBALDICT64K","GROUP","IDENTITY","IGNORE","ILIKE","IN","INITIALLY","INNER","INTO","IS","ISNULL","LANGUAGE","LEADING","LIKE","LIMIT","LOCALTIME","LOCALTIMESTAMP","LUN","LUNS","MINUS","NATURAL","NEW","NOT","NOTNULL","NULL","NULLS","OFF","OFFLINE","OFFSET","OID","OLD","ON","ONLY","OPEN","ORDER","OUTER","OVERLAPS","PARALLEL","PARTITION","PERCENT","PERMISSIONS","PLACING","PRIMARY","RECOVER","REFERENCES","REJECTLOG","RESORT","RESPECT","RESTORE","SIMILAR","SNAPSHOT","SOME","SYSTEM","TABLE","TAG","TDES","THEN","TIMESTAMP","TO","TOP","TRAILING","TRUE","UNIQUE","USING","VERBOSE","WALLET","WITHOUT","ACCEPTANYDATE","ACCEPTINVCHARS","BLANKSASNULL","DATEFORMAT","EMPTYASNULL","ENCODING","ESCAPE","EXPLICIT_IDS","FILLRECORD","IGNOREBLANKLINES","IGNOREHEADER","REMOVEQUOTES","ROUNDEC","TIMEFORMAT","TRIMBLANKS","TRUNCATECOLUMNS","COMPROWS","COMPUPDATE","MAXERROR","NOLOAD","STATUPDATE","FORMAT","CSV","DELIMITER","FIXEDWIDTH","SHAPEFILE","AVRO","JSON","PARQUET","ORC","ACCESS_KEY_ID","CREDENTIALS","ENCRYPTED","IAM_ROLE","MASTER_SYMMETRIC_KEY","SECRET_ACCESS_KEY","SESSION_TOKEN","BZIP2","GZIP","LZOP","ZSTD","MANIFEST","READRATIO","REGION","SSH","RAW","AZ64","BYTEDICT","DELTA","DELTA32K","LZO","MOSTLY8","MOSTLY16","MOSTLY32","RUNLENGTH","TEXT255","TEXT32K","CATALOG_ROLE","SECRET_ARN","EXTERNAL","AUTO","EVEN","KEY","PREDICATE","COMPRESSION"],to=["ARRAY","BIGINT","BPCHAR","CHAR","CHARACTER VARYING","CHARACTER","DECIMAL","INT","INT2","INT4","INT8","INTEGER","NCHAR","NUMERIC","NVARCHAR","SMALLINT","TEXT","VARBYTE","VARCHAR"],To=v(["SELECT [ALL | DISTINCT]"]),ro=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER BY","LIMIT","OFFSET","INSERT INTO","VALUES","SET"]),ar=v(["CREATE [TEMPORARY | TEMP | LOCAL TEMPORARY | LOCAL TEMP] TABLE [IF NOT EXISTS]"]),zt=v(["CREATE [OR REPLACE | MATERIALIZED] VIEW","UPDATE","DELETE [FROM]","DROP TABLE [IF EXISTS]","ALTER TABLE","ALTER TABLE APPEND","ADD [COLUMN]","DROP [COLUMN]","RENAME TO","RENAME COLUMN","ALTER COLUMN","TYPE","ENCODE","TRUNCATE [TABLE]","ABORT","ALTER DATABASE","ALTER DATASHARE","ALTER DEFAULT PRIVILEGES","ALTER GROUP","ALTER MATERIALIZED VIEW","ALTER PROCEDURE","ALTER SCHEMA","ALTER USER","ANALYSE","ANALYZE","ANALYSE COMPRESSION","ANALYZE COMPRESSION","BEGIN","CALL","CANCEL","CLOSE","COMMIT","COPY","CREATE DATABASE","CREATE DATASHARE","CREATE EXTERNAL FUNCTION","CREATE EXTERNAL SCHEMA","CREATE EXTERNAL TABLE","CREATE FUNCTION","CREATE GROUP","CREATE LIBRARY","CREATE MODEL","CREATE PROCEDURE","CREATE SCHEMA","CREATE USER","DEALLOCATE","DECLARE","DESC DATASHARE","DROP DATABASE","DROP DATASHARE","DROP FUNCTION","DROP GROUP","DROP LIBRARY","DROP MODEL","DROP MATERIALIZED VIEW","DROP PROCEDURE","DROP SCHEMA","DROP USER","DROP VIEW","DROP","EXECUTE","EXPLAIN","FETCH","GRANT","LOCK","PREPARE","REFRESH MATERIALIZED VIEW","RESET","REVOKE","ROLLBACK","SELECT INTO","SET SESSION AUTHORIZATION","SET SESSION CHARACTERISTICS","SHOW","SHOW EXTERNAL TABLE","SHOW MODEL","SHOW DATASHARES","SHOW PROCEDURE","SHOW TABLE","SHOW VIEW","START TRANSACTION","UNLOAD","VACUUM"]),Ro=v(["UNION [ALL]","EXCEPT","INTERSECT","MINUS"]),no=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN"]),Ao=v(["NULL AS","DATA CATALOG","HIVE METASTORE","{ROWS | RANGE} BETWEEN"]),so={name:"redshift",tokenizerOptions:{reservedSelect:To,reservedClauses:[...ro,...ar,...zt],reservedSetOperations:Ro,reservedJoins:no,reservedPhrases:Ao,reservedKeywords:Eo,reservedDataTypes:to,reservedFunctionNames:eo,stringTypes:["''-qq"],identTypes:['""-qq'],identChars:{first:"#"},paramTypes:{numbered:["$"]},operators:["^","%","@","|/","||/","&","|","~","<<",">>","||","::"]},formatOptions:{alwaysDenseOperators:["::"],onelineClauses:[...ar,...zt],tabularOnelineClauses:zt}},So=["ADD","AFTER","ALL","ALTER","ANALYZE","AND","ANTI","ANY","ARCHIVE","AS","ASC","AT","AUTHORIZATION","BETWEEN","BOTH","BUCKET","BUCKETS","BY","CACHE","CASCADE","CAST","CHANGE","CHECK","CLEAR","CLUSTER","CLUSTERED","CODEGEN","COLLATE","COLLECTION","COLUMN","COLUMNS","COMMENT","COMMIT","COMPACT","COMPACTIONS","COMPUTE","CONCATENATE","CONSTRAINT","COST","CREATE","CROSS","CUBE","CURRENT","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","DATA","DATABASE","DATABASES","DAY","DBPROPERTIES","DEFINED","DELETE","DELIMITED","DESC","DESCRIBE","DFS","DIRECTORIES","DIRECTORY","DISTINCT","DISTRIBUTE","DIV","DROP","ESCAPE","ESCAPED","EXCEPT","EXCHANGE","EXISTS","EXPORT","EXTENDED","EXTERNAL","EXTRACT","FALSE","FETCH","FIELDS","FILTER","FILEFORMAT","FIRST","FIRST_VALUE","FOLLOWING","FOR","FOREIGN","FORMAT","FORMATTED","FULL","FUNCTION","FUNCTIONS","GLOBAL","GRANT","GROUP","GROUPING","HOUR","IF","IGNORE","IMPORT","IN","INDEX","INDEXES","INNER","INPATH","INPUTFORMAT","INTERSECT","INTO","IS","ITEMS","KEYS","LAST","LAST_VALUE","LATERAL","LAZY","LEADING","LEFT","LIKE","LINES","LIST","LOCAL","LOCATION","LOCK","LOCKS","LOGICAL","MACRO","MATCHED","MERGE","MINUTE","MONTH","MSCK","NAMESPACE","NAMESPACES","NATURAL","NO","NOT","NULL","NULLS","OF","ONLY","OPTION","OPTIONS","OR","ORDER","OUT","OUTER","OUTPUTFORMAT","OVER","OVERLAPS","OVERLAY","OVERWRITE","OWNER","PARTITION","PARTITIONED","PARTITIONS","PERCENT","PLACING","POSITION","PRECEDING","PRIMARY","PRINCIPALS","PROPERTIES","PURGE","QUERY","RANGE","RECORDREADER","RECORDWRITER","RECOVER","REDUCE","REFERENCES","RENAME","REPAIR","REPLACE","RESPECT","RESTRICT","REVOKE","RIGHT","RLIKE","ROLE","ROLES","ROLLBACK","ROLLUP","ROW","ROWS","SCHEMA","SECOND","SELECT","SEMI","SEPARATED","SERDE","SERDEPROPERTIES","SESSION_USER","SETS","SHOW","SKEWED","SOME","SORT","SORTED","START","STATISTICS","STORED","STRATIFY","SUBSTR","SUBSTRING","TABLE","TABLES","TBLPROPERTIES","TEMPORARY","TERMINATED","THEN","TO","TOUCH","TRAILING","TRANSACTION","TRANSACTIONS","TRIM","TRUE","TRUNCATE","UNARCHIVE","UNBOUNDED","UNCACHE","UNIQUE","UNKNOWN","UNLOCK","UNSET","USE","USER","USING","VIEW","WINDOW","YEAR","ANALYSE","ARRAY_ZIP","COALESCE","CONTAINS","CONVERT","DAYS","DAY_HOUR","DAY_MINUTE","DAY_SECOND","DECODE","DEFAULT","DISTINCTROW","ENCODE","EXPLODE","EXPLODE_OUTER","FIXED","GREATEST","GROUP_CONCAT","HOURS","HOUR_MINUTE","HOUR_SECOND","IFNULL","LEAST","LEVEL","MINUTE_SECOND","NULLIF","OFFSET","ON","OPTIMIZE","REGEXP","SEPARATOR","SIZE","TYPE","TYPES","UNSIGNED","VARIABLES","YEAR_MONTH"],oo=["ARRAY","BIGINT","BINARY","BOOLEAN","BYTE","CHAR","DATE","DEC","DECIMAL","DOUBLE","FLOAT","INT","INTEGER","INTERVAL","LONG","MAP","NUMERIC","REAL","SHORT","SMALLINT","STRING","STRUCT","TIMESTAMP_LTZ","TIMESTAMP_NTZ","TIMESTAMP","TINYINT","VARCHAR"],Oo=["APPROX_COUNT_DISTINCT","APPROX_PERCENTILE","AVG","BIT_AND","BIT_OR","BIT_XOR","BOOL_AND","BOOL_OR","COLLECT_LIST","COLLECT_SET","CORR","COUNT","COUNT","COUNT","COUNT_IF","COUNT_MIN_SKETCH","COVAR_POP","COVAR_SAMP","EVERY","FIRST","FIRST_VALUE","GROUPING","GROUPING_ID","KURTOSIS","LAST","LAST_VALUE","MAX","MAX_BY","MEAN","MIN","MIN_BY","PERCENTILE","PERCENTILE","PERCENTILE_APPROX","SKEWNESS","STD","STDDEV","STDDEV_POP","STDDEV_SAMP","SUM","VAR_POP","VAR_SAMP","VARIANCE","CUME_DIST","DENSE_RANK","LAG","LEAD","NTH_VALUE","NTILE","PERCENT_RANK","RANK","ROW_NUMBER","ARRAY","ARRAY_CONTAINS","ARRAY_DISTINCT","ARRAY_EXCEPT","ARRAY_INTERSECT","ARRAY_JOIN","ARRAY_MAX","ARRAY_MIN","ARRAY_POSITION","ARRAY_REMOVE","ARRAY_REPEAT","ARRAY_UNION","ARRAYS_OVERLAP","ARRAYS_ZIP","FLATTEN","SEQUENCE","SHUFFLE","SLICE","SORT_ARRAY","ELEMENT_AT","ELEMENT_AT","MAP_CONCAT","MAP_ENTRIES","MAP_FROM_ARRAYS","MAP_FROM_ENTRIES","MAP_KEYS","MAP_VALUES","STR_TO_MAP","ADD_MONTHS","CURRENT_DATE","CURRENT_DATE","CURRENT_TIMESTAMP","CURRENT_TIMESTAMP","CURRENT_TIMEZONE","DATE_ADD","DATE_FORMAT","DATE_FROM_UNIX_DATE","DATE_PART","DATE_SUB","DATE_TRUNC","DATEDIFF","DAY","DAYOFMONTH","DAYOFWEEK","DAYOFYEAR","EXTRACT","FROM_UNIXTIME","FROM_UTC_TIMESTAMP","HOUR","LAST_DAY","MAKE_DATE","MAKE_DT_INTERVAL","MAKE_INTERVAL","MAKE_TIMESTAMP","MAKE_YM_INTERVAL","MINUTE","MONTH","MONTHS_BETWEEN","NEXT_DAY","NOW","QUARTER","SECOND","SESSION_WINDOW","TIMESTAMP_MICROS","TIMESTAMP_MILLIS","TIMESTAMP_SECONDS","TO_DATE","TO_TIMESTAMP","TO_UNIX_TIMESTAMP","TO_UTC_TIMESTAMP","TRUNC","UNIX_DATE","UNIX_MICROS","UNIX_MILLIS","UNIX_SECONDS","UNIX_TIMESTAMP","WEEKDAY","WEEKOFYEAR","WINDOW","YEAR","FROM_JSON","GET_JSON_OBJECT","JSON_ARRAY_LENGTH","JSON_OBJECT_KEYS","JSON_TUPLE","SCHEMA_OF_JSON","TO_JSON","ABS","ACOS","ACOSH","AGGREGATE","ARRAY_SORT","ASCII","ASIN","ASINH","ASSERT_TRUE","ATAN","ATAN2","ATANH","BASE64","BIN","BIT_COUNT","BIT_GET","BIT_LENGTH","BROUND","BTRIM","CARDINALITY","CBRT","CEIL","CEILING","CHAR_LENGTH","CHARACTER_LENGTH","CHR","CONCAT","CONCAT_WS","CONV","COS","COSH","COT","CRC32","CURRENT_CATALOG","CURRENT_DATABASE","CURRENT_USER","DEGREES","ELT","EXP","EXPM1","FACTORIAL","FIND_IN_SET","FLOOR","FORALL","FORMAT_NUMBER","FORMAT_STRING","FROM_CSV","GETBIT","HASH","HEX","HYPOT","INITCAP","INLINE","INLINE_OUTER","INPUT_FILE_BLOCK_LENGTH","INPUT_FILE_BLOCK_START","INPUT_FILE_NAME","INSTR","ISNAN","ISNOTNULL","ISNULL","JAVA_METHOD","LCASE","LEFT","LENGTH","LEVENSHTEIN","LN","LOCATE","LOG","LOG10","LOG1P","LOG2","LOWER","LPAD","LTRIM","MAP_FILTER","MAP_ZIP_WITH","MD5","MOD","MONOTONICALLY_INCREASING_ID","NAMED_STRUCT","NANVL","NEGATIVE","NVL","NVL2","OCTET_LENGTH","OVERLAY","PARSE_URL","PI","PMOD","POSEXPLODE","POSEXPLODE_OUTER","POSITION","POSITIVE","POW","POWER","PRINTF","RADIANS","RAISE_ERROR","RAND","RANDN","RANDOM","REFLECT","REGEXP_EXTRACT","REGEXP_EXTRACT_ALL","REGEXP_LIKE","REGEXP_REPLACE","REPEAT","REPLACE","REVERSE","RIGHT","RINT","ROUND","RPAD","RTRIM","SCHEMA_OF_CSV","SENTENCES","SHA","SHA1","SHA2","SHIFTLEFT","SHIFTRIGHT","SHIFTRIGHTUNSIGNED","SIGN","SIGNUM","SIN","SINH","SOUNDEX","SPACE","SPARK_PARTITION_ID","SPLIT","SQRT","STACK","SUBSTR","SUBSTRING","SUBSTRING_INDEX","TAN","TANH","TO_CSV","TRANSFORM_KEYS","TRANSFORM_VALUES","TRANSLATE","TRIM","TRY_ADD","TRY_DIVIDE","TYPEOF","UCASE","UNBASE64","UNHEX","UPPER","UUID","VERSION","WIDTH_BUCKET","XPATH","XPATH_BOOLEAN","XPATH_DOUBLE","XPATH_FLOAT","XPATH_INT","XPATH_LONG","XPATH_NUMBER","XPATH_SHORT","XPATH_STRING","XXHASH64","ZIP_WITH","CAST","COALESCE","NULLIF"],io=v(["SELECT [ALL | DISTINCT]"]),ao=v(["WITH","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","SORT BY","CLUSTER BY","DISTRIBUTE BY","LIMIT","INSERT [INTO | OVERWRITE] [TABLE]","VALUES","INSERT OVERWRITE [LOCAL] DIRECTORY","LOAD DATA [LOCAL] INPATH","[OVERWRITE] INTO TABLE"]),Ir=v(["CREATE [EXTERNAL] TABLE [IF NOT EXISTS]"]),eT=v(["CREATE [OR REPLACE] [GLOBAL TEMPORARY | TEMPORARY] VIEW [IF NOT EXISTS]","DROP TABLE [IF EXISTS]","ALTER TABLE","ADD COLUMNS","DROP {COLUMN | COLUMNS}","RENAME TO","RENAME COLUMN","ALTER COLUMN","TRUNCATE TABLE","LATERAL VIEW","ALTER DATABASE","ALTER VIEW","CREATE DATABASE","CREATE FUNCTION","DROP DATABASE","DROP FUNCTION","DROP VIEW","REPAIR TABLE","USE DATABASE","TABLESAMPLE","PIVOT","TRANSFORM","EXPLAIN","ADD FILE","ADD JAR","ANALYZE TABLE","CACHE TABLE","CLEAR CACHE","DESCRIBE DATABASE","DESCRIBE FUNCTION","DESCRIBE QUERY","DESCRIBE TABLE","LIST FILE","LIST JAR","REFRESH","REFRESH TABLE","REFRESH FUNCTION","RESET","SHOW COLUMNS","SHOW CREATE TABLE","SHOW DATABASES","SHOW FUNCTIONS","SHOW PARTITIONS","SHOW TABLE EXTENDED","SHOW TABLES","SHOW TBLPROPERTIES","SHOW VIEWS","UNCACHE TABLE"]),Io=v(["UNION [ALL | DISTINCT]","EXCEPT [ALL | DISTINCT]","INTERSECT [ALL | DISTINCT]"]),No=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN","[LEFT] {ANTI | SEMI} JOIN","NATURAL [LEFT] {ANTI | SEMI} JOIN"]),lo=v(["ON DELETE","ON UPDATE","CURRENT ROW","{ROWS | RANGE} BETWEEN"]),_o={name:"spark",tokenizerOptions:{reservedSelect:io,reservedClauses:[...ao,...Ir,...eT],reservedSetOperations:Io,reservedJoins:No,reservedPhrases:lo,supportsXor:!0,reservedKeywords:So,reservedDataTypes:oo,reservedFunctionNames:Oo,extraParens:["[]"],stringTypes:["''-bs",'""-bs',{quote:"''-raw",prefixes:["R","X"],requirePrefix:!0},{quote:'""-raw',prefixes:["R","X"],requirePrefix:!0}],identTypes:["``"],variableTypes:[{quote:"{}",prefixes:["$"],requirePrefix:!0}],operators:["%","~","^","|","&","<=>","==","!","||","->"],postProcess:Lo},formatOptions:{onelineClauses:[...Ir,...eT],tabularOnelineClauses:eT}};function Lo(E){return E.map((e,T)=>{const t=E[T-1]||tt,r=E[T+1]||tt;return FE.WINDOW(e)&&r.type==="OPEN_PAREN"?AE(EE({},e),{type:"RESERVED_FUNCTION_NAME"}):e.text==="ITEMS"&&e.type==="RESERVED_KEYWORD"&&!(t.text==="COLLECTION"&&r.text==="TERMINATED")?AE(EE({},e),{type:"IDENTIFIER",text:e.raw}):e})}var Co=["ABS","CHANGES","CHAR","COALESCE","FORMAT","GLOB","HEX","IFNULL","IIF","INSTR","LAST_INSERT_ROWID","LENGTH","LIKE","LIKELIHOOD","LIKELY","LOAD_EXTENSION","LOWER","LTRIM","NULLIF","PRINTF","QUOTE","RANDOM","RANDOMBLOB","REPLACE","ROUND","RTRIM","SIGN","SOUNDEX","SQLITE_COMPILEOPTION_GET","SQLITE_COMPILEOPTION_USED","SQLITE_OFFSET","SQLITE_SOURCE_ID","SQLITE_VERSION","SUBSTR","SUBSTRING","TOTAL_CHANGES","TRIM","TYPEOF","UNICODE","UNLIKELY","UPPER","ZEROBLOB","AVG","COUNT","GROUP_CONCAT","MAX","MIN","SUM","TOTAL","DATE","TIME","DATETIME","JULIANDAY","UNIXEPOCH","STRFTIME","row_number","rank","dense_rank","percent_rank","cume_dist","ntile","lag","lead","first_value","last_value","nth_value","ACOS","ACOSH","ASIN","ASINH","ATAN","ATAN2","ATANH","CEIL","CEILING","COS","COSH","DEGREES","EXP","FLOOR","LN","LOG","LOG","LOG10","LOG2","MOD","PI","POW","POWER","RADIANS","SIN","SINH","SQRT","TAN","TANH","TRUNC","JSON","JSON_ARRAY","JSON_ARRAY_LENGTH","JSON_ARRAY_LENGTH","JSON_EXTRACT","JSON_INSERT","JSON_OBJECT","JSON_PATCH","JSON_REMOVE","JSON_REPLACE","JSON_SET","JSON_TYPE","JSON_TYPE","JSON_VALID","JSON_QUOTE","JSON_GROUP_ARRAY","JSON_GROUP_OBJECT","JSON_EACH","JSON_TREE","CAST"],uo=["ABORT","ACTION","ADD","AFTER","ALL","ALTER","AND","ARE","ALWAYS","ANALYZE","AS","ASC","ATTACH","AUTOINCREMENT","BEFORE","BEGIN","BETWEEN","BY","CASCADE","CASE","CAST","CHECK","COLLATE","COLUMN","COMMIT","CONFLICT","CONSTRAINT","CREATE","CROSS","CURRENT","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","DATABASE","DEFAULT","DEFERRABLE","DEFERRED","DELETE","DESC","DETACH","DISTINCT","DO","DROP","EACH","ELSE","END","ESCAPE","EXCEPT","EXCLUDE","EXCLUSIVE","EXISTS","EXPLAIN","FAIL","FILTER","FIRST","FOLLOWING","FOR","FOREIGN","FROM","FULL","GENERATED","GLOB","GROUP","GROUPS","HAVING","IF","IGNORE","IMMEDIATE","IN","INDEX","INDEXED","INITIALLY","INNER","INSERT","INSTEAD","INTERSECT","INTO","IS","ISNULL","JOIN","KEY","LAST","LEFT","LIKE","LIMIT","MATCH","MATERIALIZED","NATURAL","NO","NOT","NOTHING","NOTNULL","NULL","NULLS","OF","OFFSET","ON","ONLY","OPEN","OR","ORDER","OTHERS","OUTER","OVER","PARTITION","PLAN","PRAGMA","PRECEDING","PRIMARY","QUERY","RAISE","RANGE","RECURSIVE","REFERENCES","REGEXP","REINDEX","RELEASE","RENAME","REPLACE","RESTRICT","RETURNING","RIGHT","ROLLBACK","ROW","ROWS","SAVEPOINT","SELECT","SET","TABLE","TEMP","TEMPORARY","THEN","TIES","TO","TRANSACTION","TRIGGER","UNBOUNDED","UNION","UNIQUE","UPDATE","USING","VACUUM","VALUES","VIEW","VIRTUAL","WHEN","WHERE","WINDOW","WITH","WITHOUT"],co=["ANY","ARRAY","BLOB","CHARACTER","DECIMAL","INT","INTEGER","NATIVE CHARACTER","NCHAR","NUMERIC","NVARCHAR","REAL","TEXT","VARCHAR","VARYING CHARACTER"],fo=v(["SELECT [ALL | DISTINCT]"]),Po=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","INSERT [OR ABORT | OR FAIL | OR IGNORE | OR REPLACE | OR ROLLBACK] INTO","REPLACE INTO","VALUES","SET"]),Nr=v(["CREATE [TEMPORARY | TEMP] TABLE [IF NOT EXISTS]"]),ET=v(["CREATE [TEMPORARY | TEMP] VIEW [IF NOT EXISTS]","UPDATE [OR ABORT | OR FAIL | OR IGNORE | OR REPLACE | OR ROLLBACK]","ON CONFLICT","DELETE FROM","DROP TABLE [IF EXISTS]","ALTER TABLE","ADD [COLUMN]","DROP [COLUMN]","RENAME [COLUMN]","RENAME TO","SET SCHEMA"]),Do=v(["UNION [ALL]","EXCEPT","INTERSECT"]),po=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN"]),Mo=v(["ON {UPDATE | DELETE} [SET NULL | SET DEFAULT]","{ROWS | RANGE | GROUPS} BETWEEN"]),Uo={name:"sqlite",tokenizerOptions:{reservedSelect:fo,reservedClauses:[...Po,...Nr,...ET],reservedSetOperations:Do,reservedJoins:po,reservedPhrases:Mo,reservedKeywords:uo,reservedDataTypes:co,reservedFunctionNames:Co,stringTypes:["''-qq",{quote:"''-raw",prefixes:["X"],requirePrefix:!0}],identTypes:['""-qq',"``","[]"],paramTypes:{positional:!0,numbered:["?"],named:[":","@","$"]},operators:["%","~","&","|","<<",">>","==","->","->>","||"]},formatOptions:{onelineClauses:[...Nr,...ET],tabularOnelineClauses:ET}},mo=["GROUPING","RANK","DENSE_RANK","PERCENT_RANK","CUME_DIST","ROW_NUMBER","POSITION","OCCURRENCES_REGEX","POSITION_REGEX","EXTRACT","CHAR_LENGTH","CHARACTER_LENGTH","OCTET_LENGTH","CARDINALITY","ABS","MOD","LN","EXP","POWER","SQRT","FLOOR","CEIL","CEILING","WIDTH_BUCKET","SUBSTRING","SUBSTRING_REGEX","UPPER","LOWER","CONVERT","TRANSLATE","TRANSLATE_REGEX","TRIM","OVERLAY","NORMALIZE","SPECIFICTYPE","CURRENT_DATE","CURRENT_TIME","LOCALTIME","CURRENT_TIMESTAMP","LOCALTIMESTAMP","COUNT","AVG","MAX","MIN","SUM","STDDEV_POP","STDDEV_SAMP","VAR_SAMP","VAR_POP","COLLECT","FUSION","INTERSECTION","COVAR_POP","COVAR_SAMP","CORR","REGR_SLOPE","REGR_INTERCEPT","REGR_COUNT","REGR_R2","REGR_AVGX","REGR_AVGY","REGR_SXX","REGR_SYY","REGR_SXY","PERCENTILE_CONT","PERCENTILE_DISC","CAST","COALESCE","NULLIF","ROUND","SIN","COS","TAN","ASIN","ACOS","ATAN"],ho=["ALL","ALLOCATE","ALTER","ANY","ARE","AS","ASC","ASENSITIVE","ASYMMETRIC","AT","ATOMIC","AUTHORIZATION","BEGIN","BETWEEN","BOTH","BY","CALL","CALLED","CASCADED","CAST","CHECK","CLOSE","COALESCE","COLLATE","COLUMN","COMMIT","CONDITION","CONNECT","CONSTRAINT","CORRESPONDING","CREATE","CROSS","CUBE","CURRENT","CURRENT_CATALOG","CURRENT_DEFAULT_TRANSFORM_GROUP","CURRENT_PATH","CURRENT_ROLE","CURRENT_SCHEMA","CURRENT_TRANSFORM_GROUP_FOR_TYPE","CURRENT_USER","CURSOR","CYCLE","DEALLOCATE","DAY","DECLARE","DEFAULT","DELETE","DEREF","DESC","DESCRIBE","DETERMINISTIC","DISCONNECT","DISTINCT","DROP","DYNAMIC","EACH","ELEMENT","END-EXEC","ESCAPE","EVERY","EXCEPT","EXEC","EXECUTE","EXISTS","EXTERNAL","FALSE","FETCH","FILTER","FOR","FOREIGN","FREE","FROM","FULL","FUNCTION","GET","GLOBAL","GRANT","GROUP","HAVING","HOLD","HOUR","IDENTITY","IN","INDICATOR","INNER","INOUT","INSENSITIVE","INSERT","INTERSECT","INTO","IS","LANGUAGE","LARGE","LATERAL","LEADING","LEFT","LIKE","LIKE_REGEX","LOCAL","MATCH","MEMBER","MERGE","METHOD","MINUTE","MODIFIES","MODULE","MONTH","NATURAL","NEW","NO","NONE","NOT","NULL","NULLIF","OF","OLD","ON","ONLY","OPEN","ORDER","OUT","OUTER","OVER","OVERLAPS","PARAMETER","PARTITION","PRECISION","PREPARE","PRIMARY","PROCEDURE","RANGE","READS","REAL","RECURSIVE","REF","REFERENCES","REFERENCING","RELEASE","RESULT","RETURN","RETURNS","REVOKE","RIGHT","ROLLBACK","ROLLUP","ROW","ROWS","SAVEPOINT","SCOPE","SCROLL","SEARCH","SECOND","SELECT","SENSITIVE","SESSION_USER","SET","SIMILAR","SOME","SPECIFIC","SQL","SQLEXCEPTION","SQLSTATE","SQLWARNING","START","STATIC","SUBMULTISET","SYMMETRIC","SYSTEM","SYSTEM_USER","TABLE","TABLESAMPLE","THEN","TIMEZONE_HOUR","TIMEZONE_MINUTE","TO","TRAILING","TRANSLATION","TREAT","TRIGGER","TRUE","UESCAPE","UNION","UNIQUE","UNKNOWN","UNNEST","UPDATE","USER","USING","VALUE","VALUES","WHENEVER","WINDOW","WITHIN","WITHOUT","YEAR"],Go=["ARRAY","BIGINT","BINARY LARGE OBJECT","BINARY VARYING","BINARY","BLOB","BOOLEAN","CHAR LARGE OBJECT","CHAR VARYING","CHAR","CHARACTER LARGE OBJECT","CHARACTER VARYING","CHARACTER","CLOB","DATE","DEC","DECIMAL","DOUBLE","FLOAT","INT","INTEGER","INTERVAL","MULTISET","NATIONAL CHAR VARYING","NATIONAL CHAR","NATIONAL CHARACTER LARGE OBJECT","NATIONAL CHARACTER VARYING","NATIONAL CHARACTER","NCHAR LARGE OBJECT","NCHAR VARYING","NCHAR","NCLOB","NUMERIC","SMALLINT","TIME","TIMESTAMP","VARBINARY","VARCHAR"],go=v(["SELECT [ALL | DISTINCT]"]),Ho=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY [ALL | DISTINCT]","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","FETCH {FIRST | NEXT}","INSERT INTO","VALUES","SET"]),lr=v(["CREATE [GLOBAL TEMPORARY | LOCAL TEMPORARY] TABLE"]),tT=v(["CREATE [RECURSIVE] VIEW","UPDATE","WHERE CURRENT OF","DELETE FROM","DROP TABLE","ALTER TABLE","ADD COLUMN","DROP [COLUMN]","RENAME COLUMN","RENAME TO","ALTER [COLUMN]","{SET | DROP} DEFAULT","ADD SCOPE","DROP SCOPE {CASCADE | RESTRICT}","RESTART WITH","TRUNCATE TABLE","SET SCHEMA"]),bo=v(["UNION [ALL | DISTINCT]","EXCEPT [ALL | DISTINCT]","INTERSECT [ALL | DISTINCT]"]),yo=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN"]),Bo=v(["ON {UPDATE | DELETE} [SET NULL | SET DEFAULT]","{ROWS | RANGE} BETWEEN"]),vo={name:"sql",tokenizerOptions:{reservedSelect:go,reservedClauses:[...Ho,...lr,...tT],reservedSetOperations:bo,reservedJoins:yo,reservedPhrases:Bo,reservedKeywords:ho,reservedDataTypes:Go,reservedFunctionNames:mo,stringTypes:[{quote:"''-qq-bs",prefixes:["N","U&"]},{quote:"''-raw",prefixes:["X"],requirePrefix:!0}],identTypes:['""-qq',"``"],paramTypes:{positional:!0},operators:["||"]},formatOptions:{onelineClauses:[...lr,...tT],tabularOnelineClauses:tT}},Fo=["ABS","ACOS","ALL_MATCH","ANY_MATCH","APPROX_DISTINCT","APPROX_MOST_FREQUENT","APPROX_PERCENTILE","APPROX_SET","ARBITRARY","ARRAYS_OVERLAP","ARRAY_AGG","ARRAY_DISTINCT","ARRAY_EXCEPT","ARRAY_INTERSECT","ARRAY_JOIN","ARRAY_MAX","ARRAY_MIN","ARRAY_POSITION","ARRAY_REMOVE","ARRAY_SORT","ARRAY_UNION","ASIN","ATAN","ATAN2","AT_TIMEZONE","AVG","BAR","BETA_CDF","BING_TILE","BING_TILES_AROUND","BING_TILE_AT","BING_TILE_COORDINATES","BING_TILE_POLYGON","BING_TILE_QUADKEY","BING_TILE_ZOOM_LEVEL","BITWISE_AND","BITWISE_AND_AGG","BITWISE_LEFT_SHIFT","BITWISE_NOT","BITWISE_OR","BITWISE_OR_AGG","BITWISE_RIGHT_SHIFT","BITWISE_RIGHT_SHIFT_ARITHMETIC","BITWISE_XOR","BIT_COUNT","BOOL_AND","BOOL_OR","CARDINALITY","CAST","CBRT","CEIL","CEILING","CHAR2HEXINT","CHECKSUM","CHR","CLASSIFY","COALESCE","CODEPOINT","COLOR","COMBINATIONS","CONCAT","CONCAT_WS","CONTAINS","CONTAINS_SEQUENCE","CONVEX_HULL_AGG","CORR","COS","COSH","COSINE_SIMILARITY","COUNT","COUNT_IF","COVAR_POP","COVAR_SAMP","CRC32","CUME_DIST","CURRENT_CATALOG","CURRENT_DATE","CURRENT_GROUPS","CURRENT_SCHEMA","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_TIMEZONE","CURRENT_USER","DATE","DATE_ADD","DATE_DIFF","DATE_FORMAT","DATE_PARSE","DATE_TRUNC","DAY","DAY_OF_MONTH","DAY_OF_WEEK","DAY_OF_YEAR","DEGREES","DENSE_RANK","DOW","DOY","E","ELEMENT_AT","EMPTY_APPROX_SET","EVALUATE_CLASSIFIER_PREDICTIONS","EVERY","EXP","EXTRACT","FEATURES","FILTER","FIRST_VALUE","FLATTEN","FLOOR","FORMAT","FORMAT_DATETIME","FORMAT_NUMBER","FROM_BASE","FROM_BASE32","FROM_BASE64","FROM_BASE64URL","FROM_BIG_ENDIAN_32","FROM_BIG_ENDIAN_64","FROM_ENCODED_POLYLINE","FROM_GEOJSON_GEOMETRY","FROM_HEX","FROM_IEEE754_32","FROM_IEEE754_64","FROM_ISO8601_DATE","FROM_ISO8601_TIMESTAMP","FROM_ISO8601_TIMESTAMP_NANOS","FROM_UNIXTIME","FROM_UNIXTIME_NANOS","FROM_UTF8","GEOMETRIC_MEAN","GEOMETRY_FROM_HADOOP_SHAPE","GEOMETRY_INVALID_REASON","GEOMETRY_NEAREST_POINTS","GEOMETRY_TO_BING_TILES","GEOMETRY_UNION","GEOMETRY_UNION_AGG","GREATEST","GREAT_CIRCLE_DISTANCE","HAMMING_DISTANCE","HASH_COUNTS","HISTOGRAM","HMAC_MD5","HMAC_SHA1","HMAC_SHA256","HMAC_SHA512","HOUR","HUMAN_READABLE_SECONDS","IF","INDEX","INFINITY","INTERSECTION_CARDINALITY","INVERSE_BETA_CDF","INVERSE_NORMAL_CDF","IS_FINITE","IS_INFINITE","IS_JSON_SCALAR","IS_NAN","JACCARD_INDEX","JSON_ARRAY_CONTAINS","JSON_ARRAY_GET","JSON_ARRAY_LENGTH","JSON_EXISTS","JSON_EXTRACT","JSON_EXTRACT_SCALAR","JSON_FORMAT","JSON_PARSE","JSON_QUERY","JSON_SIZE","JSON_VALUE","KURTOSIS","LAG","LAST_DAY_OF_MONTH","LAST_VALUE","LEAD","LEARN_CLASSIFIER","LEARN_LIBSVM_CLASSIFIER","LEARN_LIBSVM_REGRESSOR","LEARN_REGRESSOR","LEAST","LENGTH","LEVENSHTEIN_DISTANCE","LINE_INTERPOLATE_POINT","LINE_INTERPOLATE_POINTS","LINE_LOCATE_POINT","LISTAGG","LN","LOCALTIME","LOCALTIMESTAMP","LOG","LOG10","LOG2","LOWER","LPAD","LTRIM","LUHN_CHECK","MAKE_SET_DIGEST","MAP","MAP_AGG","MAP_CONCAT","MAP_ENTRIES","MAP_FILTER","MAP_FROM_ENTRIES","MAP_KEYS","MAP_UNION","MAP_VALUES","MAP_ZIP_WITH","MAX","MAX_BY","MD5","MERGE","MERGE_SET_DIGEST","MILLISECOND","MIN","MINUTE","MIN_BY","MOD","MONTH","MULTIMAP_AGG","MULTIMAP_FROM_ENTRIES","MURMUR3","NAN","NGRAMS","NONE_MATCH","NORMALIZE","NORMAL_CDF","NOW","NTH_VALUE","NTILE","NULLIF","NUMERIC_HISTOGRAM","OBJECTID","OBJECTID_TIMESTAMP","PARSE_DATA_SIZE","PARSE_DATETIME","PARSE_DURATION","PERCENT_RANK","PI","POSITION","POW","POWER","QDIGEST_AGG","QUARTER","RADIANS","RAND","RANDOM","RANK","REDUCE","REDUCE_AGG","REGEXP_COUNT","REGEXP_EXTRACT","REGEXP_EXTRACT_ALL","REGEXP_LIKE","REGEXP_POSITION","REGEXP_REPLACE","REGEXP_SPLIT","REGRESS","REGR_INTERCEPT","REGR_SLOPE","RENDER","REPEAT","REPLACE","REVERSE","RGB","ROUND","ROW_NUMBER","RPAD","RTRIM","SECOND","SEQUENCE","SHA1","SHA256","SHA512","SHUFFLE","SIGN","SIMPLIFY_GEOMETRY","SIN","SKEWNESS","SLICE","SOUNDEX","SPATIAL_PARTITIONING","SPATIAL_PARTITIONS","SPLIT","SPLIT_PART","SPLIT_TO_MAP","SPLIT_TO_MULTIMAP","SPOOKY_HASH_V2_32","SPOOKY_HASH_V2_64","SQRT","STARTS_WITH","STDDEV","STDDEV_POP","STDDEV_SAMP","STRPOS","ST_AREA","ST_ASBINARY","ST_ASTEXT","ST_BOUNDARY","ST_BUFFER","ST_CENTROID","ST_CONTAINS","ST_CONVEXHULL","ST_COORDDIM","ST_CROSSES","ST_DIFFERENCE","ST_DIMENSION","ST_DISJOINT","ST_DISTANCE","ST_ENDPOINT","ST_ENVELOPE","ST_ENVELOPEASPTS","ST_EQUALS","ST_EXTERIORRING","ST_GEOMETRIES","ST_GEOMETRYFROMTEXT","ST_GEOMETRYN","ST_GEOMETRYTYPE","ST_GEOMFROMBINARY","ST_INTERIORRINGN","ST_INTERIORRINGS","ST_INTERSECTION","ST_INTERSECTS","ST_ISCLOSED","ST_ISEMPTY","ST_ISRING","ST_ISSIMPLE","ST_ISVALID","ST_LENGTH","ST_LINEFROMTEXT","ST_LINESTRING","ST_MULTIPOINT","ST_NUMGEOMETRIES","ST_NUMINTERIORRING","ST_NUMPOINTS","ST_OVERLAPS","ST_POINT","ST_POINTN","ST_POINTS","ST_POLYGON","ST_RELATE","ST_STARTPOINT","ST_SYMDIFFERENCE","ST_TOUCHES","ST_UNION","ST_WITHIN","ST_X","ST_XMAX","ST_XMIN","ST_Y","ST_YMAX","ST_YMIN","SUBSTR","SUBSTRING","SUM","TAN","TANH","TDIGEST_AGG","TIMESTAMP_OBJECTID","TIMEZONE_HOUR","TIMEZONE_MINUTE","TO_BASE","TO_BASE32","TO_BASE64","TO_BASE64URL","TO_BIG_ENDIAN_32","TO_BIG_ENDIAN_64","TO_CHAR","TO_DATE","TO_ENCODED_POLYLINE","TO_GEOJSON_GEOMETRY","TO_GEOMETRY","TO_HEX","TO_IEEE754_32","TO_IEEE754_64","TO_ISO8601","TO_MILLISECONDS","TO_SPHERICAL_GEOGRAPHY","TO_TIMESTAMP","TO_UNIXTIME","TO_UTF8","TRANSFORM","TRANSFORM_KEYS","TRANSFORM_VALUES","TRANSLATE","TRIM","TRIM_ARRAY","TRUNCATE","TRY","TRY_CAST","TYPEOF","UPPER","URL_DECODE","URL_ENCODE","URL_EXTRACT_FRAGMENT","URL_EXTRACT_HOST","URL_EXTRACT_PARAMETER","URL_EXTRACT_PATH","URL_EXTRACT_PORT","URL_EXTRACT_PROTOCOL","URL_EXTRACT_QUERY","UUID","VALUES_AT_QUANTILES","VALUE_AT_QUANTILE","VARIANCE","VAR_POP","VAR_SAMP","VERSION","WEEK","WEEK_OF_YEAR","WIDTH_BUCKET","WILSON_INTERVAL_LOWER","WILSON_INTERVAL_UPPER","WITH_TIMEZONE","WORD_STEM","XXHASH64","YEAR","YEAR_OF_WEEK","YOW","ZIP","ZIP_WITH","CLASSIFIER","FIRST","LAST","MATCH_NUMBER","NEXT","PERMUTE","PREV"],Yo=["ABSENT","ADD","ADMIN","AFTER","ALL","ALTER","ANALYZE","AND","ANY","AS","ASC","AT","AUTHORIZATION","BERNOULLI","BETWEEN","BOTH","BY","CALL","CASCADE","CASE","CATALOGS","COLUMN","COLUMNS","COMMENT","COMMIT","COMMITTED","CONDITIONAL","CONSTRAINT","COPARTITION","CREATE","CROSS","CUBE","CURRENT","CURRENT_PATH","CURRENT_ROLE","DATA","DEALLOCATE","DEFAULT","DEFINE","DEFINER","DELETE","DENY","DESC","DESCRIBE","DESCRIPTOR","DISTINCT","DISTRIBUTED","DOUBLE","DROP","ELSE","EMPTY","ENCODING","END","ERROR","ESCAPE","EXCEPT","EXCLUDING","EXECUTE","EXISTS","EXPLAIN","FALSE","FETCH","FINAL","FIRST","FOLLOWING","FOR","FROM","FULL","FUNCTIONS","GRANT","GRANTED","GRANTS","GRAPHVIZ","GROUP","GROUPING","GROUPS","HAVING","IGNORE","IN","INCLUDING","INITIAL","INNER","INPUT","INSERT","INTERSECT","INTERVAL","INTO","INVOKER","IO","IS","ISOLATION","JOIN","JSON","JSON_ARRAY","JSON_OBJECT","KEEP","KEY","KEYS","LAST","LATERAL","LEADING","LEFT","LEVEL","LIKE","LIMIT","LOCAL","LOGICAL","MATCH","MATCHED","MATCHES","MATCH_RECOGNIZE","MATERIALIZED","MEASURES","NATURAL","NEXT","NFC","NFD","NFKC","NFKD","NO","NONE","NOT","NULL","NULLS","OBJECT","OF","OFFSET","OMIT","ON","ONE","ONLY","OPTION","OR","ORDER","ORDINALITY","OUTER","OUTPUT","OVER","OVERFLOW","PARTITION","PARTITIONS","PASSING","PAST","PATH","PATTERN","PER","PERMUTE","PRECEDING","PRECISION","PREPARE","PRIVILEGES","PROPERTIES","PRUNE","QUOTES","RANGE","READ","RECURSIVE","REFRESH","RENAME","REPEATABLE","RESET","RESPECT","RESTRICT","RETURNING","REVOKE","RIGHT","ROLE","ROLES","ROLLBACK","ROLLUP","ROW","ROWS","RUNNING","SCALAR","SCHEMA","SCHEMAS","SECURITY","SEEK","SELECT","SERIALIZABLE","SESSION","SET","SETS","SHOW","SKIP","SOME","START","STATS","STRING","SUBSET","SYSTEM","TABLE","TABLES","TABLESAMPLE","TEXT","THEN","TIES","TIME","TIMESTAMP","TO","TRAILING","TRANSACTION","TRUE","TYPE","UESCAPE","UNBOUNDED","UNCOMMITTED","UNCONDITIONAL","UNION","UNIQUE","UNKNOWN","UNMATCHED","UNNEST","UPDATE","USE","USER","USING","UTF16","UTF32","UTF8","VALIDATE","VALUE","VALUES","VERBOSE","VIEW","WHEN","WHERE","WINDOW","WITH","WITHIN","WITHOUT","WORK","WRAPPER","WRITE","ZONE"],Vo=["BIGINT","INT","INTEGER","SMALLINT","TINYINT","BOOLEAN","DATE","DECIMAL","REAL","DOUBLE","HYPERLOGLOG","QDIGEST","TDIGEST","P4HYPERLOGLOG","INTERVAL","TIMESTAMP","TIME","VARBINARY","VARCHAR","CHAR","ROW","ARRAY","MAP","JSON","JSON2016","IPADDRESS","GEOMETRY","UUID","SETDIGEST","JONIREGEXP","RE2JREGEXP","LIKEPATTERN","COLOR","CODEPOINTS","FUNCTION","JSONPATH"],Wo=v(["SELECT [ALL | DISTINCT]"]),wo=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY [ALL | DISTINCT]","HAVING","WINDOW","PARTITION BY","ORDER BY","LIMIT","OFFSET","FETCH {FIRST | NEXT}","INSERT INTO","VALUES","SET","MATCH_RECOGNIZE","MEASURES","ONE ROW PER MATCH","ALL ROWS PER MATCH","AFTER MATCH","PATTERN","SUBSET","DEFINE"]),_r=v(["CREATE TABLE [IF NOT EXISTS]"]),TT=v(["CREATE [OR REPLACE] [MATERIALIZED] VIEW","UPDATE","DELETE FROM","DROP TABLE [IF EXISTS]","ALTER TABLE [IF EXISTS]","ADD COLUMN [IF NOT EXISTS]","DROP COLUMN [IF EXISTS]","RENAME COLUMN [IF EXISTS]","RENAME TO","SET AUTHORIZATION [USER | ROLE]","SET PROPERTIES","EXECUTE","TRUNCATE TABLE","ALTER SCHEMA","ALTER MATERIALIZED VIEW","ALTER VIEW","CREATE SCHEMA","CREATE ROLE","DROP SCHEMA","DROP MATERIALIZED VIEW","DROP VIEW","DROP ROLE","EXPLAIN","ANALYZE","EXPLAIN ANALYZE","EXPLAIN ANALYZE VERBOSE","USE","DESCRIBE INPUT","DESCRIBE OUTPUT","REFRESH MATERIALIZED VIEW","RESET SESSION","SET SESSION","SET PATH","SET TIME ZONE","SHOW GRANTS","SHOW CREATE TABLE","SHOW CREATE SCHEMA","SHOW CREATE VIEW","SHOW CREATE MATERIALIZED VIEW","SHOW TABLES","SHOW SCHEMAS","SHOW CATALOGS","SHOW COLUMNS","SHOW STATS FOR","SHOW ROLES","SHOW CURRENT ROLES","SHOW ROLE GRANTS","SHOW FUNCTIONS","SHOW SESSION"]),$o=v(["UNION [ALL | DISTINCT]","EXCEPT [ALL | DISTINCT]","INTERSECT [ALL | DISTINCT]"]),xo=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL [INNER] JOIN","NATURAL {LEFT | RIGHT | FULL} [OUTER] JOIN"]),Xo=v(["{ROWS | RANGE | GROUPS} BETWEEN","IS [NOT] DISTINCT FROM"]),ko={name:"trino",tokenizerOptions:{reservedSelect:Wo,reservedClauses:[...wo,..._r,...TT],reservedSetOperations:$o,reservedJoins:xo,reservedPhrases:Xo,reservedKeywords:Yo,reservedDataTypes:Vo,reservedFunctionNames:Fo,extraParens:["[]","{}"],stringTypes:[{quote:"''-qq",prefixes:["U&"]},{quote:"''-raw",prefixes:["X"],requirePrefix:!0}],identTypes:['""-qq'],paramTypes:{positional:!0},operators:["%","->","=>",":","||","|","^","$"]},formatOptions:{onelineClauses:[..._r,...TT],tabularOnelineClauses:TT}},Ko=["APPROX_COUNT_DISTINCT","AVG","CHECKSUM_AGG","COUNT","COUNT_BIG","GROUPING","GROUPING_ID","MAX","MIN","STDEV","STDEVP","SUM","VAR","VARP","CUME_DIST","FIRST_VALUE","LAG","LAST_VALUE","LEAD","PERCENTILE_CONT","PERCENTILE_DISC","PERCENT_RANK","Collation - COLLATIONPROPERTY","Collation - TERTIARY_WEIGHTS","@@DBTS","@@LANGID","@@LANGUAGE","@@LOCK_TIMEOUT","@@MAX_CONNECTIONS","@@MAX_PRECISION","@@NESTLEVEL","@@OPTIONS","@@REMSERVER","@@SERVERNAME","@@SERVICENAME","@@SPID","@@TEXTSIZE","@@VERSION","CAST","CONVERT","PARSE","TRY_CAST","TRY_CONVERT","TRY_PARSE","ASYMKEY_ID","ASYMKEYPROPERTY","CERTPROPERTY","CERT_ID","CRYPT_GEN_RANDOM","DECRYPTBYASYMKEY","DECRYPTBYCERT","DECRYPTBYKEY","DECRYPTBYKEYAUTOASYMKEY","DECRYPTBYKEYAUTOCERT","DECRYPTBYPASSPHRASE","ENCRYPTBYASYMKEY","ENCRYPTBYCERT","ENCRYPTBYKEY","ENCRYPTBYPASSPHRASE","HASHBYTES","IS_OBJECTSIGNED","KEY_GUID","KEY_ID","KEY_NAME","SIGNBYASYMKEY","SIGNBYCERT","SYMKEYPROPERTY","VERIFYSIGNEDBYCERT","VERIFYSIGNEDBYASYMKEY","@@CURSOR_ROWS","@@FETCH_STATUS","CURSOR_STATUS","DATALENGTH","IDENT_CURRENT","IDENT_INCR","IDENT_SEED","IDENTITY","SQL_VARIANT_PROPERTY","@@DATEFIRST","CURRENT_TIMESTAMP","CURRENT_TIMEZONE","CURRENT_TIMEZONE_ID","DATEADD","DATEDIFF","DATEDIFF_BIG","DATEFROMPARTS","DATENAME","DATEPART","DATETIME2FROMPARTS","DATETIMEFROMPARTS","DATETIMEOFFSETFROMPARTS","DAY","EOMONTH","GETDATE","GETUTCDATE","ISDATE","MONTH","SMALLDATETIMEFROMPARTS","SWITCHOFFSET","SYSDATETIME","SYSDATETIMEOFFSET","SYSUTCDATETIME","TIMEFROMPARTS","TODATETIMEOFFSET","YEAR","JSON","ISJSON","JSON_VALUE","JSON_QUERY","JSON_MODIFY","ABS","ACOS","ASIN","ATAN","ATN2","CEILING","COS","COT","DEGREES","EXP","FLOOR","LOG","LOG10","PI","POWER","RADIANS","RAND","ROUND","SIGN","SIN","SQRT","SQUARE","TAN","CHOOSE","GREATEST","IIF","LEAST","@@PROCID","APP_NAME","APPLOCK_MODE","APPLOCK_TEST","ASSEMBLYPROPERTY","COL_LENGTH","COL_NAME","COLUMNPROPERTY","DATABASEPROPERTYEX","DB_ID","DB_NAME","FILE_ID","FILE_IDEX","FILE_NAME","FILEGROUP_ID","FILEGROUP_NAME","FILEGROUPPROPERTY","FILEPROPERTY","FILEPROPERTYEX","FULLTEXTCATALOGPROPERTY","FULLTEXTSERVICEPROPERTY","INDEX_COL","INDEXKEY_PROPERTY","INDEXPROPERTY","NEXT VALUE FOR","OBJECT_DEFINITION","OBJECT_ID","OBJECT_NAME","OBJECT_SCHEMA_NAME","OBJECTPROPERTY","OBJECTPROPERTYEX","ORIGINAL_DB_NAME","PARSENAME","SCHEMA_ID","SCHEMA_NAME","SCOPE_IDENTITY","SERVERPROPERTY","STATS_DATE","TYPE_ID","TYPE_NAME","TYPEPROPERTY","DENSE_RANK","NTILE","RANK","ROW_NUMBER","PUBLISHINGSERVERNAME","CERTENCODED","CERTPRIVATEKEY","CURRENT_USER","DATABASE_PRINCIPAL_ID","HAS_DBACCESS","HAS_PERMS_BY_NAME","IS_MEMBER","IS_ROLEMEMBER","IS_SRVROLEMEMBER","LOGINPROPERTY","ORIGINAL_LOGIN","PERMISSIONS","PWDENCRYPT","PWDCOMPARE","SESSION_USER","SESSIONPROPERTY","SUSER_ID","SUSER_NAME","SUSER_SID","SUSER_SNAME","SYSTEM_USER","USER","USER_ID","USER_NAME","ASCII","CHAR","CHARINDEX","CONCAT","CONCAT_WS","DIFFERENCE","FORMAT","LEFT","LEN","LOWER","LTRIM","NCHAR","PATINDEX","QUOTENAME","REPLACE","REPLICATE","REVERSE","RIGHT","RTRIM","SOUNDEX","SPACE","STR","STRING_AGG","STRING_ESCAPE","STUFF","SUBSTRING","TRANSLATE","TRIM","UNICODE","UPPER","$PARTITION","@@ERROR","@@IDENTITY","@@PACK_RECEIVED","@@ROWCOUNT","@@TRANCOUNT","BINARY_CHECKSUM","CHECKSUM","COMPRESS","CONNECTIONPROPERTY","CONTEXT_INFO","CURRENT_REQUEST_ID","CURRENT_TRANSACTION_ID","DECOMPRESS","ERROR_LINE","ERROR_MESSAGE","ERROR_NUMBER","ERROR_PROCEDURE","ERROR_SEVERITY","ERROR_STATE","FORMATMESSAGE","GET_FILESTREAM_TRANSACTION_CONTEXT","GETANSINULL","HOST_ID","HOST_NAME","ISNULL","ISNUMERIC","MIN_ACTIVE_ROWVERSION","NEWID","NEWSEQUENTIALID","ROWCOUNT_BIG","SESSION_CONTEXT","XACT_STATE","@@CONNECTIONS","@@CPU_BUSY","@@IDLE","@@IO_BUSY","@@PACK_SENT","@@PACKET_ERRORS","@@TIMETICKS","@@TOTAL_ERRORS","@@TOTAL_READ","@@TOTAL_WRITE","TEXTPTR","TEXTVALID","COLUMNS_UPDATED","EVENTDATA","TRIGGER_NESTLEVEL","UPDATE","COALESCE","NULLIF"],Jo=["ADD","ALL","ALTER","AND","ANY","AS","ASC","AUTHORIZATION","BACKUP","BEGIN","BETWEEN","BREAK","BROWSE","BULK","BY","CASCADE","CHECK","CHECKPOINT","CLOSE","CLUSTERED","COALESCE","COLLATE","COLUMN","COMMIT","COMPUTE","CONSTRAINT","CONTAINS","CONTAINSTABLE","CONTINUE","CONVERT","CREATE","CROSS","CURRENT","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURSOR","DATABASE","DBCC","DEALLOCATE","DECLARE","DEFAULT","DELETE","DENY","DESC","DISK","DISTINCT","DISTRIBUTED","DROP","DUMP","ERRLVL","ESCAPE","EXEC","EXECUTE","EXISTS","EXIT","EXTERNAL","FETCH","FILE","FILLFACTOR","FOR","FOREIGN","FREETEXT","FREETEXTTABLE","FROM","FULL","FUNCTION","GOTO","GRANT","GROUP","HAVING","HOLDLOCK","IDENTITY","IDENTITYCOL","IDENTITY_INSERT","IF","IN","INDEX","INNER","INSERT","INTERSECT","INTO","IS","JOIN","KEY","KILL","LEFT","LIKE","LINENO","LOAD","MERGE","NOCHECK","NONCLUSTERED","NOT","NULL","NULLIF","OF","OFF","OFFSETS","ON","OPEN","OPENDATASOURCE","OPENQUERY","OPENROWSET","OPENXML","OPTION","OR","ORDER","OUTER","OVER","PERCENT","PIVOT","PLAN","PRIMARY","PRINT","PROC","PROCEDURE","PUBLIC","RAISERROR","READ","READTEXT","RECONFIGURE","REFERENCES","REPLICATION","RESTORE","RESTRICT","RETURN","REVERT","REVOKE","RIGHT","ROLLBACK","ROWCOUNT","ROWGUIDCOL","RULE","SAVE","SCHEMA","SECURITYAUDIT","SELECT","SEMANTICKEYPHRASETABLE","SEMANTICSIMILARITYDETAILSTABLE","SEMANTICSIMILARITYTABLE","SESSION_USER","SET","SETUSER","SHUTDOWN","SOME","STATISTICS","SYSTEM_USER","TABLE","TABLESAMPLE","TEXTSIZE","THEN","TO","TOP","TRAN","TRANSACTION","TRIGGER","TRUNCATE","TRY_CONVERT","TSEQUAL","UNION","UNIQUE","UNPIVOT","UPDATE","UPDATETEXT","USE","USER","VALUES","VIEW","WAITFOR","WHERE","WHILE","WITH","WITHIN GROUP","WRITETEXT","ABSOLUTE","ACTION","ADA","ALLOCATE","ARE","ASSERTION","AT","AVG","BIT_LENGTH","BOTH","CASCADED","CAST","CATALOG","CHARACTER_LENGTH","CHAR_LENGTH","COLLATION","CONNECT","CONNECTION","CONSTRAINTS","CORRESPONDING","COUNT","DAY","DEFERRABLE","DEFERRED","DESCRIBE","DESCRIPTOR","DIAGNOSTICS","DISCONNECT","DOMAIN","END-EXEC","EXCEPTION","EXTRACT","FALSE","FIRST","FORTRAN","FOUND","GET","GLOBAL","GO","HOUR","IMMEDIATE","INCLUDE","INDICATOR","INITIALLY","INPUT","INSENSITIVE","INTERVAL","ISOLATION","LANGUAGE","LAST","LEADING","LEVEL","LOCAL","LOWER","MATCH","MAX","MIN","MINUTE","MODULE","MONTH","NAMES","NATURAL","NEXT","NO","NONE","OCTET_LENGTH","ONLY","OUTPUT","OVERLAPS","PAD","PARTIAL","PASCAL","POSITION","PREPARE","PRESERVE","PRIOR","PRIVILEGES","RELATIVE","ROWS","SCROLL","SECOND","SECTION","SESSION","SIZE","SPACE","SQL","SQLCA","SQLCODE","SQLERROR","SQLSTATE","SQLWARNING","SUBSTRING","SUM","TEMPORARY","TIMEZONE_HOUR","TIMEZONE_MINUTE","TRAILING","TRANSLATE","TRANSLATION","TRIM","TRUE","UNKNOWN","UPPER","USAGE","VALUE","WHENEVER","WORK","WRITE","YEAR","ZONE"],qo=["BINARY","BIT","CHAR","CHAR","CHARACTER","DATE","DATETIME2","DATETIMEOFFSET","DEC","DECIMAL","DOUBLE","FLOAT","INT","INTEGER","NATIONAL","NCHAR","NUMERIC","NVARCHAR","PRECISION","REAL","SMALLINT","TIME","TIMESTAMP","VARBINARY","VARCHAR"],Qo=v(["SELECT [ALL | DISTINCT]"]),Zo=v(["WITH","INTO","FROM","WHERE","GROUP BY","HAVING","WINDOW","PARTITION BY","ORDER BY","OFFSET","FETCH {FIRST | NEXT}","FOR {BROWSE | XML | JSON}","OPTION","INSERT [INTO]","VALUES","SET","MERGE [INTO]","WHEN [NOT] MATCHED [BY TARGET | BY SOURCE] [THEN]","UPDATE SET","CREATE [OR ALTER] {PROC | PROCEDURE}"]),Lr=v(["CREATE TABLE"]),rT=v(["CREATE [OR ALTER] [MATERIALIZED] VIEW","UPDATE","WHERE CURRENT OF","DELETE [FROM]","DROP TABLE [IF EXISTS]","ALTER TABLE","ADD","DROP COLUMN [IF EXISTS]","ALTER COLUMN","TRUNCATE TABLE","ADD SENSITIVITY CLASSIFICATION","ADD SIGNATURE","AGGREGATE","ANSI_DEFAULTS","ANSI_NULLS","ANSI_NULL_DFLT_OFF","ANSI_NULL_DFLT_ON","ANSI_PADDING","ANSI_WARNINGS","APPLICATION ROLE","ARITHABORT","ARITHIGNORE","ASSEMBLY","ASYMMETRIC KEY","AUTHORIZATION","AVAILABILITY GROUP","BACKUP","BACKUP CERTIFICATE","BACKUP MASTER KEY","BACKUP SERVICE MASTER KEY","BEGIN CONVERSATION TIMER","BEGIN DIALOG CONVERSATION","BROKER PRIORITY","BULK INSERT","CERTIFICATE","CLOSE MASTER KEY","CLOSE SYMMETRIC KEY","COLLATE","COLUMN ENCRYPTION KEY","COLUMN MASTER KEY","COLUMNSTORE INDEX","CONCAT_NULL_YIELDS_NULL","CONTEXT_INFO","CONTRACT","CREDENTIAL","CRYPTOGRAPHIC PROVIDER","CURSOR_CLOSE_ON_COMMIT","DATABASE","DATABASE AUDIT SPECIFICATION","DATABASE ENCRYPTION KEY","DATABASE HADR","DATABASE SCOPED CONFIGURATION","DATABASE SCOPED CREDENTIAL","DATABASE SET","DATEFIRST","DATEFORMAT","DEADLOCK_PRIORITY","DENY","DENY XML","DISABLE TRIGGER","ENABLE TRIGGER","END CONVERSATION","ENDPOINT","EVENT NOTIFICATION","EVENT SESSION","EXECUTE AS","EXTERNAL DATA SOURCE","EXTERNAL FILE FORMAT","EXTERNAL LANGUAGE","EXTERNAL LIBRARY","EXTERNAL RESOURCE POOL","EXTERNAL TABLE","FIPS_FLAGGER","FMTONLY","FORCEPLAN","FULLTEXT CATALOG","FULLTEXT INDEX","FULLTEXT STOPLIST","FUNCTION","GET CONVERSATION GROUP","GET_TRANSMISSION_STATUS","GRANT","GRANT XML","IDENTITY_INSERT","IMPLICIT_TRANSACTIONS","INDEX","LANGUAGE","LOCK_TIMEOUT","LOGIN","MASTER KEY","MESSAGE TYPE","MOVE CONVERSATION","NOCOUNT","NOEXEC","NUMERIC_ROUNDABORT","OFFSETS","OPEN MASTER KEY","OPEN SYMMETRIC KEY","PARSEONLY","PARTITION FUNCTION","PARTITION SCHEME","PROCEDURE","QUERY_GOVERNOR_COST_LIMIT","QUEUE","QUOTED_IDENTIFIER","RECEIVE","REMOTE SERVICE BINDING","REMOTE_PROC_TRANSACTIONS","RESOURCE GOVERNOR","RESOURCE POOL","RESTORE","RESTORE FILELISTONLY","RESTORE HEADERONLY","RESTORE LABELONLY","RESTORE MASTER KEY","RESTORE REWINDONLY","RESTORE SERVICE MASTER KEY","RESTORE VERIFYONLY","REVERT","REVOKE","REVOKE XML","ROLE","ROUTE","ROWCOUNT","RULE","SCHEMA","SEARCH PROPERTY LIST","SECURITY POLICY","SELECTIVE XML INDEX","SEND","SENSITIVITY CLASSIFICATION","SEQUENCE","SERVER AUDIT","SERVER AUDIT SPECIFICATION","SERVER CONFIGURATION","SERVER ROLE","SERVICE","SERVICE MASTER KEY","SETUSER","SHOWPLAN_ALL","SHOWPLAN_TEXT","SHOWPLAN_XML","SIGNATURE","SPATIAL INDEX","STATISTICS","STATISTICS IO","STATISTICS PROFILE","STATISTICS TIME","STATISTICS XML","SYMMETRIC KEY","SYNONYM","TABLE","TABLE IDENTITY","TEXTSIZE","TRANSACTION ISOLATION LEVEL","TRIGGER","TYPE","UPDATE STATISTICS","USER","WORKLOAD GROUP","XACT_ABORT","XML INDEX","XML SCHEMA COLLECTION"]),jo=v(["UNION [ALL]","EXCEPT","INTERSECT"]),zo=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","{CROSS | OUTER} APPLY"]),eO=v(["ON {UPDATE | DELETE} [SET NULL | SET DEFAULT]","{ROWS | RANGE} BETWEEN"]),EO={name:"transactsql",tokenizerOptions:{reservedSelect:Qo,reservedClauses:[...Zo,...Lr,...rT],reservedSetOperations:jo,reservedJoins:zo,reservedPhrases:eO,reservedKeywords:Jo,reservedDataTypes:qo,reservedFunctionNames:Ko,nestedBlockComments:!0,stringTypes:[{quote:"''-qq",prefixes:["N"]}],identTypes:['""-qq',"[]"],identChars:{first:"#@",rest:"#@$"},paramTypes:{named:["@"],quoted:["@"]},operators:["%","&","|","^","~","!<","!>","+=","-=","*=","/=","%=","|=","&=","^=","::",":"],propertyAccessOperators:[".."]},formatOptions:{alwaysDenseOperators:["::"],onelineClauses:[...Lr,...rT],tabularOnelineClauses:rT}},tO=["ADD","ALL","ALTER","ANALYZE","AND","AS","ASC","ASENSITIVE","BEFORE","BETWEEN","_BINARY","BOTH","BY","CALL","CASCADE","CASE","CHANGE","CHECK","COLLATE","COLUMN","CONDITION","CONSTRAINT","CONTINUE","CONVERT","CREATE","CROSS","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURSOR","DATABASE","DATABASES","DAY_HOUR","DAY_MICROSECOND","DAY_MINUTE","DAY_SECOND","DECLARE","DEFAULT","DELAYED","DELETE","DESC","DESCRIBE","DETERMINISTIC","DISTINCT","DISTINCTROW","DIV","DROP","DUAL","EACH","ELSE","ELSEIF","ENCLOSED","ESCAPED","EXCEPT","EXISTS","EXIT","EXPLAIN","EXTRA_JOIN","FALSE","FETCH","FOR","FORCE","FORCE_COMPILED_MODE","FORCE_INTERPRETER_MODE","FOREIGN","FROM","FULL","FULLTEXT","GRANT","GROUP","HAVING","HEARTBEAT_NO_LOGGING","HIGH_PRIORITY","HOUR_MICROSECOND","HOUR_MINUTE","HOUR_SECOND","IF","IGNORE","IN","INDEX","INFILE","INNER","INOUT","INSENSITIVE","INSERT","IN","_INTERNAL_DYNAMIC_TYPECAST","INTERSECT","INTERVAL","INTO","ITERATE","JOIN","KEY","KEYS","KILL","LEADING","LEAVE","LEFT","LIKE","LIMIT","LINES","LOAD","LOCALTIME","LOCALTIMESTAMP","LOCK","LOOP","LOW_PRIORITY","MATCH","MAXVALUE","MINUS","MINUTE_MICROSECOND","MINUTE_SECOND","MOD","MODIFIES","NATURAL","NO_QUERY_REWRITE","NOT","NO_WRITE_TO_BINLOG","NO_QUERY_REWRITE","NULL","ON","OPTIMIZE","OPTION","OPTIONALLY","OR","ORDER","OUT","OUTER","OUTFILE","OVER","PRIMARY","PROCEDURE","PURGE","RANGE","READ","READS","REFERENCES","REGEXP","RELEASE","RENAME","REPEAT","REPLACE","REQUIRE","RESTRICT","RETURN","REVOKE","RIGHT","RIGHT_ANTI_JOIN","RIGHT_SEMI_JOIN","RIGHT_STRAIGHT_JOIN","RLIKE","SCHEMA","SCHEMAS","SECOND_MICROSECOND","SELECT","SEMI_JOIN","SENSITIVE","SEPARATOR","SET","SHOW","SIGNAL","SPATIAL","SPECIFIC","SQL","SQL_BIG_RESULT","SQL_BUFFER_RESULT","SQL_CACHE","SQL_CALC_FOUND_ROWS","SQLEXCEPTION","SQL_NO_CACHE","SQL_NO_LOGGING","SQL_SMALL_RESULT","SQLSTATE","SQLWARNING","STRAIGHT_JOIN","TABLE","TERMINATED","THEN","TO","TRAILING","TRIGGER","TRUE","UNBOUNDED","UNDO","UNION","UNIQUE","UNLOCK","UPDATE","USAGE","USE","USING","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","_UTF8","VALUES","WHEN","WHERE","WHILE","WINDOW","WITH","WITHIN","WRITE","XOR","YEAR_MONTH","ZEROFILL"],TO=["BIGINT","BINARY","BIT","BLOB","CHAR","CHARACTER","DATETIME","DEC","DECIMAL","DOUBLE PRECISION","DOUBLE","ENUM","FIXED","FLOAT","FLOAT4","FLOAT8","INT","INT1","INT2","INT3","INT4","INT8","INTEGER","LONG","LONGBLOB","LONGTEXT","MEDIUMBLOB","MEDIUMINT","MEDIUMTEXT","MIDDLEINT","NATIONAL CHAR","NATIONAL VARCHAR","NUMERIC","PRECISION","REAL","SMALLINT","TEXT","TIME","TIMESTAMP","TINYBLOB","TINYINT","TINYTEXT","UNSIGNED","VARBINARY","VARCHAR","VARCHARACTER","YEAR"],rO=["ABS","ACOS","ADDDATE","ADDTIME","AES_DECRYPT","AES_ENCRYPT","ANY_VALUE","APPROX_COUNT_DISTINCT","APPROX_COUNT_DISTINCT_ACCUMULATE","APPROX_COUNT_DISTINCT_COMBINE","APPROX_COUNT_DISTINCT_ESTIMATE","APPROX_GEOGRAPHY_INTERSECTS","APPROX_PERCENTILE","ASCII","ASIN","ATAN","ATAN2","AVG","BIN","BINARY","BIT_AND","BIT_COUNT","BIT_OR","BIT_XOR","CAST","CEIL","CEILING","CHAR","CHARACTER_LENGTH","CHAR_LENGTH","CHARSET","COALESCE","COERCIBILITY","COLLATION","COLLECT","CONCAT","CONCAT_WS","CONNECTION_ID","CONV","CONVERT","CONVERT_TZ","COS","COT","COUNT","CUME_DIST","CURDATE","CURRENT_DATE","CURRENT_ROLE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","CURTIME","DATABASE","DATE","DATE_ADD","DATEDIFF","DATE_FORMAT","DATE_SUB","DATE_TRUNC","DAY","DAYNAME","DAYOFMONTH","DAYOFWEEK","DAYOFYEAR","DECODE","DEFAULT","DEGREES","DENSE_RANK","DIV","DOT_PRODUCT","ELT","EUCLIDEAN_DISTANCE","EXP","EXTRACT","FIELD","FIRST","FIRST_VALUE","FLOOR","FORMAT","FOUND_ROWS","FROM_BASE64","FROM_DAYS","FROM_UNIXTIME","GEOGRAPHY_AREA","GEOGRAPHY_CONTAINS","GEOGRAPHY_DISTANCE","GEOGRAPHY_INTERSECTS","GEOGRAPHY_LATITUDE","GEOGRAPHY_LENGTH","GEOGRAPHY_LONGITUDE","GEOGRAPHY_POINT","GEOGRAPHY_WITHIN_DISTANCE","GEOMETRY_AREA","GEOMETRY_CONTAINS","GEOMETRY_DISTANCE","GEOMETRY_FILTER","GEOMETRY_INTERSECTS","GEOMETRY_LENGTH","GEOMETRY_POINT","GEOMETRY_WITHIN_DISTANCE","GEOMETRY_X","GEOMETRY_Y","GREATEST","GROUPING","GROUP_CONCAT","HEX","HIGHLIGHT","HOUR","ICU_VERSION","IF","IFNULL","INET_ATON","INET_NTOA","INET6_ATON","INET6_NTOA","INITCAP","INSERT","INSTR","INTERVAL","IS","IS NULL","JSON_AGG","JSON_ARRAY_CONTAINS_DOUBLE","JSON_ARRAY_CONTAINS_JSON","JSON_ARRAY_CONTAINS_STRING","JSON_ARRAY_PUSH_DOUBLE","JSON_ARRAY_PUSH_JSON","JSON_ARRAY_PUSH_STRING","JSON_DELETE_KEY","JSON_EXTRACT_DOUBLE","JSON_EXTRACT_JSON","JSON_EXTRACT_STRING","JSON_EXTRACT_BIGINT","JSON_GET_TYPE","JSON_LENGTH","JSON_SET_DOUBLE","JSON_SET_JSON","JSON_SET_STRING","JSON_SPLICE_DOUBLE","JSON_SPLICE_JSON","JSON_SPLICE_STRING","LAG","LAST_DAY","LAST_VALUE","LCASE","LEAD","LEAST","LEFT","LENGTH","LIKE","LN","LOCALTIME","LOCALTIMESTAMP","LOCATE","LOG","LOG10","LOG2","LPAD","LTRIM","MATCH","MAX","MD5","MEDIAN","MICROSECOND","MIN","MINUTE","MOD","MONTH","MONTHNAME","MONTHS_BETWEEN","NOT","NOW","NTH_VALUE","NTILE","NULLIF","OCTET_LENGTH","PERCENT_RANK","PERCENTILE_CONT","PERCENTILE_DISC","PI","PIVOT","POSITION","POW","POWER","QUARTER","QUOTE","RADIANS","RAND","RANK","REGEXP","REPEAT","REPLACE","REVERSE","RIGHT","RLIKE","ROUND","ROW_COUNT","ROW_NUMBER","RPAD","RTRIM","SCALAR","SCHEMA","SEC_TO_TIME","SHA1","SHA2","SIGMOID","SIGN","SIN","SLEEP","SPLIT","SOUNDEX","SOUNDS LIKE","SOURCE_POS_WAIT","SPACE","SQRT","STDDEV","STDDEV_POP","STDDEV_SAMP","STR_TO_DATE","SUBDATE","SUBSTR","SUBSTRING","SUBSTRING_INDEX","SUM","SYS_GUID","TAN","TIME","TIMEDIFF","TIME_BUCKET","TIME_FORMAT","TIMESTAMP","TIMESTAMPADD","TIMESTAMPDIFF","TIME_TO_SEC","TO_BASE64","TO_CHAR","TO_DAYS","TO_JSON","TO_NUMBER","TO_SECONDS","TO_TIMESTAMP","TRIM","TRUNC","TRUNCATE","UCASE","UNHEX","UNIX_TIMESTAMP","UPDATEXML","UPPER","UTC_DATE","UTC_TIME","UTC_TIMESTAMP","UUID","VALUES","VARIANCE","VAR_POP","VAR_SAMP","VECTOR_SUB","VERSION","WEEK","WEEKDAY","WEEKOFYEAR","YEAR"],RO=v(["SELECT [ALL | DISTINCT | DISTINCTROW]"]),nO=v(["WITH","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER BY","LIMIT","OFFSET","INSERT [IGNORE] [INTO]","VALUES","REPLACE [INTO]","ON DUPLICATE KEY UPDATE","SET","CREATE [OR REPLACE] [TEMPORARY] PROCEDURE [IF NOT EXISTS]","CREATE [OR REPLACE] [EXTERNAL] FUNCTION"]),Cr=v(["CREATE [ROWSTORE] [REFERENCE | TEMPORARY | GLOBAL TEMPORARY] TABLE [IF NOT EXISTS]"]),RT=v(["CREATE VIEW","UPDATE","DELETE [FROM]","DROP [TEMPORARY] TABLE [IF EXISTS]","ALTER [ONLINE] TABLE","ADD [COLUMN]","ADD [UNIQUE] {INDEX | KEY}","DROP [COLUMN]","MODIFY [COLUMN]","CHANGE","RENAME [TO | AS]","TRUNCATE [TABLE]","ADD AGGREGATOR","ADD LEAF","AGGREGATOR SET AS MASTER","ALTER DATABASE","ALTER PIPELINE","ALTER RESOURCE POOL","ALTER USER","ALTER VIEW","ANALYZE TABLE","ATTACH DATABASE","ATTACH LEAF","ATTACH LEAF ALL","BACKUP DATABASE","BINLOG","BOOTSTRAP AGGREGATOR","CACHE INDEX","CALL","CHANGE","CHANGE MASTER TO","CHANGE REPLICATION FILTER","CHANGE REPLICATION SOURCE TO","CHECK BLOB CHECKSUM","CHECK TABLE","CHECKSUM TABLE","CLEAR ORPHAN DATABASES","CLONE","COMMIT","CREATE DATABASE","CREATE GROUP","CREATE INDEX","CREATE LINK","CREATE MILESTONE","CREATE PIPELINE","CREATE RESOURCE POOL","CREATE ROLE","CREATE USER","DEALLOCATE PREPARE","DESCRIBE","DETACH DATABASE","DETACH PIPELINE","DROP DATABASE","DROP FUNCTION","DROP INDEX","DROP LINK","DROP PIPELINE","DROP PROCEDURE","DROP RESOURCE POOL","DROP ROLE","DROP USER","DROP VIEW","EXECUTE","EXPLAIN","FLUSH","FORCE","GRANT","HANDLER","HELP","KILL CONNECTION","KILLALL QUERIES","LOAD DATA","LOAD INDEX INTO CACHE","LOAD XML","LOCK INSTANCE FOR BACKUP","LOCK TABLES","MASTER_POS_WAIT","OPTIMIZE TABLE","PREPARE","PURGE BINARY LOGS","REBALANCE PARTITIONS","RELEASE SAVEPOINT","REMOVE AGGREGATOR","REMOVE LEAF","REPAIR TABLE","REPLACE","REPLICATE DATABASE","RESET","RESET MASTER","RESET PERSIST","RESET REPLICA","RESET SLAVE","RESTART","RESTORE DATABASE","RESTORE REDUNDANCY","REVOKE","ROLLBACK","ROLLBACK TO SAVEPOINT","SAVEPOINT","SET CHARACTER SET","SET DEFAULT ROLE","SET NAMES","SET PASSWORD","SET RESOURCE GROUP","SET ROLE","SET TRANSACTION","SHOW","SHOW CHARACTER SET","SHOW COLLATION","SHOW COLUMNS","SHOW CREATE DATABASE","SHOW CREATE FUNCTION","SHOW CREATE PIPELINE","SHOW CREATE PROCEDURE","SHOW CREATE TABLE","SHOW CREATE USER","SHOW CREATE VIEW","SHOW DATABASES","SHOW ENGINE","SHOW ENGINES","SHOW ERRORS","SHOW FUNCTION CODE","SHOW FUNCTION STATUS","SHOW GRANTS","SHOW INDEX","SHOW MASTER STATUS","SHOW OPEN TABLES","SHOW PLUGINS","SHOW PRIVILEGES","SHOW PROCEDURE CODE","SHOW PROCEDURE STATUS","SHOW PROCESSLIST","SHOW PROFILE","SHOW PROFILES","SHOW RELAYLOG EVENTS","SHOW REPLICA STATUS","SHOW REPLICAS","SHOW SLAVE","SHOW SLAVE HOSTS","SHOW STATUS","SHOW TABLE STATUS","SHOW TABLES","SHOW VARIABLES","SHOW WARNINGS","SHUTDOWN","SNAPSHOT DATABASE","SOURCE_POS_WAIT","START GROUP_REPLICATION","START PIPELINE","START REPLICA","START SLAVE","START TRANSACTION","STOP GROUP_REPLICATION","STOP PIPELINE","STOP REPLICA","STOP REPLICATING","STOP SLAVE","TEST PIPELINE","UNLOCK INSTANCE","UNLOCK TABLES","USE","XA","ITERATE","LEAVE","LOOP","REPEAT","RETURN","WHILE"]),AO=v(["UNION [ALL | DISTINCT]","EXCEPT","INTERSECT","MINUS"]),sO=v(["JOIN","{LEFT | RIGHT | FULL} [OUTER] JOIN","{INNER | CROSS} JOIN","NATURAL {LEFT | RIGHT} [OUTER] JOIN","STRAIGHT_JOIN"]),SO=v(["ON DELETE","ON UPDATE","CHARACTER SET","{ROWS | RANGE} BETWEEN","IDENTIFIED BY"]),oO={name:"singlestoredb",tokenizerOptions:{reservedSelect:RO,reservedClauses:[...nO,...Cr,...RT],reservedSetOperations:AO,reservedJoins:sO,reservedPhrases:SO,reservedKeywords:tO,reservedDataTypes:TO,reservedFunctionNames:rO,stringTypes:['""-qq-bs',"''-qq-bs",{quote:"''-raw",prefixes:["B","X"],requirePrefix:!0}],identTypes:["``"],identChars:{first:"$",rest:"$",allowFirstCharNumber:!0},variableTypes:[{regex:"@@?[A-Za-z0-9_$]+"},{quote:"``",prefixes:["@"],requirePrefix:!0}],lineCommentTypes:["--","#"],operators:[":=","&","|","^","~","<<",">>","<=>","&&","||","::","::$","::%",":>","!:>","*.*"],postProcess:yt},formatOptions:{alwaysDenseOperators:["::","::$","::%"],onelineClauses:[...Cr,...RT],tabularOnelineClauses:RT}},OO=["ABS","ACOS","ACOSH","ADD_MONTHS","ALL_USER_NAMES","ANY_VALUE","APPROX_COUNT_DISTINCT","APPROX_PERCENTILE","APPROX_PERCENTILE_ACCUMULATE","APPROX_PERCENTILE_COMBINE","APPROX_PERCENTILE_ESTIMATE","APPROX_TOP_K","APPROX_TOP_K_ACCUMULATE","APPROX_TOP_K_COMBINE","APPROX_TOP_K_ESTIMATE","APPROXIMATE_JACCARD_INDEX","APPROXIMATE_SIMILARITY","ARRAY_AGG","ARRAY_APPEND","ARRAY_CAT","ARRAY_COMPACT","ARRAY_CONSTRUCT","ARRAY_CONSTRUCT_COMPACT","ARRAY_CONTAINS","ARRAY_INSERT","ARRAY_INTERSECTION","ARRAY_POSITION","ARRAY_PREPEND","ARRAY_SIZE","ARRAY_SLICE","ARRAY_TO_STRING","ARRAY_UNION_AGG","ARRAY_UNIQUE_AGG","ARRAYS_OVERLAP","AS_ARRAY","AS_BINARY","AS_BOOLEAN","AS_CHAR","AS_VARCHAR","AS_DATE","AS_DECIMAL","AS_NUMBER","AS_DOUBLE","AS_REAL","AS_INTEGER","AS_OBJECT","AS_TIME","AS_TIMESTAMP_LTZ","AS_TIMESTAMP_NTZ","AS_TIMESTAMP_TZ","ASCII","ASIN","ASINH","ATAN","ATAN2","ATANH","AUTO_REFRESH_REGISTRATION_HISTORY","AUTOMATIC_CLUSTERING_HISTORY","AVG","BASE64_DECODE_BINARY","BASE64_DECODE_STRING","BASE64_ENCODE","BIT_LENGTH","BITAND","BITAND_AGG","BITMAP_BIT_POSITION","BITMAP_BUCKET_NUMBER","BITMAP_CONSTRUCT_AGG","BITMAP_COUNT","BITMAP_OR_AGG","BITNOT","BITOR","BITOR_AGG","BITSHIFTLEFT","BITSHIFTRIGHT","BITXOR","BITXOR_AGG","BOOLAND","BOOLAND_AGG","BOOLNOT","BOOLOR","BOOLOR_AGG","BOOLXOR","BOOLXOR_AGG","BUILD_SCOPED_FILE_URL","BUILD_STAGE_FILE_URL","CASE","CAST","CBRT","CEIL","CHARINDEX","CHECK_JSON","CHECK_XML","CHR","CHAR","COALESCE","COLLATE","COLLATION","COMPLETE_TASK_GRAPHS","COMPRESS","CONCAT","CONCAT_WS","CONDITIONAL_CHANGE_EVENT","CONDITIONAL_TRUE_EVENT","CONTAINS","CONVERT_TIMEZONE","COPY_HISTORY","CORR","COS","COSH","COT","COUNT","COUNT_IF","COVAR_POP","COVAR_SAMP","CUME_DIST","CURRENT_ACCOUNT","CURRENT_AVAILABLE_ROLES","CURRENT_CLIENT","CURRENT_DATABASE","CURRENT_DATE","CURRENT_IP_ADDRESS","CURRENT_REGION","CURRENT_ROLE","CURRENT_SCHEMA","CURRENT_SCHEMAS","CURRENT_SECONDARY_ROLES","CURRENT_SESSION","CURRENT_STATEMENT","CURRENT_TASK_GRAPHS","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_TRANSACTION","CURRENT_USER","CURRENT_VERSION","CURRENT_WAREHOUSE","DATA_TRANSFER_HISTORY","DATABASE_REFRESH_HISTORY","DATABASE_REFRESH_PROGRESS","DATABASE_REFRESH_PROGRESS_BY_JOB","DATABASE_STORAGE_USAGE_HISTORY","DATE_FROM_PARTS","DATE_PART","DATE_TRUNC","DATEADD","DATEDIFF","DAYNAME","DECODE","DECOMPRESS_BINARY","DECOMPRESS_STRING","DECRYPT","DECRYPT_RAW","DEGREES","DENSE_RANK","DIV0","EDITDISTANCE","ENCRYPT","ENCRYPT_RAW","ENDSWITH","EQUAL_NULL","EXP","EXPLAIN_JSON","EXTERNAL_FUNCTIONS_HISTORY","EXTERNAL_TABLE_FILES","EXTERNAL_TABLE_FILE_REGISTRATION_HISTORY","EXTRACT","EXTRACT_SEMANTIC_CATEGORIES","FACTORIAL","FIRST_VALUE","FLATTEN","FLOOR","GENERATE_COLUMN_DESCRIPTION","GENERATOR","GET","GET_ABSOLUTE_PATH","GET_DDL","GET_IGNORE_CASE","GET_OBJECT_REFERENCES","GET_PATH","GET_PRESIGNED_URL","GET_RELATIVE_PATH","GET_STAGE_LOCATION","GETBIT","GREATEST","GROUPING","GROUPING_ID","HASH","HASH_AGG","HAVERSINE","HEX_DECODE_BINARY","HEX_DECODE_STRING","HEX_ENCODE","HLL","HLL_ACCUMULATE","HLL_COMBINE","HLL_ESTIMATE","HLL_EXPORT","HLL_IMPORT","HOUR","MINUTE","SECOND","IFF","IFNULL","ILIKE","ILIKE ANY","INFER_SCHEMA","INITCAP","INSERT","INVOKER_ROLE","INVOKER_SHARE","IS_ARRAY","IS_BINARY","IS_BOOLEAN","IS_CHAR","IS_VARCHAR","IS_DATE","IS_DATE_VALUE","IS_DECIMAL","IS_DOUBLE","IS_REAL","IS_GRANTED_TO_INVOKER_ROLE","IS_INTEGER","IS_NULL_VALUE","IS_OBJECT","IS_ROLE_IN_SESSION","IS_TIME","IS_TIMESTAMP_LTZ","IS_TIMESTAMP_NTZ","IS_TIMESTAMP_TZ","JAROWINKLER_SIMILARITY","JSON_EXTRACT_PATH_TEXT","KURTOSIS","LAG","LAST_DAY","LAST_QUERY_ID","LAST_TRANSACTION","LAST_VALUE","LEAD","LEAST","LEFT","LENGTH","LEN","LIKE","LIKE ALL","LIKE ANY","LISTAGG","LN","LOCALTIME","LOCALTIMESTAMP","LOG","LOGIN_HISTORY","LOGIN_HISTORY_BY_USER","LOWER","LPAD","LTRIM","MATERIALIZED_VIEW_REFRESH_HISTORY","MD5","MD5_HEX","MD5_BINARY","MD5_NUMBER — Obsoleted","MD5_NUMBER_LOWER64","MD5_NUMBER_UPPER64","MEDIAN","MIN","MAX","MINHASH","MINHASH_COMBINE","MOD","MODE","MONTHNAME","MONTHS_BETWEEN","NEXT_DAY","NORMAL","NTH_VALUE","NTILE","NULLIF","NULLIFZERO","NVL","NVL2","OBJECT_AGG","OBJECT_CONSTRUCT","OBJECT_CONSTRUCT_KEEP_NULL","OBJECT_DELETE","OBJECT_INSERT","OBJECT_KEYS","OBJECT_PICK","OCTET_LENGTH","PARSE_IP","PARSE_JSON","PARSE_URL","PARSE_XML","PERCENT_RANK","PERCENTILE_CONT","PERCENTILE_DISC","PI","PIPE_USAGE_HISTORY","POLICY_CONTEXT","POLICY_REFERENCES","POSITION","POW","POWER","PREVIOUS_DAY","QUERY_ACCELERATION_HISTORY","QUERY_HISTORY","QUERY_HISTORY_BY_SESSION","QUERY_HISTORY_BY_USER","QUERY_HISTORY_BY_WAREHOUSE","RADIANS","RANDOM","RANDSTR","RANK","RATIO_TO_REPORT","REGEXP","REGEXP_COUNT","REGEXP_INSTR","REGEXP_LIKE","REGEXP_REPLACE","REGEXP_SUBSTR","REGEXP_SUBSTR_ALL","REGR_AVGX","REGR_AVGY","REGR_COUNT","REGR_INTERCEPT","REGR_R2","REGR_SLOPE","REGR_SXX","REGR_SXY","REGR_SYY","REGR_VALX","REGR_VALY","REPEAT","REPLACE","REPLICATION_GROUP_REFRESH_HISTORY","REPLICATION_GROUP_REFRESH_PROGRESS","REPLICATION_GROUP_REFRESH_PROGRESS_BY_JOB","REPLICATION_GROUP_USAGE_HISTORY","REPLICATION_USAGE_HISTORY","REST_EVENT_HISTORY","RESULT_SCAN","REVERSE","RIGHT","RLIKE","ROUND","ROW_NUMBER","RPAD","RTRIM","RTRIMMED_LENGTH","SEARCH_OPTIMIZATION_HISTORY","SEQ1","SEQ2","SEQ4","SEQ8","SERVERLESS_TASK_HISTORY","SHA1","SHA1_HEX","SHA1_BINARY","SHA2","SHA2_HEX","SHA2_BINARY","SIGN","SIN","SINH","SKEW","SOUNDEX","SPACE","SPLIT","SPLIT_PART","SPLIT_TO_TABLE","SQRT","SQUARE","ST_AREA","ST_ASEWKB","ST_ASEWKT","ST_ASGEOJSON","ST_ASWKB","ST_ASBINARY","ST_ASWKT","ST_ASTEXT","ST_AZIMUTH","ST_CENTROID","ST_COLLECT","ST_CONTAINS","ST_COVEREDBY","ST_COVERS","ST_DIFFERENCE","ST_DIMENSION","ST_DISJOINT","ST_DISTANCE","ST_DWITHIN","ST_ENDPOINT","ST_ENVELOPE","ST_GEOGFROMGEOHASH","ST_GEOGPOINTFROMGEOHASH","ST_GEOGRAPHYFROMWKB","ST_GEOGRAPHYFROMWKT","ST_GEOHASH","ST_GEOMETRYFROMWKB","ST_GEOMETRYFROMWKT","ST_HAUSDORFFDISTANCE","ST_INTERSECTION","ST_INTERSECTS","ST_LENGTH","ST_MAKEGEOMPOINT","ST_GEOM_POINT","ST_MAKELINE","ST_MAKEPOINT","ST_POINT","ST_MAKEPOLYGON","ST_POLYGON","ST_NPOINTS","ST_NUMPOINTS","ST_PERIMETER","ST_POINTN","ST_SETSRID","ST_SIMPLIFY","ST_SRID","ST_STARTPOINT","ST_SYMDIFFERENCE","ST_UNION","ST_WITHIN","ST_X","ST_XMAX","ST_XMIN","ST_Y","ST_YMAX","ST_YMIN","STAGE_DIRECTORY_FILE_REGISTRATION_HISTORY","STAGE_STORAGE_USAGE_HISTORY","STARTSWITH","STDDEV","STDDEV_POP","STDDEV_SAMP","STRIP_NULL_VALUE","STRTOK","STRTOK_SPLIT_TO_TABLE","STRTOK_TO_ARRAY","SUBSTR","SUBSTRING","SUM","SYSDATE","SYSTEM$ABORT_SESSION","SYSTEM$ABORT_TRANSACTION","SYSTEM$AUTHORIZE_PRIVATELINK","SYSTEM$AUTHORIZE_STAGE_PRIVATELINK_ACCESS","SYSTEM$BEHAVIOR_CHANGE_BUNDLE_STATUS","SYSTEM$CANCEL_ALL_QUERIES","SYSTEM$CANCEL_QUERY","SYSTEM$CLUSTERING_DEPTH","SYSTEM$CLUSTERING_INFORMATION","SYSTEM$CLUSTERING_RATIO ","SYSTEM$CURRENT_USER_TASK_NAME","SYSTEM$DATABASE_REFRESH_HISTORY ","SYSTEM$DATABASE_REFRESH_PROGRESS","SYSTEM$DATABASE_REFRESH_PROGRESS_BY_JOB ","SYSTEM$DISABLE_BEHAVIOR_CHANGE_BUNDLE","SYSTEM$DISABLE_DATABASE_REPLICATION","SYSTEM$ENABLE_BEHAVIOR_CHANGE_BUNDLE","SYSTEM$ESTIMATE_QUERY_ACCELERATION","SYSTEM$ESTIMATE_SEARCH_OPTIMIZATION_COSTS","SYSTEM$EXPLAIN_JSON_TO_TEXT","SYSTEM$EXPLAIN_PLAN_JSON","SYSTEM$EXTERNAL_TABLE_PIPE_STATUS","SYSTEM$GENERATE_SAML_CSR","SYSTEM$GENERATE_SCIM_ACCESS_TOKEN","SYSTEM$GET_AWS_SNS_IAM_POLICY","SYSTEM$GET_PREDECESSOR_RETURN_VALUE","SYSTEM$GET_PRIVATELINK","SYSTEM$GET_PRIVATELINK_AUTHORIZED_ENDPOINTS","SYSTEM$GET_PRIVATELINK_CONFIG","SYSTEM$GET_SNOWFLAKE_PLATFORM_INFO","SYSTEM$GET_TAG","SYSTEM$GET_TAG_ALLOWED_VALUES","SYSTEM$GET_TAG_ON_CURRENT_COLUMN","SYSTEM$GET_TAG_ON_CURRENT_TABLE","SYSTEM$GLOBAL_ACCOUNT_SET_PARAMETER","SYSTEM$LAST_CHANGE_COMMIT_TIME","SYSTEM$LINK_ACCOUNT_OBJECTS_BY_NAME","SYSTEM$MIGRATE_SAML_IDP_REGISTRATION","SYSTEM$PIPE_FORCE_RESUME","SYSTEM$PIPE_STATUS","SYSTEM$REVOKE_PRIVATELINK","SYSTEM$REVOKE_STAGE_PRIVATELINK_ACCESS","SYSTEM$SET_RETURN_VALUE","SYSTEM$SHOW_OAUTH_CLIENT_SECRETS","SYSTEM$STREAM_GET_TABLE_TIMESTAMP","SYSTEM$STREAM_HAS_DATA","SYSTEM$TASK_DEPENDENTS_ENABLE","SYSTEM$TYPEOF","SYSTEM$USER_TASK_CANCEL_ONGOING_EXECUTIONS","SYSTEM$VERIFY_EXTERNAL_OAUTH_TOKEN","SYSTEM$WAIT","SYSTEM$WHITELIST","SYSTEM$WHITELIST_PRIVATELINK","TAG_REFERENCES","TAG_REFERENCES_ALL_COLUMNS","TAG_REFERENCES_WITH_LINEAGE","TAN","TANH","TASK_DEPENDENTS","TASK_HISTORY","TIME_FROM_PARTS","TIME_SLICE","TIMEADD","TIMEDIFF","TIMESTAMP_FROM_PARTS","TIMESTAMPADD","TIMESTAMPDIFF","TO_ARRAY","TO_BINARY","TO_BOOLEAN","TO_CHAR","TO_VARCHAR","TO_DATE","DATE","TO_DECIMAL","TO_NUMBER","TO_NUMERIC","TO_DOUBLE","TO_GEOGRAPHY","TO_GEOMETRY","TO_JSON","TO_OBJECT","TO_TIME","TIME","TO_TIMESTAMP","TO_TIMESTAMP_LTZ","TO_TIMESTAMP_NTZ","TO_TIMESTAMP_TZ","TO_VARIANT","TO_XML","TRANSLATE","TRIM","TRUNCATE","TRUNC","TRUNC","TRY_BASE64_DECODE_BINARY","TRY_BASE64_DECODE_STRING","TRY_CAST","TRY_HEX_DECODE_BINARY","TRY_HEX_DECODE_STRING","TRY_PARSE_JSON","TRY_TO_BINARY","TRY_TO_BOOLEAN","TRY_TO_DATE","TRY_TO_DECIMAL","TRY_TO_NUMBER","TRY_TO_NUMERIC","TRY_TO_DOUBLE","TRY_TO_GEOGRAPHY","TRY_TO_GEOMETRY","TRY_TO_TIME","TRY_TO_TIMESTAMP","TRY_TO_TIMESTAMP_LTZ","TRY_TO_TIMESTAMP_NTZ","TRY_TO_TIMESTAMP_TZ","TYPEOF","UNICODE","UNIFORM","UPPER","UUID_STRING","VALIDATE","VALIDATE_PIPE_LOAD","VAR_POP","VAR_SAMP","VARIANCE","VARIANCE_SAMP","VARIANCE_POP","WAREHOUSE_LOAD_HISTORY","WAREHOUSE_METERING_HISTORY","WIDTH_BUCKET","XMLGET","YEAR","YEAROFWEEK","YEAROFWEEKISO","DAY","DAYOFMONTH","DAYOFWEEK","DAYOFWEEKISO","DAYOFYEAR","WEEK","WEEK","WEEKOFYEAR","WEEKISO","MONTH","QUARTER","ZEROIFNULL","ZIPF"],iO=["ACCOUNT","ALL","ALTER","AND","ANY","AS","BETWEEN","BY","CASE","CAST","CHECK","COLUMN","CONNECT","CONNECTION","CONSTRAINT","CREATE","CROSS","CURRENT","CURRENT_DATE","CURRENT_TIME","CURRENT_TIMESTAMP","CURRENT_USER","DATABASE","DELETE","DISTINCT","DROP","ELSE","EXISTS","FALSE","FOLLOWING","FOR","FROM","FULL","GRANT","GROUP","GSCLUSTER","HAVING","ILIKE","IN","INCREMENT","INNER","INSERT","INTERSECT","INTO","IS","ISSUE","JOIN","LATERAL","LEFT","LIKE","LOCALTIME","LOCALTIMESTAMP","MINUS","NATURAL","NOT","NULL","OF","ON","OR","ORDER","ORGANIZATION","QUALIFY","REGEXP","REVOKE","RIGHT","RLIKE","ROW","ROWS","SAMPLE","SCHEMA","SELECT","SET","SOME","START","TABLE","TABLESAMPLE","THEN","TO","TRIGGER","TRUE","TRY_CAST","UNION","UNIQUE","UPDATE","USING","VALUES","VIEW","WHEN","WHENEVER","WHERE","WITH","COMMENT"],aO=["NUMBER","DECIMAL","NUMERIC","INT","INTEGER","BIGINT","SMALLINT","TINYINT","BYTEINT","FLOAT","FLOAT4","FLOAT8","DOUBLE","DOUBLE PRECISION","REAL","VARCHAR","CHAR","CHARACTER","STRING","TEXT","BINARY","VARBINARY","BOOLEAN","DATE","DATETIME","TIME","TIMESTAMP","TIMESTAMP_LTZ","TIMESTAMP_NTZ","TIMESTAMP","TIMESTAMP_TZ","VARIANT","OBJECT","ARRAY","GEOGRAPHY","GEOMETRY"],IO=v(["SELECT [ALL | DISTINCT]"]),NO=v(["WITH [RECURSIVE]","FROM","WHERE","GROUP BY","HAVING","PARTITION BY","ORDER BY","QUALIFY","LIMIT","OFFSET","FETCH [FIRST | NEXT]","INSERT [OVERWRITE] [ALL INTO | INTO | ALL | FIRST]","{THEN | ELSE} INTO","VALUES","SET","CLUSTER BY","[WITH] {MASKING POLICY | TAG | ROW ACCESS POLICY}","COPY GRANTS","USING TEMPLATE","MERGE INTO","WHEN MATCHED [AND]","THEN {UPDATE SET | DELETE}","WHEN NOT MATCHED THEN INSERT"]),ur=v(["CREATE [OR REPLACE] [VOLATILE] TABLE [IF NOT EXISTS]","CREATE [OR REPLACE] [LOCAL | GLOBAL] {TEMP|TEMPORARY} TABLE [IF NOT EXISTS]"]),nT=v(["CREATE [OR REPLACE] [SECURE] [RECURSIVE] VIEW [IF NOT EXISTS]","UPDATE","DELETE FROM","DROP TABLE [IF EXISTS]","ALTER TABLE [IF EXISTS]","RENAME TO","SWAP WITH","[SUSPEND | RESUME] RECLUSTER","DROP CLUSTERING KEY","ADD [COLUMN]","RENAME COLUMN","{ALTER | MODIFY} [COLUMN]","DROP [COLUMN]","{ADD | ALTER | MODIFY | DROP} [CONSTRAINT]","RENAME CONSTRAINT","{ADD | DROP} SEARCH OPTIMIZATION","{SET | UNSET} TAG","{ADD | DROP} ROW ACCESS POLICY","DROP ALL ROW ACCESS POLICIES","{SET | DROP} DEFAULT","{SET | DROP} NOT NULL","[SET DATA] TYPE","UNSET COMMENT","{SET | UNSET} MASKING POLICY","TRUNCATE [TABLE] [IF EXISTS]","ALTER ACCOUNT","ALTER API INTEGRATION","ALTER CONNECTION","ALTER DATABASE","ALTER EXTERNAL TABLE","ALTER FAILOVER GROUP","ALTER FILE FORMAT","ALTER FUNCTION","ALTER INTEGRATION","ALTER MASKING POLICY","ALTER MATERIALIZED VIEW","ALTER NETWORK POLICY","ALTER NOTIFICATION INTEGRATION","ALTER PIPE","ALTER PROCEDURE","ALTER REPLICATION GROUP","ALTER RESOURCE MONITOR","ALTER ROLE","ALTER ROW ACCESS POLICY","ALTER SCHEMA","ALTER SECURITY INTEGRATION","ALTER SEQUENCE","ALTER SESSION","ALTER SESSION POLICY","ALTER SHARE","ALTER STAGE","ALTER STORAGE INTEGRATION","ALTER STREAM","ALTER TAG","ALTER TASK","ALTER USER","ALTER VIEW","ALTER WAREHOUSE","BEGIN","CALL","COMMIT","COPY INTO","CREATE ACCOUNT","CREATE API INTEGRATION","CREATE CONNECTION","CREATE DATABASE","CREATE EXTERNAL FUNCTION","CREATE EXTERNAL TABLE","CREATE FAILOVER GROUP","CREATE FILE FORMAT","CREATE FUNCTION","CREATE INTEGRATION","CREATE MANAGED ACCOUNT","CREATE MASKING POLICY","CREATE MATERIALIZED VIEW","CREATE NETWORK POLICY","CREATE NOTIFICATION INTEGRATION","CREATE PIPE","CREATE PROCEDURE","CREATE REPLICATION GROUP","CREATE RESOURCE MONITOR","CREATE ROLE","CREATE ROW ACCESS POLICY","CREATE SCHEMA","CREATE SECURITY INTEGRATION","CREATE SEQUENCE","CREATE SESSION POLICY","CREATE SHARE","CREATE STAGE","CREATE STORAGE INTEGRATION","CREATE STREAM","CREATE TAG","CREATE TASK","CREATE USER","CREATE WAREHOUSE","DELETE","DESCRIBE DATABASE","DESCRIBE EXTERNAL TABLE","DESCRIBE FILE FORMAT","DESCRIBE FUNCTION","DESCRIBE INTEGRATION","DESCRIBE MASKING POLICY","DESCRIBE MATERIALIZED VIEW","DESCRIBE NETWORK POLICY","DESCRIBE PIPE","DESCRIBE PROCEDURE","DESCRIBE RESULT","DESCRIBE ROW ACCESS POLICY","DESCRIBE SCHEMA","DESCRIBE SEQUENCE","DESCRIBE SESSION POLICY","DESCRIBE SHARE","DESCRIBE STAGE","DESCRIBE STREAM","DESCRIBE TABLE","DESCRIBE TASK","DESCRIBE TRANSACTION","DESCRIBE USER","DESCRIBE VIEW","DESCRIBE WAREHOUSE","DROP CONNECTION","DROP DATABASE","DROP EXTERNAL TABLE","DROP FAILOVER GROUP","DROP FILE FORMAT","DROP FUNCTION","DROP INTEGRATION","DROP MANAGED ACCOUNT","DROP MASKING POLICY","DROP MATERIALIZED VIEW","DROP NETWORK POLICY","DROP PIPE","DROP PROCEDURE","DROP REPLICATION GROUP","DROP RESOURCE MONITOR","DROP ROLE","DROP ROW ACCESS POLICY","DROP SCHEMA","DROP SEQUENCE","DROP SESSION POLICY","DROP SHARE","DROP STAGE","DROP STREAM","DROP TAG","DROP TASK","DROP USER","DROP VIEW","DROP WAREHOUSE","EXECUTE IMMEDIATE","EXECUTE TASK","EXPLAIN","GET","GRANT OWNERSHIP","GRANT ROLE","INSERT","LIST","MERGE","PUT","REMOVE","REVOKE ROLE","ROLLBACK","SHOW COLUMNS","SHOW CONNECTIONS","SHOW DATABASES","SHOW DATABASES IN FAILOVER GROUP","SHOW DATABASES IN REPLICATION GROUP","SHOW DELEGATED AUTHORIZATIONS","SHOW EXTERNAL FUNCTIONS","SHOW EXTERNAL TABLES","SHOW FAILOVER GROUPS","SHOW FILE FORMATS","SHOW FUNCTIONS","SHOW GLOBAL ACCOUNTS","SHOW GRANTS","SHOW INTEGRATIONS","SHOW LOCKS","SHOW MANAGED ACCOUNTS","SHOW MASKING POLICIES","SHOW MATERIALIZED VIEWS","SHOW NETWORK POLICIES","SHOW OBJECTS","SHOW ORGANIZATION ACCOUNTS","SHOW PARAMETERS","SHOW PIPES","SHOW PRIMARY KEYS","SHOW PROCEDURES","SHOW REGIONS","SHOW REPLICATION ACCOUNTS","SHOW REPLICATION DATABASES","SHOW REPLICATION GROUPS","SHOW RESOURCE MONITORS","SHOW ROLES","SHOW ROW ACCESS POLICIES","SHOW SCHEMAS","SHOW SEQUENCES","SHOW SESSION POLICIES","SHOW SHARES","SHOW SHARES IN FAILOVER GROUP","SHOW SHARES IN REPLICATION GROUP","SHOW STAGES","SHOW STREAMS","SHOW TABLES","SHOW TAGS","SHOW TASKS","SHOW TRANSACTIONS","SHOW USER FUNCTIONS","SHOW USERS","SHOW VARIABLES","SHOW VIEWS","SHOW WAREHOUSES","TRUNCATE MATERIALIZED VIEW","UNDROP DATABASE","UNDROP SCHEMA","UNDROP TABLE","UNDROP TAG","UNSET","USE DATABASE","USE ROLE","USE SCHEMA","USE SECONDARY ROLES","USE WAREHOUSE"]),lO=v(["UNION [ALL]","MINUS","EXCEPT","INTERSECT"]),_O=v(["[INNER] JOIN","[NATURAL] {LEFT | RIGHT | FULL} [OUTER] JOIN","{CROSS | NATURAL} JOIN"]),LO=v(["{ROWS | RANGE} BETWEEN","ON {UPDATE | DELETE} [SET NULL | SET DEFAULT]"]),CO={name:"snowflake",tokenizerOptions:{reservedSelect:IO,reservedClauses:[...NO,...ur,...nT],reservedSetOperations:lO,reservedJoins:_O,reservedPhrases:LO,reservedKeywords:iO,reservedDataTypes:aO,reservedFunctionNames:OO,stringTypes:["$$","''-qq-bs"],identTypes:['""-qq'],variableTypes:[{regex:"[$][1-9]\\\\d*"},{regex:"[$][_a-zA-Z][_a-zA-Z0-9$]*"}],extraParens:["[]"],identChars:{rest:"$"},lineCommentTypes:["--","//"],operators:["%","::","||","=>"],propertyAccessOperators:[":"]},formatOptions:{alwaysDenseOperators:["::"],onelineClauses:[...ur,...nT],tabularOnelineClauses:nT}},ot=E=>E[E.length-1],$R=E=>E.sort((e,T)=>T.length-e.length||e.localeCompare(T)),Dt=E=>E.replace(/\\s+/gu," "),AT=E=>/\\n/.test(E),CE=E=>E.replace(/[.*+?^${}()|[\\]\\\\]/gu,"\\\\$&"),cr=/\\s+/uy,wE=E=>new RegExp(`(?:${E})`,"uy"),uO=E=>E.split("").map(e=>/ /gu.test(e)?"\\\\s+":`[${e.toUpperCase()}${e.toLowerCase()}]`).join(""),cO=E=>E+"(?:-"+E+")*",fO=({prefixes:E,requirePrefix:e})=>`(?:${E.map(uO).join("|")}${e?"":"|"})`,PO=E=>new RegExp(`(?:${E.map(CE).join("|")}).*?(?=\\r
|\\r|
|$)`,"uy"),fr=(E,e=[])=>{const T=E==="open"?0:1,t=["()",...e].map(r=>r[T]);return wE(t.map(CE).join("|"))},Pr=E=>wE(`${$R(E).map(CE).join("|")}`),DO=({rest:E,dashes:e})=>E||e?`(?![${E||""}${e?"-":""}])`:"",hE=(E,e={})=>{if(E.length===0)return/^\\b$/u;const T=DO(e),t=$R(E).map(CE).join("|").replace(/ /gu,"\\\\s+");return new RegExp(`(?:${t})${T}\\\\b`,"iuy")},sT=(E,e)=>{if(!E.length)return;const T=E.map(CE).join("|");return wE(`(?:${T})(?:${e})`)},dO=()=>{const E={"<":">","[":"]","(":")","{":"}"},e="{left}(?:(?!{right}').)*?{right}",T=Object.entries(E).map(([n,s])=>e.replace(/{left}/g,CE(n)).replace(/{right}/g,CE(s))),t=CE(Object.keys(E).join(""));return`[Qq]'(?:${String.raw`(?<tag>[^\\s${t}])(?:(?!\\k<tag>').)*?\\k<tag>`}|${T.join("|")})'`},Dr={"``":"(?:`[^`]*`)+","[]":String.raw`(?:\\[[^\\]]*\\])(?:\\][^\\]]*\\])*`,'""-qq':String.raw`(?:"[^"]*")+`,'""-bs':String.raw`(?:"[^"\\\\]*(?:\\\\.[^"\\\\]*)*")`,'""-qq-bs':String.raw`(?:"[^"\\\\]*(?:\\\\.[^"\\\\]*)*")+`,'""-raw':String.raw`(?:"[^"]*")`,"''-qq":String.raw`(?:'[^']*')+`,"''-bs":String.raw`(?:'[^'\\\\]*(?:\\\\.[^'\\\\]*)*')`,"''-qq-bs":String.raw`(?:'[^'\\\\]*(?:\\\\.[^'\\\\]*)*')+`,"''-raw":String.raw`(?:'[^']*')`,$$:String.raw`(?<tag>\\$\\w*\\$)[\\s\\S]*?\\k<tag>`,"\'\'\'..\'\'\'":String.raw`\'\'\'[^\\\\]*?(?:\\\\.[^\\\\]*?)*?\'\'\'`,'""".."""':String.raw`"""[^\\\\]*?(?:\\\\.[^\\\\]*?)*?"""`,"{}":String.raw`(?:\\{[^\\}]*\\})`,"q''":dO()},xR=E=>typeof E=="string"?Dr[E]:"regex"in E?E.regex:fO(E)+Dr[E.quote],pO=E=>wE(E.map(e=>"regex"in e?e.regex:xR(e)).join("|")),XR=E=>E.map(xR).join("|"),dr=E=>wE(XR(E)),MO=(E={})=>wE(kR(E)),kR=({first:E,rest:e,dashes:T,allowFirstCharNumber:t}={})=>{const r="\\\\p{Alphabetic}\\\\p{Mark}_",R="\\\\p{Decimal_Number}",n=CE(E??""),s=CE(e??""),S=t?`[${r}${R}${n}][${r}${R}${s}]*`:`[${r}${n}][${r}${R}${s}]*`;return T?cO(S):S};function KR(E,e){const T=E.slice(0,e).split(/\\n/);return{line:T.length,col:T[T.length-1].length+1}}var UO=class{constructor(E,e){this.rules=E,this.dialectName=e,this.input="",this.index=0}tokenize(E){this.input=E,this.index=0;const e=[];let T;for(;this.index<this.input.length;){const t=this.getWhitespace();if(this.index<this.input.length){if(T=this.getNextToken(),!T)throw this.createParseError();e.push(AE(EE({},T),{precedingWhitespace:t}))}}return e}createParseError(){const E=this.input.slice(this.index,this.index+10),{line:e,col:T}=KR(this.input,this.index);return new Error(`Parse error: Unexpected "${E}" at line ${e} column ${T}.
${this.dialectInfo()}`)}dialectInfo(){return this.dialectName==="sql"?`This likely happens because you're using the default "sql" dialect.
If possible, please select a more specific dialect (like sqlite, postgresql, etc).`:`SQL dialect used: "${this.dialectName}".`}getWhitespace(){cr.lastIndex=this.index;const E=cr.exec(this.input);if(E)return this.index+=E[0].length,E[0]}getNextToken(){for(const E of this.rules){const e=this.match(E);if(e)return e}}match(E){E.regex.lastIndex=this.index;const e=E.regex.exec(this.input);if(e){const T=e[0],t={type:E.type,raw:T,text:E.text?E.text(T):T,start:this.index};return E.key&&(t.key=E.key(T)),this.index+=T.length,t}}},pr=/\\/\\*/uy,mO=/([^/*]|\\*[^/]|\\/[^*])+/uy,hO=/\\*\\//uy,GO=class{constructor(){this.lastIndex=0}exec(E){let e="",T,t=0;if(T=this.matchSection(pr,E))e+=T,t++;else return null;for(;t>0;)if(T=this.matchSection(pr,E))e+=T,t++;else if(T=this.matchSection(hO,E))e+=T,t--;else if(T=this.matchSection(mO,E))e+=T;else return null;return[e]}matchSection(E,e){E.lastIndex=this.lastIndex;const T=E.exec(e);return T&&(this.lastIndex+=T[0].length),T?T[0]:null}},gO=class{constructor(E,e){this.cfg=E,this.dialectName=e,this.rulesBeforeParams=this.buildRulesBeforeParams(E),this.rulesAfterParams=this.buildRulesAfterParams(E)}tokenize(E,e){const T=[...this.rulesBeforeParams,...this.buildParamRules(this.cfg,e),...this.rulesAfterParams],t=new UO(T,this.dialectName).tokenize(E);return this.cfg.postProcess?this.cfg.postProcess(t):t}buildRulesBeforeParams(E){var e,T;return this.validRules([{type:"BLOCK_COMMENT",regex:/(\\/\\* *sql-formatter-disable *\\*\\/[\\s\\S]*?(?:\\/\\* *sql-formatter-enable *\\*\\/|$))/uy},{type:"BLOCK_COMMENT",regex:E.nestedBlockComments?new GO:/(\\/\\*[^]*?\\*\\/)/uy},{type:"LINE_COMMENT",regex:PO((e=E.lineCommentTypes)!=null?e:["--"])},{type:"QUOTED_IDENTIFIER",regex:dr(E.identTypes)},{type:"NUMBER",regex:/(?:0x[0-9a-fA-F]+|0b[01]+|(?:-\\s*)?[0-9]+(?:\\.[0-9]*)?(?:[eE][-+]?[0-9]+(?:\\.[0-9]+)?)?)(?![\\w\\p{Alphabetic}])/uy},{type:"RESERVED_PHRASE",regex:hE((T=E.reservedPhrases)!=null?T:[],E.identChars),text:Je},{type:"CASE",regex:/CASE\\b/iuy,text:Je},{type:"END",regex:/END\\b/iuy,text:Je},{type:"BETWEEN",regex:/BETWEEN\\b/iuy,text:Je},{type:"LIMIT",regex:E.reservedClauses.includes("LIMIT")?/LIMIT\\b/iuy:void 0,text:Je},{type:"RESERVED_CLAUSE",regex:hE(E.reservedClauses,E.identChars),text:Je},{type:"RESERVED_SELECT",regex:hE(E.reservedSelect,E.identChars),text:Je},{type:"RESERVED_SET_OPERATION",regex:hE(E.reservedSetOperations,E.identChars),text:Je},{type:"WHEN",regex:/WHEN\\b/iuy,text:Je},{type:"ELSE",regex:/ELSE\\b/iuy,text:Je},{type:"THEN",regex:/THEN\\b/iuy,text:Je},{type:"RESERVED_JOIN",regex:hE(E.reservedJoins,E.identChars),text:Je},{type:"AND",regex:/AND\\b/iuy,text:Je},{type:"OR",regex:/OR\\b/iuy,text:Je},{type:"XOR",regex:E.supportsXor?/XOR\\b/iuy:void 0,text:Je},{type:"RESERVED_FUNCTION_NAME",regex:hE(E.reservedFunctionNames,E.identChars),text:Je},{type:"RESERVED_DATA_TYPE",regex:hE(E.reservedDataTypes,E.identChars),text:Je},{type:"RESERVED_KEYWORD",regex:hE(E.reservedKeywords,E.identChars),text:Je}])}buildRulesAfterParams(E){var e,T;return this.validRules([{type:"VARIABLE",regex:E.variableTypes?pO(E.variableTypes):void 0},{type:"STRING",regex:dr(E.stringTypes)},{type:"IDENTIFIER",regex:MO(E.identChars)},{type:"DELIMITER",regex:/[;]/uy},{type:"COMMA",regex:/[,]/y},{type:"OPEN_PAREN",regex:fr("open",E.extraParens)},{type:"CLOSE_PAREN",regex:fr("close",E.extraParens)},{type:"OPERATOR",regex:Pr(["+","-","/",">","<","=","<>","<=",">=","!=",...(e=E.operators)!=null?e:[]])},{type:"ASTERISK",regex:/[*]/uy},{type:"PROPERTY_ACCESS_OPERATOR",regex:Pr([".",...(T=E.propertyAccessOperators)!=null?T:[]])}])}buildParamRules(E,e){var T,t,r,R,n;const s={named:(e==null?void 0:e.named)||((T=E.paramTypes)==null?void 0:T.named)||[],quoted:(e==null?void 0:e.quoted)||((t=E.paramTypes)==null?void 0:t.quoted)||[],numbered:(e==null?void 0:e.numbered)||((r=E.paramTypes)==null?void 0:r.numbered)||[],positional:typeof(e==null?void 0:e.positional)=="boolean"?e.positional:(R=E.paramTypes)==null?void 0:R.positional,custom:(e==null?void 0:e.custom)||((n=E.paramTypes)==null?void 0:n.custom)||[]};return this.validRules([{type:"NAMED_PARAMETER",regex:sT(s.named,kR(E.paramChars||E.identChars)),key:S=>S.slice(1)},{type:"QUOTED_PARAMETER",regex:sT(s.quoted,XR(E.identTypes)),key:S=>(({tokenKey:A,quoteChar:o})=>A.replace(new RegExp(CE("\\\\"+o),"gu"),o))({tokenKey:S.slice(2,-1),quoteChar:S.slice(-1)})},{type:"NUMBERED_PARAMETER",regex:sT(s.numbered,"[0-9]+"),key:S=>S.slice(1)},{type:"POSITIONAL_PARAMETER",regex:s.positional?/[?]/y:void 0},...s.custom.map(S=>{var A;return{type:"CUSTOM_PARAMETER",regex:wE(S.regex),key:(A=S.key)!=null?A:o=>o}})])}validRules(E){return E.filter(e=>!!e.regex)}},Je=E=>Dt(E.toUpperCase()),Mr=new Map,HO=E=>{let e=Mr.get(E);return e||(e=bO(E),Mr.set(E,e)),e},bO=E=>({tokenizer:new gO(E.tokenizerOptions,E.name),formatOptions:yO(E.formatOptions)}),yO=E=>{var e;return{alwaysDenseOperators:E.alwaysDenseOperators||[],onelineClauses:Object.fromEntries(E.onelineClauses.map(T=>[T,!0])),tabularOnelineClauses:Object.fromEntries(((e=E.tabularOnelineClauses)!=null?e:E.onelineClauses).map(T=>[T,!0]))}};function BO(E){return E.indentStyle==="tabularLeft"||E.indentStyle==="tabularRight"?" ".repeat(10):E.useTabs?"	":" ".repeat(E.tabWidth)}function ZE(E){return E.indentStyle==="tabularLeft"||E.indentStyle==="tabularRight"}var vO=class{constructor(E){this.params=E,this.index=0}get({key:E,text:e}){return this.params?E?this.params[E]:this.params[this.index++]:e}getPositionalParameterIndex(){return this.index}setPositionalParameterIndex(E){this.index=E}};function FO(E){return E.map(YO).map(VO).map(WO).map(wO).map($O)}var YO=(E,e,T)=>{if(wR(E.type)){const t=xO(T,e);if(t&&t.type==="PROPERTY_ACCESS_OPERATOR")return AE(EE({},E),{type:"IDENTIFIER",text:E.raw})}return E},VO=(E,e,T)=>{if(E.type==="RESERVED_FUNCTION_NAME"){const t=It(T,e);if(!t||!JR(t))return AE(EE({},E),{type:"RESERVED_KEYWORD"})}return E},WO=(E,e,T)=>{if(E.type==="RESERVED_DATA_TYPE"){const t=It(T,e);if(t&&JR(t))return AE(EE({},E),{type:"RESERVED_PARAMETERIZED_DATA_TYPE"})}return E},wO=(E,e,T)=>{if(E.type==="IDENTIFIER"){const t=It(T,e);if(t&&qR(t))return AE(EE({},E),{type:"ARRAY_IDENTIFIER"})}return E},$O=(E,e,T)=>{if(E.type==="RESERVED_DATA_TYPE"){const t=It(T,e);if(t&&qR(t))return AE(EE({},E),{type:"ARRAY_KEYWORD"})}return E},xO=(E,e)=>It(E,e,-1),It=(E,e,T=1)=>{let t=1;for(;E[e+t*T]&&XO(E[e+t*T]);)t++;return E[e+t*T]},JR=E=>E.type==="OPEN_PAREN"&&E.text==="(",qR=E=>E.type==="OPEN_PAREN"&&E.text==="[",XO=E=>E.type==="BLOCK_COMMENT"||E.type==="LINE_COMMENT",QR=class{constructor(E){this.tokenize=E,this.index=0,this.tokens=[],this.input=""}reset(E,e){this.input=E,this.index=0,this.tokens=this.tokenize(E)}next(){return this.tokens[this.index++]}save(){}formatError(E){const{line:e,col:T}=KR(this.input,E.start);return`Parse error at token: ${E.text} at line ${e} column ${T}`}has(E){return E in VR}};function ST(E){return E[0]}var Re=new QR(E=>[]),yE=([[E]])=>E,qe=E=>({type:"keyword",tokenType:E.type,text:E.text,raw:E.raw}),Ur=E=>({type:"data_type",text:E.text,raw:E.raw}),Qe=(E,{leading:e,trailing:T})=>(e!=null&&e.length&&(E=AE(EE({},E),{leadingComments:e})),T!=null&&T.length&&(E=AE(EE({},E),{trailingComments:T})),E),kO=(E,{leading:e,trailing:T})=>{if(e!=null&&e.length){const[t,...r]=E;E=[Qe(t,{leading:e}),...r]}if(T!=null&&T.length){const t=E.slice(0,-1),r=E[E.length-1];E=[...t,Qe(r,{trailing:T})]}return E},KO={Lexer:Re,ParserRules:[{name:"main$ebnf$1",symbols:[]},{name:"main$ebnf$1",symbols:["main$ebnf$1","statement"],postprocess:E=>E[0].concat([E[1]])},{name:"main",symbols:["main$ebnf$1"],postprocess:([E])=>{const e=E[E.length-1];return e&&!e.hasSemicolon?e.children.length>0?E:E.slice(0,-1):E}},{name:"statement$subexpression$1",symbols:[Re.has("DELIMITER")?{type:"DELIMITER"}:DELIMITER]},{name:"statement$subexpression$1",symbols:[Re.has("EOF")?{type:"EOF"}:EOF]},{name:"statement",symbols:["expressions_or_clauses","statement$subexpression$1"],postprocess:([E,[e]])=>({type:"statement",children:E,hasSemicolon:e.type==="DELIMITER"})},{name:"expressions_or_clauses$ebnf$1",symbols:[]},{name:"expressions_or_clauses$ebnf$1",symbols:["expressions_or_clauses$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"expressions_or_clauses$ebnf$2",symbols:[]},{name:"expressions_or_clauses$ebnf$2",symbols:["expressions_or_clauses$ebnf$2","clause"],postprocess:E=>E[0].concat([E[1]])},{name:"expressions_or_clauses",symbols:["expressions_or_clauses$ebnf$1","expressions_or_clauses$ebnf$2"],postprocess:([E,e])=>[...E,...e]},{name:"clause$subexpression$1",symbols:["limit_clause"]},{name:"clause$subexpression$1",symbols:["select_clause"]},{name:"clause$subexpression$1",symbols:["other_clause"]},{name:"clause$subexpression$1",symbols:["set_operation"]},{name:"clause",symbols:["clause$subexpression$1"],postprocess:yE},{name:"limit_clause$ebnf$1$subexpression$1$ebnf$1",symbols:["free_form_sql"]},{name:"limit_clause$ebnf$1$subexpression$1$ebnf$1",symbols:["limit_clause$ebnf$1$subexpression$1$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"limit_clause$ebnf$1$subexpression$1",symbols:[Re.has("COMMA")?{type:"COMMA"}:COMMA,"limit_clause$ebnf$1$subexpression$1$ebnf$1"]},{name:"limit_clause$ebnf$1",symbols:["limit_clause$ebnf$1$subexpression$1"],postprocess:ST},{name:"limit_clause$ebnf$1",symbols:[],postprocess:()=>null},{name:"limit_clause",symbols:[Re.has("LIMIT")?{type:"LIMIT"}:LIMIT,"_","expression_chain_","limit_clause$ebnf$1"],postprocess:([E,e,T,t])=>{if(t){const[r,R]=t;return{type:"limit_clause",limitKw:Qe(qe(E),{trailing:e}),offset:T,count:R}}else return{type:"limit_clause",limitKw:Qe(qe(E),{trailing:e}),count:T}}},{name:"select_clause$subexpression$1$ebnf$1",symbols:[]},{name:"select_clause$subexpression$1$ebnf$1",symbols:["select_clause$subexpression$1$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"select_clause$subexpression$1",symbols:["all_columns_asterisk","select_clause$subexpression$1$ebnf$1"]},{name:"select_clause$subexpression$1$ebnf$2",symbols:[]},{name:"select_clause$subexpression$1$ebnf$2",symbols:["select_clause$subexpression$1$ebnf$2","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"select_clause$subexpression$1",symbols:["asteriskless_free_form_sql","select_clause$subexpression$1$ebnf$2"]},{name:"select_clause",symbols:[Re.has("RESERVED_SELECT")?{type:"RESERVED_SELECT"}:RESERVED_SELECT,"select_clause$subexpression$1"],postprocess:([E,[e,T]])=>({type:"clause",nameKw:qe(E),children:[e,...T]})},{name:"select_clause",symbols:[Re.has("RESERVED_SELECT")?{type:"RESERVED_SELECT"}:RESERVED_SELECT],postprocess:([E])=>({type:"clause",nameKw:qe(E),children:[]})},{name:"all_columns_asterisk",symbols:[Re.has("ASTERISK")?{type:"ASTERISK"}:ASTERISK],postprocess:()=>({type:"all_columns_asterisk"})},{name:"other_clause$ebnf$1",symbols:[]},{name:"other_clause$ebnf$1",symbols:["other_clause$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"other_clause",symbols:[Re.has("RESERVED_CLAUSE")?{type:"RESERVED_CLAUSE"}:RESERVED_CLAUSE,"other_clause$ebnf$1"],postprocess:([E,e])=>({type:"clause",nameKw:qe(E),children:e})},{name:"set_operation$ebnf$1",symbols:[]},{name:"set_operation$ebnf$1",symbols:["set_operation$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"set_operation",symbols:[Re.has("RESERVED_SET_OPERATION")?{type:"RESERVED_SET_OPERATION"}:RESERVED_SET_OPERATION,"set_operation$ebnf$1"],postprocess:([E,e])=>({type:"set_operation",nameKw:qe(E),children:e})},{name:"expression_chain_$ebnf$1",symbols:["expression_with_comments_"]},{name:"expression_chain_$ebnf$1",symbols:["expression_chain_$ebnf$1","expression_with_comments_"],postprocess:E=>E[0].concat([E[1]])},{name:"expression_chain_",symbols:["expression_chain_$ebnf$1"],postprocess:ST},{name:"expression_chain$ebnf$1",symbols:[]},{name:"expression_chain$ebnf$1",symbols:["expression_chain$ebnf$1","_expression_with_comments"],postprocess:E=>E[0].concat([E[1]])},{name:"expression_chain",symbols:["expression","expression_chain$ebnf$1"],postprocess:([E,e])=>[E,...e]},{name:"andless_expression_chain$ebnf$1",symbols:[]},{name:"andless_expression_chain$ebnf$1",symbols:["andless_expression_chain$ebnf$1","_andless_expression_with_comments"],postprocess:E=>E[0].concat([E[1]])},{name:"andless_expression_chain",symbols:["andless_expression","andless_expression_chain$ebnf$1"],postprocess:([E,e])=>[E,...e]},{name:"expression_with_comments_",symbols:["expression","_"],postprocess:([E,e])=>Qe(E,{trailing:e})},{name:"_expression_with_comments",symbols:["_","expression"],postprocess:([E,e])=>Qe(e,{leading:E})},{name:"_andless_expression_with_comments",symbols:["_","andless_expression"],postprocess:([E,e])=>Qe(e,{leading:E})},{name:"free_form_sql$subexpression$1",symbols:["asteriskless_free_form_sql"]},{name:"free_form_sql$subexpression$1",symbols:["asterisk"]},{name:"free_form_sql",symbols:["free_form_sql$subexpression$1"],postprocess:yE},{name:"asteriskless_free_form_sql$subexpression$1",symbols:["asteriskless_andless_expression"]},{name:"asteriskless_free_form_sql$subexpression$1",symbols:["logic_operator"]},{name:"asteriskless_free_form_sql$subexpression$1",symbols:["comma"]},{name:"asteriskless_free_form_sql$subexpression$1",symbols:["comment"]},{name:"asteriskless_free_form_sql$subexpression$1",symbols:["other_keyword"]},{name:"asteriskless_free_form_sql",symbols:["asteriskless_free_form_sql$subexpression$1"],postprocess:yE},{name:"expression$subexpression$1",symbols:["andless_expression"]},{name:"expression$subexpression$1",symbols:["logic_operator"]},{name:"expression",symbols:["expression$subexpression$1"],postprocess:yE},{name:"andless_expression$subexpression$1",symbols:["asteriskless_andless_expression"]},{name:"andless_expression$subexpression$1",symbols:["asterisk"]},{name:"andless_expression",symbols:["andless_expression$subexpression$1"],postprocess:yE},{name:"asteriskless_andless_expression$subexpression$1",symbols:["atomic_expression"]},{name:"asteriskless_andless_expression$subexpression$1",symbols:["between_predicate"]},{name:"asteriskless_andless_expression$subexpression$1",symbols:["case_expression"]},{name:"asteriskless_andless_expression",symbols:["asteriskless_andless_expression$subexpression$1"],postprocess:yE},{name:"atomic_expression$subexpression$1",symbols:["array_subscript"]},{name:"atomic_expression$subexpression$1",symbols:["function_call"]},{name:"atomic_expression$subexpression$1",symbols:["property_access"]},{name:"atomic_expression$subexpression$1",symbols:["parenthesis"]},{name:"atomic_expression$subexpression$1",symbols:["curly_braces"]},{name:"atomic_expression$subexpression$1",symbols:["square_brackets"]},{name:"atomic_expression$subexpression$1",symbols:["operator"]},{name:"atomic_expression$subexpression$1",symbols:["identifier"]},{name:"atomic_expression$subexpression$1",symbols:["parameter"]},{name:"atomic_expression$subexpression$1",symbols:["literal"]},{name:"atomic_expression$subexpression$1",symbols:["data_type"]},{name:"atomic_expression$subexpression$1",symbols:["keyword"]},{name:"atomic_expression",symbols:["atomic_expression$subexpression$1"],postprocess:yE},{name:"array_subscript",symbols:[Re.has("ARRAY_IDENTIFIER")?{type:"ARRAY_IDENTIFIER"}:ARRAY_IDENTIFIER,"_","square_brackets"],postprocess:([E,e,T])=>({type:"array_subscript",array:Qe({type:"identifier",quoted:!1,text:E.text},{trailing:e}),parenthesis:T})},{name:"array_subscript",symbols:[Re.has("ARRAY_KEYWORD")?{type:"ARRAY_KEYWORD"}:ARRAY_KEYWORD,"_","square_brackets"],postprocess:([E,e,T])=>({type:"array_subscript",array:Qe(qe(E),{trailing:e}),parenthesis:T})},{name:"function_call",symbols:[Re.has("RESERVED_FUNCTION_NAME")?{type:"RESERVED_FUNCTION_NAME"}:RESERVED_FUNCTION_NAME,"_","parenthesis"],postprocess:([E,e,T])=>({type:"function_call",nameKw:Qe(qe(E),{trailing:e}),parenthesis:T})},{name:"parenthesis",symbols:[{literal:"("},"expressions_or_clauses",{literal:")"}],postprocess:([E,e,T])=>({type:"parenthesis",children:e,openParen:"(",closeParen:")"})},{name:"curly_braces$ebnf$1",symbols:[]},{name:"curly_braces$ebnf$1",symbols:["curly_braces$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"curly_braces",symbols:[{literal:"{"},"curly_braces$ebnf$1",{literal:"}"}],postprocess:([E,e,T])=>({type:"parenthesis",children:e,openParen:"{",closeParen:"}"})},{name:"square_brackets$ebnf$1",symbols:[]},{name:"square_brackets$ebnf$1",symbols:["square_brackets$ebnf$1","free_form_sql"],postprocess:E=>E[0].concat([E[1]])},{name:"square_brackets",symbols:[{literal:"["},"square_brackets$ebnf$1",{literal:"]"}],postprocess:([E,e,T])=>({type:"parenthesis",children:e,openParen:"[",closeParen:"]"})},{name:"property_access$subexpression$1",symbols:["identifier"]},{name:"property_access$subexpression$1",symbols:["array_subscript"]},{name:"property_access$subexpression$1",symbols:["all_columns_asterisk"]},{name:"property_access$subexpression$1",symbols:["parameter"]},{name:"property_access",symbols:["atomic_expression","_",Re.has("PROPERTY_ACCESS_OPERATOR")?{type:"PROPERTY_ACCESS_OPERATOR"}:PROPERTY_ACCESS_OPERATOR,"_","property_access$subexpression$1"],postprocess:([E,e,T,t,[r]])=>({type:"property_access",object:Qe(E,{trailing:e}),operator:T.text,property:Qe(r,{leading:t})})},{name:"between_predicate",symbols:[Re.has("BETWEEN")?{type:"BETWEEN"}:BETWEEN,"_","andless_expression_chain","_",Re.has("AND")?{type:"AND"}:AND,"_","andless_expression"],postprocess:([E,e,T,t,r,R,n])=>({type:"between_predicate",betweenKw:qe(E),expr1:kO(T,{leading:e,trailing:t}),andKw:qe(r),expr2:[Qe(n,{leading:R})]})},{name:"case_expression$ebnf$1",symbols:["expression_chain_"],postprocess:ST},{name:"case_expression$ebnf$1",symbols:[],postprocess:()=>null},{name:"case_expression$ebnf$2",symbols:[]},{name:"case_expression$ebnf$2",symbols:["case_expression$ebnf$2","case_clause"],postprocess:E=>E[0].concat([E[1]])},{name:"case_expression",symbols:[Re.has("CASE")?{type:"CASE"}:CASE,"_","case_expression$ebnf$1","case_expression$ebnf$2",Re.has("END")?{type:"END"}:END],postprocess:([E,e,T,t,r])=>({type:"case_expression",caseKw:Qe(qe(E),{trailing:e}),endKw:qe(r),expr:T||[],clauses:t})},{name:"case_clause",symbols:[Re.has("WHEN")?{type:"WHEN"}:WHEN,"_","expression_chain_",Re.has("THEN")?{type:"THEN"}:THEN,"_","expression_chain_"],postprocess:([E,e,T,t,r,R])=>({type:"case_when",whenKw:Qe(qe(E),{trailing:e}),thenKw:Qe(qe(t),{trailing:r}),condition:T,result:R})},{name:"case_clause",symbols:[Re.has("ELSE")?{type:"ELSE"}:ELSE,"_","expression_chain_"],postprocess:([E,e,T])=>({type:"case_else",elseKw:Qe(qe(E),{trailing:e}),result:T})},{name:"comma$subexpression$1",symbols:[Re.has("COMMA")?{type:"COMMA"}:COMMA]},{name:"comma",symbols:["comma$subexpression$1"],postprocess:([[E]])=>({type:"comma"})},{name:"asterisk$subexpression$1",symbols:[Re.has("ASTERISK")?{type:"ASTERISK"}:ASTERISK]},{name:"asterisk",symbols:["asterisk$subexpression$1"],postprocess:([[E]])=>({type:"operator",text:E.text})},{name:"operator$subexpression$1",symbols:[Re.has("OPERATOR")?{type:"OPERATOR"}:OPERATOR]},{name:"operator",symbols:["operator$subexpression$1"],postprocess:([[E]])=>({type:"operator",text:E.text})},{name:"identifier$subexpression$1",symbols:[Re.has("IDENTIFIER")?{type:"IDENTIFIER"}:IDENTIFIER]},{name:"identifier$subexpression$1",symbols:[Re.has("QUOTED_IDENTIFIER")?{type:"QUOTED_IDENTIFIER"}:QUOTED_IDENTIFIER]},{name:"identifier$subexpression$1",symbols:[Re.has("VARIABLE")?{type:"VARIABLE"}:VARIABLE]},{name:"identifier",symbols:["identifier$subexpression$1"],postprocess:([[E]])=>({type:"identifier",quoted:E.type!=="IDENTIFIER",text:E.text})},{name:"parameter$subexpression$1",symbols:[Re.has("NAMED_PARAMETER")?{type:"NAMED_PARAMETER"}:NAMED_PARAMETER]},{name:"parameter$subexpression$1",symbols:[Re.has("QUOTED_PARAMETER")?{type:"QUOTED_PARAMETER"}:QUOTED_PARAMETER]},{name:"parameter$subexpression$1",symbols:[Re.has("NUMBERED_PARAMETER")?{type:"NUMBERED_PARAMETER"}:NUMBERED_PARAMETER]},{name:"parameter$subexpression$1",symbols:[Re.has("POSITIONAL_PARAMETER")?{type:"POSITIONAL_PARAMETER"}:POSITIONAL_PARAMETER]},{name:"parameter$subexpression$1",symbols:[Re.has("CUSTOM_PARAMETER")?{type:"CUSTOM_PARAMETER"}:CUSTOM_PARAMETER]},{name:"parameter",symbols:["parameter$subexpression$1"],postprocess:([[E]])=>({type:"parameter",key:E.key,text:E.text})},{name:"literal$subexpression$1",symbols:[Re.has("NUMBER")?{type:"NUMBER"}:NUMBER]},{name:"literal$subexpression$1",symbols:[Re.has("STRING")?{type:"STRING"}:STRING]},{name:"literal",symbols:["literal$subexpression$1"],postprocess:([[E]])=>({type:"literal",text:E.text})},{name:"keyword$subexpression$1",symbols:[Re.has("RESERVED_KEYWORD")?{type:"RESERVED_KEYWORD"}:RESERVED_KEYWORD]},{name:"keyword$subexpression$1",symbols:[Re.has("RESERVED_PHRASE")?{type:"RESERVED_PHRASE"}:RESERVED_PHRASE]},{name:"keyword$subexpression$1",symbols:[Re.has("RESERVED_JOIN")?{type:"RESERVED_JOIN"}:RESERVED_JOIN]},{name:"keyword",symbols:["keyword$subexpression$1"],postprocess:([[E]])=>qe(E)},{name:"data_type$subexpression$1",symbols:[Re.has("RESERVED_DATA_TYPE")?{type:"RESERVED_DATA_TYPE"}:RESERVED_DATA_TYPE]},{name:"data_type",symbols:["data_type$subexpression$1"],postprocess:([[E]])=>Ur(E)},{name:"data_type",symbols:[Re.has("RESERVED_PARAMETERIZED_DATA_TYPE")?{type:"RESERVED_PARAMETERIZED_DATA_TYPE"}:RESERVED_PARAMETERIZED_DATA_TYPE,"_","parenthesis"],postprocess:([E,e,T])=>({type:"parameterized_data_type",dataType:Qe(Ur(E),{trailing:e}),parenthesis:T})},{name:"logic_operator$subexpression$1",symbols:[Re.has("AND")?{type:"AND"}:AND]},{name:"logic_operator$subexpression$1",symbols:[Re.has("OR")?{type:"OR"}:OR]},{name:"logic_operator$subexpression$1",symbols:[Re.has("XOR")?{type:"XOR"}:XOR]},{name:"logic_operator",symbols:["logic_operator$subexpression$1"],postprocess:([[E]])=>qe(E)},{name:"other_keyword$subexpression$1",symbols:[Re.has("WHEN")?{type:"WHEN"}:WHEN]},{name:"other_keyword$subexpression$1",symbols:[Re.has("THEN")?{type:"THEN"}:THEN]},{name:"other_keyword$subexpression$1",symbols:[Re.has("ELSE")?{type:"ELSE"}:ELSE]},{name:"other_keyword$subexpression$1",symbols:[Re.has("END")?{type:"END"}:END]},{name:"other_keyword",symbols:["other_keyword$subexpression$1"],postprocess:([[E]])=>qe(E)},{name:"_$ebnf$1",symbols:[]},{name:"_$ebnf$1",symbols:["_$ebnf$1","comment"],postprocess:E=>E[0].concat([E[1]])},{name:"_",symbols:["_$ebnf$1"],postprocess:([E])=>E},{name:"comment",symbols:[Re.has("LINE_COMMENT")?{type:"LINE_COMMENT"}:LINE_COMMENT],postprocess:([E])=>({type:"line_comment",text:E.text,precedingWhitespace:E.precedingWhitespace})},{name:"comment",symbols:[Re.has("BLOCK_COMMENT")?{type:"BLOCK_COMMENT"}:BLOCK_COMMENT],postprocess:([E])=>({type:"block_comment",text:E.text,precedingWhitespace:E.precedingWhitespace})},{name:"comment",symbols:[Re.has("DISABLE_COMMENT")?{type:"DISABLE_COMMENT"}:DISABLE_COMMENT],postprocess:([E])=>({type:"disable_comment",text:E.text,precedingWhitespace:E.precedingWhitespace})}],ParserStart:"main"},JO=KO,{Parser:qO,Grammar:QO}=zA;function ZO(E){let e={};const T=new QR(r=>[...FO(E.tokenize(r,e)),WR(r.length)]),t=new qO(QO.fromCompiled(JO),{lexer:T});return{parse:(r,R)=>{e=R;const{results:n}=t.feed(r);if(n.length===1)return n[0];throw n.length===0?new Error("Parse error: Invalid SQL"):new Error(`Parse error: Ambiguous grammar
${JSON.stringify(n,void 0,2)}`)}}}var ZR=class{constructor(E){this.indentation=E,this.items=[]}add(...E){for(const e of E)switch(e){case 0:this.items.push(0);break;case 1:this.trimHorizontalWhitespace();break;case 2:this.trimWhitespace();break;case 3:this.trimHorizontalWhitespace(),this.addNewline(3);break;case 4:this.trimHorizontalWhitespace(),this.addNewline(4);break;case 5:this.addIndentation();break;case 6:this.items.push(6);break;default:this.items.push(e)}}trimHorizontalWhitespace(){for(;jO(ot(this.items));)this.items.pop()}trimWhitespace(){for(;zO(ot(this.items));)this.items.pop()}addNewline(E){if(this.items.length>0)switch(ot(this.items)){case 3:this.items.pop(),this.items.push(E);break;case 4:break;default:this.items.push(E);break}}addIndentation(){for(let E=0;E<this.indentation.getLevel();E++)this.items.push(6)}toString(){return this.items.map(E=>this.itemToString(E)).join("")}getLayoutItems(){return this.items}itemToString(E){switch(E){case 0:return" ";case 3:case 4:return`
`;case 6:return this.indentation.getSingleIndent();default:return E}}},jO=E=>E===0||E===6,zO=E=>E===0||E===6||E===3;function mr(E,e){if(e==="standard")return E;let T=[];return E.length>=10&&E.includes(" ")&&([E,...T]=E.split(" ")),e==="tabularLeft"?E=E.padEnd(9," "):E=E.padStart(9," "),E+["",...T].join(" ")}function hr(E){return is(E)||E==="RESERVED_CLAUSE"||E==="RESERVED_SELECT"||E==="RESERVED_SET_OPERATION"||E==="RESERVED_JOIN"||E==="LIMIT"}var oT="top-level",ei="block-level",jR=class{constructor(E){this.indent=E,this.indentTypes=[]}getSingleIndent(){return this.indent}getLevel(){return this.indentTypes.length}increaseTopLevel(){this.indentTypes.push(oT)}increaseBlockLevel(){this.indentTypes.push(ei)}decreaseTopLevel(){this.indentTypes.length>0&&ot(this.indentTypes)===oT&&this.indentTypes.pop()}decreaseBlockLevel(){for(;this.indentTypes.length>0&&this.indentTypes.pop()===oT;);}},Ei=class extends ZR{constructor(E){super(new jR("")),this.expressionWidth=E,this.length=0,this.trailingSpace=!1}add(...E){if(E.forEach(e=>this.addToLength(e)),this.length>this.expressionWidth)throw new NT;super.add(...E)}addToLength(E){if(typeof E=="string")this.length+=E.length,this.trailingSpace=!1;else{if(E===4||E===3)throw new NT;E===5||E===6||E===0?this.trailingSpace||(this.length++,this.trailingSpace=!0):(E===2||E===1)&&this.trailingSpace&&(this.trailingSpace=!1,this.length--)}}},NT=class extends Error{},ti=class lT{constructor({cfg:e,dialectCfg:T,params:t,layout:r,inline:R=!1}){this.inline=!1,this.nodes=[],this.index=-1,this.cfg=e,this.dialectCfg=T,this.inline=R,this.params=t,this.layout=r}format(e){for(this.nodes=e,this.index=0;this.index<this.nodes.length;this.index++)this.formatNode(this.nodes[this.index]);return this.layout}formatNode(e){this.formatComments(e.leadingComments),this.formatNodeWithoutComments(e),this.formatComments(e.trailingComments)}formatNodeWithoutComments(e){switch(e.type){case"function_call":return this.formatFunctionCall(e);case"parameterized_data_type":return this.formatParameterizedDataType(e);case"array_subscript":return this.formatArraySubscript(e);case"property_access":return this.formatPropertyAccess(e);case"parenthesis":return this.formatParenthesis(e);case"between_predicate":return this.formatBetweenPredicate(e);case"case_expression":return this.formatCaseExpression(e);case"case_when":return this.formatCaseWhen(e);case"case_else":return this.formatCaseElse(e);case"clause":return this.formatClause(e);case"set_operation":return this.formatSetOperation(e);case"limit_clause":return this.formatLimitClause(e);case"all_columns_asterisk":return this.formatAllColumnsAsterisk(e);case"literal":return this.formatLiteral(e);case"identifier":return this.formatIdentifier(e);case"parameter":return this.formatParameter(e);case"operator":return this.formatOperator(e);case"comma":return this.formatComma(e);case"line_comment":return this.formatLineComment(e);case"block_comment":return this.formatBlockComment(e);case"disable_comment":return this.formatBlockComment(e);case"data_type":return this.formatDataType(e);case"keyword":return this.formatKeywordNode(e)}}formatFunctionCall(e){this.withComments(e.nameKw,()=>{this.layout.add(this.showFunctionKw(e.nameKw))}),this.formatNode(e.parenthesis)}formatParameterizedDataType(e){this.withComments(e.dataType,()=>{this.layout.add(this.showDataType(e.dataType))}),this.formatNode(e.parenthesis)}formatArraySubscript(e){let T;switch(e.array.type){case"data_type":T=this.showDataType(e.array);break;case"keyword":T=this.showKw(e.array);break;default:T=this.showIdentifier(e.array);break}this.withComments(e.array,()=>{this.layout.add(T)}),this.formatNode(e.parenthesis)}formatPropertyAccess(e){this.formatNode(e.object),this.layout.add(1,e.operator),this.formatNode(e.property)}formatParenthesis(e){const T=this.formatInlineExpression(e.children);T?(this.layout.add(e.openParen),this.layout.add(...T.getLayoutItems()),this.layout.add(1,e.closeParen,0)):(this.layout.add(e.openParen,3),ZE(this.cfg)?(this.layout.add(5),this.layout=this.formatSubExpression(e.children)):(this.layout.indentation.increaseBlockLevel(),this.layout.add(5),this.layout=this.formatSubExpression(e.children),this.layout.indentation.decreaseBlockLevel()),this.layout.add(3,5,e.closeParen,0))}formatBetweenPredicate(e){this.layout.add(this.showKw(e.betweenKw),0),this.layout=this.formatSubExpression(e.expr1),this.layout.add(1,0,this.showNonTabularKw(e.andKw),0),this.layout=this.formatSubExpression(e.expr2),this.layout.add(0)}formatCaseExpression(e){this.formatNode(e.caseKw),this.layout.indentation.increaseBlockLevel(),this.layout=this.formatSubExpression(e.expr),this.layout=this.formatSubExpression(e.clauses),this.layout.indentation.decreaseBlockLevel(),this.layout.add(3,5),this.formatNode(e.endKw)}formatCaseWhen(e){this.layout.add(3,5),this.formatNode(e.whenKw),this.layout=this.formatSubExpression(e.condition),this.formatNode(e.thenKw),this.layout=this.formatSubExpression(e.result)}formatCaseElse(e){this.layout.add(3,5),this.formatNode(e.elseKw),this.layout=this.formatSubExpression(e.result)}formatClause(e){this.isOnelineClause(e)?this.formatClauseInOnelineStyle(e):ZE(this.cfg)?this.formatClauseInTabularStyle(e):this.formatClauseInIndentedStyle(e)}isOnelineClause(e){return ZE(this.cfg)?this.dialectCfg.tabularOnelineClauses[e.nameKw.text]:this.dialectCfg.onelineClauses[e.nameKw.text]}formatClauseInIndentedStyle(e){this.layout.add(3,5,this.showKw(e.nameKw),3),this.layout.indentation.increaseTopLevel(),this.layout.add(5),this.layout=this.formatSubExpression(e.children),this.layout.indentation.decreaseTopLevel()}formatClauseInOnelineStyle(e){this.layout.add(3,5,this.showKw(e.nameKw),0),this.layout=this.formatSubExpression(e.children)}formatClauseInTabularStyle(e){this.layout.add(3,5,this.showKw(e.nameKw),0),this.layout.indentation.increaseTopLevel(),this.layout=this.formatSubExpression(e.children),this.layout.indentation.decreaseTopLevel()}formatSetOperation(e){this.layout.add(3,5,this.showKw(e.nameKw),3),this.layout.add(5),this.layout=this.formatSubExpression(e.children)}formatLimitClause(e){this.withComments(e.limitKw,()=>{this.layout.add(3,5,this.showKw(e.limitKw))}),this.layout.indentation.increaseTopLevel(),ZE(this.cfg)?this.layout.add(0):this.layout.add(3,5),e.offset?(this.layout=this.formatSubExpression(e.offset),this.layout.add(1,",",0),this.layout=this.formatSubExpression(e.count)):this.layout=this.formatSubExpression(e.count),this.layout.indentation.decreaseTopLevel()}formatAllColumnsAsterisk(e){this.layout.add("*",0)}formatLiteral(e){this.layout.add(e.text,0)}formatIdentifier(e){this.layout.add(this.showIdentifier(e),0)}formatParameter(e){this.layout.add(this.params.get(e),0)}formatOperator({text:e}){this.cfg.denseOperators||this.dialectCfg.alwaysDenseOperators.includes(e)?this.layout.add(1,e):e===":"?this.layout.add(1,e,0):this.layout.add(e,0)}formatComma(e){this.inline?this.layout.add(1,",",0):this.layout.add(1,",",3,5)}withComments(e,T){this.formatComments(e.leadingComments),T(),this.formatComments(e.trailingComments)}formatComments(e){e&&e.forEach(T=>{T.type==="line_comment"?this.formatLineComment(T):this.formatBlockComment(T)})}formatLineComment(e){AT(e.precedingWhitespace||"")?this.layout.add(3,5,e.text,4,5):this.layout.getLayoutItems().length>0?this.layout.add(2,0,e.text,4,5):this.layout.add(e.text,4,5)}formatBlockComment(e){e.type==="block_comment"&&this.isMultilineBlockComment(e)?(this.splitBlockComment(e.text).forEach(T=>{this.layout.add(3,5,T)}),this.layout.add(3,5)):this.layout.add(e.text,0)}isMultilineBlockComment(e){return AT(e.text)||AT(e.precedingWhitespace||"")}isDocComment(e){const T=e.split(/\\n/);return/^\\/\\*\\*?$/.test(T[0])&&T.slice(1,T.length-1).every(t=>/^\\s*\\*/.test(t))&&/^\\s*\\*\\/$/.test(ot(T))}splitBlockComment(e){return this.isDocComment(e)?e.split(/\\n/).map(T=>/^\\s*\\*/.test(T)?" "+T.replace(/^\\s*/,""):T):e.split(/\\n/).map(T=>T.replace(/^\\s*/,""))}formatSubExpression(e){return new lT({cfg:this.cfg,dialectCfg:this.dialectCfg,params:this.params,layout:this.layout,inline:this.inline}).format(e)}formatInlineExpression(e){const T=this.params.getPositionalParameterIndex();try{return new lT({cfg:this.cfg,dialectCfg:this.dialectCfg,params:this.params,layout:new Ei(this.cfg.expressionWidth),inline:!0}).format(e)}catch(t){if(t instanceof NT){this.params.setPositionalParameterIndex(T);return}else throw t}}formatKeywordNode(e){switch(e.tokenType){case"RESERVED_JOIN":return this.formatJoin(e);case"AND":case"OR":case"XOR":return this.formatLogicalOperator(e);default:return this.formatKeyword(e)}}formatJoin(e){ZE(this.cfg)?(this.layout.indentation.decreaseTopLevel(),this.layout.add(3,5,this.showKw(e),0),this.layout.indentation.increaseTopLevel()):this.layout.add(3,5,this.showKw(e),0)}formatKeyword(e){this.layout.add(this.showKw(e),0)}formatLogicalOperator(e){this.cfg.logicalOperatorNewline==="before"?ZE(this.cfg)?(this.layout.indentation.decreaseTopLevel(),this.layout.add(3,5,this.showKw(e),0),this.layout.indentation.increaseTopLevel()):this.layout.add(3,5,this.showKw(e),0):this.layout.add(this.showKw(e),3,5)}formatDataType(e){this.layout.add(this.showDataType(e),0)}showKw(e){return hr(e.tokenType)?mr(this.showNonTabularKw(e),this.cfg.indentStyle):this.showNonTabularKw(e)}showNonTabularKw(e){switch(this.cfg.keywordCase){case"preserve":return Dt(e.raw);case"upper":return e.text;case"lower":return e.text.toLowerCase()}}showFunctionKw(e){return hr(e.tokenType)?mr(this.showNonTabularFunctionKw(e),this.cfg.indentStyle):this.showNonTabularFunctionKw(e)}showNonTabularFunctionKw(e){switch(this.cfg.functionCase){case"preserve":return Dt(e.raw);case"upper":return e.text;case"lower":return e.text.toLowerCase()}}showIdentifier(e){if(e.quoted)return e.text;switch(this.cfg.identifierCase){case"preserve":return e.text;case"upper":return e.text.toUpperCase();case"lower":return e.text.toLowerCase()}}showDataType(e){switch(this.cfg.dataTypeCase){case"preserve":return Dt(e.raw);case"upper":return e.text;case"lower":return e.text.toLowerCase()}}},Ti=class{constructor(E,e){this.dialect=E,this.cfg=e,this.params=new vO(this.cfg.params)}format(E){const e=this.parse(E);return this.formatAst(e).trimEnd()}parse(E){return ZO(this.dialect.tokenizer).parse(E,this.cfg.paramTypes||{})}formatAst(E){return E.map(e=>this.formatStatement(e)).join(`
`.repeat(this.cfg.linesBetweenQueries+1))}formatStatement(E){const e=new ti({cfg:this.cfg,dialectCfg:this.dialect.formatOptions,params:this.params,layout:new ZR(new jR(BO(this.cfg)))}).format(E.children);return E.hasSemicolon&&(this.cfg.newlineBeforeSemicolon?e.add(3,";"):e.add(2,";")),e.toString()}},_T=class extends Error{};function ri(E){const e=["multilineLists","newlineBeforeOpenParen","newlineBeforeCloseParen","aliasAs","commaPosition","tabulateAlias"];for(const T of e)if(T in E)throw new _T(`${T} config is no more supported.`);if(E.expressionWidth<=0)throw new _T(`expressionWidth config must be positive number. Received ${E.expressionWidth} instead.`);return E.params&&!Ri(E.params)&&console.warn('WARNING: All "params" option values should be strings.'),E}function Ri(E){return(E instanceof Array?E:Object.values(E)).every(T=>typeof T=="string")}var zR={bigquery:"bigquery",db2:"db2",db2i:"db2i",hive:"hive",mariadb:"mariadb",mysql:"mysql",n1ql:"n1ql",plsql:"plsql",postgresql:"postgresql",redshift:"redshift",spark:"spark",sqlite:"sqlite",sql:"sql",tidb:"tidb",trino:"trino",transactsql:"transactsql",tsql:"transactsql",singlestoredb:"singlestoredb",snowflake:"snowflake"},ni=Object.keys(zR),Ai={tabWidth:2,useTabs:!1,keywordCase:"preserve",identifierCase:"preserve",dataTypeCase:"preserve",functionCase:"preserve",indentStyle:"standard",logicalOperatorNewline:"before",expressionWidth:50,linesBetweenQueries:1,denseOperators:!1,newlineBeforeSemicolon:!1},Gr=(E,e={})=>{if(typeof e.language=="string"&&!ni.includes(e.language))throw new _T(`Unsupported SQL dialect: ${e.language}`);const T=zR[e.language||"sql"];return si(E,AE(EE({},e),{dialect:YR[T]}))},si=(E,e)=>{var T=e,{dialect:t}=T,r=ts(T,["dialect"]);if(typeof E!="string")throw new Error("Invalid query argument. Expected string, instead got "+typeof E);const R=ri(EE(EE({},Ai),r));return new Ti(HO(t),R).format(E)};function gr(E,e,T){const t=E.slice();return t[1]=e[T],t}function Hr(E,e){let T,t=e[1].name+"",r,R,n,s=e[1].instantiated_value+"",S,A;return{key:E,first:null,c(){T=f("b"),r=te(t),R=te(":"),n=$(),S=te(s),A=f("br"),this.first=T},m(o,i){V(o,T,i),l(T,r),l(T,R),V(o,n,i),V(o,S,i),V(o,A,i)},p(o,i){e=o,i&1&&t!==(t=e[1].name+"")&&Le(r,t),i&1&&s!==(s=e[1].instantiated_value+"")&&Le(S,s)},d(o){o&&(Y(T),Y(n),Y(S),Y(A))}}}function Si(E){let e,T,t,r,R,n,s,S=E[0].function_name+"",A,o,i,_=[],c=new Map,P,p,C,L,I,u,H,b=E[0].description+"",M,O,N,D,B=Gr(E[0].instantiated_sql)+"",h,G,F,W,x=E[0].sql_template+"",J,oe=De(E[0].arguments);const z=Oe=>Oe[1].name;for(let Oe=0;Oe<oe.length;Oe+=1){let Ue=gr(E,oe,Oe),ye=z(Ue);c.set(ye,_[Oe]=Hr(ye,Ue))}return{c(){e=f("div"),T=f("div"),t=f("div"),t.innerHTML='<span class="inline-flex justify-center items-center size-8 rounded-full border-4 border-teal-100 bg-teal-200 text-teal-800 dark:border-teal-900 dark:bg-teal-800 dark:text-teal-400"><svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M12 22c5.523 0 10-4.477 10-10S17.523 2 12 2 2 6.477 2 12s4.477 10 10 10z"></path><path d="m9 12 2 2 4-4"></path></svg></span>',r=$(),R=f("div"),n=f("h3"),s=te("Function: "),A=te(S),o=$(),i=f("p");for(let Oe=0;Oe<_.length;Oe+=1)_[Oe].c();P=$(),p=f("div"),C=f("div"),C.innerHTML='<nav class="flex space-x-2" aria-label="Tabs"><button type="button" class="hs-tab-active:font-semibold hs-tab-active:border-blue-600 hs-tab-active:text-blue-600 py-4 px-1 inline-flex items-center gap-x-2 border-b-2 border-transparent text-sm whitespace-nowrap text-gray-500 hover:text-blue-600 disabled:opacity-50 disabled:pointer-events-none dark:text-neutral-400 dark:hover:text-blue-500 active" id="basic-tabs-item-1" data-hs-tab="#basic-tabs-1" aria-controls="basic-tabs-1" role="tab">Description</button> <button type="button" class="hs-tab-active:font-semibold hs-tab-active:border-blue-600 hs-tab-active:text-blue-600 py-4 px-1 inline-flex items-center gap-x-2 border-b-2 border-transparent text-sm whitespace-nowrap text-gray-500 hover:text-blue-600 disabled:opacity-50 disabled:pointer-events-none dark:text-neutral-400 dark:hover:text-blue-500" id="basic-tabs-item-2" data-hs-tab="#basic-tabs-2" aria-controls="basic-tabs-2" role="tab">SQL</button> <button type="button" class="hs-tab-active:font-semibold hs-tab-active:border-blue-600 hs-tab-active:text-blue-600 py-4 px-1 inline-flex items-center gap-x-2 border-b-2 border-transparent text-sm whitespace-nowrap text-gray-500 hover:text-blue-600 disabled:opacity-50 disabled:pointer-events-none dark:text-neutral-400 dark:hover:text-blue-500" id="basic-tabs-item-3" data-hs-tab="#basic-tabs-3" aria-controls="basic-tabs-3" role="tab">Template</button></nav>',L=$(),I=f("div"),u=f("div"),H=f("p"),M=te(b),O=$(),N=f("div"),D=f("p"),h=te(B),G=$(),F=f("div"),W=f("p"),J=te(x),a(t,"class","flex-shrink-0"),a(n,"class","text-gray-800 font-semibold dark:text-white"),a(i,"class","text-sm text-gray-700 dark:text-neutral-400"),a(R,"class","ms-3"),a(T,"class","flex"),a(C,"class","border-b border-gray-200 px-4 dark:border-neutral-700"),a(H,"class","text-gray-500 dark:text-neutral-400"),a(u,"id","basic-tabs-1"),a(u,"role","tabpanel"),a(u,"aria-labelledby","basic-tabs-item-1"),a(D,"class","text-gray-500 dark:text-neutral-400 font-mono text-teal-800"),a(N,"id","basic-tabs-2"),a(N,"class","hidden"),a(N,"role","tabpanel"),a(N,"aria-labelledby","basic-tabs-item-2"),a(W,"class","text-gray-500 dark:text-neutral-400 font-mono"),a(F,"id","basic-tabs-3"),a(F,"class","hidden"),a(F,"role","tabpanel"),a(F,"aria-labelledby","basic-tabs-item-3"),a(I,"class","px-4"),a(p,"class","w-full bg-white rounded-lg shadow-md dark:bg-neutral-800"),a(e,"class","bg-teal-50 border-t-2 border-teal-500 rounded-lg p-4 dark:bg-teal-800/30"),a(e,"role","alert")},m(Oe,Ue){V(Oe,e,Ue),l(e,T),l(T,t),l(T,r),l(T,R),l(R,n),l(n,s),l(n,A),l(R,o),l(R,i);for(let ye=0;ye<_.length;ye+=1)_[ye]&&_[ye].m(i,null);l(e,P),l(e,p),l(p,C),l(p,L),l(p,I),l(I,u),l(u,H),l(H,M),l(I,O),l(I,N),l(N,D),l(D,h),l(I,G),l(I,F),l(F,W),l(W,J)},p(Oe,[Ue]){Ue&1&&S!==(S=Oe[0].function_name+"")&&Le(A,S),Ue&1&&(oe=De(Oe[0].arguments),_=un(_,Ue,z,1,Oe,oe,c,i,Cn,Hr,null,gr)),Ue&1&&b!==(b=Oe[0].description+"")&&Le(M,b),Ue&1&&B!==(B=Gr(Oe[0].instantiated_sql)+"")&&Le(h,B),Ue&1&&x!==(x=Oe[0].sql_template+"")&&Le(J,x)},i:j,o:j,d(Oe){Oe&&Y(e);for(let Ue=0;Ue<_.length;Ue+=1)_[Ue].d()}}}function oi(E,e,T){let{functionData:t}=e;return E.$$set=r=>{"functionData"in r&&T(0,t=r.functionData)},[t]}class Oi extends ue{constructor(e){super(),Ce(this,e,oi,Si,_e,{functionData:0})}}function br(E,e,T){const t=E.slice();return t[12]=e[T],t}function yr(E,e,T){const t=E.slice();return t[12]=e[T],t[13]=e,t[14]=T,t}function ii(E){let e,T,t,r=E[0].function_name+"",R,n,s,S=E[0].description+"",A,o,i,_,c,P,p,C,L,I=De(E[0].arguments),u=[];for(let H=0;H<I.length;H+=1)u[H]=Br(br(E,I,H));return{c(){e=f("div"),T=f("div"),t=f("h3"),R=te(r),n=$(),s=f("p"),A=te(S),o=$(),i=f("ul");for(let H=0;H<u.length;H+=1)u[H].c();_=$(),c=f("button"),c.textContent="Delete",P=$(),p=f("button"),p.textContent="Edit",a(t,"class","text-lg font-bold text-gray-800 dark:text-white"),a(s,"class","mt-1 text-xs font-medium uppercase text-gray-500 dark:text-neutral-500"),a(i,"class","mt-1 marker:text-blue-600 list-disc ps-5 space-y-2 text-sm text-gray-600 dark:text-neutral-400"),a(T,"class","flex-grow"),a(c,"class","mt-3 inline-flex items-center gap-x-1 text-sm font-semibold rounded-lg border border-red-600 text-red-600 hover:border-red-500 hover:text-red-500 disabled:opacity-50 disabled:pointer-events-none dark:border-red-500 dark:text-red-500 dark:hover:text-red-400 dark:hover:border-red-400 py-1 px-4 max-w-fit"),a(p,"class","mt-3 inline-flex items-center gap-x-1 text-sm font-semibold rounded-lg border border-blue-600 text-blue-600 hover:border-blue-500 hover:text-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:border-blue-500 dark:text-blue-500 dark:hover:text-blue-400 dark:hover:border-blue-400 py-1 px-4 max-w-fit"),a(e,"class","p-4 flex flex-col bg-white border border-t-4 border-t-blue-600 shadow-sm rounded-xl dark:bg-neutral-900 dark:border-neutral-7000 dark:border-t-blue-500 dark:shadow-neutral-700/70")},m(H,b){V(H,e,b),l(e,T),l(T,t),l(t,R),l(T,n),l(T,s),l(s,A),l(T,o),l(T,i);for(let M=0;M<u.length;M+=1)u[M]&&u[M].m(i,null);l(e,_),l(e,c),l(e,P),l(e,p),C||(L=[Ne(c,"click",E[10]),Ne(p,"click",E[11])],C=!0)},p(H,b){if(b&1&&r!==(r=H[0].function_name+"")&&Le(R,r),b&1&&S!==(S=H[0].description+"")&&Le(A,S),b&1){I=De(H[0].arguments);let M;for(M=0;M<I.length;M+=1){const O=br(H,I,M);u[M]?u[M].p(O,b):(u[M]=Br(O),u[M].c(),u[M].m(i,null))}for(;M<u.length;M+=1)u[M].d(1);u.length=I.length}},d(H){H&&Y(e),nE(u,H),C=!1,NE(L)}}}function ai(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p,C,L,I,u,H,b,M,O,N,D,B=De(E[1].arguments),h=[];for(let G=0;G<B.length;G+=1)h[G]=vr(yr(E,B,G));return{c(){e=f("div"),T=f("form"),t=f("div"),r=f("div"),R=f("span"),R.textContent="Function Name",n=$(),s=f("input"),S=$(),A=f("div"),o=f("span"),o.textContent="Function Description",i=$(),_=f("textarea"),c=$(),P=f("div"),p=f("span"),p.textContent="SQL Template",C=$(),L=f("textarea"),I=$();for(let G=0;G<h.length;G+=1)h[G].c();u=$(),H=f("div"),b=f("button"),b.textContent="Cancel",M=$(),O=f("button"),O.textContent="Save",a(R,"class","inline-block text-sm font-bold text-gray-800 mt-2.5 dark:text-neutral-200"),a(s,"type","text"),a(s,"class","py-2 px-3 pe-11 block w-full border-gray-200 shadow-sm rounded-lg text-sm focus:border-blue-500 focus:ring-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:bg-neutral-900 dark:border-neutral-700 dark:text-neutral-400 dark:placeholder-neutral-500 dark:focus:ring-neutral-600"),a(s,"placeholder","Enter function name"),a(r,"class","space-y-2"),a(t,"class","space-y-4 sm:space-y-6"),a(o,"class","inline-block text-sm font-bold text-gray-800 mt-2.5 dark:text-neutral-200"),a(_,"class","py-2 px-3 pe-11 block w-full border-gray-200 shadow-sm rounded-lg text-sm focus:border-blue-500 focus:ring-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:bg-neutral-900 dark:border-neutral-700 dark:text-neutral-400 dark:placeholder-neutral-500 dark:focus:ring-neutral-600"),a(_,"placeholder","Enter function description"),a(A,"class","space-y-2"),a(p,"class","inline-block text-sm font-bold text-gray-800 mt-2.5 dark:text-neutral-200"),a(L,"class","py-2 px-3 pe-11 block w-full border-gray-200 shadow-sm rounded-lg text-sm focus:border-blue-500 focus:ring-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:bg-neutral-900 dark:border-neutral-700 dark:text-neutral-400 dark:placeholder-neutral-500 dark:focus:ring-neutral-600 font-mono"),a(L,"placeholder","Enter function description"),a(P,"class","space-y-2"),a(b,"class","py-2 px-3 inline-flex items-center gap-x-2 text-sm font-semibold rounded-lg border border-red-500 text-red-500 hover:border-red-400 hover:text-red-400 disabled:opacity-50 disabled:pointer-events-none"),a(O,"class","py-2 px-3 inline-flex items-center gap-x-2 text-sm font-semibold rounded-lg border border-transparent bg-blue-600 text-white hover:bg-blue-700 disabled:opacity-50 disabled:pointer-events-none"),a(H,"class","mt-3 flex gap-x-2"),a(e,"class","p-4 flex flex-col bg-white border border-t-4 border-t-red-600 shadow-sm rounded-xl dark:bg-neutral-900 dark:border-neutral-7000 dark:border-t-red-500 dark:shadow-neutral-700/70")},m(G,F){V(G,e,F),l(e,T),l(T,t),l(t,r),l(r,R),l(r,n),l(r,s),Ye(s,E[1].function_name),l(T,S),l(T,A),l(A,o),l(A,i),l(A,_),Ye(_,E[1].description),l(T,c),l(T,P),l(P,p),l(P,C),l(P,L),Ye(L,E[1].sql_template),l(T,I);for(let W=0;W<h.length;W+=1)h[W]&&h[W].m(T,null);l(e,u),l(e,H),l(H,b),l(H,M),l(H,O),N||(D=[Ne(s,"input",E[2]),Ne(_,"input",E[3]),Ne(L,"input",E[4]),Ne(b,"click",E[8]),Ne(O,"click",E[9])],N=!0)},p(G,F){if(F&2&&s.value!==G[1].function_name&&Ye(s,G[1].function_name),F&2&&Ye(_,G[1].description),F&2&&Ye(L,G[1].sql_template),F&2){B=De(G[1].arguments);let W;for(W=0;W<B.length;W+=1){const x=yr(G,B,W);h[W]?h[W].p(x,F):(h[W]=vr(x),h[W].c(),h[W].m(T,null))}for(;W<h.length;W+=1)h[W].d(1);h.length=B.length}},d(G){G&&Y(e),nE(h,G),N=!1,NE(D)}}}function Br(E){let e,T,t=E[12].name+"",r,R,n,s=E[12].description+"",S;return{c(){e=f("li"),T=f("b"),r=te(t),R=te(":"),n=$(),S=te(s)},m(A,o){V(A,e,o),l(e,T),l(T,r),l(T,R),l(e,n),l(e,S)},p(A,o){o&1&&t!==(t=A[12].name+"")&&Le(r,t),o&1&&s!==(s=A[12].description+"")&&Le(S,s)},d(A){A&&Y(e)}}}function vr(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p,C,L,I,u,H,b;function M(){E[5].call(R,E[13],E[14])}function O(){E[6].call(o,E[13],E[14])}function N(){E[7].call(p,E[13],E[14])}return{c(){e=f("div"),T=f("div"),t=f("span"),t.textContent="Argument Name",r=$(),R=f("input"),n=$(),s=f("div"),S=f("span"),S.textContent="Argument Description",A=$(),o=f("input"),i=$(),_=f("div"),c=f("span"),c.textContent="Argument Type",P=$(),p=f("select"),C=f("option"),C.textContent="String",L=f("option"),L.textContent="Numeric",I=f("option"),I.textContent="Boolean",u=$(),a(t,"class","inline-block text-sm font-bold text-gray-800 mt-2.5 dark:text-neutral-200"),a(R,"type","text"),a(R,"class","py-2 px-3 pe-11 block w-full border-gray-200 shadow-sm rounded-lg text-sm focus:border-blue-500 focus:ring-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:bg-neutral-900 dark:border-neutral-700 dark:text-neutral-400 dark:placeholder-neutral-500 dark:focus:ring-neutral-600"),a(R,"placeholder","Enter argument name"),a(T,"class","space-y-2"),a(S,"class","inline-block text-sm font-bold text-gray-800 mt-2.5 dark:text-neutral-200"),a(o,"type","text"),a(o,"class","py-2 px-3 pe-11 block w-full border-gray-200 shadow-sm rounded-lg text-sm focus:border-blue-500 focus:ring-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:bg-neutral-900 dark:border-neutral-700 dark:text-neutral-400 dark:placeholder-neutral-500 dark:focus:ring-neutral-600"),a(o,"placeholder","Enter argument description"),a(s,"class","space-y-2"),a(c,"class","inline-block text-sm font-bold text-gray-800 mt-2.5 dark:text-neutral-200"),C.__value="STRING",Ye(C,C.__value),L.__value="NUMERIC",Ye(L,L.__value),I.__value="BOOLEAN",Ye(I,I.__value),a(p,"class","py-2 px-3 pe-11 block w-full border-gray-200 shadow-sm rounded-lg text-sm focus:border-blue-500 focus:ring-blue-500 disabled:opacity-50 disabled:pointer-events-none dark:bg-neutral-900 dark:border-neutral-700 dark:text-neutral-400 dark:placeholder-neutral-500 dark:focus:ring-neutral-600"),E[12].general_type===void 0&&dt(N),a(_,"class","space-y-2"),a(e,"class","space-y-4")},m(D,B){V(D,e,B),l(e,T),l(T,t),l(T,r),l(T,R),Ye(R,E[12].name),l(e,n),l(e,s),l(s,S),l(s,A),l(s,o),Ye(o,E[12].description),l(e,i),l(e,_),l(_,c),l(_,P),l(_,p),l(p,C),l(p,L),l(p,I),WT(p,E[12].general_type,!0),l(e,u),H||(b=[Ne(R,"input",M),Ne(o,"input",O),Ne(p,"change",N)],H=!0)},p(D,B){E=D,B&2&&R.value!==E[12].name&&Ye(R,E[12].name),B&2&&o.value!==E[12].description&&Ye(o,E[12].description),B&2&&WT(p,E[12].general_type)},d(D){D&&Y(e),H=!1,NE(b)}}}function Ii(E){let e;function T(R,n){return R[1]!=null?ai:ii}let t=T(E),r=t(E);return{c(){r.c(),e=je()},m(R,n){r.m(R,n),V(R,e,n)},p(R,[n]){t===(t=T(R))&&r?r.p(R,n):(r.d(1),r=t(R),r&&(r.c(),r.m(e.parentNode,e)))},i:j,o:j,d(R){R&&Y(e),r.d(R)}}}function Ni(E,e,T){let{functionTemplate:t}=e,r=null;function R(){r.function_name=this.value,T(1,r)}function n(){r.description=this.value,T(1,r)}function s(){r.sql_template=this.value,T(1,r)}function S(p,C){p[C].name=this.value,T(1,r)}function A(p,C){p[C].description=this.value,T(1,r)}function o(p,C){p[C].general_type=On(this),T(1,r)}const i=()=>{T(1,r=null)},_=()=>{r&&(Vn(t.function_name,r),T(0,t=r),T(1,r=null))},c=()=>{t&&window.confirm("Are you sure you want to delete this function?")&&Wn(t.function_name)},P=()=>{T(1,r=structuredClone(t))};return E.$$set=p=>{"functionTemplate"in p&&T(0,t=p.functionTemplate)},[t,r,R,n,s,S,A,o,i,_,c,P]}class en extends ue{constructor(e){super(),Ce(this,e,Ni,Ii,_e,{functionTemplate:0})}}function Fr(E,e,T){const t=E.slice();return t[15]=e[T],t}function Yr(E,e,T){const t=E.slice();return t[18]=e[T],t}function Vr(E,e,T){const t=E.slice();return t[21]=e[T],t}function Wr(E,e,T){const t=E.slice();return t[21]=e[T],t}function wr(E){let e,T;return e=new ZA({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function $r(E){let e,T;return e=new Ze({props:{$$slots:{default:[_i]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108866&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function xr(E){let e,T;return e=new IE({props:{message:E[21],onSubmit:uT}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&2&&(R.message=t[21]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function li(E){let e=E[1].header+"",T,t,r,R,n=De(E[1].questions),s=[];for(let A=0;A<n.length;A+=1)s[A]=xr(Wr(E,n,A));const S=A=>y(s[A],1,1,()=>{s[A]=null});return{c(){T=te(e),t=$();for(let A=0;A<s.length;A+=1)s[A].c();r=je()},m(A,o){V(A,T,o),V(A,t,o);for(let i=0;i<s.length;i+=1)s[i]&&s[i].m(A,o);V(A,r,o),R=!0},p(A,o){if((!R||o&2)&&e!==(e=A[1].header+"")&&Le(T,e),o&2){n=De(A[1].questions);let i;for(i=0;i<n.length;i+=1){const _=Wr(A,n,i);s[i]?(s[i].p(_,o),m(s[i],1)):(s[i]=xr(_),s[i].c(),m(s[i],1),s[i].m(r.parentNode,r))}for(Ge(),i=n.length;i<s.length;i+=1)S(i);ge()}},i(A){if(!R){for(let o=0;o<n.length;o+=1)m(s[o]);R=!0}},o(A){s=s.filter(Boolean);for(let o=0;o<s.length;o+=1)y(s[o]);R=!1},d(A){A&&(Y(T),Y(t),Y(r)),nE(s,A)}}}function _i(E){let e,T;return e=new aE({props:{$$slots:{default:[li]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108866&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Li(E){let e,T;return e=new Ze({props:{$$slots:{default:[Fi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ci(E){let e,T;return e=new Ze({props:{$$slots:{default:[Vi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108888&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ui(E){let e,T;return e=new Ze({props:{$$slots:{default:[Wi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ci(E){let e,T;return e=new Ze({props:{$$slots:{default:[wi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function fi(E){let e,T;return e=new Ze({props:{$$slots:{default:[xi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Pi(E){let e,T;return e=new WE({props:{message:"Put your SQL here",$$slots:{default:[Xi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108864&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Di(E){let e,T,t,r,R,n,s,S,A,o;e=new WE({props:{message:E[18].question}}),t=new Ze({props:{$$slots:{default:[Ji]},$$scope:{ctx:E}}}),R=new Ze({props:{$$slots:{default:[qi]},$$scope:{ctx:E}}}),s=new Ze({props:{$$slots:{default:[Qi]},$$scope:{ctx:E}}});let i=E[18].summary&&Xr(E);return{c(){K(e.$$.fragment),T=$(),K(t.$$.fragment),r=$(),K(R.$$.fragment),n=$(),K(s.$$.fragment),S=$(),i&&i.c(),A=je()},m(_,c){X(e,_,c),V(_,T,c),X(t,_,c),V(_,r,c),X(R,_,c),V(_,n,c),X(s,_,c),V(_,S,c),i&&i.m(_,c),V(_,A,c),o=!0},p(_,c){const P={};c&8&&(P.message=_[18].question),e.$set(P);const p={};c&67108872&&(p.$$scope={dirty:c,ctx:_}),t.$set(p);const C={};c&67108872&&(C.$$scope={dirty:c,ctx:_}),R.$set(C);const L={};c&67108872&&(L.$$scope={dirty:c,ctx:_}),s.$set(L),_[18].summary?i?(i.p(_,c),c&8&&m(i,1)):(i=Xr(_),i.c(),m(i,1),i.m(A.parentNode,A)):i&&(Ge(),y(i,1,1,()=>{i=null}),ge())},i(_){o||(m(e.$$.fragment,_),m(t.$$.fragment,_),m(R.$$.fragment,_),m(s.$$.fragment,_),m(i),o=!0)},o(_){y(e.$$.fragment,_),y(t.$$.fragment,_),y(R.$$.fragment,_),y(s.$$.fragment,_),y(i),o=!1},d(_){_&&(Y(T),Y(r),Y(n),Y(S),Y(A)),k(e,_),k(t,_),k(R,_),k(s,_),i&&i.d(_)}}}function di(E){let e,T;return e=new Ze({props:{$$slots:{default:[zi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108873&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function pi(E){let e,T;return e=new Ze({props:{$$slots:{default:[ea]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Mi(E){let e,T;return e=new WE({props:{message:"No, the results were not correct."}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p:j,i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ui(E){let e,T;return e=new WE({props:{message:"Yes, the results were correct."}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p:j,i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function mi(E){let e,T,t=E[0].ask_results_correct&&Kr(E);return{c(){t&&t.c(),e=je()},m(r,R){t&&t.m(r,R),V(r,e,R),T=!0},p(r,R){r[0].ask_results_correct?t?(t.p(r,R),R&1&&m(t,1)):(t=Kr(r),t.c(),m(t,1),t.m(e.parentNode,e)):t&&(Ge(),y(t,1,1,()=>{t=null}),ge())},i(r){T||(m(t),T=!0)},o(r){y(t),T=!1},d(r){r&&Y(e),t&&t.d(r)}}}function hi(E){let e,T,t=E[0].ask_results_correct&&qr(E);return{c(){t&&t.c(),e=je()},m(r,R){t&&t.m(r,R),V(r,e,R),T=!0},p(r,R){r[0].ask_results_correct?t?R&1&&m(t,1):(t=qr(r),t.c(),m(t,1),t.m(e.parentNode,e)):t&&(Ge(),y(t,1,1,()=>{t=null}),ge())},i(r){T||(m(t),T=!0)},o(r){y(t),T=!1},d(r){r&&Y(e),t&&t.d(r)}}}function Gi(E){let e,T;return e=new WE({props:{message:"Change the chart based on these instructions",$$slots:{default:[Ra]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108864&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function gi(E){let e,T,t=E[0].chart&&Qr(E);return{c(){t&&t.c(),e=je()},m(r,R){t&&t.m(r,R),V(r,e,R),T=!0},p(r,R){r[0].chart?t?(t.p(r,R),R&1&&m(t,1)):(t=Qr(r),t.c(),m(t,1),t.m(e.parentNode,e)):t&&(Ge(),y(t,1,1,()=>{t=null}),ge())},i(r){T||(m(t),T=!0)},o(r){y(t),T=!1},d(r){r&&Y(e),t&&t.d(r)}}}function Hi(E){let e,T,t=E[0].table&&jr(E);return{c(){t&&t.c(),e=je()},m(r,R){t&&t.m(r,R),V(r,e,R),T=!0},p(r,R){r[0].table?t?(t.p(r,R),R&1&&m(t,1)):(t=jr(r),t.c(),m(t,1),t.m(e.parentNode,e)):t&&(Ge(),y(t,1,1,()=>{t=null}),ge())},i(r){T||(m(t),T=!0)},o(r){y(t),T=!1},d(r){r&&Y(e),t&&t.d(r)}}}function bi(E){let e,T;return e=new Ze({props:{$$slots:{default:[oa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function yi(E){let e,T,t=E[0].sql==!0&&eR(E);return{c(){t&&t.c(),e=je()},m(r,R){t&&t.m(r,R),V(r,e,R),T=!0},p(r,R){r[0].sql==!0?t?(t.p(r,R),R&1&&m(t,1)):(t=eR(r),t.c(),m(t,1),t.m(e.parentNode,e)):t&&(Ge(),y(t,1,1,()=>{t=null}),ge())},i(r){T||(m(t),T=!0)},o(r){y(t),T=!1},d(r){r&&Y(e),t&&t.d(r)}}}function Bi(E){let e,T;return e=new WE({props:{message:E[18].question}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.message=t[18].question),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function vi(E){let e=JSON.stringify(E[18])+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p(t,r){r&8&&e!==(e=JSON.stringify(t[18])+"")&&Le(T,e)},d(t){t&&Y(T)}}}function Fi(E){let e,T;return e=new aE({props:{$$slots:{default:[vi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Yi(E){let e,T=E[18].question+"",t;return{c(){e=te("I interpreted your question as: "),t=te(T)},m(r,R){V(r,e,R),V(r,t,R)},p(r,R){R&8&&T!==(T=r[18].question+"")&&Le(t,T)},d(r){r&&(Y(e),Y(t))}}}function Vi(E){let e,T,t,r,R,n;e=new aE({props:{$$slots:{default:[Yi]},$$scope:{ctx:E}}});function s(){return E[12](E[18])}return t=new IE({props:{message:"Edit Question",onSubmit:s}}),R=new IE({props:{message:"New Question",onSubmit:E[13]}}),{c(){K(e.$$.fragment),T=$(),K(t.$$.fragment),r=$(),K(R.$$.fragment)},m(S,A){X(e,S,A),V(S,T,A),X(t,S,A),V(S,r,A),X(R,S,A),n=!0},p(S,A){E=S;const o={};A&67108872&&(o.$$scope={dirty:A,ctx:E}),e.$set(o);const i={};A&24&&(i.onSubmit=s),t.$set(i)},i(S){n||(m(e.$$.fragment,S),m(t.$$.fragment,S),m(R.$$.fragment,S),n=!0)},o(S){y(e.$$.fragment,S),y(t.$$.fragment,S),y(R.$$.fragment,S),n=!1},d(S){S&&(Y(T),Y(r)),k(e,S),k(t,S),k(R,S)}}}function Wi(E){let e,T;return e=new en({props:{functionTemplate:E[18].function_template}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.functionTemplate=t[18].function_template),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function wi(E){let e,T;return e=new Oi({props:{functionData:E[18].function}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.functionData=t[18].function),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function $i(E){let e=E[18].text+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p(t,r){r&8&&e!==(e=t[18].text+"")&&Le(T,e)},d(t){t&&Y(T)}}}function xi(E){let e,T,t,r;return e=new aE({props:{$$slots:{default:[$i]},$$scope:{ctx:E}}}),t=new qA({props:{message:E[18].text}}),{c(){K(e.$$.fragment),T=$(),K(t.$$.fragment)},m(R,n){X(e,R,n),V(R,T,n),X(t,R,n),r=!0},p(R,n){const s={};n&67108872&&(s.$$scope={dirty:n,ctx:R}),e.$set(s);const S={};n&8&&(S.message=R[18].text),t.$set(S)},i(R){r||(m(e.$$.fragment,R),m(t.$$.fragment,R),r=!0)},o(R){y(e.$$.fragment,R),y(t.$$.fragment,R),r=!1},d(R){R&&Y(T),k(e,R),k(t,R)}}}function Xi(E){let e,T;return e=new $A({props:{onSubmit:vn}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p:j,i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ki(E){let e=E[18].sql+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p(t,r){r&8&&e!==(e=t[18].sql+"")&&Le(T,e)},d(t){t&&Y(T)}}}function Ki(E){let e,T;return e=new bR({props:{$$slots:{default:[ki]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ji(E){let e,T;return e=new aE({props:{$$slots:{default:[Ki]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function qi(E){let e,T;return e=new gR({props:{id:E[18].id,df:E[18].df}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.id=t[18].id),r&8&&(R.df=t[18].df),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Qi(E){let e,T;return e=new HR({props:{fig:E[18].fig}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.fig=t[18].fig),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Xr(E){let e,T;return e=new Ze({props:{$$slots:{default:[ji]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Zi(E){let e=E[18].summary+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p(t,r){r&8&&e!==(e=t[18].summary+"")&&Le(T,e)},d(t){t&&Y(T)}}}function ji(E){let e,T;return e=new aE({props:{$$slots:{default:[Zi]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function kr(E){let e,T;function t(){return E[11](E[18])}return e=new IE({props:{message:"Auto Fix",onSubmit:t}}),{c(){K(e.$$.fragment)},m(r,R){X(e,r,R),T=!0},p(r,R){E=r;const n={};R&8&&(n.onSubmit=t),e.$set(n)},i(r){T||(m(e.$$.fragment,r),T=!0)},o(r){y(e.$$.fragment,r),T=!1},d(r){k(e,r)}}}function zi(E){let e,T,t,r,R,n;e=new PT({props:{message:E[18].error}}),t=new IE({props:{message:"Manually Fix",onSubmit:E[10]}});let s=E[0].auto_fix_sql&&kr(E);return{c(){K(e.$$.fragment),T=$(),K(t.$$.fragment),r=$(),s&&s.c(),R=je()},m(S,A){X(e,S,A),V(S,T,A),X(t,S,A),V(S,r,A),s&&s.m(S,A),V(S,R,A),n=!0},p(S,A){const o={};A&8&&(o.message=S[18].error),e.$set(o),S[0].auto_fix_sql?s?(s.p(S,A),A&1&&m(s,1)):(s=kr(S),s.c(),m(s,1),s.m(R.parentNode,R)):s&&(Ge(),y(s,1,1,()=>{s=null}),ge())},i(S){n||(m(e.$$.fragment,S),m(t.$$.fragment,S),m(s),n=!0)},o(S){y(e.$$.fragment,S),y(t.$$.fragment,S),y(s),n=!1},d(S){S&&(Y(T),Y(r),Y(R)),k(e,S),k(t,S),s&&s.d(S)}}}function ea(E){let e,T;return e=new PT({props:{message:E[18].error}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.message=t[18].error),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Kr(E){let e,T;return e=new WE({props:{message:"",$$slots:{default:[ta]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108865&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ea(E){let e,T,t,r=E[0].function_generation&&Jr(E);return T=new IE({props:{message:"Yes, train as Question-SQL pair",onSubmit:E[8]}}),{c(){r&&r.c(),e=$(),K(T.$$.fragment)},m(R,n){r&&r.m(R,n),V(R,e,n),X(T,R,n),t=!0},p(R,n){R[0].function_generation?r?n&1&&m(r,1):(r=Jr(R),r.c(),m(r,1),r.m(e.parentNode,e)):r&&(Ge(),y(r,1,1,()=>{r=null}),ge())},i(R){t||(m(r),m(T.$$.fragment,R),t=!0)},o(R){y(r),y(T.$$.fragment,R),t=!1},d(R){R&&Y(e),r&&r.d(R),k(T,R)}}}function Jr(E){let e,T;return e=new IE({props:{message:"Yes, create function",onSubmit:E[7]}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ta(E){let e=fT()!==null,T,t,r,R=e&&Ea(E);return t=new IE({props:{message:"No",onSubmit:E[9]}}),{c(){R&&R.c(),T=$(),K(t.$$.fragment)},m(n,s){R&&R.m(n,s),V(n,T,s),X(t,n,s),r=!0},p(n,s){e&&R.p(n,s)},i(n){r||(m(R),m(t.$$.fragment,n),r=!0)},o(n){y(R),y(t.$$.fragment,n),r=!1},d(n){n&&Y(T),R&&R.d(n),k(t,n)}}}function qr(E){let e,T;return e=new Ze({props:{$$slots:{default:[ra]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ta(E){let e;return{c(){e=te("Were the results correct?")},m(T,t){V(T,e,t)},d(T){T&&Y(e)}}}function ra(E){let e,T;return e=new aE({props:{$$slots:{default:[Ta]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108864&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ra(E){let e,T;return e=new kA({props:{onSubmit:E[6],placeholder:"Make the line red",buttonText:"Update Chart",currentValue:""}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p:j,i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Qr(E){let e,T,t,r;e=new Ze({props:{$$slots:{default:[na]},$$scope:{ctx:E}}});let R=E[0].redraw_chart&&Zr(E);return{c(){K(e.$$.fragment),T=$(),R&&R.c(),t=je()},m(n,s){X(e,n,s),V(n,T,s),R&&R.m(n,s),V(n,t,s),r=!0},p(n,s){const S={};s&67108872&&(S.$$scope={dirty:s,ctx:n}),e.$set(S),n[0].redraw_chart?R?s&1&&m(R,1):(R=Zr(n),R.c(),m(R,1),R.m(t.parentNode,t)):R&&(Ge(),y(R,1,1,()=>{R=null}),ge())},i(n){r||(m(e.$$.fragment,n),m(R),r=!0)},o(n){y(e.$$.fragment,n),y(R),r=!1},d(n){n&&(Y(T),Y(t)),k(e,n),R&&R.d(n)}}}function na(E){let e,T;return e=new HR({props:{fig:E[18].fig}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.fig=t[18].fig),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Zr(E){let e,T;return e=new Ze({props:{$$slots:{default:[Aa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Aa(E){let e,T;return e=new IE({props:{message:"Redraw Chart",onSubmit:Fn}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p:j,i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function jr(E){let e,T;return e=new Ze({props:{$$slots:{default:[sa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function sa(E){let e,T;return e=new gR({props:{id:E[18].id,df:E[18].df}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.id=t[18].id),r&8&&(R.df=t[18].df),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function zr(E){let e,T;return e=new IE({props:{message:E[21],onSubmit:uT}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&8&&(R.message=t[21]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Sa(E){let e=E[18].header+"",T,t,r,R,n=De(E[18].questions),s=[];for(let A=0;A<n.length;A+=1)s[A]=zr(Vr(E,n,A));const S=A=>y(s[A],1,1,()=>{s[A]=null});return{c(){T=te(e),t=$();for(let A=0;A<s.length;A+=1)s[A].c();r=je()},m(A,o){V(A,T,o),V(A,t,o);for(let i=0;i<s.length;i+=1)s[i]&&s[i].m(A,o);V(A,r,o),R=!0},p(A,o){if((!R||o&8)&&e!==(e=A[18].header+"")&&Le(T,e),o&8){n=De(A[18].questions);let i;for(i=0;i<n.length;i+=1){const _=Vr(A,n,i);s[i]?(s[i].p(_,o),m(s[i],1)):(s[i]=zr(_),s[i].c(),m(s[i],1),s[i].m(r.parentNode,r))}for(Ge(),i=n.length;i<s.length;i+=1)S(i);ge()}},i(A){if(!R){for(let o=0;o<n.length;o+=1)m(s[o]);R=!0}},o(A){s=s.filter(Boolean);for(let o=0;o<s.length;o+=1)y(s[o]);R=!1},d(A){A&&(Y(T),Y(t),Y(r)),nE(s,A)}}}function oa(E){let e,T;return e=new aE({props:{$$slots:{default:[Sa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function eR(E){let e,T;return e=new Ze({props:{$$slots:{default:[aa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Oa(E){let e=E[18].text+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p(t,r){r&8&&e!==(e=t[18].text+"")&&Le(T,e)},d(t){t&&Y(T)}}}function ia(E){let e,T;return e=new bR({props:{$$slots:{default:[Oa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function aa(E){let e,T;return e=new aE({props:{$$slots:{default:[ia]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&67108872&&(R.$$scope={dirty:r,ctx:t}),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ER(E){let e,T,t,r;const R=[Bi,yi,bi,Hi,gi,Gi,hi,mi,Ui,Mi,pi,di,Di,Pi,fi,ci,ui,Ci,Li],n=[];function s(S,A){return S[18].type==="user_question"?0:S[18].type==="sql"?1:S[18].type==="question_list"?2:S[18].type==="df"?3:S[18].type==="plotly_figure"?4:S[18].type==="chart_modification"?5:S[18].type==="feedback_question"?6:S[18].type==="feedback_buttons"?7:S[18].type==="feedback_correct"?8:S[18].type==="feedback_incorrect"?9:S[18].type==="error"?10:S[18].type==="sql_error"?11:S[18].type==="question_cache"?12:S[18].type==="user_sql"?13:S[18].type==="text"?14:S[18].type==="function"?15:S[18].type==="function_template"?16:S[18].type==="rewritten_question"?17:18}return e=s(E),T=n[e]=R[e](E),{c(){T.c(),t=je()},m(S,A){n[e].m(S,A),V(S,t,A),r=!0},p(S,A){let o=e;e=s(S),e===o?n[e].p(S,A):(Ge(),y(n[o],1,1,()=>{n[o]=null}),ge(),T=n[e],T?T.p(S,A):(T=n[e]=R[e](S),T.c()),m(T,1),T.m(t.parentNode,t))},i(S){r||(m(T),r=!0)},o(S){y(T),r=!1},d(S){S&&Y(t),n[e].d(S)}}}function tR(E){let e,T;return e=new dA({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function TR(E){let e,T,t=De(E[3]),r=[];for(let n=0;n<t.length;n+=1)r[n]=RR(Fr(E,t,n));const R=n=>y(r[n],1,1,()=>{r[n]=null});return{c(){for(let n=0;n<r.length;n+=1)r[n].c();e=je()},m(n,s){for(let S=0;S<r.length;S+=1)r[S]&&r[S].m(n,s);V(n,e,s),T=!0},p(n,s){if(s&8){t=De(n[3]);let S;for(S=0;S<t.length;S+=1){const A=Fr(n,t,S);r[S]?(r[S].p(A,s),m(r[S],1)):(r[S]=RR(A),r[S].c(),m(r[S],1),r[S].m(e.parentNode,e))}for(Ge(),S=t.length;S<r.length;S+=1)R(S);ge()}},i(n){if(!T){for(let s=0;s<t.length;s+=1)m(r[s]);T=!0}},o(n){r=r.filter(Boolean);for(let s=0;s<r.length;s+=1)y(r[s]);T=!1},d(n){n&&Y(e),nE(r,n)}}}function rR(E){let e,T;function t(){return E[14](E[15])}return e=new yA({props:{message:"Re-Run SQL",onSubmit:t}}),{c(){K(e.$$.fragment)},m(r,R){X(e,r,R),T=!0},p(r,R){E=r;const n={};R&8&&(n.onSubmit=t),e.$set(n)},i(r){T||(m(e.$$.fragment,r),T=!0)},o(r){y(e.$$.fragment,r),T=!1},d(r){k(e,r)}}}function RR(E){let e,T,t=E[15].type==="question_cache"&&rR(E);return{c(){t&&t.c(),e=je()},m(r,R){t&&t.m(r,R),V(r,e,R),T=!0},p(r,R){r[15].type==="question_cache"?t?(t.p(r,R),R&8&&m(t,1)):(t=rR(r),t.c(),m(t,1),t.m(e.parentNode,e)):t&&(Ge(),y(t,1,1,()=>{t=null}),ge())},i(r){T||(m(t),T=!0)},o(r){y(t),T=!1},d(r){r&&Y(e),t&&t.d(r)}}}function Ia(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p;t=new EA({});let C=E[0].debug&&wr(),L=E[1]&&E[1].type=="question_list"&&!E[2]&&$r(E),I=De(E[3]),u=[];for(let O=0;O<I.length;O+=1)u[O]=ER(Yr(E,I,O));const H=O=>y(u[O],1,1,()=>{u[O]=null});let b=E[5]&&tR();i=new lA({});let M=E[2]&&TR(E);return P=new IA({props:{onSubmit:uT}}),{c(){e=f("div"),T=f("div"),K(t.$$.fragment),r=$(),C&&C.c(),R=$(),L&&L.c(),n=$(),s=f("ul");for(let O=0;O<u.length;O+=1)u[O].c();S=$(),b&&b.c(),A=$(),o=f("footer"),K(i.$$.fragment),_=$(),M&&M.c(),c=$(),K(P.$$.fragment),a(s,"class","mt-16 space-y-5"),a(T,"class","py-10 lg:py-14"),a(o,"class","max-w-4xl mx-auto sticky bottom-0 z-10 p-3 sm:py-6"),a(e,"class","relative h-screen w-full lg:pl-64")},m(O,N){V(O,e,N),l(e,T),X(t,T,null),l(T,r),C&&C.m(T,null),l(T,R),L&&L.m(T,null),l(T,n),l(T,s);for(let D=0;D<u.length;D+=1)u[D]&&u[D].m(s,null);l(s,S),b&&b.m(s,null),l(e,A),l(e,o),X(i,o,null),l(o,_),M&&M.m(o,null),l(o,c),X(P,o,null),p=!0},p(O,[N]){if(O[0].debug?C?N&1&&m(C,1):(C=wr(),C.c(),m(C,1),C.m(T,R)):C&&(Ge(),y(C,1,1,()=>{C=null}),ge()),O[1]&&O[1].type=="question_list"&&!O[2]?L?(L.p(O,N),N&6&&m(L,1)):(L=$r(O),L.c(),m(L,1),L.m(T,n)):L&&(Ge(),y(L,1,1,()=>{L=null}),ge()),N&25){I=De(O[3]);let D;for(D=0;D<I.length;D+=1){const B=Yr(O,I,D);u[D]?(u[D].p(B,N),m(u[D],1)):(u[D]=ER(B),u[D].c(),m(u[D],1),u[D].m(s,S))}for(Ge(),D=I.length;D<u.length;D+=1)H(D);ge()}O[5]?b?N&32&&m(b,1):(b=tR(),b.c(),m(b,1),b.m(s,null)):b&&(Ge(),y(b,1,1,()=>{b=null}),ge()),O[2]?M?(M.p(O,N),N&4&&m(M,1)):(M=TR(O),M.c(),m(M,1),M.m(o,c)):M&&(Ge(),y(M,1,1,()=>{M=null}),ge())},i(O){if(!p){m(t.$$.fragment,O),m(C),m(L);for(let N=0;N<I.length;N+=1)m(u[N]);m(b),m(i.$$.fragment,O),m(M),m(P.$$.fragment,O),p=!0}},o(O){y(t.$$.fragment,O),y(C),y(L),u=u.filter(Boolean);for(let N=0;N<u.length;N+=1)y(u[N]);y(b),y(i.$$.fragment,O),y(M),y(P.$$.fragment,O),p=!1},d(O){O&&Y(e),k(t),C&&C.d(),L&&L.d(),nE(u,O),b&&b.d(),k(i),M&&M.d(),k(P)}}}function Na(E,e,T){let t,r,R,n,s,S;return eE(E,VE,I=>T(0,t=I)),eE(E,LT,I=>T(1,r=I)),eE(E,Ht,I=>T(2,R=I)),eE(E,YE,I=>T(3,n=I)),eE(E,BE,I=>T(4,s=I)),eE(E,St,I=>T(5,S=I)),[t,r,R,n,s,S,I=>{$n(I)},()=>{Yn()},()=>{wn()},()=>{wT()},()=>{wT()},I=>{Bn(I.error)},I=>{it(),OT(BE,s=I.question,s)},()=>{it()},I=>I.type==="question_cache"?dn(I.id):void 0]}class la extends ue{constructor(e){super(),Ce(this,e,Na,Ia,_e,{})}}function _a(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p,C,L,I,u;return{c(){e=f("div"),T=f("div"),t=f("div"),r=f("div"),R=f("h3"),R.textContent="Are you sure?",n=$(),s=f("button"),s.innerHTML='<span class="sr-only">Close</span> <svg class="w-3.5 h-3.5" width="8" height="8" viewBox="0 0 8 8" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M0.258206 1.00652C0.351976 0.912791 0.479126 0.860131 0.611706 0.860131C0.744296 0.860131 0.871447 0.912791 0.965207 1.00652L3.61171 3.65302L6.25822 1.00652C6.30432 0.958771 6.35952 0.920671 6.42052 0.894471C6.48152 0.868271 6.54712 0.854471 6.61352 0.853901C6.67992 0.853321 6.74572 0.865971 6.80722 0.891111C6.86862 0.916251 6.92442 0.953381 6.97142 1.00032C7.01832 1.04727 7.05552 1.1031 7.08062 1.16454C7.10572 1.22599 7.11842 1.29183 7.11782 1.35822C7.11722 1.42461 7.10342 1.49022 7.07722 1.55122C7.05102 1.61222 7.01292 1.6674 6.96522 1.71352L4.31871 4.36002L6.96522 7.00648C7.05632 7.10078 7.10672 7.22708 7.10552 7.35818C7.10442 7.48928 7.05182 7.61468 6.95912 7.70738C6.86642 7.80018 6.74102 7.85268 6.60992 7.85388C6.47882 7.85498 6.35252 7.80458 6.25822 7.71348L3.61171 5.06702L0.965207 7.71348C0.870907 7.80458 0.744606 7.85498 0.613506 7.85388C0.482406 7.85268 0.357007 7.80018 0.264297 7.70738C0.171597 7.61468 0.119017 7.48928 0.117877 7.35818C0.116737 7.22708 0.167126 7.10078 0.258206 7.00648L2.90471 4.36002L0.258206 1.71352C0.164476 1.61976 0.111816 1.4926 0.111816 1.36002C0.111816 1.22744 0.164476 1.10028 0.258206 1.00652Z" fill="currentColor"></path></svg>',S=$(),A=f("div"),o=f("p"),i=te(E[0]),_=$(),c=f("div"),P=f("button"),P.textContent="Close",p=$(),C=f("button"),L=te(E[1]),a(R,"class","font-bold text-gray-800 dark:text-white"),a(s,"type","button"),a(s,"class","hs-dropdown-toggle inline-flex flex-shrink-0 justify-center items-center h-8 w-8 rounded-md text-gray-500 hover:text-gray-400 focus:outline-none focus:ring-2 focus:ring-gray-400 focus:ring-offset-2 focus:ring-offset-white transition-all text-sm dark:focus:ring-gray-700 dark:focus:ring-offset-gray-800"),a(s,"data-hs-overlay","#hs-vertically-centered-modal"),a(r,"class","flex justify-between items-center py-3 px-4 border-b dark:border-gray-700"),a(o,"class","text-gray-800 dark:text-gray-400"),a(A,"class","p-4 overflow-y-auto"),a(P,"type","button"),a(P,"class","hs-dropdown-toggle py-3 px-4 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700 shadow-sm align-middle hover:bg-gray-50 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white focus:ring-blue-600 transition-all text-sm dark:bg-slate-900 dark:hover:bg-slate-800 dark:border-gray-700 dark:text-gray-400 dark:hover:text-white dark:focus:ring-offset-gray-800"),a(P,"data-hs-overlay","#hs-vertically-centered-modal"),a(C,"class","py-3 px-4 inline-flex justify-center items-center gap-2 rounded-md border border-transparent font-semibold bg-blue-500 text-white hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-500 focus:ring-offset-2 transition-all text-sm dark:focus:ring-offset-gray-800"),a(c,"class","flex justify-end items-center gap-x-2 py-3 px-4 border-t dark:border-gray-700"),a(t,"class","flex flex-col bg-white border shadow-sm rounded-xl dark:bg-gray-800 dark:border-gray-700 dark:shadow-slate-700/[.7]"),a(T,"class","hs-overlay-open:mt-7 hs-overlay-open:opacity-100 hs-overlay-open:duration-500 mt-0 opacity-0 ease-out transition-all sm:max-w-lg sm:w-full m-3 sm:mx-auto min-h-[calc(100%-3.5rem)] flex items-center"),a(e,"class","hs-overlay open w-full h-full fixed top-0 left-0 z-[60] overflow-x-hidden overflow-y-auto")},m(H,b){V(H,e,b),l(e,T),l(T,t),l(t,r),l(r,R),l(r,n),l(r,s),l(t,S),l(t,A),l(A,o),l(o,i),l(t,_),l(t,c),l(c,P),l(c,p),l(c,C),l(C,L),I||(u=[Ne(s,"click",function(){zE(E[2])&&E[2].apply(this,arguments)}),Ne(P,"click",function(){zE(E[2])&&E[2].apply(this,arguments)}),Ne(C,"click",function(){zE(E[3])&&E[3].apply(this,arguments)})],I=!0)},p(H,[b]){E=H,b&1&&Le(i,E[0]),b&2&&Le(L,E[1])},i:j,o:j,d(H){H&&Y(e),I=!1,NE(u)}}}function La(E,e,T){let{message:t}=e,{buttonLabel:r}=e,{onClose:R}=e,{onConfirm:n}=e;return E.$$set=s=>{"message"in s&&T(0,t=s.message),"buttonLabel"in s&&T(1,r=s.buttonLabel),"onClose"in s&&T(2,R=s.onClose),"onConfirm"in s&&T(3,n=s.onConfirm)},[t,r,R,n]}class Ca extends ue{constructor(e){super(),Ce(this,e,La,_a,_e,{message:0,buttonLabel:1,onClose:2,onConfirm:3})}}function nR(E,e,T){const t=E.slice();return t[10]=e[T].name,t[11]=e[T].description,t[12]=e[T].example,t}function AR(E){let e,T,t,r,R,n,s,S,A,o,i,_;return o=Sn(E[7][0]),{c(){e=f("div"),T=f("div"),t=f("input"),r=$(),R=f("label"),n=f("span"),n.textContent=`${E[10]}`,s=$(),S=f("span"),S.textContent=`${E[11]}`,A=$(),a(t,"id","hs-radio-"+E[10]),t.__value=E[10],Ye(t,t.__value),a(t,"name","hs-radio-with-description"),a(t,"type","radio"),a(t,"class","border-gray-200 rounded-full text-blue-600 focus:ring-blue-500 dark:bg-gray-800 dark:border-gray-700 dark:checked:bg-blue-500 dark:checked:border-blue-500 dark:focus:ring-offset-gray-800"),a(t,"aria-describedby","hs-radio-delete-description"),a(T,"class","flex items-center h-5 mt-1"),a(n,"class","block text-sm font-semibold text-gray-800 dark:text-gray-300"),a(S,"id","hs-radio-ddl-description"),a(S,"class","block text-sm text-gray-600 dark:text-gray-500"),a(R,"for","hs-radio-"+E[10]),a(R,"class","ml-3"),a(e,"class","relative flex items-start"),o.p(t)},m(c,P){V(c,e,P),l(e,T),l(T,t),t.checked=t.__value===E[0],l(e,r),l(e,R),l(R,n),l(R,s),l(R,S),l(e,A),i||(_=Ne(t,"change",E[6]),i=!0)},p(c,P){P&1&&(t.checked=t.__value===c[0])},d(c){c&&Y(e),o.r(),i=!1,_()}}}function ua(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p,C,L,I,u,H,b,M,O,N,D,B=De(E[3]),h=[];for(let G=0;G<B.length;G+=1)h[G]=AR(nR(E,B,G));return{c(){var G;e=f("div"),T=f("div"),t=f("div"),r=f("div"),R=f("h2"),R.textContent="Add Training Data",n=$(),s=f("button"),s.innerHTML='<span class="sr-only">Close</span> <svg class="w-3.5 h-3.5" width="8" height="8" viewBox="0 0 8 8" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M0.258206 1.00652C0.351976 0.912791 0.479126 0.860131 0.611706 0.860131C0.744296 0.860131 0.871447 0.912791 0.965207 1.00652L3.61171 3.65302L6.25822 1.00652C6.30432 0.958771 6.35952 0.920671 6.42052 0.894471C6.48152 0.868271 6.54712 0.854471 6.61352 0.853901C6.67992 0.853321 6.74572 0.865971 6.80722 0.891111C6.86862 0.916251 6.92442 0.953381 6.97142 1.00032C7.01832 1.04727 7.05552 1.1031 7.08062 1.16454C7.10572 1.22599 7.11842 1.29183 7.11782 1.35822C7.11722 1.42461 7.10342 1.49022 7.07722 1.55122C7.05102 1.61222 7.01292 1.6674 6.96522 1.71352L4.31871 4.36002L6.96522 7.00648C7.05632 7.10078 7.10672 7.22708 7.10552 7.35818C7.10442 7.48928 7.05182 7.61468 6.95912 7.70738C6.86642 7.80018 6.74102 7.85268 6.60992 7.85388C6.47882 7.85498 6.35252 7.80458 6.25822 7.71348L3.61171 5.06702L0.965207 7.71348C0.870907 7.80458 0.744606 7.85498 0.613506 7.85388C0.482406 7.85268 0.357007 7.80018 0.264297 7.70738C0.171597 7.61468 0.119017 7.48928 0.117877 7.35818C0.116737 7.22708 0.167126 7.10078 0.258206 7.00648L2.90471 4.36002L0.258206 1.71352C0.164476 1.61976 0.111816 1.4926 0.111816 1.36002C0.111816 1.22744 0.164476 1.10028 0.258206 1.00652Z" fill="currentColor"></path></svg>',S=$(),A=f("span"),A.textContent="Training Data Type",o=$(),i=f("div");for(let F=0;F<h.length;F+=1)h[F].c();_=$(),c=f("div"),P=f("label"),p=te("Your "),C=te(E[0]),L=$(),I=f("div"),u=f("textarea"),b=$(),M=f("div"),O=f("button"),O.textContent="Save",a(R,"class","text-xl text-gray-800 font-bold sm:text-3xl dark:text-white"),a(s,"type","button"),a(s,"class","hs-dropdown-toggle inline-flex flex-shrink-0 justify-center items-center h-8 w-8 rounded-md text-gray-500 hover:text-gray-400 focus:outline-none focus:ring-2 focus:ring-gray-400 focus:ring-offset-2 focus:ring-offset-white transition-all text-sm dark:focus:ring-gray-700 dark:focus:ring-offset-gray-800"),a(s,"data-hs-overlay","#hs-vertically-centered-modal"),a(r,"class","flex justify-between items-center py-3 px-4 border-b dark:border-gray-700 mb-2"),a(A,"class","block mb-2 text-sm font-medium dark:text-white"),a(i,"class","grid space-y-3 mb-1"),a(P,"for","hs-feedback-post-comment-textarea-1"),a(P,"class","block mt-2 mb-2 text-sm font-medium dark:text-white"),a(u,"id","hs-feedback-post-comment-textarea-1"),a(u,"name","hs-feedback-post-comment-textarea-1"),a(u,"rows","3"),a(u,"class","py-3 px-4 block w-full border border-gray-200 rounded-md text-sm focus:border-blue-500 focus:ring-blue-500 sm:p-4 dark:bg-gray-800 dark:border-gray-700 dark:text-gray-400"),a(u,"placeholder",H=((G=E[3].find(E[8]))==null?void 0:G.example)??"No example available"),a(I,"class","mt-1"),a(c,"class","mt-2 border-t dark:border-gray-700"),a(O,"class","py-3 px-4 inline-flex justify-center items-center gap-2 rounded-md border border-transparent font-semibold bg-blue-500 text-white hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-500 focus:ring-offset-2 transition-all dark:focus:ring-offset-gray-800"),a(M,"class","mt-6 grid"),a(t,"class","mt-5 p-4 relative z-10 bg-white border rounded-xl sm:mt-10 md:p-10 dark:bg-gray-800 dark:border-gray-700"),a(T,"class","mx-auto max-w-2xl"),a(e,"class","max-w-[85rem] px-4 py-10 sm:px-6 lg:px-8 lg:py-14 mx-auto")},m(G,F){V(G,e,F),l(e,T),l(T,t),l(t,r),l(r,R),l(r,n),l(r,s),l(t,S),l(t,A),l(t,o),l(t,i);for(let W=0;W<h.length;W+=1)h[W]&&h[W].m(i,null);l(t,_),l(t,c),l(c,P),l(P,p),l(P,C),l(c,L),l(c,I),l(I,u),Ye(u,E[2]),l(t,b),l(t,M),l(M,O),N||(D=[Ne(s,"click",function(){zE(E[1])&&E[1].apply(this,arguments)}),Ne(u,"input",E[9]),Ne(O,"click",E[4])],N=!0)},p(G,[F]){var W;if(E=G,F&9){B=De(E[3]);let x;for(x=0;x<B.length;x+=1){const J=nR(E,B,x);h[x]?h[x].p(J,F):(h[x]=AR(J),h[x].c(),h[x].m(i,null))}for(;x<h.length;x+=1)h[x].d(1);h.length=B.length}F&1&&Le(C,E[0]),F&1&&H!==(H=((W=E[3].find(E[8]))==null?void 0:W.example)??"No example available")&&a(u,"placeholder",H),F&4&&Ye(u,E[2])},i:j,o:j,d(G){G&&Y(e),nE(h,G),N=!1,NE(D)}}}function ca(E,e,T){let{onDismiss:t}=e,{onTrain:r}=e,{selectedTrainingDataType:R="SQL"}=e,n=[{name:"DDL",description:"These are the CREATE TABLE statements that define your database structure.",example:"CREATE TABLE table_name (column_1 datatype, column_2 datatype, column_3 datatype);"},{name:"Documentation",description:"This can be any text-based documentation. Keep the chunks small and focused on a single topic.",example:"Our definition of ABC is XYZ."},{name:"SQL",description:"This can be any SQL statement that works. The more the merrier.",example:"SELECT column_1, column_2 FROM table_name;"}],s="";const S=()=>{r(s,R.toLowerCase())},A=[[]];function o(){R=this.__value,T(0,R)}const i=c=>c.name===R;function _(){s=this.value,T(2,s)}return E.$$set=c=>{"onDismiss"in c&&T(1,t=c.onDismiss),"onTrain"in c&&T(5,r=c.onTrain),"selectedTrainingDataType"in c&&T(0,R=c.selectedTrainingDataType)},[R,t,s,n,S,r,o,A,i,_]}class fa extends ue{constructor(e){super(),Ce(this,e,ca,ua,_e,{onDismiss:1,onTrain:5,selectedTrainingDataType:0})}}function sR(E,e,T){const t=E.slice();return t[21]=e[T],t}function SR(E,e,T){const t=E.slice();return t[24]=e[T],t}function oR(E,e,T){const t=E.slice();return t[24]=e[T],t}function iR(E){let e,T;return e=new fa({props:{onDismiss:E[13],onTrain:E[0]}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&1&&(R.onTrain=t[0]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Pa(E){let e;return{c(){e=te("Action")},m(T,t){V(T,e,t)},p:j,d(T){T&&Y(e)}}}function Da(E){let e=E[24]+"",T;return{c(){T=te(e)},m(t,r){V(t,T,r)},p:j,d(t){t&&Y(T)}}}function aR(E){let e,T,t,r;function R(S,A){return S[24]!="id"?Da:Pa}let s=R(E)(E);return{c(){e=f("th"),T=f("div"),t=f("span"),s.c(),r=$(),a(t,"class","text-xs font-semibold uppercase tracking-wide text-gray-800 dark:text-gray-200"),a(T,"class","flex items-center gap-x-2"),a(e,"scope","col"),a(e,"class","px-6 py-3 text-left")},m(S,A){V(S,e,A),l(e,T),l(T,t),s.m(t,null),l(e,r)},p(S,A){s.p(S,A)},d(S){S&&Y(e),s.d()}}}function da(E){let e,T,t;function r(){return E[18](E[21],E[24])}return{c(){e=f("button"),e.textContent="Delete",a(e,"type","button"),a(e,"class","py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border-2 border-red-200 font-semibold text-red-500 hover:text-white hover:bg-red-500 hover:border-red-500 focus:outline-none focus:ring-2 focus:ring-red-200 focus:ring-offset-2 transition-all text-sm dark:focus:ring-offset-gray-800")},m(R,n){V(R,e,n),T||(t=Ne(e,"click",r),T=!0)},p(R,n){E=R},d(R){R&&Y(e),T=!1,t()}}}function pa(E){let e,T=E[21][E[24]]+"",t;return{c(){e=f("span"),t=te(T),a(e,"class","text-gray-800 dark:text-gray-200")},m(r,R){V(r,e,R),l(e,t)},p(r,R){R&16&&T!==(T=r[21][r[24]]+"")&&Le(t,T)},d(r){r&&Y(e)}}}function IR(E){let e,T;function t(n,s){return n[24]!="id"?pa:da}let R=t(E)(E);return{c(){e=f("td"),T=f("div"),R.c(),a(T,"class","px-6 py-3"),a(e,"class","h-px w-px ")},m(n,s){V(n,e,s),l(e,T),R.m(T,null)},p(n,s){R.p(n,s)},d(n){n&&Y(e),R.d()}}}function NR(E){let e,T,t=De(E[8]),r=[];for(let R=0;R<t.length;R+=1)r[R]=IR(SR(E,t,R));return{c(){e=f("tr");for(let R=0;R<r.length;R+=1)r[R].c();T=$()},m(R,n){V(R,e,n);for(let s=0;s<r.length;s+=1)r[s]&&r[s].m(e,null);l(e,T)},p(R,n){if(n&304){t=De(R[8]);let s;for(s=0;s<t.length;s+=1){const S=SR(R,t,s);r[s]?r[s].p(S,n):(r[s]=IR(S),r[s].c(),r[s].m(e,T))}for(;s<r.length;s+=1)r[s].d(1);r.length=t.length}},d(R){R&&Y(e),nE(r,R)}}}function lR(E){let e,T;return e=new Ca({props:{message:"Are you sure you want to delete this?",buttonLabel:"Delete",onClose:E[19],onConfirm:E[20]}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&32&&(R.onClose=t[19]),r&34&&(R.onConfirm=t[20]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ma(E){let e,T,t,r,R,n,s,S,A,o,i,_,c,P,p,C,L,I,u,H,b,M,O,N,D,B,h,G=E[2]+1+"",F,W,x=Math.min(E[3],E[7].length)+"",J,oe,z,Oe,Ue,ye,TE,dE,me,fE,rE,pE,Tt,rt,Ve=E[6]&&iR(E),ME=De(E[8]),$e=[];for(let Te=0;Te<ME.length;Te+=1)$e[Te]=aR(oR(E,ME,Te));let ze=De(E[4]),ve=[];for(let Te=0;Te<ze.length;Te+=1)ve[Te]=NR(sR(E,ze,Te));let Fe=E[5]!=null&&lR(E);return{c(){Ve&&Ve.c(),e=$(),T=f("div"),t=f("div"),r=f("div"),R=f("div"),n=f("div"),s=f("div"),S=f("div"),S.innerHTML='<h2 class="text-xl font-semibold text-gray-800 dark:text-gray-200">Training Data</h2> <p class="text-sm text-gray-600 dark:text-gray-400">Add or remove training data. Good training data is the key to accuracy.</p>',A=$(),o=f("div"),i=f("div"),_=f("button"),_.textContent="View all",c=$(),P=f("button"),P.innerHTML=`<svg class="w-3 h-3" xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 16 16" fill="none"><path d="M2.63452 7.50001L13.6345 7.5M8.13452 13V2" stroke="currentColor" stroke-width="2" stroke-linecap="round"></path></svg>
                    Add training data`,p=$(),C=f("table"),L=f("thead"),I=f("tr");for(let Te=0;Te<$e.length;Te+=1)$e[Te].c();u=$(),H=f("tbody");for(let Te=0;Te<ve.length;Te+=1)ve[Te].c();b=$(),M=f("div"),O=f("div"),N=f("p"),N.textContent="Showing:",D=$(),B=f("div"),h=f("span"),F=te(G),W=te(" - "),J=te(x),oe=$(),z=f("p"),z.textContent=`of ${E[7].length}`,Oe=$(),Ue=f("div"),ye=f("div"),TE=f("button"),TE.innerHTML=`<svg class="w-3 h-3" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path fill-rule="evenodd" d="M11.354 1.646a.5.5 0 0 1 0 .708L5.707 8l5.647 5.646a.5.5 0 0 1-.708.708l-6-6a.5.5 0 0 1 0-.708l6-6a.5.5 0 0 1 .708 0z"></path></svg>
                    Prev`,dE=$(),me=f("button"),me.innerHTML=`Next
                    <svg class="w-3 h-3" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path fill-rule="evenodd" d="M4.646 1.646a.5.5 0 0 1 .708 0l6 6a.5.5 0 0 1 0 .708l-6 6a.5.5 0 0 1-.708-.708L10.293 8 4.646 2.354a.5.5 0 0 1 0-.708z"></path></svg>`,fE=$(),Fe&&Fe.c(),rE=je(),a(_,"class","py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700 shadow-sm align-middle hover:bg-gray-50 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white focus:ring-blue-600 transition-all text-sm dark:bg-slate-900 dark:hover:bg-slate-800 dark:border-gray-700 dark:text-gray-400 dark:hover:text-white dark:focus:ring-offset-gray-800"),a(P,"class","py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border border-transparent font-semibold bg-blue-500 text-white hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-500 focus:ring-offset-2 transition-all text-sm dark:focus:ring-offset-gray-800"),a(i,"class","inline-flex gap-x-2"),a(s,"class","px-6 py-4 grid gap-3 md:flex md:justify-between md:items-center border-b border-gray-200 dark:border-gray-700"),a(L,"class","bg-gray-50 dark:bg-slate-800"),a(H,"class","divide-y divide-gray-200 dark:divide-gray-700"),a(C,"class","min-w-full divide-y divide-gray-200 dark:divide-gray-700"),a(N,"class","text-sm text-gray-600 dark:text-gray-400"),a(h,"class","py-2 px-3 pr-9 block w-full border-gray-200 rounded-md text-sm focus:border-blue-500 focus:ring-blue-500 dark:bg-slate-900 dark:border-gray-700 dark:text-gray-400"),a(B,"class","max-w-sm space-y-3"),a(z,"class","text-sm text-gray-600 dark:text-gray-400"),a(O,"class","inline-flex items-center gap-x-2"),a(TE,"type","button"),a(TE,"class","py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700 shadow-sm align-middle hover:bg-gray-50 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white focus:ring-blue-600 transition-all text-sm dark:bg-slate-900 dark:hover:bg-slate-800 dark:border-gray-700 dark:text-gray-400 dark:hover:text-white dark:focus:ring-offset-gray-800"),a(me,"type","button"),a(me,"class","py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700 shadow-sm align-middle hover:bg-gray-50 focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-offset-white focus:ring-blue-600 transition-all text-sm dark:bg-slate-900 dark:hover:bg-slate-800 dark:border-gray-700 dark:text-gray-400 dark:hover:text-white dark:focus:ring-offset-gray-800"),a(ye,"class","inline-flex gap-x-2"),a(M,"class","px-6 py-4 grid gap-3 md:flex md:justify-between md:items-center border-t border-gray-200 dark:border-gray-700"),a(n,"class","bg-white border border-gray-200 rounded-xl shadow-sm overflow-hidden dark:bg-slate-900 dark:border-gray-700"),a(R,"class","p-1.5 min-w-full inline-block align-middle"),a(r,"class","-m-1.5 overflow-x-auto"),a(t,"class","flex flex-col"),a(T,"class","max-w-[85rem] px-4 py-10 sm:px-6 lg:px-8 lg:py-14 mx-auto")},m(Te,Xe){Ve&&Ve.m(Te,Xe),V(Te,e,Xe),V(Te,T,Xe),l(T,t),l(t,r),l(r,R),l(R,n),l(n,s),l(s,S),l(s,A),l(s,o),l(o,i),l(i,_),l(i,c),l(i,P),l(n,p),l(n,C),l(C,L),l(L,I);for(let se=0;se<$e.length;se+=1)$e[se]&&$e[se].m(I,null);l(C,u),l(C,H);for(let se=0;se<ve.length;se+=1)ve[se]&&ve[se].m(H,null);l(n,b),l(n,M),l(M,O),l(O,N),l(O,D),l(O,B),l(B,h),l(h,F),l(h,W),l(h,J),l(O,oe),l(O,z),l(M,Oe),l(M,Ue),l(Ue,ye),l(ye,TE),l(ye,dE),l(ye,me),V(Te,fE,Xe),Fe&&Fe.m(Te,Xe),V(Te,rE,Xe),pE=!0,Tt||(rt=[Ne(_,"click",E[11]),Ne(P,"click",E[12]),Ne(TE,"click",E[9]),Ne(me,"click",E[10])],Tt=!0)},p(Te,[Xe]){if(Te[6]?Ve?(Ve.p(Te,Xe),Xe&64&&m(Ve,1)):(Ve=iR(Te),Ve.c(),m(Ve,1),Ve.m(e.parentNode,e)):Ve&&(Ge(),y(Ve,1,1,()=>{Ve=null}),ge()),Xe&256){ME=De(Te[8]);let se;for(se=0;se<ME.length;se+=1){const gE=oR(Te,ME,se);$e[se]?$e[se].p(gE,Xe):($e[se]=aR(gE),$e[se].c(),$e[se].m(I,null))}for(;se<$e.length;se+=1)$e[se].d(1);$e.length=ME.length}if(Xe&304){ze=De(Te[4]);let se;for(se=0;se<ze.length;se+=1){const gE=sR(Te,ze,se);ve[se]?ve[se].p(gE,Xe):(ve[se]=NR(gE),ve[se].c(),ve[se].m(H,null))}for(;se<ve.length;se+=1)ve[se].d(1);ve.length=ze.length}(!pE||Xe&4)&&G!==(G=Te[2]+1+"")&&Le(F,G),(!pE||Xe&8)&&x!==(x=Math.min(Te[3],Te[7].length)+"")&&Le(J,x),Te[5]!=null?Fe?(Fe.p(Te,Xe),Xe&32&&m(Fe,1)):(Fe=lR(Te),Fe.c(),m(Fe,1),Fe.m(rE.parentNode,rE)):Fe&&(Ge(),y(Fe,1,1,()=>{Fe=null}),ge())},i(Te){pE||(m(Ve),m(Fe),pE=!0)},o(Te){y(Ve),y(Fe),pE=!1},d(Te){Te&&(Y(e),Y(T),Y(fE),Y(rE)),Ve&&Ve.d(Te),nE($e,Te),nE(ve,Te),Fe&&Fe.d(Te),Tt=!1,NE(rt)}}}function Ua(E,e,T){let{df:t}=e,{onTrain:r}=e,{removeTrainingData:R}=e,n=JSON.parse(t),s=n.length>0?Object.keys(n[0]):[],S=10,A=1,o=Math.ceil(n.length/S),i=(A-1)*S,_=A*S,c=n.slice(i,_);const P=()=>{A>1&&T(16,A--,A)},p=()=>{A<o&&T(16,A++,A)},C=()=>{T(16,A=1),T(15,S=n.length)};let L=null,I=!1;const u=()=>{T(6,I=!0)},H=()=>{T(6,I=!1)},b=(N,D)=>{T(5,L=N[D])},M=()=>{T(5,L=null)},O=()=>{L&&R(L)};return E.$$set=N=>{"df"in N&&T(14,t=N.df),"onTrain"in N&&T(0,r=N.onTrain),"removeTrainingData"in N&&T(1,R=N.removeTrainingData)},E.$$.update=()=>{E.$$.dirty&98304&&T(2,i=(A-1)*S),E.$$.dirty&98304&&T(3,_=A*S),E.$$.dirty&12&&T(4,c=n.slice(i,_)),E.$$.dirty&32768&&T(17,o=Math.ceil(n.length/S)),E.$$.dirty&196608&&console.log(A,o)},[r,R,i,_,c,L,I,n,s,P,p,C,u,H,t,S,A,o,b,M,O]}class ma extends ue{constructor(e){super(),Ce(this,e,Ua,Ma,_e,{df:14,onTrain:0,removeTrainingData:1})}}function ha(E){let e;return{c(){e=f("div"),e.innerHTML='<div class="flex flex-auto flex-col justify-center items-center p-4 md:p-5"><div class="flex justify-center"><div class="animate-spin inline-block w-6 h-6 border-[3px] border-current border-t-transparent text-blue-600 rounded-full" role="status" aria-label="loading"><span class="sr-only">Loading...</span></div></div></div>',a(e,"class","min-h-[15rem] flex flex-col bg-white border shadow-sm rounded-xl dark:bg-gray-800 dark:border-gray-700 dark:shadow-slate-700/[.7]")},m(T,t){V(T,e,t)},p:j,i:j,o:j,d(T){T&&Y(e)}}}function Ga(E){let e,T,t,r;const R=[Ha,ga],n=[];function s(S,A){return S[0].type==="df"?0:S[0].type==="error"?1:-1}return~(e=s(E))&&(T=n[e]=R[e](E)),{c(){T&&T.c(),t=je()},m(S,A){~e&&n[e].m(S,A),V(S,t,A),r=!0},p(S,A){let o=e;e=s(S),e===o?~e&&n[e].p(S,A):(T&&(Ge(),y(n[o],1,1,()=>{n[o]=null}),ge()),~e?(T=n[e],T?T.p(S,A):(T=n[e]=R[e](S),T.c()),m(T,1),T.m(t.parentNode,t)):T=null)},i(S){r||(m(T),r=!0)},o(S){y(T),r=!1},d(S){S&&Y(t),~e&&n[e].d(S)}}}function ga(E){let e,T;return e=new PT({props:{message:E[0].error}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&1&&(R.message=t[0].error),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Ha(E){let e,T;return e=new ma({props:{df:E[0].df,removeTrainingData:Un,onTrain:Hn}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&1&&(R.df=t[0].df),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ba(E){let e,T,t,r,R;const n=[Ga,ha],s=[];function S(A,o){return A[0]!==null?0:1}return t=S(E),r=s[t]=n[t](E),{c(){e=f("div"),T=f("div"),r.c(),a(T,"class","py-10 lg:py-14"),a(e,"class","relative h-screen w-full lg:pl-64")},m(A,o){V(A,e,o),l(e,T),s[t].m(T,null),R=!0},p(A,[o]){let i=t;t=S(A),t===i?s[t].p(A,o):(Ge(),y(s[i],1,1,()=>{s[i]=null}),ge(),r=s[t],r?r.p(A,o):(r=s[t]=n[t](A),r.c()),m(r,1),r.m(T,null))},i(A){R||(m(r),R=!0)},o(A){y(r),R=!1},d(A){A&&Y(e),s[t].d()}}}function ya(E,e,T){let t;return eE(E,gt,r=>T(0,t=r)),[t]}class Ba extends ue{constructor(e){super(),Ce(this,e,ya,ba,_e,{})}}function va(E){let e;return{c(){e=f("body"),e.innerHTML=`<div class="max-w-[50rem] flex flex-col mx-auto size-full"><header class="mb-auto flex justify-center z-50 w-full py-4"><nav class="px-4 sm:px-6 lg:px-8" aria-label="Global"><span class="flex-none text-xl font-semibold sm:text-3xl dark:text-white">Vanna.AI</span></nav></header> <div class="text-center py-10 px-4 sm:px-6 lg:px-8"><h1 class="block text-7xl font-bold text-gray-800 sm:text-9xl dark:text-white">No Training Data</h1> <h1 class="block text-2xl font-bold text-white">Did you read the docs?</h1> <p class="mt-3 text-gray-600 dark:text-gray-400">Oops, something went wrong.</p> <p class="text-gray-600 dark:text-gray-400">You need some training data before you can use Vanna</p> <div class="mt-5 flex flex-col justify-center items-center gap-2 sm:flex-row sm:gap-3"><a class="w-full sm:w-auto py-3 px-4 inline-flex justify-center items-center gap-x-2 text-sm font-semibold rounded-lg border border-transparent bg-blue-600 text-white hover:bg-blue-700 disabled:opacity-50 disabled:pointer-events-none dark:focus:outline-none dark:focus:ring-1 dark:focus:ring-gray-600" href="https://github.com/htmlstreamofficial/preline/tree/main/examples/html" target="_blank"><svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="16" height="16" fill="currentColor" viewBox="0 0 16 16"><path d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.012 8.012 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path></svg>
            Get the source code</a> <a class="w-full sm:w-auto py-3 px-4 inline-flex justify-center items-center gap-x-2 text-sm font-semibold rounded-lg border border-transparent text-blue-600 hover:text-blue-800 disabled:opacity-50 disabled:pointer-events-none dark:text-blue-500 dark:hover:text-blue-400 dark:focus:outline-none dark:focus:ring-1 dark:focus:ring-gray-600" href="../examples.html"><svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="m15 18-6-6 6-6"></path></svg>
            Back to examples</a></div></div> <footer class="mt-auto text-center py-5"><div class="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8"><p class="text-sm text-gray-500">© All Rights Reserved. 2022.</p></div></footer></div>`,a(e,"class","flex h-full")},m(T,t){V(T,e,t)},p:j,i:j,o:j,d(T){T&&Y(e)}}}class Fa extends ue{constructor(e){super(),Ce(this,e,null,va,_e,{})}}function Ya(E){let e,T,t;return{c(){e=f("div"),T=f("div"),t=f("div"),a(t,"class","mt-7 bg-white border border-gray-200 rounded-xl shadow-sm dark:bg-gray-800 dark:border-gray-700"),a(T,"class","w-full max-w-md mx-auto p-6"),a(e,"class","dark:bg-slate-900 bg-gray-100 flex h-screen items-center py-16")},m(r,R){V(r,e,R),l(e,T),l(T,t),t.innerHTML=E[0]},p(r,[R]){R&1&&(t.innerHTML=r[0])},i:j,o:j,d(r){r&&Y(e)}}}function Va(E,e,T){let t;return eE(E,bt,r=>T(0,t=r)),[t]}class Wa extends ue{constructor(e){super(),Ce(this,e,Va,Ya,_e,{})}}function wa(E){let e,T,t,r,R,n;return{c(){e=f("div"),T=f("div"),T.innerHTML='<h3 class="text-lg font-bold text-gray-800 dark:text-white">New</h3> <p class="mt-1 text-xs font-medium uppercase text-gray-500 dark:text-neutral-500">Create a New Function</p>',t=$(),r=f("button"),r.innerHTML=`Create
    <svg class="flex-shrink-0 size-4" xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="m9 18 6-6-6-6"></path></svg>`,a(r,"class","mt-3 inline-flex items-center gap-x-1 text-sm font-semibold rounded-lg border border-transparent text-green-600 hover:text-green-800 disabled:opacity-50 disabled:pointer-events-none dark:text-green-500 dark:hover:text-green-400"),a(e,"class","p-4 flex flex-col justify-between bg-white border border-t-4 border-t-green-600 shadow-sm rounded-xl dark:bg-neutral-900 dark:border-neutral-700 dark:border-t-green-500 dark:shadow-neutral-700/70")},m(s,S){V(s,e,S),l(e,T),l(e,t),l(e,r),R||(n=Ne(r,"click",E[0]),R=!0)},p:j,i:j,o:j,d(s){s&&Y(e),R=!1,n()}}}function $a(E){return[()=>{it()}]}class xa extends ue{constructor(e){super(),Ce(this,e,$a,wa,_e,{})}}function _R(E,e,T){const t=E.slice();return t[1]=e[T],t}function LR(E){let e,T;return e=new aE({props:{$$slots:{default:[Xa]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Xa(E){let e;return{c(){e=te("No functions found")},m(T,t){V(T,e,t)},d(T){T&&Y(e)}}}function CR(E){let e,T;return e=new en({props:{functionTemplate:E[1]}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},p(t,r){const R={};r&1&&(R.functionTemplate=t[1]),e.$set(R)},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ka(E){let e,T,t,r,R,n,s,S,A,o=E[0].length==0&&LR(E),i=De(E[0]),_=[];for(let P=0;P<i.length;P+=1)_[P]=CR(_R(E,i,P));const c=P=>y(_[P],1,1,()=>{_[P]=null});return S=new xa({}),{c(){e=f("div"),T=f("div"),t=f("h1"),t.textContent="Functions",r=$(),o&&o.c(),R=$(),n=f("div");for(let P=0;P<_.length;P+=1)_[P].c();s=$(),K(S.$$.fragment),a(t,"class","text-3xl font-bold text-gray-800 sm:text-4xl dark:text-white text-center"),a(n,"class","grid grid-cols-4 gap-4 p-4"),a(T,"class","py-10 lg:py-14"),a(e,"class","relative h-screen w-full lg:pl-64")},m(P,p){V(P,e,p),l(e,T),l(T,t),l(T,r),o&&o.m(T,null),l(T,R),l(T,n);for(let C=0;C<_.length;C+=1)_[C]&&_[C].m(n,null);l(n,s),X(S,n,null),A=!0},p(P,[p]){if(P[0].length==0?o?p&1&&m(o,1):(o=LR(P),o.c(),m(o,1),o.m(T,R)):o&&(Ge(),y(o,1,1,()=>{o=null}),ge()),p&1){i=De(P[0]);let C;for(C=0;C<i.length;C+=1){const L=_R(P,i,C);_[C]?(_[C].p(L,p),m(_[C],1)):(_[C]=CR(L),_[C].c(),m(_[C],1),_[C].m(n,s))}for(Ge(),C=i.length;C<_.length;C+=1)c(C);ge()}},i(P){if(!A){m(o);for(let p=0;p<i.length;p+=1)m(_[p]);m(S.$$.fragment,P),A=!0}},o(P){y(o),_=_.filter(Boolean);for(let p=0;p<_.length;p+=1)y(_[p]);y(S.$$.fragment,P),A=!1},d(P){P&&Y(e),o&&o.d(),nE(_,P),k(S)}}}function Ka(E,e,T){let t;return eE(E,MR,r=>T(0,t=r)),[t]}class Ja extends ue{constructor(e){super(),Ce(this,e,Ka,ka,_e,{})}}function uR(E){let e,T;return e=new Kn({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function qa(E){let e,T;return e=new aE({props:{$$slots:{default:[EI]},$$scope:{ctx:E}}}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Qa(E){let e,T;return e=new Ja({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function Za(E){let e,T;return e=new Wa({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function ja(E){let e,T;return e=new Fa({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function za(E){let e,T;return e=new Ba({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function eI(E){let e,T;return e=new la({}),{c(){K(e.$$.fragment)},m(t,r){X(e,t,r),T=!0},i(t){T||(m(e.$$.fragment,t),T=!0)},o(t){y(e.$$.fragment,t),T=!1},d(t){k(e,t)}}}function EI(E){let e;return{c(){e=te("Unknown page")},m(T,t){V(T,e,t)},d(T){T&&Y(e)}}}function tI(E){let e,T,t,r,R,n=E[0]!=="login"&&uR();const s=[eI,za,ja,Za,Qa,qa],S=[];function A(o,i){return o[0]==="chat"?0:o[0]==="training-data"?1:o[0]==="no-training-data"?2:o[0]==="login"?3:o[0]==="functions"?4:5}return t=A(E),r=S[t]=s[t](E),{c(){e=f("main"),n&&n.c(),T=$(),r.c()},m(o,i){V(o,e,i),n&&n.m(e,null),l(e,T),S[t].m(e,null),R=!0},p(o,[i]){o[0]!=="login"?n?i&1&&m(n,1):(n=uR(),n.c(),m(n,1),n.m(e,T)):n&&(Ge(),y(n,1,1,()=>{n=null}),ge());let _=t;t=A(o),t!==_&&(Ge(),y(S[_],1,1,()=>{S[_]=null}),ge(),r=S[t],r||(r=S[t]=s[t](o),r.c()),m(r,1),r.m(e,null))},i(o){R||(m(n),m(r),R=!0)},o(o){y(n),y(r),R=!1},d(o){o&&Y(e),n&&n.d(),S[t].d()}}}function TI(E,e,T){let t;return eE(E,DE,r=>T(0,t=r)),DR(async()=>{pn(),mR();const R=new URL(window.location.href).hash.slice(1);R==="training-data"?hR():R==="functions"?cT():it()}),[t]}class rI extends ue{constructor(e){super(),Ce(this,e,TI,tI,_e,{})}}new rI({target:document.getElementById("app")});
'''


================================================
FILE: src/vanna/legacy/flask/auth.py
================================================
from abc import ABC, abstractmethod

import flask


class AuthInterface(ABC):
    @abstractmethod
    def get_user(self, flask_request) -> any:
        pass

    @abstractmethod
    def is_logged_in(self, user: any) -> bool:
        pass

    @abstractmethod
    def override_config_for_user(self, user: any, config: dict) -> dict:
        pass

    @abstractmethod
    def login_form(self) -> str:
        pass

    @abstractmethod
    def login_handler(self, flask_request) -> str:
        pass

    @abstractmethod
    def callback_handler(self, flask_request) -> str:
        pass

    @abstractmethod
    def logout_handler(self, flask_request) -> str:
        pass


class NoAuth(AuthInterface):
    def get_user(self, flask_request) -> any:
        return {}

    def is_logged_in(self, user: any) -> bool:
        return True

    def override_config_for_user(self, user: any, config: dict) -> dict:
        return config

    def login_form(self) -> str:
        return ""

    def login_handler(self, flask_request) -> str:
        return "No login required"

    def callback_handler(self, flask_request) -> str:
        return "No login required"

    def logout_handler(self, flask_request) -> str:
        return "No login required"


================================================
FILE: src/vanna/legacy/google/__init__.py
================================================
from .bigquery_vector import BigQuery_VectorStore
from .gemini_chat import GoogleGeminiChat


================================================
FILE: src/vanna/legacy/google/bigquery_vector.py
================================================
import datetime
import os
import uuid
from typing import List, Optional
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel

import pandas as pd
from google.cloud import bigquery

from ..base import VannaBase


class BigQuery_VectorStore(VannaBase):
    def __init__(self, config: dict, **kwargs):
        self.config = config

        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
        self.n_results_documentation = config.get(
            "n_results_documentation", config.get("n_results", 10)
        )
        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))

        if "api_key" in config or os.getenv("GOOGLE_API_KEY"):
            """
            If Google api_key is provided through config
            or set as an environment variable, assign it.
            """
            print("Configuring genai")
            self.type = "GEMINI"
            import google.generativeai as genai

            genai.configure(api_key=config["api_key"])

            self.genai = genai
        else:
            self.type = "VERTEX_AI"
            # Authenticate using VertexAI

        if self.config.get("project_id"):
            self.project_id = self.config.get("project_id")
        else:
            self.project_id = os.getenv("GOOGLE_CLOUD_PROJECT")

        if self.project_id is None:
            raise ValueError("Project ID is not set")

        self.conn = bigquery.Client(project=self.project_id)

        dataset_name = self.config.get("bigquery_dataset_name", "vanna_managed")
        self.dataset_id = f"{self.project_id}.{dataset_name}"
        dataset = bigquery.Dataset(self.dataset_id)

        try:
            self.conn.get_dataset(self.dataset_id)  # Make an API request.
            print(f"Dataset {self.dataset_id} already exists")
        except Exception:
            # Dataset does not exist, create it
            dataset.location = "US"
            self.conn.create_dataset(dataset, timeout=30)  # Make an API request.
            print(f"Created dataset {self.dataset_id}")

        # Create a table called training_data in the dataset that contains the columns:
        # id, training_data_type, question, content, embedding, created_at

        self.table_id = f"{self.dataset_id}.training_data"
        schema = [
            bigquery.SchemaField("id", "STRING", mode="REQUIRED"),
            bigquery.SchemaField("training_data_type", "STRING", mode="REQUIRED"),
            bigquery.SchemaField("question", "STRING", mode="REQUIRED"),
            bigquery.SchemaField("content", "STRING", mode="REQUIRED"),
            bigquery.SchemaField("embedding", "FLOAT64", mode="REPEATED"),
            bigquery.SchemaField("created_at", "TIMESTAMP", mode="REQUIRED"),
        ]

        table = bigquery.Table(self.table_id, schema=schema)

        try:
            self.conn.get_table(self.table_id)  # Make an API request.
            print(f"Table {self.table_id} already exists")
        except Exception:
            # Table does not exist, create it
            self.conn.create_table(table, timeout=30)  # Make an API request.
            print(f"Created table {self.table_id}")

        # Create VECTOR INDEX IF NOT EXISTS
        # TODO: This requires 5000 rows before it can be created
        # vector_index_query = f"""
        # CREATE VECTOR INDEX IF NOT EXISTS my_index
        # ON `{self.table_id}`(embedding)
        # OPTIONS(
        #     distance_type='COSINE',
        #     index_type='IVF',
        #     ivf_options='{{"num_lists": 1000}}'
        # )
        # """

        # try:
        #     self.conn.query(vector_index_query).result()  # Make an API request.
        #     print(f"Vector index on {self.table_id} created or already exists")
        # except Exception as e:
        #     print(f"Failed to create vector index: {e}")

    def store_training_data(
        self,
        training_data_type: str,
        question: str,
        content: str,
        embedding: List[float],
        **kwargs,
    ) -> str:
        id = str(uuid.uuid4())
        created_at = datetime.datetime.now()
        self.conn.insert_rows_json(
            self.table_id,
            [
                {
                    "id": id,
                    "training_data_type": training_data_type,
                    "question": question,
                    "content": content,
                    "embedding": embedding,
                    "created_at": created_at.isoformat(),
                }
            ],
        )

        return id

    def fetch_similar_training_data(
        self, training_data_type: str, question: str, n_results, **kwargs
    ) -> pd.DataFrame:
        question_embedding = self.generate_question_embedding(question)

        query = f"""
        SELECT
            base.id as id,
            base.question as question,
            base.training_data_type as training_data_type,
            base.content as content,
            distance
        FROM
            VECTOR_SEARCH(
                TABLE `{self.table_id}`,
                'embedding',
                (SELECT * FROM UNNEST([STRUCT({question_embedding})])),
                top_k => 5,
                distance_type => 'COSINE',
                options => '{{"use_brute_force":true}}'
            )
        WHERE
            base.training_data_type = '{training_data_type}'
        """

        results = self.conn.query(query).result().to_dataframe()
        return results

    def get_embeddings(self, data: str, task: str) -> List[float]:
        embeddings = None

        if self.type == "VERTEX_AI":
            input = [TextEmbeddingInput(data, task)]
            model = TextEmbeddingModel.from_pretrained("text-embedding-004")

            result = model.get_embeddings(input)

            if len(result) > 0:
                embeddings = result[0].values
        else:
            # Use Gemini Consumer API
            result = self.genai.embed_content(
                model="models/text-embedding-004", content=data, task_type=task
            )

            if "embedding" in result:
                embeddings = result["embedding"]

        return embeddings

    def generate_question_embedding(self, data: str, **kwargs) -> List[float]:
        result = self.get_embeddings(data, "RETRIEVAL_QUERY")

        if result is not None:
            return result
        else:
            raise ValueError("No embeddings returned")

    def generate_storage_embedding(self, data: str, **kwargs) -> List[float]:
        result = self.get_embeddings(data, "RETRIEVAL_DOCUMENT")

        if result is not None:
            return result
        else:
            raise ValueError("No embeddings returned")

        # task = "RETRIEVAL_DOCUMENT"
        # inputs = [TextEmbeddingInput(data, task)]
        # embeddings = self.vertex_embedding_model.get_embeddings(inputs)

        # if len(embeddings) == 0:
        #     raise ValueError("No embeddings returned")

        # return embeddings[0].values

        return result

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        return self.generate_storage_embedding(data, **kwargs)

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        df = self.fetch_similar_training_data(
            training_data_type="sql", question=question, n_results=self.n_results_sql
        )

        # Return a list of dictionaries with only question, sql fields. The content field needs to be renamed to sql
        return df.rename(columns={"content": "sql"})[["question", "sql"]].to_dict(
            orient="records"
        )

    def get_related_ddl(self, question: str, **kwargs) -> list:
        df = self.fetch_similar_training_data(
            training_data_type="ddl", question=question, n_results=self.n_results_ddl
        )

        # Return a list of strings of the content
        return df["content"].tolist()

    def get_related_documentation(self, question: str, **kwargs) -> list:
        df = self.fetch_similar_training_data(
            training_data_type="documentation",
            question=question,
            n_results=self.n_results_documentation,
        )

        # Return a list of strings of the content
        return df["content"].tolist()

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        doc = {"question": question, "sql": sql}

        embedding = self.generate_embedding(str(doc))

        return self.store_training_data(
            training_data_type="sql",
            question=question,
            content=sql,
            embedding=embedding,
        )

    def add_ddl(self, ddl: str, **kwargs) -> str:
        embedding = self.generate_embedding(ddl)

        return self.store_training_data(
            training_data_type="ddl", question="", content=ddl, embedding=embedding
        )

    def add_documentation(self, documentation: str, **kwargs) -> str:
        embedding = self.generate_embedding(documentation)

        return self.store_training_data(
            training_data_type="documentation",
            question="",
            content=documentation,
            embedding=embedding,
        )

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        query = (
            f"SELECT id, training_data_type, question, content FROM `{self.table_id}`"
        )

        return self.conn.query(query).result().to_dataframe()

    def remove_training_data(self, id: str, **kwargs) -> bool:
        query = f"DELETE FROM `{self.table_id}` WHERE id = '{id}'"

        try:
            self.conn.query(query).result()
            return True

        except Exception as e:
            print(f"Failed to remove training data: {e}")
            return False


================================================
FILE: src/vanna/legacy/google/gemini_chat.py
================================================
import os

from ..base import VannaBase


class GoogleGeminiChat(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        # default temperature - can be overrided using config
        self.temperature = 0.7

        if "temperature" in config:
            self.temperature = config["temperature"]

        if "model_name" in config:
            model_name = config["model_name"]
        else:
            model_name = "gemini-1.5-pro"

        self.google_api_key = None

        if "api_key" in config or os.getenv("GOOGLE_API_KEY"):
            """
            If Google api_key is provided through config
            or set as an environment variable, assign it.
            """
            import google.generativeai as genai

            genai.configure(api_key=config["api_key"])
            self.chat_model = genai.GenerativeModel(model_name)
        else:
            # Authenticate using VertexAI
            import google.auth
            import vertexai
            from vertexai.generative_models import GenerativeModel

            json_file_path = config.get(
                "google_credentials"
            )  # Assuming the JSON file path is provided in the config

            if not json_file_path or not os.path.exists(json_file_path):
                raise FileNotFoundError(
                    f"JSON credentials file not found at: {json_file_path}"
                )

            try:
                # Validate and set the JSON file path for GOOGLE_APPLICATION_CREDENTIALS
                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = json_file_path

                # Initialize VertexAI with the credentials
                credentials, _ = google.auth.default()
                vertexai.init(credentials=credentials)
                self.chat_model = GenerativeModel(model_name)
            except google.auth.exceptions.DefaultCredentialsError as e:
                raise RuntimeError(f"Default credentials error: {e}")
            except google.auth.exceptions.TransportError as e:
                raise RuntimeError(f"Transport error during authentication: {e}")
            except Exception as e:
                raise RuntimeError(f"Failed to authenticate using JSON file: {e}")

    def system_message(self, message: str) -> any:
        return message

    def user_message(self, message: str) -> any:
        return message

    def assistant_message(self, message: str) -> any:
        return message

    def submit_prompt(self, prompt, **kwargs) -> str:
        response = self.chat_model.generate_content(
            prompt,
            generation_config={
                "temperature": self.temperature,
            },
        )
        return response.text


================================================
FILE: src/vanna/legacy/hf/__init__.py
================================================
from .hf import Hf


================================================
FILE: src/vanna/legacy/hf/hf.py
================================================
import re
from transformers import AutoTokenizer, AutoModelForCausalLM

from ..base import VannaBase


class Hf(VannaBase):
    def __init__(self, config=None):
        model_name_or_path = self.config.get(
            "model_name_or_path", None
        )  # e.g. meta-llama/Meta-Llama-3-8B-Instruct or local path to the model checkpoint files
        # list of quantization methods supported by transformers package: https://huggingface.co/docs/transformers/main/en/quantization/overview
        quantization_config = self.config.get("quantization_config", None)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name_or_path,
            quantization_config=quantization_config,
            device_map="auto",
        )

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def extract_sql_query(self, text):
        """
        Extracts the first SQL statement after the word 'select', ignoring case,
        matches until the first semicolon, three backticks, or the end of the string,
        and removes three backticks if they exist in the extracted string.

        Args:
        - text (str): The string to search within for an SQL statement.

        Returns:
        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
        """
        # Regular expression to find 'select' (ignoring case) and capture until ';', '```', or end of string
        pattern = re.compile(r"select.*?(?:;|```|$)", re.IGNORECASE | re.DOTALL)

        match = pattern.search(text)
        if match:
            # Remove three backticks from the matched string if they exist
            return match.group(0).replace("```", "")
        else:
            return text

    def generate_sql(self, question: str, **kwargs) -> str:
        # Use the super generate_sql
        sql = super().generate_sql(question, **kwargs)

        # Replace "\_" with "_"
        sql = sql.replace("\\_", "_")

        sql = sql.replace("\\", "")

        return self.extract_sql_query(sql)

    def submit_prompt(self, prompt, **kwargs) -> str:
        input_ids = self.tokenizer.apply_chat_template(
            prompt, add_generation_prompt=True, return_tensors="pt"
        ).to(self.model.device)

        outputs = self.model.generate(
            input_ids,
            max_new_tokens=512,
            eos_token_id=self.tokenizer.eos_token_id,
            do_sample=True,
            temperature=1,
            top_p=0.9,
        )
        response = outputs[0][input_ids.shape[-1] :]
        response = self.tokenizer.decode(response, skip_special_tokens=True)
        self.log(response)

        return response


================================================
FILE: src/vanna/legacy/local.py
================================================
from .chromadb.chromadb_vector import ChromaDB_VectorStore
from .openai.openai_chat import OpenAI_Chat


class LocalContext_OpenAI(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)


================================================
FILE: src/vanna/legacy/marqo/__init__.py
================================================
from .marqo import Marqo_VectorStore


================================================
FILE: src/vanna/legacy/marqo/marqo.py
================================================
import uuid

import marqo
import pandas as pd

from ..base import VannaBase


class Marqo_VectorStore(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        if config is not None and "marqo_url" in config:
            marqo_url = config["marqo_url"]
        else:
            marqo_url = "http://localhost:8882"

        if config is not None and "marqo_model" in config:
            marqo_model = config["marqo_model"]
        else:
            marqo_model = "hf/all_datasets_v4_MiniLM-L6"

        self.mq = marqo.Client(url=marqo_url)

        for index in ["vanna-sql", "vanna-ddl", "vanna-doc"]:
            try:
                self.mq.create_index(index, model=marqo_model)
            except Exception as e:
                print(e)
                print(f"Marqo index {index} already exists")
                pass

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        # Marqo doesn't need to generate embeddings
        pass

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        id = str(uuid.uuid4()) + "-sql"
        question_sql_dict = {
            "question": question,
            "sql": sql,
            "_id": id,
        }

        self.mq.index("vanna-sql").add_documents(
            [question_sql_dict],
            tensor_fields=["question", "sql"],
        )

        return id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        id = str(uuid.uuid4()) + "-ddl"
        ddl_dict = {
            "ddl": ddl,
            "_id": id,
        }

        self.mq.index("vanna-ddl").add_documents(
            [ddl_dict],
            tensor_fields=["ddl"],
        )
        return id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        id = str(uuid.uuid4()) + "-doc"
        doc_dict = {
            "doc": documentation,
            "_id": id,
        }

        self.mq.index("vanna-doc").add_documents(
            [doc_dict],
            tensor_fields=["doc"],
        )
        return id

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        data = []

        for hit in self.mq.index("vanna-doc").search("", limit=1000)["hits"]:
            data.append(
                {
                    "id": hit["_id"],
                    "training_data_type": "documentation",
                    "question": "",
                    "content": hit["doc"],
                }
            )

        for hit in self.mq.index("vanna-ddl").search("", limit=1000)["hits"]:
            data.append(
                {
                    "id": hit["_id"],
                    "training_data_type": "ddl",
                    "question": "",
                    "content": hit["ddl"],
                }
            )

        for hit in self.mq.index("vanna-sql").search("", limit=1000)["hits"]:
            data.append(
                {
                    "id": hit["_id"],
                    "training_data_type": "sql",
                    "question": hit["question"],
                    "content": hit["sql"],
                }
            )

        df = pd.DataFrame(data)

        return df

    def remove_training_data(self, id: str, **kwargs) -> bool:
        if id.endswith("-sql"):
            self.mq.index("vanna-sql").delete_documents(ids=[id])
            return True
        elif id.endswith("-ddl"):
            self.mq.index("vanna-ddl").delete_documents(ids=[id])
            return True
        elif id.endswith("-doc"):
            self.mq.index("vanna-doc").delete_documents(ids=[id])
            return True
        else:
            return False

    # Static method to extract the documents from the results of a query
    @staticmethod
    def _extract_documents(data) -> list:
        # Check if 'hits' key is in the dictionary and if it's a list
        if "hits" in data and isinstance(data["hits"], list):
            # Iterate over each item in 'hits'

            if len(data["hits"]) == 0:
                return []

            # If there is a "doc" key, return the value of that key
            if "doc" in data["hits"][0]:
                return [hit["doc"] for hit in data["hits"]]

            # If there is a "ddl" key, return the value of that key
            if "ddl" in data["hits"][0]:
                return [hit["ddl"] for hit in data["hits"]]

            # Otherwise return the entire hit
            return [
                {key: value for key, value in hit.items() if not key.startswith("_")}
                for hit in data["hits"]
            ]
        else:
            # Return an empty list if 'hits' is not found or not a list
            return []

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        return Marqo_VectorStore._extract_documents(
            self.mq.index("vanna-sql").search(question)
        )

    def get_related_ddl(self, question: str, **kwargs) -> list:
        return Marqo_VectorStore._extract_documents(
            self.mq.index("vanna-ddl").search(question)
        )

    def get_related_documentation(self, question: str, **kwargs) -> list:
        return Marqo_VectorStore._extract_documents(
            self.mq.index("vanna-doc").search(question)
        )


================================================
FILE: src/vanna/legacy/milvus/__init__.py
================================================
from .milvus_vector import Milvus_VectorStore


================================================
FILE: src/vanna/legacy/milvus/milvus_vector.py
================================================
import uuid
from typing import List

import pandas as pd
from pymilvus import DataType, MilvusClient, model

from ..base import VannaBase

# Setting the URI as a local file, e.g.`./milvus.db`,
# is the most convenient method, as it automatically utilizes Milvus Lite
# to store all data in this file.
#
# If you have large scale of data such as more than a million docs, we
# recommend setting up a more performant Milvus server on docker or kubernetes.
# When using this setup, please use the server URI,
# e.g.`http://localhost:19530`, as your URI.

DEFAULT_MILVUS_URI = "./milvus.db"
# DEFAULT_MILVUS_URI = "http://localhost:19530"

MAX_LIMIT_SIZE = 10_000


class Milvus_VectorStore(VannaBase):
    """
    Vectorstore implementation using Milvus - https://milvus.io/docs/quickstart.md

    Args:
        - config (dict, optional): Dictionary of `Milvus_VectorStore config` options. Defaults to `None`.
            - milvus_client: A `pymilvus.MilvusClient` instance.
            - embedding_function:
                A `milvus_model.base.BaseEmbeddingFunction` instance. Defaults to `DefaultEmbeddingFunction()`.
                For more models, please refer to:
                https://milvus.io/docs/embeddings.md
    """

    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        if "milvus_client" in config:
            self.milvus_client = config["milvus_client"]
        else:
            self.milvus_client = MilvusClient(uri=DEFAULT_MILVUS_URI)

        if "embedding_function" in config:
            self.embedding_function = config.get("embedding_function")
        else:
            self.embedding_function = model.DefaultEmbeddingFunction()
        self._embedding_dim = self.embedding_function.encode_documents(["foo"])[
            0
        ].shape[0]
        self._create_collections()
        self.n_results = config.get("n_results", 10)

    def _create_collections(self):
        self._create_sql_collection("vannasql")
        self._create_ddl_collection("vannaddl")
        self._create_doc_collection("vannadoc")

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        return self.embedding_function.encode_documents(data).tolist()

    def _create_sql_collection(self, name: str):
        if not self.milvus_client.has_collection(collection_name=name):
            vannasql_schema = MilvusClient.create_schema(
                auto_id=False,
                enable_dynamic_field=False,
            )
            vannasql_schema.add_field(
                field_name="id",
                datatype=DataType.VARCHAR,
                max_length=65535,
                is_primary=True,
            )
            vannasql_schema.add_field(
                field_name="text", datatype=DataType.VARCHAR, max_length=65535
            )
            vannasql_schema.add_field(
                field_name="sql", datatype=DataType.VARCHAR, max_length=65535
            )
            vannasql_schema.add_field(
                field_name="vector",
                datatype=DataType.FLOAT_VECTOR,
                dim=self._embedding_dim,
            )

            vannasql_index_params = self.milvus_client.prepare_index_params()
            vannasql_index_params.add_index(
                field_name="vector",
                index_name="vector",
                index_type="AUTOINDEX",
                metric_type="L2",
            )
            self.milvus_client.create_collection(
                collection_name=name,
                schema=vannasql_schema,
                index_params=vannasql_index_params,
                consistency_level="Strong",
            )

    def _create_ddl_collection(self, name: str):
        if not self.milvus_client.has_collection(collection_name=name):
            vannaddl_schema = MilvusClient.create_schema(
                auto_id=False,
                enable_dynamic_field=False,
            )
            vannaddl_schema.add_field(
                field_name="id",
                datatype=DataType.VARCHAR,
                max_length=65535,
                is_primary=True,
            )
            vannaddl_schema.add_field(
                field_name="ddl", datatype=DataType.VARCHAR, max_length=65535
            )
            vannaddl_schema.add_field(
                field_name="vector",
                datatype=DataType.FLOAT_VECTOR,
                dim=self._embedding_dim,
            )

            vannaddl_index_params = self.milvus_client.prepare_index_params()
            vannaddl_index_params.add_index(
                field_name="vector",
                index_name="vector",
                index_type="AUTOINDEX",
                metric_type="L2",
            )
            self.milvus_client.create_collection(
                collection_name=name,
                schema=vannaddl_schema,
                index_params=vannaddl_index_params,
                consistency_level="Strong",
            )

    def _create_doc_collection(self, name: str):
        if not self.milvus_client.has_collection(collection_name=name):
            vannadoc_schema = MilvusClient.create_schema(
                auto_id=False,
                enable_dynamic_field=False,
            )
            vannadoc_schema.add_field(
                field_name="id",
                datatype=DataType.VARCHAR,
                max_length=65535,
                is_primary=True,
            )
            vannadoc_schema.add_field(
                field_name="doc", datatype=DataType.VARCHAR, max_length=65535
            )
            vannadoc_schema.add_field(
                field_name="vector",
                datatype=DataType.FLOAT_VECTOR,
                dim=self._embedding_dim,
            )

            vannadoc_index_params = self.milvus_client.prepare_index_params()
            vannadoc_index_params.add_index(
                field_name="vector",
                index_name="vector",
                index_type="AUTOINDEX",
                metric_type="L2",
            )
            self.milvus_client.create_collection(
                collection_name=name,
                schema=vannadoc_schema,
                index_params=vannadoc_index_params,
                consistency_level="Strong",
            )

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        if len(question) == 0 or len(sql) == 0:
            raise Exception("pair of question and sql can not be null")
        _id = str(uuid.uuid4()) + "-sql"
        embedding = self.embedding_function.encode_documents([question])[0]
        self.milvus_client.insert(
            collection_name="vannasql",
            data={"id": _id, "text": question, "sql": sql, "vector": embedding},
        )
        return _id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        if len(ddl) == 0:
            raise Exception("ddl can not be null")
        _id = str(uuid.uuid4()) + "-ddl"
        embedding = self.embedding_function.encode_documents([ddl])[0]
        self.milvus_client.insert(
            collection_name="vannaddl",
            data={"id": _id, "ddl": ddl, "vector": embedding},
        )
        return _id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        if len(documentation) == 0:
            raise Exception("documentation can not be null")
        _id = str(uuid.uuid4()) + "-doc"
        embedding = self.embedding_function.encode_documents([documentation])[0]
        self.milvus_client.insert(
            collection_name="vannadoc",
            data={"id": _id, "doc": documentation, "vector": embedding},
        )
        return _id

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        sql_data = self.milvus_client.query(
            collection_name="vannasql",
            output_fields=["*"],
            limit=MAX_LIMIT_SIZE,
        )
        df = pd.DataFrame()
        df_sql = pd.DataFrame(
            {
                "id": [doc["id"] for doc in sql_data],
                "question": [doc["text"] for doc in sql_data],
                "content": [doc["sql"] for doc in sql_data],
            }
        )
        df = pd.concat([df, df_sql])

        ddl_data = self.milvus_client.query(
            collection_name="vannaddl",
            output_fields=["*"],
            limit=MAX_LIMIT_SIZE,
        )

        df_ddl = pd.DataFrame(
            {
                "id": [doc["id"] for doc in ddl_data],
                "question": [None for doc in ddl_data],
                "content": [doc["ddl"] for doc in ddl_data],
            }
        )
        df = pd.concat([df, df_ddl])

        doc_data = self.milvus_client.query(
            collection_name="vannadoc",
            output_fields=["*"],
            limit=MAX_LIMIT_SIZE,
        )

        df_doc = pd.DataFrame(
            {
                "id": [doc["id"] for doc in doc_data],
                "question": [None for doc in doc_data],
                "content": [doc["doc"] for doc in doc_data],
            }
        )
        df = pd.concat([df, df_doc])
        return df

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        search_params = {
            "metric_type": "L2",
            "params": {"nprobe": 128},
        }
        embeddings = self.embedding_function.encode_queries([question])
        res = self.milvus_client.search(
            collection_name="vannasql",
            anns_field="vector",
            data=embeddings,
            limit=self.n_results,
            output_fields=["text", "sql"],
            search_params=search_params,
        )
        res = res[0]

        list_sql = []
        for doc in res:
            dict = {}
            dict["question"] = doc["entity"]["text"]
            dict["sql"] = doc["entity"]["sql"]
            list_sql.append(dict)
        return list_sql

    def get_related_ddl(self, question: str, **kwargs) -> list:
        search_params = {
            "metric_type": "L2",
            "params": {"nprobe": 128},
        }
        embeddings = self.embedding_function.encode_queries([question])
        res = self.milvus_client.search(
            collection_name="vannaddl",
            anns_field="vector",
            data=embeddings,
            limit=self.n_results,
            output_fields=["ddl"],
            search_params=search_params,
        )
        res = res[0]

        list_ddl = []
        for doc in res:
            list_ddl.append(doc["entity"]["ddl"])
        return list_ddl

    def get_related_documentation(self, question: str, **kwargs) -> list:
        search_params = {
            "metric_type": "L2",
            "params": {"nprobe": 128},
        }
        embeddings = self.embedding_function.encode_queries([question])
        res = self.milvus_client.search(
            collection_name="vannadoc",
            anns_field="vector",
            data=embeddings,
            limit=self.n_results,
            output_fields=["doc"],
            search_params=search_params,
        )
        res = res[0]

        list_doc = []
        for doc in res:
            list_doc.append(doc["entity"]["doc"])
        return list_doc

    def remove_training_data(self, id: str, **kwargs) -> bool:
        if id.endswith("-sql"):
            self.milvus_client.delete(collection_name="vannasql", ids=[id])
            return True
        elif id.endswith("-ddl"):
            self.milvus_client.delete(collection_name="vannaddl", ids=[id])
            return True
        elif id.endswith("-doc"):
            self.milvus_client.delete(collection_name="vannadoc", ids=[id])
            return True
        else:
            return False


================================================
FILE: src/vanna/legacy/mistral/__init__.py
================================================
from .mistral import Mistral


================================================
FILE: src/vanna/legacy/mistral/mistral.py
================================================
import os

from mistralai import Mistral as MistralClient
from mistralai import UserMessage

from ..base import VannaBase


class Mistral(VannaBase):
    def __init__(self, config=None):
        if config is None:
            raise ValueError(
                "For Mistral, config must be provided with an api_key and model"
            )

        if "api_key" not in config:
            raise ValueError("config must contain a Mistral api_key")

        if "model" not in config:
            raise ValueError("config must contain a Mistral model")

        api_key = config["api_key"]
        model = config["model"]
        self.client = MistralClient(api_key=api_key)
        self.model = model

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def generate_sql(self, question: str, **kwargs) -> str:
        # Use the super generate_sql
        sql = super().generate_sql(question, **kwargs)

        # Replace "\_" with "_"
        sql = sql.replace("\\_", "_")

        return sql

    def submit_prompt(self, prompt, **kwargs) -> str:
        chat_response = self.client.chat.complete(
            model=self.model,
            messages=prompt,
        )

        return chat_response.choices[0].message.content


================================================
FILE: src/vanna/legacy/mock/__init__.py
================================================
from .embedding import MockEmbedding
from .llm import MockLLM
from .vectordb import MockVectorDB


================================================
FILE: src/vanna/legacy/mock/embedding.py
================================================
from typing import List

from ..base import VannaBase


class MockEmbedding(VannaBase):
    def __init__(self, config=None):
        pass

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        return [1.0, 2.0, 3.0, 4.0, 5.0]


================================================
FILE: src/vanna/legacy/mock/llm.py
================================================
from ..base import VannaBase


class MockLLM(VannaBase):
    def __init__(self, config=None):
        pass

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        return "Mock LLM response"


================================================
FILE: src/vanna/legacy/mock/vectordb.py
================================================
import pandas as pd

from ..base import VannaBase


class MockVectorDB(VannaBase):
    def __init__(self, config=None):
        pass

    def _get_id(self, value: str, **kwargs) -> str:
        # Hash the value and return the ID
        return str(hash(value))

    def add_ddl(self, ddl: str, **kwargs) -> str:
        return self._get_id(ddl)

    def add_documentation(self, doc: str, **kwargs) -> str:
        return self._get_id(doc)

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        return self._get_id(question)

    def get_related_ddl(self, question: str, **kwargs) -> list:
        return []

    def get_related_documentation(self, question: str, **kwargs) -> list:
        return []

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        return []

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        return pd.DataFrame(
            {
                "id": {
                    0: "19546-ddl",
                    1: "91597-sql",
                    2: "133976-sql",
                    3: "59851-doc",
                    4: "73046-sql",
                },
                "training_data_type": {
                    0: "ddl",
                    1: "sql",
                    2: "sql",
                    3: "documentation",
                    4: "sql",
                },
                "question": {
                    0: None,
                    1: "What are the top selling genres?",
                    2: "What are the low 7 artists by sales?",
                    3: None,
                    4: "What is the total sales for each customer?",
                },
                "content": {
                    0: "CREATE TABLE [Invoice]\n(\n    [InvoiceId] INTEGER  NOT NULL,\n    [CustomerId] INTEGER  NOT NULL,\n    [InvoiceDate] DATETIME  NOT NULL,\n    [BillingAddress] NVARCHAR(70),\n    [BillingCity] NVARCHAR(40),\n    [BillingState] NVARCHAR(40),\n    [BillingCountry] NVARCHAR(40),\n    [BillingPostalCode] NVARCHAR(10),\n    [Total] NUMERIC(10,2)  NOT NULL,\n    CONSTRAINT [PK_Invoice] PRIMARY KEY  ([InvoiceId]),\n    FOREIGN KEY ([CustomerId]) REFERENCES [Customer] ([CustomerId]) \n\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n)",
                    1: "SELECT g.Name AS Genre, SUM(il.Quantity) AS TotalSales\nFROM Genre g\nJOIN Track t ON g.GenreId = t.GenreId\nJOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY g.GenreId, g.Name\nORDER BY TotalSales DESC;",
                    2: "SELECT a.ArtistId, a.Name, SUM(il.Quantity) AS TotalSales\nFROM Artist a\nINNER JOIN Album al ON a.ArtistId = al.ArtistId\nINNER JOIN Track t ON al.AlbumId = t.AlbumId\nINNER JOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY a.ArtistId, a.Name\nORDER BY TotalSales ASC\nLIMIT 7;",
                    3: "This is a SQLite database. For dates rememeber to use SQLite syntax.",
                    4: "SELECT c.CustomerId, c.FirstName, c.LastName, SUM(i.Total) AS TotalSales\nFROM Customer c\nJOIN Invoice i ON c.CustomerId = i.CustomerId\nGROUP BY c.CustomerId, c.FirstName, c.LastName;",
                },
            }
        )

    def remove_training_data(id: str, **kwargs) -> bool:
        return True


================================================
FILE: src/vanna/legacy/ollama/__init__.py
================================================
from .ollama import Ollama


================================================
FILE: src/vanna/legacy/ollama/ollama.py
================================================
import json
import re

from httpx import Timeout

from ..base import VannaBase
from ..exceptions import DependencyError


class Ollama(VannaBase):
    def __init__(self, config=None):
        try:
            ollama = __import__("ollama")
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method, run command:"
                " \npip install ollama"
            )

        if not config:
            raise ValueError("config must contain at least Ollama model")
        if "model" not in config.keys():
            raise ValueError("config must contain at least Ollama model")
        self.host = config.get("ollama_host", "http://localhost:11434")
        self.model = config["model"]
        if ":" not in self.model:
            self.model += ":latest"

        self.ollama_timeout = config.get("ollama_timeout", 240.0)

        self.ollama_client = ollama.Client(
            self.host, timeout=Timeout(self.ollama_timeout)
        )
        self.keep_alive = config.get("keep_alive", None)
        self.ollama_options = config.get("options", {})
        self.num_ctx = self.ollama_options.get("num_ctx", 2048)
        self.__pull_model_if_ne(self.ollama_client, self.model)

    @staticmethod
    def __pull_model_if_ne(ollama_client, model):
        model_response = ollama_client.list()
        model_lists = [
            model_element["model"] for model_element in model_response.get("models", [])
        ]
        if model not in model_lists:
            ollama_client.pull(model)

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def extract_sql(self, llm_response):
        """
        Extracts the first SQL statement after the word 'select', ignoring case,
        matches until the first semicolon, three backticks, or the end of the string,
        and removes three backticks if they exist in the extracted string.

        Args:
        - llm_response (str): The string to search within for an SQL statement.

        Returns:
        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
        """
        # Remove ollama-generated extra characters
        llm_response = llm_response.replace("\\_", "_")
        llm_response = llm_response.replace("\\", "")

        # Regular expression to find ```sql' and capture until '```'
        sql = re.search(r"```sql\n((.|\n)*?)(?=;|\[|```)", llm_response, re.DOTALL)
        # Regular expression to find 'select, with (ignoring case) and capture until ';', [ (this happens in case of mistral) or end of string
        select_with = re.search(
            r"(select|(with.*?as \())(.*?)(?=;|\[|```)",
            llm_response,
            re.IGNORECASE | re.DOTALL,
        )
        if sql:
            self.log(f"Output from LLM: {llm_response} \nExtracted SQL: {sql.group(1)}")
            return sql.group(1).replace("```", "")
        elif select_with:
            self.log(
                f"Output from LLM: {llm_response} \nExtracted SQL: {select_with.group(0)}"
            )
            return select_with.group(0)
        else:
            return llm_response

    def submit_prompt(self, prompt, **kwargs) -> str:
        self.log(
            f"Ollama parameters:\n"
            f"model={self.model},\n"
            f"options={self.ollama_options},\n"
            f"keep_alive={self.keep_alive}"
        )
        self.log(f"Prompt Content:\n{json.dumps(prompt, ensure_ascii=False)}")
        response_dict = self.ollama_client.chat(
            model=self.model,
            messages=prompt,
            stream=False,
            options=self.ollama_options,
            keep_alive=self.keep_alive,
        )

        self.log(f"Ollama Response:\n{str(response_dict)}")

        return response_dict["message"]["content"]


================================================
FILE: src/vanna/legacy/openai/__init__.py
================================================
from .openai_chat import OpenAI_Chat
from .openai_embeddings import OpenAI_Embeddings


================================================
FILE: src/vanna/legacy/openai/openai_chat.py
================================================
import os

from openai import OpenAI

from ..base import VannaBase


class OpenAI_Chat(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        # default parameters - can be overrided using config
        self.temperature = 0.7

        if "temperature" in config:
            self.temperature = config["temperature"]

        if "api_type" in config:
            raise Exception(
                "Passing api_type is now deprecated. Please pass an OpenAI client instead."
            )

        if "api_base" in config:
            raise Exception(
                "Passing api_base is now deprecated. Please pass an OpenAI client instead."
            )

        if "api_version" in config:
            raise Exception(
                "Passing api_version is now deprecated. Please pass an OpenAI client instead."
            )

        if client is not None:
            self.client = client
            return

        if config is None and client is None:
            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            return

        if "api_key" in config:
            self.client = OpenAI(api_key=config["api_key"])

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        # Count the number of tokens in the message log
        # Use 4 as an approximation for the number of characters per token
        num_tokens = 0
        for message in prompt:
            num_tokens += len(message["content"]) / 4

        if kwargs.get("model", None) is not None:
            model = kwargs.get("model", None)
            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        elif kwargs.get("engine", None) is not None:
            engine = kwargs.get("engine", None)
            print(f"Using model {engine} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                engine=engine,
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        elif self.config is not None and "engine" in self.config:
            print(
                f"Using engine {self.config['engine']} for {num_tokens} tokens (approx)"
            )
            response = self.client.chat.completions.create(
                engine=self.config["engine"],
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        elif self.config is not None and "model" in self.config:
            print(
                f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
            )
            response = self.client.chat.completions.create(
                model=self.config["model"],
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        else:
            if num_tokens > 3500:
                model = "gpt-3.5-turbo-16k"
            else:
                model = "gpt-3.5-turbo"

            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )

        # Find the first response from the chatbot that has text in it (some responses may not have text)
        for choice in response.choices:
            if "text" in choice:
                return choice.text

        # If no response with text is found, return the first response's content (which may be empty)
        return response.choices[0].message.content


================================================
FILE: src/vanna/legacy/openai/openai_embeddings.py
================================================
from openai import OpenAI

from ..base import VannaBase


class OpenAI_Embeddings(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        if client is not None:
            self.client = client
            return

        if self.client is not None:
            return

        self.client = OpenAI()

        if config is None:
            return

        if "api_type" in config:
            self.client.api_type = config["api_type"]

        if "api_base" in config:
            self.client.api_base = config["api_base"]

        if "api_version" in config:
            self.client.api_version = config["api_version"]

        if "api_key" in config:
            self.client.api_key = config["api_key"]

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        if self.config is not None and "engine" in self.config:
            embedding = self.client.embeddings.create(
                engine=self.config["engine"],
                input=data,
            )
        else:
            embedding = self.client.embeddings.create(
                model="text-embedding-ada-002",
                input=data,
            )

        return embedding.get("data")[0]["embedding"]


================================================
FILE: src/vanna/legacy/opensearch/__init__.py
================================================
from .opensearch_vector import OpenSearch_VectorStore
from .opensearch_vector_semantic import OpenSearch_Semantic_VectorStore


================================================
FILE: src/vanna/legacy/opensearch/opensearch_vector.py
================================================
import base64
import uuid
from typing import List

import pandas as pd
from opensearchpy import OpenSearch

from ..base import VannaBase


class OpenSearch_VectorStore(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        document_index = "vanna_document_index"
        ddl_index = "vanna_ddl_index"
        question_sql_index = "vanna_questions_sql_index"
        if config is not None and "es_document_index" in config:
            document_index = config["es_document_index"]
        if config is not None and "es_ddl_index" in config:
            ddl_index = config["es_ddl_index"]
        if config is not None and "es_question_sql_index" in config:
            question_sql_index = config["es_question_sql_index"]

        self.document_index = document_index
        self.ddl_index = ddl_index
        self.question_sql_index = question_sql_index
        print(
            "OpenSearch_VectorStore initialized with document_index: ",
            document_index,
            " ddl_index: ",
            ddl_index,
            " question_sql_index: ",
            question_sql_index,
        )

        document_index_settings = {
            "settings": {"index": {"number_of_shards": 6, "number_of_replicas": 2}},
            "mappings": {
                "properties": {
                    "question": {
                        "type": "text",
                    },
                    "doc": {
                        "type": "text",
                    },
                }
            },
        }

        ddl_index_settings = {
            "settings": {"index": {"number_of_shards": 6, "number_of_replicas": 2}},
            "mappings": {
                "properties": {
                    "ddl": {
                        "type": "text",
                    },
                    "doc": {
                        "type": "text",
                    },
                }
            },
        }

        question_sql_index_settings = {
            "settings": {"index": {"number_of_shards": 6, "number_of_replicas": 2}},
            "mappings": {
                "properties": {
                    "question": {
                        "type": "text",
                    },
                    "sql": {
                        "type": "text",
                    },
                }
            },
        }

        if config is not None and "es_document_index_settings" in config:
            document_index_settings = config["es_document_index_settings"]
        if config is not None and "es_ddl_index_settings" in config:
            ddl_index_settings = config["es_ddl_index_settings"]
        if config is not None and "es_question_sql_index_settings" in config:
            question_sql_index_settings = config["es_question_sql_index_settings"]

        self.document_index_settings = document_index_settings
        self.ddl_index_settings = ddl_index_settings
        self.question_sql_index_settings = question_sql_index_settings

        es_urls = None
        if config is not None and "es_urls" in config:
            es_urls = config["es_urls"]

        # Host and port
        if config is not None and "es_host" in config:
            host = config["es_host"]
        else:
            host = "localhost"

        if config is not None and "es_port" in config:
            port = config["es_port"]
        else:
            port = 9200

        if config is not None and "es_ssl" in config:
            ssl = config["es_ssl"]
        else:
            ssl = False

        if config is not None and "es_verify_certs" in config:
            verify_certs = config["es_verify_certs"]
        else:
            verify_certs = False

        # Authentication
        if config is not None and "es_user" in config:
            auth = (config["es_user"], config["es_password"])
        else:
            # Default to admin:admin
            auth = None

        headers = None
        # base64 authentication
        if (
            config is not None
            and "es_encoded_base64" in config
            and "es_user" in config
            and "es_password" in config
        ):
            if config["es_encoded_base64"]:
                encoded_credentials = base64.b64encode(
                    (config["es_user"] + ":" + config["es_password"]).encode("utf-8")
                ).decode("utf-8")
                headers = {"Authorization": "Basic " + encoded_credentials}
                # remove auth from config
                auth = None

        # custom headers
        if config is not None and "es_headers" in config:
            headers = config["es_headers"]

        if config is not None and "es_timeout" in config:
            timeout = config["es_timeout"]
        else:
            timeout = 60

        if config is not None and "es_max_retries" in config:
            max_retries = config["es_max_retries"]
        else:
            max_retries = 10

        if config is not None and "es_http_compress" in config:
            es_http_compress = config["es_http_compress"]
        else:
            es_http_compress = False

        print(
            "OpenSearch_VectorStore initialized with es_urls: ",
            es_urls,
            " host: ",
            host,
            " port: ",
            port,
            " ssl: ",
            ssl,
            " verify_certs: ",
            verify_certs,
            " timeout: ",
            timeout,
            " max_retries: ",
            max_retries,
        )
        if es_urls is not None:
            # Initialize the OpenSearch client by passing a list of URLs
            self.client = OpenSearch(
                hosts=[es_urls],
                http_compress=es_http_compress,
                use_ssl=ssl,
                verify_certs=verify_certs,
                timeout=timeout,
                max_retries=max_retries,
                retry_on_timeout=True,
                http_auth=auth,
                headers=headers,
            )
        else:
            # Initialize the OpenSearch client by passing a host and port
            self.client = OpenSearch(
                hosts=[{"host": host, "port": port}],
                http_compress=es_http_compress,
                use_ssl=ssl,
                verify_certs=verify_certs,
                timeout=timeout,
                max_retries=max_retries,
                retry_on_timeout=True,
                http_auth=auth,
                headers=headers,
            )

            print("OpenSearch_VectorStore initialized with client over ")

        # 执行一个简单的查询来检查连接
        try:
            print("Connected to OpenSearch cluster:")
            info = self.client.info()
            print("OpenSearch cluster info:", info)
        except Exception as e:
            print("Error connecting to OpenSearch cluster:", e)

        # Create the indices if they don't exist
        self.create_index_if_not_exists(
            self.document_index, self.document_index_settings
        )
        self.create_index_if_not_exists(self.ddl_index, self.ddl_index_settings)
        self.create_index_if_not_exists(
            self.question_sql_index, self.question_sql_index_settings
        )

    def create_index(self):
        for index in [self.document_index, self.ddl_index, self.question_sql_index]:
            try:
                self.client.indices.create(index)
            except Exception as e:
                print("Error creating index: ", e)
                print(f"opensearch index {index} already exists")
                pass

    def create_index_if_not_exists(self, index_name: str, index_settings: dict) -> bool:
        try:
            if not self.client.indices.exists(index_name):
                print(f"Index {index_name} does not exist. Creating...")
                self.client.indices.create(index=index_name, body=index_settings)
                return True
            else:
                print(f"Index {index_name} already exists.")
                return False
        except Exception as e:
            print(f"Error creating index: {index_name} ", e)
            return False

    def add_ddl(self, ddl: str, **kwargs) -> str:
        # Assuming that you have a DDL index in your OpenSearch
        id = str(uuid.uuid4()) + "-ddl"
        ddl_dict = {"ddl": ddl}
        response = self.client.index(
            index=self.ddl_index, body=ddl_dict, id=id, **kwargs
        )
        return response["_id"]

    def add_documentation(self, doc: str, **kwargs) -> str:
        # Assuming you have a documentation index in your OpenSearch
        id = str(uuid.uuid4()) + "-doc"
        doc_dict = {"doc": doc}
        response = self.client.index(
            index=self.document_index, id=id, body=doc_dict, **kwargs
        )
        return response["_id"]

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        # Assuming you have a Questions and SQL index in your OpenSearch
        id = str(uuid.uuid4()) + "-sql"
        question_sql_dict = {"question": question, "sql": sql}
        response = self.client.index(
            index=self.question_sql_index, body=question_sql_dict, id=id, **kwargs
        )
        return response["_id"]

    def get_related_ddl(self, question: str, **kwargs) -> List[str]:
        # Assume you have some vector search mechanism associated with your data
        query = {"query": {"match": {"ddl": question}}}
        print(query)
        response = self.client.search(index=self.ddl_index, body=query, **kwargs)
        return [hit["_source"]["ddl"] for hit in response["hits"]["hits"]]

    def get_related_documentation(self, question: str, **kwargs) -> List[str]:
        query = {"query": {"match": {"doc": question}}}
        print(query)
        response = self.client.search(index=self.document_index, body=query, **kwargs)
        return [hit["_source"]["doc"] for hit in response["hits"]["hits"]]

    def get_similar_question_sql(self, question: str, **kwargs) -> List[str]:
        query = {"query": {"match": {"question": question}}}
        print(query)
        response = self.client.search(
            index=self.question_sql_index, body=query, **kwargs
        )
        return [
            (hit["_source"]["question"], hit["_source"]["sql"])
            for hit in response["hits"]["hits"]
        ]

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        # This will be a simple example pulling all data from an index
        # WARNING: Do not use this approach in production for large indices!
        data = []
        response = self.client.search(
            index=self.document_index, body={"query": {"match_all": {}}}, size=1000
        )
        print(response)
        # records = [hit['_source'] for hit in response['hits']['hits']]
        for hit in response["hits"]["hits"]:
            data.append(
                {
                    "id": hit["_id"],
                    "training_data_type": "documentation",
                    "question": "",
                    "content": hit["_source"]["doc"],
                }
            )

        response = self.client.search(
            index=self.question_sql_index, body={"query": {"match_all": {}}}, size=1000
        )
        # records = [hit['_source'] for hit in response['hits']['hits']]
        for hit in response["hits"]["hits"]:
            data.append(
                {
                    "id": hit["_id"],
                    "training_data_type": "sql",
                    "question": hit.get("_source", {}).get("question", ""),
                    "content": hit.get("_source", {}).get("sql", ""),
                }
            )

        response = self.client.search(
            index=self.ddl_index, body={"query": {"match_all": {}}}, size=1000
        )
        # records = [hit['_source'] for hit in response['hits']['hits']]
        for hit in response["hits"]["hits"]:
            data.append(
                {
                    "id": hit["_id"],
                    "training_data_type": "ddl",
                    "question": "",
                    "content": hit["_source"]["ddl"],
                }
            )

        return pd.DataFrame(data)

    def remove_training_data(self, id: str, **kwargs) -> bool:
        try:
            if id.endswith("-sql"):
                self.client.delete(index=self.question_sql_index, id=id)
                return True
            elif id.endswith("-ddl"):
                self.client.delete(index=self.ddl_index, id=id, **kwargs)
                return True
            elif id.endswith("-doc"):
                self.client.delete(index=self.document_index, id=id, **kwargs)
                return True
            else:
                return False
        except Exception as e:
            print("Error deleting training dataError deleting training data: ", e)
            return False

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        # opensearch doesn't need to generate embeddings
        pass


# OpenSearch_VectorStore.__init__(self, config={'es_urls':
# "https://opensearch-node.test.com:9200", 'es_encoded_base64': True, 'es_user':
# "admin", 'es_password': "admin", 'es_verify_certs': True})


# OpenSearch_VectorStore.__init__(self, config={'es_host':
# "https://opensearch-node.test.com", 'es_port': 9200, 'es_user': "admin",
# 'es_password': "admin", 'es_verify_certs': True})


================================================
FILE: src/vanna/legacy/opensearch/opensearch_vector_semantic.py
================================================
import json

import pandas as pd
from langchain_community.vectorstores import OpenSearchVectorSearch

from ..base import VannaBase
from ..utils import deterministic_uuid


class OpenSearch_Semantic_VectorStore(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        if config is None:
            config = {}

        if "embedding_function" in config:
            self.embedding_function = config.get("embedding_function")
        else:
            from langchain_huggingface import HuggingFaceEmbeddings

            self.embedding_function = HuggingFaceEmbeddings(
                model_name="all-MiniLM-L6-v2"
            )

        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
        self.n_results_documentation = config.get(
            "n_results_documentation", config.get("n_results", 10)
        )
        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))

        self.document_index = config.get("es_document_index", "vanna_document_index")
        self.ddl_index = config.get("es_ddl_index", "vanna_ddl_index")
        self.question_sql_index = config.get(
            "es_question_sql_index", "vanna_questions_sql_index"
        )

        self.log(
            f"OpenSearch_Semantic_VectorStore initialized with document_index: {self.document_index}, ddl_index: {self.ddl_index}, question_sql_index: {self.question_sql_index}"
        )

        es_urls = config.get("es_urls", "https://localhost:9200")
        ssl = config.get("es_ssl", True)
        verify_certs = config.get("es_verify_certs", True)

        if "es_user" in config:
            auth = (config["es_user"], config["es_password"])
        else:
            auth = None

        headers = config.get("es_headers", None)
        timeout = config.get("es_timeout", 60)
        max_retries = config.get("es_max_retries", 10)

        common_args = {
            "opensearch_url": es_urls,
            "embedding_function": self.embedding_function,
            "engine": "faiss",
            "http_auth": auth,
            "use_ssl": ssl,
            "verify_certs": verify_certs,
            "timeout": timeout,
            "max_retries": max_retries,
            "retry_on_timeout": True,
            "headers": headers,
        }

        self.documentation_store = OpenSearchVectorSearch(
            index_name=self.document_index, **common_args
        )
        self.ddl_store = OpenSearchVectorSearch(
            index_name=self.ddl_index, **common_args
        )
        self.sql_store = OpenSearchVectorSearch(
            index_name=self.question_sql_index, **common_args
        )

    def add_ddl(self, ddl: str, **kwargs) -> str:
        _id = deterministic_uuid(ddl) + "-ddl"
        self.ddl_store.add_texts(texts=[ddl], ids=[_id], **kwargs)
        return _id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        _id = deterministic_uuid(documentation) + "-doc"
        self.documentation_store.add_texts(texts=[documentation], ids=[_id], **kwargs)
        return _id

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        question_sql_json = json.dumps(
            {
                "question": question,
                "sql": sql,
            },
            ensure_ascii=False,
        )

        _id = deterministic_uuid(question_sql_json) + "-sql"
        self.sql_store.add_texts(texts=[question_sql_json], ids=[_id], **kwargs)
        return _id

    def get_related_ddl(self, question: str, **kwargs) -> list:
        documents = self.ddl_store.similarity_search(
            query=question, k=self.n_results_ddl
        )
        return [document.page_content for document in documents]

    def get_related_documentation(self, question: str, **kwargs) -> list:
        documents = self.documentation_store.similarity_search(
            query=question, k=self.n_results_documentation
        )
        return [document.page_content for document in documents]

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        documents = self.sql_store.similarity_search(
            query=question, k=self.n_results_sql
        )
        return [json.loads(document.page_content) for document in documents]

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        data = []
        query = {"query": {"match_all": {}}}

        indices = [
            {"index": self.document_index, "type": "documentation"},
            {"index": self.question_sql_index, "type": "sql"},
            {"index": self.ddl_index, "type": "ddl"},
        ]

        # Use documentation_store.client consistently for search on all indices
        opensearch_client = self.documentation_store.client

        for index_info in indices:
            index_name = index_info["index"]
            training_data_type = index_info["type"]
            scroll = "1m"  # keep scroll context for 1 minute
            response = opensearch_client.search(
                index=index_name,
                ignore_unavailable=True,
                body=query,
                scroll=scroll,
                size=1000,
            )

            scroll_id = response.get("_scroll_id")

            while scroll_id:
                hits = response["hits"]["hits"]
                if not hits:
                    break  # No more hits, exit loop

                for hit in hits:
                    source = hit["_source"]
                    if training_data_type == "sql":
                        try:
                            doc_dict = json.loads(source["text"])
                            content = doc_dict.get("sql")
                            question = doc_dict.get("question")
                        except json.JSONDecodeError as e:
                            self.log(
                                f"Skipping row with custom_id {hit['_id']} due to JSON parsing error: {e}",
                                "Error",
                            )
                            continue
                    else:  # documentation or ddl
                        content = source["text"]
                        question = None

                    data.append(
                        {
                            "id": hit["_id"],
                            "training_data_type": training_data_type,
                            "question": question,
                            "content": content,
                        }
                    )

                # Get next batch of results, using documentation_store.client.scroll
                response = opensearch_client.scroll(scroll_id=scroll_id, scroll=scroll)
                scroll_id = response.get("_scroll_id")

        return pd.DataFrame(data)

    def remove_training_data(self, id: str, **kwargs) -> bool:
        try:
            if id.endswith("-sql"):
                return self.sql_store.delete(ids=[id], **kwargs)
            elif id.endswith("-ddl"):
                return self.ddl_store.delete(ids=[id], **kwargs)
            elif id.endswith("-doc"):
                return self.documentation_store.delete(ids=[id], **kwargs)
            else:
                return False
        except Exception as e:
            self.log(
                f"Error deleting training dataError deleting training data: {e}",
                "Error",
            )
            return False

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        pass


================================================
FILE: src/vanna/legacy/oracle/__init__.py
================================================
from .oracle_vector import Oracle_VectorStore


================================================
FILE: src/vanna/legacy/oracle/oracle_vector.py
================================================
import json
import uuid
from typing import List, Optional, Tuple

import oracledb
import pandas as pd
from chromadb.utils import embedding_functions

from ..base import VannaBase

default_ef = embedding_functions.DefaultEmbeddingFunction()


class Oracle_VectorStore(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        if config is not None:
            self.embedding_function = config.get("embedding_function", default_ef)
            self.pre_delete_collection = config.get("pre_delete_collection", False)
            self.cmetadata = config.get("cmetadata", {"created_by": "oracle"})
        else:
            self.embedding_function = default_ef
            self.pre_delete_collection = False
            self.cmetadata = {"created_by": "oracle"}

        self.oracle_conn = oracledb.connect(dsn=config.get("dsn"))
        self.oracle_conn.call_timeout = 30000
        self.documentation_collection = "documentation"
        self.ddl_collection = "ddl"
        self.sql_collection = "sql"
        self.n_results = config.get("n_results", 10)
        self.n_results_ddl = config.get("n_results_ddl", self.n_results)
        self.n_results_sql = config.get("n_results_sql", self.n_results)
        self.n_results_documentation = config.get(
            "n_results_documentation", self.n_results
        )
        self.create_tables_if_not_exists()
        self.create_collections_if_not_exists(self.documentation_collection)
        self.create_collections_if_not_exists(self.ddl_collection)
        self.create_collections_if_not_exists(self.sql_collection)

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embeddings = self.embedding_function([data])
        if len(embeddings) == 1:
            return list(embeddings[0].astype(float))
        return list(embeddings.astype(float))

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        cmetadata = self.cmetadata.copy()
        collection = self.get_collection(self.sql_collection)
        question_sql_json = json.dumps(
            {
                "question": question,
                "sql": sql,
            },
            ensure_ascii=False,
        )
        id = str(uuid.uuid4())
        embeddings = self.generate_embedding(question)
        custom_id = id + "-sql"

        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(None, oracledb.DB_TYPE_VECTOR)
        cursor.execute(
            """
          INSERT INTO oracle_embedding (
              collection_id,
              embedding,
              document,
              cmetadata,
              custom_id,
              uuid
          ) VALUES (
              :1,
              TO_VECTOR(:2),
              :3,
              :4,
              :5,
              :6
          )
      """,
            [
                collection["uuid"],
                embeddings,
                question_sql_json,
                json.dumps(cmetadata),
                custom_id,
                id,
            ],
        )

        self.oracle_conn.commit()
        cursor.close()
        return id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        collection = self.get_collection(self.ddl_collection)
        question_ddl_json = json.dumps(
            {
                "question": None,
                "ddl": ddl,
            },
            ensure_ascii=False,
        )
        id = str(uuid.uuid4())
        custom_id = id + "-ddl"
        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(None, oracledb.DB_TYPE_VECTOR)
        cursor.execute(
            """
      INSERT INTO oracle_embedding (
          collection_id,
          embedding,
          document,
          cmetadata,
          custom_id,
          uuid
      ) VALUES (
          :1,
          TO_VECTOR(:2),
          :3,
          :4,
          :5,
          :6
      )
      """,
            [
                collection["uuid"],
                self.generate_embedding(ddl),
                question_ddl_json,
                json.dumps(self.cmetadata),
                custom_id,
                id,
            ],
        )
        self.oracle_conn.commit()
        cursor.close()
        return id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        collection = self.get_collection(self.documentation_collection)
        question_documentation_json = json.dumps(
            {
                "question": None,
                "documentation": documentation,
            },
            ensure_ascii=False,
        )
        id = str(uuid.uuid4())
        custom_id = id + "-doc"
        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(None, oracledb.DB_TYPE_VECTOR)
        cursor.execute(
            """
      INSERT INTO oracle_embedding (
          collection_id,
          embedding,
          document,
          cmetadata,
          custom_id,
          uuid
      ) VALUES (
          :1,
          TO_VECTOR(:2),
          :3,
          :4,
          :5,
          :6
      )
      """,
            [
                collection["uuid"],
                self.generate_embedding(documentation),
                question_documentation_json,
                json.dumps(self.cmetadata),
                custom_id,
                id,
            ],
        )
        self.oracle_conn.commit()
        cursor.close()
        return id

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        df = pd.DataFrame()

        cursor = self.oracle_conn.cursor()
        sql_collection = self.get_collection(self.sql_collection)
        cursor.execute(
            """
                    SELECT
                        document,
                        uuid
                    FROM oracle_embedding
                    WHERE
                        collection_id = :1
                  """,
            [sql_collection["uuid"]],
        )
        sql_data = cursor.fetchall()

        if sql_data is not None:
            # Extract the documents and ids
            documents = [row_data[0] for row_data in sql_data]
            ids = [row_data[1] for row_data in sql_data]

            # Create a DataFrame
            df_sql = pd.DataFrame(
                {
                    "id": ids,
                    "question": [
                        json.loads(doc)["question"]
                        if isinstance(doc, str)
                        else doc["question"]
                        for doc in documents
                    ],
                    "content": [
                        json.loads(doc)["sql"] if isinstance(doc, str) else doc["sql"]
                        for doc in documents
                    ],
                }
            )
            df_sql["training_data_type"] = "sql"
            df = pd.concat([df, df_sql])

        ddl_collection = self.get_collection(self.ddl_collection)
        cursor.execute(
            """
          SELECT
              document,
              uuid
          FROM oracle_embedding
          WHERE
              collection_id = :1
      """,
            [ddl_collection["uuid"]],
        )
        ddl_data = cursor.fetchall()

        if ddl_data is not None:
            # Extract the documents and ids
            documents = [row_data[0] for row_data in ddl_data]
            ids = [row_data[1] for row_data in ddl_data]

            # Create a DataFrame
            df_ddl = pd.DataFrame(
                {
                    "id": ids,
                    "question": [None for _ in documents],
                    "content": [
                        json.loads(doc)["ddl"] if isinstance(doc, str) else doc["ddl"]
                        for doc in documents
                    ],
                }
            )
            df_ddl["training_data_type"] = "ddl"
            df = pd.concat([df, df_ddl])

        doc_collection = self.get_collection(self.documentation_collection)
        cursor.execute(
            """
          SELECT
              document,
              uuid
          FROM oracle_embedding
          WHERE
              collection_id = :1
      """,
            [doc_collection["uuid"]],
        )
        doc_data = cursor.fetchall()

        if doc_data is not None:
            # Extract the documents and ids
            documents = [row_data[0] for row_data in doc_data]
            ids = [row_data[1] for row_data in doc_data]

            # Create a DataFrame
            df_doc = pd.DataFrame(
                {
                    "id": ids,
                    "question": [None for _ in documents],
                    "content": [
                        json.loads(doc)["documentation"]
                        if isinstance(doc, str)
                        else doc["documentation"]
                        for doc in documents
                    ],
                }
            )
            df_doc["training_data_type"] = "documentation"
            df = pd.concat([df, df_doc])

        self.oracle_conn.commit()
        cursor.close()
        return df

    def remove_training_data(self, id: str, **kwargs) -> bool:
        cursor = self.oracle_conn.cursor()
        cursor.execute(
            """
          DELETE
          FROM
              oracle_embedding
          WHERE
          uuid  = :1
      """,
            [id],
        )

        self.oracle_conn.commit()
        cursor.close()
        return True

    def update_training_data(
        self, id: str, train_type: str, question: str, **kwargs
    ) -> bool:
        print(f"{train_type=}")
        update_content = kwargs["content"]
        if train_type == "sql":
            update_json = json.dumps(
                {
                    "question": question,
                    "sql": update_content,
                }
            )
        elif train_type == "ddl":
            update_json = json.dumps(
                {
                    "question": None,
                    "ddl": update_content,
                }
            )
        elif train_type == "documentation":
            update_json = json.dumps(
                {
                    "question": None,
                    "documentation": update_content,
                }
            )
        else:
            update_json = json.dumps(
                {
                    "question": question,
                    "sql": update_content,
                }
            )
        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(oracledb.DB_TYPE_VECTOR, oracledb.DB_TYPE_JSON)
        cursor.execute(
            """
                  UPDATE
                      oracle_embedding
                  SET
                      embedding = TO_VECTOR(:1),
                      document = JSON_MERGEPATCH(document, :2)
                  WHERE
                  uuid  = :3
                  """,
            [self.generate_embedding(update_content), update_json, id],
        )

        self.oracle_conn.commit()
        cursor.close()
        return True

    @staticmethod
    def _extract_documents(query_results) -> list:
        """
        Static method to extract the documents from the results of a query.

        Args:
            query_results (pd.DataFrame): The dataframe to use.

        Returns:
            List[str] or None: The extracted documents, or an empty list or single document if an error occurred.
        """
        if query_results is None or len(query_results) == 0:
            return []

        documents = [
            json.loads(row_data[0]) if isinstance(row_data[0], str) else row_data[0]
            for row_data in query_results
        ]

        return documents

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        embeddings = self.generate_embedding(question)
        collection = self.get_collection(self.sql_collection)
        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(None, oracledb.DB_TYPE_VECTOR, oracledb.DB_TYPE_VECTOR)
        cursor.execute(
            """
      SELECT document
      FROM oracle_embedding
      WHERE collection_id = :1
      ORDER BY VECTOR_DISTANCE(embedding, TO_VECTOR(:2), COSINE)
      FETCH FIRST :3 ROWS ONLY
      """,
            [collection["uuid"], embeddings, self.n_results_sql],
        )
        results = cursor.fetchall()
        cursor.close()
        return self._extract_documents(results)

    def get_related_ddl(self, question: str, **kwargs) -> list:
        collection = self.get_collection(self.ddl_collection)
        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(None, oracledb.DB_TYPE_VECTOR)
        cursor.execute(
            """
          SELECT
              document
          FROM oracle_embedding
          WHERE
              collection_id = :1
          ORDER BY VECTOR_DISTANCE(embedding, TO_VECTOR(:2), COSINE)
          FETCH FIRST :top_k ROWS ONLY
      """,
            [collection["uuid"], self.generate_embedding(question), 100],
        )
        results = cursor.fetchall()

        self.oracle_conn.commit()
        cursor.close()
        return Oracle_VectorStore._extract_documents(results)

    def search_tables_metadata(
        self,
        engine: str = None,
        catalog: str = None,
        schema: str = None,
        table_name: str = None,
        ddl: str = None,
        size: int = 10,
        **kwargs,
    ) -> list:
        pass

    def get_related_documentation(self, question: str, **kwargs) -> list:
        collection = self.get_collection(self.documentation_collection)
        cursor = self.oracle_conn.cursor()
        cursor.setinputsizes(None, oracledb.DB_TYPE_VECTOR)
        cursor.execute(
            """
          SELECT
              document
          FROM oracle_embedding
          WHERE
              collection_id = :1
          ORDER BY VECTOR_DISTANCE(embedding, TO_VECTOR(:2), DOT)
          FETCH FIRST :top_k ROWS ONLY
      """,
            [collection["uuid"], self.generate_embedding(question), 100],
        )
        results = cursor.fetchall()

        self.oracle_conn.commit()
        cursor.close()

        return Oracle_VectorStore._extract_documents(results)

    def create_tables_if_not_exists(self) -> None:
        cursor = self.oracle_conn.cursor()
        cursor.execute(
            """
          CREATE TABLE IF NOT EXISTS oracle_collection (
              name      VARCHAR2(200) NOT NULL,
              cmetadata json NOT NULL,
              uuid      VARCHAR2(200) NOT NULL,
              CONSTRAINT oc_key_uuid PRIMARY KEY ( uuid )
          )
      """
        )

        cursor.execute(
            """
          CREATE TABLE IF NOT EXISTS oracle_embedding (
              collection_id VARCHAR2(200) NOT NULL,
              embedding     vector NOT NULL,
              document      json NOT NULL,
              cmetadata     json NOT NULL,
              custom_id     VARCHAR2(200) NOT NULL,
              uuid          VARCHAR2(200) NOT NULL,
              CONSTRAINT oe_key_uuid PRIMARY KEY ( uuid )
          )
      """
        )

        self.oracle_conn.commit()
        cursor.close()

    def create_collections_if_not_exists(
        self,
        name: str,
        cmetadata: Optional[dict] = None,
    ) -> Tuple[dict, bool]:
        """
        Get or create a collection.
        Returns [Collection, bool] where the bool is True if the collection was created.
        """
        if self.pre_delete_collection:
            self.delete_collection(name)
        created = False
        collection = self.get_collection(name)
        if collection:
            return collection, created

        cmetadata = (
            json.dumps(self.cmetadata) if cmetadata is None else json.dumps(cmetadata)
        )
        collection_id = str(uuid.uuid4())
        cursor = self.oracle_conn.cursor()
        cursor.execute(
            """
          INSERT INTO oracle_collection(name, cmetadata, uuid)
          VALUES (:1, :2, :3)
      """,
            [name, cmetadata, str(collection_id)],
        )

        self.oracle_conn.commit()
        cursor.close()

        collection = {"name": name, "cmetadata": cmetadata, "uuid": collection_id}
        created = True
        return collection, created

    def get_collection(self, name) -> Optional[dict]:
        return self.get_by_name(name)

    def get_by_name(self, name: str) -> Optional[dict]:
        cursor = self.oracle_conn.cursor()
        cursor.execute(
            """
          SELECT
              name,
              cmetadata,
              uuid
          FROM
              oracle_collection
          WHERE
              name = :1
          FETCH FIRST 1 ROWS ONLY
      """,
            [name],
        )

        for row in cursor:
            return {"name": row[0], "cmetadata": row[1], "uuid": row[2]}

        return  # type: ignore

    def delete_collection(self, name) -> None:
        collection = self.get_collection(name)
        if not collection:
            return

        cursor = self.oracle_conn.cursor()
        cursor.execute(
            """
          DELETE
          FROM
              oracle_embedding
          WHERE
          collection_id = ( SELECT uuid FROM oracle_collection WHERE name = :1 )
      """,
            [name],
        )
        cursor.execute(
            """
          DELETE
          FROM
              oracle_collection
          WHERE
              name = :1
      """,
            [name],
        )

        self.oracle_conn.commit()
        cursor.close()


================================================
FILE: src/vanna/legacy/pgvector/__init__.py
================================================
from .pgvector import PG_VectorStore


================================================
FILE: src/vanna/legacy/pgvector/pgvector.py
================================================
import ast
import json
import logging
import uuid

import pandas as pd
from langchain_core.documents import Document
from langchain_postgres.vectorstores import PGVector
from sqlalchemy import create_engine, text

from .. import ValidationError
from ..base import VannaBase
from ..types import TrainingPlan, TrainingPlanItem


class PG_VectorStore(VannaBase):
    def __init__(self, config=None):
        if not config or "connection_string" not in config:
            raise ValueError(
                "A valid 'config' dictionary with a 'connection_string' is required."
            )

        VannaBase.__init__(self, config=config)

        if config and "connection_string" in config:
            self.connection_string = config.get("connection_string")
            self.n_results = config.get("n_results", 10)

        if config and "embedding_function" in config:
            self.embedding_function = config.get("embedding_function")
        else:
            from langchain_huggingface import HuggingFaceEmbeddings

            self.embedding_function = HuggingFaceEmbeddings(
                model_name="all-MiniLM-L6-v2"
            )

        self.sql_collection = PGVector(
            embeddings=self.embedding_function,
            collection_name="sql",
            connection=self.connection_string,
        )
        self.ddl_collection = PGVector(
            embeddings=self.embedding_function,
            collection_name="ddl",
            connection=self.connection_string,
        )
        self.documentation_collection = PGVector(
            embeddings=self.embedding_function,
            collection_name="documentation",
            connection=self.connection_string,
        )

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        question_sql_json = json.dumps(
            {
                "question": question,
                "sql": sql,
            },
            ensure_ascii=False,
        )
        id = str(uuid.uuid4()) + "-sql"
        createdat = kwargs.get("createdat")
        doc = Document(
            page_content=question_sql_json,
            metadata={"id": id, "createdat": createdat},
        )
        self.sql_collection.add_documents([doc], ids=[doc.metadata["id"]])

        return id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        _id = str(uuid.uuid4()) + "-ddl"
        doc = Document(
            page_content=ddl,
            metadata={"id": _id},
        )
        self.ddl_collection.add_documents([doc], ids=[doc.metadata["id"]])
        return _id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        _id = str(uuid.uuid4()) + "-doc"
        doc = Document(
            page_content=documentation,
            metadata={"id": _id},
        )
        self.documentation_collection.add_documents([doc], ids=[doc.metadata["id"]])
        return _id

    def get_collection(self, collection_name):
        match collection_name:
            case "sql":
                return self.sql_collection
            case "ddl":
                return self.ddl_collection
            case "documentation":
                return self.documentation_collection
            case _:
                raise ValueError("Specified collection does not exist.")

    def get_similar_question_sql(self, question: str) -> list:
        documents = self.sql_collection.similarity_search(
            query=question, k=self.n_results
        )
        return [ast.literal_eval(document.page_content) for document in documents]

    def get_related_ddl(self, question: str, **kwargs) -> list:
        documents = self.ddl_collection.similarity_search(
            query=question, k=self.n_results
        )
        return [document.page_content for document in documents]

    def get_related_documentation(self, question: str, **kwargs) -> list:
        documents = self.documentation_collection.similarity_search(
            query=question, k=self.n_results
        )
        return [document.page_content for document in documents]

    def train(
        self,
        question: str | None = None,
        sql: str | None = None,
        ddl: str | None = None,
        documentation: str | None = None,
        plan: TrainingPlan | None = None,
        createdat: str | None = None,
    ):
        if question and not sql:
            raise ValidationError("Please provide a SQL query.")

        if documentation:
            logging.info(f"Adding documentation: {documentation}")
            return self.add_documentation(documentation)

        if sql and question:
            return self.add_question_sql(
                question=question, sql=sql, createdat=createdat
            )

        if ddl:
            logging.info(f"Adding ddl: {ddl}")
            return self.add_ddl(ddl)

        if plan:
            for item in plan._plan:
                if item.item_type == TrainingPlanItem.ITEM_TYPE_DDL:
                    self.add_ddl(item.item_value)
                elif item.item_type == TrainingPlanItem.ITEM_TYPE_IS:
                    self.add_documentation(item.item_value)
                elif (
                    item.item_type == TrainingPlanItem.ITEM_TYPE_SQL and item.item_name
                ):
                    self.add_question_sql(question=item.item_name, sql=item.item_value)

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        # Establishing the connection
        engine = create_engine(self.connection_string)

        # Querying the 'langchain_pg_embedding' table
        query_embedding = "SELECT cmetadata, document FROM langchain_pg_embedding"
        df_embedding = pd.read_sql(query_embedding, engine)

        # List to accumulate the processed rows
        processed_rows = []

        # Process each row in the DataFrame
        for _, row in df_embedding.iterrows():
            custom_id = row["cmetadata"]["id"]
            document = row["document"]
            training_data_type = (
                "documentation" if custom_id[-3:] == "doc" else custom_id[-3:]
            )

            if training_data_type == "sql":
                # Convert the document string to a dictionary
                try:
                    doc_dict = ast.literal_eval(document)
                    question = doc_dict.get("question")
                    content = doc_dict.get("sql")
                except (ValueError, SyntaxError):
                    logging.info(
                        f"Skipping row with custom_id {custom_id} due to parsing error."
                    )
                    continue
            elif training_data_type in ["documentation", "ddl"]:
                question = None  # Default value for question
                content = document
            else:
                # If the suffix is not recognized, skip this row
                logging.info(
                    f"Skipping row with custom_id {custom_id} due to unrecognized training data type."
                )
                continue

            # Append the processed data to the list
            processed_rows.append(
                {
                    "id": custom_id,
                    "question": question,
                    "content": content,
                    "training_data_type": training_data_type,
                }
            )

        # Create a DataFrame from the list of processed rows
        df_processed = pd.DataFrame(processed_rows)

        return df_processed

    def remove_training_data(self, id: str, **kwargs) -> bool:
        # Create the database engine
        engine = create_engine(self.connection_string)

        # SQL DELETE statement
        delete_statement = text(
            """
            DELETE FROM langchain_pg_embedding
            WHERE cmetadata ->> 'id' = :id
        """
        )

        # Connect to the database and execute the delete statement
        with engine.connect() as connection:
            # Start a transaction
            with connection.begin() as transaction:
                try:
                    result = connection.execute(delete_statement, {"id": id})
                    # Commit the transaction if the delete was successful
                    transaction.commit()
                    # Check if any row was deleted and return True or False accordingly
                    return result.rowcount > 0
                except Exception as e:
                    # Rollback the transaction in case of error
                    logging.error(f"An error occurred: {e}")
                    transaction.rollback()
                    return False

    def remove_collection(self, collection_name: str) -> bool:
        engine = create_engine(self.connection_string)

        # Determine the suffix to look for based on the collection name
        suffix_map = {"ddl": "ddl", "sql": "sql", "documentation": "doc"}
        suffix = suffix_map.get(collection_name)

        if not suffix:
            logging.info(
                "Invalid collection name. Choose from 'ddl', 'sql', or 'documentation'."
            )
            return False

        # SQL query to delete rows based on the condition
        query = text(
            f"""
            DELETE FROM langchain_pg_embedding
            WHERE cmetadata->>'id' LIKE '%{suffix}'
        """
        )

        # Execute the deletion within a transaction block
        with engine.connect() as connection:
            with connection.begin() as transaction:
                try:
                    result = connection.execute(query)
                    transaction.commit()  # Explicitly commit the transaction
                    if result.rowcount > 0:
                        logging.info(
                            f"Deleted {result.rowcount} rows from "
                            f"langchain_pg_embedding where collection is {collection_name}."
                        )
                        return True
                    else:
                        logging.info(
                            f"No rows deleted for collection {collection_name}."
                        )
                        return False
                except Exception as e:
                    logging.error(f"An error occurred: {e}")
                    transaction.rollback()  # Rollback in case of error
                    return False

    def generate_embedding(self, *args, **kwargs):
        pass


================================================
FILE: src/vanna/legacy/pinecone/__init__.py
================================================
from .pinecone_vector import PineconeDB_VectorStore

__all__ = ["PineconeDB_VectorStore"]


================================================
FILE: src/vanna/legacy/pinecone/pinecone_vector.py
================================================
import json
from typing import List

from pinecone import Pinecone, PodSpec, ServerlessSpec
import pandas as pd
from ..base import VannaBase
from ..utils import deterministic_uuid

from fastembed import TextEmbedding


class PineconeDB_VectorStore(VannaBase):
    """
    Vectorstore using PineconeDB

    Args:
        config (dict): Configuration dictionary. Defaults to {}. You must provide either a Pinecone Client or an API key in the config.
            - client (Pinecone, optional): Pinecone client. Defaults to None.
            - api_key (str, optional): Pinecone API key. Defaults to None.
            - n_results (int, optional): Number of results to return. Defaults to 10.
            - dimensions (int, optional): Dimensions of the embeddings. Defaults to 384 which coresponds to the dimensions of BAAI/bge-small-en-v1.5.
            - fastembed_model (str, optional): Fastembed model to use. Defaults to "BAAI/bge-small-en-v1.5".
            - documentation_namespace (str, optional): Namespace for documentation. Defaults to "documentation".
            - distance_metric (str, optional): Distance metric to use. Defaults to "cosine".
            - ddl_namespace (str, optional): Namespace for DDL. Defaults to "ddl".
            - sql_namespace (str, optional): Namespace for SQL. Defaults to "sql".
            - index_name (str, optional): Name of the index. Defaults to "vanna-index".
            - metadata_config (dict, optional): Metadata configuration if using a pinecone pod. Defaults to {}.
            - server_type (str, optional): Type of Pinecone server to use. Defaults to "serverless". Options are "serverless" or "pod".
            - podspec (PodSpec, optional): PodSpec configuration if using a pinecone pod. Defaults to PodSpec(environment="us-west-2", pod_type="p1.x1", metadata_config=self.metadata_config).
            - serverless_spec (ServerlessSpec, optional): ServerlessSpec configuration if using a pinecone serverless index. Defaults to ServerlessSpec(cloud="aws", region="us-west-2").
    Raises:
        ValueError: If config is None, api_key is not provided OR client is not provided, client is not an instance of Pinecone, or server_type is not "serverless" or "pod".
    """

    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)
        if config is None:
            raise ValueError(
                "config is required, pass either a Pinecone client or an API key in the config."
            )
        client = config.get("client")
        api_key = config.get("api_key")
        if not api_key and not client:
            raise ValueError(
                "api_key is required in config or pass a configured client"
            )
        if not client and api_key:
            self._client = Pinecone(api_key=api_key)
        elif not isinstance(client, Pinecone):
            raise ValueError("client must be an instance of Pinecone")
        else:
            self._client = client

        self.n_results = config.get("n_results", 10)
        self.dimensions = config.get("dimensions", 384)
        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
        self.documentation_namespace = config.get(
            "documentation_namespace", "documentation"
        )
        self.distance_metric = config.get("distance_metric", "cosine")
        self.ddl_namespace = config.get("ddl_namespace", "ddl")
        self.sql_namespace = config.get("sql_namespace", "sql")
        self.index_name = config.get("index_name", "vanna-index")
        self.metadata_config = config.get("metadata_config", {})
        self.server_type = config.get("server_type", "serverless")
        if self.server_type not in ["serverless", "pod"]:
            raise ValueError("server_type must be either 'serverless' or 'pod'")
        self.podspec = config.get(
            "podspec",
            PodSpec(
                environment="us-west-2",
                pod_type="p1.x1",
                metadata_config=self.metadata_config,
            ),
        )
        self.serverless_spec = config.get(
            "serverless_spec", ServerlessSpec(cloud="aws", region="us-west-2")
        )
        self._setup_index()

    def _set_index_host(self, host: str) -> None:
        self.Index = self._client.Index(host=host)

    def _setup_index(self) -> None:
        existing_indexes = self._get_indexes()
        if self.index_name not in existing_indexes and self.server_type == "serverless":
            self._client.create_index(
                name=self.index_name,
                dimension=self.dimensions,
                metric=self.distance_metric,
                spec=self.serverless_spec,
            )
            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
            self._set_index_host(pinecone_index_host)
        elif self.index_name not in existing_indexes and self.server_type == "pod":
            self._client.create_index(
                name=self.index_name,
                dimension=self.dimensions,
                metric=self.distance_metric,
                spec=self.podspec,
            )
            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
            self._set_index_host(pinecone_index_host)
        else:
            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
            self._set_index_host(pinecone_index_host)

    def _get_indexes(self) -> list:
        return [index["name"] for index in self._client.list_indexes()]

    def _check_if_embedding_exists(self, id: str, namespace: str) -> bool:
        fetch_response = self.Index.fetch(ids=[id], namespace=namespace)
        if fetch_response.vectors == {}:
            return False
        return True

    def add_ddl(self, ddl: str, **kwargs) -> str:
        id = deterministic_uuid(ddl) + "-ddl"
        if self._check_if_embedding_exists(id=id, namespace=self.ddl_namespace):
            print(f"DDL with id: {id} already exists in the index. Skipping...")
            return id
        self.Index.upsert(
            vectors=[(id, self.generate_embedding(ddl), {"ddl": ddl})],
            namespace=self.ddl_namespace,
        )
        return id

    def add_documentation(self, doc: str, **kwargs) -> str:
        id = deterministic_uuid(doc) + "-doc"

        if self._check_if_embedding_exists(
            id=id, namespace=self.documentation_namespace
        ):
            print(
                f"Documentation with id: {id} already exists in the index. Skipping..."
            )
            return id
        self.Index.upsert(
            vectors=[(id, self.generate_embedding(doc), {"documentation": doc})],
            namespace=self.documentation_namespace,
        )
        return id

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        question_sql_json = json.dumps(
            {
                "question": question,
                "sql": sql,
            },
            ensure_ascii=False,
        )
        id = deterministic_uuid(question_sql_json) + "-sql"
        if self._check_if_embedding_exists(id=id, namespace=self.sql_namespace):
            print(
                f"Question-SQL with id: {id} already exists in the index. Skipping..."
            )
            return id
        self.Index.upsert(
            vectors=[
                (
                    id,
                    self.generate_embedding(question_sql_json),
                    {"sql": question_sql_json},
                )
            ],
            namespace=self.sql_namespace,
        )
        return id

    def get_related_ddl(self, question: str, **kwargs) -> list:
        res = self.Index.query(
            namespace=self.ddl_namespace,
            vector=self.generate_embedding(question),
            top_k=self.n_results,
            include_values=True,
            include_metadata=True,
        )
        return [match["metadata"]["ddl"] for match in res["matches"]] if res else []

    def get_related_documentation(self, question: str, **kwargs) -> list:
        res = self.Index.query(
            namespace=self.documentation_namespace,
            vector=self.generate_embedding(question),
            top_k=self.n_results,
            include_values=True,
            include_metadata=True,
        )
        return (
            [match["metadata"]["documentation"] for match in res["matches"]]
            if res
            else []
        )

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        res = self.Index.query(
            namespace=self.sql_namespace,
            vector=self.generate_embedding(question),
            top_k=self.n_results,
            include_values=True,
            include_metadata=True,
        )
        return (
            [
                {
                    key: value
                    for key, value in json.loads(match["metadata"]["sql"]).items()
                }
                for match in res["matches"]
            ]
            if res
            else []
        )

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        # Pinecone does not support getting all vectors in a namespace, so we have to query for the top_k vectors with a dummy vector
        df = pd.DataFrame()
        namespaces = {
            "sql": self.sql_namespace,
            "ddl": self.ddl_namespace,
            "documentation": self.documentation_namespace,
        }

        for data_type, namespace in namespaces.items():
            data = self.Index.query(
                top_k=10000,  # max results that pinecone allows
                namespace=namespace,
                include_values=True,
                include_metadata=True,
                vector=[0.0] * self.dimensions,
            )

            if data is not None:
                id_list = [match["id"] for match in data["matches"]]
                content_list = [
                    match["metadata"][data_type] for match in data["matches"]
                ]
                question_list = [
                    (
                        json.loads(match["metadata"][data_type])["question"]
                        if data_type == "sql"
                        else None
                    )
                    for match in data["matches"]
                ]

                df_data = pd.DataFrame(
                    {
                        "id": id_list,
                        "question": question_list,
                        "content": content_list,
                    }
                )
                df_data["training_data_type"] = data_type
                df = pd.concat([df, df_data])

        return df

    def remove_training_data(self, id: str, **kwargs) -> bool:
        if id.endswith("-sql"):
            self.Index.delete(ids=[id], namespace=self.sql_namespace)
            return True
        elif id.endswith("-ddl"):
            self.Index.delete(ids=[id], namespace=self.ddl_namespace)
            return True
        elif id.endswith("-doc"):
            self.Index.delete(ids=[id], namespace=self.documentation_namespace)
            return True
        else:
            return False

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embedding_model = TextEmbedding(model_name=self.fastembed_model)
        embedding = next(embedding_model.embed(data))
        return embedding.tolist()


================================================
FILE: src/vanna/legacy/qdrant/__init__.py
================================================
from .qdrant import Qdrant_VectorStore

__all__ = ["Qdrant_VectorStore"]


================================================
FILE: src/vanna/legacy/qdrant/qdrant.py
================================================
from functools import cached_property
from typing import List, Tuple

import pandas as pd
from qdrant_client import QdrantClient, grpc, models

from ..base import VannaBase
from ..utils import deterministic_uuid

SCROLL_SIZE = 1000


class Qdrant_VectorStore(VannaBase):
    """
    Vectorstore implementation using Qdrant - https://qdrant.tech/

    Args:
        - config (dict, optional): Dictionary of `Qdrant_VectorStore config` options. Defaults to `{}`.
            - client: A `qdrant_client.QdrantClient` instance. Overrides other config options.
            - location: If `":memory:"` - use in-memory Qdrant instance. If `str` - use it as a `url` parameter.
            - url: Either host or str of "Optional[scheme], host, Optional[port], Optional[prefix]". Eg. `"http://localhost:6333"`.
            - prefer_grpc: If `true` - use gPRC interface whenever possible in custom methods.
            - https: If `true` - use HTTPS(SSL) protocol. Default: `None`
            - api_key: API key for authentication in Qdrant Cloud. Default: `None`
            - timeout: Timeout for REST and gRPC API requests. Defaults to 5 seconds for REST and unlimited for gRPC.
            - path: Persistence path for QdrantLocal. Default: `None`.
            - prefix: Prefix to the REST URL paths. Example: `service/v1` will result in `http://localhost:6333/service/v1/{qdrant-endpoint}`.
            - n_results: Number of results to return from similarity search. Defaults to 10.
            - fastembed_model: [Model](https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models) to use for `fastembed.TextEmbedding`.
              Defaults to `"BAAI/bge-small-en-v1.5"`.
            - collection_params: Additional parameters to pass to `qdrant_client.QdrantClient#create_collection()` method.
            - distance_metric: Distance metric to use when creating collections. Defaults to `qdrant_client.models.Distance.COSINE`.
            - documentation_collection_name: Name of the collection to store documentation. Defaults to `"documentation"`.
            - ddl_collection_name: Name of the collection to store DDL. Defaults to `"ddl"`.
            - sql_collection_name: Name of the collection to store SQL. Defaults to `"sql"`.

    Raises:
        TypeError: If config["client"] is not a `qdrant_client.QdrantClient` instance
    """

    def __init__(
        self,
        config={},
    ):
        VannaBase.__init__(self, config=config)
        client = config.get("client")

        if client is None:
            self._client = QdrantClient(
                location=config.get("location", None),
                url=config.get("url", None),
                prefer_grpc=config.get("prefer_grpc", False),
                https=config.get("https", None),
                api_key=config.get("api_key", None),
                timeout=config.get("timeout", None),
                path=config.get("path", None),
                prefix=config.get("prefix", None),
            )
        elif not isinstance(client, QdrantClient):
            raise TypeError(
                f"Unsupported client of type {client.__class__} was set in config"
            )

        else:
            self._client = client

        self.n_results = config.get("n_results", 10)
        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
        self.collection_params = config.get("collection_params", {})
        self.distance_metric = config.get("distance_metric", models.Distance.COSINE)
        self.documentation_collection_name = config.get(
            "documentation_collection_name", "documentation"
        )
        self.ddl_collection_name = config.get("ddl_collection_name", "ddl")
        self.sql_collection_name = config.get("sql_collection_name", "sql")

        self.id_suffixes = {
            self.ddl_collection_name: "ddl",
            self.documentation_collection_name: "doc",
            self.sql_collection_name: "sql",
        }

        self._setup_collections()

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        question_answer = "Question: {0}\n\nSQL: {1}".format(question, sql)
        id = deterministic_uuid(question_answer)

        self._client.upsert(
            self.sql_collection_name,
            points=[
                models.PointStruct(
                    id=id,
                    vector=self.generate_embedding(question_answer),
                    payload={
                        "question": question,
                        "sql": sql,
                    },
                )
            ],
        )

        return self._format_point_id(id, self.sql_collection_name)

    def add_ddl(self, ddl: str, **kwargs) -> str:
        id = deterministic_uuid(ddl)
        self._client.upsert(
            self.ddl_collection_name,
            points=[
                models.PointStruct(
                    id=id,
                    vector=self.generate_embedding(ddl),
                    payload={
                        "ddl": ddl,
                    },
                )
            ],
        )
        return self._format_point_id(id, self.ddl_collection_name)

    def add_documentation(self, documentation: str, **kwargs) -> str:
        id = deterministic_uuid(documentation)

        self._client.upsert(
            self.documentation_collection_name,
            points=[
                models.PointStruct(
                    id=id,
                    vector=self.generate_embedding(documentation),
                    payload={
                        "documentation": documentation,
                    },
                )
            ],
        )

        return self._format_point_id(id, self.documentation_collection_name)

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        df = pd.DataFrame()

        if sql_data := self._get_all_points(self.sql_collection_name):
            question_list = [data.payload["question"] for data in sql_data]
            sql_list = [data.payload["sql"] for data in sql_data]
            id_list = [
                self._format_point_id(data.id, self.sql_collection_name)
                for data in sql_data
            ]

            df_sql = pd.DataFrame(
                {
                    "id": id_list,
                    "question": question_list,
                    "content": sql_list,
                }
            )

            df_sql["training_data_type"] = "sql"

            df = pd.concat([df, df_sql])

        if ddl_data := self._get_all_points(self.ddl_collection_name):
            ddl_list = [data.payload["ddl"] for data in ddl_data]
            id_list = [
                self._format_point_id(data.id, self.ddl_collection_name)
                for data in ddl_data
            ]

            df_ddl = pd.DataFrame(
                {
                    "id": id_list,
                    "question": [None for _ in ddl_list],
                    "content": ddl_list,
                }
            )

            df_ddl["training_data_type"] = "ddl"

            df = pd.concat([df, df_ddl])

        if doc_data := self._get_all_points(self.documentation_collection_name):
            document_list = [data.payload["documentation"] for data in doc_data]
            id_list = [
                self._format_point_id(data.id, self.documentation_collection_name)
                for data in doc_data
            ]

            df_doc = pd.DataFrame(
                {
                    "id": id_list,
                    "question": [None for _ in document_list],
                    "content": document_list,
                }
            )

            df_doc["training_data_type"] = "documentation"

            df = pd.concat([df, df_doc])

        return df

    def remove_training_data(self, id: str, **kwargs) -> bool:
        try:
            id, collection_name = self._parse_point_id(id)
            res = self._client.delete(collection_name, points_selector=[id])
            return True
        except ValueError:
            return False

    def remove_collection(self, collection_name: str) -> bool:
        """
        This function can reset the collection to empty state.

        Args:
            collection_name (str): sql or ddl or documentation

        Returns:
            bool: True if collection is deleted, False otherwise
        """
        if collection_name in self.id_suffixes.keys():
            self._client.delete_collection(collection_name)
            self._setup_collections()
            return True
        else:
            return False

    @cached_property
    def embeddings_dimension(self):
        return len(self.generate_embedding("ABCDEF"))

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        results = self._client.query_points(
            self.sql_collection_name,
            query=self.generate_embedding(question),
            limit=self.n_results,
            with_payload=True,
        ).points

        return [dict(result.payload) for result in results]

    def get_related_ddl(self, question: str, **kwargs) -> list:
        results = self._client.query_points(
            self.ddl_collection_name,
            query=self.generate_embedding(question),
            limit=self.n_results,
            with_payload=True,
        ).points

        return [result.payload["ddl"] for result in results]

    def get_related_documentation(self, question: str, **kwargs) -> list:
        results = self._client.query_points(
            self.documentation_collection_name,
            query=self.generate_embedding(question),
            limit=self.n_results,
            with_payload=True,
        ).points

        return [result.payload["documentation"] for result in results]

    def generate_embedding(self, data: str, **kwargs) -> List[float]:
        embedding_model = self._client._get_or_init_model(
            model_name=self.fastembed_model
        )
        embedding = next(embedding_model.embed(data))

        return embedding.tolist()

    def _get_all_points(self, collection_name: str):
        results: List[models.Record] = []
        next_offset = None
        stop_scrolling = False
        while not stop_scrolling:
            records, next_offset = self._client.scroll(
                collection_name,
                limit=SCROLL_SIZE,
                offset=next_offset,
                with_payload=True,
                with_vectors=False,
            )
            stop_scrolling = next_offset is None or (
                isinstance(next_offset, grpc.PointId)
                and next_offset.num == 0
                and next_offset.uuid == ""
            )

            results.extend(records)

        return results

    def _setup_collections(self):
        if not self._client.collection_exists(self.sql_collection_name):
            self._client.create_collection(
                collection_name=self.sql_collection_name,
                vectors_config=models.VectorParams(
                    size=self.embeddings_dimension,
                    distance=self.distance_metric,
                ),
                **self.collection_params,
            )

        if not self._client.collection_exists(self.ddl_collection_name):
            self._client.create_collection(
                collection_name=self.ddl_collection_name,
                vectors_config=models.VectorParams(
                    size=self.embeddings_dimension,
                    distance=self.distance_metric,
                ),
                **self.collection_params,
            )
        if not self._client.collection_exists(self.documentation_collection_name):
            self._client.create_collection(
                collection_name=self.documentation_collection_name,
                vectors_config=models.VectorParams(
                    size=self.embeddings_dimension,
                    distance=self.distance_metric,
                ),
                **self.collection_params,
            )

    def _format_point_id(self, id: str, collection_name: str) -> str:
        return "{0}-{1}".format(id, self.id_suffixes[collection_name])

    def _parse_point_id(self, id: str) -> Tuple[str, str]:
        id, curr_suffix = id.rsplit("-", 1)
        for collection_name, suffix in self.id_suffixes.items():
            if curr_suffix == suffix:
                return id, collection_name
        raise ValueError(f"Invalid id {id}")


================================================
FILE: src/vanna/legacy/qianfan/Qianfan_Chat.py
================================================
import qianfan

from ..base import VannaBase


class Qianfan_Chat(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        if "api_key" not in config:
            raise Exception("Missing api_key in config")
        self.api_key = config["api_key"]

        if "secret_key" not in config:
            raise Exception("Missing secret_key in config")
        self.secret_key = config["secret_key"]

        # default parameters - can be overrided using config
        self.temperature = 0.9
        self.max_tokens = 1024

        if "temperature" in config:
            self.temperature = config["temperature"]

        if "max_tokens" in config:
            self.max_tokens = config["max_tokens"]

        self.model = config["model"] if "model" in config else "ERNIE-Speed"

        if client is not None:
            self.client = client
            return

        self.client = qianfan.ChatCompletion(ak=self.api_key, sk=self.secret_key)

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def get_sql_prompt(
        self,
        initial_prompt: str,
        question: str,
        question_sql_list: list,
        ddl_list: list,
        doc_list: list,
        **kwargs,
    ):
        """
        Example:
        ```python
        vn.get_sql_prompt(
            question="What are the top 10 customers by sales?",
            question_sql_list=[{"question": "What are the top 10 customers by sales?", "sql": "SELECT * FROM customers ORDER BY sales DESC LIMIT 10"}],
            ddl_list=["CREATE TABLE customers (id INT, name TEXT, sales DECIMAL)"],
            doc_list=["The customers table contains information about customers and their sales."],
        )

        ```

        This method is used to generate a prompt for the LLM to generate SQL.

        Args:
            question (str): The question to generate SQL for.
            question_sql_list (list): A list of questions and their corresponding SQL statements.
            ddl_list (list): A list of DDL statements.
            doc_list (list): A list of documentation.

        Returns:
            any: The prompt for the LLM to generate SQL.
        """

        if initial_prompt is None:
            initial_prompt = (
                f"You are a {self.dialect} expert. "
                + "Please help to generate a SQL to answer the question based on some context.Please don't give any explanation for your answer. Just only generate a SQL \n"
            )

        initial_prompt = self.add_ddl_to_prompt(
            initial_prompt, ddl_list, max_tokens=self.max_tokens
        )

        if self.static_documentation != "":
            doc_list.append(self.static_documentation)

        initial_prompt = self.add_documentation_to_prompt(
            initial_prompt, doc_list, max_tokens=self.max_tokens
        )
        message_log = []

        if question_sql_list is None or len(question_sql_list) == 0:
            initial_prompt = initial_prompt + f"question: {question}"
            message_log.append(self.user_message(initial_prompt))
        else:
            for i, example in question_sql_list:
                if example is None:
                    print("example is None")
                else:
                    if (
                        example is not None
                        and "question" in example
                        and "sql" in example
                    ):
                        if i == 0:
                            initial_prompt = (
                                initial_prompt + f"question: {example['question']}"
                            )
                            message_log.append(self.user_message(initial_prompt))
                        else:
                            message_log.append(self.user_message(example["question"]))
                        message_log.append(self.assistant_message(example["sql"]))

            message_log.append(self.user_message(question))
        return message_log

    def submit_prompt(self, prompt, **kwargs) -> str:
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        # Count the number of tokens in the message log
        # Use 4 as an approximation for the number of characters per token
        num_tokens = 0
        for message in prompt:
            num_tokens += len(message["content"]) / 4

        if kwargs.get("model", None) is not None:
            model = kwargs.get("model", None)
            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.do(
                model=self.model,
                messages=prompt,
                max_output_tokens=self.max_tokens,
                stop=None,
                temperature=self.temperature,
            )
        elif self.config is not None and "model" in self.config:
            print(
                f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
            )
            response = self.client.do(
                model=self.config.get("model"),
                messages=prompt,
                max_output_tokens=self.max_tokens,
                stop=None,
                temperature=self.temperature,
            )
        else:
            if num_tokens > 3500:
                model = "ERNIE-Speed-128K"
            else:
                model = "ERNIE-Speed-8K"

            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.do(
                model=model,
                messages=prompt,
                max_output_tokens=self.max_tokens,
                stop=None,
                temperature=self.temperature,
            )

        return response.body.get("result")


================================================
FILE: src/vanna/legacy/qianfan/Qianfan_embeddings.py
================================================
import qianfan

from ..base import VannaBase


class Qianfan_Embeddings(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        if client is not None:
            self.client = client
            return

        if "api_key" not in config:
            raise Exception("Missing api_key in config")
        self.api_key = config["api_key"]

        if "secret_key" not in config:
            raise Exception("Missing secret_key in config")
        self.secret_key = config["secret_key"]

        self.client = qianfan.Embedding(ak=self.api_key, sk=self.secret_key)

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        if self.config is not None and "model" in self.config:
            embedding = self.client.do(
                model=self.config["model"],
                input=[data],
            )
        else:
            embedding = self.client.do(
                model="bge-large-zh",
                input=[data],
            )

        return embedding.get("data")[0]["embedding"]


================================================
FILE: src/vanna/legacy/qianfan/__init__.py
================================================
from .Qianfan_Chat import Qianfan_Chat
from .Qianfan_embeddings import Qianfan_Embeddings


================================================
FILE: src/vanna/legacy/qianwen/QianwenAI_chat.py
================================================
import os

from openai import OpenAI

from ..base import VannaBase


class QianWenAI_Chat(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        # default parameters - can be overrided using config
        self.temperature = 0.7

        if "temperature" in config:
            self.temperature = config["temperature"]

        if "api_type" in config:
            raise Exception(
                "Passing api_type is now deprecated. Please pass an OpenAI client instead."
            )

        if "api_base" in config:
            raise Exception(
                "Passing api_base is now deprecated. Please pass an OpenAI client instead."
            )

        if "api_version" in config:
            raise Exception(
                "Passing api_version is now deprecated. Please pass an OpenAI client instead."
            )

        if client is not None:
            self.client = client
            return

        if config is None and client is None:
            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
            return

        if "api_key" in config:
            if "base_url" not in config:
                self.client = OpenAI(
                    api_key=config["api_key"],
                    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                )
            else:
                self.client = OpenAI(
                    api_key=config["api_key"], base_url=config["base_url"]
                )

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        # Count the number of tokens in the message log
        # Use 4 as an approximation for the number of characters per token
        num_tokens = 0
        for message in prompt:
            num_tokens += len(message["content"]) / 4

        if kwargs.get("model", None) is not None:
            model = kwargs.get("model", None)
            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        elif kwargs.get("engine", None) is not None:
            engine = kwargs.get("engine", None)
            print(f"Using model {engine} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                engine=engine,
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        elif self.config is not None and "engine" in self.config:
            print(
                f"Using engine {self.config['engine']} for {num_tokens} tokens (approx)"
            )
            response = self.client.chat.completions.create(
                engine=self.config["engine"],
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        elif self.config is not None and "model" in self.config:
            print(
                f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
            )
            response = self.client.chat.completions.create(
                model=self.config["model"],
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )
        else:
            if num_tokens > 3500:
                model = "qwen-long"
            else:
                model = "qwen-plus"

            print(f"Using model {model} for {num_tokens} tokens (approx)")
            response = self.client.chat.completions.create(
                model=model,
                messages=prompt,
                stop=None,
                temperature=self.temperature,
            )

        # Find the first response from the chatbot that has text in it (some responses may not have text)
        for choice in response.choices:
            if "text" in choice:
                return choice.text

        # If no response with text is found, return the first response's content (which may be empty)
        return response.choices[0].message.content


================================================
FILE: src/vanna/legacy/qianwen/QianwenAI_embeddings.py
================================================
from openai import OpenAI

from ..base import VannaBase


class QianWenAI_Embeddings(VannaBase):
    def __init__(self, client=None, config=None):
        VannaBase.__init__(self, config=config)

        if client is not None:
            self.client = client
            return

        if self.client is not None:
            return

        self.client = OpenAI()

        if config is None:
            return

        if "api_type" in config:
            self.client.api_type = config["api_type"]

        if "api_base" in config:
            self.client.api_base = config["api_base"]

        if "api_version" in config:
            self.client.api_version = config["api_version"]

        if "api_key" in config:
            self.client.api_key = config["api_key"]

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        if self.config is not None and "engine" in self.config:
            embedding = self.client.embeddings.create(
                engine=self.config["engine"],
                input=data,
            )
        else:
            embedding = self.client.embeddings.create(
                model="bge-large-zh",
                input=data,
            )

        return embedding.get("data")[0]["embedding"]


================================================
FILE: src/vanna/legacy/qianwen/__init__.py
================================================
from .QianwenAI_chat import QianWenAI_Chat
from .QianwenAI_embeddings import QianWenAI_Embeddings


================================================
FILE: src/vanna/legacy/remote.py
================================================
import dataclasses
import json
from io import StringIO
from typing import Callable, List, Tuple, Union

import pandas as pd
import requests

from .base import VannaBase
from .types import (
    AccuracyStats,
    ApiKey,
    DataFrameJSON,
    DataResult,
    Explanation,
    FullQuestionDocument,
    NewOrganization,
    NewOrganizationMember,
    Organization,
    OrganizationList,
    PlotlyResult,
    Question,
    QuestionCategory,
    QuestionId,
    QuestionList,
    QuestionSQLPair,
    QuestionStringList,
    SQLAnswer,
    Status,
    StatusWithId,
    StringData,
    TrainingData,
    UserEmail,
    UserOTP,
    Visibility,
)
from .vannadb import VannaDB_VectorStore


class VannaDefault(VannaDB_VectorStore):
    def __init__(self, model: str, api_key: str, config=None):
        VannaBase.__init__(self, config=config)
        VannaDB_VectorStore.__init__(
            self, vanna_model=model, vanna_api_key=api_key, config=config
        )

        self._model = model
        self._api_key = api_key

        self._endpoint = (
            "https://ask.vanna.ai/rpc"
            if config is None or "endpoint" not in config
            else config["endpoint"]
        )

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        # JSON-ify the prompt
        json_prompt = json.dumps(prompt, ensure_ascii=False)

        params = [StringData(data=json_prompt)]

        d = self._rpc_call(method="submit_prompt", params=params)

        if "result" not in d:
            return None

        # Load the result into a dataclass
        results = StringData(**d["result"])

        return results.data


================================================
FILE: src/vanna/legacy/types/__init__.py
================================================
from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Union


@dataclass
class Status:
    success: bool
    message: str


@dataclass
class StatusWithId:
    success: bool
    message: str
    id: str


@dataclass
class QuestionList:
    questions: List[FullQuestionDocument]


@dataclass
class FullQuestionDocument:
    id: QuestionId
    question: Question
    answer: SQLAnswer | None
    data: DataResult | None
    plotly: PlotlyResult | None


@dataclass
class QuestionSQLPair:
    question: str
    sql: str
    tag: Union[str, None]


@dataclass
class Organization:
    name: str
    user: str | None
    connection: Connection | None


@dataclass
class OrganizationList:
    organizations: List[str]


@dataclass
class QuestionStringList:
    questions: List[str]


@dataclass
class Visibility:
    visibility: bool


@dataclass
class UserEmail:
    email: str


@dataclass
class NewOrganization:
    org_name: str
    db_type: str


@dataclass
class NewOrganizationMember:
    org_name: str
    email: str
    is_admin: bool


@dataclass
class UserOTP:
    email: str
    otp: str


@dataclass
class ApiKey:
    key: str


@dataclass
class QuestionId:
    id: str


@dataclass
class Question:
    question: str


@dataclass
class QuestionCategory:
    question: str
    category: str

    NO_SQL_GENERATED = "No SQL Generated"
    SQL_UNABLE_TO_RUN = "SQL Unable to Run"
    BOOTSTRAP_TRAINING_QUERY = "Bootstrap Training Query"
    SQL_RAN = "SQL Ran Successfully"
    FLAGGED_FOR_REVIEW = "Flagged for Review"
    REVIEWED_AND_APPROVED = "Reviewed and Approved"
    REVIEWED_AND_REJECTED = "Reviewed and Rejected"
    REVIEWED_AND_UPDATED = "Reviewed and Updated"


@dataclass
class AccuracyStats:
    num_questions: int
    data: Dict[str, int]


@dataclass
class Followup:
    followup: str


@dataclass
class QuestionEmbedding:
    question: Question
    embedding: List[float]


@dataclass
class Connection:
    # TODO: implement
    pass


@dataclass
class SQLAnswer:
    raw_answer: str
    prefix: str
    postfix: str
    sql: str


@dataclass
class Explanation:
    explanation: str


@dataclass
class DataResult:
    question: str | None
    sql: str | None
    table_markdown: str
    error: str | None
    correction_attempts: int


@dataclass
class PlotlyResult:
    plotly_code: str


@dataclass
class WarehouseDefinition:
    name: str
    tables: List[TableDefinition]


@dataclass
class TableDefinition:
    schema_name: str
    table_name: str
    ddl: str | None
    columns: List[ColumnDefinition]


@dataclass
class ColumnDefinition:
    name: str
    type: str
    is_primary_key: bool
    is_foreign_key: bool
    foreign_key_table: str
    foreign_key_column: str


@dataclass
class Diagram:
    raw: str
    mermaid_code: str


@dataclass
class StringData:
    data: str


@dataclass
class DataFrameJSON:
    data: str


@dataclass
class TrainingData:
    questions: List[dict]
    ddl: List[str]
    documentation: List[str]


@dataclass
class TrainingPlanItem:
    item_type: str
    item_group: str
    item_name: str
    item_value: str

    def __str__(self):
        if self.item_type == self.ITEM_TYPE_SQL:
            return f"Train on SQL: {self.item_group} {self.item_name}"
        elif self.item_type == self.ITEM_TYPE_DDL:
            return f"Train on DDL: {self.item_group} {self.item_name}"
        elif self.item_type == self.ITEM_TYPE_IS:
            return f"Train on Information Schema: {self.item_group} {self.item_name}"

    ITEM_TYPE_SQL = "sql"
    ITEM_TYPE_DDL = "ddl"
    ITEM_TYPE_IS = "is"


class TrainingPlan:
    """
    A class representing a training plan. You can see what's in it, and remove items from it that you don't want trained.

    **Example:**
    ```python
    plan = vn.get_training_plan()

    plan.get_summary()
    ```

    """

    _plan: List[TrainingPlanItem]

    def __init__(self, plan: List[TrainingPlanItem]):
        self._plan = plan

    def __str__(self):
        return "\n".join(self.get_summary())

    def __repr__(self):
        return self.__str__()

    def get_summary(self) -> List[str]:
        """
        **Example:**
        ```python
        plan = vn.get_training_plan()

        plan.get_summary()
        ```

        Get a summary of the training plan.

        Returns:
            List[str]: A list of strings describing the training plan.
        """

        return [f"{item}" for item in self._plan]

    def remove_item(self, item: str):
        """
        **Example:**
        ```python
        plan = vn.get_training_plan()

        plan.remove_item("Train on SQL: What is the average salary of employees?")
        ```

        Remove an item from the training plan.

        Args:
            item (str): The item to remove.
        """
        for plan_item in self._plan:
            if str(plan_item) == item:
                self._plan.remove(plan_item)
                break


================================================
FILE: src/vanna/legacy/utils.py
================================================
import hashlib
import os
import re
import uuid
from typing import Union

from .exceptions import ImproperlyConfigured, ValidationError


def validate_config_path(path):
    if not os.path.exists(path):
        raise ImproperlyConfigured(f"No such configuration file: {path}")

    if not os.path.isfile(path):
        raise ImproperlyConfigured(f"Config should be a file: {path}")

    if not os.access(path, os.R_OK):
        raise ImproperlyConfigured(
            f"Cannot read the config file. Please grant read privileges: {path}"
        )


def sanitize_model_name(model_name):
    try:
        model_name = model_name.lower()

        # Replace spaces with a hyphen
        model_name = model_name.replace(" ", "-")

        if "-" in model_name:
            # remove double hyphones
            model_name = re.sub(r"-+", "-", model_name)
            if "_" in model_name:
                # If name contains both underscores and hyphen replace all underscores with hyphens
                model_name = re.sub(r"_", "-", model_name)

        # Remove special characters only allow underscore
        model_name = re.sub(r"[^a-zA-Z0-9-_]", "", model_name)

        # Remove hyphen or underscore if any at the last or first
        if model_name[-1] in ("-", "_"):
            model_name = model_name[:-1]
        if model_name[0] in ("-", "_"):
            model_name = model_name[1:]

        return model_name
    except Exception as e:
        raise ValidationError(e)


def deterministic_uuid(content: Union[str, bytes]) -> str:
    """Creates deterministic UUID on hash value of string or byte content.

    Args:
        content: String or byte representation of data.

    Returns:
        UUID of the content.
    """
    if isinstance(content, str):
        content_bytes = content.encode("utf-8")
    elif isinstance(content, bytes):
        content_bytes = content
    else:
        raise ValueError(f"Content type {type(content)} not supported !")

    hash_object = hashlib.sha256(content_bytes)
    hash_hex = hash_object.hexdigest()
    namespace = uuid.UUID("00000000-0000-0000-0000-000000000000")
    content_uuid = str(uuid.uuid5(namespace, hash_hex))

    return content_uuid


================================================
FILE: src/vanna/legacy/vannadb/__init__.py
================================================
from .vannadb_vector import VannaDB_VectorStore


================================================
FILE: src/vanna/legacy/vannadb/vannadb_vector.py
================================================
import dataclasses
import json
from io import StringIO

import pandas as pd
import requests

from ..advanced import VannaAdvanced
from ..base import VannaBase
from ..types import (
    DataFrameJSON,
    NewOrganization,
    OrganizationList,
    Question,
    QuestionSQLPair,
    Status,
    StatusWithId,
    StringData,
    TrainingData,
)
from ..utils import sanitize_model_name


class VannaDB_VectorStore(VannaBase, VannaAdvanced):
    def __init__(self, vanna_model: str, vanna_api_key: str, config=None):
        VannaBase.__init__(self, config=config)

        self._model = vanna_model
        self._api_key = vanna_api_key

        self._endpoint = (
            "https://ask.vanna.ai/rpc"
            if config is None or "endpoint" not in config
            else config["endpoint"]
        )
        self.related_training_data = {}
        self._graphql_endpoint = "https://functionrag.com/query"
        self._graphql_headers = {
            "Content-Type": "application/json",
            "API-KEY": self._api_key,
            "NAMESPACE": self._model,
        }

    def _rpc_call(self, method, params):
        if method != "list_orgs":
            headers = {
                "Content-Type": "application/json",
                "Vanna-Key": self._api_key,
                "Vanna-Org": self._model,
            }
        else:
            headers = {
                "Content-Type": "application/json",
                "Vanna-Key": self._api_key,
                "Vanna-Org": "demo-tpc-h",
            }

        data = {
            "method": method,
            "params": [self._dataclass_to_dict(obj) for obj in params],
        }

        response = requests.post(self._endpoint, headers=headers, data=json.dumps(data))
        return response.json()

    def _dataclass_to_dict(self, obj):
        return dataclasses.asdict(obj)

    def get_all_functions(self) -> list:
        query = """
            {
                get_all_sql_functions {
                    function_name
                    description
                    post_processing_code_template
                    arguments {
                        name
                        description
                        general_type
                        is_user_editable
                        available_values
                    }
                    sql_template
                }
            }
        """

        response = requests.post(
            self._graphql_endpoint, headers=self._graphql_headers, json={"query": query}
        )
        response_json = response.json()
        if (
            response.status_code == 200
            and "data" in response_json
            and "get_all_sql_functions" in response_json["data"]
        ):
            self.log(response_json["data"]["get_all_sql_functions"])
            resp = response_json["data"]["get_all_sql_functions"]

            print(resp)

            return resp
        else:
            raise Exception(
                f"Query failed to run by returning code of {response.status_code}. {response.text}"
            )

    def get_function(self, question: str, additional_data: dict = {}) -> dict:
        query = """
        query GetFunction($question: String!, $staticFunctionArguments: [StaticFunctionArgument]) {
            get_and_instantiate_function(question: $question, static_function_arguments: $staticFunctionArguments) {
                ... on SQLFunction {
                function_name
                description
                post_processing_code_template
                instantiated_post_processing_code
                arguments {
                    name
                    description
                    general_type
                    is_user_editable
                    instantiated_value
                    available_values
                }
                sql_template
                instantiated_sql
            }
            }
        }
        """
        static_function_arguments = [
            {"name": key, "value": str(value)} for key, value in additional_data.items()
        ]
        variables = {
            "question": question,
            "staticFunctionArguments": static_function_arguments,
        }
        response = requests.post(
            self._graphql_endpoint,
            headers=self._graphql_headers,
            json={"query": query, "variables": variables},
        )
        response_json = response.json()
        if (
            response.status_code == 200
            and "data" in response_json
            and "get_and_instantiate_function" in response_json["data"]
        ):
            self.log(response_json["data"]["get_and_instantiate_function"])
            resp = response_json["data"]["get_and_instantiate_function"]

            print(resp)

            return resp
        else:
            raise Exception(
                f"Query failed to run by returning code of {response.status_code}. {response.text}"
            )

    def create_function(
        self, question: str, sql: str, plotly_code: str, **kwargs
    ) -> dict:
        query = """
        mutation CreateFunction($question: String!, $sql: String!, $plotly_code: String!) {
            generate_and_create_sql_function(question: $question, sql: $sql, post_processing_code: $plotly_code) {
                function_name
                description
                arguments {
                    name
                    description
                    general_type
                    is_user_editable
                }
                sql_template
                post_processing_code_template
            }
        }
        """
        variables = {"question": question, "sql": sql, "plotly_code": plotly_code}
        response = requests.post(
            self._graphql_endpoint,
            headers=self._graphql_headers,
            json={"query": query, "variables": variables},
        )
        response_json = response.json()
        if (
            response.status_code == 200
            and "data" in response_json
            and response_json["data"] is not None
            and "generate_and_create_sql_function" in response_json["data"]
        ):
            resp = response_json["data"]["generate_and_create_sql_function"]

            print(resp)

            return resp
        else:
            raise Exception(
                f"Query failed to run by returning code of {response.status_code}. {response.text}"
            )

    def update_function(self, old_function_name: str, updated_function: dict) -> bool:
        """
        Update an existing SQL function based on the provided parameters.

        Args:
            old_function_name (str): The current name of the function to be updated.
            updated_function (dict): A dictionary containing the updated function details. Expected keys:
                - 'function_name': The new name of the function.
                - 'description': The new description of the function.
                - 'arguments': A list of dictionaries describing the function arguments.
                - 'sql_template': The new SQL template for the function.
                - 'post_processing_code_template': The new post-processing code template.

        Returns:
            bool: True if the function was successfully updated, False otherwise.
        """
        mutation = """
        mutation UpdateSQLFunction($input: SQLFunctionUpdate!) {
            update_sql_function(input: $input)
        }
        """

        SQLFunctionUpdate = {
            "function_name",
            "description",
            "arguments",
            "sql_template",
            "post_processing_code_template",
        }

        # Define the expected keys for each argument in the arguments list
        ArgumentKeys = {
            "name",
            "general_type",
            "description",
            "is_user_editable",
            "available_values",
        }

        # Function to validate and transform arguments
        def validate_arguments(args):
            return [
                {key: arg[key] for key in arg if key in ArgumentKeys} for arg in args
            ]

        # Keep only the keys that conform to the SQLFunctionUpdate GraphQL input type
        updated_function = {
            key: value
            for key, value in updated_function.items()
            if key in SQLFunctionUpdate
        }

        # Special handling for 'arguments' to ensure they conform to the spec
        if "arguments" in updated_function:
            updated_function["arguments"] = validate_arguments(
                updated_function["arguments"]
            )

        variables = {
            "input": {"old_function_name": old_function_name, **updated_function}
        }

        print("variables", variables)

        response = requests.post(
            self._graphql_endpoint,
            headers=self._graphql_headers,
            json={"query": mutation, "variables": variables},
        )
        response_json = response.json()
        if (
            response.status_code == 200
            and "data" in response_json
            and response_json["data"] is not None
            and "update_sql_function" in response_json["data"]
        ):
            return response_json["data"]["update_sql_function"]
        else:
            raise Exception(
                f"Mutation failed to run by returning code of {response.status_code}. {response.text}"
            )

    def delete_function(self, function_name: str) -> bool:
        mutation = """
        mutation DeleteSQLFunction($function_name: String!) {
            delete_sql_function(function_name: $function_name)
        }
        """
        variables = {"function_name": function_name}
        response = requests.post(
            self._graphql_endpoint,
            headers=self._graphql_headers,
            json={"query": mutation, "variables": variables},
        )
        response_json = response.json()
        if (
            response.status_code == 200
            and "data" in response_json
            and response_json["data"] is not None
            and "delete_sql_function" in response_json["data"]
        ):
            return response_json["data"]["delete_sql_function"]
        else:
            raise Exception(
                f"Mutation failed to run by returning code of {response.status_code}. {response.text}"
            )

    def create_model(self, model: str, **kwargs) -> bool:
        """
        **Example:**
        ```python
        success = vn.create_model("my_model")
        ```
        Create a new model.

        Args:
            model (str): The name of the model to create.

        Returns:
            bool: True if the model was created, False otherwise.
        """
        model = sanitize_model_name(model)
        params = [NewOrganization(org_name=model, db_type="")]

        d = self._rpc_call(method="create_org", params=params)

        if "result" not in d:
            return False

        status = Status(**d["result"])

        return status.success

    def get_models(self) -> list:
        """
        **Example:**
        ```python
        models = vn.get_models()
        ```

        List the models that belong to the user.

        Returns:
            List[str]: A list of model names.
        """
        d = self._rpc_call(method="list_my_models", params=[])

        if "result" not in d:
            return []

        orgs = OrganizationList(**d["result"])

        return orgs.organizations

    def generate_embedding(self, data: str, **kwargs) -> list[float]:
        # This is done server-side
        pass

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        if "tag" in kwargs:
            tag = kwargs["tag"]
        else:
            tag = "Manually Trained"

        params = [QuestionSQLPair(question=question, sql=sql, tag=tag)]

        d = self._rpc_call(method="add_sql", params=params)

        if "result" not in d:
            raise Exception("Error adding question and SQL pair", d)

        status = StatusWithId(**d["result"])

        return status.id

    def add_ddl(self, ddl: str, **kwargs) -> str:
        params = [StringData(data=ddl)]

        d = self._rpc_call(method="add_ddl", params=params)

        if "result" not in d:
            raise Exception("Error adding DDL", d)

        status = StatusWithId(**d["result"])

        return status.id

    def add_documentation(self, documentation: str, **kwargs) -> str:
        params = [StringData(data=documentation)]

        d = self._rpc_call(method="add_documentation", params=params)

        if "result" not in d:
            raise Exception("Error adding documentation", d)

        status = StatusWithId(**d["result"])

        return status.id

    def get_training_data(self, **kwargs) -> pd.DataFrame:
        params = []

        d = self._rpc_call(method="get_training_data", params=params)

        if "result" not in d:
            return None

        # Load the result into a dataclass
        training_data = DataFrameJSON(**d["result"])

        df = pd.read_json(StringIO(training_data.data))

        return df

    def remove_training_data(self, id: str, **kwargs) -> bool:
        params = [StringData(data=id)]

        d = self._rpc_call(method="remove_training_data", params=params)

        if "result" not in d:
            raise Exception("Error removing training data")

        status = Status(**d["result"])

        if not status.success:
            raise Exception(f"Error removing training data: {status.message}")

        return status.success

    def get_related_training_data_cached(self, question: str) -> TrainingData:
        params = [Question(question=question)]

        d = self._rpc_call(method="get_related_training_data", params=params)

        if "result" not in d:
            return None

        # Load the result into a dataclass
        training_data = TrainingData(**d["result"])

        self.related_training_data[question] = training_data

        return training_data

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        if question in self.related_training_data:
            training_data = self.related_training_data[question]
        else:
            training_data = self.get_related_training_data_cached(question)

        return training_data.questions

    def get_related_ddl(self, question: str, **kwargs) -> list:
        if question in self.related_training_data:
            training_data = self.related_training_data[question]
        else:
            training_data = self.get_related_training_data_cached(question)

        return training_data.ddl

    def get_related_documentation(self, question: str, **kwargs) -> list:
        if question in self.related_training_data:
            training_data = self.related_training_data[question]
        else:
            training_data = self.get_related_training_data_cached(question)

        return training_data.documentation


================================================
FILE: src/vanna/legacy/vllm/__init__.py
================================================
from .vllm import Vllm


================================================
FILE: src/vanna/legacy/vllm/vllm.py
================================================
import re

import requests

from ..base import VannaBase


class Vllm(VannaBase):
    def __init__(self, config=None):
        if config is None or "vllm_host" not in config:
            self.host = "http://localhost:8000"
        else:
            self.host = config["vllm_host"]

        if config is None or "model" not in config:
            raise ValueError("check the config for vllm")
        else:
            self.model = config["model"]

        if "auth-key" in config:
            self.auth_key = config["auth-key"]
        else:
            self.auth_key = None

        if "temperature" in config:
            self.temperature = config["temperature"]
        else:
            # default temperature - can be overrided using config
            self.temperature = 0.7

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def extract_sql_query(self, text):
        """
        Extracts the first SQL statement after the word 'select', ignoring case,
        matches until the first semicolon, three backticks, or the end of the string,
        and removes three backticks if they exist in the extracted string.

        Args:
        - text (str): The string to search within for an SQL statement.

        Returns:
        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
        """
        # Regular expression to find 'select' (ignoring case) and capture until ';', '```', or end of string
        pattern = re.compile(r"select.*?(?:;|```|$)", re.IGNORECASE | re.DOTALL)

        match = pattern.search(text)
        if match:
            # Remove three backticks from the matched string if they exist
            return match.group(0).replace("```", "")
        else:
            return text

    def generate_sql(self, question: str, **kwargs) -> str:
        # Use the super generate_sql
        sql = super().generate_sql(question, **kwargs)

        # Replace "\_" with "_"
        sql = sql.replace("\\_", "_")

        sql = sql.replace("\\", "")

        return self.extract_sql_query(sql)

    def submit_prompt(self, prompt, **kwargs) -> str:
        url = f"{self.host}/v1/chat/completions"
        data = {
            "model": self.model,
            "temperature": self.temperature,
            "stream": False,
            "messages": prompt,
        }

        if self.auth_key is not None:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.auth_key}",
            }

            response = requests.post(url, headers=headers, json=data)

        else:
            response = requests.post(url, json=data)

        response_dict = response.json()

        self.log(response.text)

        return response_dict["choices"][0]["message"]["content"]


================================================
FILE: src/vanna/legacy/weaviate/__init__.py
================================================
from .weaviate_vector import WeaviateDatabase


================================================
FILE: src/vanna/legacy/weaviate/weaviate_vector.py
================================================
import weaviate
import weaviate.classes as wvc
from fastembed import TextEmbedding

from vanna.base import VannaBase


class WeaviateDatabase(VannaBase):
    def __init__(self, config=None):
        """
        Initialize the VannaEnhanced class with the provided configuration.

        :param config: Dictionary containing configuration parameters.

        params:
        weaviate_url (str): Weaviate cluster URL while using weaviate cloud,
        weaviate_api_key (str): Weaviate API key while using weaviate cloud,
        weaviate_port (num): Weaviate port while using local weaviate,
        weaviate_grpc (num): Weaviate gRPC port while using local weaviate,
        fastembed_model (str): Fastembed model name for text embeddings. BAAI/bge-small-en-v1.5 by default.

        """
        super().__init__(config=config)

        if config is None:
            raise ValueError("config is required")

        self.n_results = config.get("n_results", 3)
        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
        self.weaviate_api_key = config.get("weaviate_api_key")
        self.weaviate_url = config.get("weaviate_url")
        self.weaviate_port = config.get("weaviate_port")
        self.weaviate_grpc_port = config.get("weaviate_grpc", 50051)

        if not self.weaviate_api_key and not self.weaviate_port:
            raise ValueError("Add proper credentials to connect to weaviate")

        self.weaviate_client = self._initialize_weaviate_client()
        self.embeddings = TextEmbedding(model_name=self.fastembed_model)

        self.training_data_cluster = {
            "sql": "SQLTrainingDataEntry",
            "ddl": "DDLEntry",
            "doc": "DocumentationEntry",
        }

        self._create_collections_if_not_exist()

    def _create_collections_if_not_exist(self):
        properties_dict = {
            self.training_data_cluster["ddl"]: [
                wvc.config.Property(
                    name="description", data_type=wvc.config.DataType.TEXT
                ),
            ],
            self.training_data_cluster["doc"]: [
                wvc.config.Property(
                    name="description", data_type=wvc.config.DataType.TEXT
                ),
            ],
            self.training_data_cluster["sql"]: [
                wvc.config.Property(name="sql", data_type=wvc.config.DataType.TEXT),
                wvc.config.Property(
                    name="natural_language_question", data_type=wvc.config.DataType.TEXT
                ),
            ],
        }

        for cluster, properties in properties_dict.items():
            if not self.weaviate_client.collections.exists(cluster):
                self.weaviate_client.collections.create(
                    name=cluster, properties=properties
                )

    def _initialize_weaviate_client(self):
        if self.weaviate_api_key:
            return weaviate.connect_to_wcs(
                cluster_url=self.weaviate_url,
                auth_credentials=weaviate.auth.AuthApiKey(self.weaviate_api_key),
                additional_config=weaviate.config.AdditionalConfig(timeout=(10, 300)),
                skip_init_checks=True,
            )
        else:
            return weaviate.connect_to_local(
                port=self.weaviate_port,
                grpc_port=self.weaviate_grpc_port,
                additional_config=weaviate.config.AdditionalConfig(timeout=(10, 300)),
                skip_init_checks=True,
            )

    def generate_embedding(self, data: str, **kwargs):
        embedding_model = TextEmbedding(model_name=self.fastembed_model)
        embedding = next(embedding_model.embed(data))
        return embedding.tolist()

    def _insert_data(self, cluster_key: str, data_object: dict, vector: list) -> str:
        self.weaviate_client.connect()
        response = self.weaviate_client.collections.get(
            self.training_data_cluster[cluster_key]
        ).data.insert(properties=data_object, vector=vector)
        self.weaviate_client.close()
        return response

    def add_ddl(self, ddl: str, **kwargs) -> str:
        data_object = {
            "description": ddl,
        }
        response = self._insert_data("ddl", data_object, self.generate_embedding(ddl))
        return f"{response}-ddl"

    def add_documentation(self, doc: str, **kwargs) -> str:
        data_object = {
            "description": doc,
        }
        response = self._insert_data("doc", data_object, self.generate_embedding(doc))
        return f"{response}-doc"

    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        data_object = {
            "sql": sql,
            "natural_language_question": question,
        }
        response = self._insert_data(
            "sql", data_object, self.generate_embedding(question)
        )
        return f"{response}-sql"

    def _query_collection(
        self, cluster_key: str, vector_input: list, return_properties: list
    ) -> list:
        self.weaviate_client.connect()
        collection = self.weaviate_client.collections.get(
            self.training_data_cluster[cluster_key]
        )
        response = collection.query.near_vector(
            near_vector=vector_input,
            limit=self.n_results,
            return_properties=return_properties,
        )
        response_list = [item.properties for item in response.objects]
        self.weaviate_client.close()
        return response_list

    def get_related_ddl(self, question: str, **kwargs) -> list:
        vector_input = self.generate_embedding(question)
        response_list = self._query_collection("ddl", vector_input, ["description"])
        return [item["description"] for item in response_list]

    def get_related_documentation(self, question: str, **kwargs) -> list:
        vector_input = self.generate_embedding(question)
        response_list = self._query_collection("doc", vector_input, ["description"])
        return [item["description"] for item in response_list]

    def get_similar_question_sql(self, question: str, **kwargs) -> list:
        vector_input = self.generate_embedding(question)
        response_list = self._query_collection(
            "sql", vector_input, ["sql", "natural_language_question"]
        )
        return [
            {"question": item["natural_language_question"], "sql": item["sql"]}
            for item in response_list
        ]

    def get_training_data(self, **kwargs) -> list:
        self.weaviate_client.connect()
        combined_response_list = []
        for collection_name in self.training_data_cluster.values():
            if self.weaviate_client.collections.exists(collection_name):
                collection = self.weaviate_client.collections.get(collection_name)
                response_list = [item.properties for item in collection.iterator()]
                combined_response_list.extend(response_list)
        self.weaviate_client.close()
        return combined_response_list

    def remove_training_data(self, id: str, **kwargs) -> bool:
        self.weaviate_client.connect()
        success = False
        if id.endswith("-sql"):
            id = id.replace("-sql", "")
            success = self.weaviate_client.collections.get(
                self.training_data_cluster["sql"]
            ).data.delete_by_id(id)
        elif id.endswith("-ddl"):
            id = id.replace("-ddl", "")
            success = self.weaviate_client.collections.get(
                self.training_data_cluster["ddl"]
            ).data.delete_by_id(id)
        elif id.endswith("-doc"):
            id = id.replace("-doc", "")
            success = self.weaviate_client.collections.get(
                self.training_data_cluster["doc"]
            ).data.delete_by_id(id)
        self.weaviate_client.close()
        return success


================================================
FILE: src/vanna/legacy/xinference/__init__.py
================================================
from .xinference import Xinference


================================================
FILE: src/vanna/legacy/xinference/xinference.py
================================================
from xinference_client.client.restful.restful_client import (
    Client,
    RESTfulChatModelHandle,
)

from ..base import VannaBase


class Xinference(VannaBase):
    def __init__(self, config=None):
        VannaBase.__init__(self, config=config)

        if not config or "base_url" not in config:
            raise ValueError("config must contain at least Xinference base_url")

        base_url = config["base_url"]
        api_key = config.get("api_key", "not empty")
        self.xinference_client = Client(base_url=base_url, api_key=api_key)

    def system_message(self, message: str) -> any:
        return {"role": "system", "content": message}

    def user_message(self, message: str) -> any:
        return {"role": "user", "content": message}

    def assistant_message(self, message: str) -> any:
        return {"role": "assistant", "content": message}

    def submit_prompt(self, prompt, **kwargs) -> str:
        if prompt is None:
            raise Exception("Prompt is None")

        if len(prompt) == 0:
            raise Exception("Prompt is empty")

        num_tokens = 0
        for message in prompt:
            num_tokens += len(message["content"]) / 4

        model_uid = kwargs.get("model_uid") or self.config.get("model_uid", None)
        if model_uid is None:
            raise ValueError("model_uid is required")

        xinference_model = self.xinference_client.get_model(model_uid)
        if isinstance(xinference_model, RESTfulChatModelHandle):
            print(f"Using model_uid {model_uid} for {num_tokens} tokens (approx)")

            response = xinference_model.chat(prompt)
            return response["choices"][0]["message"]["content"]
        else:
            raise NotImplementedError(
                f"Xinference model handle type {type(xinference_model)} is not supported, required RESTfulChatModelHandle"
            )


================================================
FILE: src/vanna/py.typed
================================================


================================================
FILE: src/vanna/servers/__init__.py
================================================
"""
Server implementations for the Vanna Agents framework.

This module provides Flask and FastAPI server factories for serving
Vanna agents over HTTP with SSE, WebSocket, and polling endpoints.
"""

from .base import ChatHandler, ChatRequest, ChatStreamChunk
from .cli.server_runner import ExampleAgentLoader

__all__ = [
    "ChatHandler",
    "ChatRequest",
    "ChatStreamChunk",
    "ExampleAgentLoader",
]


================================================
FILE: src/vanna/servers/__main__.py
================================================
"""
Entry point for running Vanna Agents servers.
"""

from .cli.server_runner import main

if __name__ == "__main__":
    main()


================================================
FILE: src/vanna/servers/base/__init__.py
================================================
"""
Base server components for the Vanna Agents framework.

This module provides framework-agnostic components for handling chat
requests and responses.
"""

from .chat_handler import ChatHandler
from .models import ChatRequest, ChatStreamChunk, ChatResponse
from .templates import INDEX_HTML

__all__ = [
    "ChatHandler",
    "ChatRequest",
    "ChatStreamChunk",
    "ChatResponse",
    "INDEX_HTML",
]


================================================
FILE: src/vanna/servers/base/chat_handler.py
================================================
"""
Framework-agnostic chat handling logic.
"""

import uuid
from typing import AsyncGenerator, List

from ...core import Agent
from .models import ChatRequest, ChatResponse, ChatStreamChunk


class ChatHandler:
    """Core chat handling logic - framework agnostic."""

    def __init__(
        self,
        agent: Agent,
    ):
        """Initialize chat handler.

        Args:
            agent: The agent to handle chat requests
        """
        self.agent = agent

    async def handle_stream(
        self, request: ChatRequest
    ) -> AsyncGenerator[ChatStreamChunk, None]:
        """Stream chat responses.

        Args:
            request: Chat request

        Yields:
            Chat stream chunks
        """
        conversation_id = request.conversation_id or self._generate_conversation_id()
        # Use request_id from client for tracking, or use the one generated internally
        request_id = request.request_id or str(uuid.uuid4())

        async for component in self.agent.send_message(
            request_context=request.request_context,
            message=request.message,
            conversation_id=conversation_id,
        ):
            yield ChatStreamChunk.from_component(component, conversation_id, request_id)

    async def handle_poll(self, request: ChatRequest) -> ChatResponse:
        """Handle polling-based chat.

        Args:
            request: Chat request

        Returns:
            Complete chat response
        """
        chunks = []
        async for chunk in self.handle_stream(request):
            chunks.append(chunk)

        return ChatResponse.from_chunks(chunks)

    def _generate_conversation_id(self) -> str:
        """Generate new conversation ID."""
        return f"conv_{uuid.uuid4().hex[:8]}"


================================================
FILE: src/vanna/servers/base/models.py
================================================
"""
Request and response models for server endpoints.
"""

import time
import uuid
from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field

from ...components import UiComponent, RichComponent
from ...core.component_manager import ComponentUpdate
from ...core.user.request_context import RequestContext


class ChatRequest(BaseModel):
    """Request model for chat endpoints."""

    message: str = Field(description="User message")
    conversation_id: Optional[str] = Field(default=None, description="Conversation ID")
    request_id: Optional[str] = Field(
        default=None, description="Request ID for tracing"
    )
    request_context: RequestContext = Field(
        default_factory=RequestContext,
        description="Request context for user resolution",
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional metadata"
    )


class ChatStreamChunk(BaseModel):
    """Single chunk in a streaming chat response."""

    rich: Dict[str, Any] = Field(description="Rich component data for advanced UIs")
    simple: Optional[Dict[str, Any]] = Field(
        default=None, description="Simple component data for basic UIs"
    )

    # Stream metadata
    conversation_id: str = Field(description="Conversation ID")
    request_id: str = Field(description="Request ID")
    timestamp: float = Field(default_factory=time.time, description="Timestamp")

    @classmethod
    def from_component(
        cls,
        component: Union[UiComponent, RichComponent],
        conversation_id: str,
        request_id: str,
    ) -> "ChatStreamChunk":
        """Create chunk from UI component or rich component."""

        if isinstance(component, UiComponent):
            # Full UiComponent with both rich and simple
            rich_data = component.rich_component.serialize_for_frontend()
            simple_data = None
            if component.simple_component:
                simple_data = component.simple_component.serialize_for_frontend()

            return cls(
                rich=rich_data,
                simple=simple_data,
                conversation_id=conversation_id,
                request_id=request_id,
            )

        # Rich component only (no simple fallback)
        rich_data = component.serialize_for_frontend()
        return cls(
            rich=rich_data,
            simple=None,
            conversation_id=conversation_id,
            request_id=request_id,
        )

    @classmethod
    def from_component_update(
        cls, update: ComponentUpdate, conversation_id: str, request_id: str
    ) -> "ChatStreamChunk":
        """Create chunk from component update."""
        update_payload = update.serialize_for_frontend()
        return cls(
            rich=update_payload,
            simple=None,  # Component updates don't have simple representations
            conversation_id=conversation_id,
            request_id=request_id,
        )


class ChatResponse(BaseModel):
    """Complete chat response for polling endpoints."""

    chunks: List[ChatStreamChunk] = Field(description="Response chunks")
    conversation_id: str = Field(description="Conversation ID")
    request_id: str = Field(description="Request ID")
    total_chunks: int = Field(description="Total number of chunks")

    @classmethod
    def from_chunks(cls, chunks: List[ChatStreamChunk]) -> "ChatResponse":
        """Create response from chunks."""
        if not chunks:
            return cls(chunks=[], conversation_id="", request_id="", total_chunks=0)

        return cls(
            chunks=chunks,
            conversation_id=chunks[0].conversation_id,
            request_id=chunks[0].request_id,
            total_chunks=len(chunks),
        )


================================================
FILE: src/vanna/servers/base/rich_chat_handler.py
================================================
# """
# Rich component-aware chat handling logic.
# """

# import uuid
# from typing import AsyncGenerator, Callable, List, Optional, Union

# from ...core import Agent, User
# from ...core.rich_components import RichComponent
# from ...core.component_manager import ComponentManager, ComponentUpdate
# from .models import ChatRequest, ChatResponse, ChatStreamChunk


# class RichChatHandler:
#     """Rich component-aware chat handling logic."""

#     def __init__(
#         self,
#         agent: Agent,
#         default_user_factory: Optional[Callable[[Optional[str]], User]] = None,
#     ):
#         """Initialize rich chat handler.

#         Args:
#             agent: The agent to handle chat requests
#             default_user_factory: Function to create default user from user_id
#         """
#         self.agent = agent
#         self.default_user_factory = default_user_factory or self._create_default_user
#         self.component_managers: dict[str, ComponentManager] = {}  # Per conversation

#     async def handle_stream(
#         self, request: ChatRequest
#     ) -> AsyncGenerator[ChatStreamChunk, None]:
#         """Stream chat responses with rich component support.

#         Args:
#             request: Chat request

#         Yields:
#             Chat stream chunks including rich component updates
#         """
#         user = self._resolve_user(request.user_id)
#         conversation_id = request.conversation_id or self._generate_conversation_id()
#         request_id = request.request_id or str(uuid.uuid4())

#         # Get or create component manager for this conversation
#         if conversation_id not in self.component_managers:
#             self.component_managers[conversation_id] = ComponentManager()

#         component_manager = self.component_managers[conversation_id]

#         async for component in self.agent.send_message(
#             conversation_id=conversation_id,
#             user=user,
#             message=request.message,
#             request_id=request_id,
#         ):
#             if isinstance(component, RichComponent):
#                 # Handle rich component through manager
#                 update = component_manager.emit(component)
#                 yield ChatStreamChunk.from_component_update(update, conversation_id, request_id)
#             else:
#                 # Handle legacy components
#                 yield ChatStreamChunk.from_component(component, conversation_id, request_id)

#     async def handle_poll(self, request: ChatRequest) -> ChatResponse:
#         """Handle polling request with rich component support.

#         Args:
#             request: Chat request

#         Returns:
#             Complete chat response with all components
#         """
#         chunks: List[ChatStreamChunk] = []

#         async for chunk in self.handle_stream(request):
#             chunks.append(chunk)

#         return ChatResponse.from_chunks(chunks)

#     def get_component_manager(self, conversation_id: str) -> Optional[ComponentManager]:
#         """Get the component manager for a conversation."""
#         return self.component_managers.get(conversation_id)

#     def get_component(self, conversation_id: str, component_id: str) -> Optional[RichComponent]:
#         """Get a specific component from a conversation."""
#         manager = self.get_component_manager(conversation_id)
#         return manager.get_component(component_id) if manager else None

#     def get_all_components(self, conversation_id: str) -> List[RichComponent]:
#         """Get all components in a conversation."""
#         manager = self.get_component_manager(conversation_id)
#         return manager.get_all_components() if manager else []

#     def update_component(
#         self,
#         conversation_id: str,
#         component_id: str,
#         **updates
#     ) -> Optional[ComponentUpdate]:
#         """Update a component in a conversation."""
#         manager = self.get_component_manager(conversation_id)
#         return manager.update_component(component_id, **updates) if manager else None

#     def remove_component(
#         self,
#         conversation_id: str,
#         component_id: str
#     ) -> Optional[ComponentUpdate]:
#         """Remove a component from a conversation."""
#         manager = self.get_component_manager(conversation_id)
#         return manager.remove_component(component_id) if manager else None

#     def clear_conversation_components(self, conversation_id: str):
#         """Clear all components for a conversation."""
#         if conversation_id in self.component_managers:
#             del self.component_managers[conversation_id]

#     def _resolve_user(self, user_id: Optional[str]) -> User:
#         """Resolve user from ID or create default."""
#         if user_id:
#             # In a real implementation, you'd fetch from a user store
#             return User(id=user_id, username=f"user_{user_id}", email="", permissions=[])

#         return self.default_user_factory(user_id)

#     def _create_default_user(self, user_id: Optional[str]) -> User:
#         """Create a default user."""
#         user_id = user_id or "anonymous"
#         return User(
#             id=user_id,
#             username=f"user_{user_id}",
#             email="",
#             permissions=[]
#         )

#     def _generate_conversation_id(self) -> str:
#         """Generate a new conversation ID."""
#         return str(uuid.uuid4())


================================================
FILE: src/vanna/servers/base/templates.py
================================================
"""
HTML templates for Vanna Agents servers.
"""

from typing import Optional


def get_vanna_component_script(
    dev_mode: bool = False,
    static_path: str = "/static",
    cdn_url: str = "https://img.vanna.ai/vanna-components.js",
) -> str:
    """Get the script tag for loading Vanna web components.

    Args:
        dev_mode: If True, load from local static files
        static_path: Path to static assets in dev mode
        cdn_url: CDN URL for production

    Returns:
        HTML script tag for loading components
    """
    if dev_mode:
        return (
            f'<script type="module" src="{static_path}/vanna-components.js"></script>'
        )
    else:
        return f'<script type="module" src="{cdn_url}"></script>'


def get_index_html(
    dev_mode: bool = False,
    static_path: str = "/static",
    cdn_url: str = "https://img.vanna.ai/vanna-components.js",
    api_base_url: str = "",
) -> str:
    """Generate index HTML with configurable component loading.

    Args:
        dev_mode: If True, load components from local static files
        static_path: Path to static assets in dev mode
        cdn_url: CDN URL for production components
        api_base_url: Base URL for API endpoints

    Returns:
        Complete HTML page as string
    """
    component_script = get_vanna_component_script(dev_mode, static_path, cdn_url)

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Vanna Agents Chat</title>
    <link rel="preconnect" href="https://fonts.googleapis.com">
    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
    <link href="https://fonts.googleapis.com/css2?family=Roboto+Slab:wght@400;500;600;700&family=Space+Grotesk:wght@300;400;500;600;700&family=Space+Mono:wght@400;700&display=swap" rel="stylesheet">
    <script src="https://cdn.tailwindcss.com"></script>
    <script>
        tailwind.config = {{
            theme: {{
                extend: {{
                    colors: {{
                        'vanna-navy': '#023d60',
                        'vanna-cream': '#e7e1cf',
                        'vanna-teal': '#15a8a8',
                        'vanna-orange': '#fe5d26',
                        'vanna-magenta': '#bf1363',
                    }},
                    fontFamily: {{
                        'sans': ['Space Grotesk', 'ui-sans-serif', 'system-ui'],
                        'serif': ['Roboto Slab', 'ui-serif', 'Georgia'],
                        'mono': ['Space Mono', 'ui-monospace', 'monospace'],
                    }}
                }}
            }}
        }}
    </script>
    <style>
        body {{
            background: linear-gradient(to bottom, #e7e1cf, #ffffff, #e7e1cf);
            min-height: 100vh;
            position: relative;
            overflow-x: hidden;
        }}

        /* Background decorations matching landing page */
        body::before {{
            content: '';
            position: fixed;
            inset: 0;
            pointer-events: none;
            z-index: 0;
            /* Radial gradients with brand colors */
            background:
                radial-gradient(circle at top left, rgba(21, 168, 168, 0.12), transparent 60%),
                radial-gradient(circle at bottom right, rgba(254, 93, 38, 0.08), transparent 65%);
        }}

        body::after {{
            content: '';
            position: fixed;
            inset: 0;
            pointer-events: none;
            z-index: 0;
            /* Dot pattern with retro computing aesthetic */
            background-image: radial-gradient(circle at 2px 2px, rgba(2, 61, 96, 0.3) 1px, transparent 0);
            background-size: 32px 32px;
            /* Grid overlay */
            background-image:
                radial-gradient(circle at 2px 2px, rgba(2, 61, 96, 0.3) 1px, transparent 0),
                linear-gradient(rgba(2, 61, 96, 0.1) 1px, transparent 1px),
                linear-gradient(90deg, rgba(2, 61, 96, 0.1) 1px, transparent 1px);
            background-size: 32px 32px, 100px 100px, 100px 100px;
        }}

        /* Ensure content is above background */
        body > * {{
            position: relative;
            z-index: 1;
        }}

        vanna-chat {{
            width: 100%;
            height: 100%;
            display: block;
        }}
    </style>
    {component_script}
</head>
<body>
    <div class="max-w-6xl mx-auto p-5">
        <!-- Header -->
        <div class="text-center mb-8">
            <h1 class="text-4xl font-bold text-vanna-navy mb-2 font-serif">Vanna Agents</h1>
            <p class="text-lg font-mono font-bold text-vanna-teal mb-4">DATA-FIRST AGENTS</p>
            <p class="text-slate-600 mb-4">Interactive AI Assistant powered by Vanna Agents Framework</p>
            <a href="javascript:window.location='view-source:'+window.location.href" class="inline-flex items-center gap-2 px-4 py-2 bg-vanna-teal text-white text-sm font-medium rounded-lg hover:bg-vanna-navy transition">
                <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                    <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M10 20l4-16m4 4l4 4-4 4M6 16l-4-4 4-4"/>
                </svg>
                View Page Source
            </a>
        </div>

        {('    <div class="bg-vanna-orange/10 border border-vanna-orange/30 rounded-lg p-3 mb-5 text-vanna-orange text-sm font-medium">📦 Development Mode: Loading components from local assets</div>' if dev_mode else "")}

        <!-- Login Form -->
        <div id="loginContainer" class="max-w-md mx-auto mb-10 bg-white p-8 rounded-xl shadow-lg border border-vanna-teal/30">
            <div class="text-center mb-6">
                <h2 class="text-2xl font-semibold text-vanna-navy mb-2 font-serif">Login to Continue</h2>
                <p class="text-sm text-slate-600">Select your email to access the chat</p>
            </div>

            <div class="mb-5">
                <label for="emailInput" class="block mb-2 text-sm font-medium text-vanna-navy">Email Address</label>
                <select
                    id="emailInput"
                    class="w-full px-4 py-3 text-sm border border-vanna-teal/30 rounded-lg focus:outline-none focus:ring-2 focus:ring-vanna-teal focus:border-transparent bg-white"
                >
                    <option value="">Select an email...</option>
                    <option value="admin@example.com">admin@example.com</option>
                    <option value="user@example.com">user@example.com</option>
                </select>
            </div>

            <button id="loginButton" class="w-full px-4 py-3 bg-vanna-teal text-white text-sm font-medium rounded-lg hover:bg-vanna-navy focus:outline-none focus:ring-2 focus:ring-vanna-teal focus:ring-offset-2 transition disabled:bg-gray-400 disabled:cursor-not-allowed">
                Continue
            </button>

            <div class="mt-5 p-3 bg-vanna-teal/10 border-l-4 border-vanna-teal rounded text-xs text-vanna-navy leading-relaxed">
                <strong>Demo Mode:</strong> This is a frontend-only authentication demo.
                Your email will be stored as a cookie and automatically sent with all API requests.
            </div>
        </div>

        <!-- Logged In Status (hidden by default) -->
        <div id="loggedInStatus" class="hidden text-center p-4 bg-vanna-teal/10 border border-vanna-teal/30 rounded-lg mb-5">
            Logged in as <span id="loggedInEmail" class="font-semibold text-vanna-navy"></span>
            <br>
            <button id="logoutButton" class="mt-2 px-3 py-1.5 bg-vanna-navy text-white text-xs rounded hover:bg-vanna-teal transition">
                Logout
            </button>
        </div>

        <!-- Chat Container (hidden by default) -->
        <div id="chatSections" class="hidden">
            <div class="bg-white rounded-xl shadow-lg h-[600px] overflow-hidden border border-vanna-teal/30">
                <vanna-chat
                    api-base="{api_base_url}"
                    sse-endpoint="{api_base_url}/api/vanna/v2/chat_sse"
                    ws-endpoint="{api_base_url}/api/vanna/v2/chat_websocket"
                    poll-endpoint="{api_base_url}/api/vanna/v2/chat_poll">
                </vanna-chat>
            </div>

            <div class="mt-8 p-5 bg-white rounded-lg shadow border border-vanna-teal/30">
                <h3 class="text-lg font-semibold text-vanna-navy mb-3 font-serif">API Endpoints</h3>
                <ul class="space-y-2">
                    <li class="p-2 bg-vanna-cream/50 rounded font-mono text-sm">
                        <span class="font-bold text-vanna-teal mr-2">POST</span>{api_base_url}/api/vanna/v2/chat_sse - Server-Sent Events streaming
                    </li>
                    <li class="p-2 bg-vanna-cream/50 rounded font-mono text-sm">
                        <span class="font-bold text-vanna-teal mr-2">WS</span>{api_base_url}/api/vanna/v2/chat_websocket - WebSocket real-time chat
                    </li>
                    <li class="p-2 bg-vanna-cream/50 rounded font-mono text-sm">
                        <span class="font-bold text-vanna-teal mr-2">POST</span>{api_base_url}/api/vanna/v2/chat_poll - Request/response polling
                    </li>
                    <li class="p-2 bg-vanna-cream/50 rounded font-mono text-sm">
                        <span class="font-bold text-vanna-teal mr-2">GET</span>{api_base_url}/health - Health check
                    </li>
                </ul>
            </div>
        </div>
    </div>

    <script>
        // Cookie helpers
        const getCookie = (name) => {{
            const value = `; ${{document.cookie}}`;
            const parts = value.split(`; ${{name}}=`);
            return parts.length === 2 ? parts.pop().split(';').shift() : null;
        }};

        const setCookie = (name, value) => {{
            const expires = new Date(Date.now() + 365 * 864e5).toUTCString();
            document.cookie = `${{name}}=${{value}}; expires=${{expires}}; path=/; SameSite=Lax`;
        }};

        const deleteCookie = (name) => {{
            document.cookie = `${{name}}=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;`;
        }};

        // Login/Logout
        document.addEventListener('DOMContentLoaded', () => {{
            const email = getCookie('vanna_email');

            // Check if already logged in
            if (email) {{
                loginContainer.classList.add('hidden');
                loggedInStatus.classList.remove('hidden');
                chatSections.classList.remove('hidden');
                loggedInEmail.textContent = email;
            }}

            // Login button
            loginButton.addEventListener('click', () => {{
                const email = emailInput.value.trim();
                if (!email) {{
                    alert('Please select an email address');
                    return;
                }}
                setCookie('vanna_email', email);
                loginContainer.classList.add('hidden');
                loggedInStatus.classList.remove('hidden');
                chatSections.classList.remove('hidden');
                loggedInEmail.textContent = email;
            }});

            // Logout button
            logoutButton.addEventListener('click', () => {{
                deleteCookie('vanna_email');
                loginContainer.classList.remove('hidden');
                loggedInStatus.classList.add('hidden');
                chatSections.classList.add('hidden');
                emailInput.value = '';
            }});

            // Enter key
            emailInput.addEventListener('keypress', (e) => {{
                if (e.key === 'Enter') loginButton.click();
            }});
        }});
    </script>

    <script>
        // Artifact demo event listener
        document.addEventListener('DOMContentLoaded', () => {{
            const vannaChat = document.querySelector('vanna-chat');

            if (vannaChat) {{
                // Add artifact event listener to demonstrate external rendering
                vannaChat.addEventListener('artifact-opened', (event) => {{
                    const {{ artifactId, type, title, trigger }} = event.detail;

                    console.log('🎨 Artifact Event:', {{ artifactId, type, title, trigger }});

                    // For demo: open all artifacts externally
                    setTimeout(() => {{
                        const newWindow = window.open('', '_blank', 'width=900,height=700');
                        if (newWindow) {{
                            newWindow.document.write(event.detail.getStandaloneHTML());
                            newWindow.document.close();
                            newWindow.document.title = title || 'Vanna Artifact';
                            console.log(`📱 Opened ${{title}} in new window`);
                        }}
                    }}, 100);

                    // Prevent default in-chat rendering
                    event.detail.preventDefault();
                    console.log('✋ Showing placeholder in chat instead of full artifact');
                }});

                console.log('🎯 Artifact demo mode: All artifacts will open externally');
            }}
        }});

        // Fallback if web component doesn't load
        if (!customElements.get('vanna-chat')) {{
            setTimeout(() => {{
                if (!customElements.get('vanna-chat')) {{
                    document.querySelector('vanna-chat').innerHTML = `
                        <div class="p-10 text-center text-gray-600">
                            <h3 class="text-xl font-semibold mb-2">Vanna Chat Component</h3>
                            <p class="mb-2">Web component failed to load. Please check your connection.</p>
                            <p class="text-sm text-gray-400">
                                {("Loading from: local static assets" if dev_mode else f"Loading from: {cdn_url}")}
                            </p>
                        </div>
                    `;
                }}
            }}, 2000);
        }}
    </script>
</body>
</html>"""


# Backward compatibility - default production HTML
INDEX_HTML = get_index_html()


================================================
FILE: src/vanna/servers/cli/__init__.py
================================================
"""
CLI components for Vanna Agents servers.
"""

from .server_runner import ExampleAgentLoader

__all__ = ["ExampleAgentLoader"]


================================================
FILE: src/vanna/servers/cli/server_runner.py
================================================
"""
CLI for running Vanna Agents servers with example agents.
"""

import importlib
import json
from typing import Dict, Optional, Any, cast, TextIO, Union

import click

from ...core import Agent


class ExampleAgentLoader:
    """Loads example agents for the CLI."""

    @staticmethod
    def list_available_examples() -> Dict[str, str]:
        """Return available examples with descriptions."""
        return {
            "mock_quickstart": "Basic agent with mock LLM service",
            "anthropic_quickstart": "Agent configured for Anthropic's Claude API",
            "openai_quickstart": "Agent configured for OpenAI's GPT models",
            "mock_custom_tool": "Agent with custom tool demonstration (mock LLM)",
            "mock_quota_example": "Agent with usage quota management (mock LLM)",
            "mock_rich_components_demo": "Rich components demonstration with cards, tasks, and progress (mock LLM)",
            "coding_agent_example": "Coding agent with file system tools (list, read, write files)",
            "email_auth_example": "Email-based authentication demonstration (mock LLM)",
            "claude_sqlite_example": "Claude agent with SQLite database querying capabilities",
            "mock_sqlite_example": "Mock agent with SQLite database demonstration",
        }

    @staticmethod
    def load_example_agent(example_name: str) -> Agent:
        """Load an example agent by name.

        Args:
            example_name: Name of the example to load

        Returns:
            Configured agent instance

        Raises:
            ValueError: If example not found or failed to load
        """
        try:
            # Import the example module
            module = importlib.import_module(f"vanna.examples.{example_name}")

            # Look for standard factory functions
            factory_functions = [
                "create_demo_agent",
                "create_agent",
                "create_basic_demo",
            ]

            for func_name in factory_functions:
                if hasattr(module, func_name):
                    factory = getattr(module, func_name)
                    return cast(Agent, factory())

            # Look for module-level agent instances
            if hasattr(module, "main_agent"):
                return cast(Agent, module.main_agent)

            raise AttributeError(f"No agent factory found in {example_name}")

        except ImportError as e:
            raise ValueError(f"Example '{example_name}' not found: {e}")
        except Exception as e:
            raise ValueError(f"Failed to load example '{example_name}': {e}")


@click.command()
@click.option(
    "--framework",
    type=click.Choice(["flask", "fastapi"]),
    default="fastapi",
    help="Web framework to use",
)
@click.option("--port", default=8000, help="Port to run server on")
@click.option("--host", default="0.0.0.0", help="Host to bind server to")
@click.option(
    "--example", help="Example agent to use (use --list-examples to see options)"
)
@click.option("--list-examples", is_flag=True, help="List available example agents")
@click.option(
    "--config", type=click.File("r"), help="JSON config file for server settings"
)
@click.option("--debug", is_flag=True, help="Enable debug mode")
@click.option(
    "--dev",
    is_flag=True,
    help="Enable development mode (load components from local assets)",
)
@click.option(
    "--static-folder", default=None, help="Static folder path for development mode"
)
@click.option(
    "--cdn-url",
    default="https://img.vanna.ai/vanna-components.js",
    help="CDN URL for web components",
)
def main(
    framework: str,
    port: int,
    host: str,
    example: Optional[str],
    list_examples: bool,
    config: Optional[click.File],
    debug: bool,
    dev: bool,
    static_folder: Optional[str],
    cdn_url: str,
) -> None:
    """Run Vanna Agents server with optional example agent."""

    if list_examples:
        click.echo("Available example agents:")
        examples = ExampleAgentLoader.list_available_examples()
        for name, description in examples.items():
            click.echo(f"  {name:20} - {description}")
        return

    # Load configuration
    server_config = {}
    if config:
        server_config = json.load(cast(TextIO, config))

    # Set default static folder based on dev mode
    if static_folder is None:
        static_folder = "frontend/webcomponent/static" if dev else "static"

    # Add CLI options to config
    server_config.update(
        {
            "dev_mode": dev,
            "static_folder": static_folder,
            "cdn_url": cdn_url,
            "api_base_url": "",  # Can be overridden in config file
        }
    )

    # Create agent
    if example:
        try:
            agent = ExampleAgentLoader.load_example_agent(example)
            click.echo(f"✓ Loaded example agent: {example}")
        except ValueError as e:
            click.echo(f"Error: {e}", err=True)
            return
    else:
        # Fallback to basic agent
        try:
            from ...agents import create_basic_agent
            from ...integrations.mock import MockLlmService

            llm_service = MockLlmService(
                response_content="Hello! I'm a Vanna Agents demo server. How can I help you?"
            )
            agent = create_basic_agent(llm_service)
            click.echo(
                "✓ Using basic demo agent (use --example to specify different agent)"
            )
        except ImportError as e:
            click.echo(f"Error: Could not create basic agent: {e}", err=True)
            return

    from ..flask.app import VannaFlaskServer
    from ..fastapi.app import VannaFastAPIServer

    # Create and run server
    server: Union[VannaFlaskServer, VannaFastAPIServer]
    if framework == "flask":
        server = VannaFlaskServer(agent, config=server_config)
        click.echo(f"🚀 Starting Flask server on http://{host}:{port}")
        if dev:
            click.echo(
                f"📦 Development mode: loading web components from ./{static_folder}/"
            )
        else:
            click.echo(f"🌍 Production mode: loading web components from CDN")
        try:
            server.run(host=host, port=port, debug=debug)
        except KeyboardInterrupt:
            click.echo("\n👋 Server stopped")
    else:
        server = VannaFastAPIServer(agent, config=server_config)
        click.echo(f"🚀 Starting FastAPI server on http://{host}:{port}")
        click.echo(f"📖 API docs available at http://{host}:{port}/docs")
        if dev:
            click.echo(
                f"📦 Development mode: loading web components from ./{static_folder}/"
            )
        else:
            click.echo(f"🌍 Production mode: loading web components from CDN")
        try:
            server.run(host=host, port=port)
        except KeyboardInterrupt:
            click.echo("\n👋 Server stopped")


if __name__ == "__main__":
    main()


================================================
FILE: src/vanna/servers/fastapi/__init__.py
================================================
"""
FastAPI server implementation for Vanna Agents.
"""

from .app import VannaFastAPIServer

__all__ = ["VannaFastAPIServer"]


================================================
FILE: src/vanna/servers/fastapi/app.py
================================================
"""
FastAPI server factory for Vanna Agents.
"""

from typing import Any, Dict, Optional

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

from ...core import Agent
from ..base import ChatHandler
from .routes import register_chat_routes


class VannaFastAPIServer:
    """FastAPI server factory for Vanna Agents."""

    def __init__(self, agent: Agent, config: Optional[Dict[str, Any]] = None):
        """Initialize FastAPI server.

        Args:
            agent: The agent to serve (must have user_resolver configured)
            config: Optional server configuration
        """
        self.agent = agent
        self.config = config or {}
        self.chat_handler = ChatHandler(agent)

    def create_app(self) -> FastAPI:
        """Create configured FastAPI app.

        Returns:
            Configured FastAPI application
        """
        # Create FastAPI app
        app_config = self.config.get("fastapi", {})
        app = FastAPI(
            title="Vanna Agents API",
            description="API server for Vanna Agents framework",
            version="0.1.0",
            **app_config,
        )

        # Configure CORS if enabled
        cors_config = self.config.get("cors", {})
        if cors_config.get("enabled", True):
            cors_params = {k: v for k, v in cors_config.items() if k != "enabled"}

            # Set sensible defaults
            cors_params.setdefault("allow_origins", ["*"])
            cors_params.setdefault("allow_credentials", True)
            cors_params.setdefault("allow_methods", ["*"])
            cors_params.setdefault("allow_headers", ["*"])

            app.add_middleware(CORSMiddleware, **cors_params)

        # Add static file serving in dev mode
        dev_mode = self.config.get("dev_mode", False)
        if dev_mode:
            static_folder = self.config.get("static_folder", "static")
            try:
                import os

                if os.path.exists(static_folder):
                    app.mount(
                        "/static", StaticFiles(directory=static_folder), name="static"
                    )
            except Exception:
                pass  # Static files not available

        # Register routes
        register_chat_routes(app, self.chat_handler, self.config)

        # Add health check
        @app.get("/health")
        async def health_check() -> Dict[str, str]:
            return {"status": "healthy", "service": "vanna"}

        return app

    def run(self, **kwargs: Any) -> None:
        """Run the FastAPI server.

        This method automatically detects if running in an async environment
        (Jupyter, Colab, IPython, etc.) and:
        - Uses appropriate async handling for existing event loops
        - Sets up port forwarding if in Google Colab
        - Displays the correct URL for accessing the app

        Args:
            **kwargs: Arguments passed to uvicorn configuration
        """
        import sys
        import asyncio
        import uvicorn

        # Check if we're in an environment with a running event loop FIRST
        in_async_env = False
        try:
            asyncio.get_running_loop()
            in_async_env = True
        except RuntimeError:
            in_async_env = False

        # If in async environment, apply nest_asyncio BEFORE creating the app
        if in_async_env:
            try:
                import nest_asyncio

                nest_asyncio.apply()
            except ImportError:
                print("Warning: nest_asyncio not installed. Installing...")
                import subprocess

                subprocess.check_call(
                    [sys.executable, "-m", "pip", "install", "nest_asyncio"]
                )
                import nest_asyncio

                nest_asyncio.apply()

        # Now create the app after nest_asyncio is applied
        app = self.create_app()

        # Set defaults
        run_kwargs = {"host": "0.0.0.0", "port": 8000, "log_level": "info", **kwargs}

        # Get the port and other config from run_kwargs
        port = run_kwargs.get("port", 8000)
        host = run_kwargs.get("host", "0.0.0.0")
        log_level = run_kwargs.get("log_level", "info")

        # Check if we're specifically in Google Colab for port forwarding
        in_colab = "google.colab" in sys.modules

        if in_colab:
            try:
                from google.colab import output

                output.serve_kernel_port_as_window(port)
                from google.colab.output import eval_js

                print("Your app is running at:")
                print(eval_js(f"google.colab.kernel.proxyPort({port})"))
            except Exception as e:
                print(f"Warning: Could not set up Colab port forwarding: {e}")
                print(f"Your app is running at: http://localhost:{port}")
        else:
            print("Your app is running at:")
            print(f"http://localhost:{port}")

        if in_async_env:
            # In Jupyter/Colab, create config with loop="asyncio" and use asyncio.run()
            # This matches the working pattern from Colab
            config = uvicorn.Config(
                app, host=host, port=port, log_level=log_level, loop="asyncio"
            )
            server = uvicorn.Server(config)
            asyncio.run(server.serve())
        else:
            # Normal execution outside of Jupyter/Colab
            uvicorn.run(app, **run_kwargs)


================================================
FILE: src/vanna/servers/fastapi/routes.py
================================================
"""
FastAPI route implementations for Vanna Agents.
"""

import json
import traceback
from typing import Any, AsyncGenerator, Dict, Optional

from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import StreamingResponse, HTMLResponse

from ..base import ChatHandler, ChatRequest, ChatResponse
from ..base.templates import get_index_html
from ...core.user.request_context import RequestContext


def register_chat_routes(
    app: FastAPI, chat_handler: ChatHandler, config: Optional[Dict[str, Any]] = None
) -> None:
    """Register chat routes on FastAPI app.

    Args:
        app: FastAPI application
        chat_handler: Chat handler instance
        config: Server configuration
    """
    config = config or {}

    @app.get("/", response_class=HTMLResponse)
    async def index() -> str:
        """Serve the main chat interface."""
        dev_mode = config.get("dev_mode", False)
        cdn_url = config.get("cdn_url", "https://img.vanna.ai/vanna-components.js")
        api_base_url = config.get("api_base_url", "")

        return get_index_html(
            dev_mode=dev_mode, cdn_url=cdn_url, api_base_url=api_base_url
        )

    @app.post("/api/vanna/v2/chat_sse")
    async def chat_sse(
        chat_request: ChatRequest, http_request: Request
    ) -> StreamingResponse:
        """Server-Sent Events endpoint for streaming chat."""
        # Extract request context for user resolution
        chat_request.request_context = RequestContext(
            cookies=dict(http_request.cookies),
            headers=dict(http_request.headers),
            remote_addr=http_request.client.host if http_request.client else None,
            query_params=dict(http_request.query_params),
            metadata=chat_request.metadata,
        )

        async def generate() -> AsyncGenerator[str, None]:
            """Generate SSE stream."""
            try:
                async for chunk in chat_handler.handle_stream(chat_request):
                    chunk_json = chunk.model_dump_json()
                    yield f"data: {chunk_json}\n\n"
                yield "data: [DONE]\n\n"
            except Exception as e:
                traceback.print_stack()
                traceback.print_exc()
                error_data = {
                    "type": "error",
                    "data": {"message": str(e)},
                    "conversation_id": chat_request.conversation_id or "",
                    "request_id": chat_request.request_id or "",
                }
                yield f"data: {json.dumps(error_data)}\n\n"

        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",  # Disable nginx buffering
            },
        )

    @app.websocket("/api/vanna/v2/chat_websocket")
    async def chat_websocket(websocket: WebSocket) -> None:
        """WebSocket endpoint for real-time chat."""
        await websocket.accept()

        try:
            while True:
                # Receive message
                try:
                    data = await websocket.receive_json()

                    # Extract request context for user resolution
                    metadata = data.get("metadata", {})
                    data["request_context"] = RequestContext(
                        cookies=dict(websocket.cookies),
                        headers=dict(websocket.headers),
                        remote_addr=websocket.client.host if websocket.client else None,
                        query_params=dict(websocket.query_params),
                        metadata=metadata,
                    )

                    chat_request = ChatRequest(**data)
                except Exception as e:
                    traceback.print_stack()
                    traceback.print_exc()
                    await websocket.send_json(
                        {
                            "type": "error",
                            "data": {"message": f"Invalid request: {str(e)}"},
                        }
                    )
                    continue

                # Stream response
                try:
                    async for chunk in chat_handler.handle_stream(chat_request):
                        await websocket.send_json(chunk.model_dump())

                    # Send completion signal
                    await websocket.send_json(
                        {
                            "type": "completion",
                            "data": {"status": "done"},
                            "conversation_id": chunk.conversation_id
                            if "chunk" in locals()
                            else "",
                            "request_id": chunk.request_id
                            if "chunk" in locals()
                            else "",
                        }
                    )

                except Exception as e:
                    traceback.print_stack()
                    traceback.print_exc()
                    await websocket.send_json(
                        {
                            "type": "error",
                            "data": {"message": str(e)},
                            "conversation_id": chat_request.conversation_id or "",
                            "request_id": chat_request.request_id or "",
                        }
                    )

        except WebSocketDisconnect:
            pass
        except Exception as e:
            traceback.print_stack()
            traceback.print_exc()
            try:
                await websocket.send_json(
                    {
                        "type": "error",
                        "data": {"message": f"WebSocket error: {str(e)}"},
                    }
                )
            except Exception:
                pass
            finally:
                await websocket.close()

    @app.post("/api/vanna/v2/chat_poll")
    async def chat_poll(
        chat_request: ChatRequest, http_request: Request
    ) -> ChatResponse:
        """Polling endpoint for chat."""
        # Extract request context for user resolution
        chat_request.request_context = RequestContext(
            cookies=dict(http_request.cookies),
            headers=dict(http_request.headers),
            remote_addr=http_request.client.host if http_request.client else None,
            query_params=dict(http_request.query_params),
            metadata=chat_request.metadata,
        )

        try:
            result = await chat_handler.handle_poll(chat_request)
            return result
        except Exception as e:
            traceback.print_stack()
            traceback.print_exc()
            raise HTTPException(status_code=500, detail=f"Chat failed: {str(e)}")


================================================
FILE: src/vanna/servers/flask/__init__.py
================================================
"""
Flask server implementation for Vanna Agents.
"""

from .app import VannaFlaskServer

__all__ = ["VannaFlaskServer"]


================================================
FILE: src/vanna/servers/flask/app.py
================================================
"""
Flask server factory for Vanna Agents.
"""

import asyncio
from typing import Any, Dict, Optional

from flask import Flask
from flask_cors import CORS

from ...core import Agent
from ..base import ChatHandler
from .routes import register_chat_routes


class VannaFlaskServer:
    """Flask server factory for Vanna Agents."""

    def __init__(self, agent: Agent, config: Optional[Dict[str, Any]] = None):
        """Initialize Flask server.

        Args:
            agent: The agent to serve (must have user_resolver configured)
            config: Optional server configuration
        """
        self.agent = agent
        self.config = config or {}
        self.chat_handler = ChatHandler(agent)

    def create_app(self) -> Flask:
        """Create configured Flask app.

        Returns:
            Configured Flask application
        """
        # Check if dev mode is enabled
        dev_mode = self.config.get("dev_mode", False)
        static_folder = self.config.get("static_folder", "static") if dev_mode else None

        app = Flask(__name__, static_folder=static_folder, static_url_path="/static")

        # Apply configuration
        app.config.update(self.config.get("flask", {}))

        # Enable CORS if configured
        cors_config = self.config.get("cors", {})
        if cors_config.get("enabled", True):
            CORS(app, **{k: v for k, v in cors_config.items() if k != "enabled"})

        # Register routes
        register_chat_routes(app, self.chat_handler, self.config)

        # Add health check
        @app.route("/health")
        def health_check() -> Dict[str, str]:
            return {"status": "healthy", "service": "vanna"}

        return app

    def run(self, **kwargs: Any) -> None:
        """Run the Flask server.

        This method automatically detects if running in an async environment
        (Jupyter, Colab, IPython, etc.) and:
        - Installs and applies nest_asyncio to handle existing event loops
        - Sets up port forwarding if in Google Colab
        - Displays the correct URL for accessing the app

        Args:
            **kwargs: Arguments passed to Flask.run()
        """
        import sys

        app = self.create_app()

        # Set defaults
        run_kwargs = {"host": "0.0.0.0", "port": 5000, "debug": False, **kwargs}

        # Get the port from run_kwargs
        port = run_kwargs.get("port", 5000)

        # Check if we're in an environment with a running event loop
        # (Jupyter, Colab, IPython, VS Code notebooks, etc.)
        in_async_env = False
        try:
            import asyncio

            try:
                asyncio.get_running_loop()
                in_async_env = True
            except RuntimeError:
                in_async_env = False
        except Exception:
            pass

        if in_async_env:
            # Apply nest_asyncio to allow nested event loops
            try:
                import nest_asyncio

                nest_asyncio.apply()
            except ImportError:
                print("Warning: nest_asyncio not installed. Installing...")
                import subprocess

                subprocess.check_call(
                    [sys.executable, "-m", "pip", "install", "nest_asyncio"]
                )
                import nest_asyncio

                nest_asyncio.apply()

        # Check if we're specifically in Google Colab for port forwarding
        in_colab = "google.colab" in sys.modules

        if in_colab:
            try:
                from google.colab import output

                output.serve_kernel_port_as_window(port)
                from google.colab.output import eval_js

                print("Your app is running at:")
                print(eval_js(f"google.colab.kernel.proxyPort({port})"))
            except Exception as e:
                print(f"Warning: Could not set up Colab port forwarding: {e}")
                print(f"Your app is running at: http://localhost:{port}")
        else:
            print("Your app is running at:")
            print(f"http://localhost:{port}")

        app.run(**run_kwargs)


================================================
FILE: src/vanna/servers/flask/routes.py
================================================
"""
Flask route implementations for Vanna Agents.
"""

import asyncio
import json
import traceback
from typing import Any, AsyncGenerator, Dict, Generator, Optional, Union

from flask import Flask, Response, jsonify, request

from ..base import ChatHandler, ChatRequest
from ..base.templates import get_index_html
from ...core.user.request_context import RequestContext


def register_chat_routes(
    app: Flask, chat_handler: ChatHandler, config: Optional[Dict[str, Any]] = None
) -> None:
    """Register chat routes on Flask app.

    Args:
        app: Flask application
        chat_handler: Chat handler instance
        config: Server configuration
    """
    config = config or {}

    @app.route("/")
    def index() -> str:
        """Serve the main chat interface."""
        dev_mode = config.get("dev_mode", False)
        cdn_url = config.get("cdn_url", "https://img.vanna.ai/vanna-components.js")
        api_base_url = config.get("api_base_url", "")

        return get_index_html(
            dev_mode=dev_mode, cdn_url=cdn_url, api_base_url=api_base_url
        )

    @app.route("/api/vanna/v2/chat_sse", methods=["POST"])
    def chat_sse() -> Union[Response, tuple[Response, int]]:
        """Server-Sent Events endpoint for streaming chat."""
        try:
            data = request.get_json()
            if not data:
                return jsonify({"error": "JSON body required"}), 400

            # Extract request context for user resolution
            data["request_context"] = RequestContext(
                cookies=dict(request.cookies),
                headers=dict(request.headers),
                remote_addr=request.remote_addr,
                query_params=dict(request.args),
            )

            chat_request = ChatRequest(**data)
        except Exception as e:
            traceback.print_stack()
            traceback.print_exc()
            return jsonify({"error": f"Invalid request: {str(e)}"}), 400

        def generate() -> Generator[str, None, None]:
            """Generate SSE stream."""
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

            try:

                async def async_generate() -> AsyncGenerator[str, None]:
                    async for chunk in chat_handler.handle_stream(chat_request):
                        chunk_json = chunk.model_dump_json()
                        yield f"data: {chunk_json}\n\n"

                gen = async_generate()
                try:
                    while True:
                        chunk = loop.run_until_complete(gen.__anext__())
                        yield chunk
                except StopAsyncIteration:
                    yield "data: [DONE]\n\n"
            finally:
                loop.close()

        return Response(
            generate(),
            mimetype="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",  # Disable nginx buffering
            },
        )

    @app.route("/api/vanna/v2/chat_websocket")
    def chat_websocket() -> tuple[Response, int]:
        """WebSocket endpoint placeholder."""
        return jsonify(
            {
                "error": "WebSocket endpoint not implemented in basic Flask example",
                "suggestion": "Use Flask-SocketIO for WebSocket support",
            }
        ), 501

    @app.route("/api/vanna/v2/chat_poll", methods=["POST"])
    def chat_poll() -> Union[Response, tuple[Response, int]]:
        """Polling endpoint for chat."""
        try:
            data = request.get_json()
            if not data:
                return jsonify({"error": "JSON body required"}), 400

            # Extract request context for user resolution
            data["request_context"] = RequestContext(
                cookies=dict(request.cookies),
                headers=dict(request.headers),
                remote_addr=request.remote_addr,
                query_params=dict(request.args),
            )

            chat_request = ChatRequest(**data)
        except Exception as e:
            traceback.print_stack()
            traceback.print_exc()
            return jsonify({"error": f"Invalid request: {str(e)}"}), 400

        # Run async handler in new event loop
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            result = loop.run_until_complete(chat_handler.handle_poll(chat_request))
            return jsonify(result.model_dump())
        except Exception as e:
            traceback.print_stack()
            traceback.print_exc()
            return jsonify({"error": f"Chat failed: {str(e)}"}), 500
        finally:
            loop.close()


================================================
FILE: src/vanna/tools/__init__.py
================================================
"""Built-in tool implementations."""

from .file_system import (
    CommandResult,
    FileSystem,
    ListFilesTool,
    LocalFileSystem,
    ReadFileTool,
    SearchFilesTool,
    WriteFileTool,
    create_file_system_tools,
)
from .python import (
    PipInstallTool,
    RunPythonFileTool,
    create_python_tools,
)
from vanna.integrations.plotly import PlotlyChartGenerator
from .run_sql import RunSqlTool
from .visualize_data import VisualizeDataTool

__all__ = [
    # File system
    "FileSystem",
    "LocalFileSystem",
    "ListFilesTool",
    "SearchFilesTool",
    "ReadFileTool",
    "WriteFileTool",
    "create_file_system_tools",
    "CommandResult",
    # Python tools
    "RunPythonFileTool",
    "PipInstallTool",
    "create_python_tools",
    # SQL
    "RunSqlTool",
    # Visualization
    "PlotlyChartGenerator",
    "VisualizeDataTool",
]


================================================
FILE: src/vanna/tools/agent_memory.py
================================================
"""
Agent memory tools.

This module provides agent memory operations through an abstract AgentMemory interface,
allowing for different implementations (local vector DB, remote cloud service, etc.).
The tools access AgentMemory via ToolContext, which is populated by the Agent.
"""

import logging
from typing import Any, Dict, List, Optional, Type
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

from vanna.core.tool import Tool, ToolContext, ToolResult
from vanna.core.agent.config import UiFeature
from vanna.capabilities.agent_memory import AgentMemory
from vanna.components import (
    UiComponent,
    StatusBarUpdateComponent,
    CardComponent,
)


class SaveQuestionToolArgsParams(BaseModel):
    """Parameters for saving question-tool-argument combinations."""

    question: str = Field(description="The original question that was asked")
    tool_name: str = Field(
        description="The name of the tool that was used successfully"
    )
    args: Dict[str, Any] = Field(
        description="The arguments that were passed to the tool"
    )


class SearchSavedCorrectToolUsesParams(BaseModel):
    """Parameters for searching saved tool usage patterns."""

    question: str = Field(
        description="The question to find similar tool usage patterns for"
    )
    limit: Optional[int] = Field(
        default=10, description="Maximum number of results to return"
    )
    similarity_threshold: Optional[float] = Field(
        default=0.7, description="Minimum similarity score for results (0.0-1.0)"
    )
    tool_name_filter: Optional[str] = Field(
        default=None, description="Filter results to specific tool name"
    )


class SaveTextMemoryParams(BaseModel):
    """Parameters for saving free-form text memories."""

    content: str = Field(description="The text content to save as a memory")


class SaveQuestionToolArgsTool(Tool[SaveQuestionToolArgsParams]):
    """Tool for saving successful question-tool-argument combinations."""

    @property
    def name(self) -> str:
        return "save_question_tool_args"

    @property
    def description(self) -> str:
        return (
            "Save a successful question-tool-argument combination for future reference"
        )

    def get_args_schema(self) -> Type[SaveQuestionToolArgsParams]:
        return SaveQuestionToolArgsParams

    async def execute(
        self, context: ToolContext, args: SaveQuestionToolArgsParams
    ) -> ToolResult:
        """Save the tool usage pattern to agent memory."""
        try:
            await context.agent_memory.save_tool_usage(
                question=args.question,
                tool_name=args.tool_name,
                args=args.args,
                context=context,
                success=True,
            )

            success_msg = (
                f"Successfully saved usage pattern for '{args.tool_name}' tool"
            )
            return ToolResult(
                success=True,
                result_for_llm=success_msg,
                ui_component=UiComponent(
                    rich_component=StatusBarUpdateComponent(
                        status="success",
                        message="Saved to memory",
                        detail=f"Saved pattern for '{args.tool_name}'",
                    ),
                    simple_component=None,
                ),
            )

        except Exception as e:
            error_message = f"Failed to save memory: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=StatusBarUpdateComponent(
                        status="error", message="Failed to save memory", detail=str(e)
                    ),
                    simple_component=None,
                ),
                error=str(e),
            )


class SearchSavedCorrectToolUsesTool(Tool[SearchSavedCorrectToolUsesParams]):
    """Tool for searching saved tool usage patterns."""

    @property
    def name(self) -> str:
        return "search_saved_correct_tool_uses"

    @property
    def description(self) -> str:
        return "Search for similar tool usage patterns based on a question"

    def get_args_schema(self) -> Type[SearchSavedCorrectToolUsesParams]:
        return SearchSavedCorrectToolUsesParams

    async def execute(
        self, context: ToolContext, args: SearchSavedCorrectToolUsesParams
    ) -> ToolResult:
        """Search for similar tool usage patterns."""
        try:
            results = await context.agent_memory.search_similar_usage(
                question=args.question,
                context=context,
                limit=args.limit or 10,
                similarity_threshold=args.similarity_threshold or 0.7,
                tool_name_filter=args.tool_name_filter,
            )

            if not results:
                no_results_msg = (
                    "No similar tool usage patterns found for this question."
                )

                # Check if user has access to detailed memory results
                ui_features_available = context.metadata.get(
                    "ui_features_available", []
                )
                show_detailed_results = (
                    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS
                    in ui_features_available
                )

                # Create UI component based on access level
                if show_detailed_results:
                    # Admin view: Show card indicating 0 results
                    ui_component = UiComponent(
                        rich_component=CardComponent(
                            title="🧠 Memory Search: 0 Results",
                            content="No similar tool usage patterns found for this question.\n\nSearched agent memory with no matches.",
                            icon="🔍",
                            status="info",
                            collapsible=True,
                            collapsed=True,
                            markdown=True,
                        ),
                        simple_component=None,
                    )
                else:
                    # Non-admin view: Simple status message
                    ui_component = UiComponent(
                        rich_component=StatusBarUpdateComponent(
                            status="idle",
                            message="No similar patterns found",
                            detail="Searched agent memory",
                        ),
                        simple_component=None,
                    )

                return ToolResult(
                    success=True,
                    result_for_llm=no_results_msg,
                    ui_component=ui_component,
                )

            # Format results for LLM
            results_text = f"Found {len(results)} similar tool usage pattern(s):\n\n"
            for i, result in enumerate(results, 1):
                memory = result.memory
                results_text += f"{i}. {memory.tool_name} (similarity: {result.similarity_score:.2f})\n"
                results_text += f"   Question: {memory.question}\n"
                results_text += f"   Args: {memory.args}\n\n"

            logger.info(f"Agent memory search results: {results_text.strip()}")

            # Check if user has access to detailed memory results
            ui_features_available = context.metadata.get("ui_features_available", [])
            show_detailed_results = (
                UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS
                in ui_features_available
            )

            # Create UI component based on access level
            if show_detailed_results:
                # Admin view: Show detailed results in collapsible card
                detailed_content = "**Retrieved memories passed to LLM:**\n\n"
                for i, result in enumerate(results, 1):
                    memory = result.memory
                    detailed_content += f"**{i}. {memory.tool_name}** (similarity: {result.similarity_score:.2f})\n"
                    detailed_content += f"- **Question:** {memory.question}\n"
                    detailed_content += f"- **Arguments:** `{memory.args}`\n"
                    if memory.timestamp:
                        detailed_content += f"- **Timestamp:** {memory.timestamp}\n"
                    if memory.memory_id:
                        detailed_content += f"- **ID:** `{memory.memory_id}`\n"
                    detailed_content += "\n"

                ui_component = UiComponent(
                    rich_component=CardComponent(
                        title=f"🧠 Memory Search: {len(results)} Result(s)",
                        content=detailed_content.strip(),
                        icon="🔍",
                        status="info",
                        collapsible=True,
                        collapsed=True,  # Start collapsed to avoid clutter
                        markdown=True,  # Render content as markdown
                    ),
                    simple_component=None,
                )
            else:
                # Non-admin view: Simple status message
                ui_component = UiComponent(
                    rich_component=StatusBarUpdateComponent(
                        status="success",
                        message=f"Found {len(results)} similar pattern(s)",
                        detail="Retrieved from agent memory",
                    ),
                    simple_component=None,
                )

            return ToolResult(
                success=True,
                result_for_llm=results_text.strip(),
                ui_component=ui_component,
            )

        except Exception as e:
            error_message = f"Failed to search memories: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=StatusBarUpdateComponent(
                        status="error", message="Failed to search memory", detail=str(e)
                    ),
                    simple_component=None,
                ),
                error=str(e),
            )


class SaveTextMemoryTool(Tool[SaveTextMemoryParams]):
    """Tool for saving free-form text memories."""

    @property
    def name(self) -> str:
        return "save_text_memory"

    @property
    def description(self) -> str:
        return "Save free-form text memory for important insights, observations, or context"

    def get_args_schema(self) -> Type[SaveTextMemoryParams]:
        return SaveTextMemoryParams

    async def execute(
        self, context: ToolContext, args: SaveTextMemoryParams
    ) -> ToolResult:
        """Save a text memory to agent memory."""
        try:
            text_memory = await context.agent_memory.save_text_memory(
                content=args.content, context=context
            )

            success_msg = (
                f"Successfully saved text memory with ID: {text_memory.memory_id}"
            )
            return ToolResult(
                success=True,
                result_for_llm=success_msg,
                ui_component=UiComponent(
                    rich_component=StatusBarUpdateComponent(
                        status="success",
                        message="Saved text memory",
                        detail=f"ID: {text_memory.memory_id}",
                    ),
                    simple_component=None,
                ),
            )

        except Exception as e:
            error_message = f"Failed to save text memory: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=StatusBarUpdateComponent(
                        status="error",
                        message="Failed to save text memory",
                        detail=str(e),
                    ),
                    simple_component=None,
                ),
                error=str(e),
            )


================================================
FILE: src/vanna/tools/file_system.py
================================================
"""
File system tools with dependency injection support.

This module provides file system operations through an abstract FileSystem interface,
allowing for different implementations (local, remote, sandboxed, etc.).
The tools accept a FileSystem instance via dependency injection.
"""

import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import Any, List, Optional, Type
import difflib
import hashlib

from pydantic import BaseModel, Field, model_validator

from vanna.core.tool import Tool, ToolContext, ToolResult
from vanna.components import (
    UiComponent,
    CardComponent,
    NotificationComponent,
    ComponentType,
    SimpleTextComponent,
)

MAX_SEARCH_FILE_BYTES = 1_000_000
FILENAME_MATCH_SNIPPET = "[filename match]"


@dataclass
class FileSearchMatch:
    """Represents a single search result within a file system."""

    path: str
    snippet: Optional[str] = None


@dataclass
class CommandResult:
    """Represents the result of executing a shell command."""

    stdout: str
    stderr: str
    returncode: int


def _make_snippet(text: str, query: str, context_window: int = 60) -> Optional[str]:
    """Return a short snippet around the first occurrence of query in text."""

    lowered = text.lower()
    index = lowered.find(query.lower())
    if index == -1:
        return None

    start = max(0, index - context_window)
    end = min(len(text), index + len(query) + context_window)
    snippet = text[start:end].replace("\n", " ").strip()

    if start > 0:
        snippet = f"…{snippet}"
    if end < len(text):
        snippet = f"{snippet}…"

    return snippet


class FileSystem(ABC):
    """Abstract base class for file system operations."""

    @abstractmethod
    async def list_files(self, directory: str, context: ToolContext) -> List[str]:
        """List files in a directory."""
        pass

    @abstractmethod
    async def read_file(self, filename: str, context: ToolContext) -> str:
        """Read the contents of a file."""
        pass

    @abstractmethod
    async def write_file(
        self, filename: str, content: str, context: ToolContext, overwrite: bool = False
    ) -> None:
        """Write content to a file."""
        pass

    @abstractmethod
    async def exists(self, path: str, context: ToolContext) -> bool:
        """Check if a file or directory exists."""
        pass

    @abstractmethod
    async def is_directory(self, path: str, context: ToolContext) -> bool:
        """Check if a path is a directory."""
        pass

    @abstractmethod
    async def search_files(
        self,
        query: str,
        context: ToolContext,
        *,
        max_results: int = 20,
        include_content: bool = False,
    ) -> List[FileSearchMatch]:
        """Search for files matching a query within the accessible namespace."""
        pass

    @abstractmethod
    async def run_bash(
        self,
        command: str,
        context: ToolContext,
        *,
        timeout: Optional[float] = None,
    ) -> CommandResult:
        """Execute a bash command within the accessible namespace."""
        pass


class LocalFileSystem(FileSystem):
    """Local file system implementation with per-user isolation."""

    def __init__(self, working_directory: str = "."):
        """Initialize with a working directory.

        Args:
            working_directory: Base directory where user-specific folders will be created
        """
        self.working_directory = Path(working_directory)

    def _get_user_directory(self, context: ToolContext) -> Path:
        """Get the user-specific directory by hashing the user ID.

        Args:
            context: Tool context containing user information

        Returns:
            Path to the user-specific directory
        """
        # Hash the user ID to create a directory name
        user_hash = hashlib.sha256(context.user.id.encode()).hexdigest()[:16]
        user_dir = self.working_directory / user_hash

        # Create the directory if it doesn't exist
        user_dir.mkdir(parents=True, exist_ok=True)

        return user_dir

    def _resolve_path(self, path: str, context: ToolContext) -> Path:
        """Resolve a path relative to the user's directory.

        Args:
            path: Path relative to user directory
            context: Tool context containing user information

        Returns:
            Absolute path within user's directory
        """
        user_dir = self._get_user_directory(context)
        resolved = user_dir / path

        # Ensure the path is within the user's directory (prevent directory traversal)
        try:
            resolved.resolve().relative_to(user_dir.resolve())
        except ValueError:
            raise PermissionError(
                f"Access denied: path '{path}' is outside user directory"
            )

        return resolved

    async def list_files(self, directory: str, context: ToolContext) -> List[str]:
        """List files in a directory within the user's isolated space."""
        directory_path = self._resolve_path(directory, context)

        if not directory_path.exists():
            raise FileNotFoundError(f"Directory '{directory}' does not exist")

        if not directory_path.is_dir():
            raise NotADirectoryError(f"'{directory}' is not a directory")

        files = []
        for item in directory_path.iterdir():
            if item.is_file():
                files.append(item.name)

        return sorted(files)

    async def read_file(self, filename: str, context: ToolContext) -> str:
        """Read the contents of a file within the user's isolated space."""
        file_path = self._resolve_path(filename, context)

        if not file_path.exists():
            raise FileNotFoundError(f"File '{filename}' does not exist")

        if not file_path.is_file():
            raise IsADirectoryError(f"'{filename}' is a directory, not a file")

        return file_path.read_text(encoding="utf-8")

    async def write_file(
        self, filename: str, content: str, context: ToolContext, overwrite: bool = False
    ) -> None:
        """Write content to a file within the user's isolated space."""
        file_path = self._resolve_path(filename, context)

        # Create parent directories if they don't exist
        file_path.parent.mkdir(parents=True, exist_ok=True)

        if file_path.exists() and not overwrite:
            raise FileExistsError(
                f"File '{filename}' already exists. Use overwrite=True to replace it."
            )

        file_path.write_text(content, encoding="utf-8")

    async def exists(self, path: str, context: ToolContext) -> bool:
        """Check if a file or directory exists within the user's isolated space."""
        try:
            resolved_path = self._resolve_path(path, context)
            return resolved_path.exists()
        except PermissionError:
            return False

    async def is_directory(self, path: str, context: ToolContext) -> bool:
        """Check if a path is a directory within the user's isolated space."""
        try:
            resolved_path = self._resolve_path(path, context)
            return resolved_path.exists() and resolved_path.is_dir()
        except PermissionError:
            return False

    async def search_files(
        self,
        query: str,
        context: ToolContext,
        *,
        max_results: int = 20,
        include_content: bool = False,
    ) -> List[FileSearchMatch]:
        """Search for files within the user's isolated space."""

        trimmed_query = query.strip()
        if not trimmed_query:
            raise ValueError("Search query must not be empty")

        user_dir = self._get_user_directory(context)
        matches: List[FileSearchMatch] = []
        query_lower = trimmed_query.lower()

        for path in user_dir.rglob("*"):
            if len(matches) >= max_results:
                break

            if not path.is_file():
                continue

            relative_path = path.relative_to(user_dir).as_posix()
            include_entry = False
            snippet: Optional[str] = None

            if query_lower in path.name.lower():
                include_entry = True
                snippet = FILENAME_MATCH_SNIPPET

            content: Optional[str] = None
            if include_content:
                try:
                    size = path.stat().st_size
                except OSError:
                    if include_entry:
                        matches.append(
                            FileSearchMatch(path=relative_path, snippet=snippet)
                        )
                    continue

                if size <= MAX_SEARCH_FILE_BYTES:
                    try:
                        content = path.read_text(encoding="utf-8")
                    except (UnicodeDecodeError, OSError):
                        content = None
                elif not include_entry:
                    # Skip oversized files if they do not match by name
                    continue

            if include_content and content is not None:
                if query_lower in content.lower():
                    snippet = _make_snippet(content, trimmed_query) or snippet
                    include_entry = True
                elif not include_entry:
                    continue

            if include_entry:
                matches.append(FileSearchMatch(path=relative_path, snippet=snippet))

        return matches

    async def run_bash(
        self,
        command: str,
        context: ToolContext,
        *,
        timeout: Optional[float] = None,
    ) -> CommandResult:
        """Execute a bash command within the user's isolated space."""

        if not command.strip():
            raise ValueError("Command must not be empty")

        user_dir = self._get_user_directory(context)

        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(user_dir),
        )

        try:
            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                process.communicate(), timeout=timeout
            )
        except asyncio.TimeoutError as exc:
            process.kill()
            await process.wait()
            raise TimeoutError(f"Command timed out after {timeout} seconds") from exc

        stdout = stdout_bytes.decode("utf-8", errors="replace")
        stderr = stderr_bytes.decode("utf-8", errors="replace")

        return CommandResult(
            stdout=stdout, stderr=stderr, returncode=process.returncode or 0
        )


class SearchFilesArgs(BaseModel):
    """Arguments for searching files."""

    query: str = Field(description="Text to search for in file names or contents")
    include_content: bool = Field(
        default=True,
        description="Whether to search within file contents in addition to file names",
    )
    max_results: int = Field(
        default=20,
        ge=1,
        le=100,
        description="Maximum number of matches to return",
    )


class SearchFilesTool(Tool[SearchFilesArgs]):
    """Tool to search for files using the injected file system implementation."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "search_files"

    @property
    def description(self) -> str:
        return "Search for files by name or content"

    def get_args_schema(self) -> Type[SearchFilesArgs]:
        return SearchFilesArgs

    async def execute(self, context: ToolContext, args: SearchFilesArgs) -> ToolResult:
        try:
            matches = await self.file_system.search_files(
                args.query,
                context,
                max_results=args.max_results,
                include_content=args.include_content,
            )
        except Exception as exc:
            error_msg = f"Error searching files: {exc}"
            return ToolResult(
                success=False,
                result_for_llm=error_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_msg,
                    ),
                    simple_component=SimpleTextComponent(text=error_msg),
                ),
                error=str(exc),
            )

        if not matches:
            message = f"No matches found for '{args.query}'."
            return ToolResult(
                success=True,
                result_for_llm=message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="info",
                        message=message,
                    ),
                    simple_component=SimpleTextComponent(text=message),
                ),
            )

        lines: List[str] = []
        for match in matches:
            snippet = match.snippet
            if snippet == FILENAME_MATCH_SNIPPET:
                snippet_text = "(matched filename)"
            elif snippet:
                snippet_text = snippet
            else:
                snippet_text = ""

            if snippet_text and len(snippet_text) > 200:
                snippet_text = f"{snippet_text[:197]}…"

            if snippet_text:
                lines.append(f"- {match.path}: {snippet_text}")
            else:
                lines.append(f"- {match.path}")

        summary = f"Found {len(matches)} match(es) for '{args.query}' (max {args.max_results})."
        content = "\n".join(lines)

        return ToolResult(
            success=True,
            result_for_llm=f"{summary}\n{content}",
            ui_component=UiComponent(
                rich_component=CardComponent(
                    type=ComponentType.CARD,
                    title=f"Search results for '{args.query}'",
                    content=content,
                ),
                simple_component=SimpleTextComponent(text=summary),
            ),
        )


class ListFilesArgs(BaseModel):
    """Arguments for listing files."""

    directory: str = Field(
        default=".", description="Directory to list (defaults to current)"
    )


class ListFilesTool(Tool[ListFilesArgs]):
    """Tool to list files in a directory using dependency injection for file system access."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        """Initialize with optional file system dependency."""
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "list_files"

    @property
    def description(self) -> str:
        return "List files in a directory"

    def get_args_schema(self) -> Type[ListFilesArgs]:
        return ListFilesArgs

    async def execute(self, context: ToolContext, args: ListFilesArgs) -> ToolResult:
        try:
            files = await self.file_system.list_files(args.directory, context)

            if not files:
                result = f"No files found in directory '{args.directory}'"
                files_list = "No files found"
            else:
                files_list = "\n".join(f"- {f}" for f in files)
                result = f"Files in '{args.directory}':\n{files_list}"

            return ToolResult(
                success=True,
                result_for_llm=result,
                ui_component=UiComponent(
                    rich_component=CardComponent(
                        type=ComponentType.CARD,
                        title=f"Files in {args.directory}",
                        content=files_list,
                    ),
                    simple_component=SimpleTextComponent(text=result),
                ),
            )
        except Exception as e:
            error_msg = f"Error listing files: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_msg,
                    ),
                    simple_component=SimpleTextComponent(text=error_msg),
                ),
                error=str(e),
            )


class ReadFileArgs(BaseModel):
    """Arguments for reading a file."""

    filename: str = Field(description="Name of the file to read")


class ReadFileTool(Tool[ReadFileArgs]):
    """Tool to read file contents using dependency injection for file system access."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        """Initialize with optional file system dependency."""
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "read_file"

    @property
    def description(self) -> str:
        return "Read the contents of a file"

    def get_args_schema(self) -> Type[ReadFileArgs]:
        return ReadFileArgs

    async def execute(self, context: ToolContext, args: ReadFileArgs) -> ToolResult:
        try:
            content = await self.file_system.read_file(args.filename, context)
            result = f"Content of '{args.filename}':\n\n{content}"

            return ToolResult(
                success=True,
                result_for_llm=result,
                ui_component=UiComponent(
                    rich_component=CardComponent(
                        type=ComponentType.CARD,
                        title=f"Contents of {args.filename}",
                        content=content,
                    ),
                    simple_component=SimpleTextComponent(
                        text=f"File content:\n{content}"
                    ),
                ),
            )
        except Exception as e:
            error_msg = f"Error reading file: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_msg,
                    ),
                    simple_component=SimpleTextComponent(text=error_msg),
                ),
                error=str(e),
            )


class WriteFileArgs(BaseModel):
    """Arguments for writing a file."""

    filename: str = Field(description="Name of the file to write")
    content: str = Field(description="Content to write to the file")
    overwrite: bool = Field(
        default=False, description="Whether to overwrite existing files"
    )


class WriteFileTool(Tool[WriteFileArgs]):
    """Tool to write content to a file using dependency injection for file system access."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        """Initialize with optional file system dependency."""
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "write_file"

    @property
    def description(self) -> str:
        return "Write content to a file"

    def get_args_schema(self) -> Type[WriteFileArgs]:
        return WriteFileArgs

    async def execute(self, context: ToolContext, args: WriteFileArgs) -> ToolResult:
        try:
            await self.file_system.write_file(
                args.filename, args.content, context, args.overwrite
            )

            success_msg = f"Successfully wrote {len(args.content)} characters to '{args.filename}'"

            return ToolResult(
                success=True,
                result_for_llm=success_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="success",
                        message=f"File '{args.filename}' written successfully",
                    ),
                    simple_component=SimpleTextComponent(
                        text=f"Wrote to {args.filename}"
                    ),
                ),
            )
        except Exception as e:
            error_msg = f"Error writing file: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_msg,
                    ),
                    simple_component=SimpleTextComponent(text=error_msg),
                ),
                error=str(e),
            )


class LineEdit(BaseModel):
    """Definition of a single line-based edit operation."""

    start_line: int = Field(
        ge=1, description="First line (1-based) affected by this edit"
    )
    end_line: Optional[int] = Field(
        default=None,
        description=(
            "Last line (1-based, inclusive) to replace. Set to start_line - 1 to insert before start_line. "
            "Defaults to start_line, replacing a single line."
        ),
    )
    new_content: str = Field(
        default="", description="Replacement text (preserves provided newlines)"
    )

    @model_validator(mode="after")
    def validate_line_range(self) -> "LineEdit":
        effective_end = self.start_line if self.end_line is None else self.end_line

        if effective_end < self.start_line - 1:
            raise ValueError("end_line must be >= start_line - 1")

        return self


class EditFileArgs(BaseModel):
    """Arguments for editing one or more sections within a file."""

    filename: str = Field(description="Path to the file to edit")
    edits: List[LineEdit] = Field(
        description="List of edits to apply. Later entries should reference higher line numbers.",
        min_length=1,
    )


class EditFileTool(Tool[EditFileArgs]):
    """Tool to apply line-based edits to an existing file."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "edit_file"

    @property
    def description(self) -> str:
        return "Modify specific lines within a file"

    def get_args_schema(self) -> Type[EditFileArgs]:
        return EditFileArgs

    async def execute(self, context: ToolContext, args: EditFileArgs) -> ToolResult:
        try:
            original_content = await self.file_system.read_file(args.filename, context)
        except Exception as exc:
            error_msg = f"Error loading file '{args.filename}': {exc}"
            return ToolResult(
                success=False,
                result_for_llm=error_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_msg,
                    ),
                    simple_component=SimpleTextComponent(text=error_msg),
                ),
                error=str(exc),
            )

        lines = original_content.splitlines(keepends=True)
        applied_edits: List[str] = []

        # Apply edits starting from the bottom so line numbers remain valid for each operation
        for edit in sorted(args.edits, key=lambda e: e.start_line, reverse=True):
            start_line = edit.start_line
            end_line = edit.end_line if edit.end_line is not None else edit.start_line

            if start_line < 1:
                return self._range_error(
                    args.filename, start_line, end_line, "start_line must be >= 1"
                )

            if end_line < start_line - 1:
                return self._range_error(
                    args.filename,
                    start_line,
                    end_line,
                    "end_line must be >= start_line - 1",
                )

            is_insertion = end_line == start_line - 1

            if not is_insertion and start_line > len(lines):
                return self._range_error(
                    args.filename,
                    start_line,
                    end_line,
                    f"start_line {start_line} is beyond the end of the file (len={len(lines)})",
                )

            if is_insertion:
                if start_line > len(lines) + 1:
                    return self._range_error(
                        args.filename,
                        start_line,
                        end_line,
                        "Cannot insert beyond one line past the end of the file",
                    )
                start_index = min(start_line - 1, len(lines))
                end_index = start_index
            else:
                if end_line > len(lines):
                    return self._range_error(
                        args.filename,
                        start_line,
                        end_line,
                        f"end_line {end_line} is beyond the end of the file (len={len(lines)})",
                    )
                start_index = start_line - 1
                end_index = end_line

            replacement_lines = edit.new_content.splitlines(keepends=True)
            lines[start_index:end_index] = replacement_lines

            if is_insertion:
                inserted_count = len(replacement_lines)
                applied_edits.append(
                    f"Inserted {inserted_count} line(s) at line {start_line}"
                )
            else:
                removed_count = end_line - start_line + 1
                applied_edits.append(
                    f"Replaced lines {start_line}-{end_line} (removed {removed_count} line(s))"
                )

        new_content = "".join(lines)

        if new_content == original_content:
            message = (
                f"No changes applied to '{args.filename}' (content already up to date)."
            )
            return ToolResult(
                success=True,
                result_for_llm=message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="info",
                        message=message,
                    ),
                    simple_component=SimpleTextComponent(text=message),
                ),
            )

        try:
            await self.file_system.write_file(
                args.filename, new_content, context, overwrite=True
            )
        except Exception as exc:
            error_msg = f"Error writing updated contents to '{args.filename}': {exc}"
            return ToolResult(
                success=False,
                result_for_llm=error_msg,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_msg,
                    ),
                    simple_component=SimpleTextComponent(text=error_msg),
                ),
                error=str(exc),
            )

        diff_lines = list(
            difflib.unified_diff(
                original_content.splitlines(),
                new_content.splitlines(),
                fromfile=f"a/{args.filename}",
                tofile=f"b/{args.filename}",
                lineterm="",
            )
        )

        diff_text = (
            "\n".join(diff_lines) if diff_lines else "(No textual diff available)"
        )
        summary = (
            f"Updated '{args.filename}' with {len(args.edits)} edit(s).\n"
            + "\n".join(reversed(applied_edits))
        )

        return ToolResult(
            success=True,
            result_for_llm=f"{summary}\n\n{diff_text}",
            ui_component=UiComponent(
                rich_component=CardComponent(
                    type=ComponentType.CARD,
                    title=f"Edited {args.filename}",
                    content=diff_text,
                ),
                simple_component=SimpleTextComponent(text=summary),
            ),
        )

    def _range_error(
        self, filename: str, start_line: int, end_line: int, message: str
    ) -> ToolResult:
        error_msg = f"Invalid edit range for '{filename}': start_line={start_line}, end_line={end_line}. {message}"
        return ToolResult(
            success=False,
            result_for_llm=error_msg,
            ui_component=UiComponent(
                rich_component=NotificationComponent(
                    type=ComponentType.NOTIFICATION,
                    level="error",
                    message=error_msg,
                ),
                simple_component=SimpleTextComponent(text=error_msg),
            ),
            error=message,
        )


# Convenience function for creating tools with default local file system
def create_file_system_tools(
    file_system: Optional[FileSystem] = None,
) -> List[Tool[Any]]:
    """Create a set of file system tools with optional dependency injection."""
    fs = file_system or LocalFileSystem()
    return [
        ListFilesTool(fs),
        SearchFilesTool(fs),
        ReadFileTool(fs),
        WriteFileTool(fs),
        EditFileTool(fs),
    ]


================================================
FILE: src/vanna/tools/python.py
================================================
"""Python-specific tooling built on top of the file system service."""

from __future__ import annotations

import shlex
import sys
from typing import Any, List, Optional, Sequence, Type

from pydantic import BaseModel, Field

from vanna.components import (
    UiComponent,
    CardComponent,
    ComponentType,
    NotificationComponent,
    SimpleTextComponent,
)
from vanna.core.tool import Tool, ToolContext, ToolResult

from .file_system import CommandResult, FileSystem, LocalFileSystem

MAX_OUTPUT_LENGTH = 4000


class RunPythonFileArgs(BaseModel):
    """Arguments required to execute a Python file."""

    filename: str = Field(
        description="Python file to execute (relative to the workspace root)"
    )
    arguments: Sequence[str] = Field(
        default_factory=list,
        description="Optional arguments to pass to the Python script",
    )
    timeout_seconds: Optional[float] = Field(
        default=None,
        ge=0,
        description="Optional timeout for the command in seconds",
    )


class RunPythonFileTool(Tool[RunPythonFileArgs]):
    """Execute a Python file using the provided file system service."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "run_python_file"

    @property
    def description(self) -> str:
        return "Execute a Python file using the workspace interpreter"

    def get_args_schema(self) -> Type[RunPythonFileArgs]:
        return RunPythonFileArgs

    async def execute(
        self, context: ToolContext, args: RunPythonFileArgs
    ) -> ToolResult:
        exists = await self.file_system.exists(args.filename, context)
        if not exists:
            message = f"Cannot execute '{args.filename}' because it does not exist."
            return _error_result(message)

        command_parts = [sys.executable, args.filename]
        command_parts.extend(args.arguments)
        command = _quote_command(command_parts)

        try:
            result = await self.file_system.run_bash(
                command,
                context,
                timeout=args.timeout_seconds,
            )
        except TimeoutError as exc:
            message = str(exc)
            return _error_result(message)

        summary = f"Executed python {args.filename} (exit code {result.returncode})."
        success = result.returncode == 0
        return _result_from_command(summary, command, result, success=success)


class PipInstallArgs(BaseModel):
    """Arguments required to run pip install."""

    packages: List[str] = Field(
        description="Packages (with optional specifiers) to install", min_length=1
    )
    upgrade: bool = Field(
        default=False,
        description="Whether to include --upgrade in the pip invocation",
    )
    extra_args: Sequence[str] = Field(
        default_factory=list,
        description="Additional arguments to pass to pip install",
    )
    timeout_seconds: Optional[float] = Field(
        default=None,
        ge=0,
        description="Optional timeout for the command in seconds",
    )


class PipInstallTool(Tool[PipInstallArgs]):
    """Install Python packages using pip inside the workspace environment."""

    def __init__(self, file_system: Optional[FileSystem] = None):
        self.file_system = file_system or LocalFileSystem()

    @property
    def name(self) -> str:
        return "pip_install"

    @property
    def description(self) -> str:
        return "Install Python packages using pip"

    def get_args_schema(self) -> Type[PipInstallArgs]:
        return PipInstallArgs

    async def execute(self, context: ToolContext, args: PipInstallArgs) -> ToolResult:
        command_parts = [sys.executable, "-m", "pip", "install"]
        if args.upgrade:
            command_parts.append("--upgrade")
        command_parts.extend(args.packages)
        command_parts.extend(args.extra_args)
        command = _quote_command(command_parts)

        try:
            result = await self.file_system.run_bash(
                command,
                context,
                timeout=args.timeout_seconds,
            )
        except TimeoutError as exc:
            return _error_result(str(exc))

        success = result.returncode == 0
        summary = (
            "pip install completed successfully"
            if success
            else f"pip install failed (exit code {result.returncode})."
        )

        return _result_from_command(summary, command, result, success=success)


def create_python_tools(file_system: Optional[FileSystem] = None) -> List[Tool[Any]]:
    """Create Python-specific tools backed by a shared file system service."""

    fs = file_system or LocalFileSystem()
    return [
        RunPythonFileTool(fs),
        PipInstallTool(fs),
    ]


def _quote_command(parts: Sequence[str]) -> str:
    return " ".join(shlex.quote(part) for part in parts)


def _truncate(text: str, limit: int = MAX_OUTPUT_LENGTH) -> str:
    if len(text) <= limit:
        return text
    return f"{text[: limit - 1]}…"


def _result_from_command(
    summary: str,
    command: str,
    result: CommandResult,
    *,
    success: bool = True,
) -> ToolResult:
    stdout = result.stdout.strip()
    stderr = result.stderr.strip()

    blocks: List[str] = [f"$ {command}"]
    if stdout:
        blocks.append("STDOUT:\n" + _truncate(stdout))
    if stderr:
        blocks.append("STDERR:\n" + _truncate(stderr))
    if not stdout and not stderr:
        blocks.append("(no output)")

    content = "\n\n".join(blocks)
    card_status = "success" if success else "error"
    component = CardComponent(
        type=ComponentType.CARD,
        title="Command Result",
        content=content,
        status=card_status,
    )

    return ToolResult(
        success=success,
        result_for_llm=f"{summary}\n\n{content}",
        ui_component=UiComponent(
            rich_component=component,
            simple_component=SimpleTextComponent(text=summary),
        ),
        error=None if success else content,
    )


def _error_result(message: str) -> ToolResult:
    return ToolResult(
        success=False,
        result_for_llm=message,
        ui_component=UiComponent(
            rich_component=NotificationComponent(
                type=ComponentType.NOTIFICATION,
                level="error",
                message=message,
            ),
            simple_component=SimpleTextComponent(text=message),
        ),
        error=message,
    )


================================================
FILE: src/vanna/tools/run_sql.py
================================================
"""Generic SQL query execution tool with dependency injection."""

from typing import Any, Dict, List, Optional, Type, cast
import uuid
from vanna.core.tool import Tool, ToolContext, ToolResult
from vanna.components import (
    UiComponent,
    DataFrameComponent,
    NotificationComponent,
    ComponentType,
    SimpleTextComponent,
)
from vanna.capabilities.sql_runner import SqlRunner, RunSqlToolArgs
from vanna.capabilities.file_system import FileSystem
from vanna.integrations.local import LocalFileSystem


class RunSqlTool(Tool[RunSqlToolArgs]):
    """Tool that executes SQL queries using an injected SqlRunner implementation."""

    def __init__(
        self,
        sql_runner: SqlRunner,
        file_system: Optional[FileSystem] = None,
        custom_tool_name: Optional[str] = None,
        custom_tool_description: Optional[str] = None,
    ):
        """Initialize the tool with a SqlRunner implementation.

        Args:
            sql_runner: SqlRunner implementation that handles actual query execution
            file_system: FileSystem implementation for saving results (defaults to LocalFileSystem)
            custom_tool_name: Optional custom name for the tool (overrides default "run_sql")
            custom_tool_description: Optional custom description for the tool (overrides default description)
        """
        self.sql_runner = sql_runner
        self.file_system = file_system or LocalFileSystem()
        self._custom_name = custom_tool_name
        self._custom_description = custom_tool_description

    @property
    def name(self) -> str:
        return self._custom_name if self._custom_name else "run_sql"

    @property
    def description(self) -> str:
        return (
            self._custom_description
            if self._custom_description
            else "Execute SQL queries against the configured database"
        )

    def get_args_schema(self) -> Type[RunSqlToolArgs]:
        return RunSqlToolArgs

    async def execute(self, context: ToolContext, args: RunSqlToolArgs) -> ToolResult:
        """Execute a SQL query using the injected SqlRunner."""
        try:
            # Use the injected SqlRunner to execute the query
            df = await self.sql_runner.run_sql(args, context)

            # Determine query type
            query_type = args.sql.strip().upper().split()[0]

            if query_type == "SELECT":
                # Handle SELECT queries with results
                if df.empty:
                    result = "Query executed successfully. No rows returned."
                    ui_component = UiComponent(
                        rich_component=DataFrameComponent(
                            rows=[],
                            columns=[],
                            title="Query Results",
                            description="No rows returned",
                        ),
                        simple_component=SimpleTextComponent(text=result),
                    )
                    metadata = {
                        "row_count": 0,
                        "columns": [],
                        "query_type": query_type,
                        "results": [],
                    }
                else:
                    # Convert DataFrame to records
                    results_data = df.to_dict("records")
                    columns = df.columns.tolist()
                    row_count = len(df)

                    # Write DataFrame to CSV file for downstream tools
                    file_id = str(uuid.uuid4())[:8]
                    filename = f"query_results_{file_id}.csv"
                    csv_content = df.to_csv(index=False)
                    await self.file_system.write_file(
                        filename, csv_content, context, overwrite=True
                    )

                    # Create result text for LLM with truncated results
                    results_preview = csv_content
                    if len(results_preview) > 1000:
                        results_preview = (
                            results_preview[:1000]
                            + "\n(Results truncated to 1000 characters. FOR LARGE RESULTS YOU DO NOT NEED TO SUMMARIZE THESE RESULTS OR PROVIDE OBSERVATIONS. THE NEXT STEP SHOULD BE A VISUALIZE_DATA CALL)"
                        )

                    result = f"{results_preview}\n\nResults saved to file: {filename}\n\n**IMPORTANT: FOR VISUALIZE_DATA USE FILENAME: {filename}**"

                    # Create DataFrame component for UI
                    dataframe_component = DataFrameComponent.from_records(
                        records=cast(List[Dict[str, Any]], results_data),
                        title="Query Results",
                        description=f"SQL query returned {row_count} rows with {len(columns)} columns",
                    )

                    ui_component = UiComponent(
                        rich_component=dataframe_component,
                        simple_component=SimpleTextComponent(text=result),
                    )

                    metadata = {
                        "row_count": row_count,
                        "columns": columns,
                        "query_type": query_type,
                        "results": results_data,
                        "output_file": filename,
                    }
            else:
                # For non-SELECT queries (INSERT, UPDATE, DELETE, etc.)
                # The SqlRunner should return a DataFrame with affected row count
                rows_affected = len(df) if not df.empty else 0
                result = (
                    f"Query executed successfully. {rows_affected} row(s) affected."
                )

                metadata = {"rows_affected": rows_affected, "query_type": query_type}
                ui_component = UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION, level="success", message=result
                    ),
                    simple_component=SimpleTextComponent(text=result),
                )

            return ToolResult(
                success=True,
                result_for_llm=result,
                ui_component=ui_component,
                metadata=metadata,
            )

        except Exception as e:
            error_message = f"Error executing query: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_message,
                    ),
                    simple_component=SimpleTextComponent(text=error_message),
                ),
                error=str(e),
                metadata={"error_type": "sql_error"},
            )


================================================
FILE: src/vanna/tools/visualize_data.py
================================================
"""Tool for visualizing DataFrame data from CSV files."""

from typing import Optional, Type
import logging
import pandas as pd
from pydantic import BaseModel, Field

from vanna.core.tool import Tool, ToolContext, ToolResult
from vanna.components import (
    UiComponent,
    ChartComponent,
    NotificationComponent,
    ComponentType,
    SimpleTextComponent,
)

from .file_system import FileSystem, LocalFileSystem
from vanna.integrations.plotly import PlotlyChartGenerator

logger = logging.getLogger(__name__)


class VisualizeDataArgs(BaseModel):
    """Arguments for visualize_data tool."""

    filename: str = Field(description="Name of the CSV file to visualize")
    title: Optional[str] = Field(
        default=None, description="Optional title for the chart"
    )


class VisualizeDataTool(Tool[VisualizeDataArgs]):
    """Tool that reads CSV files and generates visualizations using dependency injection."""

    def __init__(
        self,
        file_system: Optional[FileSystem] = None,
        plotly_generator: Optional[PlotlyChartGenerator] = None,
    ):
        """Initialize the tool with FileSystem and PlotlyChartGenerator.

        Args:
            file_system: FileSystem implementation for reading CSV files (defaults to LocalFileSystem)
            plotly_generator: PlotlyChartGenerator for creating Plotly charts (defaults to PlotlyChartGenerator())
        """
        self.file_system = file_system or LocalFileSystem()
        self.plotly_generator = plotly_generator or PlotlyChartGenerator()

    @property
    def name(self) -> str:
        return "visualize_data"

    @property
    def description(self) -> str:
        return "Create a visualization from a CSV file. The tool automatically selects an appropriate chart type based on the data."

    def get_args_schema(self) -> Type[VisualizeDataArgs]:
        return VisualizeDataArgs

    async def execute(
        self, context: ToolContext, args: VisualizeDataArgs
    ) -> ToolResult:
        """Read CSV file and generate visualization."""
        try:
            logger.info(f"Starting visualization for file: {args.filename}")

            # Read the CSV file using FileSystem
            csv_content = await self.file_system.read_file(args.filename, context)
            logger.info(f"Read {len(csv_content)} bytes from CSV file")

            # Parse CSV into DataFrame
            import io

            df = pd.read_csv(io.StringIO(csv_content))
            logger.info(
                f"Parsed DataFrame with shape {df.shape}, columns: {df.columns.tolist()}, dtypes: {df.dtypes.to_dict()}"
            )

            # Generate title
            title = args.title or f"Visualization of {args.filename}"

            # Generate chart using PlotlyChartGenerator
            logger.info("Generating chart...")
            chart_dict = self.plotly_generator.generate_chart(df, title)
            logger.info(
                f"Chart generated, type: {type(chart_dict)}, keys: {list(chart_dict.keys()) if isinstance(chart_dict, dict) else 'N/A'}"
            )

            # Create result message
            row_count = len(df)
            col_count = len(df.columns)
            result = f"Created visualization from '{args.filename}' ({row_count} rows, {col_count} columns)."

            # Create ChartComponent
            logger.info("Creating ChartComponent...")
            chart_component = ChartComponent(
                chart_type="plotly",
                data=chart_dict,
                title=title,
                config={
                    "data_shape": {"rows": row_count, "columns": col_count},
                    "source_file": args.filename,
                },
            )
            logger.info("ChartComponent created successfully")

            logger.info("Creating ToolResult...")
            tool_result = ToolResult(
                success=True,
                result_for_llm=result,
                ui_component=UiComponent(
                    rich_component=chart_component,
                    simple_component=SimpleTextComponent(text=result),
                ),
                metadata={
                    "filename": args.filename,
                    "rows": row_count,
                    "columns": col_count,
                    "chart": chart_dict,
                },
            )
            logger.info("ToolResult created successfully")
            return tool_result

        except FileNotFoundError as e:
            logger.error(f"File not found: {args.filename}", exc_info=True)
            error_message = f"File not found: {args.filename}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_message,
                    ),
                    simple_component=SimpleTextComponent(text=error_message),
                ),
                error=str(e),
                metadata={"error_type": "file_not_found"},
            )
        except pd.errors.ParserError as e:
            logger.error(f"CSV parse error for {args.filename}", exc_info=True)
            error_message = f"Failed to parse CSV file '{args.filename}': {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_message,
                    ),
                    simple_component=SimpleTextComponent(text=error_message),
                ),
                error=str(e),
                metadata={"error_type": "csv_parse_error"},
            )
        except ValueError as e:
            logger.error(f"Visualization error for {args.filename}", exc_info=True)
            error_message = f"Cannot visualize data: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_message,
                    ),
                    simple_component=SimpleTextComponent(text=error_message),
                ),
                error=str(e),
                metadata={"error_type": "visualization_error"},
            )
        except Exception as e:
            logger.error(
                f"Unexpected error creating visualization for {args.filename}",
                exc_info=True,
            )
            error_message = f"Error creating visualization: {str(e)}"
            return ToolResult(
                success=False,
                result_for_llm=error_message,
                ui_component=UiComponent(
                    rich_component=NotificationComponent(
                        type=ComponentType.NOTIFICATION,
                        level="error",
                        message=error_message,
                    ),
                    simple_component=SimpleTextComponent(text=error_message),
                ),
                error=str(e),
                metadata={"error_type": "general_error"},
            )


================================================
FILE: src/vanna/utils/__init__.py
================================================


================================================
FILE: src/vanna/web_components/__init__.py
================================================
"""
Web components for Vanna Agents.

This module provides web components built with Lit that can be embedded
in web applications to provide rich UI for Vanna agent interactions.
"""

import os
from pathlib import Path
from typing import Dict


def get_component_files() -> Dict[str, Path]:
    """Get paths to all web component files."""
    component_dir = Path(__file__).parent
    return {
        "js": component_dir / "index.js",
        "css": component_dir / "style.css",
    }


def get_component_html() -> str:
    """Get HTML template for including components."""
    files = get_component_files()

    html = """
<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Vanna AI Chat</title>
</head>
<body>
    <vanna-chat title="Vanna AI Assistant"></vanna-chat>
    <script type="module" src="{js_file}"></script>
</body>
</html>
""".format(js_file=files["js"].name)

    return html


__all__ = ["get_component_files", "get_component_html"]


================================================
FILE: tests/conftest.py
================================================
"""
Pytest configuration and shared fixtures for Vanna v2 test suite.
"""

import os
import pytest
import sqlite3
from pathlib import Path

# Configure pytest-asyncio
pytest_plugins = ["pytest_asyncio"]


def pytest_configure(config):
    """Configure pytest with custom markers."""
    config.addinivalue_line(
        "markers", "anthropic: marks tests requiring Anthropic API key"
    )
    config.addinivalue_line("markers", "openai: marks tests requiring OpenAI API key")
    config.addinivalue_line(
        "markers", "azureopenai: marks tests requiring Azure OpenAI API key"
    )
    config.addinivalue_line("markers", "gemini: marks tests requiring Google API key")
    config.addinivalue_line("markers", "ollama: marks tests requiring Ollama")
    config.addinivalue_line("markers", "chromadb: marks tests requiring ChromaDB")
    config.addinivalue_line("markers", "legacy: marks tests for LegacyVannaAdapter")


def pytest_collection_modifyitems(config, items):
    """Automatically skip tests if required API keys are missing."""
    for item in items:
        # Skip Anthropic tests if no API key
        if "anthropic" in item.keywords:
            if not os.getenv("ANTHROPIC_API_KEY"):
                item.add_marker(
                    pytest.mark.skip(
                        reason="ANTHROPIC_API_KEY environment variable not set"
                    )
                )

        # Skip OpenAI tests if no API key
        if "openai" in item.keywords:
            if not os.getenv("OPENAI_API_KEY"):
                item.add_marker(
                    pytest.mark.skip(
                        reason="OPENAI_API_KEY environment variable not set"
                    )
                )

        # Skip Azure OpenAI tests if no API key
        if "azureopenai" in item.keywords:
            if not os.getenv("AZURE_OPENAI_API_KEY"):
                item.add_marker(
                    pytest.mark.skip(
                        reason="AZURE_OPENAI_API_KEY environment variable not set"
                    )
                )

        # Skip Gemini tests if no API key
        if "gemini" in item.keywords:
            if not (os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")):
                item.add_marker(
                    pytest.mark.skip(
                        reason="GOOGLE_API_KEY or GEMINI_API_KEY environment variable not set"
                    )
                )


@pytest.fixture(scope="session")
def chinook_db(tmp_path_factory):
    """
    Downloads the Chinook SQLite database and returns a SqliteRunner.

    Uses session scope so the database is only downloaded once per test session.
    """
    import httpx
    from vanna.integrations.sqlite import SqliteRunner

    tmp_path = tmp_path_factory.mktemp("data")
    db_path = tmp_path / "Chinook.sqlite"

    # Download the database
    response = httpx.get("https://vanna.ai/Chinook.sqlite")
    response.raise_for_status()

    db_path.write_bytes(response.content)

    return SqliteRunner(database_path=str(db_path))


================================================
FILE: tests/test_agent_memory.py
================================================
"""
Integration tests for AgentMemory implementations.

These tests verify actual functionality against running services.
Tests are marked by implementation and skip if dependencies are unavailable.
"""

import pytest
import asyncio
import tempfile
import shutil
import os
from vanna.core.user import User
from vanna.core.tool import ToolContext
from vanna.integrations.local.agent_memory import DemoAgentMemory


@pytest.fixture
def test_user():
    """Test user for context."""
    return User(
        id="test_user",
        username="test",
        email="test@example.com",
        group_memberships=["user"],
    )


def create_test_context(test_user, agent_memory):
    """Helper to create test context with specific agent memory."""
    return ToolContext(
        user=test_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
        metadata={},
    )


@pytest.fixture
def chromadb_memory():
    """Create ChromaDB memory instance."""
    try:
        from vanna.integrations.chromadb import ChromaAgentMemory

        temp_dir = tempfile.mkdtemp()
        memory = ChromaAgentMemory(
            persist_directory=temp_dir, collection_name="test_memories"
        )

        yield memory

        shutil.rmtree(temp_dir, ignore_errors=True)
    except ImportError:
        pytest.skip("ChromaDB not installed")


@pytest.fixture
def qdrant_memory():
    """Create Qdrant memory instance."""
    try:
        from vanna.integrations.qdrant import QdrantAgentMemory

        temp_dir = tempfile.mkdtemp()
        memory = QdrantAgentMemory(path=temp_dir)

        yield memory

        shutil.rmtree(temp_dir, ignore_errors=True)
    except ImportError:
        pytest.skip("Qdrant not installed")


@pytest.fixture
def faiss_memory():
    """Create FAISS memory instance."""
    try:
        from vanna.integrations.faiss import FAISSAgentMemory

        temp_dir = tempfile.mkdtemp()
        memory = FAISSAgentMemory(persist_path=temp_dir)

        yield memory

        shutil.rmtree(temp_dir, ignore_errors=True)
    except ImportError:
        pytest.skip("FAISS not installed")


# Parametrized tests for local implementations
@pytest.mark.parametrize(
    "memory_fixture", ["chromadb_memory", "qdrant_memory", "faiss_memory"]
)
class TestLocalAgentMemory:
    """Tests for local AgentMemory implementations (ChromaDB, Qdrant, FAISS)."""

    @pytest.mark.asyncio
    async def test_save_and_search(self, memory_fixture, test_user, request):
        """Test saving and searching tool usage patterns."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        await memory.save_tool_usage(
            question="Show me top customers",
            tool_name="run_sql",
            args={"sql": "SELECT * FROM customers ORDER BY total_spent DESC LIMIT 10"},
            context=test_context,
            success=True,
        )

        results = await memory.search_similar_usage(
            question="Show me best customers",
            context=test_context,
            limit=5,
            similarity_threshold=0.5,
        )

        assert len(results) > 0
        assert results[0].memory.question == "Show me top customers"
        assert results[0].memory.tool_name == "run_sql"
        assert results[0].similarity_score > 0.5

    @pytest.mark.asyncio
    async def test_multiple_memories(self, memory_fixture, test_user, request):
        """Test storing and searching multiple memories."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        patterns = [
            ("Show me sales data", "run_sql", {"sql": "SELECT * FROM sales"}),
            ("Get customer list", "run_sql", {"sql": "SELECT * FROM customers"}),
            (
                "Generate sales report",
                "run_sql",
                {"sql": "SELECT date, SUM(amount) FROM sales GROUP BY date"},
            ),
        ]

        for question, tool_name, args in patterns:
            await memory.save_tool_usage(
                question=question,
                tool_name=tool_name,
                args=args,
                context=test_context,
                success=True,
            )

        results = await memory.search_similar_usage(
            question="Show me sales information",
            context=test_context,
            limit=10,
            similarity_threshold=0.3,
        )

        assert len(results) >= 2

    @pytest.mark.asyncio
    async def test_tool_filter(self, memory_fixture, test_user, request):
        """Test filtering by tool name."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        await memory.save_tool_usage(
            question="Query database",
            tool_name="run_sql",
            args={"sql": "SELECT 1"},
            context=test_context,
        )

        await memory.save_tool_usage(
            question="Search files",
            tool_name="search_files",
            args={"pattern": "*.py"},
            context=test_context,
        )

        results = await memory.search_similar_usage(
            question="Query data",
            context=test_context,
            tool_name_filter="run_sql",
            similarity_threshold=0.0,
        )

        assert len(results) > 0
        assert all(r.memory.tool_name == "run_sql" for r in results)

    @pytest.mark.asyncio
    async def test_clear_memories(self, memory_fixture, test_user, request):
        """Test clearing memories."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        await memory.save_tool_usage("Q1", "run_sql", {}, test_context)
        await memory.save_tool_usage("Q2", "run_sql", {}, test_context)
        await memory.save_tool_usage("Q3", "visualize_data", {}, test_context)

        deleted = await memory.clear_memories(context=test_context, tool_name="run_sql")

        assert deleted >= 0

    @pytest.mark.asyncio
    async def test_get_recent_memories(self, memory_fixture, test_user, request):
        """Test getting recent memories."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        await memory.save_tool_usage("Q1", "run_sql", {"sql": "SELECT 1"}, test_context)
        await asyncio.sleep(0.01)
        await memory.save_tool_usage("Q2", "run_sql", {"sql": "SELECT 2"}, test_context)
        await asyncio.sleep(0.01)
        await memory.save_tool_usage(
            "Q3", "visualize_data", {"file": "data.csv"}, test_context
        )
        await asyncio.sleep(0.01)
        await memory.save_tool_usage("Q4", "run_sql", {"sql": "SELECT 3"}, test_context)
        await asyncio.sleep(0.01)
        await memory.save_tool_usage(
            "Q5", "search_files", {"pattern": "*.py"}, test_context
        )

        recent = await memory.get_recent_memories(context=test_context, limit=3)

        assert isinstance(recent, list)
        assert len(recent) <= 3

        if recent:
            assert all(hasattr(m, "memory_id") for m in recent)

    @pytest.mark.asyncio
    async def test_delete_by_id(self, memory_fixture, test_user, request):
        """Test deleting memories by ID."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        await memory.save_tool_usage("Q1", "run_sql", {"sql": "SELECT 1"}, test_context)
        await memory.save_tool_usage("Q2", "run_sql", {"sql": "SELECT 2"}, test_context)
        await memory.save_tool_usage(
            "Q3", "visualize_data", {"file": "data.csv"}, test_context
        )

        recent = await memory.get_recent_memories(context=test_context, limit=10)

        if not recent:
            pytest.skip("Implementation doesn't support get_recent_memories")

        assert all(m.memory_id is not None for m in recent)

        memory_to_delete = recent[0]
        deleted = await memory.delete_by_id(
            context=test_context, memory_id=memory_to_delete.memory_id
        )

        assert deleted is True

        remaining = await memory.get_recent_memories(context=test_context, limit=10)
        assert all(m.memory_id != memory_to_delete.memory_id for m in remaining)

        fake_deleted = await memory.delete_by_id(
            context=test_context, memory_id="fake-id-12345"
        )
        assert fake_deleted is False

    @pytest.mark.asyncio
    async def test_save_and_search_text_memories(
        self, memory_fixture, test_user, request
    ):
        """Test saving and searching text memories."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        # Save a text memory
        text_memory = await memory.save_text_memory(
            content="The status column uses 1 for active, 0 for inactive",
            context=test_context,
        )

        assert text_memory.memory_id is not None
        assert (
            text_memory.content == "The status column uses 1 for active, 0 for inactive"
        )

        # Search for similar text memories
        results = await memory.search_text_memories(
            query="status column meaning",
            context=test_context,
            limit=5,
            similarity_threshold=0.3,
        )

        assert len(results) > 0
        assert (
            results[0].memory.content
            == "The status column uses 1 for active, 0 for inactive"
        )
        assert results[0].similarity_score > 0.3

    @pytest.mark.asyncio
    async def test_multiple_text_memories(self, memory_fixture, test_user, request):
        """Test storing and searching multiple text memories."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        text_contents = [
            "The fiscal year starts in April",
            "MRR means Monthly Recurring Revenue",
            "Always exclude test accounts where email contains 'test'",
        ]

        for content in text_contents:
            await memory.save_text_memory(content=content, context=test_context)

        # Search for fiscal year info
        results = await memory.search_text_memories(
            query="When does the fiscal year start?",
            context=test_context,
            limit=10,
            similarity_threshold=0.2,
        )

        assert len(results) >= 1
        assert any("fiscal year" in r.memory.content.lower() for r in results)

    @pytest.mark.asyncio
    async def test_get_recent_text_memories(self, memory_fixture, test_user, request):
        """Test getting recent text memories."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        await memory.save_text_memory("First text memory", test_context)
        await asyncio.sleep(0.01)
        await memory.save_text_memory("Second text memory", test_context)
        await asyncio.sleep(0.01)
        await memory.save_text_memory("Third text memory", test_context)

        recent = await memory.get_recent_text_memories(context=test_context, limit=2)

        assert isinstance(recent, list)
        assert len(recent) <= 2

        if recent:
            assert all(hasattr(m, "memory_id") for m in recent)
            assert all(hasattr(m, "content") for m in recent)

    @pytest.mark.asyncio
    async def test_delete_text_memory(self, memory_fixture, test_user, request):
        """Test deleting text memories by ID."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        text_memory = await memory.save_text_memory(
            "Test memory to delete", test_context
        )

        assert text_memory.memory_id is not None

        deleted = await memory.delete_text_memory(
            context=test_context, memory_id=text_memory.memory_id
        )

        assert deleted is True

        # Verify it's gone
        recent = await memory.get_recent_text_memories(context=test_context, limit=10)
        assert all(m.memory_id != text_memory.memory_id for m in recent)

        # Try deleting non-existent memory
        fake_deleted = await memory.delete_text_memory(
            context=test_context, memory_id="fake-text-id-12345"
        )
        assert fake_deleted is False

    @pytest.mark.asyncio
    async def test_mixed_tool_and_text_memories(
        self, memory_fixture, test_user, request
    ):
        """Test that tool memories and text memories can coexist without errors."""
        memory = request.getfixturevalue(memory_fixture)
        test_context = create_test_context(test_user, memory)

        await memory.clear_memories(context=test_context)

        # Save some tool memories
        await memory.save_tool_usage(
            question="Show me top customers",
            tool_name="run_sql",
            args={"sql": "SELECT * FROM customers"},
            context=test_context,
        )

        # Save some text memories
        await memory.save_text_memory(
            content="The fiscal year starts in April",
            context=test_context,
        )

        await memory.save_tool_usage(
            question="Get sales data",
            tool_name="run_sql",
            args={"sql": "SELECT * FROM sales"},
            context=test_context,
        )

        await memory.save_text_memory(
            content="MRR means Monthly Recurring Revenue",
            context=test_context,
        )

        # This should only return tool memories, not text memories
        tool_memories = await memory.get_recent_memories(context=test_context, limit=10)

        assert isinstance(tool_memories, list)
        assert len(tool_memories) == 2  # Should only have the 2 tool memories
        assert all(hasattr(m, "question") for m in tool_memories)
        assert all(hasattr(m, "tool_name") for m in tool_memories)

        # This should only return text memories, not tool memories
        text_memories = await memory.get_recent_text_memories(
            context=test_context, limit=10
        )

        assert isinstance(text_memories, list)
        assert len(text_memories) == 2  # Should only have the 2 text memories
        assert all(hasattr(m, "content") for m in text_memories)
        assert all(
            "fiscal year" in m.content or "MRR" in m.content for m in text_memories
        )


================================================
FILE: tests/test_agent_memory_sanity.py
================================================
"""
Sanity tests for AgentMemory implementations.

These tests verify that:
1. Each AgentMemory implementation correctly implements the AgentMemory interface
2. Imports are working correctly for all vector store modules
3. Basic class instantiation works (without requiring actual service connections)

Note: These tests do NOT execute actual vector operations against services.
They are lightweight sanity checks for the implementation structure.
"""

import pytest
from inspect import signature, iscoroutinefunction
from abc import ABC


class TestAgentMemoryInterface:
    """Test that the AgentMemory interface is properly defined."""

    def test_agent_memory_import(self):
        """Test that AgentMemory can be imported."""
        from vanna.capabilities.agent_memory import AgentMemory

        assert AgentMemory is not None

    def test_agent_memory_is_abstract(self):
        """Test that AgentMemory is an abstract base class."""
        from vanna.capabilities.agent_memory import AgentMemory

        assert issubclass(AgentMemory, ABC)

    def test_agent_memory_has_required_methods(self):
        """Test that AgentMemory defines all required abstract methods."""
        from vanna.capabilities.agent_memory import AgentMemory

        required_methods = [
            "save_tool_usage",
            "save_text_memory",
            "search_similar_usage",
            "search_text_memories",
            "get_recent_memories",
            "get_recent_text_memories",
            "delete_by_id",
            "delete_text_memory",
            "clear_memories",
        ]

        for method_name in required_methods:
            assert hasattr(AgentMemory, method_name)
            method = getattr(AgentMemory, method_name)
            assert getattr(method, "__isabstractmethod__", False), (
                f"{method_name} should be abstract"
            )

    def test_all_methods_are_async(self):
        """Test that all AgentMemory methods are async."""
        from vanna.capabilities.agent_memory import AgentMemory

        methods = [
            "save_tool_usage",
            "save_text_memory",
            "search_similar_usage",
            "search_text_memories",
            "get_recent_memories",
            "get_recent_text_memories",
            "delete_by_id",
            "delete_text_memory",
            "clear_memories",
        ]

        for method_name in methods:
            method = getattr(AgentMemory, method_name)
            assert iscoroutinefunction(method), f"{method_name} should be async"


class TestToolMemoryModel:
    """Test the ToolMemory model."""

    def test_tool_memory_import(self):
        """Test that ToolMemory can be imported."""
        from vanna.capabilities.agent_memory import ToolMemory

        assert ToolMemory is not None

    def test_tool_memory_is_pydantic_model(self):
        """Test that ToolMemory is a Pydantic model."""
        from vanna.capabilities.agent_memory import ToolMemory
        from pydantic import BaseModel

        assert issubclass(ToolMemory, BaseModel)

    def test_tool_memory_has_required_fields(self):
        """Test that ToolMemory has all required fields."""
        from vanna.capabilities.agent_memory import ToolMemory

        required_fields = ["question", "tool_name", "args"]
        optional_fields = ["memory_id", "timestamp", "success", "metadata"]

        model_fields = list(ToolMemory.model_fields.keys())

        for field in required_fields:
            assert field in model_fields, f"Required field '{field}' missing"

        for field in optional_fields:
            assert field in model_fields, f"Optional field '{field}' missing"

    def test_tool_memory_instantiation(self):
        """Test that ToolMemory can be instantiated."""
        from vanna.capabilities.agent_memory import ToolMemory

        memory = ToolMemory(
            memory_id="test-123",
            question="What is the total sales?",
            tool_name="run_sql",
            args={"sql": "SELECT SUM(amount) FROM sales"},
        )

        assert memory.memory_id == "test-123"
        assert memory.question == "What is the total sales?"
        assert memory.tool_name == "run_sql"
        assert memory.args == {"sql": "SELECT SUM(amount) FROM sales"}
        assert memory.success is True  # Default value


class TestToolMemorySearchResultModel:
    """Test the ToolMemorySearchResult model."""

    def test_memory_search_result_import(self):
        """Test that ToolMemorySearchResult can be imported."""
        from vanna.capabilities.agent_memory import ToolMemorySearchResult

        assert ToolMemorySearchResult is not None

    def test_memory_search_result_instantiation(self):
        """Test that ToolMemorySearchResult can be instantiated."""
        from vanna.capabilities.agent_memory import ToolMemorySearchResult, ToolMemory

        memory = ToolMemory(question="test", tool_name="test_tool", args={})

        result = ToolMemorySearchResult(memory=memory, similarity_score=0.95, rank=1)

        assert result.memory == memory
        assert result.similarity_score == 0.95
        assert result.rank == 1


class TestTextMemoryModel:
    """Test the TextMemory model."""

    def test_text_memory_import(self):
        """Test that TextMemory can be imported."""
        from vanna.capabilities.agent_memory import TextMemory

        assert TextMemory is not None

    def test_text_memory_is_pydantic_model(self):
        """Test that TextMemory is a Pydantic model."""
        from vanna.capabilities.agent_memory import TextMemory
        from pydantic import BaseModel

        assert issubclass(TextMemory, BaseModel)

    def test_text_memory_has_required_fields(self):
        """Test that TextMemory has all required fields."""
        from vanna.capabilities.agent_memory import TextMemory

        required_fields = ["content"]
        optional_fields = ["memory_id", "timestamp"]

        model_fields = list(TextMemory.model_fields.keys())

        for field in required_fields:
            assert field in model_fields, f"Required field '{field}' missing"

        for field in optional_fields:
            assert field in model_fields, f"Optional field '{field}' missing"

    def test_text_memory_instantiation(self):
        """Test that TextMemory can be instantiated."""
        from vanna.capabilities.agent_memory import TextMemory

        memory = TextMemory(
            memory_id="text-123", content="Remember to handle edge cases"
        )

        assert memory.memory_id == "text-123"
        assert memory.content == "Remember to handle edge cases"


class TestTextMemorySearchResultModel:
    """Test the TextMemorySearchResult model."""

    def test_text_memory_search_result_import(self):
        """Test that TextMemorySearchResult can be imported."""
        from vanna.capabilities.agent_memory import TextMemorySearchResult

        assert TextMemorySearchResult is not None

    def test_text_memory_search_result_instantiation(self):
        """Test that TextMemorySearchResult can be instantiated."""
        from vanna.capabilities.agent_memory import TextMemorySearchResult, TextMemory

        memory = TextMemory(content="Example memory")

        result = TextMemorySearchResult(memory=memory, similarity_score=0.88, rank=2)

        assert result.memory == memory
        assert result.similarity_score == 0.88
        assert result.rank == 2


class TestChromaDBAgentMemory:
    """Sanity tests for ChromaDB AgentMemory implementation."""

    def test_chromadb_import(self):
        """Test that ChromaAgentMemory can be imported."""
        try:
            from vanna.integrations.chromadb import ChromaAgentMemory

            assert ChromaAgentMemory is not None
        except ImportError:
            pytest.skip("ChromaDB not installed")

    def test_chromadb_implements_agent_memory(self):
        """Test that ChromaAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.chromadb import ChromaAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(ChromaAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("ChromaDB not installed")

    def test_chromadb_has_all_methods(self):
        """Test that ChromaAgentMemory implements all required methods."""
        try:
            from vanna.integrations.chromadb import ChromaAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(ChromaAgentMemory, method_name)
                method = getattr(ChromaAgentMemory, method_name)
                assert not getattr(method, "__isabstractmethod__", False), (
                    f"{method_name} should be implemented (not abstract)"
                )
        except ImportError:
            pytest.skip("ChromaDB not installed")

    def test_chromadb_instantiation(self):
        """Test that ChromaAgentMemory can be instantiated."""
        try:
            from vanna.integrations.chromadb import ChromaAgentMemory
            import tempfile

            temp_dir = tempfile.mkdtemp()
            memory = ChromaAgentMemory(
                persist_directory=temp_dir, collection_name="test"
            )

            assert memory is not None
            assert memory.persist_directory == temp_dir
            assert memory.collection_name == "test"
        except ImportError:
            pytest.skip("ChromaDB not installed")


class TestQdrantAgentMemory:
    """Sanity tests for Qdrant AgentMemory implementation."""

    def test_qdrant_import(self):
        """Test that QdrantAgentMemory can be imported."""
        try:
            from vanna.integrations.qdrant import QdrantAgentMemory

            assert QdrantAgentMemory is not None
        except ImportError:
            pytest.skip("Qdrant not installed")

    def test_qdrant_implements_agent_memory(self):
        """Test that QdrantAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.qdrant import QdrantAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(QdrantAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("Qdrant not installed")

    def test_qdrant_has_all_methods(self):
        """Test that QdrantAgentMemory implements all required methods."""
        try:
            from vanna.integrations.qdrant import QdrantAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(QdrantAgentMemory, method_name)
        except ImportError:
            pytest.skip("Qdrant not installed")

    def test_qdrant_instantiation(self):
        """Test that QdrantAgentMemory can be instantiated."""
        try:
            from vanna.integrations.qdrant import QdrantAgentMemory

            # In-memory mode doesn't require actual service
            memory = QdrantAgentMemory(path=":memory:")

            assert memory is not None
        except ImportError:
            pytest.skip("Qdrant not installed")


class TestPineconeAgentMemory:
    """Sanity tests for Pinecone AgentMemory implementation."""

    def test_pinecone_import(self):
        """Test that PineconeAgentMemory can be imported."""
        try:
            from vanna.integrations.pinecone import PineconeAgentMemory

            assert PineconeAgentMemory is not None
        except ImportError:
            pytest.skip("Pinecone not installed")

    def test_pinecone_implements_agent_memory(self):
        """Test that PineconeAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.pinecone import PineconeAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(PineconeAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("Pinecone not installed")

    def test_pinecone_has_all_methods(self):
        """Test that PineconeAgentMemory implements all required methods."""
        try:
            from vanna.integrations.pinecone import PineconeAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(PineconeAgentMemory, method_name)
        except ImportError:
            pytest.skip("Pinecone not installed")


class TestMilvusAgentMemory:
    """Sanity tests for Milvus AgentMemory implementation."""

    def test_milvus_import(self):
        """Test that MilvusAgentMemory can be imported."""
        try:
            from vanna.integrations.milvus import MilvusAgentMemory

            assert MilvusAgentMemory is not None
        except ImportError:
            pytest.skip("Milvus not installed")

    def test_milvus_implements_agent_memory(self):
        """Test that MilvusAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.milvus import MilvusAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(MilvusAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("Milvus not installed")

    def test_milvus_has_all_methods(self):
        """Test that MilvusAgentMemory implements all required methods."""
        try:
            from vanna.integrations.milvus import MilvusAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(MilvusAgentMemory, method_name)
        except ImportError:
            pytest.skip("Milvus not installed")


class TestWeaviateAgentMemory:
    """Sanity tests for Weaviate AgentMemory implementation."""

    def test_weaviate_import(self):
        """Test that WeaviateAgentMemory can be imported."""
        try:
            from vanna.integrations.weaviate import WeaviateAgentMemory

            assert WeaviateAgentMemory is not None
        except ImportError:
            pytest.skip("Weaviate not installed")

    def test_weaviate_implements_agent_memory(self):
        """Test that WeaviateAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.weaviate import WeaviateAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(WeaviateAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("Weaviate not installed")

    def test_weaviate_has_all_methods(self):
        """Test that WeaviateAgentMemory implements all required methods."""
        try:
            from vanna.integrations.weaviate import WeaviateAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(WeaviateAgentMemory, method_name)
        except ImportError:
            pytest.skip("Weaviate not installed")


class TestFAISSAgentMemory:
    """Sanity tests for FAISS AgentMemory implementation."""

    def test_faiss_import(self):
        """Test that FAISSAgentMemory can be imported."""
        try:
            from vanna.integrations.faiss import FAISSAgentMemory

            assert FAISSAgentMemory is not None
        except ImportError:
            pytest.skip("FAISS not installed")

    def test_faiss_implements_agent_memory(self):
        """Test that FAISSAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.faiss import FAISSAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(FAISSAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("FAISS not installed")

    def test_faiss_has_all_methods(self):
        """Test that FAISSAgentMemory implements all required methods."""
        try:
            from vanna.integrations.faiss import FAISSAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(FAISSAgentMemory, method_name)
        except ImportError:
            pytest.skip("FAISS not installed")

    def test_faiss_instantiation(self):
        """Test that FAISSAgentMemory can be instantiated."""
        try:
            from vanna.integrations.faiss import FAISSAgentMemory
            import tempfile

            temp_dir = tempfile.mkdtemp()
            memory = FAISSAgentMemory(persist_path=temp_dir)

            assert memory is not None
            assert memory.persist_path == temp_dir
        except ImportError:
            pytest.skip("FAISS not installed")


class TestOpenSearchAgentMemory:
    """Sanity tests for OpenSearch AgentMemory implementation."""

    def test_opensearch_import(self):
        """Test that OpenSearchAgentMemory can be imported."""
        try:
            from vanna.integrations.opensearch import OpenSearchAgentMemory

            assert OpenSearchAgentMemory is not None
        except ImportError:
            pytest.skip("OpenSearch not installed")

    def test_opensearch_implements_agent_memory(self):
        """Test that OpenSearchAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.opensearch import OpenSearchAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(OpenSearchAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("OpenSearch not installed")

    def test_opensearch_has_all_methods(self):
        """Test that OpenSearchAgentMemory implements all required methods."""
        try:
            from vanna.integrations.opensearch import OpenSearchAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(OpenSearchAgentMemory, method_name)
        except ImportError:
            pytest.skip("OpenSearch not installed")


class TestAzureAISearchAgentMemory:
    """Sanity tests for Azure AI Search AgentMemory implementation."""

    def test_azuresearch_import(self):
        """Test that AzureAISearchAgentMemory can be imported."""
        try:
            from vanna.integrations.azuresearch import AzureAISearchAgentMemory

            assert AzureAISearchAgentMemory is not None
        except ImportError:
            pytest.skip("Azure Search not installed")

    def test_azuresearch_implements_agent_memory(self):
        """Test that AzureAISearchAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.azuresearch import AzureAISearchAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(AzureAISearchAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("Azure Search not installed")

    def test_azuresearch_has_all_methods(self):
        """Test that AzureAISearchAgentMemory implements all required methods."""
        try:
            from vanna.integrations.azuresearch import AzureAISearchAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(AzureAISearchAgentMemory, method_name)
        except ImportError:
            pytest.skip("Azure Search not installed")


class TestMarqoAgentMemory:
    """Sanity tests for Marqo AgentMemory implementation."""

    def test_marqo_import(self):
        """Test that MarqoAgentMemory can be imported."""
        try:
            from vanna.integrations.marqo import MarqoAgentMemory

            assert MarqoAgentMemory is not None
        except ImportError:
            pytest.skip("Marqo not installed")

    def test_marqo_implements_agent_memory(self):
        """Test that MarqoAgentMemory implements AgentMemory."""
        try:
            from vanna.integrations.marqo import MarqoAgentMemory
            from vanna.capabilities.agent_memory import AgentMemory

            assert issubclass(MarqoAgentMemory, AgentMemory)
        except ImportError:
            pytest.skip("Marqo not installed")

    def test_marqo_has_all_methods(self):
        """Test that MarqoAgentMemory implements all required methods."""
        try:
            from vanna.integrations.marqo import MarqoAgentMemory

            required_methods = [
                "save_tool_usage",
                "save_text_memory",
                "search_similar_usage",
                "search_text_memories",
                "get_recent_memories",
                "get_recent_text_memories",
                "delete_by_id",
                "delete_text_memory",
                "clear_memories",
            ]

            for method_name in required_methods:
                assert hasattr(MarqoAgentMemory, method_name)
        except ImportError:
            pytest.skip("Marqo not installed")


class TestDemoAgentMemory:
    """Sanity tests for DemoAgentMemory (in-memory) implementation."""

    def test_demo_import(self):
        """Test that DemoAgentMemory can be imported."""
        from vanna.integrations.local.agent_memory import DemoAgentMemory

        assert DemoAgentMemory is not None

    def test_demo_implements_agent_memory(self):
        """Test that DemoAgentMemory implements AgentMemory."""
        from vanna.integrations.local.agent_memory import DemoAgentMemory
        from vanna.capabilities.agent_memory import AgentMemory

        assert issubclass(DemoAgentMemory, AgentMemory)

    def test_demo_has_all_methods(self):
        """Test that DemoAgentMemory implements all required methods."""
        from vanna.integrations.local.agent_memory import DemoAgentMemory

        required_methods = [
            "save_tool_usage",
            "save_text_memory",
            "search_similar_usage",
            "search_text_memories",
            "get_recent_memories",
            "get_recent_text_memories",
            "delete_by_id",
            "delete_text_memory",
            "clear_memories",
        ]

        for method_name in required_methods:
            assert hasattr(DemoAgentMemory, method_name)

    def test_demo_instantiation(self):
        """Test that DemoAgentMemory can be instantiated."""
        from vanna.integrations.local.agent_memory import DemoAgentMemory

        memory = DemoAgentMemory(max_items=100)

        assert memory is not None
        assert memory._max_items == 100


================================================
FILE: tests/test_agents.py
================================================
"""
Simple end-to-end tests for Vanna agents.

Tests use agent.send_message and validate the response components.
"""

import os
import pytest
from vanna.core.user import User
from vanna.core.user.resolver import UserResolver
from vanna.core.user.request_context import RequestContext


class SimpleUserResolver(UserResolver):
    """Simple user resolver for tests - always returns the same test user."""

    async def resolve_user(self, request_context: RequestContext) -> User:
        return User(
            id="test_user", email="test@example.com", group_memberships=["user"]
        )


def create_agent(llm_service, sql_runner):
    """Helper to create a configured agent."""
    from vanna import Agent, AgentConfig
    from vanna.core.registry import ToolRegistry
    from vanna.tools import RunSqlTool
    from vanna.integrations.local.file_system import LocalFileSystem
    from vanna.integrations.local.agent_memory import DemoAgentMemory
    from vanna.tools.agent_memory import (
        SaveQuestionToolArgsTool,
        SearchSavedCorrectToolUsesTool,
    )

    tools = ToolRegistry()

    # Add SQL tool
    db_tool = RunSqlTool(sql_runner=sql_runner, file_system=LocalFileSystem())
    tools.register_local_tool(db_tool, access_groups=["user"])

    # Add memory tools (they access agent_memory via ToolContext)
    agent_memory = DemoAgentMemory(max_items=1000)
    tools.register_local_tool(SaveQuestionToolArgsTool(), access_groups=["user"])
    tools.register_local_tool(SearchSavedCorrectToolUsesTool(), access_groups=["user"])

    return Agent(
        llm_service=llm_service,
        tool_registry=tools,
        user_resolver=SimpleUserResolver(),
        agent_memory=agent_memory,
        config=AgentConfig(),
    )


async def test_agent_top_artist(agent, expected_artist="Iron Maiden"):
    """Common test logic for testing agent responses about top artist by sales."""
    # Create a simple request context
    request_context = RequestContext(cookies={}, headers={})

    # Collect all components from the async generator
    components = []
    async for component in agent.send_message(
        request_context, "Who is the top artist by sales?"
    ):
        components.append(component)

    # Validate we got components
    assert len(components) > 0, "Should receive at least one component"

    # Print all components for debugging
    print(f"\n\n=== Received {len(components)} components ===")
    for i, component in enumerate(components):
        print(f"\nComponent {i + 1}:")
        print(f"  Type: {component.type if hasattr(component, 'type') else 'no type'}")
        if hasattr(component, "text"):
            print(
                f"  Text: {component.text[:200]}..."
                if len(component.text) > 200
                else f"  Text: {component.text}"
            )
        if hasattr(component, "content"):
            print(f"  Content: {str(component.content)[:200]}...")
        print(f"  Full: {component}")

    # Look for the expected artist in any component
    found_artist = False
    for component in components:
        # Check rich_component.content
        if hasattr(component, "rich_component") and hasattr(
            component.rich_component, "content"
        ):
            if expected_artist in component.rich_component.content:
                found_artist = True
                break
        # Check simple_component.text
        if hasattr(component, "simple_component") and hasattr(
            component.simple_component, "text"
        ):
            if expected_artist in component.simple_component.text:
                found_artist = True
                break

    assert found_artist, (
        f"Response should mention '{expected_artist}' as the top artist. Got {len(components)} components."
    )


@pytest.mark.anthropic
@pytest.mark.asyncio
async def test_anthropic_top_artist(chinook_db):
    """Test Anthropic agent finding the top artist by sales."""
    from vanna.integrations.anthropic import AnthropicLlmService

    api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = AnthropicLlmService(api_key=api_key, model="claude-sonnet-4-5")

    agent = create_agent(llm, chinook_db)
    await test_agent_top_artist(agent)


@pytest.mark.openai
@pytest.mark.asyncio
async def test_openai_top_artist(chinook_db):
    """Test OpenAI agent finding the top artist by sales."""
    from vanna.integrations.openai import OpenAILlmService

    api_key = os.getenv("OPENAI_API_KEY")
    llm = OpenAILlmService(api_key=api_key, model="gpt-5")

    agent = create_agent(llm, chinook_db)
    await test_agent_top_artist(agent)


@pytest.mark.azureopenai
@pytest.mark.asyncio
async def test_azure_openai_top_artist(chinook_db):
    """Test Azure OpenAI agent finding the top artist by sales."""
    from vanna.integrations.azureopenai import AzureOpenAILlmService

    # Get Azure OpenAI credentials from environment
    api_key = os.getenv("AZURE_OPENAI_API_KEY")
    model = os.getenv("AZURE_OPENAI_MODEL")
    azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
    api_version = os.getenv("AZURE_OPENAI_API_VERSION")

    llm = AzureOpenAILlmService(
        model=model,
        api_key=api_key,
        azure_endpoint=azure_endpoint,
        api_version=api_version,
    )

    agent = create_agent(llm, chinook_db)
    await test_agent_top_artist(agent)


# @pytest.mark.openai
# @pytest.mark.asyncio
# async def test_openai_responses_top_artist(chinook_db):
#     """Test OpenAI Responses API agent finding the top artist by sales."""
#     from vanna.integrations.openai import OpenAIResponsesService

#     api_key = os.getenv("OPENAI_API_KEY")
#     llm = OpenAIResponsesService(api_key=api_key, model="gpt-5")

#     agent = create_agent(llm, chinook_db)
#     await test_agent_top_artist(agent)


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_top_artist(chinook_db):
    """Test Ollama agent finding the top artist by sales."""
    from vanna.integrations.ollama import OllamaLlmService

    llm = OllamaLlmService(
        model="gpt-oss:20b-cloud",
        host=os.getenv("OLLAMA_HOST", "http://localhost:11434"),
    )

    agent = create_agent(llm, chinook_db)
    await test_agent_top_artist(agent)


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_top_artist(chinook_db):
    """Test Gemini agent finding the top artist by sales."""
    from vanna.integrations.google import GeminiLlmService

    # API key will be picked up from GOOGLE_API_KEY or GEMINI_API_KEY env var
    llm = GeminiLlmService(model="gemini-2.5-pro", temperature=0.0)

    agent = create_agent(llm, chinook_db)
    await test_agent_top_artist(agent)


================================================
FILE: tests/test_azureopenai_llm.py
================================================
"""
Unit tests for Azure OpenAI LLM service integration.

These tests validate the Azure OpenAI integration without making actual API calls.
"""

import pytest
from unittest.mock import Mock, patch, AsyncMock
from vanna.integrations.azureopenai import AzureOpenAILlmService
from vanna.integrations.azureopenai.llm import _is_reasoning_model


class TestReasoningModelDetection:
    """Test reasoning model detection logic."""

    def test_is_reasoning_model_o1(self):
        """Test that o1 models are detected as reasoning models."""
        assert _is_reasoning_model("o1")
        assert _is_reasoning_model("o1-mini")
        assert _is_reasoning_model("o1-preview")

    def test_is_reasoning_model_o3(self):
        """Test that o3 models are detected as reasoning models."""
        assert _is_reasoning_model("o3-mini")

    def test_is_reasoning_model_gpt5(self):
        """Test that GPT-5 series models are detected as reasoning models."""
        assert _is_reasoning_model("gpt-5")
        assert _is_reasoning_model("gpt-5-mini")
        assert _is_reasoning_model("gpt-5-nano")
        assert _is_reasoning_model("gpt-5-pro")
        assert _is_reasoning_model("gpt-5-codex")

    def test_is_not_reasoning_model(self):
        """Test that standard models are not detected as reasoning models."""
        assert not _is_reasoning_model("gpt-4")
        assert not _is_reasoning_model("gpt-4o")
        assert not _is_reasoning_model("gpt-4-turbo")
        assert not _is_reasoning_model("gpt-3.5-turbo")

    def test_case_insensitive_detection(self):
        """Test that model detection is case insensitive."""
        assert _is_reasoning_model("GPT-5")
        assert _is_reasoning_model("O1-MINI")
        assert not _is_reasoning_model("GPT-4O")


class TestAzureOpenAILlmServiceInitialization:
    """Test Azure OpenAI service initialization."""

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_with_all_params(self, mock_azure_openai):
        """Test initialization with all parameters provided."""
        service = AzureOpenAILlmService(
            model="gpt-4o",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
            api_version="2024-10-21",
        )

        assert service.model == "gpt-4o"
        assert not service._is_reasoning_model

        # Verify AzureOpenAI was called with correct params
        mock_azure_openai.assert_called_once()
        call_kwargs = mock_azure_openai.call_args[1]
        assert call_kwargs["azure_endpoint"] == "https://test.openai.azure.com"
        assert call_kwargs["api_version"] == "2024-10-21"
        assert call_kwargs["api_key"] == "test-key"

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_with_reasoning_model(self, mock_azure_openai):
        """Test initialization with a reasoning model."""
        service = AzureOpenAILlmService(
            model="gpt-5",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
        )

        assert service.model == "gpt-5"
        assert service._is_reasoning_model

    @patch.dict(
        "os.environ",
        {
            "AZURE_OPENAI_MODEL": "gpt-4o-deployment",
            "AZURE_OPENAI_API_KEY": "env-key",
            "AZURE_OPENAI_ENDPOINT": "https://env.openai.azure.com",
            "AZURE_OPENAI_API_VERSION": "2024-06-01",
        },
    )
    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_from_environment(self, mock_azure_openai):
        """Test initialization from environment variables."""
        service = AzureOpenAILlmService()

        assert service.model == "gpt-4o-deployment"

        # Verify AzureOpenAI was called with env values
        call_kwargs = mock_azure_openai.call_args[1]
        assert call_kwargs["azure_endpoint"] == "https://env.openai.azure.com"
        assert call_kwargs["api_version"] == "2024-06-01"
        assert call_kwargs["api_key"] == "env-key"

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_missing_model_raises(self, mock_azure_openai):
        """Test that missing model parameter raises ValueError."""
        with pytest.raises(ValueError, match="model parameter.*is required"):
            AzureOpenAILlmService(
                api_key="test-key",
                azure_endpoint="https://test.openai.azure.com",
            )

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_missing_endpoint_raises(self, mock_azure_openai):
        """Test that missing azure_endpoint raises ValueError."""
        with pytest.raises(ValueError, match="azure_endpoint is required"):
            AzureOpenAILlmService(
                model="gpt-4o",
                api_key="test-key",
            )

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_missing_auth_raises(self, mock_azure_openai):
        """Test that missing authentication raises ValueError."""
        with pytest.raises(ValueError, match="Authentication required"):
            AzureOpenAILlmService(
                model="gpt-4o",
                azure_endpoint="https://test.openai.azure.com",
            )

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_with_azure_ad_token_provider(self, mock_azure_openai):
        """Test initialization with Azure AD token provider."""
        mock_token_provider = Mock()
        service = AzureOpenAILlmService(
            model="gpt-4o",
            azure_endpoint="https://test.openai.azure.com",
            azure_ad_token_provider=mock_token_provider,
        )

        # Verify token provider was used instead of API key
        call_kwargs = mock_azure_openai.call_args[1]
        assert call_kwargs["azure_ad_token_provider"] == mock_token_provider
        assert "api_key" not in call_kwargs

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_init_default_api_version(self, mock_azure_openai):
        """Test that default API version is used when not specified."""
        service = AzureOpenAILlmService(
            model="gpt-4o",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
        )

        # Verify default API version (2024-10-21)
        call_kwargs = mock_azure_openai.call_args[1]
        assert call_kwargs["api_version"] == "2024-10-21"


class TestAzureOpenAILlmServicePayloadBuilding:
    """Test payload building for API requests."""

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_build_payload_includes_temperature_for_standard_model(
        self, mock_azure_openai
    ):
        """Test that temperature is included for standard models."""
        from vanna.core.llm import LlmRequest, LlmMessage
        from vanna.core.user import User

        service = AzureOpenAILlmService(
            model="gpt-4o",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
        )

        request = LlmRequest(
            messages=[LlmMessage(role="user", content="test")],
            user=User(id="test", email="test@example.com", group_memberships=[]),
            temperature=0.8,
        )

        payload = service._build_payload(request)

        assert "temperature" in payload
        assert payload["temperature"] == 0.8
        assert payload["model"] == "gpt-4o"

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_build_payload_excludes_temperature_for_reasoning_model(
        self, mock_azure_openai
    ):
        """Test that temperature is excluded for reasoning models."""
        from vanna.core.llm import LlmRequest, LlmMessage
        from vanna.core.user import User

        service = AzureOpenAILlmService(
            model="gpt-5",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
        )

        request = LlmRequest(
            messages=[LlmMessage(role="user", content="test")],
            user=User(id="test", email="test@example.com", group_memberships=[]),
            temperature=0.8,
        )

        payload = service._build_payload(request)

        assert "temperature" not in payload
        assert payload["model"] == "gpt-5"

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_build_payload_with_system_prompt(self, mock_azure_openai):
        """Test that system prompt is added to messages."""
        from vanna.core.llm import LlmRequest, LlmMessage
        from vanna.core.user import User

        service = AzureOpenAILlmService(
            model="gpt-4o",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
        )

        request = LlmRequest(
            messages=[LlmMessage(role="user", content="test")],
            user=User(id="test", email="test@example.com", group_memberships=[]),
            system_prompt="You are a helpful assistant.",
        )

        payload = service._build_payload(request)

        assert len(payload["messages"]) == 2
        assert payload["messages"][0]["role"] == "system"
        assert payload["messages"][0]["content"] == "You are a helpful assistant."

    @patch("vanna.integrations.azureopenai.llm.AzureOpenAI")
    def test_build_payload_with_tools(self, mock_azure_openai):
        """Test that tools are properly formatted in payload."""
        from vanna.core.llm import LlmRequest, LlmMessage
        from vanna.core.user import User
        from vanna.core.tool import ToolSchema

        service = AzureOpenAILlmService(
            model="gpt-4o",
            api_key="test-key",
            azure_endpoint="https://test.openai.azure.com",
        )

        tool = ToolSchema(
            name="test_tool",
            description="A test tool",
            parameters={
                "type": "object",
                "properties": {"param1": {"type": "string"}},
            },
        )

        request = LlmRequest(
            messages=[LlmMessage(role="user", content="test")],
            user=User(id="test", email="test@example.com", group_memberships=[]),
            tools=[tool],
        )

        payload = service._build_payload(request)

        assert "tools" in payload
        assert len(payload["tools"]) == 1
        assert payload["tools"][0]["type"] == "function"
        assert payload["tools"][0]["function"]["name"] == "test_tool"
        assert payload["tool_choice"] == "auto"


class TestImportError:
    """Test import error handling."""

    def test_import_error_message(self):
        """Test that helpful error message is shown when openai is not installed."""
        with patch.dict("sys.modules", {"openai": None}):
            # Force module reload to trigger import error
            import sys

            if "vanna.integrations.azureopenai.llm" in sys.modules:
                del sys.modules["vanna.integrations.azureopenai.llm"]

            with pytest.raises(
                ImportError, match="pip install 'vanna\\[azureopenai\\]'"
            ):
                from vanna.integrations.azureopenai import AzureOpenAILlmService

                AzureOpenAILlmService(
                    model="gpt-4o",
                    api_key="test",
                    azure_endpoint="https://test.openai.azure.com",
                )


================================================
FILE: tests/test_chromadb_persistence_fix.py
================================================
"""
Test for ChromaDB persistence fix.

This test verifies that ChromaDB collections can be retrieved without triggering
unnecessary embedding function initialization/model downloads.
"""

import pytest
import tempfile
import shutil
import asyncio

from vanna.integrations.chromadb import ChromaAgentMemory
from vanna.core.user import User
from vanna.core.tool import ToolContext


@pytest.fixture
def test_user():
    """Test user for context."""
    return User(
        id="test_user",
        username="test",
        email="test@example.com",
        group_memberships=["user"],
    )


def create_test_context(test_user, agent_memory):
    """Helper to create test context."""
    return ToolContext(
        user=test_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
        metadata={},
    )


@pytest.mark.asyncio
async def test_chromadb_collection_retrieval_without_embedding_function(test_user):
    """
    Test that existing ChromaDB collections can be retrieved without
    initializing the embedding function (avoiding model downloads).

    This test simulates the real-world scenario where:
    1. A collection is created with an embedding function (first app run)
    2. The app restarts and retrieves the existing collection
    3. The embedding function should NOT be re-initialized on retrieval
    """
    try:
        import chromadb
        from chromadb.config import Settings
    except ImportError:
        pytest.skip("ChromaDB not installed")

    temp_dir = tempfile.mkdtemp()

    try:
        # Session 1: Create a collection using ChromaAgentMemory (simulating first app run)
        # This will create the collection with an embedding function
        memory1 = ChromaAgentMemory(
            persist_directory=temp_dir, collection_name="test_collection"
        )

        context = create_test_context(test_user, memory1)

        # Save some memories (this will create the collection)
        # We need to add explicit embeddings to avoid model download in test environment
        collection = memory1._get_collection()
        collection.add(
            ids=["mem1", "mem2"],
            documents=["test question 1", "test question 2"],
            embeddings=[[0.1] * 384, [0.2] * 384],
            metadatas=[
                {
                    "question": "test question 1",
                    "tool_name": "test_tool",
                    "args_json": "{}",
                    "timestamp": "2024-01-01T00:00:00",
                    "success": True,
                    "metadata_json": "{}",
                },
                {
                    "question": "test question 2",
                    "tool_name": "test_tool",
                    "args_json": "{}",
                    "timestamp": "2024-01-01T00:01:00",
                    "success": True,
                    "metadata_json": "{}",
                },
            ],
        )

        # Clean up references to simulate app restart
        del collection
        del memory1

        # Session 2: Create new ChromaAgentMemory instance (simulating app restart)
        # This should retrieve the existing collection WITHOUT calling _get_embedding_function
        memory2 = ChromaAgentMemory(
            persist_directory=temp_dir, collection_name="test_collection"
        )

        # Mock _get_embedding_function to verify it's not called
        original_get_ef = memory2._get_embedding_function

        def mock_get_ef():
            pytest.fail(
                "_get_embedding_function was called when retrieving existing collection"
            )

        memory2._get_embedding_function = mock_get_ef

        # This should retrieve the existing collection without calling _get_embedding_function
        collection2 = memory2._get_collection()

        # Restore original method
        memory2._get_embedding_function = original_get_ef

        # Verify collection was retrieved successfully
        assert collection2 is not None
        assert collection2.name == "test_collection"
        assert collection2.count() == 2

        # Test that we can use public API methods on the retrieved collection
        context2 = create_test_context(test_user, memory2)
        recent = await memory2.get_recent_memories(context=context2, limit=10)
        assert len(recent) == 2
        assert recent[0].question in ["test question 1", "test question 2"]

    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


@pytest.mark.asyncio
async def test_chromadb_collection_creation_with_embedding_function():
    """
    Test that NEW ChromaDB collections are created WITH the embedding function.
    """
    try:
        from vanna.integrations.chromadb import ChromaAgentMemory
    except ImportError:
        pytest.skip("ChromaDB not installed")

    temp_dir = tempfile.mkdtemp()

    try:
        # Test: Create ChromaAgentMemory for a non-existent collection
        memory = ChromaAgentMemory(
            persist_directory=temp_dir, collection_name="new_collection"
        )

        # Track if _get_embedding_function was called
        get_ef_called = []
        original_get_ef = memory._get_embedding_function

        def tracking_get_ef():
            get_ef_called.append(True)
            return original_get_ef()

        memory._get_embedding_function = tracking_get_ef

        # This should create a new collection and SHOULD call _get_embedding_function
        collection = memory._get_collection()

        # Restore original
        memory._get_embedding_function = original_get_ef

        # Verify collection was created
        assert collection is not None
        assert collection.name == "new_collection"

        # Verify _get_embedding_function was called
        assert get_ef_called, (
            "_get_embedding_function should be called when creating new collection"
        )

    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


if __name__ == "__main__":
    # Run tests directly
    asyncio.run(test_chromadb_collection_retrieval_without_embedding_function())
    asyncio.run(test_chromadb_collection_creation_with_embedding_function())


================================================
FILE: tests/test_database_sanity.py
================================================
"""
Sanity tests for database implementations.

These tests verify that:
1. Each database implementation correctly implements the SqlRunner interface
2. Imports are working correctly for all database modules
3. Basic class instantiation works (without requiring actual database connections)

Note: These tests do NOT execute actual queries against databases.
They are lightweight sanity checks for the implementation structure.
"""

import pytest
from abc import abstractmethod
from inspect import signature, iscoroutinefunction
import pandas as pd


class TestSqlRunnerInterface:
    """Test that the SqlRunner interface is properly defined."""

    def test_sql_runner_import(self):
        """Test that SqlRunner can be imported."""
        from vanna.capabilities.sql_runner import SqlRunner

        assert SqlRunner is not None

    def test_sql_runner_is_abstract(self):
        """Test that SqlRunner is an abstract base class."""
        from vanna.capabilities.sql_runner import SqlRunner
        from abc import ABC

        assert issubclass(SqlRunner, ABC)

    def test_sql_runner_has_run_sql_method(self):
        """Test that SqlRunner defines the run_sql abstract method."""
        from vanna.capabilities.sql_runner import SqlRunner

        assert hasattr(SqlRunner, "run_sql")
        assert getattr(SqlRunner.run_sql, "__isabstractmethod__", False)

    def test_run_sql_method_signature(self):
        """Test that run_sql has the correct method signature."""
        from vanna.capabilities.sql_runner import SqlRunner

        sig = signature(SqlRunner.run_sql)
        params = list(sig.parameters.keys())

        # Should have: self, args, context
        assert len(params) == 3
        assert params[0] == "self"
        assert params[1] == "args"
        assert params[2] == "context"

    def test_run_sql_is_async(self):
        """Test that run_sql is defined as an async method."""
        from vanna.capabilities.sql_runner import SqlRunner

        # Abstract methods are wrapped, so we check if it's meant to be async
        # by looking at the method definition
        assert iscoroutinefunction(SqlRunner.run_sql)


class TestRunSqlToolArgsModel:
    """Test the RunSqlToolArgs model."""

    def test_run_sql_tool_args_import(self):
        """Test that RunSqlToolArgs can be imported."""
        from vanna.capabilities.sql_runner import RunSqlToolArgs

        assert RunSqlToolArgs is not None

    def test_run_sql_tool_args_has_sql_field(self):
        """Test that RunSqlToolArgs has a 'sql' field."""
        from vanna.capabilities.sql_runner import RunSqlToolArgs

        # Create an instance
        args = RunSqlToolArgs(sql="SELECT 1")
        assert hasattr(args, "sql")
        assert args.sql == "SELECT 1"

    def test_run_sql_tool_args_is_pydantic_model(self):
        """Test that RunSqlToolArgs is a Pydantic model."""
        from vanna.capabilities.sql_runner import RunSqlToolArgs
        from pydantic import BaseModel

        assert issubclass(RunSqlToolArgs, BaseModel)


class TestPostgresRunner:
    """Sanity tests for PostgresRunner implementation."""

    def test_postgres_runner_import(self):
        """Test that PostgresRunner can be imported."""
        from vanna.integrations.postgres import PostgresRunner

        assert PostgresRunner is not None

    def test_postgres_runner_implements_sql_runner(self):
        """Test that PostgresRunner implements SqlRunner interface."""
        from vanna.integrations.postgres import PostgresRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(PostgresRunner, SqlRunner)

    def test_postgres_runner_has_run_sql_method(self):
        """Test that PostgresRunner implements run_sql method."""
        from vanna.integrations.postgres import PostgresRunner

        assert hasattr(PostgresRunner, "run_sql")
        # Should not be abstract anymore
        assert not getattr(PostgresRunner.run_sql, "__isabstractmethod__", False)

    def test_postgres_runner_instantiation_with_connection_string(self):
        """Test that PostgresRunner can be instantiated with connection string."""
        from vanna.integrations.postgres import PostgresRunner

        # This should not raise an error (no actual connection is made in __init__)
        runner = PostgresRunner(connection_string="postgresql://user:pass@localhost/db")
        assert runner is not None
        assert runner.connection_string == "postgresql://user:pass@localhost/db"
        assert runner.connection_params is None

    def test_postgres_runner_instantiation_with_params(self):
        """Test that PostgresRunner can be instantiated with individual parameters."""
        from vanna.integrations.postgres import PostgresRunner

        runner = PostgresRunner(
            host="localhost",
            port=5432,
            database="testdb",
            user="testuser",
            password="testpass",
        )
        assert runner is not None
        assert runner.connection_string is None
        assert runner.connection_params is not None
        assert runner.connection_params["host"] == "localhost"
        assert runner.connection_params["database"] == "testdb"

    def test_postgres_runner_requires_valid_params(self):
        """Test that PostgresRunner raises error with invalid parameters."""
        from vanna.integrations.postgres import PostgresRunner

        with pytest.raises(ValueError, match="Either provide connection_string OR"):
            PostgresRunner()  # No parameters provided

    def test_postgres_runner_checks_psycopg2_import(self):
        """Test that PostgresRunner checks for psycopg2 package."""
        from vanna.integrations.postgres import PostgresRunner

        # If psycopg2 is not installed, this should raise ImportError
        # If it is installed, this should work fine
        try:
            runner = PostgresRunner(connection_string="postgresql://test")
            assert runner.psycopg2 is not None
        except ImportError as e:
            assert "psycopg2" in str(e)


class TestSqliteRunner:
    """Sanity tests for SqliteRunner implementation."""

    def test_sqlite_runner_import(self):
        """Test that SqliteRunner can be imported."""
        from vanna.integrations.sqlite import SqliteRunner

        assert SqliteRunner is not None

    def test_sqlite_runner_implements_sql_runner(self):
        """Test that SqliteRunner implements SqlRunner interface."""
        from vanna.integrations.sqlite import SqliteRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(SqliteRunner, SqlRunner)

    def test_sqlite_runner_has_run_sql_method(self):
        """Test that SqliteRunner implements run_sql method."""
        from vanna.integrations.sqlite import SqliteRunner

        assert hasattr(SqliteRunner, "run_sql")
        # Should not be abstract anymore
        assert not getattr(SqliteRunner.run_sql, "__isabstractmethod__", False)

    def test_sqlite_runner_instantiation(self):
        """Test that SqliteRunner can be instantiated with a database path."""
        from vanna.integrations.sqlite import SqliteRunner

        runner = SqliteRunner(database_path="/tmp/test.db")
        assert runner is not None
        assert runner.database_path == "/tmp/test.db"

    def test_sqlite_uses_builtin_sqlite3(self):
        """Test that SqliteRunner uses Python's built-in sqlite3 module."""
        import sqlite3
        from vanna.integrations.sqlite import SqliteRunner

        # sqlite3 should be importable (it's part of Python standard library)
        assert sqlite3 is not None


class TestLegacySqlRunner:
    """Sanity tests for LegacySqlRunner adapter."""

    def test_legacy_sql_runner_import(self):
        """Test that LegacySqlRunner can be imported."""
        from vanna.legacy.adapter import LegacySqlRunner

        assert LegacySqlRunner is not None

    def test_legacy_sql_runner_implements_sql_runner(self):
        """Test that LegacySqlRunner implements SqlRunner interface."""
        from vanna.legacy.adapter import LegacySqlRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(LegacySqlRunner, SqlRunner)

    def test_legacy_sql_runner_has_run_sql_method(self):
        """Test that LegacySqlRunner implements run_sql method."""
        from vanna.legacy.adapter import LegacySqlRunner

        assert hasattr(LegacySqlRunner, "run_sql")
        # Should not be abstract anymore
        assert not getattr(LegacySqlRunner.run_sql, "__isabstractmethod__", False)

    def test_legacy_sql_runner_instantiation(self):
        """Test that LegacySqlRunner can be instantiated with a VannaBase instance."""
        from vanna.legacy.adapter import LegacySqlRunner
        from unittest.mock import Mock

        # Create a mock VannaBase instance
        mock_vn = Mock()

        runner = LegacySqlRunner(vn=mock_vn)
        assert runner is not None
        assert runner.vn is mock_vn


class TestDatabaseIntegrationModules:
    """Test that database integration modules can be imported."""

    def test_postgres_module_import(self):
        """Test that the postgres integration module can be imported."""
        try:
            import vanna.integrations.postgres

            assert vanna.integrations.postgres is not None
        except ImportError as e:
            pytest.fail(f"Failed to import postgres module: {e}")

    def test_sqlite_module_import(self):
        """Test that the sqlite integration module can be imported."""
        try:
            import vanna.integrations.sqlite

            assert vanna.integrations.sqlite is not None
        except ImportError as e:
            pytest.fail(f"Failed to import sqlite module: {e}")

    def test_postgres_module_exports_runner(self):
        """Test that postgres module exports PostgresRunner."""
        from vanna.integrations.postgres import PostgresRunner

        assert PostgresRunner is not None

    def test_sqlite_module_exports_runner(self):
        """Test that sqlite module exports SqliteRunner."""
        from vanna.integrations.sqlite import SqliteRunner

        assert SqliteRunner is not None


class TestLegacyVannaBaseConnections:
    """Test that legacy VannaBase connection methods exist."""

    def test_vanna_base_import(self):
        """Test that VannaBase can be imported."""
        from vanna.legacy.base.base import VannaBase

        assert VannaBase is not None

    def test_vanna_base_has_connection_methods(self):
        """Test that VannaBase has various database connection methods."""
        from vanna.legacy.base.base import VannaBase

        connection_methods = [
            "connect_to_snowflake",
            "connect_to_sqlite",
            "connect_to_postgres",
            "connect_to_mysql",
            "connect_to_clickhouse",
            "connect_to_oracle",
            "connect_to_bigquery",
            "connect_to_duckdb",
            "connect_to_mssql",
            "connect_to_presto",
            "connect_to_hive",
        ]

        for method_name in connection_methods:
            assert hasattr(VannaBase, method_name), f"VannaBase missing {method_name}"

    def test_vanna_base_has_run_sql_method(self):
        """Test that VannaBase has a run_sql method."""
        from vanna.legacy.base.base import VannaBase

        assert hasattr(VannaBase, "run_sql")


class TestLegacyVannaAdapter:
    """Test the LegacyVannaAdapter."""

    def test_legacy_vanna_adapter_import(self):
        """Test that LegacyVannaAdapter can be imported."""
        from vanna.legacy.adapter import LegacyVannaAdapter

        assert LegacyVannaAdapter is not None

    def test_legacy_vanna_adapter_is_tool_registry(self):
        """Test that LegacyVannaAdapter extends ToolRegistry."""
        from vanna.legacy.adapter import LegacyVannaAdapter
        from vanna.core.registry import ToolRegistry

        assert issubclass(LegacyVannaAdapter, ToolRegistry)


class TestSnowflakeRunner:
    """Sanity tests for SnowflakeRunner implementation."""

    def test_snowflake_runner_import(self):
        """Test that SnowflakeRunner can be imported."""
        from vanna.integrations.snowflake import SnowflakeRunner

        assert SnowflakeRunner is not None

    def test_snowflake_runner_implements_sql_runner(self):
        """Test that SnowflakeRunner implements SqlRunner interface."""
        from vanna.integrations.snowflake import SnowflakeRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(SnowflakeRunner, SqlRunner)

    def test_snowflake_runner_has_run_sql_method(self):
        """Test that SnowflakeRunner implements run_sql method."""
        from vanna.integrations.snowflake import SnowflakeRunner

        assert hasattr(SnowflakeRunner, "run_sql")
        assert not getattr(SnowflakeRunner.run_sql, "__isabstractmethod__", False)

    def test_snowflake_runner_instantiation(self):
        """Test that SnowflakeRunner can be instantiated with required parameters."""
        from vanna.integrations.snowflake import SnowflakeRunner

        runner = SnowflakeRunner(
            account="test-account",
            username="test-user",
            password="test-pass",
            database="test-db",
        )
        assert runner is not None
        assert runner.account == "test-account"

    def test_snowflake_runner_key_pair_auth_with_path(self, tmp_path):
        """Test that SnowflakeRunner can be instantiated with RSA key-pair authentication using path."""
        from vanna.integrations.snowflake import SnowflakeRunner

        # Create a temporary private key file
        private_key_file = tmp_path / "test_private_key.p8"
        private_key_file.write_text(
            "-----BEGIN PRIVATE KEY-----\ntest_key_content\n-----END PRIVATE KEY-----"
        )

        runner = SnowflakeRunner(
            account="test-account",
            username="test-user",
            private_key_path=str(private_key_file),
            private_key_passphrase="test-passphrase",
            database="test-db",
        )
        assert runner is not None
        assert runner.account == "test-account"
        assert runner.username == "test-user"
        assert runner.password is None
        assert runner.private_key_path == str(private_key_file)
        assert runner.private_key_passphrase == "test-passphrase"

    def test_snowflake_runner_key_pair_auth_with_content(self):
        """Test that SnowflakeRunner can be instantiated with RSA key-pair authentication using content."""
        from vanna.integrations.snowflake import SnowflakeRunner

        private_key_content = (
            b"-----BEGIN PRIVATE KEY-----\ntest_key_content\n-----END PRIVATE KEY-----"
        )

        runner = SnowflakeRunner(
            account="test-account",
            username="test-user",
            private_key_content=private_key_content,
            database="test-db",
        )
        assert runner is not None
        assert runner.account == "test-account"
        assert runner.username == "test-user"
        assert runner.password is None
        assert runner.private_key_content == private_key_content

    def test_snowflake_runner_key_pair_auth_without_passphrase(self, tmp_path):
        """Test that SnowflakeRunner works with unencrypted private key (no passphrase)."""
        from vanna.integrations.snowflake import SnowflakeRunner

        # Create a temporary private key file
        private_key_file = tmp_path / "test_private_key_unencrypted.p8"
        private_key_file.write_text(
            "-----BEGIN PRIVATE KEY-----\ntest_key_content\n-----END PRIVATE KEY-----"
        )

        runner = SnowflakeRunner(
            account="test-account",
            username="test-user",
            private_key_path=str(private_key_file),
            database="test-db",
        )
        assert runner is not None
        assert runner.private_key_passphrase is None

    def test_snowflake_runner_missing_auth_raises_error(self):
        """Test that SnowflakeRunner raises error when no authentication method is provided."""
        from vanna.integrations.snowflake import SnowflakeRunner
        import pytest

        with pytest.raises(
            ValueError,
            match="Either password or private_key_path/private_key_content must be provided",
        ):
            SnowflakeRunner(
                account="test-account", username="test-user", database="test-db"
            )

    def test_snowflake_runner_invalid_key_path_raises_error(self):
        """Test that SnowflakeRunner raises error when private key file doesn't exist."""
        from vanna.integrations.snowflake import SnowflakeRunner
        import pytest

        with pytest.raises(FileNotFoundError, match="Private key file not found"):
            SnowflakeRunner(
                account="test-account",
                username="test-user",
                private_key_path="/nonexistent/path/to/key.p8",
                database="test-db",
            )

    def test_snowflake_runner_password_auth_backwards_compatible(self):
        """Test that SnowflakeRunner maintains backward compatibility with password auth."""
        from vanna.integrations.snowflake import SnowflakeRunner

        runner = SnowflakeRunner(
            account="test-account",
            username="test-user",
            password="test-password",
            database="test-db",
            role="test-role",
            warehouse="test-warehouse",
        )
        assert runner is not None
        assert runner.password == "test-password"
        assert runner.private_key_path is None
        assert runner.private_key_content is None


class TestMySQLRunner:
    """Sanity tests for MySQLRunner implementation."""

    def test_mysql_runner_import(self):
        """Test that MySQLRunner can be imported."""
        from vanna.integrations.mysql import MySQLRunner

        assert MySQLRunner is not None

    def test_mysql_runner_implements_sql_runner(self):
        """Test that MySQLRunner implements SqlRunner interface."""
        from vanna.integrations.mysql import MySQLRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(MySQLRunner, SqlRunner)

    def test_mysql_runner_has_run_sql_method(self):
        """Test that MySQLRunner implements run_sql method."""
        from vanna.integrations.mysql import MySQLRunner

        assert hasattr(MySQLRunner, "run_sql")
        assert not getattr(MySQLRunner.run_sql, "__isabstractmethod__", False)

    def test_mysql_runner_instantiation(self):
        """Test that MySQLRunner can be instantiated with required parameters."""
        from vanna.integrations.mysql import MySQLRunner

        runner = MySQLRunner(
            host="localhost", database="test-db", user="test-user", password="test-pass"
        )
        assert runner is not None
        assert runner.host == "localhost"


class TestClickHouseRunner:
    """Sanity tests for ClickHouseRunner implementation."""

    def test_clickhouse_runner_import(self):
        """Test that ClickHouseRunner can be imported."""
        from vanna.integrations.clickhouse import ClickHouseRunner

        assert ClickHouseRunner is not None

    def test_clickhouse_runner_implements_sql_runner(self):
        """Test that ClickHouseRunner implements SqlRunner interface."""
        from vanna.integrations.clickhouse import ClickHouseRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(ClickHouseRunner, SqlRunner)

    def test_clickhouse_runner_has_run_sql_method(self):
        """Test that ClickHouseRunner implements run_sql method."""
        from vanna.integrations.clickhouse import ClickHouseRunner

        assert hasattr(ClickHouseRunner, "run_sql")
        assert not getattr(ClickHouseRunner.run_sql, "__isabstractmethod__", False)

    def test_clickhouse_runner_instantiation(self):
        """Test that ClickHouseRunner can be instantiated with required parameters."""
        from vanna.integrations.clickhouse import ClickHouseRunner

        runner = ClickHouseRunner(
            host="localhost", database="test-db", user="test-user", password="test-pass"
        )
        assert runner is not None
        assert runner.host == "localhost"


class TestOracleRunner:
    """Sanity tests for OracleRunner implementation."""

    def test_oracle_runner_import(self):
        """Test that OracleRunner can be imported."""
        from vanna.integrations.oracle import OracleRunner

        assert OracleRunner is not None

    def test_oracle_runner_implements_sql_runner(self):
        """Test that OracleRunner implements SqlRunner interface."""
        from vanna.integrations.oracle import OracleRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(OracleRunner, SqlRunner)

    def test_oracle_runner_has_run_sql_method(self):
        """Test that OracleRunner implements run_sql method."""
        from vanna.integrations.oracle import OracleRunner

        assert hasattr(OracleRunner, "run_sql")
        assert not getattr(OracleRunner.run_sql, "__isabstractmethod__", False)

    def test_oracle_runner_instantiation(self):
        """Test that OracleRunner can be instantiated with required parameters."""
        from vanna.integrations.oracle import OracleRunner

        runner = OracleRunner(
            user="test-user", password="test-pass", dsn="localhost:1521/ORCL"
        )
        assert runner is not None
        assert runner.user == "test-user"


class TestBigQueryRunner:
    """Sanity tests for BigQueryRunner implementation."""

    def test_bigquery_runner_import(self):
        """Test that BigQueryRunner can be imported."""
        from vanna.integrations.bigquery import BigQueryRunner

        assert BigQueryRunner is not None

    def test_bigquery_runner_implements_sql_runner(self):
        """Test that BigQueryRunner implements SqlRunner interface."""
        from vanna.integrations.bigquery import BigQueryRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(BigQueryRunner, SqlRunner)

    def test_bigquery_runner_has_run_sql_method(self):
        """Test that BigQueryRunner implements run_sql method."""
        from vanna.integrations.bigquery import BigQueryRunner

        assert hasattr(BigQueryRunner, "run_sql")
        assert not getattr(BigQueryRunner.run_sql, "__isabstractmethod__", False)

    def test_bigquery_runner_instantiation(self):
        """Test that BigQueryRunner can be instantiated with required parameters."""
        from vanna.integrations.bigquery import BigQueryRunner

        runner = BigQueryRunner(project_id="test-project")
        assert runner is not None
        assert runner.project_id == "test-project"


class TestDuckDBRunner:
    """Sanity tests for DuckDBRunner implementation."""

    def test_duckdb_runner_import(self):
        """Test that DuckDBRunner can be imported."""
        from vanna.integrations.duckdb import DuckDBRunner

        assert DuckDBRunner is not None

    def test_duckdb_runner_implements_sql_runner(self):
        """Test that DuckDBRunner implements SqlRunner interface."""
        from vanna.integrations.duckdb import DuckDBRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(DuckDBRunner, SqlRunner)

    def test_duckdb_runner_has_run_sql_method(self):
        """Test that DuckDBRunner implements run_sql method."""
        from vanna.integrations.duckdb import DuckDBRunner

        assert hasattr(DuckDBRunner, "run_sql")
        assert not getattr(DuckDBRunner.run_sql, "__isabstractmethod__", False)

    def test_duckdb_runner_instantiation_memory(self):
        """Test that DuckDBRunner can be instantiated for in-memory database."""
        from vanna.integrations.duckdb import DuckDBRunner

        runner = DuckDBRunner(database_path=":memory:")
        assert runner is not None
        assert runner.database_path == ":memory:"

    def test_duckdb_runner_instantiation_file(self):
        """Test that DuckDBRunner can be instantiated with file path."""
        from vanna.integrations.duckdb import DuckDBRunner

        runner = DuckDBRunner(database_path="/tmp/test.duckdb")
        assert runner is not None
        assert runner.database_path == "/tmp/test.duckdb"


class TestMSSQLRunner:
    """Sanity tests for MSSQLRunner implementation."""

    def test_mssql_runner_import(self):
        """Test that MSSQLRunner can be imported."""
        from vanna.integrations.mssql import MSSQLRunner

        assert MSSQLRunner is not None

    def test_mssql_runner_implements_sql_runner(self):
        """Test that MSSQLRunner implements SqlRunner interface."""
        from vanna.integrations.mssql import MSSQLRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(MSSQLRunner, SqlRunner)

    def test_mssql_runner_has_run_sql_method(self):
        """Test that MSSQLRunner implements run_sql method."""
        from vanna.integrations.mssql import MSSQLRunner

        assert hasattr(MSSQLRunner, "run_sql")
        assert not getattr(MSSQLRunner.run_sql, "__isabstractmethod__", False)

    def test_mssql_runner_instantiation(self):
        """Test that MSSQLRunner can be instantiated with ODBC connection string."""
        from vanna.integrations.mssql import MSSQLRunner

        runner = MSSQLRunner(
            odbc_conn_str="Driver={SQL Server};Server=localhost;Database=test;Trusted_Connection=yes;"
        )
        assert runner is not None


class TestPrestoRunner:
    """Sanity tests for PrestoRunner implementation."""

    def test_presto_runner_import(self):
        """Test that PrestoRunner can be imported."""
        from vanna.integrations.presto import PrestoRunner

        assert PrestoRunner is not None

    def test_presto_runner_implements_sql_runner(self):
        """Test that PrestoRunner implements SqlRunner interface."""
        from vanna.integrations.presto import PrestoRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(PrestoRunner, SqlRunner)

    def test_presto_runner_has_run_sql_method(self):
        """Test that PrestoRunner implements run_sql method."""
        from vanna.integrations.presto import PrestoRunner

        assert hasattr(PrestoRunner, "run_sql")
        assert not getattr(PrestoRunner.run_sql, "__isabstractmethod__", False)

    def test_presto_runner_instantiation(self):
        """Test that PrestoRunner can be instantiated with required parameters."""
        from vanna.integrations.presto import PrestoRunner

        runner = PrestoRunner(host="localhost", user="test-user", password="test-pass")
        assert runner is not None
        assert runner.host == "localhost"


class TestHiveRunner:
    """Sanity tests for HiveRunner implementation."""

    def test_hive_runner_import(self):
        """Test that HiveRunner can be imported."""
        from vanna.integrations.hive import HiveRunner

        assert HiveRunner is not None

    def test_hive_runner_implements_sql_runner(self):
        """Test that HiveRunner implements SqlRunner interface."""
        from vanna.integrations.hive import HiveRunner
        from vanna.capabilities.sql_runner import SqlRunner

        assert issubclass(HiveRunner, SqlRunner)

    def test_hive_runner_has_run_sql_method(self):
        """Test that HiveRunner implements run_sql method."""
        from vanna.integrations.hive import HiveRunner

        assert hasattr(HiveRunner, "run_sql")
        assert not getattr(HiveRunner.run_sql, "__isabstractmethod__", False)

    def test_hive_runner_instantiation(self):
        """Test that HiveRunner can be instantiated with required parameters."""
        from vanna.integrations.hive import HiveRunner

        runner = HiveRunner(host="localhost", user="test-user", password="test-pass")
        assert runner is not None
        assert runner.host == "localhost"


================================================
FILE: tests/test_gemini_integration.py
================================================
"""
Google Gemini integration tests.

Basic unit tests for the Gemini LLM service integration.
End-to-end tests are in test_agents.py.

Note: Tests requiring API calls need GOOGLE_API_KEY environment variable.
"""

import os
import pytest
from vanna.core.llm import LlmRequest, LlmMessage
from vanna.core.tool import ToolSchema
from vanna.core.user import User


@pytest.fixture
def test_user():
    """Test user for LLM requests."""
    return User(
        id="test_user",
        username="test",
        email="test@example.com",
        group_memberships=["user"],
    )


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_import():
    """Test that Gemini integration can be imported."""
    from vanna.integrations.google import GeminiLlmService

    print("✓ GeminiLlmService imported successfully")
    assert GeminiLlmService is not None


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_initialization_without_key():
    """Test that Gemini service raises error without API key."""
    from vanna.integrations.google import GeminiLlmService

    # Clear both env vars if they exist
    old_google_key = os.environ.pop("GOOGLE_API_KEY", None)
    old_gemini_key = os.environ.pop("GEMINI_API_KEY", None)

    try:
        with pytest.raises(ValueError, match="Google API key is required"):
            llm = GeminiLlmService(model="gemini-2.5-pro")
    finally:
        # Restore the keys if they existed
        if old_google_key:
            os.environ["GOOGLE_API_KEY"] = old_google_key
        if old_gemini_key:
            os.environ["GEMINI_API_KEY"] = old_gemini_key


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_initialization():
    """Test that Gemini service can be initialized with API key."""
    from vanna.integrations.google import GeminiLlmService

    # This test will be skipped by conftest.py if GOOGLE_API_KEY is not set
    llm = GeminiLlmService(
        model="gemini-2.5-pro",
        temperature=0.7,
    )

    print(f"✓ GeminiLlmService initialized")
    print(f"  Model: {llm.model_name}")
    print(f"  Temperature: {llm.temperature}")

    assert llm.model_name == "gemini-2.5-pro"
    assert llm.temperature == 0.7


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_basic_request(test_user):
    """Test a basic request without tools."""
    from vanna.integrations.google import GeminiLlmService

    llm = GeminiLlmService(model="gemini-2.5-pro", temperature=0.0)

    request = LlmRequest(
        user=test_user,
        messages=[
            LlmMessage(role="user", content="What is 2+2? Answer with just the number.")
        ],
    )

    print(f"\n=== Basic Request Test ===")
    print(f"Sending request to Gemini...")

    response = await llm.send_request(request)

    print(f"✓ Response received")
    print(f"  Content type: {type(response.content)}")
    print(f"  Content: {response.content}")
    print(f"  Finish reason: {response.finish_reason}")
    print(f"  Tool calls: {response.tool_calls}")
    print(f"  Usage: {response.usage}")

    # Verify response structure
    assert response is not None
    assert response.content is not None
    assert isinstance(response.content, str)
    assert "4" in response.content


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_streaming_request(test_user):
    """Test streaming request."""
    from vanna.integrations.google import GeminiLlmService

    llm = GeminiLlmService(model="gemini-2.5-pro", temperature=0.0)

    request = LlmRequest(
        user=test_user,
        messages=[
            LlmMessage(role="user", content="Count from 1 to 5, one number per line.")
        ],
        stream=True,
    )

    print(f"\n=== Streaming Request Test ===")
    print(f"Streaming from Gemini...")

    chunks = []
    async for chunk in llm.stream_request(request):
        chunks.append(chunk)
        if chunk.content:
            print(f"  Chunk: {chunk.content}")

    print(f"✓ Streaming completed")
    print(f"  Total chunks: {len(chunks)}")

    # Verify we got chunks
    assert len(chunks) > 0
    # At least one chunk should have content
    assert any(c.content for c in chunks)


@pytest.mark.gemini
@pytest.mark.asyncio
async def test_gemini_validate_tools():
    """Test tool validation (does not require API key for actual calls)."""
    from vanna.integrations.google import GeminiLlmService

    # For validation testing, we need to initialize but won't make API calls
    llm = GeminiLlmService(model="gemini-2.5-pro")

    # Valid tool
    valid_tool = ToolSchema(
        name="test_tool",
        description="A test tool",
        parameters={"type": "object", "properties": {}},
    )

    # Invalid tool (no name)
    invalid_tool = ToolSchema(
        name="",
        description="Invalid tool",
        parameters={"type": "object", "properties": {}},
    )

    # Invalid tool (no description)
    invalid_tool2 = ToolSchema(
        name="test_tool_2",
        description="",
        parameters={"type": "object", "properties": {}},
    )

    errors = await llm.validate_tools([valid_tool])
    assert len(errors) == 0, "Valid tool should have no errors"

    errors = await llm.validate_tools([invalid_tool])
    assert len(errors) > 0, "Tool with empty name should have errors"
    assert any("name" in e.lower() for e in errors)

    errors = await llm.validate_tools([invalid_tool2])
    assert len(errors) > 0, "Tool with empty description should have errors"
    assert any("description" in e.lower() for e in errors)

    print("✓ Tool validation working correctly")


================================================
FILE: tests/test_legacy_adapter.py
================================================
"""
Test for LegacyVannaAdapter retrofit functionality.

This test validates that a legacy VannaBase instance can be wrapped
with LegacyVannaAdapter and used with the new Agents framework.
"""

import os
import pytest
from vanna.core.user import User
from vanna.core.user.resolver import UserResolver
from vanna.core.user.request_context import RequestContext


class SimpleUserResolver(UserResolver):
    """Simple user resolver for tests - returns test user or admin."""

    async def resolve_user(self, request_context: RequestContext) -> User:
        user_email = request_context.cookies.get("vanna_email", "test@example.com")

        if user_email == "admin@example.com":
            return User(id="admin_user", email=user_email, group_memberships=["admin"])

        return User(id=user_email, email=user_email, group_memberships=["user"])


@pytest.mark.legacy
@pytest.mark.asyncio
async def test_legacy_adapter_with_anthropic():
    """Test LegacyVannaAdapter wrapping a legacy VannaBase instance with Anthropic LLM."""
    from vanna import Agent, AgentConfig
    from vanna.legacy.adapter import LegacyVannaAdapter
    from vanna.legacy.mock import MockLLM
    from vanna.legacy.chromadb import ChromaDB_VectorStore
    from vanna.integrations.anthropic import AnthropicLlmService

    # Create a legacy VannaBase instance (using multiple inheritance like v0.x)
    class MyVanna(ChromaDB_VectorStore, MockLLM):
        def __init__(self, config=None):
            ChromaDB_VectorStore.__init__(self, config=config)
            MockLLM.__init__(self, config=config)

    vn = MyVanna()

    # Connect to the Chinook database using legacy method
    vn.connect_to_sqlite("https://vanna.ai/Chinook.sqlite")

    # Add some training data
    vn.add_question_sql(
        question="Who is the top artist by sales?",
        sql="SELECT a.Name, SUM(il.UnitPrice * il.Quantity) as total FROM Artist a JOIN Album al ON a.ArtistId = al.ArtistId JOIN Track t ON al.AlbumId = t.AlbumId JOIN InvoiceLine il ON t.TrackId = il.TrackId GROUP BY a.ArtistId ORDER BY total DESC LIMIT 1",
    )

    # Wrap legacy VannaBase with LegacyVannaAdapter
    legacy_adapter = LegacyVannaAdapter(vn)

    # Create agent with new LLM service
    api_key = os.getenv("ANTHROPIC_API_KEY")
    llm = AnthropicLlmService(api_key=api_key, model="claude-haiku-4-5")

    agent = Agent(
        llm_service=llm,
        tool_registry=legacy_adapter,  # LegacyVannaAdapter is a ToolRegistry
        agent_memory=legacy_adapter,  # LegacyVannaAdapter implements AgentMemory
        user_resolver=SimpleUserResolver(),
        config=AgentConfig(),
    )

    # Test that the agent can answer a question
    request_context = RequestContext(cookies={}, headers={})

    components = []
    async for component in agent.send_message(
        request_context, "Who is the top artist by sales?"
    ):
        components.append(component)

    # Validate we got components
    assert len(components) > 0, "Should receive at least one component"

    # Look for a successful response (either table or text mentioning an artist)
    has_response = False
    for component in components:
        if hasattr(component, "rich_component") and component.rich_component:
            has_response = True
            break
        if hasattr(component, "simple_component") and component.simple_component:
            if (
                hasattr(component.simple_component, "text")
                and component.simple_component.text
            ):
                has_response = True
                break

    assert has_response, "Should receive at least one response component"


@pytest.mark.legacy
@pytest.mark.asyncio
async def test_legacy_adapter_memory_operations():
    """Test that LegacyVannaAdapter properly implements AgentMemory interface."""
    from vanna.legacy.adapter import LegacyVannaAdapter
    from vanna.legacy.mock import MockLLM
    from vanna.legacy.chromadb import ChromaDB_VectorStore
    from vanna.core.tool import ToolContext
    from vanna.core.user import User

    # Create a legacy VannaBase instance
    class MyVanna(ChromaDB_VectorStore, MockLLM):
        def __init__(self, config=None):
            ChromaDB_VectorStore.__init__(self, config=config)
            MockLLM.__init__(self, config=config)

    vn = MyVanna()
    adapter = LegacyVannaAdapter(vn)

    # Create a properly constructed tool context
    user = User(id="test_user", email="test@example.com", group_memberships=["user"])
    context = ToolContext(
        user=user,
        conversation_id="test-conversation",
        request_id="test-request",
        agent_memory=adapter,  # Use the adapter itself as the agent_memory
    )

    # Test save_tool_usage
    await adapter.save_tool_usage(
        question="What are the total sales?",
        tool_name="run_sql",
        args={"sql": "SELECT SUM(Total) FROM Invoice"},
        context=context,
        success=True,
    )

    # Test search_similar_usage
    results = await adapter.search_similar_usage(
        question="What are sales?", context=context, limit=5
    )

    # Should find the saved usage
    assert len(results) > 0, "Should find similar tool usage"
    assert results[0].memory.question == "What are the total sales?"
    assert results[0].memory.args["sql"] == "SELECT SUM(Total) FROM Invoice"

    # Test save_text_memory
    text_memory = await adapter.save_text_memory(
        content="The Invoice table contains all sales transactions.", context=context
    )
    assert text_memory.content == "The Invoice table contains all sales transactions."

    # Test search_text_memories
    text_results = await adapter.search_text_memories(
        query="sales", context=context, limit=5
    )
    assert len(text_results) > 0, "Should find similar text memories"

    # Test get_recent_memories (uses blank string retrieval)
    recent = await adapter.get_recent_memories(context=context, limit=5)
    assert isinstance(recent, list), "Should return a list of memories"

    # Test get_recent_text_memories
    recent_text = await adapter.get_recent_text_memories(context=context, limit=5)
    assert isinstance(recent_text, list), "Should return a list of text memories"


================================================
FILE: tests/test_llm_context_enhancer.py
================================================
"""
Unit tests for LlmContextEnhancer functionality.

These tests validate that the Agent properly calls the LlmContextEnhancer methods
to enhance system prompts and user messages.
"""

import pytest
from typing import List
from vanna.core.enhancer.base import LlmContextEnhancer
from vanna.core.enhancer.default import DefaultLlmContextEnhancer
from vanna.core.user import User
from vanna.core.user.resolver import UserResolver
from vanna.core.user.request_context import RequestContext
from vanna.core.llm.models import LlmMessage
from vanna.core.llm.base import LlmService
from vanna.capabilities.agent_memory import (
    AgentMemory,
    TextMemory,
    TextMemorySearchResult,
)
from vanna.core.tool import ToolContext


class MockAgentMemory(AgentMemory):
    """Mock AgentMemory for testing."""

    def __init__(self):
        self.text_memories: List[TextMemory] = []

    async def save_tool_usage(
        self, question, tool_name, args, context, success=True, metadata=None
    ):
        pass

    async def save_text_memory(self, content, context):
        memory = TextMemory(
            memory_id=f"mem-{len(self.text_memories)}", content=content, timestamp=None
        )
        self.text_memories.append(memory)
        return memory

    async def search_similar_usage(
        self,
        question,
        context,
        *,
        limit=10,
        similarity_threshold=0.7,
        tool_name_filter=None,
    ):
        return []

    async def search_text_memories(
        self, query, context, *, limit=10, similarity_threshold=0.7
    ):
        """Return mock search results based on stored memories."""
        results = []
        for idx, memory in enumerate(self.text_memories[:limit]):
            results.append(
                TextMemorySearchResult(
                    memory=memory, similarity_score=0.9 - (idx * 0.1), rank=idx + 1
                )
            )
        return results

    async def get_recent_memories(self, context, limit=10):
        return []

    async def get_recent_text_memories(self, context, limit=10):
        return self.text_memories[:limit]

    async def delete_by_id(self, context, memory_id):
        return False

    async def delete_text_memory(self, context, memory_id):
        return False

    async def clear_memories(self, context, tool_name=None, before_date=None):
        return 0


class MockLlmService(LlmService):
    """Mock LLM service that records calls."""

    def __init__(self):
        self.requests = []  # Store full request objects
        self.next_response = "Mock response"

    async def send_request(self, request):
        """Record the call and return a mock response."""
        from vanna.core.llm.models import LlmResponse, LlmResponseMessage

        # Store the full request object
        self.requests.append(request)

        # Return a simple text response
        return LlmResponse(
            message=LlmResponseMessage(role="assistant", content=self.next_response),
            finish_reason="end_turn",
        )

    async def stream_request(self, request):
        """Mock streaming - just yield a single chunk."""
        from vanna.core.llm.models import LlmStreamChunk

        # Store the full request object
        self.requests.append(request)

        yield LlmStreamChunk(delta=self.next_response, finish_reason="end_turn")

    async def validate_tools(self, tools):
        """Mock validation - no errors."""
        return []


class SimpleUserResolver(UserResolver):
    """Simple user resolver for tests."""

    async def resolve_user(self, request_context: RequestContext) -> User:
        return User(
            id="test_user", email="test@example.com", group_memberships=["user"]
        )


class TrackingEnhancer(LlmContextEnhancer):
    """Custom enhancer that tracks when methods are called."""

    def __init__(self):
        self.enhance_system_prompt_calls = []
        self.enhance_user_messages_calls = []

    async def enhance_system_prompt(
        self, system_prompt: str, user_message: str, user: User
    ) -> str:
        """Track call and add a marker to the system prompt."""
        self.enhance_system_prompt_calls.append(
            {
                "system_prompt": system_prompt,
                "user_message": user_message,
                "user_id": user.id,
            }
        )
        return system_prompt + "\n\n[ENHANCED_SYSTEM_PROMPT]"

    async def enhance_user_messages(
        self, messages: List[LlmMessage], user: User
    ) -> List[LlmMessage]:
        """Track call and add a marker to messages."""
        self.enhance_user_messages_calls.append(
            {"message_count": len(messages), "user_id": user.id}
        )
        # Add a marker to the last user message
        if messages and messages[-1].role == "user":
            enhanced_messages = messages[:-1] + [
                LlmMessage(
                    role="user",
                    content=messages[-1].content + " [ENHANCED_USER_MESSAGE]",
                )
            ]
            return enhanced_messages
        return messages


@pytest.mark.asyncio
async def test_custom_enhancer_system_prompt_is_called():
    """Test that a custom LlmContextEnhancer.enhance_system_prompt is called by the Agent."""
    from vanna import Agent, AgentConfig
    from vanna.core.registry import ToolRegistry

    # Create mock components
    llm = MockLlmService()
    tools = ToolRegistry()
    enhancer = TrackingEnhancer()
    agent_memory = MockAgentMemory()

    # Create agent with the tracking enhancer
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=SimpleUserResolver(),
        agent_memory=agent_memory,
        llm_context_enhancer=enhancer,
        config=AgentConfig(),
    )

    # Send a message
    request_context = RequestContext(cookies={}, headers={})
    components = []
    async for component in agent.send_message(request_context, "Hello, world!"):
        components.append(component)

    # Verify enhance_system_prompt was called
    assert len(enhancer.enhance_system_prompt_calls) == 1, (
        "enhance_system_prompt should be called exactly once"
    )
    call = enhancer.enhance_system_prompt_calls[0]
    assert call["user_message"] == "Hello, world!", (
        "Should pass the user message to enhancer"
    )
    assert call["user_id"] == "test_user", "Should pass the correct user"
    assert "system_prompt" in call, "Should pass the system prompt"

    # Verify the enhanced system prompt was sent to the LLM
    assert len(llm.requests) >= 1, "LLM should be called at least once"
    first_request = llm.requests[0]
    assert first_request.system_prompt is not None, "Should have a system prompt"
    assert "[ENHANCED_SYSTEM_PROMPT]" in first_request.system_prompt, (
        f"System prompt should contain enhancement marker. Got: {first_request.system_prompt}"
    )


@pytest.mark.asyncio
async def test_custom_enhancer_user_messages_is_called():
    """Test that a custom LlmContextEnhancer.enhance_user_messages is called by the Agent."""
    from vanna import Agent, AgentConfig
    from vanna.core.registry import ToolRegistry

    # Create mock components
    llm = MockLlmService()
    tools = ToolRegistry()
    enhancer = TrackingEnhancer()
    agent_memory = MockAgentMemory()

    # Create agent with the tracking enhancer
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=SimpleUserResolver(),
        agent_memory=agent_memory,
        llm_context_enhancer=enhancer,
        config=AgentConfig(),
    )

    # Send a message
    request_context = RequestContext(cookies={}, headers={})
    components = []
    async for component in agent.send_message(request_context, "Test message"):
        components.append(component)

    # Verify enhance_user_messages was called
    assert len(enhancer.enhance_user_messages_calls) >= 1
    call = enhancer.enhance_user_messages_calls[0]
    assert call["user_id"] == "test_user"
    assert call["message_count"] > 0

    # Verify the enhanced user message was sent to the LLM
    assert len(llm.requests) >= 1
    first_request = llm.requests[0]
    user_messages = [m for m in first_request.messages if m.role == "user"]
    assert len(user_messages) >= 1, "Should have at least one user message"
    assert "[ENHANCED_USER_MESSAGE]" in user_messages[0].content, (
        f"User message should be enhanced. Got: {user_messages[0].content}"
    )


@pytest.mark.asyncio
async def test_default_enhancer_with_agent_memory():
    """Test that DefaultLlmContextEnhancer properly enhances system prompt with memories."""
    from vanna import Agent, AgentConfig
    from vanna.core.registry import ToolRegistry

    # Create mock components
    llm = MockLlmService()
    tools = ToolRegistry()
    agent_memory = MockAgentMemory()

    # Add some test memories
    user = User(id="test_user", email="test@example.com", group_memberships=["user"])
    context = ToolContext(
        user=user, conversation_id="test", request_id="test", agent_memory=agent_memory
    )

    await agent_memory.save_text_memory(
        "The database has a users table with columns: id, name, email", context
    )
    await agent_memory.save_text_memory(
        "The products table contains: product_id, name, price, category", context
    )

    # Create default enhancer with agent memory
    enhancer = DefaultLlmContextEnhancer(agent_memory=agent_memory)

    # Create agent
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=SimpleUserResolver(),
        agent_memory=agent_memory,
        llm_context_enhancer=enhancer,
        config=AgentConfig(),
    )

    # Send a message that should trigger memory retrieval
    request_context = RequestContext(cookies={}, headers={})
    components = []
    async for component in agent.send_message(
        request_context, "Show me the database schema"
    ):
        components.append(component)

    # Verify that the DefaultLlmContextEnhancer added memory context to the system prompt
    assert len(llm.requests) >= 1, "LLM should be called"
    first_request = llm.requests[0]

    # The DefaultLlmContextEnhancer should add "Relevant Context from Memory" to system prompt
    assert first_request.system_prompt is not None, "Should have a system prompt"
    assert "Relevant Context from Memory" in first_request.system_prompt, (
        f"System prompt should include memory context. Got: {first_request.system_prompt}"
    )

    # Should contain one or both of the memories we added
    assert (
        "users table" in first_request.system_prompt
        or "products table" in first_request.system_prompt
    ), (
        f"System prompt should contain our test memories. Got: {first_request.system_prompt}"
    )


@pytest.mark.asyncio
async def test_default_enhancer_without_agent_memory():
    """Test that DefaultLlmContextEnhancer works without agent memory (no enhancement)."""
    from vanna import Agent, AgentConfig
    from vanna.core.registry import ToolRegistry

    # Create mock components
    llm = MockLlmService()
    tools = ToolRegistry()

    # Create default enhancer without agent memory
    enhancer = DefaultLlmContextEnhancer(agent_memory=None)
    agent_memory = MockAgentMemory()

    # Create agent
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=SimpleUserResolver(),
        agent_memory=agent_memory,
        llm_context_enhancer=enhancer,
        config=AgentConfig(),
    )

    # Send a message
    request_context = RequestContext(cookies={}, headers={})
    components = []
    async for component in agent.send_message(request_context, "Hello"):
        components.append(component)

    # Verify the system prompt does NOT include memory context (since enhancer has no agent_memory)
    assert len(llm.requests) >= 1, "LLM should be called"
    first_request = llm.requests[0]

    # The system prompt should exist but NOT contain memory context
    if first_request.system_prompt:
        assert "Relevant Context from Memory" not in first_request.system_prompt, (
            "Should not add memory context when enhancer has no agent_memory"
        )


@pytest.mark.asyncio
async def test_no_enhancer_means_no_enhancement():
    """Test that when no enhancer is provided, no enhancement occurs."""
    from vanna import Agent, AgentConfig
    from vanna.core.registry import ToolRegistry

    # Create mock components
    llm = MockLlmService()
    tools = ToolRegistry()
    agent_memory = MockAgentMemory()

    # Create agent WITHOUT an enhancer (should use default no-op behavior)
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=SimpleUserResolver(),
        agent_memory=agent_memory,
        config=AgentConfig(),
    )

    # Send a message
    request_context = RequestContext(cookies={}, headers={})
    components = []
    async for component in agent.send_message(request_context, "Hello"):
        components.append(component)

    # Verify basic system prompt exists but has no enhancements
    assert len(llm.requests) >= 1, "LLM should be called"
    first_request = llm.requests[0]

    # Should be the basic system prompt with no enhancements
    if first_request.system_prompt:
        assert "Relevant Context from Memory" not in first_request.system_prompt, (
            "Should not add memory context when no enhancer is provided"
        )


================================================
FILE: tests/test_memory_tools.py
================================================
"""
Tests for agent memory tools, including UI feature access control.
"""

import pytest
import uuid
from vanna.tools.agent_memory import (
    SearchSavedCorrectToolUsesTool,
    SearchSavedCorrectToolUsesParams,
)
from vanna.core.tool import ToolContext
from vanna.core.user import User
from vanna.core.agent.config import UiFeature
from vanna.integrations.local.agent_memory import DemoAgentMemory
from vanna.core.rich_component import ComponentType


@pytest.fixture
def demo_agent_memory():
    """Create a demo agent memory instance."""
    return DemoAgentMemory(max_items=100)


@pytest.fixture
def admin_user():
    """Create an admin user."""
    return User(id="admin", email="admin@example.com", group_memberships=["admin"])


@pytest.fixture
def regular_user():
    """Create a regular user."""
    return User(id="user", email="user@example.com", group_memberships=["user"])


@pytest.fixture
def search_tool():
    """Create a search tool instance."""
    return SearchSavedCorrectToolUsesTool()


class TestMemoryToolDetailedResults:
    """Test memory tool detailed results feature."""

    @pytest.mark.asyncio
    async def test_admin_sees_detailed_results(
        self, search_tool, demo_agent_memory, admin_user
    ):
        """Test that admin users see detailed memory results in a collapsible card."""
        # Create context with admin user and feature enabled
        context = ToolContext(
            user=admin_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={
                "ui_features_available": [
                    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS
                ]
            },
        )

        # Save some memories
        await demo_agent_memory.save_tool_usage(
            question="What is the total sales?",
            tool_name="run_sql",
            args={"query": "SELECT SUM(total) FROM invoices"},
            context=context,
            success=True,
        )

        # Search for similar patterns
        search_params = SearchSavedCorrectToolUsesParams(
            question="What are the total sales?", limit=10, similarity_threshold=0.5
        )

        result = await search_tool.execute(context, search_params)

        # Verify result
        assert result.success is True
        assert result.ui_component is not None
        assert result.ui_component.rich_component is not None

        # Check that it's a CardComponent (detailed view)
        assert result.ui_component.rich_component.type == ComponentType.CARD

        # Check collapsible properties
        card = result.ui_component.rich_component
        assert card.collapsible is True
        assert card.collapsed is True  # Should start collapsed

        # Verify content includes detailed information
        assert "Retrieved memories passed to LLM" in card.content
        assert "run_sql" in card.content
        assert "similarity:" in card.content.lower()
        assert "Question:" in card.content
        assert "Arguments:" in card.content

    @pytest.mark.asyncio
    async def test_non_admin_sees_simple_status(
        self, search_tool, demo_agent_memory, regular_user
    ):
        """Test that non-admin users see simple status message."""
        # Create context with regular user (no detailed results feature)
        context = ToolContext(
            user=regular_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={"ui_features_available": []},  # No detailed results feature
        )

        # Save some memories
        await demo_agent_memory.save_tool_usage(
            question="What is the total sales?",
            tool_name="run_sql",
            args={"query": "SELECT SUM(total) FROM invoices"},
            context=context,
            success=True,
        )

        # Search for similar patterns
        search_params = SearchSavedCorrectToolUsesParams(
            question="What are the total sales?", limit=10, similarity_threshold=0.5
        )

        result = await search_tool.execute(context, search_params)

        # Verify result
        assert result.success is True
        assert result.ui_component is not None
        assert result.ui_component.rich_component is not None

        # Check that it's a StatusBarUpdateComponent (simple view)
        assert (
            result.ui_component.rich_component.type == ComponentType.STATUS_BAR_UPDATE
        )

        # Verify it shows success message
        status = result.ui_component.rich_component
        assert status.status == "success"
        assert "similar pattern" in status.message.lower()

    @pytest.mark.asyncio
    async def test_detailed_results_include_all_memory_fields(
        self, search_tool, demo_agent_memory, admin_user
    ):
        """Test that detailed results include all relevant memory fields."""
        # Create context with admin user and feature enabled
        context = ToolContext(
            user=admin_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={
                "ui_features_available": [
                    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS
                ]
            },
        )

        # Save a memory
        await demo_agent_memory.save_tool_usage(
            question="Show me customer names",
            tool_name="run_sql",
            args={"query": "SELECT name FROM customers"},
            context=context,
            success=True,
        )

        # Search for it
        search_params = SearchSavedCorrectToolUsesParams(
            question="Show customer names", limit=10, similarity_threshold=0.3
        )

        result = await search_tool.execute(context, search_params)

        # Verify detailed content
        card = result.ui_component.rich_component
        content = card.content

        # Check for all expected fields
        assert "Question:" in content
        assert "Show me customer names" in content
        assert "Arguments:" in content
        assert "run_sql" in content
        assert "similarity:" in content.lower()

        # Timestamp and ID should be included if available
        # (DemoAgentMemory might not set these, but the code should handle them)

    @pytest.mark.asyncio
    async def test_no_results_works_for_both_admin_and_user(
        self, search_tool, demo_agent_memory, admin_user, regular_user
    ):
        """Test that admin sees card with 0 results while regular user sees status bar."""
        # Test with admin
        admin_context = ToolContext(
            user=admin_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={
                "ui_features_available": [
                    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS
                ]
            },
        )

        search_params = SearchSavedCorrectToolUsesParams(
            question="This query will not match anything",
            limit=10,
            similarity_threshold=0.99,
        )

        admin_result = await search_tool.execute(admin_context, search_params)

        assert admin_result.success is True
        assert "No similar tool usage patterns found" in admin_result.result_for_llm
        # Admin should see a card showing 0 results
        assert admin_result.ui_component.rich_component.type == ComponentType.CARD
        assert "0 Results" in admin_result.ui_component.rich_component.title
        assert admin_result.ui_component.rich_component.collapsible is True

        # Test with regular user
        user_context = ToolContext(
            user=regular_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={"ui_features_available": []},
        )

        user_result = await search_tool.execute(user_context, search_params)

        assert user_result.success is True
        assert "No similar tool usage patterns found" in user_result.result_for_llm
        # Regular user should see a status bar update
        assert (
            user_result.ui_component.rich_component.type
            == ComponentType.STATUS_BAR_UPDATE
        )

    @pytest.mark.asyncio
    async def test_llm_result_same_for_admin_and_user(
        self, search_tool, demo_agent_memory, admin_user, regular_user
    ):
        """Test that the LLM receives the same information regardless of UI feature access."""
        # Save a memory
        admin_context = ToolContext(
            user=admin_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={
                "ui_features_available": [
                    UiFeature.UI_FEATURE_SHOW_MEMORY_DETAILED_RESULTS
                ]
            },
        )

        await demo_agent_memory.save_tool_usage(
            question="Count all records",
            tool_name="run_sql",
            args={"query": "SELECT COUNT(*) FROM table"},
            context=admin_context,
            success=True,
        )

        search_params = SearchSavedCorrectToolUsesParams(
            question="Count records", limit=10, similarity_threshold=0.3
        )

        # Get admin result
        admin_result = await search_tool.execute(admin_context, search_params)

        # Get regular user result
        user_context = ToolContext(
            user=regular_user,
            conversation_id=str(uuid.uuid4()),
            request_id=str(uuid.uuid4()),
            agent_memory=demo_agent_memory,
            metadata={"ui_features_available": []},
        )

        user_result = await search_tool.execute(user_context, search_params)

        # Both should have the same result_for_llm
        assert admin_result.result_for_llm == user_result.result_for_llm
        assert "Found" in admin_result.result_for_llm
        assert "similar tool usage pattern" in admin_result.result_for_llm


if __name__ == "__main__":
    pytest.main([__file__, "-v"])


================================================
FILE: tests/test_ollama_direct.py
================================================
"""
Direct Ollama integration tests to diagnose and verify the integration.

These tests check each aspect of the Ollama integration separately.
"""

import pytest
from vanna.core.llm import LlmRequest, LlmMessage
from vanna.core.tool import ToolSchema
from vanna.core.user import User


@pytest.fixture
def test_user():
    """Test user for LLM requests."""
    return User(
        id="test_user",
        username="test",
        email="test@example.com",
        group_memberships=["user"],
    )


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_import():
    """Test that Ollama integration can be imported."""
    try:
        from vanna.integrations.ollama import OllamaLlmService

        print("✓ OllamaLlmService imported successfully")
        assert OllamaLlmService is not None
    except ImportError as e:
        pytest.fail(f"Failed to import OllamaLlmService: {e}")


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_initialization():
    """Test that Ollama service can be initialized."""
    from vanna.integrations.ollama import OllamaLlmService

    try:
        llm = OllamaLlmService(
            model="llama3.2", host="http://localhost:11434", temperature=0.0
        )

        print(f"✓ OllamaLlmService initialized")
        print(f"  Model: {llm.model}")
        print(f"  Host: {llm.host}")
        print(f"  Temperature: {llm.temperature}")
        print(f"  Context window: {llm.num_ctx}")

        assert llm.model == "llama3.2"
        assert llm.host == "http://localhost:11434"

    except Exception as e:
        pytest.fail(f"Failed to initialize OllamaLlmService: {e}")


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_basic_request(test_user):
    """Test a basic request without tools."""
    from vanna.integrations.ollama import OllamaLlmService

    llm = OllamaLlmService(model="llama3.2", temperature=0.0)

    request = LlmRequest(
        user=test_user,
        messages=[
            LlmMessage(role="user", content="What is 2+2? Answer with just the number.")
        ],
    )

    print(f"\n=== Basic Request Test ===")
    print(f"Sending request to Ollama...")

    try:
        response = await llm.send_request(request)

        print(f"✓ Response received")
        print(f"  Content type: {type(response.content)}")
        print(f"  Content: {response.content}")
        print(f"  Finish reason: {response.finish_reason}")
        print(f"  Tool calls: {response.tool_calls}")
        print(f"  Usage: {response.usage}")

        # Verify response structure
        assert response is not None
        assert response.content is not None
        assert isinstance(response.content, str)

    except Exception as e:
        print(f"❌ Error during request: {e}")
        import traceback

        traceback.print_exc()
        raise


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_pydantic_response(test_user):
    """Test that the response is a valid Pydantic model."""
    from vanna.integrations.ollama import OllamaLlmService
    from vanna.core.llm import LlmResponse

    llm = OllamaLlmService(model="llama3.2", temperature=0.0)

    request = LlmRequest(
        user=test_user, messages=[LlmMessage(role="user", content="Say 'hello'")]
    )

    print(f"\n=== Pydantic Validation Test ===")

    try:
        response = await llm.send_request(request)

        # Verify it's a Pydantic model
        assert isinstance(response, LlmResponse)
        print(f"✓ Response is LlmResponse instance")

        # Test model_dump
        dumped = response.model_dump()
        print(f"✓ model_dump() works: {dumped}")

        # Test model_dump_json
        json_str = response.model_dump_json()
        print(f"✓ model_dump_json() works: {json_str[:100]}...")

        # Test reconstruction
        reconstructed = LlmResponse(**dumped)
        print(f"✓ Reconstruction works")
        assert reconstructed.content == response.content

    except Exception as e:
        print(f"❌ Pydantic error: {e}")
        import traceback

        traceback.print_exc()
        raise


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_streaming(test_user):
    """Test streaming responses."""
    from vanna.integrations.ollama import OllamaLlmService
    from vanna.core.llm import LlmStreamChunk

    llm = OllamaLlmService(model="llama3.2", temperature=0.0)

    request = LlmRequest(
        user=test_user, messages=[LlmMessage(role="user", content="Count from 1 to 3.")]
    )

    print(f"\n=== Streaming Test ===")

    try:
        chunks = []
        content_parts = []

        async for chunk in llm.stream_request(request):
            chunks.append(chunk)

            # Verify chunk is correct type
            assert isinstance(chunk, LlmStreamChunk)

            if chunk.content:
                content_parts.append(chunk.content)
                print(f"  Chunk {len(chunks)}: {chunk.content!r}")

        full_content = "".join(content_parts)

        print(f"✓ Streaming completed")
        print(f"  Total chunks: {len(chunks)}")
        print(f"  Full content: {full_content}")
        print(f"  Final finish_reason: {chunks[-1].finish_reason}")

        assert len(chunks) > 0
        assert chunks[-1].finish_reason is not None

    except Exception as e:
        print(f"❌ Streaming error: {e}")
        import traceback

        traceback.print_exc()
        raise


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_tool_calling_attempt(test_user):
    """Test tool calling with Ollama (may not work with all models)."""
    from vanna.integrations.ollama import OllamaLlmService

    llm = OllamaLlmService(model="llama3.2", temperature=0.0)

    tools = [
        ToolSchema(
            name="get_weather",
            description="Get the current weather for a location",
            parameters={
                "type": "object",
                "properties": {
                    "location": {"type": "string", "description": "City name"}
                },
                "required": ["location"],
            },
        )
    ]

    request = LlmRequest(
        user=test_user,
        system_prompt="You are a helpful assistant. When asked about weather, use the get_weather tool.",
        messages=[
            LlmMessage(role="user", content="What's the weather in San Francisco?")
        ],
        tools=tools,
    )

    print(f"\n=== Tool Calling Test ===")
    print(f"Model: {llm.model}")
    print(f"Tools provided: {[t.name for t in tools]}")

    try:
        response = await llm.send_request(request)

        print(f"\nResponse:")
        print(f"  Content: {response.content}")
        print(f"  Tool calls: {response.tool_calls}")
        print(f"  Finish reason: {response.finish_reason}")

        if response.tool_calls:
            print(f"\n✓ Tool calling works!")
            print(f"  Number of tool calls: {len(response.tool_calls)}")
            for i, tc in enumerate(response.tool_calls):
                print(f"  Tool call {i + 1}:")
                print(f"    ID: {tc.id}")
                print(f"    Name: {tc.name}")
                print(f"    Arguments: {tc.arguments}")
                print(f"    Arguments type: {type(tc.arguments)}")
        else:
            print(f"\n⚠️  Model returned text instead of tool calls")
            print(f"  This is expected for models without tool calling support")
            print(f"  Try using llama3.1:8b or mistral for better tool calling")

        # Test passes either way - we're just diagnosing
        assert response is not None

    except Exception as e:
        print(f"❌ Tool calling error: {e}")
        import traceback

        traceback.print_exc()
        raise


@pytest.mark.ollama
@pytest.mark.asyncio
async def test_ollama_payload_building(test_user):
    """Test that the payload is built correctly."""
    from vanna.integrations.ollama import OllamaLlmService

    llm = OllamaLlmService(model="llama3.2", temperature=0.5, num_ctx=4096)

    tools = [
        ToolSchema(
            name="test_tool",
            description="A test tool",
            parameters={"type": "object", "properties": {}},
        )
    ]

    request = LlmRequest(
        user=test_user,
        system_prompt="You are a test assistant.",
        messages=[LlmMessage(role="user", content="Hello")],
        tools=tools,
    )

    print(f"\n=== Payload Building Test ===")

    try:
        # Access the internal method to inspect payload
        payload = llm._build_payload(request)

        print(f"Built payload:")
        print(f"  Model: {payload.get('model')}")
        print(f"  Messages count: {len(payload.get('messages', []))}")
        print(f"  First message: {payload.get('messages', [{}])[0]}")
        print(f"  Options: {payload.get('options')}")
        print(f"  Tools: {payload.get('tools')}")

        # Verify payload structure
        assert payload["model"] == "llama3.2"
        assert len(payload["messages"]) == 2  # system + user
        assert payload["messages"][0]["role"] == "system"
        assert payload["messages"][1]["role"] == "user"
        assert payload["options"]["temperature"] == 0.5
        assert payload["options"]["num_ctx"] == 4096
        assert "tools" in payload
        assert len(payload["tools"]) == 1

        print(f"✓ Payload structure is correct")

    except Exception as e:
        print(f"❌ Payload building error: {e}")
        import traceback

        traceback.print_exc()
        raise


================================================
FILE: tests/test_tool_permissions.py
================================================
"""
Tests for tool access control and permissions.
"""

import pytest
from pydantic import BaseModel, Field
from typing import Type, TypeVar, Union

from vanna.core.tool import (
    Tool,
    ToolContext,
    ToolResult,
    ToolCall,
    ToolRejection,
    ToolSchema,
)
from vanna.core.user import User
from vanna.core.registry import ToolRegistry
from vanna.components import UiComponent, SimpleTextComponent
from vanna.integrations.local.agent_memory import DemoAgentMemory

T = TypeVar("T")


class SimpleToolArgs(BaseModel):
    """Simple args for testing."""

    message: str = Field(description="A message")


class MockTool(Tool[SimpleToolArgs]):
    """Mock tool for testing."""

    def __init__(self, tool_name: str = "mock_tool"):
        self._tool_name = tool_name

    @property
    def name(self) -> str:
        return self._tool_name

    @property
    def description(self) -> str:
        return f"A mock tool for testing: {self._tool_name}"

    def get_args_schema(self) -> Type[SimpleToolArgs]:
        return SimpleToolArgs

    async def execute(self, context: ToolContext, args: SimpleToolArgs) -> ToolResult:
        return ToolResult(
            success=True, result_for_llm=f"Mock tool executed: {args.message}"
        )


@pytest.fixture
def agent_memory():
    """Agent memory for testing."""
    return DemoAgentMemory(max_items=100)


@pytest.fixture
def admin_user():
    """Admin user with admin group."""
    return User(
        id="admin_1",
        username="admin",
        email="admin@example.com",
        group_memberships=["admin", "user"],
    )


@pytest.fixture
def regular_user():
    """Regular user with user group."""
    return User(
        id="user_1",
        username="user",
        email="user@example.com",
        group_memberships=["user"],
    )


@pytest.fixture
def analyst_user():
    """Analyst user with analyst group."""
    return User(
        id="analyst_1",
        username="analyst",
        email="analyst@example.com",
        group_memberships=["analyst", "user"],
    )


@pytest.fixture
def guest_user():
    """Guest user with no groups."""
    return User(
        id="guest_1", username="guest", email="guest@example.com", group_memberships=[]
    )


@pytest.mark.asyncio
async def test_tool_access_empty_groups_allows_all(regular_user, agent_memory):
    """Test that empty access_groups allows all users."""
    print("\n=== Empty Access Groups Test ===")

    registry = ToolRegistry()
    tool = MockTool("public_tool")

    # Register with empty access groups
    registry.register_local_tool(tool, access_groups=[])

    # Create context
    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    # Execute tool
    tool_call = ToolCall(
        id="call_1", name="public_tool", arguments={"message": "hello"}
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Tool executed with empty access groups")
    print(f"  Success: {result.success}")

    assert result.success is True
    assert "Mock tool executed" in result.result_for_llm


@pytest.mark.asyncio
async def test_tool_access_granted_matching_group(admin_user, agent_memory):
    """Test that user with matching group can access tool."""
    print("\n=== Access Granted Test ===")

    registry = ToolRegistry()
    tool = MockTool("admin_tool")

    # Register with admin-only access
    registry.register_local_tool(tool, access_groups=["admin"])

    context = ToolContext(
        user=admin_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1", name="admin_tool", arguments={"message": "admin action"}
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Admin user accessed admin-only tool")
    print(f"  User groups: {admin_user.group_memberships}")
    print(f"  Tool access groups: ['admin']")
    print(f"  Success: {result.success}")

    assert result.success is True


@pytest.mark.asyncio
async def test_tool_access_denied_no_matching_group(regular_user, agent_memory):
    """Test that user without matching group cannot access tool."""
    print("\n=== Access Denied Test ===")

    registry = ToolRegistry()
    tool = MockTool("admin_tool")

    # Register with admin-only access
    registry.register_local_tool(tool, access_groups=["admin"])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1", name="admin_tool", arguments={"message": "should fail"}
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Regular user denied access to admin-only tool")
    print(f"  User groups: {regular_user.group_memberships}")
    print(f"  Tool access groups: ['admin']")
    print(f"  Success: {result.success}")
    print(f"  Error: {result.error}")

    assert result.success is False
    assert "Insufficient group access" in result.result_for_llm
    assert "admin_tool" in result.result_for_llm


@pytest.mark.asyncio
async def test_tool_access_multiple_allowed_groups(
    analyst_user, admin_user, regular_user, agent_memory
):
    """Test tool with multiple allowed groups."""
    print("\n=== Multiple Allowed Groups Test ===")

    registry = ToolRegistry()
    tool = MockTool("data_tool")

    # Allow both admin and analyst groups
    registry.register_local_tool(tool, access_groups=["admin", "analyst"])

    # Test analyst can access
    analyst_context = ToolContext(
        user=analyst_user,
        conversation_id="test_conv",
        request_id="test_req_1",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(id="call_1", name="data_tool", arguments={"message": "test"})
    result = await registry.execute(tool_call, analyst_context)

    print(f"✓ Analyst accessed tool")
    assert result.success is True

    # Test admin can access
    admin_context = ToolContext(
        user=admin_user,
        conversation_id="test_conv",
        request_id="test_req_2",
        agent_memory=agent_memory,
    )

    result = await registry.execute(tool_call, admin_context)

    print(f"✓ Admin accessed tool")
    assert result.success is True

    # Test regular user cannot access
    user_context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req_3",
        agent_memory=agent_memory,
    )

    result = await registry.execute(tool_call, user_context)

    print(f"✓ Regular user denied access")
    assert result.success is False


@pytest.mark.asyncio
async def test_tool_access_guest_user_denied(guest_user, agent_memory):
    """Test that guest user with no groups cannot access restricted tools."""
    print("\n=== Guest User Denied Test ===")

    registry = ToolRegistry()
    tool = MockTool("restricted_tool")

    registry.register_local_tool(tool, access_groups=["user"])

    context = ToolContext(
        user=guest_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1", name="restricted_tool", arguments={"message": "test"}
    )
    result = await registry.execute(tool_call, context)

    print(f"✓ Guest user with no groups denied access")
    print(f"  User groups: {guest_user.group_memberships}")
    print(f"  Tool access groups: ['user']")

    assert result.success is False
    assert "Insufficient group access" in result.result_for_llm


@pytest.mark.asyncio
async def test_get_schemas_filters_by_user(admin_user, regular_user):
    """Test that get_schemas only returns tools accessible to user."""
    print("\n=== Schema Filtering Test ===")

    registry = ToolRegistry()

    # Register tools with different access levels
    registry.register_local_tool(MockTool("public_tool"), access_groups=[])
    registry.register_local_tool(MockTool("admin_tool"), access_groups=["admin"])
    registry.register_local_tool(MockTool("user_tool"), access_groups=["user"])

    # Admin user should see all tools
    admin_schemas = await registry.get_schemas(admin_user)
    admin_tool_names = [s.name for s in admin_schemas]

    print(f"✓ Admin user schemas: {admin_tool_names}")
    assert "public_tool" in admin_tool_names
    assert "admin_tool" in admin_tool_names
    assert "user_tool" in admin_tool_names
    assert len(admin_tool_names) == 3

    # Regular user should only see public and user tools
    user_schemas = await registry.get_schemas(regular_user)
    user_tool_names = [s.name for s in user_schemas]

    print(f"✓ Regular user schemas: {user_tool_names}")
    assert "public_tool" in user_tool_names
    assert "user_tool" in user_tool_names
    assert "admin_tool" not in user_tool_names
    assert len(user_tool_names) == 2


@pytest.mark.asyncio
async def test_tool_not_found(agent_memory):
    """Test execution of non-existent tool."""
    print("\n=== Tool Not Found Test ===")

    registry = ToolRegistry()

    user = User(id="user", username="user", group_memberships=["user"])
    context = ToolContext(
        user=user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1", name="nonexistent_tool", arguments={"message": "test"}
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Non-existent tool handled gracefully")
    print(f"  Error: {result.error}")

    assert result.success is False
    assert "not found" in result.result_for_llm.lower()


@pytest.mark.asyncio
async def test_duplicate_tool_registration():
    """Test that registering the same tool twice raises error."""
    print("\n=== Duplicate Registration Test ===")

    registry = ToolRegistry()
    tool = MockTool("duplicate_tool")

    # First registration should succeed
    registry.register_local_tool(tool, access_groups=[])
    print(f"✓ First registration succeeded")

    # Second registration should fail
    with pytest.raises(ValueError, match="already registered"):
        registry.register_local_tool(tool, access_groups=[])

    print(f"✓ Duplicate registration properly rejected")


@pytest.mark.asyncio
async def test_tool_access_group_intersection(admin_user, agent_memory):
    """Test that access is granted on ANY matching group (not all groups)."""
    print("\n=== Group Intersection Test ===")

    registry = ToolRegistry()
    tool = MockTool("multi_group_tool")

    # Tool requires either admin OR analyst
    registry.register_local_tool(tool, access_groups=["admin", "analyst"])

    # User has admin (but not analyst) - should still have access
    context = ToolContext(
        user=admin_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1", name="multi_group_tool", arguments={"message": "test"}
    )
    result = await registry.execute(tool_call, context)

    print(f"✓ User with ONE matching group granted access")
    print(f"  User groups: {admin_user.group_memberships}")
    print(f"  Tool requires: ['admin', 'analyst']")
    print(f"  Intersection: {set(admin_user.group_memberships) & {'admin', 'analyst'}}")

    assert result.success is True


@pytest.mark.asyncio
async def test_list_tools():
    """Test listing all registered tools."""
    print("\n=== List Tools Test ===")

    registry = ToolRegistry()

    registry.register_local_tool(MockTool("tool1"), access_groups=[])
    registry.register_local_tool(MockTool("tool2"), access_groups=["admin"])
    registry.register_local_tool(MockTool("tool3"), access_groups=["user"])

    tools = await registry.list_tools()

    print(f"✓ Listed {len(tools)} tools")
    print(f"  Tools: {tools}")

    assert len(tools) == 3
    assert "tool1" in tools
    assert "tool2" in tools
    assert "tool3" in tools


# ============================================================================
# transform_args Tests
# ============================================================================


@pytest.mark.asyncio
async def test_transform_args_default_no_transformation(regular_user, agent_memory):
    """Test that default transform_args implementation returns args unchanged."""
    print("\n=== Default transform_args (NoOp) Test ===")

    registry = ToolRegistry()
    tool = MockTool("test_tool")
    registry.register_local_tool(tool, access_groups=[])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    # Test the transform_args method directly
    original_args = SimpleToolArgs(message="original message")
    transformed_args = await registry.transform_args(
        tool=tool,
        args=original_args,
        user=regular_user,
        context=context,
    )

    print(f"✓ Default transform_args returns args unchanged")
    print(f"  Original: {original_args.message}")
    print(f"  Transformed: {transformed_args.message}")

    assert transformed_args == original_args
    assert transformed_args.message == "original message"
    assert not isinstance(transformed_args, ToolRejection)


class CustomTransformRegistry(ToolRegistry):
    """Custom registry that modifies arguments."""

    async def transform_args(
        self,
        tool: Tool[T],
        args: T,
        user: User,
        context: ToolContext,
    ) -> Union[T, ToolRejection]:
        """Custom transform that appends user info to message."""
        if isinstance(args, SimpleToolArgs):
            # Modify the args by appending user info
            modified_args = SimpleToolArgs(
                message=f"{args.message} [user: {user.username}]"
            )
            return modified_args
        return args


@pytest.mark.asyncio
async def test_transform_args_custom_modification(regular_user, agent_memory):
    """Test custom transform_args that modifies arguments."""
    print("\n=== Custom transform_args Modification Test ===")

    registry = CustomTransformRegistry()
    tool = MockTool("test_tool")
    registry.register_local_tool(tool, access_groups=[])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    # Execute tool - args should be transformed
    tool_call = ToolCall(
        id="call_1",
        name="test_tool",
        arguments={"message": "hello"},
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Tool executed with transformed args")
    print(f"  Result: {result.result_for_llm}")

    assert result.success is True
    assert "hello [user: user]" in result.result_for_llm
    assert "Mock tool executed" in result.result_for_llm


class RejectionRegistry(ToolRegistry):
    """Custom registry that rejects certain arguments."""

    async def transform_args(
        self,
        tool: Tool[T],
        args: T,
        user: User,
        context: ToolContext,
    ) -> Union[T, ToolRejection]:
        """Reject args containing 'forbidden' keyword."""
        if isinstance(args, SimpleToolArgs):
            if "forbidden" in args.message.lower():
                return ToolRejection(
                    reason=f"Message contains forbidden keyword. User '{user.username}' is not allowed to use this word."
                )
        return args


@pytest.mark.asyncio
async def test_transform_args_rejection(regular_user, agent_memory):
    """Test transform_args returning ToolRejection."""
    print("\n=== transform_args Rejection Test ===")

    registry = RejectionRegistry()
    tool = MockTool("test_tool")
    registry.register_local_tool(tool, access_groups=[])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    # Execute tool with forbidden word
    tool_call = ToolCall(
        id="call_1",
        name="test_tool",
        arguments={"message": "this is forbidden"},
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Tool execution rejected by transform_args")
    print(f"  Error: {result.error}")

    assert result.success is False
    assert "forbidden keyword" in result.result_for_llm
    assert "not allowed" in result.result_for_llm
    assert "user" in result.result_for_llm


@pytest.mark.asyncio
async def test_transform_args_allows_approved_content(regular_user, agent_memory):
    """Test that transform_args allows approved content."""
    print("\n=== transform_args Allows Approved Content Test ===")

    registry = RejectionRegistry()
    tool = MockTool("test_tool")
    registry.register_local_tool(tool, access_groups=[])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    # Execute tool with approved message
    tool_call = ToolCall(
        id="call_1",
        name="test_tool",
        arguments={"message": "this is allowed"},
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ Tool execution allowed for approved content")
    print(f"  Result: {result.result_for_llm}")

    assert result.success is True
    assert "Mock tool executed" in result.result_for_llm
    assert "this is allowed" in result.result_for_llm


class RowLevelSecurityRegistry(ToolRegistry):
    """Custom registry that applies row-level security transformations."""

    async def transform_args(
        self,
        tool: Tool[T],
        args: T,
        user: User,
        context: ToolContext,
    ) -> Union[T, ToolRejection]:
        """Apply RLS by modifying SQL queries based on user groups."""
        if isinstance(args, SimpleToolArgs):
            # Simulate SQL query transformation for RLS
            if "SELECT" in args.message.upper():
                # Add WHERE clause based on user groups
                if "admin" in user.group_memberships:
                    # Admins see everything - no modification needed
                    return args
                elif "analyst" in user.group_memberships:
                    # Analysts see filtered data
                    modified_message = args.message + " WHERE department='analytics'"
                    return SimpleToolArgs(message=modified_message)
                else:
                    # Regular users see even more restricted data
                    modified_message = args.message + f" WHERE user_id='{user.id}'"
                    return SimpleToolArgs(message=modified_message)
        return args


@pytest.mark.asyncio
async def test_transform_args_row_level_security(
    admin_user, analyst_user, regular_user, agent_memory
):
    """Test transform_args implementing row-level security."""
    print("\n=== Row-Level Security Test ===")

    registry = RowLevelSecurityRegistry()
    tool = MockTool("sql_tool")
    registry.register_local_tool(tool, access_groups=[])

    base_query = "SELECT * FROM users"

    # Test admin user - no modification
    admin_context = ToolContext(
        user=admin_user,
        conversation_id="test_conv",
        request_id="test_req_1",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1",
        name="sql_tool",
        arguments={"message": base_query},
    )

    result = await registry.execute(tool_call, admin_context)
    print(f"✓ Admin query: {result.result_for_llm}")
    assert result.success is True
    assert "WHERE" not in result.result_for_llm

    # Test analyst user - department filter
    analyst_context = ToolContext(
        user=analyst_user,
        conversation_id="test_conv",
        request_id="test_req_2",
        agent_memory=agent_memory,
    )

    result = await registry.execute(tool_call, analyst_context)
    print(f"✓ Analyst query: {result.result_for_llm}")
    assert result.success is True
    assert "WHERE department='analytics'" in result.result_for_llm

    # Test regular user - user_id filter
    regular_context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req_3",
        agent_memory=agent_memory,
    )

    result = await registry.execute(tool_call, regular_context)
    print(f"✓ Regular user query: {result.result_for_llm}")
    assert result.success is True
    assert f"WHERE user_id='{regular_user.id}'" in result.result_for_llm


@pytest.mark.asyncio
async def test_transform_args_called_during_execution(regular_user, agent_memory):
    """Test that transform_args is called during tool execution flow."""
    print("\n=== transform_args Called During Execution Test ===")

    class InstrumentedRegistry(ToolRegistry):
        def __init__(self):
            super().__init__()
            self.transform_called = False
            self.transform_call_count = 0

        async def transform_args(self, tool, args, user, context):
            self.transform_called = True
            self.transform_call_count += 1
            return await super().transform_args(tool, args, user, context)

    registry = InstrumentedRegistry()
    tool = MockTool("test_tool")
    registry.register_local_tool(tool, access_groups=[])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv",
        request_id="test_req",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1",
        name="test_tool",
        arguments={"message": "test"},
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ transform_args was called during execution")
    print(f"  Called: {registry.transform_called}")
    print(f"  Call count: {registry.transform_call_count}")

    assert result.success is True
    assert registry.transform_called is True
    assert registry.transform_call_count == 1


@pytest.mark.asyncio
async def test_transform_args_receives_correct_parameters(regular_user, agent_memory):
    """Test that transform_args receives correct parameters."""
    print("\n=== transform_args Parameter Validation Test ===")

    class ParameterCheckRegistry(ToolRegistry):
        def __init__(self):
            super().__init__()
            self.received_tool = None
            self.received_args = None
            self.received_user = None
            self.received_context = None

        async def transform_args(self, tool, args, user, context):
            self.received_tool = tool
            self.received_args = args
            self.received_user = user
            self.received_context = context
            return await super().transform_args(tool, args, user, context)

    registry = ParameterCheckRegistry()
    tool = MockTool("test_tool")
    registry.register_local_tool(tool, access_groups=[])

    context = ToolContext(
        user=regular_user,
        conversation_id="test_conv_123",
        request_id="test_req_456",
        agent_memory=agent_memory,
    )

    tool_call = ToolCall(
        id="call_1",
        name="test_tool",
        arguments={"message": "test message"},
    )

    result = await registry.execute(tool_call, context)

    print(f"✓ transform_args received correct parameters")
    print(f"  Tool name: {registry.received_tool.name}")
    print(f"  Args message: {registry.received_args.message}")
    print(f"  User: {registry.received_user.username}")
    print(f"  Context conv_id: {registry.received_context.conversation_id}")

    assert result.success is True
    assert registry.received_tool is not None
    assert registry.received_tool.name == "test_tool"
    assert registry.received_args is not None
    assert registry.received_args.message == "test message"
    assert registry.received_user is not None
    assert registry.received_user.id == regular_user.id
    assert registry.received_context is not None
    assert registry.received_context.conversation_id == "test_conv_123"


@pytest.mark.asyncio
async def test_transform_args_called_during_agent_send_message():
    """Test that transform_args is called during Agent.send_message workflow."""
    print("\n=== transform_args in Agent.send_message Test ===")

    # Import necessary components
    from vanna import Agent
    from vanna.core.user.resolver import UserResolver
    from vanna.core.user.request_context import RequestContext
    from vanna.core.llm import LlmService, LlmRequest, LlmResponse, LlmMessage
    from vanna.core.agent.config import AgentConfig

    # Create a custom registry that tracks transform_args calls
    class InstrumentedAgentRegistry(ToolRegistry):
        def __init__(self):
            super().__init__()
            self.transform_calls = []

        async def transform_args(self, tool, args, user, context):
            # Track that transform_args was called
            self.transform_calls.append(
                {
                    "tool_name": tool.name,
                    "args": args,
                    "user_id": user.id,
                    "conversation_id": context.conversation_id,
                }
            )
            # Apply a transformation to verify it's used
            if isinstance(args, SimpleToolArgs):
                modified_args = SimpleToolArgs(message=f"{args.message} [transformed]")
                return modified_args
            return args

    # Create a simple user resolver
    class TestUserResolver(UserResolver):
        async def resolve_user(self, request_context: RequestContext) -> User:
            return User(
                id="test_user_123",
                username="test_user",
                email="test@example.com",
                group_memberships=["user"],
            )

    # Create a mock LLM service that calls a tool
    class MockLlmService(LlmService):
        async def send_request(self, request: LlmRequest) -> LlmResponse:
            # Return a response that calls our mock tool
            return LlmResponse(
                content="I'll use the mock tool",
                tool_calls=[
                    ToolCall(
                        id="call_123",
                        name="mock_tool",
                        arguments={"message": "test message"},
                    )
                ],
                raw_response=None,
            )

        async def stream_request(self, request: LlmRequest):
            # Yield a single chunk with tool call
            from vanna.core.llm import LlmStreamChunk

            yield LlmStreamChunk(
                content="I'll use the mock tool",
                tool_calls=[
                    ToolCall(
                        id="call_123",
                        name="mock_tool",
                        arguments={"message": "test message"},
                    )
                ],
                raw_chunk=None,
            )

        async def validate_tools(self, tools: list[ToolSchema]) -> None:
            # No validation needed for this test
            pass

    # Set up the agent
    instrumented_registry = InstrumentedAgentRegistry()
    tool = MockTool("mock_tool")
    instrumented_registry.register_local_tool(tool, access_groups=[])

    agent_memory = DemoAgentMemory(max_items=100)

    agent = Agent(
        llm_service=MockLlmService(),
        tool_registry=instrumented_registry,
        user_resolver=TestUserResolver(),
        agent_memory=agent_memory,
        config=AgentConfig(),
    )

    # Send a message through the agent
    request_context = RequestContext(cookies={}, headers={})
    components = []

    async for component in agent.send_message(request_context, "test message"):
        components.append(component)

    print(f"✓ Agent.send_message completed")
    print(f"  Transform calls: {len(instrumented_registry.transform_calls)}")
    print(f"  Components received: {len(components)}")

    # Verify that transform_args was called
    assert len(instrumented_registry.transform_calls) > 0, (
        "transform_args should be called during Agent.send_message"
    )

    # Check the first transform call
    first_call = instrumented_registry.transform_calls[0]
    assert first_call["tool_name"] == "mock_tool"
    assert first_call["user_id"] == "test_user_123"
    assert first_call["conversation_id"] is not None

    print(f"  ✓ transform_args was called with correct parameters")
    print(f"    Tool: {first_call['tool_name']}")
    print(f"    User: {first_call['user_id']}")
    print(
        f"    Total transform_args calls: {len(instrumented_registry.transform_calls)}"
    )

    # Verify that the args passed to transform_args have the expected message
    assert first_call["args"].message == "test message"
    print(f"    Original args message: {first_call['args'].message}")

    print(
        f"  ✓ All checks passed - transform_args is called during Agent.send_message workflow"
    )


================================================
FILE: tests/test_workflow.py
================================================
"""
Tests for the default workflow handler, including memory display and deletion.
"""

import pytest
import uuid
from datetime import datetime
from vanna.core.workflow.default import DefaultWorkflowHandler
from vanna.core.user import User
from vanna.core.user.resolver import UserResolver
from vanna.core.user.request_context import RequestContext
from vanna.core.storage import Conversation
from vanna.core.tool import ToolContext
from vanna.capabilities.agent_memory import ToolMemory, TextMemory
from vanna.integrations.local.agent_memory import DemoAgentMemory
from vanna.core.rich_component import ComponentType


class SimpleUserResolver(UserResolver):
    """Simple user resolver for tests."""

    async def resolve_user(self, request_context: RequestContext) -> User:
        return User(
            id="test_user", email="test@example.com", group_memberships=["user"]
        )


class MockAgent:
    """Mock agent for testing workflow handlers."""

    def __init__(self, agent_memory=None, tool_registry=None):
        self.agent_memory = agent_memory
        self.tool_registry = tool_registry or MockToolRegistry()


class MockToolRegistry:
    """Mock tool registry for testing."""

    async def get_schemas(self, user):
        """Return mock tool schemas."""
        return []


class MockConversation:
    """Mock conversation for testing."""

    def __init__(self, conversation_id=None):
        self.id = conversation_id or str(uuid.uuid4())


@pytest.fixture
def test_user():
    """Create a test user."""
    return User(id="test_user", email="test@example.com", group_memberships=["user"])


@pytest.fixture
def admin_test_user():
    """Create an admin test user for tests that require admin access."""
    return User(id="admin_user", email="admin@example.com", group_memberships=["admin"])


@pytest.fixture
def test_conversation():
    """Create a test conversation."""
    return MockConversation()


@pytest.fixture
def workflow_handler():
    """Create a workflow handler instance."""
    return DefaultWorkflowHandler()


@pytest.fixture
def agent_with_memory():
    """Create a mock agent with memory."""
    memory = DemoAgentMemory(max_items=100)
    return MockAgent(agent_memory=memory)


@pytest.fixture
def agent_without_memory():
    """Create a mock agent without memory."""
    return MockAgent(agent_memory=None)


class TestWorkflowCommands:
    """Test basic workflow command handling."""

    @pytest.mark.asyncio
    async def test_help_command(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test that /help command returns help message (non-admin view)."""
        result = await workflow_handler.try_handle(
            agent_with_memory, test_user, test_conversation, "/help"
        )

        assert result.should_skip_llm is True
        assert len(result.components) > 0

        # Check that help content includes basic commands
        help_component = result.components[0]
        content = help_component.rich_component.content
        assert "/help" in content
        # Admin commands should NOT be visible to non-admin users
        assert "Admin Commands" not in content
        assert "/status" not in content
        assert "/memories" not in content
        assert "/delete" not in content

    @pytest.mark.asyncio
    async def test_status_command(
        self,
        workflow_handler,
        agent_with_memory,
        admin_test_user,
        test_conversation,
    ):
        """Test that /status command returns status information (admin only)."""
        result = await workflow_handler.try_handle(
            agent_with_memory, admin_test_user, test_conversation, "/status"
        )

        assert result.should_skip_llm is True
        assert len(result.components) > 0

    @pytest.mark.asyncio
    async def test_memories_command(
        self,
        workflow_handler,
        agent_with_memory,
        admin_test_user,
        test_conversation,
    ):
        """Test that /memories command returns memory list (admin only)."""
        result = await workflow_handler.try_handle(
            agent_with_memory, admin_test_user, test_conversation, "/memories"
        )

        assert result.should_skip_llm is True
        assert len(result.components) > 0

    @pytest.mark.asyncio
    async def test_unknown_command(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test that unknown commands are passed to LLM."""
        result = await workflow_handler.try_handle(
            agent_with_memory, test_user, test_conversation, "What is the weather?"
        )

        assert result.should_skip_llm is False


class TestMemoriesView:
    """Test the memories view functionality."""

    @pytest.mark.asyncio
    async def test_memories_no_agent_memory(
        self, workflow_handler, agent_without_memory, test_user, test_conversation
    ):
        """Test memories view when agent has no memory capability."""
        result = await workflow_handler._get_recent_memories(
            agent_without_memory, test_user, test_conversation
        )

        assert result.should_skip_llm is True
        assert len(result.components) == 1

        content = result.components[0].rich_component.content
        assert "No Memory System" in content

    @pytest.mark.asyncio
    async def test_memories_empty(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test memories view when no memories exist."""
        result = await workflow_handler._get_recent_memories(
            agent_with_memory, test_user, test_conversation
        )

        assert result.should_skip_llm is True
        assert len(result.components) == 1

        content = result.components[0].rich_component.content
        assert "No recent memories found" in content

    @pytest.mark.asyncio
    async def test_memories_with_tool_memories(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test memories view displays tool memories correctly."""
        # Create a context and add some tool memories
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        # Add tool memories
        await agent_with_memory.agent_memory.save_tool_usage(
            question="What is the total sales?",
            tool_name="run_sql",
            args={"query": "SELECT SUM(total) FROM invoices"},
            context=context,
            success=True,
        )

        await agent_with_memory.agent_memory.save_tool_usage(
            question="Show me customer names",
            tool_name="run_sql",
            args={"query": "SELECT name FROM customers"},
            context=context,
            success=True,
        )

        # Get memories view
        result = await workflow_handler._get_recent_memories(
            agent_with_memory, test_user, test_conversation
        )

        assert result.should_skip_llm is True
        assert len(result.components) > 1  # Header + at least one memory card

        # Check header
        header = result.components[0].rich_component.content
        assert "Recent Memories" in header
        assert "2" in header  # Should show count

        # Check that we have tool memory section
        found_tool_section = False
        for component in result.components:
            if hasattr(component.rich_component, "content"):
                if "Tool Memories" in component.rich_component.content:
                    found_tool_section = True
                    break

        assert found_tool_section

    @pytest.mark.asyncio
    async def test_memories_with_text_memories(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test memories view displays text memories correctly."""
        # Create a context and add some text memories
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        # Add text memories
        await agent_with_memory.agent_memory.save_text_memory(
            content="Remember to always check the data quality",
            context=context,
        )

        await agent_with_memory.agent_memory.save_text_memory(
            content="The fiscal year starts in April",
            context=context,
        )

        # Get memories view
        result = await workflow_handler._get_recent_memories(
            agent_with_memory, test_user, test_conversation
        )

        assert result.should_skip_llm is True
        assert len(result.components) > 1  # Header + at least one memory card

        # Check that we have text memory section
        found_text_section = False
        for component in result.components:
            if hasattr(component.rich_component, "content"):
                if "Text Memories" in component.rich_component.content:
                    found_text_section = True
                    break

        assert found_text_section

    @pytest.mark.asyncio
    async def test_memories_with_both_types(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test memories view displays both tool and text memories."""
        # Create a context
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        # Add tool memory
        await agent_with_memory.agent_memory.save_tool_usage(
            question="What is the total sales?",
            tool_name="run_sql",
            args={"query": "SELECT SUM(total) FROM invoices"},
            context=context,
            success=True,
        )

        # Add text memory
        await agent_with_memory.agent_memory.save_text_memory(
            content="Important note about the data",
            context=context,
        )

        # Get memories view
        result = await workflow_handler._get_recent_memories(
            agent_with_memory, test_user, test_conversation
        )

        assert result.should_skip_llm is True

        # Should have header + text section + text cards + tool section + tool cards
        assert len(result.components) >= 4

        # Check both sections exist
        content_str = str(
            [
                c.rich_component.content
                if hasattr(c.rich_component, "content")
                else str(c.rich_component)
                for c in result.components
            ]
        )
        assert "Text Memories" in content_str
        assert "Tool Memories" in content_str

    @pytest.mark.asyncio
    async def test_memories_have_delete_buttons(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test that memory cards include delete buttons."""
        # Create a context and add a memory
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        await agent_with_memory.agent_memory.save_tool_usage(
            question="Test question",
            tool_name="test_tool",
            args={"param": "value"},
            context=context,
            success=True,
        )

        # Get memories view
        result = await workflow_handler._get_recent_memories(
            agent_with_memory, test_user, test_conversation
        )

        # Find a card component with actions
        found_delete_button = False
        for component in result.components:
            if (
                hasattr(component.rich_component, "actions")
                and component.rich_component.actions
            ):
                for action in component.rich_component.actions:
                    if "/delete" in action.get("action", ""):
                        found_delete_button = True
                        break

        assert found_delete_button, "Should have delete button in memory cards"


class TestMemoryDeletion:
    """Test memory deletion functionality."""

    @pytest.mark.asyncio
    async def test_delete_no_agent_memory(
        self, workflow_handler, agent_without_memory, test_user, test_conversation
    ):
        """Test delete command when agent has no memory capability."""
        result = await workflow_handler._delete_memory(
            agent_without_memory, test_user, test_conversation, "test-id"
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "No Memory System" in content

    @pytest.mark.asyncio
    async def test_delete_no_id_provided(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test delete command without memory ID."""
        result = await workflow_handler._delete_memory(
            agent_with_memory, test_user, test_conversation, ""
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Invalid Command" in content
        assert "memory_id" in content

    @pytest.mark.asyncio
    async def test_delete_nonexistent_memory(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test deleting a memory that doesn't exist."""
        result = await workflow_handler._delete_memory(
            agent_with_memory, test_user, test_conversation, "nonexistent-id"
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Memory Not Found" in content

    @pytest.mark.asyncio
    async def test_delete_tool_memory_success(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test successfully deleting a tool memory."""
        # Create a context and add a memory
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        await agent_with_memory.agent_memory.save_tool_usage(
            question="Test question",
            tool_name="test_tool",
            args={"param": "value"},
            context=context,
            success=True,
        )

        # Get the memory ID
        memories = await agent_with_memory.agent_memory.get_recent_memories(
            context, limit=1
        )
        assert len(memories) == 1
        memory_id = memories[0].memory_id

        # Delete the memory
        result = await workflow_handler._delete_memory(
            agent_with_memory, test_user, test_conversation, memory_id
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Memory Deleted" in content
        assert memory_id in content

        # Verify memory is gone
        memories_after = await agent_with_memory.agent_memory.get_recent_memories(
            context, limit=10
        )
        assert len(memories_after) == 0

    @pytest.mark.asyncio
    async def test_delete_text_memory_success(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test successfully deleting a text memory."""
        # Create a context and add a text memory
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        text_memory = await agent_with_memory.agent_memory.save_text_memory(
            content="Test text memory",
            context=context,
        )

        memory_id = text_memory.memory_id

        # Delete the memory
        result = await workflow_handler._delete_memory(
            agent_with_memory, test_user, test_conversation, memory_id
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Memory Deleted" in content

        # Verify memory is gone
        text_memories_after = (
            await agent_with_memory.agent_memory.get_recent_text_memories(
                context, limit=10
            )
        )
        assert len(text_memories_after) == 0

    @pytest.mark.asyncio
    async def test_delete_command_parsing(
        self,
        workflow_handler,
        agent_with_memory,
        admin_test_user,
        test_conversation,
    ):
        """Test that /delete command is properly parsed (admin only)."""
        # Add a memory first
        context = ToolContext(
            user=admin_test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        await agent_with_memory.agent_memory.save_tool_usage(
            question="Test",
            tool_name="test",
            args={},
            context=context,
            success=True,
        )

        memories = await agent_with_memory.agent_memory.get_recent_memories(
            context, limit=1
        )
        memory_id = memories[0].memory_id

        # Test command parsing
        result = await workflow_handler.try_handle(
            agent_with_memory,
            admin_test_user,
            test_conversation,
            f"/delete {memory_id}",
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Memory Deleted" in content or "Memory Not Found" in content


class TestWorkflowComponentStructure:
    """Test the structure of components returned by workflow."""

    @pytest.mark.asyncio
    async def test_help_has_rich_component(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test that help command returns properly structured components."""
        result = await workflow_handler.try_handle(
            agent_with_memory, test_user, test_conversation, "/help"
        )

        assert len(result.components) > 0
        for component in result.components:
            assert hasattr(component, "rich_component")
            assert hasattr(component, "simple_component")

    @pytest.mark.asyncio
    async def test_memories_cards_have_proper_structure(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test that memory cards have proper structure."""
        # Add a memory
        context = ToolContext(
            user=test_user,
            conversation_id=test_conversation.id,
            request_id=str(uuid.uuid4()),
            agent_memory=agent_with_memory.agent_memory,
        )

        await agent_with_memory.agent_memory.save_tool_usage(
            question="Test",
            tool_name="test_tool",
            args={"key": "value"},
            context=context,
            success=True,
        )

        result = await workflow_handler._get_recent_memories(
            agent_with_memory, test_user, test_conversation
        )

        # Find card components
        card_found = False
        for component in result.components:
            if hasattr(component.rich_component, "type"):
                from vanna.core.rich_component import ComponentType

                if component.rich_component.type == ComponentType.CARD:
                    card_found = True
                    # Check card has required fields
                    assert hasattr(component.rich_component, "title")
                    assert hasattr(component.rich_component, "content")
                    assert hasattr(component.rich_component, "actions")
                    assert len(component.rich_component.actions) > 0

        assert card_found, "Should have at least one card component"


class TestStarterUI:
    """Test the starter UI functionality."""

    @pytest.mark.asyncio
    async def test_starter_ui_single_component(
        self, workflow_handler, agent_with_memory, test_user, test_conversation
    ):
        """Test that starter UI returns a single component."""

        # Mock tool registry with SQL tool
        class MockToolRegistryWithSQL:
            async def get_schemas(self, user):
                from vanna.core.tool import ToolSchema

                return [
                    ToolSchema(name="run_sql", description="Run SQL", parameters={})
                ]

        agent_with_memory.tool_registry = MockToolRegistryWithSQL()

        components = await workflow_handler.get_starter_ui(
            agent_with_memory, test_user, test_conversation
        )

        assert components is not None
        assert len(components) == 1, "Starter UI should return exactly one component"
        assert hasattr(components[0], "rich_component")

    @pytest.mark.asyncio
    async def test_starter_ui_user_view(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that non-admin users get a simple welcome message via RichTextComponent."""
        # Create user without admin role
        user = User(
            id="test_user", email="test@example.com", group_memberships=["user"]
        )

        # Mock tool registry with SQL tool
        class MockToolRegistryWithSQL:
            async def get_schemas(self, user):
                from vanna.core.tool import ToolSchema

                return [
                    ToolSchema(name="run_sql", description="Run SQL", parameters={})
                ]

        agent_with_memory.tool_registry = MockToolRegistryWithSQL()

        components = await workflow_handler.get_starter_ui(
            agent_with_memory, user, test_conversation
        )

        assert len(components) == 1
        rich_text = components[0].rich_component
        assert rich_text.type == ComponentType.TEXT
        assert "Welcome to Vanna AI" in rich_text.content
        # User view should be simple and not mention memory management
        assert "Memory Management" not in rich_text.content
        assert "Admin" not in rich_text.content

    @pytest.mark.asyncio
    async def test_starter_ui_admin_view(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that admin users get setup status and memory management info."""
        # Create admin user
        admin_user = User(
            id="admin_user", email="admin@example.com", group_memberships=["admin"]
        )

        # Mock tool registry with SQL and memory tools
        class MockToolRegistryComplete:
            async def get_schemas(self, user):
                from vanna.core.tool import ToolSchema

                return [
                    ToolSchema(name="run_sql", description="Run SQL", parameters={}),
                    ToolSchema(
                        name="search_saved_correct_tool_uses",
                        description="Search",
                        parameters={},
                    ),
                    ToolSchema(
                        name="save_question_tool_args",
                        description="Save",
                        parameters={},
                    ),
                ]

        agent_with_memory.tool_registry = MockToolRegistryComplete()

        components = await workflow_handler.get_starter_ui(
            agent_with_memory, admin_user, test_conversation
        )

        assert len(components) == 1
        card = components[0].rich_component
        assert card.type == ComponentType.CARD
        assert hasattr(card, "title")
        assert "Admin" in card.title  # Should indicate admin status
        assert "🔒 Admin View" in card.content  # Should show admin badge
        assert "Setup:" in card.content
        # Admin view should include memory management info
        assert "Memory Management" in card.content
        assert "As an admin" in card.content
        # Should have View Memories button
        assert hasattr(card, "actions")
        assert any(
            "View Memories" in action.get("label", "") for action in card.actions
        )

    @pytest.mark.asyncio
    async def test_starter_ui_admin_without_memory(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that admin users see setup status even without memory tools."""
        # Create admin user
        admin_user = User(
            id="admin_user", email="admin@example.com", group_memberships=["admin"]
        )

        # Mock tool registry with only SQL
        class MockToolRegistrySQL:
            async def get_schemas(self, user):
                from vanna.core.tool import ToolSchema

                return [
                    ToolSchema(name="run_sql", description="Run SQL", parameters={}),
                ]

        agent_with_memory.tool_registry = MockToolRegistrySQL()

        components = await workflow_handler.get_starter_ui(
            agent_with_memory, admin_user, test_conversation
        )

        assert len(components) == 1
        card = components[0].rich_component
        assert card.type == ComponentType.CARD
        # Admin should see setup status showing incomplete memory setup
        assert "Admin" in card.title
        assert "🔒 Admin View" in card.content
        assert "Setup:" in card.content
        assert "Memory ✗" in card.content
        # Should NOT have View Memories button since memory is not available
        memory_button_exists = any(
            "View Memories" in action.get("label", "")
            for action in (card.actions or [])
        )
        assert not memory_button_exists


class TestAdminOnlyCommands:
    """Test that admin-only commands are properly restricted."""

    @pytest.mark.asyncio
    async def test_non_admin_cannot_access_status(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that non-admin users cannot access /status command."""
        # Create non-admin user
        user = User(id="user", email="user@example.com", group_memberships=["user"])

        result = await workflow_handler.try_handle(
            agent_with_memory, user, test_conversation, "/status"
        )

        assert result.should_skip_llm is True
        assert len(result.components) == 1
        rich_text = result.components[0].rich_component
        assert "Access Denied" in rich_text.content
        assert "administrators" in rich_text.content

    @pytest.mark.asyncio
    async def test_non_admin_cannot_access_memories(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that non-admin users cannot access /memories command."""
        # Create non-admin user
        user = User(id="user", email="user@example.com", group_memberships=["user"])

        result = await workflow_handler.try_handle(
            agent_with_memory, user, test_conversation, "/memories"
        )

        assert result.should_skip_llm is True
        assert len(result.components) == 1
        rich_text = result.components[0].rich_component
        assert "Access Denied" in rich_text.content
        assert "administrators" in rich_text.content

    @pytest.mark.asyncio
    async def test_non_admin_cannot_delete_memories(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that non-admin users cannot use /delete command."""
        # Create non-admin user
        user = User(id="user", email="user@example.com", group_memberships=["user"])

        result = await workflow_handler.try_handle(
            agent_with_memory, user, test_conversation, "/delete some-memory-id"
        )

        assert result.should_skip_llm is True
        assert len(result.components) == 1
        rich_text = result.components[0].rich_component
        assert "Access Denied" in rich_text.content
        assert "administrators" in rich_text.content

    @pytest.mark.asyncio
    async def test_admin_can_access_memories(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that admin users can access /memories command."""
        # Create admin user
        admin_user = User(
            id="admin", email="admin@example.com", group_memberships=["admin"]
        )

        result = await workflow_handler.try_handle(
            agent_with_memory, admin_user, test_conversation, "/memories"
        )

        assert result.should_skip_llm is True
        # Should show memories (even if empty, no access denied)
        assert len(result.components) >= 1
        assert "Access Denied" not in result.components[0].rich_component.content

    @pytest.mark.asyncio
    async def test_help_shows_admin_commands_for_admin(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that /help shows admin commands for admin users."""
        # Create admin user
        admin_user = User(
            id="admin", email="admin@example.com", group_memberships=["admin"]
        )

        result = await workflow_handler.try_handle(
            agent_with_memory, admin_user, test_conversation, "/help"
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Admin Commands" in content
        assert "/status" in content
        assert "/memories" in content
        assert "/delete" in content

    @pytest.mark.asyncio
    async def test_help_hides_admin_commands_for_non_admin(
        self, workflow_handler, agent_with_memory, test_conversation
    ):
        """Test that /help hides admin commands for non-admin users."""
        # Create non-admin user
        user = User(id="user", email="user@example.com", group_memberships=["user"])

        result = await workflow_handler.try_handle(
            agent_with_memory, user, test_conversation, "/help"
        )

        assert result.should_skip_llm is True
        content = result.components[0].rich_component.content
        assert "Admin Commands" not in content
        # Still shows regular commands
        assert "/help" in content
        # But not admin commands
        assert "/status" not in content


if __name__ == "__main__":
    pytest.main([__file__, "-v"])


================================================
FILE: tox.ini
================================================
[tox]
envlist =
    ruff
    mypy
    py311-unit
    py311-agent-memory-sanity
    py311-anthropic
    py311-openai
    py311-legacy
    py311-chromadb
    py311-qdrant
    py311-faiss
    py311-postgres-sanity
    py311-sqlite-sanity
    py311-snowflake-sanity
    py311-mysql-sanity
    py311-clickhouse-sanity
    py311-oracle-sanity
    py311-bigquery-sanity
    py311-duckdb-sanity
    py311-mssql-sanity
    py311-presto-sanity
    py311-hive-sanity

[testenv]
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0

[testenv:py311-unit]
description = Run unit tests (no external dependencies required)
commands =
    pytest tests/test_tool_permissions.py tests/test_llm_context_enhancer.py tests/test_workflow.py tests/test_memory_tools.py -v

[testenv:py311-agent-memory-sanity]
description = Run sanity tests for all AgentMemory implementations (no actual service connections required)
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras =
    chromadb
    qdrant
commands =
    pytest tests/test_agent_memory_sanity.py -v

[testenv:py311-anthropic]
description = Test with Anthropic
extras = anthropic
passenv = ANTHROPIC_API_KEY
commands =
    python -c "from vanna.integrations.anthropic import AnthropicLlmService; print('✓ Anthropic import successful')"
    pytest tests/ -v -m anthropic

[testenv:py311-openai]
description = Test with OpenAI
extras = openai
passenv = OPENAI_API_KEY
commands =
    python -c "from vanna.integrations.openai import OpenAILlmService; print('✓ OpenAI import successful')"
    pytest tests/ -v -m openai

; [testenv:py311-gemini]
; description = Test with Google Gemini
; extras = gemini
; passenv =
;     GOOGLE_API_KEY
;     GEMINI_API_KEY
; commands =
;     python -c "from vanna.integrations.google import GeminiLlmService; print('✓ Gemini import successful')"
;     pytest tests/ -v -m gemini

; [testenv:py311-ollama]
; description = Test with Ollama
; extras = ollama
; passenv = OLLAMA_HOST
; commands =
;     python -c "from vanna.integrations.ollama import OllamaLlmService; print('✓ Ollama import successful')"
;     pytest tests/ -v -m ollama

[testenv:py311-legacy]
description = Test LegacyVannaAdapter with Anthropic
extras =
    anthropic
    chromadb
passenv = ANTHROPIC_API_KEY
commands =
    python -c "from vanna.legacy.adapter import LegacyVannaAdapter; from vanna.legacy.chromadb import ChromaDB_VectorStore; from vanna.legacy.mock import MockLLM; print('✓ Legacy adapter imports successful')"
    pytest tests/test_legacy_adapter.py -v -m legacy

[testenv:py311-chromadb]
description = Test ChromaDB AgentMemory
extras = chromadb
commands =
    pytest tests/test_agent_memory.py::TestLocalAgentMemory -k chromadb -v

[testenv:py311-qdrant]
description = Test Qdrant AgentMemory
extras = qdrant
commands =
    pytest tests/test_agent_memory.py::TestLocalAgentMemory -k qdrant -v

[testenv:py311-faiss]
description = Test FAISS AgentMemory
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
    faiss-cpu
commands =
    pytest tests/test_agent_memory.py::TestLocalAgentMemory -k faiss -v

[testenv:py311-db-sanity]
description = Run sanity tests for all database implementations (no actual DB connections required)
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
    pytest-mock>=3.10.0
extras =
    postgres
commands =
    pytest tests/test_database_sanity.py -v

[testenv:py311-postgres-sanity]
description = Sanity tests for PostgreSQL implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = postgres
commands =
    python -c "from vanna.integrations.postgres import PostgresRunner; print('✓ PostgresRunner import successful')"
    pytest tests/test_database_sanity.py::TestPostgresRunner -v

[testenv:py311-sqlite-sanity]
description = Sanity tests for SQLite implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
commands =
    python -c "from vanna.integrations.sqlite import SqliteRunner; print('✓ SqliteRunner import successful')"
    pytest tests/test_database_sanity.py::TestSqliteRunner -v

[testenv:py311-snowflake-sanity]
description = Sanity tests for Snowflake implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = snowflake
commands =
    python -c "from vanna.integrations.snowflake import SnowflakeRunner; print('✓ SnowflakeRunner import successful')"
    pytest tests/test_database_sanity.py::TestSnowflakeRunner -v

[testenv:py311-mysql-sanity]
description = Sanity tests for MySQL implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = mysql
commands =
    python -c "from vanna.integrations.mysql import MySQLRunner; print('✓ MySQLRunner import successful')"
    pytest tests/test_database_sanity.py::TestMySQLRunner -v

[testenv:py311-clickhouse-sanity]
description = Sanity tests for ClickHouse implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = clickhouse
commands =
    python -c "from vanna.integrations.clickhouse import ClickHouseRunner; print('✓ ClickHouseRunner import successful')"
    pytest tests/test_database_sanity.py::TestClickHouseRunner -v

[testenv:py311-oracle-sanity]
description = Sanity tests for Oracle implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = oracle
commands =
    python -c "from vanna.integrations.oracle import OracleRunner; print('✓ OracleRunner import successful')"
    pytest tests/test_database_sanity.py::TestOracleRunner -v

[testenv:py311-bigquery-sanity]
description = Sanity tests for BigQuery implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = bigquery
commands =
    python -c "from vanna.integrations.bigquery import BigQueryRunner; print('✓ BigQueryRunner import successful')"
    pytest tests/test_database_sanity.py::TestBigQueryRunner -v

[testenv:py311-duckdb-sanity]
description = Sanity tests for DuckDB implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = duckdb
commands =
    python -c "from vanna.integrations.duckdb import DuckDBRunner; print('✓ DuckDBRunner import successful')"
    pytest tests/test_database_sanity.py::TestDuckDBRunner -v

[testenv:py311-mssql-sanity]
description = Sanity tests for MSSQL implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = mssql
commands =
    python -c "from vanna.integrations.mssql import MSSQLRunner; print('✓ MSSQLRunner import successful')"
    pytest tests/test_database_sanity.py::TestMSSQLRunner -v

[testenv:py311-presto-sanity]
description = Sanity tests for Presto implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = presto
commands =
    python -c "from vanna.integrations.presto import PrestoRunner; print('✓ PrestoRunner import successful')"
    pytest tests/test_database_sanity.py::TestPrestoRunner -v

[testenv:py311-hive-sanity]
description = Sanity tests for Hive implementation
deps =
    pytest>=7.0.0
    pytest-asyncio>=0.21.0
extras = hive
commands =
    python -c "from vanna.integrations.hive import HiveRunner; print('✓ HiveRunner import successful')"
    pytest tests/test_database_sanity.py::TestHiveRunner -v

[testenv:ruff]
description = Check code formatting and linting with ruff (uses pyproject.toml config)
extras = dev
commands =
    ruff format --check src/vanna/ tests/
    ruff check src/vanna/ tests/

[testenv:mypy]
description = Run mypy type checking with strict mode
extras = dev
commands =
    mypy src/vanna/tools src/vanna/core src/vanna/capabilities src/vanna/agents src/vanna/utils src/vanna/web_components src/vanna/components --strict