refactor validate prompt

ddtrace/llmobs/_utils.py

@@ -1,4 +1,6 @@
+from typing import Dict
 from typing import Optional
+from typing import Union
 
 import ddtrace
 from ddtrace import Span
@@ -17,32 +19,32 @@
 log = get_logger(__name__)
 
 
-def validate_prompt(prompt: dict) -> dict:
-    validated_prompt = {}
+def validate_prompt(prompt: dict) -> Dict[str, Union[str, dict]]:
+    validated_prompt = {}  # type: Dict[str, Union[str, dict]]
     if not isinstance(prompt, dict):
         raise TypeError("Prompt must be a dictionary")
-    if prompt.get("variables"):
-        variables = prompt["variables"]
+    variables = prompt.get("variables")
+    template = prompt.get("template")
+    version = prompt.get("version")
+    prompt_id = prompt.get("id")
+    if variables is not None:
         if not isinstance(variables, dict):
             raise TypeError("Prompt variables must be a dictionary.")
         if not any(isinstance(k, str) or isinstance(v, str) for k, v in variables.items()):
             raise TypeError("Prompt variable keys and values must be strings.")
-        # for key, value in variables.items():
-        #     if not isinstance(key, str) or not isinstance(value, str):
-        #         raise TypeError("Prompt variable keys and values must be strings.")
-        validated_prompt["variables"] = prompt["variables"]
-    if prompt.get("template"):
-        if not isinstance(prompt["template"], str):
+        validated_prompt["variables"] = variables
+    if template is not None:
+        if not isinstance(template, str):
             raise TypeError("Prompt template must be a string")
-        validated_prompt["template"] = prompt["template"]
-    if prompt.get("version"):
-        if not isinstance(prompt["version"], str):
+        validated_prompt["template"] = template
+    if version is not None:
+        if not isinstance(version, str):
             raise TypeError("Prompt version must be a string.")
-        validated_prompt["version"] = prompt["version"]
-    if prompt.get("id"):
-        if not isinstance(prompt["id"], str):
+        validated_prompt["version"] = version
+    if prompt_id is not None:
+        if not isinstance(prompt_id, str):
             raise TypeError("Prompt id must be a string.")
-        validated_prompt["id"] = prompt["id"]
+        validated_prompt["id"] = prompt_id
     return validated_prompt