models

type AdvisoryLockError struct {
	Err error
}

type AppState struct {
	LLMClient     ZepLLM
	MemoryStore   MemoryStore[any]
	DocumentStore DocumentStore[any]
	UserStore     UserStore
	TaskRouter    TaskRouter
	TaskPublisher TaskPublisher
	Config        *config.Config
}

type BadRequestError struct {
	Message string
}

type CreateDocumentCollectionRequest struct {
	Name                string                 `json:"name"                 validate:"required,alphanum,min=3,max=40"`
	Description         string                 `json:"description"          validate:"omitempty,max=1000"`
	Metadata            map[string]interface{} `json:"metadata,omitempty"`
	EmbeddingDimensions int                    `json:"embedding_dimensions" validate:"required,numeric,min=8,max=2000"`
	// these needs to be pointers so that we can distinguish between false and unset when validating
	IsAutoEmbedded *bool `json:"is_auto_embedded"     validate:"required,boolean"`
}

type CreateDocumentRequest struct {
	DocumentID string                 `json:"document_id,omitempty" validate:"omitempty,printascii,max=100"`
	Content    string                 `json:"content,omitempty"     validate:"required_without=Embedding,omitempty"`
	Metadata   map[string]interface{} `json:"metadata,omitempty"`
	Embedding  []float32              `json:"embedding,omitempty"   validate:"required_without=Content,omitempty"`
}

type CreateSessionRequest struct {
	SessionID string `json:"session_id"`
	// Must be a pointer to allow for null values
	UserID   *string                `json:"user_id"`
	Metadata map[string]interface{} `json:"metadata"`
}

type CreateUserRequest struct {
	UserID    string                 `json:"user_id"`
	Email     string                 `json:"email"`
	FirstName string                 `json:"first_name"`
	LastName  string                 `json:"last_name"`
	Metadata  map[string]interface{} `json:"metadata"`
}

type DistanceFunction string

type DocEmbeddingTask struct {
	UUID uuid.UUID `json:"uuid"`
}

type DocEmbeddingUpdate struct {
	UUID           uuid.UUID `json:"uuid"`
	CollectionName string    `json:"collection_name"`
	ProcessedAt    time.Time `json:"time"`
	Embedding      []float32 `json:"embedding,omitempty" bun:"type:vector,nullzero"`
}

type Document struct {
	DocumentBase
	Embedding []float32 `bun:"type:vector,nullzero" json:"embedding,omitempty"`
}

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

type DocumentCollection struct {
	UUID                      uuid.UUID              `bun:",pk,type:uuid,default:gen_random_uuid()"                     yaml:"uuid"`
	CreatedAt                 time.Time              `bun:"type:timestamptz,nullzero,notnull,default:current_timestamp" yaml:"created_at"`
	UpdatedAt                 time.Time              `bun:"type:timestamptz,nullzero,default:current_timestamp"         yaml:"updated_at"`
	Name                      string                 `bun:",notnull,unique"                                             yaml:"name"`
	Description               string                 `bun:",notnull"                                                    yaml:"description"`
	Metadata                  map[string]interface{} `bun:"type:jsonb,nullzero,json_use_number"                         yaml:"metadata"`
	TableName                 string                 `bun:",notnull"                                                    yaml:"table_name"`
	EmbeddingModelName        string                 `bun:",notnull"                                                    yaml:"embedding_model_name"`
	EmbeddingDimensions       int                    `bun:",notnull"                                                    yaml:"embedding_dimensions"`
	IsAutoEmbedded            bool                   `bun:",notnull"                                                    yaml:"is_auto_embedded"`  // Is the collection automatically embedded by Zep?
	DistanceFunction          DistanceFunction       `bun:",notnull"                                                    yaml:"distance_function"` // Distance function to use for index
	IsNormalized              bool                   `bun:",notnull"                                                    yaml:"is_normalized"`     // Are the embeddings normalized?
	IsIndexed                 bool                   `bun:",notnull"                                                    yaml:"is_indexed"`        // Has an index been created on the collection table?
	IndexType                 IndexType              `bun:",notnull"                                                    yaml:"index_type"`        // Type of index to use
	ListCount                 int                    `bun:",notnull"                                                    yaml:"list_count"`        // Number of lists in the collection index
	ProbeCount                int                    `bun:",notnull"                                                    yaml:"probe_count"`       // Number of probes to use when searching the index
	*DocumentCollectionCounts ` yaml:"document_collection_counts,inline"`
}

type DocumentCollectionCounts struct {
	DocumentCount         int `bun:"document_count"          json:"document_count"          yaml:"document_count,omitempty"`          // Number of documents in the collection
	DocumentEmbeddedCount int `bun:"document_embedded_count" json:"document_embedded_count" yaml:"document_embedded_count,omitempty"` // Number of documents with embeddings
}

type DocumentCollectionResponse struct {
	UUID                uuid.UUID              `json:"uuid"`
	CreatedAt           time.Time              `json:"created_at"`
	UpdatedAt           time.Time              `json:"updated_at"`
	Name                string                 `json:"name"`
	Description         string                 `json:"description"`
	Metadata            map[string]interface{} `json:"metadata,omitempty"`
	EmbeddingModelName  string                 `json:"embedding_model_name,omitempty"`
	EmbeddingDimensions int                    `json:"embedding_dimensions"`
	IsAutoEmbedded      bool                   `json:"is_auto_embedded"`
	IsNormalized        bool                   `json:"is_normalized"`
	IsIndexed           bool                   `json:"is_indexed"`
	*DocumentCollectionCounts
}

type DocumentResponse struct {
	UUID       uuid.UUID              `json:"uuid"`
	CreatedAt  time.Time              `json:"created_at"`
	UpdatedAt  time.Time              `json:"updated_at"`
	DocumentID string                 `json:"document_id"`
	Content    string                 `json:"content"`
	Metadata   map[string]interface{} `json:"metadata,omitempty"`
	Embedding  []float32              `json:"embedding"`
	IsEmbedded bool                   `json:"is_embedded"`
}

type DocumentSearchPayload struct {
	CollectionName string                 `json:"collection_name"`
	Text           string                 `json:"text,omitempty"`
	Embedding      []float32              `json:"embedding,omitempty"`
	Metadata       map[string]interface{} `json:"metadata,omitempty"`
	SearchType     SearchType             `json:"search_type"`
	MMRLambda      float32                `json:"mmr_lambda,omitempty"`
}

type DocumentSearchResult struct {
	*DocumentResponse
	Score float64 `json:"score"`
}

type DocumentSearchResultPage struct {
	Results     []DocumentSearchResult `json:"results"`
	QueryVector []float32              `json:"query_vector"`
	ResultCount int                    `json:"result_count"`
	TotalPages  int                    `json:"total_pages"`
	CurrentPage int                    `json:"current_page"`
}

type DocumentStore[T any] interface {
	// CreateCollection creates a new DocumentCollection.
	// If a collection with the same name already exists, it will be overwritten.
	CreateCollection(
		ctx context.Context,
		collection DocumentCollection,
	) error
	UpdateCollection(
		ctx context.Context,
		collection DocumentCollection,
	) error
	// GetCollection retrieves a DocumentCollection by name.
	GetCollection(
		ctx context.Context,
		collectionName string,
	) (DocumentCollection, error)
	// GetCollectionList retrieves the list of DocumentCollection.
	GetCollectionList(
		ctx context.Context,
	) ([]DocumentCollection, error)
	// DeleteCollection deletes a DocumentCollection by name.
	DeleteCollection(
		ctx context.Context,
		collectionName string,
	) error
	// CreateDocuments creates a batch of Documents.
	CreateDocuments(
		ctx context.Context,
		collectionName string,
		documents []Document,
	) ([]uuid.UUID, error)
	// UpdateDocuments updates a batch of Documents.
	// The provided Document UUIDs must match existing documents.
	UpdateDocuments(
		ctx context.Context,
		collectionName string,
		documents []Document,
	) error
	// GetDocuments retrieves a Document by UUID.
	GetDocuments(
		ctx context.Context,
		collectionName string,
		uuids []uuid.UUID,
		DocumentID []string,
	) ([]Document, error)
	// DeleteDocuments deletes a Document by UUID.
	DeleteDocuments(
		ctx context.Context,
		collectionName string,
		documentUUIDs []uuid.UUID,
	) error
	// SearchCollection retrieves a collection of DocumentSearchResultPage based on the provided search query.
	// It accepts an optional limit for the total number of results, as well as parameters for pagination: pageNumber and pageSize.
	// Parameters:
	// - limit: Defines the maximum number of results returned. If it's 0, all the results will be returned.
	// - pageNumber: Specifies the current page number in the pagination scheme.
	// - pageSize: Determines the number of results per page. If it's -1, all results are returned on a single page.
	// The mmr parameter is used to enable/disable the MMR algorithm for search results.
	// The function will return the results in pages as determined by pageSize.
	SearchCollection(
		ctx context.Context,
		query *DocumentSearchPayload,
		limit int,
		pageNumber int,
		pageSize int,
	) (*DocumentSearchResultPage, error)
	// CreateCollectionIndex creates an index on the collection. Manually calling this function will drop and
	// recreate the index, if it exists.
	// force: If true, the index will be created even if there are too few documents in the collection.
	CreateCollectionIndex(ctx context.Context, collectionName string, force bool) error
	// OnStart is called when the application starts. This is a good place to initialize any resources or configs that
	// are required by the MemoryStore implementation.
	OnStart(ctx context.Context) error
	// Shutdown is called when the application is shutting down. This is a good place to clean up any resources or configs
	Shutdown(ctx context.Context) error
	// GetClient returns the underlying storage client
	GetClient() any
}

type EmbeddingModel struct {
	Service      string `json:"service"`
	Dimensions   int    `json:"dimensions"`
	IsNormalized bool   `json:"normalized"`
}

type Entity struct {
	Name    string        `json:"name"`
	Label   string        `json:"label"`
	Matches []EntityMatch `json:"matches"`
}

type EntityMatch struct {
	Start int    `json:"start"`
	End   int    `json:"end"`
	Text  string `json:"text"`
}

type EntityRequest struct {
	Texts []EntityRequestRecord `json:"texts"`
}

type EntityRequestRecord struct {
	UUID     string `json:"uuid"`
	Text     string `json:"text"`
	Language string `json:"language"`
}

type EntityResponse struct {
	Texts []EntityResponseRecord `json:"texts"`
}

type EntityResponseRecord struct {
	UUID     string   `json:"uuid"`
	Entities []Entity `json:"entities"`
}

type GetDocumentListRequest struct {
	UUIDs       []uuid.UUID `json:"uuids"        validate:"required_without=DocumentIDs"`
	DocumentIDs []string    `json:"document_ids" validate:"required_without=UUIDs"`
}

type GetDocumentRequest struct {
	UUID       uuid.UUID `json:"uuid"        validate:"required_without=DocumentID,uuid,omitempty"`
	DocumentID string    `json:"document_id" validate:"required_without=UUID,alphanum,max=40,omitempty"`
}

type IndexType string

type Intent struct {
	UUID      uuid.UUID `json:"uuid,omitempty"`
	Name      string    `json:"name,omitempty"`
	Documents []string  `json:"documents"`
}

type IntentCollection struct {
	UUID    uuid.UUID `json:"uuid,omitempty"`
	Name    string    `json:"name,omitempty"`
	Intents []Intent  `json:"intents"`
}

type IntentPromptTemplateData struct {
	Input string
}

type IntentResponse struct {
	Intent string `json:"intent"`
}

type Memory struct {
	Messages []Message              `json:"messages"`
	Summary  *Summary               `json:"summary,omitempty"`
	Metadata map[string]interface{} `json:"metadata,omitempty"`
}

type MemorySearchPayload struct {
	Text        string                 `json:"text"`
	Metadata    map[string]interface{} `json:"metadata,omitempty"`
	SearchScope SearchScope            `json:"search_scope,omitempty"`
	SearchType  SearchType             `json:"search_type,omitempty"`
	MMRLambda   float32                `json:"mmr_lambda,omitempty"`
}

type MemorySearchResult struct {
	Message   *Message               `json:"message"`
	Summary   *Summary               `json:"summary"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`
	Dist      float64                `json:"dist"`
	Embedding []float32              `json:"embedding"`
}

type MemoryStore[T any] interface {
	MemoryStorer
	MessageStorer
	SessionStorer
	SummaryStorer
	// PurgeDeleted hard deletes all deleted data in the MemoryStore.
	PurgeDeleted(ctx context.Context) error
	// Close is called when the application is shutting down. This is a good place to clean up any resources used by
	// the MemoryStore implementation.
	Close() error
}

type MemoryStorer interface {
	// 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
	GetMemory(ctx context.Context,
		sessionID string,
		lastNMessages int) (*Memory, error)
	// PutMemory stores a Memory for a given sessionID. If the SessionID doesn't exist, a new one is created.
	PutMemory(ctx context.Context,
		sessionID string,
		memoryMessages *Memory,
		skipNotify bool) error // skipNotify is used to prevent loops when calling NotifyExtractors.
	// SearchMemory retrieves a collection of SearchResults for a given sessionID and query. Currently, The
	// MemorySearchResult structure can include both Messages and Summaries.
	SearchMemory(
		ctx context.Context,
		sessionID string,
		query *MemorySearchPayload,
		limit int) ([]MemorySearchResult, error)
}

type Message struct {
	UUID       uuid.UUID              `json:"uuid"`
	CreatedAt  time.Time              `json:"created_at"`
	UpdatedAt  time.Time              `json:"updated_at"`
	Role       string                 `json:"role"`
	Content    string                 `json:"content"`
	Metadata   map[string]interface{} `json:"metadata,omitempty"`
	TokenCount int                    `json:"token_count"`
}

type MessageListResponse struct {
	Messages   []Message `json:"messages"`
	TotalCount int       `json:"total_count"`
	RowCount   int       `json:"row_count"`
}

type MessageStorer interface {
	// UpdateMessages updates a collection of Messages for a given sessionID. If includeContent is true, the
	// role and content fields are updated, too. If isPrivileged is true, the `system` key may be updated.
	UpdateMessages(
		ctx context.Context,
		sessionID string,
		messages []Message,
		isPrivileged bool,
		includeContent bool) error
	// GetMessagesByUUID retrieves messages for a given sessionID and UUID slice.
	GetMessagesByUUID(
		ctx context.Context,
		sessionID string,
		uuids []uuid.UUID,
	) ([]Message, error)
	// GetMessageList retrieves a list of messages for a given sessionID. Paginated by cursor and limit.
	GetMessageList(ctx context.Context,
		sessionID string,
		pageNumber int,
		pageSize int,
	) (*MessageListResponse, error)
	// CreateMessageEmbeddings stores a collection of TextData for a given sessionID.
	CreateMessageEmbeddings(ctx context.Context,
		sessionID string,
		embeddings []TextData) error
	// GetMessageEmbeddings retrieves a collection of TextData for a given sessionID.
	GetMessageEmbeddings(ctx context.Context,
		sessionID string) ([]TextData, error)
}

type MessageSummaryTask struct {
	UUID uuid.UUID `json:"uuid"`
}

type MessageTask struct {
	UUID uuid.UUID `json:"uuid"`
}

type NotFoundError struct {
	Resource string
}

type SearchDocumentResult struct {
	*Document
	Score float64 `json:"score" bun:"score"`
}

type SearchScope string

type SearchType string

type Session struct {
	UUID      uuid.UUID              `json:"uuid"`
	ID        int64                  `json:"id"`
	CreatedAt time.Time              `json:"created_at"`
	UpdatedAt time.Time              `json:"updated_at"`
	DeletedAt *time.Time             `json:"deleted_at"`
	SessionID string                 `json:"session_id"`
	Metadata  map[string]interface{} `json:"metadata"`
	// Must be a pointer to allow for null values
	UserID *string `json:"user_id"`
}

type SessionListResponse struct {
	Sessions   []*Session `json:"sessions"`
	TotalCount int        `json:"total_count"`
	RowCount   int        `json:"response_count"`
}

type SessionManager interface {
	Create(ctx context.Context, session *CreateSessionRequest) (*Session, error)
	Get(ctx context.Context, sessionID string) (*Session, error)
	Update(ctx context.Context, session *UpdateSessionRequest, isPrivileged bool) (*Session, error)
	Delete(ctx context.Context, sessionID string) error
	ListAll(ctx context.Context, cursor int64, limit int) ([]*Session, error)
}

type SessionStorer interface {
	// CreateSession creates a new Session for a given sessionID.
	CreateSession(
		ctx context.Context,
		session *CreateSessionRequest,
	) (*Session, error)
	// GetSession retrieves a Session for a given sessionID.
	GetSession(
		ctx context.Context,
		sessionID string,
	) (*Session, error)
	// UpdateSession updates a Session for a given sessionID. Omly the metadata is updated.
	UpdateSession(
		ctx context.Context,
		session *UpdateSessionRequest,
	) (*Session, error)
	// DeleteSession deletes all records for a given sessionID. This is a soft delete. Related Messages
	// and MessageEmbeddings are also soft deleted.
	DeleteSession(ctx context.Context, sessionID string) error
	// ListSessions returns a list of all Sessions, paginated by cursor and limit.
	ListSessions(
		ctx context.Context,
		cursor int64,
		limit int,
	) ([]*Session, error)
	// ListSessionsOrdered returns an ordered list of all Sessions, paginated by pageNumber and pageSize, and
	// the total count of all sessions.
	// orderedBy is the column to order by. asc is a boolean indicating whether to order ascending or descending.
	ListSessionsOrdered(
		ctx context.Context,
		pageNumber int,
		pageSize int,
		orderedBy string,
		asc bool,
	) (*SessionListResponse, error)
}

type Summary struct {
	UUID             uuid.UUID              `json:"uuid"`
	CreatedAt        time.Time              `json:"created_at"`
	Content          string                 `json:"content"`
	SummaryPointUUID uuid.UUID              `json:"recent_message_uuid"` // The most recent message UUID that was used to generate this summary
	Metadata         map[string]interface{} `json:"metadata,omitempty"`
	TokenCount       int                    `json:"token_count"`
}

type SummaryListResponse struct {
	Summaries  []Summary `json:"summaries"`
	TotalCount int       `json:"total_count"`
	RowCount   int       `json:"row_count"`
}

type SummaryStorer interface {
	// GetSummary retrieves the most recent Summary for a given sessionID. The Summary return includes the UUID of the
	// SummaryPoint, which the most recent Message in the collection of messages that was used to generate the Summary.
	GetSummary(ctx context.Context,
		sessionID string) (*Summary, error)
	GetSummaryByUUID(ctx context.Context,
		sessionID string,
		uuid uuid.UUID) (*Summary, error)
	// GetSummaryList retrieves a list of Summary for a given sessionID. Paginated by cursor and limit.
	GetSummaryList(ctx context.Context,
		sessionID string,
		pageNumber int,
		pageSize int,
	) (*SummaryListResponse, error)
	// CreateSummary stores a new Summary for a given sessionID.
	CreateSummary(ctx context.Context,
		sessionID string,
		summary *Summary) error
	// UpdateSummary updates the metadata for a given Summary. The Summary UUID must be set.
	UpdateSummary(ctx context.Context,
		sessionID string,
		summary *Summary,
		includeContent bool,
	) error
	// PutSummaryEmbedding stores a TextData for a given sessionID and Summary UUID.
	PutSummaryEmbedding(ctx context.Context,
		sessionID string,
		embedding *TextData) error
}

type Task interface {
	Execute(ctx context.Context, event *message.Message) error
	HandleError(err error)
}

type TaskPublisher interface {
	Publish(taskType TaskTopic, metadata map[string]string, payload any) error
	PublishMessage(metadata map[string]string, payload []MessageTask) error
	Close() error
}

type TaskRouter interface {
	Run(ctx context.Context) error
	AddTask(ctx context.Context, name string, taskType TaskTopic, task Task)
	RunHandlers(ctx context.Context) error
	IsRunning() bool
	Close() error
}

type TaskTopic string

type TextData struct {
	TextUUID  uuid.UUID `json:"uuid,omitempty"` // MemoryStore's unique ID associated with this text.
	Text      string    `json:"text"`
	Embedding []float32 `json:"embedding,omitempty"`
	Language  string    `json:"language"`
}

type TextEmbeddingCollection struct {
	UUID       uuid.UUID  `json:"uuid,omitempty"`
	Name       string     `json:"name,omitempty"`
	Embeddings []TextData `json:"documents"`
}

type UpdateDocumentCollectionRequest struct {
	Description string                 `json:"description"        validate:"max=1000"`
	Metadata    map[string]interface{} `json:"metadata,omitempty"`
}

type UpdateDocumentListRequest struct {
	UUID uuid.UUID `json:"uuid" validate:"required"`
	UpdateDocumentRequest
}

type UpdateDocumentRequest struct {
	DocumentID string                 `json:"document_id"        validate:"printascii,max=40,omitempty"`
	Metadata   map[string]interface{} `json:"metadata,omitempty" validate:"omitempty"`
}

type UpdateSessionRequest struct {
	SessionID string                 `json:"session_id"`
	Metadata  map[string]interface{} `json:"metadata"`
}

type UpdateUserRequest struct {
	UUID      uuid.UUID              `json:"uuid"`
	UserID    string                 `json:"user_id"`
	Email     string                 `json:"email"`
	FirstName string                 `json:"first_name"`
	LastName  string                 `json:"last_name"`
	Metadata  map[string]interface{} `json:"metadata"`
}

type User struct {
	UUID      uuid.UUID              `json:"uuid"`
	ID        int64                  `json:"id"`
	CreatedAt time.Time              `json:"created_at"`
	UpdatedAt time.Time              `json:"updated_at"`
	DeletedAt *time.Time             `json:"deleted_at"`
	UserID    string                 `json:"user_id"`
	Email     string                 `json:"email,omitempty"`
	FirstName string                 `json:"first_name,omitempty"`
	LastName  string                 `json:"last_name,omitempty"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`
}

type UserListResponse struct {
	Users      []*User `json:"users"`
	TotalCount int     `json:"total_count"`
	RowCount   int     `json:"row_count"`
}

type UserStore interface {
	Create(ctx context.Context, user *CreateUserRequest) (*User, error)
	Get(ctx context.Context, userID string) (*User, error)
	Update(ctx context.Context, user *UpdateUserRequest, isPrivileged bool) (*User, error)
	Delete(ctx context.Context, userID string) error
	GetSessions(ctx context.Context, userID string) ([]*Session, error)
	ListAll(ctx context.Context, cursor int64, limit int) ([]*User, error)
	ListAllOrdered(ctx context.Context,
		pageNumber int,
		pageSize int,
		orderBy string,
		asc bool,
	) (*UserListResponse, error)
}

type ZepLLM interface {
	// Call runs the LLM chat completion against the prompt
	// this version of Call uses the chat endpoint of an LLM, but
	// we pass in a simple string prompt
	Call(
		ctx context.Context,
		prompt string,
		options ...llms.CallOption,
	) (string, error)
	// EmbedTexts embeds the given texts
	EmbedTexts(ctx context.Context, texts []string) ([][]float32, error)
	// GetTokenCount returns the number of tokens in the given text
	GetTokenCount(text string) (int, error)
	// Init initializes the LLM
	Init(ctx context.Context, cfg *config.Config) error
}