diff --git a/cozeloop/_client.py b/cozeloop/_client.py index 8f67a44..a511492 100644 --- a/cozeloop/_client.py +++ b/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]: diff --git a/cozeloop/_noop.py b/cozeloop/_noop.py index 423d104..56d508d 100644 --- a/cozeloop/_noop.py +++ b/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 diff --git a/cozeloop/internal/prompt/cache.py b/cozeloop/internal/prompt/cache.py index dc3a272..7034797 100644 --- a/cozeloop/internal/prompt/cache.py +++ b/cozeloop/internal/prompt/cache.py @@ -41,15 +41,15 @@ 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) @@ -57,13 +57,13 @@ def get_all_prompt_queries(self) -> List[Tuple[str, str]]: 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}") diff --git a/cozeloop/internal/prompt/openapi.py b/cozeloop/internal/prompt/openapi.py index 7335d4b..539f73d 100644 --- a/cozeloop/internal/prompt/openapi.py +++ b/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): diff --git a/cozeloop/internal/prompt/prompt.py b/cozeloop/internal/prompt/prompt.py index fd04232..8181ba3 100644 --- a/cozeloop/internal/prompt/prompt.py +++ b/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,7 +39,7 @@ 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, @@ -47,10 +47,10 @@ def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]: 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) diff --git a/cozeloop/prompt.py b/cozeloop/prompt.py index 192be52..085b6f2 100644 --- a/cozeloop/prompt.py +++ b/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. """ diff --git a/cozeloop/spec/tracespec/span_key.py b/cozeloop/spec/tracespec/span_key.py index 3ee8002..31de969 100644 --- a/cozeloop/spec/tracespec/span_key.py +++ b/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. diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/prompt/__init__.py b/examples/prompt/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/prompt/prompt_hub/__init__.py b/examples/prompt/prompt_hub/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/prompt/prompt_hub.py b/examples/prompt/prompt_hub/prompt_hub.py similarity index 100% rename from examples/prompt/prompt_hub.py rename to examples/prompt/prompt_hub/prompt_hub.py diff --git a/examples/prompt/advance/prompt_hub_with_jinja.py b/examples/prompt/prompt_hub/prompt_hub_with_jinja.py similarity index 100% rename from examples/prompt/advance/prompt_hub_with_jinja.py rename to examples/prompt/prompt_hub/prompt_hub_with_jinja.py diff --git a/examples/prompt/prompt_hub/prompt_hub_with_label.py b/examples/prompt/prompt_hub/prompt_hub_with_label.py new file mode 100644 index 0000000..bfb8df5 --- /dev/null +++ b/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()