get prompt with label (#22)

Change-Id: Ief6540e9b9bc12aa326a00841dbd79c842a85d04

Author: kasarolzzw

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.