33
44import json
55from typing import Dict , Any , List , Optional
6-
76import pydantic
8- from jinja2 import Environment , BaseLoader , Undefined
9- from jinja2 . utils import missing , object_type_repr
7+
8+ from jinja2 import BaseLoader , Undefined
109from jinja2 .sandbox import SandboxedEnvironment
10+ from jinja2 .utils import missing , object_type_repr
1111
12- from cozeloop .spec .tracespec import PROMPT_KEY , INPUT , PROMPT_VERSION , V_SCENE_PROMPT_TEMPLATE , V_SCENE_PROMPT_HUB
1312from cozeloop .entities .prompt import (Prompt , Message , VariableDef , VariableType , TemplateType , Role ,
1413 PromptVariable )
1514from cozeloop .internal import consts
1918from cozeloop .internal .prompt .converter import _convert_prompt , _to_span_prompt_input , _to_span_prompt_output
2019from cozeloop .internal .prompt .openapi import OpenAPIClient , PromptQuery
2120from cozeloop .internal .trace .trace import TraceProvider
21+ from cozeloop .spec .tracespec import PROMPT_KEY , INPUT , PROMPT_VERSION , V_SCENE_PROMPT_TEMPLATE , V_SCENE_PROMPT_HUB , PROMPT_LABEL
2222
2323
2424class PromptProvider :
@@ -40,18 +40,18 @@ def __init__(
4040 auto_refresh = True )
4141 self .prompt_trace = prompt_trace
4242
43- def get_prompt (self , prompt_key : str , version : str = '' ) -> Optional [Prompt ]:
43+ def get_prompt (self , prompt_key : str , version : str = '' , label : str = '' ) -> Optional [Prompt ]:
4444 # Trace reporting
4545 if self .prompt_trace and self .trace_provider is not None :
4646 with self .trace_provider .start_span (consts .TRACE_PROMPT_HUB_SPAN_NAME ,
4747 consts .TRACE_PROMPT_HUB_SPAN_TYPE ,
4848 scene = V_SCENE_PROMPT_HUB ) as prompt_hub_pan :
4949 prompt_hub_pan .set_tags ({
5050 PROMPT_KEY : prompt_key ,
51- INPUT : json .dumps ({PROMPT_KEY : prompt_key , PROMPT_VERSION : version })
51+ INPUT : json .dumps ({PROMPT_KEY : prompt_key , PROMPT_VERSION : version , PROMPT_LABEL : label })
5252 })
5353 try :
54- prompt = self ._get_prompt (prompt_key , version )
54+ prompt = self ._get_prompt (prompt_key , version , label )
5555 if prompt is not None :
5656
5757 output = None
@@ -72,20 +72,20 @@ def get_prompt(self, prompt_key: str, version: str = '') -> Optional[Prompt]:
7272 prompt_hub_pan .set_error (str (e ))
7373 raise e
7474 else :
75- return self ._get_prompt (prompt_key , version )
75+ return self ._get_prompt (prompt_key , version , label )
7676
77- def _get_prompt (self , prompt_key : str , version : str ) -> Optional [Prompt ]:
77+ def _get_prompt (self , prompt_key : str , version : str , label : str = '' ) -> Optional [Prompt ]:
7878 """
7979 Get Prompt, prioritize retrieving from cache, if not found then fetch from server
8080 """
8181 # Try to get from cache
82- prompt = self .cache .get (prompt_key , version )
82+ prompt = self .cache .get (prompt_key , version , label )
8383 # If not in cache, fetch from server and cache it
8484 if prompt is None :
85- result = self .openapi_client .mpull_prompt (self .workspace_id , [PromptQuery (prompt_key = prompt_key , version = version )])
85+ result = self .openapi_client .mpull_prompt (self .workspace_id , [PromptQuery (prompt_key = prompt_key , version = version , label = label )])
8686 if result :
8787 prompt = _convert_prompt (result [0 ].prompt )
88- self .cache .set (prompt_key , version , prompt )
88+ self .cache .set (prompt_key , version , label , prompt )
8989 # object cache item should be read only
9090 return prompt .copy (deep = True )
9191
0 commit comments