feat(go/genkit): added persistent chat session and agent support

[dysrama]

func SimpleChat() {
	ctx := context.Background()
	g, err := genkit.New(nil)
	if err != nil {
		log.Fatal(err)
	}

	if err := vertexai.Init(ctx, g, nil); err != nil {
		log.Fatal(err)
	}

	m := vertexai.Model(g, "gemini-1.5-flash")

	chat, err := genkit.NewChat(
		ctx,
		g,
		genkit.WithModel(m),
		genkit.WithSystemText("You're a pirate first mate. Address the user as Captain and assist them however you can."),
		genkit.WithConfig(ai.GenerationCommonConfig{Temperature: 1.3}),
	)
	if err != nil {
		log.Fatal(err)
	}

	resp, err := chat.Send(ctx, "Hello")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Print(resp.Text())
}

func StatefulChat() {
	ctx := context.Background()
	g, err := genkit.New(nil)
	if err != nil {
		log.Fatal(err)
	}

	if err := vertexai.Init(ctx, g, nil); err != nil {
		log.Fatal(err)
	}

	m := vertexai.Model(g, "gemini-1.5-pro")

	// To include state in a session, you need to instantiate a session explicitly
	session, err := genkit.NewSession(ctx,
		genkit.WithSessionData(genkit.SessionData{
			State: map[string]any{
				"username": "Michael",
			},
		},
		),
	)
	if err != nil {
		log.Fatal(err)
	}

	// Modify the state using tools
	chatTool := genkit.DefineTool(
		g,
		"updateName",
		"use this tool to update the name of the user",
		func(ctx context.Context, input struct {
			Name string
		}) (string, error) {
			// Set name in state
			session, err := genkit.SessionFromContext(ctx)
			if err != nil {
				return "", err
			}

			err = session.UpdateState(map[string]any{
				"username": input.Name,
			})
			if err != nil {
				return "", err
			}

			return "changed username to " + input.Name, nil
		},
	)

	chat, err := genkit.NewChat(
		ctx,
		g,
		genkit.WithModel(m),
		genkit.WithSystemText("You're Kitt from Knight Rider. Address the user as Kitt would and always introduce yourself."),
		genkit.WithConfig(ai.GenerationCommonConfig{Temperature: 1}),
		genkit.WithSession(session),
		genkit.WithTools(chatTool),
	)
	if err != nil {
		log.Fatal(err)
	}

	resp, err := chat.Send(ctx, "Hello, my name is Earl")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Print(resp.Text())
	fmt.Print(session.SessionData.State["username"])
}

func MultiThreadChat() {
	ctx := context.Background()
	g, err := genkit.New(nil)
	if err != nil {
		log.Fatal(err)
	}

	if err := vertexai.Init(ctx, g, nil); err != nil {
		log.Fatal(err)
	}

	m := vertexai.Model(g, "gemini-1.5-flash")

	pirateChat, err := genkit.NewChat(
		ctx,
		g,
		genkit.WithModel(m),
		genkit.WithSystemText("You're a pirate first mate. Address the user as Captain and assist them however you can."),
		genkit.WithConfig(ai.GenerationCommonConfig{Temperature: 1.3}),
		genkit.WithThreadName("pirate"),
	)
	if err != nil {
		log.Fatal(err)
	}

	resp, err := pirateChat.Send(ctx, "Hello")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Print(resp.Text())

	lawyerChat, err := genkit.NewChat(
		ctx,
		g,
		genkit.WithModel(m),
		genkit.WithSystemText("You're a lawyer. Give unsolicited advice no matter what is asked."),
		genkit.WithConfig(ai.GenerationCommonConfig{Temperature: 1.3}),
		genkit.WithThreadName("lawyer"),
	)
	if err != nil {
		log.Fatal(err)
	}

	resp, err = lawyerChat.Send(ctx, "Hello")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Print(resp.Text())
}

func PersistentStorageChat() {
	ctx := context.Background()
	g, err := genkit.New(nil)
	if err != nil {
		log.Fatal(err)
	}

	if err := vertexai.Init(ctx, g, nil); err != nil {
		log.Fatal(err)
	}

	m := vertexai.Model(g, "gemini-1.5-pro")

	// To override default in-mem session storage
	store := &MyOwnSessionStore{
		SessionData: make(map[string]genkit.SessionData),
	}
	session, err := genkit.NewSession(ctx,
		genkit.WithSessionStore(store),
	)
	if err != nil {
		log.Fatal(err)
	}

	chat, err := genkit.NewChat(
		ctx,
		g,
		genkit.WithModel(m),
		genkit.WithSystemText("You're a helpful chatbox. Help the user."),
		genkit.WithConfig(ai.GenerationCommonConfig{Temperature: 1}),
		genkit.WithSession(session),
	)
	if err != nil {
		log.Fatal(err)
	}

	resp, err := chat.Send(ctx, "Hello, my name is Earl")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Print(resp.Text())

	// Load and use existing session
	session2, err := genkit.LoadSession(ctx, "mySessionID", store)
	if err != nil {
		log.Fatal(err)
	}

	chat2, err := genkit.NewChat(
		ctx,
		g,
		genkit.WithModel(m),
		genkit.WithSystemText("You're a helpful chatbox. Help the user."),
		genkit.WithConfig(ai.GenerationCommonConfig{Temperature: 1}),
		genkit.WithSession(session2),
	)
	if err != nil {
		log.Fatal(err)
	}

	resp, err = chat2.Send(ctx, "What's my name?")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Print(resp.Text())
}

type MyOwnSessionStore struct {
	SessionData map[string]genkit.SessionData
}

func (s *MyOwnSessionStore) Get(sessionId string) (data genkit.SessionData, err error) {
	d, err := os.ReadFile("/tmp/" + sessionId)
	if err != nil {
		return data, err
	}
	err = json.Unmarshal(d, &data)
	if err != nil {
		return data, err
	}

	if data.Threads == nil {
		data.Threads = make(map[string][]*ai.Message)
	}

	s.SessionData[sessionId] = data
	return s.SessionData[sessionId], nil
}

func (s *MyOwnSessionStore) Save(sessionId string, data genkit.SessionData) error {
	s.SessionData[sessionId] = data
	d, err := json.Marshal(data)
	if err != nil {
		return err
	}
	err = os.WriteFile("/tmp/"+sessionId, []byte(d), 0644)
	if err != nil {
		return err
	}
	return nil
}

[apascal07]

Louise, could you please a sample that uses this new API? We'll use it as the canonical one for playing around with this in Dev UI.

[dysrama]

Louise, could you please a sample that uses this new API? We'll use it as the canonical one for playing around with this in Dev UI.

Added in description, went with the same examples we have in node

[apascal07]

Can you actually add

Louise, could you please a sample that uses this new API? We'll use it as the canonical one for playing around with this in Dev UI.

Added in description, went with the same examples we have in node

Can you actually commit this to the samples directory please? We want to have a canonical example that we can run anytime.

[dysrama]

Can you actually commit this to the samples directory please? We want to have a canonical example that we can run anytime.

Yep, done

[apascal07]

We should be able to have a stricter type here like in prompt and generate cases, no?

[dysrama]

WithOutputSchema in generate also takes any and doesn't do any type checking? The schema in chat is just passed along to the generate WithOutputSchema function

[Dasio]

Any update on this? I would like to try switch from Node to Go, but this is the feature I'm missing

[apascal07]

Any update on this? I would like to try switch from Node to Go, but this is the feature I'm missing

Hi David, thanks for your interest in this! We're going to come back to this shortly and try to have it out in the next few weeks. Thanks for your patience.

[fterrag]

@apascal07 Hi there! What's needed to move this PR along?

[apascal07]

@apascal07 Hi there! What's needed to move this PR along?

We will pick work back up here very soon with a redefined design and scope.