Feat/jinja2

cozeloop/entities/prompt.py

@@ -9,6 +9,7 @@
 
 class TemplateType(str, Enum):
     NORMAL = "normal"
+    JINJA2 = "jinja2"
 
 
 class Role(str, Enum):
@@ -26,6 +27,15 @@ class ToolType(str, Enum):
 class VariableType(str, Enum):
     STRING = "string"
     PLACEHOLDER = "placeholder"
+    BOOLEAN = "boolean"
+    INTEGER = "integer"
+    FLOAT = "float"
+    OBJECT = "object"
+    ARRAY_STRING = "array<string>"
+    ARRAY_BOOLEAN = "array<boolean>"
+    ARRAY_INTEGER = "array<integer>"
+    ARRAY_FLOAT = "array<float>"
+    ARRAY_OBJECT = "array<object>"
 
 
 class ToolChoiceType(str, Enum):

cozeloop/internal/prompt/converter.py

@@ -59,7 +59,16 @@ def _convert_message(msg: OpenAPIMessage) -> EntityMessage:
 def _convert_variable_type(openapi_type: OpenAPIVariableType) -> EntityVariableType:
     type_mapping = {
         OpenAPIVariableType.STRING: EntityVariableType.STRING,
-        OpenAPIVariableType.PLACEHOLDER: EntityVariableType.PLACEHOLDER
+        OpenAPIVariableType.PLACEHOLDER: EntityVariableType.PLACEHOLDER,
+        OpenAPIVariableType.BOOLEAN: EntityVariableType.BOOLEAN,
+        OpenAPIVariableType.INTEGER: EntityVariableType.INTEGER,
+        OpenAPIVariableType.FLOAT: EntityVariableType.FLOAT,
+        OpenAPIVariableType.OBJECT: EntityVariableType.OBJECT,
+        OpenAPIVariableType.ARRAY_STRING: EntityVariableType.ARRAY_STRING,
+        OpenAPIVariableType.ARRAY_INTEGER: EntityVariableType.ARRAY_INTEGER,
+        OpenAPIVariableType.ARRAY_FLOAT: EntityVariableType.ARRAY_FLOAT,
+        OpenAPIVariableType.ARRAY_BOOLEAN: EntityVariableType.ARRAY_BOOLEAN,
+        OpenAPIVariableType.ARRAY_OBJECT: EntityVariableType.ARRAY_OBJECT
     }
     return type_mapping.get(openapi_type, EntityVariableType.STRING)  # Default to STRING type
 
@@ -122,7 +131,8 @@ def _convert_llm_config(config: OpenAPIModelConfig) -> EntityModelConfig:
 
 def _convert_template_type(openapi_template_type: OpenAPITemplateType) -> EntityTemplateType:
     template_mapping = {
-        OpenAPITemplateType.NORMAL: EntityTemplateType.NORMAL
+        OpenAPITemplateType.NORMAL: EntityTemplateType.NORMAL,
+        OpenAPITemplateType.JINJA2: EntityTemplateType.JINJA2
     }
     return template_mapping.get(openapi_template_type, EntityTemplateType.NORMAL)  # Default to NORMAL type
 

cozeloop/internal/prompt/openapi.py

@@ -14,6 +14,7 @@
 
 class TemplateType(str, Enum):
     NORMAL = "normal"
+    JINJA2 = "jinja2"
 
 
 class Role(str, Enum):
@@ -31,6 +32,15 @@ class ToolType(str, Enum):
 class VariableType(str, Enum):
     STRING = "string"
     PLACEHOLDER = "placeholder"
+    BOOLEAN = "boolean"
+    INTEGER = "integer"
+    FLOAT = "float"
+    OBJECT = "object"
+    ARRAY_STRING = "array<string>"
+    ARRAY_BOOLEAN = "array<boolean>"
+    ARRAY_INTEGER = "array<integer>"
+    ARRAY_FLOAT = "array<float>"
+    ARRAY_OBJECT = "array<object>"
 
 
 class ToolChoiceType(str, Enum):

cozeloop/internal/prompt/prompt.py

@@ -6,6 +6,7 @@
 
 from jinja2 import Environment, BaseLoader, Undefined
 from jinja2.utils import missing, object_type_repr
+from jinja2.sandbox import SandboxedEnvironment
 
 from cozeloop.spec.tracespec import PROMPT_KEY, INPUT, PROMPT_VERSION, V_SCENE_PROMPT_TEMPLATE, V_SCENE_PROMPT_HUB
 from cozeloop.entities.prompt import (Prompt, Message, VariableDef, VariableType, TemplateType, Role,
@@ -153,6 +154,27 @@ def _validate_variable_values_type(self, variable_defs: List[VariableDef], varia
             elif var_def.type == VariableType.PLACEHOLDER:
                 if not (isinstance(val, Message) or (isinstance(val, List) and all(isinstance(item, Message) for item in val))):
                     raise ValueError(f"type of variable '{var_def.key}' should be Message like object")
+            elif var_def.type == VariableType.BOOLEAN:
+                if not isinstance(val, bool):
+                    raise ValueError(f"type of variable '{var_def.key}' should be bool")
+            elif var_def.type == VariableType.INTEGER:
+                if not isinstance(val, int):
+                    raise ValueError(f"type of variable '{var_def.key}' should be int")
+            elif var_def.type == VariableType.FLOAT:
+                if not isinstance(val, float):
+                    raise ValueError(f"type of variable '{var_def.key}' should be float")
+            elif var_def.type == VariableType.ARRAY_STRING:
+                if not isinstance(val, list) or not all(isinstance(item, str) for item in val):
+                    raise ValueError(f"type of variable '{var_def.key}' should be array<string>")
+            elif var_def.type == VariableType.ARRAY_BOOLEAN:
+                if not isinstance(val, list) or not all(isinstance(item, bool) for item in val):
+                    raise ValueError(f"type of variable '{var_def.key}' should be array<boolean>")
+            elif var_def.type == VariableType.ARRAY_INTEGER:
+                if not isinstance(val, list) or not all(isinstance(item, int) for item in val):
+                    raise ValueError(f"type of variable '{var_def.key}' should be array<integer>")
+            elif var_def.type == VariableType.ARRAY_FLOAT:
+                if not isinstance(val, list) or not all(isinstance(item, float) for item in val):
+                    raise ValueError(f"type of variable '{var_def.key}' should be array<float>")
 
     def _format_normal_messages(
             self,
@@ -217,7 +239,7 @@ def _render_text_content(
     ) -> str:
         if template_type == TemplateType.NORMAL:
             # Create custom Environment using DebugUndefined to preserve original form of undefined variables
-            env = Environment(
+            env = SandboxedEnvironment(
                 loader=BaseLoader(),
                 undefined=CustomUndefined,
                 variable_start_string='{{',
@@ -230,10 +252,20 @@ def _render_text_content(
             render_vars = {k: variables.get(k, '') for k in variable_def_map.keys()}
             # Render template
             return template.render(**render_vars)
+        elif template_type == TemplateType.JINJA2:
+            return self._render_jinja2_template(template_str, variable_def_map, variables)
         else:
             raise ValueError(f"text render unsupported template type: {template_type}")
 
 
+    def _render_jinja2_template(self, template_str: str, variable_def_map: Dict[str, VariableDef],
+                                variables: Dict[str, Any]) -> str:
+        """渲染 Jinja2 模板"""
+        env = SandboxedEnvironment()
+        template = env.from_string(template_str)
+        render_vars = {k: variables[k] for k in variable_def_map.keys() if variables is not None and k in variables}
+        return template.render(**render_vars)
+
 class CustomUndefined(Undefined):
     __slots__ = ()
 

examples/prompt/advance/prompt_hub_with_jinja.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+import json
+import time
+from typing import List
+
+import cozeloop
+from cozeloop import Message
+from cozeloop.entities.prompt import Role
+from cozeloop.spec.tracespec import CALL_OPTIONS, ModelCallOption, ModelMessage, ModelInput
+
+
+def convert_model_input(messages: List[Message]) -> ModelInput:
+    model_messages = []
+    for message in messages:
+        model_messages.append(ModelMessage(
+            role=str(message.role),
+            content=message.content if message.content is not None else ""
+        ))
+
+    return ModelInput(
+        messages=model_messages
+    )
+
+
+class LLMRunner:
+    def __init__(self, client):
+        self.client = client
+
+    def llm_call(self, input_data):
+        """
+        Simulate an LLM call and set relevant span tags.
+        """
+        span = self.client.start_span("llmCall", "model")
+        try:
+            # Assuming llm is processing
+            # output = ChatOpenAI().invoke(input=input_data)
+
+            # mock resp
+            time.sleep(1)
+            output = "I'm a robot. I don't have a specific name. You can give me one."
+            input_token = 232
+            output_token = 1211
+
+            # set tag key: `input`
+            span.set_input(convert_model_input(input_data))
+            # set tag key: `output`
+            span.set_output(output)
+            # set tag key: `model_provider`, e.g., openai, etc.
+            span.set_model_provider("openai")
+            # set tag key: `start_time_first_resp`
+            # Timestamp of the first packet return from LLM, unit: microseconds.
+            # When `start_time_first_resp` is set, a tag named `latency_first_resp` calculated
+            # based on the span's StartTime will be added, meaning the latency for the first packet.
+            span.set_start_time_first_resp(int(time.time() * 1000000))
+            # set tag key: `input_tokens`. The amount of input tokens.
+            # when the `input_tokens` value is set, it will automatically sum with the `output_tokens` to calculate the `tokens` tag.
+            span.set_input_tokens(input_token)
+            # set tag key: `output_tokens`. The amount of output tokens.
+            # when the `output_tokens` value is set, it will automatically sum with the `input_tokens` to calculate the `tokens` tag.
+            span.set_output_tokens(output_token)
+            # set tag key: `model_name`, e.g., gpt-4-1106-preview, etc.
+            span.set_model_name("gpt-4-1106-preview")
+            span.set_tags({CALL_OPTIONS: ModelCallOption(
+                temperature=0.5,
+                top_p=0.5,
+                top_k=10,
+                presence_penalty=0.5,
+                frequency_penalty=0.5,
+                max_tokens=1024,
+            )})
+
+            return None
+        except Exception as e:
+            raise e
+        finally:
+            span.finish()
+
+# If you want to use the jinja templates in prompts, you can refer to the following.
+if __name__ == '__main__':
+    # 1.Create a prompt on the platform
+    # You can create a Prompt on the platform's Prompt development page (set Prompt Key to 'prompt_hub_demo'),
+    # add the following messages to the template, and submit a version.
+    # System: You are a helpful bot, the conversation topic is {{var1}}.
+    # Placeholder: placeholder1
+    # User: My question is {{var2}}
+    # Placeholder: placeholder2
+
+    # Set the following environment variables first.
+    # COZELOOP_WORKSPACE_ID=your workspace id
+    # COZELOOP_API_TOKEN=your token
+    # 2.New loop client
+    client = cozeloop.new_client(
+        # Set whether to report a trace span when get or format prompt.
+        # Default value is false.
+        prompt_trace=True)
+
+    # 3. new root span
+    rootSpan = client.start_span("root_span", "main_span")
+
+    # 4. Get the prompt
+    # If no specific version is specified, the latest version of the corresponding prompt will be obtained
+    prompt = client.get_prompt(prompt_key="prompt_hub_demo", version="0.0.1")
+    if prompt is not None:
+        # Get messages of the prompt
+        if prompt.prompt_template is not None:
+            messages = prompt.prompt_template.messages
+            print(
+                f"prompt messages: {json.dumps([message.model_dump(exclude_none=True) for message in messages], ensure_ascii=False)}")
+        # Get llm config of the prompt
+        if prompt.llm_config is not None:
+            llm_config = prompt.llm_config
+            print(f"prompt llm_config: {llm_config.model_dump_json(exclude_none=True)}")
+
+        # 5.Format messages of the prompt
+        formatted_messages = client.prompt_format(prompt, {
+            "var_string": "hi",
+            "var_int": 5,
+            "var_bool": True,
+            "var_float": 1.0,
+            "var_object": {
+                "name":    "John",
+                "age":     30,
+                "hobbies": ["reading", "coding"],
+                "address": {
+                    "city":   "bejing",
+                    "street": "123 Main",
+                },
+            },
+            "var_array_string": ["hello", "nihao"],
+            "var_array_boolean": [True, False, True],
+            "var_array_int": [1, 2, 3, 4],
+            "var_array_float": [1.0, 2.0],
+            "var_array_object": [{"key": "123"}, {"value": 100}],
+            # Placeholder variable type should be Message/List[Message]
+            "placeholder1": [Message(role=Role.USER, content="Hello!"),
+                             Message(role=Role.ASSISTANT, content="Hello!")]
+            # Other variables in the prompt template that are not provided with corresponding values will be
+            # considered as empty values.
+        })
+        print(
+            f"formatted_messages: {json.dumps([message.model_dump(exclude_none=True) for message in formatted_messages], ensure_ascii=False)}")
+
+        # 6.LLM call
+        llm_runner = LLMRunner(client)
+        llm_runner.llm_call(formatted_messages)
+
+    rootSpan.finish()
+    # 4. (optional) flush or close
+    # -- force flush, report all traces in the queue
+    # Warning! In general, this method is not needed to be call, as spans will be automatically reported in batches.
+    # Note that flush will block and wait for the report to complete, and it may cause frequent reporting,
+    # affecting performance.
+    client.flush()