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
|
||||
|
||||
- **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
|
||||
- **Evaluations**: Added `evaluations` CLI script for running benchmarks (replaces `python -m evaluations.benchmark`)
|
||||
- **Evaluations**: Added `--db` option to override evaluation database path
|
||||
|
|
|
|||
|
|
@ -30,8 +30,9 @@ embeddings:
|
|||
vector_dim: 768
|
||||
|
||||
qa:
|
||||
provider: ollama
|
||||
model: qwen3
|
||||
model:
|
||||
provider: ollama
|
||||
name: qwen3
|
||||
```
|
||||
|
||||
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
|
||||
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
|
||||
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:
|
||||
- `--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.
|
||||
|
||||
|
|
|
|||
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
|
||||
- [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
|
||||
- [Server](server.md) - File monitoring and server mode
|
||||
- [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)
|
||||
- **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
|
||||
|
||||
|
|
@ -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
|
||||
- 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
|
||||
|
||||
|
|
|
|||
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 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.
|
||||
|
|
|
|||
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
|
||||
|
||||
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
|
||||
# haiku.rag.yaml
|
||||
|
|
|
|||
|
|
@ -122,7 +122,7 @@ agui:
|
|||
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
|
||||
|
||||
|
|
|
|||
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
|
||||
|
||||
qa:
|
||||
provider: openai
|
||||
model: gpt-4o-mini # or gpt-4o, gpt-4, etc.
|
||||
model:
|
||||
provider: openai
|
||||
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):
|
||||
|
|
@ -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).
|
||||
|
||||
See [Configuration](configuration.md) for all available options.
|
||||
See [Configuration](config-index.md) for all available options.
|
||||
|
||||
## Adding the first documents
|
||||
|
||||
|
|
@ -215,4 +216,4 @@ The following people are presenting talks at PyCon Finland 2025:
|
|||
|
||||
## 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 typer
|
||||
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.evaluators import LLMJudge
|
||||
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 haiku.rag.client import HaikuRAG
|
||||
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.qa import get_qa_agent
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
load_dotenv()
|
||||
|
||||
QA_JUDGE_MODEL = "qwen3"
|
||||
|
||||
logfire.configure(send_to_logfire="if-token-present", service_name="evals")
|
||||
logfire.instrument_pydantic_ai()
|
||||
configure_cli_logging()
|
||||
|
|
@ -37,7 +35,7 @@ def build_experiment_metadata(
|
|||
dataset_key: str,
|
||||
test_cases: int,
|
||||
config: AppConfig,
|
||||
judge_model: str,
|
||||
judge_config: ModelConfig,
|
||||
) -> dict[str, Any]:
|
||||
"""Build experiment metadata for Logfire tracking."""
|
||||
return {
|
||||
|
|
@ -48,12 +46,20 @@ def build_experiment_metadata(
|
|||
"embedder_dim": config.embeddings.vector_dim,
|
||||
"chunk_size": config.processing.chunk_size,
|
||||
"context_chunk_radius": config.processing.context_chunk_radius,
|
||||
"rerank_provider": config.reranking.provider,
|
||||
"rerank_model": config.reranking.model,
|
||||
"qa_provider": config.qa.provider,
|
||||
"qa_model": config.qa.model,
|
||||
"judge_provider": "ollama",
|
||||
"judge_model": judge_model,
|
||||
"rerank_provider": config.reranking.model.provider
|
||||
if config.reranking.model
|
||||
else None,
|
||||
"rerank_model": config.reranking.model.name if config.reranking.model else None,
|
||||
"qa_provider": config.qa.model.provider,
|
||||
"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"
|
||||
|
||||
judge_config = ModelConfig(
|
||||
provider="ollama", name="gpt-oss", enable_thinking=False
|
||||
)
|
||||
experiment_metadata = build_experiment_metadata(
|
||||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
judge_model=QA_JUDGE_MODEL,
|
||||
judge_config=judge_config,
|
||||
)
|
||||
|
||||
report = await dataset.evaluate(
|
||||
|
|
@ -213,10 +222,8 @@ async def run_qa_benchmark(
|
|||
for index, doc in enumerate(corpus, start=1)
|
||||
]
|
||||
|
||||
judge_model = OpenAIChatModel(
|
||||
model_name=QA_JUDGE_MODEL,
|
||||
provider=OllamaProvider(base_url=f"{config.providers.ollama.base_url}/v1"),
|
||||
)
|
||||
judge_config = ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=False)
|
||||
judge_model = get_model(judge_config, config)
|
||||
|
||||
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
|
||||
name=spec.key,
|
||||
|
|
@ -249,7 +256,7 @@ async def run_qa_benchmark(
|
|||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
judge_model=QA_JUDGE_MODEL,
|
||||
judge_config=judge_config,
|
||||
)
|
||||
|
||||
report = await evaluation_dataset.evaluate(
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
from pydantic import BaseModel
|
||||
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.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.
|
||||
|
||||
|
|
@ -37,15 +37,15 @@ class LLMJudge:
|
|||
"""LLM-as-judge for evaluating answer equivalence using Pydantic AI."""
|
||||
|
||||
def __init__(self, model: str = "gpt-oss"):
|
||||
# Create Ollama model
|
||||
ollama_model = OpenAIChatModel(
|
||||
model_name=model,
|
||||
provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"),
|
||||
# Create model using get_model with thinking disabled
|
||||
model_config = ModelConfig(
|
||||
provider="ollama", model=model, enable_thinking=False
|
||||
)
|
||||
model_obj = get_model(model_config, Config)
|
||||
|
||||
# Create Pydantic AI agent
|
||||
self._agent = Agent(
|
||||
model=ollama_model,
|
||||
model=model_obj,
|
||||
output_type=LLMJudgeResponseSchema,
|
||||
system_prompt=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
retries=3,
|
||||
|
|
|
|||
|
|
@ -65,7 +65,7 @@ def create_a2a_app(
|
|||
broker = InMemoryBroker()
|
||||
|
||||
# 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(
|
||||
model=model,
|
||||
deps_type=AgentDependencies,
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ class AgentDeps:
|
|||
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(
|
||||
model,
|
||||
|
|
|
|||
|
|
@ -47,7 +47,7 @@ if not db_path.exists():
|
|||
|
||||
logger.info(f"Initializing research assistant with database: {db_path}")
|
||||
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
|
||||
|
|
@ -153,8 +153,8 @@ async def health_check(_: Request) -> JSONResponse:
|
|||
{
|
||||
"status": "healthy",
|
||||
"agent_model": str(agent.model),
|
||||
"research_provider": Config.research.provider,
|
||||
"research_model": Config.research.model,
|
||||
"research_provider": Config.research.model.provider,
|
||||
"research_model": Config.research.model.name,
|
||||
"db_path": str(db_path),
|
||||
"db_exists": db_path.exists(),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,18 +17,21 @@ providers:
|
|||
base_url: http://host.docker.internal:11434
|
||||
|
||||
research:
|
||||
provider: ollama
|
||||
model: gpt-oss:latest
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss:latest
|
||||
max_iterations: 3
|
||||
confidence_threshold: 0.8
|
||||
max_concurrency: 1
|
||||
|
||||
# For OpenAI:
|
||||
# research:
|
||||
# provider: openai
|
||||
# model: gpt-4o-mini
|
||||
# model:
|
||||
# provider: openai
|
||||
# name: gpt-4o-mini
|
||||
|
||||
# For Anthropic:
|
||||
# research:
|
||||
# provider: anthropic
|
||||
# model: claude-3-5-haiku-20241022
|
||||
# model:
|
||||
# provider: anthropic
|
||||
# name: claude-3-5-haiku-20241022
|
||||
|
|
|
|||
|
|
@ -6,6 +6,25 @@ from pydantic import BaseModel, Field
|
|||
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):
|
||||
data_dir: Path = Field(default_factory=get_default_data_dir)
|
||||
vacuum_retention_seconds: int = 86400
|
||||
|
|
@ -31,21 +50,30 @@ class EmbeddingsConfig(BaseModel):
|
|||
|
||||
|
||||
class RerankingConfig(BaseModel):
|
||||
provider: str = ""
|
||||
model: str = ""
|
||||
model: ModelConfig | None = None
|
||||
|
||||
|
||||
class QAConfig(BaseModel):
|
||||
provider: str = "ollama"
|
||||
model: str = "gpt-oss"
|
||||
model: ModelConfig = Field(
|
||||
default_factory=lambda: ModelConfig(
|
||||
provider="ollama",
|
||||
name="gpt-oss",
|
||||
enable_thinking=False,
|
||||
)
|
||||
)
|
||||
max_sub_questions: int = 3
|
||||
max_iterations: int = 2
|
||||
max_concurrency: int = 1
|
||||
|
||||
|
||||
class ResearchConfig(BaseModel):
|
||||
provider: str = "ollama"
|
||||
model: str = "gpt-oss"
|
||||
model: ModelConfig = Field(
|
||||
default_factory=lambda: ModelConfig(
|
||||
provider="ollama",
|
||||
name="gpt-oss",
|
||||
enable_thinking=False,
|
||||
)
|
||||
)
|
||||
max_iterations: int = 3
|
||||
confidence_threshold: float = 0.8
|
||||
max_concurrency: int = 1
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
"""Common utilities for graph implementations."""
|
||||
|
||||
from haiku.rag.graph.common.utils import get_model
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
__all__ = ["get_model"]
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from pydantic_graph.beta import StepContext
|
|||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
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.common import get_model
|
||||
from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
|
||||
|
|
@ -52,8 +52,7 @@ class GraphAgentDeps(Protocol):
|
|||
|
||||
|
||||
def create_plan_node[AgentDepsT: GraphAgentDeps](
|
||||
provider: str,
|
||||
model: str,
|
||||
model_config: ModelConfig,
|
||||
deps_type: type[AgentDepsT],
|
||||
activity_message: str = "Creating plan",
|
||||
output_retries: int | None = None,
|
||||
|
|
@ -62,8 +61,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
|
|||
"""Create a plan node for any graph.
|
||||
|
||||
Args:
|
||||
provider: Model provider (e.g., 'openai', 'anthropic')
|
||||
model: Model name
|
||||
model_config: ModelConfig with provider, model, and settings
|
||||
deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
|
||||
activity_message: Message to show during planning activity
|
||||
output_retries: Number of output retries for the agent (optional)
|
||||
|
|
@ -86,7 +84,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
|
|||
try:
|
||||
# Build agent configuration
|
||||
agent_config = {
|
||||
"model": get_model(provider, model, config),
|
||||
"model": get_model(model_config, config),
|
||||
"output_type": ResearchPlan,
|
||||
"instructions": (
|
||||
PLAN_PROMPT
|
||||
|
|
@ -141,8 +139,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
|
|||
|
||||
|
||||
def create_search_node[AgentDepsT: GraphAgentDeps](
|
||||
provider: str,
|
||||
model: str,
|
||||
model_config: ModelConfig,
|
||||
deps_type: type[AgentDepsT],
|
||||
with_step_wrapper: bool = True,
|
||||
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.
|
||||
|
||||
Args:
|
||||
provider: Model provider
|
||||
model: Model name
|
||||
model_config: ModelConfig with provider, model, and settings
|
||||
deps_type: Type of dependencies for the agent
|
||||
with_step_wrapper: Whether to wrap with agui_emitter start/finish step
|
||||
success_message_format: Format string for success activity message
|
||||
|
|
@ -186,8 +182,7 @@ def create_search_node[AgentDepsT: GraphAgentDeps](
|
|||
state,
|
||||
deps,
|
||||
sub_q,
|
||||
provider,
|
||||
model,
|
||||
model_config,
|
||||
deps_type,
|
||||
success_message_format,
|
||||
handle_exceptions,
|
||||
|
|
@ -204,8 +199,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
|
|||
state: GraphState,
|
||||
deps: GraphDeps,
|
||||
sub_q: str,
|
||||
provider: str,
|
||||
model: str,
|
||||
model_config: ModelConfig,
|
||||
deps_type: type[AgentDepsT],
|
||||
success_message_format: str,
|
||||
handle_exceptions: bool,
|
||||
|
|
@ -223,7 +217,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
|
|||
)
|
||||
|
||||
agent = Agent(
|
||||
model=get_model(provider, model, config),
|
||||
model=get_model(model_config, config),
|
||||
output_type=ToolOutput(SearchAnswer, max_retries=3),
|
||||
instructions=SEARCH_AGENT_PROMPT,
|
||||
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:
|
||||
Configured Deep QA graph
|
||||
"""
|
||||
provider = config.qa.provider
|
||||
model = config.qa.model
|
||||
model_config = config.qa.model
|
||||
g = GraphBuilder(
|
||||
state_type=DeepQAState,
|
||||
deps_type=DeepQADeps,
|
||||
|
|
@ -40,8 +39,7 @@ def build_deep_qa_graph(
|
|||
# Create and register the plan node using the factory
|
||||
plan = g.step(
|
||||
create_plan_node(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_config=model_config,
|
||||
deps_type=DeepQADependencies, # type: ignore[arg-type]
|
||||
activity_message="Planning approach",
|
||||
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
|
||||
search_one = g.step(
|
||||
create_search_node(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_config=model_config,
|
||||
deps_type=DeepQADependencies, # type: ignore[arg-type]
|
||||
with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step
|
||||
success_message_format="Answered: {sub_q}",
|
||||
|
|
@ -92,7 +89,7 @@ def build_deep_qa_graph(
|
|||
|
||||
try:
|
||||
agent = Agent(
|
||||
model=get_model(provider, model, config),
|
||||
model=get_model(model_config, config),
|
||||
output_type=DeepQAEvaluation,
|
||||
instructions=DECISION_PROMPT,
|
||||
retries=3,
|
||||
|
|
@ -173,7 +170,7 @@ def build_deep_qa_graph(
|
|||
)
|
||||
|
||||
agent = Agent(
|
||||
model=get_model(provider, model, config),
|
||||
model=get_model(model_config, config),
|
||||
output_type=DeepQAAnswer,
|
||||
instructions=prompt_template,
|
||||
retries=3,
|
||||
|
|
|
|||
|
|
@ -36,8 +36,7 @@ def build_research_graph(
|
|||
Returns:
|
||||
Configured Research graph
|
||||
"""
|
||||
provider = config.research.provider
|
||||
model = config.research.model
|
||||
model_config = config.research.model
|
||||
g = GraphBuilder(
|
||||
state_type=ResearchState,
|
||||
deps_type=ResearchDeps,
|
||||
|
|
@ -47,8 +46,7 @@ def build_research_graph(
|
|||
# Create and register the plan node using the factory
|
||||
plan = g.step(
|
||||
create_plan_node(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_config=model_config,
|
||||
deps_type=ResearchDependencies, # type: ignore[arg-type]
|
||||
activity_message="Creating research plan",
|
||||
output_retries=3,
|
||||
|
|
@ -59,8 +57,7 @@ def build_research_graph(
|
|||
# Create and register the search_one node using the factory
|
||||
search_one = g.step(
|
||||
create_search_node(
|
||||
provider=provider,
|
||||
model=model,
|
||||
model_config=model_config,
|
||||
deps_type=ResearchDependencies, # type: ignore[arg-type]
|
||||
with_step_wrapper=True,
|
||||
success_message_format="Found answer with {confidence:.0%} confidence",
|
||||
|
|
@ -99,7 +96,7 @@ def build_research_graph(
|
|||
|
||||
try:
|
||||
agent = Agent(
|
||||
model=get_model(provider, model, config),
|
||||
model=get_model(model_config, config),
|
||||
output_type=InsightAnalysis,
|
||||
instructions=INSIGHT_AGENT_PROMPT,
|
||||
retries=3,
|
||||
|
|
@ -168,7 +165,7 @@ def build_research_graph(
|
|||
|
||||
try:
|
||||
agent = Agent(
|
||||
model=get_model(provider, model, config),
|
||||
model=get_model(model_config, config),
|
||||
output_type=EvaluationResult,
|
||||
instructions=DECISION_AGENT_PROMPT,
|
||||
retries=3,
|
||||
|
|
@ -247,7 +244,7 @@ def build_research_graph(
|
|||
|
||||
try:
|
||||
agent = Agent(
|
||||
model=get_model(provider, model, config),
|
||||
model=get_model(model_config, config),
|
||||
output_type=ResearchReport,
|
||||
instructions=SYNTHESIS_AGENT_PROMPT,
|
||||
retries=3,
|
||||
|
|
|
|||
|
|
@ -21,13 +21,9 @@ def get_qa_agent(
|
|||
Returns:
|
||||
A configured QuestionAnswerAgent instance.
|
||||
"""
|
||||
provider = config.qa.provider
|
||||
model_name = config.qa.model
|
||||
|
||||
return QuestionAnswerAgent(
|
||||
client=client,
|
||||
provider=provider,
|
||||
model=model_name,
|
||||
model_config=config.qa.model,
|
||||
use_citations=use_citations,
|
||||
system_prompt=system_prompt,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,10 @@
|
|||
from pydantic import BaseModel, Field
|
||||
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.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
|
||||
|
||||
|
||||
|
|
@ -26,8 +25,7 @@ class QuestionAnswerAgent:
|
|||
def __init__(
|
||||
self,
|
||||
client: HaikuRAG,
|
||||
provider: str,
|
||||
model: str,
|
||||
model_config: ModelConfig,
|
||||
use_citations: bool = False,
|
||||
q: float = 0.0,
|
||||
system_prompt: str | None = None,
|
||||
|
|
@ -38,7 +36,7 @@ class QuestionAnswerAgent:
|
|||
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(
|
||||
model=model_obj,
|
||||
|
|
@ -66,26 +64,6 @@ class QuestionAnswerAgent:
|
|||
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:
|
||||
"""Answer a question using the RAG system."""
|
||||
deps = Dependencies(client=self._client)
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
|||
|
||||
reranker: RerankerBase | None = None
|
||||
|
||||
if config.reranking.provider == "mxbai":
|
||||
if config.reranking.model and config.reranking.model.provider == "mxbai":
|
||||
try:
|
||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
||||
|
||||
|
|
@ -33,7 +33,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
|||
except ImportError:
|
||||
reranker = None
|
||||
|
||||
elif config.reranking.provider == "cohere":
|
||||
elif config.reranking.model and config.reranking.model.provider == "cohere":
|
||||
try:
|
||||
from haiku.rag.reranking.cohere import CohereReranker
|
||||
|
||||
|
|
@ -41,20 +41,20 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
|||
except ImportError:
|
||||
reranker = None
|
||||
|
||||
elif config.reranking.provider == "vllm":
|
||||
elif config.reranking.model and config.reranking.model.provider == "vllm":
|
||||
try:
|
||||
from haiku.rag.reranking.vllm import VLLMReranker
|
||||
|
||||
reranker = VLLMReranker(config.reranking.model)
|
||||
reranker = VLLMReranker(config.reranking.model.name)
|
||||
except ImportError:
|
||||
reranker = None
|
||||
|
||||
elif config.reranking.provider == "zeroentropy":
|
||||
elif config.reranking.model and config.reranking.model.provider == "zeroentropy":
|
||||
try:
|
||||
from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
|
||||
|
||||
# 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)
|
||||
except ImportError:
|
||||
reranker = None
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from haiku.rag.store.models.chunk import Chunk
|
|||
|
||||
|
||||
class RerankerBase:
|
||||
_model: str = Config.reranking.model
|
||||
_model: str | None = Config.reranking.model.name if Config.reranking.model else None
|
||||
|
||||
async def rerank(
|
||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||
|
|
|
|||
|
|
@ -22,8 +22,9 @@ class CohereReranker(RerankerBase):
|
|||
|
||||
documents = [chunk.content for chunk in chunks]
|
||||
|
||||
model_name = self._model or "rerank-v3.5"
|
||||
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 = []
|
||||
|
|
|
|||
|
|
@ -7,9 +7,12 @@ from haiku.rag.store.models.chunk import Chunk
|
|||
|
||||
class MxBAIReranker(RerankerBase):
|
||||
def __init__(self):
|
||||
self._client = MxbaiRerankV2(
|
||||
Config.reranking.model, disable_transformers_warnings=True
|
||||
model_name = (
|
||||
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(
|
||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||
|
|
|
|||
|
|
@ -37,8 +37,9 @@ class ZeroEntropyReranker(RerankerBase):
|
|||
documents = [chunk.content for chunk in chunks]
|
||||
|
||||
# Call Zero Entropy reranking API
|
||||
model_name = self._model or "zerank-1"
|
||||
response = self._client.models.rerank(
|
||||
model=self._model,
|
||||
model=model_name,
|
||||
query=query,
|
||||
documents=documents,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4,10 +4,240 @@ import sys
|
|||
from importlib import metadata
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
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:
|
||||
"""Format bytes as human-readable string."""
|
||||
size = float(num_bytes)
|
||||
|
|
@ -137,12 +367,12 @@ def prefetch_models():
|
|||
required_models: set[str] = set()
|
||||
if Config.embeddings.provider == "ollama":
|
||||
required_models.add(Config.embeddings.model)
|
||||
if Config.qa.provider == "ollama":
|
||||
required_models.add(Config.qa.model)
|
||||
if Config.research.provider == "ollama":
|
||||
required_models.add(Config.research.model)
|
||||
if Config.reranking.provider == "ollama":
|
||||
required_models.add(Config.reranking.model)
|
||||
if Config.qa.model.provider == "ollama":
|
||||
required_models.add(Config.qa.model.name)
|
||||
if Config.research.model.provider == "ollama":
|
||||
required_models.add(Config.research.model.name)
|
||||
if Config.reranking.model and Config.reranking.model.provider == "ollama":
|
||||
required_models.add(Config.reranking.model.name)
|
||||
|
||||
if not required_models:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -59,7 +59,12 @@ nav:
|
|||
- index.md
|
||||
- Getting started: tutorial.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
|
||||
- Python: python.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):
|
||||
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)
|
||||
|
||||
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):
|
||||
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)
|
||||
|
||||
graph = build_deep_qa_graph()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,3 @@
|
|||
import asyncio
|
||||
|
||||
import pytest
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
|
|
@ -25,12 +23,6 @@ def test_build_graph_and_state():
|
|||
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
|
||||
async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
|
||||
"""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):
|
||||
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)
|
||||
|
||||
graph = build_research_graph()
|
||||
|
|
|
|||
|
|
@ -115,6 +115,7 @@ async def test_chunk_repository_crud(temp_db_path):
|
|||
)
|
||||
|
||||
created_chunk = await chunk_repo.create(chunk)
|
||||
assert isinstance(created_chunk, Chunk)
|
||||
assert created_chunk.id is not None
|
||||
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."""
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
# Mock OpenAIChatModel to return TestModel
|
||||
# Mock get_model to return TestModel
|
||||
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:
|
||||
|
|
@ -760,9 +760,9 @@ async def test_client_ask_with_cite(monkeypatch, temp_db_path):
|
|||
"""Test asking questions with citations."""
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
# Mock OpenAIChatModel to return TestModel
|
||||
# Mock get_model to return TestModel
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ def test_vllm_embedder_uses_config():
|
|||
|
||||
def test_openai_embedder_uses_config():
|
||||
"""Test that openai embedder uses the config passed to get_embedder."""
|
||||
|
||||
custom_config = AppConfig(
|
||||
embeddings=EmbeddingsConfig(
|
||||
provider="openai",
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from evaluations.evaluators import LLMJudge
|
|||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
from haiku.rag.qa.agent import QuestionAnswerAgent
|
||||
|
||||
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):
|
||||
"""Test Ollama QA with LLM judge."""
|
||||
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()
|
||||
|
||||
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):
|
||||
"""Test OpenAI QA with LLM judge."""
|
||||
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()
|
||||
|
||||
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):
|
||||
"""Test Anthropic QA with LLM judge."""
|
||||
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()
|
||||
|
||||
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):
|
||||
"""Test vLLM QA with LLM judge."""
|
||||
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()
|
||||
|
||||
doc = qa_corpus[1]
|
||||
|
|
|
|||
|
|
@ -40,17 +40,19 @@ async def test_reranker_base():
|
|||
@pytest.mark.asyncio
|
||||
async def test_mxbai_reranker():
|
||||
try:
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
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._model = "mixedbread-ai/mxbai-rerank-base-v2"
|
||||
reranked = await reranker.rerank(
|
||||
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
||||
)
|
||||
assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
|
||||
assert all(isinstance(score, float) for chunk, score in reranked)
|
||||
Config.reranking.model = ""
|
||||
Config.reranking.model = None
|
||||
|
||||
except ImportError:
|
||||
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.models import ModelConfig
|
||||
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():
|
||||
|
|
@ -119,3 +132,169 @@ Emoji test: 🚀 ✅ 📝"""
|
|||
assert "测试文档" in result_markdown
|
||||
assert "¡Hola mundo!" 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