@@ -4,6 +4,7 @@ package anthropic_test
44
55import (
66 "context"
7+ "encoding/json"
78 "errors"
89 "os"
910 "testing"
@@ -190,3 +191,168 @@ func TestBetaMessageCountTokensWithOptionalParams(t *testing.T) {
190191 t .Fatalf ("err should be nil: %s" , err .Error ())
191192 }
192193}
194+
195+ func TestBetaAccumulate (t * testing.T ) {
196+ for name , testCase := range map [string ]struct {
197+ expected anthropic.BetaMessage
198+ events []string
199+ }{
200+ "empty message" : {
201+ expected : anthropic.BetaMessage {Usage : anthropic.BetaUsage {}},
202+ events : []string {
203+ `{"type": "message_start", "message": {}}` ,
204+ `{"type: "message_stop"}` ,
205+ },
206+ },
207+ "text content block" : {
208+ events : []string {
209+ `{"type": "message_start", "message": {}}` ,
210+ `{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": "This "}}` ,
211+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "is a "}}` ,
212+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta": "text": "text block!"}}` ,
213+ `{"type": "content_block_stop", "index": 0}` ,
214+ `{"type": "message_stop"}` ,
215+ },
216+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
217+ {Type : "text" , Text : "This is a text block!" },
218+ }},
219+ },
220+ "text content block with citations" : {
221+ events : []string {
222+ `{"type": "message_start", "message": {}}` ,
223+ `{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": "1 + 1"}}` ,
224+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " = 2"}}` ,
225+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "citations_delta", "citation": {"type": "char_location", "cited_text": "1 + 1 = 2", "document_index": 0, "document_title": "Math Facts", "start_char_index": 300, "end_char_index": 310 }}}` ,
226+ `{"type": "content_block_stop", "index": 0}` ,
227+ `{"type": "message_stop"}` ,
228+ },
229+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
230+ {Type : "text" , Text : "1 + 1 = 2" , Citations : []anthropic.BetaTextCitationUnion {{
231+ Type : "char_location" ,
232+ CitedText : "1 + 1 = 2" ,
233+ DocumentIndex : 0 ,
234+ DocumentTitle : "Math Facts" ,
235+ StartCharIndex : 300 ,
236+ EndCharIndex : 310 ,
237+ }}},
238+ }},
239+ },
240+ "tool use block" : {
241+ events : []string {
242+ `{"type": "message_start", "message": {}}` ,
243+ `{"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_id", "name": "tool_name", "input": {}}}` ,
244+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"argument\":"}}` ,
245+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": " \"value\"}"}}` ,
246+ `{"type": "content_block_stop", "index": 0}` ,
247+ `{"type": "message_stop"}` ,
248+ },
249+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
250+ {Type : "tool_use" , ID : "toolu_id" , Name : "tool_name" , Input : []byte (`{"argument": "value"}` )},
251+ }},
252+ },
253+ "tool use block with no params" : {
254+ events : []string {
255+ `{"type": "message_start", "message": {}}` ,
256+ `{"type": "content_block_start": "index": 0, "content_block": {"type": "tool_use", "id": "toolu_id", "name": "tool_name", input: {}}}` ,
257+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}` ,
258+ `{"type": "content_block_stop", "index": 0}` ,
259+ `{"type": "message_stop"}` ,
260+ },
261+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
262+ {Type : "tool_use" , ID : "toolu_id" , Name : "tool_name" },
263+ }},
264+ },
265+ "server tool use block" : {
266+ events : []string {
267+ `{"type": "message_start", "message": {}}` ,
268+ `{"type": "content_block_start": "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_id", "name": "web_search", input: {}}}` ,
269+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}` ,
270+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"query\": \"weat"}}` ,
271+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "her\"}"}}` ,
272+ `{"type": "content_block_stop", "index": 0}` ,
273+ `{"type": "message_stop"}` ,
274+ },
275+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
276+ {Type : "server_tool_use" , ID : "srvtoolu_id" , Name : "web_search" , Input : []byte (`{"query": "weather"}` )},
277+ }},
278+ },
279+ "thinking block" : {
280+ events : []string {
281+ `{"type": "message_start", "message": {}}` ,
282+ `{"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": "Let me think..."}}` ,
283+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "
284+ First, let's try this..."}}` ,
285+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "
286+ Therefore, the answer is..."}}` ,
287+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "ThinkingSignature"}}` ,
288+ `{"type": "content_block_stop", "index": 0}` ,
289+ `{"type": "message_stop"}` ,
290+ },
291+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
292+ {Type : "thinking" , Thinking : "Let me think...\n First, let's try this...\n Therefore, the answer is..." , Signature : "ThinkingSignature" },
293+ }},
294+ },
295+ "redacted thinking block" : {
296+ events : []string {
297+ `{"type": "message_start", "message": {}}` ,
298+ `{"type": "content_block_start", "index": 0, "content_block": {"type": "redacted_thinking", "data": "Redacted"}}` ,
299+ `{"type": "content_block_stop", "index": 0}` ,
300+ `{"type": "message_stop"}` ,
301+ },
302+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
303+ {Type : "redacted_thinking" , Data : "Redacted" },
304+ }},
305+ },
306+ "multiple content blocks" : {
307+ events : []string {
308+ `{"type": "message_start", "message": {}}` ,
309+ `{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": "Let me look up "}}` ,
310+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "the weather for "}}` ,
311+ `{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta": "text": "you."}}` ,
312+ `{"type": "content_block_stop", "index": 0}` ,
313+ `{"type": "content_block_start", "index": 1, "content_block": {"type": "thinking", "thinking": ""}}` ,
314+ `{"type": "content_block_delta", "index": 1, "delta": {"type": "thinking_delta", "thinking": "I can look this "}}` ,
315+ `{"type": "content_block_delta", "index": 1, "delta": {"type": "thinking_delta", "thinking": "up using a tool."}}` ,
316+ `{"type": "content_block_stop", "index": 1}` ,
317+ `{"type": "content_block_start", "index": 2, "content_block": {"type": "tool_use", "id": "toolu_id", "name": "get_weather", "input": {}}}` ,
318+ `{"type": "content_block_delta", "index": 2, "delta": {"type": "input_json_delta", "partial_json": "{\"city\": "}}` ,
319+ `{"type": "content_block_delta", "index": 2, "delta": {"type": "input_json_delta", "partial_json": "\"Los Angeles\"}"}}` ,
320+ `{"type": "content_block_stop", "index": 2}` ,
321+ `{"type": "content_block_start", "index": 3, "content_block": {"type": "text", "text": ""}}` ,
322+ `{"type": "content_block_delta", "index": 3, "delta": {"type": "text_delta", "text": "The weather in Los Angeles"}}` ,
323+ `{"type": "content_block_delta", "index": 3, "delta": {"type": "text_delta", "text": " is 85 degrees Fahrenheit!"}}` ,
324+ `{"type": "content_block_stop", "index": 3"}` ,
325+ `{"type": "message_stop"}` ,
326+ },
327+ expected : anthropic.BetaMessage {Content : []anthropic.BetaContentBlockUnion {
328+ {Type : "text" , Text : "Let me look up the weather for you." },
329+ {Type : "thinking" , Thinking : "I can look this up using a tool." },
330+ {Type : "tool_use" , ID : "toolu_id" , Name : "get_weather" , Input : []byte (`{"city": "Los Angeles"}` )},
331+ {Type : "text" , Text : "The weather in Los Angeles is 85 degrees Fahrenheit!" },
332+ }},
333+ },
334+ } {
335+ t .Run (name , func (t * testing.T ) {
336+ message := anthropic.BetaMessage {}
337+ for _ , eventStr := range testCase .events {
338+ event := anthropic.BetaRawMessageStreamEventUnion {}
339+ err := (& event ).UnmarshalJSON ([]byte (eventStr ))
340+ if err != nil {
341+ t .Fatal (err )
342+ }
343+ (& message ).Accumulate (event )
344+ }
345+ marshaledMessage , err := json .Marshal (message )
346+ if err != nil {
347+ t .Fatal (err )
348+ }
349+ marshaledExpectedMessage , err := json .Marshal (testCase .expected )
350+ if err != nil {
351+ t .Fatal (err )
352+ }
353+ if string (marshaledMessage ) != string (marshaledExpectedMessage ) {
354+ t .Fatalf ("Mismatched message: expected %s but got %s" , marshaledExpectedMessage , marshaledMessage )
355+ }
356+ })
357+ }
358+ }
0 commit comments