53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
import os
|
|
from typing import TYPE_CHECKING, Optional
|
|
|
|
if TYPE_CHECKING:
|
|
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
|
|
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
|
|
|
from litellm.types.prompts.init_prompts import SupportedPromptIntegrations
|
|
|
|
from .arize_phoenix_prompt_manager import ArizePhoenixPromptManager
|
|
|
|
# Global instances
|
|
global_arize_config: Optional[dict] = None
|
|
|
|
|
|
def prompt_initializer(
|
|
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
|
|
) -> "CustomPromptManagement":
|
|
"""
|
|
Initialize a prompt from Arize Phoenix.
|
|
"""
|
|
api_key = getattr(litellm_params, "api_key", None) or os.environ.get(
|
|
"PHOENIX_API_KEY"
|
|
)
|
|
api_base = getattr(litellm_params, "api_base", None)
|
|
prompt_id = getattr(litellm_params, "prompt_id", None)
|
|
|
|
if not api_key or not api_base:
|
|
raise ValueError(
|
|
"api_key and api_base are required for Arize Phoenix prompt integration"
|
|
)
|
|
|
|
try:
|
|
arize_prompt_manager = ArizePhoenixPromptManager(
|
|
**{
|
|
"api_key": api_key,
|
|
"api_base": api_base,
|
|
"prompt_id": prompt_id,
|
|
**litellm_params.model_dump(
|
|
exclude={"api_key", "api_base", "prompt_id"}
|
|
),
|
|
},
|
|
)
|
|
|
|
return arize_prompt_manager
|
|
except Exception as e:
|
|
raise e
|
|
|
|
|
|
prompt_initializer_registry = {
|
|
SupportedPromptIntegrations.ARIZE_PHOENIX.value: prompt_initializer,
|
|
}
|