fixed liting errors

This commit is contained in:
nkoorty 2025-07-07 18:07:07 +02:00
parent 9c219db9ba
commit dc8bdae64f
2 changed files with 21 additions and 22 deletions

View File

@ -11,7 +11,6 @@ from concurrent.futures import Executor, ThreadPoolExecutor
from datetime import datetime from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, final from typing import Any, Callable, Dict, List, Optional, Tuple, Type, final
from jinja2 import Template
from jinja2.sandbox import SandboxedEnvironment from jinja2.sandbox import SandboxedEnvironment
from dbgpt._private.pydantic import ConfigDict, Field from dbgpt._private.pydantic import ConfigDict, Field
@ -40,18 +39,18 @@ class ConversableAgent(Role, Agent):
"""ConversableAgent is an agent that can communicate with other agents.""" """ConversableAgent is an agent that can communicate with other agents."""
model_config = ConfigDict(arbitrary_types_allowed=True) model_config = ConfigDict(arbitrary_types_allowed=True)
# Dangerous template patterns that could lead to code execution # Dangerous template patterns that could lead to code execution
_DANGEROUS_TEMPLATE_PATTERNS = [ _DANGEROUS_TEMPLATE_PATTERNS = [
r'\{\{.*__.*\}\}', # Double underscore methods r"\{\{.*__.*\}\}", # Double underscore methods
r'\{\{.*import.*\}\}', # Import statements r"\{\{.*import.*\}\}", # Import statements
r'\{\{.*exec.*\}\}', # Exec calls r"\{\{.*exec.*\}\}", # Exec calls
r'\{\{.*eval.*\}\}', # Eval calls r"\{\{.*eval.*\}\}", # Eval calls
r'\{\{.*open.*\}\}', # File operations r"\{\{.*open.*\}\}", # File operations
r'\{\{.*subprocess.*\}\}', # Subprocess calls r"\{\{.*subprocess.*\}\}", # Subprocess calls
r'\{\{.*os\..*\}\}', # OS module access r"\{\{.*os\..*\}\}", # OS module access
r'\{\{.*globals.*\}\}', # Globals access r"\{\{.*globals.*\}\}", # Globals access
r'\{\{.*\[.*\].*\}\}', # Bracket notation access r"\{\{.*\[.*\].*\}\}", # Bracket notation access
] ]
agent_context: Optional[AgentContext] = Field(None, description="Agent context") agent_context: Optional[AgentContext] = Field(None, description="Agent context")
@ -902,13 +901,13 @@ class ConversableAgent(Role, Agent):
"""Sanitize template parameters to prevent injection attacks.""" """Sanitize template parameters to prevent injection attacks."""
if not params: if not params:
return params return params
return self._sanitize_dict(params) return self._sanitize_dict(params)
def _sanitize_dict(self, data: Dict[str, Any]) -> Dict[str, Any]: def _sanitize_dict(self, data: Dict[str, Any]) -> Dict[str, Any]:
"""Sanitize dictionary values recursively.""" """Sanitize dictionary values recursively."""
return {key: self._sanitize_value(value) for key, value in data.items()} return {key: self._sanitize_value(value) for key, value in data.items()}
def _sanitize_value(self, value: Any) -> Any: def _sanitize_value(self, value: Any) -> Any:
"""Sanitize a single value based on its type.""" """Sanitize a single value based on its type."""
if isinstance(value, str): if isinstance(value, str):
@ -918,7 +917,7 @@ class ConversableAgent(Role, Agent):
elif isinstance(value, list): elif isinstance(value, list):
return [self._sanitize_value(item) for item in value] return [self._sanitize_value(item) for item in value]
return value return value
def _sanitize_string(self, text: str) -> str: def _sanitize_string(self, text: str) -> str:
"""Check string for dangerous template injection patterns.""" """Check string for dangerous template injection patterns."""
for pattern in self._DANGEROUS_TEMPLATE_PATTERNS: for pattern in self._DANGEROUS_TEMPLATE_PATTERNS:
@ -1100,7 +1099,7 @@ class ConversableAgent(Role, Agent):
"""Build system prompt with security controls.""" """Build system prompt with security controls."""
if self.bind_prompt: if self.bind_prompt:
return self._render_bind_prompt(resource_vars, context) return self._render_bind_prompt(resource_vars, context)
# Fallback to build_prompt with sanitized context # Fallback to build_prompt with sanitized context
sanitized_context = self.sanitize_template_params(context or {}) sanitized_context = self.sanitize_template_params(context or {})
return await self.build_prompt( return await self.build_prompt(
@ -1111,11 +1110,11 @@ class ConversableAgent(Role, Agent):
is_retry_chat=is_retry_chat, is_retry_chat=is_retry_chat,
**sanitized_context, **sanitized_context,
) )
def _render_bind_prompt( def _render_bind_prompt(
self, self,
resource_vars: Optional[Dict] = None, resource_vars: Optional[Dict] = None,
context: Optional[Dict[str, Any]] = None context: Optional[Dict[str, Any]] = None,
) -> str: ) -> str:
"""Render bind prompt template with sanitized parameters.""" """Render bind prompt template with sanitized parameters."""
prompt_param = {} prompt_param = {}
@ -1124,7 +1123,7 @@ class ConversableAgent(Role, Agent):
if context: if context:
sanitized_context = self.sanitize_template_params(context) sanitized_context = self.sanitize_template_params(context)
prompt_param.update(sanitized_context) prompt_param.update(sanitized_context)
if self.bind_prompt.template_format == "f-string": if self.bind_prompt.template_format == "f-string":
return self.bind_prompt.template.format(**prompt_param) return self.bind_prompt.template.format(**prompt_param)
elif self.bind_prompt.template_format == "jinja2": elif self.bind_prompt.template_format == "jinja2":

View File

@ -31,7 +31,7 @@ def _jinja2_formatter(template: str, **kwargs: Any) -> str:
"jinja2 not installed, which is needed to use the jinja2_formatter. " "jinja2 not installed, which is needed to use the jinja2_formatter. "
"Please install it with `pip install jinja2`." "Please install it with `pip install jinja2`."
) )
env = SandboxedEnvironment() env = SandboxedEnvironment()
return env.from_string(template).render(**kwargs) return env.from_string(template).render(**kwargs)