Merge pull request #37 from mutablelogic/v1

Adding JSON streaming

Author: djthorpe

README.md

@@ -7,7 +7,7 @@ This repository contains a generic HTTP client which can be adapted to provide:
 * Ability to send files  and data of type `multipart/form-data`
 * Ability to send data of type `application/x-www-form-urlencoded`
 * Debugging capabilities to see the request and response data
-* Streaming text events
+* Streaming text and JSON events
 
 API Documentation: https://pkg.go.dev/github.com/mutablelogic/go-client
 
@@ -159,6 +159,9 @@ modify each individual request when using the `Do` method:
 * `OptTextStreamCallback(func(TextStreamCallback) error)` allows you to set a callback
     function to process a streaming text response of type `text/event-stream`. See below for
     more details.
+* `OptJsonStreamCallback(func(any) error)` allows you to set a callback for JSON streaming
+    responses. The callback should have the signature `func(any) error`. See below for
+    more details.
 
 ## Authentication
 
@@ -191,9 +194,9 @@ You can also set the token on a per-request basis using the `OptToken` option in
 
 You can create a payload with form data:
 
-* `client.NewFormRequest(payload any, accept string)` returns a new request with a Form 
+* `client.NewFormRequest(payload any, accept string)` returns a new request with a Form
     data payload which defaults to POST.
-* `client.NewMultipartRequest(payload any, accept string)` returns a new request with 
+* `client.NewMultipartRequest(payload any, accept string)` returns a new request with
     a Multipart Form data payload which defaults to POST. This is useful for file uploads.
 
 The payload should be a `struct` where the fields are converted to form tuples. File uploads require a field of type `multipart.File`. For example,
@@ -241,9 +244,10 @@ type Unmarshaler interface {
 }
 ```
 
-## Streaming Responses
+## Text Streaming Responses
 
-The client implements a streaming text event callback which can be used to process a stream of text events, as per the [Mozilla specification](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events).
+The client implements a streaming text event callback which can be used to process a stream of text events,
+as per the [Mozilla specification](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events).
 
 In order to process streamed events, pass the `OptTextStreamCallback()` option to the request
 with a callback function, which should have the following signature:
@@ -272,3 +276,12 @@ If you return an error of type `io.EOF` from the callback, then the stream will
 Similarly, if you return any other error the stream will be closed and the error returned.
 
 Usually, you would pair this option with `OptNoTimeout` to prevent the request from timing out.
+
+## JSON Streaming Responses
+
+The client decodes JSON streaming responses by passing a callback function to the `OptJsonStreamCallback()` option.
+The callback with signature `func(any) error` is called for each JSON object in the stream, where the argument
+is the same type as the object in the request.
+
+You can return an error from the callback to stop the stream and return the error, or return `io.EOF` to stop the stream
+immediately and return success.

client.go

@@ -61,6 +61,7 @@ const (
 	PathSeparator             = string(os.PathSeparator)
 	ContentTypeAny            = "*/*"
 	ContentTypeJson           = "application/json"
+	ContentTypeJsonStream     = "application/x-ndjson"
 	ContentTypeTextXml        = "text/xml"
 	ContentTypeApplicationXml = "application/xml"
 	ContentTypeTextPlain      = "text/plain"
@@ -306,7 +307,7 @@ func do(client *http.Client, req *http.Request, accept string, strict bool, out
 
 	// Decode the body
 	switch mimetype {
-	case ContentTypeJson:
+	case ContentTypeJson, ContentTypeJsonStream:
 		// JSON decode is streamable
 		dec := json.NewDecoder(response.Body)
 		for {

cmd/agent/chat.go

@@ -0,0 +1,169 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	// Packages
+	markdown "github.com/MichaelMure/go-term-markdown"
+	agent "github.com/mutablelogic/go-client/pkg/agent"
+)
+
+/////////////////////////////////////////////////////////////////////
+// TYPES
+
+type ChatCmd struct {
+	Prompt string `arg:"" optional:"" help:"The prompt to generate a response for"`
+	Agent  string `flag:"agent" help:"The agent to use"`
+	Model  string `flag:"model" help:"The model to use"`
+	Stream bool   `flag:"stream" help:"Stream the response"`
+}
+
+/////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func (cmd *ChatCmd) Run(globals *Globals) error {
+	// Get the agent and the model
+	model_agent, model := globals.getModel(globals.ctx, cmd.Agent, cmd.Model)
+	if model_agent == nil || model == nil {
+		return fmt.Errorf("model %q not found, or not set on command line", globals.state.Model)
+	}
+
+	// Generate the options
+	opts := make([]agent.Opt, 0)
+	if cmd.Stream {
+		opts = append(opts, agent.OptStream(func(r agent.Response) {
+			fmt.Println(r)
+		}))
+	}
+
+	// Add tools
+	if tools := globals.getTools(); len(tools) > 0 {
+		opts = append(opts, agent.OptTools(tools...))
+	}
+
+	// If the prompt is empty, then we're in interative mode
+	context := []agent.Context{}
+	if cmd.Prompt == "" {
+		if globals.term == nil {
+			return fmt.Errorf("prompt is empty and not in interactive mode")
+		}
+	} else {
+		context = append(context, model_agent.UserPrompt(cmd.Prompt))
+	}
+
+FOR_LOOP:
+	for {
+		// When there is no context, create some
+		if len(context) == 0 {
+			if prompt, err := globals.term.ReadLine(model.Name() + "> "); err != nil {
+				return err
+			} else if prompt == "" {
+				break FOR_LOOP
+			} else {
+				context = append(context, model_agent.UserPrompt(prompt))
+			}
+		}
+
+		// Generate a chat completion
+		response, err := model_agent.Generate(globals.ctx, model, context, opts...)
+		if err != nil {
+			return err
+		}
+
+		// If the response is a tool call, then run the tool
+		if response.ToolCall != nil {
+			result, err := globals.runTool(globals.ctx, response.ToolCall)
+			if err != nil {
+				return err
+			}
+			response.Context = append(response.Context, result)
+		} else {
+			if globals.term != nil {
+				w, _ := globals.term.Size()
+				fmt.Println(string(markdown.Render(response.Text, w, 0)))
+			} else {
+				fmt.Println(response.Text)
+			}
+
+			// Make empty context
+			response.Context = []agent.Context{}
+		}
+
+		// Context comes from the response
+		context = response.Context
+	}
+
+	// Return success
+	return nil
+}
+
+/////////////////////////////////////////////////////////////////////
+// PRIVATE METHODS
+
+// Get the model, either from state or from the command-line flags.
+// If the model is not found, or there is another error, return nil
+func (globals *Globals) getModel(ctx context.Context, agent, model string) (agent.Agent, agent.Model) {
+	state := globals.state
+	if agent != "" {
+		state.Agent = agent
+	}
+	if model != "" {
+		state.Model = model
+	}
+
+	// Cycle through the agents and models to find the one we want
+	for _, agent := range globals.agents {
+		// Filter by agent
+		if state.Agent != "" && agent.Name() != state.Agent {
+			continue
+		}
+
+		// Retrieve the models for this agent
+		models, err := agent.Models(ctx)
+		if err != nil {
+			continue
+		}
+
+		// Filter by model
+		for _, model := range models {
+			if state.Model != "" && model.Name() != state.Model {
+				continue
+			}
+
+			// This is the model we're using....
+			state.Agent = agent.Name()
+			state.Model = model.Name()
+			return agent, model
+		}
+	}
+
+	// No model found
+	return nil, nil
+}
+
+// Get the tools
+func (globals *Globals) getTools() []agent.Tool {
+	return globals.tools
+}
+
+// Return a tool by name. If the tool is not found, return nil
+func (globals *Globals) getTool(name string) agent.Tool {
+	for _, tool := range globals.tools {
+		if tool.Name() == name {
+			return tool
+		}
+	}
+	return nil
+}
+
+// Run a tool from a tool call, and return the result
+func (globals *Globals) runTool(ctx context.Context, call *agent.ToolCall) (*agent.ToolResult, error) {
+	tool := globals.getTool(call.Name)
+	if tool == nil {
+		return nil, fmt.Errorf("tool %q not found", call.Name)
+	}
+
+	// Run the tool
+	return tool.Run(ctx, call)
+}

cmd/agent/list_agents.go

@@ -0,0 +1,31 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+/////////////////////////////////////////////////////////////////////
+// TYPES
+
+type ListAgentsCmd struct {
+}
+
+/////////////////////////////////////////////////////////////////////
+// METHODS
+
+func (cmd *ListAgentsCmd) Run(ctx *Globals) error {
+	result := make([]string, 0)
+	for _, agent := range ctx.agents {
+		result = append(result, agent.Name())
+	}
+
+	data, err := json.MarshalIndent(result, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	fmt.Println(string(data))
+
+	return nil
+}

cmd/agent/list_models.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+/////////////////////////////////////////////////////////////////////
+// TYPES
+
+type ListModelsCmd struct {
+}
+
+type modeljson struct {
+	Agent string `json:"agent"`
+	Model string `json:"model"`
+}
+
+/////////////////////////////////////////////////////////////////////
+// METHODS
+
+func (cmd *ListModelsCmd) Run(ctx *Globals) error {
+	result := make([]modeljson, 0)
+	for _, agent := range ctx.agents {
+		models, err := agent.Models(ctx.ctx)
+		if err != nil {
+			return err
+		}
+		for _, model := range models {
+			result = append(result, modeljson{Agent: agent.Name(), Model: model.Name()})
+		}
+	}
+
+	data, err := json.MarshalIndent(result, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	fmt.Println(string(data))
+
+	return nil
+}

cmd/agent/list_tools.go

@@ -0,0 +1,37 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+/////////////////////////////////////////////////////////////////////
+// TYPES
+
+type ListToolsCmd struct {
+}
+
+type tooljson struct {
+	Provider    string `json:"provider"`
+	Name        string `json:"name"`
+	Description string `json:"description"`
+}
+
+/////////////////////////////////////////////////////////////////////
+// METHODS
+
+func (cmd *ListToolsCmd) Run(ctx *Globals) error {
+	result := make([]tooljson, 0)
+	for _, tool := range ctx.tools {
+		result = append(result, tooljson{Provider: tool.Provider(), Name: tool.Name(), Description: tool.Description()})
+	}
+
+	data, err := json.MarshalIndent(result, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	fmt.Println(string(data))
+
+	return nil
+}