memorystore

package
v0.4.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: May 10, 2023 License: Apache-2.0 Imports: 20 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func CleanDB 

func CleanDB(t *testing.T, db *bun.DB)

func NewPostgresConn 

func NewPostgresConn(dsn string) *bun.DB

NewPostgresConn creates a new bun.DB connection to a postgres database using the provided DSN. The connection is configured to pool connections based on the number of PROCs available.

Types

type BaseMemoryStore 

type BaseMemoryStore[T any] struct {
	Client T
	// contains filtered or unexported fields
}

BaseMemoryStore is the base implementation of a MemoryStore. Client is the underlying datastore client, such as a database connection. The extractorObservers slice is used to store all registered Extractors.

func (*BaseMemoryStore[T]) Attach 

func (s *BaseMemoryStore[T]) Attach(observer models.Extractor)

Attach registers an Extractor to the MemoryStore

func (*BaseMemoryStore[T]) NotifyExtractors 

func (s *BaseMemoryStore[T]) NotifyExtractors(
	ctx context.Context,
	appState *models.AppState,
	eventData *models.MessageEvent,
)

NotifyExtractors notifies all registered Extractors of a new MessageEvent

type PgMessageStore 

type PgMessageStore struct {
	bun.BaseModel `bun:"table:message,alias:m"`

	// TODO: replace UUIDs with sortable ULIDs or UUIDv7s to avoid having to have both a UUID and an ID.
	// see https://blog.daveallie.com/ulid-primary-keys
	UUID uuid.UUID `bun:",pk,type:uuid,default:gen_random_uuid()"`
	// ID is used only for sorting / slicing purposes as we can't sort by CreatedAt for messages created simultaneously
	ID         int64                  `bun:",autoincrement"`
	CreatedAt  time.Time              `bun:"type:timestamptz,notnull,default:current_timestamp"`
	UpdatedAt  time.Time              `bun:"type:timestamptz,nullzero,default:current_timestamp"`
	DeletedAt  time.Time              `bun:"type:timestamptz,soft_delete,nullzero"`
	SessionID  string                 `bun:",notnull"`
	Role       string                 `bun:",notnull"`
	Content    string                 `bun:",notnull"`
	TokenCount int                    `bun:",notnull"`
	Metadata   map[string]interface{} `bun:"type:jsonb,nullzero"`
	Session    *PgSession             `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade"`
}

func (*PgMessageStore) AfterCreateTable 

func (*PgMessageStore) AfterCreateTable(
	ctx context.Context,
	query *bun.CreateTableQuery,
) error

func (*PgMessageStore) BeforeCreateTable 

func (s *PgMessageStore) BeforeCreateTable(
	_ context.Context,
	_ *bun.CreateTableQuery,
) error

type PgMessageVectorStore 

type PgMessageVectorStore struct {
	bun.BaseModel `bun:"table:message_embedding,alias:me"`

	UUID        uuid.UUID       `bun:",pk,type:uuid,default:gen_random_uuid()"`
	CreatedAt   time.Time       `bun:"type:timestamptz,notnull,default:current_timestamp"`
	UpdatedAt   time.Time       `bun:"type:timestamptz,nullzero,default:current_timestamp"`
	DeletedAt   time.Time       `bun:"type:timestamptz,soft_delete,nullzero"`
	SessionID   string          `bun:",notnull"`
	MessageUUID uuid.UUID       `bun:"type:uuid,notnull,unique"`
	Embedding   pgvector.Vector `bun:"type:vector(1536)"`
	IsEmbedded  bool            `bun:"type:bool,notnull,default:false"`
	Session     *PgSession      `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade"`
	Message     *PgMessageStore `bun:"rel:belongs-to,join:message_uuid=uuid,on_delete:cascade"`
}

PgMessageVectorStore stores the embeddings for a message. TODO: Vector dims from config

func (*PgMessageVectorStore) AfterCreateTable 

func (*PgMessageVectorStore) AfterCreateTable(
	ctx context.Context,
	query *bun.CreateTableQuery,
) error

func (*PgMessageVectorStore) BeforeCreateTable 

func (s *PgMessageVectorStore) BeforeCreateTable(
	_ context.Context,
	_ *bun.CreateTableQuery,
) error

type PgSession 

type PgSession struct {
	bun.BaseModel `bun:"table:session,alias:s"`

	UUID      uuid.UUID              `bun:",pk,type:uuid,default:gen_random_uuid()"`
	SessionID string                 `bun:",unique,notnull"`
	CreatedAt time.Time              `bun:"type:timestamptz,nullzero,notnull,default:current_timestamp"`
	UpdatedAt time.Time              `bun:"type:timestamptz,nullzero,notnull,default:current_timestamp"`
	DeletedAt time.Time              `bun:"type:timestamptz,soft_delete,nullzero"`
	Metadata  map[string]interface{} `bun:"type:jsonb,nullzero"`
}

func (*PgSession) AfterCreateTable 

func (*PgSession) AfterCreateTable(
	ctx context.Context,
	query *bun.CreateTableQuery,
) error

func (*PgSession) BeforeCreateTable 

func (s *PgSession) BeforeCreateTable(
	_ context.Context,
	_ *bun.CreateTableQuery,
) error

BeforeCreateTable is a dummy method to ensure uniform interface across all table models - used in table creation iterator

type PgSummaryStore 

type PgSummaryStore struct {
	bun.BaseModel `bun:"table:summary,alias:su"`

	UUID             uuid.UUID              `bun:",pk,type:uuid,default:gen_random_uuid()"`
	CreatedAt        time.Time              `bun:"type:timestamptz,notnull,default:current_timestamp"`
	UpdatedAt        time.Time              `bun:"type:timestamptz,nullzero,default:current_timestamp"`
	DeletedAt        time.Time              `bun:"type:timestamptz,soft_delete,nullzero"`
	SessionID        string                 `bun:",notnull"`
	Content          string                 `bun:",nullzero"` // allow null as we might want to use Metadata without a summary
	Metadata         map[string]interface{} `bun:"type:jsonb,nullzero"`
	TokenCount       int                    `bun:",notnull"`
	SummaryPointUUID uuid.UUID              `bun:"type:uuid,notnull,unique"` // the UUID of the most recent message that was used to create the summary
	Session          *PgSession             `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade"`
	Message          *PgMessageStore        `bun:"rel:belongs-to,join:summary_point_uuid=uuid,on_delete:cascade"`
}

func (*PgSummaryStore) AfterCreateTable 

func (*PgSummaryStore) AfterCreateTable(
	ctx context.Context,
	query *bun.CreateTableQuery,
) error

func (*PgSummaryStore) BeforeCreateTable 

func (s *PgSummaryStore) BeforeCreateTable(
	_ context.Context,
	_ *bun.CreateTableQuery,
) error

type PostgresMemoryStore 

type PostgresMemoryStore struct {
	BaseMemoryStore[*bun.DB]
}

func NewPostgresMemoryStore 

func NewPostgresMemoryStore(
	appState *models.AppState,
	client *bun.DB,
) (*PostgresMemoryStore, error)

NewPostgresMemoryStore returns a new PostgresMemoryStore. Use this to correctly initialize the store.

func (*PostgresMemoryStore) Close 

func (pms *PostgresMemoryStore) Close() error

func (*PostgresMemoryStore) DeleteSession 

func (pms *PostgresMemoryStore) DeleteSession(ctx context.Context, sessionID string) error

DeleteSession deletes a session from the memory store. This is a soft delete. TODO: A hard delete will be implemented as an out-of-band process or left to the implementer.

func (*PostgresMemoryStore) GetClient 

func (pms *PostgresMemoryStore) GetClient() *bun.DB

func (*PostgresMemoryStore) GetMemory 

func (pms *PostgresMemoryStore) GetMemory(
	ctx context.Context,
	appState *models.AppState,
	sessionID string,
	lastNMessages int,
) (*models.Memory, error)

GetMemory returns the memory for a given sessionID.

func (*PostgresMemoryStore) GetMessageVectors 

func (pms *PostgresMemoryStore) GetMessageVectors(ctx context.Context,
	_ *models.AppState,
	sessionID string,
	isEmbedded bool,
) ([]models.Embeddings, error)

func (*PostgresMemoryStore) GetSummary 

func (pms *PostgresMemoryStore) GetSummary(
	ctx context.Context,
	_ *models.AppState,
	sessionID string,
) (*models.Summary, error)

func (*PostgresMemoryStore) OnStart 

func (pms *PostgresMemoryStore) OnStart(
	_ context.Context,
	_ *models.AppState,
) error

func (*PostgresMemoryStore) PutMemory 

func (pms *PostgresMemoryStore) PutMemory(
	ctx context.Context,
	appState *models.AppState,
	sessionID string,
	memoryMessages *models.Memory,
) error

func (*PostgresMemoryStore) PutMessageVectors 

func (pms *PostgresMemoryStore) PutMessageVectors(ctx context.Context,
	_ *models.AppState,
	sessionID string,
	embeddings []models.Embeddings,
	isEmbedded bool,
) error

func (*PostgresMemoryStore) PutSummary 

func (pms *PostgresMemoryStore) PutSummary(
	ctx context.Context,
	_ *models.AppState,
	sessionID string,
	summary *models.Summary,
) error

func (*PostgresMemoryStore) SearchMemory 

func (pms *PostgresMemoryStore) SearchMemory(
	ctx context.Context,
	appState *models.AppState,
	sessionID string,
	query *models.SearchPayload,
	limit int,
) ([]models.SearchResult, error)

type StorageError 

type StorageError struct {
	// contains filtered or unexported fields
}

func NewStorageError 

func NewStorageError(message string, originalError error) *StorageError

func (*StorageError) Error 

func (e *StorageError) Error() string

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.