Merge pull request #163 from ggozad/feat/model-customizations
Support for per-model configuration settings such as thinking, temperature, max_tokens
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
b9f608a039
44 changed files with 1550 additions and 1036 deletions
|
|
@ -3,6 +3,10 @@
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|
||||||
|
- **Model Customization**: Added support for per-model configuration settings
|
||||||
|
- New `enable_thinking` parameter to control reasoning behavior (true/false/None)
|
||||||
|
- Support for `temperature` and `max_tokens` settings on QA and research models
|
||||||
|
- All settings apply to any provider that supports them
|
||||||
- **Database Inspector**: New `inspect` CLI command launches interactive TUI for browsing documents and chunks & searching
|
- **Database Inspector**: New `inspect` CLI command launches interactive TUI for browsing documents and chunks & searching
|
||||||
- **Evaluations**: Added `evaluations` CLI script for running benchmarks (replaces `python -m evaluations.benchmark`)
|
- **Evaluations**: Added `evaluations` CLI script for running benchmarks (replaces `python -m evaluations.benchmark`)
|
||||||
- **Evaluations**: Added `--db` option to override evaluation database path
|
- **Evaluations**: Added `--db` option to override evaluation database path
|
||||||
|
|
|
||||||
|
|
@ -30,8 +30,9 @@ embeddings:
|
||||||
vector_dim: 768
|
vector_dim: 768
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
model:
|
||||||
provider: ollama
|
provider: ollama
|
||||||
model: qwen3
|
name: qwen3
|
||||||
```
|
```
|
||||||
|
|
||||||
See [Configuration docs](https://ggozad.github.io/haiku.rag/configuration/) for all available options.
|
See [Configuration docs](https://ggozad.github.io/haiku.rag/configuration/) for all available options.
|
||||||
|
|
|
||||||
|
|
@ -69,7 +69,7 @@ haiku-rag add-src /path/to/documents/
|
||||||
```
|
```
|
||||||
|
|
||||||
!!! note
|
!!! note
|
||||||
When adding a directory, the same content filters configured for [file monitoring](configuration.md#filtering-monitored-files) are applied. This means `ignore_patterns` and `include_patterns` from your configuration will be used to filter which files are added.
|
When adding a directory, the same content filters configured for [file monitoring](processing.md#filtering-monitored-files) are applied. This means `ignore_patterns` and `include_patterns` from your configuration will be used to filter which files are added.
|
||||||
|
|
||||||
!!! note
|
!!! note
|
||||||
As you add documents to `haiku.rag` the database keeps growing. By default, LanceDB supports versioning
|
As you add documents to `haiku.rag` the database keeps growing. By default, LanceDB supports versioning
|
||||||
|
|
@ -158,7 +158,7 @@ haiku-rag research "How does haiku.rag organize and query documents?" --verbose
|
||||||
Flags:
|
Flags:
|
||||||
- `--verbose`: Show planning, searching previews, evaluation summary, and stop reason
|
- `--verbose`: Show planning, searching previews, evaluation summary, and stop reason
|
||||||
|
|
||||||
Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration.md) under the `research` section.
|
Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](config-index.md) under the `research` section.
|
||||||
|
|
||||||
When `--verbose` is set, the CLI consumes the research graph's AG-UI event stream, displaying step events and activity snapshots as agents progress through planning, search, evaluation, and synthesis. Without `--verbose`, only the final research report is displayed.
|
When `--verbose` is set, the CLI consumes the research graph's AG-UI event stream, displaying step events and activity snapshots as agents progress through planning, search, evaluation, and synthesis. Without `--verbose`, only the final research report is displayed.
|
||||||
|
|
||||||
|
|
|
||||||
193
docs/config-index.md
Normal file
193
docs/config-index.md
Normal file
|
|
@ -0,0 +1,193 @@
|
||||||
|
# Configuration
|
||||||
|
|
||||||
|
Configuration is done through YAML configuration files.
|
||||||
|
|
||||||
|
!!! note
|
||||||
|
If you create a db with certain settings and later change them, `haiku.rag` will detect incompatibilities (for example, if you change embedding provider) and will exit. You can **rebuild** the database to apply the new settings, see [Rebuild Database](./cli.md#rebuild-database).
|
||||||
|
|
||||||
|
## Getting Started
|
||||||
|
|
||||||
|
Generate a configuration file with defaults:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-rag init-config
|
||||||
|
```
|
||||||
|
|
||||||
|
This creates a `haiku.rag.yaml` file in your current directory with all available settings.
|
||||||
|
|
||||||
|
## Configuration File Locations
|
||||||
|
|
||||||
|
`haiku.rag` searches for configuration files in this order:
|
||||||
|
|
||||||
|
1. Path specified via `--config` flag: `haiku-rag --config /path/to/config.yaml <command>`
|
||||||
|
2. `./haiku.rag.yaml` (current directory)
|
||||||
|
3. Platform-specific user directory:
|
||||||
|
- **Linux**: `~/.local/share/haiku.rag/haiku.rag.yaml`
|
||||||
|
- **macOS**: `~/Library/Application Support/haiku.rag/haiku.rag.yaml`
|
||||||
|
- **Windows**: `C:/Users/<USER>/AppData/Roaming/haiku.rag/haiku.rag.yaml`
|
||||||
|
|
||||||
|
## Minimal Configuration
|
||||||
|
|
||||||
|
A minimal configuration file with defaults:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
# haiku.rag.yaml
|
||||||
|
environment: production
|
||||||
|
|
||||||
|
embeddings:
|
||||||
|
provider: ollama
|
||||||
|
model: qwen3-embedding:4b
|
||||||
|
vector_dim: 2560
|
||||||
|
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: gpt-oss
|
||||||
|
enable_thinking: false
|
||||||
|
```
|
||||||
|
|
||||||
|
## Complete Configuration Example
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
# haiku.rag.yaml
|
||||||
|
environment: production
|
||||||
|
|
||||||
|
storage:
|
||||||
|
data_dir: "" # Empty = use default platform location
|
||||||
|
vacuum_retention_seconds: 86400
|
||||||
|
|
||||||
|
monitor:
|
||||||
|
directories:
|
||||||
|
- /path/to/documents
|
||||||
|
- /another/path
|
||||||
|
ignore_patterns: [] # Gitignore-style patterns to exclude
|
||||||
|
include_patterns: [] # Gitignore-style patterns to include
|
||||||
|
|
||||||
|
lancedb:
|
||||||
|
uri: "" # Empty for local, or db://, s3://, az://, gs://
|
||||||
|
api_key: ""
|
||||||
|
region: ""
|
||||||
|
|
||||||
|
embeddings:
|
||||||
|
provider: ollama
|
||||||
|
model: qwen3-embedding:4b
|
||||||
|
vector_dim: 2560
|
||||||
|
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: "" # Empty to disable, or mxbai, cohere, zeroentropy, vllm
|
||||||
|
name: ""
|
||||||
|
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: gpt-oss
|
||||||
|
enable_thinking: false
|
||||||
|
max_sub_questions: 3
|
||||||
|
max_iterations: 2
|
||||||
|
max_concurrency: 1
|
||||||
|
|
||||||
|
research:
|
||||||
|
model:
|
||||||
|
provider: "" # Empty to use qa settings
|
||||||
|
name: ""
|
||||||
|
enable_thinking: false
|
||||||
|
max_iterations: 3
|
||||||
|
confidence_threshold: 0.8
|
||||||
|
max_concurrency: 1
|
||||||
|
|
||||||
|
search:
|
||||||
|
vector_index_metric: cosine # cosine, l2, or dot
|
||||||
|
vector_refine_factor: 30
|
||||||
|
|
||||||
|
agui:
|
||||||
|
host: "0.0.0.0"
|
||||||
|
port: 8000
|
||||||
|
cors_origins: ["*"]
|
||||||
|
cors_credentials: true
|
||||||
|
cors_methods: ["GET", "POST", "OPTIONS"]
|
||||||
|
cors_headers: ["*"]
|
||||||
|
|
||||||
|
processing:
|
||||||
|
converter: docling-local # docling-local or docling-serve
|
||||||
|
chunker: docling-local # docling-local or docling-serve
|
||||||
|
chunker_type: hybrid # hybrid or hierarchical
|
||||||
|
chunk_size: 256
|
||||||
|
context_chunk_radius: 0
|
||||||
|
chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B"
|
||||||
|
chunking_merge_peers: true
|
||||||
|
chunking_use_markdown_tables: false
|
||||||
|
markdown_preprocessor: ""
|
||||||
|
conversion_options:
|
||||||
|
do_ocr: true
|
||||||
|
force_ocr: false
|
||||||
|
ocr_lang: []
|
||||||
|
do_table_structure: true
|
||||||
|
table_mode: accurate
|
||||||
|
table_cell_matching: true
|
||||||
|
images_scale: 2.0
|
||||||
|
|
||||||
|
providers:
|
||||||
|
ollama:
|
||||||
|
base_url: http://localhost:11434
|
||||||
|
|
||||||
|
vllm:
|
||||||
|
embeddings_base_url: ""
|
||||||
|
rerank_base_url: ""
|
||||||
|
qa_base_url: ""
|
||||||
|
research_base_url: ""
|
||||||
|
|
||||||
|
docling_serve:
|
||||||
|
base_url: http://localhost:5001
|
||||||
|
api_key: ""
|
||||||
|
timeout: 300
|
||||||
|
```
|
||||||
|
|
||||||
|
## Programmatic Configuration
|
||||||
|
|
||||||
|
When using haiku.rag as a Python library, you can pass configuration directly to the `HaikuRAG` client:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from haiku.rag.config import AppConfig
|
||||||
|
from haiku.rag.config.models import ModelConfig, QAConfig, EmbeddingsConfig
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
|
||||||
|
# Create custom configuration
|
||||||
|
custom_config = AppConfig(
|
||||||
|
qa=QAConfig(
|
||||||
|
model=ModelConfig(
|
||||||
|
provider="openai",
|
||||||
|
model="gpt-4o",
|
||||||
|
temperature=0.7
|
||||||
|
)
|
||||||
|
),
|
||||||
|
embeddings=EmbeddingsConfig(
|
||||||
|
model=ModelConfig(
|
||||||
|
provider="ollama",
|
||||||
|
model="qwen3-embedding:4b"
|
||||||
|
),
|
||||||
|
vector_dim=2560
|
||||||
|
),
|
||||||
|
processing={"chunk_size": 512}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Pass configuration to the client
|
||||||
|
client = HaikuRAG(config=custom_config)
|
||||||
|
```
|
||||||
|
|
||||||
|
If you don't pass a config, the client uses the global configuration loaded from your YAML file or defaults.
|
||||||
|
|
||||||
|
This is useful for:
|
||||||
|
- Jupyter notebooks
|
||||||
|
- Python scripts
|
||||||
|
- Testing with different configurations
|
||||||
|
- Applications that need multiple clients with different configurations
|
||||||
|
|
||||||
|
## Configuration Topics
|
||||||
|
|
||||||
|
For detailed configuration of specific topics, see:
|
||||||
|
|
||||||
|
- **[Providers](providers.md)** - Model settings and provider-specific configuration (embeddings, QA, reranking)
|
||||||
|
- **[QA and Research](qa-research.md)** - Question answering and research workflow configuration
|
||||||
|
- **[Storage](storage.md)** - Database, remote storage, and vector indexing
|
||||||
|
- **[Document Processing](processing.md)** - Document conversion, chunking, and file monitoring
|
||||||
|
|
@ -1,829 +0,0 @@
|
||||||
# Configuration
|
|
||||||
|
|
||||||
Configuration is done through YAML configuration files.
|
|
||||||
|
|
||||||
!!! note
|
|
||||||
If you create a db with certain settings and later change them, `haiku.rag` will detect incompatibilities (for example, if you change embedding provider) and will exit. You can **rebuild** the database to apply the new settings, see [Rebuild Database](./cli.md#rebuild-database).
|
|
||||||
|
|
||||||
## Getting Started
|
|
||||||
|
|
||||||
Generate a configuration file with defaults:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
haiku-rag init-config
|
|
||||||
```
|
|
||||||
|
|
||||||
This creates a `haiku.rag.yaml` file in your current directory with all available settings.
|
|
||||||
|
|
||||||
## Configuration File Locations
|
|
||||||
|
|
||||||
`haiku.rag` searches for configuration files in this order:
|
|
||||||
|
|
||||||
1. Path specified via `--config` flag: `haiku-rag --config /path/to/config.yaml <command>`
|
|
||||||
2. `./haiku.rag.yaml` (current directory)
|
|
||||||
3. Platform-specific user directory:
|
|
||||||
- **Linux**: `~/.local/share/haiku.rag/haiku.rag.yaml`
|
|
||||||
- **macOS**: `~/Library/Application Support/haiku.rag/haiku.rag.yaml`
|
|
||||||
- **Windows**: `C:/Users/<USER>/AppData/Roaming/haiku.rag/haiku.rag.yaml`
|
|
||||||
|
|
||||||
## Minimal Configuration
|
|
||||||
|
|
||||||
A minimal configuration file with defaults:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
# haiku.rag.yaml
|
|
||||||
environment: production
|
|
||||||
|
|
||||||
embeddings:
|
|
||||||
provider: ollama
|
|
||||||
model: qwen3-embedding:4b
|
|
||||||
vector_dim: 2560
|
|
||||||
|
|
||||||
qa:
|
|
||||||
provider: ollama
|
|
||||||
model: gpt-oss
|
|
||||||
```
|
|
||||||
|
|
||||||
## Complete Configuration Example
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
# haiku.rag.yaml
|
|
||||||
environment: production
|
|
||||||
|
|
||||||
storage:
|
|
||||||
data_dir: "" # Empty = use default platform location
|
|
||||||
vacuum_retention_seconds: 86400
|
|
||||||
|
|
||||||
monitor:
|
|
||||||
directories:
|
|
||||||
- /path/to/documents
|
|
||||||
- /another/path
|
|
||||||
ignore_patterns: [] # Gitignore-style patterns to exclude
|
|
||||||
include_patterns: [] # Gitignore-style patterns to include
|
|
||||||
|
|
||||||
lancedb:
|
|
||||||
uri: "" # Empty for local, or db://, s3://, az://, gs://
|
|
||||||
api_key: ""
|
|
||||||
region: ""
|
|
||||||
|
|
||||||
embeddings:
|
|
||||||
provider: ollama
|
|
||||||
model: qwen3-embedding:4b
|
|
||||||
vector_dim: 2560
|
|
||||||
|
|
||||||
reranking:
|
|
||||||
provider: "" # Empty to disable, or mxbai, cohere, zeroentropy, vllm
|
|
||||||
model: ""
|
|
||||||
|
|
||||||
qa:
|
|
||||||
provider: ollama
|
|
||||||
model: gpt-oss
|
|
||||||
|
|
||||||
research:
|
|
||||||
provider: "" # Empty to use qa settings
|
|
||||||
model: ""
|
|
||||||
max_iterations: 3
|
|
||||||
confidence_threshold: 0.8
|
|
||||||
max_concurrency: 1
|
|
||||||
|
|
||||||
search:
|
|
||||||
vector_index_metric: cosine # cosine, l2, or dot
|
|
||||||
vector_refine_factor: 30
|
|
||||||
|
|
||||||
agui:
|
|
||||||
host: "0.0.0.0"
|
|
||||||
port: 8000
|
|
||||||
cors_origins: ["*"]
|
|
||||||
cors_credentials: true
|
|
||||||
cors_methods: ["GET", "POST", "OPTIONS"]
|
|
||||||
cors_headers: ["*"]
|
|
||||||
|
|
||||||
processing:
|
|
||||||
converter: docling-local # docling-local or docling-serve
|
|
||||||
chunker: docling-local # docling-local or docling-serve
|
|
||||||
chunker_type: hybrid # hybrid or hierarchical
|
|
||||||
chunk_size: 256
|
|
||||||
context_chunk_radius: 0
|
|
||||||
chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B"
|
|
||||||
chunking_merge_peers: true
|
|
||||||
chunking_use_markdown_tables: false
|
|
||||||
markdown_preprocessor: ""
|
|
||||||
conversion_options:
|
|
||||||
do_ocr: true
|
|
||||||
force_ocr: false
|
|
||||||
ocr_lang: []
|
|
||||||
do_table_structure: true
|
|
||||||
table_mode: accurate
|
|
||||||
table_cell_matching: true
|
|
||||||
images_scale: 2.0
|
|
||||||
|
|
||||||
providers:
|
|
||||||
ollama:
|
|
||||||
base_url: http://localhost:11434
|
|
||||||
|
|
||||||
vllm:
|
|
||||||
embeddings_base_url: ""
|
|
||||||
rerank_base_url: ""
|
|
||||||
qa_base_url: ""
|
|
||||||
research_base_url: ""
|
|
||||||
|
|
||||||
docling_serve:
|
|
||||||
base_url: http://localhost:5001
|
|
||||||
api_key: ""
|
|
||||||
timeout: 300
|
|
||||||
```
|
|
||||||
|
|
||||||
## Programmatic Configuration
|
|
||||||
|
|
||||||
When using haiku.rag as a Python library, you can pass configuration directly to the `HaikuRAG` client:
|
|
||||||
|
|
||||||
```python
|
|
||||||
from haiku.rag.config import AppConfig
|
|
||||||
from haiku.rag.client import HaikuRAG
|
|
||||||
|
|
||||||
# Create custom configuration
|
|
||||||
custom_config = AppConfig(
|
|
||||||
qa={"provider": "openai", "model": "gpt-4o"},
|
|
||||||
embeddings={"provider": "ollama", "model": "qwen3-embedding:4b", "vector_dim": 2560},
|
|
||||||
processing={"chunk_size": 512}
|
|
||||||
)
|
|
||||||
|
|
||||||
# Pass configuration to the client
|
|
||||||
client = HaikuRAG(config=custom_config)
|
|
||||||
```
|
|
||||||
|
|
||||||
If you don't pass a config, the client uses the global configuration loaded from your YAML file or defaults.
|
|
||||||
|
|
||||||
This is useful for:
|
|
||||||
- Jupyter notebooks
|
|
||||||
- Python scripts
|
|
||||||
- Testing with different configurations
|
|
||||||
- Applications that need multiple clients with different configurations
|
|
||||||
|
|
||||||
## File Monitoring
|
|
||||||
|
|
||||||
Set directories to monitor for automatic indexing:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
monitor:
|
|
||||||
directories:
|
|
||||||
- /path/to/documents
|
|
||||||
- /another_path/to/documents
|
|
||||||
```
|
|
||||||
|
|
||||||
### Filtering Monitored Files
|
|
||||||
|
|
||||||
Use gitignore-style patterns to control which files are monitored:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
monitor:
|
|
||||||
directories:
|
|
||||||
- /path/to/documents
|
|
||||||
|
|
||||||
# Exclude specific files or directories
|
|
||||||
ignore_patterns:
|
|
||||||
- "*draft*" # Ignore files with "draft" in the name
|
|
||||||
- "temp/" # Ignore temp directory
|
|
||||||
- "**/archive/**" # Ignore all archive directories
|
|
||||||
- "*.backup" # Ignore backup files
|
|
||||||
|
|
||||||
# Only include specific files (whitelist mode)
|
|
||||||
include_patterns:
|
|
||||||
- "*.md" # Only markdown files
|
|
||||||
- "*.pdf" # Only PDF files
|
|
||||||
- "**/docs/**" # Only files in docs directories
|
|
||||||
```
|
|
||||||
|
|
||||||
**How patterns work:**
|
|
||||||
|
|
||||||
1. **Extension filtering** - Only supported file types are considered
|
|
||||||
2. **Include patterns** - If specified, only matching files are included (whitelist)
|
|
||||||
3. **Ignore patterns** - Matching files are excluded (blacklist)
|
|
||||||
4. **Combining both** - Include patterns are applied first, then ignore patterns
|
|
||||||
|
|
||||||
**Common patterns:**
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
# Only monitor markdown documentation, but ignore drafts
|
|
||||||
monitor:
|
|
||||||
include_patterns:
|
|
||||||
- "*.md"
|
|
||||||
ignore_patterns:
|
|
||||||
- "*draft*"
|
|
||||||
- "*WIP*"
|
|
||||||
|
|
||||||
# Monitor all supported files except in specific directories
|
|
||||||
monitor:
|
|
||||||
ignore_patterns:
|
|
||||||
- "node_modules/"
|
|
||||||
- ".git/"
|
|
||||||
- "**/test/**"
|
|
||||||
- "**/temp/**"
|
|
||||||
```
|
|
||||||
|
|
||||||
Patterns follow [gitignore syntax](https://git-scm.com/docs/gitignore#_pattern_format):
|
|
||||||
|
|
||||||
- `*` matches anything except `/`
|
|
||||||
- `**` matches zero or more directories
|
|
||||||
- `?` matches any single character
|
|
||||||
- `[abc]` matches any character in the set
|
|
||||||
|
|
||||||
## Document Processing
|
|
||||||
|
|
||||||
Configure how documents are converted and chunked:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
processing:
|
|
||||||
# Chunking configuration
|
|
||||||
chunk_size: 256 # Maximum tokens per chunk
|
|
||||||
context_chunk_radius: 0 # Context radius for chunk expansion
|
|
||||||
markdown_preprocessor: "" # Optional preprocessor script
|
|
||||||
|
|
||||||
# Converter selection
|
|
||||||
converter: docling-local # docling-local or docling-serve
|
|
||||||
|
|
||||||
# Chunker selection and configuration
|
|
||||||
chunker: docling-local # docling-local or docling-serve
|
|
||||||
chunker_type: hybrid # hybrid or hierarchical
|
|
||||||
chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B" # HuggingFace model for tokenization
|
|
||||||
chunking_merge_peers: true # Merge undersized successive chunks
|
|
||||||
chunking_use_markdown_tables: false # Use markdown tables vs narrative format
|
|
||||||
|
|
||||||
# Conversion options (works with both local and remote converters)
|
|
||||||
conversion_options:
|
|
||||||
# OCR settings
|
|
||||||
do_ocr: true # Enable OCR for bitmap content
|
|
||||||
force_ocr: false # Replace existing text with OCR
|
|
||||||
ocr_lang: [] # OCR languages (e.g., ["en", "fr", "de"])
|
|
||||||
|
|
||||||
# Table extraction
|
|
||||||
do_table_structure: true # Extract table structure
|
|
||||||
table_mode: accurate # fast or accurate
|
|
||||||
table_cell_matching: true # Match table cells back to PDF cells
|
|
||||||
|
|
||||||
# Image settings
|
|
||||||
images_scale: 2.0 # Image scale factor
|
|
||||||
```
|
|
||||||
|
|
||||||
### Conversion Options
|
|
||||||
|
|
||||||
The `conversion_options` section allows fine-grained control over document conversion. These options work with both `docling-local` and `docling-serve` converters.
|
|
||||||
|
|
||||||
#### OCR Settings
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
conversion_options:
|
|
||||||
do_ocr: true # Enable OCR for bitmap/scanned content
|
|
||||||
force_ocr: false # Replace all text with OCR output
|
|
||||||
ocr_lang: [] # List of OCR languages, e.g., ["en", "fr", "de"]
|
|
||||||
```
|
|
||||||
|
|
||||||
- **do_ocr**: When `true`, applies OCR to images and scanned pages. Disable for faster processing if documents contain only native text.
|
|
||||||
- **force_ocr**: When `true`, replaces existing text layers with OCR output. Useful for documents with poor text extraction.
|
|
||||||
- **ocr_lang**: List of language codes for OCR. Empty list uses default language detection. Examples: `["en"]`, `["en", "fr", "de"]`.
|
|
||||||
|
|
||||||
#### Table Extraction
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
conversion_options:
|
|
||||||
do_table_structure: true # Extract structured table data
|
|
||||||
table_mode: accurate # fast or accurate
|
|
||||||
table_cell_matching: true # Match cells back to PDF
|
|
||||||
```
|
|
||||||
|
|
||||||
- **do_table_structure**: When `true`, extracts table structure. Disable for faster processing if tables aren't important.
|
|
||||||
- **table_mode**:
|
|
||||||
- `accurate`: Better table structure recognition (slower)
|
|
||||||
- `fast`: Faster processing with simpler table detection
|
|
||||||
- **table_cell_matching**: When `true`, matches detected table cells back to PDF cells. Disable if tables have merged cells across columns.
|
|
||||||
|
|
||||||
#### Image Settings
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
conversion_options:
|
|
||||||
images_scale: 2.0 # Image resolution scale factor
|
|
||||||
```
|
|
||||||
|
|
||||||
- **images_scale**: Scale factor for extracted images. Higher values = better quality but larger size. Typical range: 1.0-3.0.
|
|
||||||
|
|
||||||
### Local vs Remote Processing
|
|
||||||
|
|
||||||
**Local processing** (default):
|
|
||||||
|
|
||||||
- Uses `docling` library locally
|
|
||||||
- No external dependencies
|
|
||||||
- Good for development and small workloads
|
|
||||||
|
|
||||||
**Remote processing** (docling-serve):
|
|
||||||
|
|
||||||
- Offloads processing to docling-serve API
|
|
||||||
- Better for heavy workloads and production
|
|
||||||
- Requires docling-serve instance (see [Remote processing setup](remote-processing.md))
|
|
||||||
|
|
||||||
To use remote processing:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
processing:
|
|
||||||
converter: docling-serve
|
|
||||||
chunker: docling-serve
|
|
||||||
|
|
||||||
providers:
|
|
||||||
docling_serve:
|
|
||||||
base_url: http://localhost:5001
|
|
||||||
api_key: "your-api-key" # Optional
|
|
||||||
timeout: 300 # Request timeout in seconds
|
|
||||||
```
|
|
||||||
|
|
||||||
Conversion options work identically for both local and remote processing.
|
|
||||||
|
|
||||||
### Chunking Strategies
|
|
||||||
|
|
||||||
**Hybrid chunking** (default):
|
|
||||||
- Structure-aware chunking
|
|
||||||
- Respects document boundaries
|
|
||||||
- Best for most use cases
|
|
||||||
|
|
||||||
**Hierarchical chunking**:
|
|
||||||
- Creates hierarchical chunk structure
|
|
||||||
- Preserves document hierarchy
|
|
||||||
- Useful for complex documents
|
|
||||||
|
|
||||||
### Table Serialization
|
|
||||||
|
|
||||||
Control how tables are represented in chunks:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
processing:
|
|
||||||
chunking_use_markdown_tables: false # Default: narrative format
|
|
||||||
```
|
|
||||||
|
|
||||||
- `false`: Tables as narrative text ("Value A, Column 2 = Value B")
|
|
||||||
- `true`: Tables as markdown (preserves table structure)
|
|
||||||
|
|
||||||
## Embedding Providers
|
|
||||||
|
|
||||||
If you use Ollama, you can use any pulled model that supports embeddings.
|
|
||||||
|
|
||||||
### Ollama (Default)
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
embeddings:
|
|
||||||
provider: ollama
|
|
||||||
model: mxbai-embed-large
|
|
||||||
vector_dim: 1024
|
|
||||||
```
|
|
||||||
|
|
||||||
The Ollama base URL can be configured in your config file or via environment variable:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
providers:
|
|
||||||
ollama:
|
|
||||||
base_url: http://localhost:11434
|
|
||||||
```
|
|
||||||
|
|
||||||
Or via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export OLLAMA_BASE_URL=http://localhost:11434
|
|
||||||
```
|
|
||||||
|
|
||||||
If not configured, it defaults to `http://localhost:11434`.
|
|
||||||
|
|
||||||
!!! note
|
|
||||||
You can use a `.env` file in your project directory to set environment variables like `OLLAMA_BASE_URL` and API keys (e.g., `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`). These will be automatically loaded when running `haiku-rag` commands.
|
|
||||||
|
|
||||||
### VoyageAI
|
|
||||||
|
|
||||||
If you installed `haiku.rag` (full package), VoyageAI is already included. If you installed `haiku.rag-slim`, install with VoyageAI extras:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uv pip install haiku.rag-slim[voyageai]
|
|
||||||
```
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
embeddings:
|
|
||||||
provider: voyageai
|
|
||||||
model: voyage-3.5
|
|
||||||
vector_dim: 1024
|
|
||||||
```
|
|
||||||
|
|
||||||
Set your API key via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export VOYAGE_API_KEY=your-api-key
|
|
||||||
```
|
|
||||||
|
|
||||||
### OpenAI
|
|
||||||
|
|
||||||
OpenAI embeddings are included in the default installation:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
embeddings:
|
|
||||||
provider: openai
|
|
||||||
model: text-embedding-3-small # or text-embedding-3-large
|
|
||||||
vector_dim: 1536
|
|
||||||
```
|
|
||||||
|
|
||||||
Set your API key via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export OPENAI_API_KEY=your-api-key
|
|
||||||
```
|
|
||||||
|
|
||||||
### vLLM
|
|
||||||
|
|
||||||
For high-performance local inference, you can use vLLM to serve embedding models with OpenAI-compatible APIs:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
embeddings:
|
|
||||||
provider: vllm
|
|
||||||
model: mixedbread-ai/mxbai-embed-large-v1
|
|
||||||
vector_dim: 512
|
|
||||||
|
|
||||||
providers:
|
|
||||||
vllm:
|
|
||||||
embeddings_base_url: http://localhost:8000
|
|
||||||
```
|
|
||||||
|
|
||||||
**Note:** You need to run a vLLM server separately with an embedding model loaded.
|
|
||||||
|
|
||||||
## Question Answering Providers
|
|
||||||
|
|
||||||
Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used.
|
|
||||||
|
|
||||||
### Ollama (Default)
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
qa:
|
|
||||||
provider: ollama
|
|
||||||
model: gpt-oss
|
|
||||||
```
|
|
||||||
|
|
||||||
The Ollama base URL can be configured via the `OLLAMA_BASE_URL` environment variable, config file, or defaults to `http://localhost:11434`:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export OLLAMA_BASE_URL=http://localhost:11434
|
|
||||||
```
|
|
||||||
|
|
||||||
Or in your config file:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
providers:
|
|
||||||
ollama:
|
|
||||||
base_url: http://localhost:11434
|
|
||||||
```
|
|
||||||
|
|
||||||
### OpenAI
|
|
||||||
|
|
||||||
OpenAI QA is included in the default installation:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
qa:
|
|
||||||
provider: openai
|
|
||||||
model: gpt-4o-mini # or gpt-4, gpt-3.5-turbo, etc.
|
|
||||||
```
|
|
||||||
|
|
||||||
Set your API key via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export OPENAI_API_KEY=your-api-key
|
|
||||||
```
|
|
||||||
|
|
||||||
### Anthropic
|
|
||||||
|
|
||||||
Anthropic QA is included in the default installation:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
qa:
|
|
||||||
provider: anthropic
|
|
||||||
model: claude-3-5-haiku-20241022 # or claude-3-5-sonnet-20241022, etc.
|
|
||||||
```
|
|
||||||
|
|
||||||
Set your API key via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export ANTHROPIC_API_KEY=your-api-key
|
|
||||||
```
|
|
||||||
|
|
||||||
### vLLM
|
|
||||||
|
|
||||||
For high-performance local inference:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
qa:
|
|
||||||
provider: vllm
|
|
||||||
model: Qwen/Qwen3-4B # Any model with tool support in vLLM
|
|
||||||
|
|
||||||
providers:
|
|
||||||
vllm:
|
|
||||||
qa_base_url: http://localhost:8002
|
|
||||||
```
|
|
||||||
|
|
||||||
**Note:** You need to run a vLLM server separately with a model that supports tool calling loaded. Consult the specific model's documentation for proper vLLM serving configuration.
|
|
||||||
|
|
||||||
### Other Providers
|
|
||||||
|
|
||||||
Any provider supported by Pydantic AI can be used. Examples:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
# Google Gemini
|
|
||||||
qa:
|
|
||||||
provider: gemini
|
|
||||||
model: gemini-1.5-flash
|
|
||||||
|
|
||||||
# Groq
|
|
||||||
qa:
|
|
||||||
provider: groq
|
|
||||||
model: llama-3.3-70b-versatile
|
|
||||||
|
|
||||||
# Mistral
|
|
||||||
qa:
|
|
||||||
provider: mistral
|
|
||||||
model: mistral-small-latest
|
|
||||||
```
|
|
||||||
|
|
||||||
See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the complete list of supported providers and models.
|
|
||||||
|
|
||||||
## Reranking
|
|
||||||
|
|
||||||
Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (10x the requested limit) and then reranks them to return the most relevant results.
|
|
||||||
|
|
||||||
Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below.
|
|
||||||
|
|
||||||
### MixedBread AI
|
|
||||||
|
|
||||||
If you installed `haiku.rag` (full package), MxBAI is already included. If you installed `haiku.rag-slim`, add the mxbai extra:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uv pip install haiku.rag-slim[mxbai]
|
|
||||||
```
|
|
||||||
|
|
||||||
Then configure:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
reranking:
|
|
||||||
provider: mxbai
|
|
||||||
model: mixedbread-ai/mxbai-rerank-base-v2
|
|
||||||
```
|
|
||||||
|
|
||||||
### Cohere
|
|
||||||
|
|
||||||
If you installed `haiku.rag` (full package), Cohere is already included. If you installed `haiku.rag-slim`, add the cohere extra:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uv pip install haiku.rag-slim[cohere]
|
|
||||||
```
|
|
||||||
|
|
||||||
Then configure:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
reranking:
|
|
||||||
provider: cohere
|
|
||||||
model: rerank-v3.5
|
|
||||||
```
|
|
||||||
|
|
||||||
Set your API key via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export CO_API_KEY=your-api-key
|
|
||||||
```
|
|
||||||
|
|
||||||
### Zero Entropy
|
|
||||||
|
|
||||||
If you installed `haiku.rag` (full package), Zero Entropy is already included. If you installed `haiku.rag-slim`, add the zeroentropy extra:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uv pip install haiku.rag-slim[zeroentropy]
|
|
||||||
```
|
|
||||||
|
|
||||||
Then configure:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
reranking:
|
|
||||||
provider: zeroentropy
|
|
||||||
model: zerank-1 # Currently the only available model
|
|
||||||
```
|
|
||||||
|
|
||||||
Set your API key via environment variable:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
export ZEROENTROPY_API_KEY=your-api-key
|
|
||||||
```
|
|
||||||
|
|
||||||
### vLLM
|
|
||||||
|
|
||||||
For high-performance local reranking using dedicated reranking models:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
reranking:
|
|
||||||
provider: vllm
|
|
||||||
model: mixedbread-ai/mxbai-rerank-base-v2
|
|
||||||
|
|
||||||
providers:
|
|
||||||
vllm:
|
|
||||||
rerank_base_url: http://localhost:8001
|
|
||||||
```
|
|
||||||
|
|
||||||
**Note:** vLLM reranking uses the `/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded. Consult the specific model's documentation for proper vLLM serving configuration.
|
|
||||||
|
|
||||||
## Research Configuration
|
|
||||||
|
|
||||||
Configure the multi-agent research workflow:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
research:
|
|
||||||
provider: "" # Empty to use qa settings
|
|
||||||
model: "" # Empty to use qa model
|
|
||||||
max_iterations: 3 # Maximum search/evaluate cycles
|
|
||||||
confidence_threshold: 0.8 # Stop when confidence meets/exceeds this
|
|
||||||
max_concurrency: 1 # Sub-questions searched in parallel per iteration
|
|
||||||
```
|
|
||||||
|
|
||||||
- **provider/model**: LLM provider and model for research. Leave empty to use the same settings as `qa`.
|
|
||||||
- **max_iterations**: Maximum number of search/evaluate cycles before stopping (default: 3)
|
|
||||||
- **confidence_threshold**: Stop research when evaluation confidence score meets or exceeds this threshold (default: 0.8)
|
|
||||||
- **max_concurrency**: Number of sub-questions to search in parallel during each iteration (default: 1)
|
|
||||||
|
|
||||||
The research workflow plans sub-questions, searches in parallel batches, evaluates findings, and iterates until reaching the confidence threshold or max iterations.
|
|
||||||
|
|
||||||
## AG-UI Server Configuration
|
|
||||||
|
|
||||||
Configure the AG-UI HTTP server for streaming graph execution events:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
agui:
|
|
||||||
host: "0.0.0.0"
|
|
||||||
port: 8000
|
|
||||||
cors_origins: ["*"]
|
|
||||||
cors_credentials: true
|
|
||||||
cors_methods: ["GET", "POST", "OPTIONS"]
|
|
||||||
cors_headers: ["*"]
|
|
||||||
```
|
|
||||||
|
|
||||||
Start the AG-UI server with:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
haiku-rag serve --agui
|
|
||||||
```
|
|
||||||
|
|
||||||
The server exposes:
|
|
||||||
- `GET /health` - Health check endpoint
|
|
||||||
- `POST /v1/agent/stream` - Research graph streaming endpoint (Server-Sent Events)
|
|
||||||
|
|
||||||
See [Server Mode](server.md) for more details.
|
|
||||||
|
|
||||||
## Other Settings
|
|
||||||
|
|
||||||
### Database and Storage
|
|
||||||
|
|
||||||
By default, `haiku.rag` uses a local LanceDB database:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
storage:
|
|
||||||
data_dir: /path/to/data # Empty = use default platform location
|
|
||||||
```
|
|
||||||
|
|
||||||
For remote storage, use the `lancedb` settings with various backends:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
# LanceDB Cloud
|
|
||||||
lancedb:
|
|
||||||
uri: db://your-database-name
|
|
||||||
api_key: your-api-key
|
|
||||||
region: us-west-2 # optional
|
|
||||||
|
|
||||||
# Amazon S3
|
|
||||||
lancedb:
|
|
||||||
uri: s3://my-bucket/my-table
|
|
||||||
# Use AWS credentials or IAM roles
|
|
||||||
|
|
||||||
# Azure Blob Storage
|
|
||||||
lancedb:
|
|
||||||
uri: az://my-container/my-table
|
|
||||||
# Use Azure credentials
|
|
||||||
|
|
||||||
# Google Cloud Storage
|
|
||||||
lancedb:
|
|
||||||
uri: gs://my-bucket/my-table
|
|
||||||
# Use GCP credentials
|
|
||||||
|
|
||||||
# HDFS
|
|
||||||
lancedb:
|
|
||||||
uri: hdfs://namenode:port/path/to/table
|
|
||||||
```
|
|
||||||
|
|
||||||
Authentication is handled through standard cloud provider credentials (AWS CLI, Azure CLI, gcloud, etc.) or by setting `api_key` for LanceDB Cloud.
|
|
||||||
|
|
||||||
**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization is still performed locally.
|
|
||||||
|
|
||||||
#### Database Auto-creation
|
|
||||||
|
|
||||||
haiku.rag intelligently handles database creation based on operation type:
|
|
||||||
|
|
||||||
- **Write operations** (add, add-src, delete, rebuild): Automatically create the database and required tables if they don't exist
|
|
||||||
- **Read operations** (list, get, search, ask, research): Fail with a clear error if the database doesn't exist
|
|
||||||
|
|
||||||
This prevents the common mistake where a search query accidentally creates an empty database. To initialize your database, simply add your first document using `haiku-rag add` or `haiku-rag add-src`.
|
|
||||||
|
|
||||||
### Vector Indexing
|
|
||||||
|
|
||||||
Configure vector indexing behavior for efficient similarity search:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
search:
|
|
||||||
vector_index_metric: cosine # cosine, l2, or dot
|
|
||||||
vector_refine_factor: 30 # Re-ranking factor for accuracy
|
|
||||||
```
|
|
||||||
|
|
||||||
- **vector_index_metric**: Distance metric for vector similarity:
|
|
||||||
- `cosine`: Cosine similarity (default, best for most embeddings)
|
|
||||||
- `l2`: Euclidean distance
|
|
||||||
- `dot`: Dot product similarity
|
|
||||||
- **vector_refine_factor**: Improves accuracy when using a vector index by retrieving `refine_factor * limit` candidates (using approximate search) and re-ranking them with exact distances. Higher values increase accuracy but slow down queries. Default: 30
|
|
||||||
- **Only applies with a vector index** - has no effect on brute-force search, which already returns exact results
|
|
||||||
|
|
||||||
!!! note
|
|
||||||
Vector indexes are only necessary for large datasets with over 100,000 chunks. For smaller datasets, LanceDB's brute-force kNN search provides exact results with good performance. Only create an index if you notice search performance degradation on large datasets.
|
|
||||||
|
|
||||||
**Index creation:**
|
|
||||||
|
|
||||||
Vector indexes are **not created automatically** during document ingestion to avoid slowing down the process. After you've added documents (at least 256 chunks required), create the index manually:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
haiku-rag create-index
|
|
||||||
```
|
|
||||||
|
|
||||||
This command:
|
|
||||||
- Checks if you have enough data (minimum 256 chunks)
|
|
||||||
- Creates an IVF_PQ index for fast approximate nearest neighbor (ANN) search
|
|
||||||
- Uses LanceDB's automatic parameter calculation based on your dataset size and vector dimensions
|
|
||||||
|
|
||||||
**Re-indexing:**
|
|
||||||
|
|
||||||
Indexes are not automatically updated when you add new documents. After adding a significant amount of new data:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
haiku-rag create-index # Rebuilds the index with all data
|
|
||||||
```
|
|
||||||
|
|
||||||
Searches still work with stale indexes - LanceDB uses the index for old data (fast ANN) and brute-force kNN for new unindexed rows, then combines the results. However, performance degrades as more unindexed data accumulates.
|
|
||||||
|
|
||||||
For datasets with fewer than 256 chunks, searches use brute-force kNN scans (exact nearest neighbors, 100% recall) which work well for small datasets but don't scale beyond a few hundred thousand vectors.
|
|
||||||
|
|
||||||
### Document Processing
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
processing:
|
|
||||||
# Chunk size for document processing
|
|
||||||
chunk_size: 256
|
|
||||||
|
|
||||||
# Number of adjacent chunks to include before/after retrieved chunks for context
|
|
||||||
# 0 = no expansion (default), 1 = include 1 chunk before and after, etc.
|
|
||||||
# When expanded chunks overlap or are adjacent, they are automatically merged
|
|
||||||
# into single chunks with continuous content to eliminate duplication
|
|
||||||
context_chunk_radius: 0
|
|
||||||
|
|
||||||
# Optional dotted path or file path to a callable that preprocesses
|
|
||||||
# markdown content before chunking
|
|
||||||
markdown_preprocessor: ""
|
|
||||||
|
|
||||||
storage:
|
|
||||||
# Vacuum retention threshold (seconds) for automatic cleanup
|
|
||||||
# When documents are added/updated, old table versions older than this are removed
|
|
||||||
# Default: 86400 seconds (1 day, safe for concurrent connections)
|
|
||||||
# Set to 0 for aggressive cleanup (removes all old versions immediately)
|
|
||||||
vacuum_retention_seconds: 86400
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Markdown Preprocessor
|
|
||||||
|
|
||||||
Optionally preprocess Markdown before chunking by pointing to a callable that receives and returns Markdown text. This is useful for normalizing content, stripping boilerplate, or applying custom transformations before chunk boundaries are computed.
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
processing:
|
|
||||||
# A callable path in one of these formats:
|
|
||||||
# - package.module:func
|
|
||||||
# - package.module.func
|
|
||||||
# - /abs/or/relative/path/to/file.py:func
|
|
||||||
markdown_preprocessor: my_pkg.preprocess:clean_md
|
|
||||||
```
|
|
||||||
|
|
||||||
!!! note
|
|
||||||
- The function signature should be `def clean_md(text: str) -> str` or `async def clean_md(text: str) -> str`.
|
|
||||||
- If the function raises or returns a non-string, haiku.rag logs a warning and proceeds without preprocessing.
|
|
||||||
- The preprocessor affects only the chunking pipeline. The stored document content remains unchanged.
|
|
||||||
|
|
||||||
Example implementation:
|
|
||||||
|
|
||||||
```python
|
|
||||||
# my_pkg/preprocess.py
|
|
||||||
def clean_md(text: str) -> str:
|
|
||||||
# strip HTML comments and collapse multiple blank lines
|
|
||||||
lines = [line for line in text.splitlines() if not line.strip().startswith("<!--")]
|
|
||||||
out = []
|
|
||||||
for line in lines:
|
|
||||||
if line.strip() == "" and (out and out[-1] == ""):
|
|
||||||
continue
|
|
||||||
out.append(line)
|
|
||||||
return "\n".join(out)
|
|
||||||
```
|
|
||||||
|
|
@ -55,7 +55,7 @@ haiku-rag ask "Who is the author of haiku.rag?"
|
||||||
|
|
||||||
- [Getting started](tutorial.md) - Tutorial
|
- [Getting started](tutorial.md) - Tutorial
|
||||||
- [Installation](installation.md) - Install haiku.rag with different providers
|
- [Installation](installation.md) - Install haiku.rag with different providers
|
||||||
- [Configuration](configuration.md) - Environment variables and settings
|
- [Configuration](config-index.md) - Environment variables and settings
|
||||||
- [CLI](cli.md) - Command line interface usage
|
- [CLI](cli.md) - Command line interface usage
|
||||||
- [Server](server.md) - File monitoring and server mode
|
- [Server](server.md) - File monitoring and server mode
|
||||||
- [MCP](mcp.md) - Model Context Protocol integration
|
- [MCP](mcp.md) - Model Context Protocol integration
|
||||||
|
|
|
||||||
|
|
@ -43,7 +43,7 @@ The slim package has minimal dependencies and lets you install only what you nee
|
||||||
- **OpenAI** (GPT models for QA and embeddings)
|
- **OpenAI** (GPT models for QA and embeddings)
|
||||||
- **Anthropic** (Claude models for QA)
|
- **Anthropic** (Claude models for QA)
|
||||||
|
|
||||||
See [Configuration](configuration.md) for configuring providers including advanced options like vLLM.
|
See [Configuration](config-index.md) for configuring providers including advanced options like vLLM.
|
||||||
|
|
||||||
## Requirements
|
## Requirements
|
||||||
|
|
||||||
|
|
@ -71,7 +71,7 @@ When using `haiku.rag-slim`, you can skip installing the `docling` extra and ins
|
||||||
- Offloading heavy document processing to a dedicated service
|
- Offloading heavy document processing to a dedicated service
|
||||||
- Production deployments with separate processing infrastructure
|
- Production deployments with separate processing infrastructure
|
||||||
|
|
||||||
See [Remote processing](remote-processing.md) for setup instructions and [Document Processing](configuration.md#document-processing) for configuration options.
|
See [Remote processing](remote-processing.md) for setup instructions and [Document Processing](processing.md) for configuration options.
|
||||||
|
|
||||||
## Docker
|
## Docker
|
||||||
|
|
||||||
|
|
|
||||||
250
docs/processing.md
Normal file
250
docs/processing.md
Normal file
|
|
@ -0,0 +1,250 @@
|
||||||
|
# Document Processing & Monitoring
|
||||||
|
|
||||||
|
This guide covers how haiku.rag converts, chunks, and monitors documents.
|
||||||
|
|
||||||
|
## Document Processing
|
||||||
|
|
||||||
|
Configure how documents are converted and chunked:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
processing:
|
||||||
|
# Chunking configuration
|
||||||
|
chunk_size: 256 # Maximum tokens per chunk
|
||||||
|
context_chunk_radius: 0 # Context radius for chunk expansion
|
||||||
|
markdown_preprocessor: "" # Optional preprocessor script
|
||||||
|
|
||||||
|
# Converter selection
|
||||||
|
converter: docling-local # docling-local or docling-serve
|
||||||
|
|
||||||
|
# Chunker selection and configuration
|
||||||
|
chunker: docling-local # docling-local or docling-serve
|
||||||
|
chunker_type: hybrid # hybrid or hierarchical
|
||||||
|
chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B" # HuggingFace model for tokenization
|
||||||
|
chunking_merge_peers: true # Merge undersized successive chunks
|
||||||
|
chunking_use_markdown_tables: false # Use markdown tables vs narrative format
|
||||||
|
|
||||||
|
# Conversion options (works with both local and remote converters)
|
||||||
|
conversion_options:
|
||||||
|
# OCR settings
|
||||||
|
do_ocr: true # Enable OCR for bitmap content
|
||||||
|
force_ocr: false # Replace existing text with OCR
|
||||||
|
ocr_lang: [] # OCR languages (e.g., ["en", "fr", "de"])
|
||||||
|
|
||||||
|
# Table extraction
|
||||||
|
do_table_structure: true # Extract table structure
|
||||||
|
table_mode: accurate # fast or accurate
|
||||||
|
table_cell_matching: true # Match table cells back to PDF cells
|
||||||
|
|
||||||
|
# Image settings
|
||||||
|
images_scale: 2.0 # Image scale factor
|
||||||
|
```
|
||||||
|
|
||||||
|
### Conversion Options
|
||||||
|
|
||||||
|
The `conversion_options` section allows fine-grained control over document conversion. These options work with both `docling-local` and `docling-serve` converters.
|
||||||
|
|
||||||
|
#### OCR Settings
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
conversion_options:
|
||||||
|
do_ocr: true # Enable OCR for bitmap/scanned content
|
||||||
|
force_ocr: false # Replace all text with OCR output
|
||||||
|
ocr_lang: [] # List of OCR languages, e.g., ["en", "fr", "de"]
|
||||||
|
```
|
||||||
|
|
||||||
|
- **do_ocr**: When `true`, applies OCR to images and scanned pages. Disable for faster processing if documents contain only native text.
|
||||||
|
- **force_ocr**: When `true`, replaces existing text layers with OCR output. Useful for documents with poor text extraction.
|
||||||
|
- **ocr_lang**: List of language codes for OCR. Empty list uses default language detection. Examples: `["en"]`, `["en", "fr", "de"]`.
|
||||||
|
|
||||||
|
#### Table Extraction
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
conversion_options:
|
||||||
|
do_table_structure: true # Extract structured table data
|
||||||
|
table_mode: accurate # fast or accurate
|
||||||
|
table_cell_matching: true # Match cells back to PDF
|
||||||
|
```
|
||||||
|
|
||||||
|
- **do_table_structure**: When `true`, extracts table structure. Disable for faster processing if tables aren't important.
|
||||||
|
- **table_mode**:
|
||||||
|
- `accurate`: Better table structure recognition (slower)
|
||||||
|
- `fast`: Faster processing with simpler table detection
|
||||||
|
- **table_cell_matching**: When `true`, matches detected table cells back to PDF cells. Disable if tables have merged cells across columns.
|
||||||
|
|
||||||
|
#### Image Settings
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
conversion_options:
|
||||||
|
images_scale: 2.0 # Image resolution scale factor
|
||||||
|
```
|
||||||
|
|
||||||
|
- **images_scale**: Scale factor for extracted images. Higher values = better quality but larger size. Typical range: 1.0-3.0.
|
||||||
|
|
||||||
|
### Local vs Remote Processing
|
||||||
|
|
||||||
|
**Local processing** (default):
|
||||||
|
|
||||||
|
- Uses `docling` library locally
|
||||||
|
- No external dependencies
|
||||||
|
- Good for development and small workloads
|
||||||
|
|
||||||
|
**Remote processing** (docling-serve):
|
||||||
|
|
||||||
|
- Offloads processing to docling-serve API
|
||||||
|
- Better for heavy workloads and production
|
||||||
|
- Requires docling-serve instance (see [Remote processing setup](remote-processing.md))
|
||||||
|
|
||||||
|
To use remote processing:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
processing:
|
||||||
|
converter: docling-serve
|
||||||
|
chunker: docling-serve
|
||||||
|
|
||||||
|
providers:
|
||||||
|
docling_serve:
|
||||||
|
base_url: http://localhost:5001
|
||||||
|
api_key: "your-api-key" # Optional
|
||||||
|
timeout: 300 # Request timeout in seconds
|
||||||
|
```
|
||||||
|
|
||||||
|
Conversion options work identically for both local and remote processing.
|
||||||
|
|
||||||
|
### Chunking Strategies
|
||||||
|
|
||||||
|
**Hybrid chunking** (default):
|
||||||
|
- Structure-aware chunking
|
||||||
|
- Respects document boundaries
|
||||||
|
- Best for most use cases
|
||||||
|
|
||||||
|
**Hierarchical chunking**:
|
||||||
|
- Creates hierarchical chunk structure
|
||||||
|
- Preserves document hierarchy
|
||||||
|
- Useful for complex documents
|
||||||
|
|
||||||
|
### Table Serialization
|
||||||
|
|
||||||
|
Control how tables are represented in chunks:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
processing:
|
||||||
|
chunking_use_markdown_tables: false # Default: narrative format
|
||||||
|
```
|
||||||
|
|
||||||
|
- `false`: Tables as narrative text ("Value A, Column 2 = Value B")
|
||||||
|
- `true`: Tables as markdown (preserves table structure)
|
||||||
|
|
||||||
|
### Chunk Size and Context
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
processing:
|
||||||
|
# Chunk size for document processing
|
||||||
|
chunk_size: 256
|
||||||
|
|
||||||
|
# Number of adjacent chunks to include before/after retrieved chunks for context
|
||||||
|
# 0 = no expansion (default), 1 = include 1 chunk before and after, etc.
|
||||||
|
# When expanded chunks overlap or are adjacent, they are automatically merged
|
||||||
|
# into single chunks with continuous content to eliminate duplication
|
||||||
|
context_chunk_radius: 0
|
||||||
|
```
|
||||||
|
|
||||||
|
### Markdown Preprocessor
|
||||||
|
|
||||||
|
Optionally preprocess Markdown before chunking by pointing to a callable that receives and returns Markdown text. This is useful for normalizing content, stripping boilerplate, or applying custom transformations before chunk boundaries are computed.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
processing:
|
||||||
|
# A callable path in one of these formats:
|
||||||
|
# - package.module:func
|
||||||
|
# - package.module.func
|
||||||
|
# - /abs/or/relative/path/to/file.py:func
|
||||||
|
markdown_preprocessor: my_pkg.preprocess:clean_md
|
||||||
|
```
|
||||||
|
|
||||||
|
!!! note
|
||||||
|
- The function signature should be `def clean_md(text: str) -> str` or `async def clean_md(text: str) -> str`.
|
||||||
|
- If the function raises or returns a non-string, haiku.rag logs a warning and proceeds without preprocessing.
|
||||||
|
- The preprocessor affects only the chunking pipeline. The stored document content remains unchanged.
|
||||||
|
|
||||||
|
Example implementation:
|
||||||
|
|
||||||
|
```python
|
||||||
|
# my_pkg/preprocess.py
|
||||||
|
def clean_md(text: str) -> str:
|
||||||
|
# strip HTML comments and collapse multiple blank lines
|
||||||
|
lines = [line for line in text.splitlines() if not line.strip().startswith("<!--")]
|
||||||
|
out = []
|
||||||
|
for line in lines:
|
||||||
|
if line.strip() == "" and (out and out[-1] == ""):
|
||||||
|
continue
|
||||||
|
out.append(line)
|
||||||
|
return "\n".join(out)
|
||||||
|
```
|
||||||
|
|
||||||
|
## File Monitoring
|
||||||
|
|
||||||
|
Set directories to monitor for automatic indexing:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
monitor:
|
||||||
|
directories:
|
||||||
|
- /path/to/documents
|
||||||
|
- /another_path/to/documents
|
||||||
|
```
|
||||||
|
|
||||||
|
### Filtering Monitored Files
|
||||||
|
|
||||||
|
Use gitignore-style patterns to control which files are monitored:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
monitor:
|
||||||
|
directories:
|
||||||
|
- /path/to/documents
|
||||||
|
|
||||||
|
# Exclude specific files or directories
|
||||||
|
ignore_patterns:
|
||||||
|
- "*draft*" # Ignore files with "draft" in the name
|
||||||
|
- "temp/" # Ignore temp directory
|
||||||
|
- "**/archive/**" # Ignore all archive directories
|
||||||
|
- "*.backup" # Ignore backup files
|
||||||
|
|
||||||
|
# Only include specific files (whitelist mode)
|
||||||
|
include_patterns:
|
||||||
|
- "*.md" # Only markdown files
|
||||||
|
- "*.pdf" # Only PDF files
|
||||||
|
- "**/docs/**" # Only files in docs directories
|
||||||
|
```
|
||||||
|
|
||||||
|
**How patterns work:**
|
||||||
|
|
||||||
|
1. **Extension filtering** - Only supported file types are considered
|
||||||
|
2. **Include patterns** - If specified, only matching files are included (whitelist)
|
||||||
|
3. **Ignore patterns** - Matching files are excluded (blacklist)
|
||||||
|
4. **Combining both** - Include patterns are applied first, then ignore patterns
|
||||||
|
|
||||||
|
**Common patterns:**
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
# Only monitor markdown documentation, but ignore drafts
|
||||||
|
monitor:
|
||||||
|
include_patterns:
|
||||||
|
- "*.md"
|
||||||
|
ignore_patterns:
|
||||||
|
- "*draft*"
|
||||||
|
- "*WIP*"
|
||||||
|
|
||||||
|
# Monitor all supported files except in specific directories
|
||||||
|
monitor:
|
||||||
|
ignore_patterns:
|
||||||
|
- "node_modules/"
|
||||||
|
- ".git/"
|
||||||
|
- "**/test/**"
|
||||||
|
- "**/temp/**"
|
||||||
|
```
|
||||||
|
|
||||||
|
Patterns follow [gitignore syntax](https://git-scm.com/docs/gitignore#_pattern_format):
|
||||||
|
|
||||||
|
- `*` matches anything except `/`
|
||||||
|
- `**` matches zero or more directories
|
||||||
|
- `?` matches any single character
|
||||||
|
- `[abc]` matches any character in the set
|
||||||
339
docs/providers.md
Normal file
339
docs/providers.md
Normal file
|
|
@ -0,0 +1,339 @@
|
||||||
|
# Providers
|
||||||
|
|
||||||
|
haiku.rag supports multiple AI providers for embeddings, question answering, and reranking. This guide covers provider-specific configuration and setup.
|
||||||
|
|
||||||
|
!!! note
|
||||||
|
You can use a `.env` file in your project directory to set environment variables like `OLLAMA_BASE_URL` and API keys (e.g., `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`). These will be automatically loaded when running `haiku-rag` commands.
|
||||||
|
|
||||||
|
## Model Settings
|
||||||
|
|
||||||
|
Configure model behavior for `qa` and `research` workflows. These settings apply to any provider that supports them.
|
||||||
|
|
||||||
|
### Basic Settings
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: gpt-oss
|
||||||
|
temperature: 0.7
|
||||||
|
max_tokens: 500
|
||||||
|
```
|
||||||
|
|
||||||
|
**Available options:**
|
||||||
|
|
||||||
|
- **temperature**: Sampling temperature (0.0-1.0+)
|
||||||
|
- Lower (0.0-0.3): Deterministic, focused responses
|
||||||
|
- Medium (0.4-0.7): Balanced
|
||||||
|
- Higher (0.8-1.0+): Creative, varied responses
|
||||||
|
- **max_tokens**: Maximum tokens in response
|
||||||
|
- **enable_thinking**: Control reasoning behavior (see below)
|
||||||
|
|
||||||
|
### Thinking Control
|
||||||
|
|
||||||
|
The `enable_thinking` setting controls whether models use explicit reasoning steps before answering.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
enable_thinking: false # Faster responses
|
||||||
|
|
||||||
|
research:
|
||||||
|
model:
|
||||||
|
enable_thinking: true # Deeper reasoning
|
||||||
|
```
|
||||||
|
|
||||||
|
**Values:**
|
||||||
|
- `false`: Disable reasoning for faster responses
|
||||||
|
- `true`: Enable reasoning for complex tasks
|
||||||
|
- Not set: Use model defaults
|
||||||
|
|
||||||
|
**Provider support:**
|
||||||
|
|
||||||
|
See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/) for detailed provider support. haiku.rag supports thinking control for:
|
||||||
|
|
||||||
|
- **OpenAI**: Reasoning models (o1, o3, gpt-oss)
|
||||||
|
- **Anthropic**: All Claude models
|
||||||
|
- **Google**: Gemini models with thinking support
|
||||||
|
- **Groq**: Models with reasoning capabilities
|
||||||
|
- **Bedrock**: Claude, OpenAI, and Qwen models
|
||||||
|
- **Ollama**: Models supporting reasoning (gpt-oss, etc.)
|
||||||
|
- **vLLM**: Models supporting reasoning (gpt-oss, etc.)
|
||||||
|
|
||||||
|
**When to use:**
|
||||||
|
- Disable for simple queries, RAG workflows, speed-critical applications
|
||||||
|
- Enable for complex reasoning, mathematical problems, research tasks
|
||||||
|
|
||||||
|
## Embedding Providers
|
||||||
|
|
||||||
|
If you use Ollama, you can use any pulled model that supports embeddings.
|
||||||
|
|
||||||
|
### Ollama (Default)
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
embeddings:
|
||||||
|
provider: ollama
|
||||||
|
model: mxbai-embed-large
|
||||||
|
vector_dim: 1024
|
||||||
|
```
|
||||||
|
|
||||||
|
The Ollama base URL can be configured in your config file or via environment variable:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
providers:
|
||||||
|
ollama:
|
||||||
|
base_url: http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
Or via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export OLLAMA_BASE_URL=http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
If not configured, it defaults to `http://localhost:11434`.
|
||||||
|
|
||||||
|
### VoyageAI
|
||||||
|
|
||||||
|
If you installed `haiku.rag` (full package), VoyageAI is already included. If you installed `haiku.rag-slim`, install with VoyageAI extras:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv pip install haiku.rag-slim[voyageai]
|
||||||
|
```
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
embeddings:
|
||||||
|
provider: voyageai
|
||||||
|
model: voyage-3.5
|
||||||
|
vector_dim: 1024
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export VOYAGE_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
### OpenAI
|
||||||
|
|
||||||
|
OpenAI embeddings are included in the default installation:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
embeddings:
|
||||||
|
provider: openai
|
||||||
|
model: text-embedding-3-small # or text-embedding-3-large
|
||||||
|
vector_dim: 1536
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export OPENAI_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
### vLLM
|
||||||
|
|
||||||
|
For high-performance local inference, you can use vLLM to serve embedding models with OpenAI-compatible APIs:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
embeddings:
|
||||||
|
provider: vllm
|
||||||
|
model: mixedbread-ai/mxbai-embed-large-v1
|
||||||
|
vector_dim: 512
|
||||||
|
|
||||||
|
providers:
|
||||||
|
vllm:
|
||||||
|
embeddings_base_url: http://localhost:8000
|
||||||
|
```
|
||||||
|
|
||||||
|
**Note:** You need to run a vLLM server separately with an embedding model loaded.
|
||||||
|
|
||||||
|
## Question Answering Providers
|
||||||
|
|
||||||
|
Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used.
|
||||||
|
|
||||||
|
### Ollama (Default)
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: gpt-oss
|
||||||
|
```
|
||||||
|
|
||||||
|
The Ollama base URL can be configured via the `OLLAMA_BASE_URL` environment variable, config file, or defaults to `http://localhost:11434`:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export OLLAMA_BASE_URL=http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
Or in your config file:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
providers:
|
||||||
|
ollama:
|
||||||
|
base_url: http://localhost:11434
|
||||||
|
```
|
||||||
|
|
||||||
|
### OpenAI
|
||||||
|
|
||||||
|
OpenAI QA is included in the default installation:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: openai
|
||||||
|
name: gpt-4o-mini # or gpt-4, gpt-3.5-turbo, etc.
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export OPENAI_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
### Anthropic
|
||||||
|
|
||||||
|
Anthropic QA is included in the default installation:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: anthropic
|
||||||
|
name: claude-3-5-haiku-20241022 # or claude-3-5-sonnet-20241022, etc.
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export ANTHROPIC_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
### vLLM
|
||||||
|
|
||||||
|
For high-performance local inference:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: vllm
|
||||||
|
name: Qwen/Qwen3-4B # Any model with tool support in vLLM
|
||||||
|
|
||||||
|
providers:
|
||||||
|
vllm:
|
||||||
|
qa_base_url: http://localhost:8002
|
||||||
|
```
|
||||||
|
|
||||||
|
**Note:** You need to run a vLLM server separately with a model that supports tool calling loaded. Consult the specific model's documentation for proper vLLM serving configuration.
|
||||||
|
|
||||||
|
### Other Providers
|
||||||
|
|
||||||
|
Any provider supported by Pydantic AI can be used. Examples:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
# Google Gemini
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: gemini
|
||||||
|
name: gemini-1.5-flash
|
||||||
|
|
||||||
|
# Groq
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: groq
|
||||||
|
name: llama-3.3-70b-versatile
|
||||||
|
|
||||||
|
# Mistral
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: mistral
|
||||||
|
name: mistral-small-latest
|
||||||
|
```
|
||||||
|
|
||||||
|
See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the complete list of supported providers and models.
|
||||||
|
|
||||||
|
## Reranking Providers
|
||||||
|
|
||||||
|
Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (10x the requested limit) and then reranks them to return the most relevant results.
|
||||||
|
|
||||||
|
Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below.
|
||||||
|
|
||||||
|
### MixedBread AI
|
||||||
|
|
||||||
|
If you installed `haiku.rag` (full package), MxBAI is already included. If you installed `haiku.rag-slim`, add the mxbai extra:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv pip install haiku.rag-slim[mxbai]
|
||||||
|
```
|
||||||
|
|
||||||
|
Then configure:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: mxbai
|
||||||
|
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||||
|
```
|
||||||
|
|
||||||
|
### Cohere
|
||||||
|
|
||||||
|
If you installed `haiku.rag` (full package), Cohere is already included. If you installed `haiku.rag-slim`, add the cohere extra:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv pip install haiku.rag-slim[cohere]
|
||||||
|
```
|
||||||
|
|
||||||
|
Then configure:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: cohere
|
||||||
|
name: rerank-v3.5
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export CO_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
### Zero Entropy
|
||||||
|
|
||||||
|
If you installed `haiku.rag` (full package), Zero Entropy is already included. If you installed `haiku.rag-slim`, add the zeroentropy extra:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv pip install haiku.rag-slim[zeroentropy]
|
||||||
|
```
|
||||||
|
|
||||||
|
Then configure:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: zeroentropy
|
||||||
|
name: zerank-1 # Currently the only available model
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export ZEROENTROPY_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
### vLLM
|
||||||
|
|
||||||
|
For high-performance local reranking using dedicated reranking models:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: vllm
|
||||||
|
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||||
|
|
||||||
|
providers:
|
||||||
|
vllm:
|
||||||
|
rerank_base_url: http://localhost:8001
|
||||||
|
```
|
||||||
|
|
||||||
|
**Note:** vLLM reranking uses the `/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded. Consult the specific model's documentation for proper vLLM serving configuration.
|
||||||
|
|
@ -321,6 +321,6 @@ answer = await client.ask(
|
||||||
|
|
||||||
The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources. Citations prefer the document title when present, otherwise they use the URI.
|
The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources. Citations prefer the document title when present, otherwise they use the URI.
|
||||||
|
|
||||||
The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration.md)).
|
The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](config-index.md)).
|
||||||
|
|
||||||
See also: [Agents](agents.md) for details on the QA agent and the multi‑agent research workflow.
|
See also: [Agents](agents.md) for details on the QA agent and the multi‑agent research workflow.
|
||||||
|
|
|
||||||
71
docs/qa-research.md
Normal file
71
docs/qa-research.md
Normal file
|
|
@ -0,0 +1,71 @@
|
||||||
|
# QA and Research Configuration
|
||||||
|
|
||||||
|
## Question Answering Configuration
|
||||||
|
|
||||||
|
Configure the QA workflow:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
qa:
|
||||||
|
model:
|
||||||
|
provider: ollama
|
||||||
|
name: gpt-oss
|
||||||
|
enable_thinking: false
|
||||||
|
max_sub_questions: 3 # Maximum sub-questions for deep QA
|
||||||
|
max_iterations: 2 # Maximum search iterations per sub-question
|
||||||
|
max_concurrency: 1 # Sub-questions processed in parallel
|
||||||
|
```
|
||||||
|
|
||||||
|
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
|
||||||
|
- **max_sub_questions**: For deep QA mode, maximum number of sub-questions to generate (default: 3)
|
||||||
|
- **max_iterations**: Maximum search/evaluate cycles per sub-question (default: 2)
|
||||||
|
- **max_concurrency**: Number of sub-questions to process in parallel (default: 1)
|
||||||
|
|
||||||
|
Deep QA mode (`haiku-rag ask --deep`) decomposes complex questions into sub-questions, processes them in parallel batches, and synthesizes the results.
|
||||||
|
|
||||||
|
## Research Configuration
|
||||||
|
|
||||||
|
Configure the multi-agent research workflow:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
research:
|
||||||
|
model:
|
||||||
|
provider: "" # Empty to use qa settings
|
||||||
|
name: "" # Empty to use qa model
|
||||||
|
enable_thinking: false
|
||||||
|
max_iterations: 3
|
||||||
|
confidence_threshold: 0.8
|
||||||
|
max_concurrency: 1
|
||||||
|
```
|
||||||
|
|
||||||
|
- **model**: LLM configuration. Leave provider/model empty to inherit from `qa` (see [Providers](providers.md#model-settings))
|
||||||
|
- **max_iterations**: Maximum search/evaluate cycles (default: 3)
|
||||||
|
- **confidence_threshold**: Stop when confidence score meets/exceeds this (default: 0.8)
|
||||||
|
- **max_concurrency**: Sub-questions searched in parallel per iteration (default: 1)
|
||||||
|
|
||||||
|
The research workflow plans sub-questions, searches in parallel batches, evaluates findings, and iterates until reaching the confidence threshold or max iterations.
|
||||||
|
|
||||||
|
## AG-UI Server Configuration
|
||||||
|
|
||||||
|
Configure the AG-UI HTTP server for streaming graph execution events:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
agui:
|
||||||
|
host: "0.0.0.0"
|
||||||
|
port: 8000
|
||||||
|
cors_origins: ["*"]
|
||||||
|
cors_credentials: true
|
||||||
|
cors_methods: ["GET", "POST", "OPTIONS"]
|
||||||
|
cors_headers: ["*"]
|
||||||
|
```
|
||||||
|
|
||||||
|
Start the AG-UI server with:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-rag serve --agui
|
||||||
|
```
|
||||||
|
|
||||||
|
The server exposes:
|
||||||
|
- `GET /health` - Health check endpoint
|
||||||
|
- `POST /v1/agent/stream` - Research graph streaming endpoint (Server-Sent Events)
|
||||||
|
|
||||||
|
See [Server Mode](server.md) for more details.
|
||||||
|
|
@ -44,7 +44,7 @@ docker run -p 5001:5001 -e DOCLING_SERVE_ENABLE_UI=1 quay.io/docling-project/doc
|
||||||
|
|
||||||
### Configuration
|
### Configuration
|
||||||
|
|
||||||
Configure haiku.rag to use docling-serve. See the [Document Processing section in Configuration](configuration.md#document-processing) for all available options.
|
Configure haiku.rag to use docling-serve. See the [Document Processing](processing.md) guide for all available options.
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
# haiku.rag.yaml
|
# haiku.rag.yaml
|
||||||
|
|
|
||||||
|
|
@ -122,7 +122,7 @@ agui:
|
||||||
cors_credentials: true
|
cors_credentials: true
|
||||||
```
|
```
|
||||||
|
|
||||||
See [Configuration](configuration.md#ag-ui-server-configuration) for all available options.
|
See [Configuration](qa-research.md#ag-ui-server-configuration) for all available options.
|
||||||
|
|
||||||
### Using the Streaming Endpoints
|
### Using the Streaming Endpoints
|
||||||
|
|
||||||
|
|
|
||||||
103
docs/storage.md
Normal file
103
docs/storage.md
Normal file
|
|
@ -0,0 +1,103 @@
|
||||||
|
# Database and Storage
|
||||||
|
|
||||||
|
## Local Storage
|
||||||
|
|
||||||
|
By default, `haiku.rag` uses a local LanceDB database:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
storage:
|
||||||
|
data_dir: /path/to/data # Empty = use default platform location
|
||||||
|
vacuum_retention_seconds: 86400 # Cleanup threshold in seconds
|
||||||
|
```
|
||||||
|
|
||||||
|
- **data_dir**: Directory for local database storage. When empty, uses platform-specific default locations
|
||||||
|
- **vacuum_retention_seconds**: When documents are added/updated, old table versions older than this are removed. Default: 86400 seconds (1 day, safe for concurrent connections). Set to 0 for aggressive cleanup (removes all old versions immediately)
|
||||||
|
|
||||||
|
## Remote Storage
|
||||||
|
|
||||||
|
For remote storage, use the `lancedb` settings with various backends:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
# LanceDB Cloud
|
||||||
|
lancedb:
|
||||||
|
uri: db://your-database-name
|
||||||
|
api_key: your-api-key
|
||||||
|
region: us-west-2 # optional
|
||||||
|
|
||||||
|
# Amazon S3
|
||||||
|
lancedb:
|
||||||
|
uri: s3://my-bucket/my-table
|
||||||
|
# Use AWS credentials or IAM roles
|
||||||
|
|
||||||
|
# Azure Blob Storage
|
||||||
|
lancedb:
|
||||||
|
uri: az://my-container/my-table
|
||||||
|
# Use Azure credentials
|
||||||
|
|
||||||
|
# Google Cloud Storage
|
||||||
|
lancedb:
|
||||||
|
uri: gs://my-bucket/my-table
|
||||||
|
# Use GCP credentials
|
||||||
|
|
||||||
|
# HDFS
|
||||||
|
lancedb:
|
||||||
|
uri: hdfs://namenode:port/path/to/table
|
||||||
|
```
|
||||||
|
|
||||||
|
Authentication is handled through standard cloud provider credentials (AWS CLI, Azure CLI, gcloud, etc.) or by setting `api_key` for LanceDB Cloud.
|
||||||
|
|
||||||
|
**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization is still performed locally.
|
||||||
|
|
||||||
|
## Database Auto-creation
|
||||||
|
|
||||||
|
haiku.rag intelligently handles database creation based on operation type:
|
||||||
|
|
||||||
|
- **Write operations** (add, add-src, delete, rebuild): Automatically create the database and required tables if they don't exist
|
||||||
|
- **Read operations** (list, get, search, ask, research): Fail with a clear error if the database doesn't exist
|
||||||
|
|
||||||
|
This prevents the common mistake where a search query accidentally creates an empty database. To initialize your database, simply add your first document using `haiku-rag add` or `haiku-rag add-src`.
|
||||||
|
|
||||||
|
## Vector Indexing
|
||||||
|
|
||||||
|
Configure vector indexing behavior for efficient similarity search:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
search:
|
||||||
|
vector_index_metric: cosine # cosine, l2, or dot
|
||||||
|
vector_refine_factor: 30 # Re-ranking factor for accuracy
|
||||||
|
```
|
||||||
|
|
||||||
|
- **vector_index_metric**: Distance metric for vector similarity:
|
||||||
|
- `cosine`: Cosine similarity (default, best for most embeddings)
|
||||||
|
- `l2`: Euclidean distance
|
||||||
|
- `dot`: Dot product similarity
|
||||||
|
- **vector_refine_factor**: Improves accuracy when using a vector index by retrieving `refine_factor * limit` candidates (using approximate search) and re-ranking them with exact distances. Higher values increase accuracy but slow down queries. Default: 30
|
||||||
|
- **Only applies with a vector index** - has no effect on brute-force search, which already returns exact results
|
||||||
|
|
||||||
|
!!! note
|
||||||
|
Vector indexes are only necessary for large datasets with over 100,000 chunks. For smaller datasets, LanceDB's brute-force kNN search provides exact results with good performance. Only create an index if you notice search performance degradation on large datasets.
|
||||||
|
|
||||||
|
**Index creation:**
|
||||||
|
|
||||||
|
Vector indexes are **not created automatically** during document ingestion to avoid slowing down the process. After you've added documents (at least 256 chunks required), create the index manually:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-rag create-index
|
||||||
|
```
|
||||||
|
|
||||||
|
This command:
|
||||||
|
- Checks if you have enough data (minimum 256 chunks)
|
||||||
|
- Creates an IVF_PQ index for fast approximate nearest neighbor (ANN) search
|
||||||
|
- Uses LanceDB's automatic parameter calculation based on your dataset size and vector dimensions
|
||||||
|
|
||||||
|
**Re-indexing:**
|
||||||
|
|
||||||
|
Indexes are not automatically updated when you add new documents. After adding a significant amount of new data:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-rag create-index # Rebuilds the index with all data
|
||||||
|
```
|
||||||
|
|
||||||
|
Searches still work with stale indexes - LanceDB uses the index for old data (fast ANN) and brute-force kNN for new unindexed rows, then combines the results. However, performance degrades as more unindexed data accumulates.
|
||||||
|
|
||||||
|
For datasets with fewer than 256 chunks, searches use brute-force kNN scans (exact nearest neighbors, 100% recall) which work well for small datasets but don't scale beyond a few hundred thousand vectors.
|
||||||
|
|
@ -38,8 +38,9 @@ embeddings:
|
||||||
vector_dim: 1536
|
vector_dim: 1536
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
model:
|
||||||
provider: openai
|
provider: openai
|
||||||
model: gpt-4o-mini # or gpt-4o, gpt-4, etc.
|
name: gpt-4o-mini # or gpt-4o, gpt-4, etc.
|
||||||
```
|
```
|
||||||
|
|
||||||
Set your OpenAI API key as an environment variable (API keys should not be stored in the YAML file):
|
Set your OpenAI API key as an environment variable (API keys should not be stored in the YAML file):
|
||||||
|
|
@ -50,7 +51,7 @@ export OPENAI_API_KEY="<your OpenAI API key>"
|
||||||
|
|
||||||
For the list of available OpenAI models and their vector dimensions, see the [OpenAI documentation](https://platform.openai.com/docs/guides/embeddings).
|
For the list of available OpenAI models and their vector dimensions, see the [OpenAI documentation](https://platform.openai.com/docs/guides/embeddings).
|
||||||
|
|
||||||
See [Configuration](configuration.md) for all available options.
|
See [Configuration](config-index.md) for all available options.
|
||||||
|
|
||||||
## Adding the first documents
|
## Adding the first documents
|
||||||
|
|
||||||
|
|
@ -215,4 +216,4 @@ The following people are presenting talks at PyCon Finland 2025:
|
||||||
|
|
||||||
## Configuration
|
## Configuration
|
||||||
|
|
||||||
See [Configuration page](./configuration.md) for complete documentation on YAML configuration and all available options.
|
See [Configuration page](./config-index.md) for complete documentation on YAML configuration and all available options.
|
||||||
|
|
|
||||||
|
|
@ -6,8 +6,6 @@ from typing import Any, cast
|
||||||
import logfire
|
import logfire
|
||||||
import typer
|
import typer
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from pydantic_ai.models.openai import OpenAIChatModel
|
|
||||||
from pydantic_ai.providers.ollama import OllamaProvider
|
|
||||||
from pydantic_evals import Case, Dataset as EvalDataset
|
from pydantic_evals import Case, Dataset as EvalDataset
|
||||||
from pydantic_evals.evaluators import LLMJudge
|
from pydantic_evals.evaluators import LLMJudge
|
||||||
from pydantic_evals.reporting import ReportCaseFailure
|
from pydantic_evals.reporting import ReportCaseFailure
|
||||||
|
|
@ -20,13 +18,13 @@ from evaluations.evaluators import ANSWER_EQUIVALENCE_RUBRIC
|
||||||
from evaluations.prompts import WIX_SUPPORT_PROMPT
|
from evaluations.prompts import WIX_SUPPORT_PROMPT
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
|
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
from haiku.rag.logging import configure_cli_logging
|
from haiku.rag.logging import configure_cli_logging
|
||||||
from haiku.rag.qa import get_qa_agent
|
from haiku.rag.qa import get_qa_agent
|
||||||
|
from haiku.rag.utils import get_model
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
QA_JUDGE_MODEL = "qwen3"
|
|
||||||
|
|
||||||
logfire.configure(send_to_logfire="if-token-present", service_name="evals")
|
logfire.configure(send_to_logfire="if-token-present", service_name="evals")
|
||||||
logfire.instrument_pydantic_ai()
|
logfire.instrument_pydantic_ai()
|
||||||
configure_cli_logging()
|
configure_cli_logging()
|
||||||
|
|
@ -37,7 +35,7 @@ def build_experiment_metadata(
|
||||||
dataset_key: str,
|
dataset_key: str,
|
||||||
test_cases: int,
|
test_cases: int,
|
||||||
config: AppConfig,
|
config: AppConfig,
|
||||||
judge_model: str,
|
judge_config: ModelConfig,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
"""Build experiment metadata for Logfire tracking."""
|
"""Build experiment metadata for Logfire tracking."""
|
||||||
return {
|
return {
|
||||||
|
|
@ -48,12 +46,20 @@ def build_experiment_metadata(
|
||||||
"embedder_dim": config.embeddings.vector_dim,
|
"embedder_dim": config.embeddings.vector_dim,
|
||||||
"chunk_size": config.processing.chunk_size,
|
"chunk_size": config.processing.chunk_size,
|
||||||
"context_chunk_radius": config.processing.context_chunk_radius,
|
"context_chunk_radius": config.processing.context_chunk_radius,
|
||||||
"rerank_provider": config.reranking.provider,
|
"rerank_provider": config.reranking.model.provider
|
||||||
"rerank_model": config.reranking.model,
|
if config.reranking.model
|
||||||
"qa_provider": config.qa.provider,
|
else None,
|
||||||
"qa_model": config.qa.model,
|
"rerank_model": config.reranking.model.name if config.reranking.model else None,
|
||||||
"judge_provider": "ollama",
|
"qa_provider": config.qa.model.provider,
|
||||||
"judge_model": judge_model,
|
"qa_model": config.qa.model.name,
|
||||||
|
"qa_temperature": config.qa.model.temperature,
|
||||||
|
"qa_max_tokens": config.qa.model.max_tokens,
|
||||||
|
"qa_enable_thinking": config.qa.model.enable_thinking,
|
||||||
|
"judge_provider": judge_config.provider,
|
||||||
|
"judge_model": judge_config.name,
|
||||||
|
"judge_temperature": judge_config.temperature,
|
||||||
|
"judge_max_tokens": judge_config.max_tokens,
|
||||||
|
"judge_enable_thinking": judge_config.enable_thinking,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -161,11 +167,14 @@ async def run_retrieval_benchmark(
|
||||||
|
|
||||||
eval_name = name if name is not None else f"{spec.key}_retrieval_evaluation"
|
eval_name = name if name is not None else f"{spec.key}_retrieval_evaluation"
|
||||||
|
|
||||||
|
judge_config = ModelConfig(
|
||||||
|
provider="ollama", name="gpt-oss", enable_thinking=False
|
||||||
|
)
|
||||||
experiment_metadata = build_experiment_metadata(
|
experiment_metadata = build_experiment_metadata(
|
||||||
dataset_key=spec.key,
|
dataset_key=spec.key,
|
||||||
test_cases=len(cases),
|
test_cases=len(cases),
|
||||||
config=config,
|
config=config,
|
||||||
judge_model=QA_JUDGE_MODEL,
|
judge_config=judge_config,
|
||||||
)
|
)
|
||||||
|
|
||||||
report = await dataset.evaluate(
|
report = await dataset.evaluate(
|
||||||
|
|
@ -213,10 +222,8 @@ async def run_qa_benchmark(
|
||||||
for index, doc in enumerate(corpus, start=1)
|
for index, doc in enumerate(corpus, start=1)
|
||||||
]
|
]
|
||||||
|
|
||||||
judge_model = OpenAIChatModel(
|
judge_config = ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=False)
|
||||||
model_name=QA_JUDGE_MODEL,
|
judge_model = get_model(judge_config, config)
|
||||||
provider=OllamaProvider(base_url=f"{config.providers.ollama.base_url}/v1"),
|
|
||||||
)
|
|
||||||
|
|
||||||
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
|
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
|
||||||
name=spec.key,
|
name=spec.key,
|
||||||
|
|
@ -249,7 +256,7 @@ async def run_qa_benchmark(
|
||||||
dataset_key=spec.key,
|
dataset_key=spec.key,
|
||||||
test_cases=len(cases),
|
test_cases=len(cases),
|
||||||
config=config,
|
config=config,
|
||||||
judge_model=QA_JUDGE_MODEL,
|
judge_config=judge_config,
|
||||||
)
|
)
|
||||||
|
|
||||||
report = await evaluation_dataset.evaluate(
|
report = await evaluation_dataset.evaluate(
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,9 @@
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from pydantic_ai import Agent
|
from pydantic_ai import Agent
|
||||||
from pydantic_ai.models.openai import OpenAIChatModel
|
|
||||||
from pydantic_ai.providers.ollama import OllamaProvider
|
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
from haiku.rag.utils import get_model
|
||||||
|
|
||||||
ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent.
|
ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent.
|
||||||
|
|
||||||
|
|
@ -37,15 +37,15 @@ class LLMJudge:
|
||||||
"""LLM-as-judge for evaluating answer equivalence using Pydantic AI."""
|
"""LLM-as-judge for evaluating answer equivalence using Pydantic AI."""
|
||||||
|
|
||||||
def __init__(self, model: str = "gpt-oss"):
|
def __init__(self, model: str = "gpt-oss"):
|
||||||
# Create Ollama model
|
# Create model using get_model with thinking disabled
|
||||||
ollama_model = OpenAIChatModel(
|
model_config = ModelConfig(
|
||||||
model_name=model,
|
provider="ollama", model=model, enable_thinking=False
|
||||||
provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"),
|
|
||||||
)
|
)
|
||||||
|
model_obj = get_model(model_config, Config)
|
||||||
|
|
||||||
# Create Pydantic AI agent
|
# Create Pydantic AI agent
|
||||||
self._agent = Agent(
|
self._agent = Agent(
|
||||||
model=ollama_model,
|
model=model_obj,
|
||||||
output_type=LLMJudgeResponseSchema,
|
output_type=LLMJudgeResponseSchema,
|
||||||
system_prompt=ANSWER_EQUIVALENCE_RUBRIC,
|
system_prompt=ANSWER_EQUIVALENCE_RUBRIC,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
|
||||||
|
|
@ -65,7 +65,7 @@ def create_a2a_app(
|
||||||
broker = InMemoryBroker()
|
broker = InMemoryBroker()
|
||||||
|
|
||||||
# Create the agent with native search tool
|
# Create the agent with native search tool
|
||||||
model = get_model(config.qa.provider, config.qa.model)
|
model = get_model(config.qa.model, config)
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=model,
|
model=model,
|
||||||
deps_type=AgentDependencies,
|
deps_type=AgentDependencies,
|
||||||
|
|
|
||||||
|
|
@ -35,7 +35,7 @@ class AgentDeps:
|
||||||
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
|
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
|
||||||
|
|
||||||
|
|
||||||
model = get_model(Config.research.provider, Config.research.model)
|
model = get_model(Config.research.model, Config)
|
||||||
|
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model,
|
model,
|
||||||
|
|
|
||||||
|
|
@ -47,7 +47,7 @@ if not db_path.exists():
|
||||||
|
|
||||||
logger.info(f"Initializing research assistant with database: {db_path}")
|
logger.info(f"Initializing research assistant with database: {db_path}")
|
||||||
logger.info(
|
logger.info(
|
||||||
f"Research Provider: {Config.research.provider}, Model: {Config.research.model}"
|
f"Research Provider: {Config.research.model.provider}, Model: {Config.research.model.name}"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Store client reference for proper lifecycle management
|
# Store client reference for proper lifecycle management
|
||||||
|
|
@ -153,8 +153,8 @@ async def health_check(_: Request) -> JSONResponse:
|
||||||
{
|
{
|
||||||
"status": "healthy",
|
"status": "healthy",
|
||||||
"agent_model": str(agent.model),
|
"agent_model": str(agent.model),
|
||||||
"research_provider": Config.research.provider,
|
"research_provider": Config.research.model.provider,
|
||||||
"research_model": Config.research.model,
|
"research_model": Config.research.model.name,
|
||||||
"db_path": str(db_path),
|
"db_path": str(db_path),
|
||||||
"db_exists": db_path.exists(),
|
"db_exists": db_path.exists(),
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -17,18 +17,21 @@ providers:
|
||||||
base_url: http://host.docker.internal:11434
|
base_url: http://host.docker.internal:11434
|
||||||
|
|
||||||
research:
|
research:
|
||||||
|
model:
|
||||||
provider: ollama
|
provider: ollama
|
||||||
model: gpt-oss:latest
|
name: gpt-oss:latest
|
||||||
max_iterations: 3
|
max_iterations: 3
|
||||||
confidence_threshold: 0.8
|
confidence_threshold: 0.8
|
||||||
max_concurrency: 1
|
max_concurrency: 1
|
||||||
|
|
||||||
# For OpenAI:
|
# For OpenAI:
|
||||||
# research:
|
# research:
|
||||||
|
# model:
|
||||||
# provider: openai
|
# provider: openai
|
||||||
# model: gpt-4o-mini
|
# name: gpt-4o-mini
|
||||||
|
|
||||||
# For Anthropic:
|
# For Anthropic:
|
||||||
# research:
|
# research:
|
||||||
|
# model:
|
||||||
# provider: anthropic
|
# provider: anthropic
|
||||||
# model: claude-3-5-haiku-20241022
|
# name: claude-3-5-haiku-20241022
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,25 @@ from pydantic import BaseModel, Field
|
||||||
from haiku.rag.utils import get_default_data_dir
|
from haiku.rag.utils import get_default_data_dir
|
||||||
|
|
||||||
|
|
||||||
|
class ModelConfig(BaseModel):
|
||||||
|
"""Configuration for a language model.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
provider: Model provider (ollama, openai, anthropic, etc.)
|
||||||
|
name: Model name/identifier
|
||||||
|
enable_thinking: Control reasoning behavior (true/false/None for default)
|
||||||
|
temperature: Sampling temperature (0.0 to 1.0+)
|
||||||
|
max_tokens: Maximum tokens to generate
|
||||||
|
"""
|
||||||
|
|
||||||
|
provider: str = "ollama"
|
||||||
|
name: str = "gpt-oss"
|
||||||
|
|
||||||
|
enable_thinking: bool | None = None
|
||||||
|
temperature: float | None = None
|
||||||
|
max_tokens: int | None = None
|
||||||
|
|
||||||
|
|
||||||
class StorageConfig(BaseModel):
|
class StorageConfig(BaseModel):
|
||||||
data_dir: Path = Field(default_factory=get_default_data_dir)
|
data_dir: Path = Field(default_factory=get_default_data_dir)
|
||||||
vacuum_retention_seconds: int = 86400
|
vacuum_retention_seconds: int = 86400
|
||||||
|
|
@ -31,21 +50,30 @@ class EmbeddingsConfig(BaseModel):
|
||||||
|
|
||||||
|
|
||||||
class RerankingConfig(BaseModel):
|
class RerankingConfig(BaseModel):
|
||||||
provider: str = ""
|
model: ModelConfig | None = None
|
||||||
model: str = ""
|
|
||||||
|
|
||||||
|
|
||||||
class QAConfig(BaseModel):
|
class QAConfig(BaseModel):
|
||||||
provider: str = "ollama"
|
model: ModelConfig = Field(
|
||||||
model: str = "gpt-oss"
|
default_factory=lambda: ModelConfig(
|
||||||
|
provider="ollama",
|
||||||
|
name="gpt-oss",
|
||||||
|
enable_thinking=False,
|
||||||
|
)
|
||||||
|
)
|
||||||
max_sub_questions: int = 3
|
max_sub_questions: int = 3
|
||||||
max_iterations: int = 2
|
max_iterations: int = 2
|
||||||
max_concurrency: int = 1
|
max_concurrency: int = 1
|
||||||
|
|
||||||
|
|
||||||
class ResearchConfig(BaseModel):
|
class ResearchConfig(BaseModel):
|
||||||
provider: str = "ollama"
|
model: ModelConfig = Field(
|
||||||
model: str = "gpt-oss"
|
default_factory=lambda: ModelConfig(
|
||||||
|
provider="ollama",
|
||||||
|
name="gpt-oss",
|
||||||
|
enable_thinking=False,
|
||||||
|
)
|
||||||
|
)
|
||||||
max_iterations: int = 3
|
max_iterations: int = 3
|
||||||
confidence_threshold: float = 0.8
|
confidence_threshold: float = 0.8
|
||||||
max_concurrency: int = 1
|
max_concurrency: int = 1
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,5 @@
|
||||||
"""Common utilities for graph implementations."""
|
"""Common utilities for graph implementations."""
|
||||||
|
|
||||||
from haiku.rag.graph.common.utils import get_model
|
from haiku.rag.utils import get_model
|
||||||
|
|
||||||
__all__ = ["get_model"]
|
__all__ = ["get_model"]
|
||||||
|
|
|
||||||
|
|
@ -11,7 +11,7 @@ from pydantic_graph.beta import StepContext
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
from haiku.rag.config.models import AppConfig
|
from haiku.rag.config.models import AppConfig, ModelConfig
|
||||||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||||
from haiku.rag.graph.common import get_model
|
from haiku.rag.graph.common import get_model
|
||||||
from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
|
from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
|
||||||
|
|
@ -52,8 +52,7 @@ class GraphAgentDeps(Protocol):
|
||||||
|
|
||||||
|
|
||||||
def create_plan_node[AgentDepsT: GraphAgentDeps](
|
def create_plan_node[AgentDepsT: GraphAgentDeps](
|
||||||
provider: str,
|
model_config: ModelConfig,
|
||||||
model: str,
|
|
||||||
deps_type: type[AgentDepsT],
|
deps_type: type[AgentDepsT],
|
||||||
activity_message: str = "Creating plan",
|
activity_message: str = "Creating plan",
|
||||||
output_retries: int | None = None,
|
output_retries: int | None = None,
|
||||||
|
|
@ -62,8 +61,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
|
||||||
"""Create a plan node for any graph.
|
"""Create a plan node for any graph.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
provider: Model provider (e.g., 'openai', 'anthropic')
|
model_config: ModelConfig with provider, model, and settings
|
||||||
model: Model name
|
|
||||||
deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
|
deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
|
||||||
activity_message: Message to show during planning activity
|
activity_message: Message to show during planning activity
|
||||||
output_retries: Number of output retries for the agent (optional)
|
output_retries: Number of output retries for the agent (optional)
|
||||||
|
|
@ -86,7 +84,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
|
||||||
try:
|
try:
|
||||||
# Build agent configuration
|
# Build agent configuration
|
||||||
agent_config = {
|
agent_config = {
|
||||||
"model": get_model(provider, model, config),
|
"model": get_model(model_config, config),
|
||||||
"output_type": ResearchPlan,
|
"output_type": ResearchPlan,
|
||||||
"instructions": (
|
"instructions": (
|
||||||
PLAN_PROMPT
|
PLAN_PROMPT
|
||||||
|
|
@ -141,8 +139,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
|
||||||
|
|
||||||
|
|
||||||
def create_search_node[AgentDepsT: GraphAgentDeps](
|
def create_search_node[AgentDepsT: GraphAgentDeps](
|
||||||
provider: str,
|
model_config: ModelConfig,
|
||||||
model: str,
|
|
||||||
deps_type: type[AgentDepsT],
|
deps_type: type[AgentDepsT],
|
||||||
with_step_wrapper: bool = True,
|
with_step_wrapper: bool = True,
|
||||||
success_message_format: str = "Answered: {sub_q}",
|
success_message_format: str = "Answered: {sub_q}",
|
||||||
|
|
@ -152,8 +149,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps](
|
||||||
"""Create a search_one node for any graph.
|
"""Create a search_one node for any graph.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
provider: Model provider
|
model_config: ModelConfig with provider, model, and settings
|
||||||
model: Model name
|
|
||||||
deps_type: Type of dependencies for the agent
|
deps_type: Type of dependencies for the agent
|
||||||
with_step_wrapper: Whether to wrap with agui_emitter start/finish step
|
with_step_wrapper: Whether to wrap with agui_emitter start/finish step
|
||||||
success_message_format: Format string for success activity message
|
success_message_format: Format string for success activity message
|
||||||
|
|
@ -186,8 +182,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps](
|
||||||
state,
|
state,
|
||||||
deps,
|
deps,
|
||||||
sub_q,
|
sub_q,
|
||||||
provider,
|
model_config,
|
||||||
model,
|
|
||||||
deps_type,
|
deps_type,
|
||||||
success_message_format,
|
success_message_format,
|
||||||
handle_exceptions,
|
handle_exceptions,
|
||||||
|
|
@ -204,8 +199,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
|
||||||
state: GraphState,
|
state: GraphState,
|
||||||
deps: GraphDeps,
|
deps: GraphDeps,
|
||||||
sub_q: str,
|
sub_q: str,
|
||||||
provider: str,
|
model_config: ModelConfig,
|
||||||
model: str,
|
|
||||||
deps_type: type[AgentDepsT],
|
deps_type: type[AgentDepsT],
|
||||||
success_message_format: str,
|
success_message_format: str,
|
||||||
handle_exceptions: bool,
|
handle_exceptions: bool,
|
||||||
|
|
@ -223,7 +217,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
|
||||||
)
|
)
|
||||||
|
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=get_model(provider, model, config),
|
model=get_model(model_config, config),
|
||||||
output_type=ToolOutput(SearchAnswer, max_retries=3),
|
output_type=ToolOutput(SearchAnswer, max_retries=3),
|
||||||
instructions=SEARCH_AGENT_PROMPT,
|
instructions=SEARCH_AGENT_PROMPT,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
|
||||||
|
|
@ -1,48 +0,0 @@
|
||||||
"""Common utilities for all graph implementations."""
|
|
||||||
|
|
||||||
from pydantic_ai.models.openai import OpenAIChatModel
|
|
||||||
from pydantic_ai.providers.ollama import OllamaProvider
|
|
||||||
from pydantic_ai.providers.openai import OpenAIProvider
|
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.config.models import AppConfig
|
|
||||||
|
|
||||||
|
|
||||||
def get_model(
|
|
||||||
provider: str, model: str, config: AppConfig = Config
|
|
||||||
) -> OpenAIChatModel | str:
|
|
||||||
"""
|
|
||||||
Get a model instance for the specified provider and model name.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
provider: The model provider ("ollama", "vllm", or other)
|
|
||||||
model: The model name
|
|
||||||
config: AppConfig object (defaults to global Config)
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
A configured model instance
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
ValueError: If the provider is unknown
|
|
||||||
"""
|
|
||||||
if provider == "ollama":
|
|
||||||
return OpenAIChatModel(
|
|
||||||
model_name=model,
|
|
||||||
provider=OllamaProvider(base_url=f"{config.providers.ollama.base_url}/v1"),
|
|
||||||
)
|
|
||||||
elif provider == "vllm":
|
|
||||||
return OpenAIChatModel(
|
|
||||||
model_name=model,
|
|
||||||
provider=OpenAIProvider(
|
|
||||||
base_url=f"{config.providers.vllm.research_base_url or config.providers.vllm.qa_base_url}/v1",
|
|
||||||
api_key="none",
|
|
||||||
),
|
|
||||||
)
|
|
||||||
elif provider in ("openai", "anthropic", "gemini", "groq", "bedrock"):
|
|
||||||
# These providers use string format
|
|
||||||
return f"{provider}:{model}"
|
|
||||||
else:
|
|
||||||
raise ValueError(
|
|
||||||
f"Unknown model provider: {provider}. "
|
|
||||||
f"Supported providers: ollama, vllm, openai, anthropic, gemini, groq, bedrock"
|
|
||||||
)
|
|
||||||
|
|
@ -29,8 +29,7 @@ def build_deep_qa_graph(
|
||||||
Returns:
|
Returns:
|
||||||
Configured Deep QA graph
|
Configured Deep QA graph
|
||||||
"""
|
"""
|
||||||
provider = config.qa.provider
|
model_config = config.qa.model
|
||||||
model = config.qa.model
|
|
||||||
g = GraphBuilder(
|
g = GraphBuilder(
|
||||||
state_type=DeepQAState,
|
state_type=DeepQAState,
|
||||||
deps_type=DeepQADeps,
|
deps_type=DeepQADeps,
|
||||||
|
|
@ -40,8 +39,7 @@ def build_deep_qa_graph(
|
||||||
# Create and register the plan node using the factory
|
# Create and register the plan node using the factory
|
||||||
plan = g.step(
|
plan = g.step(
|
||||||
create_plan_node(
|
create_plan_node(
|
||||||
provider=provider,
|
model_config=model_config,
|
||||||
model=model,
|
|
||||||
deps_type=DeepQADependencies, # type: ignore[arg-type]
|
deps_type=DeepQADependencies, # type: ignore[arg-type]
|
||||||
activity_message="Planning approach",
|
activity_message="Planning approach",
|
||||||
output_retries=None, # Deep QA doesn't use output_retries
|
output_retries=None, # Deep QA doesn't use output_retries
|
||||||
|
|
@ -52,8 +50,7 @@ def build_deep_qa_graph(
|
||||||
# Create and register the search_one node using the factory
|
# Create and register the search_one node using the factory
|
||||||
search_one = g.step(
|
search_one = g.step(
|
||||||
create_search_node(
|
create_search_node(
|
||||||
provider=provider,
|
model_config=model_config,
|
||||||
model=model,
|
|
||||||
deps_type=DeepQADependencies, # type: ignore[arg-type]
|
deps_type=DeepQADependencies, # type: ignore[arg-type]
|
||||||
with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step
|
with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step
|
||||||
success_message_format="Answered: {sub_q}",
|
success_message_format="Answered: {sub_q}",
|
||||||
|
|
@ -92,7 +89,7 @@ def build_deep_qa_graph(
|
||||||
|
|
||||||
try:
|
try:
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=get_model(provider, model, config),
|
model=get_model(model_config, config),
|
||||||
output_type=DeepQAEvaluation,
|
output_type=DeepQAEvaluation,
|
||||||
instructions=DECISION_PROMPT,
|
instructions=DECISION_PROMPT,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
@ -173,7 +170,7 @@ def build_deep_qa_graph(
|
||||||
)
|
)
|
||||||
|
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=get_model(provider, model, config),
|
model=get_model(model_config, config),
|
||||||
output_type=DeepQAAnswer,
|
output_type=DeepQAAnswer,
|
||||||
instructions=prompt_template,
|
instructions=prompt_template,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
|
||||||
|
|
@ -36,8 +36,7 @@ def build_research_graph(
|
||||||
Returns:
|
Returns:
|
||||||
Configured Research graph
|
Configured Research graph
|
||||||
"""
|
"""
|
||||||
provider = config.research.provider
|
model_config = config.research.model
|
||||||
model = config.research.model
|
|
||||||
g = GraphBuilder(
|
g = GraphBuilder(
|
||||||
state_type=ResearchState,
|
state_type=ResearchState,
|
||||||
deps_type=ResearchDeps,
|
deps_type=ResearchDeps,
|
||||||
|
|
@ -47,8 +46,7 @@ def build_research_graph(
|
||||||
# Create and register the plan node using the factory
|
# Create and register the plan node using the factory
|
||||||
plan = g.step(
|
plan = g.step(
|
||||||
create_plan_node(
|
create_plan_node(
|
||||||
provider=provider,
|
model_config=model_config,
|
||||||
model=model,
|
|
||||||
deps_type=ResearchDependencies, # type: ignore[arg-type]
|
deps_type=ResearchDependencies, # type: ignore[arg-type]
|
||||||
activity_message="Creating research plan",
|
activity_message="Creating research plan",
|
||||||
output_retries=3,
|
output_retries=3,
|
||||||
|
|
@ -59,8 +57,7 @@ def build_research_graph(
|
||||||
# Create and register the search_one node using the factory
|
# Create and register the search_one node using the factory
|
||||||
search_one = g.step(
|
search_one = g.step(
|
||||||
create_search_node(
|
create_search_node(
|
||||||
provider=provider,
|
model_config=model_config,
|
||||||
model=model,
|
|
||||||
deps_type=ResearchDependencies, # type: ignore[arg-type]
|
deps_type=ResearchDependencies, # type: ignore[arg-type]
|
||||||
with_step_wrapper=True,
|
with_step_wrapper=True,
|
||||||
success_message_format="Found answer with {confidence:.0%} confidence",
|
success_message_format="Found answer with {confidence:.0%} confidence",
|
||||||
|
|
@ -99,7 +96,7 @@ def build_research_graph(
|
||||||
|
|
||||||
try:
|
try:
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=get_model(provider, model, config),
|
model=get_model(model_config, config),
|
||||||
output_type=InsightAnalysis,
|
output_type=InsightAnalysis,
|
||||||
instructions=INSIGHT_AGENT_PROMPT,
|
instructions=INSIGHT_AGENT_PROMPT,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
@ -168,7 +165,7 @@ def build_research_graph(
|
||||||
|
|
||||||
try:
|
try:
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=get_model(provider, model, config),
|
model=get_model(model_config, config),
|
||||||
output_type=EvaluationResult,
|
output_type=EvaluationResult,
|
||||||
instructions=DECISION_AGENT_PROMPT,
|
instructions=DECISION_AGENT_PROMPT,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
@ -247,7 +244,7 @@ def build_research_graph(
|
||||||
|
|
||||||
try:
|
try:
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=get_model(provider, model, config),
|
model=get_model(model_config, config),
|
||||||
output_type=ResearchReport,
|
output_type=ResearchReport,
|
||||||
instructions=SYNTHESIS_AGENT_PROMPT,
|
instructions=SYNTHESIS_AGENT_PROMPT,
|
||||||
retries=3,
|
retries=3,
|
||||||
|
|
|
||||||
|
|
@ -21,13 +21,9 @@ def get_qa_agent(
|
||||||
Returns:
|
Returns:
|
||||||
A configured QuestionAnswerAgent instance.
|
A configured QuestionAnswerAgent instance.
|
||||||
"""
|
"""
|
||||||
provider = config.qa.provider
|
|
||||||
model_name = config.qa.model
|
|
||||||
|
|
||||||
return QuestionAnswerAgent(
|
return QuestionAnswerAgent(
|
||||||
client=client,
|
client=client,
|
||||||
provider=provider,
|
model_config=config.qa.model,
|
||||||
model=model_name,
|
|
||||||
use_citations=use_citations,
|
use_citations=use_citations,
|
||||||
system_prompt=system_prompt,
|
system_prompt=system_prompt,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1,10 @@
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
from pydantic_ai import Agent, RunContext
|
from pydantic_ai import Agent, RunContext
|
||||||
from pydantic_ai.models.openai import OpenAIChatModel
|
|
||||||
from pydantic_ai.providers.ollama import OllamaProvider
|
|
||||||
from pydantic_ai.providers.openai import OpenAIProvider
|
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
from haiku.rag.graph.common import get_model
|
||||||
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
|
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -26,8 +25,7 @@ class QuestionAnswerAgent:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
client: HaikuRAG,
|
client: HaikuRAG,
|
||||||
provider: str,
|
model_config: ModelConfig,
|
||||||
model: str,
|
|
||||||
use_citations: bool = False,
|
use_citations: bool = False,
|
||||||
q: float = 0.0,
|
q: float = 0.0,
|
||||||
system_prompt: str | None = None,
|
system_prompt: str | None = None,
|
||||||
|
|
@ -38,7 +36,7 @@ class QuestionAnswerAgent:
|
||||||
system_prompt = (
|
system_prompt = (
|
||||||
QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
|
QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
|
||||||
)
|
)
|
||||||
model_obj = self._get_model(provider, model)
|
model_obj = get_model(model_config, Config)
|
||||||
|
|
||||||
self._agent = Agent(
|
self._agent = Agent(
|
||||||
model=model_obj,
|
model=model_obj,
|
||||||
|
|
@ -66,26 +64,6 @@ class QuestionAnswerAgent:
|
||||||
for chunk, score in expanded_results
|
for chunk, score in expanded_results
|
||||||
]
|
]
|
||||||
|
|
||||||
def _get_model(self, provider: str, model: str):
|
|
||||||
"""Get the appropriate model object for the provider."""
|
|
||||||
if provider == "ollama":
|
|
||||||
return OpenAIChatModel(
|
|
||||||
model_name=model,
|
|
||||||
provider=OllamaProvider(
|
|
||||||
base_url=f"{Config.providers.ollama.base_url}/v1"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
elif provider == "vllm":
|
|
||||||
return OpenAIChatModel(
|
|
||||||
model_name=model,
|
|
||||||
provider=OpenAIProvider(
|
|
||||||
base_url=f"{Config.providers.vllm.qa_base_url}/v1", api_key="none"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
# For all other providers, use the provider:model format
|
|
||||||
return f"{provider}:{model}"
|
|
||||||
|
|
||||||
async def answer(self, question: str) -> str:
|
async def answer(self, question: str) -> str:
|
||||||
"""Answer a question using the RAG system."""
|
"""Answer a question using the RAG system."""
|
||||||
deps = Dependencies(client=self._client)
|
deps = Dependencies(client=self._client)
|
||||||
|
|
|
||||||
|
|
@ -24,7 +24,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
|
|
||||||
reranker: RerankerBase | None = None
|
reranker: RerankerBase | None = None
|
||||||
|
|
||||||
if config.reranking.provider == "mxbai":
|
if config.reranking.model and config.reranking.model.provider == "mxbai":
|
||||||
try:
|
try:
|
||||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
from haiku.rag.reranking.mxbai import MxBAIReranker
|
||||||
|
|
||||||
|
|
@ -33,7 +33,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
except ImportError:
|
except ImportError:
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
elif config.reranking.provider == "cohere":
|
elif config.reranking.model and config.reranking.model.provider == "cohere":
|
||||||
try:
|
try:
|
||||||
from haiku.rag.reranking.cohere import CohereReranker
|
from haiku.rag.reranking.cohere import CohereReranker
|
||||||
|
|
||||||
|
|
@ -41,20 +41,20 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
except ImportError:
|
except ImportError:
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
elif config.reranking.provider == "vllm":
|
elif config.reranking.model and config.reranking.model.provider == "vllm":
|
||||||
try:
|
try:
|
||||||
from haiku.rag.reranking.vllm import VLLMReranker
|
from haiku.rag.reranking.vllm import VLLMReranker
|
||||||
|
|
||||||
reranker = VLLMReranker(config.reranking.model)
|
reranker = VLLMReranker(config.reranking.model.name)
|
||||||
except ImportError:
|
except ImportError:
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
||||||
elif config.reranking.provider == "zeroentropy":
|
elif config.reranking.model and config.reranking.model.provider == "zeroentropy":
|
||||||
try:
|
try:
|
||||||
from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
|
from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
|
||||||
|
|
||||||
# Use configured model or default to zerank-1
|
# Use configured model or default to zerank-1
|
||||||
model = config.reranking.model or "zerank-1"
|
model = config.reranking.model.name or "zerank-1"
|
||||||
reranker = ZeroEntropyReranker(model)
|
reranker = ZeroEntropyReranker(model)
|
||||||
except ImportError:
|
except ImportError:
|
||||||
reranker = None
|
reranker = None
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
class RerankerBase:
|
class RerankerBase:
|
||||||
_model: str = Config.reranking.model
|
_model: str | None = Config.reranking.model.name if Config.reranking.model else None
|
||||||
|
|
||||||
async def rerank(
|
async def rerank(
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
|
|
|
||||||
|
|
@ -22,8 +22,9 @@ class CohereReranker(RerankerBase):
|
||||||
|
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
|
model_name = self._model or "rerank-v3.5"
|
||||||
response = self._client.rerank(
|
response = self._client.rerank(
|
||||||
model=self._model, query=query, documents=documents, top_n=top_n
|
model=model_name, query=query, documents=documents, top_n=top_n
|
||||||
)
|
)
|
||||||
|
|
||||||
reranked_chunks = []
|
reranked_chunks = []
|
||||||
|
|
|
||||||
|
|
@ -7,9 +7,12 @@ from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
class MxBAIReranker(RerankerBase):
|
class MxBAIReranker(RerankerBase):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._client = MxbaiRerankV2(
|
model_name = (
|
||||||
Config.reranking.model, disable_transformers_warnings=True
|
Config.reranking.model.name
|
||||||
|
if Config.reranking.model
|
||||||
|
else "mxbai-rerank-base-v2"
|
||||||
)
|
)
|
||||||
|
self._client = MxbaiRerankV2(model_name, disable_transformers_warnings=True)
|
||||||
|
|
||||||
async def rerank(
|
async def rerank(
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
|
|
|
||||||
|
|
@ -37,8 +37,9 @@ class ZeroEntropyReranker(RerankerBase):
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
# Call Zero Entropy reranking API
|
# Call Zero Entropy reranking API
|
||||||
|
model_name = self._model or "zerank-1"
|
||||||
response = self._client.models.rerank(
|
response = self._client.models.rerank(
|
||||||
model=self._model,
|
model=model_name,
|
||||||
query=query,
|
query=query,
|
||||||
documents=documents,
|
documents=documents,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -4,10 +4,240 @@ import sys
|
||||||
from importlib import metadata
|
from importlib import metadata
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import ModuleType
|
from types import ModuleType
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
from packaging.version import Version, parse
|
from packaging.version import Version, parse
|
||||||
|
|
||||||
|
|
||||||
|
def apply_common_settings(
|
||||||
|
settings: Any | None,
|
||||||
|
settings_class: type[Any],
|
||||||
|
model_config: Any,
|
||||||
|
) -> Any | None:
|
||||||
|
"""Apply common settings (temperature, max_tokens) to model settings.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
settings: Existing settings instance or None
|
||||||
|
settings_class: Settings class to instantiate if needed
|
||||||
|
model_config: ModelConfig with temperature and max_tokens
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Updated settings instance or None if no settings to apply
|
||||||
|
"""
|
||||||
|
if model_config.temperature is None and model_config.max_tokens is None:
|
||||||
|
return settings
|
||||||
|
|
||||||
|
if settings is None:
|
||||||
|
settings_dict = settings_class()
|
||||||
|
else:
|
||||||
|
settings_dict = settings
|
||||||
|
|
||||||
|
if model_config.temperature is not None:
|
||||||
|
settings_dict["temperature"] = model_config.temperature
|
||||||
|
|
||||||
|
if model_config.max_tokens is not None:
|
||||||
|
settings_dict["max_tokens"] = model_config.max_tokens
|
||||||
|
|
||||||
|
return settings_dict
|
||||||
|
|
||||||
|
|
||||||
|
def get_model(
|
||||||
|
model_config: Any,
|
||||||
|
app_config: Any | None = None,
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
Get a model instance for the specified configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_config: ModelConfig with provider, model, and settings
|
||||||
|
app_config: AppConfig for provider base URLs (defaults to global Config)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A configured model instance
|
||||||
|
"""
|
||||||
|
from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings
|
||||||
|
from pydantic_ai.providers.ollama import OllamaProvider
|
||||||
|
from pydantic_ai.providers.openai import OpenAIProvider
|
||||||
|
|
||||||
|
if app_config is None:
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
|
||||||
|
app_config = Config
|
||||||
|
|
||||||
|
provider = model_config.provider
|
||||||
|
model = model_config.name
|
||||||
|
|
||||||
|
if provider == "ollama":
|
||||||
|
model_settings = None
|
||||||
|
|
||||||
|
# Apply thinking control for gpt-oss
|
||||||
|
if model == "gpt-oss" and model_config.enable_thinking is not None:
|
||||||
|
if model_config.enable_thinking is False:
|
||||||
|
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
|
||||||
|
else:
|
||||||
|
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="high")
|
||||||
|
|
||||||
|
model_settings = apply_common_settings(
|
||||||
|
model_settings, OpenAIChatModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return OpenAIChatModel(
|
||||||
|
model_name=model,
|
||||||
|
provider=OllamaProvider(
|
||||||
|
base_url=f"{app_config.providers.ollama.base_url}/v1"
|
||||||
|
),
|
||||||
|
settings=model_settings,
|
||||||
|
)
|
||||||
|
|
||||||
|
elif provider == "openai":
|
||||||
|
openai_settings: Any = None
|
||||||
|
|
||||||
|
# Apply thinking control
|
||||||
|
if model_config.enable_thinking is not None:
|
||||||
|
if model_config.enable_thinking is False:
|
||||||
|
openai_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
|
||||||
|
else:
|
||||||
|
openai_settings = OpenAIChatModelSettings(
|
||||||
|
openai_reasoning_effort="high"
|
||||||
|
)
|
||||||
|
|
||||||
|
openai_settings = apply_common_settings(
|
||||||
|
openai_settings, OpenAIChatModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return OpenAIChatModel(model_name=model, settings=openai_settings)
|
||||||
|
|
||||||
|
elif provider == "anthropic":
|
||||||
|
from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
|
||||||
|
|
||||||
|
anthropic_settings: Any = None
|
||||||
|
|
||||||
|
# Apply thinking control
|
||||||
|
if model_config.enable_thinking is not None:
|
||||||
|
if model_config.enable_thinking:
|
||||||
|
anthropic_settings = AnthropicModelSettings(
|
||||||
|
anthropic_thinking={"type": "enabled", "budget_tokens": 4096}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
anthropic_settings = AnthropicModelSettings(
|
||||||
|
anthropic_thinking={"type": "disabled"}
|
||||||
|
)
|
||||||
|
|
||||||
|
anthropic_settings = apply_common_settings(
|
||||||
|
anthropic_settings, AnthropicModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return AnthropicModel(model_name=model, settings=anthropic_settings)
|
||||||
|
|
||||||
|
elif provider == "gemini":
|
||||||
|
from pydantic_ai.models.google import GoogleModel, GoogleModelSettings
|
||||||
|
|
||||||
|
gemini_settings: Any = None
|
||||||
|
|
||||||
|
# Apply thinking control
|
||||||
|
if model_config.enable_thinking is not None:
|
||||||
|
gemini_settings = GoogleModelSettings(
|
||||||
|
google_thinking_config={
|
||||||
|
"include_thoughts": model_config.enable_thinking
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
gemini_settings = apply_common_settings(
|
||||||
|
gemini_settings, GoogleModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return GoogleModel(model_name=model, settings=gemini_settings)
|
||||||
|
|
||||||
|
elif provider == "groq":
|
||||||
|
from pydantic_ai.models.groq import GroqModel, GroqModelSettings
|
||||||
|
|
||||||
|
groq_settings: Any = None
|
||||||
|
|
||||||
|
# Apply thinking control
|
||||||
|
if model_config.enable_thinking is not None:
|
||||||
|
if model_config.enable_thinking:
|
||||||
|
groq_settings = GroqModelSettings(groq_reasoning_format="parsed")
|
||||||
|
else:
|
||||||
|
groq_settings = GroqModelSettings(groq_reasoning_format="hidden")
|
||||||
|
|
||||||
|
groq_settings = apply_common_settings(
|
||||||
|
groq_settings, GroqModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return GroqModel(model_name=model, settings=groq_settings)
|
||||||
|
|
||||||
|
elif provider == "bedrock":
|
||||||
|
from pydantic_ai.models.bedrock import (
|
||||||
|
BedrockConverseModel,
|
||||||
|
BedrockModelSettings,
|
||||||
|
)
|
||||||
|
|
||||||
|
bedrock_settings: Any = None
|
||||||
|
|
||||||
|
# Apply thinking control for Claude models
|
||||||
|
if model_config.enable_thinking is not None:
|
||||||
|
additional_fields: dict[str, Any] = {}
|
||||||
|
if model.startswith("anthropic.claude"):
|
||||||
|
if model_config.enable_thinking:
|
||||||
|
additional_fields = {
|
||||||
|
"thinking": {"type": "enabled", "budget_tokens": 4096}
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
additional_fields = {"thinking": {"type": "disabled"}}
|
||||||
|
elif "gpt" in model or "o1" in model or "o3" in model:
|
||||||
|
# OpenAI models on Bedrock
|
||||||
|
additional_fields = {
|
||||||
|
"reasoning_effort": "high"
|
||||||
|
if model_config.enable_thinking
|
||||||
|
else "low"
|
||||||
|
}
|
||||||
|
elif "qwen" in model:
|
||||||
|
# Qwen models on Bedrock
|
||||||
|
additional_fields = {
|
||||||
|
"reasoning_config": "high"
|
||||||
|
if model_config.enable_thinking
|
||||||
|
else "low"
|
||||||
|
}
|
||||||
|
|
||||||
|
if additional_fields:
|
||||||
|
bedrock_settings = BedrockModelSettings(
|
||||||
|
bedrock_additional_model_requests_fields=additional_fields
|
||||||
|
)
|
||||||
|
|
||||||
|
bedrock_settings = apply_common_settings(
|
||||||
|
bedrock_settings, BedrockModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return BedrockConverseModel(model_name=model, settings=bedrock_settings)
|
||||||
|
|
||||||
|
elif provider == "vllm":
|
||||||
|
vllm_settings = None
|
||||||
|
|
||||||
|
# Apply thinking control for gpt-oss
|
||||||
|
if model == "gpt-oss" and model_config.enable_thinking is not None:
|
||||||
|
if model_config.enable_thinking is False:
|
||||||
|
vllm_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
|
||||||
|
else:
|
||||||
|
vllm_settings = OpenAIChatModelSettings(openai_reasoning_effort="high")
|
||||||
|
|
||||||
|
vllm_settings = apply_common_settings(
|
||||||
|
vllm_settings, OpenAIChatModelSettings, model_config
|
||||||
|
)
|
||||||
|
|
||||||
|
return OpenAIChatModel(
|
||||||
|
model_name=model,
|
||||||
|
provider=OpenAIProvider(
|
||||||
|
base_url=f"{app_config.providers.vllm.research_base_url or app_config.providers.vllm.qa_base_url}/v1",
|
||||||
|
api_key="none",
|
||||||
|
),
|
||||||
|
settings=vllm_settings,
|
||||||
|
)
|
||||||
|
|
||||||
|
else:
|
||||||
|
# For any other provider, use string format and let Pydantic AI handle it
|
||||||
|
return f"{provider}:{model}"
|
||||||
|
|
||||||
|
|
||||||
def format_bytes(num_bytes: int) -> str:
|
def format_bytes(num_bytes: int) -> str:
|
||||||
"""Format bytes as human-readable string."""
|
"""Format bytes as human-readable string."""
|
||||||
size = float(num_bytes)
|
size = float(num_bytes)
|
||||||
|
|
@ -137,12 +367,12 @@ def prefetch_models():
|
||||||
required_models: set[str] = set()
|
required_models: set[str] = set()
|
||||||
if Config.embeddings.provider == "ollama":
|
if Config.embeddings.provider == "ollama":
|
||||||
required_models.add(Config.embeddings.model)
|
required_models.add(Config.embeddings.model)
|
||||||
if Config.qa.provider == "ollama":
|
if Config.qa.model.provider == "ollama":
|
||||||
required_models.add(Config.qa.model)
|
required_models.add(Config.qa.model.name)
|
||||||
if Config.research.provider == "ollama":
|
if Config.research.model.provider == "ollama":
|
||||||
required_models.add(Config.research.model)
|
required_models.add(Config.research.model.name)
|
||||||
if Config.reranking.provider == "ollama":
|
if Config.reranking.model and Config.reranking.model.provider == "ollama":
|
||||||
required_models.add(Config.reranking.model)
|
required_models.add(Config.reranking.model.name)
|
||||||
|
|
||||||
if not required_models:
|
if not required_models:
|
||||||
return
|
return
|
||||||
|
|
|
||||||
|
|
@ -59,7 +59,12 @@ nav:
|
||||||
- index.md
|
- index.md
|
||||||
- Getting started: tutorial.md
|
- Getting started: tutorial.md
|
||||||
- Installation: installation.md
|
- Installation: installation.md
|
||||||
- Configuration: configuration.md
|
- Configuration:
|
||||||
|
- config-index.md
|
||||||
|
- Providers: providers.md
|
||||||
|
- QA and Research: qa-research.md
|
||||||
|
- Document Processing: processing.md
|
||||||
|
- Storage: storage.md
|
||||||
- CLI: cli.md
|
- CLI: cli.md
|
||||||
- Python: python.md
|
- Python: python.md
|
||||||
- Agents: agents.md
|
- Agents: agents.md
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,10 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
|
||||||
def test_model_factory(provider, model, config=None):
|
def test_model_factory(provider, model, config=None):
|
||||||
return TestModel()
|
return TestModel()
|
||||||
|
|
||||||
monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)
|
# Patch all locations where get_model is imported
|
||||||
|
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
|
||||||
|
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
|
||||||
|
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
|
||||||
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
|
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
|
||||||
|
|
||||||
graph = build_deep_qa_graph()
|
graph = build_deep_qa_graph()
|
||||||
|
|
@ -50,7 +53,10 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
|
||||||
def test_model_factory(provider, model, config=None):
|
def test_model_factory(provider, model, config=None):
|
||||||
return TestModel()
|
return TestModel()
|
||||||
|
|
||||||
monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)
|
# Patch all locations where get_model is imported
|
||||||
|
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
|
||||||
|
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
|
||||||
|
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
|
||||||
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
|
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
|
||||||
|
|
||||||
graph = build_deep_qa_graph()
|
graph = build_deep_qa_graph()
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
import asyncio
|
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from pydantic_ai.models.test import TestModel
|
from pydantic_ai.models.test import TestModel
|
||||||
|
|
||||||
|
|
@ -25,12 +23,6 @@ def test_build_graph_and_state():
|
||||||
assert state.context.sub_questions == []
|
assert state.context.sub_questions == []
|
||||||
|
|
||||||
|
|
||||||
def test_async_loop_available():
|
|
||||||
# Ensure an event loop can be created in test env
|
|
||||||
loop = asyncio.new_event_loop()
|
|
||||||
loop.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
|
async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
|
||||||
"""Test research graph with mocked LLM using AG-UI events."""
|
"""Test research graph with mocked LLM using AG-UI events."""
|
||||||
|
|
@ -39,7 +31,10 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
|
||||||
def test_model_factory(_provider, _model, _config=None):
|
def test_model_factory(_provider, _model, _config=None):
|
||||||
return TestModel()
|
return TestModel()
|
||||||
|
|
||||||
monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)
|
# Patch all locations where get_model is imported
|
||||||
|
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
|
||||||
|
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
|
||||||
|
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
|
||||||
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
|
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
|
||||||
|
|
||||||
graph = build_research_graph()
|
graph = build_research_graph()
|
||||||
|
|
|
||||||
|
|
@ -115,6 +115,7 @@ async def test_chunk_repository_crud(temp_db_path):
|
||||||
)
|
)
|
||||||
|
|
||||||
created_chunk = await chunk_repo.create(chunk)
|
created_chunk = await chunk_repo.create(chunk)
|
||||||
|
assert isinstance(created_chunk, Chunk)
|
||||||
assert created_chunk.id is not None
|
assert created_chunk.id is not None
|
||||||
assert created_chunk.content == "Test chunk content"
|
assert created_chunk.content == "Test chunk content"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -736,9 +736,9 @@ async def test_client_ask_without_cite(monkeypatch, temp_db_path):
|
||||||
"""Test asking questions without citations."""
|
"""Test asking questions without citations."""
|
||||||
from pydantic_ai.models.test import TestModel
|
from pydantic_ai.models.test import TestModel
|
||||||
|
|
||||||
# Mock OpenAIChatModel to return TestModel
|
# Mock get_model to return TestModel
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
"haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel()
|
"haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel()
|
||||||
)
|
)
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path) as client:
|
async with HaikuRAG(temp_db_path) as client:
|
||||||
|
|
@ -760,9 +760,9 @@ async def test_client_ask_with_cite(monkeypatch, temp_db_path):
|
||||||
"""Test asking questions with citations."""
|
"""Test asking questions with citations."""
|
||||||
from pydantic_ai.models.test import TestModel
|
from pydantic_ai.models.test import TestModel
|
||||||
|
|
||||||
# Mock OpenAIChatModel to return TestModel
|
# Mock get_model to return TestModel
|
||||||
monkeypatch.setattr(
|
monkeypatch.setattr(
|
||||||
"haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel()
|
"haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel()
|
||||||
)
|
)
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path) as client:
|
async with HaikuRAG(temp_db_path) as client:
|
||||||
|
|
|
||||||
|
|
@ -55,6 +55,7 @@ def test_vllm_embedder_uses_config():
|
||||||
|
|
||||||
def test_openai_embedder_uses_config():
|
def test_openai_embedder_uses_config():
|
||||||
"""Test that openai embedder uses the config passed to get_embedder."""
|
"""Test that openai embedder uses the config passed to get_embedder."""
|
||||||
|
|
||||||
custom_config = AppConfig(
|
custom_config = AppConfig(
|
||||||
embeddings=EmbeddingsConfig(
|
embeddings=EmbeddingsConfig(
|
||||||
provider="openai",
|
provider="openai",
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ from evaluations.evaluators import LLMJudge
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
from haiku.rag.qa.agent import QuestionAnswerAgent
|
from haiku.rag.qa.agent import QuestionAnswerAgent
|
||||||
|
|
||||||
OPENAI_AVAILABLE = bool(os.getenv("OPENAI_API_KEY"))
|
OPENAI_AVAILABLE = bool(os.getenv("OPENAI_API_KEY"))
|
||||||
|
|
@ -17,7 +18,9 @@ VLLM_QA_AVAILABLE = bool(Config.providers.vllm.qa_base_url)
|
||||||
async def test_qa_ollama(qa_corpus: Dataset, temp_db_path):
|
async def test_qa_ollama(qa_corpus: Dataset, temp_db_path):
|
||||||
"""Test Ollama QA with LLM judge."""
|
"""Test Ollama QA with LLM judge."""
|
||||||
client = HaikuRAG(temp_db_path)
|
client = HaikuRAG(temp_db_path)
|
||||||
qa = QuestionAnswerAgent(client, "ollama", "qwen3")
|
qa = QuestionAnswerAgent(
|
||||||
|
client, ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=False)
|
||||||
|
)
|
||||||
llm_judge = LLMJudge()
|
llm_judge = LLMJudge()
|
||||||
|
|
||||||
doc = qa_corpus[1]
|
doc = qa_corpus[1]
|
||||||
|
|
@ -41,7 +44,7 @@ async def test_qa_ollama(qa_corpus: Dataset, temp_db_path):
|
||||||
async def test_qa_openai(qa_corpus: Dataset, temp_db_path):
|
async def test_qa_openai(qa_corpus: Dataset, temp_db_path):
|
||||||
"""Test OpenAI QA with LLM judge."""
|
"""Test OpenAI QA with LLM judge."""
|
||||||
client = HaikuRAG(temp_db_path)
|
client = HaikuRAG(temp_db_path)
|
||||||
qa = QuestionAnswerAgent(client, "openai", "gpt-4o-mini")
|
qa = QuestionAnswerAgent(client, ModelConfig(provider="openai", name="gpt-4o-mini"))
|
||||||
llm_judge = LLMJudge()
|
llm_judge = LLMJudge()
|
||||||
|
|
||||||
doc = qa_corpus[1]
|
doc = qa_corpus[1]
|
||||||
|
|
@ -65,7 +68,9 @@ async def test_qa_openai(qa_corpus: Dataset, temp_db_path):
|
||||||
async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path):
|
async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path):
|
||||||
"""Test Anthropic QA with LLM judge."""
|
"""Test Anthropic QA with LLM judge."""
|
||||||
client = HaikuRAG(temp_db_path)
|
client = HaikuRAG(temp_db_path)
|
||||||
qa = QuestionAnswerAgent(client, "anthropic", "claude-3-5-haiku-20241022")
|
qa = QuestionAnswerAgent(
|
||||||
|
client, ModelConfig(provider="anthropic", name="claude-3-5-haiku-20241022")
|
||||||
|
)
|
||||||
llm_judge = LLMJudge()
|
llm_judge = LLMJudge()
|
||||||
|
|
||||||
doc = qa_corpus[1]
|
doc = qa_corpus[1]
|
||||||
|
|
@ -89,7 +94,7 @@ async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path):
|
||||||
async def test_qa_vllm(qa_corpus: Dataset, temp_db_path):
|
async def test_qa_vllm(qa_corpus: Dataset, temp_db_path):
|
||||||
"""Test vLLM QA with LLM judge."""
|
"""Test vLLM QA with LLM judge."""
|
||||||
client = HaikuRAG(temp_db_path)
|
client = HaikuRAG(temp_db_path)
|
||||||
qa = QuestionAnswerAgent(client, "vllm", "Qwen/Qwen3-4B")
|
qa = QuestionAnswerAgent(client, ModelConfig(provider="vllm", name="Qwen/Qwen3-4B"))
|
||||||
llm_judge = LLMJudge()
|
llm_judge = LLMJudge()
|
||||||
|
|
||||||
doc = qa_corpus[1]
|
doc = qa_corpus[1]
|
||||||
|
|
|
||||||
|
|
@ -40,17 +40,19 @@ async def test_reranker_base():
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_mxbai_reranker():
|
async def test_mxbai_reranker():
|
||||||
try:
|
try:
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
from haiku.rag.reranking.mxbai import MxBAIReranker
|
||||||
|
|
||||||
Config.reranking.model = "mixedbread-ai/mxbai-rerank-base-v2"
|
Config.reranking.model = ModelConfig(
|
||||||
|
provider="mxbai", name="mixedbread-ai/mxbai-rerank-base-v2"
|
||||||
|
)
|
||||||
reranker = MxBAIReranker()
|
reranker = MxBAIReranker()
|
||||||
# reranker._model = "mixedbread-ai/mxbai-rerank-base-v2"
|
|
||||||
reranked = await reranker.rerank(
|
reranked = await reranker.rerank(
|
||||||
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
||||||
)
|
)
|
||||||
assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
|
assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
|
||||||
assert all(isinstance(score, float) for chunk, score in reranked)
|
assert all(isinstance(score, float) for chunk, score in reranked)
|
||||||
Config.reranking.model = ""
|
Config.reranking.model = None
|
||||||
|
|
||||||
except ImportError:
|
except ImportError:
|
||||||
pytest.skip("MxBAI package not installed")
|
pytest.skip("MxBAI package not installed")
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,18 @@
|
||||||
|
import importlib.util
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pydantic_ai.models.openai import OpenAIChatModel
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
from haiku.rag.converters import get_converter
|
from haiku.rag.converters import get_converter
|
||||||
|
from haiku.rag.utils import get_model
|
||||||
|
|
||||||
|
# Check for optional dependencies
|
||||||
|
HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None
|
||||||
|
HAS_GOOGLE = importlib.util.find_spec("google.generativeai") is not None
|
||||||
|
HAS_GROQ = importlib.util.find_spec("groq") is not None
|
||||||
|
HAS_BEDROCK = importlib.util.find_spec("botocore") is not None
|
||||||
|
|
||||||
|
|
||||||
def test_text_to_docling_document():
|
def test_text_to_docling_document():
|
||||||
|
|
@ -119,3 +132,169 @@ Emoji test: 🚀 ✅ 📝"""
|
||||||
assert "测试文档" in result_markdown
|
assert "测试文档" in result_markdown
|
||||||
assert "¡Hola mundo!" in result_markdown
|
assert "¡Hola mundo!" in result_markdown
|
||||||
assert "🚀" in result_markdown
|
assert "🚀" in result_markdown
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_ollama():
|
||||||
|
"""Test get_model returns OpenAIChatModel for Ollama."""
|
||||||
|
model_config = ModelConfig(provider="ollama", name="llama3")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_ollama_without_thinking():
|
||||||
|
"""Test get_model configures thinking for gpt-oss on Ollama."""
|
||||||
|
model_config = ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=False)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_ollama_with_settings():
|
||||||
|
"""Test get_model applies temperature and max_tokens for Ollama."""
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="ollama", name="llama3", temperature=0.5, max_tokens=100
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_openai():
|
||||||
|
"""Test get_model returns OpenAIChatModel for OpenAI."""
|
||||||
|
model_config = ModelConfig(provider="openai", name="gpt-4o")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_openai_with_thinking():
|
||||||
|
"""Test get_model configures thinking for OpenAI reasoning models."""
|
||||||
|
model_config = ModelConfig(provider="openai", name="o1", enable_thinking=True)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
|
||||||
|
def test_get_model_anthropic():
|
||||||
|
"""Test get_model returns AnthropicModel for Anthropic."""
|
||||||
|
from pydantic_ai.models.anthropic import AnthropicModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(provider="anthropic", name="claude-3-5-sonnet-20241022")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, AnthropicModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
|
||||||
|
def test_get_model_anthropic_with_thinking():
|
||||||
|
"""Test get_model configures thinking for Anthropic."""
|
||||||
|
from pydantic_ai.models.anthropic import AnthropicModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="anthropic",
|
||||||
|
name="claude-3-5-sonnet-20241022",
|
||||||
|
enable_thinking=True,
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, AnthropicModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
|
||||||
|
def test_get_model_gemini():
|
||||||
|
"""Test get_model returns GoogleModel for Gemini."""
|
||||||
|
from pydantic_ai.models.google import GoogleModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(provider="gemini", name="gemini-2.0-flash-exp")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, GoogleModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
|
||||||
|
def test_get_model_gemini_with_thinking():
|
||||||
|
"""Test get_model configures thinking for Gemini."""
|
||||||
|
from pydantic_ai.models.google import GoogleModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="gemini", name="gemini-2.0-flash-thinking-exp", enable_thinking=True
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, GoogleModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
|
||||||
|
def test_get_model_groq():
|
||||||
|
"""Test get_model returns GroqModel for Groq."""
|
||||||
|
from pydantic_ai.models.groq import GroqModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(provider="groq", name="llama-3.3-70b-versatile")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, GroqModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
|
||||||
|
def test_get_model_groq_with_thinking():
|
||||||
|
"""Test get_model configures thinking format for Groq."""
|
||||||
|
from pydantic_ai.models.groq import GroqModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="groq", name="llama-3.3-70b-versatile", enable_thinking=False
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, GroqModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
|
||||||
|
def test_get_model_bedrock():
|
||||||
|
"""Test get_model returns BedrockConverseModel for Bedrock."""
|
||||||
|
from pydantic_ai.models.bedrock import BedrockConverseModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="bedrock", name="anthropic.claude-3-5-sonnet-20241022-v2:0"
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, BedrockConverseModel)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
|
||||||
|
def test_get_model_bedrock_with_thinking():
|
||||||
|
"""Test get_model configures thinking for Bedrock Claude models."""
|
||||||
|
from pydantic_ai.models.bedrock import BedrockConverseModel
|
||||||
|
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="bedrock",
|
||||||
|
name="anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||||
|
enable_thinking=True,
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, BedrockConverseModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_vllm():
|
||||||
|
"""Test get_model returns OpenAIChatModel for vLLM."""
|
||||||
|
model_config = ModelConfig(provider="vllm", name="Qwen/Qwen3-4B")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_vllm_with_thinking():
|
||||||
|
"""Test get_model configures thinking for gpt-oss on vLLM."""
|
||||||
|
model_config = ModelConfig(provider="vllm", name="gpt-oss", enable_thinking=False)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_unknown_provider():
|
||||||
|
"""Test get_model returns string format for unknown providers."""
|
||||||
|
model_config = ModelConfig(provider="mistral", name="mistral-large-latest")
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, str)
|
||||||
|
assert result == "mistral:mistral-large-latest"
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_model_with_all_settings():
|
||||||
|
"""Test get_model applies all settings together."""
|
||||||
|
model_config = ModelConfig(
|
||||||
|
provider="openai",
|
||||||
|
name="gpt-4o",
|
||||||
|
enable_thinking=False,
|
||||||
|
temperature=0.7,
|
||||||
|
max_tokens=500,
|
||||||
|
)
|
||||||
|
result = get_model(model_config)
|
||||||
|
assert isinstance(result, OpenAIChatModel)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue