rag

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

type Community struct {
	ID         string         `json:"id"`
	Level      int            `json:"level"`
	Title      string         `json:"title"`
	Summary    string         `json:"summary"`
	Entities   []string       `json:"entities"`
	ParentID   string         `json:"parent_id,omitempty"`
	Children   []string       `json:"children,omitempty"`
	Properties map[string]any `json:"properties,omitempty"`
	Score      float64        `json:"score,omitempty"`
	CreatedAt  time.Time      `json:"created_at"`
	UpdatedAt  time.Time      `json:"updated_at"`
}

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

type Config struct {
	VectorRAG *VectorRAGConfig `json:"vector_rag,omitempty"`
	GraphRAG  *GraphRAGConfig  `json:"graph_rag,omitempty"`
	LightRAG  *LightRAGConfig  `json:"lightrag,omitempty"`
}

type Document struct {
	ID        string         `json:"id"`
	Content   string         `json:"content"`
	Metadata  map[string]any `json:"metadata"`
	Embedding []float32      `json:"embedding,omitempty"`
	CreatedAt time.Time      `json:"created_at"`
	UpdatedAt time.Time      `json:"updated_at"`
}

type DocumentLoader interface {
	Load(ctx context.Context) ([]Document, error)
}

type DocumentSearchResult struct {
	Document Document       `json:"document"`
	Score    float64        `json:"score"`
	Metadata map[string]any `json:"metadata,omitempty"`
}

type Embedder interface {
	EmbedDocument(ctx context.Context, text string) ([]float32, error)
	EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error)
	GetDimension() int
}

type Engine interface {
	Query(ctx context.Context, query string) (*QueryResult, error)
	QueryWithConfig(ctx context.Context, query string, config *RetrievalConfig) (*QueryResult, error)
	AddDocuments(ctx context.Context, docs []Document) error
	DeleteDocument(ctx context.Context, docID string) error
	UpdateDocument(ctx context.Context, doc Document) error
	SimilaritySearch(ctx context.Context, query string, k int) ([]Document, error)
	SimilaritySearchWithScores(ctx context.Context, query string, k int) ([]DocumentSearchResult, error)
}

type Entity struct {
	ID         string         `json:"id"`
	Type       string         `json:"type"`
	Name       string         `json:"name"`
	Properties map[string]any `json:"properties"`
	Embedding  []float32      `json:"embedding,omitempty"`
	CreatedAt  time.Time      `json:"created_at"`
	UpdatedAt  time.Time      `json:"updated_at"`
}

type GlobalRetrievalConfig struct {
	// Maximum number of communities to retrieve
	MaxCommunities int `json:"max_communities"`

	// Include community hierarchy
	IncludeHierarchy bool `json:"include_hierarchy"`

	// Weight for community relevance
	CommunityWeight float64 `json:"community_weight"`

	// Maximum hierarchy depth
	MaxHierarchyDepth int `json:"max_hierarchy_depth"`
}

type GraphQuery struct {
	EntityTypes   []string       `json:"entity_types,omitempty"`
	Relationships []string       `json:"relationships,omitempty"`
	Filters       map[string]any `json:"filters,omitempty"`
	Limit         int            `json:"limit,omitempty"`
	MaxDepth      int            `json:"max_depth,omitempty"`
	StartEntity   string         `json:"start_entity,omitempty"`
	EntityType    string         `json:"entity_type,omitempty"`
}

type GraphQueryResult struct {
	Entities      []*Entity       `json:"entities"`
	Relationships []*Relationship `json:"relationships"`
	Paths         [][]*Entity     `json:"paths,omitempty"`
	Score         float64         `json:"score"`
	Scores        []float64       `json:"scores,omitempty"`
	Metadata      map[string]any  `json:"metadata,omitempty"`
}

type GraphRAGConfig struct {
	DatabaseURL      string              `json:"database_url"`
	ModelProvider    string              `json:"model_provider"`
	EmbeddingModel   string              `json:"embedding_model"`
	ChatModel        string              `json:"chat_model"`
	EntityTypes      []string            `json:"entity_types"`
	Relationships    map[string][]string `json:"relationships"`
	MaxDepth         int                 `json:"max_depth"`
	EnableReasoning  bool                `json:"enable_reasoning"`
	ExtractionPrompt string              `json:"extraction_prompt"`
}

type HybridRetrievalConfig struct {
	// Weight for local retrieval results (0-1)
	LocalWeight float64 `json:"local_weight"`

	// Weight for global retrieval results (0-1)
	GlobalWeight float64 `json:"global_weight"`

	// Fusion method: "rrf" (reciprocal rank fusion) or "weighted"
	FusionMethod string `json:"fusion_method"`

	// RRF parameter for rank fusion
	RFFK int `json:"rrf_k"`
}

type KnowledgeGraph interface {
	AddEntity(ctx context.Context, entity *Entity) error
	AddRelationship(ctx context.Context, relationship *Relationship) error
	Query(ctx context.Context, query *GraphQuery) (*GraphQueryResult, error)
	GetRelatedEntities(ctx context.Context, entityID string, maxDepth int) ([]*Entity, error)
	GetEntity(ctx context.Context, entityID string) (*Entity, error)
}

type LLMInterface interface {
	Generate(ctx context.Context, prompt string) (string, error)
	GenerateWithConfig(ctx context.Context, prompt string, config map[string]any) (string, error)
	GenerateWithSystem(ctx context.Context, system, prompt string) (string, error)
}

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

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

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

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

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

type LightRAGConfig struct {
	// Retrieval mode: "local", "global", "hybrid", or "naive"
	Mode string `json:"mode"`

	// Local retrieval configuration
	LocalConfig LocalRetrievalConfig `json:"local_config"`

	// Global retrieval configuration
	GlobalConfig GlobalRetrievalConfig `json:"global_config"`

	// Hybrid retrieval configuration
	HybridConfig HybridRetrievalConfig `json:"hybrid_config"`

	// Chunk size for text splitting
	ChunkSize int `json:"chunk_size"`

	// Chunk overlap for text splitting
	ChunkOverlap int `json:"chunk_overlap"`

	// Threshold for entity extraction
	EntityExtractionThreshold float64 `json:"entity_extraction_threshold"`

	// Maximum number of entities to extract per chunk
	MaxEntitiesPerChunk int `json:"max_entities_per_chunk"`

	// Enable community detection for global retrieval
	EnableCommunityDetection bool `json:"enable_community_detection"`

	// Community detection algorithm: "louvain", "leiden", or "label_propagation"
	CommunityDetectionAlgorithm string `json:"community_detection_algorithm"`

	// Number of communities to return in global retrieval
	MaxCommunities int `json:"max_communities"`

	// Temperature for LLM-based operations
	Temperature float64 `json:"temperature"`

	// Custom prompt templates
	PromptTemplates map[string]string `json:"prompt_templates,omitempty"`
}

type LocalRetrievalConfig struct {
	// Maximum number of hops in the knowledge graph
	MaxHops int `json:"max_hops"`

	// Number of entities to retrieve
	TopK int `json:"top_k"`

	// Include entity descriptions
	IncludeDescriptions bool `json:"include_descriptions"`

	// Weight for entity relevance
	EntityWeight float64 `json:"entity_weight"`
}

type Metrics struct {
	TotalQueries    int64         `json:"total_queries"`
	TotalDocuments  int64         `json:"total_documents"`
	AverageLatency  time.Duration `json:"average_latency"`
	MinLatency      time.Duration `json:"min_latency"`
	MaxLatency      time.Duration `json:"max_latency"`
	LastQueryTime   time.Time     `json:"last_query_time"`
	CacheHits       int64         `json:"cache_hits"`
	CacheMisses     int64         `json:"cache_misses"`
	IndexingLatency time.Duration `json:"indexing_latency"`
}

type PipelineConfig struct {
	// Retrieval configuration
	TopK           int     // Number of documents to retrieve
	ScoreThreshold float64 // Minimum relevance score
	UseReranking   bool    // Whether to use reranking
	UseFallback    bool    // Whether to use fallback search

	// Generation configuration
	SystemPrompt     string
	IncludeCitations bool
	MaxTokens        int
	Temperature      float64

	// Components
	Loader      RAGDocumentLoader
	Splitter    RAGTextSplitter
	Embedder    Embedder
	VectorStore VectorStore
	Retriever   Retriever
	Reranker    Reranker
	LLM         llms.Model
}

type QueryResult struct {
	Query        string         `json:"query"`
	Answer       string         `json:"answer"`
	Sources      []Document     `json:"sources"`
	Context      string         `json:"context"`
	Confidence   float64        `json:"confidence"`
	ResponseTime time.Duration  `json:"response_time"`
	Metadata     map[string]any `json:"metadata"`
}

type RAGConfig struct {
	Config        `json:"config"`
	EnableCache   bool          `json:"enable_cache"`
	CacheSize     int           `json:"cache_size"`
	EnableMetrics bool          `json:"enable_metrics"`
	Debug         bool          `json:"debug"`
	Timeout       time.Duration `json:"timeout"`
}

type RAGDocument struct {
	Content   string         `json:"content"`
	Metadata  map[string]any `json:"metadata"`
	CreatedAt time.Time      `json:"created_at"`
	UpdatedAt time.Time      `json:"updated_at"`
}

type RAGDocumentLoader interface {
	Load(ctx context.Context) ([]RAGDocument, error)
}

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

type RAGState struct {
	Query              string
	Documents          []RAGDocument
	RetrievedDocuments []RAGDocument
	RankedDocuments    []DocumentSearchResult
	Context            string
	Answer             string
	Citations          []string
	Metadata           map[string]any
}

type RAGTextSplitter interface {
	SplitDocuments(documents []RAGDocument) ([]RAGDocument, error)
}

type Relationship struct {
	ID         string         `json:"id"`
	Source     string         `json:"source"`
	Target     string         `json:"target"`
	Type       string         `json:"type"`
	Properties map[string]any `json:"properties"`
	Weight     float64        `json:"weight,omitempty"`
	Confidence float64        `json:"confidence,omitempty"`
	CreatedAt  time.Time      `json:"created_at"`
}

type Reranker interface {
	Rerank(ctx context.Context, query string, documents []DocumentSearchResult) ([]DocumentSearchResult, error)
}

type RetrievalConfig struct {
	K              int            `json:"k"`
	ScoreThreshold float64        `json:"score_threshold"`
	SearchType     string         `json:"search_type"`
	Filter         map[string]any `json:"filter,omitempty"`
	IncludeScores  bool           `json:"include_scores"`
}

type Retriever interface {
	Retrieve(ctx context.Context, query string) ([]Document, error)
	RetrieveWithK(ctx context.Context, query string, k int) ([]Document, error)
	RetrieveWithConfig(ctx context.Context, query string, config *RetrievalConfig) ([]DocumentSearchResult, error)
}

type RetrieverTool struct {
	Engine  Engine
	NameVal string
	DescVal string
}

type TextSplitter interface {
	SplitText(text string) []string
	SplitDocuments(documents []Document) []Document
	JoinText(chunks []string) string
}

type VectorRAGConfig struct {
	EmbeddingModel    string          `json:"embedding_model"`
	VectorStoreType   string          `json:"vector_store_type"`
	VectorStoreConfig map[string]any  `json:"vector_store_config"`
	ChunkSize         int             `json:"chunk_size"`
	ChunkOverlap      int             `json:"chunk_overlap"`
	EnableReranking   bool            `json:"enable_reranking"`
	RetrieverConfig   RetrievalConfig `json:"retriever_config"`
}

type VectorStore interface {
	Add(ctx context.Context, documents []Document) error
	Search(ctx context.Context, query []float32, k int) ([]DocumentSearchResult, error)
	SearchWithFilter(ctx context.Context, query []float32, k int, filter map[string]any) ([]DocumentSearchResult, error)
	Delete(ctx context.Context, ids []string) error
	Update(ctx context.Context, documents []Document) error
	GetStats(ctx context.Context) (*VectorStoreStats, error)
}

type VectorStoreStats struct {
	TotalDocuments int       `json:"total_documents"`
	TotalVectors   int       `json:"total_vectors"`
	Dimension      int       `json:"dimension"`
	LastUpdated    time.Time `json:"last_updated"`
}