diff --git a/litellm/__init__.py b/litellm/__init__.py index 67f675839cd..87f075e1f13 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -30,36 +30,15 @@ from typing import ( get_args, TYPE_CHECKING, Tuple, - overload, Type, ) from litellm.types.integrations.datadog import DatadogInitParams from litellm._logging import ( set_verbose, _turn_on_debug, - verbose_logger, - json_logs, - _turn_on_json, - log_level, ) import re from litellm.constants import ( - DEFAULT_BATCH_SIZE, - DEFAULT_FLUSH_INTERVAL_SECONDS, - ROUTER_MAX_FALLBACKS, - DEFAULT_MAX_RETRIES, - DEFAULT_REPLICATE_POLLING_RETRIES, - DEFAULT_REPLICATE_POLLING_DELAY_SECONDS, - LITELLM_CHAT_PROVIDERS, - HUMANLOOP_PROMPT_CACHE_TTL_SECONDS, - OPENAI_CHAT_COMPLETION_PARAMS, - OPENAI_CHAT_COMPLETION_PARAMS as _openai_completion_params, # backwards compatibility - OPENAI_FINISH_REASONS, - OPENAI_FINISH_REASONS as _openai_finish_reasons, # backwards compatibility - openai_compatible_endpoints, - openai_compatible_providers, - openai_text_completion_compatible_providers, - _openai_like_providers, replicate_models, clarifai_models, huggingface_models, @@ -67,20 +46,15 @@ from litellm.constants import ( together_ai_models, baseten_models, WANDB_MODELS, - REPEATED_STREAMING_CHUNK_LIMIT, - request_timeout, open_ai_embedding_models, cohere_embedding_models, bedrock_embedding_models, - known_tokenizer_config, - BEDROCK_INVOKE_PROVIDERS_LITERAL, - BEDROCK_EMBEDDING_PROVIDERS_LITERAL, BEDROCK_CONVERSE_MODELS, DEFAULT_MAX_TOKENS, DEFAULT_SOFT_BUDGET, - DEFAULT_ALLOWED_FAILS, ) import httpx + # register_async_client_cleanup is lazy-loaded and called on first access litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV" @@ -145,6 +119,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "focus", "posthog", "levo", + "microsoft_purview", ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None @@ -152,7 +127,9 @@ _known_custom_logger_compatible_callbacks: List = list( get_args(_custom_logger_compatible_callbacks_literal) ) callbacks: List[ - Union[Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger"] # CustomLogger is lazy-loaded + Union[ + Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger" + ] # CustomLogger is lazy-loaded ] = [] callback_settings: Dict[str, Dict[str, Any]] = {} initialized_langfuse_clients: int = 0 @@ -162,42 +139,50 @@ prometheus_initialize_budget_metrics: Optional[bool] = False require_auth_for_metrics_endpoint: Optional[bool] = False argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload. -gcs_pub_sub_use_v1: Optional[bool] = ( - False # if you want to use v1 gcs pubsub logged payload -) -generic_api_use_v1: Optional[bool] = ( - False # if you want to use v1 generic api logged payload -) +gcs_pub_sub_use_v1: Optional[ + bool +] = False # if you want to use v1 gcs pubsub logged payload +generic_api_use_v1: Optional[ + bool +] = False # if you want to use v1 generic api logged payload argilla_transformation_object: Optional[Dict[str, Any]] = None -_async_input_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded +_async_input_callback: List[ + Union[str, Callable, "CustomLogger"] +] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. -_async_success_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded +_async_success_callback: List[ + Union[str, Callable, "CustomLogger"] +] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. -_async_failure_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded +_async_failure_callback: List[ + Union[str, Callable, "CustomLogger"] +] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] post_call_rules: List[Callable] = [] turn_off_message_logging: Optional[bool] = False -standard_logging_payload_excluded_fields: Optional[List[str]] = None # Fields to exclude from StandardLoggingPayload before callbacks receive it +standard_logging_payload_excluded_fields: Optional[ + List[str] +] = None # Fields to exclude from StandardLoggingPayload before callbacks receive it log_raw_request_response: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False filter_invalid_headers: Optional[bool] = False -add_user_information_to_llm_headers: Optional[bool] = ( - None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers -) +add_user_information_to_llm_headers: Optional[ + bool +] = None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers store_audit_logs = False # Enterprise feature, allow users to see audit logs ### end of callbacks ############# -email: Optional[str] = ( - None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -token: Optional[str] = ( - None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) +email: Optional[ + str +] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +token: Optional[ + str +] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) @@ -259,9 +244,9 @@ use_client: bool = False ssl_verify: Union[str, bool] = True ssl_security_level: Optional[str] = None ssl_certificate: Optional[str] = None -ssl_ecdh_curve: Optional[str] = ( - None # Set to 'X25519' to disable PQC and improve performance -) +ssl_ecdh_curve: Optional[ + str +] = None # Set to 'X25519' to disable PQC and improve performance disable_streaming_logging: bool = False disable_token_counter: bool = False disable_add_transform_inline_image_block: bool = False @@ -314,24 +299,20 @@ enable_loadbalancing_on_batch_endpoints: Optional[bool] = None enable_caching_on_provider_specific_optional_params: bool = ( False # feature-flag for caching on optional params - e.g. 'top_k' ) -caching: bool = ( - False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -caching_with_models: bool = ( - False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -cache: Optional["Cache"] = ( - None # cache object <- use this - https://docs.litellm.ai/docs/caching -) +caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +cache: Optional[ + "Cache" +] = None # cache object <- use this - https://docs.litellm.ai/docs/caching default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers -budget_duration: Optional[str] = ( - None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). -) +budget_duration: Optional[ + str +] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). default_soft_budget: float = ( DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 ) @@ -340,9 +321,7 @@ forward_traceparent_to_llm_provider: bool = False _current_cost = 0.0 # private variable, used if max budget is set error_logs: Dict = {} -add_function_to_prompt: bool = ( - False # if function calling not supported by api, append function call details to system prompt -) +add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt client_session: Optional[httpx.Client] = None aclient_session: Optional[httpx.AsyncClient] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' @@ -389,9 +368,7 @@ prometheus_emit_stream_label: bool = False disable_add_prefix_to_prompt: bool = ( False # used by anthropic, to disable adding prefix to prompt ) -disable_copilot_system_to_assistant: bool = ( - False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. -) +disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. public_mcp_servers: Optional[List[str]] = None public_model_groups: Optional[List[str]] = None public_agent_groups: Optional[List[str]] = None @@ -410,17 +387,13 @@ if TYPE_CHECKING: ######## Networking Settings ######## -use_aiohttp_transport: bool = ( - True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. -) +use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead disable_aiohttp_trust_env: bool = ( False # When False, aiohttp will respect HTTP(S)_PROXY env vars ) -force_ipv4: bool = ( - False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. -) +force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. network_mock: bool = False # When True, use mock transport — no real network calls ####### STOP SEQUENCE LIMIT ####### @@ -435,13 +408,13 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = ( - None # for the request overall (incl. fallbacks + model retries) -) +num_retries_per_request: Optional[ + int +] = None # for the request overall (incl. fallbacks + model retries) ####### SECRET MANAGERS ##################### -secret_manager_client: Optional[Any] = ( - None # list of instantiated key management clients - e.g. azure kv, infisical, etc. -) +secret_manager_client: Optional[ + Any +] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc. _google_kms_resource_name: Optional[str] = None _key_management_system: Optional["KeyManagementSystem"] = None # Note: KeyManagementSettings must be eagerly imported because _key_management_settings @@ -454,12 +427,12 @@ output_parse_pii: bool = False from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map model_cost = get_model_cost_map(url=model_cost_map_url) -cost_discount_config: Dict[str, float] = ( - {} -) # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount -cost_margin_config: Dict[str, Union[float, Dict[str, float]]] = ( - {} -) # Provider-specific or global cost margins. Examples: +cost_discount_config: Dict[ + str, float +] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount +cost_margin_config: Dict[ + str, Union[float, Dict[str, float]] +] = {} # Provider-specific or global cost margins. Examples: # Percentage: {"openai": 0.10} = 10% margin # Fixed: {"openai": {"fixed_amount": 0.001}} = $0.001 per request # Global: {"global": 0.05} = 5% global margin on all providers @@ -1077,7 +1050,7 @@ models_by_provider: dict = { "aws_polly": aws_polly_models, "gigachat": gigachat_models, "llamagate": llamagate_models, - "bedrock_mantle": bedrock_mantle_models + "bedrock_mantle": bedrock_mantle_models, } # mapping for those models which have larger equivalents @@ -1128,22 +1101,14 @@ openai_video_generation_models = ["sora-2"] # Import KeyManagementSettings here (before utils import) because _key_management_settings # is accessed during import time in secret_managers/main.py (via dd_tracing -> datadog -> _service_logger -> utils) from litellm.types.secret_managers.main import KeyManagementSettings + _key_management_settings: KeyManagementSettings = KeyManagementSettings() # client must be imported immediately as it's used as a decorator at function definition time -from .utils import client + # Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py # (which imports tiktoken) at import time -from .llms.custom_llm import CustomLLM -from .llms.anthropic.common_utils import AnthropicModelInfo -from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config -from .llms.deprecated_providers.palm import ( - PalmConfig, -) # here to prevent breaking changes -from .llms.deprecated_providers.aleph_alpha import AlephAlphaConfig -from .llms.gemini.common_utils import GeminiModelInfo - from .llms.vertex_ai.vertex_embeddings.transformation import ( VertexAITextEmbeddingConfig, @@ -1152,14 +1117,9 @@ from .llms.vertex_ai.vertex_embeddings.transformation import ( vertexAITextEmbeddingConfig = VertexAITextEmbeddingConfig() -from .llms.bedrock.embed.amazon_titan_v2_transformation import ( - AmazonTitanV2Config, -) -from .llms.topaz.common_utils import TopazModelInfo - # OpenAIOSeriesConfig is lazy loaded - openaiOSeriesConfig will be created on first access # OpenAIGPTConfig, OpenAIGPT5Config, etc. are lazy loaded - instances will be created on first access -from .llms.xai.common_utils import XAIModelInfo + # PublicAI now uses JSON-based configuration (see litellm/llms/openai_like/providers.json) # All remaining configs are now lazy loaded - see _lazy_imports_registry.py @@ -1171,68 +1131,7 @@ from litellm.types.utils import LlmProviders from .main import * # type: ignore # Skills API -from .skills.main import ( - create_skill, - acreate_skill, - list_skills, - alist_skills, - get_skill, - aget_skill, - delete_skill, - adelete_skill, -) -from .evals.main import ( - create_eval, - acreate_eval, - list_evals, - alist_evals, - get_eval, - aget_eval, - delete_eval, - adelete_eval, - cancel_eval, - acancel_eval, - create_run, - acreate_run, - list_runs, - alist_runs, - get_run, - aget_run, - delete_run, - adelete_run, - cancel_run, - acancel_run, -) from .integrations import * -from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients -from .exceptions import ( - AuthenticationError, - InvalidRequestError, - BadRequestError, - ImageFetchError, - NotFoundError, - PermissionDeniedError, - RateLimitError, - ServiceUnavailableError, - BadGatewayError, - OpenAIError, - ContextWindowExceededError, - ContentPolicyViolationError, - BudgetExceededError, - APIError, - Timeout, - APIConnectionError, - UnsupportedParamsError, - APIResponseValidationError, - UnprocessableEntityError, - InternalServerError, - JSONSchemaValidationError, - LITELLM_EXCEPTION_TYPES, - MockException, -) -from .budget_manager import BudgetManager -from .proxy.proxy_cli import run_server -from .router import Router from .assistants.main import * from .batches.main import * from .images.main import * @@ -1241,46 +1140,19 @@ from .batch_completion.main import * # type: ignore from .rerank_api.main import * from .llms.anthropic.experimental_pass_through.messages.handler import * from .responses.main import * + # Interactions API is available as litellm.interactions module # Usage: litellm.interactions.create(), litellm.interactions.get(), etc. -from . import interactions -from .skills.main import ( - create_skill, - acreate_skill, - list_skills, - alist_skills, - get_skill, - aget_skill, - delete_skill, - adelete_skill, -) from .containers.main import * from .ocr.main import * from .rag.main import * from .search.main import * -from .realtime_api.main import _arealtime -from .responses.main import _aresponses_websocket from .fine_tuning.main import * from .files.main import * -from .vector_store_files.main import ( - acreate as avector_store_file_create, - adelete as avector_store_file_delete, - alist as avector_store_file_list, - aretrieve as avector_store_file_retrieve, - aretrieve_content as avector_store_file_content, - aupdate as avector_store_file_update, - create as vector_store_file_create, - delete as vector_store_file_delete, - list as vector_store_file_list, - retrieve as vector_store_file_retrieve, - retrieve_content as vector_store_file_content, - update as vector_store_file_update, -) from .scheduler import * ### ADAPTERS ### from .types.adapter import AdapterItem -import litellm.anthropic_interface as anthropic adapters: List[AdapterItem] = [] @@ -1294,26 +1166,22 @@ vector_store_registry: Optional[VectorStoreRegistry] = None vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None ### RAG ### -from . import rag ### CUSTOM LLMs ### from .types.llms.custom_llm import CustomLLMItem custom_provider_map: List[CustomLLMItem] = [] -_custom_providers: List[str] = ( - [] -) # internal helper util, used to track names of custom providers -disable_hf_tokenizer_download: Optional[bool] = ( - None # disable huggingface tokenizer download. Defaults to openai clk100 -) +_custom_providers: List[ + str +] = [] # internal helper util, used to track names of custom providers +disable_hf_tokenizer_download: Optional[ + bool +] = None # disable huggingface tokenizer download. Defaults to openai clk100 global_disable_no_log_param: bool = False ### CLI UTILITIES ### -from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key ### PASSTHROUGH ### -from .passthrough import allm_passthrough_route, llm_passthrough_route -from .google_genai import agenerate_content ### GLOBAL CONFIG ### global_bitbucket_config: Optional[Dict[str, Any]] = None @@ -1344,131 +1212,323 @@ if TYPE_CHECKING: from litellm.caching.caching import Cache # Type stubs for lazy-loaded configs to help mypy - from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig as AmazonConverseConfig - from .llms.openai_like.chat.handler import OpenAILikeChatConfig as OpenAILikeChatConfig - from .llms.galadriel.chat.transformation import GaladrielChatConfig as GaladrielChatConfig + from .llms.bedrock.chat.converse_transformation import ( + AmazonConverseConfig as AmazonConverseConfig, + ) + from .llms.openai_like.chat.handler import ( + OpenAILikeChatConfig as OpenAILikeChatConfig, + ) + from .llms.galadriel.chat.transformation import ( + GaladrielChatConfig as GaladrielChatConfig, + ) from .llms.github.chat.transformation import GithubChatConfig as GithubChatConfig - from .llms.azure_ai.anthropic.transformation import AzureAnthropicConfig as AzureAnthropicConfig + from .llms.azure_ai.anthropic.transformation import ( + AzureAnthropicConfig as AzureAnthropicConfig, + ) from .llms.bytez.chat.transformation import BytezChatConfig as BytezChatConfig - from .llms.compactifai.chat.transformation import CompactifAIChatConfig as CompactifAIChatConfig + from .llms.compactifai.chat.transformation import ( + CompactifAIChatConfig as CompactifAIChatConfig, + ) from .llms.empower.chat.transformation import EmpowerChatConfig as EmpowerChatConfig from .llms.minimax.chat.transformation import MinimaxChatConfig as MinimaxChatConfig - from .llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig as AiohttpOpenAIChatConfig - from .llms.huggingface.chat.transformation import HuggingFaceChatConfig as HuggingFaceChatConfig - from .llms.huggingface.embedding.transformation import HuggingFaceEmbeddingConfig as HuggingFaceEmbeddingConfig + from .llms.aiohttp_openai.chat.transformation import ( + AiohttpOpenAIChatConfig as AiohttpOpenAIChatConfig, + ) + from .llms.huggingface.chat.transformation import ( + HuggingFaceChatConfig as HuggingFaceChatConfig, + ) + from .llms.huggingface.embedding.transformation import ( + HuggingFaceEmbeddingConfig as HuggingFaceEmbeddingConfig, + ) from .llms.oobabooga.chat.transformation import OobaboogaConfig as OobaboogaConfig from .llms.maritalk import MaritalkConfig as MaritalkConfig - from .llms.openrouter.chat.transformation import OpenrouterConfig as OpenrouterConfig + from .llms.openrouter.chat.transformation import ( + OpenrouterConfig as OpenrouterConfig, + ) from .llms.datarobot.chat.transformation import DataRobotConfig as DataRobotConfig from .llms.anthropic.chat.transformation import AnthropicConfig as AnthropicConfig - from .llms.anthropic.completion.transformation import AnthropicTextConfig as AnthropicTextConfig + from .llms.anthropic.completion.transformation import ( + AnthropicTextConfig as AnthropicTextConfig, + ) from .llms.groq.stt.transformation import GroqSTTConfig as GroqSTTConfig from .llms.triton.completion.transformation import TritonConfig as TritonConfig - from .llms.triton.completion.transformation import TritonGenerateConfig as TritonGenerateConfig - from .llms.triton.completion.transformation import TritonInferConfig as TritonInferConfig - from .llms.triton.embedding.transformation import TritonEmbeddingConfig as TritonEmbeddingConfig - from .llms.huggingface.rerank.transformation import HuggingFaceRerankConfig as HuggingFaceRerankConfig - from .llms.databricks.chat.transformation import DatabricksConfig as DatabricksConfig - from .llms.databricks.embed.transformation import DatabricksEmbeddingConfig as DatabricksEmbeddingConfig + from .llms.triton.completion.transformation import ( + TritonGenerateConfig as TritonGenerateConfig, + ) + from .llms.triton.completion.transformation import ( + TritonInferConfig as TritonInferConfig, + ) + from .llms.triton.embedding.transformation import ( + TritonEmbeddingConfig as TritonEmbeddingConfig, + ) + from .llms.huggingface.rerank.transformation import ( + HuggingFaceRerankConfig as HuggingFaceRerankConfig, + ) + from .llms.databricks.chat.transformation import ( + DatabricksConfig as DatabricksConfig, + ) + from .llms.databricks.embed.transformation import ( + DatabricksEmbeddingConfig as DatabricksEmbeddingConfig, + ) from .llms.predibase.chat.transformation import PredibaseConfig as PredibaseConfig from .llms.replicate.chat.transformation import ReplicateConfig as ReplicateConfig from .llms.snowflake.chat.transformation import SnowflakeConfig as SnowflakeConfig - from .llms.cohere.rerank.transformation import CohereRerankConfig as CohereRerankConfig - from .llms.cohere.rerank_v2.transformation import CohereRerankV2Config as CohereRerankV2Config - from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig as AzureAIRerankConfig - from .llms.infinity.rerank.transformation import InfinityRerankConfig as InfinityRerankConfig - from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig as JinaAIRerankConfig - from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig as DeepinfraRerankConfig - from .llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig as HostedVLLMRerankConfig - from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig as NvidiaNimRerankConfig - from .llms.nvidia_nim.rerank.ranking_transformation import NvidiaNimRankingConfig as NvidiaNimRankingConfig - from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as VertexAIRerankConfig - from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig as FireworksAIRerankConfig - from .llms.voyage.rerank.transformation import VoyageRerankConfig as VoyageRerankConfig - from .llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig as IBMWatsonXRerankConfig + from .llms.cohere.rerank.transformation import ( + CohereRerankConfig as CohereRerankConfig, + ) + from .llms.cohere.rerank_v2.transformation import ( + CohereRerankV2Config as CohereRerankV2Config, + ) + from .llms.azure_ai.rerank.transformation import ( + AzureAIRerankConfig as AzureAIRerankConfig, + ) + from .llms.infinity.rerank.transformation import ( + InfinityRerankConfig as InfinityRerankConfig, + ) + from .llms.jina_ai.rerank.transformation import ( + JinaAIRerankConfig as JinaAIRerankConfig, + ) + from .llms.deepinfra.rerank.transformation import ( + DeepinfraRerankConfig as DeepinfraRerankConfig, + ) + from .llms.hosted_vllm.rerank.transformation import ( + HostedVLLMRerankConfig as HostedVLLMRerankConfig, + ) + from .llms.nvidia_nim.rerank.transformation import ( + NvidiaNimRerankConfig as NvidiaNimRerankConfig, + ) + from .llms.nvidia_nim.rerank.ranking_transformation import ( + NvidiaNimRankingConfig as NvidiaNimRankingConfig, + ) + from .llms.vertex_ai.rerank.transformation import ( + VertexAIRerankConfig as VertexAIRerankConfig, + ) + from .llms.fireworks_ai.rerank.transformation import ( + FireworksAIRerankConfig as FireworksAIRerankConfig, + ) + from .llms.voyage.rerank.transformation import ( + VoyageRerankConfig as VoyageRerankConfig, + ) + from .llms.watsonx.rerank.transformation import ( + IBMWatsonXRerankConfig as IBMWatsonXRerankConfig, + ) from .llms.clarifai.chat.transformation import ClarifaiConfig as ClarifaiConfig from .llms.ai21.chat.transformation import AI21ChatConfig as AI21ChatConfig from .llms.meta_llama.chat.transformation import LlamaAPIConfig as LlamaAPIConfig - from .llms.together_ai.completion.transformation import TogetherAITextCompletionConfig as TogetherAITextCompletionConfig - from .llms.cloudflare.chat.transformation import CloudflareChatConfig as CloudflareChatConfig + from .llms.together_ai.completion.transformation import ( + TogetherAITextCompletionConfig as TogetherAITextCompletionConfig, + ) + from .llms.cloudflare.chat.transformation import ( + CloudflareChatConfig as CloudflareChatConfig, + ) from .llms.novita.chat.transformation import NovitaConfig as NovitaConfig from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig from .llms.ollama.completion.transformation import OllamaConfig as OllamaConfig - from .llms.sagemaker.completion.transformation import SagemakerConfig as SagemakerConfig - from .llms.sagemaker.chat.transformation import SagemakerChatConfig as SagemakerChatConfig + from .llms.sagemaker.completion.transformation import ( + SagemakerConfig as SagemakerConfig, + ) + from .llms.sagemaker.chat.transformation import ( + SagemakerChatConfig as SagemakerChatConfig, + ) from .llms.cohere.chat.transformation import CohereChatConfig as CohereChatConfig - from .llms.anthropic.experimental_pass_through.messages.transformation import AnthropicMessagesConfig as AnthropicMessagesConfig - from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig + from .llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig as AnthropicMessagesConfig, + ) + from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig, + ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig - from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig as VertexGeminiConfig - from .llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig as GoogleAIStudioGeminiConfig - from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig as VertexAIAnthropicConfig - from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import VertexAILlama3Config as VertexAILlama3Config - from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import VertexAIAi21Config as VertexAIAi21Config - from .llms.bedrock.chat.invoke_handler import AmazonCohereChatConfig as AmazonCohereChatConfig - from .llms.bedrock.common_utils import AmazonBedrockGlobalConfig as AmazonBedrockGlobalConfig - from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import AmazonAI21Config as AmazonAI21Config - from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import AmazonInvokeNovaConfig as AmazonInvokeNovaConfig - from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import AmazonQwen2Config as AmazonQwen2Config - from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import AmazonQwen3Config as AmazonQwen3Config - from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import AmazonAnthropicConfig as AmazonAnthropicConfig - from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig - from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import AmazonCohereConfig as AmazonCohereConfig - from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import AmazonLlamaConfig as AmazonLlamaConfig - from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import AmazonDeepSeekR1Config as AmazonDeepSeekR1Config - from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import AmazonMistralConfig as AmazonMistralConfig - from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import AmazonMoonshotConfig as AmazonMoonshotConfig - from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import AmazonTitanConfig as AmazonTitanConfig - from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig - from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig as AmazonInvokeConfig - from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig - from .llms.bedrock.image_generation.amazon_stability1_transformation import AmazonStabilityConfig as AmazonStabilityConfig - from .llms.bedrock.image_generation.amazon_stability3_transformation import AmazonStability3Config as AmazonStability3Config - from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig as AmazonNovaCanvasConfig - from .llms.bedrock.embed.amazon_titan_g1_transformation import AmazonTitanG1Config as AmazonTitanG1Config - from .llms.bedrock.embed.amazon_titan_multimodal_transformation import AmazonTitanMultimodalEmbeddingG1Config as AmazonTitanMultimodalEmbeddingG1Config - from .llms.cohere.chat.v2_transformation import CohereV2ChatConfig as CohereV2ChatConfig - from .llms.bedrock.embed.cohere_transformation import BedrockCohereEmbeddingConfig as BedrockCohereEmbeddingConfig - from .llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig as TwelveLabsMarengoEmbeddingConfig - from .llms.bedrock.embed.amazon_nova_transformation import AmazonNovaEmbeddingConfig as AmazonNovaEmbeddingConfig - from .llms.openai.openai import OpenAIConfig as OpenAIConfig, MistralEmbeddingConfig as MistralEmbeddingConfig - from .llms.openai.image_variations.transformation import OpenAIImageVariationConfig as OpenAIImageVariationConfig - from .llms.deepgram.audio_transcription.transformation import DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig - from .llms.topaz.image_variations.transformation import TopazImageVariationConfig as TopazImageVariationConfig - from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig as OpenAITextCompletionConfig + from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig as VertexGeminiConfig, + ) + from .llms.gemini.chat.transformation import ( + GoogleAIStudioGeminiConfig as GoogleAIStudioGeminiConfig, + ) + from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( + VertexAIAnthropicConfig as VertexAIAnthropicConfig, + ) + from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import ( + VertexAILlama3Config as VertexAILlama3Config, + ) + from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( + VertexAIAi21Config as VertexAIAi21Config, + ) + from .llms.bedrock.chat.invoke_handler import ( + AmazonCohereChatConfig as AmazonCohereChatConfig, + ) + from .llms.bedrock.common_utils import ( + AmazonBedrockGlobalConfig as AmazonBedrockGlobalConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import ( + AmazonAI21Config as AmazonAI21Config, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import ( + AmazonInvokeNovaConfig as AmazonInvokeNovaConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import ( + AmazonQwen2Config as AmazonQwen2Config, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( + AmazonQwen3Config as AmazonQwen3Config, + ) + from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import ( + AmazonAnthropicConfig as AmazonAnthropicConfig, + ) + from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import ( + AmazonCohereConfig as AmazonCohereConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import ( + AmazonLlamaConfig as AmazonLlamaConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import ( + AmazonDeepSeekR1Config as AmazonDeepSeekR1Config, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import ( + AmazonMistralConfig as AmazonMistralConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig as AmazonMoonshotConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import ( + AmazonTitanConfig as AmazonTitanConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import ( + AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig, + ) + from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig as AmazonInvokeConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig, + ) + from .llms.bedrock.image_generation.amazon_stability1_transformation import ( + AmazonStabilityConfig as AmazonStabilityConfig, + ) + from .llms.bedrock.image_generation.amazon_stability3_transformation import ( + AmazonStability3Config as AmazonStability3Config, + ) + from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import ( + AmazonNovaCanvasConfig as AmazonNovaCanvasConfig, + ) + from .llms.bedrock.embed.amazon_titan_g1_transformation import ( + AmazonTitanG1Config as AmazonTitanG1Config, + ) + from .llms.bedrock.embed.amazon_titan_multimodal_transformation import ( + AmazonTitanMultimodalEmbeddingG1Config as AmazonTitanMultimodalEmbeddingG1Config, + ) + from .llms.cohere.chat.v2_transformation import ( + CohereV2ChatConfig as CohereV2ChatConfig, + ) + from .llms.bedrock.embed.cohere_transformation import ( + BedrockCohereEmbeddingConfig as BedrockCohereEmbeddingConfig, + ) + from .llms.bedrock.embed.twelvelabs_marengo_transformation import ( + TwelveLabsMarengoEmbeddingConfig as TwelveLabsMarengoEmbeddingConfig, + ) + from .llms.bedrock.embed.amazon_nova_transformation import ( + AmazonNovaEmbeddingConfig as AmazonNovaEmbeddingConfig, + ) + from .llms.openai.openai import ( + OpenAIConfig as OpenAIConfig, + MistralEmbeddingConfig as MistralEmbeddingConfig, + ) + from .llms.openai.image_variations.transformation import ( + OpenAIImageVariationConfig as OpenAIImageVariationConfig, + ) + from .llms.deepgram.audio_transcription.transformation import ( + DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig, + ) + from .llms.topaz.image_variations.transformation import ( + TopazImageVariationConfig as TopazImageVariationConfig, + ) + from litellm.llms.openai.completion.transformation import ( + OpenAITextCompletionConfig as OpenAITextCompletionConfig, + ) from .llms.groq.chat.transformation import GroqChatConfig as GroqChatConfig - from .llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig as BedrockMantleChatConfig + from .llms.bedrock_mantle.chat.transformation import ( + BedrockMantleChatConfig as BedrockMantleChatConfig, + ) from .llms.a2a.chat.transformation import A2AConfig as A2AConfig - from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig - from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig - from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig - from .llms.perplexity.embedding.transformation import PerplexityEmbeddingConfig as PerplexityEmbeddingConfig - from .llms.azure_ai.chat.transformation import AzureAIStudioConfig as AzureAIStudioConfig + from .llms.voyage.embedding.transformation import ( + VoyageEmbeddingConfig as VoyageEmbeddingConfig, + ) + from .llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig, + ) + from .llms.infinity.embedding.transformation import ( + InfinityEmbeddingConfig as InfinityEmbeddingConfig, + ) + from .llms.perplexity.embedding.transformation import ( + PerplexityEmbeddingConfig as PerplexityEmbeddingConfig, + ) + from .llms.azure_ai.chat.transformation import ( + AzureAIStudioConfig as AzureAIStudioConfig, + ) from .llms.mistral.chat.transformation import MistralConfig as MistralConfig - from .llms.openai.responses.transformation import OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig - from .llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig as AzureOpenAIResponsesAPIConfig - from .llms.azure.responses.o_series_transformation import AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig - from .llms.xai.responses.transformation import XAIResponsesAPIConfig as XAIResponsesAPIConfig - from .llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig - from .llms.volcengine.responses.transformation import VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig - from .llms.manus.responses.transformation import ManusResponsesAPIConfig as ManusResponsesAPIConfig - from .llms.perplexity.responses.transformation import PerplexityResponsesConfig as PerplexityResponsesConfig - from .llms.databricks.responses.transformation import DatabricksResponsesAPIConfig as DatabricksResponsesAPIConfig - from .llms.openrouter.responses.transformation import OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig - from .llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig - from .llms.openai.chat.o_series_transformation import OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config - from .llms.anthropic.skills.transformation import AnthropicSkillsConfig as AnthropicSkillsConfig - from .llms.base_llm.skills.transformation import BaseSkillsAPIConfig as BaseSkillsAPIConfig - from .llms.gradient_ai.chat.transformation import GradientAIConfig as GradientAIConfig + from .llms.openai.responses.transformation import ( + OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig, + ) + from .llms.azure.responses.transformation import ( + AzureOpenAIResponsesAPIConfig as AzureOpenAIResponsesAPIConfig, + ) + from .llms.azure.responses.o_series_transformation import ( + AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig, + ) + from .llms.xai.responses.transformation import ( + XAIResponsesAPIConfig as XAIResponsesAPIConfig, + ) + from .llms.litellm_proxy.responses.transformation import ( + LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig, + ) + from .llms.volcengine.responses.transformation import ( + VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig, + ) + from .llms.manus.responses.transformation import ( + ManusResponsesAPIConfig as ManusResponsesAPIConfig, + ) + from .llms.perplexity.responses.transformation import ( + PerplexityResponsesConfig as PerplexityResponsesConfig, + ) + from .llms.databricks.responses.transformation import ( + DatabricksResponsesAPIConfig as DatabricksResponsesAPIConfig, + ) + from .llms.openrouter.responses.transformation import ( + OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig, + ) + from .llms.gemini.interactions.transformation import ( + GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, + ) + from .llms.anthropic.skills.transformation import ( + AnthropicSkillsConfig as AnthropicSkillsConfig, + ) + from .llms.base_llm.skills.transformation import ( + BaseSkillsAPIConfig as BaseSkillsAPIConfig, + ) + from .llms.gradient_ai.chat.transformation import ( + GradientAIConfig as GradientAIConfig, + ) from .llms.openai.chat.gpt_transformation import OpenAIGPTConfig as OpenAIGPTConfig - from .llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config as OpenAIGPT5Config - from .llms.openai.transcriptions.whisper_transformation import OpenAIWhisperAudioTranscriptionConfig as OpenAIWhisperAudioTranscriptionConfig - from .llms.openai.transcriptions.gpt_transformation import OpenAIGPTAudioTranscriptionConfig as OpenAIGPTAudioTranscriptionConfig - from .llms.openai.chat.gpt_audio_transformation import OpenAIGPTAudioConfig as OpenAIGPTAudioConfig + from .llms.openai.chat.gpt_5_transformation import ( + OpenAIGPT5Config as OpenAIGPT5Config, + ) + from .llms.openai.transcriptions.whisper_transformation import ( + OpenAIWhisperAudioTranscriptionConfig as OpenAIWhisperAudioTranscriptionConfig, + ) + from .llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTAudioTranscriptionConfig as OpenAIGPTAudioTranscriptionConfig, + ) + from .llms.openai.chat.gpt_audio_transformation import ( + OpenAIGPTAudioConfig as OpenAIGPTAudioConfig, + ) from .llms.nvidia_nim.chat.transformation import NvidiaNimConfig as NvidiaNimConfig - from .llms.nvidia_nim.embed import NvidiaNimEmbeddingConfig as NvidiaNimEmbeddingConfig + from .llms.nvidia_nim.embed import ( + NvidiaNimEmbeddingConfig as NvidiaNimEmbeddingConfig, + ) # Type stubs for lazy-loaded config instances openaiOSeriesConfig: OpenAIOSeriesConfig @@ -1480,21 +1540,47 @@ if TYPE_CHECKING: # Import config classes that need type stubs (for mypy) - import with _ prefix to avoid circular reference from .llms.vllm.completion.transformation import VLLMConfig as _VLLMConfig - from .llms.deepseek.chat.transformation import DeepSeekChatConfig as _DeepSeekChatConfig - from .llms.sap.chat.transformation import GenAIHubOrchestrationConfig as _GenAIHubOrchestrationConfig - from .llms.sap.embed.transformation import GenAIHubEmbeddingConfig as _GenAIHubEmbeddingConfig - from .llms.azure.chat.o_series_transformation import AzureOpenAIO1Config as _AzureOpenAIO1Config - from .llms.perplexity.chat.transformation import PerplexityChatConfig as _PerplexityChatConfig + from .llms.deepseek.chat.transformation import ( + DeepSeekChatConfig as _DeepSeekChatConfig, + ) + from .llms.sap.chat.transformation import ( + GenAIHubOrchestrationConfig as _GenAIHubOrchestrationConfig, + ) + from .llms.sap.embed.transformation import ( + GenAIHubEmbeddingConfig as _GenAIHubEmbeddingConfig, + ) + from .llms.azure.chat.o_series_transformation import ( + AzureOpenAIO1Config as _AzureOpenAIO1Config, + ) + from .llms.perplexity.chat.transformation import ( + PerplexityChatConfig as _PerplexityChatConfig, + ) from .llms.nscale.chat.transformation import NscaleConfig as _NscaleConfig - from .llms.watsonx.chat.transformation import IBMWatsonXChatConfig as _IBMWatsonXChatConfig - from .llms.watsonx.completion.transformation import IBMWatsonXAIConfig as _IBMWatsonXAIConfig - from .llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfig as _LiteLLMProxyChatConfig + from .llms.watsonx.chat.transformation import ( + IBMWatsonXChatConfig as _IBMWatsonXChatConfig, + ) + from .llms.watsonx.completion.transformation import ( + IBMWatsonXAIConfig as _IBMWatsonXAIConfig, + ) + from .llms.litellm_proxy.chat.transformation import ( + LiteLLMProxyChatConfig as _LiteLLMProxyChatConfig, + ) from .llms.deepinfra.chat.transformation import DeepInfraConfig as _DeepInfraConfig - from .llms.llamafile.chat.transformation import LlamafileChatConfig as _LlamafileChatConfig - from .llms.lm_studio.chat.transformation import LMStudioChatConfig as _LMStudioChatConfig - from .llms.lm_studio.embed.transformation import LmStudioEmbeddingConfig as _LmStudioEmbeddingConfig - from .llms.watsonx.embed.transformation import IBMWatsonXEmbeddingConfig as _IBMWatsonXEmbeddingConfig - from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig as _VertexGeminiConfig + from .llms.llamafile.chat.transformation import ( + LlamafileChatConfig as _LlamafileChatConfig, + ) + from .llms.lm_studio.chat.transformation import ( + LMStudioChatConfig as _LMStudioChatConfig, + ) + from .llms.lm_studio.embed.transformation import ( + LmStudioEmbeddingConfig as _LmStudioEmbeddingConfig, + ) + from .llms.watsonx.embed.transformation import ( + IBMWatsonXEmbeddingConfig as _IBMWatsonXEmbeddingConfig, + ) + from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig as _VertexGeminiConfig, + ) # Type stubs for lazy-loaded config classes (to help mypy understand types) VLLMConfig: Type[_VLLMConfig] @@ -1514,56 +1600,121 @@ if TYPE_CHECKING: IBMWatsonXEmbeddingConfig: Type[_IBMWatsonXEmbeddingConfig] VertexAIConfig: Type[_VertexGeminiConfig] # Alias for VertexGeminiConfig - from .llms.featherless_ai.chat.transformation import FeatherlessAIConfig as FeatherlessAIConfig + from .llms.featherless_ai.chat.transformation import ( + FeatherlessAIConfig as FeatherlessAIConfig, + ) from .llms.cerebras.chat import CerebrasConfig as CerebrasConfig from .llms.baseten.chat import BasetenConfig as BasetenConfig from .llms.sambanova.chat import SambanovaConfig as SambanovaConfig - from .llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig as SambaNovaEmbeddingConfig - from .llms.fireworks_ai.chat.transformation import FireworksAIConfig as FireworksAIConfig - from .llms.fireworks_ai.completion.transformation import FireworksAITextCompletionConfig as FireworksAITextCompletionConfig - from .llms.fireworks_ai.audio_transcription.transformation import FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig - from .llms.fireworks_ai.embed.fireworks_ai_transformation import FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig - from .llms.friendliai.chat.transformation import FriendliaiChatConfig as FriendliaiChatConfig - from .llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig as JinaAIEmbeddingConfig + from .llms.sambanova.embedding.transformation import ( + SambaNovaEmbeddingConfig as SambaNovaEmbeddingConfig, + ) + from .llms.fireworks_ai.chat.transformation import ( + FireworksAIConfig as FireworksAIConfig, + ) + from .llms.fireworks_ai.completion.transformation import ( + FireworksAITextCompletionConfig as FireworksAITextCompletionConfig, + ) + from .llms.fireworks_ai.audio_transcription.transformation import ( + FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig, + ) + from .llms.fireworks_ai.embed.fireworks_ai_transformation import ( + FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig, + ) + from .llms.friendliai.chat.transformation import ( + FriendliaiChatConfig as FriendliaiChatConfig, + ) + from .llms.jina_ai.embedding.transformation import ( + JinaAIEmbeddingConfig as JinaAIEmbeddingConfig, + ) from .llms.xai.chat.transformation import XAIChatConfig as XAIChatConfig from .llms.zai.chat.transformation import ZAIChatConfig as ZAIChatConfig from .llms.aiml.chat.transformation import AIMLChatConfig as AIMLChatConfig - from .llms.volcengine.chat.transformation import VolcEngineChatConfig as VolcEngineChatConfig, VolcEngineChatConfig as VolcEngineConfig - from .llms.codestral.completion.transformation import CodestralTextCompletionConfig as CodestralTextCompletionConfig - from .llms.azure.azure import AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig + from .llms.codestral.completion.transformation import ( + CodestralTextCompletionConfig as CodestralTextCompletionConfig, + ) + from .llms.azure.azure import ( + AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig, + ) from .llms.heroku.chat.transformation import HerokuChatConfig as HerokuChatConfig from .llms.cometapi.chat.transformation import CometAPIConfig as CometAPIConfig - from .llms.azure.chat.gpt_transformation import AzureOpenAIConfig as AzureOpenAIConfig - from .llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config as AzureOpenAIGPT5Config - from .llms.azure.completion.transformation import AzureOpenAITextConfig as AzureOpenAITextConfig - from .llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig as HostedVLLMChatConfig - from .llms.hosted_vllm.embedding.transformation import HostedVLLMEmbeddingConfig as HostedVLLMEmbeddingConfig - from .llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig - from .llms.github_copilot.chat.transformation import GithubCopilotConfig as GithubCopilotConfig - from .llms.github_copilot.responses.transformation import GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig - from .llms.github_copilot.embedding.transformation import GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig + from .llms.azure.chat.gpt_transformation import ( + AzureOpenAIConfig as AzureOpenAIConfig, + ) + from .llms.azure.chat.gpt_5_transformation import ( + AzureOpenAIGPT5Config as AzureOpenAIGPT5Config, + ) + from .llms.azure.completion.transformation import ( + AzureOpenAITextConfig as AzureOpenAITextConfig, + ) + from .llms.hosted_vllm.chat.transformation import ( + HostedVLLMChatConfig as HostedVLLMChatConfig, + ) + from .llms.hosted_vllm.embedding.transformation import ( + HostedVLLMEmbeddingConfig as HostedVLLMEmbeddingConfig, + ) + from .llms.hosted_vllm.responses.transformation import ( + HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig, + ) + from .llms.github_copilot.chat.transformation import ( + GithubCopilotConfig as GithubCopilotConfig, + ) + from .llms.github_copilot.responses.transformation import ( + GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig, + ) + from .llms.github_copilot.embedding.transformation import ( + GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig, + ) from .llms.chatgpt.chat.transformation import ChatGPTConfig as ChatGPTConfig - from .llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig + from .llms.chatgpt.responses.transformation import ( + ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig, + ) from .llms.gigachat.chat.transformation import GigaChatConfig as GigaChatConfig - from .llms.gigachat.embedding.transformation import GigaChatEmbeddingConfig as GigaChatEmbeddingConfig + from .llms.gigachat.embedding.transformation import ( + GigaChatEmbeddingConfig as GigaChatEmbeddingConfig, + ) from .llms.nebius.chat.transformation import NebiusConfig as NebiusConfig from .llms.wandb.chat.transformation import WandbConfig as WandbConfig - from .llms.dashscope.chat.transformation import DashScopeChatConfig as DashScopeChatConfig - from .llms.moonshot.chat.transformation import MoonshotChatConfig as MoonshotChatConfig - from .llms.docker_model_runner.chat.transformation import DockerModelRunnerChatConfig as DockerModelRunnerChatConfig + from .llms.dashscope.chat.transformation import ( + DashScopeChatConfig as DashScopeChatConfig, + ) + from .llms.moonshot.chat.transformation import ( + MoonshotChatConfig as MoonshotChatConfig, + ) + from .llms.docker_model_runner.chat.transformation import ( + DockerModelRunnerChatConfig as DockerModelRunnerChatConfig, + ) from .llms.v0.chat.transformation import V0ChatConfig as V0ChatConfig from .llms.oci.chat.transformation import OCIChatConfig as OCIChatConfig from .llms.morph.chat.transformation import MorphChatConfig as MorphChatConfig from .llms.ragflow.chat.transformation import RAGFlowConfig as RAGFlowConfig - from .llms.lambda_ai.chat.transformation import LambdaAIChatConfig as LambdaAIChatConfig - from .llms.hyperbolic.chat.transformation import HyperbolicChatConfig as HyperbolicChatConfig - from .llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig as VercelAIGatewayConfig - from .llms.ovhcloud.chat.transformation import OVHCloudChatConfig as OVHCloudChatConfig - from .llms.ovhcloud.embedding.transformation import OVHCloudEmbeddingConfig as OVHCloudEmbeddingConfig - from .llms.cometapi.embed.transformation import CometAPIEmbeddingConfig as CometAPIEmbeddingConfig - from .llms.lemonade.chat.transformation import LemonadeChatConfig as LemonadeChatConfig - from .llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig - from .llms.amazon_nova.chat.transformation import AmazonNovaChatConfig as AmazonNovaChatConfig + from .llms.lambda_ai.chat.transformation import ( + LambdaAIChatConfig as LambdaAIChatConfig, + ) + from .llms.hyperbolic.chat.transformation import ( + HyperbolicChatConfig as HyperbolicChatConfig, + ) + from .llms.vercel_ai_gateway.chat.transformation import ( + VercelAIGatewayConfig as VercelAIGatewayConfig, + ) + from .llms.ovhcloud.chat.transformation import ( + OVHCloudChatConfig as OVHCloudChatConfig, + ) + from .llms.ovhcloud.embedding.transformation import ( + OVHCloudEmbeddingConfig as OVHCloudEmbeddingConfig, + ) + from .llms.cometapi.embed.transformation import ( + CometAPIEmbeddingConfig as CometAPIEmbeddingConfig, + ) + from .llms.lemonade.chat.transformation import ( + LemonadeChatConfig as LemonadeChatConfig, + ) + from .llms.snowflake.embedding.transformation import ( + SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig, + ) + from .llms.amazon_nova.chat.transformation import ( + AmazonNovaChatConfig as AmazonNovaChatConfig, + ) from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES from litellm.types.utils import ( @@ -1624,10 +1775,10 @@ if TYPE_CHECKING: # Bedrock tool name mappings instance (lazy-loaded) from litellm.caching.caching import InMemoryCache + bedrock_tool_name_mappings: InMemoryCache # Azure exception class (lazy-loaded) - from litellm.llms.azure.common_utils import AzureOpenAIError # Secret manager types (lazy-loaded) from litellm.types.secret_managers.main import ( @@ -1642,11 +1793,15 @@ if TYPE_CHECKING: from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams # Logging callback manager class and instance (lazy-loaded) - from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager + from litellm.litellm_core_utils.logging_callback_manager import ( + LoggingCallbackManager, + ) + logging_callback_manager: LoggingCallbackManager # provider_list is lazy-loaded from litellm.types.utils import LlmProviders + provider_list: List[Union[LlmProviders, str]] # Note: AmazonConverseConfig and OpenAILikeChatConfig are imported above in TYPE_CHECKING block @@ -1663,7 +1818,7 @@ _async_client_cleanup_registered = False if os.getenv("LITELLM_DISABLE_LAZY_LOADING", "").lower() in ("1", "true", "yes", "on"): # Load encoding at import time (pre-#18070 behavior) # This ensures encoding is initialized before VCR starts recording - from .main import encoding + pass def __getattr__(name: str) -> Any: @@ -1671,7 +1826,10 @@ def __getattr__(name: str) -> Any: global _async_client_cleanup_registered # Register async client cleanup on first access (only once) if not _async_client_cleanup_registered: - from litellm.llms.custom_httpx.async_client_cleanup import register_async_client_cleanup + from litellm.llms.custom_httpx.async_client_cleanup import ( + register_async_client_cleanup, + ) + register_async_client_cleanup() _async_client_cleanup_registered = True @@ -1688,36 +1846,45 @@ def __getattr__(name: str) -> Any: # Lazy load encoding from main.py to avoid heavy tiktoken import if name == "encoding": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "encoding" not in _globals: from .main import encoding as _encoding + _globals["encoding"] = _encoding return _globals["encoding"] # Lazy load bedrock_tool_name_mappings instance if name == "bedrock_tool_name_mappings": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "bedrock_tool_name_mappings" not in _globals: - from .llms.bedrock.chat.invoke_handler import bedrock_tool_name_mappings as _bedrock_tool_name_mappings + from .llms.bedrock.chat.invoke_handler import ( + bedrock_tool_name_mappings as _bedrock_tool_name_mappings, + ) + _globals["bedrock_tool_name_mappings"] = _bedrock_tool_name_mappings return _globals["bedrock_tool_name_mappings"] # Lazy load AzureOpenAIError exception class if name == "AzureOpenAIError": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "AzureOpenAIError" not in _globals: from .llms.azure.common_utils import AzureOpenAIError as _AzureOpenAIError + _globals["AzureOpenAIError"] = _AzureOpenAIError return _globals["AzureOpenAIError"] # Lazy load openaiOSeriesConfig instance if name == "openaiOSeriesConfig": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() if "openaiOSeriesConfig" not in _globals: # Import the config class and instantiate it @@ -1735,6 +1902,7 @@ def __getattr__(name: str) -> Any: } if name in _config_instances: from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() if name not in _globals: # Import the config class and instantiate it @@ -1749,17 +1917,20 @@ def __getattr__(name: str) -> Any: # Lazy load provider_list if name == "provider_list": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "provider_list" not in _globals: # LlmProviders is eagerly imported above, so we can import it directly from litellm.types.utils import LlmProviders + _globals["provider_list"] = list(LlmProviders) return _globals["provider_list"] # Lazy load priority_reservation_settings instance if name == "priority_reservation_settings": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "priority_reservation_settings" not in _globals: @@ -1771,6 +1942,7 @@ def __getattr__(name: str) -> Any: # Lazy load logging_callback_manager instance if name == "logging_callback_manager": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "logging_callback_manager" not in _globals: @@ -1782,43 +1954,41 @@ def __getattr__(name: str) -> Any: # Lazy load _service_logger module if name == "_service_logger": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "_service_logger" not in _globals: # Import the module lazily import litellm._service_logger + _globals["_service_logger"] = litellm._service_logger return _globals["_service_logger"] # Lazy load evals module functions - if name in ["acreate_eval", "alist_evals", "aget_eval", "aupdate_eval", "adelete_eval", "acancel_eval", - "create_eval", "list_evals", "get_eval", "update_eval", "delete_eval", "cancel_eval", - "acreate_run", "alist_runs", "aget_run", "acancel_run", "adelete_run", - "create_run", "list_runs", "get_run", "cancel_run", "delete_run"]: - from litellm.evals.main import ( - acreate_eval, - alist_evals, - aget_eval, - aupdate_eval, - adelete_eval, - acancel_eval, - create_eval, - list_evals, - get_eval, - update_eval, - delete_eval, - cancel_eval, - acreate_run, - alist_runs, - aget_run, - acancel_run, - adelete_run, - create_run, - list_runs, - get_run, - cancel_run, - delete_run, - ) + if name in [ + "acreate_eval", + "alist_evals", + "aget_eval", + "aupdate_eval", + "adelete_eval", + "acancel_eval", + "create_eval", + "list_evals", + "get_eval", + "update_eval", + "delete_eval", + "cancel_eval", + "acreate_run", + "alist_runs", + "aget_run", + "acancel_run", + "adelete_run", + "create_run", + "list_runs", + "get_run", + "cancel_run", + "delete_run", + ]: return locals()[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index c2b0c4ddce9..5b568fbadd6 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -454,5 +454,38 @@ } }, "description": "SQS Queue (AWS) Logging Integration" + }, + { + "id": "microsoft_purview", + "displayName": "Microsoft Purview", + "logo": "microsoft.svg", + "supports_key_team_logging": false, + "dynamic_params": { + "MICROSOFT_PURVIEW_TENANT_ID": { + "type": "text", + "ui_name": "Azure Tenant ID", + "description": "Azure Active Directory Tenant ID", + "required": true + }, + "MICROSOFT_PURVIEW_CLIENT_ID": { + "type": "password", + "ui_name": "App Client ID", + "description": "App Registration Client ID with Content.Process.All permission", + "required": true + }, + "MICROSOFT_PURVIEW_CLIENT_SECRET": { + "type": "password", + "ui_name": "App Client Secret", + "description": "App Registration Client Secret", + "required": true + }, + "MICROSOFT_PURVIEW_APP_ID": { + "type": "text", + "ui_name": "App ID (GUID)", + "description": "Registered Application GUID in Purview policy location", + "required": false + } + }, + "description": "Microsoft Purview AI Compliance & Audit Logging Integration" } -] +] \ No newline at end of file diff --git a/litellm/integrations/microsoft_purview/__init__.py b/litellm/integrations/microsoft_purview/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/integrations/microsoft_purview/microsoft_purview.py b/litellm/integrations/microsoft_purview/microsoft_purview.py new file mode 100644 index 00000000000..032a4c7c9a9 --- /dev/null +++ b/litellm/integrations/microsoft_purview/microsoft_purview.py @@ -0,0 +1,424 @@ +""" +Microsoft Purview Integration - sends LLM prompts & responses to the Microsoft Graph +processContent API for compliance, DLP, and audit tracking. + +Reference API: https://learn.microsoft.com/en-us/graph/api/userdatasecurityandgovernance-processcontent +""" + +import asyncio +import os +import traceback +from collections import defaultdict +from typing import List, Optional, Dict, Any +from datetime import datetime, timezone + + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.utils import StandardLoggingPayload +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + +class MicrosoftPurviewLogger(CustomBatchLogger): + """ + Logger that sends LLM interactions to Microsoft Purview via the + Microsoft Graph processContent API for compliance and audit. + """ + + def __init__( + self, + tenant_id: Optional[str] = None, + client_id: Optional[str] = None, + client_secret: Optional[str] = None, + app_name: Optional[str] = None, + app_version: Optional[str] = None, + app_id: Optional[str] = None, + default_user_id: Optional[str] = None, + graph_api_version: str = "v1.0", + log_prompts: bool = True, + log_responses: bool = True, + **kwargs, + ): + """ + Initialize Microsoft Purview logger using the Graph API + """ + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + + self.tenant_id = ( + tenant_id + or os.getenv("MICROSOFT_PURVIEW_TENANT_ID") + or os.getenv("AZURE_TENANT_ID") + ) + self.client_id = ( + client_id + or os.getenv("MICROSOFT_PURVIEW_CLIENT_ID") + or os.getenv("AZURE_CLIENT_ID") + ) + self.client_secret = ( + client_secret + or os.getenv("MICROSOFT_PURVIEW_CLIENT_SECRET") + or os.getenv("AZURE_CLIENT_SECRET") + ) + + self.app_name = app_name or os.getenv( + "MICROSOFT_PURVIEW_APP_NAME", "LiteLLM Proxy" + ) + self.app_version = app_version or os.getenv( + "MICROSOFT_PURVIEW_APP_VERSION", getattr(litellm, "_version", "0.0.0") + ) + self.app_id = app_id or os.getenv("MICROSOFT_PURVIEW_APP_ID") + + self.default_user_id = default_user_id or os.getenv( + "MICROSOFT_PURVIEW_DEFAULT_USER_ID", "lite-llm-unknown-user" + ) + self.graph_api_version = graph_api_version + + # Boolean flags controls what is captured in logs + self.log_prompts = log_prompts + self.log_responses = log_responses + + if not self.tenant_id: + raise ValueError( + "MICROSOFT_PURVIEW_TENANT_ID is required to use Microsoft Purview integration" + ) + if not self.client_id: + raise ValueError( + "MICROSOFT_PURVIEW_CLIENT_ID is required to use Microsoft Purview integration" + ) + if not self.client_secret: + raise ValueError( + "MICROSOFT_PURVIEW_CLIENT_SECRET is required to use Microsoft Purview integration" + ) + + # OAuth2 scope for Microsoft Graph + self.oauth_scope = "https://graph.microsoft.com/.default" + self.oauth_token: Optional[str] = None + self.oauth_token_expires_at: Optional[float] = None + + self.flush_lock = asyncio.Lock() + super().__init__(**kwargs, flush_lock=self.flush_lock) + asyncio.create_task(self.periodic_flush()) + self.log_queue: List[StandardLoggingPayload] = [] + + async def _get_oauth_token(self) -> str: + """ + Get OAuth2 Bearer token for Microsoft Graph + """ + import time + + if ( + self.oauth_token + and self.oauth_token_expires_at + and time.time() < self.oauth_token_expires_at - 60 + ): # Refresh 60 seconds before expiry + return self.oauth_token + + assert self.tenant_id is not None, "tenant_id is required" + assert self.client_id is not None, "client_id is required" + assert self.client_secret is not None, "client_secret is required" + + token_url = ( + f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + ) + + token_data = { + "client_id": self.client_id, + "client_secret": self.client_secret, + "scope": self.oauth_scope, + "grant_type": "client_credentials", + } + + response = await self.async_httpx_client.post( + url=token_url, + data=token_data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + + if response.status_code != 200: + raise Exception( + f"Failed to get OAuth2 token: {response.status_code} - {response.text}" + ) + + token_response = response.json() + self.oauth_token = token_response.get("access_token") + expires_in = token_response.get("expires_in", 3600) + + if not self.oauth_token: + raise Exception("OAuth2 token response did not contain access_token") + + self.oauth_token_expires_at = time.time() + expires_in + return self.oauth_token + + def _extract_user_id(self, payload: StandardLoggingPayload) -> str: + """Get the user identity to map to the user-scoped Purview API""" + metadata = payload.get("metadata", {}) or {} + + user_id = metadata.get("user_api_key_user_id") + if user_id: + return str(user_id) + + end_user = payload.get("end_user") + if end_user: + return str(end_user) + + # fallback if not available + return self.default_user_id + + def _serialize_messages(self, messages: Any) -> str: + """Serialize prompts to a string, limiting total size if necessary""" + if isinstance(messages, str): + text = messages + elif isinstance(messages, list): + try: + # Try to extract just text for readability in Purview + parts = [] + for msg in messages: + if isinstance(msg, dict): + role = msg.get("role", "user") + content = msg.get("content", "") + if isinstance(content, str): + parts.append(f"[{role}]: {content}") + else: + parts.append(f"[{role}]: {safe_dumps(content)}") + else: + parts.append(str(msg)) + text = "\n\n".join(parts) + except Exception: + text = safe_dumps(messages) + else: + text = safe_dumps(messages) + + return text + + def _extract_response_text(self, payload: StandardLoggingPayload) -> str: + """Extract the model response""" + response = payload.get("response", {}) + if not response: + return "" + + if isinstance(response, str): + return response + + try: + choices = response.get("choices", []) + if choices and len(choices) > 0: + message = choices[0].get("message", {}) + content = message.get("content") + if content: + return str(content) + except Exception: + pass + + return safe_dumps(response) + + def _format_time(self, timestamp: Any) -> str: + """Format timestamp to ISO 8601 strictly""" + try: + if timestamp: + dt = datetime.fromtimestamp(timestamp, tz=timezone.utc) + else: + dt = datetime.now(timezone.utc) + # Purview requires exactly yYYY-MM-DDThh:mm:ss format + return dt.strftime("%Y-%m-%dT%H:%M:%S") + except Exception: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S") + + def _build_process_content_request(self, payload: StandardLoggingPayload) -> dict: + """Assemble the Microsoft Graph API request body for processContent""" + entries = [] + trace_id = payload.get("trace_id", "") or "purview-unknown-trace" + + # 1. Add User Prompts + if self.log_prompts: + messages = payload.get("messages", []) + prompt_text = self._serialize_messages(messages) + if prompt_text: + entries.append( + { + "@odata.type": "microsoft.graph.processConversationMetadata", + "identifier": f"{trace_id}-prompt", + "content": { + "@odata.type": "microsoft.graph.textContent", + "data": prompt_text, + }, + "name": "LLM Prompt", + "correlationId": trace_id, + "sequenceNumber": 0, + "isTruncated": False, + "createdDateTime": self._format_time(payload.get("startTime")), + "modifiedDateTime": self._format_time(payload.get("startTime")), + } + ) + + # 2. Add AI Response + if self.log_responses: + response_text = self._extract_response_text(payload) + if response_text: + entries.append( + { + "@odata.type": "microsoft.graph.processConversationMetadata", + "identifier": f"{trace_id}-response", + "content": { + "@odata.type": "microsoft.graph.textContent", + "data": response_text, + }, + "name": "LLM Response", + "correlationId": trace_id, + "sequenceNumber": 1, + "isTruncated": False, + "createdDateTime": self._format_time(payload.get("endTime")), + "modifiedDateTime": self._format_time(payload.get("endTime")), + } + ) + + # If nothing to send based on configuration or empty payload + if not entries: + return {} + + req_body = { + "contentToProcess": { + "contentEntries": entries, + "activityMetadata": { + "activity": "uploadText" # Or "downloadText". "uploadText" represents generating intent + receiving content + }, + "integratedAppMetadata": { + "name": self.app_name, + "version": self.app_version, + }, + } + } + + # Add protectedAppMetadata if an app_id was specified (needed for full mapping in Purview) + if self.app_id: + req_body["contentToProcess"]["protectedAppMetadata"] = { + "name": self.app_name, + "version": self.app_version, + "applicationLocation": { + "@odata.type": "microsoft.graph.policyLocationApplication", + "value": self.app_id, + }, + } + + return req_body + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """Async log success events to Microsoft Purview API queue""" + try: + verbose_logger.debug( + "Microsoft Purview: Queueing success log for model %s", + kwargs.get("model"), + ) + standard_logging_payload = kwargs.get("standard_logging_object", None) + + if standard_logging_payload is None: + return + + self.log_queue.append(standard_logging_payload) + + if len(self.log_queue) >= self.batch_size: + await self.async_send_batch() + + except Exception as e: + verbose_logger.exception( + f"Microsoft Purview Success Logging Error - {str(e)}\n{traceback.format_exc()}" + ) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + """Async log failure events to Microsoft Purview API queue""" + try: + verbose_logger.debug( + "Microsoft Purview: Queueing failure log for model %s", + kwargs.get("model"), + ) + standard_logging_payload = kwargs.get("standard_logging_object", None) + + if standard_logging_payload is None: + return + + self.log_queue.append(standard_logging_payload) + + if len(self.log_queue) >= self.batch_size: + await self.async_send_batch() + + except Exception as e: + verbose_logger.exception( + f"Microsoft Purview Failure Logging Error - {str(e)}\n{traceback.format_exc()}" + ) + + async def async_send_batch(self): + """ + Sends the batch of logs to Microsoft Graph Process Content API + """ + try: + if not self.log_queue: + return + + verbose_logger.debug( + "Microsoft Purview - about to flush %s events", len(self.log_queue) + ) + + # 1. Group payloads by user id (Since Graph API is per-user) + groups: Dict[str, list] = defaultdict(list) + for payload in self.log_queue: + user_id = self._extract_user_id(payload) + req_body = self._build_process_content_request(payload) + if req_body: + groups[user_id].append(req_body) + + if not groups: + self.log_queue.clear() + return + + # 2. Get OAuth2 Token + bearer_token = await self._get_oauth_token() + headers = { + "Authorization": f"Bearer {bearer_token}", + "Content-Type": "application/json", + } + + # 3. Fire requests concurrently + # Although they belong to different users, we loop through. Graph API doesn't support batching processContent inside a single call currently. + # We process them in parallel + tasks = [] + for user_id, requests in groups.items(): + api_endpoint = f"https://graph.microsoft.com/{self.graph_api_version}/users/{user_id}/dataSecurityAndGovernance/processContent" + + for req_body in requests: + tasks.append( + self.async_httpx_client.post( + url=api_endpoint, json=req_body, headers=headers + ) + ) + + responses = await asyncio.gather(*tasks, return_exceptions=True) + + for index, response in enumerate(responses): + if isinstance(response, Exception): + verbose_logger.error( + "Microsoft Purview Graph API encountered error: %s", + str(response), + ) + elif response.status_code not in [200, 202, 204]: + verbose_logger.error( + "Microsoft Purview Graph API error: status_code=%s, response=%s", + response.status_code, + response.text, + ) + + verbose_logger.debug( + "Microsoft Purview: Flushed %s processContent calls", len(tasks) + ) + + except Exception as e: + verbose_logger.exception( + f"Microsoft Purview Error sending batch API - {str(e)}\n{traceback.format_exc()}" + ) + finally: + self.log_queue.clear() diff --git a/litellm/integrations/microsoft_purview/types.py b/litellm/integrations/microsoft_purview/types.py new file mode 100644 index 00000000000..5e5523871e4 --- /dev/null +++ b/litellm/integrations/microsoft_purview/types.py @@ -0,0 +1,61 @@ +from typing import TypedDict, List, Optional + + +class TextContent(TypedDict): + odata_type: str # @odata.type -> "microsoft.graph.textContent" + data: str + + +class ProcessConversationMetadata(TypedDict): + odata_type: str # @odata.type -> "microsoft.graph.processConversationMetadata" + identifier: str + content: TextContent + name: str + correlationId: str + sequenceNumber: int + isTruncated: bool + createdDateTime: str + modifiedDateTime: str + + +class ActivityMetadata(TypedDict): + activity: str # e.g., "uploadText", "downloadText" + + +class OperatingSystemSpecifications(TypedDict): + operatingSystemPlatform: str + operatingSystemVersion: str + + +class DeviceMetadata(TypedDict, total=False): + deviceType: str + operatingSystemSpecifications: OperatingSystemSpecifications + ipAddress: str + + +class PolicyLocationApplication(TypedDict): + odata_type: str # @odata.type -> "microsoft.graph.policyLocationApplication" + value: str + + +class ProtectedApplicationMetadata(TypedDict, total=False): + name: str + version: str + applicationLocation: PolicyLocationApplication + + +class IntegratedApplicationMetadata(TypedDict): + name: str + version: str + + +class ProcessContentRequest(TypedDict): + contentEntries: List[ProcessConversationMetadata] + activityMetadata: ActivityMetadata + deviceMetadata: Optional[DeviceMetadata] + protectedAppMetadata: Optional[ProtectedApplicationMetadata] + integratedAppMetadata: IntegratedApplicationMetadata + + +class ProcessContentRequestBody(TypedDict): + contentToProcess: ProcessContentRequest diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6f587abcdf1..61b6963d0c6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -130,6 +130,7 @@ from ..integrations.argilla import ArgillaLogger from ..integrations.arize.arize_phoenix import ArizePhoenixLogger from ..integrations.athina import AthinaLogger from ..integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger +from ..integrations.microsoft_purview.microsoft_purview import MicrosoftPurviewLogger from ..integrations.azure_storage.azure_storage import AzureBlobStorageLogger from ..integrations.custom_prompt_management import CustomPromptManagement from ..integrations.datadog.datadog import DataDogLogger @@ -352,9 +353,9 @@ class Logging(LiteLLMLoggingBaseClass): ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[Any] = ( - [] - ) # for generating complete stream response + self.sync_streaming_chunks: List[ + Any + ] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks @@ -746,9 +747,9 @@ class Logging(LiteLLMLoggingBaseClass): prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details["prompt_integration"] = ( - logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = logger.__class__.__name__ return logger except Exception: # If check fails, continue to next logger @@ -816,9 +817,9 @@ class Logging(LiteLLMLoggingBaseClass): if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( non_default_params ): - self.model_call_details["prompt_integration"] = ( - anthropic_cache_control_logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = anthropic_cache_control_logger.__class__.__name__ return anthropic_cache_control_logger ######################################################### @@ -830,9 +831,9 @@ class Logging(LiteLLMLoggingBaseClass): internal_usage_cache=None, llm_router=None, ) - self.model_call_details["prompt_integration"] = ( - vector_store_custom_logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = vector_store_custom_logger.__class__.__name__ # Add to global callbacks so post-call hooks are invoked if ( vector_store_custom_logger @@ -892,9 +893,9 @@ class Logging(LiteLLMLoggingBaseClass): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"]["api_base"] = ( - self._get_masked_api_base(additional_args.get("api_base", "")) - ) + self.model_call_details["litellm_params"][ + "api_base" + ] = self._get_masked_api_base(additional_args.get("api_base", "")) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 # Log the exact input to the LLM API @@ -923,10 +924,10 @@ class Logging(LiteLLMLoggingBaseClass): try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata["raw_request"] = ( - "redacted by litellm. \ + _metadata[ + "raw_request" + ] = "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" - ) else: curl_command = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -937,34 +938,34 @@ class Logging(LiteLLMLoggingBaseClass): _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, - ) + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, ) except Exception as e: - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - error=str(e), - ) + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + error=str(e), ) - _metadata["raw_request"] = ( - "Unable to Log \ + _metadata[ + "raw_request" + ] = "Unable to Log \ raw request: {}".format( - str(e) - ) + str(e) ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: @@ -1265,13 +1266,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[MCPPostCallResponseObject] = ( - await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, - ) + response: Optional[ + MCPPostCallResponseObject + ] = await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1466,9 +1467,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info return None try: @@ -1494,9 +1495,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info return None @@ -1652,9 +1653,9 @@ class Logging(LiteLLMLoggingBaseClass): result=logging_result ) - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload(logging_result, start_time, end_time) - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload(logging_result, start_time, end_time) if ( standard_logging_payload := self.model_call_details.get( @@ -1732,9 +1733,9 @@ class Logging(LiteLLMLoggingBaseClass): end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details["completion_start_time"] = ( - self.completion_start_time - ) + self.model_call_details[ + "completion_start_time" + ] = self.completion_start_time self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1771,10 +1772,10 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - result, start_time, end_time - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload( + result, start_time, end_time ) if ( standard_logging_payload := self.model_call_details.get( @@ -1783,9 +1784,9 @@ class Logging(LiteLLMLoggingBaseClass): ) is not None: emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: - self.model_call_details["standard_logging_object"] = ( - standard_logging_object - ) + self.model_call_details[ + "standard_logging_object" + ] = standard_logging_object else: self.model_call_details["response_cost"] = None @@ -1943,17 +1944,17 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( "Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details["complete_streaming_response"] = ( - complete_streaming_response - ) - self.model_call_details["response_cost"] = ( - self._response_cost_calculator(result=complete_streaming_response) - ) + self.model_call_details[ + "complete_streaming_response" + ] = complete_streaming_response + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator(result=complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time ) if ( standard_logging_payload := self.model_call_details.get( @@ -2287,10 +2288,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2314,10 +2315,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] @@ -2456,9 +2457,9 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details["async_complete_streaming_response"] = ( - complete_streaming_response - ) + self.model_call_details[ + "async_complete_streaming_response" + ] = complete_streaming_response try: if self.model_call_details.get("cache_hit", False) is True: @@ -2469,10 +2470,10 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details["response_cost"] = ( - self._response_cost_calculator( - result=complete_streaming_response - ) + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator( + result=complete_streaming_response ) verbose_logger.debug( @@ -2485,10 +2486,10 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = None ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time ) # print standard logging payload @@ -2515,9 +2516,9 @@ class Logging(LiteLLMLoggingBaseClass): # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload(result, start_time, end_time) - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload(result, start_time, end_time) # print standard logging payload if ( @@ -2760,18 +2761,18 @@ class Logging(LiteLLMLoggingBaseClass): ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, ) return start_time, end_time @@ -3678,6 +3679,14 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _azure_sentinel_logger = AzureSentinelLogger() _in_memory_loggers.append(_azure_sentinel_logger) return _azure_sentinel_logger # type: ignore + elif logging_integration == "microsoft_purview": + for callback in _in_memory_loggers: + if isinstance(callback, MicrosoftPurviewLogger): + return callback # type: ignore + + _microsoft_purview_logger = MicrosoftPurviewLogger() + _in_memory_loggers.append(_microsoft_purview_logger) + return _microsoft_purview_logger # type: ignore elif logging_integration == "gcs_bucket": for callback in _in_memory_loggers: if isinstance(callback, GCSBucketLogger): @@ -3735,9 +3744,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 service_name=arize_config.project_name, ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" for callback in _in_memory_loggers: if ( isinstance(callback, ArizeLogger) @@ -3763,13 +3772,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={arize_phoenix_config.project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={arize_phoenix_config.project_name}" # Set Phoenix project name from environment variable phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) @@ -3777,19 +3786,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={phoenix_project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={phoenix_project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={phoenix_project_name}" # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - arize_phoenix_config.otlp_auth_headers - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = arize_phoenix_config.otlp_auth_headers for callback in _in_memory_loggers: if ( @@ -3965,9 +3974,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"api_key={os.getenv('LANGTRACE_API_KEY')}" - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetry) @@ -4284,6 +4293,10 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, AzureSentinelLogger): return callback + elif logging_integration == "microsoft_purview": + for callback in _in_memory_loggers: + if isinstance(callback, MicrosoftPurviewLogger): + return callback elif logging_integration == "gcs_bucket": for callback in _in_memory_loggers: if isinstance(callback, GCSBucketLogger): @@ -4881,10 +4894,10 @@ class StandardLoggingPayloadSetup: for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params["additional_headers"] = ( - StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] - ) + clean_hidden_params[ + "additional_headers" + ] = StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -5036,7 +5049,6 @@ class StandardLoggingPayloadSetup: dynamic_litellm_session_id = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id = litellm_params.get("litellm_trace_id") - # Note: we recommend using `litellm_session_id` for session tracking # `litellm_trace_id` is an internal litellm param if dynamic_litellm_session_id: @@ -5507,9 +5519,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[k] = ( - "scrubbed_by_litellm_for_sensitive_keys" - ) + cleaned_user_api_key_metadata[ + k + ] = "scrubbed_by_litellm_for_sensitive_keys" else: cleaned_user_api_key_metadata[k] = v diff --git a/tests/litellm/integrations/test_microsoft_purview.py b/tests/litellm/integrations/test_microsoft_purview.py new file mode 100644 index 00000000000..b87fb1cafba --- /dev/null +++ b/tests/litellm/integrations/test_microsoft_purview.py @@ -0,0 +1,257 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock + +from litellm.integrations.microsoft_purview.microsoft_purview import ( + MicrosoftPurviewLogger, +) + + +@pytest.fixture +def valid_env_vars(monkeypatch): + monkeypatch.setenv("MICROSOFT_PURVIEW_TENANT_ID", "test-tenant") + monkeypatch.setenv("MICROSOFT_PURVIEW_CLIENT_ID", "test-client-id") + monkeypatch.setenv("MICROSOFT_PURVIEW_CLIENT_SECRET", "test-secret") + monkeypatch.setenv("MICROSOFT_PURVIEW_APP_NAME", "test-app") + monkeypatch.setenv("MICROSOFT_PURVIEW_APP_VERSION", "1.0.0") + monkeypatch.setenv("MICROSOFT_PURVIEW_APP_ID", "test-app-id") + + +@pytest.mark.asyncio +async def test_init_with_all_env_vars(valid_env_vars): + logger = MicrosoftPurviewLogger() + assert logger.tenant_id == "test-tenant" + assert logger.client_id == "test-client-id" + assert logger.client_secret == "test-secret" + assert logger.app_name == "test-app" + assert logger.app_version == "1.0.0" + assert logger.app_id == "test-app-id" + assert logger.oauth_scope == "https://graph.microsoft.com/.default" + + +def test_init_missing_tenant_id_raises(monkeypatch): + monkeypatch.delenv("MICROSOFT_PURVIEW_TENANT_ID", raising=False) + monkeypatch.delenv("AZURE_TENANT_ID", raising=False) + monkeypatch.setenv("MICROSOFT_PURVIEW_CLIENT_ID", "test-client-id") + monkeypatch.setenv("MICROSOFT_PURVIEW_CLIENT_SECRET", "test-secret") + with pytest.raises( + ValueError, + match="MICROSOFT_PURVIEW_TENANT_ID is required to use Microsoft Purview integration", + ): + MicrosoftPurviewLogger() + + +def test_init_missing_client_id_raises(monkeypatch): + monkeypatch.setenv("MICROSOFT_PURVIEW_TENANT_ID", "test-tenant") + monkeypatch.delenv("MICROSOFT_PURVIEW_CLIENT_ID", raising=False) + monkeypatch.delenv("AZURE_CLIENT_ID", raising=False) + monkeypatch.setenv("MICROSOFT_PURVIEW_CLIENT_SECRET", "test-secret") + with pytest.raises( + ValueError, + match="MICROSOFT_PURVIEW_CLIENT_ID is required to use Microsoft Purview integration", + ): + MicrosoftPurviewLogger() + + +def test_init_missing_client_secret_raises(monkeypatch): + monkeypatch.setenv("MICROSOFT_PURVIEW_TENANT_ID", "test-tenant") + monkeypatch.setenv("MICROSOFT_PURVIEW_CLIENT_ID", "test-client-id") + monkeypatch.delenv("MICROSOFT_PURVIEW_CLIENT_SECRET", raising=False) + monkeypatch.delenv("AZURE_CLIENT_SECRET", raising=False) + with pytest.raises( + ValueError, + match="MICROSOFT_PURVIEW_CLIENT_SECRET is required to use Microsoft Purview integration", + ): + MicrosoftPurviewLogger() + + +@pytest.mark.asyncio +async def test_extract_user_id_from_metadata(valid_env_vars): + logger = MicrosoftPurviewLogger(default_user_id="default-user") + + # Priority 1: metadata.user_api_key_user_id + payload1 = {"metadata": {"user_api_key_user_id": "test-user-1"}} + assert logger._extract_user_id(payload1) == "test-user-1" + + # Priority 2: end_user + payload2 = {"end_user": "test-user-2"} + assert logger._extract_user_id(payload2) == "test-user-2" + + # Priority 3: fallback + payload3 = {} + assert logger._extract_user_id(payload3) == "default-user" + + +@pytest.mark.asyncio +async def test_serialize_messages(valid_env_vars): + logger = MicrosoftPurviewLogger() + + # List format + messages = [ + {"role": "user", "content": "hello world"}, + {"role": "assistant", "content": "hi"}, + ] + result = logger._serialize_messages(messages) + assert "[user]: hello world" in result + assert "[assistant]: hi" in result + + # String format + assert logger._serialize_messages("hello world") == "hello world" + + +@pytest.mark.asyncio +async def test_extract_response_text(valid_env_vars): + logger = MicrosoftPurviewLogger() + + # Standard format + payload = { + "response": {"choices": [{"message": {"content": "this is a response"}}]} + } + assert logger._extract_response_text(payload) == "this is a response" + + # String format + payload_str = {"response": "just a string response"} + assert logger._extract_response_text(payload_str) == "just a string response" + + +@pytest.mark.asyncio +async def test_build_process_content_request(valid_env_vars): + logger = MicrosoftPurviewLogger() + + payload = { + "trace_id": "test-trace-123", + "startTime": 1700000000, + "endTime": 1700000010, + "messages": [{"role": "user", "content": "What is 2+2?"}], + "response": {"choices": [{"message": {"content": "4"}}]}, + } + + req = logger._build_process_content_request(payload) + + assert "contentToProcess" in req + content_to_process = req["contentToProcess"] + + assert "contentEntries" in content_to_process + assert len(content_to_process["contentEntries"]) == 2 + + prompt_entry = content_to_process["contentEntries"][0] + assert prompt_entry["identifier"] == "test-trace-123-prompt" + assert prompt_entry["name"] == "LLM Prompt" + assert prompt_entry["content"]["data"] == "[user]: What is 2+2?" + assert prompt_entry["sequenceNumber"] == 0 + assert prompt_entry["correlationId"] == "test-trace-123" + + response_entry = content_to_process["contentEntries"][1] + assert response_entry["identifier"] == "test-trace-123-response" + assert response_entry["name"] == "LLM Response" + assert response_entry["content"]["data"] == "4" + assert response_entry["sequenceNumber"] == 1 + assert response_entry["correlationId"] == "test-trace-123" + + # Verify metadata + assert content_to_process["integratedAppMetadata"]["name"] == "test-app" + assert ( + content_to_process["protectedAppMetadata"]["applicationLocation"]["value"] + == "test-app-id" + ) + + +@pytest.mark.asyncio +async def test_async_log_success_event_queues(valid_env_vars): + logger = MicrosoftPurviewLogger(batch_size=5) + + kwargs = {"model": "gpt-4", "standard_logging_object": {"trace_id": "1"}} + + await logger.async_log_success_event(kwargs, None, None, None) + assert len(logger.log_queue) == 1 + assert logger.log_queue[0]["trace_id"] == "1" + + +@pytest.mark.asyncio +async def test_async_log_failure_event_queues(valid_env_vars): + logger = MicrosoftPurviewLogger(batch_size=5) + + kwargs = {"model": "gpt-4", "standard_logging_object": {"trace_id": "2"}} + + await logger.async_log_failure_event(kwargs, None, None, None) + assert len(logger.log_queue) == 1 + assert logger.log_queue[0]["trace_id"] == "2" + + +@pytest.mark.asyncio +async def test_async_send_batch_success(valid_env_vars, monkeypatch): + logger = MicrosoftPurviewLogger() + + # Add dummy payload to queue + payload = { + "trace_id": "test-trace-123", + "metadata": {"user_api_key_user_id": "user-A"}, + "messages": ["test message"], + } + logger.log_queue.append(payload) + + # Mock token + async def mock_get_token(): + return "fake-token" + + logger._get_oauth_token = mock_get_token + + # Mock the http client + mock_post = AsyncMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_post.return_value = mock_response + logger.async_httpx_client.post = mock_post + + await logger.async_send_batch() + + # Check that HTTP post was called correctly + assert mock_post.called + assert len(logger.log_queue) == 0 + call_kwargs = mock_post.call_args[1] + assert "url" in call_kwargs + assert "users/user-A/dataSecurityAndGovernance/processContent" in call_kwargs["url"] + assert "Bearer fake-token" in call_kwargs["headers"]["Authorization"] + + +@pytest.mark.asyncio +async def test_oauth_token_caching(valid_env_vars): + logger = MicrosoftPurviewLogger() + + import time + + logger.oauth_token = "cached-token" + logger.oauth_token_expires_at = time.time() + 3600 + + mock_post = AsyncMock() + logger.async_httpx_client.post = mock_post + + token = await logger._get_oauth_token() + + assert token == "cached-token" + assert not mock_post.called + + +@pytest.mark.asyncio +async def test_oauth_token_refresh(valid_env_vars): + logger = MicrosoftPurviewLogger() + + # Expired token + import time + + logger.oauth_token = "expired-token" + logger.oauth_token_expires_at = time.time() - 3600 + + # Mock token response + mock_post = AsyncMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"access_token": "new-token", "expires_in": 3600} + mock_post.return_value = mock_response + logger.async_httpx_client.post = mock_post + + token = await logger._get_oauth_token() + + assert token == "new-token" + assert logger.oauth_token == "new-token" + assert mock_post.called + assert logger.oauth_token_expires_at > time.time()