[feat][prompt] support get prompt by label

cozeloop/_client.py

@@ -259,10 +259,10 @@ def close(self):
         self._trace_provider.close_trace()
         self._closed = True
 
-    def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]:
+    def get_prompt(self, prompt_key: str, version: str = '', label: str = '') -> Optional[Prompt]:
         if self._closed:
             raise ClientClosedError()
-        return self._prompt_provider.get_prompt(prompt_key, version)
+        return self._prompt_provider.get_prompt(prompt_key, version, label)
 
     def prompt_format(self, prompt: Prompt, variables: Dict[str, PromptVariable]) -> List[Message]:
         if self._closed:
@@ -360,8 +360,8 @@ def close():
     return get_default_client().close()
 
 
-def get_prompt(prompt_key: str, version: str = '') -> Prompt:
-    return get_default_client().get_prompt(prompt_key, version)
+def get_prompt(prompt_key: str, version: str = '', label: str = '') -> Prompt:
+    return get_default_client().get_prompt(prompt_key, version, label)
 
 
 def prompt_format(prompt: Prompt, variables: Dict[str, Any]) -> List[Message]:

cozeloop/_noop.py

@@ -27,7 +27,7 @@ def workspace_id(self) -> str:
     def close(self):
         logger.warning(f"Noop client not supported. {self.new_exception}")
 
-    def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]:
+    def get_prompt(self, prompt_key: str, version: str = '', label: str = '') -> Optional[Prompt]:
         logger.warning(f"Noop client not supported. {self.new_exception}")
         raise self.new_exception
 

cozeloop/internal/prompt/cache.py

@@ -41,29 +41,29 @@ def __init__(self, workspace_id: str,
         if auto_refresh and self.openapi_client is not None:
             self._start_refresh_task()
 
-    def get(self, prompt_key: str, version: str) -> Optional['Prompt']:
-        cache_key = self._get_cache_key(prompt_key, version)
+    def get(self, prompt_key: str, version: str, label: str = '') -> Optional['Prompt']:
+        cache_key = self._get_cache_key(prompt_key, version, label)
         return self.cache.get(cache_key)
 
-    def set(self, prompt_key: str, version: str, value: 'Prompt') -> None:
-        cache_key = self._get_cache_key(prompt_key, version)
+    def set(self, prompt_key: str, version: str, label: str, value: 'Prompt') -> None:
+        cache_key = self._get_cache_key(prompt_key, version, label)
         self.cache[cache_key] = value
 
-    def get_all_prompt_queries(self) -> List[Tuple[str, str]]:
+    def get_all_prompt_queries(self) -> List[Tuple[str, str, str]]:
         result = []
         for cache_key in self.cache.keys():
             parsed = self._parse_cache_key(cache_key)
             if parsed:
                 result.append(parsed)
         return result
 
-    def _get_cache_key(self, prompt_key: str, version: str) -> str:
-        return f"prompt_hub:{prompt_key}:{version}"
+    def _get_cache_key(self, prompt_key: str, version: str, label: str = '') -> str:
+        return f"prompt_hub:{prompt_key}:{version}:{label}"
 
-    def _parse_cache_key(self, cache_key: str) -> Optional[Tuple[str, str]]:
+    def _parse_cache_key(self, cache_key: str) -> Optional[Tuple[str, str, str]]:
         parts = cache_key.split(':')
-        if len(parts) == 3:
-            return parts[1], parts[2]
+        if len(parts) == 4:
+            return parts[1], parts[2], parts[3]
         return None
 
     def _start_refresh_task(self):
@@ -91,13 +91,13 @@ def _refresh_all_prompts(self):
         """Refresh all cached prompts"""
         try:
             # Get all cached prompt_keys and versions
-            key_pairs = self.get_all_prompt_queries()
-            queries = [PromptQuery(prompt_key=prompt_key, version=version) for prompt_key, version in key_pairs]
+            key_tuples = self.get_all_prompt_queries()
+            queries = [PromptQuery(prompt_key=prompt_key, version=version, label=label) for prompt_key, version, label in key_tuples]
             try:
                 results = self.openapi_client.mpull_prompt(self.workspace_id, queries)
                 for result in results:
-                    prompt_key, version = result.query.prompt_key, result.query.version
-                    self.set(prompt_key, version, _convert_prompt(result.prompt))
+                    prompt_key, version, label = result.query.prompt_key, result.query.version, result.query.label
+                    self.set(prompt_key, version, label, _convert_prompt(result.prompt))
             except Exception as e:
                 logger.error(f"Error refreshing prompts: {e}")
 

cozeloop/internal/prompt/openapi.py

@@ -103,6 +103,7 @@ class Prompt(BaseModel):
 class PromptQuery(BaseModel):
     prompt_key: str
     version: Optional[str] = None
+    label: Optional[str] = None
 
 
 class MPullPromptRequest(BaseModel):

cozeloop/internal/prompt/prompt.py

@@ -4,11 +4,10 @@
 import json
 from typing import Dict, Any, List, Optional
 
-from jinja2 import Environment, BaseLoader, Undefined
-from jinja2.utils import missing, object_type_repr
+from jinja2 import BaseLoader, Undefined
 from jinja2.sandbox import SandboxedEnvironment
+from jinja2.utils import missing, object_type_repr
 
-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,
                                       PromptVariable)
 from cozeloop.internal import consts
@@ -18,6 +17,7 @@
 from cozeloop.internal.prompt.converter import _convert_prompt, _to_span_prompt_input, _to_span_prompt_output
 from cozeloop.internal.prompt.openapi import OpenAPIClient, PromptQuery
 from cozeloop.internal.trace.trace import TraceProvider
+from cozeloop.spec.tracespec import PROMPT_KEY, INPUT, PROMPT_VERSION, V_SCENE_PROMPT_TEMPLATE, V_SCENE_PROMPT_HUB, PROMPT_LABEL
 
 
 class PromptProvider:
@@ -39,18 +39,18 @@ def __init__(
                                  auto_refresh=True)
         self.prompt_trace = prompt_trace
 
-    def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]:
+    def get_prompt(self, prompt_key: str, version: str = '', label: str = '') -> Optional[Prompt]:
         # Trace reporting
         if self.prompt_trace and self.trace_provider is not None:
             with self.trace_provider.start_span(consts.TRACE_PROMPT_HUB_SPAN_NAME,
                                                 consts.TRACE_PROMPT_HUB_SPAN_TYPE,
                                                 scene=V_SCENE_PROMPT_HUB) as prompt_hub_pan:
                 prompt_hub_pan.set_tags({
                     PROMPT_KEY: prompt_key,
-                    INPUT: json.dumps({PROMPT_KEY: prompt_key, PROMPT_VERSION: version})
+                    INPUT: json.dumps({PROMPT_KEY: prompt_key, PROMPT_VERSION: version, PROMPT_LABEL: label})
                 })
                 try:
-                    prompt = self._get_prompt(prompt_key, version)
+                    prompt = self._get_prompt(prompt_key, version, label)
                     if prompt is not None:
                         prompt_hub_pan.set_tags({
                             PROMPT_VERSION: prompt.version,
@@ -65,20 +65,20 @@ def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]:
                     prompt_hub_pan.set_error(str(e))
                     raise e
         else:
-            return self._get_prompt(prompt_key, version)
+            return self._get_prompt(prompt_key, version, label)
 
-    def _get_prompt(self, prompt_key: str, version: str) -> Optional[Prompt]:
+    def _get_prompt(self, prompt_key: str, version: str, label: str = '') -> Optional[Prompt]:
         """
         Get Prompt, prioritize retrieving from cache, if not found then fetch from server
         """
         # Try to get from cache
-        prompt = self.cache.get(prompt_key, version)
+        prompt = self.cache.get(prompt_key, version, label)
         # If not in cache, fetch from server and cache it
         if prompt is None:
-            result = self.openapi_client.mpull_prompt(self.workspace_id, [PromptQuery(prompt_key=prompt_key, version=version)])
+            result = self.openapi_client.mpull_prompt(self.workspace_id, [PromptQuery(prompt_key=prompt_key, version=version, label=label)])
             if result:
                 prompt = _convert_prompt(result[0].prompt)
-                self.cache.set(prompt_key, version, prompt)
+                self.cache.set(prompt_key, version, label, prompt)
         # object cache item should be read only
         return prompt.copy(deep=True)
 

cozeloop/prompt.py

@@ -13,12 +13,13 @@ class PromptClient(ABC):
     """
 
     @abstractmethod
-    def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]:
+    def get_prompt(self, prompt_key: str, version: str = '', label: str = '') -> Optional[Prompt]:
         """
         Get a prompt by prompt key and version.
 
         :param prompt_key: A unique key for retrieving the prompt.
         :param version: The version of the prompt. Defaults to empty, which represents fetching the latest version.
+        :param label: The label of the prompt. Defaults to empty.
         :return: An instance of `entity.Prompt` if found, or None.
         """
 

cozeloop/spec/tracespec/span_key.py

@@ -19,6 +19,7 @@
 PROMPT_PROVIDER = "prompt_provider" # Prompt providers, such as Loop, Langsmith, etc.
 PROMPT_KEY = "prompt_key"
 PROMPT_VERSION = "prompt_version"
+PROMPT_LABEL = "prompt_label"
 
 # Internal experimental field.
 # It is not recommended to use for the time being. Instead, use the corresponding Set method.

examples/prompt/prompt_hub/prompt_hub_with_label.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+import json
+import os
+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:
+    """将 cozeloop 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:
+    """LLM 运行器，用于模拟 LLM 调用并设置相关的 span 标签"""
+
+    def __init__(self, client):
+        self.client = client
+
+    def llm_call(self, input_data):
+        """
+        模拟 LLM 调用并设置相关的 span 标签
+        """
+        span = self.client.start_span("llmCall", "model")
+        try:
+            # 模拟 LLM 处理过程
+            # output = ChatOpenAI().invoke(input=input_data)
+
+            # 模拟响应
+            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
+
+            # 设置 span 标签
+            span.set_input(convert_model_input(input_data))
+            span.set_output(output)
+            span.set_model_provider("openai")
+            span.set_start_time_first_resp(int(time.time() * 1000000))
+            span.set_input_tokens(input_token)
+            span.set_output_tokens(output_token)
+            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 output
+        except Exception as e:
+            raise e
+        finally:
+            span.finish()
+
+
+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_label_demo'),
+    # add the following messages to the template, and submit a version with label.
+    # 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)
+
+    try:
+        # 3. Get the prompt by prompt_key and label
+        # Note: When version is specified, label will be ignored.
+        prompt = client.get_prompt(prompt_key="prompt_hub_label_demo", label="production")
+        if prompt is not None:
+            print(f"Got prompt by label: {prompt.prompt_key}")
+
+            # 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)}")
+
+            # 4.Format messages of the prompt
+            formatted_messages = client.prompt_format(prompt, {
+                # Normal variable type should be string
+                "var1": "artificial intelligence",
+                "var2": "What is the weather like?",
+                # Placeholder variable type should be Message/List[Message]
+                "placeholder1": [Message(role=Role.USER, content="Hello!"),
+                                 Message(role=Role.ASSISTANT, content="Hello!")],
+                "placeholder2": [Message(role=Role.USER, content="Nice to meet you!")]
+                # 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)}")
+
+            # 5. Use LLM Runner to call LLM with formatted messages
+            llm_runner = LLMRunner(client)
+            result = llm_runner.llm_call(formatted_messages)
+            print(f"LLM response: {result}")
+
+        else:
+            print("Prompt not found with the specified label")
+
+    finally:
+        # 6. (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()