Skip to content

Commit 4bc28ea

Browse files
alexkromanclaude
andauthored
Type llm messages with ChatCompletionMessageParam instead of suppressing (#161)
The gateway `complete()` call carried a bare `# type: ignore[arg-type]` because `build_messages` returned `list[dict[str, str]]`, which isn't assignable to the OpenAI SDK's `Iterable[ChatCompletionMessageParam]`. Type the messages list and both function boundaries with the SDK's own param type so the `create()` call type-checks without the suppression. Tighten one build_messages test to compare the whole dict (matching its siblings), avoiding a not-required-key subscript on the now-typed union. --------- Co-authored-by: Claude <noreply@anthropic.com>
1 parent 910b0b7 commit 4bc28ea

2 files changed

Lines changed: 8 additions & 8 deletions

File tree

aai_cli/core/llm.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@
22

33
import json
44
from collections.abc import Sequence
5-
from typing import TYPE_CHECKING, Any
5+
from typing import TYPE_CHECKING
66

77
from aai_cli.core import environments
88
from aai_cli.core.errors import APIError, UsageError
99

1010
if TYPE_CHECKING:
1111
from openai import OpenAI
12-
from openai.types.chat import ChatCompletion
12+
from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
1313

1414
# The LLM Gateway is OpenAI-compatible, so we talk to it through the OpenAI SDK
1515
# pointed at the active environment's gateway base (see _client / code_gen).
@@ -76,7 +76,7 @@ def build_messages(
7676
system: str | None = None,
7777
transcript_id: str | None = None,
7878
transcript_text: str | None = None,
79-
) -> list[dict[str, str]]:
79+
) -> list[ChatCompletionMessageParam]:
8080
"""Assemble the chat `messages` array for a transcript transform or plain prompt.
8181
8282
With a `transcript_id`, the gateway injects the transcript server-side, so we
@@ -88,7 +88,7 @@ def build_messages(
8888
content = f"{prompt}\n\nTranscript:\n{transcript_text}"
8989
else:
9090
content = prompt
91-
messages: list[dict[str, str]] = []
91+
messages: list[ChatCompletionMessageParam] = []
9292
if system:
9393
messages.append({"role": "system", "content": system})
9494
messages.append({"role": "user", "content": content})
@@ -130,7 +130,7 @@ def complete(
130130
api_key: str,
131131
*,
132132
model: str,
133-
messages: list[dict[str, str]],
133+
messages: list[ChatCompletionMessageParam],
134134
max_tokens: int = DEFAULT_MAX_TOKENS,
135135
transcript_id: str | None = None,
136136
extra: dict[str, object] | None = None,
@@ -153,7 +153,7 @@ def complete(
153153
try:
154154
return client.chat.completions.create(
155155
model=model,
156-
messages=messages, # type: ignore[arg-type]
156+
messages=messages,
157157
max_tokens=max_tokens,
158158
extra_body=extra_body or None,
159159
)
@@ -182,7 +182,7 @@ def content_of(response: ChatCompletion) -> str:
182182
return content or ""
183183

184184

185-
def usage_of(response: ChatCompletion) -> dict[str, Any] | None:
185+
def usage_of(response: ChatCompletion) -> dict[str, object] | None:
186186
"""Return the token-usage block as a plain dict, if present."""
187187
usage = response.usage
188188
if usage is None:

tests/test_llm.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ def test_build_messages_transcript_id_uses_tag():
164164

165165
def test_build_messages_inline_text():
166166
msgs = llm.build_messages("summarize", transcript_text="hello world")
167-
assert msgs[0]["content"] == "summarize\n\nTranscript:\nhello world"
167+
assert msgs == [{"role": "user", "content": "summarize\n\nTranscript:\nhello world"}]
168168

169169

170170
def test_build_messages_with_system_prompt():

0 commit comments

Comments
 (0)