refactor(agents): no need to wrap them in another struct

2025-04-17 10:36:11 +01:00
parent fa486153b4
commit e42aa75639
6 changed files with 66 additions and 118 deletions
--- a/backend/agents/contact_agent.go
+++ b/backend/agents/contact_agent.go
@@ -3,11 +3,9 @@ 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"
@@ -81,12 +79,6 @@ const contactTools = `
 ]
 `

-type ContactAgent struct {
-	client client.AgentClient
-
-	contactModel models.ContactModel
-}
-
 type listContactsArguments struct{}
 type createContactsArguments struct {
 	Name        string  `json:"name"`
@@ -98,23 +90,14 @@ type linkContactArguments struct {
 	ContactID string `json:"contactId"`
 }

-func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
-	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
-		ReportTimestamp: true,
-		TimeFormat:      time.Kitchen,
-		Prefix:          "Contacts 👥",
-	}))
+func NewContactAgent(log *log.Logger, contactModel models.ContactModel) (client.AgentClient, error) {
+	agentClient, err := client.CreateAgentClient(log)
 	if err != nil {
-		return ContactAgent{}, err
-	}
-
-	agent := ContactAgent{
-		client:       agentClient,
-		contactModel: contactModel,
+		return client.AgentClient{}, err
 	}

 	agentClient.ToolHandler.AddTool("listContacts", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		return agent.contactModel.List(context.Background(), info.UserId)
+		return contactModel.List(context.Background(), info.UserId)
 	})

 	agentClient.ToolHandler.AddTool("createContact", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
@@ -126,7 +109,7 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {

 		ctx := context.Background()

-		contact, err := agent.contactModel.Save(ctx, info.UserId, model.Contacts{
+		contact, err := contactModel.Save(ctx, info.UserId, model.Contacts{
 			Name:        args.Name,
 			PhoneNumber: args.PhoneNumber,
 			Email:       args.Email,
@@ -136,7 +119,7 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
 			return model.Contacts{}, err
 		}

-		_, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contact.ID)
+		_, err = contactModel.SaveToImage(ctx, info.ImageId, contact.ID)
 		if err != nil {
 			return model.Contacts{}, err
 		}
@@ -158,7 +141,7 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
 			return "", err
 		}

-		_, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contactUuid)
+		_, err = contactModel.SaveToImage(ctx, info.ImageId, contactUuid)
 		if err != nil {
 			return "", err
 		}
@@ -166,5 +149,5 @@ func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
 		return "Saved", nil
 	})

-	return agent, nil
+	return agentClient, nil
 }
--- a/backend/agents/event_agent.go
+++ b/backend/agents/event_agent.go
@@ -3,7 +3,6 @@ package agents
 import (
 	"context"
 	"encoding/json"
-	"os"
 	"screenmark/screenmark/.gen/haystack/haystack/model"
 	"screenmark/screenmark/agents/client"
 	"screenmark/screenmark/models"
@@ -98,14 +97,6 @@ const eventTools = `
    }
 ]`

-type EventAgent struct {
-	client client.AgentClient
-
-	eventsModel models.EventModel
-
-	locationAgent LocationAgent
-}
-
 type listEventArguments struct{}
 type createEventArguments struct {
 	Name          string  `json:"name"`
@@ -117,25 +108,15 @@ type linkEventArguments struct {
 	EventID string `json:"eventId"`
 }

-func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) (EventAgent, error) {
-	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
-		ReportTimestamp: true,
-		TimeFormat:      time.Kitchen,
-		Prefix:          "Events 📍",
-	}))
+func NewEventAgent(log *log.Logger, eventsModel models.EventModel, locationAgent client.AgentClient) (client.AgentClient, error) {
+	agentClient, err := client.CreateAgentClient(log)

 	if err != nil {
-		return EventAgent{}, err
-	}
-
-	agent := EventAgent{
-		client:        agentClient,
-		eventsModel:   eventsModel,
-		locationAgent: locationAgent,
+		return client.AgentClient{}, err
 	}

 	agentClient.ToolHandler.AddTool("listEvents", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		return agent.eventsModel.List(context.Background(), info.UserId)
+		return eventsModel.List(context.Background(), info.UserId)
 	})

 	agentClient.ToolHandler.AddTool("createEvent", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
@@ -159,7 +140,7 @@ func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) (
 			return model.Events{}, err
 		}

-		events, err := agent.eventsModel.Save(ctx, info.UserId, model.Events{
+		events, err := eventsModel.Save(ctx, info.UserId, model.Events{
 			Name:          args.Name,
 			StartDateTime: &startTime,
 			EndDateTime:   &endTime,
@@ -169,7 +150,7 @@ func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) (
 			return model.Events{}, err
 		}

-		_, err = agent.eventsModel.SaveToImage(ctx, info.ImageId, events.ID)
+		_, err = eventsModel.SaveToImage(ctx, info.ImageId, events.ID)
 		if err != nil {
 			return model.Events{}, err
 		}
@@ -191,16 +172,16 @@ func NewEventAgent(eventsModel models.EventModel, locationAgent LocationAgent) (
 			return "", err
 		}

-		agent.eventsModel.SaveToImage(ctx, info.ImageId, contactUuid)
+		eventsModel.SaveToImage(ctx, info.ImageId, contactUuid)
 		return "Saved", nil
 	})

 	agentClient.ToolHandler.AddTool("getEventLocationId", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
 		query := "Can you get me the ID of the location present in this image?"
-		locationAgent.client.RunAgent(locationPrompt, locationTools, "finish", &query, info.UserId, info.ImageId, info.ImageName, *info.Image)
+		locationAgent.RunAgent(locationPrompt, locationTools, "finish", &query, info.UserId, info.ImageId, info.ImageName, *info.Image)

-		return locationAgent.client.Reply, nil
+		return locationAgent.Reply, nil
 	})

-	return agent, nil
+	return agentClient, nil
 }
--- a/backend/agents/location_agent.go
+++ b/backend/agents/location_agent.go
@@ -3,11 +3,9 @@ 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"
@@ -97,12 +95,6 @@ const locationTools = `
    }
 ]`

-type LocationAgent struct {
-	client client.AgentClient
-
-	locationModel models.LocationModel
-}
-
 type listLocationArguments struct{}
 type createLocationArguments struct {
 	Name    string  `json:"name"`
@@ -112,24 +104,15 @@ type linkLocationArguments struct {
 	LocationID string `json:"locationId"`
 }

-func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) {
-	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
-		ReportTimestamp: true,
-		TimeFormat:      time.Kitchen,
-		Prefix:          "Locations 📍",
-	}))
+func NewLocationAgent(log *log.Logger, locationModel models.LocationModel) (client.AgentClient, error) {
+	agentClient, err := client.CreateAgentClient(log)

 	if err != nil {
-		return LocationAgent{}, err
-	}
-
-	agent := LocationAgent{
-		client:        agentClient,
-		locationModel: locationModel,
+		return client.AgentClient{}, err
 	}

 	agentClient.ToolHandler.AddTool("listLocations", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		return agent.locationModel.List(context.Background(), info.UserId)
+		return locationModel.List(context.Background(), info.UserId)
 	})

 	agentClient.ToolHandler.AddTool("createLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
@@ -141,7 +124,7 @@ func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error)

 		ctx := context.Background()

-		location, err := agent.locationModel.Save(ctx, info.UserId, model.Locations{
+		location, err := locationModel.Save(ctx, info.UserId, model.Locations{
 			Name:    args.Name,
 			Address: args.Address,
 		})
@@ -150,7 +133,7 @@ func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error)
 			return model.Locations{}, err
 		}

-		_, err = agent.locationModel.SaveToImage(ctx, info.ImageId, location.ID)
+		_, err = locationModel.SaveToImage(ctx, info.ImageId, location.ID)
 		if err != nil {
 			return model.Locations{}, err
 		}
@@ -172,14 +155,14 @@ func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error)
 			return "", err
 		}

-		agent.locationModel.SaveToImage(ctx, info.ImageId, contactUuid)
+		locationModel.SaveToImage(ctx, info.ImageId, contactUuid)
 		return "Saved", nil
 	})

 	agentClient.ToolHandler.AddTool("reply", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		agent.client.Log.Debug(args)
+		agentClient.Log.Debug(args)
 		return "ok", nil
 	})

-	return agent, nil
+	return agentClient, nil
 }
--- a/backend/agents/note_agent.go
+++ b/backend/agents/note_agent.go
@@ -2,11 +2,9 @@ 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"
@@ -70,12 +68,8 @@ func (agent NoteAgent) GetNotes(userId uuid.UUID, imageId uuid.UUID, imageName s
 	return nil
 }

-func NewNoteAgent(noteModel models.NoteModel) (NoteAgent, error) {
-	client, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
-		ReportTimestamp: true,
-		TimeFormat:      time.Kitchen,
-		Prefix:          "Notes 📝",
-	}))
+func NewNoteAgent(log *log.Logger, noteModel models.NoteModel) (NoteAgent, error) {
+	client, err := client.CreateAgentClient(log)
 	if err != nil {
 		return NoteAgent{}, err
 	}
--- a/backend/agents/orchestrator.go
+++ b/backend/agents/orchestrator.go
@@ -2,9 +2,7 @@ package agents

 import (
 	"errors"
-	"os"
 	"screenmark/screenmark/agents/client"
-	"time"

 	"github.com/charmbracelet/log"
 )
@@ -114,15 +112,11 @@ type Status struct {
 	Ok bool `json:"ok"`
 }

-func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locationAgent LocationAgent, eventAgent EventAgent, imageName string, imageData []byte) (OrchestratorAgent, error) {
-	agent, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
-		ReportTimestamp: true,
-		TimeFormat:      time.Kitchen,
-		Prefix:          "Orchestrator 🎼",
-	}))
+func NewOrchestratorAgent(log *log.Logger, noteAgent NoteAgent, contactAgent client.AgentClient, locationAgent client.AgentClient, eventAgent client.AgentClient, imageName string, imageData []byte) (client.AgentClient, error) {
+	agent, err := client.CreateAgentClient(log)

 	if err != nil {
-		return OrchestratorAgent{}, err
+		return client.AgentClient{}, err
 	}

 	agent.ToolHandler.AddTool("noteAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
@@ -134,7 +128,7 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati
 	})

 	agent.ToolHandler.AddTool("contactAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		go contactAgent.client.RunAgent(contactPrompt, contactTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData)
+		go contactAgent.RunAgent(contactPrompt, contactTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData)

 		return Status{
 			Ok: true,
@@ -142,7 +136,7 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati
 	})

 	agent.ToolHandler.AddTool("locationAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		go locationAgent.client.RunAgent(locationPrompt, locationTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData)
+		go locationAgent.RunAgent(locationPrompt, locationTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData)

 		return Status{
 			Ok: true,
@@ -150,7 +144,7 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati
 	})

 	agent.ToolHandler.AddTool("eventAgent", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
-		go eventAgent.client.RunAgent(eventPrompt, eventTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData)
+		go agent.RunAgent(eventPrompt, eventTools, "finish", nil, info.UserId, info.ImageId, imageName, imageData)

 		return Status{
 			Ok: true,
@@ -165,7 +159,5 @@ func NewOrchestratorAgent(noteAgent NoteAgent, contactAgent ContactAgent, locati
 		}, errors.New("Finished! Kinda bad return type but...")
 	})

-	return OrchestratorAgent{
-		Client: agent,
-	}, nil
+	return agent, nil
 }