Compare commits
6 Commits
133cca2551
...
191ac2603f
| Author | SHA1 | Date | |
|---|---|---|---|
| 191ac2603f | |||
| 1dcd79ef96 | |||
| cb7a6c5cb5 | |||
| 36e310e646 | |||
| bb16cd6927 | |||
| 8072fb3dd0 |
@@ -6,3 +6,6 @@ venv/
|
||||
readme.md-
|
||||
*.bak
|
||||
hardware_state.json
|
||||
.env
|
||||
secrets/
|
||||
searxng/
|
||||
|
||||
@@ -14,7 +14,7 @@ from fastapi.responses import HTMLResponse, JSONResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR
|
||||
from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR, RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER, RAG_EVICTION_BATCH
|
||||
from db import init_db
|
||||
from hardware import assess_hardware
|
||||
from memory import get_memory_count
|
||||
@@ -37,6 +37,7 @@ import routers.completions as completions
|
||||
import routers.upload as upload
|
||||
import routers.ingest as ingest
|
||||
import routers.hardware as hardware
|
||||
import routers.rag_admin as rag_admin
|
||||
|
||||
# --- Logging ---
|
||||
log = logging.getLogger("jarvischat")
|
||||
@@ -56,6 +57,18 @@ async def lifespan(app: FastAPI):
|
||||
init_db()
|
||||
log.info(f"Memory system: {get_memory_count()} memories loaded")
|
||||
await assess_hardware()
|
||||
|
||||
if RAG_MAX_VECTORS > 0:
|
||||
if RAG_EVICTION_HIGH_WATER <= RAG_EVICTION_LOW_WATER:
|
||||
log.warning(
|
||||
f"RAG_EVICTION_HIGH_WATER={RAG_EVICTION_HIGH_WATER} <= "
|
||||
f"RAG_EVICTION_LOW_WATER={RAG_EVICTION_LOW_WATER} — eviction will never fire"
|
||||
)
|
||||
if RAG_EVICTION_BATCH <= 0:
|
||||
log.warning(f"RAG_EVICTION_BATCH={RAG_EVICTION_BATCH} clamped to 1")
|
||||
else:
|
||||
log.warning("RAG_MAX_VECTORS <= 0 — RAG eviction disabled")
|
||||
|
||||
yield
|
||||
log.info("JarvisChat shutting down")
|
||||
|
||||
@@ -146,6 +159,7 @@ for router_module in [
|
||||
auth_router, conversations.router, memories.router, models.router,
|
||||
presets.router, profile.router, settings.router, skills.router,
|
||||
chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router,
|
||||
rag_admin.router,
|
||||
]:
|
||||
app.include_router(router_module)
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ import logging
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
|
||||
VERSION = "v0.11.0"
|
||||
VERSION = "v0.13.0"
|
||||
OLLAMA_BASE = os.environ.get("OLLAMA_BASE", "http://localhost:11434")
|
||||
LLAMA_SERVER_BASE = os.environ.get("LLAMA_SERVER_BASE", "http://192.168.50.108:8081")
|
||||
SEARXNG_BASE = "http://localhost:8888"
|
||||
@@ -50,9 +50,21 @@ BODY_LIMIT_PROFILE_BYTES = 256 * 1024
|
||||
UPLOAD_DIR = "/tmp/jarvischat_uploads"
|
||||
MAX_UPLOAD_BYTES = 20 * 1024 * 1024
|
||||
SUPPORTED_UPLOAD_TYPES = {"text/plain", "text/markdown", "application/pdf", "application/json", "text/x-python", "text/html", "image/png", "image/jpeg", "image/gif", "image/svg+xml", "image/webp"}
|
||||
QDRANT_URL = "http://192.168.50.108:6333"
|
||||
RAG_COLLECTION = "jarvis_rag"
|
||||
UPLOAD_CONTEXT_EXPIRY_HOURS = 1
|
||||
BODY_LIMIT_UPLOAD_BYTES = MAX_UPLOAD_BYTES
|
||||
|
||||
# --- RAG eviction ---
|
||||
RAG_MAX_VECTORS = 50000
|
||||
RAG_EVICTION_HIGH_WATER = 0.80
|
||||
RAG_EVICTION_LOW_WATER = 0.20
|
||||
RAG_EVICTION_BATCH = 1000
|
||||
RAG_PINNED_SOURCES = ["upload", "profile"]
|
||||
RAG_GRACE_HOURS = 1
|
||||
RAG_ACCESS_WEIGHT = 1.0
|
||||
RAG_AGE_WEIGHT = 0.1
|
||||
|
||||
MAX_CHAT_MESSAGE_CHARS = 8000
|
||||
MAX_SEARCH_QUERY_CHARS = 500
|
||||
MAX_PROFILE_CHARS = 32000
|
||||
|
||||
@@ -0,0 +1,681 @@
|
||||
# Docker Distribution — Architecture & Planning
|
||||
|
||||
> **Part of B3 (v1.0 gate).** This document catalogs every service, volume, port, configuration, and decision needed to ship jarvisChat as a `docker compose` stack. It also defines extraction (setup) and back-out (uninstall) procedures so nothing is lost when reality disagrees with the plan.
|
||||
|
||||
## 1. Stack Overview
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────┐
|
||||
│ docker compose stack │
|
||||
│ │
|
||||
│ ┌────────────┐ ┌──────────┐ ┌────────────────────┐ │
|
||||
│ │ SearXNG │ │ Qdrant │ │ RabbitMQ │ │
|
||||
│ │ :8888 │ │ :6333 │ │ :5672 / :15672 │ │
|
||||
│ └──────┬──────┘ └────┬─────┘ └────────┬───────────┘ │
|
||||
│ │ │ │ │
|
||||
│ ▼ ▼ ▼ │
|
||||
│ ┌──────────────────────────────────────────────────┐ │
|
||||
│ │ jarvisChat (FastAPI) │ │
|
||||
│ │ :8080 (HTTP) │ │
|
||||
│ │ │ │
|
||||
│ │ SQLite ◄── jarvischat.db (volume) │ │
|
||||
│ │ Uploads ◄── /app/uploads (volume) │ │
|
||||
│ └──────────┬──────────────┬───────────────────────┘ │
|
||||
│ │ │ │
|
||||
│ ▼ ▼ │
|
||||
│ ┌──────────────┐ ┌──────────────┐ │
|
||||
│ │ llama-server │ │ Ollama │ │
|
||||
│ │ :8081 │ │ :11434 │ │
|
||||
│ │ (GPU/RPC) │ │ (embeddings) │ │
|
||||
│ └──────────────┘ └──────────────┘ │
|
||||
└─────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Service roles
|
||||
|
||||
| Service | Image | Role |
|
||||
|---------|-------|------|
|
||||
| **jarvisChat** | Custom `Dockerfile` | FastAPI app serving UI + API |
|
||||
| **SearXNG** | `searxng/searxng:latest` | Privacy-respecting web search |
|
||||
| **Qdrant** | `qdrant/qdrant:latest` | Vector database for RAG |
|
||||
| **RabbitMQ** | `rabbitmq:4-management` | Message broker for AMQP cluster |
|
||||
| **llama-server** | `ghcr.io/ggml-org/llama.cpp:server` | LLM inference (OpenAI-compat API) |
|
||||
| **Ollama** | `ollama/ollama:latest` | Embeddings for RAG chunk vectors |
|
||||
|
||||
### Non-containerized (host-level)
|
||||
|
||||
| Component | Reason |
|
||||
|-----------|--------|
|
||||
| AMD GPU driver + ROCm | Kernel access required for GPU compute |
|
||||
| llama.cpp RPC workers | Runs on *other* hosts — not on the Docker host |
|
||||
| `rocm-smi` | Hardware stats — not needed for core function |
|
||||
| `psutil` | Already inside the container via pip |
|
||||
|
||||
---
|
||||
|
||||
## 2. Service Catalog
|
||||
|
||||
### 2.1 jarvisChat (FastAPI app)
|
||||
|
||||
**Image:** `jarvischat:latest` (built from `Dockerfile`)
|
||||
|
||||
**Ports:**
|
||||
| Container | Host | Purpose |
|
||||
|-----------|------|---------|
|
||||
| 8080 | 8080 | HTTP API + UI |
|
||||
|
||||
**Volumes:**
|
||||
| Container path | Type | Purpose |
|
||||
|----------------|------|---------|
|
||||
| `/app/jarvischat.db` | named volume `jarvischat_data` | SQLite database |
|
||||
| `/app/uploads` | named volume `jarvischat_uploads` | Uploaded files |
|
||||
| `/app/hardware_state.json` | (inside volume) | Cached hardware probe |
|
||||
|
||||
**Dependencies:** Wait for SearXNG, Qdrant, RabbitMQ, llama-server, Ollama before serving.
|
||||
|
||||
**Restart:** `unless-stopped`
|
||||
|
||||
**Healthcheck:** `curl -f http://localhost:8080/`
|
||||
|
||||
### 2.2 SearXNG
|
||||
|
||||
**Image:** `searxng/searxng:latest`
|
||||
|
||||
**Ports:**
|
||||
| Container | Host | Purpose |
|
||||
|-----------|------|---------|
|
||||
| 8080 | 8888 | Search API |
|
||||
|
||||
**Volumes:**
|
||||
| Container path | Type | Purpose |
|
||||
|----------------|------|---------|
|
||||
| `/etc/searxng` | named volume `searxng_config` | `settings.yml` |
|
||||
|
||||
**Environment:**
|
||||
```env
|
||||
SEARXNG_BASE_URL=https://localhost:8888
|
||||
```
|
||||
|
||||
**Config override (`/etc/searxng/settings.yml`):**
|
||||
```yaml
|
||||
search:
|
||||
safe_search: 0
|
||||
autocomplete: ""
|
||||
server:
|
||||
secret_key: ${SEARXNG_SECRET_KEY}
|
||||
limiter: false
|
||||
image_proxy: false
|
||||
method: GET
|
||||
port: 8080
|
||||
bind_address: "0.0.0.0"
|
||||
```
|
||||
|
||||
**Restart:** `unless-stopped`
|
||||
|
||||
### 2.3 Qdrant
|
||||
|
||||
**Image:** `qdrant/qdrant:latest`
|
||||
|
||||
**Ports:**
|
||||
| Container | Host | Purpose |
|
||||
|-----------|------|---------|
|
||||
| 6333 | 6333 | HTTP API |
|
||||
| 6334 | — | gRPC (internal only) |
|
||||
|
||||
**Volumes:**
|
||||
| Container path | Type | Purpose |
|
||||
|----------------|------|---------|
|
||||
| `/qdrant/storage` | named volume `qdrant_storage` | Vector index data |
|
||||
|
||||
**Environment:**
|
||||
```env
|
||||
QDRANT__SERVICE__GRPC_PORT=6334
|
||||
```
|
||||
|
||||
**Restart:** `unless-stopped`
|
||||
|
||||
### 2.4 RabbitMQ
|
||||
|
||||
**Image:** `rabbitmq:4-management`
|
||||
|
||||
**Ports:**
|
||||
| Container | Host | Purpose |
|
||||
|-----------|------|---------|
|
||||
| 5672 | 5672 | AMQP messaging |
|
||||
| 15672 | — | Management UI (internal only) |
|
||||
|
||||
**Volumes:**
|
||||
| Container path | Type | Purpose |
|
||||
|----------------|------|---------|
|
||||
| `/var/lib/rabbitmq` | named volume `rabbitmq_data` | Message store |
|
||||
|
||||
**Environment:**
|
||||
```env
|
||||
RABBITMQ_DEFAULT_USER=jarvischat
|
||||
RABBITMQ_DEFAULT_PASS_FILE=/run/secrets/rabbitmq_password
|
||||
RABBITMQ_DEFAULT_VHOST=/
|
||||
```
|
||||
|
||||
**Restart:** `unless-stopped`
|
||||
|
||||
### 2.5 llama-server
|
||||
|
||||
**Image:** `ghcr.io/ggml-org/llama.cpp:server`
|
||||
|
||||
**Ports:**
|
||||
| Container | Host | Purpose |
|
||||
|-----------|------|---------|
|
||||
| 8081 | 8081 | OpenAI-compat API |
|
||||
|
||||
**Volumes:**
|
||||
| Container path | Type | Purpose |
|
||||
|----------------|------|---------|
|
||||
| `/models` | bind mount `./models` | Model GGUF files |
|
||||
|
||||
**Environment:**
|
||||
```env
|
||||
LLAMA_ARG_MODEL=/models/<model-file>
|
||||
LLAMA_ARG_N_GPU_LAYERS=0 # set >0 for GPU offload
|
||||
LLAMA_ARG_MAIN_GPU=0
|
||||
LLAMA_ARG_CTX_SIZE=4096
|
||||
LLAMA_ARG_HOST=0.0.0.0
|
||||
LLAMA_ARG_PORT=8081
|
||||
LLAMA_ARG_EMBEDDINGS=1
|
||||
LLAMA_ARG_LOGPROBS=1
|
||||
LLAMA_ARG_RPC= # optional: comma-separated RPC endpoints
|
||||
```
|
||||
|
||||
**Restart:** `unless-stopped`
|
||||
|
||||
**Healthcheck:** `curl -f http://localhost:8081/health`
|
||||
|
||||
**Notes:**
|
||||
- Models directory bind mount — user places `.gguf` files in `./models/` on the host
|
||||
- RPC offload to other machines (e.g., `10.0.0.50:50052,10.0.0.51:50052`)
|
||||
- If no GPU, set `LLAMA_ARG_N_GPU_LAYERS=0` for CPU-only
|
||||
- `LLAMA_ARG_EMBEDDINGS=1` required for perplexity scoring
|
||||
- `LLAMA_ARG_LOGPROBS=1` required for auto-search trigger
|
||||
|
||||
### 2.6 Ollama
|
||||
|
||||
**Image:** `ollama/ollama:latest`
|
||||
|
||||
**Ports:**
|
||||
| Container | Host | Purpose |
|
||||
|-----------|------|---------|
|
||||
| 11434 | 11434 | Embeddings API |
|
||||
|
||||
**Volumes:**
|
||||
| Container path | Type | Purpose |
|
||||
|----------------|------|---------|
|
||||
| `/root/.ollama` | named volume `ollama_models` | Pulled model blobs |
|
||||
|
||||
**Restart:** `unless-stopped`
|
||||
|
||||
**Notes:**
|
||||
- Used exclusively for embeddings (`/api/embeddings`), not inference
|
||||
- Typically needs a small model like `all-minilm:latest` or `nomic-embed-text:latest`
|
||||
- Consider replacing Ollama with llama-server's built-in embedding if it supports the same model — would remove one container
|
||||
|
||||
---
|
||||
|
||||
## 3. Configuration Management
|
||||
|
||||
### 3.1 `.env` file (generated by setup wizard)
|
||||
|
||||
```env
|
||||
# --- Secrets (auto-generated, change before production) ---
|
||||
JARVISCHAT_ADMIN_PIN=
|
||||
JARVISCHAT_COMPLETIONS_API_KEY=
|
||||
JARVISCHAT_ALLOW_DEFAULT_PIN=false
|
||||
RABBITMQ_PASSWORD=
|
||||
SEARXNG_SECRET_KEY=
|
||||
|
||||
# --- Host discovery (auto-detected by setup wizard) ---
|
||||
LLAMA_SERVER_BASE=http://llama-server:8081
|
||||
OLLAMA_BASE=http://ollama:11434
|
||||
SEARXNG_BASE=http://searxng:8888
|
||||
QDRANT_URL=http://qdrant:6333
|
||||
RABBITMQ_HOST=rabbitmq
|
||||
RABBITMQ_PORT=5672
|
||||
|
||||
# --- Performance tuning (calculated by setup wizard) ---
|
||||
RAG_MAX_VECTORS=50000
|
||||
RAG_EVICTION_HIGH_WATER=0.80
|
||||
RAG_EVICTION_LOW_WATER=0.20
|
||||
RAG_EVICTION_BATCH=1000
|
||||
|
||||
# --- llama-server options ---
|
||||
LLAMA_MODEL=llama3.1-8b-instruct.Q4_K_M.gguf
|
||||
LLAMA_N_GPU_LAYERS=0
|
||||
LLAMA_RPC_ENDPOINTS=
|
||||
LLAMA_CTX_SIZE=4096
|
||||
|
||||
# --- Ollama ---
|
||||
OLLAMA_EMBED_MODEL=all-minilm:latest
|
||||
|
||||
# --- Network ---
|
||||
JARVISCHAT_ALLOWED_CIDRS=127.0.0.0/8,::1/128,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16
|
||||
JARVISCHAT_TRUSTED_ORIGINS=
|
||||
JARVISCHAT_TRUST_X_FORWARDED_FOR=false
|
||||
```
|
||||
|
||||
### 3.2 Mapping of config.py → .env variable
|
||||
|
||||
Every config.py default that references an external service must accept a matching env var at runtime:
|
||||
|
||||
| config.py constant | .env variable | Service |
|
||||
|-------------------|---------------|---------|
|
||||
| `LLAMA_SERVER_BASE` | `LLAMA_SERVER_BASE` | llama-server |
|
||||
| `OLLAMA_BASE` | `OLLAMA_BASE` | Ollama |
|
||||
| `SEARXNG_BASE` | `SEARXNG_BASE` | SearXNG |
|
||||
| `QDRANT_URL` | `QDRANT_URL` | Qdrant |
|
||||
| `COMPLETIONS_API_KEY` | `JARVISCHAT_COMPLETIONS_API_KEY` | — |
|
||||
| `ALLOWED_CIDRS_RAW` | `JARVISCHAT_ALLOWED_CIDRS` | — |
|
||||
| `TRUST_X_FORWARDED_FOR` | `JARVISCHAT_TRUST_X_FORWARDED_FOR` | — |
|
||||
| `TRUSTED_ORIGINS` | `JARVISCHAT_TRUSTED_ORIGINS` | — |
|
||||
| `RAG_MAX_VECTORS` | `RAG_MAX_VECTORS` | — (calc'd from RAM) |
|
||||
|
||||
### 3.3 Secrets management
|
||||
|
||||
| Secret | Generated by | Stored in | Mounted to |
|
||||
|--------|-------------|-----------|------------|
|
||||
| `JARVISCHAT_ADMIN_PIN` | User prompt | `.env` | jarvisChat container |
|
||||
| `JARVISCHAT_COMPLETIONS_API_KEY` | Auto-generated, shown to user | `.env` | jarvisChat container |
|
||||
| `RABBITMQ_PASSWORD` | Auto-generated | `.env` + Docker secret | RabbitMQ container |
|
||||
| `SEARXNG_SECRET_KEY` | Auto-generated | `.env` | SearXNG container |
|
||||
|
||||
**Docker secrets approach:** Use `secrets:` in compose file for RabbitMQ password (mounted as file) rather than passing via env var, since `settings.yml` in SearXNG and RabbitMQ config can reference file-based secrets without env-var leakage.
|
||||
|
||||
### 3.4 Dockerfile for jarvisChat
|
||||
|
||||
```dockerfile
|
||||
FROM python:3.13-slim-bookworm AS builder
|
||||
WORKDIR /app
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
FROM python:3.13-slim-bookworm
|
||||
WORKDIR /app
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
curl && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
COPY --from=builder /usr/local/lib/python3.13/site-packages /usr/local/lib/python3.13/site-packages
|
||||
COPY --from=builder /usr/local/bin /usr/local/bin
|
||||
COPY . .
|
||||
|
||||
EXPOSE 8080
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=15s --retries=3 \
|
||||
CMD curl -f http://localhost:8080/ || exit 1
|
||||
|
||||
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8080"]
|
||||
```
|
||||
|
||||
**Multi-stage rationale:** First stage compiles/bundles packages (wheels), final stage is minimal. Devs can skip builder with `--target builder` for live-reload with volume mount.
|
||||
|
||||
---
|
||||
|
||||
## 4. docker-compose.yml structure
|
||||
|
||||
```yaml
|
||||
services:
|
||||
jarvischat:
|
||||
build: .
|
||||
ports: ["8080:8080"]
|
||||
volumes:
|
||||
- jarvischat_data:/app/jarvischat.db
|
||||
- jarvischat_uploads:/app/uploads
|
||||
env_file: .env
|
||||
depends_on:
|
||||
searxng: { condition: service_started }
|
||||
qdrant: { condition: service_started }
|
||||
rabbitmq: { condition: service_healthy }
|
||||
llama-server: { condition: service_healthy }
|
||||
ollama: { condition: service_started }
|
||||
restart: unless-stopped
|
||||
|
||||
searxng:
|
||||
image: searxng/searxng:latest
|
||||
ports: ["8888:8080"]
|
||||
volumes:
|
||||
- ./searxng/settings.yml:/etc/searxng/settings.yml:ro
|
||||
- searxng_config:/etc/searxng
|
||||
env_file: .env
|
||||
restart: unless-stopped
|
||||
|
||||
qdrant:
|
||||
image: qdrant/qdrant:latest
|
||||
ports: ["6333:6333"]
|
||||
volumes:
|
||||
- qdrant_storage:/qdrant/storage
|
||||
restart: unless-stopped
|
||||
|
||||
rabbitmq:
|
||||
image: rabbitmq:4-management
|
||||
ports: ["5672:5672"]
|
||||
volumes:
|
||||
- rabbitmq_data:/var/lib/rabbitmq
|
||||
env_file: .env
|
||||
secrets:
|
||||
- rabbitmq_password
|
||||
healthcheck:
|
||||
test: ["CMD", "rabbitmq-diagnostics", "check_port_connectivity"]
|
||||
interval: 15s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
restart: unless-stopped
|
||||
|
||||
llama-server:
|
||||
image: ghcr.io/ggml-org/llama.cpp:server
|
||||
ports: ["8081:8081"]
|
||||
volumes:
|
||||
- ./models:/models:ro
|
||||
env_file: .env
|
||||
command: >
|
||||
--model /models/${LLAMA_MODEL}
|
||||
--host 0.0.0.0 --port 8081
|
||||
--ctx-size ${LLAMA_CTX_SIZE:-4096}
|
||||
--n-gpu-layers ${LLAMA_N_GPU_LAYERS:-0}
|
||||
--embeddings
|
||||
--logprobs
|
||||
${LLAMA_RPC_ENDPOINTS:+--rpc ${LLAMA_RPC_ENDPOINTS}}
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:8081/health"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 5
|
||||
start_period: 60s
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
restart: unless-stopped
|
||||
|
||||
ollama:
|
||||
image: ollama/ollama:latest
|
||||
ports: ["11434:11434"]
|
||||
volumes:
|
||||
- ollama_models:/root/.ollama
|
||||
healthcheck:
|
||||
test: ["CMD", "ollama", "list"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 3
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
jarvischat_data:
|
||||
jarvischat_uploads:
|
||||
searxng_config:
|
||||
qdrant_storage:
|
||||
rabbitmq_data:
|
||||
ollama_models:
|
||||
|
||||
secrets:
|
||||
rabbitmq_password:
|
||||
file: ./secrets/rabbitmq_password.txt
|
||||
```
|
||||
|
||||
**Notes:**
|
||||
- GPU reservations use `resources.reservations.devices` — this is compose v3.8+. For AMD GPUs, replace `driver: nvidia` with `driver: amd` (experimental Docker support). For hosts without GPU, omit the `deploy` block entirely.
|
||||
- The `deploy` block only applies when deployed as a swarm stack. For `docker compose`, GPU access may need `--gpus all` or `device_requests` in config. Verify compatibility.
|
||||
- SearXNG config file (`settings.yml`) is bind-mounted read-only from the host repo clone — the setup wizard should generate this file.
|
||||
|
||||
---
|
||||
|
||||
## 5. Networking
|
||||
|
||||
### 5.1 Internal communication (compose network)
|
||||
|
||||
| From | To | Port | Protocol |
|
||||
|------|----|------|----------|
|
||||
| jarvisChat | llama-server | 8081 | HTTP |
|
||||
| jarvisChat | Ollama | 11434 | HTTP |
|
||||
| jarvisChat | SearXNG | 8080 | HTTP |
|
||||
| jarvisChat | Qdrant | 6333 | HTTP |
|
||||
| jarvisChat | RabbitMQ | 5672 | AMQP |
|
||||
| RabbitMQ | (cluster peers) | 4369 | EPMD |
|
||||
| RabbitMQ | (cluster peers) | 25672 | Inter-node |
|
||||
|
||||
### 5.2 Exposed ports (host-facing)
|
||||
|
||||
| Port | Service | Should expose? | Notes |
|
||||
|------|---------|---------------|-------|
|
||||
| 8080 | jarvisChat | ✅ Required | UI + API |
|
||||
| 8888 | SearXNG | Optional | Only if user wants standalone search |
|
||||
| 6333 | Qdrant | Optional | Only for external tooling |
|
||||
| 5672 | RabbitMQ | Optional | Only for remote AMQP clients |
|
||||
| 15672 | RabbitMQ mgmt | ❌ Internal | Healthcheck only |
|
||||
| 8081 | llama-server | Optional | Only for external tooling |
|
||||
| 11434 | Ollama | Optional | Only for external tooling |
|
||||
|
||||
**Design decision:** By default, only port 8080 (jarvisChat) is published. All other services remain on the internal compose network. Advanced users can opt-in by uncommenting `ports:` blocks.
|
||||
|
||||
### 5.3 Reverse proxy consideration
|
||||
|
||||
For production, a reverse proxy (Caddy, nginx, Traefik) should sit in front:
|
||||
|
||||
```yaml
|
||||
# Optional — compose profile: "proxy"
|
||||
caddy:
|
||||
image: caddy:latest
|
||||
ports: ["80:80", "443:443"]
|
||||
volumes:
|
||||
- ./Caddyfile:/etc/caddy/Caddyfile:ro
|
||||
- caddy_data:/data
|
||||
```
|
||||
|
||||
This is out of scope for v1.0 but documented for future.
|
||||
|
||||
---
|
||||
|
||||
## 6. Setup Wizard (Extraction)
|
||||
|
||||
`setup.sh` — idempotent, interactive, runs on first boot.
|
||||
|
||||
### Flow
|
||||
|
||||
```
|
||||
1. CHECK: Is .env present?
|
||||
├── YES → skip to step 7 (or ask to regenerate)
|
||||
└── NO → continue
|
||||
|
||||
2. INTRO: Print banner, explain what's about to happen
|
||||
|
||||
3. PROBE: Run hardware assessment
|
||||
├── psutil → RAM total, CPU count
|
||||
├── rocm-smi → VRAM (optional, best-effort)
|
||||
└── nvidia-smi → VRAM (optional, best-effort)
|
||||
|
||||
4. NETWORK: Ask for
|
||||
├── Hostname / LAN IP for this machine
|
||||
├── Admin PIN (4 digits, or accept auto-generated)
|
||||
└── (Optional) RPC endpoints for GPU offload
|
||||
|
||||
5. CALCULATE:
|
||||
├── RAG_MAX_VECTORS = max(1000, int(available_ram_gb * 100_000))
|
||||
├── LLAMA_N_GPU_LAYERS = 0 (CPU default; offer GPU detection)
|
||||
├── LLAMA_MODEL = default gguf filename
|
||||
└── RABBITMQ_PASSWORD = openssl rand -hex 20
|
||||
|
||||
6. GENERATE:
|
||||
├── .env file from template
|
||||
├── ./secrets/rabbitmq_password.txt
|
||||
├── ./searxng/settings.yml (with generated secret_key)
|
||||
└── ./models/README.txt (instructions for placing .gguf)
|
||||
|
||||
7. VERIFY:
|
||||
├── docker and docker compose plugin installed
|
||||
├── docker compose version >= 2.x
|
||||
├── SUCCESS → "Run: docker compose up -d"
|
||||
└── FAILURE → show diagnostics and links
|
||||
|
||||
8. EXTRACT model:
|
||||
├── Prompt for download URL or local path
|
||||
├── Offer to pull from HuggingFace if huggingface-cli available
|
||||
└── Guides user to place file in ./models/
|
||||
```
|
||||
|
||||
### What setup.sh creates on disk
|
||||
|
||||
```
|
||||
./docker-deploy/
|
||||
├── .env # All env vars (SECRET — add to .gitignore)
|
||||
├── docker-compose.yml # Compose stack definition
|
||||
├── Dockerfile # jarvisChat image build
|
||||
├── secrets/
|
||||
│ └── rabbitmq_password.txt # RabbitMQ password file
|
||||
├── searxng/
|
||||
│ └── settings.yml # SearXNG config with generated secret_key
|
||||
├── models/
|
||||
│ ├── README.txt # Instructions for model placement
|
||||
│ └── <model>.gguf # (user-provided)
|
||||
└── setup.log # Wizard run log
|
||||
```
|
||||
|
||||
### Idempotency
|
||||
|
||||
Re-running `setup.sh`:
|
||||
- With `.env` present: ask "Regenerate? This will overwrite existing config."
|
||||
- Without `.env`: fresh run
|
||||
- Never overwrites `./models/*.gguf` files
|
||||
- Never touches running containers — only modifies files on disk
|
||||
|
||||
---
|
||||
|
||||
## 7. Back-out Procedure (Uninstall)
|
||||
|
||||
`teardown.sh` — returns the host system to its pre-install state.
|
||||
|
||||
### What gets removed
|
||||
|
||||
| Item | Removal method |
|
||||
|------|---------------|
|
||||
| Docker containers | `docker compose down -v` |
|
||||
| Docker images | `docker rmi jarvischat:latest` (ask about other images) |
|
||||
| Docker volumes | `docker volume rm jarvischat_data ...` (prompt first) |
|
||||
| Network `jarvischat_default` | Removed with compose |
|
||||
| `.env` file | `rm .env` |
|
||||
| `secrets/` directory | `rm -rf secrets/` |
|
||||
| `searxng/` directory | `rm -rf searxng/` |
|
||||
| `setup.log` | `rm setup.log` |
|
||||
| `hardware_state.json` | `rm hardware_state.json` |
|
||||
|
||||
### What is preserved (by default)
|
||||
|
||||
| Item | Reason |
|
||||
|------|--------|
|
||||
| `./models/*.gguf` | User data — prompt for deletion |
|
||||
| `jarvischat.db` (in volume) | Prompt: "Keep database snapshot?" |
|
||||
| `./uploads/` (in volume) | Prompt: "Keep uploaded files?" |
|
||||
| Docker Engine itself | Not installed by this project — leave it |
|
||||
|
||||
### Script flow
|
||||
|
||||
```
|
||||
1. CHECK: docker compose file exists?
|
||||
├── NO → warn, continue
|
||||
└── YES → docker compose down -v
|
||||
|
||||
2. CHECK: .env exists?
|
||||
├── NO → skip
|
||||
└── YES → ask: "Remove .env?" (default no)
|
||||
|
||||
3. ASK: "Remove secrets/ and searxng/ directories?" (default no)
|
||||
|
||||
4. ASK: "Remove Docker images? (y/N)" (default no)
|
||||
├── Y → docker rmi jarvischat:latest
|
||||
├── Y → docker image ls | grep searxng/qdrant/rabbitmq → prompt per image
|
||||
└── N → skip
|
||||
|
||||
5. ASK: "Keep database volume snapshot? (Y/n)" (default yes)
|
||||
├── N → docker volume rm jarvischat_data
|
||||
└── Y → leave volume (can be reattached later)
|
||||
|
||||
6. ASK: "Remove model files from ./models/? (y/N)" (default no)
|
||||
|
||||
7. CLEANUP generated artifacts:
|
||||
├── rm -f setup.log
|
||||
├── rm -f hardware_state.json
|
||||
└── rm -f docker-compose.yml
|
||||
|
||||
8. SUMMARY:
|
||||
├── "Docker stack removed"
|
||||
├── "Persistent data preserved at: <paths>"
|
||||
└── "Models kept at: ./models/"
|
||||
```
|
||||
|
||||
### Partial rollback
|
||||
|
||||
If the setup wizard fails mid-way, a partial rollback is better than leaving detritus:
|
||||
|
||||
| Failure point | Clean up |
|
||||
|--------------|----------|
|
||||
| After .env, before compose | `rm .env; rm -rf secrets/ searxng/` |
|
||||
| After compose, before first `up` | `rm docker-compose.yml; rm -rf *` |
|
||||
| After `up` but before healthcheck | `docker compose down -v; rm -rf ./*` |
|
||||
|
||||
`setup.sh` should trap EXIT on failure and prompt: "Clean up partial install? [y/N]"
|
||||
|
||||
---
|
||||
|
||||
## 8. Open Decisions
|
||||
|
||||
| Decision | Options | Priority |
|
||||
|----------|---------|----------|
|
||||
| **Ollama vs llama-server embeddings** | Both work. Keep both for now — remove Ollama if llama-server handles embeddings. Reduce containers = simpler. | Medium |
|
||||
| **GPU support in compose** | NVIDIA: well-supported. AMD: requires `--device=/dev/kfd --device=/dev/dri` and ROCm image. Document both. | High |
|
||||
| **RabbitMQ clustering vs single node** | Single node in v1.0. Clustering docs for multi-host later. | Low |
|
||||
| **SearXNG config management** | Bind-mount a generated `settings.yml`, or let container create default and post-process. Bind-mount is cleaner. | Medium |
|
||||
| **Reverse proxy** | Caddy is simplest for auto-HTTPS. Out of scope for v1.0 but design for it. | Low |
|
||||
| **Healthcheck strategy** | `depends_on` with `condition: service_healthy` is the safest approach but increases startup time. Acceptable. | Medium |
|
||||
| **Database migration** | SQLite file in volume — no migration needed for v1.0 format. If schema changes post-v1.0, need a migration container. | Low |
|
||||
| **Linux vs macOS vs Windows** | Linux-primary. macOS may work with changes (no rocm-smi). Windows via WSL2 only. | Low |
|
||||
| **LLM model download** | HuggingFace CLI integration in setup.sh, or manual download. Manual is simpler. | Low |
|
||||
| **Dockerfile optimization** | Pin pip hashes, use `--no-cache-dir`, consider `slim` vs `alpine`. Alpine has musl compatibility issues with psutil. Stay with slim. | Medium |
|
||||
|
||||
## 9. Checklist (pre-v1.0 gate)
|
||||
|
||||
- [ ] `Dockerfile` written and builds clean
|
||||
- [ ] `docker-compose.yml` boots all containers
|
||||
- [ ] jarvisChat container reaches all services (env vars resolve correctly)
|
||||
- [ ] SearXNG settings.yml generated correctly by setup.sh
|
||||
- [ ] RabbitMQ password secret mounted correctly
|
||||
- [ ] GPU (NVIDIA) passes through to llama-server container
|
||||
- [ ] GPU (AMD) passes through to llama-server container (or documented limitation)
|
||||
- [ ] `.env.example` checked in (no real secrets)
|
||||
- [ ] `setup.sh` written, idempotent, tested on clean Debian
|
||||
- [ ] `teardown.sh` written, tested, doesn't delete models without confirmation
|
||||
- [ ] `docker compose up -d` works without any manual steps beyond setup.sh
|
||||
- [ ] `docker compose down -v` followed by `setup.sh && docker compose up -d` = fresh stack
|
||||
- [ ] Healthchecks prevent serving before dependencies are ready
|
||||
- [ ] v1.0 release tag created
|
||||
|
||||
---
|
||||
|
||||
## 10. Files to create for B3
|
||||
|
||||
```
|
||||
docker.md ← this file (planning doc)
|
||||
Dockerfile ← jarvisChat image
|
||||
docker-compose.yml ← full stack
|
||||
.env.example ← template without secrets
|
||||
setup.sh ← extraction wizard
|
||||
teardown.sh ← back-out utility
|
||||
searxng/
|
||||
settings.yml ← SearXNG config (generated by setup.sh)
|
||||
secrets/
|
||||
rabbitmq_password.txt ← generated by setup.sh
|
||||
models/
|
||||
README.txt ← instructions for placing .gguf
|
||||
```
|
||||
+268
@@ -0,0 +1,268 @@
|
||||
"""
|
||||
JarvisChat — Score-based RAG vector eviction with hysteresis.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
import httpx
|
||||
|
||||
from config import (
|
||||
QDRANT_URL, RAG_COLLECTION,
|
||||
RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER,
|
||||
RAG_EVICTION_BATCH, RAG_PINNED_SOURCES, RAG_GRACE_HOURS,
|
||||
RAG_ACCESS_WEIGHT, RAG_AGE_WEIGHT,
|
||||
)
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
|
||||
eviction_lock = asyncio.Lock()
|
||||
EVICTION_LOG: list[dict] = []
|
||||
|
||||
|
||||
async def _update_retrieval_count(point_id: str, current_count: int = 0):
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
payload = {
|
||||
"retrieval_count": current_count + 1,
|
||||
"last_accessed": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
resp = await client.put(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/payload",
|
||||
json={"points": [point_id], "payload": payload},
|
||||
timeout=5.0,
|
||||
)
|
||||
if resp.status_code not in (200, 201):
|
||||
log.warning(f"Failed to increment retrieval count for {point_id}: {resp.status_code}")
|
||||
except Exception as e:
|
||||
log.warning(f"Error incrementing retrieval count for {point_id}: {e}")
|
||||
|
||||
|
||||
async def get_collection_count() -> int:
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
resp = await client.get(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}",
|
||||
timeout=10.0,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
return resp.json().get("result", {}).get("vectors_count", 0)
|
||||
except Exception as e:
|
||||
log.warning(f"get_collection_count error: {e}")
|
||||
return 0
|
||||
|
||||
|
||||
async def get_collection_stats() -> dict:
|
||||
count = await get_collection_count()
|
||||
high_water_pct = int(RAG_EVICTION_HIGH_WATER * 100)
|
||||
low_water_pct = int(RAG_EVICTION_LOW_WATER * 100)
|
||||
percent_full = round((count / RAG_MAX_VECTORS) * 100, 1) if RAG_MAX_VECTORS > 0 else 0
|
||||
return {
|
||||
"vector_count": count,
|
||||
"max_vectors": RAG_MAX_VECTORS,
|
||||
"high_water_mark": int(RAG_MAX_VECTORS * RAG_EVICTION_HIGH_WATER),
|
||||
"low_water_mark": int(RAG_MAX_VECTORS * RAG_EVICTION_LOW_WATER),
|
||||
"high_water_pct": high_water_pct,
|
||||
"low_water_pct": low_water_pct,
|
||||
"percent_full": percent_full,
|
||||
"pinned_sources": list(RAG_PINNED_SOURCES),
|
||||
}
|
||||
|
||||
|
||||
async def evict_batch(batch_size: int) -> int:
|
||||
filter_conditions = {
|
||||
"must_not": [
|
||||
{"match": {"key": "source", "value": src}}
|
||||
for src in RAG_PINNED_SOURCES
|
||||
]
|
||||
}
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
scroll_resp = await client.post(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll",
|
||||
json={
|
||||
"filter": filter_conditions,
|
||||
"limit": min(batch_size * 10, 10000),
|
||||
"with_payload": True,
|
||||
"with_vector": False,
|
||||
},
|
||||
timeout=30.0,
|
||||
)
|
||||
if scroll_resp.status_code != 200:
|
||||
log.warning(f"Eviction scroll failed: {scroll_resp.status_code}")
|
||||
return 0
|
||||
|
||||
points = scroll_resp.json().get("result", {}).get("points", [])
|
||||
if not points:
|
||||
return 0
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
scored = []
|
||||
for p in points:
|
||||
payload = p.get("payload", {})
|
||||
date_str = payload.get("ingest_date") or payload.get("upload_date", "")
|
||||
if date_str:
|
||||
age_hours = (now - datetime.fromisoformat(date_str)).total_seconds() / 3600
|
||||
else:
|
||||
age_hours = 999999
|
||||
|
||||
if age_hours < RAG_GRACE_HOURS:
|
||||
continue
|
||||
|
||||
retrieval_count = payload.get("retrieval_count", 0) or 0
|
||||
score = retrieval_count * RAG_ACCESS_WEIGHT + age_hours * RAG_AGE_WEIGHT
|
||||
last_accessed = payload.get("last_accessed", date_str)
|
||||
scored.append((score, last_accessed, p["id"]))
|
||||
|
||||
if not scored:
|
||||
log.warning("No evictable vectors found (all pinned or newborn)")
|
||||
return 0
|
||||
|
||||
scored.sort(key=lambda x: (x[0], x[1]))
|
||||
to_delete = [p[2] for p in scored[:batch_size]]
|
||||
if not to_delete:
|
||||
return 0
|
||||
|
||||
delete_resp = await client.post(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/delete",
|
||||
json={"points": to_delete},
|
||||
timeout=30.0,
|
||||
)
|
||||
if delete_resp.status_code not in (200, 201):
|
||||
log.warning(f"Eviction delete failed: {delete_resp.status_code}")
|
||||
return 0
|
||||
|
||||
return len(to_delete)
|
||||
except Exception as e:
|
||||
log.warning(f"evict_batch error: {e}")
|
||||
return 0
|
||||
|
||||
|
||||
async def maybe_evict() -> int:
|
||||
if RAG_MAX_VECTORS <= 0:
|
||||
return 0
|
||||
effective_batch = max(RAG_EVICTION_BATCH, 1)
|
||||
|
||||
async with eviction_lock:
|
||||
count = await get_collection_count()
|
||||
threshold_high = int(RAG_MAX_VECTORS * RAG_EVICTION_HIGH_WATER)
|
||||
threshold_low = int(RAG_MAX_VECTORS * RAG_EVICTION_LOW_WATER)
|
||||
|
||||
if count < threshold_high:
|
||||
return 0
|
||||
|
||||
total_evicted = 0
|
||||
while count >= threshold_low:
|
||||
if total_evicted > 0 and count < threshold_low:
|
||||
break
|
||||
deleted = await evict_batch(effective_batch)
|
||||
if deleted == 0:
|
||||
break
|
||||
total_evicted += deleted
|
||||
count -= deleted
|
||||
if count < threshold_high and total_evicted > 0:
|
||||
break
|
||||
if count < threshold_low:
|
||||
break
|
||||
|
||||
if total_evicted > 0:
|
||||
entry = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"count": total_evicted,
|
||||
"remaining": count,
|
||||
}
|
||||
EVICTION_LOG.append(entry)
|
||||
if len(EVICTION_LOG) > 1000:
|
||||
EVICTION_LOG.pop(0)
|
||||
log.info(f"Evicted {total_evicted} vectors ({count} remaining)")
|
||||
|
||||
return total_evicted
|
||||
|
||||
|
||||
async def get_rag_operational_stats() -> dict:
|
||||
stats = await get_collection_stats()
|
||||
now = datetime.now(timezone.utc)
|
||||
cutoff_1m = now - timedelta(minutes=1)
|
||||
cutoff_5m = now - timedelta(minutes=5)
|
||||
cutoff_30m = now - timedelta(minutes=30)
|
||||
|
||||
eviction_1m = sum(
|
||||
e["count"] for e in EVICTION_LOG
|
||||
if datetime.fromisoformat(e["timestamp"]) > cutoff_1m
|
||||
)
|
||||
eviction_5m = sum(
|
||||
e["count"] for e in EVICTION_LOG
|
||||
if datetime.fromisoformat(e["timestamp"]) > cutoff_5m
|
||||
)
|
||||
eviction_30m = sum(
|
||||
e["count"] for e in EVICTION_LOG
|
||||
if datetime.fromisoformat(e["timestamp"]) > cutoff_30m
|
||||
)
|
||||
|
||||
pinned_count = 0
|
||||
avg_retrieval_count = 0.0
|
||||
at_risk_count = 0
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
pinned_filter = {
|
||||
"should": [
|
||||
{"match": {"key": "source", "value": src}}
|
||||
for src in RAG_PINNED_SOURCES
|
||||
]
|
||||
}
|
||||
pinned_resp = await client.post(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll",
|
||||
json={"filter": pinned_filter, "limit": 10000, "with_payload": True, "with_vector": False},
|
||||
timeout=10.0,
|
||||
)
|
||||
if pinned_resp.status_code == 200:
|
||||
pinned_count = len(pinned_resp.json().get("result", {}).get("points", []))
|
||||
|
||||
nonpinned_filter = {
|
||||
"must_not": [
|
||||
{"match": {"key": "source", "value": src}}
|
||||
for src in RAG_PINNED_SOURCES
|
||||
]
|
||||
}
|
||||
np_resp = await client.post(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll",
|
||||
json={"filter": nonpinned_filter, "limit": 10000, "with_payload": True, "with_vector": False},
|
||||
timeout=10.0,
|
||||
)
|
||||
if np_resp.status_code == 200:
|
||||
points = np_resp.json().get("result", {}).get("points", [])
|
||||
if points:
|
||||
retrievals = []
|
||||
scored = []
|
||||
for p in points:
|
||||
payload = p.get("payload", {})
|
||||
rc = payload.get("retrieval_count", 0) or 0
|
||||
retrievals.append(rc)
|
||||
date_str = payload.get("ingest_date") or payload.get("upload_date", "")
|
||||
if date_str:
|
||||
age_hours = (now - datetime.fromisoformat(date_str)).total_seconds() / 3600
|
||||
else:
|
||||
age_hours = 999999
|
||||
score = rc * RAG_ACCESS_WEIGHT + age_hours * RAG_AGE_WEIGHT
|
||||
last_accessed = payload.get("last_accessed", date_str)
|
||||
scored.append((score, last_accessed))
|
||||
|
||||
avg_retrieval_count = round(sum(retrievals) / len(retrievals), 2)
|
||||
|
||||
scored.sort(key=lambda x: (x[0], x[1]))
|
||||
at_risk_threshold = max(1, len(scored) // 10)
|
||||
at_risk_count = at_risk_threshold
|
||||
except Exception as e:
|
||||
log.warning(f"RAG operational stats scroll error: {e}")
|
||||
|
||||
stats.update({
|
||||
"grace_hours": RAG_GRACE_HOURS,
|
||||
"eviction_counts_last_1m": eviction_1m,
|
||||
"eviction_counts_last_5m": eviction_5m,
|
||||
"eviction_counts_last_30m": eviction_30m,
|
||||
"pinned_count": pinned_count,
|
||||
"avg_retrieval_count": avg_retrieval_count,
|
||||
"at_risk_count": at_risk_count,
|
||||
})
|
||||
return stats
|
||||
@@ -1,26 +1,31 @@
|
||||
"""
|
||||
JarvisChat - RAG pipeline: Qdrant vector search + system prompt assembly.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
|
||||
from eviction import _update_retrieval_count
|
||||
from db import get_db, get_setting, list_skills_with_state, format_active_skills_prompt
|
||||
from memory import search_memories
|
||||
from config import MAX_SKILL_PROMPT_CHARS
|
||||
from config import MAX_SKILL_PROMPT_CHARS, QDRANT_URL, RAG_COLLECTION
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
|
||||
QDRANT_URL = "http://192.168.50.108:6333"
|
||||
EMBED_URL = "http://192.168.50.210:11434"
|
||||
EMBED_MODEL = "mxbai-embed-large"
|
||||
RAG_COLLECTION = "jarvis_rag"
|
||||
RAG_SCORE_THRESHOLD = 0.25
|
||||
|
||||
# Re-export eviction symbols for backward compatibility
|
||||
from eviction import ( # noqa: E402
|
||||
maybe_evict, get_rag_operational_stats, EVICTION_LOG,
|
||||
get_collection_count, get_collection_stats, evict_batch,
|
||||
)
|
||||
|
||||
|
||||
def chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list:
|
||||
words = text.split()
|
||||
# Approximate token count as len(words) * 1.3
|
||||
target_words = int(chunk_size / 1.3)
|
||||
overlap_words = int(overlap / 1.3)
|
||||
if not words:
|
||||
@@ -54,7 +59,13 @@ async def query_rag(query: str, limit: int = 3) -> list:
|
||||
)
|
||||
if search_resp.status_code != 200:
|
||||
return []
|
||||
return search_resp.json().get("result", [])
|
||||
results = search_resp.json().get("result", [])
|
||||
for r in results:
|
||||
pid = r.get("id")
|
||||
if pid:
|
||||
current = r.get("payload", {}).get("retrieval_count", 0) or 0
|
||||
asyncio.ensure_future(_update_retrieval_count(pid, current))
|
||||
return results
|
||||
except Exception as e:
|
||||
log.warning(f"RAG query error: {e}")
|
||||
return []
|
||||
|
||||
@@ -7,6 +7,7 @@ from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from config import COMPLETIONS_API_KEY
|
||||
from eviction import maybe_evict
|
||||
from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
@@ -61,4 +62,9 @@ async def ingest_content(request: Request):
|
||||
else:
|
||||
log.warning(f"Ingest Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}")
|
||||
|
||||
if ingested > 0:
|
||||
evicted = await maybe_evict()
|
||||
if evicted:
|
||||
log.info(f"Evicted {evicted} vectors after ingest")
|
||||
|
||||
return {"chunks_ingested": ingested, "source": source, "message": f"Ingested {ingested} chunks from {source}"}
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
"""JarvisChat routers — RAG corpus management admin endpoints."""
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from eviction import get_rag_operational_stats, EVICTION_LOG
|
||||
from rag import QDRANT_URL, RAG_COLLECTION
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/api/rag/stats")
|
||||
async def rag_stats(request: Request):
|
||||
if getattr(request.state, "session_role", "none") != "admin":
|
||||
raise HTTPException(status_code=403, detail="Admin PIN required for this action")
|
||||
stats = await get_rag_operational_stats()
|
||||
stats["eviction_log_size"] = len(EVICTION_LOG)
|
||||
return stats
|
||||
|
||||
|
||||
@router.post("/api/rag/flush")
|
||||
async def rag_flush(request: Request):
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
scroll_resp = await client.post(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/scroll",
|
||||
json={"limit": 10000, "with_payload": False, "with_vector": False},
|
||||
timeout=30.0,
|
||||
)
|
||||
if scroll_resp.status_code != 200:
|
||||
return JSONResponse(status_code=502, content={"detail": f"Qdrant scroll failed: {scroll_resp.status_code}"})
|
||||
|
||||
all_points = scroll_resp.json().get("result", {}).get("points", [])
|
||||
point_ids = [p["id"] for p in all_points]
|
||||
|
||||
if not point_ids:
|
||||
return {"deleted_count": 0, "collection": RAG_COLLECTION, "status": "flushed"}
|
||||
|
||||
delete_resp = await client.post(
|
||||
f"{QDRANT_URL}/collections/{RAG_COLLECTION}/points/delete",
|
||||
json={"points": point_ids},
|
||||
timeout=30.0,
|
||||
)
|
||||
if delete_resp.status_code not in (200, 201):
|
||||
return JSONResponse(status_code=502, content={"detail": f"Qdrant delete failed: {delete_resp.status_code}"})
|
||||
|
||||
EVICTION_LOG.clear()
|
||||
log.warning(f"RAG collection '{RAG_COLLECTION}' flushed ({len(point_ids)} points deleted)")
|
||||
|
||||
return {
|
||||
"deleted_count": len(point_ids),
|
||||
"collection": RAG_COLLECTION,
|
||||
"status": "flushed",
|
||||
}
|
||||
except Exception as e:
|
||||
log.warning(f"RAG flush error: {e}")
|
||||
return JSONResponse(status_code=502, content={"detail": f"RAG flush error: {e}"})
|
||||
@@ -10,6 +10,7 @@ from fastapi.responses import JSONResponse
|
||||
|
||||
from config import UPLOAD_DIR, MAX_UPLOAD_BYTES, SUPPORTED_UPLOAD_TYPES, UPLOAD_CONTEXT_EXPIRY_HOURS
|
||||
from db import get_db, insert_upload_context, list_upload_context_by_conversation, delete_upload_context_by_id
|
||||
from eviction import maybe_evict
|
||||
from rag import chunk_text, QDRANT_URL, EMBED_URL, EMBED_MODEL, RAG_COLLECTION
|
||||
|
||||
log = logging.getLogger("jarvischat")
|
||||
@@ -87,6 +88,10 @@ async def upload_file(
|
||||
else:
|
||||
log.warning(f"Qdrant upsert failed for chunk {i}: {upsert_resp.status_code}")
|
||||
result["chunks_ingested"] = ingested
|
||||
if ingested > 0:
|
||||
evicted = await maybe_evict()
|
||||
if evicted:
|
||||
log.info(f"Evicted {evicted} vectors after upload")
|
||||
|
||||
if mode in ("context", "both"):
|
||||
expires = (datetime.now(timezone.utc) + timedelta(hours=UPLOAD_CONTEXT_EXPIRY_HOURS)).isoformat()
|
||||
|
||||
@@ -0,0 +1,409 @@
|
||||
import asyncio
|
||||
import os
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import app
|
||||
import config
|
||||
import db
|
||||
import rag
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
def make_client(tmp_path: Path) -> TestClient:
|
||||
os.environ["JARVISCHAT_ADMIN_PIN"] = "1234"
|
||||
db.DB_PATH = tmp_path / "jarvischat-rag-mgmt.db"
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
db.init_db()
|
||||
return TestClient(app.app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
def _admin_headers(client: TestClient) -> dict:
|
||||
login = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"})
|
||||
sid = login.json()["session_id"]
|
||||
return {"X-Session-ID": sid, "Origin": "http://testserver"}
|
||||
|
||||
|
||||
def _guest_headers(client: TestClient) -> dict:
|
||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
|
||||
return {"X-Session-ID": sid, "Origin": "http://testserver"}
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status, json_data=None):
|
||||
self.status_code = status
|
||||
self._json = json_data or {}
|
||||
|
||||
def json(self):
|
||||
return self._json
|
||||
|
||||
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, *a, **kw):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
pass
|
||||
|
||||
async def get(self, url, **kw):
|
||||
if "/collections/jarvis_rag" in url:
|
||||
return FakeResponse(200, {"result": {"vectors_count": 123}})
|
||||
return FakeResponse(200)
|
||||
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": []}})
|
||||
if "/points/delete" in url:
|
||||
return FakeResponse(200)
|
||||
return FakeResponse(200)
|
||||
|
||||
async def put(self, url, **kw):
|
||||
return FakeResponse(200)
|
||||
|
||||
|
||||
def _old_ts(hours_ago: float = 24) -> str:
|
||||
return (datetime.now(timezone.utc) - timedelta(hours=hours_ago)).isoformat()
|
||||
|
||||
|
||||
def _young_ts(hours_ago: float = 0.1) -> str:
|
||||
return (datetime.now(timezone.utc) - timedelta(hours=hours_ago)).isoformat()
|
||||
|
||||
|
||||
# ---------- get_collection_count ----------
|
||||
|
||||
def test_get_collection_count(monkeypatch):
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient())
|
||||
count = asyncio.run(rag.get_collection_count())
|
||||
assert count == 123
|
||||
|
||||
|
||||
# ---------- get_collection_stats ----------
|
||||
|
||||
def test_get_collection_stats_shape(monkeypatch):
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient())
|
||||
stats = asyncio.run(rag.get_collection_stats())
|
||||
assert stats["vector_count"] == 123
|
||||
assert stats["max_vectors"] == 50000
|
||||
assert stats["high_water_mark"] == 40000
|
||||
assert stats["low_water_mark"] == 10000
|
||||
assert stats["high_water_pct"] == 80
|
||||
assert stats["low_water_pct"] == 20
|
||||
assert 0 < stats["percent_full"] < 1
|
||||
assert "upload" in stats["pinned_sources"]
|
||||
assert "profile" in stats["pinned_sources"]
|
||||
|
||||
|
||||
# ---------- evict_batch ----------
|
||||
|
||||
def test_evict_batch_excludes_pinned_sources(monkeypatch):
|
||||
"""Pinned sources ('upload', 'profile') should be in the must_not scroll filter."""
|
||||
old = _old_ts(48)
|
||||
# Only non-pinned points are returned (real Qdrant would honour must_not filter)
|
||||
scroll_points = [
|
||||
{"id": "old-data", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}},
|
||||
]
|
||||
|
||||
class ScrollClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
must_not = kw.get("json", {}).get("filter", {}).get("must_not", [])
|
||||
pinned_values = [m["match"]["value"] for m in must_not]
|
||||
assert "upload" in pinned_values
|
||||
assert "profile" in pinned_values
|
||||
return FakeResponse(200, {"result": {"points": scroll_points}})
|
||||
if "/points/delete" in url:
|
||||
return FakeResponse(200)
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: ScrollClient())
|
||||
deleted = asyncio.run(rag.evict_batch(10))
|
||||
assert deleted == 1
|
||||
|
||||
|
||||
def test_evict_batch_respects_grace_period(monkeypatch):
|
||||
"""Vectors younger than RAG_GRACE_HOURS should be skipped."""
|
||||
old = _old_ts(48)
|
||||
young = _young_ts(0.1)
|
||||
scroll_points = [
|
||||
{"id": "mature", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}},
|
||||
{"id": "newborn", "payload": {"source": "terminal", "ingest_date": young, "retrieval_count": 0}},
|
||||
]
|
||||
|
||||
class GraceClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": scroll_points}})
|
||||
if "/points/delete" in url:
|
||||
deleted = kw.get("json", {}).get("points", [])
|
||||
assert "newborn" not in deleted
|
||||
return FakeResponse(200)
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: GraceClient())
|
||||
deleted = asyncio.run(rag.evict_batch(10))
|
||||
assert deleted == 1
|
||||
|
||||
|
||||
def test_evict_batch_respects_batch_size(monkeypatch):
|
||||
"""Only up to batch_size vectors should be deleted per call."""
|
||||
old = _old_ts(48)
|
||||
points = [{"id": f"p{i}", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}} for i in range(50)]
|
||||
|
||||
class BatchClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/delete" in url:
|
||||
deleted = kw.get("json", {}).get("points", [])
|
||||
assert len(deleted) == 10
|
||||
return FakeResponse(200)
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": points}})
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: BatchClient())
|
||||
deleted = asyncio.run(rag.evict_batch(10))
|
||||
assert deleted == 10
|
||||
|
||||
|
||||
def test_evict_batch_all_pinned_returns_zero(monkeypatch):
|
||||
"""If scroll returns nothing (all points filtered by must_not), evict_batch returns 0."""
|
||||
class EmptyClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": []}})
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: EmptyClient())
|
||||
deleted = asyncio.run(rag.evict_batch(10))
|
||||
assert deleted == 0
|
||||
|
||||
|
||||
def test_evict_batch_scores_lowest_first(monkeypatch):
|
||||
"""Vectors with lower scores should be evicted first."""
|
||||
old = _old_ts(48)
|
||||
points = [
|
||||
{"id": "high-score", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 100}},
|
||||
{"id": "low-score", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}},
|
||||
{"id": "mid-score", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 50}},
|
||||
]
|
||||
|
||||
class ScoreClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/delete" in url:
|
||||
deleted = kw.get("json", {}).get("points", [])
|
||||
assert "low-score" in deleted
|
||||
assert "high-score" not in deleted
|
||||
assert "mid-score" not in deleted
|
||||
return FakeResponse(200)
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": points}})
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: ScoreClient())
|
||||
deleted = asyncio.run(rag.evict_batch(1))
|
||||
assert deleted == 1
|
||||
|
||||
|
||||
# ---------- maybe_evict ----------
|
||||
|
||||
def test_maybe_evict_below_high_water(monkeypatch):
|
||||
"""When count is below high-water mark, eviction should not fire."""
|
||||
class LowCountClient(FakeAsyncClient):
|
||||
async def get(self, url, **kw):
|
||||
# 30000 < 40000 high water
|
||||
return FakeResponse(200, {"result": {"vectors_count": 30000}})
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: LowCountClient())
|
||||
rag.EVICTION_LOG.clear()
|
||||
evicted = asyncio.run(rag.maybe_evict())
|
||||
assert evicted == 0
|
||||
assert len(rag.EVICTION_LOG) == 0
|
||||
|
||||
|
||||
def test_maybe_evict_at_high_water(monkeypatch):
|
||||
"""When count reaches high-water mark, eviction should fire."""
|
||||
class HighCountClient(FakeAsyncClient):
|
||||
def __init__(self, *a, **kw):
|
||||
super().__init__()
|
||||
self.call_count = 0
|
||||
|
||||
async def get(self, url, **kw):
|
||||
return FakeResponse(200, {"result": {"vectors_count": 45000}})
|
||||
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
old = _old_ts(48)
|
||||
points = [{"id": f"evict-me-{i}", "payload": {"source": "terminal", "ingest_date": old, "retrieval_count": 0}} for i in range(100)]
|
||||
return FakeResponse(200, {"result": {"points": points}})
|
||||
if "/points/delete" in url:
|
||||
return FakeResponse(200)
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: HighCountClient())
|
||||
rag.EVICTION_LOG.clear()
|
||||
evicted = asyncio.run(rag.maybe_evict())
|
||||
assert evicted > 0
|
||||
assert len(rag.EVICTION_LOG) == 1
|
||||
entry = rag.EVICTION_LOG[0]
|
||||
assert "timestamp" in entry
|
||||
assert entry["count"] > 0
|
||||
|
||||
|
||||
def test_maybe_evict_zero_config_disabled(monkeypatch):
|
||||
"""RAG_MAX_VECTORS <= 0 should disable eviction."""
|
||||
orig = config.RAG_MAX_VECTORS
|
||||
try:
|
||||
config.RAG_MAX_VECTORS = 0
|
||||
evicted = asyncio.run(rag.maybe_evict())
|
||||
assert evicted == 0
|
||||
finally:
|
||||
config.RAG_MAX_VECTORS = orig
|
||||
|
||||
|
||||
def test_maybe_evict_all_pinned_breaks(monkeypatch):
|
||||
"""Above high water but only pinned points exist → eviction breaks with 0 deleted."""
|
||||
class AllPinnedClient(FakeAsyncClient):
|
||||
async def get(self, url, **kw):
|
||||
return FakeResponse(200, {"result": {"vectors_count": 45000}})
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": []}})
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: AllPinnedClient())
|
||||
rag.EVICTION_LOG.clear()
|
||||
evicted = asyncio.run(rag.maybe_evict())
|
||||
assert evicted == 0
|
||||
assert len(rag.EVICTION_LOG) == 0
|
||||
|
||||
|
||||
# ---------- get_rag_operational_stats ----------
|
||||
|
||||
def test_rag_operational_stats_shape(monkeypatch):
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient())
|
||||
rag.EVICTION_LOG.clear()
|
||||
rag.EVICTION_LOG.append({
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"count": 500,
|
||||
"remaining": 40000,
|
||||
})
|
||||
stats = asyncio.run(rag.get_rag_operational_stats())
|
||||
assert stats["vector_count"] == 123
|
||||
assert stats["grace_hours"] == 1
|
||||
assert "eviction_counts_last_1m" in stats
|
||||
assert "eviction_counts_last_5m" in stats
|
||||
assert "eviction_counts_last_30m" in stats
|
||||
assert stats["eviction_counts_last_1m"] == 500
|
||||
|
||||
|
||||
# ---------- GET /api/rag/stats ----------
|
||||
|
||||
def test_rag_stats_endpoint(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.get("/api/rag/stats", headers=_admin_headers(client))
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["vector_count"] == 123
|
||||
assert data["max_vectors"] == 50000
|
||||
assert "high_water_mark" in data
|
||||
assert "low_water_mark" in data
|
||||
assert data["high_water_pct"] == 80
|
||||
assert data["low_water_pct"] == 20
|
||||
assert "percent_full" in data
|
||||
assert data["pinned_sources"] == ["upload", "profile"]
|
||||
assert data["grace_hours"] == 1
|
||||
assert "eviction_counts_last_1m" in data
|
||||
assert "eviction_counts_last_5m" in data
|
||||
assert "eviction_counts_last_30m" in data
|
||||
assert "pinned_count" in data
|
||||
assert "avg_retrieval_count" in data
|
||||
assert "at_risk_count" in data
|
||||
assert "eviction_log_size" in data
|
||||
|
||||
|
||||
def test_rag_stats_requires_admin(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.get("/api/rag/stats", headers=_guest_headers(client))
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ---------- POST /api/rag/flush ----------
|
||||
|
||||
def test_rag_flush_endpoint(tmp_path, monkeypatch):
|
||||
class FlushClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": [{"id": "a"}, {"id": "b"}]}})
|
||||
if "/points/delete" in url:
|
||||
return FakeResponse(200)
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FlushClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post("/api/rag/flush", headers=_admin_headers(client))
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["status"] == "flushed"
|
||||
assert data["deleted_count"] == 2
|
||||
assert data["collection"] == rag.RAG_COLLECTION
|
||||
|
||||
|
||||
def test_rag_flush_requires_admin(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: FakeAsyncClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post("/api/rag/flush", headers=_guest_headers(client))
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
def test_rag_flush_empty_collection(tmp_path, monkeypatch):
|
||||
class EmptyClient(FakeAsyncClient):
|
||||
async def post(self, url, **kw):
|
||||
if "/points/scroll" in url:
|
||||
return FakeResponse(200, {"result": {"points": []}})
|
||||
return FakeResponse(200)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: EmptyClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post("/api/rag/flush", headers=_admin_headers(client))
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["deleted_count"] == 0
|
||||
assert data["status"] == "flushed"
|
||||
|
||||
|
||||
# ---------- Race lock ----------
|
||||
|
||||
def test_eviction_lock_prevents_concurrent_eviction(monkeypatch):
|
||||
"""Concurrent calls to maybe_evict should queue; only one evicts."""
|
||||
call_order = []
|
||||
|
||||
async def slow_get_collection_count():
|
||||
call_order.append("count")
|
||||
return 45000
|
||||
|
||||
async def slow_evict_batch(bs):
|
||||
call_order.append("evict")
|
||||
await asyncio.sleep(0.05)
|
||||
return 500
|
||||
|
||||
monkeypatch.setattr(rag, "get_collection_count", slow_get_collection_count)
|
||||
monkeypatch.setattr(rag, "evict_batch", slow_evict_batch)
|
||||
rag.EVICTION_LOG.clear()
|
||||
|
||||
async def run_concurrent():
|
||||
r1, r2 = await asyncio.gather(rag.maybe_evict(), rag.maybe_evict())
|
||||
return r1, r2
|
||||
|
||||
r1, r2 = asyncio.run(run_concurrent())
|
||||
# First call evicted, second found count already below high water or lock serialized
|
||||
assert r1 >= 0
|
||||
assert r2 >= 0
|
||||
Reference in New Issue
Block a user