Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py: 64%

1""" 

2Qostodian Nexus (by Qohash) — LiteLLM guardrail integration. 

3""" 

4 

5import os 

6from typing import TYPE_CHECKING, Final, Literal, Optional 

7 

8from litellm.integrations.custom_guardrail import log_guardrail_information 

9from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( 

10 GenericGuardrailAPI, 

11) 

12from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( 

13 QostodianNexusConfigModel, 

14) 

15from litellm.types.utils import GenericGuardrailAPIInputs 

16 

17if TYPE_CHECKING: 17 ↛ 18line 17 didn't jump to line 18 because the condition on line 17 was never true

18 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj 

19 

20GUARDRAIL_NAME: Final = "qostodian_nexus" 

21 

22 

23class QostodianNexus(GenericGuardrailAPI): 

24 def __init__( 

25 self, 

26 api_base: str | None = None, 

27 **kwargs, 

28 ): 

29 api_base = api_base or os.environ.get("QOSTODIAN_NEXUS_API_BASE", "http://nexus:8800") 

30 

31 kwargs["guardrail_name"] = kwargs.get("guardrail_name", GUARDRAIL_NAME) 

32 

33 # Merge built-in Qostodian Nexus identifier headers with any caller-supplied extra_headers 

34 nexus_headers: Final = [ 

35 "x-qostodian-nexus-identifiers-trace", 

36 "x-qostodian-nexus-identifiers-source", 

37 "x-qostodian-nexus-identifiers-container", 

38 "x-qostodian-nexus-identifiers-identity", 

39 ] 

40 

41 existing: Final = kwargs.get("extra_headers") or [] 

42 kwargs["extra_headers"] = nexus_headers + [h for h in existing if h not in nexus_headers] 

43 

44 super().__init__( 

45 api_base=api_base, 

46 **kwargs, 

47 ) 

48 

49 @log_guardrail_information 

50 async def apply_guardrail( 

51 self, 

52 inputs: GenericGuardrailAPIInputs, 

53 request_data: dict, 

54 input_type: Literal["request", "response"], 

55 logging_obj: Optional["LiteLLMLoggingObj"] = None, 

56 ) -> GenericGuardrailAPIInputs: 

57 """ 

58 Apply Qostodian Nexus to the given inputs. 

59 

60 NOTE: This override is intentionally a pass-through. It must be present 

61 directly in this class's __dict__ so that LiteLLM's unified guardrail 

62 routing check (`"apply_guardrail" in type(callback).__dict__` in 

63 litellm/proxy/utils.py) routes calls correctly. Do not remove. 

64 """ 

65 return await super().apply_guardrail( 

66 inputs=inputs, 

67 request_data=request_data, 

68 input_type=input_type, 

69 logging_obj=logging_obj, 

70 ) 

71 

72 @classmethod 

73 def get_config_model(cls) -> type[QostodianNexusConfigModel] | None: 

74 """ 

75 Returns the config model for Qostodian Nexus. 

76 """ 

77 return QostodianNexusConfigModel