fix(.): update(features): add support for dynamic variables in templates

patched-codes · Sep 23, 2024 · ec92cbf · ec92cbf
1 parent 4ae41f1
commit ec92cbf
Show file tree

Hide file tree

Showing 5 changed files with 146 additions and 54 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Released]
 
+## [0.6.11]
+    - Support of dynamic variable such as {{input:var1}} in template
+    - Fix. Only a variable one time
+    - Update and improve price table
+
 ## [0.6.9]
     - Improve display --tokens
 

diff --git a/README.md b/README.md
@@ -363,6 +363,81 @@ This command will analyze your project, count the tokens, and provide a detailed
 
 
 
+
+## Feature Highlight: Dynamic Variable Extraction for Prompt Generation
+
+`code2prompt` offers a powerful feature for dynamic variable extraction from templates, allowing for interactive and customizable prompt generation. Using the syntax `{{input:variable_name}}`, you can easily define variables that will prompt users for input during execution. 
+
+This is particularly useful for creating flexible templates for various purposes, such as generating AI prompts for Chrome extensions. Here's an example:
+
+```jinja2
+# AI Prompt Generator for Chrome Extension
+
+Generate a prompt for an AI to create a Chrome extension with the following specifications:
+
+Extension Name: {{input:extension_name}}
+Main Functionality: {{input:main_functionality}}
+Target Audience: {{input:target_audience}}
+
+## Prompt:
+
+You are an experienced Chrome extension developer. Create a detailed plan for a Chrome extension named "{{input:extension_name}}" that {{input:main_functionality}}. This extension is designed for {{input:target_audience}}.
+
+Your response should include:
+
+1. A brief description of the extension's purpose and functionality
+2. Key features (at least 3)
+3. User interface design considerations
+4. Potential challenges in development and how to overcome them
+5. Security and privacy considerations
+6. A basic code structure for the main components (manifest.json, background script, content script, etc.)
+
+Ensure that your plan is detailed, technically sound, and tailored to the needs of {{input:target_audience}}.
+
+Start from this codebase:
+
+---- 
+
+## The codebase:
+
+<codebase>
+
+<toc>
+## Table of Contents
+
+{% for file in files %}{{ file.path }}
+{% endfor %}
+</toc>
+
+<code>
+{% for file in files %}
+## {{ file.path }}
+
+```{{ file.language }}
+{{ file.content }}
+```
+
+{% endfor %}
+</code>
+
+</codebase>
+
+
+```
+
+When you run `code2prompt` with this template, it will automatically detect the `{{input:variable_name}}` patterns and prompt the user to provide values for each variable (extension_name, main_functionality, and target_audience). This allows for flexible and interactive prompt generation, making it easy to create customized AI prompts for various Chrome extension ideas.
+
+For example, if a user inputs:
+- Extension Name: "ProductivityBoost"
+- Main Functionality: "tracks time spent on different websites and provides productivity insights"
+- Target Audience: "professionals working from home"
+
+The tool will generate a tailored prompt for an AI to create a detailed plan for this specific Chrome extension. This feature is particularly useful for developers, product managers, or anyone looking to quickly generate customized AI prompts for various projects or ideas.
+
+
+
+
+
 ## Configuration File
 
 Code2Prompt supports a `.code2promptrc` configuration file in JSON format for setting default options. Place this file in your project or home directory.

diff --git a/code2prompt/core/template_processor.py b/code2prompt/core/template_processor.py
@@ -4,33 +4,17 @@
 import re
 
 def load_template(template_path):
-    """
-    Load a Jinja2 template from a file.
-
-    Args:
-        template_path (str): Path to the template file.
-
-    Returns:
-        str: The contents of the template file.
-    """
+    """ Load a Jinja2 template from a file. """
     try:
-        with open(template_path, 'r') as file:
+        with open(template_path, 'r', encoding='utf-8') as file:
             return file.read()
     except IOError as e:
-        raise IOError(f"Error loading template file: {e}")
-
+        raise IOError(f"Error loading template file: {e}") from e
 
 def get_user_inputs(template_content):
-    """
-    Extract user-defined variables from the template and prompt for input.
-    Args:
-        template_content (str): The contents of the template file.
-    Returns:
-        dict: A dictionary of user-defined variables and their values.
-    """
-    # Use a regex pattern that allows for whitespace and special characters in variable names
-    # This pattern matches anything between {{ and }} that's not a curly brace
-    pattern = r'\{\{\s*([^{}]+?)\s*\}\}'
+    """ Extract user-defined variables from the template and prompt for input. """
+    # Use a regex pattern that excludes Jinja execute blocks and matches the new input syntax
+    pattern = r'{{\s*input:([^{}]+?)\s*}}'
     user_vars = re.findall(pattern, template_content)
     user_inputs = {}
 
@@ -44,19 +28,12 @@ def get_user_inputs(template_content):
     return user_inputs
 
 def process_template(template_content, files_data, user_inputs):
-    """
-    Process the Jinja2 template with the given data and user inputs.
-
-    Args:
-        template_content (str): The contents of the template file.
-        files_data (list): List of processed file data.
-        user_inputs (dict): Dictionary of user-defined variables and their values.
-
-    Returns:
-        str: The processed template content.
-    """
+    """ Process the Jinja2 template with the given data and user inputs. """
     try:
-        template = Template(template_content)
+        # Replace {{input:variable}} with {{variable}} for Jinja2 processing
+        processed_content = re.sub(r'{{\s*input:([^{}]+?)\s*}}', r'{{\1}}', template_content)
+
+        template = Template(processed_content)
         return template.render(files=files_data, **user_inputs)
     except Exception as e:
-        raise ValueError(f"Error processing template: {e}")
+        raise ValueError(f"Error processing template: {e}") from e
diff --git a/code2prompt/main.py b/code2prompt/main.py
@@ -12,7 +12,7 @@
 from code2prompt.utils.logging_utils import setup_logger, log_token_count, log_error, log_info
 from code2prompt.utils.price_calculator import load_token_prices, calculate_prices
 
-VERSION = "0.6.10"
+VERSION = "0.6.11"
 
 DEFAULT_OPTIONS = {
     "path": [],
@@ -226,7 +226,7 @@ def display_price_table(options, token_count):
 
     headers = ["Provider", "Model", "Price for 1K Input Tokens", "Number of Input Tokens", "Total Price"]
     table = tabulate(table_data, headers=headers, tablefmt="grid")
-    log_info("\n✨ Estimated Token Prices: (All prices are in USD, it is an estimate as the current token implementation is based on OpenAI's GPT-3)")
+    log_info("\n✨ Estimated Token Prices: (All prices are in USD, it is an estimate as the current token implementation is based on OpenAI's Tokenizer)")
     log_info("\n")
     log_info(table)
     log_info("\n📝 Note: The prices are based on the token count and the provider's pricing model.")

diff --git a/tests/test_template_processor.py b/tests/test_template_processor.py
@@ -9,67 +9,102 @@ def mock_prompt():
 
 def test_get_user_inputs_single_variable(mock_prompt):
     mock_prompt.return_value = "test_value"
-    template_content = "This is a {{variable}} test."
+    template_content = "This is a {{input:variable}} test."
     result = get_user_inputs(template_content)
     assert result == {"variable": "test_value"}
     mock_prompt.assert_called_once_with("Enter value for variable: ")
 
 def test_get_user_inputs_multiple_variables(mock_prompt):
     mock_prompt.side_effect = ["value1", "value2"]
-    template_content = "{{var1}} and {{var2}} are two variables."
+    template_content = "{{input:var1}} and {{input:var2}} are two variables."
     result = get_user_inputs(template_content)
     assert result == {"var1": "value1", "var2": "value2"}
     assert mock_prompt.call_count == 2
 
 def test_get_user_inputs_duplicate_variables(mock_prompt):
     mock_prompt.return_value = "repeated_value"
-    template_content = "{{var}} appears twice: {{var}}"
+    template_content = "{{input:var}} appears twice: {{input:var}}"
     result = get_user_inputs(template_content)
     assert result == {"var": "repeated_value"}
     mock_prompt.assert_called_once_with("Enter value for var: ")
 
 def test_get_user_inputs_no_variables(mock_prompt):
-    template_content = "This template has no variables."
+    template_content = "This template has no input variables."
     result = get_user_inputs(template_content)
     assert result == {}
     mock_prompt.assert_not_called()
 
 def test_get_user_inputs_whitespace_in_variable_names(mock_prompt):
     mock_prompt.side_effect = ["value1", "value2"]
-    template_content = "{{ var1 }} and {{  var2  }} have whitespace."
+    template_content = "{{ input:var1 }} and {{  input:var2  }} have whitespace."
     result = get_user_inputs(template_content)
     assert result == {"var1": "value1", "var2": "value2"}
     assert mock_prompt.call_count == 2
 
 def test_get_user_inputs_case_sensitivity(mock_prompt):
     mock_prompt.side_effect = ["value1", "value2"]
-    template_content = "{{VAR}} and {{var}} are different."
+    template_content = "{{input:VAR}} and {{input:var}} are different."
     result = get_user_inputs(template_content)
     assert result == {"VAR": "value1", "var": "value2"}
     assert mock_prompt.call_count == 2
 
 def test_get_user_inputs_special_characters(mock_prompt):
     mock_prompt.return_value = "special_value"
-    template_content = "This is a {{special!@#$%^&*()_+}} variable."
+    template_content = "This is a {{input:special!@#$%^&*()_+}} variable."
     result = get_user_inputs(template_content)
     assert result == {"special!@#$%^&*()_+": "special_value"}
     mock_prompt.assert_called_once_with("Enter value for special!@#$%^&*()_+: ")
 
 def test_get_user_inputs_empty_variable_name(mock_prompt):
-    template_content = "This has an {{}} empty variable name."
+    template_content = "This has an {{input:}} empty variable name."
     result = get_user_inputs(template_content)
     assert result == {}
     mock_prompt.assert_not_called()
 
-#def test_get_user_inputs_nested_variables(mock_prompt):
-#    mock_prompt.side_effect = ["outer", "inner"]
-#    template_content = "Nested {{outer{{inner}}}} variables."
-#    result = get_user_inputs(template_content)
-#    assert result == {"outer": "outer", "inner": "inner"}
-#    assert mock_prompt.call_count == 2
-
 def test_get_user_inputs_malformed_variables(mock_prompt):
-    template_content = "Malformed {{var} and {var}} variables."
+    template_content = "Malformed {{input:var} and {input:var}} variables."
     result = get_user_inputs(template_content)
     assert result == {}
-    mock_prompt.assert_not_called()
+    mock_prompt.assert_not_called()
+
+def test_get_user_inputs_ignore_jinja_execute_blocks(mock_prompt):
+    template_content = """
+    {% if condition %}
+        {{var}}
+    {% endif %}
+    {{input:user_var}}
+    {% for item in items %}
+        {{item}}
+    {% endfor %}
+    """
+    mock_prompt.return_value = "user_value"
+    result = get_user_inputs(template_content)
+    assert result == {"user_var": "user_value"}
+    mock_prompt.assert_called_once_with("Enter value for user_var: ")
+
+def test_get_user_inputs_mixed_variables(mock_prompt):
+    template_content = """
+    Regular variable: {{var}}
+    Input variable: {{input:user_var}}
+    {% if condition %}
+        Jinja block variable: {{block_var}}
+    {% endif %}
+    Another input: {{input:another_var}}
+    """
+    mock_prompt.side_effect = ["user_value", "another_value"]
+    result = get_user_inputs(template_content)
+    assert result == {"user_var": "user_value", "another_var": "another_value"}
+    assert mock_prompt.call_count == 2
+
+def test_get_user_inputs_nested_jinja_blocks(mock_prompt):
+    template_content = """
+    {% if outer_condition %}
+        {% for item in items %}
+            {{input:user_var}}
+        {% endfor %}
+    {% endif %}
+    """
+    mock_prompt.return_value = "user_value"
+    result = get_user_inputs(template_content)
+    assert result == {"user_var": "user_value"}
+    mock_prompt.assert_called_once_with("Enter value for user_var: ")