drop the old monitor, rename serve→mcp, add e2e tests
This commit is contained in:
parent
1ca3c25a83
commit
7ea61a7b10
27 changed files with 539 additions and 1807 deletions
|
|
@ -1,8 +1,17 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
### 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.
|
- 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
|
## [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.
|
**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
|
```bash
|
||||||
# MCP server only (HTTP transport)
|
# HTTP transport on port 8001
|
||||||
haiku-rag serve --mcp
|
haiku-rag mcp
|
||||||
|
|
||||||
# MCP server (stdio transport)
|
# stdio transport (for Claude Desktop)
|
||||||
haiku-rag serve --mcp --stdio
|
haiku-rag mcp --stdio
|
||||||
|
|
||||||
# File monitoring only
|
# Custom port
|
||||||
haiku-rag serve --monitor
|
haiku-rag mcp --port 9000
|
||||||
|
|
||||||
# Both services
|
# Read-only mode (no write tools)
|
||||||
haiku-rag serve --monitor --mcp
|
haiku-rag --read-only 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
|
|
||||||
```
|
```
|
||||||
|
|
||||||
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
|
## Settings
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -59,12 +59,14 @@ storage:
|
||||||
data_dir: "" # Empty = use default platform location
|
data_dir: "" # Empty = use default platform location
|
||||||
vacuum_retention_seconds: 86400
|
vacuum_retention_seconds: 86400
|
||||||
|
|
||||||
monitor:
|
ingester:
|
||||||
directories:
|
sources:
|
||||||
- /path/to/documents
|
- type: fs
|
||||||
- /another/path
|
id: local-docs
|
||||||
ignore_patterns: [] # Gitignore-style patterns to exclude
|
root: /path/to/documents
|
||||||
include_patterns: [] # Gitignore-style patterns to include
|
ignore_patterns: [] # Gitignore-style patterns to exclude
|
||||||
|
include_patterns: [] # Gitignore-style patterns to include
|
||||||
|
delete_orphans: true
|
||||||
|
|
||||||
lancedb:
|
lancedb:
|
||||||
uri: "" # Empty for local, or db://, s3://, az://, gs://
|
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).
|
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:
|
For automatic ingestion of local directories, S3 buckets, or HTTP
|
||||||
|
sources (with filtering, retries, and a dead-letter queue), see the
|
||||||
```yaml
|
[Ingester](../ingester.md) page.
|
||||||
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://…`).
|
|
||||||
|
|
|
||||||
|
|
@ -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":
|
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.
|
- **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 serve --read-only --mcp`, the chat TUI, etc. They never see the documents bucket.
|
- **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.
|
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
|
```bash
|
||||||
# Default streamable HTTP transport on port 8001
|
# Default streamable HTTP transport on port 8001
|
||||||
haiku-rag serve --mcp
|
haiku-rag mcp
|
||||||
|
|
||||||
# Custom port
|
# Custom port
|
||||||
haiku-rag serve --mcp --mcp-port 9000
|
haiku-rag mcp --port 9000
|
||||||
|
|
||||||
# stdio transport (for Claude Desktop)
|
# stdio transport (for Claude Desktop)
|
||||||
haiku-rag serve --mcp --stdio
|
haiku-rag mcp --stdio
|
||||||
|
|
||||||
# Read-only mode (excludes write tools)
|
# 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.
|
**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": {
|
"mcpServers": {
|
||||||
"haiku-rag": {
|
"haiku-rag": {
|
||||||
"command": "haiku-rag",
|
"command": "haiku-rag",
|
||||||
"args": ["serve", "--mcp", "--stdio"]
|
"args": ["mcp", "--stdio"]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -44,7 +44,7 @@ With a custom database path:
|
||||||
"mcpServers": {
|
"mcpServers": {
|
||||||
"haiku-rag": {
|
"haiku-rag": {
|
||||||
"command": "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)
|
- `document` (optional): Document title/ID to pre-load (can repeat)
|
||||||
- Best for aggregation, computation, and multi-document analysis
|
- Best for aggregation, computation, and multi-document analysis
|
||||||
|
|
||||||
## Running with Other Services
|
## Continuous ingestion
|
||||||
|
|
||||||
Combine MCP with file monitoring:
|
For continuous document ingestion (filesystem watch, S3 polling, HTTP
|
||||||
|
sources, a job queue with retries), run [`haiku-ingester`](ingester.md)
|
||||||
```bash
|
as a separate process against the same LanceDB.
|
||||||
# MCP + file monitoring
|
|
||||||
haiku-rag serve --mcp --monitor
|
|
||||||
```
|
|
||||||
|
|
||||||
See [Server Mode](server.md) for details on file monitoring.
|
|
||||||
|
|
|
||||||
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:
|
To use a skill from Claude Desktop or another MCP-aware client, run the MCP server:
|
||||||
|
|
||||||
```bash
|
```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).
|
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:
|
To call the skill from Claude Desktop (or any MCP client), run the MCP server:
|
||||||
|
|
||||||
```bash
|
```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.
|
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
|
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
|
## Usage
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,9 +6,12 @@ environment: production
|
||||||
storage:
|
storage:
|
||||||
data_dir: /data
|
data_dir: /data
|
||||||
|
|
||||||
monitor:
|
ingester:
|
||||||
directories:
|
sources:
|
||||||
- /docs
|
- type: fs
|
||||||
|
id: docs
|
||||||
|
root: /docs
|
||||||
|
delete_orphans: true
|
||||||
|
|
||||||
# Remote document processing with docling-serve
|
# Remote document processing with docling-serve
|
||||||
processing:
|
processing:
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,3 @@
|
||||||
import asyncio
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
@ -20,7 +19,6 @@ from rich.progress import (
|
||||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, Config
|
||||||
from haiku.rag.mcp import create_mcp_server
|
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.chunk import SearchType
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
|
|
@ -727,77 +725,26 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
self.console.print(content)
|
self.console.print(content)
|
||||||
self.console.rule()
|
self.console.rule()
|
||||||
|
|
||||||
async def serve(
|
async def run_mcp(
|
||||||
self,
|
self,
|
||||||
enable_monitor: bool = True,
|
transport: str | None = None,
|
||||||
enable_mcp: bool = True,
|
port: int = 8001,
|
||||||
mcp_transport: str | None = None,
|
|
||||||
mcp_port: int = 8001,
|
|
||||||
):
|
):
|
||||||
"""Start the server with selected services."""
|
"""Run the MCP server until interrupted."""
|
||||||
async with HaikuRAG(
|
async with HaikuRAG(
|
||||||
self.db_path,
|
self.db_path,
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
) as client:
|
):
|
||||||
tasks = []
|
server = create_mcp_server(
|
||||||
|
self.db_path, config=self.config, read_only=self.read_only
|
||||||
# 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
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Wait for any task to complete (or KeyboardInterrupt)
|
if transport == "stdio":
|
||||||
await asyncio.gather(*tasks)
|
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:
|
except KeyboardInterrupt:
|
||||||
pass
|
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(
|
@_cli.command(
|
||||||
"serve",
|
"mcp",
|
||||||
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
|
help="Run the MCP server. For continuous ingestion, use haiku-ingester serve.",
|
||||||
)
|
)
|
||||||
def serve(
|
def mcp(
|
||||||
db: Path | None = typer.Option(
|
db: Path | None = typer.Option(
|
||||||
None,
|
None,
|
||||||
"--db",
|
"--db",
|
||||||
help="Path to the LanceDB database file",
|
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(
|
stdio: bool = typer.Option(
|
||||||
False,
|
False,
|
||||||
"--stdio",
|
"--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,
|
8001,
|
||||||
"--mcp-port",
|
"--port",
|
||||||
help="Port to bind MCP server to (ignored with --stdio)",
|
help="Port to bind MCP server to (ignored with --stdio)",
|
||||||
),
|
),
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Start the server with selected services."""
|
"""Run the MCP server."""
|
||||||
# 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)
|
|
||||||
|
|
||||||
app = create_app(db) # pragma: no cover
|
app = create_app(db) # pragma: no cover
|
||||||
|
|
||||||
transport = "stdio" if stdio else None # pragma: no cover
|
transport = "stdio" if stdio else None # pragma: no cover
|
||||||
|
|
||||||
asyncio.run( # pragma: no cover
|
asyncio.run( # pragma: no cover
|
||||||
app.serve(
|
app.run_mcp(transport=transport, port=port)
|
||||||
enable_monitor=monitor,
|
|
||||||
enable_mcp=mcp,
|
|
||||||
mcp_transport=transport,
|
|
||||||
mcp_port=mcp_port,
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -328,11 +328,11 @@ async def create_document_from_source(
|
||||||
)
|
)
|
||||||
from haiku.rag.ingester.sources.filter import FileFilter
|
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] = []
|
documents: list[Document] = []
|
||||||
filter = FileFilter(
|
filter = FileFilter()
|
||||||
ignore_patterns=client._config.monitor.ignore_patterns or None,
|
|
||||||
include_patterns=client._config.monitor.include_patterns or None,
|
|
||||||
)
|
|
||||||
for child in local_path.rglob("*"):
|
for child in local_path.rglob("*"):
|
||||||
if child.is_file() and filter.include_file(str(child)):
|
if child.is_file() and filter.include_file(str(child)):
|
||||||
doc = await create_document_from_source(
|
doc = await create_document_from_source(
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,6 @@ from haiku.rag.config.models import (
|
||||||
IngesterConfig,
|
IngesterConfig,
|
||||||
LanceDBConfig,
|
LanceDBConfig,
|
||||||
ModelConfig,
|
ModelConfig,
|
||||||
MonitorConfig,
|
|
||||||
OllamaConfig,
|
OllamaConfig,
|
||||||
ProcessingConfig,
|
ProcessingConfig,
|
||||||
PromptsConfig,
|
PromptsConfig,
|
||||||
|
|
@ -24,7 +23,6 @@ from haiku.rag.config.models import (
|
||||||
QueueConfig,
|
QueueConfig,
|
||||||
RerankingConfig,
|
RerankingConfig,
|
||||||
RetryPolicyConfig,
|
RetryPolicyConfig,
|
||||||
S3MonitorEntry,
|
|
||||||
S3SourceConfig,
|
S3SourceConfig,
|
||||||
SourceConfig,
|
SourceConfig,
|
||||||
StorageConfig,
|
StorageConfig,
|
||||||
|
|
@ -44,7 +42,6 @@ __all__ = [
|
||||||
"IngesterConfig",
|
"IngesterConfig",
|
||||||
"LanceDBConfig",
|
"LanceDBConfig",
|
||||||
"ModelConfig",
|
"ModelConfig",
|
||||||
"MonitorConfig",
|
|
||||||
"OllamaConfig",
|
"OllamaConfig",
|
||||||
"ProcessingConfig",
|
"ProcessingConfig",
|
||||||
"PromptsConfig",
|
"PromptsConfig",
|
||||||
|
|
@ -53,7 +50,6 @@ __all__ = [
|
||||||
"QueueConfig",
|
"QueueConfig",
|
||||||
"RerankingConfig",
|
"RerankingConfig",
|
||||||
"RetryPolicyConfig",
|
"RetryPolicyConfig",
|
||||||
"S3MonitorEntry",
|
|
||||||
"S3SourceConfig",
|
"S3SourceConfig",
|
||||||
"SourceConfig",
|
"SourceConfig",
|
||||||
"StorageConfig",
|
"StorageConfig",
|
||||||
|
|
|
||||||
|
|
@ -57,23 +57,6 @@ class StorageConfig(BaseModel):
|
||||||
vacuum_retention_seconds: int = 86400
|
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):
|
class LanceDBConfig(BaseModel):
|
||||||
uri: str = ""
|
uri: str = ""
|
||||||
api_key: str = ""
|
api_key: str = ""
|
||||||
|
|
@ -358,7 +341,6 @@ class IngesterConfig(BaseModel):
|
||||||
class AppConfig(BaseModel):
|
class AppConfig(BaseModel):
|
||||||
environment: str = "production"
|
environment: str = "production"
|
||||||
storage: StorageConfig = Field(default_factory=StorageConfig)
|
storage: StorageConfig = Field(default_factory=StorageConfig)
|
||||||
monitor: MonitorConfig = Field(default_factory=MonitorConfig)
|
|
||||||
lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig)
|
lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig)
|
||||||
embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
|
embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
|
||||||
reranking: RerankingConfig = Field(default_factory=RerankingConfig)
|
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)}
|
uris = {e.uri async for e in src.discover(since=None)}
|
||||||
assert (fs_root / "b.txt").as_uri() not in uris
|
assert (fs_root / "b.txt").as_uri() not in uris
|
||||||
assert (fs_root / "a.md").as_uri() 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"}
|
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:
|
class TestRebuildValidation:
|
||||||
def test_embed_only_and_rechunk_mutually_exclusive(self):
|
def test_embed_only_and_rechunk_mutually_exclusive(self):
|
||||||
result = runner.invoke(
|
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)
|
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()
|
@pytest.mark.vcr()
|
||||||
async def test_client_create_document_from_url(temp_db_path):
|
async def test_client_create_document_from_url(temp_db_path):
|
||||||
"""Test creating a document from a URL."""
|
"""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:
|
# Stop after:
|
||||||
# docker compose -f tests/docker/docker-compose.s3.yml down -v
|
# docker compose -f tests/docker/docker-compose.s3.yml down -v
|
||||||
|
|
||||||
import importlib.util
|
|
||||||
import socket
|
import socket
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
|
|
@ -11,11 +10,9 @@ import pytest
|
||||||
|
|
||||||
from haiku.rag.app import HaikuRAGApp
|
from haiku.rag.app import HaikuRAGApp
|
||||||
from haiku.rag.client import HaikuRAG
|
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
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
HAS_OBSTORE = importlib.util.find_spec("obstore") is not None
|
|
||||||
|
|
||||||
S3_ENDPOINT = "http://localhost:8333"
|
S3_ENDPOINT = "http://localhost:8333"
|
||||||
S3_BUCKET = "test-bucket"
|
S3_BUCKET = "test-bucket"
|
||||||
S3_STORAGE_OPTIONS = {
|
S3_STORAGE_OPTIONS = {
|
||||||
|
|
@ -152,151 +149,3 @@ async def test_app_info_empty_db(tmp_path, capsys):
|
||||||
|
|
||||||
out = capsys.readouterr().out
|
out = capsys.readouterr().out
|
||||||
assert "Database is empty" in 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
|
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:
|
class TestValidateConfigCompatibility:
|
||||||
"""Tests for validate_config_compatibility method."""
|
"""Tests for validate_config_compatibility method."""
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -35,8 +35,8 @@ nav = [
|
||||||
{ Tuning = "tuning.md" },
|
{ Tuning = "tuning.md" },
|
||||||
] },
|
] },
|
||||||
{ Production = [
|
{ Production = [
|
||||||
{ Server = "server.md" },
|
|
||||||
{ MCP = "mcp.md" },
|
{ MCP = "mcp.md" },
|
||||||
|
{ Ingester = "ingester.md" },
|
||||||
{ "Remote processing" = "remote-processing.md" },
|
{ "Remote processing" = "remote-processing.md" },
|
||||||
] },
|
] },
|
||||||
{ Develop = [
|
{ Develop = [
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue