fix types

Arize-ai · Feb 10, 2025 · 29a2fd1 · 29a2fd1
1 parent b20a7fc
commit 29a2fd1
Showing 1 changed file with 3 additions and 10 deletions.
diff --git a/src/phoenix/server/api/helpers/prompts/models.py b/src/phoenix/server/api/helpers/prompts/models.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Any, Literal, Mapping, Optional, Union
+from typing import Any, Literal, Mapping, Optional, Union, get_args
 
 from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import Annotated, TypeAlias, TypeGuard, assert_never
@@ -405,6 +405,7 @@ class PromptGeminiInvocationParameters(PromptModel):
 PromptInvocationParameters: TypeAlias = Annotated[
     Union[
         PromptOpenAIInvocationParameters,
+        PromptAzureOpenAIInvocationParameters,
         PromptAnthropicInvocationParameters,
         PromptGeminiInvocationParameters,
     ],
@@ -415,15 +416,7 @@ class PromptGeminiInvocationParameters(PromptModel):
 def is_prompt_invocation_parameters(
     invocation_parameters: Any,
 ) -> TypeGuard[PromptInvocationParameters]:
-    return isinstance(
-        invocation_parameters,
-        (
-            PromptOpenAIInvocationParameters,
-            PromptAzureOpenAIInvocationParameters,
-            PromptAnthropicInvocationParameters,
-            PromptGeminiInvocationParameters,
-        ),
-    )
+    return isinstance(invocation_parameters, get_args(PromptInvocationParameters.__origin__))
 
 
 class PromptInvocationParametersWrapper(PromptModel):