Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 54 additions & 0 deletions apps/common/init/init_template.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
# coding=utf-8
"""
@project: MaxKB
@Author:虎虎
@file: init_jinja.py
@date:2025/12/1 17:16
@desc:
"""
from typing import Any

from jinja2.sandbox import SandboxedEnvironment
from langchain_core.prompts.string import DEFAULT_FORMATTER_MAPPING, _HAS_JINJA2


def jinja2_formatter(template: str, /, **kwargs: Any) -> str:
"""Format a template using jinja2.
*Security warning*:
As of LangChain 0.0.329, this method uses Jinja2's
SandboxedEnvironment by default. However, this sand-boxing should
be treated as a best-effort approach rather than a guarantee of security.
Do not accept jinja2 templates from untrusted sources as they may lead
to arbitrary Python code execution.
https://jinja.palletsprojects.com/en/3.1.x/sandbox/
Args:
template: The template string.
**kwargs: The variables to format the template with.
Returns:
The formatted string.
Raises:
ImportError: If jinja2 is not installed.
"""
if not _HAS_JINJA2:
msg = (
"jinja2 not installed, which is needed to use the jinja2_formatter. "
"Please install it with `pip install jinja2`."
"Please be cautious when using jinja2 templates. "
"Do not expand jinja2 templates using unverified or user-controlled "
"inputs as that can result in arbitrary Python code execution."
)
raise ImportError(msg)

# Use a restricted sandbox that blocks ALL attribute/method access
# Only simple variable lookups like {{variable}} are allowed
# Attribute access like {{variable.attr}} or {{variable.method()}} is blocked
return SandboxedEnvironment().from_string(template).render(**kwargs)


def run():
DEFAULT_FORMATTER_MAPPING['jinja2'] = jinja2_formatter
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code looks largely correct and performs its intended function. Here are a few suggestions for improvement:

  1. Import Statements: Ensure that _HAS_JINJA2 is imported within the jinja2_formatter function to avoid undefined variable errors.

  2. Function Return Type Annotation: Since you're explicitly returning a string in the render() method call, you might consider adding it to the return type annotation of the jinja2_formatter function.

Here’s an updated version with these improvements:

+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author:虎虎
+    @file: init_jinja.py
+    @date:2025/12/1 17:16
+    @desc:
+"""
+from __future__ import annotations
+
+from typing import Any, Callable

+import jinja2.sandbox
+from langchain_core.prompts.string import DEFAULT_FORMATTER_MAPPING, _HAS_JINJA2


def jinja2_formatter(template: str, /, **kwargs: Any) -> str:
    """Format a template using jinja2.
    
    *Security warning*:
        As of LangChain 0.0.329, this method uses Jinja2's
        SandboxedEnvironment by default. However, this sand-boxing should
        be treated as a best-effort approach rather than a guarantee of security.
        Do not accept jinja2 templates from untrusted sources as they may lead
        to arbitrary Python code execution.

        https://jinja.palletsprojects.com/en/3.1.x/sandbox/

    Args:
        template: The template string.
        **kwargs: The variables to format the template with.

    Returns:
        The formatted string.

    Raises:
        ImportError: If jinja2 is not installed.
    """
    if not _HAS_JINJA2:
        msg = (
            "jinja2 not installed, which is needed to use the jinja2_formatter. "
            "Please install it with `pip install jinja2`."
            "Please be cautious when using jinja2 templates. "
            "Do not expand jinja2 templates using unverified or user-controlled "
            "inputs as that can result in arbitrary Python code execution."
        )
        raise ImportError(msg)

    # Use a restricted sandbox that blocks ALL attribute/method access
    # Only simple variable lookups like {{variable}} are allowed
    # Attribute access like {{variable.attr}} or {{variable.method()}} is blocked
    env = jinja2.sandbox.SandboxedEnvironment()
    return env.from_string(template).render(**kwargs)


run(): ...

Key Changes:

  1. Imports: Added __future__ import annotations to indicate type hints and ensured _HAS_JINJA2 is used properly within the function.
  2. Return Type Annotation: Specified -> str in the return statement of the jinja2_formatter function to clarify the expected output type.
  3. Variable Scope: Imported Callable[Any] at the top to specify the argument types more clearly; although this doesn't change the current implementation but good practice.

These changes make the code cleaner and clearer while ensuring proper error handling and type safety.

3 changes: 2 additions & 1 deletion apps/maxkb/wsgi/web.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,10 @@ def __call__(self, name, *args, **kwargs):
def post_handler():
from common.database_model_manage.database_model_manage import DatabaseModelManage
from common import event

from common.init import init_template
event.run()
DatabaseModelManage.init()
init_template.run()


# 启动后处理函数
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided Python code looks mostly fine, but there are a few minor improvements that could be made:

  1. Use __init__.py: Ensure that your directories have an __init__.py file to make them packages so they can be imported correctly.

  2. Function Naming: It's good practice to use descriptive function names. The function post_handler is understandable, but event.run() seems like it should be more specific about its role (e.g., run_event_handlers, execute_event_logic).

  3. Imports Order: Consider organizing imports based on readability and dependencies. For example, you might group related import statements together or sort them alphabetically within groups.

  4. Comments: Provide comments explaining the purpose of each section. This can help others understand what the code does without needing to read through the entire function.

Here's a slightly optimized version with these improvements applied:

from __future__ import annotations

# Initialize necessary modules at startup
def initialize():
    """
    Executes initialization tasks for the application.
    """
    from common.database_model_manage.database_model_manage import DatabaseModelManage
    DatabaseModelManage.initialize()

    from common.event_handler import run_event_handlers  # Assuming this exists
    run_event_handlers()

    from common.template_loader import load_templates  # Example usage; adjust as needed
    load_templates()


class PostProcessor:
    def __init__(self):
        pass

    def __call__(self, name, *args, **kwargs):
        """
        Performs post-processing actions if called.
        """
        # Add actual processing logic here
        print(f"Performing post-processing for: {name}")


# Start up handler
initialize()

# Placeholder for further startup logic

Explanation of Changes:

  • Initialize Function Splitting: Introduced initialize() function to encapsulate all initialization steps, making the main.py cleaner.

  • Comments Enhancements: Added comments to explain key parts of the initialize() function and improve overall clarity.

These changes make the code more modular and easier to maintain.

Expand Down
Loading