feat(log): pretty logging agent responses and tool calls

backend/agents/client/client.go

@@ -4,10 +4,11 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
-	"fmt"
 	"io"
 	"net/http"
 	"os"
+
+	"github.com/charmbracelet/log"
 )
 
 type ResponseFormat struct {
@@ -69,12 +70,14 @@ type AgentClient struct {
 
 	ToolHandler ToolsHandlers
 
+	Log *log.Logger
+
 	Do func(req *http.Request) (*http.Response, error)
 }
 
 const OPENAI_API_KEY = "OPENAI_API_KEY"
 
-func CreateAgentClient() (AgentClient, error) {
+func CreateAgentClient(log *log.Logger) (AgentClient, error) {
 	apiKey := os.Getenv(OPENAI_API_KEY)
 
 	if len(apiKey) == 0 {
@@ -89,6 +92,8 @@ func CreateAgentClient() (AgentClient, error) {
 			return client.Do(req)
 		},
 
+		Log: log,
+
 		ToolHandler: ToolsHandlers{
 			handlers: map[string]ToolHandler{},
 		},
@@ -128,8 +133,6 @@ func (client AgentClient) Request(req *AgentRequestBody) (AgentResponse, error)
 		return AgentResponse{}, err
 	}
 
-	fmt.Println(string(response))
-
 	agentResponse := AgentResponse{}
 	err = json.Unmarshal(response, &agentResponse)
 
@@ -141,6 +144,25 @@ func (client AgentClient) Request(req *AgentRequestBody) (AgentResponse, error)
 		return AgentResponse{}, errors.New("Unsupported. We currently only accept 1 choice from AI.")
 	}
 
+	client.Log.SetLevel(log.DebugLevel)
+
+	msg := agentResponse.Choices[0].Message
+
+	if len(msg.Content) > 0 {
+		client.Log.Debugf("Content: %s", msg.Content)
+	}
+
+	if msg.ToolCalls != nil && len(*msg.ToolCalls) > 0 {
+		client.Log.Debugf("Tool Call: %s", (*msg.ToolCalls)[0].Function.Name)
+
+		prettyJson, err := json.MarshalIndent((*msg.ToolCalls)[0].Function.Arguments, "", "	")
+		if err != nil {
+			return AgentResponse{}, err
+		}
+
+		client.Log.Debugf("Arguments: %s", string(prettyJson))
+	}
+
 	req.Chat.AddAiResponse(agentResponse.Choices[0].Message)
 
 	return agentResponse, nil

backend/agents/contact_agent.go

@@ -3,10 +3,13 @@ package agents
 import (
 	"context"
 	"encoding/json"
+	"os"
 	"screenmark/screenmark/.gen/haystack/haystack/model"
 	"screenmark/screenmark/agents/client"
 	"screenmark/screenmark/models"
+	"time"
 
+	"github.com/charmbracelet/log"
 	"github.com/google/uuid"
 )
 
@@ -130,7 +133,11 @@ func (agent ContactAgent) GetContacts(userId uuid.UUID, imageId uuid.UUID, image
 }
 
 func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
-	agentClient, err := client.CreateAgentClient()
+	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
+		ReportTimestamp: true,
+		TimeFormat:      time.Kitchen,
+		Prefix:          "Contacts 👥",
+	}))
 	if err != nil {
 		return ContactAgent{}, err
 	}

backend/agents/event_location_agent.go

@@ -3,11 +3,13 @@ package agents
 import (
 	"context"
 	"encoding/json"
+	"os"
 	"screenmark/screenmark/.gen/haystack/haystack/model"
 	"screenmark/screenmark/agents/client"
 	"screenmark/screenmark/models"
 	"time"
 
+	"github.com/charmbracelet/log"
 	"github.com/google/uuid"
 )
 
@@ -186,7 +188,11 @@ func (agent EventLocationAgent) GetLocations(userId uuid.UUID, imageId uuid.UUID
 }
 
 func NewLocationEventAgent(locationModel models.LocationModel, eventModel models.EventModel, contactModel models.ContactModel) (EventLocationAgent, error) {
-	agentClient, err := client.CreateAgentClient()
+	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
+		ReportTimestamp: true,
+		TimeFormat:      time.Kitchen,
+		Prefix:          "Location & Events 📅",
+	}))
 	if err != nil {
 		return EventLocationAgent{}, err
 	}

backend/agents/note_agent.go

@@ -2,10 +2,13 @@ package agents
 
 import (
 	"context"
+	"os"
 	"screenmark/screenmark/.gen/haystack/haystack/model"
 	"screenmark/screenmark/agents/client"
 	"screenmark/screenmark/models"
+	"time"
 
+	"github.com/charmbracelet/log"
 	"github.com/google/uuid"
 )
 
@@ -66,7 +69,11 @@ func (agent NoteAgent) GetNotes(userId uuid.UUID, imageId uuid.UUID, imageName s
 }
 
 func NewNoteAgent(noteModel models.NoteModel) (NoteAgent, error) {
-	client, err := client.CreateAgentClient()
+	client, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
+		ReportTimestamp: true,
+		TimeFormat:      time.Kitchen,
+		Prefix:          "Notes 📝",
+	}))
 	if err != nil {
 		return NoteAgent{}, err
 	}

backend/agents/orchestrator.go

@@ -4,8 +4,11 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"os"
 	"screenmark/screenmark/agents/client"
+	"time"
 
+	"github.com/charmbracelet/log"
 	"github.com/google/uuid"
 )
 
@@ -95,6 +98,8 @@ const MY_TOOLS = `
 
 type OrchestratorAgent struct {
 	client client.AgentClient
+
+	log log.Logger
 }
 
 type Status struct {
@@ -147,7 +152,12 @@ func (agent OrchestratorAgent) Orchestrate(userId uuid.UUID, imageId uuid.UUID,
 }
 
 func NewOrchestratorAgent(eventLocationAgent EventLocationAgent, noteAgent NoteAgent, contactAgent ContactAgent, imageName string, imageData []byte) (OrchestratorAgent, error) {
-	agent, err := client.CreateAgentClient()
+	agent, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
+		ReportTimestamp: true,
+		TimeFormat:      time.Kitchen,
+		Prefix:          "Orchestrator 🎼",
+	}))
+
 	if err != nil {
 		return OrchestratorAgent{}, err
 	}