Coverage for src/chat_limiter/types.py: 93%

1""" 

2Type definitions for chat completion requests and responses. 

3""" 

4 

5from dataclasses import dataclass 

6from enum import Enum 

7from typing import Any 

8 

9from pydantic import BaseModel 

10 

11 

12class MessageRole(str, Enum): 

13 """Message roles supported across providers.""" 

14 

15 USER = "user" 

16 ASSISTANT = "assistant" 

17 SYSTEM = "system" 

18 

19 

20@dataclass 

21class Message: 

22 """A chat message that works across all providers.""" 

23 

24 role: MessageRole 

25 content: str 

26 

27 

28class ChatCompletionRequest(BaseModel): 

29 """High-level chat completion request.""" 

30 

31 model: str 

32 messages: list[Message] 

33 max_tokens: int | None = None 

34 temperature: float | None = None 

35 top_p: float | None = None 

36 stop: str | list[str] | None = None 

37 stream: bool = False 

38 

39 # Provider-specific parameters (will be filtered per provider) 

40 frequency_penalty: float | None = None # OpenAI 

41 presence_penalty: float | None = None # OpenAI 

42 top_k: int | None = None # Anthropic 

43 

44 

45@dataclass 

46class Usage: 

47 """Token usage information.""" 

48 

49 prompt_tokens: int 

50 completion_tokens: int 

51 total_tokens: int 

52 

53 

54@dataclass 

55class Choice: 

56 """A completion choice.""" 

57 

58 index: int 

59 message: Message 

60 finish_reason: str | None = None 

61 

62 

63@dataclass 

64class ChatCompletionResponse: 

65 """High-level chat completion response.""" 

66 

67 id: str 

68 model: str 

69 choices: list[Choice] 

70 usage: Usage | None = None 

71 created: int | None = None 

72 

73 # Provider-specific metadata 

74 provider: str | None = None 

75 raw_response: dict[str, Any] | None = None 

76 

77 

78# Model mappings for each provider 

79OPENAI_MODELS = { 

80 "gpt-4o", 

81 "gpt-4o-mini", 

82 "gpt-4-turbo", 

83 "gpt-4", 

84 "gpt-3.5-turbo", 

85 "gpt-3.5-turbo-16k", 

86} 

87 

88ANTHROPIC_MODELS = { 

89 "claude-3-5-sonnet-20241022", 

90 "claude-3-5-haiku-20241022", 

91 "claude-3-opus-20240229", 

92 "claude-3-sonnet-20240229", 

93 "claude-3-haiku-20240307", 

94} 

95 

96OPENROUTER_MODELS = { 

97 # OpenAI models via OpenRouter 

98 "openai/gpt-4o", 

99 "openai/gpt-4o-mini", 

100 "openai/gpt-4-turbo", 

101 "openai/gpt-3.5-turbo", 

102 

103 # Anthropic models via OpenRouter 

104 "anthropic/claude-3-5-sonnet", 

105 "anthropic/claude-3-opus", 

106 "anthropic/claude-3-sonnet", 

107 "anthropic/claude-3-haiku", 

108 

109 # Other providers via OpenRouter 

110 "meta-llama/llama-3.1-405b-instruct", 

111 "meta-llama/llama-3.1-70b-instruct", 

112 "google/gemini-pro", 

113 "cohere/command-r-plus", 

114} 

115 

116ALL_MODELS = OPENAI_MODELS | ANTHROPIC_MODELS | OPENROUTER_MODELS 

117 

118 

119def detect_provider_from_model(model: str, use_dynamic_discovery: bool = False, api_keys: dict[str, str] | None = None) -> str | None: 

120 """ 

121 Detect provider from model name. 

122 

123 Args: 

124 model: The model name to check 

125 use_dynamic_discovery: Whether to use live API queries for model discovery 

126 api_keys: Dictionary of API keys for dynamic discovery 

127 

128 Returns: 

129 Provider name or None if not found 

130 """ 

131 # First try pattern-based detection for common cases 

132 if "/" in model: # OpenRouter format 

133 return "openrouter" 

134 

135 # Check hardcoded lists for fast lookup 

136 if model in OPENAI_MODELS: 

137 return "openai" 

138 elif model in ANTHROPIC_MODELS: 

139 return "anthropic" 

140 elif model in OPENROUTER_MODELS: 

141 return "openrouter" 

142 

143 # If dynamic discovery is enabled and we have API keys, try that 

144 if use_dynamic_discovery and api_keys: 

145 from .models import detect_provider_from_model_sync 

146 result = detect_provider_from_model_sync(model, api_keys) 

147 return result.found_provider 

148 

149 return None