ai-robot-core/ai-service/app/services/prompt/variable_resolver.py

145 lines
4.0 KiB
Python

"""
Variable resolver for prompt templates.
[AC-AISVC-56] Built-in and custom variable replacement engine.
"""
import logging
import re
from datetime import datetime
from typing import Any
logger = logging.getLogger(__name__)
VARIABLE_PATTERN = re.compile(r"\{\{(\w+)\}\}")
BUILTIN_VARIABLES = {
"persona_name": "小N",
"current_time": lambda: datetime.now().strftime("%Y-%m-%d %H:%M"),
"channel_type": "default",
"tenant_name": "平台",
"session_id": "",
}
class VariableResolver:
"""
[AC-AISVC-56] Variable replacement engine for prompt templates.
Supports:
- Built-in variables: persona_name, current_time, channel_type, tenant_name, session_id
- Custom variables: defined in template with defaults
"""
def __init__(
self,
channel_type: str = "default",
tenant_name: str = "平台",
session_id: str = "",
):
self._context = {
"channel_type": channel_type,
"tenant_name": tenant_name,
"session_id": session_id,
}
def resolve(
self,
template: str,
variables: list[dict[str, Any]] | None = None,
extra_context: dict[str, Any] | None = None,
) -> str:
"""
Resolve all {{variable}} placeholders in the template.
Args:
template: Template string with {{variable}} placeholders
variables: Custom variable definitions from template
extra_context: Additional context for resolution
Returns:
Template with all variables replaced
"""
context = self._build_context(variables, extra_context)
def replace_var(match: re.Match) -> str:
var_name = match.group(1)
if var_name in context:
value = context[var_name]
if callable(value):
return str(value())
return str(value)
logger.warning(f"Unknown variable in template: {var_name}")
return match.group(0)
resolved = VARIABLE_PATTERN.sub(replace_var, template)
return resolved
def _build_context(
self,
variables: list[dict[str, Any]] | None,
extra_context: dict[str, Any] | None,
) -> dict[str, Any]:
"""Build the complete context for variable resolution."""
context = {}
for key, value in BUILTIN_VARIABLES.items():
if key in self._context:
context[key] = self._context[key]
else:
context[key] = value
if variables:
for var in variables:
name = var.get("name")
default = var.get("default", "")
if name:
context[name] = default
if extra_context:
context.update(extra_context)
return context
def extract_variables(self, template: str) -> list[str]:
"""
Extract all variable names from a template.
Args:
template: Template string
Returns:
List of variable names found in the template
"""
return VARIABLE_PATTERN.findall(template)
def validate_variables(
self,
template: str,
defined_variables: list[dict[str, Any]] | None,
) -> dict[str, Any]:
"""
Validate that all variables in template are defined.
Args:
template: Template string
defined_variables: Variables defined in template metadata
Returns:
Dict with 'valid' boolean and 'missing' list
"""
used_vars = set(self.extract_variables(template))
builtin_vars = set(BUILTIN_VARIABLES.keys())
defined_names = set()
if defined_variables:
defined_names = {v.get("name") for v in defined_variables if v.get("name")}
available_vars = builtin_vars | defined_names
missing = used_vars - available_vars
return {
"valid": len(missing) == 0,
"missing": list(missing),
"used_variables": list(used_vars),
}