Added chat completions

Author: djthorpe

pkg/client/transport.go

@@ -19,6 +19,11 @@ type logtransport struct {
 	v bool
 }
 
+type readwrapper struct {
+	r    io.ReadCloser
+	data bytes.Buffer
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE
 
@@ -58,27 +63,72 @@ func (transport *logtransport) RoundTrip(req *http.Request) (*http.Response, err
 	defer func() {
 		fmt.Fprintln(transport.w, "  Took", time.Since(then).Milliseconds(), "ms")
 	}()
+
+	// Wrap the request
+	req.Body = &readwrapper{r: req.Body}
+
+	// Perform the roundtrip
 	resp, err := transport.RoundTripper.RoundTrip(req)
 	if err != nil {
 		fmt.Fprintln(transport.w, "error:", err)
-	} else {
-		fmt.Fprintln(transport.w, "response:", resp.Status)
-		for k, v := range resp.Header {
-			fmt.Fprintf(transport.w, "  <= %v: %q\n", k, v)
+		return resp, err
+	}
+
+	// If verbose is switched on, output the payload
+	if transport.v {
+		data, err := req.Body.(*readwrapper).as(req.Header.Get("Content-Type"))
+		if err == nil {
+			fmt.Fprintln(transport.w, "   ", string(data))
 		}
-		// If verbose is switched on, read the body
-		if transport.v && resp.Body != nil {
-			contentType := resp.Header.Get("Content-Type")
-			if contentType == ContentTypeJson || contentType == ContentTypeTextPlain {
-				defer resp.Body.Close()
-				body, err := io.ReadAll(resp.Body)
-				if err == nil {
-					fmt.Fprintln(transport.w, "    ", string(body))
-				}
-				resp.Body = io.NopCloser(bytes.NewReader(body))
+	}
+
+	fmt.Fprintln(transport.w, "response:", resp.Status)
+	for k, v := range resp.Header {
+		fmt.Fprintf(transport.w, "  <= %v: %q\n", k, v)
+	}
+
+	// If verbose is switched on, read the body
+	if transport.v && resp.Body != nil {
+		contentType := resp.Header.Get("Content-Type")
+		if contentType == ContentTypeJson || contentType == ContentTypeTextPlain {
+			defer resp.Body.Close()
+			body, err := io.ReadAll(resp.Body)
+			if err == nil {
+				fmt.Fprintln(transport.w, "    ", string(body))
 			}
+			resp.Body = io.NopCloser(bytes.NewReader(body))
 		}
 	}
 
+	// Return success
 	return resp, err
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// PRIVATE METHODS
+
+func (w *readwrapper) Read(b []byte) (n int, err error) {
+	n, err = w.r.Read(b)
+	if err == nil {
+		_, err = w.data.Write(b[:n])
+	}
+	return n, err
+}
+
+func (w *readwrapper) Close() error {
+	return w.r.Close()
+}
+
+func (w *readwrapper) as(mimetype string) ([]byte, error) {
+	switch mimetype {
+	case ContentTypeJson:
+		dest := bytes.NewBuffer(nil)
+		if err := json.Indent(dest, w.data.Bytes(), "    ", "  "); err != nil {
+			return nil, err
+		} else {
+			return dest.Bytes(), nil
+		}
+	default:
+		return w.data.Bytes(), nil
+	}
+}

pkg/openai/chat.go

@@ -9,14 +9,36 @@ const (
 	defaultChatCompletion = "gpt-3.5-turbo"
 )
 
+///////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS
+
+func NewUserMessage(text string) Message {
+	return Message{
+		Role: "user", Content: &text,
+	}
+}
+
+func NewSystemMessage(text string) Message {
+	return Message{
+		Role: "system", Content: &text,
+	}
+}
+
+func NewAssistantMessage(text string) Message {
+	return Message{
+		Role: "assistant", Content: &text,
+	}
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // API CALLS
 
 // Chat creates a model response for the given chat conversation.
-func (c *Client) Chat(opts ...Opt) (Chat, error) {
+func (c *Client) Chat(messages []Message, opts ...Opt) (Chat, error) {
 	// Create the request
 	var request reqChat
 	request.Model = defaultChatCompletion
+	request.Messages = messages
 	for _, opt := range opts {
 		if err := opt(&request); err != nil {
 			return Chat{}, err

pkg/openai/chat_test.go

@@ -1,6 +1,7 @@
 package openai_test
 
 import (
+	"encoding/json"
 	"os"
 	"testing"
 
@@ -15,7 +16,49 @@ func Test_chat_001(t *testing.T) {
 	assert.NoError(err)
 	assert.NotNil(client)
 
-	response, err := client.Chat()
+	response, err := client.Chat([]openai.Message{
+		openai.NewUserMessage("What would be the best app to use to get the weather in berlin today?"),
+	})
 	assert.NoError(err)
 	assert.NotNil(response)
+	assert.NotEmpty(response)
+
+	data, err := json.MarshalIndent(response, "", "  ")
+	assert.NoError(err)
+	t.Log(string(data))
+
+}
+
+func Test_chat_002(t *testing.T) {
+	assert := assert.New(t)
+	client, err := openai.New(GetApiKey(t), opts.OptTrace(os.Stderr, true))
+	assert.NoError(err)
+	assert.NotNil(client)
+
+	response, err := client.Chat([]openai.Message{
+		openai.NewUserMessage("What is the weather in berlin today?"),
+	}, openai.OptFunction("get_weather", "Get the weather in a specific city and country", openai.ToolParameter{
+		Name:        "city",
+		Type:        "string",
+		Description: "The city to get the weather for",
+		Required:    true,
+	}, openai.ToolParameter{
+		Name:        "country",
+		Type:        "string",
+		Description: "The country to get the weather for",
+		Required:    true,
+	}, openai.ToolParameter{
+		Name:        "time",
+		Type:        "string",
+		Description: "When to get the weather for. If not specified, defaults to the current time",
+		Required:    true,
+	}))
+	assert.NoError(err)
+	assert.NotNil(response)
+	assert.NotEmpty(response)
+
+	data, err := json.MarshalIndent(response, "", "  ")
+	assert.NoError(err)
+	t.Log(string(data))
+
 }