vtk-prompt 0.0.3__py3-none-any.whl

vtk_prompt/__init__.py

@@ -0,0 +1,25 @@
+"""VTK-Prompt - CLI tool for generating VTK visualizations using LLMs.
+
+This package provides tools to generate VTK Python code and XML files using
+LLMs (Anthropic Claude, OpenAI GPT, or NVIDIA NIM models). It also includes
+Retrieval-Augmented Generation (RAG) capabilities to improve code generation
+by providing relevant examples from the VTK examples corpus.
+
+Main components:
+- vtk-prompt: Generate and run VTK Python code
+- gen-vtk-file: Generate VTK XML files
+- vtk-build-rag: Build a RAG database from VTK examples
+- vtk-test-rag: Test the RAG database with queries
+"""
+
+try:
+    from importlib.metadata import version, PackageNotFoundError
+except ImportError:
+    from importlib_metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("vtk-prompt")
+except PackageNotFoundError:
+    __version__ = "unknown"
+__author__ = "Vicente Adolfo Bolea Sanchez"
+__email__ = "vicente.bolea@kitware.com"

vtk_prompt/build_rag_db.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+
+import argparse
+import importlib.util
+import sys
+from pathlib import Path
+
+
+def setup_rag_path():
+    """Add rag-components to the Python path.
+
+    Returns:
+        The path to the rag-components directory
+    """
+    script_dir = Path(__file__).resolve().parent
+    project_root = script_dir.parent.parent
+    rag_path = str(project_root / "rag-components")
+
+    if rag_path not in sys.path:
+        sys.path.append(rag_path)
+
+    return rag_path
+
+
+def check_dependencies():
+    """Check if required dependencies are installed.
+
+    Returns:
+        True if all dependencies are installed, False otherwise
+    """
+    required_modules = ["chromadb", "sentence_transformers", "tree_sitter_languages"]
+
+    for module in required_modules:
+        if importlib.util.find_spec(module) is None:
+            print(f"Missing required dependency: {module}")
+            return False
+
+    return True
+
+
+def main():
+    """Build a RAG database from VTK example files."""
+    parser = argparse.ArgumentParser(description="Build RAG database for VTK examples")
+    parser.add_argument(
+        "--examples-dir",
+        type=str,
+        default="data/examples",
+        help="Directory containing VTK examples (default: data/examples)",
+    )
+    parser.add_argument(
+        "--database",
+        type=str,
+        default="./db/codesage-codesage-large-v2",
+        help="Database path (default: ./db/codesage-codesage-large-v2)",
+    )
+    parser.add_argument(
+        "--collection-name",
+        type=str,
+        default="vtk-examples",
+        help="Collection name in the database (default: vtk-examples)",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="codesage/codesage-large-v2",
+        help="Embedding model name (default: codesage/codesage-large-v2)",
+    )
+    parser.add_argument(
+        "--language",
+        type=str,
+        default="python",
+        help="Language of the examples (default: python)",
+    )
+
+    args = parser.parse_args()
+
+    # Check dependencies
+    if not check_dependencies():
+        print("Please install the required dependencies with:")
+        print('pip install -e ".[rag]"')
+        sys.exit(1)
+
+    # Setup RAG path
+    rag_path = setup_rag_path()
+
+    # Import populate_db from rag-components
+    try:
+        sys.path.insert(0, rag_path)
+        from populate_db import fill_database
+    except ImportError as e:
+        print(f"Failed to import from rag-components: {e}")
+        print(
+            "Make sure the rag-components directory exists and contains the required files."
+        )
+        sys.exit(1)
+
+    # Check if examples directory exists
+    examples_dir = Path(args.examples_dir)
+    if not examples_dir.exists() or not examples_dir.is_dir():
+        print(
+            f"Error: Examples directory '{args.examples_dir}' does not exist or is not a directory"
+        )
+        sys.exit(1)
+
+    # Get all Python files in the examples directory
+    files = list(examples_dir.glob("**/*.py"))
+    if not files:
+        print(f"No Python files found in '{args.examples_dir}'")
+        sys.exit(1)
+
+    print(f"Found {len(files)} Python files in '{args.examples_dir}'")
+
+    # Create database directory if it doesn't exist
+    database_dir = Path(args.database).parent
+    database_dir.mkdir(parents=True, exist_ok=True)
+
+    # Build the RAG database
+    print(
+        f"Building RAG database at '{args.database}' using embedding model '{args.model}'..."
+    )
+    try:
+        fill_database(
+            files=files,
+            database_path=args.database,
+            embedding_model=args.model,
+            language=args.language,
+            collection_name=args.collection_name,
+        )
+
+        print(f"Successfully built RAG database at '{args.database}'")
+        print("You can now use the RAG database with vtk-prompt by running:")
+        print(
+            f'vtk-prompt "your query" -r --database {args.database} '
+            f"--collection {args.collection_name}"
+        )
+
+    except Exception as e:
+        print(f"Error building RAG database: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

vtk_prompt/generate_files.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+
+from pathlib import Path
+import os
+import json
+import sys
+import openai
+import click
+
+# Import our template system
+from .prompts import (
+    get_vtk_xml_context,
+    get_xml_role,
+)
+
+
+class VTKXMLGenerator:
+    """OpenAI client for VTK XML file generation."""
+
+    def __init__(self, api_key=None, base_url=None):
+        """Initialize the VTK XML generator.
+
+        Args:
+            api_key: OpenAI API key (defaults to OPENAI_API_KEY env var)
+            base_url: Optional custom API endpoint URL
+        """
+        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        self.base_url = base_url
+
+        if not self.api_key:
+            raise ValueError(
+                "No API key provided. Set OPENAI_API_KEY or pass api_key parameter."
+            )
+
+        self.client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    def generate_xml(self, message, model, max_tokens=4000, temperature=0.7):
+        """Generate VTK XML content from a description."""
+        examples_path = Path("data/examples/index.json")
+        if examples_path.exists():
+            _ = " ".join(json.loads(examples_path.read_text()).keys())
+        else:
+            _ = ""
+
+        context = get_vtk_xml_context(message)
+
+        response = self.client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": get_xml_role()},
+                {"role": "user", "content": context},
+            ],
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+
+        if hasattr(response, "choices") and len(response.choices) > 0:
+            content = response.choices[0].message.content or "No content in response"
+            finish_reason = response.choices[0].finish_reason
+
+            if finish_reason == "length":
+                raise ValueError(
+                    f"Output was truncated due to max_tokens limit ({max_tokens}). "
+                    "Please increase max_tokens."
+                )
+
+            return content
+
+        return "No response generated"
+
+
+# Legacy function wrapper for backwards compatibility
+def openai_query(message, model, api_key, max_tokens, temperature=0.7, base_url=None):
+    """Legacy wrapper for VTK XML generation."""
+    generator = VTKXMLGenerator(api_key, base_url)
+    return generator.generate_xml(message, model, max_tokens, temperature)
+
+
+@click.command()
+@click.argument("input_string")
+@click.option(
+    "--provider",
+    type=click.Choice(["openai", "anthropic", "gemini", "nim"]),
+    default="openai",
+    help="LLM provider to use",
+)
+@click.option("-m", "--model", default="gpt-4o", help="Model to use for generation")
+@click.option(
+    "-t", "--token", required=True, help="API token for the selected provider"
+)
+@click.option("--base-url", help="Base URL for API (auto-detected or custom)")
+@click.option(
+    "-k",
+    "--max-tokens",
+    type=int,
+    default=4000,
+    help="Maximum number of tokens to generate",
+)
+@click.option(
+    "--temperature",
+    type=float,
+    default=0.7,
+    help="Temperature for generation (0.0-2.0)",
+)
+@click.option(
+    "-o", "--output", help="Output file path (if not specified, output to stdout)"
+)
+def main(
+    input_string, provider, model, token, base_url, max_tokens, temperature, output
+):
+    """Generate VTK XML file content using LLMs.
+
+    INPUT_STRING: Description of the VTK file to generate
+    """
+
+    # Set default base URLs
+    if not base_url:
+        base_urls = {
+            "anthropic": "https://api.anthropic.com/v1",
+            "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/",
+            "nim": "https://integrate.api.nvidia.com/v1",
+        }
+        base_url = base_urls.get(provider)
+
+    # Set default models based on provider
+    if model == "gpt-4o":
+        default_models = {
+            "anthropic": "claude-3-5-sonnet-20241022",
+            "gemini": "gemini-1.5-pro",
+            "nim": "meta/llama3-70b-instruct",
+        }
+        model = default_models.get(provider, model)
+
+    # Initialize the VTK XML generator
+    try:
+        generator = VTKXMLGenerator(token, base_url)
+    except ValueError as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    # Generate the VTK XML content
+    try:
+        xml_content = generator.generate_xml(
+            input_string, model, max_tokens, temperature
+        )
+    except ValueError as e:
+        if "max_tokens" in str(e):
+            print(f"\nError: {e}", file=sys.stderr)
+            print(f"Current max_tokens: {max_tokens}", file=sys.stderr)
+            print("Try increasing with: --max-tokens <higher_number>", file=sys.stderr)
+        else:
+            print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    # Validate XML structure (basic check)
+    if xml_content.strip().startswith("<?xml") and "</VTKFile>" in xml_content:
+        # Output to file or stdout
+        if output:
+            with open(output, "w") as f:
+                f.write(xml_content)
+            print(f"VTK XML content written to {output}")
+        else:
+            print(xml_content)
+    else:
+        print("Warning: Generated content may not be valid VTK XML", file=sys.stderr)
+        if output:
+            with open(output, "w") as f:
+                f.write(xml_content)
+            print(f"Content written to {output} (please verify)")
+        else:
+            print(xml_content)
+
+
+if __name__ == "__main__":
+    main()

vtk_prompt/prompt.py

@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+
+import ast
+import os
+import sys
+import openai
+import click
+from dataclasses import dataclass
+
+from .prompts import (
+    get_no_rag_context,
+    get_rag_context,
+    get_python_role,
+)
+
+
+@dataclass
+class VTKPromptClient:
+    """OpenAI client for VTK code generation."""
+
+    collection_name: str = "vtk-examples"
+    database_path: str = "./db/codesage-codesage-large-v2"
+    verbose: bool = False
+
+    def validate_code_syntax(self, code_string):
+        """Validate Python code syntax using AST."""
+        try:
+            ast.parse(code_string)
+            return True, None
+        except SyntaxError as e:
+            return False, f"Syntax error: {e.msg} at line {e.lineno}"
+        except Exception as e:
+            return False, f"AST parsing error: {str(e)}"
+
+    def run_code(self, code_string):
+        """Execute VTK code using exec() after AST validation."""
+        is_valid, error_msg = self.validate_code_syntax(code_string)
+        if not is_valid:
+            print(f"Code validation failed: {error_msg}")
+            if self.verbose:
+                print("Generated code:")
+                print(code_string)
+            return None
+
+        if self.verbose:
+            print(code_string)
+
+        try:
+            exec(code_string, globals(), {})
+        except Exception as e:
+            print(f"Error executing code: {e}")
+            if not self.verbose:
+                print(code_string)
+            return None
+
+    def query(
+        self,
+        message,
+        api_key,
+        model="gpt-4o",
+        base_url=None,
+        max_tokens=1000,
+        temperature=0.1,
+        top_k=5,
+        rag=False,
+        retry_attempts=1,
+    ):
+        """Generate VTK code with optional RAG enhancement and retry logic.
+
+        Args:
+            message: The user query
+            api_key: API key for the service
+            model: Model name to use
+            base_url: API base URL
+            max_tokens: Maximum tokens to generate
+            temperature: Temperature for generation
+            top_k: Number of RAG examples to retrieve
+            rag: Whether to use RAG enhancement
+            retry_attempts: Number of times to retry if AST validation fails
+        """
+        if not api_key:
+            api_key = os.environ.get("OPENAI_API_KEY")
+
+        if not api_key:
+            raise ValueError(
+                "No API key provided. Set OPENAI_API_KEY or pass api_key parameter."
+            )
+
+        # Create client with current parameters
+        client = openai.OpenAI(api_key=api_key, base_url=base_url)
+
+        if rag:
+            from .rag_chat_wrapper import (
+                check_rag_components_available,
+                get_rag_snippets,
+            )
+
+            if not check_rag_components_available():
+                raise ValueError("RAG components not available")
+
+            rag_snippets = get_rag_snippets(
+                message,
+                collection_name=self.collection_name,
+                database_path=self.database_path,
+                top_k=top_k,
+            )
+
+            if not rag_snippets:
+                raise ValueError("Failed to load RAG snippets")
+
+            context_snippets = "\n\n".join(rag_snippets["code_snippets"])
+            context = get_rag_context(message, context_snippets)
+
+            if self.verbose:
+                print("CONTEXT: " + context)
+                references = rag_snippets.get("references")
+                if references:
+                    print("Using examples from:")
+                    for ref in references:
+                        print(f"- {ref}")
+        else:
+            context = get_no_rag_context(message)
+            if self.verbose:
+                print("CONTEXT: " + context)
+
+        messages = [
+            {"role": "system", "content": get_python_role()},
+            {"role": "user", "content": context},
+        ]
+
+        # Retry loop for AST validation
+        for attempt in range(retry_attempts):
+            if self.verbose:
+                if attempt > 0:
+                    print(f"Retry attempt {attempt + 1}/{retry_attempts}")
+                print(f"Making request with model: {model}, temperature: {temperature}")
+                for i, msg in enumerate(messages):
+                    print(f"Message {i} ({msg['role']}): {msg['content'][:100]}...")
+
+            response = client.chat.completions.create(
+                model=model,
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+            )
+
+            if hasattr(response, "choices") and len(response.choices) > 0:
+                content = (
+                    response.choices[0].message.content or "No content in response"
+                )
+                finish_reason = response.choices[0].finish_reason
+
+                if finish_reason == "length":
+                    raise ValueError(
+                        f"Output was truncated due to max_tokens limit ({max_tokens}). Please increase max_tokens."
+                    )
+
+                generated_code = None
+                if "import vtk" not in content:
+                    generated_code = "import vtk\n" + content
+                else:
+                    pos = content.find("import vtk")
+                    if pos != -1:
+                        generated_code = content[pos:]
+                    else:
+                        generated_code = content
+
+                is_valid, error_msg = self.validate_code_syntax(generated_code)
+                if is_valid:
+                    return generated_code, response.usage
+
+                elif attempt < retry_attempts - 1:  # Don't print on last attempt
+                    if self.verbose:
+                        print(f"AST validation failed: {error_msg}. Retrying...")
+                    # Add error feedback to context for retry
+                    messages.append({"role": "assistant", "content": content})
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": (
+                                f"The generated code has a syntax error: {error_msg}. "
+                                "Please fix the syntax and generate valid Python code."
+                            ),
+                        }
+                    )
+                else:
+                    # Last attempt failed
+                    if self.verbose:
+                        print(f"Final attempt failed AST validation: {error_msg}")
+                    return (
+                        generated_code,
+                        response.usage,
+                    )  # Return anyway, let caller handle
+            else:
+                if attempt == retry_attempts - 1:
+                    return "No response generated", response.usage
+
+        return "No response generated"
+
+
+@click.command()
+@click.argument("input_string")
+@click.option(
+    "--provider",
+    type=click.Choice(["openai", "anthropic", "gemini", "nim"]),
+    default="openai",
+    help="LLM provider to use",
+)
+@click.option("-m", "--model", default="gpt-4o", help="Model name to use")
+@click.option(
+    "-k", "--max-tokens", type=int, default=1000, help="Max # of tokens to generate"
+)
+@click.option(
+    "--temperature",
+    type=float,
+    default=0.7,
+    help="Temperature for generation (0.0-2.0)",
+)
+@click.option(
+    "-t", "--token", required=True, help="API token for the selected provider"
+)
+@click.option("--base-url", help="Base URL for API (auto-detected or custom)")
+@click.option("-r", "--rag", is_flag=True, help="Use RAG to improve code generation")
+@click.option("-v", "--verbose", is_flag=True, help="Show generated source code")
+@click.option("--collection", default="vtk-examples", help="Collection name for RAG")
+@click.option(
+    "--database",
+    default="./db/codesage-codesage-large-v2",
+    help="Database path for RAG",
+)
+@click.option(
+    "--top-k", type=int, default=5, help="Number of examples to retrieve from RAG"
+)
+@click.option(
+    "--retry-attempts",
+    type=int,
+    default=1,
+    help="Number of times to retry if AST validation fails",
+)
+def main(
+    input_string,
+    provider,
+    model,
+    max_tokens,
+    temperature,
+    token,
+    base_url,
+    rag,
+    verbose,
+    collection,
+    database,
+    top_k,
+    retry_attempts,
+):
+    """Generate and execute VTK code using LLMs.
+
+    INPUT_STRING: The code description to generate VTK code for
+    """
+
+    # Set default base URLs
+    if base_url is None:
+        base_urls = {
+            "anthropic": "https://api.anthropic.com/v1",
+            "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/",
+            "nim": "https://integrate.api.nvidia.com/v1",
+        }
+        base_url = base_urls.get(provider)
+
+    # Set default models based on provider
+    if model == "gpt-4o":
+        default_models = {
+            "anthropic": "claude-3-5-sonnet-20241022",
+            "gemini": "gemini-1.5-pro",
+            "nim": "meta/llama3-70b-instruct",
+        }
+        model = default_models.get(provider, model)
+
+    try:
+        client = VTKPromptClient(
+            collection_name=collection,
+            database_path=database,
+            verbose=verbose,
+        )
+        generated_code, usage = client.query(
+            input_string,
+            api_key=token,
+            model=model,
+            base_url=base_url,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            rag=rag,
+            retry_attempts=retry_attempts,
+        )
+
+        if verbose and usage is not None:
+            print(
+                f"Used tokens: input={usage.prompt_tokens} output={usage.completion_tokens}"
+            )
+
+        client.run_code(generated_code)
+
+    except ValueError as e:
+        if "RAG components" in str(e):
+            print("rag_components not found", file=sys.stderr)
+            sys.exit(1)
+        elif "Failed to load RAG snippets" in str(e):
+            print("failed to load rag snippets", file=sys.stderr)
+            sys.exit(2)
+
+        elif "max_tokens" in str(e):
+            print(f"\nError: {e}", file=sys.stderr)
+            print(f"Current max_tokens: {max_tokens}", file=sys.stderr)
+            print("Try increasing with: --max-tokens <higher_number>", file=sys.stderr)
+            sys.exit(3)
+
+        else:
+            print(f"Error: {e}", file=sys.stderr)
+            sys.exit(4)
+
+
+if __name__ == "__main__":
+    main()