orchestration

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/bububa/atomic-agents/agents"
	"github.com/bububa/atomic-agents/components"
	"github.com/bububa/atomic-agents/components/systemprompt/cot"
	"github.com/bububa/atomic-agents/examples"
	"github.com/bububa/atomic-agents/schema"
	"github.com/bububa/atomic-agents/tools"
	"github.com/bububa/atomic-agents/tools/calculator"
	"github.com/bububa/atomic-agents/tools/orchestration"
	"github.com/bububa/atomic-agents/tools/searxng"
	"github.com/bububa/instructor-go"
)

// Input schema for the Orchestrator Agent. Contains the user's message to be processed.
type Input struct {
	schema.Base
	// ChatMessage The user's input message to be analyzed and responded to.
	ChatMessage string `json:"chat_message" jsonschema:"title=chat_message,description=The user's message to be analyzed and responded to."`
}

type ToolType string

const (
	SearchTool     ToolType = "search"
	CalculatorTool ToolType = "calculator"
)

// Output Combined output schema for the Orchestrator Agent. Contains the tool to use and its parameters.
type Output struct {
	schema.Base
	// Tool The tool to use: 'search' or 'calculator'
	Tool ToolType `json:"tool" jsonschema:"title=tool,enum=search,enum=calculator,description=The tool to use: 'search' or 'calculator'"`
	// SearchParameters the parameters for the search tool
	SearchParameters *searxng.Input `json:"search_parameters" jsonschema:"title=search_parameters,description=The parameters for the search tool. Should only and must have value if tool is 'search'"`
	// CalculatorParameters the parameters for the calculator tool
	CalculatorParameters *calculator.Input `json:"calculator_parameters" jsonschema:"title=calculator_parameters,description=The parameters for the calculator tool. Should only have and must value if tool is 'calculator'"`
}

// FinalAnswer Schema for the final answer generated by the Orchestrator Agent.
type FinalAnswer struct {
	schema.Base
	// FinalAnswer The final answer generated based on the tool output and user query.
	FinalAnswer string `json:"final_answer" jsonschema:"title=final_answer,description=The final answer generated based on the tool output and user query."`
}

type ContextProvider struct{}

func (p *ContextProvider) Title() string {
	return "Current Date"
}

func (p *ContextProvider) Info() string {
	return fmt.Sprintf("Current date in format YYYY-MM-DD:%s", time.Now().Format("2006-01-02"))
}

func main() {
	mockPort := 8080
	mockSearchURL := fmt.Sprintf("http://localhost:%d", mockPort)
	mockQuery := "query with max results"
	mockResult := searxng.Output{
		Results: []searxng.SearchResultItem{
			{Title: "Result with Metadata", URL: "https://example.com/metadata", Content: "Content with metadata", Query: mockQuery, Metadata: "2021-01-01"},
			{Title: "Result with Published Date", Content: "Content with published date", URL: "https://example.com/published-data", Query: mockQuery, PublishedDate: "2022-01-01"},
			{Title: "Result without dates", Content: "Content without dates", URL: "https://example.com/no-dates", Query: mockQuery},
		},
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	srv := startSearxngServer(mockPort, &mockResult)
	defer srv.Shutdown(ctx)
	mem := components.NewMemory(10)

	systemPromptGenerator := cot.New(
		cot.WithBackground([]string{
			"- You are an Orchestrator Agent that decides between using a search tool or a calculator tool based on user input.",
			"- Use the search tool for queries requiring factual information, current events, or specific data.",
			"- Use the calculator tool for mathematical calculations and expressions.",
		}),
		cot.WithOutputInstructs([]string{
			"- Analyze the input to determine whether it requires a web search or a calculation.",
			"- For search queries, use the 'search' tool and provide 1-3 relevant search queries.",
			"- For calculations, use the 'calculator' tool and provide the mathematical expression to evaluate.",
			"- When uncertain, prefer using the search tool.",
			"- Format the output using the appropriate schema.",
		}),
		cot.WithContextProviders(new(ContextProvider)),
	)
	searchTool := searxng.New(searxng.WithBaseURL(mockSearchURL), searxng.WithMaxResults(3))
	calculatorTool := calculator.New()

	// example inputs
	inputs := []string{
		"Who won the Nobel Prize in Physics in 2024?",
		"Please calculate the sine of pi/3 to the third power",
	}
	finalOutput := new(FinalAnswer)
	finalAgent := agents.NewToolAgent[Input, Output, FinalAnswer](
		agents.WithClient(examples.NewInstructor(instructor.ProviderOpenAI)),
		agents.WithMemory(mem),
		agents.WithModel(os.Getenv("OPENAI_MODEL")),
		agents.WithSystemPromptGenerator(systemPromptGenerator),
		agents.WithTemperature(0.5),
		agents.WithMaxTokens(1000),
	)
	toolSelector := func(req *Output) (tools.AnonymousTool, any, error) {
		switch req.Tool {
		case SearchTool:
			return searchTool, req.SearchParameters, nil
		case CalculatorTool:
			return calculatorTool, req.CalculatorParameters, nil
		default:
			return nil, nil, errors.New("unknown tool")
		}
	}
	finalAgent.SetTool(orchestration.New(toolSelector))
	for _, userInput := range inputs {
		input := Input{
			ChatMessage: userInput,
		}
		if err := finalAgent.Run(ctx, &input, finalOutput, nil); err != nil {
			fmt.Println(err)
			return
		}
		fmt.Printf("Agent: %s\n", finalOutput.FinalAnswer)
		finalAgent.ResetMemory()
	}
	// Outputs:
	// Agent: Final Answer is I will search for the winner of the Nobel Prize in Physics in 2024.
	// Agent: The sine of pi/3 is 0.86602540378. Therefore, (sin(pi/3))^3 = 0.86602540378^3 = 0.64951905284.
}

Output: