158 lines
3.6 KiB
Go
158 lines
3.6 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"screenmark/screenmark/.gen/haystack/haystack/model"
|
|
"screenmark/screenmark/agents/client"
|
|
"screenmark/screenmark/models"
|
|
"strings"
|
|
|
|
"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.
|
|
|
|
You must respond in json format, do not add backticks to the json. ONLY valid json.
|
|
`
|
|
|
|
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
|
|
|
|
stackModel models.StackModel
|
|
}
|
|
|
|
func (agent *CreateListAgent) CreateList(log *log.Logger, userID uuid.UUID, stackID uuid.UUID, title string, userReq string) error {
|
|
request := client.AgentRequestBody{
|
|
Model: "policy/images",
|
|
Temperature: 0.3,
|
|
ResponseFormat: client.ResponseFormat{
|
|
Type: "json_schema",
|
|
JsonSchema: listJsonSchema,
|
|
},
|
|
Chat: &client.Chat{
|
|
Messages: make([]client.ChatMessage, 0),
|
|
},
|
|
}
|
|
|
|
request.Chat.AddSystem(agent.client.Options.SystemPrompt)
|
|
|
|
req := fmt.Sprintf("List title: %s | Users list description: %s", title, userReq)
|
|
|
|
request.Chat.AddUser(req)
|
|
|
|
resp, err := agent.client.Request(&request)
|
|
if err != nil {
|
|
return fmt.Errorf("request: %w", err)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
|
|
content := resp.Choices[0].Message.Content
|
|
|
|
if strings.HasPrefix(content, "```json") {
|
|
content = content[len("```json") : len(content)-3]
|
|
}
|
|
|
|
log.Info("", "res", content)
|
|
|
|
var createListArgs createNewListArguments
|
|
err = json.Unmarshal([]byte(content), &createListArgs)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
schemaItems := make([]model.SchemaItems, 0)
|
|
for _, field := range createListArgs.Fields {
|
|
schemaItems = append(schemaItems, model.SchemaItems{
|
|
StackID: stackID,
|
|
|
|
Item: field.Name,
|
|
Description: field.Description,
|
|
|
|
Value: "string", // keep it simple for now.
|
|
})
|
|
}
|
|
|
|
err = agent.stackModel.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
|
|
}
|