golightrag

type Document struct {
	ID      string
	Content string
}

type DocumentHandler interface {
	// ChunksDocument splits a document's content into smaller, manageable chunks.
	// It returns a slice of Source objects representing the document chunks,
	// without assigning IDs (IDs will be generated in the Insert function).
	ChunksDocument(content string) ([]Source, error)
	// EntityExtractionPromptData returns the data needed to generate prompts for extracting
	// entities and relationships from text content.
	// The implementation doesn't need to fill the Input field, as it will be filled in the
	// Insert function.
	EntityExtractionPromptData() EntityExtractionPromptData
	// MaxRetries determines the maximum number of retries allowed for the Chat function.
	// This is especially used when extracting entities and relationships from text content,
	// due to the incorrect format that sometimes LLM returns.
	MaxRetries() int
	// ConcurrencyCount determines the number of concurrent requests to the LLM.
	ConcurrencyCount() int
	// BackoffDuration determines the backoff duration between retries.
	BackoffDuration() time.Duration
	// GleanCount returns the maximum number of additional extraction attempts
	// to perform after the initial entity extraction to find entities that might
	// have been missed.
	GleanCount() int
	// MaxSummariesTokenLength returns the maximum token length allowed for entity
	// and relationship descriptions before they need to be summarized by the LLM.
	MaxSummariesTokenLength() int
}

type EntityContext struct {
	Name        string
	Type        string
	Description string
	RefCount    int
	CreatedAt   time.Time
}

type EntityExtractionPromptData struct {
	Goal        string
	EntityTypes []string
	Language    string
	Examples    []EntityExtractionPromptExample

	Input string
}

type EntityExtractionPromptEntityOutput struct {
	Name        string
	Type        string
	Description string
}

type EntityExtractionPromptExample struct {
	EntityTypes          []string
	Text                 string
	EntitiesOutputs      []EntityExtractionPromptEntityOutput
	RelationshipsOutputs []EntityExtractionPromptRelationshipOutput
}

type EntityExtractionPromptRelationshipOutput struct {
	SourceEntity string
	TargetEntity string
	Description  string
	Keywords     []string
	Strength     float64
}

type GraphEntity struct {
	Name         string `json:"entity_name"`
	Type         string `json:"entity_type"`
	Descriptions string `json:"entity_description"`
	SourceIDs    string
	CreatedAt    time.Time
}

type GraphRelationship struct {
	SourceEntity string   `json:"source_entity"`
	TargetEntity string   `json:"target_entity"`
	Weight       float64  `json:"relationship_strength"`
	Descriptions string   `json:"relationship_description"`
	Keywords     []string `json:"relationship_keywords"`
	SourceIDs    string
	CreatedAt    time.Time
}

type GraphStorage interface {
	// GraphEntity retrieves a single entity by name from the graph storage.
	// Returns ErrEntityNotFound if the entity doesn't exist.
	GraphEntity(name string) (GraphEntity, error)
	// GraphRelationship retrieves a relationship between sourceEntity and targetEntity.
	// Returns ErrRelationshipNotFound if the relationship doesn't exist.
	GraphRelationship(sourceEntity, targetEntity string) (GraphRelationship, error)

	// GraphUpsertEntity creates a new entity or updates an existing entity in the graph storage.
	// If the entity already exists, it should merge the new data with existing data.
	GraphUpsertEntity(entity GraphEntity) error
	// GraphUpsertRelationship creates a new relationship or updates an existing relationship
	// between two entities in the graph storage.
	// If the relationship already exists, it should merge the new data with existing data.
	GraphUpsertRelationship(relationship GraphRelationship) error

	// GraphEntities batch retrieves multiple entities by their names.
	// Returns a map with entity names as keys and entity objects as values.
	// If an entity doesn't exist, it should be omitted from the result map.
	GraphEntities(names []string) (map[string]GraphEntity, error)
	// GraphRelationships batch retrieves multiple relationships by their source-target pairs.
	// Returns a map with composite keys (formatted as "source-target") as keys and
	// relationship objects as values.
	// If a relationship doesn't exist, it should be omitted from the result map.
	GraphRelationships(pairs [][2]string) (map[string]GraphRelationship, error)

	// GraphCountEntitiesRelationships counts the number of relationships each entity has.
	// Returns a map with entity names as keys and relationship counts as values.
	// This is used to determine entity importance during queries.
	GraphCountEntitiesRelationships(names []string) (map[string]int, error)
	// GraphRelatedEntities finds entities directly connected to the specified entities.
	// Returns a map with entity names as keys and slices of directly connected entities as values.
	// Used to expand the context during queries.
	GraphRelatedEntities(names []string) (map[string][]GraphEntity, error)
}

type KeyValueStorage interface {
	// KVSource retrieves a source document chunk by its ID.
	// Returns an error if the source doesn't exist or can't be retrieved.
	KVSource(id string) (Source, error)
	// KVUpsertSources creates or updates multiple source document chunks at once.
	// Each source should be stored with its ID as the key.
	// This is called during document processing to store chunked documents.
	KVUpsertSources(sources []Source) error
}

type KeywordExtractionPromptData struct {
	Goal     string
	Examples []KeywordExtractionPromptExample

	Query   string
	History string
}

type KeywordExtractionPromptExample struct {
	Query             string
	LowLevelKeywords  []string
	HighLevelKeywords []string
}

type LLM interface {
	// Chat sends messages to the LLM and returns the response.
	// A message with an even index is guaranteed to be sent by the user, while the odd index is
	// sent by the assistant.
	Chat(messages []string) (string, error)
}

type QueryConversation struct {
	Message string
	Role    string
}

type QueryHandler interface {
	// KeywordExtractionPromptData returns the data needed to generate prompts for extracting
	// keywords from user queries and conversation history.
	// The implementation doesn't need to fill the Query and History fields, as they will be filled
	// in the Query function.
	KeywordExtractionPromptData() KeywordExtractionPromptData
}

type QueryResult struct {
	GlobalEntities      []EntityContext
	GlobalRelationships []RelationshipContext
	GlobalSources       []SourceContext
	LocalEntities       []EntityContext
	LocalRelationships  []RelationshipContext
	LocalSources        []SourceContext
}

type RelationshipContext struct {
	Source      string
	Target      string
	Keywords    string
	Description string
	Weight      float64
	RefCount    int
	CreatedAt   time.Time
}

type Source struct {
	ID         string
	Content    string
	TokenSize  int
	OrderIndex int
}

type SourceContext struct {
	Content  string
	RefCount int
}

type Storage interface {
	GraphStorage
	VectorStorage
	KeyValueStorage
}

type VectorStorage interface {
	// VectorQueryEntity performs a semantic search for entities based on the provided keywords.
	// Returns a slice of entity names that semantically match the keywords.
	// The results should be ordered by relevance.
	VectorQueryEntity(keywords string) ([]string, error)
	// VectorQueryRelationship performs a semantic search for relationships based on the provided keywords.
	// Returns a slice of source-target entity name pairs that semantically match the keywords.
	// The results should be ordered by relevance.
	VectorQueryRelationship(keywords string) ([][2]string, error)

	// VectorUpsertEntity creates or updates the vector representation of an entity.
	// The content parameter should contain the text used for semantic matching.
	// This typically includes the entity name and description.
	VectorUpsertEntity(name, content string) error
	// VectorUpsertRelationship creates or updates the vector representation of a relationship.
	// The content parameter should contain the text used for semantic matching.
	// This typically includes keywords, descriptions, and entity names.
	VectorUpsertRelationship(source, target, content string) error
}