c3b374f043
- Added LocalPath method to the Manager and Storage interfaces to retrieve the absolute path and content type of files. - Enhanced local and S3 storage implementations to support LocalPath, including handling gzipped files and content type detection. - Introduced comprehensive tests for LocalPath functionality, covering various file types, non-existent files, and gzipped content. - Updated AddFile API to utilize LocalPath for retrieving file information, improving error handling and response consistency.
386 lines
11 KiB
Go
386 lines
11 KiB
Go
package kb
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/yaoapp/gou/graphrag/types"
|
|
"github.com/yaoapp/yao/kb"
|
|
"github.com/yaoapp/yao/kb/providers/factory"
|
|
kbtypes "github.com/yaoapp/yao/kb/types"
|
|
)
|
|
|
|
/*
|
|
Usage Examples:
|
|
|
|
1. AddFile API (converter will be auto-detected based on file info):
|
|
{
|
|
"collection_id": "my_collection",
|
|
"file_id": "uploaded_file_123",
|
|
"chunking": {
|
|
"provider_id": "text_splitter",
|
|
"option_id": "default"
|
|
},
|
|
"embedding": {
|
|
"provider_id": "openai",
|
|
"option_id": "text-embedding-3-small"
|
|
},
|
|
"doc_id": "document_001",
|
|
"metadata": {
|
|
"source": "research_paper"
|
|
}
|
|
}
|
|
|
|
2. AddText API:
|
|
{
|
|
"collection_id": "my_collection",
|
|
"text": "This is the text content to be processed.",
|
|
"chunking": {
|
|
"provider_id": "text_splitter"
|
|
},
|
|
"embedding": {
|
|
"provider_id": "openai"
|
|
}
|
|
}
|
|
|
|
3. AddSegments API:
|
|
{
|
|
"collection_id": "my_collection",
|
|
"doc_id": "document_001",
|
|
"segment_texts": [
|
|
{"text": "First segment", "metadata": {"page": 1}},
|
|
{"text": "Second segment", "metadata": {"page": 2}}
|
|
],
|
|
"embedding": {
|
|
"provider_id": "openai",
|
|
"option_id": "text-embedding-3-small"
|
|
}
|
|
}
|
|
|
|
Note:
|
|
- If no option_id is specified, the default option from provider configuration will be selected
|
|
- For AddFile API, converter will be auto-detected based on filename and content_type obtained from GetFileInfo(file_id)
|
|
- ToUpsertOptions() can be called without parameters, or with filename and contentType for converter auto-detection
|
|
*/
|
|
|
|
// ProviderConfig represents a provider configuration that can be specified in two ways:
|
|
// 1. ProviderID + OptionID (option will be looked up from provider)
|
|
// 2. ProviderID + Option (option is provided directly)
|
|
type ProviderConfig struct {
|
|
ProviderID string `json:"provider_id" binding:"required"`
|
|
OptionID string `json:"option_id,omitempty"`
|
|
Option *kbtypes.ProviderOption `json:"option,omitempty"`
|
|
}
|
|
|
|
// BaseUpsertRequest contains common fields for all upsert operations
|
|
type BaseUpsertRequest struct {
|
|
// Collection ID - this will be mapped to UpsertOptions.CollectionID
|
|
CollectionID string `json:"collection_id" binding:"required"`
|
|
|
|
// Provider configurations
|
|
Chunking *ProviderConfig `json:"chunking" binding:"required"`
|
|
Embedding *ProviderConfig `json:"embedding" binding:"required"`
|
|
Extraction *ProviderConfig `json:"extraction,omitempty"`
|
|
Fetcher *ProviderConfig `json:"fetcher,omitempty"`
|
|
Converter *ProviderConfig `json:"converter,omitempty"`
|
|
|
|
// Upsert options
|
|
DocID string `json:"doc_id,omitempty"`
|
|
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
|
}
|
|
|
|
// AddFileRequest represents the request for AddFile API
|
|
type AddFileRequest struct {
|
|
BaseUpsertRequest
|
|
FileID string `json:"file_id" binding:"required"`
|
|
Uploader string `json:"uploader,omitempty"` // The name of the uploader, e.g. "s3", "local", "webdav", etc.
|
|
}
|
|
|
|
// AddTextRequest represents the request for AddText API
|
|
type AddTextRequest struct {
|
|
BaseUpsertRequest
|
|
Text string `json:"text" binding:"required"`
|
|
}
|
|
|
|
// AddURLRequest represents the request for AddURL API
|
|
type AddURLRequest struct {
|
|
BaseUpsertRequest
|
|
URL string `json:"url" binding:"required"`
|
|
}
|
|
|
|
// AddSegmentsRequest represents the request for AddSegments API
|
|
type AddSegmentsRequest struct {
|
|
BaseUpsertRequest
|
|
SegmentTexts []types.SegmentText `json:"segment_texts" binding:"required"`
|
|
}
|
|
|
|
// UpdateSegmentsRequest represents the request for UpdateSegments API
|
|
type UpdateSegmentsRequest struct {
|
|
BaseUpsertRequest
|
|
SegmentTexts []types.SegmentText `json:"segment_texts" binding:"required"`
|
|
}
|
|
|
|
// resolveProviderOption resolves a ProviderConfig to a *kbtypes.ProviderOption
|
|
// If OptionID is provided, it looks up the option from the provider
|
|
// If Option is provided directly, it uses the Option field
|
|
// If neither is provided, it selects the default option from provider's Options
|
|
func resolveProviderOption(config *ProviderConfig) (*kbtypes.ProviderOption, error) {
|
|
if config == nil {
|
|
return nil, fmt.Errorf("provider config is required")
|
|
}
|
|
|
|
if config.ProviderID == "" {
|
|
return nil, fmt.Errorf("provider_id is required")
|
|
}
|
|
|
|
// If Option is provided directly, use it
|
|
if config.Option != nil {
|
|
return config.Option, nil
|
|
}
|
|
|
|
// Get the provider from KB instance
|
|
if kb.Instance == nil {
|
|
return nil, fmt.Errorf("KB instance is not initialized")
|
|
}
|
|
|
|
// Find the provider in KB config
|
|
var provider *kbtypes.Provider
|
|
kbConfig := kb.Instance.(*kb.KnowledgeBase).Config
|
|
|
|
// Check all provider types to find the matching provider
|
|
allProviders := [][]*kbtypes.Provider{
|
|
kbConfig.Chunkings,
|
|
kbConfig.Embeddings,
|
|
kbConfig.Converters,
|
|
kbConfig.Extractors,
|
|
kbConfig.Fetchers,
|
|
}
|
|
|
|
for _, providers := range allProviders {
|
|
for _, p := range providers {
|
|
if p.ID == config.ProviderID {
|
|
provider = p
|
|
break
|
|
}
|
|
}
|
|
if provider != nil {
|
|
break
|
|
}
|
|
}
|
|
|
|
if provider == nil {
|
|
return nil, fmt.Errorf("provider %s not found", config.ProviderID)
|
|
}
|
|
|
|
// If OptionID is provided, look it up from the provider
|
|
if config.OptionID != "" {
|
|
option, exists := provider.GetOption(config.OptionID)
|
|
if !exists {
|
|
return nil, fmt.Errorf("option %s not found in provider %s", config.OptionID, config.ProviderID)
|
|
}
|
|
return option, nil
|
|
}
|
|
|
|
// If no option specified, try to find the default option
|
|
if provider.Options != nil {
|
|
for _, option := range provider.Options {
|
|
if option.Default {
|
|
return option, nil
|
|
}
|
|
}
|
|
// If no default option found but options exist, return the first one
|
|
if len(provider.Options) > 0 {
|
|
return provider.Options[0], nil
|
|
}
|
|
}
|
|
|
|
return nil, fmt.Errorf("no option specified and no default option found for provider %s", config.ProviderID)
|
|
}
|
|
|
|
// ToUpsertOptions converts BaseUpsertRequest to types.UpsertOptions
|
|
// Optional parameters: filename, contentType (for converter auto-detection)
|
|
func (r *BaseUpsertRequest) ToUpsertOptions(fileInfo ...string) (*types.UpsertOptions, error) {
|
|
var filename, contentType string
|
|
if len(fileInfo) >= 1 {
|
|
filename = fileInfo[0]
|
|
}
|
|
if len(fileInfo) >= 2 {
|
|
contentType = fileInfo[1]
|
|
}
|
|
|
|
options := &types.UpsertOptions{
|
|
CollectionID: r.CollectionID, // Collection ID maps to CollectionID
|
|
DocID: r.DocID,
|
|
Metadata: r.Metadata,
|
|
}
|
|
|
|
// Resolve and create chunking provider
|
|
chunkingOption, err := resolveProviderOption(r.Chunking)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to resolve chunking provider: %w", err)
|
|
}
|
|
|
|
chunking, err := factory.MakeChunking(r.Chunking.ProviderID, chunkingOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create chunking provider: %w", err)
|
|
}
|
|
options.Chunking = chunking
|
|
|
|
// Get chunking options
|
|
chunkingOpts, err := factory.ChunkingOptions(r.Chunking.ProviderID, chunkingOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get chunking options: %w", err)
|
|
}
|
|
options.ChunkingOptions = chunkingOpts
|
|
|
|
// Resolve and create embedding provider
|
|
embeddingOption, err := resolveProviderOption(r.Embedding)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to resolve embedding provider: %w", err)
|
|
}
|
|
|
|
embedding, err := factory.MakeEmbedding(r.Embedding.ProviderID, embeddingOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create embedding provider: %w", err)
|
|
}
|
|
options.Embedding = embedding
|
|
|
|
// Optional providers
|
|
if r.Extraction != nil {
|
|
extractionOption, err := resolveProviderOption(r.Extraction)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to resolve extraction provider: %w", err)
|
|
}
|
|
|
|
extraction, err := factory.MakeExtractor(r.Extraction.ProviderID, extractionOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create extraction provider: %w", err)
|
|
}
|
|
options.Extraction = extraction
|
|
}
|
|
|
|
if r.Fetcher != nil {
|
|
fetcherOption, err := resolveProviderOption(r.Fetcher)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to resolve fetcher provider: %w", err)
|
|
}
|
|
|
|
fetcher, err := factory.MakeFetcher(r.Fetcher.ProviderID, fetcherOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create fetcher provider: %w", err)
|
|
}
|
|
options.Fetcher = fetcher
|
|
}
|
|
|
|
// Handle converter - auto-detect if not specified
|
|
if r.Converter != nil {
|
|
// User specified converter
|
|
converterOption, err := resolveProviderOption(r.Converter)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to resolve converter provider: %w", err)
|
|
}
|
|
|
|
converter, err := factory.MakeConverter(r.Converter.ProviderID, converterOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create converter provider: %w", err)
|
|
}
|
|
options.Converter = converter
|
|
} else if filename != "" || contentType != "" {
|
|
// Auto-detect converter based on filename and content type
|
|
matched, converterID, err := factory.AutoDetectConverter(filename, contentType)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to auto-detect converter: %w", err)
|
|
}
|
|
|
|
if matched {
|
|
// Find the provider to get default option
|
|
converterConfig := &ProviderConfig{
|
|
ProviderID: converterID,
|
|
}
|
|
|
|
converterOption, err := resolveProviderOption(converterConfig)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to resolve auto-detected converter provider: %w", err)
|
|
}
|
|
|
|
converter, err := factory.MakeConverter(converterID, converterOption)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create auto-detected converter provider: %w", err)
|
|
}
|
|
options.Converter = converter
|
|
}
|
|
}
|
|
|
|
return options, nil
|
|
}
|
|
|
|
// Validate validates the common fields
|
|
func (r *BaseUpsertRequest) Validate() error {
|
|
if r.CollectionID == "" {
|
|
return fmt.Errorf("collection_id is required")
|
|
}
|
|
if r.Chunking == nil {
|
|
return fmt.Errorf("chunking provider is required")
|
|
}
|
|
if r.Embedding == nil {
|
|
return fmt.Errorf("embedding provider is required")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Validate validates the AddFileRequest fields
|
|
func (r *AddFileRequest) Validate() error {
|
|
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if r.FileID == "" {
|
|
return fmt.Errorf("file_id is required")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Validate validates the AddTextRequest fields
|
|
func (r *AddTextRequest) Validate() error {
|
|
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if r.Text == "" {
|
|
return fmt.Errorf("text is required")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Validate validates the AddURLRequest fields
|
|
func (r *AddURLRequest) Validate() error {
|
|
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if r.URL == "" {
|
|
return fmt.Errorf("url is required")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Validate validates the AddSegmentsRequest fields
|
|
func (r *AddSegmentsRequest) Validate() error {
|
|
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if len(r.SegmentTexts) == 0 {
|
|
return fmt.Errorf("segment_texts is required")
|
|
}
|
|
if r.DocID == "" {
|
|
return fmt.Errorf("doc_id is required for AddSegments operation")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Validate validates the UpdateSegmentsRequest fields
|
|
func (r *UpdateSegmentsRequest) Validate() error {
|
|
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
|
return err
|
|
}
|
|
if len(r.SegmentTexts) == 0 {
|
|
return fmt.Errorf("segment_texts is required")
|
|
}
|
|
return nil
|
|
}
|