Skip to content

Commit ba703a2

Browse files
authored
Fix AI ext openai streaming (#8553)
Should fix issue reported on discord: https://discord.com/channels/841451783728529451/1356441709856292914/1356625059585986722 When translating the openai response (https://platform.openai.com/docs/api-reference/chat-streaming/streaming) to the anthropic style streaming response (https://docs.anthropic.com/en/api/messages-streaming#event-types), we're not sending a `content_block_start` event before the `content_block_delta` events for the text response as required. Also the default prompt doesn't contain a user message, so the user's query never actually gets sent to the llm, just the system message with the context.
1 parent 227475d commit ba703a2

File tree

3 files changed

+30
-6
lines changed

3 files changed

+30
-6
lines changed

edb/buildmeta.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@
5757
# The merge conflict there is a nice reminder that you probably need
5858
# to write a patch in edb/pgsql/patches.py, and then you should preserve
5959
# the old value.
60-
EDGEDB_CATALOG_VERSION = 2025_04_07_00_00
60+
EDGEDB_CATALOG_VERSION = 2025_04_08_19_20
6161
EDGEDB_MAJOR_VERSION = 7
6262

6363

edb/lib/ext/ai.edgeql

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -656,6 +656,13 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
656656
knowledge, answer the user query."
657657
),
658658
}),
659+
(insert ext::ai::ChatPromptMessage {
660+
participant_role := ext::ai::ChatParticipantRole.User,
661+
content := (
662+
"Query: {query}\n\
663+
Answer: "
664+
),
665+
})
659666
}
660667
};
661668

edb/server/protocol/ai_ext.py

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1421,13 +1421,28 @@ async def _start_openai_like_chat(
14211421
+ b'data: ' + event_data + b'\n\n'
14221422
)
14231423
protocol.write_raw(event)
1424+
1425+
event_data = json.dumps({
1426+
"type": "content_block_start",
1427+
"index": 0,
1428+
"content_block": {
1429+
"type": "text",
1430+
"text": ""
1431+
}
1432+
}).encode("utf-8")
1433+
event = (
1434+
b'event: content_block_start\n'
1435+
+ b'data: ' + event_data + b'\n\n'
1436+
)
1437+
protocol.write_raw(event)
1438+
14241439
# if there's only one openai tool call it shows up here
14251440
if tool_calls:
14261441
for tool_call in tool_calls:
14271442
tool_index = tool_call["index"]
14281443
event_data = json.dumps({
14291444
"type": "content_block_start",
1430-
"index": tool_call["index"],
1445+
"index": tool_call["index"] + 1,
14311446
"content_block": {
14321447
"id": tool_call["id"],
14331448
"type": "tool_use",
@@ -1458,14 +1473,14 @@ async def _start_openai_like_chat(
14581473
+ b'data: { \
14591474
"type": "content_block_stop",'
14601475
+ b'"index": '
1461-
+ str(currentIndex - 1).encode()
1476+
+ str(currentIndex).encode()
14621477
+ b'}\n\n'
14631478
)
14641479
protocol.write_raw(event)
14651480

14661481
event_data = json.dumps({
14671482
"type": "content_block_start",
1468-
"index": currentIndex,
1483+
"index": currentIndex + 1,
14691484
"content_block": {
14701485
"id": tool_call.get("id"),
14711486
"type": "tool_use",
@@ -1483,7 +1498,7 @@ async def _start_openai_like_chat(
14831498
else:
14841499
event_data = json.dumps({
14851500
"type": "content_block_delta",
1486-
"index": currentIndex,
1501+
"index": currentIndex + 1,
14871502
"delta": {
14881503
"type": "tool_call_delta",
14891504
"args":
@@ -1497,7 +1512,9 @@ async def _start_openai_like_chat(
14971512
protocol.write_raw(event)
14981513
elif finish_reason := data.get("finish_reason"):
14991514
index = (
1500-
tool_index if finish_reason == "tool_calls" else 0
1515+
tool_index + 1
1516+
if finish_reason == "tool_calls"
1517+
else 0
15011518
)
15021519
event = (
15031520
b'event: content_block_stop\n'

0 commit comments

Comments
 (0)