Haystack/backend/agents/create_list_agent.go

package agents

import (
	"context"
	"encoding/json"
	"fmt"
	"screenmark/screenmark/.gen/haystack/haystack/model"
	"screenmark/screenmark/agents/client"
	"screenmark/screenmark/models"

	"github.com/charmbracelet/log"
	"github.com/google/uuid"
)

const createListAgentPrompt = `
You are an agent who's job is to produce a reasonable output for an unstructured input.

Your job is to create lists for the user, the user will give you a title and some fields they want
as part of the list. Your job is to take these fields, adjust their names so they have good names,
and add a good description for each one.

You can add fields if you think they make a lot of sense.
You can remove fields if they are not correct, but be sure before you do this.
`

const listJsonSchema = `
{
  "type": "object",
  "properties": {
	"title": {
		"type": "string",
		"description": "the title of the list"
	},
	"description": {
		"type": "string",
		"description": "the description of the list"
	},
    "fields": {
      "type": "array",
      "items": {
        "type": "object",
        "properties": {
          "name": {
            "type": "string",
            "description": "The name of the field."
          },
          "description": {
            "type": "string",
            "description": "A description of the field."
          }
        },
        "required": [
          "name",
          "description"
        ]
      },
      "description": "An array of field objects."
    }
  },
  "required": [
    "fields"
  ]
}
`

type createNewListArguments struct {
	Title       string `json:"title"`
	Description string `json:"description"`

	Fields []struct {
		Name        string `json:"name"`
		Description string `json:"description"`
	} `json:"fields"`
}

type CreateListAgent struct {
	client client.AgentClient

	listModel models.StackModel
}

func (agent *CreateListAgent) CreateList(log *log.Logger, userID uuid.UUID, userReq string) error {
	request := client.AgentRequestBody{
		Model:       "policy/images",
		Temperature: 0.3,
		ResponseFormat: client.ResponseFormat{
			Type:       "json_object",
			JsonSchema: listJsonSchema,
		},
		Chat: &client.Chat{
			Messages: make([]client.ChatMessage, 0),
		},
	}

	request.Chat.AddSystem(agent.client.Options.SystemPrompt)
	request.Chat.AddUser(userReq)

	resp, err := agent.client.Request(&request)
	if err != nil {
		return fmt.Errorf("request: %w", err)
	}

	ctx := context.Background()

	structuredOutput := resp.Choices[0].Message.Content

	var createListArgs createNewListArguments
	err = json.Unmarshal([]byte(structuredOutput), &createListArgs)
	if err != nil {
		return err
	}

	schemaItems := make([]model.SchemaItems, 0)
	for _, field := range createListArgs.Fields {
		schemaItems = append(schemaItems, model.SchemaItems{
			Item:        field.Name,
			Description: field.Description,

			Value: "string", // keep it simple for now.
		})
	}

	_, err = agent.listModel.Save(ctx, userID, createListArgs.Title, createListArgs.Description)
	if err != nil {
		return fmt.Errorf("creating list agent, saving list: %w", err)
	}

	err = agent.listModel.SaveItems(ctx, schemaItems)
	if err != nil {
		return fmt.Errorf("creating list agent, saving items: %w", err)
	}

	return nil
}

func NewCreateListAgent(log *log.Logger, listModel models.StackModel) CreateListAgent {
	client := client.CreateAgentClient(client.CreateAgentClientOptions{
		SystemPrompt: createListAgentPrompt,
		Log:          log,
	})

	agent := CreateListAgent{
		client,
		listModel,
	}

	return agent
}