Coverage for agentos/models/router.py: 85%

1""" 

2AgentOS v1.2.7 — Minimal ModelRouter for CodeAgent. 

3 

4Lightweight LLM call wrapper using httpx to OpenAI-compatible endpoints. 

5Designed as a self-contained module with zero internal dependencies. 

6""" 

7 

8from __future__ import annotations 

9 

10import os 

11from dataclasses import dataclass, field 

12from typing import Any 

13 

14import httpx 

15 

16from agentos.tools.base import ToolCall 

17 

18 

19@dataclass 

20class ModelResponse: 

21 """LLM 响应：文本内容 + 函数调用列表。""" 

22 

23 content: str 

24 tool_calls: list[ToolCall] = field(default_factory=list) 

25 

26 @property 

27 def has_tool_calls(self) -> bool: 

28 return len(self.tool_calls) > 0 

29 

30 

31@dataclass 

32class ModelSpec: 

33 """单个模型的规格定义。""" 

34 

35 provider: str 

36 model_id: str 

37 context_window: int = 128_000 

38 api_key: str = "" 

39 base_url: str = "" 

40 cost_per_1m_input: float = 0.0 

41 cost_per_1m_output: float = 0.0 

42 

43 

44class AllModelsFailed(Exception): # noqa: N818 

45 """所有模型均失败异常。""" 

46 

47 

48 

49@dataclass 

50class ModelConfig: 

51 """模型路由配置。""" 

52 

53 default_model: str = "gpt-4o-mini" 

54 fallback_chain: list[str] = field(default_factory=list) 

55 models: dict[str, ModelSpec] = field(default_factory=dict) 

56 max_retries: int = 3 

57 request_timeout: int = 120 

58 

59 

60RECOMMENDED_CONFIG = ModelConfig( 

61 default_model="gpt-4o-mini", 

62 fallback_chain=["gpt-4o", "claude-3.5-sonnet"], 

63 models={ 

64 "gpt-4o-mini": ModelSpec(provider="openai", model_id="gpt-4o-mini", context_window=128_000), 

65 "gpt-4o": ModelSpec(provider="openai", model_id="gpt-4o", context_window=128_000), 

66 "claude-3.5-sonnet": ModelSpec( 

67 provider="anthropic", model_id="claude-3.5-sonnet", context_window=200_000 

68 ), 

69 }, 

70) 

71 

72 

73@dataclass 

74class ModelRouter: 

75 """Minimal LLM router for code generation tasks.""" 

76 

77 api_key: str = field(default_factory=lambda: os.environ.get("OPENAI_API_KEY", "")) 

78 base_url: str = field( 

79 default_factory=lambda: os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") 

80 ) 

81 

82 async def chat( 

83 self, 

84 model: str, 

85 messages: list[dict[str, Any]], 

86 temperature: float = 0.0, 

87 max_tokens: int = 2048, 

88 ) -> str: 

89 """Send a chat request and return text content.""" 

90 async with httpx.AsyncClient(timeout=120) as client: 

91 resp = await client.post( 

92 f"{self.base_url}/chat/completions", 

93 headers={ 

94 "Authorization": f"Bearer {self.api_key}", 

95 "Content-Type": "application/json", 

96 }, 

97 json={ 

98 "model": model, 

99 "messages": messages, 

100 "temperature": temperature, 

101 "max_tokens": max_tokens, 

102 }, 

103 ) 

104 resp.raise_for_status() 

105 data = resp.json() 

106 return data["choices"][0]["message"]["content"]