diff --git a/sdk/eidolon_ai_sdk/agent/retriever_agent/hyde_question_transformer.py b/sdk/eidolon_ai_sdk/agent/retriever_agent/hyde_question_transformer.py index 028dfc3fd..42a84b3b5 100644 --- a/sdk/eidolon_ai_sdk/agent/retriever_agent/hyde_question_transformer.py +++ b/sdk/eidolon_ai_sdk/agent/retriever_agent/hyde_question_transformer.py @@ -1,6 +1,7 @@ from typing import List -from jinja2 import Environment, StrictUndefined +from jinja2 import StrictUndefined +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import Field from eidolon_ai_sdk.agent.retriever_agent.question_transformer import QuestionTransformerSpec, QuestionTransformer @@ -22,7 +23,7 @@ def __init__(self, **kwargs): async def transform(self, apu: APU, process_id: str, question: str) -> List[str]: thread = apu.new_thread(process_id) - env = Environment(undefined=StrictUndefined) + env = ImmutableSandboxedEnvironment(undefined=StrictUndefined) userPrompt = env.from_string(self.spec.prompt).render(question=question) response = await thread.run_request(prompts=[UserTextAPUMessage(prompt=userPrompt)], output_format="str") diff --git a/sdk/eidolon_ai_sdk/agent/retriever_agent/multi_question_transformer.py b/sdk/eidolon_ai_sdk/agent/retriever_agent/multi_question_transformer.py index 91d39531e..8e449d3e3 100644 --- a/sdk/eidolon_ai_sdk/agent/retriever_agent/multi_question_transformer.py +++ b/sdk/eidolon_ai_sdk/agent/retriever_agent/multi_question_transformer.py @@ -1,6 +1,7 @@ from typing import List -from jinja2 import Environment, StrictUndefined +from jinja2 import StrictUndefined +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import Field, BaseModel from eidolon_ai_sdk.agent.retriever_agent.question_transformer import QuestionTransformerSpec, QuestionTransformer @@ -31,7 +32,7 @@ def __init__(self, **kwargs): async def transform(self, apu: APU, process_id: str, question: str) -> List[str]: thread = apu.new_thread(process_id) - env = Environment(undefined=StrictUndefined) + env = ImmutableSandboxedEnvironment(undefined=StrictUndefined) userPrompt = env.from_string(self.spec.prompt).render( question=question, number_to_generate=self.spec.number_to_generate ) diff --git a/sdk/eidolon_ai_sdk/agent/simple_agent.py b/sdk/eidolon_ai_sdk/agent/simple_agent.py index 2b99ffb9b..eedc3545e 100644 --- a/sdk/eidolon_ai_sdk/agent/simple_agent.py +++ b/sdk/eidolon_ai_sdk/agent/simple_agent.py @@ -6,6 +6,7 @@ from fastapi import Body from jinja2 import Environment, meta, StrictUndefined +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import field_validator, model_validator, BaseModel, Field from pydantic_core import to_jsonable_python, SchemaError @@ -138,7 +139,7 @@ async def create_process(self, process_id: str): t = default_apu.main_thread(process_id) template_args = {k: v.instantiate().value for k, v in self.prompt_templates.items()} - env = Environment(undefined=StrictUndefined) + env = ImmutableSandboxedEnvironment(undefined=StrictUndefined) system_prompt = env.from_string(self.system_prompt).render(**template_args) await t.set_boot_messages(prompts=[SystemAPUMessage(prompt=system_prompt)]) @@ -224,7 +225,7 @@ async def action_fn(process_id, action, **kwargs): if isinstance(request_body, dict): template_args.update(request_body) - env = Environment(undefined=StrictUndefined) + env = ImmutableSandboxedEnvironment(undefined=StrictUndefined) text_message = UserTextAPUMessage(prompt=env.from_string(action.user_prompt).render(**template_args)) if execute_on_apu and execute_on_apu in apus: diff --git a/sdk/eidolon_ai_sdk/agent/sql_agent/agent.py b/sdk/eidolon_ai_sdk/agent/sql_agent/agent.py index 41b979de3..ffe91f3ed 100644 --- a/sdk/eidolon_ai_sdk/agent/sql_agent/agent.py +++ b/sdk/eidolon_ai_sdk/agent/sql_agent/agent.py @@ -1,6 +1,7 @@ from typing import Literal, Optional, cast -from jinja2 import Environment, StrictUndefined, Template +from jinja2 import StrictUndefined, Template +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, Field from eidolon_ai_client.events import ( @@ -72,7 +73,7 @@ def __init__(self, **kwargs): else: raise ValueError("APU does not have logic units") - environment = Environment(undefined=StrictUndefined) + environment = ImmutableSandboxedEnvironment(undefined=StrictUndefined) self._error_prompt = environment.from_string(self.spec.error_prompt) self._system_prompt = environment.from_string(self.spec.system_prompt) self._user_prompt = environment.from_string(self.spec.user_prompt) diff --git a/sdk/eidolon_ai_sdk/agent/tot_agent/checker.py b/sdk/eidolon_ai_sdk/agent/tot_agent/checker.py index a7d010957..1234fb2b9 100644 --- a/sdk/eidolon_ai_sdk/agent/tot_agent/checker.py +++ b/sdk/eidolon_ai_sdk/agent/tot_agent/checker.py @@ -1,6 +1,7 @@ from typing import List -from jinja2 import StrictUndefined, Environment +from jinja2 import StrictUndefined +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import Field, BaseModel from eidolon_ai_sdk.apu.apu import APU @@ -43,7 +44,7 @@ async def evaluate( """ checker_prompt = ( - Environment(undefined=StrictUndefined) + ImmutableSandboxedEnvironment(undefined=StrictUndefined) .from_string(self.spec.prompt) .render( problem=problem_description, diff --git a/sdk/eidolon_ai_sdk/agent/tot_agent/thought_generators.py b/sdk/eidolon_ai_sdk/agent/tot_agent/thought_generators.py index ca55fd7ea..8c26d6e28 100644 --- a/sdk/eidolon_ai_sdk/agent/tot_agent/thought_generators.py +++ b/sdk/eidolon_ai_sdk/agent/tot_agent/thought_generators.py @@ -6,10 +6,12 @@ non-repeating thoughts, which are crucial for problem-solving tasks that require exploration. """ + from abc import abstractmethod from typing import Any, Dict, List, Callable, Awaitable -from jinja2 import StrictUndefined, Environment +from jinja2 import StrictUndefined +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import Field, BaseModel from eidolon_ai_sdk.apu.agent_io import SystemAPUMessage, UserTextAPUMessage, APUMessageTypes @@ -36,7 +38,7 @@ class ThoughtGenerationStrategy(Specable[TGSConfig]): """ spec: TGSConfig - env = Environment(undefined=StrictUndefined) + env = ImmutableSandboxedEnvironment(undefined=StrictUndefined) def __init__(self, spec): super().__init__(spec) diff --git a/sdk/eidolon_ai_sdk/agent/tot_agent/tot_agent.py b/sdk/eidolon_ai_sdk/agent/tot_agent/tot_agent.py index 212837eb6..b38567526 100644 --- a/sdk/eidolon_ai_sdk/agent/tot_agent/tot_agent.py +++ b/sdk/eidolon_ai_sdk/agent/tot_agent/tot_agent.py @@ -4,7 +4,8 @@ from typing import List, Dict, Any, Literal, Optional, Type, Union from fastapi import HTTPException -from jinja2 import StrictUndefined, Environment +from jinja2 import StrictUndefined +from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import Field, BaseModel from eidolon_ai_sdk.agent.agent import register_program @@ -121,7 +122,11 @@ async def question(self, process_id, body) -> TotResponse: # override to run the tree of thoughts algorithm in a separate thread thoughts_path: List[str] = [] level = 0 - question = Environment(undefined=StrictUndefined).from_string(self.spec.user_prompt).render(**body.model_dump()) + question = ( + ImmutableSandboxedEnvironment(undefined=StrictUndefined) + .from_string(self.spec.user_prompt) + .render(**body.model_dump()) + ) async def exec_request( _boot_messages: List[LLMMessage],