Coverage for agentos/models/backends/openai.py: 29%

1"""OpenAI backend for AgentOS. 

2 

3Supports OpenAI, Azure OpenAI, and any OpenAI-compatible API (DeepSeek, Groq, etc.). 

4""" 

5 

6import os 

7from collections.abc import AsyncIterator 

8from dataclasses import dataclass, field 

9from typing import Any 

10 

11 

12@dataclass 

13class OpenAIConfig: 

14 """Configuration for OpenAI backend.""" 

15 

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

17 base_url: str = field( 

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

19 ) 

20 model: str = "gpt-4o" 

21 temperature: float = 0.7 

22 max_tokens: int = 4096 

23 top_p: float = 1.0 

24 frequency_penalty: float = 0.0 

25 presence_penalty: float = 0.0 

26 timeout: int = 60 

27 max_retries: int = 3 

28 organization: str = "" 

29 

30 

31class OpenAIClient: 

32 """OpenAI-compatible LLM client. 

33 

34 Works with: 

35 - OpenAI (GPT-4o, GPT-4, GPT-3.5) 

36 - Azure OpenAI 

37 - DeepSeek 

38 - Groq 

39 - Together AI 

40 - Any OpenAI-compatible endpoint 

41 """ 

42 

43 def __init__(self, config: OpenAIConfig | None = None): 

44 self.config = config or OpenAIConfig() 

45 self._client = None 

46 self._async_client = None 

47 

48 @property 

49 def headers(self) -> dict[str, str]: 

50 h = { 

51 "Authorization": f"Bearer {self.config.api_key}", 

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

53 } 

54 if self.config.organization: 

55 h["OpenAI-Organization"] = self.config.organization 

56 return h 

57 

58 @property 

59 def _chat_url(self) -> str: 

60 return f"{self.config.base_url}/chat/completions" 

61 

62 async def _async_request(self, messages: list[dict], **kwargs) -> dict: 

63 import httpx 

64 

65 timeout = httpx.Timeout(self.config.timeout) 

66 payload = { 

67 "model": self.config.model, 

68 "messages": messages, 

69 "temperature": kwargs.get("temperature", self.config.temperature), 

70 "max_tokens": kwargs.get("max_tokens", self.config.max_tokens), 

71 "top_p": kwargs.get("top_p", self.config.top_p), 

72 "frequency_penalty": self.config.frequency_penalty, 

73 "presence_penalty": self.config.presence_penalty, 

74 } 

75 

76 if kwargs.get("tools"): 

77 payload["tools"] = kwargs["tools"] 

78 payload["tool_choice"] = kwargs.get("tool_choice", "auto") 

79 

80 if kwargs.get("response_format"): 

81 payload["response_format"] = kwargs["response_format"] 

82 

83 for attempt in range(self.config.max_retries): 

84 try: 

85 async with httpx.AsyncClient() as client: 

86 resp = await client.post( 

87 self._chat_url, 

88 json=payload, 

89 headers=self.headers, 

90 timeout=timeout, 

91 ) 

92 resp.raise_for_status() 

93 return resp.json() 

94 except httpx.HTTPStatusError as e: 

95 if attempt == self.config.max_retries - 1: 

96 raise 

97 if e.response.status_code >= 500: 

98 import asyncio 

99 

100 await asyncio.sleep(2**attempt) 

101 continue 

102 raise 

103 

104 async def chat( 

105 self, 

106 messages: list[dict[str, str]], 

107 system: str | None = None, 

108 **kwargs, 

109 ) -> dict[str, Any]: 

110 """Send a chat completion request. 

111 

112 Args: 

113 messages: List of message dicts with 'role' and 'content'. 

114 system: Optional system prompt. 

115 **kwargs: Override config parameters. 

116 

117 Returns: 

118 Response dict with 'content', 'role', 'usage', 'model'. 

119 """ 

120 msgs = messages.copy() 

121 if system: 

122 msgs.insert(0, {"role": "system", "content": system}) 

123 

124 result = await self._async_request(msgs, **kwargs) 

125 choice = result["choices"][0] 

126 message = choice.get("message", {}) 

127 

128 tool_calls = message.get("tool_calls", []) 

129 

130 return { 

131 "content": message.get("content", ""), 

132 "role": message.get("role", "assistant"), 

133 "usage": result.get("usage", {}), 

134 "model": result.get("model", self.config.model), 

135 "finish_reason": choice.get("finish_reason", ""), 

136 "tool_calls": [ 

137 { 

138 "id": tc.get("id", ""), 

139 "name": tc.get("function", {}).get("name", ""), 

140 "arguments": tc.get("function", {}).get("arguments", "{}"), 

141 } 

142 for tc in tool_calls 

143 ], 

144 } 

145 

146 async def chat_stream( 

147 self, 

148 messages: list[dict[str, str]], 

149 system: str | None = None, 

150 **kwargs, 

151 ) -> AsyncIterator[dict[str, Any]]: 

152 """Stream chat completion tokens. 

153 

154 Yields dicts with 'delta', 'finish_reason', 'tool_call_delta'. 

155 """ 

156 import httpx 

157 

158 msgs = messages.copy() 

159 if system: 

160 msgs.insert(0, {"role": "system", "content": system}) 

161 

162 payload = { 

163 "model": self.config.model, 

164 "messages": msgs, 

165 "temperature": kwargs.get("temperature", self.config.temperature), 

166 "max_tokens": kwargs.get("max_tokens", self.config.max_tokens), 

167 "top_p": kwargs.get("top_p", self.config.top_p), 

168 "stream": True, 

169 } 

170 

171 timeout = httpx.Timeout(self.config.timeout * 2) 

172 async with httpx.AsyncClient() as client: 

173 async with client.stream( 

174 "POST", 

175 self._chat_url, 

176 json=payload, 

177 headers=self.headers, 

178 timeout=timeout, 

179 ) as response: 

180 response.raise_for_status() 

181 async for line in response.aiter_lines(): 

182 if line.startswith("data: "): 

183 data = line[6:].strip() 

184 if data == "[DONE]": 

185 break 

186 import json 

187 

188 try: 

189 chunk = json.loads(data) 

190 choice = chunk["choices"][0] 

191 delta = choice.get("delta", {}) 

192 yield { 

193 "delta": delta.get("content", ""), 

194 "finish_reason": choice.get("finish_reason"), 

195 "tool_call_delta": delta.get("tool_calls"), 

196 } 

197 except (json.JSONDecodeError, KeyError): 

198 continue 

199 

200 def sync_chat( 

201 self, 

202 messages: list[dict[str, str]], 

203 system: str | None = None, 

204 **kwargs, 

205 ) -> dict[str, Any]: 

206 """Synchronous chat completion.""" 

207 import asyncio 

208 

209 loop = asyncio.new_event_loop() 

210 try: 

211 return loop.run_until_complete(self.chat(messages, system, **kwargs)) 

212 finally: 

213 loop.close()