Coverage for packages/core/src/langgate/core/models.py: 100%

1"""Core domain data models for LangGate.""" 

2 

3from datetime import UTC, datetime 

4from decimal import Decimal 

5from typing import Annotated, Any, NewType 

6 

7from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator 

8 

9from langgate.core.fields import NormalizedDecimal 

10 

11# Type aliases for flexibility while maintaining naming compatibility 

12ServiceProviderId = NewType("ServiceProviderId", str) 

13# Model provider might differ from the inference service provider 

14# The service provider is not intended to be exposed to external consumers of the registry 

15# The service provider is used by the proxy for routing requests to the correct service 

16ModelProviderId = NewType("ModelProviderId", str) 

17 

18# Common model providers for convenience 

19MODEL_PROVIDER_OPENAI = ModelProviderId("openai") 

20MODEL_PROVIDER_ANTHROPIC = ModelProviderId("anthropic") 

21MODEL_PROVIDER_META = ModelProviderId("meta") 

22MODEL_PROVIDER_GOOGLE = ModelProviderId("google") 

23MODEL_PROVIDER_DEEPSEEK = ModelProviderId("deepseek") 

24MODEL_PROVIDER_MISTRAL = ModelProviderId("mistralai") 

25MODEL_PROVIDER_ALIBABA = ModelProviderId("alibaba") 

26MODEL_PROVIDER_XAI = ModelProviderId("xai") 

27MODEL_PROVIDER_COHERE = ModelProviderId("cohere") 

28MODEL_PROVIDER_ELEUTHERIA = ModelProviderId("eleutheria") 

29 

30 

31class ServiceProvider(BaseModel): 

32 """Information about a service provider (API service).""" 

33 

34 id: ServiceProviderId 

35 base_url: str 

36 api_key: SecretStr 

37 default_params: dict[str, Any] = Field(default_factory=dict) 

38 

39 

40class ModelProviderBase(BaseModel): 

41 id: ModelProviderId | None = None 

42 name: str | None = None 

43 description: str | None = None 

44 

45 model_config = ConfigDict(extra="allow") 

46 

47 

48class ModelProvider(ModelProviderBase): 

49 """Information about a model provider (creator).""" 

50 

51 id: ModelProviderId = Field(default=...) 

52 name: str = Field(default=...) 

53 

54 

55class ContextWindow(BaseModel): 

56 """Context window information for a model.""" 

57 

58 max_input_tokens: int = 0 

59 max_output_tokens: int = 0 

60 

61 model_config = ConfigDict(extra="allow") 

62 

63 

64class ModelCapabilities(BaseModel): 

65 """Capabilities of a language model.""" 

66 

67 supports_tools: bool | None = None 

68 supports_parallel_tool_calls: bool | None = None 

69 supports_vision: bool | None = None 

70 supports_audio_input: bool | None = None 

71 supports_audio_output: bool | None = None 

72 supports_prompt_caching: bool | None = None 

73 supports_reasoning: bool | None = None 

74 supports_response_schema: bool | None = None 

75 supports_system_messages: bool | None = None 

76 

77 model_config = ConfigDict(extra="allow") 

78 

79 

80TokenCost = Annotated[NormalizedDecimal, "TokenCost"] 

81Percentage = Annotated[NormalizedDecimal, "Percentage"] 

82TokenUsage = Annotated[NormalizedDecimal, "TokenUsage"] 

83 

84 

85class ModelCost(BaseModel): 

86 """Cost information for a language model.""" 

87 

88 input_cost_per_token: TokenCost = Field(default_factory=lambda: Decimal()) 

89 output_cost_per_token: TokenCost = Field(default_factory=lambda: Decimal()) 

90 input_cost_per_token_batches: TokenCost | None = None 

91 output_cost_per_token_batches: TokenCost | None = None 

92 cache_read_input_token_cost: TokenCost | None = None 

93 

94 model_config = ConfigDict(extra="allow") 

95 

96 

97class LLMInfoBase(BaseModel): 

98 id: str | None = None 

99 name: str | None = None 

100 provider_id: ModelProviderId | None = None 

101 

102 description: str | None = None 

103 costs: ModelCost | None = None 

104 capabilities: ModelCapabilities | None = None 

105 context_window: ContextWindow | None = None 

106 

107 model_config = ConfigDict(extra="allow") 

108 

109 

110class LLMInfo(LLMInfoBase): 

111 """Information about a language model.""" 

112 

113 id: str = Field(default=...) # "gpt-4o" 

114 name: str = Field(default=...) 

115 provider_id: ModelProviderId = Field(default=...) 

116 

117 provider: ModelProvider # Who created it (shown to users) 

118 description: str | None = None 

119 costs: ModelCost = Field(default_factory=ModelCost) 

120 capabilities: ModelCapabilities = Field(default_factory=ModelCapabilities) 

121 context_window: ContextWindow = Field(default_factory=ContextWindow) 

122 updated_dt: datetime = Field(default_factory=lambda: datetime.now(UTC)) 

123 

124 @model_validator(mode="after") 

125 def _validate_provider_id(self): 

126 self.provider_id = self.provider.id 

127 return self