diff --git a/agent/context.go b/agent/context.go new file mode 100644 index 0000000..c2bb555 --- /dev/null +++ b/agent/context.go @@ -0,0 +1,214 @@ +package agent + +import ( + "context" + nhttp "net/http" + "net/url" + "strconv" + "time" + + "github.com/pkg/errors" + + "github.com/longbridge/openapi-go/config" + httplib "github.com/longbridge/openapi-go/http" +) + +// requestTimeout bounds blocking Conversation/Continue calls. The shared +// http.Client default (see httplib.DefaultTimeout, 15s) is tuned for fast +// REST calls: in blocking mode the server holds the connection silent until +// the whole LLM turn is done, and that can legitimately take longer. +const requestTimeout = 120 * time.Second + +// listTimeout bounds the plain REST calls (Workspaces, Agents). Since +// AgentContext's underlying *http.Client has no overall timeout (see +// NewFromCfg), these would otherwise hang forever on a stuck connection; +// they're ordinary fast REST calls, so the shared default is enough. +const listTimeout = httplib.DefaultTimeout + +// AgentContext is a client for the Longbridge AI Agent conversation OpenAPI. +// +// Example: +// +// conf, err := config.NewFromEnv() +// actx, err := agent.NewFromCfg(conf) +// resp, err := actx.Conversation(context.Background(), "ag_7d3f9b2c", "How has Tesla stock performed recently?", "") +type AgentContext struct { + httpClient *httplib.Client +} + +// NewFromCfg creates an AgentContext from a *config.Config. +// +// Both blocking and streamed conversation calls can legitimately run far +// longer than a typical REST call (a full LLM turn), so unless cfg.Client is +// already set, AgentContext uses its own underlying *http.Client with no +// overall timeout: blocking calls (Conversation, Continue) are instead +// bounded by a dedicated per-request timeout (see requestTimeout), and +// streamed calls (ConversationStream, ContinueStream) are left unbounded — +// only ctx cancellation stops them. +func NewFromCfg(cfg *config.Config) (*AgentContext, error) { + agentCfg := *cfg + if agentCfg.Client == nil { + agentCfg.Client = &nhttp.Client{ + Transport: &nhttp.Transport{IdleConnTimeout: 60 * time.Second}, + } + } + httpClient, err := httplib.NewFromCfg(&agentCfg) + if err != nil { + return nil, errors.Wrap(err, "create http client error") + } + return &AgentContext{httpClient: httpClient}, nil +} + +// NewFromEnv returns an AgentContext configured from environment variables. +func NewFromEnv() (*AgentContext, error) { + cfg, err := config.NewFormEnv() + if err != nil { + return nil, errors.Wrap(err, "load config from env error") + } + return NewFromCfg(cfg) +} + +// Workspaces lists the Workspaces the current account belongs to. A +// Workspace is the organizational unit for Agents: find the target +// Workspace here first, then use Agents to list the Agents available in it. +// +// Path: GET /v1/ai/workspaces +func (c *AgentContext) Workspaces(ctx context.Context) (*WorkspacesResponse, error) { + var resp WorkspacesResponse + if err := c.httpClient.Get(ctx, "/v1/ai/workspaces", url.Values{}, &resp, httplib.WithRequestTimeout(listTimeout)); err != nil { + return nil, err + } + return &resp, nil +} + +// Agents lists the Agents in the specified Workspace. opts is optional; pass +// nil to use the server defaults (page 1, limit 20, no name filter). +// +// The returned Agent.UID is the identifier used by Conversation; only Agents +// with IsPublished set can start conversations. +// +// Path: GET /v1/ai/workspaces/{id}/agents +func (c *AgentContext) Agents(ctx context.Context, workspaceID string, opts *GetAgentsOptions) (*AgentsResponse, error) { + q := url.Values{} + if opts != nil { + if opts.Page > 0 { + q.Set("page", strconv.FormatInt(int64(opts.Page), 10)) + } + if opts.Limit > 0 { + q.Set("limit", strconv.FormatInt(int64(opts.Limit), 10)) + } + if opts.Name != "" { + q.Set("name", opts.Name) + } + } + path := "/v1/ai/workspaces/" + url.PathEscape(workspaceID) + "/agents" + var resp AgentsResponse + if err := c.httpClient.Get(ctx, path, q, &resp, httplib.WithRequestTimeout(listTimeout)); err != nil { + return nil, err + } + return &resp, nil +} + +// Conversation asks a question to the specified Agent, blocking until the +// run succeeds, is interrupted, or fails. +// +// chatUID identifies an existing conversation to continue asking in; pass an +// empty string to start a new one. +// +// The Agent generates the answer using capabilities such as market data and +// account access. When the Agent needs more information or confirmation +// from you, the run is interrupted (Status is ConversationStatusInterrupted) +// — send your answers via Continue to resume it. +// +// Path: POST /v1/ai/agents/{id}/conversations +func (c *AgentContext) Conversation(ctx context.Context, agentID, query, chatUID string) (*ConversationResponse, error) { + path := "/v1/ai/agents/" + url.PathEscape(agentID) + "/conversations" + body := conversationBody{Query: query, ChatUID: chatUID} + var resp ConversationResponse + if err := c.httpClient.Post(ctx, path, body, &resp, httplib.WithRequestTimeout(requestTimeout)); err != nil { + return nil, err + } + return &resp, nil +} + +// Continue resumes an interrupted conversation (a Conversation or Continue +// call that returned Status ConversationStatusInterrupted), blocking until +// the run succeeds, is interrupted again, or fails. +// +// chatUID and messageID come from the interrupted response: chatUID is its +// ChatUID, and messageID is its Interrupt.MessageID (as a string). answers +// is keyed by Interrupt.ToolCallID; see AnswersByToolCall's docs for its +// shape. Every question the interrupt asked must be answered. +// +// A single round of conversation may be interrupted multiple times: if the +// run returns ConversationStatusInterrupted again after continuing, call +// Continue again with the new Interrupt. +// +// Path: POST /v1/ai/agents/{id}/conversations/{chat_uid}/messages/{message_id}/continue +func (c *AgentContext) Continue(ctx context.Context, agentID, chatUID, messageID string, answers AnswersByToolCall) (*ConversationResponse, error) { + path := continuePath(agentID, chatUID, messageID) + body := continueBody{AnswersByToolCall: answers} + var resp ConversationResponse + if err := c.httpClient.Post(ctx, path, body, &resp, httplib.WithRequestTimeout(requestTimeout)); err != nil { + return nil, err + } + return &resp, nil +} + +// ConversationStream starts a conversation with the specified Agent, +// returning a *ConversationStream of run-progress events delivered over SSE +// as they arrive, instead of blocking for the final result the way +// Conversation does. +// +// A ConversationStreamEvent of type *WorkflowFinishedEvent carries the run's +// outcome, but it isn't necessarily the last event — the server may still +// emit a few more housekeeping events (e.g. a *ChatTitleUpdatedEvent) before +// closing the connection, so keep calling Next until it returns false +// rather than stopping as soon as you see it. The caller must Close the +// returned stream once done with it. +// +// An interrupted run does not emit a *WorkflowFinishedEvent at all: it emits +// a *HumanInteractionRequiredEvent instead, followed directly by +// *ChatFinishedEvent — send answers via ContinueStream (or Continue) using +// its Interrupt details to resume. +// +// chatUID identifies an existing conversation to continue asking in; pass an +// empty string to start a new one. +// +// Path: POST /v1/ai/agents/{id}/conversations (Accept: text/event-stream) +func (c *AgentContext) ConversationStream(ctx context.Context, agentID, query, chatUID string) (*ConversationStream, error) { + path := "/v1/ai/agents/" + url.PathEscape(agentID) + "/conversations" + body := conversationBody{Query: query, ChatUID: chatUID} + rc, err := c.httpClient.CallSSE(ctx, "POST", path, body) + if err != nil { + return nil, err + } + return newConversationStream(rc, nil), nil +} + +// ContinueStream resumes an interrupted conversation, returning a +// *ConversationStream of run-progress events delivered over SSE as they +// arrive, instead of blocking for the final result the way Continue does. +// See ConversationStream's docs for how to drain it correctly. The caller +// must Close the returned stream once done with it. +// +// Path: POST /v1/ai/agents/{id}/conversations/{chat_uid}/messages/{message_id}/continue (Accept: text/event-stream) +func (c *AgentContext) ContinueStream(ctx context.Context, agentID, chatUID, messageID string, answers AnswersByToolCall) (*ConversationStream, error) { + path := continuePath(agentID, chatUID, messageID) + body := continueBody{AnswersByToolCall: answers} + rc, err := c.httpClient.CallSSE(ctx, "POST", path, body) + if err != nil { + return nil, err + } + // Unlike a brand-new conversation, chat_uid/message_id are already known + // here — seed them so the final ConversationResponse carries them even + // if the server doesn't re-emit a ChatStartedEvent for a continued run. + started := &conversationStarted{ChatUID: chatUID, MessageID: messageID} + return newConversationStream(rc, started), nil +} + +func continuePath(agentID, chatUID, messageID string) string { + return "/v1/ai/agents/" + url.PathEscape(agentID) + + "/conversations/" + url.PathEscape(chatUID) + + "/messages/" + url.PathEscape(messageID) + "/continue" +} diff --git a/agent/stream.go b/agent/stream.go new file mode 100644 index 0000000..e0cd69c --- /dev/null +++ b/agent/stream.go @@ -0,0 +1,227 @@ +package agent + +import ( + "bufio" + "encoding/json" + "io" + "strings" +) + +// conversationStarted holds the chat_uid/message_id captured from an +// earlier ChatStartedEvent (nil if one hasn't been observed yet), threaded +// through to WorkflowFinishedEvent, which doesn't repeat them. +type conversationStarted struct { + ChatUID string + MessageID string +} + +// ConversationStream is an open server-sent-events stream of +// ConversationStreamEvent values from AgentContext.ConversationStream or +// AgentContext.ContinueStream. +// +// Call Next repeatedly to advance through the stream; it returns false once +// the stream ends or an error occurs — call Err to tell the two apart. Call +// Close once done, whether or not the stream was drained to the end. +// +// Usage: +// +// stream, err := actx.ConversationStream(ctx, agentID, query, "") +// if err != nil { +// // handle err +// } +// defer stream.Close() +// for stream.Next() { +// switch e := stream.Event().(type) { +// case *agent.MessageEvent: +// fmt.Print(e.Text) +// case *agent.WorkflowFinishedEvent: +// result = e.ConversationResponse +// case *agent.HumanInteractionRequiredEvent: +// interrupt = e.Interrupt +// } +// } +// if err := stream.Err(); err != nil { +// // handle err +// } +type ConversationStream struct { + body io.ReadCloser + scanner *bufio.Scanner + started *conversationStarted + event ConversationStreamEvent + err error +} + +func newConversationStream(body io.ReadCloser, started *conversationStarted) *ConversationStream { + scanner := bufio.NewScanner(body) + // A single "message" event can carry a sizeable chunk of the answer; + // grow well past bufio.Scanner's 64KB default max line length. + scanner.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + return &ConversationStream{body: body, scanner: scanner, started: started} +} + +// Next advances the stream to the next event. It returns false once the +// stream ends or an error occurs; call Err to distinguish the two. +func (s *ConversationStream) Next() bool { + if s.err != nil { + return false + } + payload, err := s.readEventData() + if err != nil { + if err != io.EOF { + s.err = err + } + return false + } + event, err := s.decodeEvent(payload) + if err != nil { + s.err = err + return false + } + s.event = event + return true +} + +// Event returns the event Next just advanced to. Only valid after a call to +// Next that returned true. +func (s *ConversationStream) Event() ConversationStreamEvent { + return s.event +} + +// Err returns the first error encountered while reading the stream, if any. +func (s *ConversationStream) Err() error { + return s.err +} + +// Close closes the underlying connection. Safe to call more than once. +func (s *ConversationStream) Close() error { + return s.body.Close() +} + +// readEventData reads raw SSE lines up to the next blank line, returning the +// joined "data:" payload. It returns io.EOF once the stream ends cleanly +// with no more data. +func (s *ConversationStream) readEventData() (string, error) { + var dataLines []string + for s.scanner.Scan() { + line := s.scanner.Text() + if line == "" { + if len(dataLines) > 0 { + return strings.Join(dataLines, "\n"), nil + } + continue + } + // Other SSE fields (event:, id:, retry:, or ":" comments) carry no + // information the API uses — the real event type is nested inside + // the "data:" JSON payload's own "event" field (see sseEnvelope). + if rest, ok := cutPrefix(line, "data:"); ok { + dataLines = append(dataLines, strings.TrimPrefix(rest, " ")) + } + } + if err := s.scanner.Err(); err != nil { + return "", err + } + if len(dataLines) > 0 { + return strings.Join(dataLines, "\n"), nil + } + return "", io.EOF +} + +func cutPrefix(s, prefix string) (string, bool) { + if !strings.HasPrefix(s, prefix) { + return s, false + } + return s[len(prefix):], true +} + +// sseEnvelope is the outer JSON object carried by every "data:" line. The +// SSE transport's own event name is always "message" — the real event type +// is this envelope's "event" field instead. +type sseEnvelope struct { + Event string `json:"event"` + Data json.RawMessage `json:"data"` +} + +// decodeEvent parses one SSE frame's data payload into a +// ConversationStreamEvent, threading chat_uid/message_id captured from an +// earlier ChatStartedEvent into WorkflowFinishedEvent (which doesn't repeat +// them). +func (s *ConversationStream) decodeEvent(payload string) (ConversationStreamEvent, error) { + var env sseEnvelope + if err := json.Unmarshal([]byte(payload), &env); err != nil { + return nil, err + } + + switch env.Event { + case "chat_started": + var e ChatStartedEvent + if err := json.Unmarshal(env.Data, &e); err != nil { + return nil, err + } + s.started = &conversationStarted{ChatUID: e.ChatUID, MessageID: e.MessageID} + return &e, nil + + case "workflow_started": + var e WorkflowStartedEvent + if err := json.Unmarshal(env.Data, &e); err != nil { + return nil, err + } + return &e, nil + + case "message": + var e MessageEvent + if err := json.Unmarshal(env.Data, &e); err != nil { + return nil, err + } + return &e, nil + + case "ping": + return &PingEvent{}, nil + + case "chat_finished": + var e ChatFinishedEvent + if err := json.Unmarshal(env.Data, &e); err != nil { + return nil, err + } + return &e, nil + + case "chat_title_updated": + var e ChatTitleUpdatedEvent + if err := json.Unmarshal(env.Data, &e); err != nil { + return nil, err + } + return &e, nil + + case "human_interaction_required": + var e Interrupt + if err := json.Unmarshal(env.Data, &e); err != nil { + return nil, err + } + return &HumanInteractionRequiredEvent{Interrupt: &e}, nil + + case "workflow_finished": + var payload workflowFinishedPayload + if err := json.Unmarshal(env.Data, &payload); err != nil { + return nil, err + } + answer := "" + if payload.Outputs.Answer != nil { + answer = *payload.Outputs.Answer + } + resp := &ConversationResponse{ + Status: payload.Status, + Answer: answer, + References: payload.Outputs.References, + ElapsedTime: payload.ElapsedTime, + Interrupt: payload.Outputs.Interrupt, + Error: payload.Outputs.Error, + } + if s.started != nil { + resp.ChatUID = s.started.ChatUID + resp.MessageID = s.started.MessageID + } + return &WorkflowFinishedEvent{ConversationResponse: resp}, nil + + default: + return &OtherEvent{Event: env.Event, Data: env.Data}, nil + } +} diff --git a/agent/stream_test.go b/agent/stream_test.go new file mode 100644 index 0000000..8e489ee --- /dev/null +++ b/agent/stream_test.go @@ -0,0 +1,266 @@ +package agent + +import ( + "io" + "strings" + "testing" +) + +// sseFrame formats a raw SSE frame the way the server sends it: an +// "event: message" line (the SSE transport's own event name, always +// "message") followed by a "data:" line carrying the real event envelope. +func sseFrame(data string) string { + return "event: message\ndata: " + data + "\n\n" +} + +// The `data:` payloads of the three example SSE frames from +// https://open.longbridge.com/en/docs/ai/chat/conversation, plus the four +// event types the docs don't mention (workflow_started, ping, +// chat_finished, chat_title_updated). +const ( + chatStartedFrame = `{"event":"chat_started","workflow_run_id":"wr_1","data":{"chat_uid":"ct_9f2c1a5b","message_id":42}}` + workflowStartedFrame = `{"event":"workflow_started","workflow_run_id":"wr_1","data":{"hit_cache":false,"inputs":{"chat_id":834552,"chat_uid":"ct_9f2c1a5b","message_id":42,"query":"How has Tesla stock performed recently?"},"started_at":1784545150,"workflow_id":176476}}` + messageFrame = `{"event":"message","workflow_run_id":"wr_1","data":{"text":"Tesla"}}` + pingFrame = `{"event":"ping","workflow_run_id":"wr_1","data":null}` + chatFinishedFrame = `{"event":"chat_finished","workflow_run_id":"wr_1","data":{"chat_id":834552,"chat_uid":"ct_9f2c1a5b","error":"","error_message":"","message_id":42}}` + workflowFinishedFrame = `{"event":"workflow_finished","workflow_run_id":"wr_1","data":{"status":"succeeded","elapsed_time":3.21,"outputs":{"answer":"Tesla (TSLA.US) recently..."}}}` + chatTitleUpdatedFrame = `{"event":"chat_title_updated","workflow_run_id":"wr_1","data":{"chat_id":834552,"chat_uid":"ct_9f2c1a5b","source":"ai_generated","title":"Tesla stock performance","updated_at":1784546957}}` + + // From https://open.longbridge.com/en/docs/ai/chat/events: the interrupt + // event fired mid-run instead of workflow_finished. + humanInteractionRequiredFrame = `{"event":"human_interaction_required","workflow_run_id":"wr_1","data":{"node_id":"n_ask_human","tool_call_id":"call_abc123","questions":[{"question":"Which time range would you like to check?","options":[{"description":"Past week"},{"description":"Past month"}],"multi_select":false}],"message_id":43,"chat_id":1001}}` +) + +func newTestStream(frames ...string) *ConversationStream { + body := io.NopCloser(strings.NewReader(strings.Join(frames, ""))) + return newConversationStream(body, nil) +} + +func TestConversationStreamFullSequence(t *testing.T) { + // Exercises the fuller, real-world sequence beyond the docs' three-event + // example, including chat_title_updated arriving *after* + // workflow_finished (observed live against the real API). + stream := newTestStream( + sseFrame(chatStartedFrame), + sseFrame(workflowStartedFrame), + sseFrame(messageFrame), + sseFrame(pingFrame), + sseFrame(chatFinishedFrame), + sseFrame(workflowFinishedFrame), + sseFrame(chatTitleUpdatedFrame), + ) + defer stream.Close() + + if !stream.Next() { + t.Fatalf("expected chat_started event, got err: %v", stream.Err()) + } + chatStarted, ok := stream.Event().(*ChatStartedEvent) + if !ok { + t.Fatalf("expected *ChatStartedEvent, got %T", stream.Event()) + } + if chatStarted.ChatUID != "ct_9f2c1a5b" || chatStarted.MessageID != "42" { + t.Errorf("unexpected ChatStartedEvent: %+v", chatStarted) + } + + if !stream.Next() { + t.Fatalf("expected workflow_started event, got err: %v", stream.Err()) + } + workflowStarted, ok := stream.Event().(*WorkflowStartedEvent) + if !ok { + t.Fatalf("expected *WorkflowStartedEvent, got %T", stream.Event()) + } + if workflowStarted.HitCache { + t.Error("expected HitCache = false") + } + if workflowStarted.Inputs.ChatUID != "ct_9f2c1a5b" || workflowStarted.Inputs.MessageID != "42" { + t.Errorf("unexpected WorkflowStartedEvent.Inputs: %+v", workflowStarted.Inputs) + } + if workflowStarted.WorkflowID != 176476 { + t.Errorf("WorkflowID = %d, want 176476", workflowStarted.WorkflowID) + } + + if !stream.Next() { + t.Fatalf("expected message event, got err: %v", stream.Err()) + } + message, ok := stream.Event().(*MessageEvent) + if !ok { + t.Fatalf("expected *MessageEvent, got %T", stream.Event()) + } + if message.Text != "Tesla" { + t.Errorf("Text = %q, want %q", message.Text, "Tesla") + } + + if !stream.Next() { + t.Fatalf("expected ping event, got err: %v", stream.Err()) + } + if _, ok := stream.Event().(*PingEvent); !ok { + t.Fatalf("expected *PingEvent, got %T", stream.Event()) + } + + if !stream.Next() { + t.Fatalf("expected chat_finished event, got err: %v", stream.Err()) + } + chatFinished, ok := stream.Event().(*ChatFinishedEvent) + if !ok { + t.Fatalf("expected *ChatFinishedEvent, got %T", stream.Event()) + } + if chatFinished.ChatUID != "ct_9f2c1a5b" || chatFinished.MessageID != "42" { + t.Errorf("unexpected ChatFinishedEvent: %+v", chatFinished) + } + if chatFinished.Error != "" || chatFinished.ErrorMessage != "" { + t.Errorf("expected empty Error/ErrorMessage, got %+v", chatFinished) + } + + if !stream.Next() { + t.Fatalf("expected workflow_finished event, got err: %v", stream.Err()) + } + workflowFinished, ok := stream.Event().(*WorkflowFinishedEvent) + if !ok { + t.Fatalf("expected *WorkflowFinishedEvent, got %T", stream.Event()) + } + if workflowFinished.ChatUID != "ct_9f2c1a5b" || workflowFinished.MessageID != "42" { + t.Errorf("expected chat_uid/message_id threaded from chat_started, got %+v", workflowFinished.ConversationResponse) + } + if workflowFinished.Status != ConversationStatusSucceeded { + t.Errorf("Status = %v, want %v", workflowFinished.Status, ConversationStatusSucceeded) + } + if workflowFinished.Answer != "Tesla (TSLA.US) recently..." { + t.Errorf("Answer = %q", workflowFinished.Answer) + } + + // Arrives *after* workflow_finished in this (real, observed) ordering. + if !stream.Next() { + t.Fatalf("expected chat_title_updated event, got err: %v", stream.Err()) + } + chatTitleUpdated, ok := stream.Event().(*ChatTitleUpdatedEvent) + if !ok { + t.Fatalf("expected *ChatTitleUpdatedEvent, got %T", stream.Event()) + } + if chatTitleUpdated.ChatUID != "ct_9f2c1a5b" || chatTitleUpdated.Source != "ai_generated" || chatTitleUpdated.Title != "Tesla stock performance" { + t.Errorf("unexpected ChatTitleUpdatedEvent: %+v", chatTitleUpdated) + } + + if stream.Next() { + t.Fatalf("expected stream to end, got event %T", stream.Event()) + } + if err := stream.Err(); err != nil { + t.Errorf("unexpected error at end of stream: %v", err) + } +} + +func TestConversationStreamInterruptedSequence(t *testing.T) { + // Per the docs, an interrupted run emits human_interaction_required and + // then chat_finished directly — no workflow_finished in between. + stream := newTestStream( + sseFrame(chatStartedFrame), + sseFrame(workflowStartedFrame), + sseFrame(humanInteractionRequiredFrame), + sseFrame(chatFinishedFrame), + ) + defer stream.Close() + + if !stream.Next() { + t.Fatalf("expected chat_started event, got err: %v", stream.Err()) + } + if _, ok := stream.Event().(*ChatStartedEvent); !ok { + t.Fatalf("expected *ChatStartedEvent, got %T", stream.Event()) + } + + if !stream.Next() { + t.Fatalf("expected workflow_started event, got err: %v", stream.Err()) + } + if _, ok := stream.Event().(*WorkflowStartedEvent); !ok { + t.Fatalf("expected *WorkflowStartedEvent, got %T", stream.Event()) + } + + if !stream.Next() { + t.Fatalf("expected human_interaction_required event, got err: %v", stream.Err()) + } + interrupted, ok := stream.Event().(*HumanInteractionRequiredEvent) + if !ok { + t.Fatalf("expected *HumanInteractionRequiredEvent, got %T", stream.Event()) + } + if interrupted.NodeID != "n_ask_human" || interrupted.ToolCallID != "call_abc123" { + t.Errorf("unexpected HumanInteractionRequiredEvent: %+v", interrupted.Interrupt) + } + if interrupted.MessageID != 43 || interrupted.ChatID != 1001 { + t.Errorf("unexpected MessageID/ChatID: %+v", interrupted.Interrupt) + } + if len(interrupted.Questions) != 1 || interrupted.Questions[0].Question != "Which time range would you like to check?" { + t.Fatalf("unexpected Questions: %+v", interrupted.Questions) + } + if len(interrupted.Questions[0].Options) != 2 || interrupted.Questions[0].MultiSelect { + t.Errorf("unexpected Questions[0]: %+v", interrupted.Questions[0]) + } + + if !stream.Next() { + t.Fatalf("expected chat_finished event, got err: %v", stream.Err()) + } + if _, ok := stream.Event().(*ChatFinishedEvent); !ok { + t.Fatalf("expected *ChatFinishedEvent, got %T", stream.Event()) + } + + if stream.Next() { + t.Fatalf("expected stream to end, got event %T", stream.Event()) + } + if err := stream.Err(); err != nil { + t.Errorf("unexpected error at end of stream: %v", err) + } +} + +func TestConversationStreamUnknownEventFallsBackToOther(t *testing.T) { + stream := newTestStream(sseFrame(`{"event":"some_future_event","data":{"foo":"bar"}}`)) + defer stream.Close() + + if !stream.Next() { + t.Fatalf("expected an event, got err: %v", stream.Err()) + } + other, ok := stream.Event().(*OtherEvent) + if !ok { + t.Fatalf("expected *OtherEvent, got %T", stream.Event()) + } + if other.Event != "some_future_event" { + t.Errorf("Event = %q, want %q", other.Event, "some_future_event") + } + if !strings.Contains(string(other.Data), `"foo":"bar"`) { + t.Errorf("Data = %s", other.Data) + } +} + +func TestConversationStreamWithoutTrailingBlankLine(t *testing.T) { + // The final frame of a real stream may not be followed by a blank line + // before the connection closes; the reader must still surface it. + body := io.NopCloser(strings.NewReader("event: message\ndata: " + messageFrame)) + stream := newConversationStream(body, nil) + defer stream.Close() + + if !stream.Next() { + t.Fatalf("expected an event, got err: %v", stream.Err()) + } + if _, ok := stream.Event().(*MessageEvent); !ok { + t.Fatalf("expected *MessageEvent, got %T", stream.Event()) + } + if stream.Next() { + t.Fatalf("expected stream to end, got event %T", stream.Event()) + } +} + +func TestContinueStreamSeedsStartedInfo(t *testing.T) { + // ContinueStream already knows chat_uid/message_id from the caller + // (unlike a brand-new conversation), so WorkflowFinishedEvent should + // carry them even without an intervening ChatStartedEvent. + body := io.NopCloser(strings.NewReader(sseFrame(workflowFinishedFrame))) + stream := newConversationStream(body, &conversationStarted{ChatUID: "ct_9f2c1a5b", MessageID: "43"}) + defer stream.Close() + + if !stream.Next() { + t.Fatalf("expected an event, got err: %v", stream.Err()) + } + workflowFinished, ok := stream.Event().(*WorkflowFinishedEvent) + if !ok { + t.Fatalf("expected *WorkflowFinishedEvent, got %T", stream.Event()) + } + if workflowFinished.ChatUID != "ct_9f2c1a5b" || workflowFinished.MessageID != "43" { + t.Errorf("unexpected ConversationResponse: %+v", workflowFinished.ConversationResponse) + } +} diff --git a/agent/types.go b/agent/types.go new file mode 100644 index 0000000..2c5c20b --- /dev/null +++ b/agent/types.go @@ -0,0 +1,472 @@ +// Package agent provides a client for the Longbridge AI Agent conversation +// OpenAPI. It covers listing Workspaces and Agents, starting and continuing +// conversations, and streaming run progress over SSE. +// +// Reference: https://open.longbridge.com/en/docs/ai/chat/conversation +package agent + +import ( + "bytes" + "encoding/json" +) + +// ─── Workspaces ─────────────────────────────────────────────────────────── + +// Workspace is a Workspace the current account belongs to. +type Workspace struct { + // ID is the Workspace ID. + ID string `json:"id"` + // Name is the Workspace name. + Name string `json:"name"` + // CreatedAt is the creation time, Unix timestamp in seconds. + CreatedAt int64 `json:"created_at"` + // UpdatedAt is the last updated time, Unix timestamp in seconds. + UpdatedAt int64 `json:"updated_at"` +} + +// WorkspacesResponse is the response for AgentContext.Workspaces. +type WorkspacesResponse struct { + // Workspaces the current account belongs to. + Workspaces []Workspace `json:"workspaces"` +} + +// ─── Agents ─────────────────────────────────────────────────────────────── + +// Agent is an Agent in a Workspace. +type Agent struct { + // UID is the Agent UID, used as the path parameter of + // AgentContext.Conversation. + UID string `json:"uid"` + // Name is the Agent name. + Name string `json:"name"` + // Description is the Agent description. + Description string `json:"description"` + // Mode is the Agent mode, e.g. "chat". + Mode string `json:"mode"` + // Icon is the icon URL. + Icon string `json:"icon"` + // IsPublished reports whether the Agent is published; only published + // Agents can start conversations. + IsPublished bool `json:"is_published"` + // PublishedAt is the publish time, Unix timestamp in seconds; 0 if + // unpublished. + PublishedAt int64 `json:"published_at"` + // CreatedAt is the creation time, Unix timestamp in seconds. + CreatedAt int64 `json:"created_at"` + // UpdatedAt is the last updated time, Unix timestamp in seconds. + UpdatedAt int64 `json:"updated_at"` +} + +// AgentsResponse is the response for AgentContext.Agents. +type AgentsResponse struct { + // Agents is the Agent list. + Agents []Agent `json:"agents"` + // Total is the total number of matching Agents. + Total int32 `json:"total"` +} + +// GetAgentsOptions holds the optional query parameters for +// AgentContext.Agents. The zero value uses the server defaults (page 1, +// limit 20, no name filter). +type GetAgentsOptions struct { + // Page is the page number, starting at 1. 0 uses the server default. + Page int32 + // Limit is the page size. 0 uses the server default. + Limit int32 + // Name fuzzy-searches by Agent name. Empty omits the filter. + Name string +} + +// ─── Conversation ───────────────────────────────────────────────────────── + +// ConversationStatus is the final run status of a conversation. +type ConversationStatus string + +const ( + // ConversationStatusSucceeded means the run completed successfully. + ConversationStatusSucceeded ConversationStatus = "succeeded" + // ConversationStatusInterrupted means the run is paused, waiting for + // AgentContext.Continue. + ConversationStatusInterrupted ConversationStatus = "interrupted" + // ConversationStatusFailed means the run failed. + ConversationStatusFailed ConversationStatus = "failed" + // ConversationStatusStopped means the run was stopped. + ConversationStatusStopped ConversationStatus = "stopped" +) + +// Reference is a source referenced by the answer. +type Reference struct { + // Index is the reference index. + Index int32 `json:"index"` + // Title is the reference title. + Title string `json:"title"` + // URL is the reference URL. + URL string `json:"url"` +} + +// QuestionOption is one option of a Question. +type QuestionOption struct { + // Description is the option text. + Description string `json:"description"` +} + +// Question is one question the Agent needs you to answer to resume an +// interrupted conversation. +type Question struct { + // Question is the question text. + Question string `json:"question"` + // Options are the answer options; empty means free-form answer. + Options []QuestionOption `json:"options"` + // MultiSelect reports whether multiple options may be selected. + MultiSelect bool `json:"multi_select"` +} + +// Interrupt is present when a conversation run is interrupted, waiting for +// AgentContext.Continue. +type Interrupt struct { + // NodeID is the ID of the node that triggered the interrupt. + NodeID string `json:"node_id"` + // ToolCallID is the tool call ID of this inquiry; used as the outer key + // of AnswersByToolCall when continuing. + ToolCallID string `json:"tool_call_id"` + // Questions you need to answer. + Questions []Question `json:"questions"` + // MessageID is the ID of the paused message. + MessageID int64 `json:"message_id"` + // ChatID is the ID of the owning conversation. + ChatID int64 `json:"chat_id"` +} + +// ConversationError is present when a conversation run failed. +type ConversationError struct { + // Code is the error code. + Code int32 `json:"code"` + // Message is the error message. + Message string `json:"message"` +} + +// AnswersByToolCall is answers keyed by Interrupt.ToolCallID, each value +// being a map of question text to answer, used as the request body of +// AgentContext.Continue and AgentContext.ContinueStream. +// +// The outer key corresponds to Interrupt.ToolCallID. The inner key is the +// Question.Question text, and the value is the answer selected or typed for +// it. Every question the interrupt asked must be answered. +type AnswersByToolCall map[string]map[string]string + +// ConversationResponse is the response for AgentContext.Conversation and +// AgentContext.Continue, and the outcome carried by WorkflowFinishedEvent +// when streaming. +type ConversationResponse struct { + // ChatUID identifies the conversation, used for follow-up questions and + // troubleshooting. + ChatUID string `json:"chat_uid"` + // MessageID is the message ID of this round. + MessageID string `json:"message_id"` + // Status is the final run status. + Status ConversationStatus `json:"status"` + // Answer is the final answer text; valid when Status is + // ConversationStatusSucceeded. + Answer string `json:"answer"` + // References are sources referenced by the answer; nil if none. + References []Reference `json:"references"` + // ElapsedTime is the run duration in seconds. + ElapsedTime float64 `json:"elapsed_time"` + // Interrupt is present only when Status is ConversationStatusInterrupted. + Interrupt *Interrupt `json:"interrupt"` + // Error is present only when the run failed. + Error *ConversationError `json:"error"` +} + +// UnmarshalJSON accepts message_id as either a JSON string or a JSON +// number: the blocking response's top-level message_id is a quoted string, +// but callers building a ConversationResponse from streamed events may feed +// it the numeric form seen in some SSE payloads (see ChatStartedEvent). +func (r *ConversationResponse) UnmarshalJSON(data []byte) error { + var raw struct { + ChatUID string `json:"chat_uid"` + MessageID json.RawMessage `json:"message_id"` + Status ConversationStatus `json:"status"` + Answer string `json:"answer"` + References []Reference `json:"references"` + ElapsedTime float64 `json:"elapsed_time"` + Interrupt *Interrupt `json:"interrupt"` + Error *ConversationError `json:"error"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + messageID, err := unmarshalStringOrInt(raw.MessageID) + if err != nil { + return err + } + r.ChatUID = raw.ChatUID + r.MessageID = messageID + r.Status = raw.Status + r.Answer = raw.Answer + r.References = raw.References + r.ElapsedTime = raw.ElapsedTime + r.Interrupt = raw.Interrupt + r.Error = raw.Error + return nil +} + +// conversationBody is the request body of AgentContext.Conversation and +// AgentContext.ConversationStream. +type conversationBody struct { + Query string `json:"query"` + ChatUID string `json:"chat_uid,omitempty"` +} + +// continueBody is the request body of AgentContext.Continue and +// AgentContext.ContinueStream. +type continueBody struct { + AnswersByToolCall AnswersByToolCall `json:"answers_by_tool_call"` +} + +// ─── Streaming events ───────────────────────────────────────────────────── + +// ConversationStreamEvent is implemented by every event observed while +// draining a *ConversationStream. Use a type switch to handle it: +// +// switch e := stream.Event().(type) { +// case *agent.MessageEvent: +// fmt.Print(e.Text) +// case *agent.WorkflowFinishedEvent: +// result = e.ConversationResponse +// case *agent.HumanInteractionRequiredEvent: +// interrupt = e.Interrupt +// } +type ConversationStreamEvent interface { + conversationStreamEvent() +} + +// ChatStartedEvent reports that the run has started. +type ChatStartedEvent struct { + // ChatUID identifies the conversation. + ChatUID string `json:"chat_uid"` + // MessageID is the message ID of this round. + MessageID string `json:"message_id"` +} + +func (*ChatStartedEvent) conversationStreamEvent() {} + +// UnmarshalJSON accepts message_id as either a JSON string or a JSON +// number: the docs' SSE example encodes it as a raw number here, unlike the +// blocking response's top-level message_id, which is a quoted string. +func (e *ChatStartedEvent) UnmarshalJSON(data []byte) error { + var raw struct { + ChatUID string `json:"chat_uid"` + MessageID json.RawMessage `json:"message_id"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + messageID, err := unmarshalStringOrInt(raw.MessageID) + if err != nil { + return err + } + e.ChatUID = raw.ChatUID + e.MessageID = messageID + return nil +} + +// WorkflowStartedInputs echoes the run's inputs, carried by +// WorkflowStartedEvent. +type WorkflowStartedInputs struct { + // ChatID is the ID of the owning conversation. + ChatID int64 `json:"chat_id"` + // ChatUID identifies the conversation. + ChatUID string `json:"chat_uid"` + // MessageID is the message ID of this round. + MessageID string `json:"message_id"` + // Query is the question that was asked. + Query string `json:"query"` +} + +// UnmarshalJSON accepts message_id as either a JSON string or a JSON +// number; see ChatStartedEvent.UnmarshalJSON. +func (i *WorkflowStartedInputs) UnmarshalJSON(data []byte) error { + var raw struct { + ChatID int64 `json:"chat_id"` + ChatUID string `json:"chat_uid"` + MessageID json.RawMessage `json:"message_id"` + Query string `json:"query"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + messageID, err := unmarshalStringOrInt(raw.MessageID) + if err != nil { + return err + } + i.ChatID = raw.ChatID + i.ChatUID = raw.ChatUID + i.MessageID = messageID + i.Query = raw.Query + return nil +} + +// WorkflowStartedEvent is observed right after ChatStartedEvent on every run +// seen so far. +type WorkflowStartedEvent struct { + // HitCache reports whether this run's answer was served from a cache. + HitCache bool `json:"hit_cache"` + // Inputs echoes the run's inputs. + Inputs WorkflowStartedInputs `json:"inputs"` + // StartedAt is a Unix timestamp in seconds. + StartedAt int64 `json:"started_at"` + // WorkflowID is an internal workflow run ID. + WorkflowID int64 `json:"workflow_id"` +} + +func (*WorkflowStartedEvent) conversationStreamEvent() {} + +// MessageEvent is an incremental piece of the answer. +type MessageEvent struct { + // Text is the incremental answer text. + Text string `json:"text"` +} + +func (*MessageEvent) conversationStreamEvent() {} + +// PingEvent is a heartbeat with no payload, observed at arbitrary points in +// the stream (including in between MessageEvent chunks). +type PingEvent struct{} + +func (*PingEvent) conversationStreamEvent() {} + +// ChatFinishedEvent is observed once all MessageEvents for this round have +// been sent, shortly before WorkflowFinishedEvent. +type ChatFinishedEvent struct { + // ChatID is the ID of the owning conversation. + ChatID int64 `json:"chat_id"` + // ChatUID identifies the conversation. + ChatUID string `json:"chat_uid"` + // MessageID is the message ID of this round. + MessageID string `json:"message_id"` + // Error has been empty in every run observed so far. + Error string `json:"error"` + // ErrorMessage has been empty in every run observed so far. + ErrorMessage string `json:"error_message"` +} + +func (*ChatFinishedEvent) conversationStreamEvent() {} + +// UnmarshalJSON accepts message_id as either a JSON string or a JSON +// number; see ChatStartedEvent.UnmarshalJSON. +func (e *ChatFinishedEvent) UnmarshalJSON(data []byte) error { + var raw struct { + ChatID int64 `json:"chat_id"` + ChatUID string `json:"chat_uid"` + MessageID json.RawMessage `json:"message_id"` + Error string `json:"error"` + ErrorMessage string `json:"error_message"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + messageID, err := unmarshalStringOrInt(raw.MessageID) + if err != nil { + return err + } + e.ChatID = raw.ChatID + e.ChatUID = raw.ChatUID + e.MessageID = messageID + e.Error = raw.Error + e.ErrorMessage = raw.ErrorMessage + return nil +} + +// WorkflowFinishedEvent reports that the run finished (succeeded, failed, or +// stopped), carrying the run's outcome. Not emitted when the run is +// interrupted instead — see HumanInteractionRequiredEvent. Not necessarily +// the last event of the stream — the server may still emit a few more +// housekeeping events (e.g. ChatTitleUpdatedEvent) before actually closing +// the connection, so keep draining the stream until Next returns false +// rather than stopping as soon as this is seen. +type WorkflowFinishedEvent struct { + *ConversationResponse +} + +func (*WorkflowFinishedEvent) conversationStreamEvent() {} + +// HumanInteractionRequiredEvent reports that the run is paused, waiting for +// more information from the caller — send answers via AgentContext.Continue +// or AgentContext.ContinueStream to resume it. An interrupted run does not +// emit WorkflowFinishedEvent; this event carries the interrupt details +// instead, and the stream ends with ChatFinishedEvent right after it. +type HumanInteractionRequiredEvent struct { + *Interrupt +} + +func (*HumanInteractionRequiredEvent) conversationStreamEvent() {} + +// ChatTitleUpdatedEvent reports the server auto-generating a short title for +// the conversation, as a UI convenience. Can arrive before or after +// WorkflowFinishedEvent; not tied to the run's outcome. +type ChatTitleUpdatedEvent struct { + // ChatID is the ID of the owning conversation. + ChatID int64 `json:"chat_id"` + // ChatUID identifies the conversation. + ChatUID string `json:"chat_uid"` + // Source describes where the title came from, e.g. "ai_generated". + Source string `json:"source"` + // Title is the new (possibly truncated) title. + Title string `json:"title"` + // UpdatedAt is a Unix timestamp in seconds. + UpdatedAt int64 `json:"updated_at"` +} + +func (*ChatTitleUpdatedEvent) conversationStreamEvent() {} + +// OtherEvent is an event type not recognized by this SDK version, carried as +// raw JSON so callers aren't broken by future additions to the API. Event +// is the SSE envelope's discriminator string, so callers can at least tell +// these apart instead of getting an opaque blob. +type OtherEvent struct { + // Event is the SSE envelope's "event" field (the event type name). + Event string + // Data is the SSE envelope's "data" field. + Data json.RawMessage +} + +func (*OtherEvent) conversationStreamEvent() {} + +// workflowOutputs are the outputs of a workflow_finished SSE event. +type workflowOutputs struct { + // Answer is the final answer text; present when the run succeeded. + Answer *string `json:"answer"` + // References are sources referenced by the answer. + References []Reference `json:"references"` + // Interrupt is present only when the status is "interrupted". + Interrupt *Interrupt `json:"interrupt"` + // Error is present only when the run failed. + Error *ConversationError `json:"error"` +} + +// workflowFinishedPayload is the payload of a workflow_finished SSE event. +type workflowFinishedPayload struct { + Status ConversationStatus `json:"status"` + ElapsedTime float64 `json:"elapsed_time"` + Outputs workflowOutputs `json:"outputs"` +} + +// unmarshalStringOrInt decodes data as either a JSON string or a JSON +// number, returning its value as a string either way. Some SSE event +// payloads encode message_id as a raw JSON number, unlike the blocking +// response's top-level message_id, which is a quoted string. +func unmarshalStringOrInt(data []byte) (string, error) { + data = bytes.TrimSpace(data) + if len(data) == 0 || string(data) == "null" { + return "", nil + } + if data[0] == '"' { + var s string + if err := json.Unmarshal(data, &s); err != nil { + return "", err + } + return s, nil + } + return string(data), nil +} diff --git a/agent/types_test.go b/agent/types_test.go new file mode 100644 index 0000000..bfd09d5 --- /dev/null +++ b/agent/types_test.go @@ -0,0 +1,179 @@ +package agent + +import ( + "encoding/json" + "testing" +) + +// The `data` payload of the "Run succeeded" example from +// https://open.longbridge.com/en/docs/ai/chat/conversation +const succeededJSON = `{ + "chat_uid": "ct_9f2c1a5b", + "message_id": "42", + "status": "succeeded", + "answer": "Tesla (TSLA.US) recently...", + "references": [ + { "index": 1, "title": "...", "url": "..." } + ], + "elapsed_time": 3.21 +}` + +// The `data` payload of the "Run interrupted" example from the same page. +const interruptedJSON = `{ + "chat_uid": "ct_9f2c1a5b", + "message_id": "43", + "status": "interrupted", + "answer": "", + "references": null, + "elapsed_time": 1.05, + "interrupt": { + "node_id": "n_ask_human", + "tool_call_id": "call_abc123", + "questions": [ + { + "question": "Which time range would you like to check?", + "options": [ + { "description": "Past week" }, + { "description": "Past month" } + ], + "multi_select": false + } + ], + "message_id": 43, + "chat_id": 1001 + } +}` + +func TestUnmarshalSucceededConversationResponse(t *testing.T) { + var resp ConversationResponse + if err := json.Unmarshal([]byte(succeededJSON), &resp); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if resp.ChatUID != "ct_9f2c1a5b" { + t.Errorf("ChatUID = %q", resp.ChatUID) + } + if resp.MessageID != "42" { + t.Errorf("MessageID = %q", resp.MessageID) + } + if resp.Status != ConversationStatusSucceeded { + t.Errorf("Status = %v", resp.Status) + } + if resp.Answer != "Tesla (TSLA.US) recently..." { + t.Errorf("Answer = %q", resp.Answer) + } + if len(resp.References) != 1 || resp.References[0].Index != 1 { + t.Errorf("References = %+v", resp.References) + } + if resp.ElapsedTime != 3.21 { + t.Errorf("ElapsedTime = %v", resp.ElapsedTime) + } + if resp.Interrupt != nil { + t.Errorf("Interrupt = %+v, want nil", resp.Interrupt) + } + if resp.Error != nil { + t.Errorf("Error = %+v, want nil", resp.Error) + } +} + +func TestUnmarshalInterruptedConversationResponse(t *testing.T) { + var resp ConversationResponse + if err := json.Unmarshal([]byte(interruptedJSON), &resp); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if resp.Status != ConversationStatusInterrupted { + t.Errorf("Status = %v", resp.Status) + } + if resp.References != nil { + t.Errorf("References = %+v, want nil", resp.References) + } + interrupt := resp.Interrupt + if interrupt == nil { + t.Fatal("Interrupt = nil, want non-nil") + } + if interrupt.NodeID != "n_ask_human" { + t.Errorf("NodeID = %q", interrupt.NodeID) + } + if interrupt.ToolCallID != "call_abc123" { + t.Errorf("ToolCallID = %q", interrupt.ToolCallID) + } + if interrupt.MessageID != 43 { + t.Errorf("MessageID = %d", interrupt.MessageID) + } + if interrupt.ChatID != 1001 { + t.Errorf("ChatID = %d", interrupt.ChatID) + } + if len(interrupt.Questions) != 1 { + t.Fatalf("Questions = %+v", interrupt.Questions) + } + q := interrupt.Questions[0] + if len(q.Options) != 2 { + t.Errorf("Options = %+v", q.Options) + } + if q.MultiSelect { + t.Error("MultiSelect = true, want false") + } +} + +func TestUnmarshalWorkspacesResponse(t *testing.T) { + const j = `{ + "workspaces": [ + { "id": "1001", "name": "My Workspace", "created_at": 1742000000, "updated_at": 1742001000 } + ] + }` + var resp WorkspacesResponse + if err := json.Unmarshal([]byte(j), &resp); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if len(resp.Workspaces) != 1 || resp.Workspaces[0].ID != "1001" { + t.Errorf("Workspaces = %+v", resp.Workspaces) + } +} + +func TestUnmarshalAgentsResponse(t *testing.T) { + const j = `{ + "agents": [ + { + "uid": "ag_7d3f9b2c", + "name": "US Stock Analyst", + "description": "Answers US stock questions with market and fundamental data", + "mode": "chat", + "icon": "https://cdn.longbridge.com/icons/agent.png", + "is_published": true, + "published_at": 1742000000, + "created_at": 1741000000, + "updated_at": 1742001000 + } + ], + "total": 12 + }` + var resp AgentsResponse + if err := json.Unmarshal([]byte(j), &resp); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if resp.Total != 12 { + t.Errorf("Total = %d", resp.Total) + } + if len(resp.Agents) != 1 || resp.Agents[0].UID != "ag_7d3f9b2c" { + t.Errorf("Agents = %+v", resp.Agents) + } + if !resp.Agents[0].IsPublished { + t.Error("IsPublished = false, want true") + } +} + +func TestContinueBodyMarshal(t *testing.T) { + body := continueBody{AnswersByToolCall: AnswersByToolCall{ + "call_abc123": {"Which time range would you like to check?": "Past month"}, + }} + b, err := json.Marshal(body) + if err != nil { + t.Fatalf("marshal error: %v", err) + } + var round continueBody + if err := json.Unmarshal(b, &round); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if round.AnswersByToolCall["call_abc123"]["Which time range would you like to check?"] != "Past month" { + t.Errorf("round-tripped body = %+v", round) + } +} diff --git a/http/client.go b/http/client.go index 6754dcd..21af944 100644 --- a/http/client.go +++ b/http/client.go @@ -39,6 +39,9 @@ type RequestOptions struct { // Request Header Header nhttp.Header body interface{} + // Timeout overrides the request timeout for a single Call. Ignored by + // CallSSE (see its docs). + Timeout time.Duration } // RequestOption use to set addition info to request @@ -60,6 +63,18 @@ func WithBody(v interface{}) RequestOption { } } +// WithRequestTimeout overrides the request timeout for a single Call, in +// place of the timeout the underlying *http.Client was configured with. +// Ignored by CallSSE, whose whole point is to keep receiving events for as +// long as the underlying operation takes; only ctx cancellation bounds it. +func WithRequestTimeout(d time.Duration) RequestOption { + return func(o *RequestOptions) { + if d > 0 { + o.Timeout = d + } + } +} + // region returns the data center region derived from the client's credentials. func (c *Client) region() dcRegion { if c.opts.OAuthClient != nil { @@ -132,19 +147,11 @@ func (c *Client) GetOTPV2(ctx context.Context, ropts ...RequestOption) (string, return res.Otp, nil } -// Call will send request with signature to http server -func (c *Client) Call(ctx context.Context, method, path string, queryParams interface{}, body interface{}, resp interface{}, ropts ...RequestOption) (err error) { - var ( - br io.Reader - bb []byte - httpResp *nhttp.Response - rb []byte - ) - - ro := &RequestOptions{} - for _, opt := range ropts { - opt(ro) - } +// buildRequest builds and signs the underlying *http.Request shared by Call +// and CallSSE: marshals the body, resolves auth (API key or OAuth) and the +// DC region, sets headers and query params, and signs the request. +func (c *Client) buildRequest(ctx context.Context, method, path string, queryParams interface{}, body interface{}, ro *RequestOptions) (req *nhttp.Request, bb []byte, err error) { + var br io.Reader if body == nil && ro.body != nil { body = ro.body @@ -153,14 +160,14 @@ func (c *Client) Call(ctx context.Context, method, path string, queryParams inte if body != nil { bb, err = json.Marshal(body) if err != nil { - return err + return nil, nil, err } br = bytes.NewBuffer(bb) } - req, err := nhttp.NewRequestWithContext(ctx, method, c.opts.URL+path, br) + req, err = nhttp.NewRequestWithContext(ctx, method, c.opts.URL+path, br) if err != nil { - return err + return nil, nil, err } appKey := c.opts.AppKey @@ -170,7 +177,7 @@ func (c *Client) Call(ctx context.Context, method, path string, queryParams inte if c.opts.OAuthClient != nil { token, err := c.opts.OAuthClient.AccessToken(ctx) if err != nil { - return err + return nil, nil, err } // Derive DC region from the token prefix ("us_" → US, otherwise AP), // then strip the prefix so only the bare token is sent to the gateway. @@ -207,7 +214,7 @@ func (c *Client) Call(ctx context.Context, method, path string, queryParams inte vals, ok := queryParams.(url.Values) if !ok { if vals, err = query.Values(queryParams); err != nil { - return + return nil, nil, err } } req.URL.RawQuery = vals.Encode() @@ -215,15 +222,62 @@ func (c *Client) Call(ctx context.Context, method, path string, queryParams inte // set signature (no-op for OAuth when appSecret is empty) signature(req, appSecret, bb) + return req, bb, nil +} + +// readErrorBody reads a non-streaming error response body and converts it +// to an *ApiError. Shared by Call's non-2xx path and CallSSE's fallback path +// (a request that asked for text/event-stream but got a one-shot JSON error +// body instead, e.g. for a 4xx/5xx before the server committed to SSE). +func readErrorBody(httpResp *nhttp.Response) error { + defer httpResp.Body.Close() + rb, err := io.ReadAll(httpResp.Body) + if err != nil { + return err + } + log.Debugf("http call response body:%s", rb) + + apiResp := &apiResponse{} + if v := httpResp.Header.Get("x-trace-id"); v != "" { + apiResp.TraceID = v + } + if isJSON(httpResp.Header.Get("content-type")) { + if err = jsonUnmarshal(bytes.NewReader(rb), apiResp); err != nil { + return err + } + } else { + apiResp.Message = string(rb) + } + return NewError(httpResp.StatusCode, apiResp) +} + +// Call will send request with signature to http server +func (c *Client) Call(ctx context.Context, method, path string, queryParams interface{}, body interface{}, resp interface{}, ropts ...RequestOption) (err error) { + ro := &RequestOptions{} + for _, opt := range ropts { + opt(ro) + } + if ro.Timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, ro.Timeout) + defer cancel() + } + + req, bb, err := c.buildRequest(ctx, method, path, queryParams, body, ro) + if err != nil { + return err + } + log.Debugf("http call method:%v url:%v body:%v", req.Method, req.URL, string(bb)) - httpResp, err = c.httpClient.Do(req) + httpResp, err := c.httpClient.Do(req) if err != nil { return err } log.Debugf("http call response headers:%v", httpResp.Header) defer httpResp.Body.Close() - if rb, err = io.ReadAll(httpResp.Body); err != nil { + rb, err := io.ReadAll(httpResp.Body) + if err != nil { return err } log.Debugf("http call response body:%s", rb) @@ -256,6 +310,43 @@ func (c *Client) Call(ctx context.Context, method, path string, queryParams inte return nil } +// CallSSE sends a request with signature to the http server and, once the +// server responds with a 200 text/event-stream, returns the raw response +// body for the caller to read as a server-sent-events stream — instead of +// buffering the whole response the way Call does. The caller must Close the +// returned io.ReadCloser once done with it. +// +// A per-request timeout set via WithRequestTimeout is ignored here: an +// event-stream response can legitimately keep delivering events for as long +// as the underlying operation takes, so only ctx cancellation bounds it. +func (c *Client) CallSSE(ctx context.Context, method, path string, body interface{}, ropts ...RequestOption) (io.ReadCloser, error) { + ro := &RequestOptions{} + for _, opt := range ropts { + opt(ro) + } + + req, bb, err := c.buildRequest(ctx, method, path, nil, body, ro) + if err != nil { + return nil, err + } + req.Header.Set("accept", "text/event-stream") + + log.Debugf("http call method:%v url:%v body:%v", req.Method, req.URL, string(bb)) + httpResp, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + log.Debugf("http call response headers:%v", httpResp.Header) + + if httpResp.StatusCode == nhttp.StatusOK && strings.Contains(httpResp.Header.Get("content-type"), "text/event-stream") { + return httpResp.Body, nil + } + + // Error responses (and any other non-SSE response) are still a one-shot + // body, not SSE. + return nil, readErrorBody(httpResp) +} + func isJSON(ct string) bool { return strings.Contains(ct, "application/json") } diff --git a/http/client_test.go b/http/client_test.go new file mode 100644 index 0000000..9f363d2 --- /dev/null +++ b/http/client_test.go @@ -0,0 +1,118 @@ +package http + +import ( + "context" + "errors" + "io" + nhttp "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func newTestClient(t *testing.T, url string) *Client { + t.Helper() + c, err := New( + WithURL(url), + WithAppKey("test_app_key"), + WithAppSecret("test_app_secret"), + WithAccessToken("test_access_token"), + ) + if err != nil { + t.Fatalf("New() error: %v", err) + } + return c +} + +func TestCallSSESuccess(t *testing.T) { + srv := httptest.NewServer(nhttp.HandlerFunc(func(w nhttp.ResponseWriter, r *nhttp.Request) { + if got := r.Header.Get("accept"); got != "text/event-stream" { + t.Errorf("Accept header = %q, want text/event-stream", got) + } + if got := r.Header.Get("x-api-key"); got != "test_app_key" { + t.Errorf("x-api-key header = %q", got) + } + w.Header().Set("content-type", "text/event-stream") + w.WriteHeader(200) + _, _ = w.Write([]byte("event: message\ndata: {\"hello\":\"world\"}\n\n")) + })) + defer srv.Close() + + c := newTestClient(t, srv.URL) + rc, err := c.CallSSE(context.Background(), "POST", "/v1/stream", map[string]string{"query": "hi"}) + if err != nil { + t.Fatalf("CallSSE error: %v", err) + } + defer rc.Close() + + b, err := io.ReadAll(rc) + if err != nil { + t.Fatalf("read error: %v", err) + } + if !strings.Contains(string(b), `{"hello":"world"}`) { + t.Errorf("body = %s", b) + } +} + +func TestCallSSEErrorResponse(t *testing.T) { + srv := httptest.NewServer(nhttp.HandlerFunc(func(w nhttp.ResponseWriter, r *nhttp.Request) { + w.Header().Set("content-type", "application/json") + w.WriteHeader(400) + _, _ = w.Write([]byte(`{"code":123,"message":"bad request"}`)) + })) + defer srv.Close() + + c := newTestClient(t, srv.URL) + _, err := c.CallSSE(context.Background(), "POST", "/v1/stream", map[string]string{"query": "hi"}) + if err == nil { + t.Fatal("expected error, got nil") + } + apiErr, ok := err.(*ApiError) + if !ok { + t.Fatalf("expected *ApiError, got %T: %v", err, err) + } + if apiErr.Code != 123 || apiErr.Message != "bad request" { + t.Errorf("unexpected ApiError: %+v", apiErr) + } +} + +func TestCallWithRequestTimeout(t *testing.T) { + srv := httptest.NewServer(nhttp.HandlerFunc(func(w nhttp.ResponseWriter, r *nhttp.Request) { + time.Sleep(200 * time.Millisecond) + w.Header().Set("content-type", "application/json") + w.WriteHeader(200) + _, _ = w.Write([]byte(`{"code":0,"message":"success","data":{}}`)) + })) + defer srv.Close() + + c := newTestClient(t, srv.URL) + var resp map[string]interface{} + err := c.Call(context.Background(), "GET", "/v1/slow", nil, nil, &resp, WithRequestTimeout(20*time.Millisecond)) + if err == nil { + t.Fatal("expected timeout error, got nil") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected context.DeadlineExceeded, got: %v", err) + } +} + +func TestCallSuccess(t *testing.T) { + srv := httptest.NewServer(nhttp.HandlerFunc(func(w nhttp.ResponseWriter, r *nhttp.Request) { + w.Header().Set("content-type", "application/json") + w.WriteHeader(200) + _, _ = w.Write([]byte(`{"code":0,"message":"success","data":{"ok":true}}`)) + })) + defer srv.Close() + + c := newTestClient(t, srv.URL) + var resp struct { + OK bool `json:"ok"` + } + if err := c.Call(context.Background(), "GET", "/v1/ok", nil, nil, &resp); err != nil { + t.Fatalf("Call error: %v", err) + } + if !resp.OK { + t.Errorf("resp.OK = false, want true") + } +}