loop/loop.go

package loop

import (
	_ "embed"
	"encoding/json"
	"fmt"
	"sync"

	"dodo.cloud/neo/tools"

	"github.com/anthropics/anthropic-sdk-go"
	"github.com/invopop/jsonschema"
)

//go:embed CONVERSATION_RULES.md
var systemPrompt string

type Message struct {
	Author   string `json:"author"`
	Contents string `json:"contents"`
	Done     bool   `json:"done"`
}

type Conversation struct {
	Messages []Message `json:"message"`
}

var todo *ToDo

func Run(pr PromptReader, client *Client, tools tools.Registry) error {
	RegisterToDoTools(tools)
	fmt.Printf("YOU: ")
	prompt, err := pr.Read()
	if err != nil {
		return err
	}
	todo = &ToDo{}
	todo.ID = "1"
	todo.Title = prompt
	todo.AssignedTo = "assistant"
	agents := []Agent{
		&UserAgent{pr},
		&AnthropicAgent{tools, client.c},
	}
	var wg sync.WaitGroup
	for _, a := range agents {
		wg.Add(1)
		go func() {
			if err := a.Run(todo); err != nil {
				panic(err)
			}
			wg.Done()
		}()
	}
	wg.Wait()
	return nil
}

// func Loop(todo *ToDo, item *ToDo, pr PromptReader, client *Client, reg tools.Registry) error {
// 	messages := []anthropic.MessageParam{
// 		anthropic.NewUserMessage(anthropic.NewTextBlock(string(k))),
// 		anthropic.NewUserMessage(anthropic.NewTextBlock(fmt.Sprintf("Work on TODO item: %s", item.ID))),
// 	}
// 	for {
// 		fmt.Println(todo.String())
// 		if item.AssignedTo == "user" {
// 			fmt.Printf("YOU %s: ", item.ID)
// 			prompt, err := pr.Read()
// 			if err != nil {
// 				return err
// 			}
// 			item.Discussion = append(item.Discussion, Comment{
// 				Author:  "user",
// 				Comment: prompt,
// 			})
// 			item.AssignedTo = "assistant"
// 			messages = append(messages, anthropic.NewUserMessage(anthropic.NewTextBlock(prompt)))
// 		} else {
// 			if len(messages) == 0 {
// 				messages = append(messages, anthropic.NewUserMessage(anthropic.NewTextBlock(fmt.Sprintf("Work on TODO item with id"))))
// 			}
// 		}

// 	}
// }

func pickToDoItem(todo *ToDo) *ToDo {
	if todo.Done {
		return nil
	}
	for _, i := range todo.Items {
		if ret := pickToDoItem(i); ret != nil {
			return ret
		}
	}
	return todo
}

func findItemByID(todo *ToDo, id string) *ToDo {
	if todo.ID == id {
		return todo
	}
	for _, i := range todo.Items {
		if ret := findItemByID(i, id); ret != nil {
			return ret
		}
	}
	return nil
}

type ToDoItem struct {
	ParentID    string `json:"parentId"`
	Title       string `json:"title"`
	Description string `json:"description"`
	AssignedTo  string `json:"assignedTo"`
}

type ToDoAddItemArgs struct {
	Items []ToDoItem `json:"items"`
}

func ToDoAddItem(args ToDoAddItemArgs) (string, error) {
	for _, td := range args.Items {
		item := findItemByID(todo, td.ParentID)
		if item == nil {
			return "error", fmt.Errorf("TODO item with given id not found: %s", td.ParentID)
		}
		id := fmt.Sprintf("%s.%d", item.ID, len(item.Items)+1)
		item.Items = append(item.Items, &ToDo{
			ID:          id,
			Title:       td.Title,
			Description: td.Description,
			AssignedTo:  td.AssignedTo,
		})
	}
	return "done", nil
}

type ToDoMarkItemDoneArgs struct {
	ID string `json:"id"`
}

func ToDoMarkItemDone(args ToDoMarkItemDoneArgs) (string, error) {
	item := findItemByID(todo, args.ID)
	if item == nil {
		return "error", fmt.Errorf("TODO item with given id not found: %s", args.ID)
	}
	item.Done = true
	return "done", nil
}

type ToDoItemAddCommentArgs struct {
	ID       string `json:"id"`
	Comment  string `json:"comment"`
	AssignTo string `json:"assignTo"`
}

func ToDoItemAddComment(args ToDoItemAddCommentArgs) (string, error) {
	item := findItemByID(todo, args.ID)
	if item == nil {
		return "error", fmt.Errorf("TODO item with given id not found: %s", args.ID)
	}
	if len(item.Discussion) == 0 {
		return "error", fmt.Errorf("You shall never initiate a discussion, if you want to clarify something create a TODO item for it.")
	}
	item.Discussion = append(item.Discussion, Comment{
		Author:  "assistant",
		Comment: args.Comment,
	})
	item.AssignedTo = args.AssignTo
	return "done", nil
}

func RegisterToDoTools(reg tools.Registry) {
	reg.Add(tools.NewFuncTool("todo_item_add", ToDoAddItem, "Add new ToDo item."))
	reg.Add(tools.NewFuncTool("todo_item_mark_done", ToDoMarkItemDone, "Marks ToDo item with given ID as done."))
	reg.Add(tools.NewFuncTool("todo_item_add_comment", ToDoItemAddComment, "Adds discussion comment to given ToDo item"))
}

func GetToolSchema(schema *jsonschema.Schema) (anthropic.ToolInputSchemaParam, error) {
	schemaBytes, err := json.Marshal(schema)
	if err != nil {
		return anthropic.ToolInputSchemaParam{}, err
	}
	var schemaMap map[string]any
	if err := json.Unmarshal(schemaBytes, &schemaMap); err != nil {
		return anthropic.ToolInputSchemaParam{}, err
	}

	inputSchema, err := parseSchemaMap(schemaMap)
	if err != nil {
		return anthropic.ToolInputSchemaParam{}, err
	}
	return inputSchema, nil
}

func parseSchemaMap(s map[string]any) (anthropic.ToolInputSchemaParam, error) {
	bytes, err := json.Marshal(s)
	if err != nil {
		return anthropic.ToolInputSchemaParam{}, fmt.Errorf("failed to marshal schema: %w", err)
	}

	var schema anthropic.ToolInputSchemaParam
	if err := json.Unmarshal(bytes, &schema); err != nil {
		return anthropic.ToolInputSchemaParam{}, fmt.Errorf("failed to unmarshal schema: %w", err)
	}

	return schema, nil
}