postgres

package
v0.25.0 Latest Latest
Warning

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

 Go to latest Published: Mar 25, 2024 License: Apache-2.0 Imports: 38 Imported by: 0

Documentation

Index

Constants

View Source 
const DefaultDocEmbeddingChunkSize = 1000
View Source 
const DefaultDocumentSearchLimit = 20
View Source 
const DefaultEFSearch = 100
View Source 
const DefaultMMRLambda = 0.5
View Source 
const DefaultMMRMultiplier = 2
View Source 
const DefaultMemorySearchLimit = 10
View Source 
const EmbeddingColName = "embedding"
View Source 
const IndexTimeout = 1 * time.Hour
View Source 
const MaxParallelWorkersPerGather = 4
View Source 
const MinRowsForIndex = 10000

MinRowsForIndex is the minimum number of rows required to create an index. The pgvector docs recommend creating the index after a representative sample of data is loaded. This is a guesstimate.

Variables

View Source 
var IndexMutexMap = make(map[string]*sync.Mutex)

IndexMutexMap stores a mutex for each collection.

Functions

func CleanDB 

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

func CreateSchema added in v0.11.0 

func CreateSchema(
	ctx context.Context,
	appState *models.AppState,
	db *bun.DB,
) error

CreateSchema creates the db schema if it does not exist.

func GenerateFixtureData added in v0.11.0 

func GenerateFixtureData(fixtureCount int, outputDir string)

func LoadFixtures added in v0.11.0 

func LoadFixtures(
	ctx context.Context,
	appState *models.AppState,
	db *bun.DB,
	fixturePath string,
) error

func MigrateEmbeddingDims added in v0.17.0 

func MigrateEmbeddingDims(
	ctx context.Context,
	db *bun.DB,
	tableName string,
	dimensions int,
) error

MigrateEmbeddingDims drops the old embedding column and creates a new one with the correct dimensions.

func NewPostgresConn 

func NewPostgresConn(appState *models.AppState) (*bun.DB, error)

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.

func NewPostgresConnForQueue added in v0.17.0 

func NewPostgresConnForQueue(appState *models.AppState) (*sql.DB, error)

NewPostgresConnForQueue creates a new pgx connection to a postgres database using the provided DSN. This connection is intended to be used for queueing tasks.

Types

type CustomRandSource added in v0.12.0 

type CustomRandSource struct {
	rand.Source
}

func (*CustomRandSource) Intn added in v0.12.0 

func (s *CustomRandSource) Intn(n int) int

Intn Override

type DocumentCollectionDAO 

type DocumentCollectionDAO struct {
	models.DocumentCollection
	// contains filtered or unexported fields
}

func NewDocumentCollectionDAO 

func NewDocumentCollectionDAO(
	appState *models.AppState,
	db *bun.DB,
	collection models.DocumentCollection,
) *DocumentCollectionDAO

func (*DocumentCollectionDAO) Create 

func (dc *DocumentCollectionDAO) Create(
	ctx context.Context,
) error

Create inserts a collection into the collections table and creates a table for the collection's documents.

func (*DocumentCollectionDAO) CreateDocuments 

func (dc *DocumentCollectionDAO) CreateDocuments(
	ctx context.Context,
	documents []models.Document,
) ([]uuid.UUID, error)

CreateDocuments inserts the given documents into the given collection.

func (*DocumentCollectionDAO) Delete 

func (dc *DocumentCollectionDAO) Delete(ctx context.Context) error

Delete deletes a collection from the collections table and drops the collection's document table.

func (*DocumentCollectionDAO) DeleteDocumentsByUUID 

func (dc *DocumentCollectionDAO) DeleteDocumentsByUUID(
	ctx context.Context,
	documentUUIDs []uuid.UUID,
) error

DeleteDocumentsByUUID deletes a single document from a collection in the SqlDB, identified by its UUID.

func (*DocumentCollectionDAO) GetAll 

func (dc *DocumentCollectionDAO) GetAll(
	ctx context.Context,
) ([]models.DocumentCollection, error)

GetAll returns a list of all collections from the collections table.

func (*DocumentCollectionDAO) GetByName 

func (dc *DocumentCollectionDAO) GetByName(
	ctx context.Context,
) error

GetByName returns a collection from the collections table by name.

func (*DocumentCollectionDAO) GetCollectionCounts added in v0.12.0 

func (dc *DocumentCollectionDAO) GetCollectionCounts(
	ctx context.Context,
) (models.DocumentCollectionCounts, error)

GetCollectionCounts returns the number of documents and embedded documents

func (*DocumentCollectionDAO) GetDocuments 

func (dc *DocumentCollectionDAO) GetDocuments(
	ctx context.Context,
	limit int,
	uuids []uuid.UUID,
	documentIDs []string,
) ([]models.Document, error)

GetDocuments retrieves documents. If `documents` is non-Nil, it will use the document UUIDs to retrieve these documents. Otherwise, it will retrieve all documents. If limit is greater than 0, it will only retrieve limit many documents.

func (*DocumentCollectionDAO) SearchDocuments 

func (dc *DocumentCollectionDAO) SearchDocuments(ctx context.Context,
	query *models.DocumentSearchPayload,
	limit int,
	pageNumber int,
	pageSize int) (*models.DocumentSearchResultPage, error)

SearchDocuments searches for documents in a collection. Currently pagination is not supported.

func (*DocumentCollectionDAO) Update 

func (dc *DocumentCollectionDAO) Update(
	ctx context.Context,
) error

Update updates a collection in the collections table.

func (*DocumentCollectionDAO) UpdateDocuments 

func (dc *DocumentCollectionDAO) UpdateDocuments(
	ctx context.Context,
	documents []models.Document,
) error

UpdateDocuments updates the document_id, metadata, and embedding columns of the given documents in the given collection. The documents must have non-nil uuids.

**IMPORTANT:** We determine which columns to update based on the fields that are non-zero in the given documents. This means that all documents must have data for the same fields. If a document is missing data for a field, there could be data loss.

type DocumentCollectionSchema 

type DocumentCollectionSchema struct {
	bun.BaseModel             `bun:"table:document_collection,alias:dc" yaml:"-"`
	models.DocumentCollection `                                         yaml:",inline"`
}

DocumentCollectionSchema represents the schema for the DocumentCollectionDAO table.

func (*DocumentCollectionSchema) AfterCreateTable 

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

func (*DocumentCollectionSchema) BeforeAppendModel added in v0.12.0 

func (s *DocumentCollectionSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type DocumentSchemaTemplate 

type DocumentSchemaTemplate struct {
	bun.BaseModel `bun:"table:document,alias:d"`
	models.DocumentBase
}

DocumentSchemaTemplate represents the schema template for Document tables. TextData is manually added when createDocumentTable is run in order to set the correct dimensions. This means the embedding is not returned when querying using bun.

type DocumentStore 

type DocumentStore struct {
	store.BaseDocumentStore[*bun.DB]
	// contains filtered or unexported fields
}

func NewDocumentStore 

func NewDocumentStore(
	ctx context.Context,
	appState *models.AppState,
	client *bun.DB,
) (*DocumentStore, error)

NewDocumentStore returns a new DocumentStore. Use this to correctly initialize the store.

func (*DocumentStore) CreateCollection 

func (ds *DocumentStore) CreateCollection(
	ctx context.Context,
	collection models.DocumentCollection,
) error

func (*DocumentStore) CreateCollectionIndex 

func (ds *DocumentStore) CreateCollectionIndex(
	ctx context.Context,
	collectionName string,
	force bool,
) error

func (*DocumentStore) CreateDocuments 

func (ds *DocumentStore) CreateDocuments(
	ctx context.Context,
	collectionName string,
	documents []models.Document,
) ([]uuid.UUID, error)

func (*DocumentStore) DeleteCollection 

func (ds *DocumentStore) DeleteCollection(
	ctx context.Context,
	collectionName string,
) error

func (*DocumentStore) DeleteDocuments 

func (ds *DocumentStore) DeleteDocuments(
	ctx context.Context,
	collectionName string,
	documentUUID []uuid.UUID,
) error

func (*DocumentStore) GetClient 

func (ds *DocumentStore) GetClient() any

func (*DocumentStore) GetCollection 

func (ds *DocumentStore) GetCollection(
	ctx context.Context,
	collectionName string,
) (models.DocumentCollection, error)

func (*DocumentStore) GetCollectionList 

func (ds *DocumentStore) GetCollectionList(
	ctx context.Context,
) ([]models.DocumentCollection, error)

func (*DocumentStore) GetDocuments 

func (ds *DocumentStore) GetDocuments(
	ctx context.Context,
	collectionName string,
	uuids []uuid.UUID,
	documentIDs []string,
) ([]models.Document, error)

func (*DocumentStore) OnStart 

func (ds *DocumentStore) OnStart(
	_ context.Context,
) error

func (*DocumentStore) SearchCollection 

func (ds *DocumentStore) SearchCollection(
	ctx context.Context,
	query *models.DocumentSearchPayload,
	limit int,
	pageNumber int,
	pageSize int,
) (*models.DocumentSearchResultPage, error)

func (*DocumentStore) Shutdown 

func (ds *DocumentStore) Shutdown(_ context.Context) error

func (*DocumentStore) UpdateCollection 

func (ds *DocumentStore) UpdateCollection(
	ctx context.Context,
	collection models.DocumentCollection,
) error

func (*DocumentStore) UpdateDocuments 

func (ds *DocumentStore) UpdateDocuments(
	ctx context.Context,
	collectionName string,
	documents []models.Document,
) error

type FixtureModel added in v0.11.0 

type FixtureModel[T Row] struct {
	Model string `yaml:"model"`
	Rows  []T    `yaml:"rows"`
}

type Fixtures added in v0.11.0 

type Fixtures[T Row] []FixtureModel[T]

type IndexStatus added in v0.13.0 

type IndexStatus struct {
	Phase       string `bun:"phase"`
	TuplesTotal int    `bun:"tuples_total"`
	TuplesDone  int    `bun:"tuples_done"`
}

type JSONQuery 

type JSONQuery struct {
	JSONPath string       `json:"jsonpath"`
	And      []*JSONQuery `json:"and,omitempty"`
	Or       []*JSONQuery `json:"or,omitempty"`
}

type MemoryDAO added in v0.20.0 

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

MemoryDAO is a data access object for Memory. A Memory is an overlay over Messages and Summaries. It is used to retrieve a set of messages and a summary for a given sessionID, to store a new set of messages from a chat client, and to search for messages and summaries.

func NewMemoryDAO added in v0.20.0 

func NewMemoryDAO(db *bun.DB, appState *models.AppState, sessionID string) (*MemoryDAO, error)

NewMemoryDAO creates a new MemoryDAO.

func (*MemoryDAO) Create added in v0.20.0 

func (m *MemoryDAO) Create(
	ctx context.Context,
	memoryMessages *models.Memory,
	skipNotify bool,
) error

Create stores a Memory for a given sessionID. If the SessionID doesn't exist, a new one is created. If skipNotify is true, the new messages will not be published to the message queue router.

func (*MemoryDAO) Get added in v0.20.0 

func (m *MemoryDAO) Get(
	ctx context.Context,
	lastNMessages int,
) (*models.Memory, error)

Get returns the most recent Summary and a list of messages for a given sessionID. Get returns:

  • the most recent Summary, if one exists
  • the lastNMessages messages, if lastNMessages > 0
  • all messages since the last SummaryPoint, if lastNMessages == 0
  • if no Summary (and no SummaryPoint) exists and lastNMessages == 0, returns all undeleted messages up to the configured message window

func (*MemoryDAO) Search added in v0.20.0 

func (m *MemoryDAO) Search(
	ctx context.Context,
	query *models.MemorySearchPayload,
	limit int,
) ([]models.MemorySearchResult, error)

type MessageDAO added in v0.20.0 

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

func NewMessageDAO added in v0.20.0 

func NewMessageDAO(db *bun.DB, appState *models.AppState, sessionID string) (*MessageDAO, error)

func (*MessageDAO) Create added in v0.20.0 

func (dao *MessageDAO) Create(
	ctx context.Context,
	message *models.Message,
) (*models.Message, error)

Create creates a new message for a session. Create does not create a session if it does not exist.

func (*MessageDAO) CreateEmbeddings added in v0.20.0 

func (dao *MessageDAO) CreateEmbeddings(
	ctx context.Context,
	embeddings []models.TextData,
) error

CreateEmbeddings saves message embeddings for a set of given messages

func (*MessageDAO) CreateMany added in v0.20.0 

func (dao *MessageDAO) CreateMany(
	ctx context.Context,
	messages []models.Message,
) ([]models.Message, error)

CreateMany creates a batch of messages for a session.

func (*MessageDAO) Delete added in v0.20.0 

func (dao *MessageDAO) Delete(ctx context.Context, messageUUID uuid.UUID) error

func (*MessageDAO) Get added in v0.20.0 

func (dao *MessageDAO) Get(ctx context.Context, messageUUID uuid.UUID) (*models.Message, error)

Get retrieves a message by its UUID.

func (*MessageDAO) GetEmbedding added in v0.20.0 

func (dao *MessageDAO) GetEmbedding(ctx context.Context, messageUUID uuid.UUID) (*models.TextData, error)

func (*MessageDAO) GetEmbeddingListBySession added in v0.20.0 

func (dao *MessageDAO) GetEmbeddingListBySession(ctx context.Context) ([]models.TextData, error)

GetEmbeddingListBySession retrieves all message embeddings for a session.

func (*MessageDAO) GetLastN added in v0.20.0 

func (dao *MessageDAO) GetLastN(
	ctx context.Context,
	lastNMessages int,
	beforeUUID uuid.UUID,
) ([]models.Message, error)

GetLastN retrieves the last N messages for a session. If uuid is provided, it will get the last N messages before and including the provided beforeUUID. Results are returned in ascending order of creation

func (*MessageDAO) GetListBySession added in v0.20.0 

func (dao *MessageDAO) GetListBySession(
	ctx context.Context,
	currentPage int,
	pageSize int) (*models.MessageListResponse, error)

GetListBySession retrieves a list of messages for a session. The list is paginated.

func (*MessageDAO) GetListByUUID added in v0.20.0 

func (dao *MessageDAO) GetListByUUID(
	ctx context.Context,
	messageUUIDs []uuid.UUID,
) ([]models.Message, error)

GetListByUUID retrieves a list of messages by their UUIDs. Does not reorder the messages.

func (*MessageDAO) GetSinceLastSummary added in v0.20.0 

func (dao *MessageDAO) GetSinceLastSummary(
	ctx context.Context,
	lastSummary *models.Summary,
	memoryWindow int,
) ([]models.Message, error)

GetSinceLastSummary retrieves messages since the last summary point, limited by the memory window. If there is no last summary point, all messages are returned, limited by the memory window. Results are returned in ascending order of creation

func (*MessageDAO) Update added in v0.20.0 

func (dao *MessageDAO) Update(ctx context.Context,
	message *models.Message,
	includeContent bool,
	isPrivileged bool) error

Update updates a message by its UUID. Metadata is updated via a merge. If includeContent is true, the content and role fields are updated, too.

func (*MessageDAO) UpdateMany added in v0.20.0 

func (dao *MessageDAO) UpdateMany(ctx context.Context,
	messages []models.Message,
	includeContent bool,
	isPrivileged bool) error

UpdateMany updates a batch of messages by their UUIDs. Metadata is updated via a merge.

type MessageStoreSchema 

type MessageStoreSchema struct {
	bun.BaseModel `bun:"table:message,alias:m" yaml:"-"`

	UUID uuid.UUID `bun:",pk,type:uuid,default:gen_random_uuid()"                     yaml:"uuid"`
	// ID is used only for sorting / slicing purposes as we can't sort by CreatedAt for messages created simultaneously
	ID         int64                  `bun:",autoincrement"                                              yaml:"id,omitempty"`
	CreatedAt  time.Time              `bun:"type:timestamptz,notnull,default:current_timestamp"          yaml:"created_at,omitempty"`
	UpdatedAt  time.Time              `bun:"type:timestamptz,nullzero,default:current_timestamp"         yaml:"updated_at,omitempty"`
	DeletedAt  time.Time              `bun:"type:timestamptz,soft_delete,nullzero"                       yaml:"deleted_at,omitempty"`
	SessionID  string                 `bun:",notnull"                                                    yaml:"session_id,omitempty"`
	Role       string                 `bun:",notnull"                                                    yaml:"role,omitempty"`
	Content    string                 `bun:",notnull"                                                    yaml:"content,omitempty"`
	TokenCount int                    `bun:",notnull"                                                    yaml:"token_count,omitempty"`
	Metadata   map[string]interface{} `bun:"type:jsonb,nullzero,json_use_number"                         yaml:"metadata,omitempty"`
	Session    *SessionSchema         `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade" yaml:"-"`
}

func (*MessageStoreSchema) AfterCreateTable 

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

func (*MessageStoreSchema) BeforeAppendModel added in v0.12.0 

func (s *MessageStoreSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type MessageVectorStoreSchema 

type MessageVectorStoreSchema 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     *SessionSchema      `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade"`
	Message     *MessageStoreSchema `bun:"rel:belongs-to,join:message_uuid=uuid,on_delete:cascade"`
}

MessageVectorStoreSchema stores the embeddings for a message.

func (*MessageVectorStoreSchema) AfterCreateTable 

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

func (*MessageVectorStoreSchema) BeforeAppendModel added in v0.12.0 

func (s *MessageVectorStoreSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type PostgresMemoryStore 

type PostgresMemoryStore struct {
	store.BaseMemoryStore[*bun.DB]
	SessionStore *SessionDAO
	// contains filtered or unexported fields
}

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) CreateMessageEmbeddings added in v0.20.0 

func (pms *PostgresMemoryStore) CreateMessageEmbeddings(ctx context.Context,
	sessionID string,
	embeddings []models.TextData,
) error

func (*PostgresMemoryStore) CreateSession added in v0.11.0 

func (pms *PostgresMemoryStore) CreateSession(
	ctx context.Context,
	session *models.CreateSessionRequest,
) (*models.Session, error)

CreateSession creates or updates a Session for a given sessionID.

func (*PostgresMemoryStore) CreateSummary added in v0.20.0 

func (pms *PostgresMemoryStore) CreateSummary(
	ctx context.Context,
	sessionID string,
	summary *models.Summary,
) 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.

func (*PostgresMemoryStore) GetClient 

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

func (*PostgresMemoryStore) GetMemory 

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

GetMemory returns the most recent Summary and a list of messages for a given sessionID. GetMemory returns:

  • the most recent Summary, if one exists
  • the lastNMessages messages, if lastNMessages > 0
  • all messages since the last SummaryPoint, if lastNMessages == 0
  • if no Summary (and no SummaryPoint) exists and lastNMessages == 0, returns all undeleted messages up to the configured message window

func (*PostgresMemoryStore) GetMessageEmbeddings added in v0.17.0 

func (pms *PostgresMemoryStore) GetMessageEmbeddings(ctx context.Context,
	sessionID string,
) ([]models.TextData, error)

func (*PostgresMemoryStore) GetMessageList added in v0.12.0 

func (pms *PostgresMemoryStore) GetMessageList(
	ctx context.Context,
	sessionID string,
	pageNumber int,
	pageSize int,
) (*models.MessageListResponse, error)

GetMessageList retrieves a list of messages for a given sessionID. Paginated by cursor and limit.

func (*PostgresMemoryStore) GetMessagesByUUID added in v0.17.0 

func (pms *PostgresMemoryStore) GetMessagesByUUID(
	ctx context.Context,
	sessionID string,
	uuids []uuid.UUID,
) ([]models.Message, error)

func (*PostgresMemoryStore) GetSession 

func (pms *PostgresMemoryStore) GetSession(
	ctx context.Context,
	sessionID string,
) (*models.Session, error)

GetSession retrieves a Session for a given sessionID.

func (*PostgresMemoryStore) GetSummary 

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

func (*PostgresMemoryStore) GetSummaryByUUID added in v0.17.0 

func (pms *PostgresMemoryStore) GetSummaryByUUID(
	ctx context.Context,
	sessionID string,
	uuid uuid.UUID) (*models.Summary, error)

func (*PostgresMemoryStore) GetSummaryList added in v0.12.0 

func (pms *PostgresMemoryStore) GetSummaryList(
	ctx context.Context,
	sessionID string,
	pageNumber int,
	pageSize int,
) (*models.SummaryListResponse, error)

func (*PostgresMemoryStore) ListSessions added in v0.11.0 

func (pms *PostgresMemoryStore) ListSessions(
	ctx context.Context,
	cursor int64,
	limit int,
) ([]*models.Session, error)

ListSessions returns a list of all Sessions.

func (*PostgresMemoryStore) ListSessionsOrdered added in v0.12.0 

func (pms *PostgresMemoryStore) ListSessionsOrdered(
	ctx context.Context,
	pageNumber int,
	pageSize int,
	orderedBy string,
	asc bool,
) (*models.SessionListResponse, error)

ListSessionsOrdered returns an ordered list of all Sessions, paginated by pageNumber and pageSize. orderedBy is the column to order by. asc is a boolean indicating whether to order ascending or descending.

func (*PostgresMemoryStore) OnStart 

func (pms *PostgresMemoryStore) OnStart(
	ctx context.Context,
) error

func (*PostgresMemoryStore) PurgeDeleted 

func (pms *PostgresMemoryStore) PurgeDeleted(ctx context.Context) error

func (*PostgresMemoryStore) PutMemory 

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

func (*PostgresMemoryStore) PutSummaryEmbedding added in v0.17.0 

func (pms *PostgresMemoryStore) PutSummaryEmbedding(
	ctx context.Context,
	sessionID string,
	embedding *models.TextData,
) error

func (*PostgresMemoryStore) SearchMemory 

func (pms *PostgresMemoryStore) SearchMemory(
	ctx context.Context,
	sessionID string,
	query *models.MemorySearchPayload,
	limit int,
) ([]models.MemorySearchResult, error)

func (*PostgresMemoryStore) UpdateMessages added in v0.20.0 

func (pms *PostgresMemoryStore) UpdateMessages(
	ctx context.Context,
	sessionID string,
	messages []models.Message,
	isPrivileged bool,
	includeContent bool,
) error

func (*PostgresMemoryStore) UpdateSession added in v0.11.0 

func (pms *PostgresMemoryStore) UpdateSession(
	ctx context.Context,
	session *models.UpdateSessionRequest,
) (*models.Session, error)

UpdateSession creates or updates a Session for a given sessionID.

func (*PostgresMemoryStore) UpdateSummary added in v0.20.0 

func (pms *PostgresMemoryStore) UpdateSummary(ctx context.Context,
	sessionID string,
	summary *models.Summary,
	metadataOnly bool,
) error

type Row added in v0.11.0 

type Row interface {
	UserSchema | SessionSchema | DocumentCollectionSchema | MessageStoreSchema
}

type SessionDAO added in v0.11.0 

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

SessionDAO implements the SessionManager interface.

func NewSessionDAO added in v0.11.0 

func NewSessionDAO(db *bun.DB) *SessionDAO

NewSessionDAO is a constructor for the SessionDAO struct. It takes a pointer to a bun.DB instance and returns a pointer to a new SessionDAO instance.

func (*SessionDAO) Create added in v0.11.0 

func (dao *SessionDAO) Create(
	ctx context.Context,
	session *models.CreateSessionRequest,
) (*models.Session, error)

Create creates a new session in the database. It takes a context and a pointer to a CreateSessionRequest struct. It returns a pointer to the created Session struct or an error if the creation fails.

func (*SessionDAO) Delete added in v0.11.0 

func (dao *SessionDAO) Delete(ctx context.Context, sessionID string) error

Delete soft-deletes a session from the database by its sessionID. It also soft-deletes all messages, message embeddings, and summaries associated with the session.

func (*SessionDAO) Get added in v0.11.0 

func (dao *SessionDAO) Get(ctx context.Context, sessionID string) (*models.Session, error)

Get retrieves a session from the database by its sessionID. It takes a context and a session ID string. It returns a pointer to the retrieved Session struct or an error if the retrieval fails.

func (*SessionDAO) ListAll added in v0.11.0 

func (dao *SessionDAO) ListAll(
	ctx context.Context,
	cursor int64,
	limit int,
) ([]*models.Session, error)

ListAll retrieves all sessions from the database. It takes a context, a cursor int64, and a limit int. It returns a slice of pointers to Session structs or an error if the retrieval fails.

func (*SessionDAO) ListAllOrdered added in v0.12.0 

func (dao *SessionDAO) ListAllOrdered(
	ctx context.Context,
	pageNumber int,
	pageSize int,
	orderBy string,
	asc bool,
) (*models.SessionListResponse, error)

func (*SessionDAO) Update added in v0.11.0 

func (dao *SessionDAO) Update(
	ctx context.Context,
	session *models.UpdateSessionRequest,
	isPrivileged bool,
) (*models.Session, error)

Update updates a session in the database. It takes a context, a pointer to a UpdateSessionRequest struct, and a boolean indicating whether the caller is privileged. It returns an error if the update fails. Note: Update will update soft-deleted sessions and undelete them. Messages and message embeddings are not undeleted.

type SessionSchema 

type SessionSchema struct {
	bun.BaseModel `bun:"table:session,alias:s" yaml:"-"`

	UUID      uuid.UUID              `bun:",pk,type:uuid,default:gen_random_uuid()"                     yaml:"uuid,omitempty"`
	ID        int64                  `bun:",autoincrement"                                              yaml:"id,omitempty"` // used as a cursor for pagination
	SessionID string                 `bun:",unique,notnull"                                             yaml:"session_id,omitempty"`
	CreatedAt time.Time              `bun:"type:timestamptz,nullzero,notnull,default:current_timestamp" yaml:"created_at,omitempty"`
	UpdatedAt time.Time              `bun:"type:timestamptz,nullzero,notnull,default:current_timestamp" yaml:"updated_at,omitempty"`
	DeletedAt time.Time              `bun:"type:timestamptz,soft_delete,nullzero"                       yaml:"deleted_at,omitempty"`
	Metadata  map[string]interface{} `bun:"type:jsonb,nullzero,json_use_number"                         yaml:"metadata,omitempty"`
	// UserUUID must be pointer type in order to be nullable
	UserID *string     `bun:","                                                           yaml:"user_id,omitempty"`
	User   *UserSchema `bun:"rel:belongs-to,join:user_id=user_id,on_delete:cascade"       yaml:"-"`
}

func (*SessionSchema) AfterCreateTable 

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

func (*SessionSchema) BeforeAppendModel added in v0.12.0 

func (s *SessionSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type SummaryDAO added in v0.20.0 

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

func NewSummaryDAO added in v0.20.0 

func NewSummaryDAO(db *bun.DB, appState *models.AppState, sessionID string) (*SummaryDAO, error)

NewSummaryDAO creates a new SummaryDAO.

func (*SummaryDAO) Create added in v0.20.0 

func (s *SummaryDAO) Create(
	ctx context.Context,
	summary *models.Summary,
) (*models.Summary, error)

Create stores a new summary for a session. The SummaryPointUUID is the UUID of the most recent message in the session when the summary was created.

func (*SummaryDAO) Get added in v0.20.0 

func (s *SummaryDAO) Get(ctx context.Context) (*models.Summary, error)

Get returns the most recent summary for a session

func (*SummaryDAO) GetByUUID added in v0.20.0 

func (s *SummaryDAO) GetByUUID(
	ctx context.Context,
	uuid uuid.UUID) (*models.Summary, error)

GetByUUID returns a summary by UUID

func (*SummaryDAO) GetEmbeddings added in v0.20.0 

func (s *SummaryDAO) GetEmbeddings(
	ctx context.Context,
) ([]models.TextData, error)

GetEmbeddings all summary embeddings for a session. Note: Does not return the summary content.

func (*SummaryDAO) GetList added in v0.20.0 

func (s *SummaryDAO) GetList(ctx context.Context,
	currentPage int,
	pageSize int,
) (*models.SummaryListResponse, error)

GetList returns a list of summaries for a session

func (*SummaryDAO) PutEmbedding added in v0.20.0 

func (s *SummaryDAO) PutEmbedding(
	ctx context.Context,
	embedding *models.TextData,
) error

PutEmbedding stores a summary embedding

func (*SummaryDAO) Update added in v0.20.0 

func (s *SummaryDAO) Update(
	ctx context.Context,
	summary *models.Summary,
	includeContent bool,
) (*models.Summary, error)

type SummaryStoreSchema 

type SummaryStoreSchema struct {
	bun.BaseModel `bun:"table:summary,alias:su" ,yaml:"-"`

	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,json_use_number"`
	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          *SessionSchema         `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade"`
	Message          *MessageStoreSchema    `bun:"rel:belongs-to,join:summary_point_uuid=uuid,on_delete:cascade"`
}

func (*SummaryStoreSchema) AfterCreateTable 

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

func (*SummaryStoreSchema) BeforeAppendModel added in v0.12.0 

func (s *SummaryStoreSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type SummaryVectorStoreSchema added in v0.17.0 

type SummaryVectorStoreSchema struct {
	bun.BaseModel `bun:"table:summary_embedding,alias:se" yaml:"-"`

	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"`
	SummaryUUID uuid.UUID           `bun:"type:uuid,notnull,unique"`
	Embedding   pgvector.Vector     `bun:"type:vector(1536)"`
	IsEmbedded  bool                `bun:"type:bool,notnull,default:false"`
	Summary     *SummaryStoreSchema `bun:"rel:belongs-to,join:summary_uuid=uuid,on_delete:cascade"`
	Session     *SessionSchema      `bun:"rel:belongs-to,join:session_id=session_id,on_delete:cascade"`
}

func (*SummaryVectorStoreSchema) AfterCreateTable added in v0.17.0 

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

func (*SummaryVectorStoreSchema) BeforeAppendModel added in v0.17.0 

func (s *SummaryVectorStoreSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type UserSchema added in v0.11.0 

type UserSchema struct {
	bun.BaseModel `bun:"table:users,alias:u" yaml:"-"`

	UUID      uuid.UUID              `bun:",pk,type:uuid,default:gen_random_uuid()"             yaml:"uuid,omitempty"`
	ID        int64                  `bun:",autoincrement"                                      yaml:"id,omitempty"` // used as a cursor for pagination
	CreatedAt time.Time              `bun:"type:timestamptz,notnull,default:current_timestamp"  yaml:"created_at,omitempty"`
	UpdatedAt time.Time              `bun:"type:timestamptz,nullzero,default:current_timestamp" yaml:"updated_at,omitempty"`
	DeletedAt time.Time              `bun:"type:timestamptz,soft_delete,nullzero"               yaml:"deleted_at,omitempty"`
	UserID    string                 `bun:",unique,notnull"                                     yaml:"user_id,omitempty"`
	Email     string                 `bun:","                                                   yaml:"email,omitempty"`
	FirstName string                 `bun:","                                                   yaml:"first_name,omitempty"`
	LastName  string                 `bun:","                                                   yaml:"last_name,omitempty"`
	Metadata  map[string]interface{} `bun:"type:jsonb,nullzero,json_use_number"                 yaml:"metadata,omitempty"`
}

func (*UserSchema) AfterCreateTable added in v0.11.0 

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

func (*UserSchema) BeforeAppendModel added in v0.12.0 

func (u *UserSchema) BeforeAppendModel(_ context.Context, query bun.Query) error

type UserStoreDAO added in v0.11.0 

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

func NewUserStoreDAO added in v0.11.0 

func NewUserStoreDAO(db *bun.DB) *UserStoreDAO

func (*UserStoreDAO) Create added in v0.11.0 

func (dao *UserStoreDAO) Create(
	ctx context.Context,
	user *models.CreateUserRequest,
) (*models.User, error)

Create creates a new user.

func (*UserStoreDAO) Delete added in v0.11.0 

func (dao *UserStoreDAO) Delete(ctx context.Context, userID string) error

Delete deletes a user.

func (*UserStoreDAO) Get added in v0.11.0 

func (dao *UserStoreDAO) Get(ctx context.Context, userID string) (*models.User, error)

Get gets a user by UserID.

func (*UserStoreDAO) GetSessions added in v0.11.0 

func (dao *UserStoreDAO) GetSessions(
	ctx context.Context,
	userID string,
) ([]*models.Session, error)

GetSessions gets all sessions for a user.

func (*UserStoreDAO) ListAll added in v0.11.0 

func (dao *UserStoreDAO) ListAll(
	ctx context.Context,
	cursor int64,
	limit int,
) ([]*models.User, error)

ListAll lists all users. The cursor is used to paginate results.

func (*UserStoreDAO) ListAllOrdered added in v0.12.0 

func (dao *UserStoreDAO) ListAllOrdered(
	ctx context.Context,
	pageNumber int,
	pageSize int,
	orderBy string,
	asc bool,
) (*models.UserListResponse, error)

func (*UserStoreDAO) Update added in v0.11.0 

func (dao *UserStoreDAO) Update(
	ctx context.Context,
	user *models.UpdateUserRequest,
	isPrivileged bool,
) (*models.User, error)

Update updates a user.

type VectorColIndex 

type VectorColIndex struct {
	Collection models.DocumentCollection
	ColName    string
	RowCount   int
	ListCount  int
	ProbeCount int
	// contains filtered or unexported fields
}

func NewVectorColIndex 

func NewVectorColIndex(
	ctx context.Context,
	appState *models.AppState,
	collection models.DocumentCollection,
) (*VectorColIndex, error)

func (*VectorColIndex) CalculateListCount 

func (vci *VectorColIndex) CalculateListCount() error

CalculateListCount calculates the number of lists to use for the index.

func (*VectorColIndex) CalculateProbes 

func (vci *VectorColIndex) CalculateProbes() error

func (*VectorColIndex) CountRows 

func (vci *VectorColIndex) CountRows(ctx context.Context) error

func (*VectorColIndex) CreateIndex 

func (vci *VectorColIndex) CreateIndex(_ context.Context, force bool) error

Source Files

View all Source files

Directories

Path Synopsis

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.