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

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{pr, 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 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 ToDoSubItem struct {
	Title       string `json:"title" jsonschema:"title=title,description=high level title of the TODO item,required"`
	Description string `json:"description" jsonschema:"title=description,description=detailed description what this TODO item is about"`
	AssignedTo  string `json:"assignedTo" jsonschema:"title=assigned to,description=name of the person who shall work on this utem"`
}

type ToDoItem struct {
	ParentID    string        `json:"parentId" jsonschema:"title=parent item id,description=ID of the parent TODO item this one shall be added to as a sub-item"`
	Title       string        `json:"title" jsonschema:"title=title,description=high level title of the TODO item,required"`
	Description string        `json:"description" jsonschema:"title=description,description=detailed description what this TODO item is about"`
	AssignedTo  string        `json:"assignedTo" jsonschema:"title=assigned to,description=name of the person who shall work on this utem"`
	Items       []ToDoSubItem `json:"items" jsonschema:"title=sub items,description:array of sub-items"`
	Parallel    bool          `json:"parallel" jsonschema:"title=parallel,description=if true sub-items may be worked on in parallel and there is no depencency between them, otherwise they shall be worked on sequentially"`
}

type ToDoAddItemArgs struct {
	Items []ToDoItem `json:"items" jsonschema:"title=items,description="items to add to the TODO list,required""`
}

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,
			Parallel:    td.Parallel,
		})
		ti := item.Items[len(item.Items)-1]
		for _, s := range td.Items {
			sid := fmt.Sprintf("%s.%d", id, len(ti.Items)+1)
			ti.Items = append(ti.Items, &ToDo{
				ID:          sid,
				Title:       s.Title,
				Description: s.Description,
				AssignedTo:  s.AssignedTo,
			})
		}
	}
	return "done", nil
}

type ToDoMarkItemDoneArgs struct {
	ID      string `json:"id" jsonschema:"title=id,description=id of the TODO item to mark as DONE,required"`
	Summary string `json:"summary" jsonschem:"title=summary,description=detailed summary of current item and all of it's sub-item trees.,required"`
}

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
	item.Summary = args.Summary
	return "done", nil
}

type ToDoItemAddCommentArgs struct {
	ID       string `json:"id" jsonschema:"title=id,description=id of the TODO item to add comment to,required"`
	Comment  string `json:"comment" jsonschema:"title=comment,description=actual comment text,required"`
	AssignTo string `json:"assignTo" jsonschema:"title=assigned to,description=name of the person who shall be assigned to TODO item with given ID, if empty assignment does not chage"`
}

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,
	})
	if args.AssignTo != "" {
		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
}