rag

type ContextChunk struct {
	EpisodeID   string                 `json:"episode_id"`
	Text        string                 `json:"text"`
	Score       float32                `json:"score"` // Similarity score (cosine distance)
	StartDate   time.Time              `json:"start_date"`
	EndDate     time.Time              `json:"end_date"`
	Authors     []string               `json:"authors"`
	CommitCount int                    `json:"commit_count"`
	FileCount   int                    `json:"file_count"`
	Metadata    map[string]interface{} `json:"metadata,omitempty"`
}

type Embedder interface {
	// Embed generates embeddings for the provided texts
	Embed(ctx context.Context, texts []string) ([]EmbeddingRecord, error)
}

type EmbeddingRecord struct {
	Text      string    `json:"text"`
	Embedding []float32 `json:"embedding"`
	Index     int       `json:"index"`
	Model     string    `json:"model"`
}

type EpisodeRecord struct {
	EpisodeID   string
	Text        string
	Embedding   []float32
	StartDate   time.Time
	EndDate     time.Time
	Authors     []string
	CommitCount int
	FileCount   int
}

type EpisodeSummary struct {
	EpisodeID   string    `json:"episode_id"`
	Title       string    `json:"title,omitempty"`
	Summary     string    `json:"summary"`
	StartDate   time.Time `json:"start_date,omitempty"`
	EndDate     time.Time `json:"end_date,omitempty"`
	Authors     []string  `json:"authors,omitempty"`
	CommitCount int       `json:"commit_count"`
	FileCount   int       `json:"file_count"`
}

type IndexOptions struct {
	// BatchSize determines how many episodes to embed at once
	BatchSize int

	// ForceReindex will delete and re-insert episodes even if they exist
	ForceReindex bool

	// SkipExisting will check if episode already exists and skip if present
	SkipExisting bool
}

type MilvusConfig struct {
	Address        string // Milvus server address (e.g., "localhost:19530")
	CollectionName string // Name of the collection
	Dimension      int    // Vector dimension (e.g., 3072 for text-embedding-3-large)
	IndexType      string // Index type (default: "HNSW")
	MetricType     string // Similarity metric (default: "COSINE")

	// HNSW index parameters
	M              int // HNSW M parameter (default: 16)
	EfConstruction int // HNSW efConstruction (default: 256)
}

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

type OpenAIEmbedder struct {
	Model     string
	Dimension int
	// contains filtered or unexported fields
}

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

type SearchOptions struct {
	EpisodeIDs []string               `json:"episode_ids,omitempty"` // Filter by specific episode IDs
	Repository string                 `json:"repository,omitempty"`  // Filter by repository name
	Metadata   map[string]interface{} `json:"metadata,omitempty"`    // Additional metadata filters
}

type VectorStore interface {
	// Insert efficiently inserts multiple episodes in a single operation
	Insert(ctx context.Context, episodes []EpisodeRecord) error

	// Flush ensures all pending data is persisted
	Flush(ctx context.Context) error

	// Search performs top-K similarity search with optional filtering
	Search(ctx context.Context, queryVector []float32, topK int, opts *SearchOptions) ([]ContextChunk, error)

	// Query checks which episode IDs exist in the store
	// Returns a map where keys are episode IDs and values indicate existence
	Query(ctx context.Context, episodeIDs []string) (map[string]bool, error)

	// Delete removes records by episode IDs
	Delete(ctx context.Context, episodeIDs []string) error

	// GetStats returns collection statistics (record count, index status, etc.)
	GetStats(ctx context.Context) (map[string]interface{}, error)

	// Close releases resources and closes connections
	Close() error
}