VTK Prompt

A command-line interface and web-based UI for generating VTK visualization code using Large Language Models (Anthropic Claude, OpenAI GPT, NVIDIA NIM, and local models).

Features
- Multiple LLM providers: Anthropic Claude, OpenAI GPT, NVIDIA NIM, and local models
- Interactive web UI with live VTK rendering
- Retrieval-Augmented Generation (RAG) with VTK examples database
- Real-time visualization of generated code
- Token usage tracking and cost monitoring
- CLI and Python API for integration
Installation
From PyPI (Stable)
pip install vtk-prompt
From TestPyPI (Latest Development)
pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ vtk-prompt
From Source
git clone https://github.com/vicentebolea/vtk-prompt.git
cd vtk-prompt
pip install -e .
Quick Start
1. Set up API keys
export ANTHROPIC_API_KEY="your-anthropic-key"
export OPENAI_API_KEY="your-openai-key"
2. Launch Web UI (Recommended)
vtk-prompt-ui
Access the UI at http://localhost:8080
3. Or use CLI
vtk-prompt "Create a red sphere"
vtk-prompt "Create a sphere with custom resolution" --rag
vtk-prompt "Create a blue cube" --provider openai
vtk-prompt "Create a cone" --provider nim --token YOUR_NIM_TOKEN
Usage
Web UI Features
The web interface provides:
- Model selection: Choose between Claude models (Haiku, Sonnet 4) and other providers
- Token control: Adjust maximum tokens for responses
- Usage tracking: Real-time display of input/output tokens and costs
- RAG integration: Toggle retrieval-augmented generation for better results
- Live preview: See VTK visualizations rendered in real-time
- Code export: View, edit, and copy generated VTK code
- Local & cloud support: Both cloud APIs and local model endpoints
Command Line Interface
vtk-prompt "Create a red sphere"
vtk-prompt "Create a textured cone with 32 resolution" \
--provider anthropic \
--model claude-3-5-sonnet-20241022 \
--max-tokens 4000 \
--rag \
--verbose
vtk-prompt "Create a blue cube" --provider openai --model gpt-4o
vtk-prompt "Create a cylinder" --provider nim --model meta/llama3-70b-instruct
RAG (Retrieval Augmented Generation)
Enhance code generation with relevant VTK examples:
- Build RAG database (one-time setup):
vtk-build-rag
- Test RAG system (optional):
vtk-test-rag "How to create a cube in VTK"
vtk-prompt "Create a vtkSphereSource with texture mapping" --rag
Python API
from vtk_prompt import VTKPromptClient
client = VTKPromptClient()
code = client.generate_code("Create a red sphere")
print(code)
Configuration
Environment Variables
ANTHROPIC_API_KEY
- Anthropic Claude API key
OPENAI_API_KEY
- OpenAI API key (also used for NVIDIA NIM)
Supported Providers & Models
Custom/Local Models
You can use local models via OpenAI-compatible APIs:
vtk-prompt "Create a sphere" \
--provider custom \
--base-url http://localhost:11434/v1 \
--model llama2
vtk-prompt "Create a cube" \
--provider custom \
--base-url http://localhost:1234/v1 \
--model local-model
Development
Setting up development environment
git clone https://github.com/vicentebolea/vtk-prompt.git
cd vtk-prompt
pip install -e ".[all]"
Running tests
black src/
flake8 src/
vtk-prompt --help
vtk-prompt-ui --help
Building package
python -m build
CLI Reference
usage: vtk-prompt [-h] [--provider {anthropic,openai,nim,custom}]
[-m MODEL] [-k MAX_TOKENS] [-t TOKEN] [--base-url BASE_URL]
[-r] [-v] [--collection COLLECTION] [--database DATABASE]
[--top-k TOP_K] input_string
Generate VTK visualization code using Large Language Models
positional arguments:
input_string Description of the VTK visualization to generate
options:
-h, --help Show this help message and exit
-m MODEL, --model MODEL
Model name to use
-k MAX_TOKENS, --max-tokens MAX_TOKENS
Maximum number of tokens to generate
-t TOKEN, --token TOKEN
API token (defaults to environment variable)
--base-url BASE_URL Base URL for API (for custom/local models)
-r, --rag Use Retrieval Augmented Generation
-v, --verbose Show generated source code
--provider {anthropic,openai,nim,custom}
LLM provider to use
RAG Options:
--collection COLLECTION
Collection name for RAG (default: vtk-examples)
--database DATABASE Database path for RAG (default: ./db/codesage-codesage-large-v2)
--top-k TOP_K Number of examples to retrieve (default: 5)
Available Commands
vtk-prompt
- Main CLI for code generation
vtk-prompt-ui
- Launch web interface
vtk-build-rag
- Build RAG database from VTK examples
vtk-test-rag
- Test RAG functionality
gen-vtk-file
- Generate VTK XML files
rag-chat
- Interactive RAG chat interface
Contributing
Contributions are welcome! Please feel free to submit a Pull Request.
- Fork the repository
- Create your feature branch (
git checkout -b feature/amazing-feature
)
- Commit your changes (
git commit -m 'Add some amazing feature'
)
- Push to the branch (
git push origin feature/amazing-feature
)
- Open a Pull Request
License
This project is licensed under the MIT License - see the LICENSE file for details.
Architecture
- Core: Python package with CLI and API
- UI: Trame-based web interface with VTK rendering
- RAG: ChromaDB + Llama Index for code example retrieval
- Providers: Unified interface for multiple LLM APIs
Links
Made with care for the VTK and scientific visualization community.