package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"

	"github.com/sashabaranov/go-openai"
	"github.com/sashabaranov/go-openai/jsonschema"
	"github.com/zenmodel/zenmodel"
)

func main() {
	bp := zenmodel.NewBrainPrint()
	// add neuron
	bp.AddNeuron("llm", chatLLM)
	bp.AddNeuron("action", callTools)

	/* This example omits error handling */
	// add entry link
	_, _ = bp.AddEntryLink("llm")

	// add link
	continueLink, _ := bp.AddLink("llm", "action")
	_, _ = bp.AddLink("action", "llm")

	// add end link
	endLink, _ := bp.AddEndLink("llm")

	// add link to cast group of a neuron
	_ = bp.AddLinkToCastGroup("llm", "continue", continueLink)
	_ = bp.AddLinkToCastGroup("llm", "end", endLink)
	// bind cast group select function for neuron
	_ = bp.BindCastGroupSelectFunc("llm", llmNext)

	// build brain
	brain := bp.Build()
	// set memory and trig all entry links
	_ = brain.EntryWithMemory(
		"messages", []openai.ChatCompletionMessage{{Role: openai.ChatMessageRoleUser, Content: "What is the weather in Boston today?"}})

	// block process util brain sleeping
	brain.Wait()

	messages, _ := json.Marshal(brain.GetMemory("messages"))
	fmt.Printf("messages: %s\n", messages)
}

// describe the function & its inputs
var tools = []openai.Tool{
	{
		Type: openai.ToolTypeFunction,
		Function: &openai.FunctionDefinition{
			Name:        "get_current_weather",
			Description: "Get the current weather in a given location",
			Parameters: jsonschema.Definition{
				Type: jsonschema.Object,
				Properties: map[string]jsonschema.Definition{
					"location": {
						Type:        jsonschema.String,
						Description: "The city and state, e.g. San Francisco, CA",
					},
					"unit": {
						Type: jsonschema.String,
						Enum: []string{"celsius", "fahrenheit"},
					},
				},
				Required: []string{"location"},
			},
		},
	},
}

func chatLLM(b zenmodel.BrainRuntime) error {
	fmt.Println("run here chatLLM...")

	// get need info form memory
	messages, _ := b.GetMemory("messages").([]openai.ChatCompletionMessage)

	ctx := context.Background()
	client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
	resp, err := client.CreateChatCompletion(ctx,
		openai.ChatCompletionRequest{
			Model:    openai.GPT3Dot5Turbo0125,
			Messages: messages,
			Tools:    tools,
		},
	)
	if err != nil || len(resp.Choices) != 1 {
		return fmt.Errorf("Completion error: err:%v len(choices):%v\n", err,
			len(resp.Choices))
	}

	msg := resp.Choices[0].Message
	fmt.Printf("LLM response: %+v\n", msg)
	messages = append(messages, msg)
	_ = b.SetMemory("messages", messages)

	return nil
}

func callTools(b zenmodel.BrainRuntime) error {
	fmt.Println("run here callTools...")

	// get need info form memory
	messages, _ := b.GetMemory("messages").([]openai.ChatCompletionMessage)
	lastMsg := messages[len(messages)-1]

	for _, call := range lastMsg.ToolCalls {
		if call.Function.Name == "get_current_weather" {
			// 根据 call.Function.Name 和 call.Function.Arguments 发起函数调用，此处模拟调用并 mock 调用结果
			fmt.Printf("call tool [%s] with arguments [%s]\n", call.Function.Name, call.Function.Arguments)
			// mock 返回 toolCalledResp
			toolCalledResp := "Sunny and 80 degrees."

			messages = append(messages, openai.ChatCompletionMessage{
				Role:       openai.ChatMessageRoleTool,
				Content:    toolCalledResp,
				Name:       lastMsg.ToolCalls[0].Function.Name,
				ToolCallID: lastMsg.ToolCalls[0].ID,
			})
		}
	}
	_ = b.SetMemory("messages", messages)

	return nil
}

func llmNext(b zenmodel.BrainRuntime) string {
	if !b.ExistMemory("messages") {
		return "end"
	}
	messages, _ := b.GetMemory("messages").([]openai.ChatCompletionMessage)
	lastMsg := messages[len(messages)-1]
	if len(lastMsg.ToolCalls) == 0 { // no need to call any tools
		return "end"
	}

	return "continue"
}

package main

import (
	"fmt"

	"github.com/zenmodel/zenmodel"
)

func main() {
	bp := zenmodel.NewBrainPrint()
	bp.AddNeuron("n1", fn1)
	bp.AddNeuron("n2", fn2)
	_, err := bp.AddLink("n1", "n2")
	if err != nil {
		fmt.Printf("add link error: %s\n", err)
		return
	}
	_, err = bp.AddEntryLink("n1")
	if err != nil {
		fmt.Printf("add entry link error: %s\n", err)
		return
	}

	brain := bp.Build()

	_ = brain.Entry()

	brain.Wait()

	name := brain.GetMemory("name").(string)
	fmt.Printf("result: my name is %s.\n", name)
}

func fn1(b zenmodel.BrainRuntime) error {
	fmt.Println("start fn1 ..............")

	if err := b.SetMemory("name", "Clay"); err != nil {
		return err
	}

	return nil
}

func fn2(b zenmodel.BrainRuntime) error {
	fmt.Println("start fn2 ..............")

	firstName := b.GetMemory("name").(string)

	name := firstName + " Zhang"
	if err := b.SetMemory("name", name); err != nil {
		return err
	}
	return nil
}

package main

import (
	"fmt"

	"github.com/zenmodel/zenmodel"
)

var (
	NeuronLeader     = "Leader"
	NeuronRD         = "RD"
	NeuronQA         = "QA"
	DecisionRD       = "RD"
	DecisionQA       = "QA"
	DecisionResponse = "Response"
)

func main() {
	bp := zenmodel.NewBrainPrint()

	bp.AddNeuron(NeuronLeader, LeaderProcess)
	bp.AddNeuron(NeuronQA, QAProcess)
	bp.AddNeuronWithProcessor(NeuronRD, NewRDProcessor())

	_, _ = bp.AddEntryLink(NeuronLeader)
	// leader out-link
	rdLink, _ := bp.AddLink(NeuronLeader, NeuronRD)
	qaLink, _ := bp.AddLink(NeuronLeader, NeuronQA)
	endLink, _ := bp.AddEndLink(NeuronLeader)

	// leader in-link
	_, _ = bp.AddLink(NeuronRD, NeuronLeader)
	_, _ = bp.AddLink(NeuronQA, NeuronLeader)

	_ = bp.AddLinkToCastGroup(NeuronLeader, DecisionRD, rdLink)
	_ = bp.AddLinkToCastGroup(NeuronLeader, DecisionQA, qaLink)
	_ = bp.AddLinkToCastGroup(NeuronLeader, DecisionResponse, endLink)
	_ = bp.BindCastGroupSelectFunc(NeuronLeader, func(b zenmodel.BrainRuntime) string {
		return b.GetMemory(memKeyDecision).(string)
	})

	brain := bp.Build()
	_ = brain.EntryWithMemory(memKeyDemand, "Help me write a function `func Add (x, y int) int` with golang to implement addition, and implement unit test in a separate _test .go file, at least 3 test cases are required")
	brain.Wait()
	fmt.Printf("Response: %s\n", brain.GetMemory(memKeyResponse).(string))
}

package main

import "testing"

func TestAdd(t *testing.T) {
		cases := []struct {
				x, y, expected int
		}{
				{1, 2, 3},
				{-1, 1, 0},
				{0, 0, 0},
		}

		for _, c := range cases {
				result := Add(c.x, c.y)
				if result != c.expected {
						t.Errorf("Add(%d, %d) == %d, expected %d", c.x, c.y, result, c.expected)
				}
		}
}

#go test -v -run .
=== RUN   TestAdd

package main

import (
	"fmt"
	"strings"

	"github.com/zenmodel/zenmodel"
	"github.com/zenmodel/zenmodel-contrib/processor/go_code_tester"
)

const (
	memKeyDemand   = "demand"
	memKeyResponse = "response"
	memKeyTask     = "task"
	memKeyFeedback = "feedback"
	memKeyDecision = "decision"
)

func LeaderProcess(b zenmodel.BrainRuntime) error {
	// if it has no task, disassemble task from demand
	if !b.ExistMemory(memKeyTask) {
		task := rephraseTaskFromDemand(b.GetMemory(memKeyDemand).(string))
		_ = b.SetMemory(memKeyTask, task)
		_ = b.SetMemory(memKeyDecision, DecisionRD)

		return nil
	}
	switch b.GetMemory(memKeyFeedback).(string) {
	case NeuronRD: // feedback from RD
		_ = b.SetMemory(memKeyDecision, DecisionQA) // pass to QA
	case NeuronQA: // feedback from QA
		ok := readTestReport(b.GetMemory(memKeyGoTestResult).(string))
		if !ok {
			// test result not ok, resend to RD
			_ = b.SetMemory(memKeyDecision, DecisionRD)
		} else {
			// pretty response from codes
			resp := genResponse(b)
			_ = b.SetMemory(memKeyResponse, resp)
			_ = b.SetMemory(memKeyDecision, DecisionResponse)
		}
	default:
		return fmt.Errorf("unknown feedback: %v\n", b.GetMemory(memKeyFeedback))
	}

	return nil
}

func rephraseTaskFromDemand(demand string) string {
	// TODO maybe use completion LLM to rephrase demand to task
	task := demand

	return task
}

func readTestReport(testResult string) bool {
	return !strings.Contains(testResult, "FAIL")
}

func genResponse(b zenmodel.BrainRuntime) string {
	codes := b.GetMemory(memKeyCodes).(*go_code_tester.Codes).String()
	testReport := b.GetMemory(memKeyGoTestResult).(string)

	var builder strings.Builder
	builder.WriteString("Dear Boss:  \n")
	builder.WriteString("After the efforts of our RD team and QA team, the final codes and test report are produced as follows:\n\n")
	builder.WriteString("==========\n\nCodes:\n\n")
	builder.WriteString(codes)
	builder.WriteString("==========\n\nTest Report:\n\n")
	builder.WriteString("```shell\n")
	builder.WriteString(testReport)
	builder.WriteString("```")
	builder.WriteString("\n")

	return builder.String()
}

package main

import (
	"fmt"
	"os"

	"github.com/sashabaranov/go-openai"
	"github.com/zenmodel/zenmodel"
	"github.com/zenmodel/zenmodel-contrib/processor/go_code_tester"
	"github.com/zenmodel/zenmodel-contrib/processor/openai_structured_output"
)

var (
	memKeyCodes = (&go_code_tester.Codes{}).FunctionName()
)

func NewRDProcessor() *CoderProcessor {
	return &CoderProcessor{
		clientConfig: openai.DefaultConfig(os.Getenv("OPENAI_API_KEY")),
		requestConfig: openai_structured_output.RequestConfig{
			Model:       openai.GPT3Dot5Turbo0125,
			Temperature: 0.7,
			Stream:      false,
		},
	}
}

type CoderProcessor struct {
	clientConfig  openai.ClientConfig
	client        *openai.Client
	requestConfig openai_structured_output.RequestConfig
}

func (p *CoderProcessor) Process(b zenmodel.BrainRuntime) error {
	var prompt string
	if !b.ExistMemory(memKeyCodes) {
		// read task, write code
		prompt = fmt.Sprintf(`{{.%s}}`, memKeyTask)
	} else {
		// read task, old code and test result, write code
		prompt = fmt.Sprintf(`{{.%s}}

My code is as follows:

%s

test result is as follows:

%s

Help me correct my code.
`, memKeyTask, b.GetMemory(memKeyCodes).(*go_code_tester.Codes).String(), memKeyGoTestResult)
	}

	structuredOutput := p.newStructuredOutputProcessor(prompt)
	if err := structuredOutput.Process(b); err != nil {
		return err
	}

	if err := b.SetMemory(memKeyFeedback, b.GetCurrentNeuronID()); err != nil {
		return err
	}

	return nil
}

func (p *CoderProcessor) DeepCopy() zenmodel.Processor {
	return &CoderProcessor{
		requestConfig: p.requestConfig,
		clientConfig:  p.clientConfig,
		client:        nil,
	}
}

func (p *CoderProcessor) WithClientConfig(clientConfig openai.ClientConfig) *CoderProcessor {
	p.clientConfig = clientConfig
	return p
}

func (p *CoderProcessor) WithClient(client *openai.Client) *CoderProcessor {
	p.client = client
	return p
}

func (p *CoderProcessor) WithRequestConfig(requestConfig openai_structured_output.RequestConfig) *CoderProcessor {
	p.requestConfig = requestConfig
	return p
}

func (p *CoderProcessor) newStructuredOutputProcessor(prompt string) *openai_structured_output.OpenAIStructuredOutputProcessor {
	processor := openai_structured_output.NewProcessor()
	_ = processor.WithPromptTemplate(prompt)
	_ = processor.WithOutputStructDefinition(go_code_tester.Codes{}, (go_code_tester.Codes{}).FunctionName(), (go_code_tester.Codes{}).FunctionDescription())
	return processor.WithClientConfig(p.clientConfig).WithRequestConfig(p.requestConfig)
}

package main

import (
	"github.com/zenmodel/zenmodel"
	"github.com/zenmodel/zenmodel-contrib/processor/go_code_tester"
)

const (
	memKeyGoTestResult = "go_test_result"
)

func QAProcess(b zenmodel.BrainRuntime) error {
	p := go_code_tester.NewProcessor().WithTestCodeKeep(true)
	if err := p.Process(b); err != nil {
		return err
	}

	if err := b.SetMemory(memKeyFeedback, b.GetCurrentNeuronID()); err != nil {
		return err
	}

	return nil
}

package main

import (
	"fmt"

	"github.com/zenmodel/zenmodel"
)

var (
	entryInput, entryPoetry, entryJoke string
)

func main() {
	bp := zenmodel.NewBrainPrint()
	bp.AddNeuron("input", inputFn)
	bp.AddNeuron("poetry-template", poetryFn)
	bp.AddNeuron("joke-template", jokeFn)
	bp.AddNeuron("generate", genFn)

	inputIn, _ := bp.AddLink("input", "generate")
	poetryIn, _ := bp.AddLink("poetry-template", "generate")
	jokeIn, _ := bp.AddLink("joke-template", "generate")

	entryInput, _ = bp.AddEntryLink("input")
	entryPoetry, _ = bp.AddEntryLink("poetry-template")
	entryJoke, _ = bp.AddEntryLink("joke-template")

	_ = bp.AddTriggerGroup("generate", inputIn, poetryIn)
	_ = bp.AddTriggerGroup("generate", inputIn, jokeIn)

	brain := bp.Build()

	// case 1: entry poetry and input
	// expect: generate poetry
	_ = brain.TrigLinks(entryPoetry)
	_ = brain.TrigLinks(entryInput)

	// case 2:entry joke and input
	// expect: generate joke
	//_ = brain.TrigLinks(entryJoke)
	//_ = brain.TrigLinks(entryInput)

	// case 3: entry poetry and joke
	// expect: keep blocking and waiting for any trigger group triggered
	//_ = brain.TrigLinks(entryPoetry)
	//_ = brain.TrigLinks(entryJoke)

	// case 4: entry only poetry
	// expect: keep blocking and waiting for any trigger group triggered
	//_ = brain.TrigLinks(entryPoetry)

	// case 5: entry all
	// expect: The first done trigger group triggered activates the generated Neuron,
	// and the trigger group triggered later does not activate the generated Neuron again.
	//_ = brain.Entry()

	brain.Wait()
}

func inputFn(b zenmodel.BrainRuntime) error {
	_ = b.SetMemory("input", "orange")
	return nil
}

func poetryFn(b zenmodel.BrainRuntime) error {
	_ = b.SetMemory("template", "poetry")
	return nil
}

func jokeFn(b zenmodel.BrainRuntime) error {
	_ = b.SetMemory("template", "joke")
	return nil
}

func genFn(b zenmodel.BrainRuntime) error {
	input := b.GetMemory("input").(string)
	tpl := b.GetMemory("template").(string)
	fmt.Printf("Generating %s for %s\n", tpl, input)
	return nil
}

package main

import (
	"fmt"

	"github.com/zenmodel/zenmodel"
)

func main() {
	bp := zenmodel.NewBrainPrint()
	bp.AddNeuron("condition", func(runtime zenmodel.BrainRuntime) error {
		return nil // do nothing
	})
	bp.AddNeuron("cell-phone", func(runtime zenmodel.BrainRuntime) error {
		fmt.Printf("Run here: Cell Phone\n")
		return nil
	})
	bp.AddNeuron("laptop", func(runtime zenmodel.BrainRuntime) error {
		fmt.Printf("Run here: Laptop\n")
		return nil
	})
	bp.AddNeuron("ps5", func(runtime zenmodel.BrainRuntime) error {
		fmt.Printf("Run here: PS5\n")
		return nil
	})
	bp.AddNeuron("tv", func(runtime zenmodel.BrainRuntime) error {
		fmt.Printf("Run here: TV\n")
		return nil
	})
	bp.AddNeuron("printer", func(runtime zenmodel.BrainRuntime) error {
		fmt.Printf("Run here: Printer\n")
		return nil
	})

	cellPhone, _ := bp.AddLink("condition", "cell-phone")
	laptop, _ := bp.AddLink("condition", "laptop")
	ps5, _ := bp.AddLink("condition", "ps5")
	tv, _ := bp.AddLink("condition", "tv")
	printer, _ := bp.AddLink("condition", "printer")
	// add entry link
	_, _ = bp.AddEntryLink("condition")

	/*
	   Category 1: Electronics
	   - Cell Phone
	   - Laptop
	   - PS5

	   Category 2: Entertainment Devices
	   - Cell Phone
	   - PS5
	   - TV

	   Category 3: Office Devices
	   - Laptop
	   - Printer
	   - Cell Phone
	*/
	_ = bp.AddLinkToCastGroup("condition", "electronics",
		cellPhone, laptop, ps5)
	_ = bp.AddLinkToCastGroup("condition",
		"entertainment-devices",
		cellPhone, ps5, tv)
	_ = bp.AddLinkToCastGroup(
		"condition", "office-devices",
		laptop, printer, cellPhone)

	_ = bp.BindCastGroupSelectFunc("condition", func(brain zenmodel.BrainRuntime) string {
		return brain.GetMemory("category").(string)
	})

	brain := bp.Build()

	_ = brain.EntryWithMemory("category", "electronics")
	//_ = brain.EntryWithMemory("category", "entertainment-devices")
	//_ = brain.EntryWithMemory("category", "office-devices")
	//_ = brain.EntryWithMemory("category", "NOT-Defined")

	brain.Wait()
}

package main

import (
	"fmt"

	"github.com/zenmodel/zenmodel"
)

func main() {
	bp := zenmodel.NewBrainPrint()
	bp.AddNeuron("nested", nestedBrain)
	_, _ = bp.AddEntryLink("nested")

	brain := bp.Build()
	_ = brain.Entry()
	brain.Wait()

	fmt.Printf("nested result: %s\n", brain.GetMemory("nested_result").(string))
}

func nestedBrain(outerBrain zenmodel.BrainRuntime) error {
	bp := zenmodel.NewBrainPrint()
	bp.AddNeuron("run", func(curBrain zenmodel.BrainRuntime) error {
		_ = curBrain.SetMemory("result", fmt.Sprintf("run here neuron: %s.%s", outerBrain.GetCurrentNeuronID(), curBrain.GetCurrentNeuronID()))
		return nil
	})
	_, _ = bp.AddEntryLink("run")

	brain := bp.Build()

	// run nested brain
	_ = brain.Entry()
	brain.Wait()
	// get nested brain result
	result := brain.GetMemory("result").(string)
	// pass nested brain result to outer brain
	_ = outerBrain.SetMemory("nested_result", result)

	return nil
}