drop the old monitor, rename serve→mcp, add e2e tests
This commit is contained in:
Yiorgis Gozadinos
parent
1ca3c25a83
commit
7ea61a7b10
27 changed files with 539 additions and 1807 deletions
|
|
@ -1,8 +1,17 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- New `haiku-ingester` service for continuous document ingestion: persistent SQLite job queue, async worker pool with retries and a dead-letter queue, FS/HTTP/S3 source adapters with per-source circuit breakers, and a FastAPI control plane (`/health`, `/jobs`, `/sources`, `/dlq`). Configured under `ingester:` in `haiku.rag.yaml`. Shipped behind the `[ingester]` extra. See [docs/ingester.md](docs/ingester.md).
|
||||
|
||||
### Removed
|
||||
|
||||
- File monitor (`haiku.rag.monitor` module, `MonitorConfig`, `S3MonitorEntry`, `AppConfig.monitor`). The `--monitor` flag on `haiku-rag serve` is gone — continuous ingestion now lives in `haiku-ingester serve`. Migrate `monitor.directories` to `ingester.sources[type=fs]` and `monitor.s3` to `ingester.sources[type=s3]`; the `delete_orphans` / `ignore_patterns` / `include_patterns` keys keep their meaning on the per-source entry.
|
||||
|
||||
### Changed
|
||||
|
||||
- `haiku-rag serve` renamed to `haiku-rag mcp` (only MCP is left). `--mcp-port` renamed to `--port`. Update any `claude_desktop_config.json` from `["serve", "--mcp", "--stdio"]` to `["mcp", "--stdio"]`.
|
||||
- Drop `list_documents` and `get_document` from the default RAG skill's tool set; the skill now exposes only `search` and `cite`. Both tools dumped unbounded content into the agent's context (full document lists, full document bodies) and `get_document` returned no chunk_ids so its output was structurally uncitable. The analysis skill already covers these uses programmatically — `await list_documents()` and `Path('/documents/{id}/content.txt').read_text()` inside `execute_code`. The tool branches remain in `create_skill_tools` and the `skill_generator` `AVAILABLE_TOOLS` set so users can still opt in when building custom skills.
|
||||
|
||||
## [0.48.1] - 2026-05-21
|
||||
|
|
|
|||
29
docs/cli.md
29
docs/cli.md
|
|
@ -397,30 +397,25 @@ haiku-rag vacuum
|
|||
|
||||
**Automatic Cleanup:** Vacuum runs automatically in the background after document operations. By default, it removes versions older than 1 day (configurable via `storage.vacuum_retention_seconds`), preserving recent versions for concurrent connections. Manual vacuum can be useful for cleanup after bulk operations or to free disk space immediately.
|
||||
|
||||
## Server
|
||||
## MCP Server
|
||||
|
||||
Start services (requires at least one flag):
|
||||
```bash
|
||||
# MCP server only (HTTP transport)
|
||||
haiku-rag serve --mcp
|
||||
# HTTP transport on port 8001
|
||||
haiku-rag mcp
|
||||
|
||||
# MCP server (stdio transport)
|
||||
haiku-rag serve --mcp --stdio
|
||||
# stdio transport (for Claude Desktop)
|
||||
haiku-rag mcp --stdio
|
||||
|
||||
# File monitoring only
|
||||
haiku-rag serve --monitor
|
||||
# Custom port
|
||||
haiku-rag mcp --port 9000
|
||||
|
||||
# Both services
|
||||
haiku-rag serve --monitor --mcp
|
||||
|
||||
# Custom MCP port
|
||||
haiku-rag serve --mcp --mcp-port 9000
|
||||
|
||||
# Read-only mode (excludes write MCP tools, disables monitor)
|
||||
haiku-rag --read-only serve --mcp
|
||||
# Read-only mode (no write tools)
|
||||
haiku-rag --read-only mcp
|
||||
```
|
||||
|
||||
See [Server Mode](server.md) for details on available services.
|
||||
See [MCP Server](server.md) for details. For continuous document
|
||||
ingestion (filesystem watch, S3 polling, HTTP sources), use
|
||||
`haiku-ingester serve`.
|
||||
|
||||
## Settings
|
||||
|
||||
|
|
|
|||
|
|
@ -59,12 +59,14 @@ 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
|
||||
ingester:
|
||||
sources:
|
||||
- type: fs
|
||||
id: local-docs
|
||||
root: /path/to/documents
|
||||
ignore_patterns: [] # Gitignore-style patterns to exclude
|
||||
include_patterns: [] # Gitignore-style patterns to include
|
||||
delete_orphans: true
|
||||
|
||||
lancedb:
|
||||
uri: "" # Empty for local, or db://, s3://, az://, gs://
|
||||
|
|
|
|||
|
|
@ -291,96 +291,8 @@ Explicit titles passed via `title=` parameter always take precedence and are nev
|
|||
|
||||
To generate titles for existing untitled documents, use [`rebuild --title-only`](../cli.md#rebuild-database).
|
||||
|
||||
## File Monitoring
|
||||
## Continuous ingestion
|
||||
|
||||
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
|
||||
|
||||
### S3 / Object Storage Sources
|
||||
|
||||
In addition to local directories, the watcher can poll S3-compatible buckets (AWS S3, SeaweedFS, MinIO, Cloudflare R2, etc.). Install the `[s3]` extra and configure one or more entries under `monitor.s3`:
|
||||
|
||||
```yaml
|
||||
monitor:
|
||||
s3:
|
||||
- uri: s3://my-bucket/incoming/
|
||||
poll_interval: 300 # seconds between sweeps; default 300
|
||||
include_patterns: ["*.pdf", "*.md"]
|
||||
ignore_patterns: ["draft*"]
|
||||
delete_orphans: true
|
||||
storage_options:
|
||||
endpoint: http://seaweed:8333
|
||||
aws_access_key_id: ${AWS_KEY}
|
||||
aws_secret_access_key: ${AWS_SECRET}
|
||||
region: us-east-1
|
||||
allow_http: "true"
|
||||
```
|
||||
|
||||
Each entry is independent: own poll interval, own include/ignore patterns, own `delete_orphans` setting, own credentials. Omit `storage_options` to fall back to the AWS default credential chain (env vars, IAM role, AWS profile).
|
||||
|
||||
The dict shape matches `lancedb.storage_options`. The same Rust `object_store` library is used by both, so credentials configured for the LanceDB backend can be copy-pasted here.
|
||||
|
||||
See [Server Mode → S3 / Object Storage Monitoring](../server.md#s3-object-storage-monitoring) for behaviour details (ETag-based change detection, orphan-deletion scope, CLI `add-src s3://…`).
|
||||
For automatic ingestion of local directories, S3 buckets, or HTTP
|
||||
sources (with filtering, retries, and a dead-letter queue), see the
|
||||
[Ingester](../ingester.md) page.
|
||||
|
|
|
|||
|
|
@ -108,8 +108,8 @@ LanceDB on S3 supports **exactly one writer + N readers** per database URI. Mult
|
|||
|
||||
The recommended layout for production is "different buckets, same account, separate IAM roles per process":
|
||||
|
||||
- **Ingestion process** — IAM role with `s3:Get/List` on the documents bucket and `s3:Get/Put/Delete` on the LanceDB bucket. Runs `haiku-rag serve --monitor` (with `monitor.s3` entries pointing at the documents bucket). Exactly one such process per LanceDB URI.
|
||||
- **Consumer processes** (1..N) — IAM role with `s3:Get/List` on the LanceDB bucket only. Run `haiku-rag serve --read-only --mcp`, the chat TUI, etc. They never see the documents bucket.
|
||||
- **Ingestion process** — IAM role with `s3:Get/List` on the documents bucket and `s3:Get/Put/Delete` on the LanceDB bucket. Runs `haiku-ingester serve` (with `ingester.sources[type=s3]` pointing at the documents bucket). Exactly one such process per LanceDB URI.
|
||||
- **Consumer processes** (1..N) — IAM role with `s3:Get/List` on the LanceDB bucket only. Run `haiku-rag --read-only mcp`, the chat TUI, etc. They never see the documents bucket.
|
||||
|
||||
Each process picks up its own credentials from the AWS default chain (env vars, IAM instance role, AWS profile), so no credentials are hard-coded in the configuration files.
|
||||
|
||||
|
|
|
|||
275
docs/ingester.md
Normal file
275
docs/ingester.md
Normal file
|
|
@ -0,0 +1,275 @@
|
|||
# Ingester
|
||||
|
||||
The ingester is a long-running service that watches sources for
|
||||
changes and feeds documents into haiku.rag's LanceDB. It runs as a
|
||||
separate process (`haiku-ingester serve`), owns its own SQLite job
|
||||
queue, and exposes a small HTTP control plane for operations.
|
||||
|
||||
Use the ingester when:
|
||||
|
||||
- you have a corpus you want to keep in sync continuously
|
||||
- documents arrive over time from filesystem, S3, or HTTP sources
|
||||
- you want retry + dead-letter behavior, not "fire and forget"
|
||||
|
||||
For one-off ingestion, the `haiku-rag add-src` CLI is enough — see
|
||||
[CLI → Add Documents](cli.md).
|
||||
|
||||
## Install
|
||||
|
||||
The ingester ships behind an optional extra:
|
||||
|
||||
```bash
|
||||
pip install 'haiku.rag-slim[ingester]'
|
||||
# or, for the full package:
|
||||
pip install 'haiku.rag[ingester]'
|
||||
```
|
||||
|
||||
That pulls `fastapi`, `uvicorn`, `aiosqlite`, and the `[s3]` extra.
|
||||
The production binary is `haiku-ingester`.
|
||||
|
||||
## Configure sources
|
||||
|
||||
Add an `ingester:` block to your `haiku.rag.yaml`. The minimum is a
|
||||
single source:
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
sources:
|
||||
- type: fs
|
||||
id: local-docs
|
||||
root: /Users/you/docs
|
||||
delete_orphans: true
|
||||
```
|
||||
|
||||
### Filesystem
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
sources:
|
||||
- type: fs
|
||||
id: local-docs # optional; auto-derives from root
|
||||
root: /Users/you/docs
|
||||
poll_interval_s: 300
|
||||
delete_orphans: true
|
||||
ignore_patterns: ["**/.git/**", "**/node_modules/**"]
|
||||
include_patterns: ["*.md", "*.pdf"] # optional whitelist
|
||||
```
|
||||
|
||||
Uses `watchfiles` for push events plus a periodic sweep that catches
|
||||
anything the OS dropped between starts. Patterns follow
|
||||
[gitignore syntax](https://git-scm.com/docs/gitignore#_pattern_format).
|
||||
|
||||
### S3 / object storage
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
sources:
|
||||
- type: s3
|
||||
id: corp-docs
|
||||
uri: s3://my-bucket/incoming/
|
||||
poll_interval_s: 300
|
||||
delete_orphans: true
|
||||
ignore_patterns: ["draft*"]
|
||||
include_patterns: ["*.pdf", "*.md"]
|
||||
storage_options:
|
||||
endpoint: http://seaweed:8333 # omit for AWS default chain
|
||||
aws_access_key_id: ${AWS_KEY}
|
||||
aws_secret_access_key: ${AWS_SECRET}
|
||||
region: us-east-1
|
||||
allow_http: "true"
|
||||
```
|
||||
|
||||
ETags are the cheap-skip key. Each sweep lists the prefix, compares
|
||||
the listed ETag against the document's stored `metadata["etag"]`, and
|
||||
only fetches keys whose ETag has changed. If the bytes turn out to
|
||||
match the stored MD5 (multipart re-upload landing a new ETag on the
|
||||
same content), only the etag is refreshed — no re-chunk.
|
||||
|
||||
`storage_options` follows the same convention as `lancedb.storage_options` —
|
||||
the dict is passed straight to obstore (the Rust `object_store` library
|
||||
LanceDB uses internally), so credentials configured for the LanceDB
|
||||
backend can be copy-pasted here.
|
||||
|
||||
### HTTP
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
sources:
|
||||
- type: http
|
||||
id: arxiv
|
||||
urls:
|
||||
- https://arxiv.org/pdf/2301.12345.pdf
|
||||
headers:
|
||||
Authorization: Bearer ${SOME_TOKEN}
|
||||
poll_interval_s: 86400
|
||||
```
|
||||
|
||||
HTTP is pull-based with HEAD-driven change detection. A `410 Gone`
|
||||
response from a configured URL triggers a delete event; other failure
|
||||
statuses fall through to UPSERT-with-no-revision so the worker can
|
||||
GET and decide.
|
||||
|
||||
## Workers and retry
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
workers:
|
||||
worker_count: 4
|
||||
max_concurrent: 4
|
||||
poll_idle_interval_s: 1.0
|
||||
claim_timeout_s: 1800
|
||||
reaper_interval_s: 60
|
||||
retry:
|
||||
max_attempts: 5
|
||||
base_delay_s: 2.0
|
||||
max_delay_s: 300.0
|
||||
jitter: 0.25 # ±25%
|
||||
```
|
||||
|
||||
The worker pool runs `worker_count` async workers behind a shared
|
||||
`max_concurrent` semaphore. Jobs that hit a `TransientError` are
|
||||
rescheduled with exponential backoff plus jitter, up to `max_attempts`,
|
||||
then land in the dead-letter queue. `PermanentError` (unsupported
|
||||
extension, 4xx HTTP except 408/429, etc.) skips retry entirely.
|
||||
|
||||
A reaper task resets jobs whose `claimed_at` is older than
|
||||
`claim_timeout_s` so a crashed worker doesn't strand its job.
|
||||
|
||||
**Per-source override.** A source can opt out of the global retry
|
||||
policy:
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
sources:
|
||||
- type: http
|
||||
id: flaky-api
|
||||
urls: [...]
|
||||
retry:
|
||||
max_attempts: 10
|
||||
base_delay_s: 10
|
||||
```
|
||||
|
||||
## Circuit breaker
|
||||
|
||||
After N consecutive `discover()` failures, a source's circuit breaker
|
||||
opens and polling pauses for a cooldown. Other sources keep running.
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
sources:
|
||||
- type: http
|
||||
id: rate-limited
|
||||
urls: [...]
|
||||
circuit_breaker:
|
||||
failure_threshold: 5
|
||||
cooldown_s: 600
|
||||
```
|
||||
|
||||
## Run it
|
||||
|
||||
```bash
|
||||
haiku-ingester serve # workers + pollers + API
|
||||
haiku-ingester serve --no-api # workers + pollers only
|
||||
haiku-ingester serve --db /path.lancedb # explicit DB
|
||||
```
|
||||
|
||||
The service blocks until SIGINT or SIGTERM. Shutdown drains the API
|
||||
server, then pollers, then in-flight workers.
|
||||
|
||||
### Single-writer constraint
|
||||
|
||||
LanceDB supports exactly one writer + N readers per database URI. Run
|
||||
exactly one `haiku-ingester serve` against a given LanceDB. Multiple
|
||||
MCP servers or read-only consumers against the same DB are fine.
|
||||
|
||||
## HTTP control plane
|
||||
|
||||
By default the ingester exposes a FastAPI control plane on
|
||||
`127.0.0.1:8765`. Set `ingester.api.auth_token` to require a Bearer
|
||||
token; without one the API stays open and the service logs a warning.
|
||||
|
||||
| Method | Path | Purpose |
|
||||
|---|---|---|
|
||||
| `GET` | `/health` | liveness + queue counts |
|
||||
| `GET` | `/sources` | configured pollers + last-poll time + breaker state |
|
||||
| `POST` | `/sources/{id}/refresh` | force an out-of-band sweep |
|
||||
| `GET` | `/jobs` | filtered list (`status`, `source_id`, `uri`, `limit`, `offset`) |
|
||||
| `GET` | `/jobs/{id}` | one job |
|
||||
| `POST` | `/jobs/{id}/retry` | reset attempts to 0, status to queued |
|
||||
| `DELETE` | `/jobs/{id}` | cancel a queued/claimed job |
|
||||
| `GET` | `/dlq` | dead jobs |
|
||||
| `POST` | `/dlq/{id}/retry` | resurrect from DLQ |
|
||||
|
||||
OpenAPI docs at `http://localhost:8765/docs`.
|
||||
|
||||
```yaml
|
||||
ingester:
|
||||
api:
|
||||
enabled: true
|
||||
host: 127.0.0.1
|
||||
port: 8765
|
||||
auth_token: ${INGESTER_TOKEN} # null → unauthenticated
|
||||
```
|
||||
|
||||
## Operating
|
||||
|
||||
### Smoke-test a single URI
|
||||
|
||||
`run-once` bypasses the queue and runs a single Job through the
|
||||
pipeline. Useful for sanity-checking a source before starting the
|
||||
service.
|
||||
|
||||
```bash
|
||||
haiku-ingester run-once /path/to/test.pdf
|
||||
haiku-ingester run-once https://example.com/spec.pdf
|
||||
haiku-ingester run-once s3://my-bucket/key.pdf
|
||||
```
|
||||
|
||||
Exit codes: `0` success, `1` transient error, `2` permanent error.
|
||||
|
||||
### The queue
|
||||
|
||||
The ingester's SQLite queue lives at
|
||||
`~/Library/Application Support/haiku.rag/ingester.db` on macOS
|
||||
(platform user data dir; configurable via `ingester.queue.path`). It's
|
||||
created automatically by `serve`.
|
||||
|
||||
For ops setup you can pre-create it:
|
||||
|
||||
```bash
|
||||
haiku-ingester queue init # create the DB and schema
|
||||
haiku-ingester queue migrate # apply pending schema changes
|
||||
```
|
||||
|
||||
### Logs
|
||||
|
||||
The service logs via Python `logging` to stderr through a Rich handler.
|
||||
A typical run looks like:
|
||||
|
||||
```
|
||||
INFO Ingester running: 4 worker(s), 1 source(s)
|
||||
INFO API listening on 127.0.0.1:8765
|
||||
INFO Swept local-docs: 142 upsert, 0 delete, 8 unchanged
|
||||
INFO Processing upsert file:///.../a.md (job 5d9a...)
|
||||
INFO Job 5d9a... succeeded in 0.34s: file:///.../a.md
|
||||
```
|
||||
|
||||
When `LOGFIRE_TOKEN` is set, spans are also shipped to Logfire.
|
||||
|
||||
### Operating against the API
|
||||
|
||||
```bash
|
||||
TOKEN=$INGESTER_TOKEN # omit -H entirely if no token configured
|
||||
|
||||
curl http://localhost:8765/health
|
||||
curl -H "Authorization: Bearer $TOKEN" http://localhost:8765/sources
|
||||
curl -H "Authorization: Bearer $TOKEN" 'http://localhost:8765/jobs?status=dead'
|
||||
|
||||
# Force a poll now
|
||||
curl -H "Authorization: Bearer $TOKEN" -X POST \
|
||||
http://localhost:8765/sources/local-docs/refresh
|
||||
|
||||
# Resurrect a dead job
|
||||
curl -H "Authorization: Bearer $TOKEN" -X POST \
|
||||
http://localhost:8765/jobs/<id>/retry
|
||||
```
|
||||
25
docs/mcp.md
25
docs/mcp.md
|
|
@ -8,16 +8,16 @@ The MCP server supports Streamable HTTP and stdio transports:
|
|||
|
||||
```bash
|
||||
# Default streamable HTTP transport on port 8001
|
||||
haiku-rag serve --mcp
|
||||
haiku-rag mcp
|
||||
|
||||
# Custom port
|
||||
haiku-rag serve --mcp --mcp-port 9000
|
||||
haiku-rag mcp --port 9000
|
||||
|
||||
# stdio transport (for Claude Desktop)
|
||||
haiku-rag serve --mcp --stdio
|
||||
haiku-rag mcp --stdio
|
||||
|
||||
# Read-only mode (excludes write tools)
|
||||
haiku-rag --read-only serve --mcp --stdio
|
||||
haiku-rag --read-only mcp --stdio
|
||||
```
|
||||
|
||||
**Read-only mode:** When `--read-only` is specified, write tools (`add_document_from_file`, `add_document_from_url`, `add_document_from_text`, `delete_document`) are not registered. Only search and query tools remain available.
|
||||
|
|
@ -31,7 +31,7 @@ Add to your Claude Desktop configuration (`claude_desktop_config.json`):
|
|||
"mcpServers": {
|
||||
"haiku-rag": {
|
||||
"command": "haiku-rag",
|
||||
"args": ["serve", "--mcp", "--stdio"]
|
||||
"args": ["mcp", "--stdio"]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -44,7 +44,7 @@ With a custom database path:
|
|||
"mcpServers": {
|
||||
"haiku-rag": {
|
||||
"command": "haiku-rag",
|
||||
"args": ["serve", "--mcp", "--stdio", "--db", "/path/to/database.lancedb"]
|
||||
"args": ["mcp", "--stdio", "--db", "/path/to/database.lancedb"]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -108,13 +108,8 @@ After restarting Claude Desktop, you can ask Claude to search your documents, ad
|
|||
- `document` (optional): Document title/ID to pre-load (can repeat)
|
||||
- Best for aggregation, computation, and multi-document analysis
|
||||
|
||||
## Running with Other Services
|
||||
## Continuous ingestion
|
||||
|
||||
Combine MCP with file monitoring:
|
||||
|
||||
```bash
|
||||
# MCP + file monitoring
|
||||
haiku-rag serve --mcp --monitor
|
||||
```
|
||||
|
||||
See [Server Mode](server.md) for details on file monitoring.
|
||||
For continuous document ingestion (filesystem watch, S3 polling, HTTP
|
||||
sources, a job queue with retries), run [`haiku-ingester`](ingester.md)
|
||||
as a separate process against the same LanceDB.
|
||||
|
|
|
|||
172
docs/server.md
172
docs/server.md
|
|
@ -1,172 +0,0 @@
|
|||
# Server Mode
|
||||
|
||||
The server provides automatic file monitoring and MCP functionality.
|
||||
|
||||
## Starting the Server
|
||||
|
||||
The `serve` command requires at least one service flag. You can enable file monitoring, MCP server, or both:
|
||||
|
||||
### MCP Server Only
|
||||
|
||||
```bash
|
||||
haiku-rag serve --mcp
|
||||
```
|
||||
|
||||
Transport options:
|
||||
- Default - Streamable HTTP transport on port 8001
|
||||
- `--stdio` - Standard input/output transport
|
||||
- `--mcp-port` - Custom port (default: 8001)
|
||||
|
||||
### File Monitoring Only
|
||||
|
||||
```bash
|
||||
haiku-rag serve --monitor
|
||||
```
|
||||
|
||||
### Both Services
|
||||
|
||||
```bash
|
||||
haiku-rag serve --monitor --mcp
|
||||
```
|
||||
|
||||
This will start file monitoring and MCP server on port 8001.
|
||||
|
||||
## File Monitoring
|
||||
|
||||
Configure directories to monitor in your `haiku.rag.yaml` (see [Document Processing](configuration/processing.md#file-monitoring) for all options):
|
||||
|
||||
```yaml
|
||||
monitor:
|
||||
directories:
|
||||
- /path/to/documents
|
||||
- /another/path
|
||||
```
|
||||
|
||||
Then start the server:
|
||||
|
||||
```bash
|
||||
haiku-rag serve --monitor
|
||||
```
|
||||
|
||||
### Monitoring Features
|
||||
|
||||
- **Startup**: Scans all monitored directories and adds new files
|
||||
- **File Added/Modified**: Automatically parses and updates documents
|
||||
- **File Deleted**: Removes corresponding documents from database
|
||||
|
||||
### Filtering Files
|
||||
|
||||
You can filter which files to monitor using gitignore-style patterns:
|
||||
|
||||
```yaml
|
||||
monitor:
|
||||
directories:
|
||||
- /path/to/documents
|
||||
|
||||
# Ignore patterns (exclude files)
|
||||
ignore_patterns:
|
||||
- "*draft*" # Ignore draft files
|
||||
- "temp/" # Ignore temp directory
|
||||
- "**/archive/**" # Ignore archive directories
|
||||
|
||||
# Include patterns (whitelist files)
|
||||
include_patterns:
|
||||
- "*.md" # Only markdown files
|
||||
- "**/docs/**" # Files in docs directories
|
||||
```
|
||||
|
||||
**Pattern behavior:**
|
||||
- Extension filtering is applied first (only supported file types)
|
||||
- Include patterns create a whitelist (if specified)
|
||||
- Ignore patterns exclude files
|
||||
- Both can be combined for fine-grained control
|
||||
|
||||
### Supported Formats
|
||||
|
||||
The file monitor processes documents using [Docling](https://github.com/DS4SD/docling), which supports:
|
||||
|
||||
**Documents:**
|
||||
- PDF (`.pdf`) - with OCR support for scanned documents
|
||||
- Microsoft Word (`.docx`)
|
||||
- Microsoft Excel (`.xlsx`)
|
||||
- Microsoft PowerPoint (`.pptx`)
|
||||
- HTML (`.html`, `.htm`)
|
||||
- Markdown (`.md`)
|
||||
- Quarto Markdown (`.qmd`)
|
||||
- R Markdown (`.rmd`)
|
||||
- LaTeX (`.tex`, `.latex`)
|
||||
- AsciiDoc (`.adoc`, `.asciidoc`)
|
||||
|
||||
**Data formats:**
|
||||
- CSV (`.csv`)
|
||||
- JSON (`.json`)
|
||||
- XML (`.xml`)
|
||||
|
||||
**Images (via OCR):**
|
||||
- PNG (`.png`)
|
||||
- JPEG (`.jpg`, `.jpeg`)
|
||||
- TIFF (`.tiff`, `.tif`)
|
||||
- BMP (`.bmp`)
|
||||
|
||||
**Code files:**
|
||||
- Python (`.py`)
|
||||
- JavaScript (`.js`)
|
||||
- TypeScript (`.ts`)
|
||||
- PlantUML (`.puml`, `.plantuml`, `.pu`)
|
||||
- And other text-based code files
|
||||
|
||||
**Plain text:**
|
||||
- Text files (`.txt`)
|
||||
- RST (`.rst`)
|
||||
|
||||
URLs are also supported - the content is fetched and converted to markdown.
|
||||
|
||||
## S3 / Object Storage Monitoring
|
||||
|
||||
The server can also poll S3-compatible object storage (AWS S3, SeaweedFS, MinIO, Cloudflare R2, etc.) for new, modified, and deleted objects, treating each one as a document source.
|
||||
|
||||
Install the optional `[s3]` extra:
|
||||
|
||||
```bash
|
||||
pip install haiku.rag-slim[s3]
|
||||
# or, for the full package:
|
||||
pip install haiku.rag[s3]
|
||||
```
|
||||
|
||||
Configure one or more S3 sources under `monitor.s3` in `haiku.rag.yaml`:
|
||||
|
||||
```yaml
|
||||
monitor:
|
||||
s3:
|
||||
- uri: s3://my-bucket/incoming/
|
||||
poll_interval: 300 # seconds between sweeps; default 300
|
||||
include_patterns: ["*.pdf", "*.md"]
|
||||
ignore_patterns: ["draft*"]
|
||||
delete_orphans: true
|
||||
storage_options:
|
||||
endpoint: http://seaweed:8333
|
||||
aws_access_key_id: ${AWS_KEY}
|
||||
aws_secret_access_key: ${AWS_SECRET}
|
||||
region: us-east-1
|
||||
allow_http: "true"
|
||||
```
|
||||
|
||||
Then start the server with `--monitor` (the same flag enables both local-directory and S3 watchers):
|
||||
|
||||
```bash
|
||||
haiku-rag serve --monitor
|
||||
```
|
||||
|
||||
Each entry in `monitor.s3` runs as its own polling task. On every sweep the watcher lists all objects under the configured prefix, compares each object's S3 ETag against the document's stored `metadata["etag"]`, and only re-fetches keys whose ETag has changed. When the bytes turn out to match the stored MD5 (e.g. the same file was re-uploaded with a different multipart chunk size), the watcher refreshes the etag and skips re-chunking. Otherwise the document is downloaded, chunked, and re-embedded.
|
||||
|
||||
### Credentials
|
||||
|
||||
`storage_options` follows the same convention as `lancedb.storage_options`. The dict is passed straight to obstore (the same Rust `object_store` library LanceDB uses internally), so any keys you've configured there work here too. When `storage_options` is omitted, the watcher falls back to the AWS default credential chain (environment variables, IAM instance role, AWS profile).
|
||||
|
||||
### Orphan deletion scope
|
||||
|
||||
`delete_orphans: true` is per-entry: a watcher only removes documents whose URI starts with that entry's `s3://bucket/prefix/`. Documents from other buckets, prefixes, or local-file sources are never touched.
|
||||
|
||||
## One-off ingestion
|
||||
|
||||
`s3://` URIs are also a first-class source for `haiku-rag add-src` and the MCP `add_document_from_url` tool. See [CLI → Add Documents](cli.md#add-documents).
|
||||
|
|
@ -87,7 +87,7 @@ See the [Web application](../apps.md) reference implementation.
|
|||
To use a skill from Claude Desktop or another MCP-aware client, run the MCP server:
|
||||
|
||||
```bash
|
||||
haiku-rag serve --mcp --stdio
|
||||
haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
The server exposes the skill tools (search, ask, analyze) over MCP. See [MCP](../mcp.md).
|
||||
|
|
|
|||
|
|
@ -144,7 +144,7 @@ See the [Web application](../apps.md) reference implementation for the full Star
|
|||
To call the skill from Claude Desktop (or any MCP client), run the MCP server:
|
||||
|
||||
```bash
|
||||
haiku-rag serve --mcp --stdio
|
||||
haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
The exposed `ask_question` tool runs this skill. See [MCP](../mcp.md) for the configuration block.
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ The docker-compose.yml mounts three volumes:
|
|||
cp haiku.rag.yaml.example haiku.rag.yaml
|
||||
```
|
||||
|
||||
The example config sets `monitor.directories: [/docs]` - this is the **container path**, not your host path. Documents placed in `./docs` on your host will appear at `/docs` inside the container.
|
||||
The example config sets `ingester.sources[0].root: /docs` - this is the **container path**, not your host path. Documents placed in `./docs` on your host will appear at `/docs` inside the container.
|
||||
|
||||
## Usage
|
||||
|
||||
|
|
|
|||
|
|
@ -6,9 +6,12 @@ environment: production
|
|||
storage:
|
||||
data_dir: /data
|
||||
|
||||
monitor:
|
||||
directories:
|
||||
- /docs
|
||||
ingester:
|
||||
sources:
|
||||
- type: fs
|
||||
id: docs
|
||||
root: /docs
|
||||
delete_orphans: true
|
||||
|
||||
# Remote document processing with docling-serve
|
||||
processing:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
|
@ -20,7 +19,6 @@ from rich.progress import (
|
|||
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||
from haiku.rag.config import AppConfig, Config
|
||||
from haiku.rag.mcp import create_mcp_server
|
||||
from haiku.rag.monitor import FileWatcher, S3Watcher
|
||||
from haiku.rag.store.models.chunk import SearchType
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
|
|
@ -727,77 +725,26 @@ class HaikuRAGApp: # pragma: no cover
|
|||
self.console.print(content)
|
||||
self.console.rule()
|
||||
|
||||
async def serve(
|
||||
async def run_mcp(
|
||||
self,
|
||||
enable_monitor: bool = True,
|
||||
enable_mcp: bool = True,
|
||||
mcp_transport: str | None = None,
|
||||
mcp_port: int = 8001,
|
||||
transport: str | None = None,
|
||||
port: int = 8001,
|
||||
):
|
||||
"""Start the server with selected services."""
|
||||
"""Run the MCP server until interrupted."""
|
||||
async with HaikuRAG(
|
||||
self.db_path,
|
||||
config=self.config,
|
||||
read_only=self.read_only,
|
||||
before=self.before,
|
||||
) as client:
|
||||
tasks = []
|
||||
|
||||
# Start file monitor if enabled (not available in read-only mode)
|
||||
if enable_monitor:
|
||||
if self.read_only:
|
||||
logger.warning(
|
||||
"File monitor disabled: cannot monitor files in read-only mode"
|
||||
)
|
||||
else:
|
||||
monitor = FileWatcher(client=client, config=self.config)
|
||||
monitor_task = asyncio.create_task(monitor.observe())
|
||||
tasks.append(monitor_task)
|
||||
|
||||
if self.config.monitor.s3:
|
||||
from haiku.rag.converters import get_converter
|
||||
|
||||
supported_extensions = get_converter(
|
||||
self.config
|
||||
).supported_extensions
|
||||
for entry in self.config.monitor.s3:
|
||||
s3_watcher = S3Watcher(
|
||||
client=client,
|
||||
entry=entry,
|
||||
supported_extensions=supported_extensions,
|
||||
)
|
||||
tasks.append(asyncio.create_task(s3_watcher.observe()))
|
||||
|
||||
# Start MCP server if enabled
|
||||
if enable_mcp:
|
||||
server = create_mcp_server(
|
||||
self.db_path, config=self.config, read_only=self.read_only
|
||||
)
|
||||
|
||||
async def run_mcp():
|
||||
if mcp_transport == "stdio":
|
||||
await server.run_stdio_async()
|
||||
else:
|
||||
logger.info(f"Starting MCP server on port {mcp_port}")
|
||||
await server.run_http_async(
|
||||
transport="streamable-http", port=mcp_port
|
||||
)
|
||||
|
||||
mcp_task = asyncio.create_task(run_mcp())
|
||||
tasks.append(mcp_task)
|
||||
|
||||
if not tasks:
|
||||
logger.warning("No services enabled")
|
||||
return
|
||||
|
||||
):
|
||||
server = create_mcp_server(
|
||||
self.db_path, config=self.config, read_only=self.read_only
|
||||
)
|
||||
try:
|
||||
# Wait for any task to complete (or KeyboardInterrupt)
|
||||
await asyncio.gather(*tasks)
|
||||
if transport == "stdio":
|
||||
await server.run_stdio_async()
|
||||
else:
|
||||
logger.info(f"Starting MCP server on port {port}")
|
||||
await server.run_http_async(transport="streamable-http", port=port)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
# Cancel all tasks
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
# Wait for cancellation
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
|
|
|||
|
|
@ -671,59 +671,33 @@ def chat( # pragma: no cover
|
|||
|
||||
|
||||
@_cli.command(
|
||||
"serve",
|
||||
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
|
||||
"mcp",
|
||||
help="Run the MCP server. For continuous ingestion, use haiku-ingester serve.",
|
||||
)
|
||||
def serve(
|
||||
def mcp(
|
||||
db: Path | None = typer.Option(
|
||||
None,
|
||||
"--db",
|
||||
help="Path to the LanceDB database file",
|
||||
),
|
||||
monitor: bool = typer.Option(
|
||||
False,
|
||||
"--monitor",
|
||||
help="Enable file monitoring",
|
||||
),
|
||||
mcp: bool = typer.Option(
|
||||
False,
|
||||
"--mcp",
|
||||
help="Enable MCP server",
|
||||
),
|
||||
stdio: bool = typer.Option(
|
||||
False,
|
||||
"--stdio",
|
||||
help="Run MCP server on stdio Transport (requires --mcp)",
|
||||
help="Run MCP server on stdio Transport",
|
||||
),
|
||||
mcp_port: int = typer.Option(
|
||||
port: int = typer.Option(
|
||||
8001,
|
||||
"--mcp-port",
|
||||
"--port",
|
||||
help="Port to bind MCP server to (ignored with --stdio)",
|
||||
),
|
||||
) -> None:
|
||||
"""Start the server with selected services."""
|
||||
# Require at least one service flag
|
||||
if not (monitor or mcp):
|
||||
typer.echo(
|
||||
"Error: At least one service flag (--monitor or --mcp) must be specified"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
if stdio and not mcp:
|
||||
typer.echo("Error: --stdio requires --mcp")
|
||||
raise typer.Exit(1)
|
||||
|
||||
"""Run the MCP server."""
|
||||
app = create_app(db) # pragma: no cover
|
||||
|
||||
transport = "stdio" if stdio else None # pragma: no cover
|
||||
|
||||
asyncio.run( # pragma: no cover
|
||||
app.serve(
|
||||
enable_monitor=monitor,
|
||||
enable_mcp=mcp,
|
||||
mcp_transport=transport,
|
||||
mcp_port=mcp_port,
|
||||
)
|
||||
app.run_mcp(transport=transport, port=port)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -328,11 +328,11 @@ async def create_document_from_source(
|
|||
)
|
||||
from haiku.rag.ingester.sources.filter import FileFilter
|
||||
|
||||
# One-shot CLI directory ingest uses the converter's supported
|
||||
# extensions but no include/ignore patterns. For pattern-based
|
||||
# filtering use `haiku-ingester serve` with an FS source.
|
||||
documents: list[Document] = []
|
||||
filter = FileFilter(
|
||||
ignore_patterns=client._config.monitor.ignore_patterns or None,
|
||||
include_patterns=client._config.monitor.include_patterns or None,
|
||||
)
|
||||
filter = FileFilter()
|
||||
for child in local_path.rglob("*"):
|
||||
if child.is_file() and filter.include_file(str(child)):
|
||||
doc = await create_document_from_source(
|
||||
|
|
|
|||
|
|
@ -15,7 +15,6 @@ from haiku.rag.config.models import (
|
|||
IngesterConfig,
|
||||
LanceDBConfig,
|
||||
ModelConfig,
|
||||
MonitorConfig,
|
||||
OllamaConfig,
|
||||
ProcessingConfig,
|
||||
PromptsConfig,
|
||||
|
|
@ -24,7 +23,6 @@ from haiku.rag.config.models import (
|
|||
QueueConfig,
|
||||
RerankingConfig,
|
||||
RetryPolicyConfig,
|
||||
S3MonitorEntry,
|
||||
S3SourceConfig,
|
||||
SourceConfig,
|
||||
StorageConfig,
|
||||
|
|
@ -44,7 +42,6 @@ __all__ = [
|
|||
"IngesterConfig",
|
||||
"LanceDBConfig",
|
||||
"ModelConfig",
|
||||
"MonitorConfig",
|
||||
"OllamaConfig",
|
||||
"ProcessingConfig",
|
||||
"PromptsConfig",
|
||||
|
|
@ -53,7 +50,6 @@ __all__ = [
|
|||
"QueueConfig",
|
||||
"RerankingConfig",
|
||||
"RetryPolicyConfig",
|
||||
"S3MonitorEntry",
|
||||
"S3SourceConfig",
|
||||
"SourceConfig",
|
||||
"StorageConfig",
|
||||
|
|
|
|||
|
|
@ -57,23 +57,6 @@ class StorageConfig(BaseModel):
|
|||
vacuum_retention_seconds: int = 86400
|
||||
|
||||
|
||||
class S3MonitorEntry(BaseModel):
|
||||
uri: str
|
||||
storage_options: dict[str, str] = Field(default_factory=dict)
|
||||
poll_interval: int = 300
|
||||
ignore_patterns: list[str] = []
|
||||
include_patterns: list[str] = []
|
||||
delete_orphans: bool = False
|
||||
|
||||
|
||||
class MonitorConfig(BaseModel):
|
||||
directories: list[Path] = []
|
||||
ignore_patterns: list[str] = []
|
||||
include_patterns: list[str] = []
|
||||
delete_orphans: bool = False
|
||||
s3: list[S3MonitorEntry] = []
|
||||
|
||||
|
||||
class LanceDBConfig(BaseModel):
|
||||
uri: str = ""
|
||||
api_key: str = ""
|
||||
|
|
@ -358,7 +341,6 @@ class IngesterConfig(BaseModel):
|
|||
class AppConfig(BaseModel):
|
||||
environment: str = "production"
|
||||
storage: StorageConfig = Field(default_factory=StorageConfig)
|
||||
monitor: MonitorConfig = Field(default_factory=MonitorConfig)
|
||||
lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig)
|
||||
embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
|
||||
reranking: RerankingConfig = Field(default_factory=RerankingConfig)
|
||||
|
|
|
|||
|
|
@ -1,261 +0,0 @@
|
|||
import asyncio
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from watchfiles import Change, awatch
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig, Config, S3MonitorEntry
|
||||
from haiku.rag.ingester.sources.filter import FileFilter
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.utils import escape_sql_string
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["FileFilter", "FileWatcher", "S3Watcher"]
|
||||
|
||||
|
||||
class FileWatcher:
|
||||
def __init__(
|
||||
self,
|
||||
client: HaikuRAG,
|
||||
config: AppConfig = Config,
|
||||
):
|
||||
from haiku.rag.converters import get_converter
|
||||
|
||||
self.paths = config.monitor.directories
|
||||
self.client = client
|
||||
self.ignore_patterns = config.monitor.ignore_patterns or None
|
||||
self.include_patterns = config.monitor.include_patterns or None
|
||||
self.delete_orphans = config.monitor.delete_orphans
|
||||
self.supported_extensions = get_converter(config).supported_extensions
|
||||
|
||||
async def observe(self):
|
||||
if not self.paths:
|
||||
logger.warning("No directories configured for monitoring")
|
||||
return
|
||||
|
||||
# Validate all paths exist before attempting to watch
|
||||
missing_paths = [p for p in self.paths if not Path(p).exists()]
|
||||
if missing_paths:
|
||||
raise FileNotFoundError(
|
||||
f"Monitor directories do not exist: {missing_paths}. "
|
||||
"Check your haiku.rag.yaml configuration."
|
||||
)
|
||||
|
||||
logger.info(f"Watching files in {self.paths}")
|
||||
filter = FileFilter(
|
||||
ignore_patterns=self.ignore_patterns,
|
||||
include_patterns=self.include_patterns,
|
||||
supported_extensions=self.supported_extensions,
|
||||
)
|
||||
await self.refresh()
|
||||
|
||||
async for changes in awatch(*self.paths, watch_filter=filter):
|
||||
await self.handler(changes)
|
||||
|
||||
async def handler(self, changes: set[tuple[Change, str]]):
|
||||
for change, path in changes:
|
||||
if change == Change.added or change == Change.modified:
|
||||
await self._upsert_document(Path(path))
|
||||
elif change == Change.deleted:
|
||||
await self._delete_document(Path(path))
|
||||
|
||||
async def refresh(self):
|
||||
# Delete orphaned documents in background if enabled
|
||||
if self.delete_orphans:
|
||||
logger.info("Starting orphan cleanup in background")
|
||||
asyncio.create_task(self._delete_orphans())
|
||||
|
||||
# Create filter to apply same logic as observe()
|
||||
filter = FileFilter(
|
||||
ignore_patterns=self.ignore_patterns,
|
||||
include_patterns=self.include_patterns,
|
||||
supported_extensions=self.supported_extensions,
|
||||
)
|
||||
|
||||
for path in self.paths:
|
||||
for f in Path(path).rglob("**/*"):
|
||||
if f.is_file() and f.suffix in self.supported_extensions:
|
||||
# Apply pattern filters
|
||||
if filter(Change.added, str(f)):
|
||||
await self._upsert_document(f)
|
||||
|
||||
async def _upsert_document(self, file: Path) -> Document | None:
|
||||
try:
|
||||
uri = file.as_uri()
|
||||
existing_doc = await self.client.get_document_by_uri(uri)
|
||||
|
||||
result = await self.client.create_document_from_source(str(file))
|
||||
doc = result if isinstance(result, Document) else result[0]
|
||||
|
||||
if existing_doc:
|
||||
# Check if document was actually updated by comparing updated_at timestamps
|
||||
if doc.updated_at > existing_doc.updated_at:
|
||||
logger.info(f"Updated document {existing_doc.id} from {file}")
|
||||
else:
|
||||
logger.info(
|
||||
f"Skipped unchanged document {existing_doc.id} from {file}"
|
||||
)
|
||||
else:
|
||||
logger.info(f"Created new document {doc.id} from {file}")
|
||||
|
||||
return doc
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to upsert document from {file}: {e}")
|
||||
return None
|
||||
|
||||
async def _delete_orphans(self):
|
||||
"""Delete documents whose source files no longer exist."""
|
||||
try:
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
# Create filter to apply same include/exclude logic
|
||||
filter = FileFilter(
|
||||
ignore_patterns=self.ignore_patterns,
|
||||
include_patterns=self.include_patterns,
|
||||
)
|
||||
|
||||
all_docs = await self.client.list_documents()
|
||||
|
||||
for doc in all_docs:
|
||||
if not doc.uri or not doc.id:
|
||||
continue
|
||||
|
||||
# Only check file:// URIs
|
||||
parsed = urlparse(doc.uri)
|
||||
if parsed.scheme != "file":
|
||||
continue
|
||||
|
||||
# Convert URI to Path, decoding URL-encoded characters (like %20 for spaces)
|
||||
file_path = Path(unquote(parsed.path))
|
||||
|
||||
# Check if file exists
|
||||
if not file_path.exists():
|
||||
# Check if file is within monitored directories
|
||||
is_monitored = any(
|
||||
file_path.is_relative_to(monitored_path)
|
||||
for monitored_path in self.paths
|
||||
)
|
||||
|
||||
# Check if file would have been included by filters
|
||||
if is_monitored and filter.include_file(str(file_path)):
|
||||
await self.client.delete_document(doc.id)
|
||||
logger.info(
|
||||
f"Deleted orphaned document {doc.id} for {file_path}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete orphaned documents: {e}")
|
||||
|
||||
async def _delete_document(self, file: Path):
|
||||
try:
|
||||
uri = file.as_uri()
|
||||
existing_doc = await self.client.get_document_by_uri(uri)
|
||||
|
||||
if existing_doc and existing_doc.id:
|
||||
await self.client.delete_document(existing_doc.id)
|
||||
logger.info(f"Deleted document {existing_doc.id} for {file}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete document for {file}: {e}")
|
||||
|
||||
|
||||
class S3Watcher:
|
||||
"""Polls an S3 prefix and keeps documents in sync with the index.
|
||||
|
||||
Uses ListObjectsV2 ETags as the cheap-skip key. When a key's listing
|
||||
ETag differs from the stored `metadata["etag"]`, delegates to
|
||||
`client.create_document_from_source` which performs the full
|
||||
HeadObject + GetObject + MD5 compare two-stage detection.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: HaikuRAG,
|
||||
entry: S3MonitorEntry,
|
||||
supported_extensions: list[str],
|
||||
) -> None:
|
||||
from haiku.rag.s3 import make_s3_store
|
||||
|
||||
parsed = urlparse(entry.uri)
|
||||
if not parsed.netloc:
|
||||
raise ValueError(f"Invalid S3 monitor URI: {entry.uri}")
|
||||
|
||||
self.client = client
|
||||
self.entry = entry
|
||||
self.bucket = parsed.netloc
|
||||
self.prefix = parsed.path.lstrip("/")
|
||||
self.uri_prefix = f"s3://{self.bucket}/{self.prefix}"
|
||||
self._make_s3_store = make_s3_store
|
||||
self.filter = FileFilter(
|
||||
ignore_patterns=entry.ignore_patterns or None,
|
||||
include_patterns=entry.include_patterns or None,
|
||||
supported_extensions=supported_extensions,
|
||||
)
|
||||
|
||||
async def observe(self) -> None:
|
||||
logger.info(
|
||||
f"Watching S3 {self.entry.uri} (poll_interval={self.entry.poll_interval}s)"
|
||||
)
|
||||
await self.refresh()
|
||||
while True:
|
||||
await asyncio.sleep(self.entry.poll_interval)
|
||||
try:
|
||||
await self.refresh()
|
||||
except Exception as e:
|
||||
logger.error(f"S3 watcher refresh failed for {self.entry.uri}: {e}")
|
||||
|
||||
async def refresh(self) -> None:
|
||||
import obstore # type: ignore[import-not-found]
|
||||
|
||||
uris_seen: dict[str, str] = {}
|
||||
store = self._make_s3_store(self.bucket, self.entry.storage_options)
|
||||
|
||||
async for batch in obstore.list(store, prefix=self.prefix or None):
|
||||
for obj in batch:
|
||||
key = obj["path"]
|
||||
if not self.filter.include_file(key):
|
||||
continue
|
||||
uri = f"s3://{self.bucket}/{key}"
|
||||
uris_seen[uri] = (obj.get("e_tag") or "").strip('"')
|
||||
|
||||
existing_etags = await self._existing_etags_under_prefix()
|
||||
|
||||
for uri, etag in uris_seen.items():
|
||||
if existing_etags.get(uri) == etag:
|
||||
continue
|
||||
await self._upsert_object(uri)
|
||||
|
||||
if self.entry.delete_orphans:
|
||||
await self._delete_orphans(set(uris_seen.keys()), existing_etags)
|
||||
|
||||
async def _existing_etags_under_prefix(self) -> dict[str, str]:
|
||||
safe_prefix = escape_sql_string(self.uri_prefix)
|
||||
docs = await self.client.list_documents(filter=f"uri LIKE '{safe_prefix}%'")
|
||||
return {
|
||||
doc.uri: (doc.metadata or {}).get("etag", "") for doc in docs if doc.uri
|
||||
}
|
||||
|
||||
async def _upsert_object(self, uri: str) -> Document | None:
|
||||
try:
|
||||
result = await self.client.create_document_from_source(
|
||||
uri, storage_options=self.entry.storage_options
|
||||
)
|
||||
doc = result if isinstance(result, Document) else result[0]
|
||||
logger.info(f"Upserted document {doc.id} from {uri}")
|
||||
return doc
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to upsert document from {uri}: {e}")
|
||||
return None
|
||||
|
||||
async def _delete_orphans(
|
||||
self, uris_seen: set[str], existing_etags: dict[str, str]
|
||||
) -> None:
|
||||
for uri in existing_etags.keys() - uris_seen:
|
||||
try:
|
||||
doc = await self.client.get_document_by_uri(uri)
|
||||
if doc and doc.id:
|
||||
await self.client.delete_document(doc.id)
|
||||
logger.info(f"Deleted orphaned document {doc.id} for {uri}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete orphan {uri}: {e}")
|
||||
|
|
@ -152,10 +152,3 @@ async def test_fs_source_discover_respects_include_patterns(fs_root: Path):
|
|||
uris = {e.uri async for e in src.discover(since=None)}
|
||||
assert (fs_root / "b.txt").as_uri() not in uris
|
||||
assert (fs_root / "a.md").as_uri() in uris
|
||||
|
||||
|
||||
def test_filefilter_backward_compatible_reexport():
|
||||
from haiku.rag.ingester.sources.filter import FileFilter as IngesterFileFilter
|
||||
from haiku.rag.monitor import FileFilter as MonitorFileFilter
|
||||
|
||||
assert MonitorFileFilter is IngesterFileFilter
|
||||
|
|
|
|||
183
tests/ingester/test_serve_integration.py
Normal file
183
tests/ingester/test_serve_integration.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
"""End-to-end ingester tests: poller -> queue -> worker -> sync_state."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import aiosqlite
|
||||
import pytest
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import FSSourceConfig
|
||||
from haiku.rag.ingester.pollers.manager import PollerManager
|
||||
from haiku.rag.ingester.queue.migrations import apply_migrations
|
||||
from haiku.rag.ingester.queue.models import JobOp
|
||||
from haiku.rag.ingester.queue.repository import JobRepo, SyncStateRepo
|
||||
from haiku.rag.ingester.workers.pool import WorkerPool
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def conn(tmp_path):
|
||||
path = tmp_path / "queue.db"
|
||||
connection = await aiosqlite.connect(str(path))
|
||||
connection.row_factory = aiosqlite.Row
|
||||
await apply_migrations(connection)
|
||||
yield connection
|
||||
await connection.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def jobs(conn):
|
||||
return JobRepo(conn)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sync(conn):
|
||||
return SyncStateRepo(conn)
|
||||
|
||||
|
||||
async def _wait_for(predicate, *, timeout: float = 5.0, interval: float = 0.05):
|
||||
"""Poll `predicate` until it returns truthy or `timeout` elapses."""
|
||||
deadline = asyncio.get_running_loop().time() + timeout
|
||||
while True:
|
||||
result = (
|
||||
await predicate() if asyncio.iscoroutinefunction(predicate) else predicate()
|
||||
)
|
||||
if result:
|
||||
return result
|
||||
if asyncio.get_running_loop().time() >= deadline:
|
||||
raise AssertionError(f"predicate never became truthy within {timeout}s")
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
|
||||
def _mock_client(docs_root) -> AsyncMock:
|
||||
"""A HaikuRAG mock that returns a fresh Document for each URI it's asked
|
||||
to ingest, mirroring real metadata shape (contentType + md5)."""
|
||||
client = AsyncMock(spec=HaikuRAG)
|
||||
|
||||
counter = {"n": 0}
|
||||
|
||||
async def _fake_create(uri, *_, metadata=None, **__):
|
||||
counter["n"] += 1
|
||||
return Document(
|
||||
id=f"doc-{counter['n']}",
|
||||
content="x",
|
||||
uri=uri,
|
||||
metadata={"contentType": "text/markdown", "md5": f"md5-{counter['n']}"},
|
||||
)
|
||||
|
||||
client.create_document_from_source.side_effect = _fake_create
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_initial_sweep_lands_succeeded_jobs(tmp_path, jobs, sync):
|
||||
"""PollerManager + WorkerPool together: a file on disk at startup becomes
|
||||
a succeeded queue row and a sync_state entry."""
|
||||
(tmp_path / "a.md").write_text("hello")
|
||||
(tmp_path / "b.md").write_text("world")
|
||||
|
||||
client = _mock_client(tmp_path)
|
||||
cfg = FSSourceConfig(
|
||||
type="fs",
|
||||
id="local",
|
||||
root=tmp_path,
|
||||
poll_interval_s=60.0,
|
||||
)
|
||||
|
||||
manager = PollerManager(
|
||||
configs=[cfg],
|
||||
job_repo=jobs,
|
||||
sync_repo=sync,
|
||||
supported_extensions=[".md"],
|
||||
)
|
||||
pool = WorkerPool(
|
||||
client=client,
|
||||
job_repo=jobs,
|
||||
sync_repo=sync,
|
||||
worker_count=2,
|
||||
max_concurrent=2,
|
||||
poll_idle_interval_s=0.05,
|
||||
)
|
||||
|
||||
await pool.start()
|
||||
await manager.start()
|
||||
try:
|
||||
|
||||
async def _two_succeeded() -> bool:
|
||||
counts = await jobs.counts_by_status()
|
||||
return counts.get("succeeded", 0) == 2
|
||||
|
||||
await _wait_for(_two_succeeded, timeout=5.0)
|
||||
finally:
|
||||
await manager.stop()
|
||||
await pool.stop()
|
||||
|
||||
counts = await jobs.counts_by_status()
|
||||
assert counts.get("succeeded", 0) == 2
|
||||
assert counts.get("queued", 0) == 0
|
||||
assert counts.get("dead", 0) == 0
|
||||
|
||||
# The worker called create_document_from_source exactly twice — once per file.
|
||||
assert client.create_document_from_source.await_count == 2
|
||||
ingested_uris = {
|
||||
call.args[0] for call in client.create_document_from_source.await_args_list
|
||||
}
|
||||
assert ingested_uris == {
|
||||
(tmp_path / "a.md").as_uri(),
|
||||
(tmp_path / "b.md").as_uri(),
|
||||
}
|
||||
|
||||
# sync_state holds last_seen_at + content_hash for each URI.
|
||||
row_a = await sync.get_row("local", (tmp_path / "a.md").as_uri())
|
||||
row_b = await sync.get_row("local", (tmp_path / "b.md").as_uri())
|
||||
assert row_a is not None and row_a.content_hash and row_a.last_ingested_at
|
||||
assert row_b is not None and row_b.content_hash and row_b.last_ingested_at
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_watchfiles_push_event_lands_as_job(tmp_path, jobs, sync):
|
||||
"""FSPoller's watchfiles loop: a file *added* after startup should land
|
||||
as a queued job without waiting for the periodic sweep. No worker pool
|
||||
here — we're only asserting that watchfiles surfaces the event to the
|
||||
poller, which enqueues."""
|
||||
cfg = FSSourceConfig(
|
||||
type="fs",
|
||||
id="local",
|
||||
root=tmp_path,
|
||||
# poll_interval is far in the future so the periodic sweep CAN'T be
|
||||
# what picks up the new file — only watchfiles can.
|
||||
poll_interval_s=3600.0,
|
||||
)
|
||||
|
||||
manager = PollerManager(
|
||||
configs=[cfg],
|
||||
job_repo=jobs,
|
||||
sync_repo=sync,
|
||||
supported_extensions=[".md"],
|
||||
)
|
||||
|
||||
await manager.start()
|
||||
try:
|
||||
# Initial sweep saw an empty dir — give it a moment to settle, then
|
||||
# write a new file. watchfiles polls fs every ~50ms by default.
|
||||
async def _initial_sweep_done() -> bool:
|
||||
return manager.pollers[0].last_polled_at is not None
|
||||
|
||||
await _wait_for(_initial_sweep_done, timeout=5.0)
|
||||
assert await jobs.counts_by_status() == {}
|
||||
|
||||
(tmp_path / "new.md").write_text("after startup")
|
||||
|
||||
async def _one_queued() -> bool:
|
||||
queued = await jobs.list_jobs(source_id="local")
|
||||
return any(j.uri == (tmp_path / "new.md").as_uri() for j in queued)
|
||||
|
||||
await _wait_for(_one_queued, timeout=5.0)
|
||||
finally:
|
||||
await manager.stop()
|
||||
|
||||
queued = await jobs.list_jobs(source_id="local")
|
||||
assert len(queued) == 1
|
||||
assert queued[0].op is JobOp.UPSERT
|
||||
assert queued[0].uri == (tmp_path / "new.md").as_uri()
|
||||
|
|
@ -65,18 +65,6 @@ class TestParseMetaOptions:
|
|||
assert result == {"equation": "a=b+c"}
|
||||
|
||||
|
||||
class TestServeValidation:
|
||||
def test_no_flags_fails(self):
|
||||
result = runner.invoke(cli, ["serve"])
|
||||
assert result.exit_code == 1
|
||||
assert "At least one service flag" in result.output
|
||||
|
||||
def test_stdio_without_mcp_fails(self):
|
||||
result = runner.invoke(cli, ["serve", "--stdio", "--monitor"])
|
||||
assert result.exit_code == 1
|
||||
assert "--stdio requires --mcp" in result.output
|
||||
|
||||
|
||||
class TestRebuildValidation:
|
||||
def test_embed_only_and_rechunk_mutually_exclusive(self):
|
||||
result = runner.invoke(
|
||||
|
|
|
|||
|
|
@ -383,55 +383,6 @@ async def test_client_create_document_from_directory(temp_db_path):
|
|||
assert not any("unsupported.xyz" in uri for uri in uris)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_client_create_document_from_directory_with_filters(
|
||||
monkeypatch, temp_db_path
|
||||
):
|
||||
"""Test creating documents from a directory with ignore and include patterns."""
|
||||
# Mock config to have ignore and include patterns
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.client.Config.monitor.ignore_patterns", ["**/ignore_me/**", "*.log"]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.client.Config.monitor.include_patterns", ["**/include/**/*.txt"]
|
||||
)
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
test_dir = Path(temp_dir) / "test_docs"
|
||||
test_dir.mkdir()
|
||||
|
||||
# Create files in include directory - should be included
|
||||
include_dir = test_dir / "include"
|
||||
include_dir.mkdir()
|
||||
(include_dir / "doc1.txt").write_text("Content of doc1")
|
||||
(include_dir / "doc2.txt").write_text("Content of doc2")
|
||||
|
||||
# Create files outside include directory - should be excluded by include pattern
|
||||
(test_dir / "doc3.txt").write_text("Content of doc3")
|
||||
|
||||
# Create files in ignore directory - should be excluded by ignore pattern
|
||||
ignore_dir = test_dir / "ignore_me"
|
||||
ignore_dir.mkdir()
|
||||
(ignore_dir / "doc4.txt").write_text("Content of doc4")
|
||||
|
||||
# Create log file - should be excluded by ignore pattern
|
||||
(test_dir / "debug.log").write_text("log content")
|
||||
|
||||
result = await client.create_document_from_source(test_dir)
|
||||
|
||||
assert isinstance(result, list)
|
||||
# Should only include doc1.txt and doc2.txt from include directory
|
||||
assert len(result) == 2
|
||||
|
||||
uris = [doc.uri for doc in result if doc.uri]
|
||||
assert any("doc1.txt" in uri for uri in uris)
|
||||
assert any("doc2.txt" in uri for uri in uris)
|
||||
assert not any("doc3.txt" in uri for uri in uris)
|
||||
assert not any("doc4.txt" in uri for uri in uris)
|
||||
assert not any("debug.log" in uri for uri in uris)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_client_create_document_from_url(temp_db_path):
|
||||
"""Test creating a document from a URL."""
|
||||
|
|
|
|||
|
|
@ -1,476 +0,0 @@
|
|||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig, MonitorConfig
|
||||
from haiku.rag.monitor import FileWatcher
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_upsert_document():
|
||||
"""Test FileWatcher._upsert_document method."""
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir) / "test.txt"
|
||||
temp_path.write_text("Test content for file watcher")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
mock_doc = Document(id="1", content="Test content", uri=temp_path.as_uri())
|
||||
mock_client.create_document_from_source.return_value = mock_doc
|
||||
mock_client.get_document_by_uri.return_value = None # No existing document
|
||||
|
||||
test_config = AppConfig(monitor=MonitorConfig(directories=[temp_path.parent]))
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
result = await watcher._upsert_document(temp_path)
|
||||
|
||||
assert result is not None
|
||||
assert result.id == "1"
|
||||
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
||||
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_upsert_existing_document():
|
||||
"""Test FileWatcher._upsert_document with existing document."""
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir) / "test.txt"
|
||||
temp_path.write_text("Test content for file watcher")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
existing_doc = Document(id="1", content="Old content", uri=temp_path.as_uri())
|
||||
updated_doc = Document(
|
||||
id="1", content="Updated content", uri=temp_path.as_uri()
|
||||
)
|
||||
|
||||
mock_client.get_document_by_uri.return_value = existing_doc
|
||||
mock_client.create_document_from_source.return_value = updated_doc
|
||||
|
||||
test_config = AppConfig(monitor=MonitorConfig(directories=[temp_path.parent]))
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
result = await watcher._upsert_document(temp_path)
|
||||
|
||||
assert result is not None
|
||||
assert result.content == "Updated content"
|
||||
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
||||
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_delete_document():
|
||||
"""Test FileWatcher._delete_document method."""
|
||||
temp_path = Path("/tmp/test_file.txt")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
existing_doc = Document(id="1", content="Content to delete", uri=temp_path.as_uri())
|
||||
mock_client.get_document_by_uri.return_value = existing_doc
|
||||
mock_client.delete_document.return_value = True
|
||||
|
||||
test_config = AppConfig(monitor=MonitorConfig(directories=[temp_path.parent]))
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
await watcher._delete_document(temp_path)
|
||||
|
||||
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
||||
mock_client.delete_document.assert_called_once_with("1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_delete_nonexistent_document():
|
||||
"""Test FileWatcher._delete_document with non-existent document."""
|
||||
temp_path = Path("/tmp/nonexistent_file.txt")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
mock_client.get_document_by_uri.return_value = None
|
||||
|
||||
test_config = AppConfig(monitor=MonitorConfig(directories=[temp_path.parent]))
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
await watcher._delete_document(temp_path)
|
||||
|
||||
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
||||
mock_client.delete_document.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_filter_ignore_patterns():
|
||||
"""Test FileFilter with ignore patterns."""
|
||||
from watchfiles import Change
|
||||
|
||||
from haiku.rag.monitor import FileFilter
|
||||
|
||||
filter = FileFilter(ignore_patterns=["*draft*.md", "temp/", "**/archive/**"])
|
||||
|
||||
# Should ignore draft markdown files
|
||||
assert not filter(Change.added, "/path/to/draft-post.md")
|
||||
|
||||
# Should ignore files in temp/ directory
|
||||
assert not filter(Change.added, "/path/temp/notes.txt")
|
||||
|
||||
# Should ignore files in archive directories
|
||||
assert not filter(Change.added, "/path/to/archive/old.pdf")
|
||||
|
||||
# Should NOT ignore regular markdown files
|
||||
assert filter(Change.added, "/path/to/readme.md")
|
||||
|
||||
# Should NOT ignore files outside temp/
|
||||
assert filter(Change.added, "/path/to/notes.txt")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_filter_include_patterns():
|
||||
"""Test FileFilter with include patterns (whitelist mode)."""
|
||||
from watchfiles import Change
|
||||
|
||||
from haiku.rag.monitor import FileFilter
|
||||
|
||||
filter = FileFilter(include_patterns=["*.md", "**/docs/**"])
|
||||
|
||||
# Should include .md files
|
||||
assert filter(Change.added, "/path/to/file.md")
|
||||
|
||||
# Should include files in docs/ directory
|
||||
assert filter(Change.added, "/path/to/docs/guide.txt")
|
||||
|
||||
# Should NOT include .txt files outside docs/
|
||||
assert not filter(Change.added, "/path/to/file.txt")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_filter_combined_patterns():
|
||||
"""Test FileFilter with both include and ignore patterns."""
|
||||
from watchfiles import Change
|
||||
|
||||
from haiku.rag.monitor import FileFilter
|
||||
|
||||
# Include all markdown files, but ignore drafts
|
||||
filter = FileFilter(
|
||||
include_patterns=["*.md"], ignore_patterns=["*draft*.md", "archive/"]
|
||||
)
|
||||
|
||||
# Should include regular .md files
|
||||
assert filter(Change.added, "/path/to/readme.md")
|
||||
|
||||
# Should ignore draft .md files (ignore takes precedence after include)
|
||||
assert not filter(Change.added, "/path/to/draft-post.md")
|
||||
|
||||
# Should ignore .md files in archive/ directory
|
||||
assert not filter(Change.added, "/path/archive/old.md")
|
||||
|
||||
# Should NOT include .txt files (not in include patterns)
|
||||
assert not filter(Change.added, "/path/to/file.txt")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_filter_extension_check():
|
||||
"""Test that FileFilter still respects extension filtering."""
|
||||
from watchfiles import Change
|
||||
|
||||
from haiku.rag.monitor import FileFilter
|
||||
|
||||
filter = FileFilter()
|
||||
|
||||
# Should include files with supported extensions
|
||||
assert filter(Change.added, "/path/to/document.pdf")
|
||||
assert filter(Change.added, "/path/to/notes.md")
|
||||
|
||||
# Should not include files with unsupported extensions
|
||||
assert not filter(Change.added, "/path/to/file.xyz")
|
||||
assert not filter(Change.added, "/path/to/binary.bin")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_with_ignore_patterns():
|
||||
"""Test FileWatcher respects ignore patterns from config."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
draft_file = temp_path / "draft.md"
|
||||
readme_file = temp_path / "readme.md"
|
||||
|
||||
draft_file.write_text("Draft content")
|
||||
readme_file.write_text("Readme content")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
mock_doc = Document(id="1", content="Readme", uri=readme_file.as_uri())
|
||||
mock_client.create_document_from_source.return_value = mock_doc
|
||||
mock_client.get_document_by_uri.return_value = None
|
||||
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(directories=[temp_path], ignore_patterns=["draft*"])
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Run refresh which should only process readme.md, not draft.md
|
||||
await watcher.refresh()
|
||||
|
||||
# Should have only called for the readme file
|
||||
assert mock_client.create_document_from_source.call_count == 1
|
||||
mock_client.create_document_from_source.assert_called_with(str(readme_file))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_with_include_patterns():
|
||||
"""Test FileWatcher respects include patterns from config."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
md_file = temp_path / "readme.md"
|
||||
pdf_file = temp_path / "document.pdf"
|
||||
py_file = temp_path / "script.py"
|
||||
|
||||
md_file.write_text("Markdown content")
|
||||
pdf_file.write_text("PDF content")
|
||||
py_file.write_text("Python content")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
mock_doc = Document(id="1", content="Markdown", uri=md_file.as_uri())
|
||||
mock_client.create_document_from_source.return_value = mock_doc
|
||||
mock_client.get_document_by_uri.return_value = None
|
||||
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(directories=[temp_path], include_patterns=["*.md"])
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Run refresh which should only process .md file, not .pdf or .py
|
||||
await watcher.refresh()
|
||||
|
||||
# Should have only called for the .md file
|
||||
assert mock_client.create_document_from_source.call_count == 1
|
||||
mock_client.create_document_from_source.assert_called_with(str(md_file))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_skips_unchanged_document():
|
||||
"""Test FileWatcher returns existing document when content hasn't changed."""
|
||||
from datetime import datetime
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
test_file = temp_path / "test.txt"
|
||||
test_content = "Test content"
|
||||
test_file.write_text(test_content)
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
# Existing document with a timestamp
|
||||
now = datetime.now()
|
||||
existing_doc = Document(
|
||||
id="1",
|
||||
content=test_content,
|
||||
uri=test_file.as_uri(),
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
mock_client.get_document_by_uri.return_value = existing_doc
|
||||
# Client returns same document with same timestamp (unchanged)
|
||||
mock_client.create_document_from_source.return_value = existing_doc
|
||||
|
||||
test_config = AppConfig(monitor=MonitorConfig(directories=[temp_path]))
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
result = await watcher._upsert_document(test_file)
|
||||
|
||||
assert result is not None
|
||||
assert result.id == "1"
|
||||
# Verify timestamp hasn't changed (document wasn't updated)
|
||||
assert result.updated_at == now
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_deletes_orphans():
|
||||
"""Test FileWatcher deletes documents whose files no longer exist."""
|
||||
import asyncio
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
existing_file = temp_path / "exists.txt"
|
||||
existing_file.write_text("Existing file")
|
||||
|
||||
# Create a document for a file that doesn't exist
|
||||
orphan_uri = (temp_path / "deleted.txt").as_uri()
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
orphan_doc = Document(id="orphan-1", content="Orphaned content", uri=orphan_uri)
|
||||
existing_doc = Document(
|
||||
id="existing-1", content="Existing content", uri=existing_file.as_uri()
|
||||
)
|
||||
|
||||
# Mock list_documents to return both documents
|
||||
mock_client.list_documents.return_value = [orphan_doc, existing_doc]
|
||||
mock_client.get_document_by_uri.return_value = None
|
||||
mock_client.create_document_from_source.return_value = existing_doc
|
||||
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(directories=[temp_path], delete_orphans=True)
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Run refresh which should delete orphan and process existing file
|
||||
await watcher.refresh()
|
||||
|
||||
# Give background task time to complete
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Should have deleted the orphan document
|
||||
mock_client.delete_document.assert_called_once_with("orphan-1")
|
||||
# Should have processed the existing file
|
||||
mock_client.create_document_from_source.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_skips_orphan_deletion_when_disabled():
|
||||
"""Test FileWatcher does not delete orphans when delete_orphans is False."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create a document for a file that doesn't exist
|
||||
orphan_uri = (temp_path / "deleted.txt").as_uri()
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
orphan_doc = Document(id="orphan-1", content="Orphaned content", uri=orphan_uri)
|
||||
|
||||
# Mock list_documents to return orphan document
|
||||
mock_client.list_documents.return_value = [orphan_doc]
|
||||
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(directories=[temp_path], delete_orphans=False)
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Run refresh
|
||||
await watcher.refresh()
|
||||
|
||||
# Should NOT have deleted the orphan document
|
||||
mock_client.delete_document.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_orphan_deletion_respects_patterns():
|
||||
"""Test orphan deletion respects include/ignore patterns."""
|
||||
import asyncio
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create documents for files that don't exist
|
||||
ignored_orphan_uri = (temp_path / "draft.md").as_uri()
|
||||
excluded_orphan_uri = (temp_path / "file.pdf").as_uri()
|
||||
included_orphan_uri = (temp_path / "readme.md").as_uri()
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
|
||||
ignored_doc = Document(
|
||||
id="ignored-1", content="Ignored", uri=ignored_orphan_uri
|
||||
)
|
||||
excluded_doc = Document(
|
||||
id="excluded-1", content="Excluded", uri=excluded_orphan_uri
|
||||
)
|
||||
included_doc = Document(
|
||||
id="included-1", content="Included", uri=included_orphan_uri
|
||||
)
|
||||
|
||||
# Mock list_documents to return all orphan documents
|
||||
mock_client.list_documents.return_value = [
|
||||
ignored_doc,
|
||||
excluded_doc,
|
||||
included_doc,
|
||||
]
|
||||
|
||||
# Config with patterns: only .md files, but exclude draft*
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(
|
||||
directories=[temp_path],
|
||||
delete_orphans=True,
|
||||
include_patterns=["*.md"],
|
||||
ignore_patterns=["draft*"],
|
||||
)
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Run refresh
|
||||
await watcher.refresh()
|
||||
|
||||
# Give background task time to complete
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Should only delete the included orphan (readme.md)
|
||||
# - draft.md matches ignore pattern -> NOT deleted
|
||||
# - file.pdf doesn't match include pattern -> NOT deleted
|
||||
# - readme.md matches include and not ignored -> DELETED
|
||||
mock_client.delete_document.assert_called_once_with("included-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_orphan_handles_spaces_in_filenames():
|
||||
"""Test orphan deletion correctly handles files with spaces in names."""
|
||||
import asyncio
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
# Create a file with spaces that exists
|
||||
existing_file = temp_path / "my file with spaces.txt"
|
||||
existing_file.write_text("Existing file")
|
||||
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
# Document with URI that has URL-encoded spaces (%20)
|
||||
existing_doc = Document(
|
||||
id="existing-1", content="Existing", uri=existing_file.as_uri()
|
||||
)
|
||||
|
||||
# Mock list_documents to return document with encoded spaces
|
||||
mock_client.list_documents.return_value = [existing_doc]
|
||||
mock_client.get_document_by_uri.return_value = None
|
||||
mock_client.create_document_from_source.return_value = existing_doc
|
||||
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(directories=[temp_path], delete_orphans=True)
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Run refresh
|
||||
await watcher.refresh()
|
||||
|
||||
# Give background task time to complete
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Should NOT delete the document since file exists
|
||||
mock_client.delete_document.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_observe_raises_on_missing_paths():
|
||||
"""Test observe() raises FileNotFoundError when directories don't exist."""
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
|
||||
test_config = AppConfig(
|
||||
monitor=MonitorConfig(
|
||||
directories=[Path("/nonexistent/path/that/does/not/exist")]
|
||||
)
|
||||
)
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
with pytest.raises(FileNotFoundError) as exc_info:
|
||||
await watcher.observe()
|
||||
|
||||
assert "Monitor directories do not exist" in str(exc_info.value)
|
||||
assert "haiku.rag.yaml" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_watcher_observe_returns_early_when_no_directories():
|
||||
"""Test observe() returns early when no directories are configured."""
|
||||
mock_client = AsyncMock(spec=HaikuRAG)
|
||||
|
||||
test_config = AppConfig(monitor=MonitorConfig(directories=[]))
|
||||
watcher = FileWatcher(client=mock_client, config=test_config)
|
||||
|
||||
# Should return without error when no directories configured
|
||||
await watcher.observe()
|
||||
|
||||
# No documents should have been processed
|
||||
mock_client.create_document_from_source.assert_not_called()
|
||||
|
|
@ -3,7 +3,6 @@
|
|||
# Stop after:
|
||||
# docker compose -f tests/docker/docker-compose.s3.yml down -v
|
||||
|
||||
import importlib.util
|
||||
import socket
|
||||
from uuid import uuid4
|
||||
|
||||
|
|
@ -11,11 +10,9 @@ import pytest
|
|||
|
||||
from haiku.rag.app import HaikuRAGApp
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig, S3MonitorEntry
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
HAS_OBSTORE = importlib.util.find_spec("obstore") is not None
|
||||
|
||||
S3_ENDPOINT = "http://localhost:8333"
|
||||
S3_BUCKET = "test-bucket"
|
||||
S3_STORAGE_OPTIONS = {
|
||||
|
|
@ -152,151 +149,3 @@ async def test_app_info_empty_db(tmp_path, capsys):
|
|||
|
||||
out = capsys.readouterr().out
|
||||
assert "Database is empty" in out
|
||||
|
||||
|
||||
# ----------------------- S3 watcher integration tests ----------------------- #
|
||||
# These exercise S3Watcher against the live SeaweedFS instance. Documents are
|
||||
# uploaded as raw S3 objects under a unique per-test prefix; the watcher's
|
||||
# refresh() is invoked directly so tests stay deterministic. LanceDB stays
|
||||
# local — these tests verify the watcher path, not LanceDB-on-S3.
|
||||
|
||||
|
||||
_obstore_required = pytest.mark.skipif(
|
||||
not HAS_OBSTORE,
|
||||
reason="obstore not installed (uv sync --extra s3)",
|
||||
)
|
||||
|
||||
|
||||
def _watcher_store():
|
||||
from haiku.rag.s3 import make_s3_store
|
||||
|
||||
return make_s3_store(S3_BUCKET, S3_STORAGE_OPTIONS)
|
||||
|
||||
|
||||
async def _put_object(prefix: str, key: str, body: bytes) -> None:
|
||||
import obstore
|
||||
|
||||
await obstore.put_async(_watcher_store(), f"{prefix}/{key}", body)
|
||||
|
||||
|
||||
async def _delete_object(prefix: str, key: str) -> None:
|
||||
import obstore
|
||||
|
||||
await obstore.delete_async(_watcher_store(), f"{prefix}/{key}")
|
||||
|
||||
|
||||
def _watcher_entry(prefix: str, **overrides) -> S3MonitorEntry:
|
||||
return S3MonitorEntry(
|
||||
uri=overrides.pop("uri", f"s3://{S3_BUCKET}/{prefix}/"),
|
||||
storage_options=overrides.pop("storage_options", S3_STORAGE_OPTIONS),
|
||||
include_patterns=overrides.pop("include_patterns", ["*.txt"]),
|
||||
delete_orphans=overrides.pop("delete_orphans", False),
|
||||
poll_interval=overrides.pop("poll_interval", 60),
|
||||
**overrides,
|
||||
)
|
||||
|
||||
|
||||
@_obstore_required
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_initial_sweep(tmp_path):
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
prefix = f"watcher-init-{uuid4().hex[:8]}"
|
||||
await _put_object(prefix, "alpha.txt", b"alpha content")
|
||||
await _put_object(prefix, "beta.txt", b"beta content")
|
||||
|
||||
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_watcher_entry(prefix),
|
||||
supported_extensions=[".txt", ".md", ".pdf"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
docs = await rag.list_documents()
|
||||
uris = sorted(d.uri or "" for d in docs)
|
||||
|
||||
assert uris == [
|
||||
f"s3://{S3_BUCKET}/{prefix}/alpha.txt",
|
||||
f"s3://{S3_BUCKET}/{prefix}/beta.txt",
|
||||
]
|
||||
|
||||
|
||||
@_obstore_required
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_detects_new_object(tmp_path):
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
prefix = f"watcher-new-{uuid4().hex[:8]}"
|
||||
await _put_object(prefix, "first.txt", b"first content")
|
||||
|
||||
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_watcher_entry(prefix),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
assert await rag.count_documents() == 1
|
||||
|
||||
await _put_object(prefix, "second.txt", b"second content")
|
||||
await watcher.refresh()
|
||||
assert await rag.count_documents() == 2
|
||||
|
||||
|
||||
@_obstore_required
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_detects_modified_object(tmp_path):
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
prefix = f"watcher-mod-{uuid4().hex[:8]}"
|
||||
uri = f"s3://{S3_BUCKET}/{prefix}/file.txt"
|
||||
|
||||
await _put_object(prefix, "file.txt", b"original content")
|
||||
|
||||
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_watcher_entry(prefix),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
first = await rag.get_document_by_uri(uri)
|
||||
assert first is not None
|
||||
first_md5 = first.metadata["md5"]
|
||||
|
||||
await _put_object(prefix, "file.txt", b"new content body")
|
||||
await watcher.refresh()
|
||||
|
||||
second = await rag.get_document_by_uri(uri)
|
||||
assert second is not None
|
||||
assert second.id == first.id
|
||||
assert second.metadata["md5"] != first_md5
|
||||
assert "new content body" in second.content
|
||||
|
||||
|
||||
@_obstore_required
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_orphan_deletion(tmp_path):
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
prefix = f"watcher-orphan-{uuid4().hex[:8]}"
|
||||
await _put_object(prefix, "kept.txt", b"keep me")
|
||||
await _put_object(prefix, "doomed.txt", b"will be deleted")
|
||||
|
||||
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_watcher_entry(prefix, delete_orphans=True),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
assert await rag.count_documents() == 2
|
||||
|
||||
await _delete_object(prefix, "doomed.txt")
|
||||
await watcher.refresh()
|
||||
|
||||
docs = await rag.list_documents()
|
||||
assert len(docs) == 1
|
||||
assert docs[0].uri == f"s3://{S3_BUCKET}/{prefix}/kept.txt"
|
||||
|
|
|
|||
|
|
@ -1,403 +0,0 @@
|
|||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig, MonitorConfig, S3MonitorEntry
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def s3_listing(monkeypatch):
|
||||
"""Patch `obstore.list_obs` with an async-iterator returning controllable batches.
|
||||
|
||||
Returns `(set_batches, list_mock)`. `set_batches([[meta, ...], ...])`
|
||||
seeds the next call's pages.
|
||||
"""
|
||||
import obstore
|
||||
|
||||
batches: list[list[MagicMock]] = []
|
||||
|
||||
def list_obs(_store, *_, **__):
|
||||
async def _iter():
|
||||
for batch in batches:
|
||||
yield batch
|
||||
|
||||
return _iter()
|
||||
|
||||
list_mock = MagicMock(side_effect=list_obs)
|
||||
monkeypatch.setattr(obstore, "list", list_mock)
|
||||
|
||||
def set_batches(new_batches):
|
||||
batches.clear()
|
||||
batches.extend(new_batches)
|
||||
|
||||
return set_batches, list_mock
|
||||
|
||||
|
||||
def _meta(path: str, etag: str) -> dict:
|
||||
# Real obstore ObjectMeta is a TypedDict; raw S3 ETags include quotes.
|
||||
return {
|
||||
"path": path,
|
||||
"e_tag": f'"{etag}"',
|
||||
"size": 0,
|
||||
"last_modified": None,
|
||||
}
|
||||
|
||||
|
||||
def _entry(**kwargs) -> S3MonitorEntry:
|
||||
return S3MonitorEntry(
|
||||
uri=kwargs.pop("uri", "s3://my-bucket/incoming/"),
|
||||
poll_interval=kwargs.pop("poll_interval", 60),
|
||||
delete_orphans=kwargs.pop("delete_orphans", False),
|
||||
ignore_patterns=kwargs.pop("ignore_patterns", []),
|
||||
include_patterns=kwargs.pop("include_patterns", []),
|
||||
storage_options=kwargs.pop("storage_options", {}),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _doc(uri: str, etag: str, doc_id: str | None = None) -> Document:
|
||||
return Document(
|
||||
id=doc_id or uri,
|
||||
content="...",
|
||||
uri=uri,
|
||||
metadata={"etag": etag, "md5": "deadbeef"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_refresh_upserts_new_objects(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[_meta("incoming/a.txt", "abc"), _meta("incoming/b.txt", "def")]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = []
|
||||
rag.create_document_from_source.return_value = Document(
|
||||
id="x", content="...", uri="s3://my-bucket/incoming/a.txt"
|
||||
)
|
||||
|
||||
watcher = S3Watcher(
|
||||
client=rag, entry=_entry(), supported_extensions=[".txt", ".md", ".pdf"]
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
assert rag.create_document_from_source.await_count == 2
|
||||
called_uris = {c.args[0] for c in rag.create_document_from_source.await_args_list}
|
||||
assert called_uris == {
|
||||
"s3://my-bucket/incoming/a.txt",
|
||||
"s3://my-bucket/incoming/b.txt",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_skips_unchanged_etag(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[_meta("incoming/a.txt", "abc")]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = [_doc("s3://my-bucket/incoming/a.txt", "abc")]
|
||||
|
||||
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
||||
await watcher.refresh()
|
||||
|
||||
rag.create_document_from_source.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_upserts_when_etag_differs(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[_meta("incoming/a.txt", "new")]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = [_doc("s3://my-bucket/incoming/a.txt", "old")]
|
||||
rag.create_document_from_source.return_value = Document(
|
||||
id="x", content="...", uri="s3://my-bucket/incoming/a.txt"
|
||||
)
|
||||
|
||||
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
||||
await watcher.refresh()
|
||||
|
||||
rag.create_document_from_source.assert_awaited_once_with(
|
||||
"s3://my-bucket/incoming/a.txt", storage_options={}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_strips_etag_quotes(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[_meta("incoming/a.txt", "abc")]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = [
|
||||
_doc("s3://my-bucket/incoming/a.txt", "abc") # already stripped in storage
|
||||
]
|
||||
|
||||
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
||||
await watcher.refresh()
|
||||
|
||||
rag.create_document_from_source.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_deletes_orphans_when_enabled(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[_meta("incoming/a.txt", "abc")]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
a_doc = _doc("s3://my-bucket/incoming/a.txt", "abc", doc_id="a-id")
|
||||
orphan = _doc("s3://my-bucket/incoming/old.txt", "stale", doc_id="orphan-id")
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = [a_doc, orphan]
|
||||
rag.get_document_by_uri.return_value = orphan
|
||||
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_entry(delete_orphans=True),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
rag.delete_document.assert_awaited_once_with("orphan-id")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_does_not_delete_orphans_when_disabled(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
orphan = _doc("s3://my-bucket/incoming/old.txt", "stale", doc_id="orphan-id")
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = [orphan]
|
||||
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_entry(delete_orphans=False),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
rag.delete_document.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_orphan_scope_is_per_entry(s3_listing):
|
||||
"""A doc under a different bucket prefix must not be touched."""
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = [] # filter scopes to my-bucket
|
||||
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_entry(delete_orphans=True),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
rag.list_documents.assert_awaited_once()
|
||||
filter_kwarg = rag.list_documents.await_args.kwargs["filter"]
|
||||
assert filter_kwarg == "uri LIKE 's3://my-bucket/incoming/%'"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_applies_include_and_ignore_patterns(s3_listing):
|
||||
set_batches, _ = s3_listing
|
||||
set_batches(
|
||||
[
|
||||
[
|
||||
_meta("incoming/keep.md", "1"),
|
||||
_meta("incoming/draft.md", "2"),
|
||||
_meta("incoming/skip.txt", "3"),
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = []
|
||||
rag.create_document_from_source.return_value = Document(
|
||||
id="x", content="...", uri="s3://my-bucket/incoming/keep.md"
|
||||
)
|
||||
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_entry(include_patterns=["*.md"], ignore_patterns=["draft*"]),
|
||||
supported_extensions=[".md", ".txt"],
|
||||
)
|
||||
await watcher.refresh()
|
||||
|
||||
assert rag.create_document_from_source.await_count == 1
|
||||
assert (
|
||||
rag.create_document_from_source.await_args.args[0]
|
||||
== "s3://my-bucket/incoming/keep.md"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_observe_survives_transient_list_failure(s3_listing):
|
||||
"""First refresh succeeds; second refresh raises; loop survives and recovers."""
|
||||
set_batches, list_mock = s3_listing
|
||||
|
||||
pages_initial = [[_meta("incoming/a.txt", "abc")]]
|
||||
pages_after = [[_meta("incoming/a.txt", "abc")]]
|
||||
|
||||
paginate_calls = {"n": 0}
|
||||
|
||||
def list_obs_side_effect(_store, *_, **__):
|
||||
paginate_calls["n"] += 1
|
||||
if paginate_calls["n"] == 2:
|
||||
raise RuntimeError("transient list failure")
|
||||
|
||||
async def _iter():
|
||||
for batch in pages_after if paginate_calls["n"] > 1 else pages_initial:
|
||||
yield batch
|
||||
|
||||
return _iter()
|
||||
|
||||
list_mock.side_effect = list_obs_side_effect
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = []
|
||||
rag.create_document_from_source.return_value = Document(
|
||||
id="x", content="...", uri="s3://my-bucket/incoming/a.txt"
|
||||
)
|
||||
|
||||
watcher = S3Watcher(
|
||||
client=rag,
|
||||
entry=_entry(poll_interval=0),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
task = asyncio.create_task(watcher.observe())
|
||||
|
||||
for _ in range(20):
|
||||
await asyncio.sleep(0)
|
||||
if paginate_calls["n"] >= 3:
|
||||
break
|
||||
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert paginate_calls["n"] >= 3 # loop kept going past the failure
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_invalid_uri_rejected():
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
with pytest.raises(ValueError, match="Invalid S3 monitor URI"):
|
||||
S3Watcher(
|
||||
client=rag,
|
||||
entry=S3MonitorEntry(uri="s3://"),
|
||||
supported_extensions=[".txt"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_watcher_upsert_failure_does_not_abort_sweep(s3_listing):
|
||||
"""A failing upsert doesn't propagate; the refresh keeps processing siblings."""
|
||||
set_batches, _ = s3_listing
|
||||
set_batches([[_meta("incoming/bad.txt", "abc"), _meta("incoming/good.txt", "def")]])
|
||||
|
||||
from haiku.rag.monitor import S3Watcher
|
||||
|
||||
rag = AsyncMock(spec=HaikuRAG)
|
||||
rag.list_documents.return_value = []
|
||||
|
||||
good_doc = Document(
|
||||
id="good-id", content="...", uri="s3://my-bucket/incoming/good.txt"
|
||||
)
|
||||
|
||||
async def maybe_fail(uri, **_):
|
||||
if uri.endswith("bad.txt"):
|
||||
raise RuntimeError("boom")
|
||||
return good_doc
|
||||
|
||||
rag.create_document_from_source.side_effect = maybe_fail
|
||||
|
||||
watcher = S3Watcher(client=rag, entry=_entry(), supported_extensions=[".txt"])
|
||||
|
||||
# The failing upsert must not propagate out of refresh().
|
||||
await watcher.refresh()
|
||||
|
||||
# Both objects were attempted — the first failure didn't abort the sibling.
|
||||
assert rag.create_document_from_source.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_serve_starts_one_s3_task_per_entry(monkeypatch, s3_listing):
|
||||
"""`serve` wires one S3Watcher task per MonitorConfig.s3 entry."""
|
||||
from haiku.rag import app as app_module
|
||||
|
||||
original_create_task = asyncio.create_task
|
||||
|
||||
def tracking_create_task(coro, *args, **kwargs):
|
||||
return original_create_task(coro, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(app_module.asyncio, "create_task", tracking_create_task)
|
||||
|
||||
config = AppConfig(
|
||||
monitor=MonitorConfig(
|
||||
s3=[
|
||||
S3MonitorEntry(uri="s3://bucket-a/x/"),
|
||||
S3MonitorEntry(uri="s3://bucket-b/y/"),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
fw_observe_calls = {"n": 0}
|
||||
|
||||
async def fake_fw_observe(self):
|
||||
fw_observe_calls["n"] += 1
|
||||
|
||||
monkeypatch.setattr(app_module.FileWatcher, "observe", fake_fw_observe)
|
||||
|
||||
s3_observe_calls = {"n": 0}
|
||||
|
||||
async def fake_s3_observe(self):
|
||||
s3_observe_calls["n"] += 1
|
||||
|
||||
monkeypatch.setattr(app_module.S3Watcher, "observe", fake_s3_observe)
|
||||
|
||||
class _Conv:
|
||||
supported_extensions = [".txt"]
|
||||
|
||||
monkeypatch.setattr("haiku.rag.converters.get_converter", lambda cfg: _Conv())
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = Path(tmp) / "db.lancedb"
|
||||
app = app_module.HaikuRAGApp(db_path=db_path, config=config)
|
||||
async with HaikuRAG(db_path, config=config, create=True):
|
||||
pass # create the database
|
||||
|
||||
await app.serve(enable_monitor=True, enable_mcp=False)
|
||||
|
||||
assert fw_observe_calls["n"] == 1
|
||||
assert s3_observe_calls["n"] == 2
|
||||
|
|
@ -44,16 +44,6 @@ async def test_settings_save_and_retrieve(temp_db_path):
|
|||
Config.processing.chunk_size = original_chunk_size
|
||||
|
||||
|
||||
def test_monitor_filter_patterns_config():
|
||||
"""Test that monitor filter patterns are available in config."""
|
||||
assert hasattr(Config.monitor, "ignore_patterns")
|
||||
assert hasattr(Config.monitor, "include_patterns")
|
||||
assert hasattr(Config.monitor, "directories")
|
||||
assert isinstance(Config.monitor.ignore_patterns, list)
|
||||
assert isinstance(Config.monitor.include_patterns, list)
|
||||
assert isinstance(Config.monitor.directories, list)
|
||||
|
||||
|
||||
class TestValidateConfigCompatibility:
|
||||
"""Tests for validate_config_compatibility method."""
|
||||
|
||||
|
|
|
|||
|
|
@ -35,8 +35,8 @@ nav = [
|
|||
{ Tuning = "tuning.md" },
|
||||
] },
|
||||
{ Production = [
|
||||
{ Server = "server.md" },
|
||||
{ MCP = "mcp.md" },
|
||||
{ Ingester = "ingester.md" },
|
||||
{ "Remote processing" = "remote-processing.md" },
|
||||
] },
|
||||
{ Develop = [
|
||||
|
|
|
|||
Loading…
Reference in a new issue