1"""OpenAI LLM provider implementation.

2

3 This module provides OpenAI API integration for dataknobs-llm, supporting:

4- GPT-4, GPT-3.5-turbo, and other OpenAI chat models

5- Text embeddings (ada-002, etc.)

6- Function calling / tool use

7- Streaming responses

8- JSON mode for structured outputs

9- Vision models (GPT-4V)

10

11 The OpenAIProvider uses the official OpenAI Python SDK and supports all

12 standard OpenAI API parameters.

13

14 Example:

15 ```python

16 from dataknobs_llm.llm.providers import OpenAIProvider

17 from dataknobs_llm.llm.base import LLMConfig

18

19 # Create provider

20 config = LLMConfig(

21 provider="openai",

22 model="gpt-4",

23 api_key="sk-...", # or set OPENAI_API_KEY env var

24 temperature=0.7,

25 max_tokens=500

26 )

27

28 async with OpenAIProvider(config) as llm:

29 # Simple completion

30 response = await llm.complete("What is Python?")

31 print(response.content)

32

33 # Streaming

34 async for chunk in llm.stream_complete("Tell a story"):

35 print(chunk.delta, end="", flush=True)

36

37 # Embeddings

38 embedding = await llm.embed("sample text")

39 print(f"Dimensions: {len(embedding)}")

40 ```

41

42 See Also:

43 - OpenAI API Documentation: https://platform.openai.com/docs

44 - openai Python package: https://github.com/openai/openai-python

45"""

46

47 import os

48 from typing import TYPE_CHECKING, Any, Dict, List, Union, AsyncIterator

49

50 from ..base import (

51 LLMConfig, LLMMessage, LLMResponse, LLMStreamResponse,

52 AsyncLLMProvider, ModelCapability,

53 LLMAdapter, normalize_llm_config

54)

55 from dataknobs_llm.prompts import AsyncPromptBuilder

56

57 if TYPE_CHECKING:

58 from dataknobs_config.config import Config

59

60

61 class OpenAIAdapter(LLMAdapter):

62 """Adapter for OpenAI API format."""

63

64 def adapt_messages(self, messages: List[LLMMessage]) -> List[Dict[str, Any]]:

65 """Convert messages to OpenAI format."""

66 adapted = []

67 for msg in messages:

68 message = {

69 'role': msg.role,

70 'content': msg.content

71 }

72 if msg.name:

73 message['name'] = msg.name

74 if msg.function_call:

75 message['function_call'] = msg.function_call

76 adapted.append(message)

77 return adapted

78

79 def adapt_response(self, response: Any) -> LLMResponse:

80 """Convert OpenAI response to standard format."""

81 choice = response.choices[0]

82 message = choice.message

83

84 return LLMResponse(

85 content=message.content or '',

86 model=response.model,

87 finish_reason=choice.finish_reason,

88 usage={

89 'prompt_tokens': response.usage.prompt_tokens,

90 'completion_tokens': response.usage.completion_tokens,

91 'total_tokens': response.usage.total_tokens

92 } if response.usage else None,

93 function_call=message.function_call if hasattr(message, 'function_call') else None

94 )

95

96 def adapt_config(self, config: LLMConfig) -> Dict[str, Any]:

97 """Convert config to OpenAI parameters."""

98 params = {

99 'model': config.model,

100 'temperature': config.temperature,

101 'top_p': config.top_p,

102 'frequency_penalty': config.frequency_penalty,

103 'presence_penalty': config.presence_penalty,

104 }

105

106 if config.max_tokens:

107 params['max_tokens'] = config.max_tokens

108 if config.stop_sequences:

109 params['stop'] = config.stop_sequences

110 if config.seed:

111 params['seed'] = config.seed

112 if config.logit_bias:

113 params['logit_bias'] = config.logit_bias

114 if config.user_id:

115 params['user'] = config.user_id

116 if config.response_format == 'json':

117 params['response_format'] = {'type': 'json_object'}

118 if config.functions:

119 params['functions'] = config.functions

120 if config.function_call:

121 params['function_call'] = config.function_call

122

123 return params

124

125

126 class OpenAIProvider(AsyncLLMProvider):

127 """OpenAI LLM provider with full API support.

128

129 Provides async access to OpenAI's chat, completion, embedding, and

130 function calling APIs. Supports all GPT models including GPT-4, GPT-3.5,

131 and specialized models (vision, embeddings).

132

133 Features:

134 - Full GPT-4 and GPT-3.5-turbo support

135 - Streaming responses for real-time output

136 - Function calling for tool use

137 - JSON mode for structured outputs

138 - Embeddings for semantic search

139 - Custom API endpoints (e.g., Azure OpenAI)

140 - Automatic retry with rate limiting

141 - Cost tracking

142

143 Example:

144 ```python

145 from dataknobs_llm.llm.providers import OpenAIProvider

146 from dataknobs_llm.llm.base import LLMConfig, LLMMessage

147

148 # Basic usage

149 config = LLMConfig(

150 provider="openai",

151 model="gpt-4",

152 api_key="sk-...",

153 temperature=0.7

154 )

155

156 async with OpenAIProvider(config) as llm:

157 # Simple question

158 response = await llm.complete("Explain async/await")

159 print(response.content)

160

161 # Multi-turn conversation

162 messages = [

163 LLMMessage(role="system", content="You are a coding tutor"),

164 LLMMessage(role="user", content="How do I use asyncio?")

165 ]

166 response = await llm.complete(messages)

167

168 # JSON mode for structured output

169 json_config = LLMConfig(

170 provider="openai",

171 model="gpt-4",

172 response_format="json",

173 system_prompt="Return JSON only"

174 )

175

176 llm = OpenAIProvider(json_config)

177 await llm.initialize()

178 response = await llm.complete(

179 "List 3 Python libraries as JSON: {name, description}"

180 )

181 import json

182 data = json.loads(response.content)

183

184 # With Azure OpenAI

185 azure_config = LLMConfig(

186 provider="openai",

187 model="gpt-4",

188 api_base="https://your-resource.openai.azure.com/",

189 api_key="azure-key"

190 )

191

192 # Function calling

193 functions = [{

194 "name": "search",

195 "description": "Search for information",

196 "parameters": {

197 "type": "object",

198 "properties": {

199 "query": {"type": "string"}

200 }

201 }

202 }]

203

204 response = await llm.function_call(messages, functions)

205 if response.function_call:

206 print(f"Call: {response.function_call['name']}")

207 ```

208

209 Args:

210 config: LLMConfig, dataknobs Config, or dict with provider settings

211 prompt_builder: Optional AsyncPromptBuilder for prompt rendering

212

213 Attributes:

214 adapter (OpenAIAdapter): Format adapter for OpenAI API

215 _client: OpenAI AsyncOpenAI client instance

216

217 See Also:

218 LLMConfig: Configuration options

219 AsyncLLMProvider: Base provider interface

220 OpenAIAdapter: Format conversion

221 """

222

223 def __init__(

224 self,

225 config: Union[LLMConfig, "Config", Dict[str, Any]],

226 prompt_builder: AsyncPromptBuilder | None = None

227 ):

228 # Normalize config first

229 llm_config = normalize_llm_config(config)

230 super().__init__(llm_config, prompt_builder=prompt_builder)

231 self.adapter = OpenAIAdapter()

232

233 async def initialize(self) -> None:

234 """Initialize OpenAI client."""

235 try:

236 import openai

237

238 api_key = self.config.api_key or os.environ.get('OPENAI_API_KEY')

239 if not api_key:

240 raise ValueError("OpenAI API key not provided")

241

242 self._client = openai.AsyncOpenAI(

243 api_key=api_key,

244 base_url=self.config.api_base,

245 timeout=self.config.timeout

246 )

247 self._is_initialized = True

248 except ImportError as e:

249 raise ImportError("openai package not installed. Install with: pip install openai") from e

250

251 async def close(self) -> None:

252 """Close OpenAI client."""

253 if self._client:

254 await self._client.close() # type: ignore[unreachable]

255 self._is_initialized = False

256

257 async def validate_model(self) -> bool:

258 """Validate model availability."""

259 try:

260 # List available models

261 models = await self._client.models.list()

262 model_ids = [m.id for m in models.data]

263 return self.config.model in model_ids

264 except Exception:

265 return False

266

267 def get_capabilities(self) -> List[ModelCapability]:

268 """Get OpenAI model capabilities."""

269 capabilities = [

270 ModelCapability.TEXT_GENERATION,

271 ModelCapability.CHAT,

272 ModelCapability.STREAMING

273 ]

274

275 if 'gpt-4' in self.config.model or 'gpt-3.5' in self.config.model:

276 capabilities.extend([

277 ModelCapability.FUNCTION_CALLING,

278 ModelCapability.JSON_MODE

279 ])

280

281 if 'vision' in self.config.model:

282 capabilities.append(ModelCapability.VISION)

283

284 if 'embedding' in self.config.model:

285 capabilities.append(ModelCapability.EMBEDDINGS)

286

287 return capabilities

288

289 async def complete(

290 self,

291 messages: Union[str, List[LLMMessage]],

292 config_overrides: Dict[str, Any] | None = None,

293 **kwargs

294 ) -> LLMResponse:

295 """Generate completion.

296

297 Args:

298 messages: Input messages or prompt

299 config_overrides: Optional dict to override config fields (model,

300 temperature, max_tokens, top_p, stop_sequences, seed)

301 **kwargs: Additional provider-specific parameters

302 """

303 if not self._is_initialized:

304 await self.initialize()

305

306 # Get runtime config (with overrides applied if provided)

307 runtime_config = self._get_runtime_config(config_overrides)

308

309 # Convert string to message list

310 if isinstance(messages, str):

311 messages = [LLMMessage(role='user', content=messages)]

312

313 # Add system prompt if configured

314 if runtime_config.system_prompt and messages[0].role != 'system':

315 messages.insert(0, LLMMessage(role='system', content=runtime_config.system_prompt))

316

317 # Adapt messages and config

318 adapted_messages = self.adapter.adapt_messages(messages)

319 params = self.adapter.adapt_config(runtime_config)

320 params.update(kwargs)

321

322 # Make API call

323 response = await self._client.chat.completions.create(

324 messages=adapted_messages,

325 **params

326 )

327

328 return self.adapter.adapt_response(response)

329

330 async def stream_complete(

331 self,

332 messages: Union[str, List[LLMMessage]],

333 config_overrides: Dict[str, Any] | None = None,

334 **kwargs

335 ) -> AsyncIterator[LLMStreamResponse]:

336 """Generate streaming completion.

337

338 Args:

339 messages: Input messages or prompt

340 config_overrides: Optional dict to override config fields (model,

341 temperature, max_tokens, top_p, stop_sequences, seed)

342 **kwargs: Additional provider-specific parameters

343 """

344 if not self._is_initialized:

345 await self.initialize()

346

347 # Get runtime config (with overrides applied if provided)

348 runtime_config = self._get_runtime_config(config_overrides)

349

350 # Convert string to message list

351 if isinstance(messages, str):

352 messages = [LLMMessage(role='user', content=messages)]

353

354 # Add system prompt if configured

355 if runtime_config.system_prompt and messages[0].role != 'system':

356 messages.insert(0, LLMMessage(role='system', content=runtime_config.system_prompt))

357

358 # Adapt messages and config

359 adapted_messages = self.adapter.adapt_messages(messages)

360 params = self.adapter.adapt_config(runtime_config)

361 params['stream'] = True

362 params.update(kwargs)

363

364 # Stream API call

365 stream = await self._client.chat.completions.create(

366 messages=adapted_messages,

367 **params

368 )

369

370 async for chunk in stream:

371 if chunk.choices[0].delta.content:

372 yield LLMStreamResponse(

373 delta=chunk.choices[0].delta.content,

374 is_final=chunk.choices[0].finish_reason is not None,

375 finish_reason=chunk.choices[0].finish_reason

376 )

377

378 async def embed(

379 self,

380 texts: Union[str, List[str]],

381 **kwargs

382 ) -> Union[List[float], List[List[float]]]:

383 """Generate embeddings."""

384 if not self._is_initialized:

385 await self.initialize()

386

387 if isinstance(texts, str):

388 texts = [texts]

389 single = True

390 else:

391 single = False

392

393 response = await self._client.embeddings.create(

394 input=texts,

395 model=self.config.model or 'text-embedding-ada-002'

396 )

397

398 embeddings = [e.embedding for e in response.data]

399 return embeddings[0] if single else embeddings

400

401 async def function_call(

402 self,

403 messages: List[LLMMessage],

404 functions: List[Dict[str, Any]],

405 **kwargs

406 ) -> LLMResponse:

407 """Execute function calling."""

408 if not self._is_initialized:

409 await self.initialize()

410

411 # Add system prompt if configured

412 if self.config.system_prompt and messages[0].role != 'system':

413 messages.insert(0, LLMMessage(role='system', content=self.config.system_prompt))

414

415 # Adapt messages and config

416 adapted_messages = self.adapter.adapt_messages(messages)

417 params = self.adapter.adapt_config(self.config)

418 params['functions'] = functions

419 params['function_call'] = kwargs.get('function_call', 'auto')

420 params.update(kwargs)

421

422 # Make API call

423 response = await self._client.chat.completions.create(

424 messages=adapted_messages,

425 **params

426 )

427

428 return self.adapter.adapt_response(response)