Merge pull request #1048 from trheyi/main
Remove unused provider.go file and enhance chunking options in chunki…
This commit is contained in:
Max
GitHub
commit
c301784d89
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