agent_etc

Author: Gekko0114

examples/normal/main.go

@@ -3,13 +3,13 @@ package main
 import (
 	"fmt"
 	"golangchain/pkg/lib"
-	"golangchain/pkg/openai"
+	"golangchain/pkg/llm"
 	"golangchain/pkg/parser"
 	"golangchain/pkg/prompt"
 )
 
 func main() {
-	llm, err := openai.NewChatOpenAI("gpt-3.5-turbo")
+	llm, err := llm.NewChatOpenAI("gpt-3.5-turbo")
 	if err != nil {
 		fmt.Println(err)
 	}

pkg/agent/agent.go

@@ -2,26 +2,72 @@ package agent
 
 import (
 	"fmt"
+	"golangchain/pkg/lib"
+	"golangchain/pkg/parser"
+	"golangchain/pkg/prompt"
 	"strings"
 )
 
-type Agent struct{}
+type Agent struct {
+	LLMChain  *lib.Pipeline
+	Prompt    *prompt.ChatPromptTemplate
+	UserInput string
+}
 
-func NewAgent() *Agent {
-	return &Agent{}
+func NewAgent(tools map[string]Tool, llm lib.Runnable) (*Agent, error) {
+	parser := parser.NewStrOutputParser()
+	prompt, err := CreatePrompt(tools)
+	if err != nil {
+		return nil, err
+	}
+	pl := lib.NewPipeline()
+	pl.Pipe(prompt).Pipe(llm).Pipe(parser)
+	agent := &Agent{
+		LLMChain: pl,
+		Prompt:   prompt,
+	}
+
+	return agent, nil
 }
 
-func (a *Agent) CreatePrompt(tools []Tool) {
+func CreatePrompt(tools map[string]Tool) (*prompt.ChatPromptTemplate, error) {
 	var toolStrings []string
 	var toolNames []string
 	for _, tool := range tools {
-		toolStrings = append(toolStrings, fmt.Sprintf("%s: %s", tool.name, tool.description))
-		toolNames = append(toolNames, tool.name)
+		toolStrings = append(toolStrings, fmt.Sprintf("%s: %s", tool.name(), tool.description()))
+		toolNames = append(toolNames, tool.name())
 	}
 	toolStringsJoined := strings.Join(toolStrings, "\n")
 	toolNamesJoined := strings.Join(toolNames, ",")
 	formatInstructions := strings.Replace(FORMAT_INSTRUCTIONS, "{{.ToolNames}}", toolNamesJoined, 1)
-	fmt.Println(toolStringsJoined)
-	fmt.Println(formatInstructions)
+	instructions := []string{
+		SYSTEM_MESSAGE_PREFIX,
+		toolStringsJoined,
+		formatInstructions,
+		SYSTEM_MESSAGE_SUFFIX,
+	}
+	instruction := strings.Join(instructions, "\n\n")
+	prompt, err := prompt.NewChatPromptTemplate(instruction, HUMAN_MESSAGE)
+	if err != nil {
+		return nil, err
+	}
+	return prompt, nil
+}
 
+func (a *Agent) Plan(intermediateSteps []string) (any, error) {
+	agentScratchpad := strings.Join(intermediateSteps, "\n")
+	m := map[string]string{
+		"Input":            a.UserInput,
+		"Agent_scratchpad": agentScratchpad,
+	}
+	prompt, err := a.Prompt.Invoke(m)
+	if err != nil {
+		return nil, err
+	}
+	agentDecision, err := a.LLMChain.Invoke(prompt)
+	if err != nil {
+		return nil, err
+	}
+	fmt.Println(agentDecision)
+	return agentDecision, err
 }

pkg/agent/agent_executor.go

@@ -7,33 +7,67 @@ import (
 
 type AgentExecutor struct {
 	Agent         *Agent
-	Tools         []Tool
+	Tools         map[string]Tool
 	MaxIterations int
 }
 
-func InitializeAgent(tools []Tool, llm *lib.Runnable) *AgentExecutor {
-	agent := NewAgent()
+func InitializeAgent(tools map[string]Tool, llm lib.Runnable) (*AgentExecutor, error) {
+	agent, err := NewAgent(tools, llm)
+	if err != nil {
+		return nil, err
+	}
 
 	return &AgentExecutor{
 		Agent:         agent,
 		Tools:         tools,
 		MaxIterations: 15,
-	}
+	}, nil
 }
 
 func (a *AgentExecutor) Invoke(input any) (any, error) {
-	return nil, nil
+	switch input.(type) {
+	case string:
+		a.Agent.UserInput = input.(string)
+	default:
+		return nil, nil
+	}
+	output, err := a.call()
+	if err != nil {
+		return nil, err
+	}
+	return output, nil
 }
 
-func (a *AgentExecutor) TakeNextStep() any {
-	return nil
+func (a *AgentExecutor) takeNextStep(intermediateSteps []string) (any, error) {
+	var results []any
+	output, err := a.Agent.Plan(intermediateSteps)
+	if err != nil {
+		return nil, err
+	}
+
+	action := output.(string)
+	var observation any
+	if tool, ok := a.Tools[action]; ok {
+		observation, err = tool.run("aaaaa")
+		if err != nil {
+			return nil, err
+		}
+	}
+	results = append(results, observation)
+
+	return results, nil
 }
 
 func (a *AgentExecutor) call() (any, error) {
 	iterations := 0
+	var intermediateSteps []string
 	var nextStepOutput any
 	for a.MaxIterations > iterations {
-		nextStepOutput := a.TakeNextStep()
+		nextStepOutput, err := a.takeNextStep(intermediateSteps)
+		if err != nil {
+			return nil, err
+		}
+		// TODO: fix
 		fmt.Println(nextStepOutput)
 		iterations += 1
 	}

pkg/agent/prompt.go

@@ -26,4 +26,4 @@ Observation: the result of the action
 Thought: I now know the final answer
 Final Answer: the final answer to the original input question`
 const SYSTEM_MESSAGE_SUFFIX = "Begin! Reminder to always use the exact characters `Final Answer` when responding."
-const HUMAN_MESSAGE = "{input}\n\n{agent_scratchpad}"
+const HUMAN_MESSAGE = "{{ .Input }}\n\n{{ .Agent_scratchpad }}"

pkg/agent/tool.go

@@ -1,27 +1,76 @@
 package agent
 
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"os"
+)
+
 var AvailableTools = map[string]Tool{
-	"llm-math": {
-		name:        "Calculator",
-		description: "Useful for when you need to answer questions about math.",
-	},
-	"serpapi": {
-		name:        "Search",
-		description: "A search engine. Useful for when you need to answer questions about current events. Input should be a search query.",
-	}}
-
-type Tool struct {
-	name        string
-	description string
-}
-
-func LoadTools(tools []string) []Tool {
-	confirmedTools := []Tool{}
+	"serpapi": &SerpAPI{}}
+
+type Tool interface {
+	name() string
+	description() string
+	run(input any) (any, error)
+}
+
+func LoadTools(tools []string) map[string]Tool {
+	confirmedTools := map[string]Tool{}
 	for _, toolname := range tools {
 		if tool, ok := AvailableTools[toolname]; ok {
 			confirmedTool := tool
-			confirmedTools = append(confirmedTools, confirmedTool)
+			confirmedTools[toolname] = confirmedTool
 		}
 	}
 	return confirmedTools
 }
+
+type OrganicResult struct {
+	Title string `json:"title"`
+	Link  string `json:"link"`
+}
+
+type SerpApiResponse struct {
+	OrganicResults []OrganicResult `json:"organic_results"`
+}
+
+type SerpAPI struct {
+}
+
+func (s *SerpAPI) name() string {
+	return "Search"
+}
+
+func (s *SerpAPI) description() string {
+	return "A search engine. Useful for when you need to answer questions about current events. Input should be a search query."
+}
+
+func (s *SerpAPI) run(input any) (any, error) {
+	apiKey := os.Getenv("SERPAPI_API_KEY")
+	searchQuery := input.(string)
+	serpApiURL := "https://serpapi.com/search"
+
+	params := url.Values{}
+	params.Add("engine", "google")
+	params.Add("q", searchQuery)
+	params.Add("api_key", apiKey)
+
+	resp, err := http.Get(fmt.Sprintf("%s?%s", serpApiURL, params.Encode()))
+	if err != nil {
+		return nil, err
+	}
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	var response SerpApiResponse
+	if err := json.Unmarshal(body, &response); err != nil {
+		return nil, err
+	}
+	return response, nil
+}

pkg/openai/openai.go renamed to pkg/llm/openai.go

@@ -1,4 +1,4 @@
-package openai
+package llm
 
 import (
 	"bytes"

pkg/openai/openai_test.go renamed to pkg/llm/openai_test.go

@@ -1,4 +1,4 @@
-package openai
+package llm
 
 import (
 	"net/http"

pkg/openai/struct.go renamed to pkg/llm/struct.go

@@ -1,4 +1,4 @@
-package openai
+package llm
 
 // message represents a single message in the chat
 type Message struct {

pkg/parser/stroutputparser.go

@@ -1,7 +1,7 @@
 package parser
 
 import (
-	"golangchain/pkg/openai"
+	"golangchain/pkg/llm"
 )
 
 type StrOutputParser struct {
@@ -13,7 +13,7 @@ func NewStrOutputParser() *StrOutputParser {
 
 func (p *StrOutputParser) Invoke(input any) (any, error) {
 	var output string
-	res, ok := input.(*openai.Response)
+	res, ok := input.(*llm.Response)
 	if ok {
 		output = res.Choices[0].Message.Content
 	} else {

pkg/parser/stroutputparser_test.go

@@ -1,7 +1,7 @@
 package parser
 
 import (
-	"golangchain/pkg/openai"
+	"golangchain/pkg/llm"
 	"testing"
 )
 
@@ -13,9 +13,9 @@ func TestInvoke(t *testing.T) {
 	}{
 		{
 			name: "template is string",
-			input: &openai.Response{
-				Choices: []openai.Choice{
-					{Message: openai.Message{Role: "assistant", Content: "This is the test response"}},
+			input: &llm.Response{
+				Choices: []llm.Choice{
+					{Message: llm.Message{Role: "assistant", Content: "This is the test response"}},
 				},
 			},
 			expected: "This is the test response",

pkg/prompt/chat_prompt.go

@@ -0,0 +1,59 @@
+package prompt
+
+import (
+	"bytes"
+	"fmt"
+	"golangchain/pkg/llm"
+	"text/template"
+)
+
+type ChatTemplate struct {
+	system *template.Template
+	human  *template.Template
+}
+
+type ChatPromptTemplate struct {
+	template *ChatTemplate
+}
+
+func NewChatPromptTemplate(system string, human string) (*ChatPromptTemplate, error) {
+	sysTmpl, err := template.New("system").Parse(system)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse template: %w", err)
+	}
+	humanTmpl, err := template.New("human").Parse(human)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse template: %w", err)
+	}
+	Chat := &ChatTemplate{
+		system: sysTmpl,
+		human:  humanTmpl,
+	}
+	return &ChatPromptTemplate{
+		template: Chat,
+	}, nil
+}
+
+func (t *ChatPromptTemplate) Invoke(input any) (any, error) {
+	var sbuf bytes.Buffer
+	if err := t.template.system.Execute(&sbuf, input); err != nil {
+		return "", err
+	}
+	sMes := llm.Message{
+		Role:    "system",
+		Content: sbuf.String(),
+	}
+	var hbuf bytes.Buffer
+	if err := t.template.human.Execute(&hbuf, input); err != nil {
+		return "", err
+	}
+	hMes := llm.Message{
+		Role:    "user",
+		Content: hbuf.String(),
+	}
+	messages := []llm.Message{
+		sMes,
+		hMes,
+	}
+	return messages, nil
+}