33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159 | class OpenAIChat(ChatLLM):
"""Chat model backed by the OpenAI Python SDK.
Args:
model: Any model the endpoint serves (e.g. `"gpt-5-nano"`).
api_key: Credential for the endpoint. Falls back to the `OPENAI_API_KEY`
environment variable when not given, so a deployment can keep its
secret under its own name and pass it here.
api_base: Endpoint URL, for any OpenAI-compatible server (vLLM, a
proxy, …). Falls back to `OPENAI_BASE_URL`, then to the SDK's
default of `https://api.openai.com/v1`.
"""
def __init__(
self,
model: str = "gpt-5-nano",
*,
api_key: str | None = None,
api_base: str | None = None,
) -> None:
self.model: str = model
self.api_key: str | None = (
api_key if api_key is not None else os.getenv("OPENAI_API_KEY")
)
self.api_base: str | None = (
api_base if api_base is not None else os.getenv("OPENAI_BASE_URL")
)
self.client = OpenAI(api_key=self.api_key, base_url=self.api_base)
self.async_client = AsyncOpenAI(api_key=self.api_key, base_url=self.api_base)
@staticmethod
def _to_openai_messages(
messages: list[Message], tools: list | None = None
) -> list[ChatCompletionMessageParam]:
"""Converts the internal Message objects into the shape expected by the OpenAI Python SDK. Casting just to make hinters happy."""
return cast(
list[ChatCompletionMessageParam],
[m.to_provider_dict() for m in messages],
)
def invoke(
self,
messages: list[Message],
*,
ctx: InvocationContext | None = None,
tools: list | None = None,
response_format: dict | type[BaseModel] | None = None,
) -> ChatLLMResponse:
openai_messages = self._to_openai_messages(messages)
with _translate_errors(self.model):
response = self.client.chat.completions.create(
model=self.model,
messages=openai_messages,
**_drop_none(
tools=tools,
response_format=response_format,
),
)
content = response.choices[0].message.content or ""
return ChatLLMResponse(message=Message(role="assistant", content=content))
async def ainvoke(
self,
messages: list[Message],
*,
ctx: InvocationContext | None = None,
tools: list | None = None,
response_format: dict | type[BaseModel] | None = None,
) -> ChatLLMResponse:
openai_messages = self._to_openai_messages(messages)
with _translate_errors(self.model):
response = await self.async_client.chat.completions.create(
model=self.model,
messages=openai_messages,
**_drop_none(
tools=tools,
response_format=response_format,
),
)
content = response.choices[0].message.content or ""
return ChatLLMResponse(message=Message(role="assistant", content=content))
def stream(
self,
messages: list[Message],
*,
ctx: InvocationContext | None = None,
tools: list | None = None,
response_format: dict | type[BaseModel] | None = None,
) -> AsyncIterator[ChatLLMStreamChunk]:
openai_messages = self._to_openai_messages(messages)
async def gen() -> AsyncIterator[ChatLLMStreamChunk]:
with _translate_errors(self.model):
stream = await self.async_client.chat.completions.create(
model=self.model,
messages=openai_messages,
stream=True,
**_drop_none(
tools=tools,
response_format=response_format,
),
)
async for chunk in stream:
choice_delta = chunk.choices[0].delta
delta = choice_delta.content
if delta:
yield ChatLLMStreamChunk(
delta=MessageChunk(role="assistant", delta=delta)
)
for tc in getattr(choice_delta, "tool_calls", None) or []:
yield ChatLLMStreamChunk(
delta=MessageChunk(role="assistant", delta=""),
tool_calls={
"index": tc.index,
"id": tc.id,
"name": getattr(tc.function, "name", None),
"arguments": getattr(tc.function, "arguments", None)
or "",
},
)
return gen()
|