Update typing and add new fields to prompt_template.py

promptlayer/types/prompt_template.py

@@ -1,4 +1,6 @@
-from typing import Dict, Literal, Sequence, TypedDict, Union
+from typing import Dict, List, Literal, Sequence, TypedDict, Union
+
+from typing_extensions import NotRequired
 
 
 class GetPromptTemplate(TypedDict, total=False):
@@ -41,30 +43,34 @@ class FunctionCall(TypedDict, total=False):
 
 class SystemMessage(TypedDict, total=False):
     role: Literal["system"]
-    template_format: Union[TemplateFormat, None]
+    input_variables: NotRequired[List[str]]
+    template_format: NotRequired[TemplateFormat]
     content: Sequence[Content]
-    name: Union[str, None]
+    name: NotRequired[str]
 
 
 class UserMessage(TypedDict, total=False):
     role: Literal["user"]
-    template_format: Union[TemplateFormat, None]
+    input_variables: NotRequired[List[str]]
+    template_format: NotRequired[TemplateFormat]
     content: Sequence[Content]
-    name: Union[str, None]
+    name: NotRequired[str]
 
 
 class AssistantMessage(TypedDict, total=False):
     role: Literal["assistant"]
-    template_format: Union[TemplateFormat, None]
-    content: Union[Sequence[Content], None]
-    function_call: Union[FunctionCall, None]
-    name: Union[str, None]
+    input_variables: NotRequired[List[str]]
+    template_format: NotRequired[TemplateFormat]
+    content: NotRequired[Sequence[Content]]
+    function_call: NotRequired[FunctionCall]
+    name: NotRequired[str]
 
 
 class FunctionMessage(TypedDict, total=False):
     role: Literal["function"]
-    template_format: Union[TemplateFormat, None]
-    content: Union[Sequence[Content], None]
+    input_variables: NotRequired[List[str]]
+    template_format: NotRequired[TemplateFormat]
+    content: NotRequired[Sequence[Content]]
     name: str
 
 
@@ -77,24 +83,51 @@ class ChatFunctionCall(TypedDict, total=False):
 
 class CompletionPromptTemplate(TypedDict, total=False):
     type: Literal["completion"]
-    template_format: Union[TemplateFormat, None]
+    template_format: NotRequired[TemplateFormat]
     content: Sequence[Content]
+    input_variables: NotRequired[List[str]]
 
 
 class ChatPromptTemplate(TypedDict, total=False):
     type: Literal["chat"]
     messages: Sequence[Message]
-    functions: Union[Sequence[Function], None]
-    function_call: Union[Literal["auto", "none"], ChatFunctionCall, None]
+    functions: NotRequired[Sequence[Function]]
+    function_call: NotRequired[Union[Literal["auto", "none"], ChatFunctionCall]]
+    input_variables: NotRequired[List[str]]
 
 
 PromptTemplate = Union[CompletionPromptTemplate, ChatPromptTemplate]
 
 
-class PublishPromptTemplate(TypedDict, total=False):
+class Model(TypedDict, total=False):
+    provider: str
+    name: str
+    parameters: Dict[str, object]
+
+
+class Metadata(TypedDict, total=False):
+    model: NotRequired[Model]
+
+
+class BasePromptTemplate(TypedDict, total=False):
     prompt_name: str
+    tags: List[str]
+
+
+class PromptVersion(TypedDict, total=False):
     prompt_template: PromptTemplate
+    commit_message: NotRequired[str]
+    metadata: NotRequired[Metadata]
+
 
+class PublishPromptTemplate(BasePromptTemplate, PromptVersion):
+    pass
 
-class PublishPromptTemplateResponse(PublishPromptTemplate):
+
+class PublishPromptTemplateResponse(TypedDict, total=False):
     id: int
+    prompt_name: str
+    tags: List[str]
+    prompt_template: PromptTemplate
+    commit_message: NotRequired[str]
+    metadata: NotRequired[Metadata]

promptlayer/utils.py

@@ -563,7 +563,7 @@ def publish_prompt_template(
         response = requests.post(
             f"{URL_API_PROMPTLAYER}/rest/prompt-templates",
             headers={"X-API-KEY": get_api_key()},
-            json=body,
+            json={"prompt_template": {**body}, "prompt_version": {**body}},
         )
         if response.status_code == 400:
             raise Exception(

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "promptlayer"
-version = "0.4.6"
+version = "0.4.7"
 description = "PromptLayer is a platform for prompt engineering and tracks your LLM requests."
 authors = ["Magniv <[email protected]>"]
 license = "Apache-2.0"