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

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 content: str 

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

24 

25 @property 

26 def has_tool_calls(self) -> bool: 

27 return len(self.tool_calls) > 0 

28 

29 

30@dataclass 

31class ModelSpec: 

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

33 provider: str 

34 model_id: str 

35 context_window: int = 128_000 

36 api_key: str = "" 

37 base_url: str = "" 

38 cost_per_1m_input: float = 0.0 

39 cost_per_1m_output: float = 0.0 

40 

41 

42class AllModelsFailed(Exception): 

43 

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

45 

46 pass 

47 

48 

49@dataclass 

50class ModelConfig: 

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

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

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

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

55 max_retries: int = 3 

56 request_timeout: int = 120 

57 

58 

59RECOMMENDED_CONFIG = ModelConfig( 

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

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

62 models={ 

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

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

65 "claude-3.5-sonnet": ModelSpec(provider="anthropic", model_id="claude-3.5-sonnet", context_window=200_000), 

66 }, 

67) 

68 

69 

70@dataclass 

71class ModelRouter: 

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

73 

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

75 base_url: str = field(default_factory=lambda: os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")) 

76 

77 async def chat( 

78 self, 

79 model: str, 

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

81 temperature: float = 0.0, 

82 max_tokens: int = 2048, 

83 ) -> str: 

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

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

86 resp = await client.post( 

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

88 headers={ 

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

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

91 }, 

92 json={ 

93 "model": model, 

94 "messages": messages, 

95 "temperature": temperature, 

96 "max_tokens": max_tokens, 

97 }, 

98 ) 

99 resp.raise_for_status() 

100 data = resp.json() 

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