message sse stuff is done

examples/messages_stream/message_stream.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+ "bufio"
+ "context"
+ "fmt"
+ "io"
+ "log"
+ "os"
+ "strings"
+
+ "github.com/dleviminzi/anthrogo"
+)
+
+func main() {
+ c, err := anthrogo.NewClient()
+ if err != nil {
+ log.Fatal(err)
+ os.Exit(1)
+ }
+
+ systemPrompt := "you are an expert in all things bananas"
+
+ // Read user input for the prompt
+ reader := bufio.NewReader(os.Stdin)
+ fmt.Print("Enter your prompt: ")
+ userPrompt, _ := reader.ReadString('\n')
+ userPrompt = strings.TrimSuffix(userPrompt, "\n")
+
+ r, _, err := c.MessageStreamRequest(context.Background(), anthrogo.MessagePayload{
+ Model: anthrogo.ModelClaude3Opus,
+ Messages: []anthrogo.Message{{
+ Role: anthrogo.RoleTypeUser,
+ Content: []anthrogo.MessageContent{{
+ Type: anthrogo.ContentTypeText,
+ Text: &userPrompt,
+ }},
+ }},
+ System: &systemPrompt,
+ MaxTokens: 1000,
+ })
+ if err != nil {
+ log.Fatal(err)
+ os.Exit(1)
+ }
+ defer r.Close()
+
+ // Create an SSEDecoder
+ decoder := anthrogo.NewMessageSSEDecoder(r)
+ for {
+ message, err := decoder.Decode(anthrogo.DecodeOptions{ContentOnly: true})
+ if err != nil {
+ if err == io.EOF {
+ break
+ }
+ fmt.Print(err)
+ continue
+ }
+
+ if message.Event == "message_stop" {
+ break
+ }
+
+ fmt.Print(message.Data.Content)
+ }
+}

message_sse_decoder.go

@@ -8,12 +8,22 @@ import (
  "strings"
 )
 
-type MessageSSEPayload struct {
- Event string `json:"event"`
- Data interface{} `json:"data"`
+// MessageEventPayload is the decoded event from anthropic
+type MessageEventPayload struct {
+ Event string
+ Data EventData
 }
 
-type MessageStartData struct {
+// EventData contains content which will be whatever the model output
+// and Data which is the full data from the event
+type EventData struct {
+ Content string
+ Data any
+}
+
+// MessageStart is one of the data types for events and it represents the start of a
+// a stream of messages. It contains metadata about the request.
+type MessageStart struct {
  Type string `json:"type"`
  Message struct {
  ID string `json:"id"`
@@ -30,7 +40,8 @@ type MessageStartData struct {
  } `json:"message"`
 }
 
-type ContentBlockStartData struct {
+// ContentBlockStart marks the start of a new content block in the message stream.
+type ContentBlockStart struct {
  Type string `json:"type"`
  Index int `json:"index"`
  ContentBlock struct {
@@ -39,38 +50,42 @@ type ContentBlockStartData struct {
  } `json:"content_block"`
 }
 
+// PingData is a ping event
 type PingData struct {
  Type string `json:"type"`
 }
 
-type ContentBlockDeltaData struct {
- Type string `json:"type"`
- Index int `json:"index"`
- Delta TextDelta `json:"delta"`
-}
-
-type TextDelta struct {
- Type string `json:"type"`
- Text string `json:"text"`
+// ContentBlockDelta carries new content for a content block in the message stream.
+type ContentBlockDelta struct {
+ Type string `json:"type"`
+ Index int `json:"index"`
+ Delta struct {
+ Type string `json:"type"`
+ Text string `json:"text"`
+ } `json:"delta"`
 }
 
-type ContentBlockStopData struct {
+// ContentBlockStop marks the end of a content block in the message stream.
+type ContentBlockStop struct {
  Type string `json:"type"`
  Index int `json:"index"`
 }
 
-type MessageDeltaData struct {
+// MessageDelta events indicate top-level changes to the final message.
+type MessageDelta struct {
  Type string `json:"type"`
  Delta interface{} `json:"delta"`
  Usage struct {
  OutputTokens int `json:"output_tokens"`
  } `json:"usage"`
 }
 
+// MessageStopData is the final event in a message stream.
 type MessageStopData struct {
  Type string `json:"type"`
 }
 
+// ErrorData is the event type for errors.
 type ErrorData struct {
  Type string `json:"type"`
  Error struct {
@@ -79,19 +94,41 @@ type ErrorData struct {
  } `json:"error"`
 }
 
+// MessageEvent is the event type for messages. It contains the message payload
+// and an error if one occurred.
+type MessageEvent struct {
+ Message *MessageEventPayload
+ Err *error
+}
+
+// MessageSSEDecoder is a decoder for the SSE stream from the message endpoint.
 type MessageSSEDecoder struct {
  reader *bufio.Reader
  content []string
 }
 
-func NewSSEDecoder(reader io.Reader) *MessageSSEDecoder {
+// DecodeOptions are options for decoding the SSE stream.
+type DecodeOptions struct {
+ ContentOnly bool
+}
+
+// NewMessageSSEDecoder creates a new MessageSSEDecoder.
+func NewMessageSSEDecoder(reader io.Reader) *MessageSSEDecoder {
  return &MessageSSEDecoder{
  reader: bufio.NewReader(reader),
  content: make([]string, 0),
  }
 }
 
-func (d *MessageSSEDecoder) Decode() (*MessageSSEPayload, error) {
+// Decode reads the next event from the SSE stream.
+func (d *MessageSSEDecoder) Decode(opts ...DecodeOptions) (*MessageEventPayload, error) {
+ var options DecodeOptions
+ if len(opts) > 1 {
+ return nil, fmt.Errorf("too many options provided, expected at most one")
+ } else if len(opts) == 1 {
+ options = opts[0]
+ }
+
  line, err := d.reader.ReadString('\n')
  if err != nil {
  if err == io.EOF {
@@ -103,7 +140,7 @@ func (d *MessageSSEDecoder) Decode() (*MessageSSEPayload, error) {
  line = strings.TrimSpace(line)
  if line == "" {
  // Recursively call Decode to read the next event
- return d.Decode()
+ return d.Decode(opts...)
  }
 
  parts := strings.SplitN(line, ":", 2)
@@ -115,38 +152,24 @@ func (d *MessageSSEDecoder) Decode() (*MessageSSEPayload, error) {
  value := strings.TrimSpace(parts[1])
 
  if field == "event" {
- data := d.decodeData(value)
-
- // Update the content array based on the event type
- switch value {
- case "content_block_start":
- contentBlockStartData := data.(ContentBlockStartData)
- index := contentBlockStartData.Index
- if index >= len(d.content) {
- d.content = append(d.content, make([]string, index-len(d.content)+1)...)
- }
- d.content[index] = contentBlockStartData.ContentBlock.Text
- case "content_block_delta":
- contentBlockDeltaData := data.(ContentBlockDeltaData)
- index := contentBlockDeltaData.Index
- if index >= len(d.content) {
- d.content = append(d.content, make([]string, index-len(d.content)+1)...)
- }
- d.content[index] += contentBlockDeltaData.Delta.Text
+ data, err := d.decodeData(value)
+ if err != nil {
+ return nil, err
  }
 
- return &MessageSSEPayload{
- Event: value,
- Data: data,
- }, nil
+ if data.Content != "" || !options.ContentOnly || value == "message_stop" {
+ return &MessageEventPayload{
+ Event: value,
+ Data: data,
+ }, nil
+ }
  }
  // Recursively call Decode to read the next event if we didn't have one here
- return d.Decode()
+ return d.Decode(opts...)
 }
 
-// TODO: check for errors and return them here.
-func (d *MessageSSEDecoder) decodeData(event string) any {
- var data any
+func (d *MessageSSEDecoder) decodeData(event string) (EventData, error) {
+ var eventData EventData
 
  for {
  line, err := d.reader.ReadString('\n')
@@ -164,40 +187,75 @@ func (d *MessageSSEDecoder) decodeData(event string) any {
 
  switch event {
  case "message_start":
- var messageStartData MessageStartData
- json.Unmarshal([]byte(jsonData), &messageStartData)
- data = messageStartData
+ var messageStartData MessageStart
+ err := json.Unmarshal([]byte(jsonData), &messageStartData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = messageStartData
  case "content_block_start":
- var contentBlockStartData ContentBlockStartData
- json.Unmarshal([]byte(jsonData), &contentBlockStartData)
- data = contentBlockStartData
+ var contentBlockStartData ContentBlockStart
+ err := json.Unmarshal([]byte(jsonData), &contentBlockStartData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = contentBlockStartData
+ eventData.Content = contentBlockStartData.ContentBlock.Text
+ d.updateContent(contentBlockStartData.Index, contentBlockStartData.ContentBlock.Text)
  case "ping":
  var pingData PingData
- json.Unmarshal([]byte(jsonData), &pingData)
- data = pingData
+ err := json.Unmarshal([]byte(jsonData), &pingData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = pingData
  case "content_block_delta":
- var contentBlockDeltaData ContentBlockDeltaData
- json.Unmarshal([]byte(jsonData), &contentBlockDeltaData)
- data = contentBlockDeltaData
+ var contentBlockDeltaData ContentBlockDelta
+ err := json.Unmarshal([]byte(jsonData), &contentBlockDeltaData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = contentBlockDeltaData
+ eventData.Content = contentBlockDeltaData.Delta.Text
+ d.updateContent(contentBlockDeltaData.Index, contentBlockDeltaData.Delta.Text)
  case "content_block_stop":
- var contentBlockStopData ContentBlockStopData
- json.Unmarshal([]byte(jsonData), &contentBlockStopData)
- data = contentBlockStopData
+ var contentBlockStopData ContentBlockStop
+ err := json.Unmarshal([]byte(jsonData), &contentBlockStopData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = contentBlockStopData
  case "message_delta":
- var messageDeltaData MessageDeltaData
- json.Unmarshal([]byte(jsonData), &messageDeltaData)
- data = messageDeltaData
+ var messageDeltaData MessageDelta
+ err := json.Unmarshal([]byte(jsonData), &messageDeltaData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = messageDeltaData
  case "message_stop":
  var messageStopData MessageStopData
- json.Unmarshal([]byte(jsonData), &messageStopData)
- data = messageStopData
+ err := json.Unmarshal([]byte(jsonData), &messageStopData)
+ if err != nil {
+ return eventData, err
+ }
+ eventData.Data = messageStopData
  case "error":
  var errorData ErrorData
- json.Unmarshal([]byte(jsonData), &errorData)
- data = errorData
+ err := json.Unmarshal([]byte(jsonData), &errorData)
+ if err != nil {
+ return eventData, err
+ }
+ return eventData, fmt.Errorf("error(%s) - %s", errorData.Error.Type, errorData.Error.Message)
  }
  }
  }
 
- return data
+ return eventData, nil
+}
+
+func (d *MessageSSEDecoder) updateContent(index int, content string) {
+ if index >= len(d.content) {
+ d.content = append(d.content, make([]string, index-len(d.content)+1)...)
+ }
+ d.content[index] += content
 }