-
-
Notifications
You must be signed in to change notification settings - Fork 735
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
memory: Add Zep Cloud integration (#746)
* memory: Add Zep Cloud integration
- Loading branch information
1 parent
ade0795
commit 89195f8
Showing
9 changed files
with
492 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "github.com/getzep/zep-go" | ||
| zepClient "github.com/getzep/zep-go/client" | ||
| zepOption "github.com/getzep/zep-go/option" | ||
| "github.com/tmc/langchaingo/chains" | ||
| "github.com/tmc/langchaingo/llms/openai" | ||
| zepLangchainMemory "github.com/tmc/langchaingo/memory/zep" | ||
| "os" | ||
| ) | ||
|
|
||
| func main() { | ||
| ctx := context.Background() | ||
|
|
||
| client := zepClient.NewClient(zepOption.WithAPIKey(os.Getenv("ZEP_API_KEY"))) | ||
| sessionID := os.Getenv("ZEP_SESSION_ID") | ||
|
|
||
| llm, err := openai.New() | ||
| if err != nil { | ||
| fmt.Printf("Error: %s\n", err) | ||
| } | ||
|
|
||
| c := chains.NewConversation( | ||
| llm, | ||
| zepLangchainMemory.NewMemory( | ||
| client, | ||
| sessionID, | ||
| zepLangchainMemory.WithMemoryType(zep.MemoryGetRequestMemoryTypePerpetual), | ||
| zepLangchainMemory.WithReturnMessages(true), | ||
| zepLangchainMemory.WithAIPrefix("Robot"), | ||
| zepLangchainMemory.WithHumanPrefix("Joe"), | ||
| ), | ||
| ) | ||
| res, err := chains.Run(ctx, c, "Hi! I'm John Doe") | ||
| if err != nil { | ||
| fmt.Printf("Error: %s\n", err) | ||
| } | ||
| fmt.Printf("Response: %s\n", res) | ||
|
|
||
| res, err = chains.Run(ctx, c, "What is my name?") | ||
| if err != nil { | ||
| fmt.Printf("Error: %s\n", err) | ||
| } | ||
| fmt.Printf("Response: %s\n", res) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,164 @@ | ||
| package zep | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| "github.com/getzep/zep-go" | ||
| zepClient "github.com/getzep/zep-go/client" | ||
| "github.com/rs/zerolog/log" | ||
| "github.com/tmc/langchaingo/llms" | ||
| "github.com/tmc/langchaingo/schema" | ||
| ) | ||
|
|
||
| // ChatMessageHistory is a struct that stores chat messages. | ||
| type ChatMessageHistory struct { | ||
| ZepClient *zepClient.Client | ||
| SessionID string | ||
| MemoryType zep.MemoryGetRequestMemoryType | ||
| HumanPrefix string | ||
| AIPrefix string | ||
| } | ||
|
|
||
| // Statically assert that ZepChatMessageHistory implement the chat message history interface. | ||
| var _ schema.ChatMessageHistory = &ChatMessageHistory{} | ||
|
|
||
| // NewZepChatMessageHistory creates a new ZepChatMessageHistory using chat message options. | ||
| func NewZepChatMessageHistory(zep *zepClient.Client, sessionID string, options ...ChatMessageHistoryOption) *ChatMessageHistory { | ||
| messageHistory := applyZepChatHistoryOptions(options...) | ||
| messageHistory.ZepClient = zep | ||
| messageHistory.SessionID = sessionID | ||
| return messageHistory | ||
| } | ||
|
|
||
| func (h *ChatMessageHistory) messagesFromZepMessages(zepMessages []*zep.Message) []llms.ChatMessage { | ||
| var chatMessages []llms.ChatMessage | ||
| for _, zepMessage := range zepMessages { | ||
| switch *zepMessage.RoleType { // nolint We do not store other message types in zep memory | ||
| case zep.RoleTypeUserRole: | ||
| chatMessages = append(chatMessages, llms.HumanChatMessage{Content: *zepMessage.Content}) | ||
| case zep.RoleTypeAssistantRole: | ||
| chatMessages = append(chatMessages, llms.AIChatMessage{Content: *zepMessage.Content}) | ||
| case zep.RoleTypeToolRole: | ||
| case zep.RoleTypeFunctionRole: | ||
| chatMessages = append(chatMessages, llms.ToolChatMessage{Content: *zepMessage.Content}) | ||
| default: | ||
| log.Print(fmt.Errorf("unknown role: %s", *zepMessage.RoleType)) | ||
| continue | ||
| } | ||
| } | ||
| return chatMessages | ||
| } | ||
|
|
||
| func (h *ChatMessageHistory) messagesToZepMessages(messages []llms.ChatMessage) []*zep.Message { | ||
| var zepMessages []*zep.Message //nolint We don't know the final size of the messages as some might be skipped due to unsupported role. | ||
| for _, m := range messages { | ||
| zepMessage := zep.Message{ | ||
| Content: zep.String(m.GetContent()), | ||
| } | ||
| switch m.GetType() { // nolint We only expect to bring these three types into chat history | ||
| case llms.ChatMessageTypeHuman: | ||
| zepMessage.RoleType = zep.RoleTypeUserRole.Ptr() | ||
| if h.HumanPrefix != "" { | ||
| zepMessage.Role = zep.String(h.HumanPrefix) | ||
| } | ||
| case llms.ChatMessageTypeAI: | ||
| zepMessage.RoleType = zep.RoleTypeAssistantRole.Ptr() | ||
| if h.AIPrefix != "" { | ||
| zepMessage.Role = zep.String(h.AIPrefix) | ||
| } | ||
| case llms.ChatMessageTypeFunction: | ||
| zepMessage.RoleType = zep.RoleTypeFunctionRole.Ptr() | ||
| case llms.ChatMessageTypeTool: | ||
| zepMessage.RoleType = zep.RoleTypeToolRole.Ptr() | ||
| default: | ||
| log.Print(fmt.Errorf("unknown role: %s", *zepMessage.RoleType)) | ||
| continue | ||
| } | ||
| zepMessages = append(zepMessages, &zepMessage) | ||
| } | ||
| return zepMessages | ||
| } | ||
|
|
||
| // Messages returns all messages stored. | ||
| func (h *ChatMessageHistory) Messages(ctx context.Context) ([]llms.ChatMessage, error) { | ||
| memory, err := h.ZepClient.Memory.Get(ctx, h.SessionID, &zep.MemoryGetRequest{ | ||
| MemoryType: h.MemoryType.Ptr(), | ||
| }) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| messages := h.messagesFromZepMessages(memory.Messages) | ||
| zepFacts := memory.Facts | ||
| systemPromptContent := "" | ||
| for _, fact := range zepFacts { | ||
| systemPromptContent += fmt.Sprintf("%s\n", fact) | ||
| } | ||
| if memory.Summary != nil && memory.Summary.Content != nil { | ||
| systemPromptContent += fmt.Sprintf("%s\n", *memory.Summary.Content) | ||
| } | ||
| if systemPromptContent != "" { | ||
| // Add system prompt to the beginning of the messages. | ||
| messages = append( | ||
| []llms.ChatMessage{ | ||
| llms.SystemChatMessage{ | ||
| Content: systemPromptContent, | ||
| }, | ||
| }, | ||
| messages..., | ||
| ) | ||
| } | ||
| return messages, nil | ||
| } | ||
|
|
||
| // AddAIMessage adds an AIMessage to the chat message history. | ||
| func (h *ChatMessageHistory) AddAIMessage(ctx context.Context, text string) error { | ||
| _, err := h.ZepClient.Memory.Add(ctx, h.SessionID, &zep.AddMemoryRequest{ | ||
| Messages: h.messagesToZepMessages( | ||
| []llms.ChatMessage{ | ||
| llms.AIChatMessage{Content: text}, | ||
| }, | ||
| ), | ||
| }) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // AddUserMessage adds a user to the chat message history. | ||
| func (h *ChatMessageHistory) AddUserMessage(ctx context.Context, text string) error { | ||
| _, err := h.ZepClient.Memory.Add(ctx, h.SessionID, &zep.AddMemoryRequest{ | ||
| Messages: h.messagesToZepMessages( | ||
| []llms.ChatMessage{ | ||
| llms.HumanChatMessage{Content: text}, | ||
| }, | ||
| ), | ||
| }) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func (h *ChatMessageHistory) Clear(ctx context.Context) error { | ||
| _, err := h.ZepClient.Memory.Delete(ctx, h.SessionID) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func (h *ChatMessageHistory) AddMessage(ctx context.Context, message llms.ChatMessage) error { | ||
| _, err := h.ZepClient.Memory.Add(ctx, h.SessionID, &zep.AddMemoryRequest{ | ||
| Messages: h.messagesToZepMessages([]llms.ChatMessage{message}), | ||
| }) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func (*ChatMessageHistory) SetMessages(_ context.Context, _ []llms.ChatMessage) error { | ||
| return nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| package zep | ||
|
|
||
| import "github.com/getzep/zep-go" | ||
|
|
||
| // ChatMessageHistoryOption is a function for creating new chat message history | ||
| // with other than the default values. | ||
| type ChatMessageHistoryOption func(m *ChatMessageHistory) | ||
|
|
||
| // WithChatHistoryMemoryType specifies zep memory type. | ||
| func WithChatHistoryMemoryType(memoryType zep.MemoryGetRequestMemoryType) ChatMessageHistoryOption { | ||
| return func(b *ChatMessageHistory) { | ||
| b.MemoryType = memoryType | ||
| } | ||
| } | ||
|
|
||
| // WithChatHistoryHumanPrefix is an option for specifying the human prefix. Will be passed as role for the message to zep. | ||
| func WithChatHistoryHumanPrefix(humanPrefix string) ChatMessageHistoryOption { | ||
| return func(b *ChatMessageHistory) { | ||
| b.HumanPrefix = humanPrefix | ||
| } | ||
| } | ||
|
|
||
| // WithChatHistoryAIPrefix is an option for specifying the AI prefix. Will be passed as role for the message to zep. | ||
| func WithChatHistoryAIPrefix(aiPrefix string) ChatMessageHistoryOption { | ||
| return func(b *ChatMessageHistory) { | ||
| b.AIPrefix = aiPrefix | ||
| } | ||
| } | ||
|
|
||
| func applyZepChatHistoryOptions(options ...ChatMessageHistoryOption) *ChatMessageHistory { | ||
| h := &ChatMessageHistory{ | ||
| MemoryType: zep.MemoryGetRequestMemoryTypePerpetual, | ||
| } | ||
|
|
||
| for _, option := range options { | ||
| option(h) | ||
| } | ||
|
|
||
| return h | ||
| } |
Oops, something went wrong.