memory

type Category struct {
	ID sop.UUID `json:"id"`
	// ParentIDs points to parent Categories, allowing a Directed Acyclic Graph (DAG) / Polyhierarchy.
	// Use-case: A category like "Database Migrations" can belong to both "Release Management"
	// and "Database Administration". Adding the explicit UseCase here gives the LLM
	// the ability to review, validate, and improve these graph edges during deep sleep cycles.
	ParentIDs       []CategoryParent `json:"parents,omitempty"`
	CenterVector    []float32        `json:"center_vector"`               // Mathematical center of this chunk/category
	ChildrenIDs     []sop.UUID       `json:"children_ids,omitempty"`      // IDs of Sub-Categories
	Radius          float32          `json:"radius,omitempty"`            // Size of the cluster
	ItemCount       int              `json:"item_count,omitempty"`        // Number of vectors/items in this bucket
	Name            string           `json:"name,omitempty"`              // Human-readable concept name
	Path            string           `json:"path,omitempty"`              // Full contextual taxonomy path (e.g. "tools / execute_script")
	Description     string           `json:"description,omitempty"`       // Broader context
	SummaryMaxCount int              `json:"summary_max_count,omitempty"` // Maximum number of summaries for items in this category
	VectorHash      string           `json:"vector_hash,omitempty"`       // Hash of EmbedderName + Content to deduplicate vectorization
}

type CategoryParent struct {
	ParentID sop.UUID `json:"parent_id"`
	UseCase  string   `json:"use_case,omitempty"` // The explicit justification (can be empty "" for the primary/obvious parent)
}

type ChunkData struct {
	Text        string   `json:"text,omitempty"`        // Small snippet or chunk directly answerable
	Description string   `json:"description,omitempty"` // Contextual description or rationale
	DocumentID  sop.UUID `json:"document_id,omitempty"` // Pointer to the heavyweight Document/Blob (can be NilUUID)
}

type Database interface {
	ai.Database
	OpenKnowledgeBase(ctx context.Context, name string, tx sop.Transaction, llm ai.Generator, embedder ai.Embeddings, documentMode bool, enableTextSearch ...bool) (*KnowledgeBase[map[string]any], error)
	NewBtree(ctx context.Context, name string, t sop.Transaction) (btree.BtreeInterface[string, any], error)
}

type DistanceKey struct {
	ParentID sop.UUID // NilUUID for Level 1 Macro-Categories
	Distance float32  // Mathematical distance relative to the bounding anchor
	ID       sop.UUID // ID of the Category
}

type DocIDs []string

type Document struct {
	ID          sop.UUID `json:"id"`
	Title       string   `json:"title,omitempty"`
	URL         string   `json:"url,omitempty"`
	Source      string   `json:"source,omitempty"`
	ContentType string   `json:"content_type,omitempty"` // e.g. "text/markdown", "text/plain"
	Content     string   `json:"content,omitempty"`
	Data        []byte   `json:"data,omitempty"` // For pure blobs, pdf binaries, etc
}

type ExportData[T any] struct {
	Config     *KnowledgeBaseConfig `json:"config,omitempty"`
	Categories []*Category          `json:"categories"`
	Documents  []*Document          `json:"documents,omitempty"`
	Items      []ExportItem[T]      `json:"items"`
}

type ExportItem[T any] struct {
	CategoryPath     string      `json:"category"`
	DocID            DocIDs      `json:"doc_id,omitempty"`
	Data             T           `json:"data"`
	Summaries        []string    `json:"summaries,omitempty"`
	SummariesVectors [][]float32 `json:"summaries_vectors,omitempty"`
	Positions        []VectorKey `json:"positions,omitempty"`
	VectorHash       string      `json:"vector_hash,omitempty"`
}

type Item[T any] struct {
	ID         sop.UUID    `json:"id"`
	CategoryID sop.UUID    `json:"category_id"`
	DocID      DocIDs      `json:"doc_id,omitempty"`      // UUID of uploaded documents OR string URI for external docs
	Summaries  []string    `json:"summaries,omitempty"`   // 1 or more distinct, clean sentences for vector indexing
	Data       T           `json:"data"`                  // The application data or structured thought
	Positions  []VectorKey `json:"positions,omitempty"`   // Direct links to its Vectors for O(1) cleanup during Category moves
	VectorHash string      `json:"vector_hash,omitempty"` // Hash of EmbedderName + Content to avoid re-vectorizing unchanged items
}

type ItemKey struct {
	CategoryID sop.UUID
	ItemID     sop.UUID
}

type KBDigestHit struct {
	DocID      []string
	Score      float32
	Category   string
	Text       string
	Query      string
	SearchType string
}

type KBDigestRequest struct {
	Queries            []string
	PerQueryLimit      int
	MaxResults         int
	MinScore           float32
	UseClosestCategory bool
	KeywordFallback    bool
}

type KnowledgeBase[T any] struct {
	Store   MemoryStore[T]
	Manager *MemoryManager[T]

	// MaxMathCategoryDistance specifies the max Euclidean distance to cluster centroids
	// to avoid calling the LLM for category categorization. Set to 0.0 or less to disable
	// and always rely on "pristine" LLM categorization.
	MaxMathCategoryDistance float32
	// contains filtered or unexported fields
}

type KnowledgeBaseConfig struct {
	Type                string            `json:"type,omitempty"`
	IsPersona           bool              `json:"is_persona,omitempty"`
	IsExclusive         bool              `json:"is_exclusive,omitempty"`
	Description         string            `json:"description,omitempty"`
	SystemPrompt        string            `json:"system_prompt,omitempty"`
	Embedder            string            `json:"embedder,omitempty"`
	EmbedderDimension   int               `json:"embedder_dimension,omitempty"`
	AllowAutoEnrichment bool              `json:"allowAutoEnrichment,omitempty"`
	AllowedTools        []string          `json:"allowed_tools,omitempty"`
	ToolQueries         []PathSearchParam `json:"tool_queries,omitempty"`
	LastModified        int64             `json:"last_modified,omitempty"`   // Unix timestamp
	LastVectorized      int64             `json:"last_vectorized,omitempty"` // Unix timestamp
	RoutingPrefix       string            `json:"routing_prefix,omitempty"`
	DomainReference     []float32         `json:"domain_reference,omitempty"`
	// DocumentMode flags whether this KB operates in traditional payload mode (Item.Data holds data),
	// or in decoupled RAG references mode where Item.Data points to the canonical large Document(MD).
	DocumentMode      bool `json:"document_mode,omitempty"`
	TextSearchEnabled bool `json:"text_search_enabled,omitempty"` // Controls if keyword/BM25 search is indexed and available
}

type LLM[T any] interface {
	// GenerateCategory invokes the model to synthesize a new Category
	// based off the underlying data structure's payload.
	GenerateCategory(ctx context.Context, payload T) (*Category, error)
}

type ListCategoriesParam struct {
	ParentPath string   `json:"parent_path"` // Optional: restrict list to a specific parent node
	ParentID   sop.UUID `json:"parent_id"`   // Optional: explicit restriction by ID
	Limit      int      `json:"limit"`
	Offset     int      `json:"offset"`
}

type ListItemsParam struct {
	CategoryPath string   `json:"category_path"` // Optional: restrict list to a specific category
	CategoryID   sop.UUID `json:"category_id"`   // Optional: restrict list by explicit ID
	Limit        int      `json:"limit"`
	Offset       int      `json:"offset"`
}

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

type MemoryStore[T any] interface {

	// Returns this Memory Store's name.
	Name() string

	// Upsert adds or updates a single item in the store.
	Upsert(ctx context.Context, item Item[T], vec []float32) error

	// UpsertByCategoryPath explicitly assigns a category ignoring spatial routing.
	UpsertByCategoryPath(ctx context.Context, categoryName string, item Item[T], vecs [][]float32) error

	// UpsertByCategoryID inserts data bypassing Category lookup.
	// vecs are DocumentTexts (768 dim), classificationVecs are CategoryTexts (256 dim) for DistanceToCategory.
	UpsertByCategoryID(ctx context.Context, catID sop.UUID, catCenterVector []float32, item Item[T], vecs [][]float32, classificationVecs [][]float32) error

	// UpsertBatch adds or updates multiple items in the store efficiently.
	UpsertBatch(ctx context.Context, items []Item[T], vecs [][]float32) error

	// FindClosestCategory explores the category tree using spatial distance to find the closest matching category.
	FindClosestCategory(ctx context.Context, vector []float32) (*Category, float32, error)

	// SemanticCategoryByPath resolves a category path expressed as pre-embedded vectors into the
	// closest matching Categories at each level of the hierarchy.
	//
	// Given a path "a/b/c" whose parts have been embedded into vectors [va, vb, vc]:
	//   - Level 0 (root): searches CategoriesByDistance with ParentID=NilUUID, anchor=DomainReference
	//   - Level N:        searches CategoriesByDistance with ParentID=prev.ID, anchor=prev.CenterVector
	// The function keeps all best-distance ties per level and returns the final best candidates.
	SemanticCategoryByPath(ctx context.Context, pathVectors [][]float32) ([]*Category, error)

	// Get retrieves a item by its logical ID.
	Get(ctx context.Context, key ItemKey) (*Item[T], error)
	// Delete removes an item by its logical ID.
	Delete(ctx context.Context, key ItemKey) error

	// Query searches for the nearest neighbors to the given vector coordinates.
	// filters is a function that returns true if the item should be included.
	Query(ctx context.Context, vec []float32, opts *SearchOptions[T]) ([]ai.Hit[T], error)

	// QueryItems searches stored items for the already-resolved category.
	// This lets the KnowledgeBase path reuse resolved categories instead of resolving them again.
	QueryItems(ctx context.Context, vec []float32, category *Category, opts *SearchOptions[T]) ([]ai.Hit[T], error)

	// QueryBatch searches for the nearest neighbors for a slice of query vectors.
	QueryBatch(ctx context.Context, vecs [][]float32, opts *SearchOptions[T]) ([][]ai.Hit[T], error)

	// QueryText performs a BM25 or keyword text search on the stored text representation of the thoughts.
	QueryText(ctx context.Context, text string, opts *SearchOptions[T]) ([]ai.Hit[T], error)

	// QueryTextBatch performs a BM25 or keyword text search for an array of queries.
	QueryTextBatch(ctx context.Context, texts []string, opts *SearchOptions[T]) ([][]ai.Hit[T], error)

	// Count returns the total number of items in the store.
	Count(ctx context.Context) (int64, error)

	// Categories returns a B-Tree interface to manually read/update hierarchical categories.
	Categories(ctx context.Context) (btree.BtreeInterface[sop.UUID, *Category], error)

	// CategoriesByPath returns a B-Tree interface to manually read/update path-indexed categories.
	CategoriesByPath(ctx context.Context) (btree.BtreeInterface[string, sop.UUID], error)

	// CategoriesByDistance returns a B-Tree interface for reading/updating distance-indexed categories.
	CategoriesByDistance(ctx context.Context) (btree.BtreeInterface[DistanceKey, byte], error)

	// AddCategory adds a new category to the store dynamically.
	// This allows for runtime expansion of the concept space without full rebalancing.
	AddCategory(ctx context.Context, c *Category) (sop.UUID, error)

	// AddCategoryParent connects an existing category to an additional parent, supporting
	// the polyhierarchy DAG structure. This is often leveraged during LLM Sleep Cycles.
	AddCategoryParent(ctx context.Context, categoryID sop.UUID, parent CategoryParent) error

	// Consolidate reads accumulated vectors from short-term memory (TempVectors),
	// dynamically routes them into existing Categories using AssignAndIndex logic,
	// and clears them from short-term memory.
	Consolidate(ctx context.Context) error

	// UpdateEmbedderInfo updates the configuration defining which embedder was used
	// to index the vectors, persisting it in the system configuration of the store.
	UpdateEmbedderInfo(ctx context.Context, provider string, model string, dimensions int) error

	// SetDomainReference sets the anchor vector for O(log N) category indexing.
	SetDomainReference(vec []float32)

	// DomainReference returns the anchor vector for O(log N) category indexing.
	DomainReference() []float32

	// SetLLM sets the LLM interface used to generate categories dynamically.
	SetLLM(llm LLM[T])

	// Vectors returns the Vectors B-Tree for advanced manipulation (Mathematical layout).
	Vectors(ctx context.Context) (btree.BtreeInterface[VectorKey, Vector], error)

	// Content returns the Content B-Tree for advanced manipulation (The actual Item Data).
	Items(ctx context.Context) (btree.BtreeInterface[ItemKey, Item[T]], error)

	// Documents returns the Documents B-Tree for reading the raw canonical documents.
	Documents(ctx context.Context) (btree.BtreeInterface[sop.UUID, Document], error)

	// UpsertDocument adds or updates a full canonical document.
	UpsertDocument(ctx context.Context, doc Document) error

	// GetDocument retrieves a full document by its ID.
	GetDocument(ctx context.Context, id sop.UUID) (*Document, error)

	// Version returns the Vector store's version number, which is a unix elapsed time.
	Version(ctx context.Context) (int64, error)
}

type MemoryUnit struct {
	AgentID    string
	UserID     string
	AllowedKBs []string // LTM scoping boundaries

	// Tracks the last time an episode was logged to STM for idle sleep cycles
	LastEpisodeTS atomic.Int64
	// contains filtered or unexported fields
}

type PathSearchParam struct {
	CategoryPath string `json:"category_path"` // e.g. "Root/Engineering/Architecture" (semantic or lexical)
	SearchText   string `json:"search_text"`   // Text to prefix search on item content/title
}

type Preference struct {
	Key          string   `json:"key"`
	BoolValue    *bool    `json:"bool_value,omitempty"`
	StringValue  string   `json:"string_value,omitempty"`
	NumberValue  *float64 `json:"number_value,omitempty"`
	UpdatedAtUTC int64    `json:"updated_at_utc,omitempty"`
	Source       string   `json:"source,omitempty"`
}

type SearchOptions[T any] struct {
	Limit int
	// CategoryPath can serve as a cheaper SearchByPath-style alternative to TextSearch
	// when the use-case has a stable, meaningful category taxonomy to route through.
	CategoryPath string
	// CategoryVector can be used to search for items within a given category.
	// The Category whose CenterVector is closest to this vector will be used as the search Category.
	CategoryVector []float32
	Filter         func(T) bool
}

type SearchRequest[T any] struct {
	Text           string
	Vector         []float32
	CategoryPath   string
	CategoryVector []float32
	Limit          int
	Filter         func(T) bool
}

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

type Thought[T any] struct {
	Summaries    []string
	CategoryPath string
	DocID        DocIDs
	Data         T
	Vectors      [][]float32
	Positions    []VectorKey
	VectorHash   string
}

type UpsertCategoryParam struct {
	ParentPaths []string   `json:"parent_paths"` // Declarative DAG routing (e.g. ["Root/Engineering", "Root/DevOps"])
	ParentIDs   []sop.UUID `json:"parent_ids"`   // Explicit DAG routing fallback
	Category    *Category  `json:"category"`     // Pointer to the Category to upsert
}

type UpsertItemParam[T any] struct {
	CategoryPath string      `json:"category_path"` // e.g. "Root/Engineering/Architecture"
	CategoryID   sop.UUID    `json:"category_id"`   // Direct ID fallback if path is empty
	Item         *Item[T]    `json:"item"`          // Pointer to avoid heavy allocation during batch
	Vectors      [][]float32 `json:"vectors"`       // Optional explicit embeddings
}

type Vector struct {
	ID         sop.UUID  `json:"id"`
	Data       []float32 `json:"data"`        // Math coordinate
	ItemID     sop.UUID  `json:"item_id"`     // Points to the actual Item
	CategoryID sop.UUID  `json:"category_id"` // Critical for category-partitioned semantic searches
}

type VectorKey struct {
	CategoryID         sop.UUID // Points to the hierarchical Category ID
	DistanceToCategory float32
	VectorID           sop.UUID // Points to the specific Vector ID
}