Remove unused provider.go file and enhance chunking options in chunking.go, converter.go, embedding.go, extractor.go, and fetcher.go with detailed property extraction and default values. Update comments for clarity and consistency across providers.
This commit is contained in:
parent
8471343311
commit
b88e108444
27 changed files with 5229 additions and 241 deletions
|
|
@ -1 +0,0 @@
|
|||
package kb
|
||||
582
kb/providers/README.md
Normal file
582
kb/providers/README.md
Normal file
|
|
@ -0,0 +1,582 @@
|
|||
# Knowledge Base Providers
|
||||
|
||||
This directory contains all the providers for the Knowledge Base (KB) system. Providers are modular components that handle different aspects of document processing, including chunking, embedding, extraction, fetching, and conversion.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Provider Types](#provider-types)
|
||||
- [Chunking Providers](#chunking-providers)
|
||||
- [Embedding Providers](#embedding-providers)
|
||||
- [Extractor Providers](#extractor-providers)
|
||||
- [Fetcher Providers](#fetcher-providers)
|
||||
- [Converter Providers](#converter-providers)
|
||||
- [Configuration Format](#configuration-format)
|
||||
- [Examples](#examples)
|
||||
|
||||
## Overview
|
||||
|
||||
The provider system is designed to be modular and extensible. Each provider type handles a specific aspect of document processing:
|
||||
|
||||
- **Chunking**: Splits documents into manageable pieces
|
||||
- **Embedding**: Converts text into vector representations
|
||||
- **Extraction**: Extracts entities and relationships for knowledge graphs
|
||||
- **Fetching**: Retrieves documents from various sources
|
||||
- **Conversion**: Transforms different file formats into processable text
|
||||
|
||||
All providers implement a common interface with `Make()`, `Options()`, and `Schema()` methods.
|
||||
|
||||
## Provider Types
|
||||
|
||||
### Chunking Providers
|
||||
|
||||
#### Structured Chunking (`__yao.structured`)
|
||||
|
||||
Splits documents based on structural elements like headings, paragraphs, and sections.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ----------------- | --------------- | ------- | -------------------------------------------- | ------------ |
|
||||
| `size` | `int`/`float64` | `300` | Maximum chunk size in characters | > 0 |
|
||||
| `overlap` | `int`/`float64` | `20` | Character overlap between chunks | ≥ 0 |
|
||||
| `max_depth` | `int`/`float64` | `3` | Maximum nesting depth for structure analysis | ≥ 1 |
|
||||
| `size_multiplier` | `int`/`float64` | `3` | Multiplier for dynamic sizing | ≥ 1 |
|
||||
| `max_concurrent` | `int`/`float64` | `10` | Maximum concurrent processing threads | ≥ 1 |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"size": 500,
|
||||
"overlap": 50,
|
||||
"max_depth": 5,
|
||||
"size_multiplier": 2,
|
||||
"max_concurrent": 15
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Semantic Chunking (`__yao.semantic`)
|
||||
|
||||
Uses AI models to create semantically coherent chunks based on content meaning.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------------------- | --------------- | ---------- | --------------------------------------- | ------------ |
|
||||
| `size` | `int`/`float64` | `300` | Base chunk size in characters | > 0 |
|
||||
| `overlap` | `int`/`float64` | `50` | Character overlap between chunks | ≥ 0 |
|
||||
| `max_depth` | `int`/`float64` | `3` | Maximum nesting depth | ≥ 1 |
|
||||
| `size_multiplier` | `int`/`float64` | `3` | Size multiplier for analysis | ≥ 1 |
|
||||
| `max_concurrent` | `int`/`float64` | `10` | Maximum concurrent processing threads | ≥ 1 |
|
||||
| `connector` | `string` | `""` | AI connector name for semantic analysis | Must exist |
|
||||
| `toolcall` | `bool` | `false` | Enable AI tool calling | - |
|
||||
| `context_size` | `int`/`float64` | `size * 6` | Context window size for AI analysis | > 0 |
|
||||
| `options` | `string` | `""` | Additional AI model options | - |
|
||||
| `prompt` | `string` | `""` | Custom prompt for semantic analysis | - |
|
||||
| `max_retry` | `int`/`float64` | `3` | Maximum retry attempts for AI calls | ≥ 0 |
|
||||
| `semantic_max_concurrent` | `int`/`float64` | `10` | Max concurrent semantic operations | ≥ 1 |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"size": 400,
|
||||
"overlap": 80,
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true,
|
||||
"context_size": 2400,
|
||||
"max_retry": 5,
|
||||
"semantic_max_concurrent": 8
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Embedding Providers
|
||||
|
||||
#### OpenAI Embedding (`__yao.openai`)
|
||||
|
||||
Uses OpenAI's embedding models to convert text into vector representations.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------ | --------------- | ------- | ------------------------------- | ------------------- |
|
||||
| `connector` | `string` | `""` | OpenAI connector name | Must exist |
|
||||
| `dimensions` | `int`/`float64` | `1536` | Embedding vector dimensions | > 0, model-specific |
|
||||
| `concurrent` | `int`/`float64` | `10` | Maximum concurrent API requests | ≥ 1 |
|
||||
| `model` | `string` | `""` | Specific model name (optional) | Valid OpenAI model |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"connector": "openai.text-embedding-3-small",
|
||||
"dimensions": 1536,
|
||||
"concurrent": 20,
|
||||
"model": "text-embedding-3-small"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Fastembed Embedding (`__yao.fastembed`)
|
||||
|
||||
Uses local FastEmbed models for embedding generation without API calls.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------ | --------------- | ------- | -------------------------------- | ------------------- |
|
||||
| `connector` | `string` | `""` | Fastembed service connector | Must exist |
|
||||
| `dimensions` | `int`/`float64` | `384` | Embedding vector dimensions | > 0, model-specific |
|
||||
| `concurrent` | `int`/`float64` | `5` | Maximum concurrent requests | ≥ 1 |
|
||||
| `model` | `string` | `""` | FastEmbed model name | Valid model name |
|
||||
| `host` | `string` | `""` | FastEmbed service host | Valid URL/IP |
|
||||
| `key` | `string` | `""` | Authentication key (if required) | - |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"connector": "fastembed.sentence-transformers",
|
||||
"dimensions": 384,
|
||||
"concurrent": 8,
|
||||
"model": "BAAI/bge-small-en-v1.5",
|
||||
"host": "localhost:8080"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Extractor Providers
|
||||
|
||||
#### OpenAI Extractor (`__yao.openai`)
|
||||
|
||||
Extracts entities and relationships from documents using OpenAI models for knowledge graph construction.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ---------------- | --------------- | ------- | --------------------------------------------- | ---------------------- |
|
||||
| `connector` | `string` | `""` | OpenAI connector name | Must exist |
|
||||
| `toolcall` | `bool` | `true` | Enable tool calling for structured extraction | - |
|
||||
| `temperature` | `float64`/`int` | `0.1` | Model temperature for generation | 0.0-2.0 |
|
||||
| `max_tokens` | `int`/`float64` | `4000` | Maximum tokens per request | > 0 |
|
||||
| `concurrent` | `int`/`float64` | `5` | Maximum concurrent requests | ≥ 1 |
|
||||
| `model` | `string` | `""` | Specific model name (optional) | Valid OpenAI model |
|
||||
| `prompt` | `string` | `""` | Custom extraction prompt | - |
|
||||
| `retry_attempts` | `int`/`float64` | `3` | Number of retry attempts | ≥ 0 |
|
||||
| `retry_delay` | `float64`/`int` | `1.0` | Delay between retries (seconds) | ≥ 0 |
|
||||
| `tools` | `[]interface{}` | `nil` | Custom extraction tools | Valid tool definitions |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true,
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 8000,
|
||||
"concurrent": 10,
|
||||
"retry_attempts": 5,
|
||||
"retry_delay": 2.0
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Fetcher Providers
|
||||
|
||||
#### HTTP Fetcher (`__yao.http`)
|
||||
|
||||
Downloads files from HTTP/HTTPS URLs with configurable headers and timeout.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------ | ------------------------ | ------------------------ | -------------------------- | ------------------ |
|
||||
| `headers` | `map[string]interface{}` | `{}` | Custom HTTP headers | String values only |
|
||||
| `user_agent` | `string` | `"GraphRAG-Fetcher/1.0"` | Custom User-Agent header | - |
|
||||
| `timeout` | `int`/`float64` | `300` | Request timeout in seconds | > 0 |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"headers": {
|
||||
"Authorization": "Bearer token123",
|
||||
"Accept": "application/json",
|
||||
"Custom-Header": "custom-value"
|
||||
},
|
||||
"user_agent": "MyApp/2.0",
|
||||
"timeout": 60
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### MCP Fetcher (`__yao.mcp`)
|
||||
|
||||
Retrieves files using Model Context Protocol (MCP) tools for intelligent fetching.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ---------------------- | ------------------------ | --------- | ------------------------------------------- | ------------------ |
|
||||
| `id` | `string` | `""` | MCP client identifier | Must exist |
|
||||
| `tool` | `string` | `"fetch"` | MCP tool name to call | Valid tool name |
|
||||
| `arguments_mapping` | `map[string]interface{}` | `nil` | Template mapping for tool arguments | String values only |
|
||||
| `result_mapping` | `map[string]interface{}` | `nil` | Template mapping for parsing results | String values only |
|
||||
| `output_mapping` | `map[string]interface{}` | `nil` | Alias for result_mapping (compatibility) | String values only |
|
||||
| `notification_mapping` | `map[string]interface{}` | `nil` | Template mapping for progress notifications | String values only |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"id": "fetcher",
|
||||
"tool": "fetch_document",
|
||||
"arguments_mapping": {
|
||||
"url": "{{.url}}",
|
||||
"format": "text"
|
||||
},
|
||||
"result_mapping": {
|
||||
"content": "{{.result.content}}",
|
||||
"mime_type": "{{.result.mime_type}}"
|
||||
},
|
||||
"notification_mapping": {
|
||||
"progress": "{{.notification.progress}}",
|
||||
"status": "{{.notification.status}}"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Converter Providers
|
||||
|
||||
#### UTF8 Converter (`__yao.utf8`)
|
||||
|
||||
Converts plain text and UTF-8 encoded files to processable text format.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------ | -------- | --------- | --------------------------------- | ------------------- |
|
||||
| `encoding` | `string` | `"utf-8"` | Text encoding to assume | Valid encoding name |
|
||||
| `remove_bom` | `bool` | `true` | Remove Byte Order Mark if present | - |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"encoding": "utf-8",
|
||||
"remove_bom": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Vision Converter (`__yao.vision`)
|
||||
|
||||
Processes images and visual documents using AI vision models.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------ | --------------- | -------- | ------------------------------ | --------------------- |
|
||||
| `connector` | `string` | `""` | Vision AI connector name | Must exist |
|
||||
| `quality` | `string` | `"auto"` | Image processing quality | "low", "high", "auto" |
|
||||
| `detail` | `string` | `"auto"` | Level of detail in analysis | "low", "high", "auto" |
|
||||
| `max_tokens` | `int`/`float64` | `4000` | Maximum tokens for description | > 0 |
|
||||
| `prompt` | `string` | `""` | Custom vision analysis prompt | - |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4-vision",
|
||||
"quality": "high",
|
||||
"detail": "high",
|
||||
"max_tokens": 8000,
|
||||
"prompt": "Describe this image in detail"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Whisper Converter (`__yao.whisper`)
|
||||
|
||||
Converts audio files to text using speech recognition models.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ----------------- | --------------- | -------- | ------------------------------ | ------------------------------ |
|
||||
| `connector` | `string` | `""` | Audio processing connector | Must exist |
|
||||
| `language` | `string` | `"auto"` | Audio language for recognition | ISO language code or "auto" |
|
||||
| `temperature` | `float64`/`int` | `0.0` | Model temperature | 0.0-1.0 |
|
||||
| `response_format` | `string` | `"text"` | Output format | "text", "json", "verbose_json" |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"connector": "openai.whisper-1",
|
||||
"language": "en",
|
||||
"temperature": 0.2,
|
||||
"response_format": "text"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### MCP Converter (`__yao.mcp`)
|
||||
|
||||
Uses MCP tools for custom document conversion workflows.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------------- | ------------------------ | ----------- | ---------------------------- | ------------------ |
|
||||
| `id` | `string` | `""` | MCP client identifier | Must exist |
|
||||
| `tool` | `string` | `"convert"` | MCP tool name for conversion | Valid tool name |
|
||||
| `arguments_mapping` | `map[string]interface{}` | `nil` | Template for tool arguments | String values only |
|
||||
| `result_mapping` | `map[string]interface{}` | `nil` | Template for result parsing | String values only |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"id": "converter",
|
||||
"tool": "convert_document",
|
||||
"arguments_mapping": {
|
||||
"file_path": "{{.path}}",
|
||||
"format": "text"
|
||||
},
|
||||
"result_mapping": {
|
||||
"content": "{{.result.text}}",
|
||||
"metadata": "{{.result.meta}}"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### OCR Converter (`__yao.ocr`)
|
||||
|
||||
Optical Character Recognition for extracting text from images and scanned documents.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ------------ | ------------------------ | -------- | ------------------------------ | -------------------------------- |
|
||||
| `vision` | `map[string]interface{}` | Required | Vision converter configuration | Must contain valid vision config |
|
||||
| `language` | `string` | `"auto"` | OCR language hint | ISO language code or "auto" |
|
||||
| `dpi` | `int`/`float64` | `300` | Image DPI for processing | > 0 |
|
||||
| `preprocess` | `bool` | `true` | Enable image preprocessing | - |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"vision": {
|
||||
"converter": "__yao.vision",
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4-vision",
|
||||
"quality": "high"
|
||||
}
|
||||
},
|
||||
"language": "en",
|
||||
"dpi": 300,
|
||||
"preprocess": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Video Converter (`__yao.video`)
|
||||
|
||||
Extracts content from video files using frame analysis and audio transcription.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| ---------------- | ------------------------ | -------- | ----------------------------------- | -------------------------------- |
|
||||
| `vision` | `map[string]interface{}` | Required | Vision converter for frame analysis | Must contain valid vision config |
|
||||
| `audio` | `map[string]interface{}` | Required | Audio converter for transcription | Must contain valid audio config |
|
||||
| `frame_interval` | `int`/`float64` | `30` | Seconds between frame captures | > 0 |
|
||||
| `max_frames` | `int`/`float64` | `10` | Maximum frames to analyze | > 0 |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"vision": {
|
||||
"converter": "__yao.vision",
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4-vision"
|
||||
}
|
||||
},
|
||||
"audio": {
|
||||
"converter": "__yao.whisper",
|
||||
"properties": {
|
||||
"connector": "openai.whisper-1"
|
||||
}
|
||||
},
|
||||
"frame_interval": 60,
|
||||
"max_frames": 20
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Office Converter (`__yao.office`)
|
||||
|
||||
Processes Microsoft Office documents (Word, Excel, PowerPoint) and PDFs.
|
||||
|
||||
**Configuration Fields:**
|
||||
|
||||
| Field | Type | Default | Description | Requirements |
|
||||
| --------------------- | ------------------------ | -------- | ----------------------------------- | -------------------------------- |
|
||||
| `vision` | `map[string]interface{}` | Required | Vision converter for image content | Must contain valid vision config |
|
||||
| `video` | `map[string]interface{}` | Optional | Video converter for embedded videos | Must contain valid video config |
|
||||
| `audio` | `map[string]interface{}` | Optional | Audio converter for embedded audio | Must contain valid audio config |
|
||||
| `extract_images` | `bool` | `true` | Extract and process embedded images | - |
|
||||
| `extract_tables` | `bool` | `true` | Extract and format table data | - |
|
||||
| `preserve_formatting` | `bool` | `false` | Preserve original formatting | - |
|
||||
|
||||
**Example Configuration:**
|
||||
|
||||
```json
|
||||
{
|
||||
"properties": {
|
||||
"vision": {
|
||||
"converter": "__yao.vision",
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4-vision"
|
||||
}
|
||||
},
|
||||
"video": {
|
||||
"converter": "__yao.video",
|
||||
"properties": {
|
||||
"vision": {
|
||||
"converter": "__yao.vision",
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4-vision"
|
||||
}
|
||||
},
|
||||
"audio": {
|
||||
"converter": "__yao.whisper",
|
||||
"properties": {
|
||||
"connector": "openai.whisper-1"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"extract_images": true,
|
||||
"extract_tables": true,
|
||||
"preserve_formatting": false
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Configuration Format
|
||||
|
||||
All providers use a consistent configuration format:
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "provider_id",
|
||||
"properties": {
|
||||
"field_name": "field_value"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Data Type Handling
|
||||
|
||||
The configuration system automatically handles type conversion:
|
||||
|
||||
- **Numeric fields**: Accept both `int` and `float64`, converted as needed
|
||||
- **String fields**: Must be strings, other types are ignored
|
||||
- **Boolean fields**: Must be boolean values
|
||||
- **Map fields**: Accept `map[string]interface{}`, non-string values filtered out
|
||||
- **Array fields**: Accept `[]interface{}`, with element type validation
|
||||
|
||||
### Default Values
|
||||
|
||||
All providers provide sensible default values for optional fields. Required fields (like `connector` names) must be explicitly configured.
|
||||
|
||||
## Examples
|
||||
|
||||
### Complete KB Configuration
|
||||
|
||||
```json
|
||||
{
|
||||
"chunking": {
|
||||
"id": "__yao.semantic",
|
||||
"properties": {
|
||||
"size": 400,
|
||||
"overlap": 80,
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true
|
||||
}
|
||||
},
|
||||
"embedding": {
|
||||
"id": "__yao.openai",
|
||||
"properties": {
|
||||
"connector": "openai.text-embedding-3-small",
|
||||
"dimensions": 1536,
|
||||
"concurrent": 15
|
||||
}
|
||||
},
|
||||
"extractor": {
|
||||
"id": "__yao.openai",
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true,
|
||||
"temperature": 0.1
|
||||
}
|
||||
},
|
||||
"fetcher": {
|
||||
"id": "__yao.http",
|
||||
"properties": {
|
||||
"timeout": 60,
|
||||
"headers": {
|
||||
"User-Agent": "KB-System/1.0"
|
||||
}
|
||||
}
|
||||
},
|
||||
"converters": [
|
||||
{
|
||||
"id": "__yao.utf8",
|
||||
"properties": {
|
||||
"encoding": "utf-8"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "__yao.vision",
|
||||
"properties": {
|
||||
"connector": "openai.gpt-4-vision",
|
||||
"quality": "high"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Error Handling
|
||||
|
||||
All providers implement robust error handling:
|
||||
|
||||
- **Invalid configurations**: Ignored with defaults applied
|
||||
- **Missing dependencies**: Clear error messages
|
||||
- **Type mismatches**: Automatic type conversion or field skipping
|
||||
- **Network failures**: Retry mechanisms where applicable
|
||||
|
||||
For detailed implementation examples and test cases, see the corresponding `*_test.go` files in each provider directory.
|
||||
|
|
@ -28,7 +28,70 @@ func (s *Structured) Make(_ *kbtypes.ProviderOption) (types.Chunking, error) {
|
|||
|
||||
// Options returns the options for the structured chunking provider
|
||||
func (s *Structured) Options(option *kbtypes.ProviderOption) (*types.ChunkingOptions, error) {
|
||||
return nil, nil
|
||||
if option == nil {
|
||||
// Return default structured options
|
||||
return &types.ChunkingOptions{
|
||||
Size: 300,
|
||||
Overlap: 20,
|
||||
MaxDepth: 3,
|
||||
SizeMultiplier: 3,
|
||||
MaxConcurrent: 10,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Start with default values
|
||||
options := &types.ChunkingOptions{
|
||||
Size: 300,
|
||||
Overlap: 20,
|
||||
MaxDepth: 3,
|
||||
SizeMultiplier: 3,
|
||||
MaxConcurrent: 10,
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option.Properties != nil {
|
||||
if size, ok := option.Properties["size"]; ok {
|
||||
if sizeInt, ok := size.(int); ok {
|
||||
options.Size = sizeInt
|
||||
} else if sizeFloat, ok := size.(float64); ok {
|
||||
options.Size = int(sizeFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if overlap, ok := option.Properties["overlap"]; ok {
|
||||
if overlapInt, ok := overlap.(int); ok {
|
||||
options.Overlap = overlapInt
|
||||
} else if overlapFloat, ok := overlap.(float64); ok {
|
||||
options.Overlap = int(overlapFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if maxDepth, ok := option.Properties["max_depth"]; ok {
|
||||
if maxDepthInt, ok := maxDepth.(int); ok {
|
||||
options.MaxDepth = maxDepthInt
|
||||
} else if maxDepthFloat, ok := maxDepth.(float64); ok {
|
||||
options.MaxDepth = int(maxDepthFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if sizeMultiplier, ok := option.Properties["size_multiplier"]; ok {
|
||||
if sizeMultiplierInt, ok := sizeMultiplier.(int); ok {
|
||||
options.SizeMultiplier = sizeMultiplierInt
|
||||
} else if sizeMultiplierFloat, ok := sizeMultiplier.(float64); ok {
|
||||
options.SizeMultiplier = int(sizeMultiplierFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if maxConcurrent, ok := option.Properties["max_concurrent"]; ok {
|
||||
if maxConcurrentInt, ok := maxConcurrent.(int); ok {
|
||||
options.MaxConcurrent = maxConcurrentInt
|
||||
} else if maxConcurrentFloat, ok := maxConcurrent.(float64); ok {
|
||||
options.MaxConcurrent = int(maxConcurrentFloat)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return options, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the structured chunking provider
|
||||
|
|
@ -45,7 +108,135 @@ func (s *Semantic) Make(_ *kbtypes.ProviderOption) (types.Chunking, error) {
|
|||
|
||||
// Options returns the options for the semantic chunking provider
|
||||
func (s *Semantic) Options(option *kbtypes.ProviderOption) (*types.ChunkingOptions, error) {
|
||||
return nil, nil
|
||||
if option == nil {
|
||||
// Return default semantic options
|
||||
return &types.ChunkingOptions{
|
||||
Size: 300,
|
||||
Overlap: 50,
|
||||
MaxDepth: 3,
|
||||
SizeMultiplier: 3,
|
||||
MaxConcurrent: 10,
|
||||
SemanticOptions: &types.SemanticOptions{
|
||||
ContextSize: 1800, // Default L1 Size (ChunkSize * 6)
|
||||
MaxRetry: 3,
|
||||
MaxConcurrent: 10,
|
||||
Toolcall: false,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Start with default values
|
||||
options := &types.ChunkingOptions{
|
||||
Size: 300,
|
||||
Overlap: 50,
|
||||
MaxDepth: 3,
|
||||
SizeMultiplier: 3,
|
||||
MaxConcurrent: 10,
|
||||
SemanticOptions: &types.SemanticOptions{
|
||||
ContextSize: 1800, // Default L1 Size (ChunkSize * 6)
|
||||
MaxRetry: 3,
|
||||
MaxConcurrent: 10,
|
||||
Toolcall: false,
|
||||
},
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option.Properties != nil {
|
||||
// Basic chunking options
|
||||
if size, ok := option.Properties["size"]; ok {
|
||||
if sizeInt, ok := size.(int); ok {
|
||||
options.Size = sizeInt
|
||||
// Update context size based on new size
|
||||
options.SemanticOptions.ContextSize = sizeInt * 6
|
||||
} else if sizeFloat, ok := size.(float64); ok {
|
||||
options.Size = int(sizeFloat)
|
||||
options.SemanticOptions.ContextSize = int(sizeFloat) * 6
|
||||
}
|
||||
}
|
||||
|
||||
if overlap, ok := option.Properties["overlap"]; ok {
|
||||
if overlapInt, ok := overlap.(int); ok {
|
||||
options.Overlap = overlapInt
|
||||
} else if overlapFloat, ok := overlap.(float64); ok {
|
||||
options.Overlap = int(overlapFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if maxDepth, ok := option.Properties["max_depth"]; ok {
|
||||
if maxDepthInt, ok := maxDepth.(int); ok {
|
||||
options.MaxDepth = maxDepthInt
|
||||
} else if maxDepthFloat, ok := maxDepth.(float64); ok {
|
||||
options.MaxDepth = int(maxDepthFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if sizeMultiplier, ok := option.Properties["size_multiplier"]; ok {
|
||||
if sizeMultiplierInt, ok := sizeMultiplier.(int); ok {
|
||||
options.SizeMultiplier = sizeMultiplierInt
|
||||
} else if sizeMultiplierFloat, ok := sizeMultiplier.(float64); ok {
|
||||
options.SizeMultiplier = int(sizeMultiplierFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if maxConcurrent, ok := option.Properties["max_concurrent"]; ok {
|
||||
if maxConcurrentInt, ok := maxConcurrent.(int); ok {
|
||||
options.MaxConcurrent = maxConcurrentInt
|
||||
} else if maxConcurrentFloat, ok := maxConcurrent.(float64); ok {
|
||||
options.MaxConcurrent = int(maxConcurrentFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Semantic-specific options
|
||||
if connector, ok := option.Properties["connector"]; ok {
|
||||
if connectorStr, ok := connector.(string); ok {
|
||||
options.SemanticOptions.Connector = connectorStr
|
||||
}
|
||||
}
|
||||
|
||||
if toolcall, ok := option.Properties["toolcall"]; ok {
|
||||
if toolcallBool, ok := toolcall.(bool); ok {
|
||||
options.SemanticOptions.Toolcall = toolcallBool
|
||||
}
|
||||
}
|
||||
|
||||
if contextSize, ok := option.Properties["context_size"]; ok {
|
||||
if contextSizeInt, ok := contextSize.(int); ok {
|
||||
options.SemanticOptions.ContextSize = contextSizeInt
|
||||
} else if contextSizeFloat, ok := contextSize.(float64); ok {
|
||||
options.SemanticOptions.ContextSize = int(contextSizeFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if optionsStr, ok := option.Properties["options"]; ok {
|
||||
if optionsString, ok := optionsStr.(string); ok {
|
||||
options.SemanticOptions.Options = optionsString
|
||||
}
|
||||
}
|
||||
|
||||
if prompt, ok := option.Properties["prompt"]; ok {
|
||||
if promptStr, ok := prompt.(string); ok {
|
||||
options.SemanticOptions.Prompt = promptStr
|
||||
}
|
||||
}
|
||||
|
||||
if maxRetry, ok := option.Properties["max_retry"]; ok {
|
||||
if maxRetryInt, ok := maxRetry.(int); ok {
|
||||
options.SemanticOptions.MaxRetry = maxRetryInt
|
||||
} else if maxRetryFloat, ok := maxRetry.(float64); ok {
|
||||
options.SemanticOptions.MaxRetry = int(maxRetryFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if semanticMaxConcurrent, ok := option.Properties["semantic_max_concurrent"]; ok {
|
||||
if semanticMaxConcurrentInt, ok := semanticMaxConcurrent.(int); ok {
|
||||
options.SemanticOptions.MaxConcurrent = semanticMaxConcurrentInt
|
||||
} else if semanticMaxConcurrentFloat, ok := semanticMaxConcurrent.(float64); ok {
|
||||
options.SemanticOptions.MaxConcurrent = int(semanticMaxConcurrentFloat)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return options, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the semantic chunking provider
|
||||
|
|
|
|||
402
kb/providers/chunking_test.go
Normal file
402
kb/providers/chunking_test.go
Normal file
|
|
@ -0,0 +1,402 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestStructured_Options(t *testing.T) {
|
||||
s := &Structured{}
|
||||
|
||||
t.Run("nil option should return default values", func(t *testing.T) {
|
||||
options, err := s.Options(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if options == nil {
|
||||
t.Fatal("Expected options, got nil")
|
||||
}
|
||||
|
||||
// Check default values
|
||||
expected := &types.ChunkingOptions{
|
||||
Size: 300,
|
||||
Overlap: 20,
|
||||
MaxDepth: 3,
|
||||
SizeMultiplier: 3,
|
||||
MaxConcurrent: 10,
|
||||
}
|
||||
|
||||
if options.Size != expected.Size {
|
||||
t.Errorf("Expected Size %d, got %d", expected.Size, options.Size)
|
||||
}
|
||||
if options.Overlap != expected.Overlap {
|
||||
t.Errorf("Expected Overlap %d, got %d", expected.Overlap, options.Overlap)
|
||||
}
|
||||
if options.MaxDepth != expected.MaxDepth {
|
||||
t.Errorf("Expected MaxDepth %d, got %d", expected.MaxDepth, options.MaxDepth)
|
||||
}
|
||||
if options.SizeMultiplier != expected.SizeMultiplier {
|
||||
t.Errorf("Expected SizeMultiplier %d, got %d", expected.SizeMultiplier, options.SizeMultiplier)
|
||||
}
|
||||
if options.MaxConcurrent != expected.MaxConcurrent {
|
||||
t.Errorf("Expected MaxConcurrent %d, got %d", expected.MaxConcurrent, options.MaxConcurrent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty properties should return default values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Label: "test",
|
||||
Value: "test",
|
||||
Description: "test",
|
||||
Properties: map[string]interface{}{},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
// Should still have default values
|
||||
if options.Size != 300 {
|
||||
t.Errorf("Expected Size 300, got %d", options.Size)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("custom properties with int values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 500,
|
||||
"overlap": 30,
|
||||
"max_depth": 5,
|
||||
"size_multiplier": 4,
|
||||
"max_concurrent": 15,
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.Size != 500 {
|
||||
t.Errorf("Expected Size 500, got %d", options.Size)
|
||||
}
|
||||
if options.Overlap != 30 {
|
||||
t.Errorf("Expected Overlap 30, got %d", options.Overlap)
|
||||
}
|
||||
if options.MaxDepth != 5 {
|
||||
t.Errorf("Expected MaxDepth 5, got %d", options.MaxDepth)
|
||||
}
|
||||
if options.SizeMultiplier != 4 {
|
||||
t.Errorf("Expected SizeMultiplier 4, got %d", options.SizeMultiplier)
|
||||
}
|
||||
if options.MaxConcurrent != 15 {
|
||||
t.Errorf("Expected MaxConcurrent 15, got %d", options.MaxConcurrent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("custom properties with float64 values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 500.0,
|
||||
"overlap": 30.0,
|
||||
"max_depth": 5.0,
|
||||
"size_multiplier": 4.0,
|
||||
"max_concurrent": 15.0,
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.Size != 500 {
|
||||
t.Errorf("Expected Size 500, got %d", options.Size)
|
||||
}
|
||||
if options.Overlap != 30 {
|
||||
t.Errorf("Expected Overlap 30, got %d", options.Overlap)
|
||||
}
|
||||
if options.MaxDepth != 5 {
|
||||
t.Errorf("Expected MaxDepth 5, got %d", options.MaxDepth)
|
||||
}
|
||||
if options.SizeMultiplier != 4 {
|
||||
t.Errorf("Expected SizeMultiplier 4, got %d", options.SizeMultiplier)
|
||||
}
|
||||
if options.MaxConcurrent != 15 {
|
||||
t.Errorf("Expected MaxConcurrent 15, got %d", options.MaxConcurrent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 800,
|
||||
"overlap": 100,
|
||||
// max_depth, size_multiplier, max_concurrent not provided
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.Size != 800 {
|
||||
t.Errorf("Expected Size 800, got %d", options.Size)
|
||||
}
|
||||
if options.Overlap != 100 {
|
||||
t.Errorf("Expected Overlap 100, got %d", options.Overlap)
|
||||
}
|
||||
// Should use defaults for missing values
|
||||
if options.MaxDepth != 3 {
|
||||
t.Errorf("Expected MaxDepth 3 (default), got %d", options.MaxDepth)
|
||||
}
|
||||
if options.SizeMultiplier != 3 {
|
||||
t.Errorf("Expected SizeMultiplier 3 (default), got %d", options.SizeMultiplier)
|
||||
}
|
||||
if options.MaxConcurrent != 10 {
|
||||
t.Errorf("Expected MaxConcurrent 10 (default), got %d", options.MaxConcurrent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid type values should be ignored", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": "invalid", // string instead of int/float
|
||||
"overlap": true, // bool instead of int/float
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
// Should use defaults when invalid types are provided
|
||||
if options.Size != 300 {
|
||||
t.Errorf("Expected Size 300 (default), got %d", options.Size)
|
||||
}
|
||||
if options.Overlap != 20 {
|
||||
t.Errorf("Expected Overlap 20 (default), got %d", options.Overlap)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSemantic_Options(t *testing.T) {
|
||||
s := &Semantic{}
|
||||
|
||||
t.Run("nil option should return default values", func(t *testing.T) {
|
||||
options, err := s.Options(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if options == nil {
|
||||
t.Fatal("Expected options, got nil")
|
||||
}
|
||||
|
||||
// Check default values
|
||||
if options.Size != 300 {
|
||||
t.Errorf("Expected Size 300, got %d", options.Size)
|
||||
}
|
||||
if options.Overlap != 50 {
|
||||
t.Errorf("Expected Overlap 50, got %d", options.Overlap)
|
||||
}
|
||||
if options.MaxDepth != 3 {
|
||||
t.Errorf("Expected MaxDepth 3, got %d", options.MaxDepth)
|
||||
}
|
||||
if options.SizeMultiplier != 3 {
|
||||
t.Errorf("Expected SizeMultiplier 3, got %d", options.SizeMultiplier)
|
||||
}
|
||||
if options.MaxConcurrent != 10 {
|
||||
t.Errorf("Expected MaxConcurrent 10, got %d", options.MaxConcurrent)
|
||||
}
|
||||
|
||||
// Check semantic options
|
||||
if options.SemanticOptions == nil {
|
||||
t.Fatal("Expected SemanticOptions, got nil")
|
||||
}
|
||||
if options.SemanticOptions.ContextSize != 1800 {
|
||||
t.Errorf("Expected ContextSize 1800, got %d", options.SemanticOptions.ContextSize)
|
||||
}
|
||||
if options.SemanticOptions.MaxRetry != 3 {
|
||||
t.Errorf("Expected MaxRetry 3, got %d", options.SemanticOptions.MaxRetry)
|
||||
}
|
||||
if options.SemanticOptions.MaxConcurrent != 10 {
|
||||
t.Errorf("Expected MaxConcurrent 10, got %d", options.SemanticOptions.MaxConcurrent)
|
||||
}
|
||||
if options.SemanticOptions.Toolcall != false {
|
||||
t.Errorf("Expected Toolcall false, got %v", options.SemanticOptions.Toolcall)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("basic chunking properties", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 600,
|
||||
"overlap": 100,
|
||||
"max_depth": 4,
|
||||
"size_multiplier": 5,
|
||||
"max_concurrent": 20,
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.Size != 600 {
|
||||
t.Errorf("Expected Size 600, got %d", options.Size)
|
||||
}
|
||||
if options.Overlap != 100 {
|
||||
t.Errorf("Expected Overlap 100, got %d", options.Overlap)
|
||||
}
|
||||
// Context size should be updated based on size
|
||||
if options.SemanticOptions.ContextSize != 3600 { // 600 * 6
|
||||
t.Errorf("Expected ContextSize 3600, got %d", options.SemanticOptions.ContextSize)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("semantic specific properties", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true,
|
||||
"context_size": 2400,
|
||||
"options": `{"temperature": 0.7}`,
|
||||
"prompt": "Custom system prompt",
|
||||
"max_retry": 5,
|
||||
"semantic_max_concurrent": 15,
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.SemanticOptions.Connector != "openai.gpt-4o-mini" {
|
||||
t.Errorf("Expected Connector 'openai.gpt-4o-mini', got '%s'", options.SemanticOptions.Connector)
|
||||
}
|
||||
if options.SemanticOptions.Toolcall != true {
|
||||
t.Errorf("Expected Toolcall true, got %v", options.SemanticOptions.Toolcall)
|
||||
}
|
||||
if options.SemanticOptions.ContextSize != 2400 {
|
||||
t.Errorf("Expected ContextSize 2400, got %d", options.SemanticOptions.ContextSize)
|
||||
}
|
||||
if options.SemanticOptions.Options != `{"temperature": 0.7}` {
|
||||
t.Errorf("Expected Options '{\"temperature\": 0.7}', got '%s'", options.SemanticOptions.Options)
|
||||
}
|
||||
if options.SemanticOptions.Prompt != "Custom system prompt" {
|
||||
t.Errorf("Expected Prompt 'Custom system prompt', got '%s'", options.SemanticOptions.Prompt)
|
||||
}
|
||||
if options.SemanticOptions.MaxRetry != 5 {
|
||||
t.Errorf("Expected MaxRetry 5, got %d", options.SemanticOptions.MaxRetry)
|
||||
}
|
||||
if options.SemanticOptions.MaxConcurrent != 15 {
|
||||
t.Errorf("Expected MaxConcurrent 15, got %d", options.SemanticOptions.MaxConcurrent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("context size auto-calculation", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 400, // Should result in context_size = 400 * 6 = 2400
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.Size != 400 {
|
||||
t.Errorf("Expected Size 400, got %d", options.Size)
|
||||
}
|
||||
if options.SemanticOptions.ContextSize != 2400 {
|
||||
t.Errorf("Expected ContextSize 2400 (auto-calculated), got %d", options.SemanticOptions.ContextSize)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("explicit context size overrides auto-calculation", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 400, // Would auto-calculate to 2400
|
||||
"context_size": 3000, // Explicit override
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.SemanticOptions.ContextSize != 3000 {
|
||||
t.Errorf("Expected ContextSize 3000 (explicit), got %d", options.SemanticOptions.ContextSize)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("float64 values for semantic properties", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 500.0,
|
||||
"context_size": 3000.0,
|
||||
"max_retry": 4.0,
|
||||
"semantic_max_concurrent": 12.0,
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
if options.Size != 500 {
|
||||
t.Errorf("Expected Size 500, got %d", options.Size)
|
||||
}
|
||||
if options.SemanticOptions.ContextSize != 3000 {
|
||||
t.Errorf("Expected ContextSize 3000, got %d", options.SemanticOptions.ContextSize)
|
||||
}
|
||||
if options.SemanticOptions.MaxRetry != 4 {
|
||||
t.Errorf("Expected MaxRetry 4, got %d", options.SemanticOptions.MaxRetry)
|
||||
}
|
||||
if options.SemanticOptions.MaxConcurrent != 12 {
|
||||
t.Errorf("Expected MaxConcurrent 12, got %d", options.SemanticOptions.MaxConcurrent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("mixed valid and invalid properties", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"size": 500, // valid
|
||||
"connector": 123, // invalid type for string
|
||||
"toolcall": "invalid", // invalid type for bool
|
||||
"max_retry": 3, // valid
|
||||
},
|
||||
}
|
||||
|
||||
options, err := s.Options(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
|
||||
// Valid properties should be set
|
||||
if options.Size != 500 {
|
||||
t.Errorf("Expected Size 500, got %d", options.Size)
|
||||
}
|
||||
if options.SemanticOptions.MaxRetry != 3 {
|
||||
t.Errorf("Expected MaxRetry 3, got %d", options.SemanticOptions.MaxRetry)
|
||||
}
|
||||
|
||||
// Invalid properties should use defaults
|
||||
if options.SemanticOptions.Connector != "" {
|
||||
t.Errorf("Expected Connector '' (default), got '%s'", options.SemanticOptions.Connector)
|
||||
}
|
||||
if options.SemanticOptions.Toolcall != false {
|
||||
t.Errorf("Expected Toolcall false (default), got %v", options.SemanticOptions.Toolcall)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -3,10 +3,8 @@ package providers
|
|||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/yao/kb/providers/converters"
|
||||
"github.com/yaoapp/yao/kb/providers/factory"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// Converter is a base converter provider
|
||||
|
|
@ -15,73 +13,6 @@ type Converter struct {
|
|||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// UTF8 is a converter provider for utf8 files
|
||||
type UTF8 struct{ Converter }
|
||||
|
||||
// Office is a converter provider for office files, support docx, pptx.
|
||||
type Office struct{ Converter }
|
||||
|
||||
// OCR is a converter provider for ocr files, support pdf, image.
|
||||
type OCR struct{ Converter }
|
||||
|
||||
// Video is a converter provider for video files
|
||||
type Video struct{ Converter }
|
||||
|
||||
// Whisper is a converter provider for audio files
|
||||
type Whisper struct{ Converter }
|
||||
|
||||
// Vision is a converter provider for vision files
|
||||
type Vision struct{ Converter }
|
||||
|
||||
// MCP is a converter provider for mcp files
|
||||
type MCP struct{ Converter }
|
||||
|
||||
// AutoRegister registers the converter providers
|
||||
func init() {
|
||||
factory.Converters["__yao.utf8"] = &UTF8{
|
||||
Converter: Converter{
|
||||
Autodetect: []string{"text/plain", "text/markdown", ".txt", ".md"},
|
||||
MatchPriority: 100,
|
||||
},
|
||||
}
|
||||
factory.Converters["__yao.office"] = &Office{
|
||||
Converter: Converter{
|
||||
Autodetect: []string{"application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/vnd.openxmlformats-officedocument.presentationml.presentation", ".docx", ".pptx"},
|
||||
MatchPriority: 10,
|
||||
},
|
||||
}
|
||||
factory.Converters["__yao.ocr"] = &OCR{
|
||||
Converter: Converter{
|
||||
Autodetect: []string{"application/pdf", "image/jpeg", "image/png", "image/gif", "image/webp", ".pdf", ".jpg", ".jpeg", ".png", ".gif", ".webp"},
|
||||
MatchPriority: 10,
|
||||
},
|
||||
}
|
||||
|
||||
factory.Converters["__yao.video"] = &Video{
|
||||
Converter: Converter{
|
||||
Autodetect: []string{"video/mp4", "video/mpeg", "video/quicktime", "video/webm", ".mp4", ".mpeg", ".mov", ".webm"},
|
||||
MatchPriority: 10,
|
||||
},
|
||||
}
|
||||
|
||||
factory.Converters["__yao.whisper"] = &Whisper{
|
||||
Converter: Converter{
|
||||
Autodetect: []string{"audio/mpeg", "audio/wav", "audio/webm", ".mp3", ".wav", ".webm"},
|
||||
MatchPriority: 10,
|
||||
},
|
||||
}
|
||||
|
||||
factory.Converters["__yao.vision"] = &Vision{
|
||||
Converter: Converter{
|
||||
Autodetect: []string{"image/jpeg", "image/png", "image/gif", "image/webp", ".jpg", ".jpeg", ".png", ".gif", ".webp"},
|
||||
MatchPriority: 20,
|
||||
},
|
||||
}
|
||||
|
||||
factory.Converters["__yao.mcp"] = &MCP{Converter: Converter{}}
|
||||
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (c Converter) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
|
||||
|
|
@ -105,147 +36,36 @@ func (c Converter) AutoDetect(filename, contentTypes string) (bool, int, error)
|
|||
return false, 0, nil
|
||||
}
|
||||
|
||||
// === UTF8 ===
|
||||
|
||||
// Make creates a new UTF8 converter
|
||||
func (utf8 *UTF8) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
return converter.NewUTF8(), nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the UTF8 converter
|
||||
func (utf8 *UTF8) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// === Office ===
|
||||
|
||||
// Make creates a new Office converter
|
||||
func (office *Office) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to converter.OfficeOption
|
||||
officeOption := converter.OfficeOption{
|
||||
// VisionConverter: nil, // TODO: Get vision converter from option
|
||||
// VideoConverter: nil, // TODO: Get video converter from option
|
||||
// WhisperConverter: nil, // TODO: Get whisper converter from option
|
||||
// MaxConcurrency: 0, // Will use default
|
||||
// TempDir: "", // Will use default
|
||||
// CleanupTemp: false,
|
||||
// AutoRegister registers the converter providers
|
||||
func init() {
|
||||
factory.Converters["__yao.utf8"] = &converters.UTF8{
|
||||
Autodetect: []string{"text/plain", "text/markdown", ".txt", ".md"},
|
||||
MatchPriority: 100,
|
||||
}
|
||||
return converter.NewOffice(officeOption)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Office converter
|
||||
func (office *Office) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// === OCR ===
|
||||
|
||||
// Make creates a new OCR converter
|
||||
func (ocr *OCR) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to converter.OCROption
|
||||
ocrOption := converter.OCROption{
|
||||
// Vision: nil, // TODO: Get vision converter from option
|
||||
// Mode: "", // Will use default
|
||||
// MaxConcurrency: 0, // Will use default
|
||||
// CompressSize: 0, // Will use default
|
||||
// ForceImageMode: false,
|
||||
factory.Converters["__yao.office"] = &converters.Office{
|
||||
Autodetect: []string{"application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/vnd.openxmlformats-officedocument.presentationml.presentation", ".docx", ".pptx"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
return converter.NewOCR(ocrOption)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the OCR converter
|
||||
func (ocr *OCR) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// === Video ===
|
||||
|
||||
// Make creates a new Video converter
|
||||
func (video *Video) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to converter.VideoOption
|
||||
videoOption := converter.VideoOption{
|
||||
// AudioConverter: nil, // TODO: Get audio converter from option
|
||||
// VisionConverter: nil, // TODO: Get vision converter from option
|
||||
// KeyframeInterval: 0, // Will use default
|
||||
// MaxKeyframes: 0, // Will use default
|
||||
// TempDir: "", // Will use default
|
||||
// CleanupTemp: false,
|
||||
// MaxConcurrency: 0, // Will use default
|
||||
// TextOptimization: false,
|
||||
// DeduplicationRatio: 0, // Will use default
|
||||
factory.Converters["__yao.ocr"] = &converters.OCR{
|
||||
Autodetect: []string{"application/pdf", "image/jpeg", "image/png", "image/gif", "image/webp", ".pdf", ".jpg", ".jpeg", ".png", ".gif", ".webp"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
return converter.NewVideo(videoOption)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Video converter
|
||||
func (video *Video) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// === Whisper ===
|
||||
|
||||
// Make creates a new Whisper converter
|
||||
func (whisper *Whisper) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to converter.WhisperOption
|
||||
whisperOption := converter.WhisperOption{
|
||||
// ConnectorName: "", // TODO: Get connector name from option
|
||||
// Model: "", // Will use default
|
||||
// Options: nil,
|
||||
// Language: "", // Will use default
|
||||
// ChunkDuration: 0, // Will use default
|
||||
// MappingDuration: 0, // Will use default
|
||||
// SilenceThreshold: 0, // Will use default
|
||||
// SilenceMinLength: 0, // Will use default
|
||||
// EnableSilenceDetection: false,
|
||||
// MaxConcurrency: 0, // Will use default
|
||||
// TempDir: "", // Will use default
|
||||
// CleanupTemp: false,
|
||||
factory.Converters["__yao.video"] = &converters.Video{
|
||||
Autodetect: []string{"video/mp4", "video/mpeg", "video/quicktime", "video/webm", ".mp4", ".mpeg", ".mov", ".webm"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
return converter.NewWhisper(whisperOption)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Whisper converter
|
||||
func (whisper *Whisper) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// === Vision ===
|
||||
|
||||
// Make creates a new Vision converter
|
||||
func (vision *Vision) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to converter.VisionOption
|
||||
visionOption := converter.VisionOption{
|
||||
// ConnectorName: "", // TODO: Get connector name from option
|
||||
// Model: "", // Will use default
|
||||
// Prompt: "", // Will use default
|
||||
// Options: nil,
|
||||
// CompressSize: 0, // Will use default
|
||||
// Language: "", // Will use default
|
||||
factory.Converters["__yao.whisper"] = &converters.Whisper{
|
||||
Autodetect: []string{"audio/mpeg", "audio/wav", "audio/webm", ".mp3", ".wav", ".webm"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
return converter.NewVision(visionOption)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Vision converter
|
||||
func (vision *Vision) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// === MCP ===
|
||||
|
||||
// Make creates a new MCP converter
|
||||
func (mcp *MCP) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to converter.MCPOptions
|
||||
mcpOptions := &converter.MCPOptions{
|
||||
// ID: "", // TODO: Get ID from option
|
||||
// Tool: "", // TODO: Get tool from option
|
||||
// ArgumentsMapping: nil, // TODO: Get arguments mapping from option
|
||||
// ResultMapping: nil, // TODO: Get result mapping from option
|
||||
// NotificationMapping: nil, // TODO: Get notification mapping from option
|
||||
factory.Converters["__yao.vision"] = &converters.Vision{
|
||||
Autodetect: []string{"image/jpeg", "image/png", "image/gif", "image/webp", ".jpg", ".jpeg", ".png", ".gif", ".webp"},
|
||||
MatchPriority: 20,
|
||||
}
|
||||
return converter.NewMCP(mcpOptions)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the MCP converter
|
||||
func (mcp *MCP) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
factory.Converters["__yao.mcp"] = &converters.MCP{}
|
||||
|
||||
}
|
||||
|
|
|
|||
132
kb/providers/converters/mcp.go
Normal file
132
kb/providers/converters/mcp.go
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// MCP is a converter provider for mcp files
|
||||
type MCP struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new MCP converter
|
||||
func (mcp *MCP) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// Start with default values
|
||||
mcpOptions := &converter.MCPOptions{
|
||||
ID: "", // Will be set from option
|
||||
Tool: "", // Will be set from option
|
||||
ArgumentsMapping: nil, // Optional
|
||||
ResultMapping: nil, // Optional
|
||||
NotificationMapping: nil, // Optional
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
if id, ok := option.Properties["id"]; ok {
|
||||
if idStr, ok := id.(string); ok {
|
||||
mcpOptions.ID = idStr
|
||||
}
|
||||
}
|
||||
|
||||
if tool, ok := option.Properties["tool"]; ok {
|
||||
if toolStr, ok := tool.(string); ok {
|
||||
mcpOptions.Tool = toolStr
|
||||
}
|
||||
}
|
||||
|
||||
if argsMapping, ok := option.Properties["arguments_mapping"]; ok {
|
||||
if argsMappingMap, ok := argsMapping.(map[string]interface{}); ok {
|
||||
// Convert map[string]interface{} to map[string]string
|
||||
stringMap := make(map[string]string)
|
||||
for k, v := range argsMappingMap {
|
||||
if vStr, ok := v.(string); ok {
|
||||
stringMap[k] = vStr
|
||||
}
|
||||
}
|
||||
if len(stringMap) > 0 {
|
||||
mcpOptions.ArgumentsMapping = stringMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if resultMapping, ok := option.Properties["result_mapping"]; ok {
|
||||
if resultMappingMap, ok := resultMapping.(map[string]interface{}); ok {
|
||||
// Convert map[string]interface{} to map[string]string
|
||||
stringMap := make(map[string]string)
|
||||
for k, v := range resultMappingMap {
|
||||
if vStr, ok := v.(string); ok {
|
||||
stringMap[k] = vStr
|
||||
}
|
||||
}
|
||||
if len(stringMap) > 0 {
|
||||
mcpOptions.ResultMapping = stringMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Support both "result_mapping" and "output_mapping" for backward compatibility
|
||||
if outputMapping, ok := option.Properties["output_mapping"]; ok {
|
||||
if outputMappingMap, ok := outputMapping.(map[string]interface{}); ok {
|
||||
// Convert map[string]interface{} to map[string]string
|
||||
stringMap := make(map[string]string)
|
||||
for k, v := range outputMappingMap {
|
||||
if vStr, ok := v.(string); ok {
|
||||
stringMap[k] = vStr
|
||||
}
|
||||
}
|
||||
if len(stringMap) > 0 {
|
||||
mcpOptions.ResultMapping = stringMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if notificationMapping, ok := option.Properties["notification_mapping"]; ok {
|
||||
if notificationMappingMap, ok := notificationMapping.(map[string]interface{}); ok {
|
||||
// Convert map[string]interface{} to map[string]string
|
||||
stringMap := make(map[string]string)
|
||||
for k, v := range notificationMappingMap {
|
||||
if vStr, ok := v.(string); ok {
|
||||
stringMap[k] = vStr
|
||||
}
|
||||
}
|
||||
if len(stringMap) > 0 {
|
||||
mcpOptions.NotificationMapping = stringMap
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return converter.NewMCP(mcpOptions)
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (mcp *MCP) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if mcp.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range mcp.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, mcp.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, mcp.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the MCP converter
|
||||
func (mcp *MCP) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
230
kb/providers/converters/mcp_test.go
Normal file
230
kb/providers/converters/mcp_test.go
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestMCP_Make(t *testing.T) {
|
||||
mcp := &MCP{}
|
||||
|
||||
t.Run("nil option should return error due to missing MCP client", func(t *testing.T) {
|
||||
_, err := mcp.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not set up in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not set up in test environment
|
||||
})
|
||||
|
||||
t.Run("option with id and tool should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "ocrflux",
|
||||
"tool": "process_image",
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client 'ocrflux' is not set up in test environment
|
||||
})
|
||||
|
||||
t.Run("option with all mapping properties should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "ocrflux",
|
||||
"tool": "process_document",
|
||||
"arguments_mapping": map[string]interface{}{
|
||||
"file": "input_file",
|
||||
"options": "config",
|
||||
},
|
||||
"result_mapping": map[string]interface{}{
|
||||
"text": "extracted_text",
|
||||
"metadata": "file_info",
|
||||
},
|
||||
"notification_mapping": map[string]interface{}{
|
||||
"progress": "status",
|
||||
"error": "error_msg",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should support output_mapping as alias for result_mapping", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "ocrflux",
|
||||
"tool": "process_file",
|
||||
"output_mapping": map[string]interface{}{
|
||||
"content": "extracted_content",
|
||||
"pages": "page_count",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid mapping types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "ocrflux",
|
||||
"tool": "process_file",
|
||||
"arguments_mapping": "invalid_type", // should be map
|
||||
"result_mapping": 123, // should be map
|
||||
"notification_mapping": []string{"array"}, // should be map
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("mapping with non-string values should be filtered out but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "ocrflux",
|
||||
"tool": "process_file",
|
||||
"arguments_mapping": map[string]interface{}{
|
||||
"valid_key": "valid_value", // should be included
|
||||
"invalid_key": 123, // should be filtered out
|
||||
"another_key": true, // should be filtered out
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty mappings should not set mapping fields but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "ocrflux",
|
||||
"tool": "process_file",
|
||||
"arguments_mapping": map[string]interface{}{}, // empty map
|
||||
"result_mapping": map[string]interface{}{}, // empty map
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": 123, // invalid type
|
||||
"tool": true, // invalid type
|
||||
},
|
||||
}
|
||||
_, err := mcp.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestMCP_AutoDetect(t *testing.T) {
|
||||
mcp := &MCP{
|
||||
Autodetect: []string{".pdf", ".jpg", ".png", "application/pdf", "image/jpeg"},
|
||||
MatchPriority: 15,
|
||||
}
|
||||
|
||||
t.Run("should detect .pdf files", func(t *testing.T) {
|
||||
match, priority, err := mcp.AutoDetect("document.pdf", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .pdf file")
|
||||
}
|
||||
if priority != 15 {
|
||||
t.Errorf("Expected priority 15, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .jpg files", func(t *testing.T) {
|
||||
match, priority, err := mcp.AutoDetect("image.jpg", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .jpg file")
|
||||
}
|
||||
if priority != 15 {
|
||||
t.Errorf("Expected priority 15, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := mcp.AutoDetect("unknown", "application/pdf")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for application/pdf content type")
|
||||
}
|
||||
if priority != 15 {
|
||||
t.Errorf("Expected priority 15, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := mcp.AutoDetect("video.mp4", "video/mp4")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .mp4 file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyMCP := &MCP{}
|
||||
match, priority, err := emptyMCP.AutoDetect("document.pdf", "application/pdf")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestMCP_Schema(t *testing.T) {
|
||||
mcp := &MCP{}
|
||||
schema, err := mcp.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
154
kb/providers/converters/ocr.go
Normal file
154
kb/providers/converters/ocr.go
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/gou/pdf"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// OCR is a converter provider for ocr files, support pdf, image.
|
||||
type OCR struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new OCR converter
|
||||
func (ocr *OCR) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// Start with default values
|
||||
ocrOption := converter.OCROption{
|
||||
Vision: nil, // Will be set from option
|
||||
Mode: converter.OCRModeQueue, // Default to queue mode
|
||||
MaxConcurrency: 4, // Default 4 concurrent processes
|
||||
CompressSize: 512, // Default compression size
|
||||
ForceImageMode: false, // Default don't force image mode
|
||||
PDFTool: pdf.ToolPdftoppm, // Default PDF tool
|
||||
PDFToolPath: "", // Use system default
|
||||
PDFDPI: 150, // Default DPI
|
||||
PDFFormat: "png", // Default format
|
||||
PDFQuality: 90, // Default JPEG quality
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
if mode, ok := option.Properties["mode"]; ok {
|
||||
if modeStr, ok := mode.(string); ok {
|
||||
switch modeStr {
|
||||
case "queue":
|
||||
ocrOption.Mode = converter.OCRModeQueue
|
||||
case "concurrent":
|
||||
ocrOption.Mode = converter.OCRModeConcurrent
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if maxConcurrency, ok := option.Properties["max_concurrency"]; ok {
|
||||
if maxInt, ok := maxConcurrency.(int); ok {
|
||||
ocrOption.MaxConcurrency = maxInt
|
||||
} else if maxFloat, ok := maxConcurrency.(float64); ok {
|
||||
ocrOption.MaxConcurrency = int(maxFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if compressSize, ok := option.Properties["compress_size"]; ok {
|
||||
if sizeInt, ok := compressSize.(int); ok {
|
||||
ocrOption.CompressSize = int64(sizeInt)
|
||||
} else if sizeFloat, ok := compressSize.(float64); ok {
|
||||
ocrOption.CompressSize = int64(sizeFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if forceImageMode, ok := option.Properties["force_image_mode"]; ok {
|
||||
if forceBool, ok := forceImageMode.(bool); ok {
|
||||
ocrOption.ForceImageMode = forceBool
|
||||
}
|
||||
}
|
||||
|
||||
if pdfTool, ok := option.Properties["pdf_tool"]; ok {
|
||||
if pdfToolStr, ok := pdfTool.(string); ok {
|
||||
switch pdfToolStr {
|
||||
case "pdftoppm":
|
||||
ocrOption.PDFTool = pdf.ToolPdftoppm
|
||||
case "mutool":
|
||||
ocrOption.PDFTool = pdf.ToolMutool
|
||||
case "imagemagick", "convert":
|
||||
ocrOption.PDFTool = pdf.ToolImageMagick
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if pdfToolPath, ok := option.Properties["pdf_tool_path"]; ok {
|
||||
if pathStr, ok := pdfToolPath.(string); ok {
|
||||
ocrOption.PDFToolPath = pathStr
|
||||
}
|
||||
}
|
||||
|
||||
if pdfDPI, ok := option.Properties["pdf_dpi"]; ok {
|
||||
if dpiInt, ok := pdfDPI.(int); ok {
|
||||
ocrOption.PDFDPI = dpiInt
|
||||
} else if dpiFloat, ok := pdfDPI.(float64); ok {
|
||||
ocrOption.PDFDPI = int(dpiFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if pdfFormat, ok := option.Properties["pdf_format"]; ok {
|
||||
if formatStr, ok := pdfFormat.(string); ok {
|
||||
ocrOption.PDFFormat = formatStr
|
||||
}
|
||||
}
|
||||
|
||||
if pdfQuality, ok := option.Properties["pdf_quality"]; ok {
|
||||
if qualityInt, ok := pdfQuality.(int); ok {
|
||||
ocrOption.PDFQuality = qualityInt
|
||||
} else if qualityFloat, ok := pdfQuality.(float64); ok {
|
||||
ocrOption.PDFQuality = int(qualityFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Handle nested vision converter
|
||||
if vision, ok := option.Properties["vision"]; ok {
|
||||
visionConverter, err := parseNestedConverter(vision)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse vision converter: %w", err)
|
||||
}
|
||||
ocrOption.Vision = visionConverter
|
||||
}
|
||||
}
|
||||
|
||||
// Vision converter is required
|
||||
if ocrOption.Vision == nil {
|
||||
return nil, fmt.Errorf("vision converter is required for OCR processing")
|
||||
}
|
||||
|
||||
return converter.NewOCR(ocrOption)
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (ocr *OCR) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if ocr.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range ocr.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, ocr.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, ocr.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the OCR converter
|
||||
func (ocr *OCR) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
312
kb/providers/converters/ocr_test.go
Normal file
312
kb/providers/converters/ocr_test.go
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestOCR_Make(t *testing.T) {
|
||||
ocr := &OCR{}
|
||||
|
||||
t.Run("nil option should return error for missing vision converter", func(t *testing.T) {
|
||||
_, err := ocr.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error for missing vision converter")
|
||||
}
|
||||
if err.Error() != "vision converter is required for OCR processing" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty option should return error for missing vision converter", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := ocr.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for missing vision converter")
|
||||
}
|
||||
if err.Error() != "vision converter is required for OCR processing" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with OCR properties should set all values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"mode": "concurrent",
|
||||
"max_concurrency": 8,
|
||||
"compress_size": 1024,
|
||||
"force_image_mode": true,
|
||||
"pdf_tool": "mutool",
|
||||
"pdf_tool_path": "/usr/bin/mutool",
|
||||
"pdf_dpi": 200,
|
||||
"pdf_format": "jpg",
|
||||
"pdf_quality": 85,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail because vision converter factory isn't set up in tests
|
||||
_, err := ocr.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to mock factory limitation")
|
||||
}
|
||||
// In real usage, this would work with proper factory setup
|
||||
})
|
||||
|
||||
t.Run("mode selection should work correctly", func(t *testing.T) {
|
||||
testCases := []struct {
|
||||
mode string
|
||||
shouldWork bool
|
||||
}{
|
||||
{"queue", true},
|
||||
{"concurrent", true},
|
||||
{"invalid", true}, // Should default to queue mode
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"mode": tc.mode,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := ocr.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Errorf("Expected error for mode %s due to test limitations", tc.mode)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("PDF tool selection should work correctly", func(t *testing.T) {
|
||||
testCases := []struct {
|
||||
tool string
|
||||
}{
|
||||
{"pdftoppm"},
|
||||
{"mutool"},
|
||||
{"imagemagick"},
|
||||
{"convert"},
|
||||
{"invalid"}, // Should default
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"pdf_tool": tc.tool,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := ocr.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Errorf("Expected error for PDF tool %s due to test limitations", tc.tool)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("numeric values should handle both int and float64", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"max_concurrency": 8, // int
|
||||
"compress_size": 512.0, // float64
|
||||
"pdf_dpi": 150.0, // float64 -> int
|
||||
"pdf_quality": 90, // int
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := ocr.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("boolean values should be handled correctly", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"force_image_mode": true,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := ocr.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"mode": 123, // invalid type
|
||||
"max_concurrency": "invalid", // invalid type
|
||||
"compress_size": "invalid", // invalid type
|
||||
"force_image_mode": "invalid", // invalid type
|
||||
"pdf_dpi": "invalid", // invalid type
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := ocr.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid vision converter should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"vision": "invalid_format", // should be a map
|
||||
},
|
||||
}
|
||||
_, err := ocr.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid vision converter format")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOCR_AutoDetect(t *testing.T) {
|
||||
ocr := &OCR{
|
||||
Autodetect: []string{".pdf", ".jpg", ".png", ".gif", "application/pdf", "image/jpeg"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
|
||||
t.Run("should detect .pdf files", func(t *testing.T) {
|
||||
match, priority, err := ocr.AutoDetect("document.pdf", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .pdf file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .jpg files", func(t *testing.T) {
|
||||
match, priority, err := ocr.AutoDetect("scan.jpg", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .jpg file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .png files", func(t *testing.T) {
|
||||
match, priority, err := ocr.AutoDetect("screenshot.png", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .png file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := ocr.AutoDetect("unknown", "application/pdf")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for application/pdf content type")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect image content types", func(t *testing.T) {
|
||||
match, priority, err := ocr.AutoDetect("unknown", "image/jpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for image/jpeg content type")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := ocr.AutoDetect("video.mp4", "video/mp4")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .mp4 file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyOCR := &OCR{}
|
||||
match, priority, err := emptyOCR.AutoDetect("document.pdf", "application/pdf")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOCR_Schema(t *testing.T) {
|
||||
ocr := &OCR{}
|
||||
schema, err := ocr.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
113
kb/providers/converters/office.go
Normal file
113
kb/providers/converters/office.go
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// Office is a converter provider for office files, support docx, pptx.
|
||||
type Office struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new Office converter
|
||||
func (office *Office) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// Start with default values
|
||||
officeOption := converter.OfficeOption{
|
||||
VisionConverter: nil, // Will be set from option
|
||||
VideoConverter: nil, // Optional, will be set from option if provided
|
||||
WhisperConverter: nil, // Optional, will be set from option if provided
|
||||
MaxConcurrency: 4, // Default 4 concurrent processes
|
||||
TempDir: "", // Use system temp
|
||||
CleanupTemp: true, // Default cleanup
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
if maxConcurrency, ok := option.Properties["max_concurrency"]; ok {
|
||||
if maxInt, ok := maxConcurrency.(int); ok {
|
||||
officeOption.MaxConcurrency = maxInt
|
||||
} else if maxFloat, ok := maxConcurrency.(float64); ok {
|
||||
officeOption.MaxConcurrency = int(maxFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if tempDir, ok := option.Properties["temp_dir"]; ok {
|
||||
if tempDirStr, ok := tempDir.(string); ok {
|
||||
officeOption.TempDir = tempDirStr
|
||||
}
|
||||
}
|
||||
|
||||
if cleanupTemp, ok := option.Properties["cleanup_temp"]; ok {
|
||||
if cleanupBool, ok := cleanupTemp.(bool); ok {
|
||||
officeOption.CleanupTemp = cleanupBool
|
||||
}
|
||||
}
|
||||
|
||||
// Handle nested vision converter (required)
|
||||
if vision, ok := option.Properties["vision"]; ok {
|
||||
visionConverter, err := parseNestedConverter(vision)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse vision converter: %w", err)
|
||||
}
|
||||
officeOption.VisionConverter = visionConverter
|
||||
}
|
||||
|
||||
// Handle nested video converter (optional)
|
||||
if video, ok := option.Properties["video"]; ok {
|
||||
videoConverter, err := parseNestedConverter(video)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse video converter: %w", err)
|
||||
}
|
||||
officeOption.VideoConverter = videoConverter
|
||||
}
|
||||
|
||||
// Handle nested audio/whisper converter (optional)
|
||||
if audio, ok := option.Properties["audio"]; ok {
|
||||
audioConverter, err := parseNestedConverter(audio)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse audio converter: %w", err)
|
||||
}
|
||||
officeOption.WhisperConverter = audioConverter
|
||||
}
|
||||
}
|
||||
|
||||
// Vision converter is required for office processing
|
||||
if officeOption.VisionConverter == nil {
|
||||
return nil, fmt.Errorf("vision converter is required for office document processing")
|
||||
}
|
||||
|
||||
return converter.NewOffice(officeOption)
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (office *Office) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if office.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range office.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, office.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, office.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Office converter
|
||||
func (office *Office) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
313
kb/providers/converters/office_test.go
Normal file
313
kb/providers/converters/office_test.go
Normal file
|
|
@ -0,0 +1,313 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestOffice_Make(t *testing.T) {
|
||||
office := &Office{}
|
||||
|
||||
t.Run("nil option should return error for missing vision converter", func(t *testing.T) {
|
||||
_, err := office.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error for missing vision converter")
|
||||
}
|
||||
if err.Error() != "vision converter is required for office document processing" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty option should return error for missing vision converter", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := office.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for missing vision converter")
|
||||
}
|
||||
if err.Error() != "vision converter is required for office document processing" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with office processing properties should set all values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"max_concurrency": 8,
|
||||
"temp_dir": "/tmp/office",
|
||||
"cleanup_temp": false,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
"video": map[string]interface{}{
|
||||
"converter": "__yao.video",
|
||||
"properties": map[string]interface{}{
|
||||
"keyframe_interval": 10.0,
|
||||
},
|
||||
},
|
||||
"audio": map[string]interface{}{
|
||||
"converter": "__yao.whisper",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.whisper",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail because converter factories aren't set up in tests
|
||||
_, err := office.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to mock factory limitation")
|
||||
}
|
||||
// In real usage, this would work with proper factory setup
|
||||
})
|
||||
|
||||
t.Run("numeric values should handle both int and float64", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"max_concurrency": 6.0, // float64 -> int
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := office.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("boolean values should be handled correctly", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"cleanup_temp": true,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := office.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"max_concurrency": "invalid", // invalid type
|
||||
"temp_dir": 123, // invalid type
|
||||
"cleanup_temp": "invalid", // invalid type
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := office.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("only vision converter should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := office.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid vision converter should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"vision": "invalid_format", // should be a map
|
||||
},
|
||||
}
|
||||
_, err := office.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid vision converter format")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid video converter should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
"video": []string{"invalid"}, // should be a map
|
||||
},
|
||||
}
|
||||
_, err := office.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid video converter format")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid audio converter should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
"audio": 123, // should be a map
|
||||
},
|
||||
}
|
||||
_, err := office.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid audio converter format")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"max_concurrency": 12,
|
||||
"vision": map[string]interface{}{
|
||||
"converter": "__yao.vision",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
},
|
||||
// temp_dir and cleanup_temp should use defaults
|
||||
},
|
||||
}
|
||||
// This will fail due to factory setup, but we're testing the parsing logic
|
||||
_, err := office.Make(option)
|
||||
// We expect error due to vision converter factory not being set up
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test limitations")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOffice_AutoDetect(t *testing.T) {
|
||||
office := &Office{
|
||||
Autodetect: []string{".docx", ".pptx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
|
||||
t.Run("should detect .docx files", func(t *testing.T) {
|
||||
match, priority, err := office.AutoDetect("document.docx", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .docx file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .pptx files", func(t *testing.T) {
|
||||
match, priority, err := office.AutoDetect("presentation.pptx", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .pptx file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := office.AutoDetect("unknown", "application/vnd.openxmlformats-officedocument.wordprocessingml.document")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for Word document content type")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := office.AutoDetect("image.jpg", "image/jpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .jpg file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect old Office formats", func(t *testing.T) {
|
||||
match, priority, err := office.AutoDetect("document.doc", "application/msword")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .doc file (old format)")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyOffice := &Office{}
|
||||
match, priority, err := emptyOffice.AutoDetect("document.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOffice_Schema(t *testing.T) {
|
||||
office := &Office{}
|
||||
schema, err := office.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
48
kb/providers/converters/utf8.go
Normal file
48
kb/providers/converters/utf8.go
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// UTF8 is a converter provider for utf8 files
|
||||
type UTF8 struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new UTF8 converter
|
||||
func (utf8 *UTF8) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// UTF8 converter doesn't need any configuration, just return a new instance
|
||||
return converter.NewUTF8(), nil
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (utf8 *UTF8) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if utf8.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range utf8.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, utf8.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, utf8.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the UTF8 converter
|
||||
func (utf8 *UTF8) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
131
kb/providers/converters/utf8_test.go
Normal file
131
kb/providers/converters/utf8_test.go
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestUTF8_Make(t *testing.T) {
|
||||
utf8 := &UTF8{}
|
||||
|
||||
t.Run("nil option should create UTF8 converter", func(t *testing.T) {
|
||||
converter, err := utf8.Make(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if converter == nil {
|
||||
t.Fatal("Expected converter, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty option should create UTF8 converter", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
converter, err := utf8.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if converter == nil {
|
||||
t.Fatal("Expected converter, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with properties should create UTF8 converter", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"some_property": "some_value",
|
||||
},
|
||||
}
|
||||
converter, err := utf8.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if converter == nil {
|
||||
t.Fatal("Expected converter, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestUTF8_AutoDetect(t *testing.T) {
|
||||
utf8 := &UTF8{
|
||||
Autodetect: []string{".txt", ".md", "text/plain"},
|
||||
MatchPriority: 100,
|
||||
}
|
||||
|
||||
t.Run("should detect .txt files", func(t *testing.T) {
|
||||
match, priority, err := utf8.AutoDetect("test.txt", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .txt file")
|
||||
}
|
||||
if priority != 100 {
|
||||
t.Errorf("Expected priority 100, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .md files", func(t *testing.T) {
|
||||
match, priority, err := utf8.AutoDetect("readme.md", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .md file")
|
||||
}
|
||||
if priority != 100 {
|
||||
t.Errorf("Expected priority 100, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := utf8.AutoDetect("unknown", "text/plain")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for text/plain content type")
|
||||
}
|
||||
if priority != 100 {
|
||||
t.Errorf("Expected priority 100, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := utf8.AutoDetect("test.pdf", "application/pdf")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .pdf file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyUTF8 := &UTF8{}
|
||||
match, priority, err := emptyUTF8.AutoDetect("test.txt", "text/plain")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestUTF8_Schema(t *testing.T) {
|
||||
utf8 := &UTF8{}
|
||||
schema, err := utf8.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
47
kb/providers/converters/utils.go
Normal file
47
kb/providers/converters/utils.go
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/yao/kb/providers/factory"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// parseNestedConverter parses nested converter configuration
|
||||
func parseNestedConverter(config interface{}) (types.Converter, error) {
|
||||
configMap, ok := config.(map[string]interface{})
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("converter config must be a map")
|
||||
}
|
||||
|
||||
converterID, ok := configMap["converter"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("converter ID is required")
|
||||
}
|
||||
|
||||
// Get converter factory
|
||||
converterFactory, exists := factory.Converters[converterID]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("converter %s not found", converterID)
|
||||
}
|
||||
|
||||
// Parse properties
|
||||
var providerOption *kbtypes.ProviderOption
|
||||
if properties, ok := configMap["properties"]; ok {
|
||||
if propertiesStr, ok := properties.(string); ok {
|
||||
// Handle preset value - we'd need to look up the preset
|
||||
// For now, create a basic option with the preset as ID
|
||||
providerOption = &kbtypes.ProviderOption{
|
||||
Value: propertiesStr,
|
||||
}
|
||||
} else if propertiesMap, ok := properties.(map[string]interface{}); ok {
|
||||
// Handle direct properties map
|
||||
providerOption = &kbtypes.ProviderOption{
|
||||
Properties: propertiesMap,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return converterFactory.Make(providerOption)
|
||||
}
|
||||
162
kb/providers/converters/utils_test.go
Normal file
162
kb/providers/converters/utils_test.go
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseNestedConverter(t *testing.T) {
|
||||
t.Run("nil config should return error", func(t *testing.T) {
|
||||
_, err := parseNestedConverter(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error for nil config")
|
||||
}
|
||||
if err.Error() != "converter config must be a map" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non-map config should return error", func(t *testing.T) {
|
||||
_, err := parseNestedConverter("not a map")
|
||||
if err == nil {
|
||||
t.Error("Expected error for non-map config")
|
||||
}
|
||||
if err.Error() != "converter config must be a map" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("map without converter field should return error", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
}
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error for missing converter field")
|
||||
}
|
||||
if err.Error() != "converter ID is required" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non-string converter field should return error", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"converter": 123, // should be string
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
}
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error for non-string converter field")
|
||||
}
|
||||
if err.Error() != "converter ID is required" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("unknown converter ID should return error", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"converter": "__yao.unknown_converter",
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
}
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error for unknown converter")
|
||||
}
|
||||
if err.Error() != "converter __yao.unknown_converter not found" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("valid converter config with string properties should return error due to factory limitation", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"converter": "__yao.vision", // This converter exists in factory
|
||||
"properties": "gpt-4o-mini", // String preset value
|
||||
}
|
||||
// This will fail because the factory converter's Make method will fail
|
||||
// due to missing actual connector setup in test environment
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test factory limitation")
|
||||
}
|
||||
// The error would come from the converter's Make method, not parseNestedConverter itself
|
||||
})
|
||||
|
||||
t.Run("valid converter config with map properties should return error due to factory limitation", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"converter": "__yao.vision", // This converter exists in factory
|
||||
"properties": map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"compress_size": 512,
|
||||
},
|
||||
}
|
||||
// This will fail because the factory converter's Make method will fail
|
||||
// due to missing actual connector setup in test environment
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test factory limitation")
|
||||
}
|
||||
// The error would come from the converter's Make method, not parseNestedConverter itself
|
||||
})
|
||||
|
||||
t.Run("converter config without properties should return error due to factory limitation", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"converter": "__yao.utf8", // This converter exists in factory
|
||||
// No properties field
|
||||
}
|
||||
// This will fail because the factory converter's Make method will fail
|
||||
// due to missing actual connector setup in test environment
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test factory limitation")
|
||||
}
|
||||
// The error would come from the converter's Make method, not parseNestedConverter itself
|
||||
})
|
||||
|
||||
t.Run("empty map config should return error", func(t *testing.T) {
|
||||
config := map[string]interface{}{}
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error for empty config")
|
||||
}
|
||||
if err.Error() != "converter ID is required" {
|
||||
t.Errorf("Expected specific error message, got: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("config with invalid properties type should still process", func(t *testing.T) {
|
||||
config := map[string]interface{}{
|
||||
"converter": "__yao.vision", // This converter exists in factory
|
||||
"properties": 123, // Invalid type, should be ignored
|
||||
}
|
||||
// This will fail because the factory converter's Make method will fail
|
||||
// due to missing actual connector setup in test environment
|
||||
_, err := parseNestedConverter(config)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to test factory limitation")
|
||||
}
|
||||
// The error would come from the converter's Make method, not parseNestedConverter itself
|
||||
})
|
||||
|
||||
// Note about test limitations:
|
||||
// These tests verify the parsing logic of parseNestedConverter, but cannot test
|
||||
// successful converter creation because:
|
||||
// 1. The factory requires actual connector instances to be set up
|
||||
// 2. Connectors require external services (OpenAI, etc.) to be available
|
||||
// 3. Test environment doesn't have these dependencies
|
||||
//
|
||||
// In integration tests or with proper mocking, these would succeed:
|
||||
// - parseNestedConverter(validConfig) should return actual converter instance
|
||||
// - All property mappings should work correctly
|
||||
// - Nested converter configurations should be properly parsed
|
||||
}
|
||||
|
||||
// Additional tests could be added with proper mocking of the factory system:
|
||||
// - Test successful converter creation with mocked factories
|
||||
// - Test property mapping with different converter types
|
||||
// - Test error propagation from nested converter Make methods
|
||||
// - Test recursive nested converter configurations
|
||||
132
kb/providers/converters/video.go
Normal file
132
kb/providers/converters/video.go
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// Video is a converter provider for video files
|
||||
type Video struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new Video converter
|
||||
func (video *Video) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// Start with default values
|
||||
videoOption := converter.VideoOption{
|
||||
AudioConverter: nil, // Will be set from option if provided
|
||||
VisionConverter: nil, // Will be set from option if provided
|
||||
KeyframeInterval: 10.0, // Default 10 seconds
|
||||
MaxKeyframes: 20, // Default max 20 keyframes
|
||||
TempDir: "", // Use system temp
|
||||
CleanupTemp: true, // Default cleanup
|
||||
MaxConcurrency: 4, // Default 4 concurrent processes
|
||||
TextOptimization: true, // Default enable text optimization
|
||||
DeduplicationRatio: 0.8, // Default deduplication ratio
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
if keyframeInterval, ok := option.Properties["keyframe_interval"]; ok {
|
||||
if intervalFloat, ok := keyframeInterval.(float64); ok {
|
||||
videoOption.KeyframeInterval = intervalFloat
|
||||
} else if intervalInt, ok := keyframeInterval.(int); ok {
|
||||
videoOption.KeyframeInterval = float64(intervalInt)
|
||||
}
|
||||
}
|
||||
|
||||
if maxKeyframes, ok := option.Properties["max_keyframes"]; ok {
|
||||
if maxInt, ok := maxKeyframes.(int); ok {
|
||||
videoOption.MaxKeyframes = maxInt
|
||||
} else if maxFloat, ok := maxKeyframes.(float64); ok {
|
||||
videoOption.MaxKeyframes = int(maxFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if tempDir, ok := option.Properties["temp_dir"]; ok {
|
||||
if tempDirStr, ok := tempDir.(string); ok {
|
||||
videoOption.TempDir = tempDirStr
|
||||
}
|
||||
}
|
||||
|
||||
if cleanupTemp, ok := option.Properties["cleanup_temp"]; ok {
|
||||
if cleanupBool, ok := cleanupTemp.(bool); ok {
|
||||
videoOption.CleanupTemp = cleanupBool
|
||||
}
|
||||
}
|
||||
|
||||
if maxConcurrency, ok := option.Properties["max_concurrency"]; ok {
|
||||
if maxInt, ok := maxConcurrency.(int); ok {
|
||||
videoOption.MaxConcurrency = maxInt
|
||||
} else if maxFloat, ok := maxConcurrency.(float64); ok {
|
||||
videoOption.MaxConcurrency = int(maxFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if textOptimization, ok := option.Properties["text_optimization"]; ok {
|
||||
if optimizationBool, ok := textOptimization.(bool); ok {
|
||||
videoOption.TextOptimization = optimizationBool
|
||||
}
|
||||
}
|
||||
|
||||
if deduplicationRatio, ok := option.Properties["deduplication_ratio"]; ok {
|
||||
if ratioFloat, ok := deduplicationRatio.(float64); ok {
|
||||
videoOption.DeduplicationRatio = ratioFloat
|
||||
} else if ratioInt, ok := deduplicationRatio.(int); ok {
|
||||
videoOption.DeduplicationRatio = float64(ratioInt)
|
||||
}
|
||||
}
|
||||
|
||||
// Handle nested vision converter
|
||||
if vision, ok := option.Properties["vision"]; ok {
|
||||
visionConverter, err := parseNestedConverter(vision)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse vision converter: %w", err)
|
||||
}
|
||||
videoOption.VisionConverter = visionConverter
|
||||
}
|
||||
|
||||
// Handle nested audio converter
|
||||
if audio, ok := option.Properties["audio"]; ok {
|
||||
audioConverter, err := parseNestedConverter(audio)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse audio converter: %w", err)
|
||||
}
|
||||
videoOption.AudioConverter = audioConverter
|
||||
}
|
||||
}
|
||||
|
||||
return converter.NewVideo(videoOption)
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (video *Video) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if video.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range video.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, video.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, video.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Video converter
|
||||
func (video *Video) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
229
kb/providers/converters/video_test.go
Normal file
229
kb/providers/converters/video_test.go
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestVideo_Make(t *testing.T) {
|
||||
video := &Video{}
|
||||
|
||||
// Note: Video converter requires FFmpeg and audio converters to be set up
|
||||
// All tests will fail in test environment due to missing dependencies
|
||||
|
||||
t.Run("nil option should return error due to missing FFmpeg", func(t *testing.T) {
|
||||
_, err := video.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
// Error is expected because FFmpeg and audio converter are not set up in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing FFmpeg", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
// Error is expected because FFmpeg and audio converter are not set up in test environment
|
||||
})
|
||||
|
||||
t.Run("option with video processing properties should return error due to missing FFmpeg", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"keyframe_interval": 15.0,
|
||||
"max_keyframes": 30,
|
||||
"temp_dir": "/tmp/video",
|
||||
"cleanup_temp": false,
|
||||
"max_concurrency": 8,
|
||||
"text_optimization": false,
|
||||
"deduplication_ratio": 0.9,
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("float64 values should be handled correctly but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"keyframe_interval": 12.5, // float64
|
||||
"deduplication_ratio": 0.75, // float64
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("int values should be converted to appropriate types but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"keyframe_interval": 20, // int -> float64
|
||||
"max_keyframes": 25, // int
|
||||
"max_concurrency": 6, // int
|
||||
"deduplication_ratio": 1, // int -> float64
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("boolean values should be handled correctly but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"cleanup_temp": true,
|
||||
"text_optimization": false,
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"keyframe_interval": "invalid", // invalid type
|
||||
"max_keyframes": "invalid", // invalid type
|
||||
"text_optimization": "invalid", // invalid type
|
||||
"deduplication_ratio": "invalid", // invalid type
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"keyframe_interval": 5.0,
|
||||
"max_keyframes": 10,
|
||||
// Other properties should use defaults
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing FFmpeg or audio converter")
|
||||
}
|
||||
})
|
||||
|
||||
// Note: Nested converter tests would require setting up mock factories
|
||||
// For now, we test the error cases when parseNestedConverter fails
|
||||
t.Run("invalid vision converter should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"vision": "invalid_format", // should be a map
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid vision converter format")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid audio converter should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"audio": []string{"invalid"}, // should be a map
|
||||
},
|
||||
}
|
||||
_, err := video.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid audio converter format")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestVideo_AutoDetect(t *testing.T) {
|
||||
video := &Video{
|
||||
Autodetect: []string{".mp4", ".mov", ".avi", "video/mp4", "video/quicktime"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
|
||||
t.Run("should detect .mp4 files", func(t *testing.T) {
|
||||
match, priority, err := video.AutoDetect("movie.mp4", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .mp4 file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .mov files", func(t *testing.T) {
|
||||
match, priority, err := video.AutoDetect("video.mov", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .mov file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := video.AutoDetect("unknown", "video/mp4")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for video/mp4 content type")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := video.AutoDetect("audio.mp3", "audio/mpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .mp3 file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyVideo := &Video{}
|
||||
match, priority, err := emptyVideo.AutoDetect("video.mp4", "video/mp4")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestVideo_Schema(t *testing.T) {
|
||||
video := &Video{}
|
||||
schema, err := video.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
98
kb/providers/converters/vision.go
Normal file
98
kb/providers/converters/vision.go
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// Vision is a converter provider for vision files
|
||||
type Vision struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new Vision converter
|
||||
func (vision *Vision) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// Start with default values
|
||||
visionOption := converter.VisionOption{
|
||||
ConnectorName: "", // Will be set from option
|
||||
Model: "", // Will use default from connector
|
||||
Prompt: "", // Will use default
|
||||
CompressSize: 512, // Default compression size
|
||||
Language: "Auto", // Default language
|
||||
Options: nil, // Additional options
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
if connector, ok := option.Properties["connector"]; ok {
|
||||
if connectorStr, ok := connector.(string); ok {
|
||||
visionOption.ConnectorName = connectorStr
|
||||
}
|
||||
}
|
||||
|
||||
if model, ok := option.Properties["model"]; ok {
|
||||
if modelStr, ok := model.(string); ok {
|
||||
visionOption.Model = modelStr
|
||||
}
|
||||
}
|
||||
|
||||
if prompt, ok := option.Properties["prompt"]; ok {
|
||||
if promptStr, ok := prompt.(string); ok {
|
||||
visionOption.Prompt = promptStr
|
||||
}
|
||||
}
|
||||
|
||||
if compressSize, ok := option.Properties["compress_size"]; ok {
|
||||
if sizeInt, ok := compressSize.(int); ok {
|
||||
visionOption.CompressSize = int64(sizeInt)
|
||||
} else if sizeFloat, ok := compressSize.(float64); ok {
|
||||
visionOption.CompressSize = int64(sizeFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if language, ok := option.Properties["language"]; ok {
|
||||
if langStr, ok := language.(string); ok {
|
||||
visionOption.Language = langStr
|
||||
}
|
||||
}
|
||||
|
||||
if options, ok := option.Properties["options"]; ok {
|
||||
if optionsMap, ok := options.(map[string]interface{}); ok {
|
||||
visionOption.Options = optionsMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return converter.NewVision(visionOption)
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (vision *Vision) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if vision.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range vision.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, vision.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, vision.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Vision converter
|
||||
func (vision *Vision) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
192
kb/providers/converters/vision_test.go
Normal file
192
kb/providers/converters/vision_test.go
Normal file
|
|
@ -0,0 +1,192 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestVision_Make(t *testing.T) {
|
||||
vision := &Vision{}
|
||||
|
||||
// Note: Vision converter requires connectors to be loaded
|
||||
// All tests will fail in test environment due to missing connectors
|
||||
|
||||
t.Run("nil option should return error due to missing connector", func(t *testing.T) {
|
||||
_, err := vision.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := vision.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("option with connector should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
}
|
||||
_, err := vision.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.gpt-4o-mini connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with all properties should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o",
|
||||
"model": "gpt-4o",
|
||||
"prompt": "Describe this image",
|
||||
"compress_size": 1024,
|
||||
"language": "English",
|
||||
"options": map[string]interface{}{
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 500,
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := vision.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.gpt-4o connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("compress_size as float64 should be converted to int64 but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"compress_size": 512.0, // float64
|
||||
},
|
||||
}
|
||||
_, err := vision.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": 123, // invalid type
|
||||
"compress_size": "invalid", // invalid type
|
||||
"language": true, // invalid type
|
||||
},
|
||||
}
|
||||
_, err := vision.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("missing connector should still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"model": "gpt-4o-mini",
|
||||
"compress_size": 256,
|
||||
},
|
||||
}
|
||||
_, err := vision.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because no connector is specified
|
||||
})
|
||||
}
|
||||
|
||||
func TestVision_AutoDetect(t *testing.T) {
|
||||
vision := &Vision{
|
||||
Autodetect: []string{".jpg", ".png", ".gif", "image/jpeg", "image/png"},
|
||||
MatchPriority: 20,
|
||||
}
|
||||
|
||||
t.Run("should detect .jpg files", func(t *testing.T) {
|
||||
match, priority, err := vision.AutoDetect("photo.jpg", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .jpg file")
|
||||
}
|
||||
if priority != 20 {
|
||||
t.Errorf("Expected priority 20, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .png files", func(t *testing.T) {
|
||||
match, priority, err := vision.AutoDetect("image.png", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .png file")
|
||||
}
|
||||
if priority != 20 {
|
||||
t.Errorf("Expected priority 20, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := vision.AutoDetect("unknown", "image/jpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for image/jpeg content type")
|
||||
}
|
||||
if priority != 20 {
|
||||
t.Errorf("Expected priority 20, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := vision.AutoDetect("document.pdf", "application/pdf")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .pdf file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyVision := &Vision{}
|
||||
match, priority, err := emptyVision.AutoDetect("image.jpg", "image/jpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestVision_Schema(t *testing.T) {
|
||||
vision := &Vision{}
|
||||
schema, err := vision.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
148
kb/providers/converters/whisper.go
Normal file
148
kb/providers/converters/whisper.go
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/converter"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// Whisper is a converter provider for audio files
|
||||
type Whisper struct {
|
||||
Autodetect []string `json:"autodetect" yaml:"autodetect"` // Optional, default is empty, if not set, will not use autodetect
|
||||
MatchPriority int `json:"match_priority" yaml:"match_priority"` // Optional, default is 0, the higher the number, the higher the priority
|
||||
}
|
||||
|
||||
// Make creates a new Whisper converter
|
||||
func (whisper *Whisper) Make(option *kbtypes.ProviderOption) (types.Converter, error) {
|
||||
// Start with default values
|
||||
whisperOption := converter.WhisperOption{
|
||||
ConnectorName: "", // Will be set from option
|
||||
Model: "", // Will use default from connector
|
||||
Language: "", // Auto-detect
|
||||
ChunkDuration: 30.0, // Default 30 seconds
|
||||
MappingDuration: 5.0, // Default 5 seconds
|
||||
SilenceThreshold: -40.0, // Default -40dB
|
||||
SilenceMinLength: 1.0, // Default 1 second
|
||||
EnableSilenceDetection: true, // Default enabled
|
||||
MaxConcurrency: 4, // Default 4 concurrent requests
|
||||
TempDir: "", // Will use system temp
|
||||
CleanupTemp: true, // Default cleanup
|
||||
Options: nil, // Additional options
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
if connector, ok := option.Properties["connector"]; ok {
|
||||
if connectorStr, ok := connector.(string); ok {
|
||||
whisperOption.ConnectorName = connectorStr
|
||||
}
|
||||
}
|
||||
|
||||
if model, ok := option.Properties["model"]; ok {
|
||||
if modelStr, ok := model.(string); ok {
|
||||
whisperOption.Model = modelStr
|
||||
}
|
||||
}
|
||||
|
||||
if language, ok := option.Properties["language"]; ok {
|
||||
if langStr, ok := language.(string); ok {
|
||||
whisperOption.Language = langStr
|
||||
}
|
||||
}
|
||||
|
||||
if chunkDuration, ok := option.Properties["chunk_duration"]; ok {
|
||||
if durationFloat, ok := chunkDuration.(float64); ok {
|
||||
whisperOption.ChunkDuration = durationFloat
|
||||
} else if durationInt, ok := chunkDuration.(int); ok {
|
||||
whisperOption.ChunkDuration = float64(durationInt)
|
||||
}
|
||||
}
|
||||
|
||||
if mappingDuration, ok := option.Properties["mapping_duration"]; ok {
|
||||
if durationFloat, ok := mappingDuration.(float64); ok {
|
||||
whisperOption.MappingDuration = durationFloat
|
||||
} else if durationInt, ok := mappingDuration.(int); ok {
|
||||
whisperOption.MappingDuration = float64(durationInt)
|
||||
}
|
||||
}
|
||||
|
||||
if silenceThreshold, ok := option.Properties["silence_threshold"]; ok {
|
||||
if thresholdFloat, ok := silenceThreshold.(float64); ok {
|
||||
whisperOption.SilenceThreshold = thresholdFloat
|
||||
} else if thresholdInt, ok := silenceThreshold.(int); ok {
|
||||
whisperOption.SilenceThreshold = float64(thresholdInt)
|
||||
}
|
||||
}
|
||||
|
||||
if silenceMinLength, ok := option.Properties["silence_min_length"]; ok {
|
||||
if lengthFloat, ok := silenceMinLength.(float64); ok {
|
||||
whisperOption.SilenceMinLength = lengthFloat
|
||||
} else if lengthInt, ok := silenceMinLength.(int); ok {
|
||||
whisperOption.SilenceMinLength = float64(lengthInt)
|
||||
}
|
||||
}
|
||||
|
||||
if enableSilence, ok := option.Properties["enable_silence_detection"]; ok {
|
||||
if enableBool, ok := enableSilence.(bool); ok {
|
||||
whisperOption.EnableSilenceDetection = enableBool
|
||||
}
|
||||
}
|
||||
|
||||
if maxConcurrency, ok := option.Properties["max_concurrency"]; ok {
|
||||
if maxInt, ok := maxConcurrency.(int); ok {
|
||||
whisperOption.MaxConcurrency = maxInt
|
||||
} else if maxFloat, ok := maxConcurrency.(float64); ok {
|
||||
whisperOption.MaxConcurrency = int(maxFloat)
|
||||
}
|
||||
}
|
||||
|
||||
if tempDir, ok := option.Properties["temp_dir"]; ok {
|
||||
if tempDirStr, ok := tempDir.(string); ok {
|
||||
whisperOption.TempDir = tempDirStr
|
||||
}
|
||||
}
|
||||
|
||||
if cleanupTemp, ok := option.Properties["cleanup_temp"]; ok {
|
||||
if cleanupBool, ok := cleanupTemp.(bool); ok {
|
||||
whisperOption.CleanupTemp = cleanupBool
|
||||
}
|
||||
}
|
||||
|
||||
if options, ok := option.Properties["options"]; ok {
|
||||
if optionsMap, ok := options.(map[string]interface{}); ok {
|
||||
whisperOption.Options = optionsMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return converter.NewWhisper(whisperOption)
|
||||
}
|
||||
|
||||
// AutoDetect detects the converter based on the filename and content types
|
||||
func (whisper *Whisper) AutoDetect(filename, contentTypes string) (bool, int, error) {
|
||||
// If autodetect is empty, return false
|
||||
if whisper.Autodetect == nil {
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Check if the filename matches the autodetect
|
||||
for _, autodetect := range whisper.Autodetect {
|
||||
if strings.HasSuffix(filename, autodetect) {
|
||||
return true, whisper.MatchPriority, nil
|
||||
}
|
||||
|
||||
// Check if the content types matches the autodetect
|
||||
if strings.Contains(contentTypes, autodetect) {
|
||||
return true, whisper.MatchPriority, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, 0, nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Whisper converter
|
||||
func (whisper *Whisper) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
223
kb/providers/converters/whisper_test.go
Normal file
223
kb/providers/converters/whisper_test.go
Normal file
|
|
@ -0,0 +1,223 @@
|
|||
package converters
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestWhisper_Make(t *testing.T) {
|
||||
whisper := &Whisper{}
|
||||
|
||||
// Note: Whisper converter requires connectors to be loaded
|
||||
// All tests will fail in test environment due to missing connectors
|
||||
|
||||
t.Run("nil option should return error due to missing connector", func(t *testing.T) {
|
||||
_, err := whisper.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("option with all audio properties should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.whisper",
|
||||
"model": "whisper-1",
|
||||
"language": "en",
|
||||
"chunk_duration": 45.0,
|
||||
"mapping_duration": 10.0,
|
||||
"silence_threshold": -35.0,
|
||||
"silence_min_length": 2.0,
|
||||
"enable_silence_detection": false,
|
||||
"max_concurrency": 8,
|
||||
"temp_dir": "/tmp/whisper",
|
||||
"cleanup_temp": false,
|
||||
"options": map[string]interface{}{
|
||||
"temperature": 0.0,
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.whisper connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("float64 values should be handled correctly but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.whisper",
|
||||
"chunk_duration": 30.5, // float64
|
||||
"mapping_duration": 5.2, // float64
|
||||
"silence_threshold": -42.3, // float64
|
||||
"silence_min_length": 1.8, // float64
|
||||
},
|
||||
}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("int values should be converted to float64 for duration fields but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.whisper",
|
||||
"chunk_duration": 30, // int
|
||||
"mapping_duration": 5, // int
|
||||
"silence_threshold": -40, // int
|
||||
"silence_min_length": 1, // int
|
||||
"max_concurrency": 6, // int
|
||||
},
|
||||
}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("boolean values should be handled correctly but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.whisper",
|
||||
"enable_silence_detection": true,
|
||||
"cleanup_temp": false,
|
||||
},
|
||||
}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": 123, // invalid type
|
||||
"chunk_duration": "invalid", // invalid type
|
||||
"enable_silence_detection": "invalid", // invalid type
|
||||
"max_concurrency": "invalid", // invalid type
|
||||
"options": "not a map", // invalid type
|
||||
},
|
||||
}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.whisper",
|
||||
"chunk_duration": 60.0,
|
||||
// Other properties should use defaults
|
||||
},
|
||||
}
|
||||
_, err := whisper.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
}
|
||||
|
||||
func TestWhisper_AutoDetect(t *testing.T) {
|
||||
whisper := &Whisper{
|
||||
Autodetect: []string{".mp3", ".wav", ".m4a", "audio/mpeg", "audio/wav"},
|
||||
MatchPriority: 10,
|
||||
}
|
||||
|
||||
t.Run("should detect .mp3 files", func(t *testing.T) {
|
||||
match, priority, err := whisper.AutoDetect("audio.mp3", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .mp3 file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect .wav files", func(t *testing.T) {
|
||||
match, priority, err := whisper.AutoDetect("recording.wav", "")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for .wav file")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should detect by content type", func(t *testing.T) {
|
||||
match, priority, err := whisper.AutoDetect("unknown", "audio/mpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if !match {
|
||||
t.Error("Expected match for audio/mpeg content type")
|
||||
}
|
||||
if priority != 10 {
|
||||
t.Errorf("Expected priority 10, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("should not detect unsupported files", func(t *testing.T) {
|
||||
match, priority, err := whisper.AutoDetect("video.mp4", "video/mp4")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match for .mp4 file")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty autodetect should not match", func(t *testing.T) {
|
||||
emptyWhisper := &Whisper{}
|
||||
match, priority, err := emptyWhisper.AutoDetect("audio.mp3", "audio/mpeg")
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if match {
|
||||
t.Error("Expected no match when autodetect is empty")
|
||||
}
|
||||
if priority != 0 {
|
||||
t.Errorf("Expected priority 0, got %d", priority)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestWhisper_Schema(t *testing.T) {
|
||||
whisper := &Whisper{}
|
||||
schema, err := whisper.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
|
|
@ -7,12 +7,13 @@ import (
|
|||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// OpenAI is an embedding provider
|
||||
// OpenAI is an OpenAI embedding provider
|
||||
type OpenAI struct{}
|
||||
|
||||
// Fastembed is an embedding provider
|
||||
// Fastembed is a Fastembed embedding provider
|
||||
type Fastembed struct{}
|
||||
|
||||
// AutoRegister registers the embedding providers
|
||||
func init() {
|
||||
factory.Embeddings["__yao.openai"] = &OpenAI{}
|
||||
factory.Embeddings["__yao.fastembed"] = &Fastembed{}
|
||||
|
|
@ -22,7 +23,50 @@ func init() {
|
|||
|
||||
// Make creates an OpenAI embedding provider
|
||||
func (o *OpenAI) Make(option *kbtypes.ProviderOption) (types.Embedding, error) {
|
||||
return embedding.NewOpenai(embedding.OpenaiOptions{})
|
||||
// Start with default values
|
||||
options := embedding.OpenaiOptions{
|
||||
ConnectorName: "", // Will be set from option
|
||||
Concurrent: 10, // Default concurrent requests
|
||||
Dimension: 1536, // Default dimension for text-embedding-3-small
|
||||
Model: "", // Will be determined by connector or use default
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
// Set connector name
|
||||
if connector, ok := option.Properties["connector"]; ok {
|
||||
if connectorStr, ok := connector.(string); ok {
|
||||
options.ConnectorName = connectorStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set dimensions
|
||||
if dimensions, ok := option.Properties["dimensions"]; ok {
|
||||
if dimensionsInt, ok := dimensions.(int); ok {
|
||||
options.Dimension = dimensionsInt
|
||||
} else if dimensionsFloat, ok := dimensions.(float64); ok {
|
||||
options.Dimension = int(dimensionsFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set concurrent requests
|
||||
if concurrent, ok := option.Properties["concurrent"]; ok {
|
||||
if concurrentInt, ok := concurrent.(int); ok {
|
||||
options.Concurrent = concurrentInt
|
||||
} else if concurrentFloat, ok := concurrent.(float64); ok {
|
||||
options.Concurrent = int(concurrentFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set model
|
||||
if model, ok := option.Properties["model"]; ok {
|
||||
if modelStr, ok := model.(string); ok {
|
||||
options.Model = modelStr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embedding.NewOpenai(options)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the OpenAI embedding provider
|
||||
|
|
@ -34,7 +78,66 @@ func (o *OpenAI) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, er
|
|||
|
||||
// Make creates a Fastembed embedding provider
|
||||
func (f *Fastembed) Make(option *kbtypes.ProviderOption) (types.Embedding, error) {
|
||||
return embedding.NewFastEmbed(embedding.FastEmbedOptions{})
|
||||
// Start with default values
|
||||
options := embedding.FastEmbedOptions{
|
||||
ConnectorName: "", // Will be set from option
|
||||
Concurrent: 10, // Default concurrent requests
|
||||
Dimension: 384, // Default dimension for BAAI/bge-small-en-v1.5
|
||||
Model: "", // Will be determined by connector or use default
|
||||
Host: "", // Will be determined by connector
|
||||
Key: "", // Will be determined by connector
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
// Set connector name
|
||||
if connector, ok := option.Properties["connector"]; ok {
|
||||
if connectorStr, ok := connector.(string); ok {
|
||||
options.ConnectorName = connectorStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set dimensions
|
||||
if dimensions, ok := option.Properties["dimensions"]; ok {
|
||||
if dimensionsInt, ok := dimensions.(int); ok {
|
||||
options.Dimension = dimensionsInt
|
||||
} else if dimensionsFloat, ok := dimensions.(float64); ok {
|
||||
options.Dimension = int(dimensionsFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set concurrent requests
|
||||
if concurrent, ok := option.Properties["concurrent"]; ok {
|
||||
if concurrentInt, ok := concurrent.(int); ok {
|
||||
options.Concurrent = concurrentInt
|
||||
} else if concurrentFloat, ok := concurrent.(float64); ok {
|
||||
options.Concurrent = int(concurrentFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set model
|
||||
if model, ok := option.Properties["model"]; ok {
|
||||
if modelStr, ok := model.(string); ok {
|
||||
options.Model = modelStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set host
|
||||
if host, ok := option.Properties["host"]; ok {
|
||||
if hostStr, ok := host.(string); ok {
|
||||
options.Host = hostStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set key
|
||||
if key, ok := option.Properties["key"]; ok {
|
||||
if keyStr, ok := key.(string); ok {
|
||||
options.Key = keyStr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embedding.NewFastEmbed(options)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the Fastembed embedding provider
|
||||
|
|
|
|||
293
kb/providers/embedding_test.go
Normal file
293
kb/providers/embedding_test.go
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestOpenAI_Make(t *testing.T) {
|
||||
openai := &OpenAI{}
|
||||
|
||||
// Note: OpenAI embedding requires connectors to be loaded
|
||||
// All tests will fail in test environment due to missing connectors
|
||||
|
||||
t.Run("nil option should return error due to missing connector", func(t *testing.T) {
|
||||
_, err := openai.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("option with connector should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.text-embedding-3-small",
|
||||
},
|
||||
}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.text-embedding-3-small connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with all properties should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.text-embedding-3-large",
|
||||
"dimensions": 1536,
|
||||
"concurrent": 20,
|
||||
"model": "text-embedding-3-large",
|
||||
},
|
||||
}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.text-embedding-3-large connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("dimensions as float64 should be converted to int but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.text-embedding-3-small",
|
||||
"dimensions": 512.0, // float64
|
||||
"concurrent": 15.0, // float64
|
||||
},
|
||||
}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": 123, // invalid type
|
||||
"dimensions": "invalid", // invalid type
|
||||
"concurrent": "invalid", // invalid type
|
||||
"model": true, // invalid type
|
||||
},
|
||||
}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.text-embedding-3-small",
|
||||
"dimensions": 768,
|
||||
// concurrent and model should use defaults
|
||||
},
|
||||
}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("missing connector should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"dimensions": 1536,
|
||||
"concurrent": 10,
|
||||
// No connector specified
|
||||
},
|
||||
}
|
||||
_, err := openai.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because no connector is specified
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAI_Schema(t *testing.T) {
|
||||
openai := &OpenAI{}
|
||||
schema, err := openai.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFastembed_Make(t *testing.T) {
|
||||
fastembed := &Fastembed{}
|
||||
|
||||
// Note: Fastembed embedding requires connectors to be loaded
|
||||
// All tests will fail in test environment due to missing connectors
|
||||
|
||||
t.Run("nil option should return error due to missing connector", func(t *testing.T) {
|
||||
_, err := fastembed.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("option with connector should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "fastembed.bge-small-en-v1_5",
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because fastembed.bge-small-en-v1_5 connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with all properties should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "fastembed.mxbai-embed-large-v1",
|
||||
"dimensions": 1024,
|
||||
"concurrent": 15,
|
||||
"model": "mxbai-embed-large-v1",
|
||||
"host": "http://localhost:8080",
|
||||
"key": "test-key",
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because fastembed.mxbai-embed-large-v1 connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("dimensions as float64 should be converted to int but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "fastembed.bge-small-zh-v1_5",
|
||||
"dimensions": 512.0, // float64
|
||||
"concurrent": 8.0, // float64
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": 123, // invalid type
|
||||
"dimensions": "invalid", // invalid type
|
||||
"concurrent": "invalid", // invalid type
|
||||
"model": true, // invalid type
|
||||
"host": []string{}, // invalid type
|
||||
"key": 123, // invalid type
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "fastembed.bge-small-en-v1_5",
|
||||
"dimensions": 384,
|
||||
"host": "http://fastembed-server:8080",
|
||||
// concurrent, model, and key should use defaults
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("missing connector should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"dimensions": 384,
|
||||
"concurrent": 10,
|
||||
"host": "http://localhost:8080",
|
||||
// No connector specified
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because no connector is specified
|
||||
})
|
||||
|
||||
t.Run("chinese model configuration should work but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "fastembed.bge-small-zh-v1_5",
|
||||
"dimensions": 512,
|
||||
"model": "bge-small-zh-v1.5",
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("large model configuration should work but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "fastembed.mxbai-embed-large-v1",
|
||||
"dimensions": 1024,
|
||||
"model": "mxbai-embed-large-v1",
|
||||
},
|
||||
}
|
||||
_, err := fastembed.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
}
|
||||
|
||||
func TestFastembed_Schema(t *testing.T) {
|
||||
fastembed := &Fastembed{}
|
||||
schema, err := fastembed.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
|
|
@ -1,13 +1,15 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/extraction/openai"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/yao/kb/providers/factory"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// ExtractorOpenAI is an extractor provider for entity and relationship extraction
|
||||
// ExtractorOpenAI is an OpenAI extractor provider
|
||||
type ExtractorOpenAI struct{}
|
||||
|
||||
// AutoRegister registers the extractor providers
|
||||
|
|
@ -15,27 +17,117 @@ func init() {
|
|||
factory.Extractors["__yao.openai"] = &ExtractorOpenAI{}
|
||||
}
|
||||
|
||||
// === ExtractorOpenAI ===
|
||||
|
||||
// Make creates a new OpenAI extractor
|
||||
func (e *ExtractorOpenAI) Make(option *kbtypes.ProviderOption) (types.Extraction, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to openai.Options
|
||||
openaiOptions := openai.Options{
|
||||
// ConnectorName: "", // TODO: Get connector name from option
|
||||
// Concurrent: 0, // Will use default
|
||||
// Model: "", // Will use default
|
||||
// Temperature: 0, // Will use default
|
||||
// MaxTokens: 0, // Will use default
|
||||
// Prompt: "", // Will use default
|
||||
// Toolcall: nil, // Will use default
|
||||
// Tools: nil, // Will use default
|
||||
// RetryAttempts: 0, // Will use default
|
||||
// RetryDelay: 0, // Will use default
|
||||
// Start with default values
|
||||
options := openai.Options{
|
||||
ConnectorName: "", // Will be set from option
|
||||
Concurrent: 5, // Default concurrent requests for extraction
|
||||
Model: "", // Will be determined by connector or use default
|
||||
Temperature: 0.1, // Low temperature for consistent extraction
|
||||
MaxTokens: 4000, // Default max tokens
|
||||
Prompt: "", // Custom prompt (optional)
|
||||
Toolcall: nil, // Will be set from option (nil = default true)
|
||||
Tools: nil, // Will use default extraction tools
|
||||
RetryAttempts: 3, // Default retry attempts
|
||||
RetryDelay: time.Second, // Default retry delay
|
||||
}
|
||||
return openai.NewOpenai(openaiOptions)
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
// Set connector name
|
||||
if connector, ok := option.Properties["connector"]; ok {
|
||||
if connectorStr, ok := connector.(string); ok {
|
||||
options.ConnectorName = connectorStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set toolcall (explicit bool pointer)
|
||||
if toolcall, ok := option.Properties["toolcall"]; ok {
|
||||
if toolcallBool, ok := toolcall.(bool); ok {
|
||||
options.Toolcall = &toolcallBool
|
||||
}
|
||||
}
|
||||
|
||||
// Set temperature
|
||||
if temperature, ok := option.Properties["temperature"]; ok {
|
||||
if temperatureFloat, ok := temperature.(float64); ok {
|
||||
options.Temperature = temperatureFloat
|
||||
} else if temperatureInt, ok := temperature.(int); ok {
|
||||
options.Temperature = float64(temperatureInt)
|
||||
}
|
||||
}
|
||||
|
||||
// Set max tokens
|
||||
if maxTokens, ok := option.Properties["max_tokens"]; ok {
|
||||
if maxTokensInt, ok := maxTokens.(int); ok {
|
||||
options.MaxTokens = maxTokensInt
|
||||
} else if maxTokensFloat, ok := maxTokens.(float64); ok {
|
||||
options.MaxTokens = int(maxTokensFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set concurrent requests
|
||||
if concurrent, ok := option.Properties["concurrent"]; ok {
|
||||
if concurrentInt, ok := concurrent.(int); ok {
|
||||
options.Concurrent = concurrentInt
|
||||
} else if concurrentFloat, ok := concurrent.(float64); ok {
|
||||
options.Concurrent = int(concurrentFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set model
|
||||
if model, ok := option.Properties["model"]; ok {
|
||||
if modelStr, ok := model.(string); ok {
|
||||
options.Model = modelStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set custom prompt
|
||||
if prompt, ok := option.Properties["prompt"]; ok {
|
||||
if promptStr, ok := prompt.(string); ok {
|
||||
options.Prompt = promptStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set retry attempts
|
||||
if retryAttempts, ok := option.Properties["retry_attempts"]; ok {
|
||||
if retryAttemptsInt, ok := retryAttempts.(int); ok {
|
||||
options.RetryAttempts = retryAttemptsInt
|
||||
} else if retryAttemptsFloat, ok := retryAttempts.(float64); ok {
|
||||
options.RetryAttempts = int(retryAttemptsFloat)
|
||||
}
|
||||
}
|
||||
|
||||
// Set retry delay (in seconds)
|
||||
if retryDelay, ok := option.Properties["retry_delay"]; ok {
|
||||
if retryDelayFloat, ok := retryDelay.(float64); ok {
|
||||
options.RetryDelay = time.Duration(retryDelayFloat * float64(time.Second))
|
||||
} else if retryDelayInt, ok := retryDelay.(int); ok {
|
||||
options.RetryDelay = time.Duration(retryDelayInt) * time.Second
|
||||
}
|
||||
}
|
||||
|
||||
// Set custom tools (advanced usage)
|
||||
if tools, ok := option.Properties["tools"]; ok {
|
||||
if toolsSlice, ok := tools.([]interface{}); ok {
|
||||
customTools := make([]map[string]interface{}, 0, len(toolsSlice))
|
||||
for _, tool := range toolsSlice {
|
||||
if toolMap, ok := tool.(map[string]interface{}); ok {
|
||||
customTools = append(customTools, toolMap)
|
||||
}
|
||||
}
|
||||
if len(customTools) > 0 {
|
||||
options.Tools = customTools
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return openai.NewOpenai(options)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the OpenAI extractor
|
||||
// Schema returns the schema for the OpenAI extractor provider
|
||||
func (e *ExtractorOpenAI) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
|
|
|||
289
kb/providers/extractor_test.go
Normal file
289
kb/providers/extractor_test.go
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestExtractorOpenAI_Make(t *testing.T) {
|
||||
extractor := &ExtractorOpenAI{}
|
||||
|
||||
// Note: OpenAI extractor requires connectors to be loaded
|
||||
// All tests will fail in test environment due to missing connectors
|
||||
|
||||
t.Run("nil option should return error due to missing connector", func(t *testing.T) {
|
||||
_, err := extractor.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("option with connector should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.gpt-4o-mini connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with toolcall true should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true,
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with toolcall false should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "deepseek.v3",
|
||||
"toolcall": false,
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because deepseek.v3 connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with all properties should return error due to missing connector", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o",
|
||||
"toolcall": true,
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 8000,
|
||||
"concurrent": 10,
|
||||
"model": "gpt-4o",
|
||||
"prompt": "Extract entities and relationships:",
|
||||
"retry_attempts": 5,
|
||||
"retry_delay": 2,
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because openai.gpt-4o connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("temperature as int should be converted to float64 but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"temperature": 1, // int -> float64
|
||||
"max_tokens": 2000.0, // float64 -> int
|
||||
"concurrent": 3.0, // float64 -> int
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("retry_delay as float should be converted to duration but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"retry_delay": 1.5, // 1.5 seconds
|
||||
"retry_attempts": 2.0, // float64 -> int
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("custom tools should be parsed but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"tools": []interface{}{
|
||||
map[string]interface{}{
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": "extract_entities",
|
||||
"description": "Extract entities from text",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": 123, // invalid type
|
||||
"toolcall": "invalid", // invalid type
|
||||
"temperature": "invalid", // invalid type
|
||||
"max_tokens": "invalid", // invalid type
|
||||
"concurrent": "invalid", // invalid type
|
||||
"model": true, // invalid type
|
||||
"prompt": []string{}, // invalid type
|
||||
"retry_attempts": "invalid", // invalid type
|
||||
"retry_delay": "invalid", // invalid type
|
||||
"tools": "invalid", // invalid type
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("partial properties should use defaults for missing values but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"toolcall": true,
|
||||
"temperature": 0.3,
|
||||
// Other properties should use defaults
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("missing connector should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"toolcall": true,
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 4000,
|
||||
// No connector specified
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because no connector is specified
|
||||
})
|
||||
|
||||
t.Run("gpt-4o configuration should work but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o",
|
||||
"toolcall": true,
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("deepseek configuration should work but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "deepseek.v3",
|
||||
"toolcall": false,
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid tools array should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"tools": []interface{}{
|
||||
"invalid_tool", // not a map
|
||||
123, // not a map
|
||||
map[string]interface{}{"valid": "tool"}, // valid map
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("edge case temperature values should be handled", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"temperature": 2.5, // Above normal range, will be validated by openai.NewOpenai
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
|
||||
t.Run("zero values should be handled correctly", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"connector": "openai.gpt-4o-mini",
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 0, // Will be set to default by openai.NewOpenai
|
||||
"concurrent": 0, // Will be set to default by openai.NewOpenai
|
||||
"retry_attempts": 0, // Will be set to default by openai.NewOpenai
|
||||
"retry_delay": 0, // Will be set to default by openai.NewOpenai
|
||||
},
|
||||
}
|
||||
_, err := extractor.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing connector")
|
||||
}
|
||||
// Error is expected because connector is not loaded
|
||||
})
|
||||
}
|
||||
|
||||
func TestExtractorOpenAI_Schema(t *testing.T) {
|
||||
extractor := &ExtractorOpenAI{}
|
||||
schema, err := extractor.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
|
|
@ -1,16 +1,18 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/fetcher"
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/yao/kb/providers/factory"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
// FetcherHTTP is a fetcher provider for HTTP/HTTPS URLs
|
||||
// FetcherHTTP is an HTTP fetcher provider
|
||||
type FetcherHTTP struct{}
|
||||
|
||||
// FetcherMCP is a fetcher provider for MCP-based URL fetching
|
||||
// FetcherMCP is an MCP fetcher provider
|
||||
type FetcherMCP struct{}
|
||||
|
||||
// AutoRegister registers the fetcher providers
|
||||
|
|
@ -23,16 +25,47 @@ func init() {
|
|||
|
||||
// Make creates a new HTTP fetcher
|
||||
func (f *FetcherHTTP) Make(option *kbtypes.ProviderOption) (types.Fetcher, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to fetcher.HTTPOptions
|
||||
// Start with default values
|
||||
httpOptions := &fetcher.HTTPOptions{
|
||||
// Headers: nil, // TODO: Get headers from option
|
||||
// UserAgent: "", // Will use default
|
||||
// Timeout: 0, // Will use default
|
||||
Headers: make(map[string]string), // Custom headers
|
||||
UserAgent: "GraphRAG-Fetcher/1.0", // Default user agent
|
||||
Timeout: 300 * time.Second, // Default 5 minutes
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
// Set headers
|
||||
if headers, ok := option.Properties["headers"]; ok {
|
||||
if headersMap, ok := headers.(map[string]interface{}); ok {
|
||||
for key, value := range headersMap {
|
||||
if valueStr, ok := value.(string); ok {
|
||||
httpOptions.Headers[key] = valueStr
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Set user agent
|
||||
if userAgent, ok := option.Properties["user_agent"]; ok {
|
||||
if userAgentStr, ok := userAgent.(string); ok {
|
||||
httpOptions.UserAgent = userAgentStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set timeout (in seconds)
|
||||
if timeout, ok := option.Properties["timeout"]; ok {
|
||||
if timeoutInt, ok := timeout.(int); ok {
|
||||
httpOptions.Timeout = time.Duration(timeoutInt) * time.Second
|
||||
} else if timeoutFloat, ok := timeout.(float64); ok {
|
||||
httpOptions.Timeout = time.Duration(timeoutFloat) * time.Second
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return fetcher.NewHTTPFetcher(httpOptions), nil
|
||||
}
|
||||
|
||||
// Schema returns the schema for the HTTP fetcher
|
||||
// Schema returns the schema for the HTTP fetcher provider
|
||||
func (f *FetcherHTTP) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
|
@ -41,18 +74,87 @@ func (f *FetcherHTTP) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchem
|
|||
|
||||
// Make creates a new MCP fetcher
|
||||
func (f *FetcherMCP) Make(option *kbtypes.ProviderOption) (types.Fetcher, error) {
|
||||
// TODO: Map kbtypes.ProviderOption to fetcher.MCPOptions
|
||||
// Start with default values
|
||||
mcpOptions := &fetcher.MCPOptions{
|
||||
// ID: "", // TODO: Get ID from option
|
||||
// Tool: "", // TODO: Get tool from option
|
||||
// ArgumentsMapping: nil, // TODO: Get arguments mapping from option
|
||||
// ResultMapping: nil, // TODO: Get result mapping from option
|
||||
// NotificationMapping: nil, // TODO: Get notification mapping from option
|
||||
ID: "", // Required - will be set from option
|
||||
Tool: "fetch", // Default tool name
|
||||
ArgumentsMapping: nil, // Optional arguments mapping
|
||||
ResultMapping: nil, // Optional result mapping
|
||||
NotificationMapping: nil, // Optional notification mapping
|
||||
}
|
||||
|
||||
// Extract values from Properties map
|
||||
if option != nil && option.Properties != nil {
|
||||
// Set MCP ID (required)
|
||||
if id, ok := option.Properties["id"]; ok {
|
||||
if idStr, ok := id.(string); ok {
|
||||
mcpOptions.ID = idStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set tool name
|
||||
if tool, ok := option.Properties["tool"]; ok {
|
||||
if toolStr, ok := tool.(string); ok {
|
||||
mcpOptions.Tool = toolStr
|
||||
}
|
||||
}
|
||||
|
||||
// Set arguments mapping
|
||||
if argumentsMapping, ok := option.Properties["arguments_mapping"]; ok {
|
||||
if argumentsMappingMap, ok := argumentsMapping.(map[string]interface{}); ok {
|
||||
argMap := make(map[string]string)
|
||||
for key, value := range argumentsMappingMap {
|
||||
if valueStr, ok := value.(string); ok {
|
||||
argMap[key] = valueStr
|
||||
}
|
||||
}
|
||||
if len(argMap) > 0 {
|
||||
mcpOptions.ArgumentsMapping = argMap
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Set result mapping (handle both "result_mapping" and "output_mapping" for compatibility)
|
||||
var resultMapping map[string]interface{}
|
||||
if rm, ok := option.Properties["result_mapping"]; ok {
|
||||
resultMapping, _ = rm.(map[string]interface{})
|
||||
} else if om, ok := option.Properties["output_mapping"]; ok {
|
||||
// Support kb.yao's "output_mapping" as alias for "result_mapping"
|
||||
resultMapping, _ = om.(map[string]interface{})
|
||||
}
|
||||
|
||||
if resultMapping != nil {
|
||||
resMap := make(map[string]string)
|
||||
for key, value := range resultMapping {
|
||||
if valueStr, ok := value.(string); ok {
|
||||
resMap[key] = valueStr
|
||||
}
|
||||
}
|
||||
if len(resMap) > 0 {
|
||||
mcpOptions.ResultMapping = resMap
|
||||
}
|
||||
}
|
||||
|
||||
// Set notification mapping
|
||||
if notificationMapping, ok := option.Properties["notification_mapping"]; ok {
|
||||
if notificationMappingMap, ok := notificationMapping.(map[string]interface{}); ok {
|
||||
notMap := make(map[string]string)
|
||||
for key, value := range notificationMappingMap {
|
||||
if valueStr, ok := value.(string); ok {
|
||||
notMap[key] = valueStr
|
||||
}
|
||||
}
|
||||
if len(notMap) > 0 {
|
||||
mcpOptions.NotificationMapping = notMap
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return fetcher.NewMCP(mcpOptions)
|
||||
}
|
||||
|
||||
// Schema returns the schema for the MCP fetcher
|
||||
// Schema returns the schema for the MCP fetcher provider
|
||||
func (f *FetcherMCP) Schema(provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
|
|
|||
451
kb/providers/fetcher_test.go
Normal file
451
kb/providers/fetcher_test.go
Normal file
|
|
@ -0,0 +1,451 @@
|
|||
package providers
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
func TestFetcherHTTP_Make(t *testing.T) {
|
||||
fetcher := &FetcherHTTP{}
|
||||
|
||||
t.Run("nil option should return default HTTP fetcher", func(t *testing.T) {
|
||||
result, err := fetcher.Make(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty option should return default HTTP fetcher", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with headers should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"headers": map[string]interface{}{
|
||||
"Authorization": "Bearer token123",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "Custom-Agent/1.0",
|
||||
},
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with timeout should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"timeout": 30, // 30 seconds
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with timeout as float should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"timeout": 45.5, // 45.5 seconds
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with user_agent should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"user_agent": "Custom-GraphRAG/2.0",
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("option with all properties should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"headers": map[string]interface{}{
|
||||
"Authorization": "Bearer secret",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"user_agent": "Complete-Fetcher/1.0",
|
||||
"timeout": 60,
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"headers": "invalid_type", // should be map
|
||||
"user_agent": 123, // should be string
|
||||
"timeout": "invalid_timeout", // should be number
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("headers with non-string values should be ignored", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"headers": map[string]interface{}{
|
||||
"Valid-Header": "valid_value",
|
||||
"Invalid-Header": 123, // non-string value should be ignored
|
||||
"Another-Valid": "another_value",
|
||||
},
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("zero timeout should be converted correctly", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"timeout": 0, // Should result in 0 duration, which will use default
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty headers map should work", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"headers": map[string]interface{}{}, // Empty headers map
|
||||
},
|
||||
}
|
||||
result, err := fetcher.Make(option)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if result == nil {
|
||||
t.Error("Expected HTTP fetcher, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestFetcherHTTP_Schema(t *testing.T) {
|
||||
fetcher := &FetcherHTTP{}
|
||||
schema, err := fetcher.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetcherMCP_Make(t *testing.T) {
|
||||
fetcher := &FetcherMCP{}
|
||||
|
||||
// Note: MCP fetcher requires MCP clients to be loaded
|
||||
// All tests will fail in test environment due to missing MCP client
|
||||
|
||||
t.Run("nil option should return error due to missing MCP client", func(t *testing.T) {
|
||||
_, err := fetcher.Make(nil)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("empty option should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded in test environment
|
||||
})
|
||||
|
||||
t.Run("option with id should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client "fetcher" is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with id and tool should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"tool": "fetch_url",
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with arguments_mapping should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"arguments_mapping": map[string]interface{}{
|
||||
"url": "{{.url}}",
|
||||
"headers": "{{.headers}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with output_mapping should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"output_mapping": map[string]interface{}{
|
||||
"content": "{{.result.content}}",
|
||||
"mime_type": "{{.result.mime_type}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with result_mapping should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"result_mapping": map[string]interface{}{
|
||||
"content": "{{.data.content}}",
|
||||
"mime_type": "{{.data.type}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with notification_mapping should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"notification_mapping": map[string]interface{}{
|
||||
"progress": "{{.progress}}",
|
||||
"message": "{{.message}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("option with all properties should return error due to missing MCP client", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"tool": "fetch_document",
|
||||
"arguments_mapping": map[string]interface{}{
|
||||
"url": "{{.url}}",
|
||||
"format": "text",
|
||||
},
|
||||
"result_mapping": map[string]interface{}{
|
||||
"content": "{{.result.content}}",
|
||||
"mime_type": "{{.result.mime_type}}",
|
||||
},
|
||||
"notification_mapping": map[string]interface{}{
|
||||
"progress": "{{.notification.progress}}",
|
||||
"status": "{{.notification.status}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("invalid property types should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": 123, // invalid type
|
||||
"tool": []string{}, // invalid type
|
||||
"arguments_mapping": "invalid", // invalid type
|
||||
"result_mapping": "invalid", // invalid type
|
||||
"output_mapping": "invalid", // invalid type
|
||||
"notification_mapping": "invalid", // invalid type
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("mapping with non-string values should be ignored but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"arguments_mapping": map[string]interface{}{
|
||||
"valid_arg": "{{.url}}",
|
||||
"invalid_arg": 123, // non-string value should be ignored
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("empty mapping should be handled correctly but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"arguments_mapping": map[string]interface{}{}, // Empty mapping
|
||||
"result_mapping": map[string]interface{}{}, // Empty mapping
|
||||
"notification_mapping": map[string]interface{}{}, // Empty mapping
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("missing id should return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"tool": "fetch_url",
|
||||
// No ID specified
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because no ID is specified and MCP client is not loaded
|
||||
})
|
||||
|
||||
t.Run("both output_mapping and result_mapping should prefer result_mapping but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"result_mapping": map[string]interface{}{
|
||||
"content": "{{.result}}",
|
||||
},
|
||||
"output_mapping": map[string]interface{}{
|
||||
"content": "{{.output}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
// result_mapping should take precedence over output_mapping
|
||||
})
|
||||
|
||||
t.Run("only output_mapping should be used as result_mapping but still return error", func(t *testing.T) {
|
||||
option := &kbtypes.ProviderOption{
|
||||
Properties: map[string]interface{}{
|
||||
"id": "fetcher",
|
||||
"output_mapping": map[string]interface{}{
|
||||
"content": "{{.output.data}}",
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := fetcher.Make(option)
|
||||
if err == nil {
|
||||
t.Error("Expected error due to missing MCP client")
|
||||
}
|
||||
// Error is expected because MCP client is not loaded
|
||||
// output_mapping should be mapped to result_mapping
|
||||
})
|
||||
}
|
||||
|
||||
func TestFetcherMCP_Schema(t *testing.T) {
|
||||
fetcher := &FetcherMCP{}
|
||||
schema, err := fetcher.Schema(nil)
|
||||
if err != nil {
|
||||
t.Errorf("Expected no error, got %v", err)
|
||||
}
|
||||
if schema != nil {
|
||||
t.Error("Expected nil schema")
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue