diff --git a/integration/zep/model.go b/integration/zep/model.go new file mode 100644 index 0000000..cb2d36d --- /dev/null +++ b/integration/zep/model.go @@ -0,0 +1,74 @@ +package zep + +// Message represents a message in a conversation. +type Message struct { + // The content of the message. + Content string `json:"content,omitempty"` + // The timestamp of when the message was created. + CreatedAt string `json:"created_at,omitempty"` + // Metadata associated with the message. + Metadata interface{} `json:"metadata,omitempty"` + // The role of the sender of the message (e.g., "user", "assistant"). + Role string `json:"role,omitempty"` + // The number of tokens in the message. + TokenCount int `json:"token_count,omitempty"` + // The unique identifier of the message. + UUID string `json:"uuid,omitempty"` +} + +// SearchPayload represents a search payload for querying memory. +type SearchPayload struct { + // Metadata associated with the search query. + Meta map[string]interface{} `json:"meta,omitempty"` + // The text of the search query. + Text string `json:"text,omitempty"` +} + +// SearchResult represents a search result from querying memory. +type SearchResult struct { + // The distance metric of the search result. + Dist float64 `json:"dist,omitempty"` + // The message associated with the search result. + Message Message `json:"message,omitempty"` + // Metadata associated with the search result. + Meta interface{} `json:"meta,omitempty"` + // The summary of the search result. + Summary Summary `json:"summary,omitempty"` +} + +// Summary represents a summary of a conversation. +type Summary struct { + // The content of the summary. + Content string `json:"content,omitempty"` + // The timestamp of when the summary was created. + CreatedAt string `json:"created_at,omitempty"` + // // Metadata associated with the summary. + Metadata interface{} `json:"metadata,omitempty"` + // The unique identifier of the most recent message in the conversation. + RecentMessageUUID string `json:"recent_message_uuid,omitempty"` + // The number of tokens in the summary. + TokenCount int `json:"token_count,omitempty"` + // The unique identifier of the summary. + UUID string `json:"uuid,omitempty"` +} + +// Represents a memory object with messages, metadata, and other attributes. +type Memory struct { + // A list of message objects, where each message contains a role and content. + Messages []Message `json:"messages,omitempty"` + // A dictionary containing metadata associated with the memory. + Metadata interface{} `json:"metadata,omitempty"` + // A Summary object. + Summary Summary `json:"summary,omitempty"` + // A unique identifier for the memory. + UUID string `json:"uuid,omitempty"` + // The timestamp when the memory was created. + CreatedAt string `json:"created_at,omitempty"` + // The token count of the memory. + TokenCount int `json:"token_count,omitempty"` +} + +type APIError struct { + Code int `json:"code,omitempty"` + Message string `json:"message,omitempty"` +} diff --git a/integration/zep/zep.go b/integration/zep/zep.go new file mode 100644 index 0000000..3cae20c --- /dev/null +++ b/integration/zep/zep.go @@ -0,0 +1,138 @@ +package zep + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" +) + +type Client struct { + baseURL string + version string + httpClient *http.Client +} + +func New(baseURL string) *Client { + return &Client{ + baseURL: baseURL, + version: "v1", + httpClient: http.DefaultClient, + } +} + +// GetMemory retrieves memory for a specific session.. +func (c *Client) GetMemory(ctx context.Context, sessionID string) (*Memory, error) { + reqURL := fmt.Sprintf("%s/api/%s/sessions/%s/memory", c.baseURL, c.version, sessionID) + + res, err := c.doRequest(ctx, http.MethodGet, reqURL, nil) + if err != nil { + return nil, err + } + defer res.Body.Close() + + body, err := io.ReadAll(res.Body) + if err != nil { + return nil, err + } + + if res.StatusCode != http.StatusOK { + var apiErr APIError + if err := json.Unmarshal(body, &apiErr); err != nil { + return nil, err + } + + return nil, fmt.Errorf("api error: %d - %s", apiErr.Code, apiErr.Message) + } + + memory := Memory{} + if err := json.Unmarshal(body, &memory); err != nil { + return nil, err + } + + return &memory, nil +} + +// AddMemory adds a new memory to a specific session. +func (c *Client) AddMemory(ctx context.Context, sessionID string, memory *Memory) (string, error) { + reqURL := fmt.Sprintf("%s/api/%s/sessions/%s/memory", c.baseURL, c.version, sessionID) + + res, err := c.doRequest(ctx, http.MethodPost, reqURL, memory) + if err != nil { + return "", err + } + defer res.Body.Close() + + body, err := io.ReadAll(res.Body) + if err != nil { + return "", err + } + + return string(body), nil +} + +// DeleteMemory deletes the memory of a specific session. +func (c *Client) DeleteMemory(ctx context.Context, sessionID string) (string, error) { + reqURL := fmt.Sprintf("%s/api/%s/sessions/%s/memory", c.baseURL, c.version, sessionID) + + res, err := c.doRequest(ctx, http.MethodDelete, reqURL, nil) + if err != nil { + return "", err + } + defer res.Body.Close() + + body, err := io.ReadAll(res.Body) + if err != nil { + return "", err + } + + return string(body), nil +} + +// SearchMessages searches memory of a specific session based on search payload provided. +func (c *Client) SearchMessages(ctx context.Context, sessionID string, payload *SearchPayload) (*SearchResult, error) { + reqURL := fmt.Sprintf("%s/api/%s/sessions/%s/search", c.baseURL, c.version, sessionID) + + res, err := c.doRequest(ctx, http.MethodPost, reqURL, payload) + if err != nil { + return nil, err + } + defer res.Body.Close() + + body, err := io.ReadAll(res.Body) + if err != nil { + return nil, err + } + + result := SearchResult{} + if err := json.Unmarshal(body, &result); err != nil { + return nil, err + } + + return &result, nil +} + +func (c *Client) doRequest(ctx context.Context, method string, url string, payload any) (*http.Response, error) { + var body io.Reader + + if payload != nil { + b, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + body = bytes.NewReader(b) + } + + httpReq, err := http.NewRequestWithContext(ctx, method, url, body) + if err != nil { + return nil, err + } + + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json") + + return c.httpClient.Do(httpReq) +}