atrium-rag 0.1.0__tar.gz

atrium_rag-0.1.0/.env.docker

@@ -0,0 +1,68 @@
+# Configuration used by `docker compose`.
+#
+# This file is the container-side counterpart to your local `.env`:
+# the hostnames differ (a container reaches services by service name,
+# and the host machine by `host.docker.internal`), everything else is
+# the same.
+#
+# `.env` on your machine is for running atrium directly (laptop, CI).
+# This file is for `docker compose up`. Keep API_KEY empty here unless
+# you are putting the stack behind a network you control.
+
+# --- Vector database ------------------------------------------------------
+# `qdrant` is the service name in docker-compose.yml. Use `pgvector`
+# with the pgvector profile instead.
+VECTOR_DB=qdrant
+QDRANT_HOST=qdrant
+QDRANT_PORT=6333
+PG_CONN=postgresql://user:pass@pgvector:5432/db
+
+# --- Embedding provider ---------------------------------------------------
+# Ollama runs on the host, not in the compose network.
+# `host.docker.internal` resolves on Docker Desktop (macOS + Windows).
+# On Linux, add `extra_hosts: ["host.docker.internal:host-gateway"]`
+# to the api service, or point this at a reachable Ollama host.
+EMBED_BASE_URL=http://host.docker.internal:11434
+EMBED_MODEL=bge-large
+EMBED_DIM=1024
+
+# --- Optional features ----------------------------------------------------
+OCR_ENABLED=0
+OCR_PROVIDER=tesseract
+OCR_LANG=eng
+TESSERACT_CMD=
+
+RERANK_ENABLED=0
+RERANK_PROVIDER=bge
+RERANK_MODEL=
+RERANK_BACKEND=
+RERANK_DEVICE=
+RERANK_MAX_LENGTH=512
+FLASHRANK_CACHE_DIR=
+COHERE_API_KEY=
+
+VISION_ENABLED=0
+VISION_MODEL=minicpm-v4.6:1b
+VISION_TIMEOUT_SECONDS=180
+
+# --- Retrieval / server ---------------------------------------------------
+DEFAULT_TOP_K=8
+DENSE_CONFIDENT_GAP=0.05
+DENSE_MIN_SCORE=0.55
+LOG_LEVEL=INFO
+# The job queue lives on the shared `atrium-data` volume so the API
+# and every worker see the same file.
+JOB_DB_PATH=/data/jobs.db
+JOB_STALE_SECONDS=1800
+INGEST_CONCURRENCY=4
+WORKER_POLL_BASE=1.0
+WORKER_POLL_MAX=10.0
+WORKER_JOB_TIMEOUT=0
+RETRIEVAL_LAB_MAX_SESSIONS=16
+
+# --- Limits / security ----------------------------------------------------
+MAX_UPLOAD_BYTES=104857600
+ALLOWED_MIME_TYPES=
+ALLOW_CHUNKED_UPLOADS=0
+API_KEY=
+CORS_ALLOW_ORIGINS=

atrium_rag-0.1.0/.env.example

@@ -0,0 +1,169 @@
+# atrium configuration template.
+#
+#   cp .env.example .env      # then edit
+#   atrium init                # or let the wizard write it for you
+#
+# Every variable here is read by atrium/main.py, atrium/worker.py and
+# the CLI. Lines starting with # are ignored; everything after a
+# value's whitespace-# is a comment.
+#
+# Precedence: a real environment variable beats this file, and this
+# file beats the built-in default. So `API_KEY=... uvicorn
+# atrium.main:app` wins over whatever is written here.
+
+# --- Vector database ------------------------------------------------------
+# Choose one. Both backends implement the same retrieval protocol.
+#   qdrant    — Qdrant
+#   pgvector  — Postgres + pgvector
+VECTOR_DB=sqlite
+
+# Qdrant. Ignored unless VECTOR_DB=qdrant.
+# Inside Docker Compose the hostname is the service name (`qdrant`);
+# on a laptop it is `localhost`.
+QDRANT_HOST=localhost
+QDRANT_PORT=6333
+
+# Postgres connection string. Ignored unless VECTOR_DB=pgvector.
+# Format: postgresql://user:password@host:port/database
+PG_CONN=postgresql://user:pass@pgvector:5432/db
+
+# --- Embedding provider ---------------------------------------------------
+# Any server that speaks the OpenAI /v1/embeddings API: Ollama, OpenAI,
+# vLLM, LM Studio, TEI, Together. The base URL includes the /v1.
+EMBED_BASE_URL=http://localhost:11434/v1
+
+# The key for that endpoint. Blank for local servers; required by
+# hosted providers.
+EMBED_API_KEY=
+
+# "http" for any OpenAI-compatible endpoint, "fastembed" for the
+# in-process ONNX embedder from the `local` extra (no server, no key).
+EMBED_PROVIDER=http
+
+# Embedding model + dimension. Defaults match the published 19-set
+# benchmark headline (bge-large / 1024). Override with care — promote
+# embedder changes only after re-running the benchmark; see README.
+EMBED_MODEL=bge-large
+EMBED_DIM=1024
+
+# --- Limits ---------------------------------------------------------------
+# Ceilings on what one request and one document may cost. The upload
+# cap bounds the bytes; these bound what those bytes turn into.
+MAX_CHUNKS_PER_DOCUMENT=50000
+MAX_PAGES_PER_DOCUMENT=10000
+MAX_FILES_PER_REQUEST=32
+# The preview endpoint returns the document's structure, not the corpus.
+MAX_PREVIEW_CHUNKS=500
+# The retrieval lab holds a document's chunks and embeddings in memory.
+MAX_LAB_CHUNKS_PER_SESSION=5000
+MAX_LAB_TOTAL_CHUNKS=50000
+
+# --- Worker ---------------------------------------------------------------
+# How many times a job may be attempted before it is given up on.
+MAX_JOB_ATTEMPTS=3
+# How long a finished job row is kept (its payload is dropped at once).
+JOB_RETENTION_SECONDS=86400
+
+# --- Where the local files live -------------------------------------------
+# The API and the worker must agree on this, or they use different job
+# stores and different corpora. Relative paths resolve against it.
+ATRIUM_DATA_DIR=
+
+# --- The address `atrium serve` binds -------------------------------------
+# Loopback by default. Set 0.0.0.0 to serve beyond this machine.
+ATRIUM_BIND=127.0.0.1
+
+# --- Local vector store ----------------------------------------------------
+# The SQLite file used when VECTOR_DB=sqlite. Nothing else is needed:
+# no server, no Docker.
+LOCAL_DB_PATH=atrium_local.db
+
+# --- Vision (optional) ----------------------------------------------------
+# VISION_ENABLED=1 captions image-only pages and diagrams with a
+# vision model during parsing. Costs a lot of time on document-heavy
+# corpora; off by default.
+VISION_ENABLED=0
+VISION_MODEL=minicpm-v4.6:1b
+VISION_TIMEOUT_SECONDS=180
+
+# --- OCR (optional) -------------------------------------------------------
+# OCR_ENABLED=1 turns on OCR for scanned pages and image-only docs.
+# Requires `pip install atrium[ocr]` and (for Tesseract) the
+# Tesseract binary on PATH — see README.
+OCR_ENABLED=0
+OCR_PROVIDER=tesseract   # or rapidocr
+OCR_LANG=eng
+# Path to the Tesseract binary. Empty = look it up on PATH.
+TESSERACT_CMD=
+
+# --- Reranking (optional) -------------------------------------------------
+# RERANK_ENABLED=1 turns on cross-encoder reranking for /search.
+# Requires `pip install atrium[rerank]` (or [flashrank]).
+RERANK_ENABLED=0
+RERANK_PROVIDER=bge   # bge | flashrank | cohere
+RERANK_MODEL=
+RERANK_DEVICE=        # bge only: cpu | cuda | mps
+RERANK_BACKEND=       # cohere | flashrank | bge-local
+RERANK_MAX_LENGTH=512
+FLASHRANK_CACHE_DIR=
+COHERE_API_KEY=
+
+# --- Retrieval ------------------------------------------------------------
+# Default top-k for /search and `atrium search`. The v0.0.1 headline
+# uses 8 — the top-k sweep on the 19-fixture benchmark showed k=8
+# captures 5 more queries than k=5 for 1.6x the scoring cost.
+DEFAULT_TOP_K=8
+
+# Thresholds for the dense-confidence gate. The gate decides when the
+# optional BM25 sparse ranker is consulted on top of the dense one.
+# Raising either one makes the gate stricter (more BM25 consulted).
+DENSE_CONFIDENT_GAP=0.05
+DENSE_MIN_SCORE=0.55
+
+# --- Server ---------------------------------------------------------------
+# LOG_LEVEL=DEBUG|INFO|WARNING|ERROR
+LOG_LEVEL=INFO
+
+# Where the async-job queue lives. A single SQLite file, shared by the
+# API and every worker. To run workers on other hosts, put this file on
+# a shared volume (NFS, a mounted PVC) — there is no Postgres job store.
+JOB_DB_PATH=jobs.db
+
+# How long a job may sit in `running` before a worker hands it back to
+# the queue (seconds). Set it above the slowest document you expect:
+# a 500-page PDF can take minutes, and reclaiming it early means two
+# workers ingest the same document.
+JOB_STALE_SECONDS=1800
+
+# How many documents may be ingested at once.
+INGEST_CONCURRENCY=4
+
+# Worker poll backoff (seconds) and per-job timeout. 0 = no timeout.
+WORKER_POLL_BASE=1.0
+WORKER_POLL_MAX=10.0
+WORKER_JOB_TIMEOUT=0
+
+# How many retrieval-lab sessions to keep in RAM.
+RETRIEVAL_LAB_MAX_SESSIONS=16
+
+# --- Limits ---------------------------------------------------------------
+# Maximum upload size in bytes (default 100 MB). Uploads larger than
+# this are rejected with 413 before being read into memory.
+MAX_UPLOAD_BYTES=104857600
+
+# Allowed MIME types for upload (comma-separated). Empty = allow any
+# detected by the built-in format detector. Set explicitly in prod for
+# tighter validation.
+ALLOWED_MIME_TYPES=
+
+# Accept uploads without a Content-Length header.
+ALLOW_CHUNKED_UPLOADS=0
+
+# --- API auth (optional) --------------------------------------------------
+# Set API_KEY to require `Authorization: Bearer <key>` on all routes
+# except /health. Empty = no auth (dev only).
+API_KEY=
+
+# --- CORS -----------------------------------------------------------------
+# Comma-separated origin allowlist. Empty = same-origin only.
+CORS_ALLOW_ORIGINS=

atrium_rag-0.1.0/Dockerfile

@@ -0,0 +1,67 @@
+# Atrium — container image.
+#
+# The image installs from pyproject.toml, not requirements.txt: the
+# distribution metadata is the source of truth for what "atrium" is,
+# and requirements.txt is a dev convenience that had already drifted
+# from it (pinned floors vs. ranges).
+#
+#   docker build -t atrium-rag .
+#   docker build -t atrium-rag --build-arg EXTRAS="ocr,rerank" .
+#
+# The runtime image needs libmagic (python-magic's native library) and
+# nothing else from the OS. Tesseract is deliberately NOT installed:
+# it is a host-level binary, the `ocr` extra is opt-in, and shipping it
+# would add ~30 MB to every image for a feature most deployments do
+# not use. If you want OCR in the container, either use the
+# `rapidocr` provider (pure pip) or build a derived image that adds
+# `apt-get install tesseract-ocr`.
+
+FROM python:3.12-slim AS base
+
+# libmagic: python-magic's native library. Without it the MIME
+# detector falls back to the extension table, which is fine but
+# weaker.
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends libmagic1 \
+    && rm -rf /var/lib/apt/lists/*
+
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1
+
+WORKDIR /app
+
+# Optional extras, e.g. --build-arg EXTRAS="ocr,rerank".
+ARG EXTRAS=""
+
+# Dependencies first: this layer is cached until pyproject.toml
+# changes, so a code-only edit doesn't reinstall the world.
+COPY pyproject.toml README.md LICENSE MANIFEST.in ./
+COPY atrium ./atrium
+RUN if [ -n "$EXTRAS" ]; then \
+        pip install --no-cache-dir ".[$EXTRAS]"; \
+    else \
+        pip install --no-cache-dir .; \
+    fi
+
+# Non-root. The API writes one file (the SQLite job queue) and needs
+# a writable HOME for model caches.
+RUN useradd --create-home --uid 10001 atrium \
+    && mkdir -p /data \
+    && chown -R atrium:atrium /app /data
+USER atrium
+
+# The job queue lives here. Mount a volume to keep job history across
+# container restarts, and to share the queue with a worker container.
+ENV JOB_DB_PATH=/data/jobs.db \
+    OLLAMA_BASE_URL=http://host.docker.internal:11434 \
+    QDRANT_HOST=qdrant
+
+VOLUME ["/data"]
+EXPOSE 8000
+
+HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \
+    CMD ["python", "-m", "atrium.healthcheck"]
+
+CMD ["uvicorn", "atrium.main:app", "--host", "0.0.0.0", "--port", "8000"]

atrium_rag-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Atrium contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

atrium_rag-0.1.0/MANIFEST.in

@@ -0,0 +1,42 @@
+# Files to include in the source distribution (sdist) but not
+# covered by setuptools' defaults.
+
+# README + LICENSE.
+include README.md
+include LICENSE
+
+# Operational files referenced in the docs.
+include Dockerfile
+include docker-compose.yml
+include Makefile
+include .env.example
+include .env.docker
+
+# Type marker so mypy/PEP 561 understands the package.
+include atrium/py.typed
+
+# Exclude dev/test/build artifacts from sdist.
+exclude .github
+prune tests
+exclude scripts
+exclude docs
+exclude .venv
+exclude venv
+exclude __pycache__
+exclude .pytest_cache
+exclude .mypy_cache
+exclude .ruff_cache
+exclude .git
+exclude .gitignore
+exclude .dockerignore
+exclude .DS_Store
+exclude *.pyc
+exclude *.pyo
+recursive-exclude * __pycache__
+recursive-exclude * *.pyc
+
+# The README's images travel with it: the PyPI page renders the README
+# from the sdist, and a relative path to a file that is not there is a
+# broken image on the front page of the project.
+include docs/images/*.png
+include docs/images/*.gif

atrium_rag-0.1.0/Makefile

@@ -0,0 +1,59 @@
+# Development tasks for atrium.
+#
+#   make install     install the package + dev extras
+#   make test        run the unit suite
+#   make lint        ruff
+#   make typecheck   mypy (advisory — see pyproject.toml)
+#   make build       build the wheel + sdist
+#   make shot        regenerate the README benchmark image
+#   make docker      build the container image
+
+PYTHON ?= python3
+VENV   ?= venv
+PIP    := $(VENV)/bin/pip
+PYTEST := $(VENV)/bin/pytest
+RUFF   := $(VENV)/bin/ruff
+MYPY   := $(VENV)/bin/mypy
+
+.PHONY: install test lint format typecheck build shot smoke docker clean
+
+install:
+	$(PIP) install -e ".[dev]"
+
+test:
+	$(PYTEST) tests/ -q
+
+lint:
+	$(RUFF) check atrium tests scripts
+
+format:
+	$(RUFF) check atrium tests scripts --fix
+
+typecheck:
+	$(MYPY) atrium
+
+build:
+	$(VENV)/bin/python -m build
+	$(VENV)/bin/twine check dist/*
+
+# Install the built wheel into a throwaway environment and run the
+# smoke test against it — the artifact a user gets, not the source
+# tree. This is what CI's "Install the built artifact" job does.
+smoke:
+	$(VENV)/bin/python -m build --wheel
+	rm -rf /tmp/atrium-smoke
+	$(PYTHON) -m venv /tmp/atrium-smoke
+	/tmp/atrium-smoke/bin/pip install --quiet dist/*.whl
+	cd /tmp && /tmp/atrium-smoke/bin/python $(CURDIR)/scripts/smoke_test.py
+
+# Regenerate docs/images/benchmark.png from the published run's raw
+# JSON. The README image must stay in step with the data behind it.
+shot:
+	$(VENV)/bin/python scripts/render_benchmark_shot.py
+
+docker:
+	docker build -t atrium-rag .
+
+clean:
+	rm -rf build dist *.egg-info
+	find . -name '__pycache__' -type d -prune -exec rm -rf {} +