Compare commits
23
Commits
v0.13.0
...
settings-e
No files matched your search
@@ -39,6 +39,14 @@ jobs:
|
||||
- name: Build
|
||||
run: go build ./...
|
||||
|
||||
# llmgw-kb runs on people's own machines, Windows included, which no
|
||||
# test here runs on: building it for each catches a Unix-only call.
|
||||
- name: Build llmgw-kb for every system
|
||||
run: |
|
||||
for target in windows/amd64 windows/arm64 darwin/amd64 darwin/arm64 linux/amd64 linux/arm64; do
|
||||
CGO_ENABLED=0 GOOS="${target%/*}" GOARCH="${target#*/}" go build -o /dev/null ./cmd/llmgw-kb
|
||||
done
|
||||
|
||||
# The Terraform provider is a module of its own, which ./... above does
|
||||
# not reach. Its tests check its client against docs/admin-api.openapi.json,
|
||||
# so a change to the admin API that breaks the provider fails here. The
|
||||
|
||||
@@ -55,6 +55,21 @@ jobs:
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# llmgw-kb, which people run on their own machines to put their
|
||||
# documents into a knowledge base: Windows too, which llmgw itself
|
||||
# does not build for.
|
||||
for target in windows/amd64 windows/arm64 darwin/amd64 darwin/arm64 linux/amd64 linux/arm64; do
|
||||
os="${target%/*}"; arch="${target#*/}"
|
||||
ext=""; [ "$os" = windows ] && ext=".exe"
|
||||
name="llmgw-kb-${version}-${os}-${arch}${ext}"
|
||||
echo "building $name"
|
||||
CGO_ENABLED=0 GOOS="$os" GOARCH="$arch" go build \
|
||||
-trimpath \
|
||||
-ldflags "-s -w -X gitea.mynodes.xyz/ilker/llmgw/internal/version.Version=${version}" \
|
||||
-o "dist/${name}" ./cmd/llmgw-kb
|
||||
(cd dist && sha256sum "${name}" >> checksums.txt)
|
||||
done
|
||||
|
||||
# The reverse proxy configurations, so a release page carries
|
||||
# everything install.md asks for and the repository is not needed.
|
||||
for f in deploy/proxy/Caddyfile deploy/proxy/nginx.conf; do
|
||||
@@ -91,6 +106,7 @@ jobs:
|
||||
# test the cross-compiled ones, which is why the native build is
|
||||
# checked rather than assumed.
|
||||
./dist/llmgw-${version}-linux-amd64 --version
|
||||
./dist/llmgw-kb-${version}-linux-amd64 --version
|
||||
cat dist/checksums.txt
|
||||
|
||||
# Signed, so a machine that trusts the key built into it installs
|
||||
|
||||
+161
-35
@@ -807,14 +807,34 @@ Patroni configuration. The firewall's allowed set is the cluster's membership
|
||||
removed changes every machine's rules. Retires `firewall.Admit` and
|
||||
`firewall_networks` in server.yaml (imported into the store).
|
||||
|
||||
### E3. agent.yaml down to bootstrap (M)
|
||||
### E3. agent.yaml down to bootstrap (M) — done 2026-10-08 (branch `settings-e`)
|
||||
agent.yaml keeps the control planes' addresses, its certificate and its id.
|
||||
Labels, fabrics, engines, heartbeat and supervisor settings move to the
|
||||
store per node and reach the agent over its stream; a change to the file is
|
||||
noticed and the agent restarts itself.
|
||||
Built as a section of the cluster's settings rather than a store of its own,
|
||||
so history, audit, `ctl config`, the console and the portal page cover it:
|
||||
`agents.default` and `agents.nodes.<id>` (config.AgentSettings; bootstrap
|
||||
keys node_id/server(s)/hostname/mgmt_addr/tls refused; each node's merge
|
||||
loaded as an agent would, so a bad fabric is refused on save). A live key:
|
||||
the control plane sends ControlMessage.AgentSettings on register and, a
|
||||
second apart per node, on change; the agent compares with what it started on,
|
||||
checks it loads beside its agent.yaml (else logs and keeps running), keeps it
|
||||
in /var/lib/llmgw-agent/agent-settings.yaml and restarts (engines adopted).
|
||||
At start: agent.yaml < kept settings < LLMGW_ env; kept settings that no
|
||||
longer load are set aside, never a reason not to start. agent.yaml is
|
||||
watched every 10 s and restarted on when it changes and loads. Left: file
|
||||
keys are not imported into the settings (they keep working under them);
|
||||
`node add` still writes a full agent.yaml; labels have two homes now
|
||||
(`ctl label`, which wins, and agents.*.labels).
|
||||
|
||||
### E4. Web settings page (S)
|
||||
### E4. Web settings page (S) — done 2026-10-08 (branch `settings-e`)
|
||||
The same document, edited on /admin/settings, for whoever has no terminal.
|
||||
Built: administrators only (staff roles get 404: webhook secrets), checked by
|
||||
the same function as PUT /api/v1/settings, made from a base version or
|
||||
refused on a newer one (the edit kept in the editor), history with view and
|
||||
roll back as a new version, audit with the size not the body, English and
|
||||
Turkish, a Settings link in the menu.
|
||||
|
||||
---
|
||||
|
||||
@@ -829,47 +849,134 @@ ACL yet); a personal key or session names a project or team and gets that
|
||||
chain only, never the union of its person's memberships; PostgreSQL store
|
||||
only, off on SQLite with a sentence saying why.
|
||||
|
||||
### F1. Store: collections, grants, chunks (M)
|
||||
- `pgvector` ≥ 0.8 installed on every Patroni member by `server init`
|
||||
(`internal/cli/serverinit_postgres.go`), `CREATE EXTENSION` in the
|
||||
migration; the SQLite migration has no tables, and the store reports the
|
||||
feature unavailable.
|
||||
- Tables: `kb_collections` (id, org_id, owner `team:…`/`project:…`/`user:…`,
|
||||
name, embed_model, dims, created_at), `kb_grants` (scope, scope_id,
|
||||
collection, perm read|write — like `guard_attachments`), `kb_documents`
|
||||
(name, digest, bytes, state, chunk count), `kb_chunks` (org_id,
|
||||
collection, doc_id, ord, body, tsv, embedding) `PARTITION BY LIST
|
||||
(org_id)`, a partition per organization, HNSW and GIN indexes per
|
||||
partition.
|
||||
- Row-level security on `kb_chunks` against `current_setting('llmgw.org')`
|
||||
and `llmgw.collections`, retrieval through its own role.
|
||||
- Rules, refused before anything is written: a grant never crosses
|
||||
organizations; owner and grants resolve to existing scopes; deleting a
|
||||
project/team/org and `users forget` delete their collections (`forget.go`).
|
||||
- Done: store tests on PostgreSQL, including an RLS test that a query
|
||||
without the settings returns no rows.
|
||||
### F1. Store: collections, grants, chunks (M) — done 2026-10-08 (branch `rag`)
|
||||
Built (`internal/store/knowledge.go`, migration 0035 PostgreSQL only):
|
||||
- `kb_collections` (owner scope/id, org, embed_model, dims), `kb_grants`
|
||||
(read|write, any scope, never across organizations; a key needs a
|
||||
project; a person needs a membership in the organization),
|
||||
`kb_documents` (digest: the same again is no change, a new one replaces
|
||||
in one transaction).
|
||||
- `kb_chunks` is created when the store opens, not in a migration (a
|
||||
migration needing an extension would stop the control plane; one that
|
||||
skipped it would never run again): partitioned by organization (hashed
|
||||
partition name), `tsv` generated (`simple`), GIN on the parent, one
|
||||
partial HNSW index per collection on its organization's partition
|
||||
(`embedding::vector(dims)`, `WHERE collection = '<id>'`). EXPLAIN on
|
||||
20k rows: partition pruned, that index used, policy applied.
|
||||
- RLS FORCEd on `kb_chunks` against `llmgw.kb_org`/`llmgw.kb_collections`
|
||||
(set with set_config per transaction), one role — the table owner — not
|
||||
a separate one. Chunks are deleted by the foreign-key cascade from their
|
||||
document, which RLS does not block.
|
||||
- `CollectionsFor(Effective, write)`: owner or grant anywhere on the chain;
|
||||
no organization (personal key without a project) reads nothing.
|
||||
- `SearchChunks`: per collection nearest-neighbour (literal id and cast so
|
||||
the partial index is used) plus a full-text arm with any word
|
||||
(`a | b`, lowercased by the database, not Go), fused by reciprocal rank.
|
||||
- Deleting a project/team/org that owns collections is refused like one
|
||||
with members (changed from "deleted with them": every other delete in the
|
||||
tree refuses); DeleteUser/ForgetUser delete the person's collections;
|
||||
deleting an organization drops its partition; grants to a deleted scope go.
|
||||
- `server init` installs `postgresql-N-pgvector` / `pgvector_N` and enables
|
||||
it; pgha installs it on every member at setup and enables it on the
|
||||
leader on every run. Neither stops on failure.
|
||||
Left: a cluster set up before (the PoC) gets no package from a rerun of
|
||||
pgha on a running cluster — install by hand or add an upgrade step;
|
||||
`server restore`/`CopyTo` do not carry kb tables (SQLite → PostgreSQL only,
|
||||
and SQLite has none); tests need `LLMGW_TEST_POSTGRES` pointing at a
|
||||
pgvector database as a non-superuser (a superuser bypasses RLS).
|
||||
|
||||
### F2. Ingest (M)
|
||||
- `llmgw ctl kb create|list|delete|grant|revoke|docs`, `ctl kb add
|
||||
<collection> <files…>`, admin API under `/api/v1/kb`, audit events.
|
||||
- Text from PDF (text layer), DOCX, HTML, Markdown, plain text — in Go, no
|
||||
Python. Split on headings/paragraphs, ~500 tokens with overlap. Embedded
|
||||
through the router at the bronze band, accounted to the collection's
|
||||
owner. One transaction per document; same digest skipped, changed digest
|
||||
replaces. OCR out of scope.
|
||||
- Console: a Knowledge screen (collections, documents, grants; add, grant,
|
||||
revoke, delete); `docs/console.md`, `docs/knowledge.md`.
|
||||
### F2. Ingest on the user's machine (L) — done 2026-10-08 (branch `rag`)
|
||||
Uploading whole files to the control plane was dropped: a PDF is 20–100×
|
||||
its text, and the work belongs where the files are. Instead a small client,
|
||||
`llmgw-kb`, run by the person on their own machine with their own key:
|
||||
- **Its own binary** (`cmd/llmgw-kb`), for Windows, macOS and Linux on amd64
|
||||
and arm64. `llmgw` itself does not build on Windows (agent syscalls), and
|
||||
employees' laptops should not carry the server, agent and setup screens. A
|
||||
test keeps server packages out of its imports.
|
||||
- `llmgw-kb sync <collection> <dir|file…> [--include/--exclude globs]
|
||||
[--dry-run] [--delete]`: scan → local catalog (path, size, mtime, sha256;
|
||||
cache under os.UserCacheDir, digests not recomputed for unchanged files)
|
||||
→ the collection's manifest from the gateway (name → digest only) → diff
|
||||
→ for new/changed files only: extract, chunk, embed through the gateway's
|
||||
`/v1/embeddings` with the person's key (its accounting, limits, guard and
|
||||
band apply), and send name, digest, chunk texts and vectors; `--delete`
|
||||
removes documents whose files are gone. `--dry-run` lists what would leave
|
||||
the machine and how much. A file that cannot be read (locked, unsupported,
|
||||
no text) is skipped with the reason; the run goes on.
|
||||
- **Formats, as many as possible.** Built in, pure Go: PDF (text layer),
|
||||
DOCX/XLSX/PPTX, ODT/ODS/ODP, RTF, EPUB, HTML/XHTML, EML (text and HTML
|
||||
parts), Jupyter notebooks, subtitles (SRT/VTT), CSV/TSV, JSON/YAML/XML,
|
||||
Markdown/RST/AsciiDoc/Org, plain text and source code; ZIP archives
|
||||
recursed into. With a helper found on the machine: legacy DOC/XLS/PPT, RTF
|
||||
and anything else LibreOffice opens (`soffice --headless --convert-to`),
|
||||
images and scanned PDFs through Tesseract (`tur+eng`). Text in UTF-8,
|
||||
UTF-16 (BOM) or Windows-1254/1252 is converted to UTF-8.
|
||||
- **Across systems:** document names relative to the synced root, `/`
|
||||
separated, Unicode NFC (macOS hands out NFD), a warning for two names that
|
||||
differ only in case; CRLF → LF; key from `LLMGW_API_KEY` or `llmgw-kb
|
||||
login` (os.UserConfigDir, 0600 where that means something); the system's
|
||||
CA store, `--ca` for another.
|
||||
- **Gateway endpoints** (client listener, person's key):
|
||||
`GET /v1/kb/collections` (what this key may read/write),
|
||||
`GET /v1/kb/collections/{id}/manifest`, `PUT
|
||||
/v1/kb/collections/{id}/documents` (one document per request: name,
|
||||
digest, bytes, chunks with text and embedding; dims checked, chunk count
|
||||
and size bounded), `DELETE …/documents?name=`. Write needs a write grant
|
||||
or ownership on the key's chain (`CollectionsFor(…, write)`); otherwise
|
||||
404.
|
||||
- **Administration** stays with operators: `llmgw ctl kb
|
||||
create|list|delete|grant|revoke|docs`, admin API `/api/v1/kb`, audit,
|
||||
console screen; `docs/knowledge.md`.
|
||||
- Trade-off accepted: vectors come from the client, so a writer could store
|
||||
vectors that do not match their text — only in collections they may
|
||||
already write.
|
||||
Built: `internal/kb/extract` (formats above; helpers found on PATH and in
|
||||
each system's install places), `internal/kb/chunk` (≈1800 characters under
|
||||
the heading path, a paragraph of overlap, sections kept apart once a chunk
|
||||
is a third full), `internal/kb/kbsync` (ledger in SQLite at
|
||||
os.UserConfigDir/llmgw-kb/ledger.db: documents keyed by gateway, collection
|
||||
and name with file, size, mtime, digest, state, reason; runs), `cmd/llmgw-kb`
|
||||
(login, collections, sync, status, history, formats), `internal/router/kb.go`
|
||||
(the four gateway endpoints, `X-Llmgw-Project` for personal keys, audit),
|
||||
`internal/control/adminapi/kb.go` + `llmgw ctl kb` + console Knowledge
|
||||
screen, `docs/knowledge.md`; CI builds the client for all six targets and
|
||||
releases attach it. Tested: every format from generated fixtures and a
|
||||
Turkish PDF (OCR path with Tesseract here), sync against a fake gateway and
|
||||
end to end against the router on pgvector.
|
||||
Tried on the PoC 2026-10-08 (control planes 140 and 150 on a dev build of
|
||||
this branch; their previous binary is /usr/local/bin/llmgw.pre-rag): Ubuntu's
|
||||
pgvector 0.6.0 installed on 149/150 and enabled; bge-m3's weights were never
|
||||
in the mirror (fetched from CompendiumLabs/bge-m3-gguf@1864f47f) and it is
|
||||
held at 1 replica; tree Demo/ik/belgeler and Demo/satis/crm with a key each;
|
||||
collection Demo/politikalar filled from a Mac through an SSH tunnel with ten
|
||||
Turkish files in nine formats (9 sent in 5 s, the .doc skipped for want of
|
||||
LibreOffice; a rerun sent nothing in 0 s; a changed file alone was sent).
|
||||
Nine Turkish questions each found their page first once search scored by
|
||||
similarity plus word coverage (rank fusion had got two wrong). The other
|
||||
team's key saw no collection and got 404 by id; with a read grant it read
|
||||
and could not write; revoked, nothing again. On the live store the llmgw
|
||||
role (not superuser, no BYPASSRLS) sees 0 chunks with no scope set and 0
|
||||
with the organization alone.
|
||||
Left: the gateway's collection list shows the owner as team:ID, not its
|
||||
path; no LibreOffice on the dev machine, so the legacy
|
||||
Office path is untested beyond its error; the client is unsigned (Gatekeeper
|
||||
and SmartScreen warn); Outlook `.msg` is not read; one document over ~60 MiB
|
||||
of JSON (a few thousand pages) is refused rather than split.
|
||||
PostgreSQL tests of store, router and adminapi share LLMGW_TEST_POSTGRES and
|
||||
the store's TRUNCATE each test; run them with `go test -p 1` or one package
|
||||
at a time, or a router test loses its rows mid-run.
|
||||
|
||||
### F3. Retrieval on the request path (M)
|
||||
### F3. Retrieval on the request path (M) — done 2026-10-08 (branch `rag`)
|
||||
- `llmgw.knowledge: [collections]` in the body (stripped before the
|
||||
engine), or default collections per project/key.
|
||||
- Allowed set from the caller's chain (§7.9.2), cached like limits; grants
|
||||
add up down the tree. Personal keys and sessions need `x-llmgw-project`
|
||||
(or the playground's picker) naming a scope the person is a member of.
|
||||
- A collection not allowed answers 404, exactly as a missing one.
|
||||
- Hybrid search with the filter inside the query (`hnsw.iterative_scan`,
|
||||
`WHERE org_id = $1 AND collection = ANY($2)`), full-text arm merged by
|
||||
reciprocal rank, `/v1/rerank` when a reranker is in the catalog. Top
|
||||
- Hybrid search with the filter inside the query (per collection, its own
|
||||
partial index; pgvector 0.6 is enough —
|
||||
scored by similarity plus the share of the question's words held (built
|
||||
in F1/F2 as SearchChunks), `/v1/rerank` when a reranker is in the catalog. Top
|
||||
chunks into a system message with document and position;
|
||||
`x-llmgw-sources` lists them.
|
||||
- Leaks closed: answer-cache key includes the allowed set; per-organization
|
||||
@@ -879,6 +986,25 @@ only, off on SQLite with a sentence saying why.
|
||||
(including via cache), that a personal key without a project is refused,
|
||||
metrics (retrieval latency, chunks returned), docs.
|
||||
|
||||
Built (`internal/router/rag.go`): `X-Llmgw-Knowledge` (header, as the
|
||||
other llmgw extensions are, not a body field), `X-Llmgw-Knowledge-K`;
|
||||
collections from the key's chain (kbEffective, shared with the kb
|
||||
endpoints), 404 for one not readable; the question embedded through the
|
||||
ordinary path on the caller's principal; SearchChunks; passages trimmed to
|
||||
the model's context; added to the one leading system message (or a new
|
||||
one); `x-llmgw-sources`; before guardInput; never cached; metrics. Default collections built (kb_defaults on org/team/project/key, narrowest
|
||||
wins, 30 s cache, `ctl kb default`, console action; never refuse); the
|
||||
portal chat's Documents picker (collection, default or none; sources listed
|
||||
under the answer; a header the page sends itself is dropped). `cache_salt` per organization on requests given passages, sent to vLLM
|
||||
replicas only (chosen replica's engine). Checked on the PoC 2026-10-08:
|
||||
vLLM 0.30.0 (cuda01) with TinyLlama-1.1B AWQ run by hand beside the 35B in
|
||||
the card's last 2.3 GB — a 668-token prompt reused 656 tokens under the same
|
||||
salt and 0 under another's, with identical text (llmgw cannot place it
|
||||
there: placement takes whole GPUs). Live on the PoC: a team default answers
|
||||
a request without the header, `none` turns it off, an explicit header with
|
||||
k=2 gives two sources and no more, llmgw-kb lists owners by path. Left: returned passages
|
||||
are not audited per request (the header and the kept request record them).
|
||||
|
||||
### F4. Adapter training (L, after F1–F3)
|
||||
A run-to-completion job placed by the control plane like a replica and
|
||||
supervised by the agent like an engine, in a Python environment the agent
|
||||
|
||||
@@ -18,11 +18,86 @@ is only released once its **exit criterion** is met, not when its code lands.
|
||||
| v0.11.0 | — | High availability from setup, the setup and console screens, usage sums and per-tenant alerts | released |
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
Until v1.0.0 the minor version is the breaking one: configuration formats and Go
|
||||
APIs change between minors without a deprecation period.
|
||||
|
||||
## Unreleased
|
||||
|
||||
### Added
|
||||
|
||||
- **The cluster's settings on the portal** (BACKLOG E4): `/admin/settings`
|
||||
for administrators. It edits the document `llmgw ctl config edit` opens:
|
||||
checked as a control plane would load it, refused when made from a
|
||||
version since replaced, with the history to view and roll back. Changes
|
||||
are audited by size, never by contents.
|
||||
- **The nodes' agent settings in the cluster's settings** (BACKLOG E3):
|
||||
`agents.default` and `agents.nodes.<id>` hold everything agent.yaml can
|
||||
say but the agent's id, control planes, address and certificate. They are
|
||||
sent to the agents as they change; an agent starts again on new settings,
|
||||
or on an edited agent.yaml, with its engines serving throughout. Settings
|
||||
that would not load are refused on save, and set aside on a node rather
|
||||
than keep its agent from starting. Nodes need this release to follow
|
||||
them.
|
||||
|
||||
## v0.14.0 — 2026-10-08
|
||||
|
||||
Knowledge bases: a company's documents, put in by people from their own
|
||||
machines, kept apart per tenant, and answered from.
|
||||
|
||||
**Upgrading.**
|
||||
- **pgvector on a cluster set up before this release.** Knowledge bases
|
||||
need pgvector, which `server init` and the high-availability database
|
||||
setup now install, but a rerun on a running cluster does not:
|
||||
1. Install `postgresql-<version>-pgvector` on every database machine.
|
||||
2. As postgres, run `CREATE EXTENSION vector` in llmgw's database.
|
||||
3. Restart the control planes.
|
||||
Until then everything else works, and `llmgw ctl kb list` says why
|
||||
knowledge bases are off. On a SQLite store they stay off.
|
||||
- **An embedding model.** A collection needs a catalog model with
|
||||
`capabilities.embeddings`, for example bge-m3, whose weights are in the
|
||||
mirror.
|
||||
- **The client.** People download `llmgw-kb` for their system from this
|
||||
release.
|
||||
|
||||
### Added
|
||||
|
||||
- **Collections and who may read them.** A collection is owned by an
|
||||
organization, team, project or person. Grants add up down the tenancy
|
||||
tree and never cross an organization. A collection a key may not read is
|
||||
answered as one that does not exist. Chunks are kept in PostgreSQL with
|
||||
pgvector, partitioned by organization, with an HNSW index per collection,
|
||||
and row-level security on the table as a second line behind every
|
||||
query's own filter. A project, team or organization that owns collections
|
||||
is not deleted; a person's collections go with their account.
|
||||
- **`llmgw-kb`**, its own download for Windows, macOS and Linux, run by
|
||||
people with their own key. It reads documents where they are and sends
|
||||
text and vectors, never the files, for what changed only; a SQLite ledger
|
||||
on the machine remembers what it did.
|
||||
- It reads PDF, Word, Excel, PowerPoint, OpenDocument, RTF, EPUB, HTML,
|
||||
MHTML, mail with attachments, notebooks, subtitles, ZIP archives, and
|
||||
text and source code in any common encoding.
|
||||
- Legacy Office files are read with LibreOffice, and scanned pages with
|
||||
Tesseract, when those are installed.
|
||||
- `--dry-run` says what would leave the machine; `--delete` removes only
|
||||
what that machine sent.
|
||||
- **Answers from documents.** A chat request names collections in
|
||||
`X-Llmgw-Knowledge`, or its part of the tree has a default (`llmgw ctl kb
|
||||
default`). The best passages, chosen by meaning plus a little for the
|
||||
question's words, are added to the system message, numbered, and
|
||||
`x-llmgw-sources` says which were used. This happens before the
|
||||
guardrails judge the request, and the answer is never served from the
|
||||
answer cache. On vLLM replicas the organization's `cache_salt` keeps
|
||||
tenants' prefix caches apart. It works through the Responses and Messages
|
||||
APIs too.
|
||||
- **Administration**: `llmgw ctl kb` (create, list, grant, revoke, default,
|
||||
docs, rm, search, delete), `/api/v1/kb`, the console's Knowledge screen,
|
||||
and a Documents picker on the portal's chat page.
|
||||
- Metrics `llmgw_kb_retrievals_total` and `llmgw_kb_retrieval_seconds`;
|
||||
[docs/knowledge.md](docs/knowledge.md).
|
||||
|
||||
## v0.13.0 — 2026-10-08
|
||||
|
||||
The whole system stopped, started and restarted from one command, a
|
||||
|
||||
@@ -6,6 +6,7 @@ LDFLAGS := -X '$(MODULE)/internal/version.Version=$(VERSION)'
|
||||
|
||||
build:
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw
|
||||
go build -ldflags "$(LDFLAGS)" -o bin/llmgw-kb ./cmd/llmgw-kb
|
||||
|
||||
test:
|
||||
go test ./...
|
||||
|
||||
@@ -3,8 +3,9 @@
|
||||
A single-binary LLM gateway for clusters with multiple InfiniBand fabrics,
|
||||
NVLink domains and heterogeneous GPUs.
|
||||
|
||||
It serves LLMs running on your own GPU servers (vLLM / SGLang) behind one
|
||||
OpenAI-compatible endpoint. Design document: [llmgw-plan.md](llmgw-plan.md).
|
||||
It serves LLMs running on your own GPU and CPU servers (vLLM, SGLang,
|
||||
llama.cpp, TensorRT-LLM) behind one OpenAI- and Anthropic-compatible
|
||||
endpoint, and answers from your own documents. Design document: [llmgw-plan.md](llmgw-plan.md).
|
||||
|
||||
## Getting started
|
||||
|
||||
@@ -37,6 +38,9 @@ Releases track the roadmap phases; see [CHANGELOG.md](CHANGELOG.md).
|
||||
| v0.9.0 | — | Managing people and the tenancy tree | released |
|
||||
| v0.10.0 | — | InfiniBand fabrics and NVLink, doctor, node add, engines installed by the nodes | released |
|
||||
| v0.11.0 | — | High availability from setup, the setup and console screens, usage sums and per-tenant alerts | released |
|
||||
| v0.12.0 | — | The cluster's settings in the store, in force when changed; the firewall switch; wipe | released |
|
||||
| v0.13.0 | — | The whole system stopped and started from one command; Qwen3Guard judges, on a GPU and CPUs together | released |
|
||||
| v0.14.0 | — | Knowledge bases: documents from people's own machines, kept apart per tenant, answered from | released |
|
||||
| v1.0.0 | 5 | Tenancy, prefix affinity, benchmarking | planned |
|
||||
|
||||
As of v0.6.0 a cluster runs unattended: agents report their hardware, the
|
||||
@@ -44,38 +48,60 @@ control plane places replicas — including ones larger than a single machine
|
||||
models load when someone asks for them and stop when nobody does, and keys,
|
||||
budgets and rate limits decide who may use any of it.
|
||||
|
||||
Since v0.11.0, not yet released (see [CHANGELOG.md](CHANGELOG.md),
|
||||
"Unreleased"):
|
||||
Since v0.11.0 (each release's detail is in [CHANGELOG.md](CHANGELOG.md)):
|
||||
|
||||
- **Security hardening** — signed releases only, a token on the admin API,
|
||||
engines not run as root, TLS for the agent port, etcd and Patroni, a
|
||||
firewall for llmgw's internal ports, encrypted backups, bounded
|
||||
connections, sign-in throttling, a Content-Security-Policy.
|
||||
- **Guardrails** — policies that check requests and answers (PII, secrets,
|
||||
prompt injection, deny patterns, a judge model such as Llama Guard), in
|
||||
observe or enforce mode, attached to any part of the tenancy tree:
|
||||
- **Knowledge bases** (v0.14.0). People put a company's documents in from
|
||||
their own machines with `llmgw-kb`, a separate download for Windows,
|
||||
macOS and Linux that reads PDF, Office, OpenDocument, mail and many more
|
||||
formats and sends text and vectors, never the files. Collections are kept
|
||||
apart per tenant, down to row-level security in PostgreSQL. A chat
|
||||
request, or a part of the tree by default, is answered from them with
|
||||
its sources: [docs/knowledge.md](docs/knowledge.md).
|
||||
- **One command for the whole system** (v0.13.0): `llmgw stop`, `start` and
|
||||
`restart`. Qwen3Guard is the guardrails' judge, on a GPU and on CPUs
|
||||
together.
|
||||
- **Settings that apply when changed** (v0.12.0). The cluster's
|
||||
configuration lives in the store, the same on every control plane and in
|
||||
force within seconds (`llmgw ctl config`, the console's Settings). The
|
||||
firewall can be switched on and off for the whole cluster.
|
||||
- **Security hardening** (v0.11.2): signed releases only, a token on the
|
||||
admin API, engines not run as root, TLS for the agent port, etcd and
|
||||
Patroni, a firewall for llmgw's internal ports, encrypted backups,
|
||||
bounded connections, sign-in throttling, a Content-Security-Policy.
|
||||
- **Guardrails** (v0.11.2). Policies check requests and answers for PII,
|
||||
secrets, prompt injection, deny patterns, and with a judge model; they
|
||||
observe or enforce, and attach to any part of the tenancy tree:
|
||||
[docs/guardrails.md](docs/guardrails.md).
|
||||
- **Administration** — an audit log of every change (`ctl audit`), staff
|
||||
roles (viewer, operator, billing, admin), single sign-on through OpenID
|
||||
Connect and provisioning through SCIM 2.0, key lifetime limits and
|
||||
rotation with overlap (`keys ttl`, `keys rotate`), usage export as CSV
|
||||
(`ctl usage export`).
|
||||
- **Notifications** — webhooks, signed or Slack-shaped, for a model parked,
|
||||
a node lost, a budget spent and new weights published (`notify.webhooks`,
|
||||
`ctl notify test`), and Alertmanager receiver templates.
|
||||
- **Traffic and models** — request tags, canaries of new weights with
|
||||
automatic rollback, quality suites gating updates, a response cache,
|
||||
kept and replayed requests, batches, Anthropic's Messages API, more of
|
||||
OpenAI's (structured output, images, rerank, audio, moderations),
|
||||
scaling by the clock, SLO-aware admission, TensorRT-LLM, MIG, models
|
||||
imported from Ollama, routing by what each vLLM replica's cache holds,
|
||||
LoRA adapters loaded on demand, other llmgw clusters as providers, and
|
||||
placement that knows racks, PDUs and power caps.
|
||||
- **Operations** — `server restore` and off-machine backup copies, secrets
|
||||
in Vault, forgetting a person and keeping a project's requests in the
|
||||
cluster, offline installs from a bundle, a Helm chart, tracing from the
|
||||
control plane to the engines, the admin API in OpenAPI with a Terraform
|
||||
provider, and the portal in English and Turkish with its docs at /docs.
|
||||
- **Administration** (v0.11.2):
|
||||
- an audit log of every change (`ctl audit`);
|
||||
- staff roles (viewer, operator, billing, admin);
|
||||
- single sign-on through OpenID Connect and provisioning through SCIM
|
||||
2.0;
|
||||
- key lifetime limits, and rotation with overlap (`keys ttl`, `keys
|
||||
rotate`);
|
||||
- usage export as CSV (`ctl usage export`).
|
||||
- **Notifications** (v0.11.2): webhooks, signed or Slack-shaped, for a
|
||||
model parked, a node lost, a budget spent and new weights published, and
|
||||
Alertmanager receiver templates.
|
||||
- **Traffic and models** (v0.11.2):
|
||||
- request tags, canaries of new weights with automatic rollback, and
|
||||
quality suites gating updates;
|
||||
- a response cache, kept and replayed requests, and batches;
|
||||
- Anthropic's Messages API, and more of OpenAI's (structured output,
|
||||
images, rerank, audio, moderations);
|
||||
- scaling by the clock, and SLO-aware admission;
|
||||
- TensorRT-LLM, MIG, and models imported from Ollama;
|
||||
- routing by what each vLLM replica's cache holds, and LoRA adapters
|
||||
loaded on demand;
|
||||
- other llmgw clusters as providers, and placement that knows racks,
|
||||
PDUs and power caps.
|
||||
- **Operations** (v0.11.2):
|
||||
- `server restore`, off-machine backup copies, and secrets in Vault;
|
||||
- forgetting a person, and keeping a project's requests in the cluster;
|
||||
- offline installs from a bundle, and a Helm chart;
|
||||
- tracing from the control plane to the engines;
|
||||
- the admin API in OpenAPI, with a Terraform provider;
|
||||
- the portal in English and Turkish, with its docs at /docs.
|
||||
|
||||
Every one of these is in the console as well as on the command line:
|
||||
[docs/console.md](docs/console.md).
|
||||
@@ -124,6 +150,9 @@ Every page is also in the binary: `llmgw docs`, or **Manual** in the console.
|
||||
official OpenAI and Anthropic SDKs, and llmgw's extensions
|
||||
- [docs/guardrails.md](docs/guardrails.md) — checking requests and answers:
|
||||
policies, detectors, a judge model, and where each applies
|
||||
- [docs/knowledge.md](docs/knowledge.md) — knowledge bases: collections,
|
||||
who may read them, `llmgw-kb` on people's machines, and answers from
|
||||
documents
|
||||
- [docs/choosing-models.md](docs/choosing-models.md) — which models to run,
|
||||
how many of each, and where the weights should live
|
||||
|
||||
@@ -152,7 +181,7 @@ Every page is also in the binary: `llmgw docs`, or **Manual** in the console.
|
||||
## Building
|
||||
|
||||
```sh
|
||||
make build # bin/llmgw
|
||||
make build # bin/llmgw and bin/llmgw-kb
|
||||
make test # unit tests
|
||||
make race # with the race detector
|
||||
```
|
||||
@@ -169,6 +198,7 @@ llmgw ctl cordon gpu-a-01 # stop new placement, leave what runs
|
||||
sudo llmgw doctor # on a node: what its GPUs, NVLink and fabric need
|
||||
llmgw ctl audit --since 24h # who changed what
|
||||
llmgw ctl usage export --from 2026-09-01 --to 2026-10-01 --by project,model
|
||||
llmgw ctl kb list # knowledge-base collections, grants, defaults
|
||||
```
|
||||
|
||||
Bare `llmgw` opens the console, which does all of this from one screen.
|
||||
|
||||
+299
-226
File diff suppressed because it is too large.
Load diff
@@ -73,6 +73,7 @@ message ControlMessage {
|
||||
ControlPlanes control_planes = 11;
|
||||
Firewall firewall = 13;
|
||||
RestartAgent restart_agent = 14;
|
||||
AgentSettings agent_settings = 15;
|
||||
}
|
||||
// trace is the W3C trace context (traceparent, and tracestate when there
|
||||
// is one) of the control plane's span that sent this command, so what the
|
||||
@@ -103,6 +104,17 @@ message Firewall {
|
||||
int32 engine_ports_to = 4;
|
||||
}
|
||||
|
||||
// AgentSettings is what the cluster's settings say of this node's agent
|
||||
// (agents.default, with agents.nodes.<id> over it), sent when it registers
|
||||
// and whenever they change. The agent keeps it beside agent.yaml, which
|
||||
// holds only what it needs to start; over that file, it is what the agent
|
||||
// runs on. A change from what it runs on has the agent start again on it,
|
||||
// its engines serving throughout.
|
||||
message AgentSettings {
|
||||
// yaml is agent.yaml's keys, as YAML; empty for none.
|
||||
string yaml = 1;
|
||||
}
|
||||
|
||||
// ControlPlanes is every control plane an agent may dial, sent by the leader
|
||||
// when an agent registers and whenever the list changes. The agent keeps it
|
||||
// beside its configuration, so a second control plane added after the node
|
||||
|
||||
@@ -0,0 +1,467 @@
|
||||
// Command llmgw-kb puts a person's documents into an llmgw knowledge base
|
||||
// from their own machine (BACKLOG F2): it reads the files there, cuts them
|
||||
// into chunks, has the gateway embed them with the person's key, and sends
|
||||
// the text and vectors — only for what changed since the last run, which a
|
||||
// ledger on the machine remembers.
|
||||
//
|
||||
// Its own binary, for Windows, macOS and Linux: it carries none of the
|
||||
// server, agent or setup code that llmgw does.
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"golang.org/x/term"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/extract"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/version"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if err := root().Execute(); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "error:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// config is what `login` keeps: the gateway, the key, and for a personal key
|
||||
// the project it acts in.
|
||||
type config struct {
|
||||
Gateway string `json:"gateway"`
|
||||
Key string `json:"key"`
|
||||
Project string `json:"project,omitempty"`
|
||||
CA string `json:"ca,omitempty"`
|
||||
}
|
||||
|
||||
func configPath() (string, error) {
|
||||
dir, err := os.UserConfigDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Join(dir, "llmgw-kb", "config.json"), nil
|
||||
}
|
||||
|
||||
// loadConfig reads the saved configuration, with the environment over it:
|
||||
// LLMGW_URL, LLMGW_API_KEY, LLMGW_PROJECT.
|
||||
func loadConfig() (config, error) {
|
||||
var c config
|
||||
if p, err := configPath(); err == nil {
|
||||
if b, err := os.ReadFile(p); err == nil {
|
||||
if err := json.Unmarshal(b, &c); err != nil {
|
||||
return c, fmt.Errorf("%s: %w", p, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
for env, dst := range map[string]*string{"LLMGW_URL": &c.Gateway, "LLMGW_API_KEY": &c.Key, "LLMGW_PROJECT": &c.Project} {
|
||||
if v := strings.TrimSpace(os.Getenv(env)); v != "" {
|
||||
*dst = v
|
||||
}
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func client(flags *globalFlags) (*kbsync.Client, error) {
|
||||
c, err := loadConfig()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if flags.gateway != "" {
|
||||
c.Gateway = flags.gateway
|
||||
}
|
||||
if flags.project != "" {
|
||||
c.Project = flags.project
|
||||
}
|
||||
if flags.ca != "" {
|
||||
c.CA = flags.ca
|
||||
}
|
||||
if c.Gateway == "" || c.Key == "" {
|
||||
return nil, errors.New("no gateway or key: run `llmgw-kb login`, or set LLMGW_URL and LLMGW_API_KEY")
|
||||
}
|
||||
h, err := kbsync.NewHTTP(c.CA)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h}, nil
|
||||
}
|
||||
|
||||
type globalFlags struct {
|
||||
gateway, project, ca, ledger string
|
||||
}
|
||||
|
||||
func (g *globalFlags) openLedger(ctx context.Context) (*kbsync.Ledger, error) {
|
||||
p := g.ledger
|
||||
if p == "" {
|
||||
var err error
|
||||
if p, err = kbsync.DefaultLedgerPath(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return kbsync.OpenLedger(ctx, p)
|
||||
}
|
||||
|
||||
func root() *cobra.Command {
|
||||
g := &globalFlags{}
|
||||
cmd := &cobra.Command{
|
||||
Use: "llmgw-kb",
|
||||
Short: "Put your documents into an llmgw knowledge base, from this machine",
|
||||
Long: `llmgw-kb reads documents on this machine, cuts them into chunks, has the
|
||||
gateway embed them with your key, and sends the text and vectors to a
|
||||
knowledge base. Files are never uploaded. Run it as often as you like: a
|
||||
ledger on this machine remembers what was sent, and only what changed is
|
||||
read and sent again.`,
|
||||
Version: version.Version,
|
||||
SilenceUsage: true,
|
||||
SilenceErrors: true,
|
||||
}
|
||||
cmd.PersistentFlags().StringVar(&g.gateway, "gateway", "", "the gateway's address (default: from login, or LLMGW_URL)")
|
||||
cmd.PersistentFlags().StringVar(&g.project, "project", "", "the project a personal key acts in (default: from login, or LLMGW_PROJECT)")
|
||||
cmd.PersistentFlags().StringVar(&g.ca, "ca", "", "a PEM file of certificates to trust besides the system's")
|
||||
cmd.PersistentFlags().StringVar(&g.ledger, "ledger", "", "the ledger file (default: in your configuration directory)")
|
||||
cmd.AddCommand(loginCmd(g), collectionsCmd(g), syncCmd(g), statusCmd(g), historyCmd(g), formatsCmd())
|
||||
return cmd
|
||||
}
|
||||
|
||||
func loginCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "login",
|
||||
Short: "Save the gateway's address and your key",
|
||||
Long: `Asks for the gateway's address and your API key, checks them, and saves them
|
||||
in your configuration directory, readable by you only. The key is typed
|
||||
without being shown, so it does not end up in your shell's history.`,
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
c, _ := loadConfig()
|
||||
in := bufio.NewReader(os.Stdin)
|
||||
ask := func(q, def string) string {
|
||||
if def != "" {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s [%s]: ", q, def)
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s: ", q)
|
||||
}
|
||||
s, _ := in.ReadString('\n')
|
||||
if s = strings.TrimSpace(s); s == "" {
|
||||
return def
|
||||
}
|
||||
return s
|
||||
}
|
||||
c.Gateway = strings.TrimRight(ask("Gateway address (https://…)", or(g.gateway, c.Gateway)), "/")
|
||||
fmt.Fprint(cmd.OutOrStdout(), "API key: ")
|
||||
var key string
|
||||
if term.IsTerminal(int(os.Stdin.Fd())) {
|
||||
b, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||
fmt.Fprintln(cmd.OutOrStdout())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key = strings.TrimSpace(string(b))
|
||||
} else {
|
||||
key, _ = in.ReadString('\n')
|
||||
key = strings.TrimSpace(key)
|
||||
}
|
||||
if key != "" {
|
||||
c.Key = key
|
||||
}
|
||||
c.Project = ask("Project, for a personal key (empty for a project's key)", or(g.project, c.Project))
|
||||
if g.ca != "" {
|
||||
c.CA = g.ca
|
||||
}
|
||||
h, err := kbsync.NewHTTP(c.CA)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cl := &kbsync.Client{Base: c.Gateway, Key: c.Key, Project: c.Project, HTTP: h, Wait: time.Second}
|
||||
ctx, cancel := context.WithTimeout(cmd.Context(), 30*time.Second)
|
||||
defer cancel()
|
||||
cs, err := cl.Collections(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("the gateway refused: %w", err)
|
||||
}
|
||||
p, err := configPath()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil {
|
||||
return err
|
||||
}
|
||||
b, _ := json.MarshalIndent(c, "", " ")
|
||||
if err := os.WriteFile(p, b, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "Saved in %s. This key reaches %d collection(s); `llmgw-kb collections` lists them.\n", p, len(cs))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func or(a, b string) string {
|
||||
if a != "" {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func collectionsCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "collections",
|
||||
Short: "List the collections your key may read, and which it may add to",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cs, err := c.Collections(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(cs) == 0 {
|
||||
fmt.Fprintln(cmd.OutOrStdout(), "No collections. An administrator creates them and grants access (llmgw ctl kb).")
|
||||
return nil
|
||||
}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "NAME\tACCESS\tOWNER\tEMBEDDED BY")
|
||||
for _, x := range cs {
|
||||
access := "read"
|
||||
if x.Writable {
|
||||
access = "read, write"
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%s\t%s (%d)\n", x.Name, access, or(x.OwnerName, x.Owner), x.EmbedModel, x.Dims)
|
||||
}
|
||||
return tw.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func syncCmd(g *globalFlags) *cobra.Command {
|
||||
var o kbsync.Options
|
||||
var verbose bool
|
||||
var ocr string
|
||||
var maxMiB int64
|
||||
cmd := &cobra.Command{
|
||||
Use: "sync COLLECTION PATH...",
|
||||
Short: "Send a collection what changed in these folders and files",
|
||||
Long: `Reads the files under each PATH that the collection does not already have as
|
||||
they are, and sends their text and vectors. Run it again whenever you like;
|
||||
unchanged files are not even read.
|
||||
|
||||
Documents are named by their path under PATH, with forward slashes; with
|
||||
several PATHs, each folder's own name comes first.
|
||||
|
||||
--dry-run reads and cuts everything, sends nothing, and says what would
|
||||
leave this machine
|
||||
--delete removes documents this machine sent whose files are gone
|
||||
(only under these PATHs; nothing anybody else sent)
|
||||
--retry reads again files skipped before (after installing LibreOffice
|
||||
or Tesseract, say)
|
||||
|
||||
Without --include, every file of a known format is read (llmgw-kb formats);
|
||||
hidden files, version control folders and Office lock files are not.`,
|
||||
Args: cobra.MinimumNArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt)
|
||||
defer stop()
|
||||
l, err := g.openLedger(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
o.Collection, o.Roots = args[0], args[1:]
|
||||
o.Extract = extract.Options{Helpers: extract.FindHelpers(), OCRLanguages: ocr, MaxBytes: maxMiB << 20}
|
||||
out := cmd.OutOrStdout()
|
||||
c.Waiting = func(why string, d time.Duration) {
|
||||
fmt.Fprintf(out, " .. waiting %s: %s\n", d, why)
|
||||
}
|
||||
start := time.Now()
|
||||
rep, err := kbsync.Sync(ctx, c, l, o, func(e kbsync.Event) {
|
||||
switch e.Kind {
|
||||
case "unchanged":
|
||||
if !verbose {
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(out, " = %s\n", e.Name)
|
||||
case "sent", "would-send":
|
||||
fmt.Fprintf(out, " %-8s %s (%d chunks, %s)\n", label(e.Kind), e.Name, e.Chunks, size(e.Bytes))
|
||||
default:
|
||||
line := fmt.Sprintf(" %-8s %s", label(e.Kind), e.Name)
|
||||
if e.Detail != "" {
|
||||
line += ": " + e.Detail
|
||||
}
|
||||
fmt.Fprintln(out, line)
|
||||
}
|
||||
})
|
||||
if rep.Collection.Name != "" {
|
||||
verb := "sent"
|
||||
if o.DryRun {
|
||||
verb = "would send"
|
||||
}
|
||||
fmt.Fprintf(out, "\n%s: %d files; %s %d (%d chunks, %s), %d unchanged, %d skipped, %d failed, %d deleted, in %s\n",
|
||||
rep.Collection.Name, rep.Files, verb, rep.Sent, rep.Chunks, size(rep.Bytes), rep.Unchanged,
|
||||
rep.Skipped, rep.Failed, rep.Deleted, time.Since(start).Round(time.Second))
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if rep.Failed > 0 {
|
||||
return fmt.Errorf("%d document(s) failed; run again to retry them", rep.Failed)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
f := cmd.Flags()
|
||||
f.StringSliceVar(&o.Include, "include", nil, "read only files matching these globs (\"*.pdf,docs/**/*.md\")")
|
||||
f.StringSliceVar(&o.Exclude, "exclude", nil, "never read files matching these globs")
|
||||
f.BoolVar(&o.Hidden, "hidden", false, "read hidden files and folders too")
|
||||
f.BoolVar(&o.DryRun, "dry-run", false, "send nothing; say what would be sent")
|
||||
f.BoolVar(&o.Delete, "delete", false, "remove documents this machine sent whose files are gone")
|
||||
f.BoolVar(&o.Retry, "retry", false, "read again files skipped before")
|
||||
f.IntVar(&o.Workers, "workers", 4, "files read and sent at once")
|
||||
f.StringVar(&ocr, "ocr-languages", "tur+eng", "Tesseract's languages for images and scanned pages")
|
||||
f.Int64Var(&maxMiB, "max-size", 256, "skip files larger than this many MiB")
|
||||
f.BoolVarP(&verbose, "verbose", "v", false, "list unchanged files too")
|
||||
return cmd
|
||||
}
|
||||
|
||||
func label(kind string) string {
|
||||
switch kind {
|
||||
case "would-send":
|
||||
return "send"
|
||||
case "would-delete":
|
||||
return "delete"
|
||||
}
|
||||
return kind
|
||||
}
|
||||
|
||||
func size(b int64) string {
|
||||
switch {
|
||||
case b >= 1<<20:
|
||||
return fmt.Sprintf("%.1f MiB", float64(b)/(1<<20))
|
||||
case b >= 1<<10:
|
||||
return fmt.Sprintf("%.0f KiB", float64(b)/(1<<10))
|
||||
}
|
||||
return fmt.Sprintf("%d B", b)
|
||||
}
|
||||
|
||||
func statusCmd(g *globalFlags) *cobra.Command {
|
||||
var all bool
|
||||
return &cobra.Command{
|
||||
Use: "status COLLECTION",
|
||||
Short: "Show what this machine has sent to a collection, and what it skipped",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
c, err := client(g)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
m, err := c.Manifest(cmd.Context(), args[0])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
l, err := g.openLedger(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
es, err := l.Entries(cmd.Context(), strings.TrimRight(c.Base, "/"), m.Collection.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
counts := map[kbsync.State]int{}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "STATE\tDOCUMENT\tCHUNKS\tWHEN\tWHY")
|
||||
for _, e := range es {
|
||||
counts[e.State]++
|
||||
if e.State == kbsync.Sent && !all {
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%d\t%s\t%s\n", e.State, e.Name, e.Chunks, e.At.Format("2006-01-02 15:04"), e.Reason)
|
||||
}
|
||||
_ = tw.Flush()
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "\n%s holds %d documents; from this machine %d sent, %d skipped, %d failed.",
|
||||
m.Collection.Name, len(m.Documents), counts[kbsync.Sent], counts[kbsync.Skipped], counts[kbsync.Failed])
|
||||
if !all {
|
||||
fmt.Fprint(cmd.OutOrStdout(), " (--all lists the sent ones too)")
|
||||
}
|
||||
fmt.Fprintln(cmd.OutOrStdout())
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func historyCmd(g *globalFlags) *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "history",
|
||||
Short: "List the latest syncs from this machine",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
l, err := g.openLedger(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer l.Close()
|
||||
runs, err := l.Runs(cmd.Context(), 20)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tw := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(tw, "STARTED\tTOOK\tCOLLECTION\tSENT\tUNCHANGED\tSKIPPED\tFAILED\tDELETED")
|
||||
for _, r := range runs {
|
||||
took := "unfinished"
|
||||
if !r.Finished.IsZero() {
|
||||
took = r.Finished.Sub(r.Started).Round(time.Second).String()
|
||||
}
|
||||
fmt.Fprintf(tw, "%s\t%s\t%s\t%d\t%d\t%d\t%d\t%d\n", r.Started.Format("2006-01-02 15:04"), took,
|
||||
r.Collection, r.Sent, r.Unchanged, r.Skipped, r.Failed, r.Delete)
|
||||
}
|
||||
return tw.Flush()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func formatsCmd() *cobra.Command {
|
||||
return &cobra.Command{
|
||||
Use: "formats",
|
||||
Short: "List the formats read, and the helpers found on this machine",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
builtIn, helper := extract.Supported()
|
||||
out := cmd.OutOrStdout()
|
||||
fmt.Fprintf(out, "Read here:\n%s\n\n", wrap(builtIn))
|
||||
fmt.Fprintf(out, "Read with a helper program:\n%s\n\n", wrap(helper))
|
||||
fmt.Fprintln(out, "Helpers:")
|
||||
for _, l := range extract.FindHelpers().Describe() {
|
||||
fmt.Fprintln(out, " "+l)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func wrap(words []string) string {
|
||||
var b strings.Builder
|
||||
line := " "
|
||||
for _, w := range words {
|
||||
if len(line)+len(w) > 78 {
|
||||
b.WriteString(strings.TrimRight(line, " ") + "\n")
|
||||
line = " "
|
||||
}
|
||||
line += w + " "
|
||||
}
|
||||
b.WriteString(strings.TrimRight(line, " "))
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The client goes onto employees' laptops: none of the server, agent, store
|
||||
// or setup code may come with it, and it must build where llmgw does not.
|
||||
func TestTheClientCarriesNoServerCode(t *testing.T) {
|
||||
out, err := exec.Command("go", "list", "-deps", ".").Output()
|
||||
if err != nil {
|
||||
t.Skipf("go list: %v", err)
|
||||
}
|
||||
for _, dep := range strings.Fields(string(out)) {
|
||||
for _, banned := range []string{"/internal/store", "/internal/router", "/internal/agent", "/internal/control",
|
||||
"/internal/cli", "/internal/guard", "github.com/jackc/pgx"} {
|
||||
if strings.Contains(dep, banned) {
|
||||
t.Errorf("llmgw-kb depends on %s", dep)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -39,7 +39,7 @@ func (d Doc) Source() string {
|
||||
// than with either reader, because the console's manual and the gateway's
|
||||
// docs portal both list them, and two lists drift.
|
||||
var docOrder = []string{"console", "quick-install", "install", "node-onboarding", "onboarding",
|
||||
"clients", "guardrails", "choosing-models", "fabrics", "facility-aware-placement", "data",
|
||||
"clients", "guardrails", "knowledge", "choosing-models", "fabrics", "facility-aware-placement", "data",
|
||||
"reference", "development"}
|
||||
|
||||
// Manual is every document, in reading order: those docOrder names, then any
|
||||
|
||||
@@ -53,6 +53,9 @@
|
||||
{
|
||||
"name": "jobs"
|
||||
},
|
||||
{
|
||||
"name": "kb"
|
||||
},
|
||||
{
|
||||
"name": "metrics"
|
||||
},
|
||||
@@ -1069,6 +1072,381 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb": {
|
||||
"get": {
|
||||
"operationId": "getKb",
|
||||
"summary": "Every knowledge-base collection, with its documents' count and grants",
|
||||
"description": "Says why there are none when the store cannot hold them (SQLite, or PostgreSQL without pgvector).",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBView"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb/collections": {
|
||||
"post": {
|
||||
"operationId": "postKbCollections",
|
||||
"summary": "Make a collection, owned by an organization, team, project or person",
|
||||
"description": "The embedding model's dimension is asked of the model when dims is not given, so the model must be able to serve.",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBCreateRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"201": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBCollectionView"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb/collections/{id}": {
|
||||
"delete": {
|
||||
"operationId": "deleteKbCollectionsId",
|
||||
"summary": "Remove a collection with its documents and grants",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "The collection's ID, its name when that is unique, or ORG/NAME.",
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"additionalProperties": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb/collections/{id}/documents": {
|
||||
"get": {
|
||||
"operationId": "getKbCollectionsIdDocuments",
|
||||
"summary": "A collection's documents",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "The collection's ID, its name when that is unique, or ORG/NAME.",
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/KBDocumentView"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
},
|
||||
"delete": {
|
||||
"operationId": "deleteKbCollectionsIdDocuments",
|
||||
"summary": "Remove one document from a collection",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "The collection's ID, its name when that is unique, or ORG/NAME.",
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "name",
|
||||
"in": "query",
|
||||
"required": true,
|
||||
"description": "The document's name, as listed.",
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"additionalProperties": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb/defaults": {
|
||||
"post": {
|
||||
"operationId": "postKbDefaults",
|
||||
"summary": "Set the collections a part of the tree answers from when a request names none",
|
||||
"description": "On an organization, team, project or key; the narrowest one with a default wins. No collections clears it. A request's own X-Llmgw-Knowledge overrides it, and X-Llmgw-Knowledge: none turns it off.",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBDefaultRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBDefaultView"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb/grants": {
|
||||
"post": {
|
||||
"operationId": "postKbGrants",
|
||||
"summary": "Let a scope, and everything below it, read or write a collection",
|
||||
"description": "Never outside the collection's organization.",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBGrantRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBGrantView"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
},
|
||||
"delete": {
|
||||
"operationId": "deleteKbGrants",
|
||||
"summary": "Take a grant away",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBGrantRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBGrantView"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/kb/search": {
|
||||
"post": {
|
||||
"operationId": "postKbSearch",
|
||||
"summary": "Try a search of a collection, as an administrator",
|
||||
"description": "No tenancy applies: this is for checking what a collection holds and finds.",
|
||||
"tags": [
|
||||
"kb"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"$ref": "#/components/parameters/Actor"
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/KBSearchRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "OK",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/KBHitView"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"401": {
|
||||
"$ref": "#/components/responses/Unauthorized"
|
||||
},
|
||||
"default": {
|
||||
"$ref": "#/components/responses/Error"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/metrics/range": {
|
||||
"post": {
|
||||
"operationId": "postMetricsRange",
|
||||
@@ -4606,6 +4984,249 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBCollectionView": {
|
||||
"type": "object",
|
||||
"description": "KBCollectionView is a collection as an administrator sees it.",
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"org": {
|
||||
"type": "string"
|
||||
},
|
||||
"org_name": {
|
||||
"type": "string"
|
||||
},
|
||||
"owner": {
|
||||
"type": "string"
|
||||
},
|
||||
"owner_name": {
|
||||
"type": "string"
|
||||
},
|
||||
"embed_model": {
|
||||
"type": "string"
|
||||
},
|
||||
"dims": {
|
||||
"type": "integer"
|
||||
},
|
||||
"documents": {
|
||||
"type": "integer"
|
||||
},
|
||||
"chunks": {
|
||||
"type": "integer"
|
||||
},
|
||||
"grants": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/KBGrantView"
|
||||
},
|
||||
"description": "Grants are who else may read or write it."
|
||||
},
|
||||
"created_at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"created_by": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBCreateRequest": {
|
||||
"type": "object",
|
||||
"description": "KBCreateRequest makes a collection.",
|
||||
"properties": {
|
||||
"owner": {
|
||||
"type": "string",
|
||||
"description": "Owner is org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT (or the ID after the colon), or user:USER-ID."
|
||||
},
|
||||
"org": {
|
||||
"type": "string",
|
||||
"description": "Org is the organization a person's collection belongs to; only for user: owners."
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"description": "Model is the catalog model that embeds its documents."
|
||||
},
|
||||
"dims": {
|
||||
"type": "integer",
|
||||
"description": "Dims is the model's embedding dimension; asked of the model when 0."
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBDefaultRequest": {
|
||||
"type": "object",
|
||||
"description": "KBDefaultRequest sets the collections a part of the tree answers from when a request names none; none clears it.",
|
||||
"properties": {
|
||||
"target": {
|
||||
"type": "string",
|
||||
"description": "Target is org:, team:, project: (by path or ID) or key:KEY-ID."
|
||||
},
|
||||
"collections": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Collections are IDs, names when unique, or ORG/NAME."
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBDefaultView": {
|
||||
"type": "object",
|
||||
"description": "KBDefaultView is one part of the tree's default collections.",
|
||||
"properties": {
|
||||
"target": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"collections": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Collections are ORG/NAME, in order."
|
||||
},
|
||||
"at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"by": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBDocumentView": {
|
||||
"type": "object",
|
||||
"description": "KBDocumentView is one document in a collection.",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"digest": {
|
||||
"type": "string"
|
||||
},
|
||||
"bytes": {
|
||||
"type": "integer",
|
||||
"format": "int64"
|
||||
},
|
||||
"chunks": {
|
||||
"type": "integer"
|
||||
},
|
||||
"added_at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"added_by": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBGrantRequest": {
|
||||
"type": "object",
|
||||
"description": "KBGrantRequest lets a scope read or write a collection, or with DELETE takes that away.",
|
||||
"properties": {
|
||||
"collection": {
|
||||
"type": "string",
|
||||
"description": "Collection is its ID, its name when that is unique, or ORG/NAME."
|
||||
},
|
||||
"target": {
|
||||
"type": "string",
|
||||
"description": "Target is org:, team:, project:, key: or user: as for guard policies."
|
||||
},
|
||||
"perm": {
|
||||
"type": "string",
|
||||
"description": "Perm is read or write; ignored by DELETE."
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBGrantView": {
|
||||
"type": "object",
|
||||
"description": "KBGrantView is one grant.",
|
||||
"properties": {
|
||||
"target": {
|
||||
"type": "string"
|
||||
},
|
||||
"name": {
|
||||
"type": "string"
|
||||
},
|
||||
"perm": {
|
||||
"type": "string"
|
||||
},
|
||||
"at": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"by": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBHitView": {
|
||||
"type": "object",
|
||||
"description": "KBHitView is one chunk a search found.",
|
||||
"properties": {
|
||||
"document": {
|
||||
"type": "string"
|
||||
},
|
||||
"ord": {
|
||||
"type": "integer"
|
||||
},
|
||||
"score": {
|
||||
"type": "number"
|
||||
},
|
||||
"text": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBSearchRequest": {
|
||||
"type": "object",
|
||||
"description": "KBSearchRequest tries a search, as an administrator: no tenancy applies.",
|
||||
"properties": {
|
||||
"collection": {
|
||||
"type": "string"
|
||||
},
|
||||
"query": {
|
||||
"type": "string"
|
||||
},
|
||||
"k": {
|
||||
"type": "integer"
|
||||
}
|
||||
}
|
||||
},
|
||||
"KBView": {
|
||||
"type": "object",
|
||||
"description": "KBView is every collection, or why there can be none.",
|
||||
"properties": {
|
||||
"available": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": "Reason is why not, when not."
|
||||
},
|
||||
"collections": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/KBCollectionView"
|
||||
}
|
||||
},
|
||||
"defaults": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/KBDefaultView"
|
||||
},
|
||||
"description": "Defaults are the collections parts of the tree answer from when a request names none."
|
||||
}
|
||||
}
|
||||
},
|
||||
"Live": {
|
||||
"type": "object",
|
||||
"description": "Live is one of the cluster's adapters as its node last reported it.",
|
||||
|
||||
@@ -164,6 +164,15 @@ which): requests, tokens, errors and spend.
|
||||
team or project is guaranteed and may at most hold, and what it holds now.
|
||||
**A** adds one; on a row, change or remove it.
|
||||
|
||||
**Knowledge** lists the knowledge-base collections: each one's owner, the
|
||||
model that embeds it, how many documents and chunks it holds and who else it
|
||||
is granted to. **A** makes one. On a row: its **Documents**, **Grant…** a
|
||||
project, team, organization, person or key read or write access, **Revoke**
|
||||
a grant, **Answer from it by default…** for a part of the tree, **Try a
|
||||
search…** as an administrator, or **Delete…** it. People
|
||||
fill collections from their own machines with `llmgw-kb`
|
||||
([knowledge.md](knowledge.md)). Off on a SQLite store, and says why.
|
||||
|
||||
**Mirror** is what the model mirror holds, one row for each version of each
|
||||
model — a file, or the set of files a model of several is: shards, a vision
|
||||
projector, a config and tokenizer — with how far its download has got and
|
||||
|
||||
+35
-2
@@ -258,8 +258,14 @@ llmgw ctl config history # who changed what, when, and why
|
||||
llmgw ctl config rollback 4 # back to version 4, as a new version
|
||||
```
|
||||
|
||||
or on the console's **Settings** screen. A change is checked before it is
|
||||
kept — one that would not load is refused — and recorded in the audit log.
|
||||
or on the console's **Settings** screen, or in a browser on the portal's
|
||||
**Settings** page (`/admin/settings`, for administrators only: the document
|
||||
can hold webhook secrets). The page edits the same document, lists its
|
||||
history, and puts an earlier version back. A change is checked before it is
|
||||
kept, so one that would not load is refused. A change made from a version
|
||||
someone else has replaced since is refused too, rather than overwriting
|
||||
theirs. Every change is recorded in the audit log, with the document's size
|
||||
but not its contents.
|
||||
Guardrails change live; anything else has every control plane start again on
|
||||
the new settings, the standbys first and the leader half a minute later, so
|
||||
one of them is always serving. Engines keep running through it.
|
||||
@@ -269,6 +275,33 @@ cluster keys — with a `baseline` guard policy, if the file had none — and
|
||||
from then on those keys in any server.yaml are not read (the server says so
|
||||
when it starts).
|
||||
|
||||
**The nodes' agents** are set the same way, under `agents`. `agents.default`
|
||||
is for every node, and `agents.nodes.<node-id>` is for one node, laid over
|
||||
the default:
|
||||
|
||||
```
|
||||
llmgw ctl config set agents.default.heartbeat_interval 5s
|
||||
llmgw ctl config set agents.nodes.cuda01.labels.gpu rtx6000
|
||||
llmgw ctl config edit # fabrics, engines, supervisor, discovery…
|
||||
```
|
||||
|
||||
Each node's `agent.yaml` then needs only what its agent must know before it
|
||||
can be told anything: `node_id`, `server` or `servers`, `hostname`,
|
||||
`mgmt_addr` and `tls`. Those keys are refused under `agents`.
|
||||
|
||||
- **Precedence.** Anything else in the file is still read, with the
|
||||
cluster's settings over it.
|
||||
- **Applying a change.** The control plane sends each node its settings, a
|
||||
second apart. An agent whose settings changed keeps them beside its
|
||||
`agent.yaml` (in `/var/lib/llmgw-agent/agent-settings.yaml`) and starts
|
||||
again on them, its engines serving throughout.
|
||||
- **Bad settings.** Settings that would not load are refused when they are
|
||||
saved. If they somehow do not load beside one node's own file, that agent
|
||||
keeps running on what it has and logs why.
|
||||
- **Editing `agent.yaml` itself.** The agent notices the change within ten
|
||||
seconds and starts again on it. If the edited file does not load, the
|
||||
agent reports it and keeps running on what it had.
|
||||
|
||||
**This machine's own keys** stay in `/etc/llmgw/server.yaml`: `listen`,
|
||||
`admin_listen`, `metrics_listen`, `agent_listen`, `agent_tls`, `auth.store`,
|
||||
`secrets`, `tracing`, `monitoring`, `mirror.api` and `mirror.tls`, the backup
|
||||
|
||||
@@ -0,0 +1,349 @@
|
||||
# Knowledge bases
|
||||
|
||||
A knowledge base is a collection of a company's documents that the cluster
|
||||
can answer from. People put their documents into it from their own machines
|
||||
with `llmgw-kb`. The tool reads the files on that machine, cuts them into
|
||||
passages ("chunks"), has the gateway embed each passage with the person's own
|
||||
key, and sends the text and the vectors. The files themselves never leave
|
||||
the machine.
|
||||
|
||||
## Why retrieval, and not training
|
||||
|
||||
Retrieval is the tool for *knowledge*: policies, procedures, contracts. The
|
||||
answer can name its source, a changed document is sent again in minutes, and
|
||||
a passage only ever reaches people allowed to read its collection.
|
||||
|
||||
Training a model on the documents (LoRA) changes how it *writes*, not what it
|
||||
reliably *knows*. It has no access control: whatever goes into the weights
|
||||
is available to everyone who may use them.
|
||||
|
||||
## What it needs
|
||||
|
||||
- **The PostgreSQL store with pgvector.** On a SQLite store knowledge bases
|
||||
are off, and `llmgw ctl kb list` and the console say why.
|
||||
- `llmgw server init` and the high-availability database setup install
|
||||
pgvector and enable it.
|
||||
- A cluster set up before this build needs it by hand, once:
|
||||
1. Install the package on every database machine:
|
||||
`sudo apt-get install postgresql-16-pgvector` (use the server's
|
||||
version number).
|
||||
2. On the leader, enable it in llmgw's database:
|
||||
`sudo -u postgres psql -d llmgw -c 'CREATE EXTENSION vector'`.
|
||||
3. Restart the control planes: `llmgw restart`.
|
||||
- **An embedding model in the catalog**, with `capabilities.embeddings`.
|
||||
Two good choices handle Turkish and run on CPUs: `BAAI/bge-m3` (1024
|
||||
dimensions) and `Qwen/Qwen3-Embedding-0.6B`.
|
||||
|
||||
Add it as any model is added (**Serve a model** in the console, or the
|
||||
catalog), with its entry saying that it embeds:
|
||||
|
||||
```yaml
|
||||
capabilities:
|
||||
embeddings: true
|
||||
```
|
||||
|
||||
A collection stays with the model it was made with: vectors from two
|
||||
models cannot be compared.
|
||||
|
||||
## For administrators
|
||||
|
||||
### Make a collection
|
||||
|
||||
A collection belongs to one part of the tenancy tree: an organization, a
|
||||
team, a project or a person. Whoever is in the owner, at any level below it,
|
||||
may read it and add to it.
|
||||
|
||||
```sh
|
||||
llmgw ctl kb create team:Acme/hr policies bge-m3
|
||||
llmgw ctl kb create user:USER-ID notes bge-m3 --org Acme
|
||||
```
|
||||
|
||||
The model's dimension is asked of the model itself, so the model has to be
|
||||
able to serve when the collection is made. `--dims 1024` gives it instead.
|
||||
|
||||
In the console: **Knowledge**, then **A**.
|
||||
|
||||
### Who may read and add
|
||||
|
||||
The owner, and everything below it, may read and add. A **grant** gives the
|
||||
same to another part of the same organization, to a person, or to one key:
|
||||
|
||||
```sh
|
||||
llmgw ctl kb grant Acme/policies project:Acme/sales/crm # read
|
||||
llmgw ctl kb grant Acme/policies team:Acme/legal --write # read and add
|
||||
llmgw ctl kb revoke Acme/policies project:Acme/sales/crm
|
||||
```
|
||||
|
||||
How access works:
|
||||
|
||||
- **Grants add up down the tree.** A grant to a team reaches its projects
|
||||
and their keys.
|
||||
- **Nothing crosses organizations.** A grant to another organization's
|
||||
team, project, key or person is refused.
|
||||
- **A person's own key** (a key with no project) reads nothing by itself.
|
||||
The person names the project they act in (`llmgw-kb login` asks for it,
|
||||
or `--project`) and must be a member of it. Joining a second team
|
||||
therefore does not quietly widen what every key of theirs can read.
|
||||
- **No access looks like no collection.** A collection a key may not read
|
||||
is answered exactly as one that does not exist, so even its name does not
|
||||
leak.
|
||||
|
||||
### Answer from collections by default
|
||||
|
||||
Clients that cannot add a header still get answers from the documents when
|
||||
their part of the tree has a default:
|
||||
|
||||
```sh
|
||||
llmgw ctl kb default team:Acme/hr Acme/policies Acme/forms
|
||||
llmgw ctl kb default key:KEY-ID Acme/policies
|
||||
llmgw ctl kb default team:Acme/hr --clear
|
||||
```
|
||||
|
||||
- **Which default applies.** The narrowest one: a key's over its project's,
|
||||
a project's over its team's, a team's over its organization's.
|
||||
- **A request can override it.** A request that names collections uses
|
||||
those instead, and `X-Llmgw-Knowledge: none` turns the default off for
|
||||
that request.
|
||||
- **A default never refuses a request.** A collection the caller may not
|
||||
read is left out for them. If the search fails (the embedding model is
|
||||
down, say), the request is answered without passages and the failure is
|
||||
logged.
|
||||
- **Changes take effect within 30 seconds.**
|
||||
|
||||
In the console: **Knowledge**, Enter on a collection, then **Answer from it
|
||||
by default…**.
|
||||
|
||||
### Look after it
|
||||
|
||||
```sh
|
||||
llmgw ctl kb list # collections, owners, documents, grants
|
||||
llmgw ctl kb docs Acme/policies # its documents, who added each and when
|
||||
llmgw ctl kb rm Acme/policies old/form.pdf # one document out
|
||||
llmgw ctl kb search Acme/policies how many days of annual leave
|
||||
llmgw ctl kb delete Acme/policies # the collection, its documents and grants
|
||||
```
|
||||
|
||||
Rules that apply:
|
||||
|
||||
- A project, team or organization that owns a collection is not deleted
|
||||
until its collections are. This is the same rule as for one with members.
|
||||
- A person's own collections go with their account (`users delete`,
|
||||
`users forget`).
|
||||
- Every document added or removed is in the audit log (`llmgw ctl audit`),
|
||||
by name. The text is never recorded.
|
||||
|
||||
## For everyone: `llmgw-kb`
|
||||
|
||||
### Get it
|
||||
|
||||
Download `llmgw-kb` for your system from the release page, next to `llmgw`:
|
||||
|
||||
| System | File |
|
||||
|---|---|
|
||||
| Windows (Intel/AMD) | `llmgw-kb-VERSION-windows-amd64.exe` |
|
||||
| Windows (ARM) | `llmgw-kb-VERSION-windows-arm64.exe` |
|
||||
| macOS (Apple silicon) | `llmgw-kb-VERSION-darwin-arm64` |
|
||||
| macOS (Intel) | `llmgw-kb-VERSION-darwin-amd64` |
|
||||
| Linux | `llmgw-kb-VERSION-linux-amd64` / `-arm64` |
|
||||
|
||||
Rename it to `llmgw-kb` (`llmgw-kb.exe` on Windows) and put it somewhere on
|
||||
your PATH. You do not need to install anything else.
|
||||
|
||||
On macOS, a downloaded program that is not signed is stopped the first time.
|
||||
Either right-click it and choose Open, or run
|
||||
`xattr -d com.apple.quarantine llmgw-kb`.
|
||||
|
||||
### Log in once
|
||||
|
||||
```sh
|
||||
llmgw-kb login
|
||||
```
|
||||
|
||||
It asks for three things, checks them, and saves them where only you can
|
||||
read them:
|
||||
|
||||
- **The gateway's address.**
|
||||
- **Your API key.** It is typed without being shown.
|
||||
- **A project.** Needed only for a personal key.
|
||||
|
||||
Instead of logging in you can set `LLMGW_URL`, `LLMGW_API_KEY` and
|
||||
`LLMGW_PROJECT`.
|
||||
|
||||
```sh
|
||||
llmgw-kb collections # what your key may read, and which you may add to
|
||||
```
|
||||
|
||||
### Send your documents
|
||||
|
||||
```sh
|
||||
llmgw-kb sync policies ~/Documents/HR
|
||||
llmgw-kb sync policies "C:\Users\ayse\Documents\İK" --include "*.pdf,*.docx"
|
||||
llmgw-kb sync policies ./handbook ./forms --exclude "drafts/**"
|
||||
```
|
||||
|
||||
Run it as often as you like. The first run sends everything. Later runs
|
||||
send only what changed, and do not even read files whose size and time have
|
||||
not changed.
|
||||
|
||||
| Option | What it does |
|
||||
|---|---|
|
||||
| `--dry-run` | Reads and cuts everything, sends nothing, and lists what would leave the machine and how much. |
|
||||
| `--delete` | Also removes documents whose files are gone. Only those this machine sent, from these folders; never anything someone else sent. |
|
||||
| `--retry` | Reads again files skipped before, for example after installing LibreOffice. |
|
||||
| `--include`, `--exclude` | Glob patterns. `*.pdf` matches a file name; `reports/**/*.xlsx` matches a path. |
|
||||
| `--hidden` | Reads hidden files and folders too (left out by default). |
|
||||
| `-v` | Lists unchanged files too. |
|
||||
|
||||
Documents are named by their path under the folder, with forward slashes,
|
||||
on every system. With several folders, each folder's own name comes first.
|
||||
|
||||
```sh
|
||||
llmgw-kb status policies # what this machine sent and skipped, and why
|
||||
llmgw-kb history # the latest runs
|
||||
llmgw-kb formats # what can be read here
|
||||
```
|
||||
|
||||
### What it reads
|
||||
|
||||
| Read here | Formats |
|
||||
|---|---|
|
||||
| Documents | PDF (text), Word (`.docx`), Excel (`.xlsx`), PowerPoint (`.pptx`, with speaker notes), OpenDocument (`.odt`, `.ods`, `.odp`), RTF, EPUB |
|
||||
| Web and mail | HTML, MHTML, e-mail (`.eml`, with its attachments) |
|
||||
| Text | Markdown, plain text, CSV, JSON, YAML, XML, subtitles, notebooks (`.ipynb`), source code |
|
||||
| Archives | ZIP, with every file inside it read in turn |
|
||||
|
||||
| Read with a helper on your machine | Helper |
|
||||
|---|---|
|
||||
| Word, Excel and PowerPoint 97–2003 (`.doc`, `.xls`, `.ppt`), Pages, Numbers, Keynote and others | [LibreOffice](https://www.libreoffice.org) |
|
||||
| Images and scanned PDFs | [Tesseract](https://tesseract-ocr.github.io) (with the Turkish language data), plus Poppler's `pdftoppm` for PDFs |
|
||||
| PDFs the built-in reader cannot parse | Poppler's `pdftotext` |
|
||||
|
||||
`llmgw-kb formats` says which helpers it found. Text in UTF-8, UTF-16, or
|
||||
the Windows Turkish and Western code pages is read correctly.
|
||||
|
||||
A file that cannot be read is skipped with the reason, and the run goes on.
|
||||
Reasons include a locked Office file, a format with no helper, and a
|
||||
scanned page with no OCR.
|
||||
|
||||
### What leaves your machine
|
||||
|
||||
For each new or changed document, two things leave the machine, both over
|
||||
the gateway's TLS connection:
|
||||
|
||||
- **The passages, to be embedded.** They go through `/v1/embeddings` like
|
||||
any request of yours. Your key's limits, budget and guardrails apply.
|
||||
- **The document itself.** Its name, its file's SHA-256 digest and size,
|
||||
and the passages with their vectors.
|
||||
|
||||
The file itself is never sent. `--dry-run` shows exactly what would be.
|
||||
|
||||
### The ledger
|
||||
|
||||
`llmgw-kb` keeps a small SQLite database of what it did with every
|
||||
document: the file it came from, its size, time and digest, and whether it
|
||||
was sent, skipped or failed, with the reason.
|
||||
|
||||
| System | Where it is |
|
||||
|---|---|
|
||||
| Windows | `%AppData%\llmgw-kb\ledger.db` |
|
||||
| macOS | `~/Library/Application Support/llmgw-kb/ledger.db` |
|
||||
| Linux | `~/.config/llmgw-kb/ledger.db` |
|
||||
|
||||
Deleting the ledger is safe: the next run hashes every file once, finds what
|
||||
the collection already has, and sends nothing that is unchanged. `--ledger`
|
||||
uses another file.
|
||||
|
||||
## Asking questions of a collection
|
||||
|
||||
A chat request names the collections to answer from in a header. Any client
|
||||
that can add a header can do it: the OpenAI and Anthropic SDKs, Continue, or
|
||||
curl.
|
||||
|
||||
```sh
|
||||
curl https://llm.example.com/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LLMGW_API_KEY" \
|
||||
-H "X-Llmgw-Knowledge: politikalar" \
|
||||
-d '{"model": "qwen2.5-coder-14b-instruct",
|
||||
"messages": [{"role": "user", "content": "Kaç gün yıllık iznim var?"}]}'
|
||||
```
|
||||
|
||||
```python
|
||||
client.chat.completions.create(
|
||||
model="qwen2.5-coder-14b-instruct",
|
||||
messages=[{"role": "user", "content": "Kaç gün yıllık iznim var?"}],
|
||||
extra_headers={"X-Llmgw-Knowledge": "politikalar"},
|
||||
)
|
||||
```
|
||||
|
||||
What the gateway does with such a request:
|
||||
|
||||
1. **It finds the collections.** It looks up the ones named, by name or ID,
|
||||
among those the key may read. A collection it may not read is answered
|
||||
404, the same as one that does not exist.
|
||||
2. **It embeds the question.** The last user message is embedded by the
|
||||
collection's model, on the caller's own account: their limits, budget and
|
||||
guardrails apply.
|
||||
3. **It searches.** The collections are searched by meaning and by words
|
||||
(below).
|
||||
4. **It adds the passages.** The best passages are numbered and added to the
|
||||
system message, with an instruction to answer from them, cite them as
|
||||
[1], [2], and say so when they do not hold the answer. The system message
|
||||
is created if the request has none. Passages that would not fit the
|
||||
model's context are left out, the weakest first.
|
||||
5. **It judges and sends the request.** The request, passages included, then
|
||||
goes through the guardrails and to the model as usual. A document
|
||||
carrying instructions aimed at its reader is judged like anything else
|
||||
the model reads.
|
||||
|
||||
On the portal's chat page, **Documents** picks a collection the person may
|
||||
read in the chosen project, the default, or none. Under each answer, the
|
||||
sources it was given are listed by number.
|
||||
|
||||
How it is shaped:
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| `X-Llmgw-Knowledge` | Collections, comma separated; `none` for no passages, whatever the default. |
|
||||
| `X-Llmgw-Knowledge-K` | How many passages, 1 to 20; 5 by default. |
|
||||
| `X-Llmgw-Project` | For a personal key: the project it acts in. |
|
||||
| `x-llmgw-sources` (answer) | The passages given, as `document#chunk`, in their numbered order; empty when nothing matched. |
|
||||
|
||||
Behaviour to know about:
|
||||
|
||||
- **No passages matched.** The model is told that nothing in the documents
|
||||
answers the question, and asked to say so rather than answer from general
|
||||
knowledge.
|
||||
- **Other APIs.** Requests through `/v1/responses` and `/v1/messages`
|
||||
(Codex, Claude Code) work the same way, since they are answered as chat
|
||||
completions.
|
||||
- **Separate prefix caches.** On vLLM replicas, a request given passages
|
||||
carries a `cache_salt` derived from its organization. One organization's
|
||||
passages then never share cached blocks with another's, so how fast an
|
||||
answer starts says nothing about what another tenant asked.
|
||||
- **No answer cache.** Answers from knowledge bases are never served from
|
||||
the answer cache, because two keys of one project can read different
|
||||
collections.
|
||||
|
||||
### How a passage is chosen
|
||||
|
||||
Each passage scores the following, and the highest scores win:
|
||||
|
||||
- its similarity in meaning to the question (cosine, 0 to 1), plus
|
||||
- up to 0.3 for the share of the question's words it contains.
|
||||
|
||||
Words almost every passage has ("ne", "için", "the", "of") are not counted.
|
||||
The meaning does most of the work: "how much leave do fathers get" finds the
|
||||
paternity page, not the annual-leave page that shares more of its words.
|
||||
The word match helps with what meaning cannot carry, such as a form number,
|
||||
a product code or a name.
|
||||
|
||||
Metrics: `llmgw_kb_retrievals_total{result}` (found, nothing, refused) and
|
||||
`llmgw_kb_retrieval_seconds`.
|
||||
|
||||
## How it is kept apart
|
||||
|
||||
Passages live in PostgreSQL, partitioned by organization, with an index per
|
||||
collection. Every search and every write names the organization and the
|
||||
collections the caller may use. On top of that, row-level security on the
|
||||
table checks the same thing again, so a query that forgot its filter
|
||||
returns nothing rather than another tenant's passages.
|
||||
@@ -26,6 +26,8 @@ rather than repeated.
|
||||
| `queue_capacity_per_band` | `int` | [Admission](#admission) |
|
||||
| `queue_capacity_total` | `int` | [Admission](#admission) |
|
||||
| `interval` | `time.Duration` | [Admission](#admission) |
|
||||
| `default` | `map[string]any` | [AgentSettings](#agentsettings) |
|
||||
| `nodes` | `map[string]map[string]any` | [AgentSettings](#agentsettings) |
|
||||
| `enabled` | `bool` | [Auth](#auth) |
|
||||
| `store` | `string` | [Auth](#auth) |
|
||||
| `cache_ttl` | `time.Duration` | [Auth](#auth) |
|
||||
@@ -294,6 +296,7 @@ rather than repeated.
|
||||
| `mail` | `mail.Config` | [Server](#server) |
|
||||
| `serving` | `ServingSettings` | [Server](#server) |
|
||||
| `firewall` | `FirewallSettings` | [Server](#server) |
|
||||
| `agents` | `AgentSettings` | [Server](#server) |
|
||||
| `firewall_networks` | `[]string` | [Server](#server) |
|
||||
| `backup_recipients` | `[]string` | [Server](#server) |
|
||||
| `backup_copy_to` | `[]string` | [Server](#server) |
|
||||
@@ -469,6 +472,21 @@ Zero takes half of all the bands' capacities together.
|
||||
|
||||
`time.Duration`
|
||||
|
||||
## AgentSettings
|
||||
|
||||
### `default`
|
||||
|
||||
`map[string]any`
|
||||
|
||||
Default is agent.yaml's keys for every node.
|
||||
|
||||
### `nodes`
|
||||
|
||||
`map[string]map[string]any`
|
||||
|
||||
Nodes are agent.yaml's keys for one node each, by node id, over
|
||||
Default.
|
||||
|
||||
## Auth
|
||||
|
||||
### `enabled`
|
||||
@@ -2520,6 +2538,18 @@ Firewall turns llmgw's firewall on every machine on or off.
|
||||
|
||||
Takes a [FirewallSettings](#firewallsettings) block.
|
||||
|
||||
### `agents`
|
||||
|
||||
`AgentSettings`
|
||||
|
||||
Agents are the nodes' agent settings: agent.yaml's keys but those an
|
||||
agent needs to start, for every node (default) and for one
|
||||
(nodes.<id>). Sent to the agents as they change; an agent whose
|
||||
settings changed starts again on them, its engines serving
|
||||
throughout.
|
||||
|
||||
Takes a [AgentSettings](#agentsettings) block.
|
||||
|
||||
### `firewall_networks`
|
||||
|
||||
`[]string`
|
||||
|
||||
@@ -18,6 +18,7 @@ require (
|
||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||
github.com/knadh/koanf/providers/file v1.2.1
|
||||
github.com/knadh/koanf/v2 v2.3.6
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0
|
||||
github.com/mattn/go-isatty v0.0.24
|
||||
github.com/prometheus/client_golang v1.24.1
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
@@ -34,6 +35,7 @@ require (
|
||||
golang.org/x/net v0.59.0
|
||||
golang.org/x/oauth2 v0.37.0
|
||||
golang.org/x/sync v0.23.0
|
||||
golang.org/x/term v0.46.0
|
||||
golang.org/x/text v0.42.0
|
||||
google.golang.org/protobuf v1.36.12
|
||||
modernc.org/sqlite v1.59.0
|
||||
|
||||
@@ -95,6 +95,8 @@ github.com/knadh/koanf/v2 v2.3.6 h1:JoQPSJmvS4aP0xNc8xMDr5tcrkSEInL23/Il7pITAKo=
|
||||
github.com/knadh/koanf/v2 v2.3.6/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28=
|
||||
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
|
||||
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0 h1:7Q+xNAZFmnfYOMweHN3c/PDFUKKfY1pVJ26K++QvVfU=
|
||||
github.com/ledongthuc/pdf v0.0.0-20260907135840-6c8c28e0e8a0/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||
|
||||
@@ -136,6 +136,10 @@ type Agent struct {
|
||||
// by mu.
|
||||
kv *kvevents.Mirror
|
||||
kvSubs map[string]*kvevents.Subscriber
|
||||
|
||||
// settingsCheck says whether the cluster's settings for this agent
|
||||
// would load beside its agent.yaml; see SetSettingsCheck.
|
||||
settingsCheck func([]byte) error
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -880,6 +884,8 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
|
||||
a.setControlPlanes(m.ControlPlanes)
|
||||
case *llmgwv1.ControlMessage_Firewall:
|
||||
go a.applyFirewall(m.Firewall)
|
||||
case *llmgwv1.ControlMessage_AgentSettings:
|
||||
go a.applySettings(m.AgentSettings)
|
||||
case *llmgwv1.ControlMessage_RestartAgent:
|
||||
go func() {
|
||||
a.log.Info("restarting, as the control plane asks; the engines keep serving",
|
||||
|
||||
@@ -103,6 +103,11 @@ type Config struct {
|
||||
Tracing obs.Tracing `koanf:"tracing"`
|
||||
|
||||
Version string
|
||||
|
||||
// Settings is the cluster's settings for this agent it was started on,
|
||||
// as the control plane sent them: what a newly sent version is
|
||||
// compared with.
|
||||
Settings string `koanf:"-"`
|
||||
}
|
||||
|
||||
// DefaultConfig returns sensible defaults.
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
// The cluster's settings for this agent (BACKLOG E3): what agent.yaml can
|
||||
// say but the agent's id, control planes, address and certificate, kept by
|
||||
// the control planes and sent over the stream. Kept here beside agent.yaml,
|
||||
// which then holds only what the agent needs to start; at start the two are
|
||||
// read together, these over the file.
|
||||
|
||||
// SettingsFile keeps the settings the control plane last sent.
|
||||
var SettingsFile = "/var/lib/llmgw-agent/agent-settings.yaml"
|
||||
|
||||
// ReadSettings is what SettingsFile holds, nil when nothing.
|
||||
func ReadSettings() []byte {
|
||||
b, err := os.ReadFile(SettingsFile)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// SetSettingsCheck says how to tell whether settings would load beside this
|
||||
// agent's own file; settings that would not are not taken.
|
||||
func (a *Agent) SetSettingsCheck(f func([]byte) error) { a.settingsCheck = f }
|
||||
|
||||
// applySettings takes the settings the control plane sent: nothing when they
|
||||
// are what this agent runs on, and otherwise kept and started again on.
|
||||
func (a *Agent) applySettings(m *llmgwv1.AgentSettings) {
|
||||
next := m.GetYaml()
|
||||
if next == a.cfg.Settings {
|
||||
return
|
||||
}
|
||||
if a.settingsCheck != nil {
|
||||
if err := a.settingsCheck([]byte(next)); err != nil {
|
||||
// Kept running on what it has: settings that will not load
|
||||
// here must not take the node out of service.
|
||||
a.log.Error("the cluster's settings for this agent do not load beside its agent.yaml; running on what it has",
|
||||
"err", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(SettingsFile), 0o755); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
tmp := SettingsFile + ".new"
|
||||
if err := os.WriteFile(tmp, []byte(next), 0o600); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
if err := os.Rename(tmp, SettingsFile); err != nil {
|
||||
a.log.Error("could not keep the cluster's settings for this agent", "err", err)
|
||||
return
|
||||
}
|
||||
a.log.Info("the cluster's settings for this agent changed; starting again on them, the engines keep serving")
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
}
|
||||
|
||||
// WatchFile starts the agent again when its configuration file changes and
|
||||
// still loads: agent.yaml is read only at start, and a change to it that
|
||||
// waited for somebody to restart the agent was a change nobody saw take
|
||||
// effect.
|
||||
func (a *Agent) WatchFile(ctx context.Context, path string, loads func() error) {
|
||||
if path == "" {
|
||||
return
|
||||
}
|
||||
sum := func() [32]byte {
|
||||
b, _ := os.ReadFile(path) //nolint:gosec // the agent's own configuration
|
||||
return sha256.Sum256(b)
|
||||
}
|
||||
was := sum()
|
||||
t := time.NewTicker(watchEvery)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
now := sum()
|
||||
if now == was {
|
||||
continue
|
||||
}
|
||||
if err := loads(); err != nil {
|
||||
// Once per change: an editor saving halfway must not stop the
|
||||
// agent, and a broken file is reported, not restarted into.
|
||||
a.log.Error("agent.yaml changed and does not load; running on what it had", "path", path, "err", err)
|
||||
was = now
|
||||
continue
|
||||
}
|
||||
a.log.Info("agent.yaml changed; starting again on it, the engines keep serving", "path", path)
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
exitForUpgrade()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// watchEvery is how often agent.yaml is looked at.
|
||||
var watchEvery = 10 * time.Second
|
||||
@@ -0,0 +1,88 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
func settingsAgent(t *testing.T, running string) (*Agent, chan struct{}) {
|
||||
t.Helper()
|
||||
oldFile, oldExit := SettingsFile, exitForUpgrade
|
||||
SettingsFile = filepath.Join(t.TempDir(), "agent-settings.yaml")
|
||||
exited := make(chan struct{}, 1)
|
||||
exitForUpgrade = func() { exited <- struct{}{} }
|
||||
t.Cleanup(func() { SettingsFile, exitForUpgrade = oldFile, oldExit })
|
||||
cfg := Config{Server: "127.0.0.1:7443", NodeID: "n1", WeightsCache: t.TempDir(), EnginesDir: t.TempDir(), Settings: running}
|
||||
cfg.Superv.WorkDir = t.TempDir()
|
||||
a, err := New(cfg, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return a, exited
|
||||
}
|
||||
|
||||
// exitedSoon is whether the agent exited within a second: it waits a
|
||||
// moment for its log line to leave first.
|
||||
func exitedSoon(ch chan struct{}) bool {
|
||||
select {
|
||||
case <-ch:
|
||||
return true
|
||||
case <-time.After(time.Second):
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnAgentOnItsSettingsStaysAndOneOffThemStartsAgain(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "heartbeat_interval: 5s\n")
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "heartbeat_interval: 5s\n"})
|
||||
if exitedSoon(exited) || ReadSettings() != nil {
|
||||
t.Error("settings it already runs on restarted it")
|
||||
}
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "heartbeat_interval: 2s\n"})
|
||||
if !exitedSoon(exited) || string(ReadSettings()) != "heartbeat_interval: 2s\n" {
|
||||
t.Errorf("new settings: kept %q", ReadSettings())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsThatWouldNotLoadAreNotTaken(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "")
|
||||
a.SetSettingsCheck(func([]byte) error { return errors.New("fabric ib-a lists no HCAs") })
|
||||
a.applySettings(&llmgwv1.AgentSettings{Yaml: "fabrics: [{name: ib-a}]\n"})
|
||||
if exitedSoon(exited) || ReadSettings() != nil {
|
||||
t.Error("settings that do not load were taken")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAChangedAgentYamlThatLoadsStartsTheAgentAgain(t *testing.T) {
|
||||
a, exited := settingsAgent(t, "")
|
||||
old := watchEvery
|
||||
watchEvery = 10 * time.Millisecond
|
||||
t.Cleanup(func() { watchEvery = old })
|
||||
p := filepath.Join(t.TempDir(), "agent.yaml")
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\n"), 0o600)
|
||||
var loads atomic.Bool
|
||||
check := func() error {
|
||||
if loads.Load() {
|
||||
return nil
|
||||
}
|
||||
return errors.New("broken")
|
||||
}
|
||||
go a.WatchFile(t.Context(), p, check)
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
// Saved broken: reported, not restarted into.
|
||||
_ = os.WriteFile(p, []byte("node_id: [\n"), 0o600)
|
||||
if exitedSoon(exited) {
|
||||
t.Fatal("a broken agent.yaml restarted the agent")
|
||||
}
|
||||
loads.Store(true)
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\nlabels: {a: b}\n"), 0o600)
|
||||
if !exitedSoon(exited) {
|
||||
t.Error("a changed agent.yaml did not restart the agent")
|
||||
}
|
||||
}
|
||||
+22
-1
@@ -40,10 +40,23 @@ func newAgentCmd() *cobra.Command {
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
log := logger()
|
||||
|
||||
cfg, err := config.LoadAgent(configPath)
|
||||
// The cluster's settings for this agent over its file, as the
|
||||
// control plane last sent them. Settings that no longer load
|
||||
// beside the file are set aside rather than keeping the agent
|
||||
// from starting: the control plane sends them again, and says
|
||||
// why they are refused.
|
||||
settings := agent.ReadSettings()
|
||||
cfg, err := config.LoadAgentWith(configPath, settings)
|
||||
if err != nil && settings != nil {
|
||||
log.Warn("the cluster's settings kept for this agent do not load beside its agent.yaml; starting without them",
|
||||
"err", err)
|
||||
settings = nil
|
||||
cfg, err = config.LoadAgent(configPath)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cfg.Settings = string(settings)
|
||||
if server != "" {
|
||||
cfg.Server = server
|
||||
}
|
||||
@@ -89,9 +102,17 @@ func newAgentCmd() *cobra.Command {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
a.SetSettingsCheck(func(b []byte) error {
|
||||
_, err := config.LoadAgentWith(configPath, b)
|
||||
return err
|
||||
})
|
||||
|
||||
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
go a.WatchFile(ctx, configPath, func() error {
|
||||
_, err := config.LoadAgentWith(configPath, agent.ReadSettings())
|
||||
return err
|
||||
})
|
||||
|
||||
log.Info("agent starting",
|
||||
"node", cfg.NodeID, "server", cfg.Server,
|
||||
|
||||
@@ -241,6 +241,8 @@ var liveSettings = []string{
|
||||
"serving",
|
||||
// Sent to the nodes and this machine's firewall unit as they change.
|
||||
"firewall", "firewall_networks",
|
||||
// Sent to the nodes' agents as they change.
|
||||
"agents",
|
||||
// Read by server init, node add and backups, not by the server.
|
||||
"backup_recipients", "backup_copy_to",
|
||||
}
|
||||
|
||||
@@ -180,6 +180,7 @@ func newCtlCmd() *cobra.Command {
|
||||
newCapacityCmd(client, &asJSON),
|
||||
newQuotaCmd(client, &asJSON),
|
||||
newGuardCmd(client, &asJSON),
|
||||
newKBCmd(client, &asJSON),
|
||||
newConfigCmd(client, &asJSON),
|
||||
newAuditCmd(client, &asJSON),
|
||||
newNotifyCmd(client, &asJSON),
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
func newKBCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
|
||||
c := &cobra.Command{
|
||||
Use: "kb",
|
||||
Short: "Knowledge bases: collections of documents requests are answered from",
|
||||
Long: "A collection belongs to an organization, team, project or person, and is\n" +
|
||||
"embedded by one model of the catalog. People put documents into it from\n" +
|
||||
"their own machines with llmgw-kb (its own download, for Windows, macOS and\n" +
|
||||
"Linux), which sends text and vectors, never the files.\n\n" +
|
||||
"Who may read or add to a collection: its owner and everything below it,\n" +
|
||||
"and whoever it is granted to (grant), never outside its organization.\n" +
|
||||
"Targets are written org:Acme, team:Acme/hr, project:Acme/hr/docs (or with\n" +
|
||||
"IDs), key:KEY-ID or user:USER-ID. A collection is named by its ID, its name\n" +
|
||||
"when that is unique, or ORG/NAME.\n\n" +
|
||||
"Needs the PostgreSQL store with pgvector; `llmgw ctl kb list` says why not.",
|
||||
}
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "list",
|
||||
Short: "Every collection: owner, model, documents and grants",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: func(cmd *cobra.Command, _ []string) error {
|
||||
var v adminapi.KBView
|
||||
if err := client().get(cmd.Context(), "/api/v1/kb", &v); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if !v.Available {
|
||||
fmt.Fprintln(out, "Knowledge bases are off: "+v.Reason)
|
||||
return nil
|
||||
}
|
||||
if len(v.Collections) == 0 {
|
||||
fmt.Fprintln(out, "No collections yet (llmgw ctl kb create).")
|
||||
return nil
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "COLLECTION\tOWNER\tMODEL\tDOCUMENTS\tCHUNKS\tGRANTED TO")
|
||||
for _, c := range v.Collections {
|
||||
var grants []string
|
||||
for _, g := range c.Grants {
|
||||
grants = append(grants, or(g.Name, g.Target)+" ("+g.Perm+")")
|
||||
}
|
||||
fmt.Fprintf(w, "%s/%s\t%s\t%s (%d)\t%d\t%d\t%s\n", or(c.OrgName, c.Org), c.Name, or(c.OwnerName, c.Owner),
|
||||
c.EmbedModel, c.Dims, c.Documents, c.Chunks, strings.Join(grants, ", "))
|
||||
}
|
||||
if err := w.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(v.Defaults) > 0 {
|
||||
fmt.Fprintln(out, "\nAnswered from by default, when a request names no collection:")
|
||||
w = tab()
|
||||
for _, d := range v.Defaults {
|
||||
fmt.Fprintf(w, " %s\t%s\n", or(d.Name, d.Target), strings.Join(d.Collections, ", "))
|
||||
}
|
||||
return w.Flush()
|
||||
}
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
var org string
|
||||
var dims int
|
||||
create := &cobra.Command{
|
||||
Use: "create OWNER NAME MODEL",
|
||||
Short: "Make a collection owned by an organization, team, project or person",
|
||||
Long: "MODEL is a catalog model with capabilities.embeddings. Its dimension is asked\n" +
|
||||
"of it, so it has to be able to serve now; --dims gives it instead.",
|
||||
Example: " llmgw ctl kb create team:Acme/hr policies bge-m3\n" +
|
||||
" llmgw ctl kb create user:USER-ID notes bge-m3 --org Acme",
|
||||
Args: cobra.ExactArgs(3),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
body, _ := json.Marshal(adminapi.KBCreateRequest{Owner: args[0], Name: args[1], Model: args[2], Org: org, Dims: dims})
|
||||
var v adminapi.KBCollectionView
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/kb/collections", json.RawMessage(body), &v); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "collection %s (%s), embedded by %s in %d dimensions\n", v.Name, v.ID, v.EmbedModel, v.Dims)
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "people with access fill it with: llmgw-kb sync %s FOLDER\n", v.Name)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
create.Flags().StringVar(&org, "org", "", "the organization of a person's collection")
|
||||
create.Flags().IntVar(&dims, "dims", 0, "the model's embedding dimension, instead of asking it")
|
||||
c.AddCommand(create)
|
||||
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "delete COLLECTION",
|
||||
Short: "Remove a collection, its documents and its grants",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if err := client().do(cmd.Context(), "DELETE", "/api/v1/kb/collections/"+url.PathEscape(args[0]), nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "collection %s removed\n", args[0])
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
for _, verb := range []struct {
|
||||
use, short, method string
|
||||
perm bool
|
||||
}{
|
||||
{"grant", "Let a scope and everything below it read a collection, or with --write add to it", "POST", true},
|
||||
{"revoke", "Take a grant away", "DELETE", false},
|
||||
} {
|
||||
var write bool
|
||||
cmd := &cobra.Command{
|
||||
Use: verb.use + " COLLECTION TARGET",
|
||||
Short: verb.short,
|
||||
Example: " llmgw ctl kb " + verb.use + " Acme/policies project:Acme/sales/crm",
|
||||
Args: cobra.ExactArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
req := adminapi.KBGrantRequest{Collection: args[0], Target: args[1], Perm: "read"}
|
||||
if write {
|
||||
req.Perm = "write"
|
||||
}
|
||||
body, _ := json.Marshal(req)
|
||||
if err := client().do(cmd.Context(), verb.method, "/api/v1/kb/grants", json.RawMessage(body), nil); err != nil {
|
||||
return err
|
||||
}
|
||||
if verb.perm {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s may %s %s\n", args[1], req.Perm, args[0])
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s no longer granted %s\n", args[1], args[0])
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if verb.perm {
|
||||
cmd.Flags().BoolVar(&write, "write", false, "add and remove documents too")
|
||||
}
|
||||
c.AddCommand(cmd)
|
||||
}
|
||||
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "docs COLLECTION",
|
||||
Short: "A collection's documents",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
var docs []adminapi.KBDocumentView
|
||||
if err := client().get(cmd.Context(), "/api/v1/kb/collections/"+url.PathEscape(args[0])+"/documents", &docs); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(docs)
|
||||
}
|
||||
w := tab()
|
||||
fmt.Fprintln(w, "DOCUMENT\tCHUNKS\tSIZE\tADDED\tBY")
|
||||
for _, d := range docs {
|
||||
fmt.Fprintf(w, "%s\t%d\t%s\t%s\t%s\n", d.Name, d.Chunks, humanBytes(d.Bytes), d.AddedAt.Local().Format("2006-01-02 15:04"), d.AddedBy)
|
||||
}
|
||||
return w.Flush()
|
||||
},
|
||||
})
|
||||
c.AddCommand(&cobra.Command{
|
||||
Use: "rm COLLECTION DOCUMENT",
|
||||
Short: "Remove one document from a collection",
|
||||
Args: cobra.ExactArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
path := "/api/v1/kb/collections/" + url.PathEscape(args[0]) + "/documents?name=" + url.QueryEscape(args[1])
|
||||
if err := client().do(cmd.Context(), "DELETE", path, nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s removed from %s\n", args[1], args[0])
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
var clear bool
|
||||
def := &cobra.Command{
|
||||
Use: "default TARGET [COLLECTION...]",
|
||||
Short: "Answer a part of the tree's chat requests from collections, when they name none",
|
||||
Long: "Requests from TARGET and everything below it are then answered from these\n" +
|
||||
"collections without naming them (X-Llmgw-Knowledge). The narrowest default\n" +
|
||||
"wins: a key's over its project's, a project's over its team's. A request's own\n" +
|
||||
"header overrides it, and X-Llmgw-Knowledge: none turns it off. A collection a\n" +
|
||||
"caller may not read is left out for them, and a search that fails answers\n" +
|
||||
"without passages rather than refusing. In force within 30 seconds.",
|
||||
Example: " llmgw ctl kb default team:Acme/hr Acme/policies Acme/forms\n" +
|
||||
" llmgw ctl kb default team:Acme/hr --clear",
|
||||
Args: cobra.MinimumNArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if clear != (len(args) == 1) {
|
||||
return fmt.Errorf("name collections, or --clear")
|
||||
}
|
||||
body, _ := json.Marshal(adminapi.KBDefaultRequest{Target: args[0], Collections: args[1:]})
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/kb/defaults", json.RawMessage(body), nil); err != nil {
|
||||
return err
|
||||
}
|
||||
if clear {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s answers from no collection by default\n", args[0])
|
||||
} else {
|
||||
fmt.Fprintf(cmd.OutOrStdout(), "%s answers from %s by default\n", args[0], strings.Join(args[1:], ", "))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
def.Flags().BoolVar(&clear, "clear", false, "no default any more")
|
||||
c.AddCommand(def)
|
||||
|
||||
var k int
|
||||
search := &cobra.Command{
|
||||
Use: "search COLLECTION QUERY...",
|
||||
Short: "Try a search of a collection, as an administrator (no tenancy applies)",
|
||||
Args: cobra.MinimumNArgs(2),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
body, _ := json.Marshal(adminapi.KBSearchRequest{Collection: args[0], Query: strings.Join(args[1:], " "), K: k})
|
||||
var hits []adminapi.KBHitView
|
||||
if err := client().do(cmd.Context(), "POST", "/api/v1/kb/search", json.RawMessage(body), &hits); err != nil {
|
||||
return err
|
||||
}
|
||||
if *asJSON {
|
||||
return json.NewEncoder(cmd.OutOrStdout()).Encode(hits)
|
||||
}
|
||||
out := cmd.OutOrStdout()
|
||||
if len(hits) == 0 {
|
||||
fmt.Fprintln(out, "Nothing found.")
|
||||
}
|
||||
for i, h := range hits {
|
||||
fmt.Fprintf(out, "%d. %s #%d\n%s\n\n", i+1, h.Document, h.Ord, indent(strings.TrimSpace(h.Text), " "))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
search.Flags().IntVarP(&k, "k", "k", 5, "how many chunks")
|
||||
c.AddCommand(search)
|
||||
return c
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"crypto/rand"
|
||||
"crypto/tls"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
@@ -262,6 +263,15 @@ func newServerCmd() *cobra.Command {
|
||||
return &llmgwv1.Firewall{Enabled: c.Firewall.Enabled, Networks: c.FirewallNetworks,
|
||||
EnginePortsFrom: int32(from), EnginePortsTo: int32(to)} //nolint:gosec // ports
|
||||
})
|
||||
// What each node's agent runs on, as the settings say now.
|
||||
svc.SetAgentSettings(func(node string) string {
|
||||
b, err := config.Marshal(current.Load().Agents.For(node))
|
||||
if err != nil {
|
||||
log.Warn("could not write a node's agent settings", "node", node, "err", err)
|
||||
return ""
|
||||
}
|
||||
return string(b)
|
||||
})
|
||||
ctl.Reconciler().SetNotifier(notifier)
|
||||
admin = ctl
|
||||
agentPath, agentHandler = svc.Handler()
|
||||
@@ -374,10 +384,15 @@ func newServerCmd() *cobra.Command {
|
||||
api.SetAgentRestarter(agentSvc.RestartAgents)
|
||||
}
|
||||
api.SetAudit(st)
|
||||
// Knowledge bases, with the cluster's own models to embed a
|
||||
// probe (a collection's dimension) and an operator's search.
|
||||
api.SetKB(&adminapi.KBAdmin{Store: st, Embed: embedder(rt.Replay)})
|
||||
// The cluster's settings, each change checked against this
|
||||
// machine's own keys before it is kept.
|
||||
api.SetSettings(&adminapi.Settings{Store: st, LocalKeys: config.LocalKeys(),
|
||||
Check: func(body []byte) error { return checkClusterSettings(configPath, body) }})
|
||||
// And the same document on the portal, for administrators.
|
||||
rt.SetSettingsCheck(func(body []byte) error { return checkClusterSettings(configPath, body) }, config.LocalKeys())
|
||||
}
|
||||
// Catalog changes need somewhere to persist, so this is wired
|
||||
// only when there is a store. Without one, the endpoints say so
|
||||
@@ -713,6 +728,10 @@ func newServerCmd() *cobra.Command {
|
||||
agentSvc.BroadcastFirewall()
|
||||
}
|
||||
}
|
||||
if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "agents") }) && agentSvc != nil {
|
||||
// A moment apart per node, so off this goroutine.
|
||||
go agentSvc.BroadcastAgentSettings()
|
||||
}
|
||||
if slices.ContainsFunc(changed, func(c string) bool { return strings.HasPrefix(c, "guard") }) {
|
||||
if gate == nil {
|
||||
// Built at start, and only when there are policies.
|
||||
@@ -1183,3 +1202,40 @@ func federationID(ctx context.Context, st *storepkg.DB) (string, error) {
|
||||
// more of a secret than needed.
|
||||
return secret[:16], nil
|
||||
}
|
||||
|
||||
// embedder has a model embed one text through the gateway itself, the way a
|
||||
// kept request is replayed: for the administrator's knowledge-base
|
||||
// endpoints, not for anybody's traffic.
|
||||
func embedder(replay func(ctx context.Context, path, model string, body []byte) (int, []byte, time.Duration, error)) func(context.Context, string, string) ([]float32, error) {
|
||||
return func(ctx context.Context, model, text string) ([]float32, error) {
|
||||
body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}})
|
||||
status, out, _, err := replay(ctx, "/v1/embeddings", model, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
msg := strings.TrimSpace(string(out))
|
||||
var e struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(out, &e) == nil && e.Error.Message != "" {
|
||||
msg = e.Error.Message
|
||||
}
|
||||
return nil, fmt.Errorf("%d: %s", status, msg)
|
||||
}
|
||||
var res struct {
|
||||
Data []struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(out, &res); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(res.Data) != 1 || len(res.Data[0].Embedding) == 0 {
|
||||
return nil, fmt.Errorf("%s answered no embedding", model)
|
||||
}
|
||||
return res.Data[0].Embedding, nil
|
||||
}
|
||||
}
|
||||
@@ -370,6 +370,8 @@ func (in *serverInit) setupPostgres(ctx context.Context) (string, error) {
|
||||
in.done("PostgreSQL sized for a machine it shares (shared_buffers 96MB, 40 connections)")
|
||||
}
|
||||
|
||||
in.setupPgvector(ctx, discovery.DiscoverHost())
|
||||
|
||||
// Proof the address works, before anything is moved into it.
|
||||
address := "@" + in.dsnPath()
|
||||
s, err := store.Open(ctx, address)
|
||||
@@ -482,8 +484,17 @@ func lastLines(b []byte, n int) string {
|
||||
|
||||
// psql runs one statement as the postgres superuser, over the local socket.
|
||||
func (in *serverInit) psql(ctx context.Context, query string) (string, error) {
|
||||
cmd := exec.CommandContext(ctx, "runuser", "-u", "postgres", "--",
|
||||
in.pg.tool("psql"), "-v", "ON_ERROR_STOP=1", "-qAtX", "-c", query)
|
||||
return in.psqlIn(ctx, "", query)
|
||||
}
|
||||
|
||||
// psqlIn runs a statement as postgres in one database, or in postgres's
|
||||
// default one when db is "".
|
||||
func (in *serverInit) psqlIn(ctx context.Context, db, query string) (string, error) {
|
||||
args := []string{"-v", "ON_ERROR_STOP=1", "-qAtX", "-c", query}
|
||||
if db != "" {
|
||||
args = append(args, "-d", db)
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, "runuser", append([]string{"-u", "postgres", "--", in.pg.tool("psql")}, args...)...)
|
||||
// postgres cannot enter root's home, and psql says so on every call.
|
||||
cmd.Dir = "/"
|
||||
out, err := cmd.CombinedOutput()
|
||||
@@ -538,3 +549,50 @@ func (in *serverInit) stopService(ctx context.Context) (bool, error) {
|
||||
in.done("llmgw-server stopped while its store moves")
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// pgvectorPackages are the package names of pgvector for a PostgreSQL major
|
||||
// version: Debian's and Ubuntu's (and PGDG's apt repository) name it after the
|
||||
// server, PGDG's rpm repository after the version alone.
|
||||
func pgvectorPackages(major int) map[string][]string {
|
||||
return map[string][]string{
|
||||
"deb": {fmt.Sprintf("postgresql-%d-pgvector", major)},
|
||||
"rpm": {fmt.Sprintf("pgvector_%d", major)},
|
||||
}
|
||||
}
|
||||
|
||||
// setupPgvector installs pgvector and enables it in llmgw's database, for
|
||||
// knowledge bases. Not a reason to stop: without it everything else works,
|
||||
// and the store says why knowledge bases are off when somebody asks for one.
|
||||
func (in *serverInit) setupPgvector(ctx context.Context, h discovery.Host) {
|
||||
warn := func(format string, a ...any) {
|
||||
fmt.Fprintf(in.out, " !! pgvector: "+format+"; knowledge bases stay off\n", a...)
|
||||
}
|
||||
if on, err := in.psqlIn(ctx, pgDatabase, "SELECT 1 FROM pg_extension WHERE extname = 'vector'"); err == nil && on == "1" {
|
||||
in.kept("pgvector in database %s", pgDatabase)
|
||||
return
|
||||
}
|
||||
have, err := in.psqlBool(ctx, "SELECT 1 FROM pg_available_extensions WHERE name = 'vector'")
|
||||
if err != nil {
|
||||
warn("%v", err)
|
||||
return
|
||||
}
|
||||
if !have {
|
||||
num, err := in.psql(ctx, "SHOW server_version_num")
|
||||
n, convErr := strconv.Atoi(num)
|
||||
if err != nil || convErr != nil || h.Family() == "" {
|
||||
warn("cannot tell which package to install (version %q, %s)", num, h.OSID)
|
||||
return
|
||||
}
|
||||
pkgs := pgvectorPackages(n / 10000)
|
||||
pi := &provision.Installer{Host: h, Log: slog.Default()}
|
||||
if err := pi.InstallPackages(ctx, pkgs); err != nil {
|
||||
warn("installing %s: %v", strings.Join(pkgs[h.Family()], " "), err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := in.psqlIn(ctx, pgDatabase, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
|
||||
warn("%v", err)
|
||||
return
|
||||
}
|
||||
in.done("pgvector enabled in database %s", pgDatabase)
|
||||
}
|
||||
@@ -45,6 +45,7 @@ const (
|
||||
cUsage
|
||||
cQuotas
|
||||
cGuard
|
||||
cKB
|
||||
cFederation
|
||||
cMirror
|
||||
cJobs
|
||||
@@ -68,7 +69,7 @@ const (
|
||||
cCount
|
||||
)
|
||||
|
||||
var consoleLabels = [cCount]string{"Overview", "Monitor", "Nodes", "GPUs", "Models", "Replicas", "Usage", "Quotas", "Guardrails", "Federation", "Mirror", "Jobs", "Logs", "Accounts", "API keys", "Organizations", "Audit", "Manual",
|
||||
var consoleLabels = [cCount]string{"Overview", "Monitor", "Nodes", "GPUs", "Models", "Replicas", "Usage", "Quotas", "Guardrails", "Knowledge", "Federation", "Mirror", "Jobs", "Logs", "Accounts", "API keys", "Organizations", "Audit", "Manual",
|
||||
"server.yaml", "Add a node", "Add a control plane", "Replicate the store", "Serve a model", "Serve the guardrails' judge", "Upgrade the cluster",
|
||||
"Restart all replicas", "Back up now", "Wipe this machine", "Leave"}
|
||||
|
||||
@@ -177,6 +178,7 @@ type consoleModel struct {
|
||||
quotas []adminapi.QuotaView
|
||||
quotasErr error
|
||||
guard guardState
|
||||
kb kbState
|
||||
|
||||
st *storeHolder
|
||||
people peopleData
|
||||
@@ -323,6 +325,9 @@ func (m consoleModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
if m.menu == cGuard {
|
||||
cmds = append(cmds, m.fetchGuard())
|
||||
}
|
||||
if m.menu == cKB {
|
||||
cmds = append(cmds, m.fetchKB())
|
||||
}
|
||||
if m.menu == cMirror {
|
||||
cmds = append(cmds, m.fetchMirror())
|
||||
}
|
||||
@@ -381,6 +386,8 @@ func (m consoleModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.people = peopleData(msg)
|
||||
case consoleQuotasMsg:
|
||||
m.quotas, m.quotasErr = msg.rows, msg.err
|
||||
case consoleKBMsg:
|
||||
m.kb = kbState{view: msg.view, err: msg.err, loaded: true}
|
||||
case consoleGuardMsg:
|
||||
window := m.guard.window
|
||||
m.guard = guardState{view: msg.view, events: msg.events, err: msg.err, window: window}
|
||||
@@ -539,6 +546,8 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.menuD = m.quotaMenu(m.quotas[m.row])
|
||||
case cGuard:
|
||||
m.menuD = m.guardAttachmentMenu(m.guard.view.Attachments[m.row])
|
||||
case cKB:
|
||||
m.menuD = m.kbMenu(m.kb.view.Collections[m.row])
|
||||
case cMirror:
|
||||
m.menuD = m.mirrorMenu(m.mirrorGroups()[m.row])
|
||||
case cJobs:
|
||||
@@ -596,6 +605,10 @@ func (m consoleModel) key(k tea.KeyMsg) (tea.Model, tea.Cmd) {
|
||||
m.quotaForm(adminapi.QuotaView{})
|
||||
case cGuard:
|
||||
m.guardAttachForm()
|
||||
case cKB:
|
||||
if m.kb.view.Available {
|
||||
m.kbCreateForm()
|
||||
}
|
||||
case cModels:
|
||||
m.addModelForm()
|
||||
case cMirror:
|
||||
@@ -651,6 +664,8 @@ func (m consoleModel) onEnter() tea.Cmd {
|
||||
return m.fetchQuotas()
|
||||
case cGuard:
|
||||
return m.fetchGuard()
|
||||
case cKB:
|
||||
return m.fetchKB()
|
||||
case cMirror:
|
||||
return m.fetchMirror()
|
||||
case cMonitor:
|
||||
@@ -1027,6 +1042,8 @@ func (m consoleModel) rowCount() int {
|
||||
return len(m.quotas)
|
||||
case cGuard:
|
||||
return len(m.guard.view.Attachments)
|
||||
case cKB:
|
||||
return len(m.kb.view.Collections)
|
||||
case cMirror:
|
||||
return len(m.mirrorGroups())
|
||||
case cJobs:
|
||||
@@ -1260,6 +1277,8 @@ func (m consoleModel) panel() (title, body, keys string) {
|
||||
return m.quotasPanel(pw)
|
||||
case cGuard:
|
||||
return m.guardPanel(pw)
|
||||
case cKB:
|
||||
return m.kbPanel(pw)
|
||||
case cAccounts, cKeys, cOrgs:
|
||||
return m.peoplePanel(pw)
|
||||
case cManual:
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
// Knowledge bases, as the console shows them: each collection with its
|
||||
// owner, model, documents and grants; Enter for what can be done to one, A
|
||||
// to make one. People fill them from their own machines with llmgw-kb.
|
||||
|
||||
type kbState struct {
|
||||
view adminapi.KBView
|
||||
err error
|
||||
// loaded is set once an answer came, so "none yet" is not shown while
|
||||
// asking.
|
||||
loaded bool
|
||||
}
|
||||
|
||||
type consoleKBMsg struct {
|
||||
view adminapi.KBView
|
||||
err error
|
||||
}
|
||||
|
||||
func (m consoleModel) fetchKB() tea.Cmd {
|
||||
c := m.client
|
||||
return func() tea.Msg {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
var msg consoleKBMsg
|
||||
msg.err = c.get(ctx, "/api/v1/kb", &msg.view)
|
||||
return msg
|
||||
}
|
||||
}
|
||||
|
||||
func (m consoleModel) kbPanel(pw int) (string, string, string) {
|
||||
keys := "→ or Enter open · ↑↓ choose · Q leave"
|
||||
if m.inPanel {
|
||||
keys = "↑↓ collection · Enter actions · A new collection · ← back · Q leave"
|
||||
}
|
||||
k := m.kb
|
||||
var b strings.Builder
|
||||
b.WriteString("\n " + sDim.Render("Documents requests are answered from; people add them with llmgw-kb.") + "\n\n")
|
||||
switch {
|
||||
case k.err != nil:
|
||||
b.WriteString(" " + sBad.Render(k.err.Error()))
|
||||
return "Knowledge", b.String(), keys
|
||||
case !k.loaded:
|
||||
b.WriteString(" " + sDim.Render("asking the control plane…"))
|
||||
return "Knowledge", b.String(), keys
|
||||
case !k.view.Available:
|
||||
b.WriteString(" " + sBad.Render("Off: ") + k.view.Reason)
|
||||
return "Knowledge", b.String(), keys
|
||||
case len(k.view.Collections) == 0:
|
||||
b.WriteString(" " + sDim.Render("No collections yet. A to make one."))
|
||||
return "Knowledge", b.String(), keys
|
||||
}
|
||||
var rows [][]string
|
||||
for _, c := range k.view.Collections {
|
||||
var grants []string
|
||||
for _, g := range c.Grants {
|
||||
grants = append(grants, or(g.Name, g.Target)+" "+g.Perm)
|
||||
}
|
||||
rows = append(rows, []string{or(c.OrgName, c.Org) + "/" + c.Name, or(c.OwnerName, c.Owner),
|
||||
fmt.Sprintf("%s (%d)", c.EmbedModel, c.Dims), fmt.Sprint(c.Documents), fmt.Sprint(c.Chunks), dash(strings.Join(grants, ", "))})
|
||||
}
|
||||
b.WriteString(table([]string{"COLLECTION", "OWNER", "MODEL", "DOCUMENTS", "CHUNKS", "GRANTED TO"}, rows, m.row, pw, m.inPanel))
|
||||
if len(k.view.Defaults) > 0 {
|
||||
b.WriteString("\n\n " + sTitle.Render("BY DEFAULT") + " " + sDim.Render("when a request names no collection") + "\n")
|
||||
var drows [][]string
|
||||
for _, d := range k.view.Defaults {
|
||||
drows = append(drows, []string{or(d.Name, d.Target), strings.Join(d.Collections, ", ")})
|
||||
}
|
||||
b.WriteString(table([]string{"FOR", "ANSWERED FROM"}, drows, -1, pw, false))
|
||||
}
|
||||
b.WriteString("\n\n " + sDim.Render("On a person's machine: llmgw-kb sync COLLECTION FOLDER"))
|
||||
return "Knowledge", b.String(), keys
|
||||
}
|
||||
|
||||
// kbRef is how the admin API is told which collection: by its ID.
|
||||
func kbRef(c adminapi.KBCollectionView) string { return url.PathEscape(c.ID) }
|
||||
|
||||
// kbCreateForm makes a collection.
|
||||
func (m *consoleModel) kbCreateForm() {
|
||||
model := ""
|
||||
for _, md := range m.data.models {
|
||||
n := strings.ToLower(md.Name)
|
||||
if strings.Contains(n, "embed") || strings.Contains(n, "bge") || strings.Contains(n, "e5") {
|
||||
model = md.Name
|
||||
break
|
||||
}
|
||||
}
|
||||
m.openForm(&formDialog{title: "New knowledge-base collection",
|
||||
note: "Owned by part of the tree: it and everything below it may read and add to it.\n" +
|
||||
"The model must embed (capabilities.embeddings) and be able to serve now.",
|
||||
fields: []formField{
|
||||
choiceF("owner", []string{"project", "team", "org", "user"}, "team"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
|
||||
textF("org", "", "for a person's collection only"),
|
||||
textF("name", "", "how people name it to llmgw-kb"),
|
||||
textF("model", model, "the catalog model that embeds it"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
for i, what := range map[int]string{1: "which " + v[0], 3: "a name", 4: "a model"} {
|
||||
if strings.TrimSpace(v[i]) == "" {
|
||||
return nil, fmt.Errorf("%s?", what)
|
||||
}
|
||||
}
|
||||
req := adminapi.KBCreateRequest{Owner: v[0] + ":" + v[1], Org: v[2], Name: v[3], Model: v[4]}
|
||||
return m.call("New collection "+v[3], http.MethodPost, "/api/v1/kb/collections", req, func(raw json.RawMessage) string {
|
||||
var c adminapi.KBCollectionView
|
||||
_ = json.Unmarshal(raw, &c)
|
||||
return fmt.Sprintf("Collection %s, embedded by %s in %d dimensions.\n\nPeople with access fill it with:\n llmgw-kb sync %s FOLDER",
|
||||
c.Name, c.EmbedModel, c.Dims, c.Name)
|
||||
}), nil
|
||||
}})
|
||||
}
|
||||
|
||||
// kbMenu is what can be done to one collection.
|
||||
func (m consoleModel) kbMenu(c adminapi.KBCollectionView) *menuDialog {
|
||||
name := or(c.OrgName, c.Org) + "/" + c.Name
|
||||
items := []menuItem{
|
||||
{"Documents", func(m *consoleModel) tea.Cmd { return m.kbDocuments(c) }},
|
||||
{"Grant…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Grant " + name,
|
||||
note: "It and everything below it may then read, or with write add to, this collection.\nNever outside " + or(c.OrgName, c.Org) + ".",
|
||||
fields: []formField{
|
||||
choiceF("to", []string{"project", "team", "org", "user", "key"}, "project"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
|
||||
choiceF("may", []string{"read", "write"}, "read"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if strings.TrimSpace(v[1]) == "" {
|
||||
return nil, fmt.Errorf("which %s?", v[0])
|
||||
}
|
||||
req := adminapi.KBGrantRequest{Collection: c.ID, Target: v[0] + ":" + strings.TrimSpace(v[1]), Perm: v[2]}
|
||||
return m.call("Grant "+c.Name, http.MethodPost, "/api/v1/kb/grants", req, nil), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
}
|
||||
for _, g := range c.Grants {
|
||||
items = append(items, menuItem{"Revoke " + or(g.Name, g.Target) + " (" + g.Perm + ")…", func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Revoke a grant on " + name,
|
||||
text: or(g.Name, g.Target) + " may then no longer " + g.Perm + " it, unless it owns it.",
|
||||
yes: func(m *consoleModel) tea.Cmd {
|
||||
return m.call("Revoke", http.MethodDelete, "/api/v1/kb/grants",
|
||||
adminapi.KBGrantRequest{Collection: c.ID, Target: g.Target}, nil)
|
||||
}}
|
||||
return nil
|
||||
}})
|
||||
}
|
||||
items = append(items,
|
||||
menuItem{"Answer from it by default…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Answer from " + name + " by default",
|
||||
note: "Chat requests from there that name no collection are answered from this one.\n" +
|
||||
"Replaces that part's default; llmgw ctl kb default sets several.",
|
||||
fields: []formField{
|
||||
choiceF("for", []string{"project", "team", "org", "key"}, "team"),
|
||||
textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, or the ID"),
|
||||
},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if strings.TrimSpace(v[1]) == "" {
|
||||
return nil, fmt.Errorf("which %s?", v[0])
|
||||
}
|
||||
req := adminapi.KBDefaultRequest{Target: v[0] + ":" + strings.TrimSpace(v[1]), Collections: []string{c.ID}}
|
||||
return m.call("Default "+c.Name, http.MethodPost, "/api/v1/kb/defaults", req, nil), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
menuItem{"Try a search…", func(m *consoleModel) tea.Cmd {
|
||||
m.openForm(&formDialog{title: "Search " + name,
|
||||
note: "As an administrator: no tenancy applies. Shows what a question would find.",
|
||||
fields: []formField{textF("question", "", ""), textF("results", "5", "how many chunks")},
|
||||
submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
|
||||
if v[0] == "" {
|
||||
return nil, fmt.Errorf("what question?")
|
||||
}
|
||||
k, _ := strconv.Atoi(v[1])
|
||||
req := adminapi.KBSearchRequest{Collection: c.ID, Query: v[0], K: k}
|
||||
return m.call("Search "+c.Name, http.MethodPost, "/api/v1/kb/search", req, func(raw json.RawMessage) string {
|
||||
var hits []adminapi.KBHitView
|
||||
_ = json.Unmarshal(raw, &hits)
|
||||
if len(hits) == 0 {
|
||||
return "Nothing found."
|
||||
}
|
||||
var b strings.Builder
|
||||
for i, h := range hits {
|
||||
text := h.Text
|
||||
if r := []rune(text); len(r) > 400 {
|
||||
text = string(r[:400]) + "…"
|
||||
}
|
||||
fmt.Fprintf(&b, "%d. %s #%d\n%s\n\n", i+1, h.Document, h.Ord, indent(strings.TrimSpace(text), " "))
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
}), nil
|
||||
}})
|
||||
return nil
|
||||
}},
|
||||
menuItem{"Delete…", func(m *consoleModel) tea.Cmd {
|
||||
m.confirmD = &confirmDialog{title: "Delete " + name,
|
||||
text: fmt.Sprintf("Its %d documents and %d chunks, and every grant, go with it. People's own files are not touched.",
|
||||
c.Documents, c.Chunks),
|
||||
yes: func(m *consoleModel) tea.Cmd {
|
||||
return m.call("Delete "+c.Name, http.MethodDelete, "/api/v1/kb/collections/"+kbRef(c), nil, nil)
|
||||
}}
|
||||
return nil
|
||||
}},
|
||||
)
|
||||
return &menuDialog{title: name, items: items}
|
||||
}
|
||||
|
||||
// kbDocuments shows a collection's documents in a dialog.
|
||||
func (m *consoleModel) kbDocuments(c adminapi.KBCollectionView) tea.Cmd {
|
||||
return m.call("Documents of "+c.Name, http.MethodGet, "/api/v1/kb/collections/"+kbRef(c)+"/documents", nil,
|
||||
func(raw json.RawMessage) string {
|
||||
var docs []adminapi.KBDocumentView
|
||||
_ = json.Unmarshal(raw, &docs)
|
||||
if len(docs) == 0 {
|
||||
return "None yet. On a person's machine: llmgw-kb sync " + c.Name + " FOLDER"
|
||||
}
|
||||
var b strings.Builder
|
||||
for i, d := range docs {
|
||||
if i == 40 {
|
||||
fmt.Fprintf(&b, "… and %d more (llmgw ctl kb docs %s)", len(docs)-40, c.ID)
|
||||
break
|
||||
}
|
||||
fmt.Fprintf(&b, "%-48s %5d chunks %s %s\n", d.Name, d.Chunks, d.AddedAt.Local().Format("2006-01-02 15:04"), d.AddedBy)
|
||||
}
|
||||
return strings.TrimRight(b.String(), "\n")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
tea "github.com/charmbracelet/bubbletea"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
|
||||
)
|
||||
|
||||
func TestACollectionIsMadeAndGrantedFromTheConsole(t *testing.T) {
|
||||
m, s := consoleWith(t, cKB)
|
||||
m.data.models = append(m.data.models, adminapi.ModelView{Name: "bge-m3"})
|
||||
m.kb = kbState{loaded: true, view: adminapi.KBView{Available: true, Collections: []adminapi.KBCollectionView{{
|
||||
ID: "kb_aaaaaaaaaaaaaaaa", Name: "policies", Org: "o1", OrgName: "Acme", Owner: "team:t1", OwnerName: "Acme / hr",
|
||||
EmbedModel: "bge-m3", Dims: 1024, Documents: 3, Chunks: 40,
|
||||
Grants: []adminapi.KBGrantView{{Target: "project:p9", Name: "Acme/sales/crm", Perm: "read"}},
|
||||
}}}}
|
||||
var tm tea.Model = m
|
||||
view := tm.View()
|
||||
for _, want := range []string{"Knowledge", "Acme/policies", "Acme / hr", "bge-m3 (1024)", "Acme/s", "llmgw-kb sync COLLECTION FOLDER"} {
|
||||
if !strings.Contains(view, want) {
|
||||
t.Errorf("missing %q in:\n%s", want, view)
|
||||
}
|
||||
}
|
||||
|
||||
// A: the form, prefilled with the embedding model.
|
||||
tm = press(tm, "a")
|
||||
f := tm.(consoleModel).formD
|
||||
if f == nil || f.fields[4].value() != "bge-m3" {
|
||||
t.Fatalf("no form, or the model not offered: %+v", f)
|
||||
}
|
||||
f.fields[1].text.SetValue("Acme/hr")
|
||||
f.fields[3].text.SetValue("handbook")
|
||||
tm = pressRun(tm, "tab", "tab", "tab", "tab", "tab", "enter")
|
||||
var req adminapi.KBCreateRequest
|
||||
got := s.last()
|
||||
if !strings.HasPrefix(got, "POST /api/v1/kb/collections ") ||
|
||||
json.Unmarshal([]byte(strings.SplitN(got, " ", 3)[2]), &req) != nil ||
|
||||
req.Owner != "team:Acme/hr" || req.Name != "handbook" || req.Model != "bge-m3" {
|
||||
t.Fatalf("sent %q", got)
|
||||
}
|
||||
|
||||
// Enter on a row: grant to a project for writing.
|
||||
m.formD = nil
|
||||
tm = m
|
||||
tm = press(tm, "enter")
|
||||
if md := tm.(consoleModel).menuD; md == nil || !strings.Contains(md.items[2].label, "Revoke Acme/sales/crm") {
|
||||
t.Fatalf("menu %+v", md)
|
||||
}
|
||||
tm = press(tm, "down", "enter")
|
||||
f = tm.(consoleModel).formD
|
||||
f.fields[1].text.SetValue("Acme/sales/crm")
|
||||
f.fields[2].opt = 1 // write
|
||||
tm = pressRun(tm, "tab", "tab", "tab", "enter")
|
||||
if got := s.last(); got != `POST /api/v1/kb/grants {"collection":"kb_aaaaaaaaaaaaaaaa","target":"project:Acme/sales/crm","perm":"write"}` {
|
||||
t.Fatalf("sent %q", got)
|
||||
}
|
||||
_ = tm
|
||||
}
|
||||
@@ -18,7 +18,12 @@ import (
|
||||
//
|
||||
// An empty path loads the defaults, which with --server and --node-id on the
|
||||
// command line is enough for a smoke test on a node with no IB.
|
||||
func LoadAgent(path string) (agent.Config, error) {
|
||||
func LoadAgent(path string) (agent.Config, error) { return LoadAgentWith(path, nil) }
|
||||
|
||||
// LoadAgentWith is LoadAgent with the cluster's settings for this node
|
||||
// (agents.default and agents.nodes.<id>, as YAML) laid over the file: the
|
||||
// file is what the agent needs to start, the settings what it runs on.
|
||||
func LoadAgentWith(path string, settings []byte) (agent.Config, error) {
|
||||
cfg := agent.DefaultConfig()
|
||||
k := koanf.New(".")
|
||||
|
||||
@@ -30,6 +35,11 @@ func LoadAgent(path string) (agent.Config, error) {
|
||||
return cfg, fmt.Errorf("parsing %s: %w", path, err)
|
||||
}
|
||||
}
|
||||
if len(settings) > 0 {
|
||||
if err := k.Load(bytesProvider(settings), yaml.Parser()); err != nil {
|
||||
return cfg, fmt.Errorf("parsing the cluster's settings for this agent: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
err := k.Load(env.Provider(".", env.Opt{
|
||||
Prefix: "LLMGW_",
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// AgentSettings are the nodes' agent settings, kept with the cluster's
|
||||
// (BACKLOG E3): everything agent.yaml can say except what an agent needs to
|
||||
// start — its id, the control planes to dial, its address and its
|
||||
// certificate. Default applies to every node, and Nodes, by node id, over
|
||||
// it; both over the node's own agent.yaml.
|
||||
type AgentSettings struct {
|
||||
// Default is agent.yaml's keys for every node.
|
||||
Default map[string]any `koanf:"default"`
|
||||
// Nodes are agent.yaml's keys for one node each, by node id, over
|
||||
// Default.
|
||||
Nodes map[string]map[string]any `koanf:"nodes"`
|
||||
}
|
||||
|
||||
// agentBootstrapKeys are what an agent needs before it can be told anything:
|
||||
// they stay in agent.yaml.
|
||||
var agentBootstrapKeys = []string{"node_id", "server", "servers", "hostname", "mgmt_addr", "tls"}
|
||||
|
||||
// For is the settings for one node: Default with the node's own over it.
|
||||
// Empty when there are none.
|
||||
func (a AgentSettings) For(node string) map[string]any {
|
||||
out := clone(stringKeys(a.Default))
|
||||
if n, ok := a.Nodes[node]; ok {
|
||||
merge(out, stringKeys(n))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Check refuses settings that set an agent's bootstrap keys, or that would
|
||||
// not load as an agent's configuration for some node.
|
||||
func (a AgentSettings) Check() error {
|
||||
check := func(where string, m map[string]any) error {
|
||||
for _, k := range agentBootstrapKeys {
|
||||
if _, ok := m[k]; ok {
|
||||
return fmt.Errorf("%s.%s stays in each node's agent.yaml: an agent needs it before it can be told anything", where, k)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := check("agents.default", stringKeys(a.Default)); err != nil {
|
||||
return err
|
||||
}
|
||||
nodes := make([]string, 0, len(a.Nodes))
|
||||
for n := range a.Nodes {
|
||||
nodes = append(nodes, n)
|
||||
}
|
||||
sort.Strings(nodes)
|
||||
for _, n := range nodes {
|
||||
if err := check("agents.nodes."+n, stringKeys(a.Nodes[n])); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Each node's settings as its agent would load them, over nothing: what
|
||||
// its own file adds cannot be known here, and what is wrong in the
|
||||
// settings alone is wrong on every node.
|
||||
for _, n := range append([]string{""}, nodes...) {
|
||||
b, err := Marshal(a.For(n))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := LoadAgentWith("", b); err != nil {
|
||||
where := "agents.default"
|
||||
if n != "" {
|
||||
where = "agents.nodes." + n
|
||||
}
|
||||
return fmt.Errorf("%s: %w", where, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AgentBootstrapKeys are the keys agent.yaml keeps, for saying so.
|
||||
func AgentBootstrapKeys() []string { return slices.Clone(agentBootstrapKeys) }
|
||||
|
||||
// stringKeys turns YAML's map[any]any, wherever it hides, into the
|
||||
// map[string]any the rest of this package merges.
|
||||
func stringKeys(m map[string]any) map[string]any {
|
||||
out := make(map[string]any, len(m))
|
||||
for k, v := range m {
|
||||
out[k] = normalize(v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalize(v any) any {
|
||||
switch t := v.(type) {
|
||||
case map[string]any:
|
||||
return stringKeys(t)
|
||||
case map[any]any:
|
||||
m := make(map[string]any, len(t))
|
||||
for k, vv := range t {
|
||||
m[fmt.Sprint(k)] = normalize(vv)
|
||||
}
|
||||
return m
|
||||
case []any:
|
||||
out := make([]any, len(t))
|
||||
for i, vv := range t {
|
||||
out[i] = normalize(vv)
|
||||
}
|
||||
return out
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestANodesAgentSettingsAreTheDefaultWithItsOwnOver(t *testing.T) {
|
||||
a := AgentSettings{
|
||||
Default: map[string]any{"heartbeat_interval": "5s", "supervisor": map[string]any{"work_dir": "/srv/w", "x": 1}},
|
||||
Nodes: map[string]map[string]any{
|
||||
"cuda01": {"supervisor": map[string]any{"work_dir": "/fast/w"}, "labels": map[string]any{"gpu": "rtx6000"}},
|
||||
},
|
||||
}
|
||||
got := a.For("cuda01")
|
||||
sup := got["supervisor"].(map[string]any)
|
||||
if sup["work_dir"] != "/fast/w" || sup["x"] != 1 || got["heartbeat_interval"] != "5s" || got["labels"] == nil {
|
||||
t.Errorf("cuda01: %v", got)
|
||||
}
|
||||
if got := a.For("node01"); got["labels"] != nil || got["supervisor"].(map[string]any)["work_dir"] != "/srv/w" {
|
||||
t.Errorf("node01: %v", got)
|
||||
}
|
||||
if got := (AgentSettings{}).For("x"); len(got) != 0 {
|
||||
t.Errorf("none: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentSettingsKeepBootstrapKeysInTheFileAndMustLoad(t *testing.T) {
|
||||
for _, a := range []AgentSettings{
|
||||
{Default: map[string]any{"servers": []any{"x:7443"}}},
|
||||
{Nodes: map[string]map[string]any{"n1": {"node_id": "other"}}},
|
||||
{Nodes: map[string]map[string]any{"n1": {"tls": map[string]any{"ca": "/x"}}}},
|
||||
} {
|
||||
if err := a.Check(); err == nil || !strings.Contains(err.Error(), "stays in each node's agent.yaml") {
|
||||
t.Errorf("%+v: %v", a, err)
|
||||
}
|
||||
}
|
||||
// A fabric with no adapters would not load on the node it names.
|
||||
bad := AgentSettings{Nodes: map[string]map[string]any{"n1": {"fabrics": []any{map[string]any{"name": "ib-a"}}}}}
|
||||
if err := bad.Check(); err == nil || !strings.Contains(err.Error(), "agents.nodes.n1") || !strings.Contains(err.Error(), "no HCAs") {
|
||||
t.Errorf("a fabric without HCAs: %v", err)
|
||||
}
|
||||
if err := (AgentSettings{Default: map[string]any{"heartbeat_interval": "3s"}}).Check(); err != nil {
|
||||
t.Errorf("a plain setting: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheClusterSettingsCarryAgentsAndRefuseBadOnes(t *testing.T) {
|
||||
good, err := Parse([]byte("agents:\n default:\n heartbeat_interval: 3s\n nodes:\n cuda01:\n labels: {gpu: rtx6000}\n"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := FromMap(good)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.Agents.For("cuda01")["labels"] == nil {
|
||||
t.Errorf("agents not read: %+v", cfg.Agents)
|
||||
}
|
||||
bad, _ := Parse([]byte("agents:\n default:\n mgmt_addr: 10.0.0.1\n"))
|
||||
if _, err := FromMap(bad); err == nil {
|
||||
t.Error("a bootstrap key in the cluster's settings was taken")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSettingsAreLaidOverTheAgentsFile(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "agent.yaml")
|
||||
_ = os.WriteFile(p, []byte("node_id: n1\nserver: cp:7443\nheartbeat_interval: 9s\nlabels: {rack: r1}\n"), 0o600)
|
||||
cfg, err := LoadAgentWith(p, []byte("heartbeat_interval: 2s\nlabels: {gpu: a100}\n"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.NodeID != "n1" || cfg.HeartbeatInterval != 2*time.Second || cfg.Labels["rack"] != "r1" || cfg.Labels["gpu"] != "a100" {
|
||||
t.Errorf("loaded %+v", cfg)
|
||||
}
|
||||
}
|
||||
@@ -158,6 +158,13 @@ type Server struct {
|
||||
// Firewall turns llmgw's firewall on every machine on or off.
|
||||
Firewall FirewallSettings `koanf:"firewall"`
|
||||
|
||||
// Agents are the nodes' agent settings: agent.yaml's keys but those an
|
||||
// agent needs to start, for every node (default) and for one
|
||||
// (nodes.<id>). Sent to the agents as they change; an agent whose
|
||||
// settings changed starts again on them, its engines serving
|
||||
// throughout.
|
||||
Agents AgentSettings `koanf:"agents"`
|
||||
|
||||
// FirewallNetworks are networks, besides those each machine reaches the
|
||||
// cluster on, that llmgw's internal ports answer: a monitoring host on
|
||||
// another subnet, nodes behind a router. Written into every machine's
|
||||
@@ -887,6 +894,9 @@ func load(b []byte) (Server, error) {
|
||||
// Validate rejects configurations that would otherwise fail confusingly at
|
||||
// request time.
|
||||
func (s Server) Validate() error {
|
||||
if err := s.Agents.Check(); err != nil {
|
||||
return err
|
||||
}
|
||||
for i, o := range s.EngineOverrides {
|
||||
if o.GPUClass == "" && len(o.Nodes) == 0 && len(o.Labels) == 0 {
|
||||
return fmt.Errorf("engine_overrides[%d] picks no nodes: give it a gpu_class, nodes or labels", i)
|
||||
|
||||
@@ -68,7 +68,9 @@ type API struct {
|
||||
// nil without a store.
|
||||
federation *federation.Service
|
||||
// guard is the guardrail policies' library and attachments.
|
||||
guard *GuardAdmin
|
||||
guard *GuardAdmin
|
||||
// kb is knowledge bases; nil without a store.
|
||||
kb *KBAdmin
|
||||
settings *Settings
|
||||
// restartAgents tells every connected agent to start again.
|
||||
restartAgents func(reason string) []string
|
||||
@@ -114,6 +116,7 @@ func (a *API) Register(mux Mux) {
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/scale", a.scaleModel)
|
||||
mux.HandleFunc("POST /api/v1/models/{name}/resources", a.setResources)
|
||||
a.registerQuotas(mux)
|
||||
a.registerKB(mux)
|
||||
a.registerGuard(mux)
|
||||
a.registerSettings(mux)
|
||||
mux.HandleFunc("POST /api/v1/agents/restart", a.restartAgentsHandler)
|
||||
|
||||
@@ -0,0 +1,545 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// Knowledge bases, administered (BACKLOG F): collections made, granted,
|
||||
// emptied and removed here; documents put into them by people with
|
||||
// llmgw-kb, through the gateway.
|
||||
|
||||
// KBStore is the part of the store knowledge bases need.
|
||||
type KBStore interface {
|
||||
KnowledgeAvailable() error
|
||||
Collections(ctx context.Context) ([]store.KBCollection, error)
|
||||
Collection(ctx context.Context, id string) (store.KBCollection, error)
|
||||
CreateCollection(ctx context.Context, c store.KBCollection) (store.KBCollection, error)
|
||||
DeleteCollection(ctx context.Context, id string) error
|
||||
Grant(ctx context.Context, g store.KBGrant) error
|
||||
Revoke(ctx context.Context, scope store.Scope, scopeID, collection string) (bool, error)
|
||||
Grants(ctx context.Context, collection string) ([]store.KBGrant, error)
|
||||
Documents(ctx context.Context, collection string) ([]store.KBDocument, error)
|
||||
DeleteDocument(ctx context.Context, collection, name string) (bool, error)
|
||||
SearchChunks(ctx context.Context, q store.KBSearch) ([]store.KBHit, error)
|
||||
SetKBDefault(ctx context.Context, d store.KBDefault) error
|
||||
KBDefaults(ctx context.Context) ([]store.KBDefault, error)
|
||||
}
|
||||
|
||||
// KBAdmin is what the knowledge-base endpoints need: the store, and a way to
|
||||
// have a model embed a text (to learn its dimension, and to try a search).
|
||||
type KBAdmin struct {
|
||||
Store KBStore
|
||||
Embed func(ctx context.Context, model, text string) ([]float32, error)
|
||||
}
|
||||
|
||||
// SetKB enables the knowledge-base endpoints.
|
||||
func (a *API) SetKB(k *KBAdmin) { a.kb = k }
|
||||
|
||||
// KBView is every collection, or why there can be none.
|
||||
type KBView struct {
|
||||
Available bool `json:"available"`
|
||||
// Reason is why not, when not.
|
||||
Reason string `json:"reason,omitempty"`
|
||||
Collections []KBCollectionView `json:"collections"`
|
||||
// Defaults are the collections parts of the tree answer from when a
|
||||
// request names none.
|
||||
Defaults []KBDefaultView `json:"defaults"`
|
||||
}
|
||||
|
||||
// KBDefaultView is one part of the tree's default collections.
|
||||
type KBDefaultView struct {
|
||||
Target string `json:"target"`
|
||||
Name string `json:"name,omitempty"`
|
||||
// Collections are ORG/NAME, in order.
|
||||
Collections []string `json:"collections"`
|
||||
At time.Time `json:"at"`
|
||||
By string `json:"by,omitempty"`
|
||||
}
|
||||
|
||||
// KBDefaultRequest sets the collections a part of the tree answers from
|
||||
// when a request names none; none clears it.
|
||||
type KBDefaultRequest struct {
|
||||
// Target is org:, team:, project: (by path or ID) or key:KEY-ID.
|
||||
Target string `json:"target"`
|
||||
// Collections are IDs, names when unique, or ORG/NAME.
|
||||
Collections []string `json:"collections"`
|
||||
}
|
||||
|
||||
// KBCollectionView is a collection as an administrator sees it.
|
||||
type KBCollectionView struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Org string `json:"org"`
|
||||
OrgName string `json:"org_name,omitempty"`
|
||||
Owner string `json:"owner"`
|
||||
OwnerName string `json:"owner_name,omitempty"`
|
||||
EmbedModel string `json:"embed_model"`
|
||||
Dims int `json:"dims"`
|
||||
Documents int `json:"documents"`
|
||||
Chunks int `json:"chunks"`
|
||||
// Grants are who else may read or write it.
|
||||
Grants []KBGrantView `json:"grants"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CreatedBy string `json:"created_by,omitempty"`
|
||||
}
|
||||
|
||||
// KBGrantView is one grant.
|
||||
type KBGrantView struct {
|
||||
Target string `json:"target"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Perm string `json:"perm"`
|
||||
At time.Time `json:"at"`
|
||||
By string `json:"by,omitempty"`
|
||||
}
|
||||
|
||||
// KBCreateRequest makes a collection.
|
||||
type KBCreateRequest struct {
|
||||
// Owner is org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT (or the ID
|
||||
// after the colon), or user:USER-ID.
|
||||
Owner string `json:"owner"`
|
||||
// Org is the organization a person's collection belongs to; only for
|
||||
// user: owners.
|
||||
Org string `json:"org,omitempty"`
|
||||
Name string `json:"name"`
|
||||
// Model is the catalog model that embeds its documents.
|
||||
Model string `json:"model"`
|
||||
// Dims is the model's embedding dimension; asked of the model when 0.
|
||||
Dims int `json:"dims,omitempty"`
|
||||
}
|
||||
|
||||
// KBGrantRequest lets a scope read or write a collection, or with DELETE
|
||||
// takes that away.
|
||||
type KBGrantRequest struct {
|
||||
// Collection is its ID, its name when that is unique, or ORG/NAME.
|
||||
Collection string `json:"collection"`
|
||||
// Target is org:, team:, project:, key: or user: as for guard policies.
|
||||
Target string `json:"target"`
|
||||
// Perm is read or write; ignored by DELETE.
|
||||
Perm string `json:"perm,omitempty"`
|
||||
}
|
||||
|
||||
// KBDocumentView is one document in a collection.
|
||||
type KBDocumentView struct {
|
||||
Name string `json:"name"`
|
||||
Digest string `json:"digest"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Chunks int `json:"chunks"`
|
||||
AddedAt time.Time `json:"added_at"`
|
||||
AddedBy string `json:"added_by,omitempty"`
|
||||
}
|
||||
|
||||
// KBSearchRequest tries a search, as an administrator: no tenancy applies.
|
||||
type KBSearchRequest struct {
|
||||
Collection string `json:"collection"`
|
||||
Query string `json:"query"`
|
||||
K int `json:"k,omitempty"`
|
||||
}
|
||||
|
||||
// KBHitView is one chunk a search found.
|
||||
type KBHitView struct {
|
||||
Document string `json:"document"`
|
||||
Ord int `json:"ord"`
|
||||
Score float64 `json:"score"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
func (a *API) registerKB(mux Mux) {
|
||||
mux.HandleFunc("GET /api/v1/kb", a.kbList)
|
||||
mux.HandleFunc("POST /api/v1/kb/collections", a.kbCreate)
|
||||
mux.HandleFunc("DELETE /api/v1/kb/collections/{id}", a.kbDelete)
|
||||
mux.HandleFunc("GET /api/v1/kb/collections/{id}/documents", a.kbDocuments)
|
||||
mux.HandleFunc("DELETE /api/v1/kb/collections/{id}/documents", a.kbDeleteDocument)
|
||||
mux.HandleFunc("POST /api/v1/kb/grants", a.kbGrant(true))
|
||||
mux.HandleFunc("DELETE /api/v1/kb/grants", a.kbGrant(false))
|
||||
mux.HandleFunc("POST /api/v1/kb/search", a.kbSearch)
|
||||
mux.HandleFunc("POST /api/v1/kb/defaults", a.kbSetDefault)
|
||||
}
|
||||
|
||||
// kbReady answers for a gateway without knowledge bases, and says whether
|
||||
// it did.
|
||||
func (a *API) kbReady(w http.ResponseWriter) bool {
|
||||
if a.kb == nil {
|
||||
writeError(w, http.StatusNotImplemented, "knowledge bases need the store, which this control plane runs without")
|
||||
return false
|
||||
}
|
||||
if err := a.kb.Store.KnowledgeAvailable(); err != nil {
|
||||
writeError(w, http.StatusNotImplemented, err.Error())
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (a *API) names(ctx context.Context) map[string]string {
|
||||
if a.quotas == nil {
|
||||
return map[string]string{}
|
||||
}
|
||||
n, err := a.scopeNames(ctx)
|
||||
if err != nil {
|
||||
return map[string]string{}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (a *API) kbList(w http.ResponseWriter, r *http.Request) {
|
||||
if a.kb == nil {
|
||||
writeJSON(w, http.StatusOK, KBView{Reason: "knowledge bases need the store", Collections: []KBCollectionView{}, Defaults: []KBDefaultView{}})
|
||||
return
|
||||
}
|
||||
if err := a.kb.Store.KnowledgeAvailable(); err != nil {
|
||||
writeJSON(w, http.StatusOK, KBView{Reason: err.Error(), Collections: []KBCollectionView{}, Defaults: []KBDefaultView{}})
|
||||
return
|
||||
}
|
||||
cs, err := a.kb.Store.Collections(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
grants, err := a.kb.Store.Grants(r.Context(), "")
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
names := a.names(r.Context())
|
||||
v := KBView{Available: true, Collections: []KBCollectionView{}, Defaults: []KBDefaultView{}}
|
||||
for _, c := range cs {
|
||||
cv := KBCollectionView{ID: c.ID, Name: c.Name, Org: c.OrgID, OrgName: names["org:"+c.OrgID],
|
||||
Owner: string(c.OwnerScope) + ":" + c.OwnerID, OwnerName: names[string(c.OwnerScope)+":"+c.OwnerID],
|
||||
EmbedModel: c.EmbedModel, Dims: c.Dims, Grants: []KBGrantView{}, CreatedAt: c.CreatedAt, CreatedBy: c.CreatedBy}
|
||||
if docs, err := a.kb.Store.Documents(r.Context(), c.ID); err == nil {
|
||||
cv.Documents = len(docs)
|
||||
for _, d := range docs {
|
||||
cv.Chunks += d.Chunks
|
||||
}
|
||||
}
|
||||
for _, g := range grants {
|
||||
if g.Collection == c.ID {
|
||||
key := string(g.Scope) + ":" + g.ScopeID
|
||||
cv.Grants = append(cv.Grants, KBGrantView{Target: key, Name: names[key], Perm: string(g.Perm), At: g.AddedAt, By: g.AddedBy})
|
||||
}
|
||||
}
|
||||
v.Collections = append(v.Collections, cv)
|
||||
}
|
||||
label := map[string]string{}
|
||||
for _, c := range cs {
|
||||
label[c.ID] = or(names["org:"+c.OrgID], c.OrgID) + "/" + c.Name
|
||||
}
|
||||
defaults, err := a.kb.Store.KBDefaults(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
for _, d := range defaults {
|
||||
key := string(d.Scope) + ":" + d.ScopeID
|
||||
dv := KBDefaultView{Target: key, Name: names[key], Collections: []string{}, At: d.AddedAt, By: d.AddedBy}
|
||||
for _, id := range d.Collections {
|
||||
dv.Collections = append(dv.Collections, or(label[id], id))
|
||||
}
|
||||
v.Defaults = append(v.Defaults, dv)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, v)
|
||||
}
|
||||
|
||||
// kbTarget reads a grant's or an owner's target: org:, team: and project:
|
||||
// by path or ID, key: and user: by ID.
|
||||
func (a *API) kbTarget(ctx context.Context, target string) (store.Scope, string, error) {
|
||||
kind, what, ok := strings.Cut(strings.TrimSpace(target), ":")
|
||||
what = strings.TrimSpace(what)
|
||||
if !ok || what == "" {
|
||||
return "", "", fmt.Errorf("%q is not org:, team:, project:, key: or user: and what", target)
|
||||
}
|
||||
switch store.Scope(kind) {
|
||||
case store.ScopeKey, store.ScopeUser:
|
||||
return store.Scope(kind), what, nil
|
||||
case store.ScopeOrg, store.ScopeTeam, store.ScopeProject:
|
||||
if a.quotas == nil {
|
||||
return store.Scope(kind), what, nil
|
||||
}
|
||||
scope, id, _, err := a.resolveScope(ctx, target)
|
||||
return scope, id, err
|
||||
}
|
||||
return "", "", fmt.Errorf("%q: org, team, project, key or user", kind)
|
||||
}
|
||||
|
||||
// kbFind reads a collection by ID, by its name when that is unique, or by
|
||||
// ORG/NAME.
|
||||
func (a *API) kbFind(ctx context.Context, ref string) (store.KBCollection, error) {
|
||||
ref = strings.TrimSpace(ref)
|
||||
if c, err := a.kb.Store.Collection(ctx, ref); err == nil {
|
||||
return c, nil
|
||||
}
|
||||
cs, err := a.kb.Store.Collections(ctx)
|
||||
if err != nil {
|
||||
return store.KBCollection{}, err
|
||||
}
|
||||
names := a.names(ctx)
|
||||
org, name, scoped := strings.Cut(ref, "/")
|
||||
var found []store.KBCollection
|
||||
for _, c := range cs {
|
||||
switch {
|
||||
case scoped && strings.EqualFold(c.Name, name) &&
|
||||
(c.OrgID == org || strings.EqualFold(names["org:"+c.OrgID], org)):
|
||||
found = append(found, c)
|
||||
case !scoped && strings.EqualFold(c.Name, ref):
|
||||
found = append(found, c)
|
||||
}
|
||||
}
|
||||
switch len(found) {
|
||||
case 0:
|
||||
return store.KBCollection{}, fmt.Errorf("no collection %q", ref)
|
||||
case 1:
|
||||
return found[0], nil
|
||||
}
|
||||
return store.KBCollection{}, fmt.Errorf("%q names a collection in more than one organization; write ORG/%s", ref, ref)
|
||||
}
|
||||
|
||||
func (a *API) kbCreate(w http.ResponseWriter, r *http.Request) {
|
||||
if !a.kbReady(w) {
|
||||
return
|
||||
}
|
||||
var in KBCreateRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
scope, id, err := a.kbTarget(r.Context(), in.Owner)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
org := ""
|
||||
if scope == store.ScopeUser {
|
||||
if in.Org == "" {
|
||||
writeError(w, http.StatusBadRequest, "a person's collection names its organization (org)")
|
||||
return
|
||||
}
|
||||
if _, org, err = a.kbTarget(r.Context(), "org:"+in.Org); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
m, ok := a.ctl.Catalog().Resolve(in.Model)
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "no model "+in.Model+" in the catalog")
|
||||
return
|
||||
}
|
||||
if !m.Capabilities.Embeddings {
|
||||
writeError(w, http.StatusBadRequest, m.Name+" does not embed (capabilities.embeddings); a collection needs an embedding model")
|
||||
return
|
||||
}
|
||||
dims := in.Dims
|
||||
if dims == 0 {
|
||||
if a.kb.Embed == nil {
|
||||
writeError(w, http.StatusBadRequest, "give the model's dimension (dims): nothing here can ask it")
|
||||
return
|
||||
}
|
||||
// Asked of the model: a dimension typed by hand and wrong would be
|
||||
// found only when the first document is refused.
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
v, err := a.kb.Embed(ctx, m.Name, "llmgw")
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "asking "+m.Name+" for its dimension: "+err.Error()+
|
||||
" (is it running? or give dims)")
|
||||
return
|
||||
}
|
||||
dims = len(v)
|
||||
}
|
||||
c, err := a.kb.Store.CreateCollection(r.Context(), store.KBCollection{OwnerScope: scope, OwnerID: id, OrgID: org,
|
||||
Name: in.Name, EmbedModel: m.Name, Dims: dims, CreatedBy: actorOf(r)})
|
||||
if err != nil {
|
||||
code := http.StatusBadRequest
|
||||
if errors.Is(err, store.ErrNotFound) {
|
||||
code = http.StatusNotFound
|
||||
}
|
||||
writeError(w, code, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusCreated, KBCollectionView{ID: c.ID, Name: c.Name, Org: c.OrgID,
|
||||
Owner: string(c.OwnerScope) + ":" + c.OwnerID, EmbedModel: c.EmbedModel, Dims: c.Dims,
|
||||
Grants: []KBGrantView{}, CreatedAt: c.CreatedAt, CreatedBy: c.CreatedBy})
|
||||
}
|
||||
|
||||
func (a *API) kbCollection(w http.ResponseWriter, r *http.Request) (store.KBCollection, bool) {
|
||||
if !a.kbReady(w) {
|
||||
return store.KBCollection{}, false
|
||||
}
|
||||
c, err := a.kbFind(r.Context(), r.PathValue("id"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return store.KBCollection{}, false
|
||||
}
|
||||
return c, true
|
||||
}
|
||||
|
||||
func (a *API) kbDelete(w http.ResponseWriter, r *http.Request) {
|
||||
c, ok := a.kbCollection(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := a.kb.Store.DeleteCollection(r.Context(), c.ID); err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"id": c.ID, "name": c.Name, "deleted": true})
|
||||
}
|
||||
|
||||
func (a *API) kbDocuments(w http.ResponseWriter, r *http.Request) {
|
||||
c, ok := a.kbCollection(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
docs, err := a.kb.Store.Documents(r.Context(), c.ID)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
out := []KBDocumentView{}
|
||||
for _, d := range docs {
|
||||
out = append(out, KBDocumentView{Name: d.Name, Digest: d.Digest, Bytes: d.Bytes, Chunks: d.Chunks,
|
||||
AddedAt: d.AddedAt, AddedBy: d.AddedBy})
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) kbDeleteDocument(w http.ResponseWriter, r *http.Request) {
|
||||
c, ok := a.kbCollection(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
name := r.URL.Query().Get("name")
|
||||
gone, err := a.kb.Store.DeleteDocument(r.Context(), c.ID, name)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
if !gone {
|
||||
writeError(w, http.StatusNotFound, "no document "+name+" in "+c.Name)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"name": name, "deleted": true})
|
||||
}
|
||||
|
||||
func (a *API) kbGrant(grant bool) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if !a.kbReady(w) {
|
||||
return
|
||||
}
|
||||
var in KBGrantRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
c, err := a.kbFind(r.Context(), in.Collection)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
scope, id, err := a.kbTarget(r.Context(), in.Target)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
if !grant {
|
||||
gone, err := a.kb.Store.Revoke(r.Context(), scope, id, c.ID)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
if !gone {
|
||||
writeError(w, http.StatusNotFound, "no such grant")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, KBGrantView{Target: string(scope) + ":" + id})
|
||||
return
|
||||
}
|
||||
perm := store.KBPerm(strings.ToLower(strings.TrimSpace(in.Perm)))
|
||||
if perm == "" {
|
||||
perm = store.KBRead
|
||||
}
|
||||
if err := a.kb.Store.Grant(r.Context(), store.KBGrant{Scope: scope, ScopeID: id, Collection: c.ID,
|
||||
Perm: perm, AddedBy: actorOf(r)}); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, KBGrantView{Target: string(scope) + ":" + id, Perm: string(perm), At: time.Now(), By: actorOf(r)})
|
||||
}
|
||||
}
|
||||
|
||||
func (a *API) kbSearch(w http.ResponseWriter, r *http.Request) {
|
||||
if !a.kbReady(w) {
|
||||
return
|
||||
}
|
||||
var in KBSearchRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
c, err := a.kbFind(r.Context(), in.Collection)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
k := in.K
|
||||
if k <= 0 || k > 50 {
|
||||
k = 5
|
||||
}
|
||||
q := store.KBSearch{OrgID: c.OrgID, Collections: []store.KBCollection{c}, Text: in.Query, K: k}
|
||||
// Said, not swallowed: a search that quietly lost its meaning half
|
||||
// looks like a collection that does not know the answer.
|
||||
if a.kb.Embed != nil {
|
||||
v, err := a.kb.Embed(r.Context(), c.EmbedModel, in.Query)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, c.EmbedModel+" did not embed the question: "+err.Error())
|
||||
return
|
||||
}
|
||||
q.Embeddings = map[string][]float32{c.EmbedModel: v}
|
||||
}
|
||||
hits, err := a.kb.Store.SearchChunks(r.Context(), q)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
out := []KBHitView{}
|
||||
for _, h := range hits {
|
||||
out = append(out, KBHitView{Document: h.Document, Ord: h.Ord, Score: h.Score, Text: h.Body})
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (a *API) kbSetDefault(w http.ResponseWriter, r *http.Request) {
|
||||
if !a.kbReady(w) {
|
||||
return
|
||||
}
|
||||
var in KBDefaultRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
scope, id, err := a.kbTarget(r.Context(), in.Target)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
d := store.KBDefault{Scope: scope, ScopeID: id, AddedBy: actorOf(r)}
|
||||
for _, ref := range in.Collections {
|
||||
c, err := a.kbFind(r.Context(), ref)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, err.Error())
|
||||
return
|
||||
}
|
||||
d.Collections = append(d.Collections, c.ID)
|
||||
}
|
||||
if err := a.kb.Store.SetKBDefault(r.Context(), d); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, KBDefaultView{Target: string(scope) + ":" + id, Collections: in.Collections, At: time.Now(), By: actorOf(r)})
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
package adminapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func kbModels() []catalog.Model {
|
||||
return append(testModels(), catalog.Model{Name: "bge-m3", Source: "/e",
|
||||
Capabilities: catalog.Capabilities{Embeddings: true},
|
||||
Profiles: []catalog.Profile{{Name: "cpu", GPUClass: "h100-80g", TP: 1}}})
|
||||
}
|
||||
|
||||
func TestKnowledgeBasesListSaysWhyThereAreNone(t *testing.T) {
|
||||
a, _ := setup(t, kbModels())
|
||||
s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
a.SetKB(&KBAdmin{Store: s})
|
||||
var v KBView
|
||||
if code := call(t, a, "GET", "/api/v1/kb", "", &v); code != http.StatusOK || v.Available || !strings.Contains(v.Reason, "SQLite") {
|
||||
t.Errorf("got %d %+v", code, v)
|
||||
}
|
||||
if code := call(t, a, "POST", "/api/v1/kb/collections", `{"owner":"org:x","name":"n","model":"bge-m3"}`, nil); code != http.StatusNotImplemented {
|
||||
t.Errorf("creating on SQLite: %d", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledgeBasesAdministered(t *testing.T) {
|
||||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||||
if dsn == "" {
|
||||
t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
|
||||
}
|
||||
if u, err := url.Parse(dsn); err != nil || !strings.Contains(u.Path, "test") {
|
||||
t.Fatalf("refusing a database whose name does not say test")
|
||||
}
|
||||
ctx := context.Background()
|
||||
s, err := store.OpenPostgres(ctx, dsn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.KnowledgeAvailable(); err != nil {
|
||||
t.Skip(err)
|
||||
}
|
||||
a, _ := setup(t, kbModels())
|
||||
a.SetQuotas(s)
|
||||
probes := 0
|
||||
a.SetKB(&KBAdmin{Store: s, Embed: func(_ context.Context, model, _ string) ([]float32, error) {
|
||||
probes++
|
||||
return make([]float32, 4), nil
|
||||
}})
|
||||
sfx := strings.ToLower(store.NewID(""))
|
||||
org, _ := s.CreateOrg(ctx, "Acme"+sfx)
|
||||
team, _ := s.CreateTeam(ctx, org, "hr")
|
||||
_, _ = s.CreateProject(ctx, team, "docs")
|
||||
other, _ := s.CreateOrg(ctx, "Other"+sfx)
|
||||
|
||||
// Owner by path; the dimension asked of the model.
|
||||
var c KBCollectionView
|
||||
rec := func(method, path, body string, out any) int { return call(t, a, method, path, body, out) }
|
||||
body := `{"owner":"team:Acme` + sfx + `/hr","name":"politikalar","model":"bge-m3"}`
|
||||
if code := callCreated(t, a, body, &c); code != http.StatusCreated || c.Dims != 4 || probes != 1 || c.Org != org {
|
||||
t.Fatalf("create: %d %+v (probes %d)", code, c, probes)
|
||||
}
|
||||
// A model that does not embed is refused before anything is written.
|
||||
if code := rec("POST", "/api/v1/kb/collections", `{"owner":"org:`+org+`","name":"x","model":"llama"}`, nil); code != http.StatusBadRequest {
|
||||
t.Errorf("a chat model as the embedder: %d", code)
|
||||
}
|
||||
// Grants by path, never across organizations.
|
||||
if code := rec("POST", "/api/v1/kb/grants", `{"collection":"Acme`+sfx+`/politikalar","target":"project:Acme`+sfx+`/hr/docs","perm":"write"}`, nil); code != http.StatusOK {
|
||||
t.Errorf("grant: %d", code)
|
||||
}
|
||||
if code := rec("POST", "/api/v1/kb/grants", `{"collection":"`+c.ID+`","target":"org:`+other+`"}`, nil); code != http.StatusBadRequest {
|
||||
t.Errorf("a grant to another organization: %d", code)
|
||||
}
|
||||
var v KBView
|
||||
rec("GET", "/api/v1/kb", "", &v)
|
||||
var got *KBCollectionView
|
||||
for i := range v.Collections {
|
||||
if v.Collections[i].ID == c.ID {
|
||||
got = &v.Collections[i]
|
||||
}
|
||||
}
|
||||
if got == nil || len(got.Grants) != 1 || got.Grants[0].Perm != "write" || !strings.Contains(got.OwnerName, "hr") {
|
||||
t.Fatalf("list: %+v", got)
|
||||
}
|
||||
// By ORG/NAME too.
|
||||
if code := rec("DELETE", "/api/v1/kb/grants", `{"collection":"Acme`+sfx+`/politikalar","target":"project:Acme`+sfx+`/hr/docs"}`, nil); code != http.StatusOK {
|
||||
t.Errorf("revoke: %d", code)
|
||||
}
|
||||
if code := rec("DELETE", "/api/v1/kb/collections/"+c.ID, "", nil); code != http.StatusOK {
|
||||
t.Errorf("delete: %d", code)
|
||||
}
|
||||
}
|
||||
|
||||
func callCreated(t *testing.T, a *API, body string, out any) int {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
a.Register(mux)
|
||||
r, _ := http.NewRequest("POST", "/api/v1/kb/collections", strings.NewReader(body))
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
w := &recorder{h: http.Header{}}
|
||||
mux.ServeHTTP(w, r)
|
||||
if w.code == http.StatusCreated {
|
||||
_ = json.Unmarshal(w.body, out)
|
||||
} else {
|
||||
t.Logf("create answered %d %s", w.code, w.body)
|
||||
}
|
||||
return w.code
|
||||
}
|
||||
|
||||
type recorder struct {
|
||||
h http.Header
|
||||
code int
|
||||
body []byte
|
||||
}
|
||||
|
||||
func (r *recorder) Header() http.Header { return r.h }
|
||||
func (r *recorder) Write(b []byte) (int, error) {
|
||||
if r.code == 0 {
|
||||
r.code = http.StatusOK
|
||||
}
|
||||
r.body = append(r.body, b...)
|
||||
return len(b), nil
|
||||
}
|
||||
func (r *recorder) WriteHeader(c int) { r.code = c }
|
||||
@@ -66,6 +66,7 @@ var PathParams = map[string]string{
|
||||
"quotas/{scope}": "The scope: org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT, or their IDs after the colon.",
|
||||
"federation/{peer}": "The federation peer's name, as `llmgw ctl federation list` lists it.",
|
||||
"holds/{domain}": "The domain as level:name, rack:r12-07, as `llmgw ctl facility` lists it.",
|
||||
"collections/{id}": "The collection's ID, its name when that is unique, or ORG/NAME.",
|
||||
}
|
||||
|
||||
// Routes is every endpoint Register mounts.
|
||||
@@ -263,6 +264,34 @@ var Routes = []Route{
|
||||
{Method: "DELETE", Path: "/api/v1/guard/attachments", Summary: "Detach a guardrail policy",
|
||||
Request: GuardAttachmentRequest{}, Response: guard.Attachment{}},
|
||||
|
||||
// Knowledge bases: collections of documents, which people fill with
|
||||
// llmgw-kb from their own machines through the gateway.
|
||||
{Method: "GET", Path: "/api/v1/kb", Summary: "Every knowledge-base collection, with its documents' count and grants",
|
||||
Description: "Says why there are none when the store cannot hold them (SQLite, or PostgreSQL without pgvector).",
|
||||
Response: KBView{}},
|
||||
{Method: "POST", Path: "/api/v1/kb/collections", Summary: "Make a collection, owned by an organization, team, project or person",
|
||||
Description: "The embedding model's dimension is asked of the model when dims is not given, so the model must be able to serve.",
|
||||
Request: KBCreateRequest{}, Response: KBCollectionView{}, Status: http.StatusCreated},
|
||||
{Method: "DELETE", Path: "/api/v1/kb/collections/{id}", Summary: "Remove a collection with its documents and grants",
|
||||
Response: map[string]any{}},
|
||||
{Method: "GET", Path: "/api/v1/kb/collections/{id}/documents", Summary: "A collection's documents",
|
||||
Response: []KBDocumentView{}},
|
||||
{Method: "DELETE", Path: "/api/v1/kb/collections/{id}/documents", Summary: "Remove one document from a collection",
|
||||
Query: []Param{{Name: "name", Description: "The document's name, as listed.", Required: true}},
|
||||
Response: map[string]any{}},
|
||||
{Method: "POST", Path: "/api/v1/kb/grants", Summary: "Let a scope, and everything below it, read or write a collection",
|
||||
Description: "Never outside the collection's organization.",
|
||||
Request: KBGrantRequest{}, Response: KBGrantView{}},
|
||||
{Method: "DELETE", Path: "/api/v1/kb/grants", Summary: "Take a grant away",
|
||||
Request: KBGrantRequest{}, Response: KBGrantView{}},
|
||||
{Method: "POST", Path: "/api/v1/kb/defaults", Summary: "Set the collections a part of the tree answers from when a request names none",
|
||||
Description: "On an organization, team, project or key; the narrowest one with a default wins. No collections clears it. " +
|
||||
"A request's own X-Llmgw-Knowledge overrides it, and X-Llmgw-Knowledge: none turns it off.",
|
||||
Request: KBDefaultRequest{}, Response: KBDefaultView{}},
|
||||
{Method: "POST", Path: "/api/v1/kb/search", Summary: "Try a search of a collection, as an administrator",
|
||||
Description: "No tenancy applies: this is for checking what a collection holds and finds.",
|
||||
Request: KBSearchRequest{}, Response: []KBHitView{}},
|
||||
|
||||
{Method: "POST", Path: "/api/v1/agents/restart", Summary: "Have every node's agent start again; their engines keep serving",
|
||||
Response: AgentRestartView{}},
|
||||
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
// SetAgentSettings gives the service each node's agent settings from the
|
||||
// cluster's (BACKLOG E3), as YAML, read at each use so a change is what is
|
||||
// sent next.
|
||||
func (s *Service) SetAgentSettings(f func(node string) string) {
|
||||
s.mu.Lock()
|
||||
s.agentSettings = f
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Service) agentSettingsMsg(node string) *llmgwv1.ControlMessage {
|
||||
s.mu.RLock()
|
||||
f := s.agentSettings
|
||||
s.mu.RUnlock()
|
||||
if f == nil {
|
||||
return nil
|
||||
}
|
||||
return &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_AgentSettings{
|
||||
AgentSettings: &llmgwv1.AgentSettings{Yaml: f(node)}}}
|
||||
}
|
||||
|
||||
// agentSettingsGap is the pause between telling one node and the next: an
|
||||
// agent whose settings changed starts again, and a whole cluster's agents
|
||||
// doing so in the same second would all re-register at once.
|
||||
var agentSettingsGap = time.Second
|
||||
|
||||
// BroadcastAgentSettings tells every connected agent what the cluster's
|
||||
// settings now say it runs on, a moment apart. A node not connected is told
|
||||
// when it registers.
|
||||
func (s *Service) BroadcastAgentSettings() {
|
||||
for _, node := range s.ConnectedNodes() {
|
||||
msg := s.agentSettingsMsg(node)
|
||||
if msg == nil {
|
||||
return
|
||||
}
|
||||
if err := s.send(node, msg); err != nil {
|
||||
s.log.Warn("could not tell a node its agent settings; it is told when it next registers",
|
||||
"node", node, "err", err)
|
||||
continue
|
||||
}
|
||||
time.Sleep(agentSettingsGap)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
)
|
||||
|
||||
func TestEachNodeIsToldItsOwnAgentSettings(t *testing.T) {
|
||||
s := New(inventory.NewStore(), nil, false, slog.New(slog.DiscardHandler))
|
||||
if s.agentSettingsMsg("n1") != nil {
|
||||
t.Error("told settings with nothing to read them from")
|
||||
}
|
||||
s.SetAgentSettings(func(node string) string {
|
||||
if node == "cuda01" {
|
||||
return "labels:\n gpu: rtx6000\n"
|
||||
}
|
||||
return ""
|
||||
})
|
||||
if got := s.agentSettingsMsg("cuda01").GetAgentSettings().GetYaml(); got != "labels:\n gpu: rtx6000\n" {
|
||||
t.Errorf("cuda01 told %q", got)
|
||||
}
|
||||
// None is said too, so an agent whose settings were removed drops them.
|
||||
if m := s.agentSettingsMsg("node01"); m == nil || m.GetAgentSettings().GetYaml() != "" {
|
||||
t.Errorf("node01 told %v", m)
|
||||
}
|
||||
}
|
||||
@@ -79,6 +79,9 @@ type Service struct {
|
||||
controlPlanes func(context.Context) ([]string, error)
|
||||
// firewall is what nodes are told of llmgw's firewall; see SetFirewall.
|
||||
firewall func() *llmgwv1.Firewall
|
||||
// agentSettings is what each node's agent is told to run on; see
|
||||
// SetAgentSettings.
|
||||
agentSettings func(node string) string
|
||||
|
||||
// logWaits are the log requests awaiting a reply, by request id.
|
||||
logMu sync.Mutex
|
||||
@@ -296,6 +299,13 @@ func (s *Service) Connect(
|
||||
return err
|
||||
}
|
||||
}
|
||||
// What the cluster's settings say this agent runs on. An agent already
|
||||
// on them ignores it; one that is not starts again on them.
|
||||
if msg := s.agentSettingsMsg(nodeID); msg != nil {
|
||||
if err := stream.Send(msg); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Recover replica state from what this node is already running, before any
|
||||
// reconciliation pass can conclude the cluster is empty and place
|
||||
|
||||
@@ -19,6 +19,10 @@ var Reference = []Block{
|
||||
{Key: "queue_capacity_total", Type: "int", Doc: "QueueCapacityTotal bounds everything queued across bands. Reached, a\nrequest is not refused while a lower band has any queued: the newest of\nthe lowest band is turned away instead, with a 503 and Retry-After.\nZero takes half of all the bands' capacities together.", Nested: ""},
|
||||
{Key: "interval", Type: "time.Duration", Doc: "", Nested: ""},
|
||||
}},
|
||||
{Package: "config", Name: "AgentSettings", Title: "AgentSettings", Doc: "", Fields: []Field{
|
||||
{Key: "default", Type: "map[string]any", Doc: "Default is agent.yaml's keys for every node.", Nested: ""},
|
||||
{Key: "nodes", Type: "map[string]map[string]any", Doc: "Nodes are agent.yaml's keys for one node each, by node id, over\nDefault.", Nested: ""},
|
||||
}},
|
||||
{Package: "config", Name: "Auth", Title: "Auth", Doc: "", Fields: []Field{
|
||||
{Key: "enabled", Type: "bool", Doc: "", Nested: ""},
|
||||
{Key: "store", Type: "string", Doc: "Store is the SQLite file holding keys, tenancy and usage.", Nested: ""},
|
||||
@@ -372,6 +376,7 @@ var Reference = []Block{
|
||||
{Key: "mail", Type: "mail.Config", Doc: "Mail is where the gateway's few emails go out from: a link to set a\nforgotten password. Unset, none is sent and people ask an operator.", Nested: "mail.Config"},
|
||||
{Key: "serving", Type: "ServingSettings", Doc: "Serving stops and starts the whole cluster's service (llmgw stop, llmgw\nstart).", Nested: "config.ServingSettings"},
|
||||
{Key: "firewall", Type: "FirewallSettings", Doc: "Firewall turns llmgw's firewall on every machine on or off.", Nested: "config.FirewallSettings"},
|
||||
{Key: "agents", Type: "AgentSettings", Doc: "Agents are the nodes' agent settings: agent.yaml's keys but those an\nagent needs to start, for every node (default) and for one\n(nodes.<id>). Sent to the agents as they change; an agent whose\nsettings changed starts again on them, its engines serving\nthroughout.", Nested: "config.AgentSettings"},
|
||||
{Key: "firewall_networks", Type: "[]string", Doc: "FirewallNetworks are networks, besides those each machine reaches the\ncluster on, that llmgw's internal ports answer: a monitoring host on\nanother subnet, nodes behind a router. Written into every machine's\nfirewall by `server init`, `node add` and `server add`.", Nested: ""},
|
||||
{Key: "backup_recipients", Type: "[]string", Doc: "BackupRecipients encrypt every archive `llmgw server backup` writes, to\nwhoever holds one of their keys: age recipients (age1…) or SSH public\nkeys (ssh-ed25519 …, ssh-rsa …), as `age` takes them. An archive holds\nthe certificate authority's key and every secret in the store, and is\ncopied off the machine; encrypted, the copy is safe wherever it lands.\nEmpty writes archives unencrypted, readable by root only.", Nested: ""},
|
||||
{Key: "backup_copy_to", Type: "[]string", Doc: "BackupCopyTo is where each archive is copied once written, so it\nsurvives this machine: a directory (a mounted share, another disk) or\nssh://user@host[:port]/path, reached with BackupCopyKey and checked\nagainst /etc/llmgw/backup_known_hosts, which learns a host the first\ntime. The same number of archives is kept there as here. Refused\nwithout backup_recipients: an archive leaves the machine encrypted.", Nested: ""},
|
||||
|
||||
@@ -32,6 +32,7 @@ var Sources = []string{
|
||||
"internal/config/oidc.go",
|
||||
"internal/config/firewall.go",
|
||||
"internal/config/serving.go",
|
||||
"internal/config/agents.go",
|
||||
"internal/mail/mail.go",
|
||||
"internal/notify/notify.go",
|
||||
"internal/secrets/secrets.go",
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
// Package chunk splits extracted text into the pieces a knowledge base keeps
|
||||
// and searches (BACKLOG F2).
|
||||
//
|
||||
// A chunk is a few paragraphs under the headings they sit beneath: about 500
|
||||
// tokens, which an embedding model reads whole and a prompt can carry several
|
||||
// of. Paragraphs are not cut while they fit; a section starts a new chunk
|
||||
// when the current one is already a fair size, so a chunk rarely mixes two
|
||||
// topics; and the last paragraph of a chunk is repeated at the start of the
|
||||
// next, so an answer that spans the boundary is found whole in one of them.
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Options size the chunks. Zero values take the defaults.
|
||||
type Options struct {
|
||||
// MaxChars is the most characters in a chunk's body: 1800, about 500
|
||||
// tokens of Turkish or English.
|
||||
MaxChars int
|
||||
// Overlap is the most characters of the previous chunk repeated at the
|
||||
// start of the next: 300.
|
||||
Overlap int
|
||||
}
|
||||
|
||||
func (o Options) max() int {
|
||||
if o.MaxChars > 0 {
|
||||
return o.MaxChars
|
||||
}
|
||||
return 1800
|
||||
}
|
||||
|
||||
func (o Options) overlap() int {
|
||||
if o.Overlap >= 0 && o.MaxChars > 0 {
|
||||
return min(o.Overlap, o.max()/2)
|
||||
}
|
||||
return 300
|
||||
}
|
||||
|
||||
// Chunk is one piece: the headings above it and its text.
|
||||
type Chunk struct {
|
||||
// Headings is the path of headings the chunk starts under, widest first.
|
||||
Headings []string
|
||||
Body string
|
||||
}
|
||||
|
||||
// Text is what is stored and embedded: the heading path, then the body. The
|
||||
// path is part of it so that "how many days?" under "Annual leave" finds the
|
||||
// paragraph that says "twenty days" without naming leave.
|
||||
func (c Chunk) Text() string {
|
||||
if len(c.Headings) == 0 {
|
||||
return c.Body
|
||||
}
|
||||
return strings.Join(c.Headings, " › ") + "\n\n" + c.Body
|
||||
}
|
||||
|
||||
// Split cuts Markdown-shaped text into chunks.
|
||||
func Split(text string, o Options) []Chunk {
|
||||
max, overlap := o.max(), o.overlap()
|
||||
var out []Chunk
|
||||
type head struct {
|
||||
level int
|
||||
title string
|
||||
}
|
||||
var stack []head // the headings in force, widest first
|
||||
titles := func() []string {
|
||||
out := make([]string, len(stack))
|
||||
for i, h := range stack {
|
||||
out[i] = h.title
|
||||
}
|
||||
return out
|
||||
}
|
||||
var cur []string // blocks of the chunk being built
|
||||
var curHeads []string
|
||||
size := 0
|
||||
// carried is the overlap at the start of cur: it does not count as the
|
||||
// chunk having content of its own.
|
||||
carried := 0
|
||||
|
||||
emit := func(keepOverlap bool) {
|
||||
if size-carried <= 0 {
|
||||
cur, size, carried = nil, 0, 0
|
||||
return
|
||||
}
|
||||
out = append(out, Chunk{Headings: curHeads, Body: strings.Join(cur, "\n\n")})
|
||||
last := cur[len(cur)-1]
|
||||
cur, size, carried = nil, 0, 0
|
||||
if keepOverlap && overlap > 0 && runes(last) <= overlap {
|
||||
cur, size, carried = []string{last}, runes(last), runes(last)
|
||||
}
|
||||
}
|
||||
add := func(b string) {
|
||||
if len(cur) == 0 || size == carried {
|
||||
curHeads = titles()
|
||||
}
|
||||
if size > 0 && size+2+runes(b) > max {
|
||||
emit(true)
|
||||
curHeads = titles()
|
||||
}
|
||||
cur = append(cur, b)
|
||||
size += runes(b) + 2
|
||||
}
|
||||
|
||||
for _, b := range blocks(text) {
|
||||
if lvl, title := heading(b); lvl > 0 {
|
||||
// A new section starts a new chunk once this one is worth
|
||||
// keeping by itself; otherwise the heading goes into it.
|
||||
if size-carried >= max/3 {
|
||||
emit(false)
|
||||
}
|
||||
for len(stack) > 0 && stack[len(stack)-1].level >= lvl {
|
||||
stack = stack[:len(stack)-1]
|
||||
}
|
||||
stack = append(stack, head{lvl, title})
|
||||
if size-carried > 0 {
|
||||
add(b)
|
||||
} else {
|
||||
// Nothing of its own yet: the chunk starts under this heading.
|
||||
cur, size, carried = nil, 0, 0
|
||||
curHeads = titles()
|
||||
}
|
||||
continue
|
||||
}
|
||||
for _, piece := range fit(b, max) {
|
||||
add(piece)
|
||||
}
|
||||
}
|
||||
emit(false)
|
||||
return out
|
||||
}
|
||||
|
||||
func runes(s string) int { return utf8.RuneCountInString(s) }
|
||||
|
||||
// heading reads a Markdown ATX heading line: its level and title.
|
||||
func heading(b string) (int, string) {
|
||||
if strings.Contains(b, "\n") {
|
||||
return 0, ""
|
||||
}
|
||||
n := 0
|
||||
for n < len(b) && b[n] == '#' {
|
||||
n++
|
||||
}
|
||||
if n == 0 || n > 6 || n >= len(b) || b[n] != ' ' {
|
||||
return 0, ""
|
||||
}
|
||||
return n, strings.TrimSpace(strings.TrimRight(strings.TrimSpace(b[n:]), "#"))
|
||||
}
|
||||
|
||||
// blocks splits text on blank lines, keeping fenced code blocks whole.
|
||||
func blocks(text string) []string {
|
||||
var out []string
|
||||
var b []string
|
||||
fenced := false
|
||||
flush := func() {
|
||||
s := strings.TrimSpace(strings.Join(b, "\n"))
|
||||
if s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
b = nil
|
||||
}
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
t := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(t, "```") || strings.HasPrefix(t, "~~~") {
|
||||
fenced = !fenced
|
||||
}
|
||||
if !fenced && t == "" {
|
||||
flush()
|
||||
continue
|
||||
}
|
||||
// A heading line is a block of its own even without blank lines
|
||||
// around it.
|
||||
if !fenced && len(b) > 0 {
|
||||
if lvl, _ := heading(t); lvl > 0 {
|
||||
flush()
|
||||
}
|
||||
}
|
||||
b = append(b, strings.TrimRight(line, " \t"))
|
||||
if !fenced {
|
||||
if lvl, _ := heading(t); lvl > 0 {
|
||||
flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return out
|
||||
}
|
||||
|
||||
// fit cuts a block longer than max: at line ends, then sentence ends, then
|
||||
// spaces, and only as a last resort mid-word.
|
||||
func fit(b string, max int) []string {
|
||||
if runes(b) <= max {
|
||||
return []string{b}
|
||||
}
|
||||
for _, split := range []func(string) []string{lines, sentences, words} {
|
||||
parts := split(b)
|
||||
if len(parts) < 2 {
|
||||
continue
|
||||
}
|
||||
var out []string
|
||||
var cur strings.Builder
|
||||
for _, p := range parts {
|
||||
if cur.Len() > 0 && runes(cur.String())+runes(p) > max {
|
||||
out = append(out, strings.TrimSpace(cur.String()))
|
||||
cur.Reset()
|
||||
}
|
||||
if runes(p) > max {
|
||||
out = append(out, fit(strings.TrimSpace(p), max)...)
|
||||
continue
|
||||
}
|
||||
cur.WriteString(p)
|
||||
}
|
||||
if strings.TrimSpace(cur.String()) != "" {
|
||||
out = append(out, strings.TrimSpace(cur.String()))
|
||||
}
|
||||
return out
|
||||
}
|
||||
// One word longer than a chunk: a hash, a line of base64.
|
||||
r := []rune(b)
|
||||
var out []string
|
||||
for len(r) > max {
|
||||
out, r = append(out, string(r[:max])), r[max:]
|
||||
}
|
||||
return append(out, string(r))
|
||||
}
|
||||
|
||||
// lines splits keeping each line's newline.
|
||||
func lines(s string) []string { return strings.SplitAfter(s, "\n") }
|
||||
|
||||
// sentences splits after ., ? or ! followed by a space, keeping both.
|
||||
func sentences(s string) []string {
|
||||
var out []string
|
||||
start := 0
|
||||
r := []rune(s)
|
||||
for i := 0; i < len(r)-1; i++ {
|
||||
if (r[i] == '.' || r[i] == '?' || r[i] == '!' || r[i] == '…') && unicode.IsSpace(r[i+1]) {
|
||||
out = append(out, string(r[start:i+2]))
|
||||
start = i + 2
|
||||
}
|
||||
}
|
||||
if start < len(r) {
|
||||
out = append(out, string(r[start:]))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func words(s string) []string { return strings.SplitAfter(s, " ") }
|
||||
@@ -0,0 +1,106 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func para(word string, n int) string { return strings.TrimSpace(strings.Repeat(word+" ", n)) }
|
||||
|
||||
func TestShortTextIsOneChunkUnderItsHeadings(t *testing.T) {
|
||||
cs := Split("# İzin\n\n## Yıllık\n\nYirmi gün.\n\nBayramlar ayrı.\n", Options{})
|
||||
if len(cs) != 1 {
|
||||
t.Fatalf("chunks %+v", cs)
|
||||
}
|
||||
if got := cs[0].Text(); got != "İzin › Yıllık\n\nYirmi gün.\n\nBayramlar ayrı." {
|
||||
t.Errorf("text %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunksStayUnderTheLimitAndOverlap(t *testing.T) {
|
||||
var b strings.Builder
|
||||
for i := range 20 {
|
||||
b.WriteString(para("kelime"+string(rune('a'+i)), 30) + "\n\n")
|
||||
}
|
||||
cs := Split(b.String(), Options{MaxChars: 600, Overlap: 300})
|
||||
if len(cs) < 5 {
|
||||
t.Fatalf("only %d chunks", len(cs))
|
||||
}
|
||||
for i, c := range cs {
|
||||
if n := utf8.RuneCountInString(c.Body); n > 600 {
|
||||
t.Errorf("chunk %d has %d characters", i, n)
|
||||
}
|
||||
if i > 0 {
|
||||
prev := strings.Split(cs[i-1].Body, "\n\n")
|
||||
if !strings.HasPrefix(c.Body, prev[len(prev)-1]) {
|
||||
t.Errorf("chunk %d does not start with the end of chunk %d", i, i-1)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestASectionStartsAChunkOnceTheLastIsWorthKeeping(t *testing.T) {
|
||||
text := "# A\n\n" + para("a", 400) + "\n\n# B\n\nkısa\n\n# C\n\nde kısa\n"
|
||||
cs := Split(text, Options{MaxChars: 1800})
|
||||
if len(cs) != 2 {
|
||||
t.Fatalf("chunks %+v", cs)
|
||||
}
|
||||
if cs[1].Headings[0] != "B" || !strings.Contains(cs[1].Body, "# C") || !strings.Contains(cs[1].Body, "de kısa") {
|
||||
t.Errorf("small sections were not kept together: %+v", cs[1])
|
||||
}
|
||||
// The overlap does not run into a new section.
|
||||
if strings.HasPrefix(cs[1].Body, "a a") {
|
||||
t.Errorf("section B starts with section A's text")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeadingLevelsReplaceTheirSiblings(t *testing.T) {
|
||||
text := "## A\n\n### A1\n\n" + para("x", 300) + "\n\n## B\n\n" + para("y", 300) + "\n"
|
||||
cs := Split(text, Options{MaxChars: 700})
|
||||
last := cs[len(cs)-1]
|
||||
if strings.Join(last.Headings, "/") != "B" {
|
||||
t.Errorf("headings after a sibling: %v", last.Headings)
|
||||
}
|
||||
if strings.Join(cs[0].Headings, "/") != "A/A1" {
|
||||
t.Errorf("first headings %v", cs[0].Headings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLongParagraphsAreCutAtSentencesThenWords(t *testing.T) {
|
||||
long := strings.Repeat("Bu bir cümledir. ", 200)
|
||||
for _, c := range Split(long, Options{MaxChars: 500}) {
|
||||
if utf8.RuneCountInString(c.Body) > 500 {
|
||||
t.Errorf("chunk of %d", utf8.RuneCountInString(c.Body))
|
||||
}
|
||||
if !strings.HasSuffix(c.Body, ".") {
|
||||
t.Errorf("cut mid-sentence: …%q", c.Body[len(c.Body)-20:])
|
||||
}
|
||||
}
|
||||
word := strings.Repeat("ş", 1200)
|
||||
cs := Split(word, Options{MaxChars: 500})
|
||||
total := 0
|
||||
for _, c := range cs {
|
||||
total += utf8.RuneCountInString(c.Body)
|
||||
if !utf8.ValidString(c.Body) {
|
||||
t.Error("a multibyte character was cut")
|
||||
}
|
||||
}
|
||||
if total != 1200 {
|
||||
t.Errorf("a long word lost characters: %d", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodeFencesStayWhole(t *testing.T) {
|
||||
text := "Örnek:\n\n```go\nfunc a() {\n\n\treturn\n}\n```\n\nSon.\n"
|
||||
cs := Split(text, Options{})
|
||||
if len(cs) != 1 || !strings.Contains(cs[0].Body, "func a() {\n\n\treturn\n}") {
|
||||
t.Errorf("chunks %+v", cs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNothingToSayIsNoChunk(t *testing.T) {
|
||||
if cs := Split("# Başlık\n\n## Alt\n", Options{}); len(cs) != 0 {
|
||||
t.Errorf("headings alone made %+v", cs)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/quotedprintable"
|
||||
"net/mail"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// maxArchiveMembers is the most members of one archive read; past it the
|
||||
// rest are reported skipped.
|
||||
const maxArchiveMembers = 10000
|
||||
|
||||
// maxArchiveBytes is the most an archive's members are inflated to in all.
|
||||
const maxArchiveBytes = 2 << 30
|
||||
|
||||
// zipParts reads every member of a zip archive this can read, each a part
|
||||
// named by its path inside.
|
||||
func zipParts(ctx context.Context, name string, b []byte, o Options, depth int) ([]Part, error) {
|
||||
zr, err := openZip(b)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var parts []Part
|
||||
var total int64
|
||||
for i, f := range zr.File {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return parts, err
|
||||
}
|
||||
if f.FileInfo().IsDir() || skipName(f.Name) {
|
||||
continue
|
||||
}
|
||||
mname := memberName(f.Name)
|
||||
if i >= maxArchiveMembers {
|
||||
parts = append(parts, Part{Name: mname, Err: fmt.Errorf("%w: past the first %d members", ErrUnsupported, maxArchiveMembers)})
|
||||
continue
|
||||
}
|
||||
if !Known(f.Name) {
|
||||
continue
|
||||
}
|
||||
if int64(f.UncompressedSize64) > o.maxBytes() {
|
||||
parts = append(parts, Part{Name: mname, Err: fmt.Errorf("%w: %d MiB is over the limit", ErrUnsupported, f.UncompressedSize64>>20)})
|
||||
continue
|
||||
}
|
||||
total += int64(f.UncompressedSize64)
|
||||
if total > maxArchiveBytes {
|
||||
parts = append(parts, Part{Name: mname, Err: fmt.Errorf("%w: the archive inflates past %d GiB", ErrUnsupported, maxArchiveBytes>>30)})
|
||||
break
|
||||
}
|
||||
data, err := readMember(f)
|
||||
if err != nil {
|
||||
parts = append(parts, Part{Name: mname, Err: err})
|
||||
continue
|
||||
}
|
||||
inner, err := extractNested(ctx, mname, data, o, depth+1)
|
||||
parts = append(parts, nest(mname, inner, err)...)
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return nil, ErrNoText
|
||||
}
|
||||
return parts, nil
|
||||
}
|
||||
|
||||
func extractNested(ctx context.Context, name string, data []byte, o Options, depth int) ([]Part, error) {
|
||||
return extractDepth(ctx, name, data, o, depth)
|
||||
}
|
||||
|
||||
// nest names the parts of a member after it: "a/b.pdf", "a/c.zip!/d.txt".
|
||||
func nest(name string, parts []Part, err error) []Part {
|
||||
if err != nil {
|
||||
return []Part{{Name: name, Err: err}}
|
||||
}
|
||||
out := make([]Part, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if p.Name == "" {
|
||||
p.Name = name
|
||||
} else {
|
||||
p.Name = name + "!/" + p.Name
|
||||
}
|
||||
out = append(out, p)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// memberName is a member's path made safe to name a document: forward
|
||||
// slashes, no leading slash or dot segments.
|
||||
func memberName(n string) string {
|
||||
n = strings.ReplaceAll(n, `\`, "/")
|
||||
n = path.Clean("/" + n)
|
||||
return strings.TrimPrefix(n, "/")
|
||||
}
|
||||
|
||||
// skipName is what systems leave in archives that is not anybody's
|
||||
// document: macOS resource forks, Office lock files, thumbnails.
|
||||
func skipName(n string) bool {
|
||||
base := path.Base(strings.ReplaceAll(n, `\`, "/"))
|
||||
return strings.HasPrefix(n, "__MACOSX/") || strings.HasPrefix(base, "._") || strings.HasPrefix(base, "~$") ||
|
||||
base == ".DS_Store" || base == "Thumbs.db" || base == "desktop.ini"
|
||||
}
|
||||
|
||||
// mailParts reads an e-mail (.eml): the message — its headers and the plain
|
||||
// text, or the HTML when there is no plain text — and each attachment this
|
||||
// can read, named by its file name.
|
||||
func mailParts(ctx context.Context, name string, b []byte, o Options, depth int) ([]Part, error) {
|
||||
m, err := mail.ReadMessage(bytes.NewReader(b))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dec := new(mime.WordDecoder)
|
||||
dec.CharsetReader = func(cs string, in io.Reader) (io.Reader, error) {
|
||||
all, err := io.ReadAll(in)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return strings.NewReader(decodeCharset(all, cs)), nil
|
||||
}
|
||||
header := func(k string) string {
|
||||
v := m.Header.Get(k)
|
||||
if d, err := dec.DecodeHeader(v); err == nil {
|
||||
return d
|
||||
}
|
||||
return v
|
||||
}
|
||||
var plain, htmlBody []string
|
||||
var attachments []Part
|
||||
err = walkMIME(m.Header.Get("Content-Type"), m.Header.Get("Content-Transfer-Encoding"), m.Body, 0,
|
||||
func(ctype, filename string, params map[string]string, body []byte) {
|
||||
if filename != "" {
|
||||
if d, err := dec.DecodeHeader(filename); err == nil {
|
||||
filename = d
|
||||
}
|
||||
if Known(filename) {
|
||||
inner, err := extractNested(ctx, filename, body, o, depth+1)
|
||||
attachments = append(attachments, nest(memberName(filename), inner, err)...)
|
||||
}
|
||||
return
|
||||
}
|
||||
switch ctype {
|
||||
case "text/plain":
|
||||
plain = append(plain, decodeCharset(body, params["charset"]))
|
||||
case "text/html":
|
||||
if t, err := htmlFrom(strings.NewReader(decodeCharset(body, params["charset"]))); err == nil {
|
||||
htmlBody = append(htmlBody, t)
|
||||
}
|
||||
case "message/rfc822":
|
||||
if inner, err := mailParts(ctx, "", body, o, depth+1); err == nil {
|
||||
attachments = append(attachments, nest("message", inner, nil)...)
|
||||
}
|
||||
}
|
||||
})
|
||||
if err != nil && len(plain) == 0 && len(htmlBody) == 0 {
|
||||
return nil, err
|
||||
}
|
||||
var text strings.Builder
|
||||
if s := header("Subject"); s != "" {
|
||||
text.WriteString("# " + s + "\n\n")
|
||||
}
|
||||
for _, k := range []string{"From", "To", "Cc", "Date"} {
|
||||
if v := header(k); v != "" {
|
||||
text.WriteString(k + ": " + v + "\n")
|
||||
}
|
||||
}
|
||||
text.WriteString("\n")
|
||||
body := plain
|
||||
if len(body) == 0 {
|
||||
body = htmlBody
|
||||
}
|
||||
for _, p := range body {
|
||||
text.WriteString(p + "\n\n")
|
||||
}
|
||||
parts := []Part{{Text: Clean(text.String())}}
|
||||
parts = append(parts, attachments...)
|
||||
return parts, nil
|
||||
}
|
||||
|
||||
// transferDecoder undoes a MIME part's Content-Transfer-Encoding.
|
||||
func transferDecoder(cte string, r io.Reader) io.Reader {
|
||||
switch strings.ToLower(strings.TrimSpace(cte)) {
|
||||
case "base64":
|
||||
return base64.NewDecoder(base64.StdEncoding, newlineStripper{r})
|
||||
case "quoted-printable":
|
||||
return quotedprintable.NewReader(r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// newlineStripper drops CR and LF, which base64 in mail is wrapped with and
|
||||
// encoding/base64 does not skip in every position.
|
||||
type newlineStripper struct{ r io.Reader }
|
||||
|
||||
func (n newlineStripper) Read(p []byte) (int, error) {
|
||||
for {
|
||||
k, err := n.r.Read(p)
|
||||
j := 0
|
||||
for _, c := range p[:k] {
|
||||
if c != '\r' && c != '\n' {
|
||||
p[j] = c
|
||||
j++
|
||||
}
|
||||
}
|
||||
if j > 0 || err != nil {
|
||||
if errors.Is(err, io.EOF) && j > 0 {
|
||||
return j, nil
|
||||
}
|
||||
return j, err
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
// Package extract turns documents into plain text for knowledge bases
|
||||
// (BACKLOG F2), on the machine that holds them.
|
||||
//
|
||||
// The text it produces is Markdown-shaped: headings as "# ", "## ", list
|
||||
// items as "- ", paragraphs separated by a blank line. The chunker splits on
|
||||
// exactly those, so a format that knows its headings (DOCX styles, HTML h1–h6,
|
||||
// ODT outline levels, slides, sheets) says so, and its chunks carry them.
|
||||
//
|
||||
// Built-in formats are pure Go, so the client builds for every system without
|
||||
// a C toolchain. What Go cannot read — Word 97 files, scanned pages — goes to
|
||||
// a helper when one is on the machine (LibreOffice, Tesseract, Poppler), and
|
||||
// is otherwise skipped with a sentence saying which helper would read it.
|
||||
package extract
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// ErrUnsupported is a file this cannot read, with the reason wrapped around
|
||||
// it ("a .doc needs LibreOffice").
|
||||
var ErrUnsupported = errors.New("unsupported")
|
||||
|
||||
// ErrNoText is a file that was read and held no text: a scanned PDF without
|
||||
// an OCR helper, an empty sheet.
|
||||
var ErrNoText = errors.New("no text in it")
|
||||
|
||||
// Part is the text of one document. A file is one part; an archive or a mail
|
||||
// with attachments is several, each named by its path inside, after the
|
||||
// outer file's name and "!/".
|
||||
type Part struct {
|
||||
// Name is "" for the file itself, or the path inside it.
|
||||
Name string
|
||||
Text string
|
||||
// Err is why this part was skipped, when it was; Text is then empty.
|
||||
Err error
|
||||
}
|
||||
|
||||
// Options are the limits and helpers extraction uses.
|
||||
type Options struct {
|
||||
Helpers Helpers
|
||||
// MaxBytes is the largest file, or member of an archive, read; 0 is
|
||||
// 256 MiB.
|
||||
MaxBytes int64
|
||||
// OCRLanguages are Tesseract's languages; "" is "tur+eng".
|
||||
OCRLanguages string
|
||||
}
|
||||
|
||||
func (o Options) maxBytes() int64 {
|
||||
if o.MaxBytes > 0 {
|
||||
return o.MaxBytes
|
||||
}
|
||||
return 256 << 20
|
||||
}
|
||||
|
||||
// File extracts a file on disk.
|
||||
func File(ctx context.Context, name string, o Options) ([]Part, error) {
|
||||
st, err := os.Stat(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if st.Size() > o.maxBytes() {
|
||||
return nil, fmt.Errorf("%w: %d MiB is over the %d MiB limit", ErrUnsupported, st.Size()>>20, o.maxBytes()>>20)
|
||||
}
|
||||
data, err := os.ReadFile(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return Bytes(ctx, filepath.Base(name), data, o)
|
||||
}
|
||||
|
||||
// Bytes extracts a document held in memory; name gives its format.
|
||||
func Bytes(ctx context.Context, name string, data []byte, o Options) ([]Part, error) {
|
||||
return extractDepth(ctx, name, data, o, 0)
|
||||
}
|
||||
|
||||
// maxDepth is how many archives deep members are followed: a zip in a mail
|
||||
// in a zip, and no further.
|
||||
const maxDepth = 3
|
||||
|
||||
func extractDepth(ctx context.Context, name string, data []byte, o Options, depth int) ([]Part, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ext := Ext(name)
|
||||
if f, ok := containers[ext]; ok {
|
||||
if depth >= maxDepth {
|
||||
return nil, fmt.Errorf("%w: archives nested more than %d deep", ErrUnsupported, maxDepth)
|
||||
}
|
||||
return f(ctx, name, data, o, depth)
|
||||
}
|
||||
text, err := one(ctx, name, ext, data, o)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
text = Clean(text)
|
||||
if strings.TrimSpace(text) == "" {
|
||||
return nil, ErrNoText
|
||||
}
|
||||
return []Part{{Text: text}}, nil
|
||||
}
|
||||
|
||||
// one extracts a single document.
|
||||
func one(ctx context.Context, name, ext string, data []byte, o Options) (string, error) {
|
||||
if f, ok := builtin[ext]; ok {
|
||||
text, err := f(data)
|
||||
if err == nil && strings.TrimSpace(text) != "" {
|
||||
return text, nil
|
||||
}
|
||||
// A PDF the Go reader cannot parse, or one with no text layer, is
|
||||
// what the helpers are for.
|
||||
if alt, ok := fallback[ext]; ok {
|
||||
if t, aerr := alt(ctx, name, data, o); aerr == nil {
|
||||
return t, nil
|
||||
} else if err == nil {
|
||||
err = aerr
|
||||
}
|
||||
}
|
||||
if err == nil {
|
||||
err = ErrNoText
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if alt, ok := fallback[ext]; ok {
|
||||
return alt(ctx, name, data, o)
|
||||
}
|
||||
if looksLikeText(data) {
|
||||
return decodeText(data), nil
|
||||
}
|
||||
if o.Helpers.Office != "" {
|
||||
return o.Helpers.officeToText(ctx, name, data, o)
|
||||
}
|
||||
return "", fmt.Errorf("%w: no reader for %q files (LibreOffice, if installed, may read it)", ErrUnsupported, ext)
|
||||
}
|
||||
|
||||
// Ext is a name's extension, lower case and without the dot.
|
||||
func Ext(name string) string {
|
||||
return strings.TrimPrefix(strings.ToLower(path.Ext(strings.ReplaceAll(name, `\`, "/"))), ".")
|
||||
}
|
||||
|
||||
// Supported lists the extensions read without a helper, and those read with
|
||||
// one, for --help and the docs.
|
||||
func Supported() (builtIn, withHelper []string) {
|
||||
for e := range builtin {
|
||||
builtIn = append(builtIn, e)
|
||||
}
|
||||
for e := range containers {
|
||||
builtIn = append(builtIn, e)
|
||||
}
|
||||
for e := range textExts {
|
||||
builtIn = append(builtIn, e)
|
||||
}
|
||||
for e := range fallback {
|
||||
if _, ok := builtin[e]; !ok {
|
||||
withHelper = append(withHelper, e)
|
||||
}
|
||||
}
|
||||
sort.Strings(builtIn)
|
||||
sort.Strings(withHelper)
|
||||
return builtIn, withHelper
|
||||
}
|
||||
|
||||
// Known reports whether a name's extension is one this reads, with or without
|
||||
// a helper, or one of the text formats.
|
||||
func Known(name string) bool {
|
||||
e := Ext(name)
|
||||
_, a := builtin[e]
|
||||
_, b := containers[e]
|
||||
_, c := fallback[e]
|
||||
return a || b || c || textExts[e]
|
||||
}
|
||||
|
||||
var (
|
||||
builtin map[string]func([]byte) (string, error)
|
||||
containers map[string]func(context.Context, string, []byte, Options, int) ([]Part, error)
|
||||
fallback map[string]func(context.Context, string, []byte, Options) (string, error)
|
||||
)
|
||||
|
||||
func init() {
|
||||
builtin = map[string]func([]byte) (string, error){
|
||||
"pdf": pdfText,
|
||||
"html": htmlText, "htm": htmlText, "xhtml": htmlText, "mht": mhtmlText, "mhtml": mhtmlText,
|
||||
"docx": docxText, "docm": docxText, "dotx": docxText, "dotm": docxText,
|
||||
"xlsx": xlsxText, "xlsm": xlsxText, "xltx": xlsxText, "xltm": xlsxText,
|
||||
"pptx": pptxText, "pptm": pptxText, "ppsx": pptxText, "ppsm": pptxText, "potx": pptxText,
|
||||
"odt": odfText, "ott": odfText, "ods": odfText, "ots": odfText, "odp": odfText, "otp": odfText,
|
||||
"fodt": flatODFText, "fods": flatODFText, "fodp": flatODFText,
|
||||
"rtf": rtfText,
|
||||
"epub": epubText,
|
||||
"ipynb": notebookText,
|
||||
"srt": subtitleText, "vtt": subtitleText,
|
||||
}
|
||||
for e := range textExts {
|
||||
builtin[e] = func(b []byte) (string, error) { return decodeText(b), nil }
|
||||
}
|
||||
containers = map[string]func(context.Context, string, []byte, Options, int) ([]Part, error){
|
||||
"zip": zipParts,
|
||||
"eml": mailParts,
|
||||
}
|
||||
fallback = map[string]func(context.Context, string, []byte, Options) (string, error){
|
||||
"pdf": pdfWithHelpers,
|
||||
"rtf": viaOffice,
|
||||
}
|
||||
for _, e := range officeOnly {
|
||||
fallback[e] = viaOffice
|
||||
}
|
||||
for _, e := range imageExts {
|
||||
fallback[e] = viaOCR
|
||||
}
|
||||
}
|
||||
|
||||
// textExts are formats that are text already: read as they are, decoded.
|
||||
var textExts = map[string]bool{
|
||||
"txt": true, "text": true, "md": true, "markdown": true, "mdx": true, "rst": true, "adoc": true,
|
||||
"asciidoc": true, "org": true, "tex": true, "log": true, "csv": true, "tsv": true, "json": true,
|
||||
"jsonl": true, "ndjson": true, "yaml": true, "yml": true, "toml": true, "ini": true, "cfg": true,
|
||||
"conf": true, "properties": true, "env": true, "xml": true, "svg": true, "sql": true,
|
||||
"go": true, "py": true, "js": true, "mjs": true, "cjs": true, "ts": true, "tsx": true, "jsx": true,
|
||||
"java": true, "kt": true, "kts": true, "scala": true, "c": true, "h": true, "cc": true, "cpp": true,
|
||||
"cxx": true, "hpp": true, "hh": true, "cs": true, "fs": true, "vb": true, "rs": true, "swift": true,
|
||||
"m": true, "mm": true, "rb": true, "php": true, "pl": true, "pm": true, "lua": true, "r": true,
|
||||
"jl": true, "dart": true, "ex": true, "exs": true, "erl": true, "hs": true, "clj": true,
|
||||
"sh": true, "bash": true, "zsh": true, "fish": true, "ps1": true, "psm1": true, "bat": true,
|
||||
"cmd": true, "groovy": true, "gradle": true, "tf": true, "hcl": true, "proto": true,
|
||||
"graphql": true, "gql": true, "css": true, "scss": true, "sass": true, "less": true,
|
||||
"vue": true, "svelte": true, "dockerfile": true, "makefile": true, "mk": true, "cmake": true,
|
||||
"bib": true, "nfo": true, "diff": true, "patch": true,
|
||||
}
|
||||
|
||||
// officeOnly are formats only LibreOffice reads here: the binary Office
|
||||
// formats before 2007, and others it opens.
|
||||
var officeOnly = []string{
|
||||
"doc", "dot", "xls", "xlt", "xlsb", "ppt", "pps", "pot", "wpd", "wps", "wri",
|
||||
"pages", "numbers", "key", "sxw", "sxc", "sxi", "hwp", "pub", "vsd", "vsdx",
|
||||
}
|
||||
|
||||
// imageExts are read by OCR.
|
||||
var imageExts = []string{"png", "jpg", "jpeg", "tif", "tiff", "bmp", "gif", "webp", "jp2", "pnm"}
|
||||
|
||||
// looksLikeText reports whether bytes are text: no NUL in the first 8 KiB
|
||||
// (UTF-16 has its BOM checked first), and mostly printable.
|
||||
func looksLikeText(b []byte) bool {
|
||||
if hasUTF16BOM(b) {
|
||||
return true
|
||||
}
|
||||
head := b[:min(len(b), 8192)]
|
||||
if bytes.IndexByte(head, 0) >= 0 {
|
||||
return false
|
||||
}
|
||||
if utf8.Valid(head) {
|
||||
return true
|
||||
}
|
||||
// A single-byte code page: allow it when control characters are rare.
|
||||
ctl := 0
|
||||
for _, c := range head {
|
||||
if c < 0x20 && c != '\n' && c != '\r' && c != '\t' && c != '\f' {
|
||||
ctl++
|
||||
}
|
||||
}
|
||||
return ctl*100 < len(head)
|
||||
}
|
||||
|
||||
var (
|
||||
spaceRun = regexp.MustCompile(`[ \t\x{00A0}]+\n`)
|
||||
blankRun = regexp.MustCompile(`\n{3,}`)
|
||||
)
|
||||
|
||||
// Clean normalises extracted text: LF line ends, no trailing spaces, no more
|
||||
// than one blank line in a row, no control characters but tab and newline.
|
||||
func Clean(s string) string {
|
||||
s = strings.ReplaceAll(s, "\r\n", "\n")
|
||||
s = strings.ReplaceAll(s, "\r", "\n")
|
||||
s = strings.Map(func(r rune) rune {
|
||||
switch {
|
||||
case r == '\n' || r == '\t':
|
||||
return r
|
||||
case r == '\f' || r == '\v':
|
||||
return '\n'
|
||||
case r < 0x20 || r == 0x7f || r == utf8.RuneError || r == '\uFEFF':
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, s)
|
||||
s = spaceRun.ReplaceAllString(s, "\n")
|
||||
s = blankRun.ReplaceAllString(s, "\n\n")
|
||||
return strings.TrimSpace(s) + "\n"
|
||||
}
|
||||
@@ -0,0 +1,323 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"golang.org/x/text/encoding/charmap"
|
||||
"golang.org/x/text/encoding/unicode"
|
||||
)
|
||||
|
||||
// zipOf builds a zip from name → content, in the order given.
|
||||
func zipOf(t *testing.T, files ...string) []byte {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
w := zip.NewWriter(&buf)
|
||||
for i := 0; i+1 < len(files); i += 2 {
|
||||
f, err := w.Create(files[i])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := f.Write([]byte(files[i+1])); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := w.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func extractOne(t *testing.T, name string, data []byte) string {
|
||||
t.Helper()
|
||||
parts, err := Bytes(context.Background(), name, data, Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
if len(parts) != 1 || parts[0].Err != nil {
|
||||
t.Fatalf("%s: parts %+v", name, parts)
|
||||
}
|
||||
return parts[0].Text
|
||||
}
|
||||
|
||||
func wants(t *testing.T, got string, want ...string) {
|
||||
t.Helper()
|
||||
for _, w := range want {
|
||||
if !strings.Contains(got, w) {
|
||||
t.Errorf("missing %q in:\n%s", w, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTextInEveryEncodingBecomesUTF8(t *testing.T) {
|
||||
const s = "Yıllık izin: Şubat'ta ğüşİöç"
|
||||
u16, _ := unicode.UTF16(unicode.LittleEndian, unicode.UseBOM).NewEncoder().Bytes([]byte(s))
|
||||
cp, _ := charmap.Windows1254.NewEncoder().Bytes([]byte(s))
|
||||
for name, b := range map[string][]byte{
|
||||
"utf8": []byte(s),
|
||||
"utf8-bom": append([]byte{0xEF, 0xBB, 0xBF}, s...),
|
||||
"utf16": u16,
|
||||
"cp1254": cp,
|
||||
"crlf": []byte(strings.ReplaceAll(s+"\nikinci", "\n", "\r\n")),
|
||||
} {
|
||||
got := extractOne(t, "a.txt", b)
|
||||
if !strings.HasPrefix(got, s) {
|
||||
t.Errorf("%s: got %q", name, got)
|
||||
}
|
||||
if strings.Contains(got, "\r") {
|
||||
t.Errorf("%s: kept a CR", name)
|
||||
}
|
||||
}
|
||||
// macOS writes names and sometimes text decomposed; the text comes out
|
||||
// composed, so the same word is the same bytes everywhere.
|
||||
if got := extractOne(t, "a.txt", []byte("Şirket")); !strings.HasPrefix(got, "Şirket") {
|
||||
t.Errorf("NFD text came out as %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTMLKeepsHeadingsAndDropsScripts(t *testing.T) {
|
||||
page := `<html><head><title>T</title><script>var x="gizli";</script></head><body>
|
||||
<nav>menü</nav><h1>İzin Politikası</h1><p>Yıllık <b>izin</b> yirmi gün.</p>
|
||||
<ul><li>Bir</li><li>İki</li></ul><table><tr><td>a</td><td>b</td></tr></table>
|
||||
<p><b>Mer</b>haba</p></body></html>`
|
||||
got := extractOne(t, "a.html", []byte(page))
|
||||
wants(t, got, "# İzin Politikası", "Yıllık izin yirmi gün.", "- Bir", "- İki", "a\tb", "Merhaba")
|
||||
for _, no := range []string{"gizli", "menü"} {
|
||||
if strings.Contains(got, no) {
|
||||
t.Errorf("%q leaked into:\n%s", no, got)
|
||||
}
|
||||
}
|
||||
// Declared charset honoured.
|
||||
l5, _ := charmap.ISO8859_9.NewEncoder().Bytes([]byte(`<meta charset="iso-8859-9"><p>Şirket</p>`))
|
||||
wants(t, extractOne(t, "b.htm", l5), "Şirket")
|
||||
}
|
||||
|
||||
func TestWordDocument(t *testing.T) {
|
||||
doc := `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body>
|
||||
<w:p><w:pPr><w:pStyle w:val="Balk1"/></w:pPr><w:r><w:t>Yıllık İzin</w:t></w:r></w:p>
|
||||
<w:p><w:r><w:t xml:space="preserve">Çalışanlar </w:t></w:r><w:r><w:t>yirmi gün</w:t></w:r><w:del><w:r><w:delText>otuz</w:delText></w:r></w:del><w:r><w:tab/><w:t>izin alır.</w:t></w:r></w:p>
|
||||
<w:p><w:pPr><w:numPr><w:ilvl w:val="0"/></w:numPr></w:pPr><w:r><w:t>Madde bir</w:t></w:r></w:p>
|
||||
<w:tbl><w:tr><w:tc><w:p><w:r><w:t>Hücre1</w:t></w:r></w:p></w:tc><w:tc><w:p><w:r><w:t>Hücre2</w:t></w:r></w:p></w:tc></w:tr></w:tbl>
|
||||
<w:p><w:pPr><w:pStyle w:val="Heading2"/></w:pPr><w:r><w:t>Masraflar</w:t></w:r></w:p>
|
||||
</w:body></w:document>`
|
||||
got := extractOne(t, "a.docx", zipOf(t, "word/document.xml", doc,
|
||||
"word/footnotes.xml", `<w:footnotes xmlns:w="x"><w:footnote><w:p><w:r><w:t>Dipnot</w:t></w:r></w:p></w:footnote></w:footnotes>`))
|
||||
wants(t, got, "# Yıllık İzin", "Çalışanlar yirmi gün\tizin alır.", "- Madde bir", "Hücre1\tHücre2", "## Masraflar", "Dipnot")
|
||||
if strings.Contains(got, "otuz") {
|
||||
t.Errorf("a tracked deletion is in the text:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkbook(t *testing.T) {
|
||||
got := extractOne(t, "a.xlsx", zipOf(t,
|
||||
"xl/workbook.xml", `<workbook xmlns:r="r"><sheets><sheet name="Bütçe" sheetId="1" r:id="rId1"/><sheet name="Boş" sheetId="2" r:id="rId2"/></sheets></workbook>`,
|
||||
"xl/_rels/workbook.xml.rels", `<Relationships><Relationship Id="rId1" Target="worksheets/sheet1.xml"/><Relationship Id="rId2" Target="worksheets/sheet2.xml"/></Relationships>`,
|
||||
"xl/sharedStrings.xml", `<sst><si><t>Kalem</t></si><si><r><t>Tu</t></r><r><t>tar</t></r></si><si><t>Kira</t></si></sst>`,
|
||||
"xl/worksheets/sheet1.xml", `<worksheet><sheetData><row><c t="s"><v>0</v></c><c t="s"><v>1</v></c></row><row><c t="s"><v>2</v></c><c><v>1500.5</v></c><c t="inlineStr"><is><t>not</t></is></c></row></sheetData></worksheet>`,
|
||||
"xl/worksheets/sheet2.xml", `<worksheet><sheetData/></worksheet>`,
|
||||
))
|
||||
wants(t, got, "## Bütçe", "Kalem\tTutar", "Kira\t1500.5\tnot")
|
||||
}
|
||||
|
||||
func TestPresentationInItsOwnOrderWithNotes(t *testing.T) {
|
||||
slide := func(text string) string {
|
||||
return `<p:sld xmlns:a="a" xmlns:p="p"><p:cSld><p:spTree><p:sp><p:txBody><a:p><a:r><a:t>` + text +
|
||||
`</a:t></a:r></a:p><a:p><a:fld type="slidenum"><a:t>9</a:t></a:fld></a:p></p:txBody></p:sp></p:spTree></p:cSld></p:sld>`
|
||||
}
|
||||
got := extractOne(t, "a.pptx", zipOf(t,
|
||||
"ppt/presentation.xml", `<p:presentation xmlns:p="p" xmlns:r="r"><p:sldIdLst><p:sldId id="256" r:id="rId3"/><p:sldId id="257" r:id="rId2"/></p:sldIdLst></p:presentation>`,
|
||||
"ppt/_rels/presentation.xml.rels", `<Relationships><Relationship Id="rId2" Target="slides/slide1.xml"/><Relationship Id="rId3" Target="slides/slide2.xml"/></Relationships>`,
|
||||
"ppt/slides/slide1.xml", slide("İkinci slayt"),
|
||||
"ppt/slides/slide2.xml", slide("Birinci slayt"),
|
||||
"ppt/slides/_rels/slide2.xml.rels", `<Relationships><Relationship Id="rId1" Type="http://x/notesSlide" Target="../notesSlides/notesSlide7.xml"/></Relationships>`,
|
||||
"ppt/notesSlides/notesSlide7.xml", slide("Konuşmacı notu"),
|
||||
))
|
||||
first, second := strings.Index(got, "Birinci slayt"), strings.Index(got, "İkinci slayt")
|
||||
if first < 0 || second < 0 || first > second {
|
||||
t.Errorf("slides out of the presentation's order:\n%s", got)
|
||||
}
|
||||
wants(t, got, "## Slide 1", "Notes:", "Konuşmacı notu")
|
||||
if strings.Contains(got, "9") {
|
||||
t.Errorf("a slide number field is in the text:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenDocument(t *testing.T) {
|
||||
content := `<office:document-content xmlns:office="o" xmlns:text="t" xmlns:table="tb"><office:body><office:text>
|
||||
<text:h text:outline-level="2">Bölüm</text:h><text:p>Bir<text:s text:c="2"/>iki<text:tab/>üç</text:p>
|
||||
<text:list><text:list-item><text:p>madde</text:p></text:list-item></text:list>
|
||||
<table:table table:name="Tablo1"><table:table-row><table:table-cell><text:p>x</text:p></table:table-cell><table:table-cell><text:p>y</text:p></table:table-cell></table:table-row></table:table>
|
||||
</office:text></office:body></office:document-content>`
|
||||
got := extractOne(t, "a.odt", zipOf(t, "mimetype", "application/vnd.oasis.opendocument.text", "content.xml", content))
|
||||
wants(t, got, "## Bölüm", "Bir iki\tüç", "- madde", "## Tablo1", "x\ty")
|
||||
wants(t, extractOne(t, "a.fodt", []byte(content)), "## Bölüm")
|
||||
}
|
||||
|
||||
func TestRichText(t *testing.T) {
|
||||
rtf := `{\rtf1\ansi\ansicpg1254\deff0{\fonttbl{\f0 Calibri;}}{\colortbl;\red0\green0\blue0;}` +
|
||||
`{\*\generator Riched20;}\pard Y\'fdll\'fdk izin\par ` +
|
||||
`\u350?irket \u-3913?\'3f ok\tab sonra\par{\header gizli}\cell b\row }`
|
||||
got := extractOne(t, "a.rtf", []byte(rtf))
|
||||
wants(t, got, "Yıllık izin", "Şirket", "ok\tsonra")
|
||||
for _, no := range []string{"Calibri", "Riched", "gizli"} {
|
||||
if strings.Contains(got, no) {
|
||||
t.Errorf("%q leaked into:\n%s", no, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEbookInSpineOrder(t *testing.T) {
|
||||
got := extractOne(t, "a.epub", zipOf(t,
|
||||
"mimetype", "application/epub+zip",
|
||||
"META-INF/container.xml", `<container><rootfiles><rootfile full-path="OEBPS/content.opf"/></rootfiles></container>`,
|
||||
"OEBPS/content.opf", `<package><manifest><item id="c1" href="b.xhtml"/><item id="c2" href="a.xhtml"/></manifest><spine><itemref idref="c2"/><itemref idref="c1"/></spine></package>`,
|
||||
"OEBPS/a.xhtml", `<html><body><h1>Önce</h1></body></html>`,
|
||||
"OEBPS/b.xhtml", `<html><body><h1>Sonra</h1></body></html>`,
|
||||
))
|
||||
if a, b := strings.Index(got, "Önce"), strings.Index(got, "Sonra"); a < 0 || b < a {
|
||||
t.Errorf("chapters out of spine order:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotebookAndSubtitles(t *testing.T) {
|
||||
nb := `{"metadata":{"kernelspec":{"language":"python"}},"cells":[{"cell_type":"markdown","source":["# Analiz\n","Veri"]},
|
||||
{"cell_type":"code","source":"print(1)","outputs":[{"text":"çıktı"}]}]}`
|
||||
got := extractOne(t, "a.ipynb", []byte(nb))
|
||||
wants(t, got, "# Analiz", "```python\nprint(1)\n```")
|
||||
if strings.Contains(got, "çıktı") {
|
||||
t.Errorf("an output is in the text:\n%s", got)
|
||||
}
|
||||
srt := "1\n00:00:01,000 --> 00:00:02,000\n<i>Merhaba</i>\n\n2\n00:00:02,000 --> 00:00:03,000\nMerhaba\nDünya\n"
|
||||
if got := extractOne(t, "a.srt", []byte(srt)); got != "Merhaba\nDünya\n" {
|
||||
t.Errorf("subtitles came out as %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFTextLayer(t *testing.T) {
|
||||
b, err := os.ReadFile("testdata/turkish.pdf")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wants(t, extractOne(t, "a.pdf", b), "Yıllık İzin Politikası", "Masraflar finans ekibine iletilir.")
|
||||
// Garbage is a failure, not a panic.
|
||||
if _, err := Bytes(context.Background(), "bad.pdf", []byte("%PDF-1.4 nonsense"), Options{}); err == nil {
|
||||
t.Error("a broken PDF was read")
|
||||
}
|
||||
}
|
||||
|
||||
func TestArchivesAreReadMemberByMember(t *testing.T) {
|
||||
inner := zipOf(t, "deep/c.md", "# Derin")
|
||||
archive := zipOf(t,
|
||||
"docs/a.txt", "Birinci",
|
||||
`docs\b.html`, "<p>İkinci</p>",
|
||||
"__MACOSX/docs/._a.txt", "çöp",
|
||||
"docs/~$kilit.docx", "çöp",
|
||||
"inner.zip", string(inner),
|
||||
"image.bin", "\x00\x01",
|
||||
)
|
||||
parts, err := Bytes(context.Background(), "arsiv.zip", archive, Options{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := map[string]string{}
|
||||
for _, p := range parts {
|
||||
if p.Err != nil {
|
||||
t.Errorf("%s: %v", p.Name, p.Err)
|
||||
}
|
||||
got[p.Name] = p.Text
|
||||
}
|
||||
if len(got) != 3 || !strings.Contains(got["docs/a.txt"], "Birinci") || !strings.Contains(got["docs/b.html"], "İkinci") ||
|
||||
!strings.Contains(got["inner.zip!/deep/c.md"], "# Derin") {
|
||||
t.Errorf("archive parts %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMailWithAttachments(t *testing.T) {
|
||||
att := zipOf(t, "x.txt", "ekteki metin")
|
||||
enc := func(b []byte) string {
|
||||
s := base64.StdEncoding.EncodeToString(b)
|
||||
var lines []string
|
||||
for len(s) > 76 {
|
||||
lines, s = append(lines, s[:76]), s[76:]
|
||||
}
|
||||
return strings.Join(append(lines, s), "\r\n")
|
||||
}
|
||||
msg := "From: =?UTF-8?B?QXnFn2U=?= <a@example.com>\r\n" +
|
||||
"Subject: =?iso-8859-9?Q?=DEubat_raporu?=\r\n" +
|
||||
"Content-Type: multipart/mixed; boundary=B\r\n\r\n" +
|
||||
"--B\r\nContent-Type: multipart/alternative; boundary=C\r\n\r\n" +
|
||||
"--C\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Transfer-Encoding: quoted-printable\r\n\r\n" +
|
||||
"Rapor ekte, te=C5=9Fekk=C3=BCrler.\r\n" +
|
||||
"--C\r\nContent-Type: text/html; charset=utf-8\r\n\r\n<p>HTML hali</p>\r\n--C--\r\n" +
|
||||
"--B\r\nContent-Type: application/zip; name=\"ek.zip\"\r\nContent-Transfer-Encoding: base64\r\n" +
|
||||
"Content-Disposition: attachment; filename=\"ek.zip\"\r\n\r\n" + enc(att) + "\r\n--B--\r\n"
|
||||
parts, err := Bytes(context.Background(), "a.eml", []byte(msg), Options{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("parts %+v", parts)
|
||||
}
|
||||
wants(t, parts[0].Text, "# Şubat raporu", "From: Ayşe", "Rapor ekte, teşekkürler.")
|
||||
if strings.Contains(parts[0].Text, "HTML hali") {
|
||||
t.Errorf("the HTML alternative was read beside the plain text:\n%s", parts[0].Text)
|
||||
}
|
||||
if parts[1].Name != "ek.zip!/x.txt" || !strings.Contains(parts[1].Text, "ekteki metin") {
|
||||
t.Errorf("attachment %+v", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnreadableFilesSayWhy(t *testing.T) {
|
||||
_, err := Bytes(context.Background(), "a.doc", []byte{0xD0, 0xCF, 0x11, 0xE0, 0}, Options{})
|
||||
if !errors.Is(err, ErrUnsupported) || !strings.Contains(err.Error(), "LibreOffice") {
|
||||
t.Errorf("a .doc without LibreOffice gave %v", err)
|
||||
}
|
||||
_, err = Bytes(context.Background(), "a.png", []byte{0x89, 'P', 'N', 'G', 0}, Options{})
|
||||
if !errors.Is(err, ErrUnsupported) || !strings.Contains(err.Error(), "Tesseract") {
|
||||
t.Errorf("an image without Tesseract gave %v", err)
|
||||
}
|
||||
_, err = Bytes(context.Background(), "a.bin", []byte{0, 1, 2, 3}, Options{})
|
||||
if !errors.Is(err, ErrUnsupported) {
|
||||
t.Errorf("binary gave %v", err)
|
||||
}
|
||||
_, err = Bytes(context.Background(), "a.txt", []byte(" \n"), Options{})
|
||||
if !errors.Is(err, ErrNoText) {
|
||||
t.Errorf("an empty file gave %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHelpersWhenInstalled(t *testing.T) {
|
||||
h := FindHelpers()
|
||||
ctx := context.Background()
|
||||
if h.Tesseract != "" && h.PdfToPPM != "" {
|
||||
// The text-layer PDF read by OCR instead: what a scanned one gets.
|
||||
b, _ := os.ReadFile("testdata/turkish.pdf")
|
||||
got, err := pdfWithHelpers(ctx, "a.pdf", b, Options{Helpers: Helpers{Tesseract: h.Tesseract, PdfToPPM: h.PdfToPPM}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wants(t, got, "Masraflar")
|
||||
} else {
|
||||
t.Log("no Tesseract with pdftoppm here; OCR not tried")
|
||||
}
|
||||
if h.Office == "" {
|
||||
t.Log("no LibreOffice here; legacy Office formats not tried")
|
||||
return
|
||||
}
|
||||
// A Word document converted to the legacy format and back would need a
|
||||
// fixture made by Word; an ODT through the Office path proves the
|
||||
// conversion runs and its result is read.
|
||||
content := `<office:document-content xmlns:office="urn:oasis:names:tc:opendocument:xmlns:office:1.0" xmlns:text="urn:oasis:names:tc:opendocument:xmlns:text:1.0"><office:body><office:text><text:p>LibreOffice yolu</text:p></office:text></office:body></office:document-content>`
|
||||
got, err := h.officeToText(ctx, "a.fodt", []byte(`<?xml version="1.0"?>`+content), Options{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wants(t, got, "LibreOffice yolu")
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Helpers are programs on this machine that read what Go here cannot. Any
|
||||
// may be "", and the formats that need it are then skipped with a reason.
|
||||
type Helpers struct {
|
||||
// Office is LibreOffice's soffice: the binary Office formats and others.
|
||||
Office string
|
||||
// Tesseract reads images and scanned pages.
|
||||
Tesseract string
|
||||
// PdfToText and PdfToPPM are Poppler's: a second PDF reader, and pages
|
||||
// rendered for Tesseract.
|
||||
PdfToText string
|
||||
PdfToPPM string
|
||||
}
|
||||
|
||||
// helperTimeout bounds one helper run; a conversion that hangs is a file
|
||||
// skipped, not a sync stopped.
|
||||
const helperTimeout = 3 * time.Minute
|
||||
|
||||
// FindHelpers looks for the helpers on PATH and where their installers put
|
||||
// them on each system.
|
||||
func FindHelpers() Helpers {
|
||||
return Helpers{
|
||||
Office: find("soffice", officePlaces()...),
|
||||
Tesseract: find("tesseract", tesseractPlaces()...),
|
||||
PdfToText: find("pdftotext"),
|
||||
PdfToPPM: find("pdftoppm"),
|
||||
}
|
||||
}
|
||||
|
||||
// Describe says which helpers were found, and what each missing one would
|
||||
// read.
|
||||
func (h Helpers) Describe() []string {
|
||||
line := func(have, name, reads string) string {
|
||||
if have != "" {
|
||||
return fmt.Sprintf("%-11s %s", name, have)
|
||||
}
|
||||
return fmt.Sprintf("%-11s not found — %s are skipped", name, reads)
|
||||
}
|
||||
return []string{
|
||||
line(h.Office, "LibreOffice", "Word 97–2003, Excel 97–2003, PowerPoint 97–2003, Pages, Numbers, Keynote and similar files"),
|
||||
line(h.Tesseract, "Tesseract", "images and scanned PDFs"),
|
||||
line(h.PdfToText, "pdftotext", "PDFs the built-in reader cannot parse"),
|
||||
line(h.PdfToPPM, "pdftoppm", "scanned PDFs (Tesseract needs their pages as images)"),
|
||||
}
|
||||
}
|
||||
|
||||
func find(name string, places ...string) string {
|
||||
exe := name
|
||||
if runtime.GOOS == "windows" {
|
||||
exe += ".exe"
|
||||
}
|
||||
if p, err := exec.LookPath(exe); err == nil {
|
||||
return p
|
||||
}
|
||||
if runtime.GOOS != "windows" {
|
||||
if p, err := exec.LookPath(name); err == nil {
|
||||
return p
|
||||
}
|
||||
}
|
||||
for _, p := range places {
|
||||
if matches, _ := filepath.Glob(p); len(matches) > 0 {
|
||||
sort.Strings(matches)
|
||||
return matches[len(matches)-1]
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func officePlaces() []string {
|
||||
switch runtime.GOOS {
|
||||
case "windows":
|
||||
return []string{
|
||||
`C:\Program Files\LibreOffice*\program\soffice.exe`,
|
||||
`C:\Program Files (x86)\LibreOffice*\program\soffice.exe`,
|
||||
}
|
||||
case "darwin":
|
||||
return []string{"/Applications/LibreOffice.app/Contents/MacOS/soffice"}
|
||||
}
|
||||
return []string{"/usr/bin/libreoffice", "/usr/lib/libreoffice/program/soffice", "/opt/libreoffice*/program/soffice",
|
||||
"/snap/bin/libreoffice"}
|
||||
}
|
||||
|
||||
func tesseractPlaces() []string {
|
||||
switch runtime.GOOS {
|
||||
case "windows":
|
||||
return []string{`C:\Program Files\Tesseract-OCR\tesseract.exe`, `C:\Program Files (x86)\Tesseract-OCR\tesseract.exe`}
|
||||
case "darwin":
|
||||
return []string{"/opt/homebrew/bin/tesseract", "/usr/local/bin/tesseract"}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// workDir is a temporary directory for one helper run, removed after.
|
||||
func workDir() (string, func(), error) {
|
||||
d, err := os.MkdirTemp("", "llmgw-kb-")
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
return d, func() { _ = os.RemoveAll(d) }, nil
|
||||
}
|
||||
|
||||
// writeInput puts a document where a helper can read it, under a safe name
|
||||
// that keeps its extension.
|
||||
func writeInput(dir, name string, data []byte) (string, error) {
|
||||
p := filepath.Join(dir, "input."+Ext(name))
|
||||
return p, os.WriteFile(p, data, 0o600)
|
||||
}
|
||||
|
||||
func run(ctx context.Context, dir string, bin string, args ...string) ([]byte, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, helperTimeout)
|
||||
defer cancel()
|
||||
cmd := exec.CommandContext(ctx, bin, args...)
|
||||
cmd.Dir = dir
|
||||
out, err := cmd.Output()
|
||||
if ctx.Err() != nil {
|
||||
return nil, fmt.Errorf("%s took over %s", filepath.Base(bin), helperTimeout)
|
||||
}
|
||||
if err != nil {
|
||||
var ee *exec.ExitError
|
||||
if errors.As(err, &ee) && len(ee.Stderr) > 0 {
|
||||
return nil, fmt.Errorf("%s: %v: %s", filepath.Base(bin), err, lastLine(ee.Stderr))
|
||||
}
|
||||
return nil, fmt.Errorf("%s: %w", filepath.Base(bin), err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func lastLine(b []byte) string {
|
||||
lines := strings.Split(strings.TrimSpace(string(b)), "\n")
|
||||
return strings.TrimSpace(lines[len(lines)-1])
|
||||
}
|
||||
|
||||
// fileURL is a path as LibreOffice takes its profile directory: a file URL,
|
||||
// with a drive letter on Windows.
|
||||
func fileURL(p string) string {
|
||||
p = filepath.ToSlash(p)
|
||||
if !strings.HasPrefix(p, "/") {
|
||||
p = "/" + p
|
||||
}
|
||||
return "file://" + p
|
||||
}
|
||||
|
||||
var spreadsheets = map[string]bool{"xls": true, "xlt": true, "xlsb": true, "numbers": true, "sxc": true}
|
||||
var presentations = map[string]bool{"ppt": true, "pps": true, "pot": true, "key": true, "sxi": true}
|
||||
|
||||
func viaOffice(ctx context.Context, name string, data []byte, o Options) (string, error) {
|
||||
if o.Helpers.Office == "" {
|
||||
return "", fmt.Errorf("%w: a .%s file needs LibreOffice, which is not installed here", ErrUnsupported, Ext(name))
|
||||
}
|
||||
return o.Helpers.officeToText(ctx, name, data, o)
|
||||
}
|
||||
|
||||
// officeToText has LibreOffice convert a document to the matching Office
|
||||
// Open XML format, then reads that: headings, sheets and slides survive,
|
||||
// which a conversion to plain text would lose.
|
||||
func (h Helpers) officeToText(ctx context.Context, name string, data []byte, _ Options) (string, error) {
|
||||
dir, done, err := workDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer done()
|
||||
in, err := writeInput(dir, name, data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
ext := Ext(name)
|
||||
to, read := "docx", docxText
|
||||
switch {
|
||||
case spreadsheets[ext]:
|
||||
to, read = "xlsx", xlsxText
|
||||
case presentations[ext]:
|
||||
to, read = "pptx", pptxText
|
||||
}
|
||||
// A profile of its own: two conversions at once with the user's
|
||||
// profile wait on each other's lock, or fail.
|
||||
profile := fileURL(filepath.Join(dir, "profile"))
|
||||
if _, err := run(ctx, dir, h.Office, "-env:UserInstallation="+profile, "--headless", "--norestore",
|
||||
"--convert-to", to, "--outdir", dir, in); err != nil {
|
||||
return "", err
|
||||
}
|
||||
converted, err := os.ReadFile(filepath.Join(dir, "input."+to))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("LibreOffice did not convert it: %w", err)
|
||||
}
|
||||
return read(converted)
|
||||
}
|
||||
|
||||
func viaOCR(ctx context.Context, name string, data []byte, o Options) (string, error) {
|
||||
if o.Helpers.Tesseract == "" {
|
||||
return "", fmt.Errorf("%w: an image needs Tesseract to be read, which is not installed here", ErrUnsupported)
|
||||
}
|
||||
dir, done, err := workDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer done()
|
||||
in, err := writeInput(dir, name, data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return ocr(ctx, dir, o, in)
|
||||
}
|
||||
|
||||
func ocr(ctx context.Context, dir string, o Options, file string) (string, error) {
|
||||
langs := o.OCRLanguages
|
||||
if langs == "" {
|
||||
langs = "tur+eng"
|
||||
}
|
||||
out, err := run(ctx, dir, o.Helpers.Tesseract, file, "stdout", "-l", langs)
|
||||
if err != nil && langs != "eng" {
|
||||
// A Tesseract without the Turkish data still reads English.
|
||||
out, err = run(ctx, dir, o.Helpers.Tesseract, file, "stdout", "-l", "eng")
|
||||
}
|
||||
return string(out), err
|
||||
}
|
||||
|
||||
// maxOCRPages bounds how many pages of one scanned PDF are read.
|
||||
const maxOCRPages = 500
|
||||
|
||||
// pdfWithHelpers reads a PDF the built-in reader could not: with pdftotext,
|
||||
// and when that finds no text either, page by page through Tesseract.
|
||||
func pdfWithHelpers(ctx context.Context, name string, data []byte, o Options) (string, error) {
|
||||
h := o.Helpers
|
||||
if h.PdfToText == "" && (h.Tesseract == "" || h.PdfToPPM == "") {
|
||||
return "", fmt.Errorf("%w: no text layer the built-in reader can use; Poppler (pdftotext) or "+
|
||||
"Tesseract with pdftoppm would read it", ErrNoText)
|
||||
}
|
||||
dir, done, err := workDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer done()
|
||||
in, err := writeInput(dir, name, data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if h.PdfToText != "" {
|
||||
if out, err := run(ctx, dir, h.PdfToText, "-enc", "UTF-8", in, "-"); err == nil && strings.TrimSpace(string(out)) != "" {
|
||||
return string(out), nil
|
||||
}
|
||||
}
|
||||
if h.Tesseract == "" || h.PdfToPPM == "" {
|
||||
return "", fmt.Errorf("%w: a scanned PDF needs Tesseract and pdftoppm", ErrNoText)
|
||||
}
|
||||
if _, err := run(ctx, dir, h.PdfToPPM, "-r", "300", "-gray", "-png", "-l", fmt.Sprint(maxOCRPages), in,
|
||||
filepath.Join(dir, "page")); err != nil {
|
||||
return "", err
|
||||
}
|
||||
pages, _ := filepath.Glob(filepath.Join(dir, "page-*.png"))
|
||||
// pdftoppm pads page numbers to the width of the last one, so the
|
||||
// names sort in page order.
|
||||
sort.Strings(pages)
|
||||
var all strings.Builder
|
||||
for _, p := range pages {
|
||||
t, err := ocr(ctx, dir, o, p)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
all.WriteString(t)
|
||||
all.WriteString("\n\n")
|
||||
}
|
||||
return all.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/mail"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/html"
|
||||
"golang.org/x/net/html/charset"
|
||||
)
|
||||
|
||||
// out builds Markdown-shaped text: paragraphs, headings, list items, table
|
||||
// rows, with whitespace collapsed inside them.
|
||||
type out struct {
|
||||
b strings.Builder
|
||||
// line is the text of the line being built, whitespace collapsed.
|
||||
line strings.Builder
|
||||
prefix string
|
||||
}
|
||||
|
||||
// text adds running text; runs of whitespace become one space, and none is
|
||||
// added where the source had none ("<b>Mer</b>haba" is one word).
|
||||
func (o *out) text(s string) {
|
||||
if s == "" {
|
||||
return
|
||||
}
|
||||
r := []rune(s)
|
||||
lead, trail := isSpace(r[0]), isSpace(r[len(r)-1])
|
||||
for i, f := range strings.FieldsFunc(s, isSpace) {
|
||||
if i > 0 || lead {
|
||||
o.space()
|
||||
}
|
||||
o.line.WriteString(f)
|
||||
}
|
||||
if trail {
|
||||
o.space()
|
||||
}
|
||||
}
|
||||
|
||||
// space adds one space unless the line is empty or already ends in one.
|
||||
func (o *out) space() {
|
||||
l := o.line.String()
|
||||
if l != "" && !strings.HasSuffix(l, " ") && !strings.HasSuffix(l, "\t") {
|
||||
o.line.WriteByte(' ')
|
||||
}
|
||||
}
|
||||
|
||||
// raw adds text as it is, for preformatted blocks.
|
||||
func (o *out) raw(s string) { o.line.WriteString(s) }
|
||||
|
||||
// tab separates cells.
|
||||
func (o *out) tab() {
|
||||
o.trimLine()
|
||||
o.line.WriteString("\t")
|
||||
}
|
||||
|
||||
func (o *out) trimLine() {
|
||||
s := strings.TrimRight(o.line.String(), " ")
|
||||
o.line.Reset()
|
||||
o.line.WriteString(s)
|
||||
}
|
||||
|
||||
// newline ends the line being built, if it has anything.
|
||||
func (o *out) newline() {
|
||||
s := strings.TrimSpace(o.line.String())
|
||||
o.line.Reset()
|
||||
if s == "" {
|
||||
return
|
||||
}
|
||||
o.b.WriteString(o.prefix)
|
||||
o.b.WriteString(s)
|
||||
o.b.WriteByte('\n')
|
||||
o.prefix = ""
|
||||
}
|
||||
|
||||
// para ends a paragraph: the line, and a blank line after it.
|
||||
func (o *out) para() {
|
||||
o.newline()
|
||||
if o.b.Len() > 0 && !strings.HasSuffix(o.b.String(), "\n\n") {
|
||||
o.b.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
|
||||
// heading starts a heading of a level, 1 to 6.
|
||||
func (o *out) heading(level int) {
|
||||
o.para()
|
||||
level = max(1, min(level, 6))
|
||||
o.prefix = strings.Repeat("#", level) + " "
|
||||
}
|
||||
|
||||
// item starts a list item.
|
||||
func (o *out) item() {
|
||||
o.newline()
|
||||
o.prefix = "- "
|
||||
}
|
||||
|
||||
func (o *out) String() string {
|
||||
o.newline()
|
||||
return o.b.String()
|
||||
}
|
||||
|
||||
func isSpace(r rune) bool {
|
||||
return r == ' ' || r == '\t' || r == '\n' || r == '\r' || r == '\f' || r == ' '
|
||||
}
|
||||
|
||||
// htmlSkip are elements whose content is not text a reader sees.
|
||||
var htmlSkip = map[string]bool{
|
||||
"script": true, "style": true, "noscript": true, "template": true, "svg": true,
|
||||
"head": true, "nav": true, "iframe": true, "object": true, "canvas": true, "button": true, "select": true,
|
||||
}
|
||||
|
||||
// htmlBlock are elements that end a paragraph.
|
||||
var htmlBlock = map[string]bool{
|
||||
"p": true, "div": true, "section": true, "article": true, "main": true, "header": true,
|
||||
"footer": true, "aside": true, "blockquote": true, "pre": true, "table": true, "ul": true,
|
||||
"ol": true, "dl": true, "figure": true, "figcaption": true, "address": true, "hr": true,
|
||||
"form": true, "fieldset": true, "details": true, "summary": true, "body": true, "dt": true, "dd": true,
|
||||
"caption": true,
|
||||
}
|
||||
|
||||
// htmlText reads an HTML page, in the charset it declares.
|
||||
func htmlText(b []byte) (string, error) {
|
||||
r, err := charset.NewReader(bytes.NewReader(b), "text/html")
|
||||
if err != nil {
|
||||
r = bytes.NewReader(b)
|
||||
}
|
||||
return htmlFrom(r)
|
||||
}
|
||||
|
||||
// htmlFrom reads HTML already decoded to UTF-8.
|
||||
func htmlFrom(r io.Reader) (string, error) {
|
||||
z := html.NewTokenizer(r)
|
||||
var o out
|
||||
skip, pre := 0, 0
|
||||
for {
|
||||
tt := z.Next()
|
||||
switch tt {
|
||||
case html.ErrorToken:
|
||||
if z.Err() == io.EOF {
|
||||
return o.String(), nil
|
||||
}
|
||||
return o.String(), z.Err()
|
||||
case html.TextToken:
|
||||
if skip > 0 {
|
||||
continue
|
||||
}
|
||||
if pre > 0 {
|
||||
o.raw(string(z.Text()))
|
||||
} else {
|
||||
o.text(string(z.Text()))
|
||||
}
|
||||
case html.StartTagToken, html.SelfClosingTagToken, html.EndTagToken:
|
||||
name, _ := z.TagName()
|
||||
tag := string(name)
|
||||
end := tt == html.EndTagToken
|
||||
if htmlSkip[tag] {
|
||||
if tt == html.StartTagToken {
|
||||
skip++
|
||||
} else if end && skip > 0 {
|
||||
skip--
|
||||
}
|
||||
continue
|
||||
}
|
||||
if skip > 0 {
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case len(tag) == 2 && tag[0] == 'h' && tag[1] >= '1' && tag[1] <= '6':
|
||||
if end {
|
||||
o.para()
|
||||
} else {
|
||||
o.heading(int(tag[1] - '0'))
|
||||
}
|
||||
case tag == "title" && !end:
|
||||
// Inside head, which is skipped; a title outside it is a
|
||||
// heading.
|
||||
o.heading(1)
|
||||
case tag == "li":
|
||||
if !end {
|
||||
o.item()
|
||||
} else {
|
||||
o.newline()
|
||||
}
|
||||
case tag == "br":
|
||||
o.newline()
|
||||
case tag == "tr":
|
||||
o.newline()
|
||||
case tag == "td" || tag == "th":
|
||||
if end {
|
||||
o.tab()
|
||||
}
|
||||
case tag == "pre":
|
||||
if end {
|
||||
pre--
|
||||
} else {
|
||||
pre++
|
||||
}
|
||||
o.para()
|
||||
case htmlBlock[tag]:
|
||||
o.para()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// mhtmlText reads a web page saved as one file (MHTML): a MIME message whose
|
||||
// HTML part is the page.
|
||||
func mhtmlText(b []byte) (string, error) {
|
||||
m, err := mail.ReadMessage(bytes.NewReader(b))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var best string
|
||||
err = walkMIME(m.Header.Get("Content-Type"), m.Header.Get("Content-Transfer-Encoding"), m.Body, 0,
|
||||
func(ctype, _ string, params map[string]string, body []byte) {
|
||||
if best != "" {
|
||||
return
|
||||
}
|
||||
switch ctype {
|
||||
case "text/html":
|
||||
best, _ = htmlFrom(strings.NewReader(decodeCharset(body, params["charset"])))
|
||||
}
|
||||
})
|
||||
return best, err
|
||||
}
|
||||
|
||||
// walkMIME calls leaf for every non-multipart part of a MIME body, decoded
|
||||
// from its transfer encoding.
|
||||
func walkMIME(contentType, cte string, body io.Reader, depth int,
|
||||
leaf func(ctype, filename string, params map[string]string, body []byte),
|
||||
) error {
|
||||
if depth > 8 {
|
||||
return nil
|
||||
}
|
||||
ctype, params, err := mime.ParseMediaType(contentType)
|
||||
if err != nil {
|
||||
ctype, params = "text/plain", map[string]string{}
|
||||
}
|
||||
if strings.HasPrefix(ctype, "multipart/") {
|
||||
mr := multipart.NewReader(body, params["boundary"])
|
||||
for {
|
||||
p, err := mr.NextRawPart()
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := walkMIME(p.Header.Get("Content-Type"), p.Header.Get("Content-Transfer-Encoding"), p, depth+1,
|
||||
func(ct, fn string, pr map[string]string, b []byte) {
|
||||
if fn == "" {
|
||||
fn = filenameOf(p.Header.Get("Content-Disposition"))
|
||||
}
|
||||
leaf(ct, fn, pr, b)
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
data, err := io.ReadAll(transferDecoder(cte, body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
leaf(ctype, params["name"], params, data)
|
||||
return nil
|
||||
}
|
||||
|
||||
func filenameOf(disposition string) string {
|
||||
_, p, err := mime.ParseMediaType(disposition)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return p["filename"]
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/ledongthuc/pdf"
|
||||
)
|
||||
|
||||
// epubText reads an e-book: its chapters in reading order (the spine).
|
||||
func epubText(b []byte) (string, error) {
|
||||
zr, err := openZip(b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
container, _ := member(zr, "META-INF/container.xml")
|
||||
if container == nil {
|
||||
return "", errors.New("not an EPUB: no META-INF/container.xml")
|
||||
}
|
||||
var opf string
|
||||
_ = walkXML(container, func(e xmlEvent) {
|
||||
if e.start && e.name == "rootfile" && opf == "" {
|
||||
opf = e.get("full-path")
|
||||
}
|
||||
})
|
||||
pkg, _ := member(zr, opf)
|
||||
if pkg == nil {
|
||||
return "", fmt.Errorf("not an EPUB: no package document %q", opf)
|
||||
}
|
||||
dir := path.Dir(opf)
|
||||
items := map[string]string{}
|
||||
var spine []string
|
||||
_ = walkXML(pkg, func(e xmlEvent) {
|
||||
switch {
|
||||
case e.start && e.name == "item":
|
||||
items[e.get("id")] = path.Clean(path.Join(dir, e.get("href")))
|
||||
case e.start && e.name == "itemref":
|
||||
spine = append(spine, e.get("idref"))
|
||||
}
|
||||
})
|
||||
var all strings.Builder
|
||||
for _, id := range spine {
|
||||
page, _ := member(zr, items[id])
|
||||
if page == nil {
|
||||
continue
|
||||
}
|
||||
t, err := htmlText(page)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
all.WriteString(t)
|
||||
all.WriteString("\n\n")
|
||||
}
|
||||
return all.String(), nil
|
||||
}
|
||||
|
||||
// notebookText reads a Jupyter notebook: Markdown cells as they are, code
|
||||
// cells fenced; outputs left out, since they are mostly numbers and images.
|
||||
func notebookText(b []byte) (string, error) {
|
||||
var nb struct {
|
||||
Cells []struct {
|
||||
CellType string `json:"cell_type"`
|
||||
Source json.RawMessage `json:"source"`
|
||||
} `json:"cells"`
|
||||
Metadata struct {
|
||||
Kernel struct {
|
||||
Language string `json:"language"`
|
||||
} `json:"kernelspec"`
|
||||
} `json:"metadata"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &nb); err != nil {
|
||||
return "", err
|
||||
}
|
||||
var out strings.Builder
|
||||
for _, c := range nb.Cells {
|
||||
src := jsonText(c.Source)
|
||||
if strings.TrimSpace(src) == "" {
|
||||
continue
|
||||
}
|
||||
switch c.CellType {
|
||||
case "code":
|
||||
out.WriteString("```" + nb.Metadata.Kernel.Language + "\n" + strings.TrimRight(src, "\n") + "\n```\n\n")
|
||||
default:
|
||||
out.WriteString(strings.TrimRight(src, "\n") + "\n\n")
|
||||
}
|
||||
}
|
||||
return out.String(), nil
|
||||
}
|
||||
|
||||
// jsonText is a notebook source: a string, or a list of lines.
|
||||
func jsonText(raw json.RawMessage) string {
|
||||
var s string
|
||||
if json.Unmarshal(raw, &s) == nil {
|
||||
return s
|
||||
}
|
||||
var lines []string
|
||||
if json.Unmarshal(raw, &lines) == nil {
|
||||
return strings.Join(lines, "")
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
var (
|
||||
subtitleTiming = regexp.MustCompile(`^\d{1,2}:\d{2}(:\d{2})?[.,]\d{3}\s+-->`)
|
||||
subtitleIndex = regexp.MustCompile(`^\d+$`)
|
||||
subtitleTag = regexp.MustCompile(`<[^>]+>`)
|
||||
)
|
||||
|
||||
// subtitleText reads SRT and WebVTT: the spoken lines, without numbers,
|
||||
// timings and styling.
|
||||
func subtitleText(b []byte) (string, error) {
|
||||
var out strings.Builder
|
||||
last := ""
|
||||
for _, line := range strings.Split(decodeText(b), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
switch {
|
||||
case line == "", line == "WEBVTT", strings.HasPrefix(line, "NOTE"),
|
||||
subtitleIndex.MatchString(line), subtitleTiming.MatchString(line):
|
||||
continue
|
||||
}
|
||||
line = strings.TrimSpace(subtitleTag.ReplaceAllString(line, ""))
|
||||
// Captions repeat a line across cues as it scrolls.
|
||||
if line == "" || line == last {
|
||||
continue
|
||||
}
|
||||
last = line
|
||||
out.WriteString(line + "\n")
|
||||
}
|
||||
return out.String(), nil
|
||||
}
|
||||
|
||||
// pdfText reads a PDF's text layer, page by page, rows in reading order.
|
||||
// A scanned PDF has none; pdfWithHelpers is then tried.
|
||||
func pdfText(b []byte) (text string, err error) {
|
||||
// The reader panics on some malformed files rather than failing; a bad
|
||||
// PDF is one skipped file, not a stopped sync.
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
text, err = "", fmt.Errorf("the PDF could not be parsed: %v", r)
|
||||
}
|
||||
}()
|
||||
r, err := pdf.NewReader(bytes.NewReader(b), int64(len(b)))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var o out
|
||||
for i := 1; i <= r.NumPage(); i++ {
|
||||
p := r.Page(i)
|
||||
if p.V.IsNull() {
|
||||
continue
|
||||
}
|
||||
rows, err := p.GetTextByRow()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, row := range rows {
|
||||
var line strings.Builder
|
||||
for _, w := range row.Content {
|
||||
line.WriteString(w.S)
|
||||
}
|
||||
o.text(line.String())
|
||||
o.newline()
|
||||
}
|
||||
o.para()
|
||||
}
|
||||
return o.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// odfText reads an OpenDocument file (LibreOffice's ODT, ODS, ODP): its
|
||||
// content.xml.
|
||||
func odfText(b []byte) (string, error) {
|
||||
zr, err := openZip(b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
content, err := member(zr, "content.xml")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if content == nil {
|
||||
return "", errors.New("not an OpenDocument file: no content.xml")
|
||||
}
|
||||
return odfXML(content)
|
||||
}
|
||||
|
||||
// flatODFText reads the single-file XML form (.fodt, .fods, .fodp).
|
||||
func flatODFText(b []byte) (string, error) { return odfXML(b) }
|
||||
|
||||
// odfXML reads OpenDocument's text model, which all three kinds share:
|
||||
// text:h with an outline level is a heading, text:p a paragraph, list items,
|
||||
// table rows and cells, slides (draw:page) and sheets (table:table) named as
|
||||
// headings.
|
||||
func odfXML(b []byte) (string, error) {
|
||||
var o out
|
||||
skip, slide := 0, 0
|
||||
// Paragraphs inside a cell or a list item do not end the row or the list.
|
||||
inCell, inItem := 0, 0
|
||||
err := walkXML(b, func(e xmlEvent) {
|
||||
switch {
|
||||
case e.start:
|
||||
switch e.name {
|
||||
case "tracked-changes", "notes-configuration", "annotation", "frame-name":
|
||||
skip++
|
||||
case "h":
|
||||
lvl, _ := strconv.Atoi(e.get("outline-level"))
|
||||
o.heading(max(lvl, 1))
|
||||
case "list-item":
|
||||
inItem++
|
||||
o.item()
|
||||
case "table-cell":
|
||||
inCell++
|
||||
case "tab":
|
||||
o.raw("\t")
|
||||
case "s":
|
||||
n, _ := strconv.Atoi(e.get("c"))
|
||||
for range max(n, 1) {
|
||||
o.raw(" ")
|
||||
}
|
||||
case "line-break":
|
||||
o.newline()
|
||||
case "page":
|
||||
slide++
|
||||
o.heading(2)
|
||||
name := e.get("name")
|
||||
if name == "" || isDefaultPageName(name) {
|
||||
name = fmt.Sprintf("Slide %d", slide)
|
||||
}
|
||||
o.text(name)
|
||||
o.para()
|
||||
case "table":
|
||||
if n := e.get("name"); n != "" {
|
||||
o.heading(2)
|
||||
o.text(n)
|
||||
o.para()
|
||||
}
|
||||
}
|
||||
case e.end:
|
||||
switch e.name {
|
||||
case "tracked-changes", "notes-configuration", "annotation", "frame-name":
|
||||
skip--
|
||||
case "h", "p":
|
||||
switch {
|
||||
case inCell > 0:
|
||||
o.text(" ")
|
||||
case inItem > 0:
|
||||
o.newline()
|
||||
default:
|
||||
o.para()
|
||||
}
|
||||
case "list-item":
|
||||
inItem--
|
||||
o.newline()
|
||||
case "table-cell":
|
||||
inCell--
|
||||
o.tab()
|
||||
case "table-row":
|
||||
o.newline()
|
||||
case "table", "page":
|
||||
o.para()
|
||||
}
|
||||
default:
|
||||
if skip == 0 {
|
||||
o.raw(e.text)
|
||||
}
|
||||
}
|
||||
})
|
||||
return o.String(), err
|
||||
}
|
||||
|
||||
// isDefaultPageName is LibreOffice's own name for an unnamed slide, "page1".
|
||||
func isDefaultPageName(s string) bool {
|
||||
if len(s) < 5 || s[:4] != "page" {
|
||||
return false
|
||||
}
|
||||
_, err := strconv.Atoi(s[4:])
|
||||
return err == nil
|
||||
}
|
||||
@@ -0,0 +1,463 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"encoding/xml"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"path"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// maxMember is the most one member of an Office file or EPUB is inflated
|
||||
// to: a zip bomb stops here rather than at the machine's memory.
|
||||
const maxMember = 512 << 20
|
||||
|
||||
func openZip(b []byte) (*zip.Reader, error) {
|
||||
return zip.NewReader(bytes.NewReader(b), int64(len(b)))
|
||||
}
|
||||
|
||||
// member reads one file inside a zip, by its exact name; nil when absent.
|
||||
func member(zr *zip.Reader, name string) ([]byte, error) {
|
||||
for _, f := range zr.File {
|
||||
if f.Name == name {
|
||||
return readMember(f)
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func readMember(f *zip.File) ([]byte, error) {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rc.Close()
|
||||
b, err := io.ReadAll(io.LimitReader(rc, maxMember+1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(b) > maxMember {
|
||||
return nil, fmt.Errorf("%s inflates past %d MiB", f.Name, maxMember>>20)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// relID is an element's r:id, the namespaced id that names a relationship —
|
||||
// not its own plain id, which sldId has as well.
|
||||
func (e xmlEvent) relID() string {
|
||||
for _, a := range e.attr {
|
||||
if a.Name.Local == "id" && a.Name.Space != "" {
|
||||
return a.Value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// xmlEvent is one step of walking an XML document: a start, an end, or
|
||||
// character data, with the element's local name and attributes.
|
||||
type xmlEvent struct {
|
||||
start, end bool
|
||||
name string
|
||||
attr []xml.Attr
|
||||
text string
|
||||
}
|
||||
|
||||
func (e xmlEvent) get(local string) string {
|
||||
for _, a := range e.attr {
|
||||
if a.Name.Local == local {
|
||||
return a.Value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// walkXML calls f for each element start and end and each run of character
|
||||
// data. Namespaces are dropped: OOXML and ODF are read by local names.
|
||||
func walkXML(b []byte, f func(xmlEvent)) error {
|
||||
d := xml.NewDecoder(bytes.NewReader(b))
|
||||
d.Strict = false
|
||||
d.CharsetReader = func(cs string, in io.Reader) (io.Reader, error) {
|
||||
all, err := io.ReadAll(in)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return strings.NewReader(decodeCharset(all, cs)), nil
|
||||
}
|
||||
for {
|
||||
tok, err := d.Token()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
switch t := tok.(type) {
|
||||
case xml.StartElement:
|
||||
f(xmlEvent{start: true, name: t.Name.Local, attr: t.Attr})
|
||||
case xml.EndElement:
|
||||
f(xmlEvent{end: true, name: t.Name.Local})
|
||||
case xml.CharData:
|
||||
f(xmlEvent{text: string(t)})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// headingLevel reads a heading level out of a paragraph style: Word's
|
||||
// "Heading1", "Heading 2", "heading3", its Turkish "Balk1" (the id Word gives
|
||||
// "Başlık 1"), "Title" and "Subtitle".
|
||||
func headingLevel(style string) int {
|
||||
s := strings.ToLower(strings.ReplaceAll(style, " ", ""))
|
||||
switch s {
|
||||
case "title", "konubaşlığı", "konubal":
|
||||
return 1
|
||||
case "subtitle", "altkonubaşlığı", "altkonubal":
|
||||
return 2
|
||||
}
|
||||
for _, p := range []string{"heading", "balk", "başlık", "berschrift", "titre", "título", "titolo", "kop"} {
|
||||
if rest, ok := strings.CutPrefix(s, p); ok {
|
||||
if n, err := strconv.Atoi(rest); err == nil && n >= 1 && n <= 9 {
|
||||
return min(n, 6)
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// docxText reads a Word document: the body, then footnotes and endnotes.
|
||||
// Paragraph styles give headings; tables are rows of tab-separated cells.
|
||||
func docxText(b []byte) (string, error) {
|
||||
zr, err := openZip(b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
body, err := member(zr, "word/document.xml")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if body == nil {
|
||||
return "", errors.New("not a Word document: no word/document.xml")
|
||||
}
|
||||
var o out
|
||||
if err := wordML(body, &o); err != nil {
|
||||
return "", err
|
||||
}
|
||||
for _, extra := range []string{"word/footnotes.xml", "word/endnotes.xml"} {
|
||||
if x, _ := member(zr, extra); x != nil {
|
||||
_ = wordML(x, &o)
|
||||
}
|
||||
}
|
||||
return o.String(), nil
|
||||
}
|
||||
|
||||
func wordML(b []byte, o *out) error {
|
||||
inText, inDel, inCell := false, 0, 0
|
||||
level, list := 0, false
|
||||
return walkXML(b, func(e xmlEvent) {
|
||||
switch {
|
||||
case e.start:
|
||||
switch e.name {
|
||||
case "p":
|
||||
level, list = 0, false
|
||||
case "pStyle":
|
||||
level = headingLevel(e.get("val"))
|
||||
if strings.Contains(strings.ToLower(e.get("val")), "list") {
|
||||
list = true
|
||||
}
|
||||
case "numPr":
|
||||
list = true
|
||||
case "t":
|
||||
inText = true
|
||||
if level > 0 && inCell == 0 && o.line.Len() == 0 && o.prefix == "" {
|
||||
o.heading(level)
|
||||
} else if list && inCell == 0 && o.line.Len() == 0 && o.prefix == "" {
|
||||
o.item()
|
||||
}
|
||||
case "delText":
|
||||
// Tracked deletions are not the document's text.
|
||||
inDel++
|
||||
case "tc":
|
||||
inCell++
|
||||
case "tab":
|
||||
o.raw("\t")
|
||||
case "br", "cr":
|
||||
o.newline()
|
||||
}
|
||||
case e.end:
|
||||
switch e.name {
|
||||
case "t":
|
||||
inText = false
|
||||
case "delText":
|
||||
inDel--
|
||||
case "p":
|
||||
switch {
|
||||
case inCell > 0:
|
||||
// A cell's paragraphs stay in its row.
|
||||
o.text(" ")
|
||||
case list:
|
||||
o.newline()
|
||||
default:
|
||||
o.para()
|
||||
}
|
||||
case "tc":
|
||||
inCell--
|
||||
o.tab()
|
||||
case "tr":
|
||||
o.newline()
|
||||
case "tbl":
|
||||
o.para()
|
||||
}
|
||||
default:
|
||||
if inText && inDel == 0 {
|
||||
o.raw(e.text)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// xlsxText reads a workbook: each sheet under its name as a heading, each
|
||||
// row a line of tab-separated cells, shared strings resolved.
|
||||
func xlsxText(b []byte) (string, error) {
|
||||
zr, err := openZip(b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var shared []string
|
||||
if ss, _ := member(zr, "xl/sharedStrings.xml"); ss != nil {
|
||||
var cur strings.Builder
|
||||
inT, inRPh := false, 0
|
||||
_ = walkXML(ss, func(e xmlEvent) {
|
||||
switch {
|
||||
case e.start && e.name == "si":
|
||||
cur.Reset()
|
||||
case e.end && e.name == "si":
|
||||
shared = append(shared, cur.String())
|
||||
case e.start && e.name == "rPh":
|
||||
inRPh++ // phonetic guides, not the cell's text
|
||||
case e.end && e.name == "rPh":
|
||||
inRPh--
|
||||
case e.start && e.name == "t":
|
||||
inT = true
|
||||
case e.end && e.name == "t":
|
||||
inT = false
|
||||
case !e.start && !e.end && inT && inRPh == 0:
|
||||
cur.WriteString(e.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
sheets, err := xlsxSheets(zr)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var o out
|
||||
for _, sh := range sheets {
|
||||
data, err := member(zr, sh.path)
|
||||
if err != nil || data == nil {
|
||||
continue
|
||||
}
|
||||
o.heading(2)
|
||||
o.text(sh.name)
|
||||
o.para()
|
||||
var cell strings.Builder
|
||||
cellType, inV, inIs := "", false, false
|
||||
first := true
|
||||
_ = walkXML(data, func(e xmlEvent) {
|
||||
switch {
|
||||
case e.start && e.name == "row":
|
||||
first = true
|
||||
case e.end && e.name == "row":
|
||||
o.newline()
|
||||
case e.start && e.name == "c":
|
||||
cellType = e.get("t")
|
||||
cell.Reset()
|
||||
case e.end && e.name == "c":
|
||||
v := cell.String()
|
||||
if cellType == "s" {
|
||||
if i, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && i >= 0 && i < len(shared) {
|
||||
v = shared[i]
|
||||
}
|
||||
}
|
||||
v = strings.TrimSpace(strings.ReplaceAll(v, "\n", " "))
|
||||
if v != "" {
|
||||
if !first {
|
||||
o.tab()
|
||||
}
|
||||
o.text(v)
|
||||
first = false
|
||||
}
|
||||
case e.start && e.name == "v":
|
||||
inV = true
|
||||
case e.end && e.name == "v":
|
||||
inV = false
|
||||
case e.start && e.name == "is":
|
||||
inIs = true
|
||||
case e.end && e.name == "is":
|
||||
inIs = false
|
||||
case !e.start && !e.end && (inV || inIs):
|
||||
cell.WriteString(e.text)
|
||||
}
|
||||
})
|
||||
o.para()
|
||||
}
|
||||
return o.String(), nil
|
||||
}
|
||||
|
||||
type sheetRef struct{ name, path string }
|
||||
|
||||
// xlsxSheets lists a workbook's sheets in their order, with the part each is
|
||||
// kept in.
|
||||
func xlsxSheets(zr *zip.Reader) ([]sheetRef, error) {
|
||||
wb, err := member(zr, "xl/workbook.xml")
|
||||
if err != nil || wb == nil {
|
||||
return nil, errors.New("not a workbook: no xl/workbook.xml")
|
||||
}
|
||||
rels := map[string]string{}
|
||||
if r, _ := member(zr, "xl/_rels/workbook.xml.rels"); r != nil {
|
||||
_ = walkXML(r, func(e xmlEvent) {
|
||||
if e.start && e.name == "Relationship" {
|
||||
rels[e.get("Id")] = e.get("Target")
|
||||
}
|
||||
})
|
||||
}
|
||||
var out []sheetRef
|
||||
_ = walkXML(wb, func(e xmlEvent) {
|
||||
if !e.start || e.name != "sheet" {
|
||||
return
|
||||
}
|
||||
target := rels[e.relID()]
|
||||
if target == "" {
|
||||
return
|
||||
}
|
||||
p := path.Clean(path.Join("xl", target))
|
||||
if strings.HasPrefix(target, "/") {
|
||||
p = strings.TrimPrefix(target, "/")
|
||||
}
|
||||
out = append(out, sheetRef{name: e.get("name"), path: p})
|
||||
})
|
||||
return out, nil
|
||||
}
|
||||
|
||||
var slideNum = regexp.MustCompile(`(\d+)\.xml$`)
|
||||
|
||||
func numbered(zr *zip.Reader, prefix string) []*zip.File {
|
||||
var fs []*zip.File
|
||||
for _, f := range zr.File {
|
||||
if strings.HasPrefix(f.Name, prefix) && slideNum.MatchString(f.Name) && !strings.Contains(f.Name[len(prefix):], "/") {
|
||||
fs = append(fs, f)
|
||||
}
|
||||
}
|
||||
n := func(f *zip.File) int {
|
||||
v, _ := strconv.Atoi(slideNum.FindStringSubmatch(f.Name)[1])
|
||||
return v
|
||||
}
|
||||
sort.Slice(fs, func(i, j int) bool { return n(fs[i]) < n(fs[j]) })
|
||||
return fs
|
||||
}
|
||||
|
||||
// rels reads a part's relationships: id → target, and type → target, with
|
||||
// targets resolved against the part's directory.
|
||||
func rels(zr *zip.Reader, part string) (byID, byType map[string]string) {
|
||||
byID, byType = map[string]string{}, map[string]string{}
|
||||
dir, base := path.Split(part)
|
||||
r, _ := member(zr, dir+"_rels/"+base+".rels")
|
||||
if r == nil {
|
||||
return byID, byType
|
||||
}
|
||||
_ = walkXML(r, func(e xmlEvent) {
|
||||
if !e.start || e.name != "Relationship" || e.get("TargetMode") == "External" {
|
||||
return
|
||||
}
|
||||
t := e.get("Target")
|
||||
if strings.HasPrefix(t, "/") {
|
||||
t = strings.TrimPrefix(t, "/")
|
||||
} else {
|
||||
t = path.Clean(path.Join(dir, t))
|
||||
}
|
||||
byID[e.get("Id")] = t
|
||||
ty := e.get("Type")
|
||||
byType[ty[strings.LastIndex(ty, "/")+1:]] = t
|
||||
})
|
||||
return byID, byType
|
||||
}
|
||||
|
||||
// pptxText reads a presentation: each slide, in the presentation's order, as
|
||||
// a heading, its text, and its speaker notes.
|
||||
func pptxText(b []byte) (string, error) {
|
||||
zr, err := openZip(b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var order []string
|
||||
if pres, _ := member(zr, "ppt/presentation.xml"); pres != nil {
|
||||
ids, _ := rels(zr, "ppt/presentation.xml")
|
||||
_ = walkXML(pres, func(e xmlEvent) {
|
||||
if e.start && e.name == "sldId" {
|
||||
if t := ids[e.relID()]; t != "" {
|
||||
order = append(order, t)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
if len(order) == 0 {
|
||||
for _, f := range numbered(zr, "ppt/slides/slide") {
|
||||
order = append(order, f.Name)
|
||||
}
|
||||
}
|
||||
if len(order) == 0 {
|
||||
return "", errors.New("not a presentation: no slides")
|
||||
}
|
||||
var o out
|
||||
drawing := func(d []byte) {
|
||||
inT, inField := false, 0
|
||||
_ = walkXML(d, func(e xmlEvent) {
|
||||
switch {
|
||||
case e.start && e.name == "fld":
|
||||
// Slide numbers and dates, not the slide's text.
|
||||
inField++
|
||||
case e.end && e.name == "fld":
|
||||
inField--
|
||||
case e.start && e.name == "t":
|
||||
inT = true
|
||||
case e.end && e.name == "t":
|
||||
inT = false
|
||||
case e.end && e.name == "p":
|
||||
o.newline()
|
||||
case e.start && e.name == "br":
|
||||
o.newline()
|
||||
case e.end && e.name == "sp":
|
||||
o.para()
|
||||
case e.end && e.name == "tc":
|
||||
o.tab()
|
||||
case e.end && e.name == "tr":
|
||||
o.newline()
|
||||
case !e.start && !e.end && inT && inField == 0:
|
||||
o.raw(e.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
for i, part := range order {
|
||||
d, err := member(zr, part)
|
||||
if err != nil || d == nil {
|
||||
continue
|
||||
}
|
||||
o.heading(2)
|
||||
o.text(fmt.Sprintf("Slide %d", i+1))
|
||||
o.para()
|
||||
drawing(d)
|
||||
if _, byType := rels(zr, part); byType["notesSlide"] != "" {
|
||||
if nd, _ := member(zr, byType["notesSlide"]); nd != nil {
|
||||
o.para()
|
||||
o.text("Notes:")
|
||||
o.newline()
|
||||
drawing(nd)
|
||||
}
|
||||
}
|
||||
o.para()
|
||||
}
|
||||
return o.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf16"
|
||||
|
||||
"golang.org/x/text/encoding/charmap"
|
||||
"golang.org/x/text/unicode/norm"
|
||||
)
|
||||
|
||||
// rtfSkip are destinations whose content is not the document's text.
|
||||
var rtfSkip = map[string]bool{
|
||||
"fonttbl": true, "colortbl": true, "stylesheet": true, "info": true, "pict": true,
|
||||
"header": true, "headerl": true, "headerr": true, "headerf": true, "footer": true,
|
||||
"footerl": true, "footerr": true, "footerf": true, "object": true, "xmlnstbl": true,
|
||||
"listtable": true, "listoverridetable": true, "rsidtbl": true, "generator": true,
|
||||
"themedata": true, "colorschememapping": true, "latentstyles": true, "datastore": true,
|
||||
"filetbl": true, "revtbl": true, "pgdsctbl": true, "fldinst": true, "bkmkstart": true,
|
||||
"bkmkend": true, "shppict": true, "nonshppict": true, "datafield": true, "blipuid": true,
|
||||
}
|
||||
|
||||
// rtfText reads Rich Text Format: text, paragraphs, tabs and table cells,
|
||||
// characters in the document's code page (\ansicpg) or as \u escapes;
|
||||
// font tables, pictures, headers and other destinations skipped.
|
||||
func rtfText(b []byte) (string, error) {
|
||||
s := string(b)
|
||||
if !strings.HasPrefix(strings.TrimSpace(s), "{\\rtf") {
|
||||
return "", errors.New("not RTF: no {\\rtf header")
|
||||
}
|
||||
type state struct {
|
||||
skip bool
|
||||
uc int // how many fallback characters follow a \u
|
||||
}
|
||||
cp := charmap.Windows1252
|
||||
stack := []state{{uc: 1}}
|
||||
var o out
|
||||
var pending []byte // \'hh bytes, decoded in runs so a code page's pairs stay together
|
||||
var surrogate rune
|
||||
flush := func() {
|
||||
if len(pending) == 0 {
|
||||
return
|
||||
}
|
||||
d, _ := cp.NewDecoder().Bytes(pending)
|
||||
o.raw(string(d))
|
||||
pending = pending[:0]
|
||||
}
|
||||
skipChars := 0
|
||||
cur := func() *state { return &stack[len(stack)-1] }
|
||||
for i := 0; i < len(s); i++ {
|
||||
c := s[i]
|
||||
switch c {
|
||||
case '{':
|
||||
flush()
|
||||
stack = append(stack, *cur())
|
||||
case '}':
|
||||
flush()
|
||||
if len(stack) > 1 {
|
||||
stack = stack[:len(stack)-1]
|
||||
}
|
||||
case '\\':
|
||||
if i+1 >= len(s) {
|
||||
break
|
||||
}
|
||||
n := s[i+1]
|
||||
switch {
|
||||
case n == '\'' && i+3 < len(s):
|
||||
v, err := strconv.ParseUint(s[i+2:i+4], 16, 8)
|
||||
i += 3
|
||||
if err != nil || cur().skip {
|
||||
continue
|
||||
}
|
||||
if skipChars > 0 {
|
||||
skipChars--
|
||||
continue
|
||||
}
|
||||
pending = append(pending, byte(v))
|
||||
case n == '*':
|
||||
i++
|
||||
cur().skip = true
|
||||
case n == '\\' || n == '{' || n == '}':
|
||||
i++
|
||||
if !cur().skip {
|
||||
flush()
|
||||
o.raw(string(n))
|
||||
}
|
||||
case n == '~':
|
||||
i++
|
||||
if !cur().skip {
|
||||
o.raw(" ")
|
||||
}
|
||||
case n == '-' || n == '_':
|
||||
i++
|
||||
case n == '\n' || n == '\r':
|
||||
i++
|
||||
if !cur().skip {
|
||||
flush()
|
||||
o.para()
|
||||
}
|
||||
case isLetter(n):
|
||||
j := i + 1
|
||||
for j < len(s) && isLetter(s[j]) {
|
||||
j++
|
||||
}
|
||||
word := s[i+1 : j]
|
||||
k := j
|
||||
if k < len(s) && (s[k] == '-' || (s[k] >= '0' && s[k] <= '9')) {
|
||||
k++
|
||||
for k < len(s) && s[k] >= '0' && s[k] <= '9' {
|
||||
k++
|
||||
}
|
||||
}
|
||||
arg, hasArg := 0, k > j
|
||||
if hasArg {
|
||||
arg, _ = strconv.Atoi(s[j:k])
|
||||
}
|
||||
if k < len(s) && s[k] == ' ' {
|
||||
k++
|
||||
}
|
||||
i = k - 1
|
||||
flush()
|
||||
if rtfSkip[word] {
|
||||
cur().skip = true
|
||||
continue
|
||||
}
|
||||
switch word {
|
||||
case "ansicpg":
|
||||
cp = codePage(arg)
|
||||
case "uc":
|
||||
cur().uc = arg
|
||||
case "u":
|
||||
if cur().skip {
|
||||
continue
|
||||
}
|
||||
r := rune(arg)
|
||||
if r < 0 {
|
||||
r += 65536
|
||||
}
|
||||
switch {
|
||||
case utf16.IsSurrogate(r) && r < 0xDC00:
|
||||
surrogate = r
|
||||
case utf16.IsSurrogate(r):
|
||||
o.raw(string(utf16.DecodeRune(surrogate, r)))
|
||||
surrogate = 0
|
||||
default:
|
||||
o.raw(string(r))
|
||||
}
|
||||
skipChars = cur().uc
|
||||
case "par", "sect", "page":
|
||||
if !cur().skip {
|
||||
o.para()
|
||||
}
|
||||
case "line":
|
||||
if !cur().skip {
|
||||
o.newline()
|
||||
}
|
||||
case "tab":
|
||||
if !cur().skip {
|
||||
o.raw("\t")
|
||||
}
|
||||
case "cell":
|
||||
if !cur().skip {
|
||||
o.tab()
|
||||
}
|
||||
case "row":
|
||||
if !cur().skip {
|
||||
o.newline()
|
||||
}
|
||||
case "emdash":
|
||||
o.raw("—")
|
||||
case "endash":
|
||||
o.raw("–")
|
||||
case "bullet":
|
||||
o.raw("•")
|
||||
case "lquote", "rquote":
|
||||
o.raw("'")
|
||||
case "ldblquote", "rdblquote":
|
||||
o.raw("\"")
|
||||
}
|
||||
default:
|
||||
i++
|
||||
}
|
||||
case '\r', '\n':
|
||||
// Line breaks in RTF source are not text.
|
||||
default:
|
||||
if cur().skip {
|
||||
continue
|
||||
}
|
||||
if skipChars > 0 {
|
||||
skipChars--
|
||||
continue
|
||||
}
|
||||
if c >= 0x80 {
|
||||
// Unescaped 8-bit text, as some writers leave it.
|
||||
pending = append(pending, c)
|
||||
continue
|
||||
}
|
||||
flush()
|
||||
o.raw(string(c))
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return norm.NFC.String(o.String()), nil
|
||||
}
|
||||
|
||||
func isLetter(c byte) bool { return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') }
|
||||
BIN
Binary file not shown.
@@ -0,0 +1,95 @@
|
||||
package extract
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"golang.org/x/net/html/charset"
|
||||
"golang.org/x/text/encoding/charmap"
|
||||
"golang.org/x/text/encoding/unicode"
|
||||
"golang.org/x/text/unicode/norm"
|
||||
)
|
||||
|
||||
func hasUTF16BOM(b []byte) bool {
|
||||
return len(b) >= 2 && ((b[0] == 0xFF && b[1] == 0xFE) || (b[0] == 0xFE && b[1] == 0xFF))
|
||||
}
|
||||
|
||||
// decodeText turns a text file into UTF-8, NFC.
|
||||
//
|
||||
// UTF-8 is taken as it is (its BOM dropped); UTF-16 by its BOM, which is how
|
||||
// Windows' Notepad writes "Unicode". Anything else is a single-byte code
|
||||
// page, and which one cannot be known from the bytes; Windows-1254 is
|
||||
// assumed, since it is what a Turkish Windows writes and it agrees with
|
||||
// Windows-1252 on every letter but six (ğĞıİşŞ in place of ðÐýÝþÞ) — which
|
||||
// are rare in the western languages 1252 serves.
|
||||
func decodeText(b []byte) string {
|
||||
var s string
|
||||
switch {
|
||||
case bytes.HasPrefix(b, []byte{0xEF, 0xBB, 0xBF}):
|
||||
s = string(b[3:])
|
||||
case hasUTF16BOM(b):
|
||||
out, err := unicode.UTF16(unicode.LittleEndian, unicode.ExpectBOM).NewDecoder().Bytes(b)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
s = string(out)
|
||||
case utf8.Valid(b):
|
||||
s = string(b)
|
||||
default:
|
||||
out, err := charmap.Windows1254.NewDecoder().Bytes(b)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
s = string(out)
|
||||
}
|
||||
return norm.NFC.String(strings.ReplaceAll(s, "\r\n", "\n"))
|
||||
}
|
||||
|
||||
// decodeCharset decodes bytes declared to be in a charset, as mail and web
|
||||
// pages declare theirs, by any of the WHATWG labels; an unknown or missing
|
||||
// one is treated as decodeText would.
|
||||
func decodeCharset(b []byte, cs string) string {
|
||||
cs = strings.ToLower(strings.TrimSpace(cs))
|
||||
if cs == "" || cs == "utf-8" || cs == "utf8" || cs == "us-ascii" {
|
||||
return decodeText(b)
|
||||
}
|
||||
enc, _ := charset.Lookup(cs)
|
||||
if enc == nil {
|
||||
return decodeText(b)
|
||||
}
|
||||
out, err := enc.NewDecoder().Bytes(b)
|
||||
if err != nil {
|
||||
return decodeText(b)
|
||||
}
|
||||
return norm.NFC.String(strings.ReplaceAll(string(out), "\r\n", "\n"))
|
||||
}
|
||||
|
||||
// codePage is the decoder for an RTF \ansicpgN or a Windows code page number.
|
||||
func codePage(n int) *charmap.Charmap {
|
||||
switch n {
|
||||
case 1250:
|
||||
return charmap.Windows1250
|
||||
case 1251:
|
||||
return charmap.Windows1251
|
||||
case 1252:
|
||||
return charmap.Windows1252
|
||||
case 1253:
|
||||
return charmap.Windows1253
|
||||
case 1254:
|
||||
return charmap.Windows1254
|
||||
case 1255:
|
||||
return charmap.Windows1255
|
||||
case 1256:
|
||||
return charmap.Windows1256
|
||||
case 1257:
|
||||
return charmap.Windows1257
|
||||
case 1258:
|
||||
return charmap.Windows1258
|
||||
case 850:
|
||||
return charmap.CodePage850
|
||||
case 437:
|
||||
return charmap.CodePage437
|
||||
}
|
||||
return charmap.Windows1252
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
package kbsync
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Client speaks to a gateway's knowledge-base and embeddings endpoints with
|
||||
// a person's key.
|
||||
type Client struct {
|
||||
// Base is the gateway's address, "https://llm.example.com".
|
||||
Base string
|
||||
Key string
|
||||
// Project is the project a personal key acts in.
|
||||
Project string
|
||||
HTTP *http.Client
|
||||
// Wait is how long a request waits out a model loading or a busy
|
||||
// gateway (429/503 with Retry-After) before giving up; 0 is 15 minutes.
|
||||
Wait time.Duration
|
||||
// Waiting, when set, is told each time a request waits, and why.
|
||||
Waiting func(why string, d time.Duration)
|
||||
}
|
||||
|
||||
// NewHTTP is an HTTP client trusting the system's certificates, and those in
|
||||
// caFile when one is given.
|
||||
func NewHTTP(caFile string) (*http.Client, error) {
|
||||
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||
if caFile != "" {
|
||||
pem, err := os.ReadFile(caFile)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pool, err := x509.SystemCertPool()
|
||||
if err != nil || pool == nil {
|
||||
pool = x509.NewCertPool()
|
||||
}
|
||||
if !pool.AppendCertsFromPEM(pem) {
|
||||
return nil, fmt.Errorf("no certificate in %s", caFile)
|
||||
}
|
||||
tr.TLSClientConfig = &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12}
|
||||
}
|
||||
return &http.Client{Transport: tr, Timeout: 10 * time.Minute}, nil
|
||||
}
|
||||
|
||||
// Collection is a collection as the gateway shows it to this key.
|
||||
type Collection struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Owner string `json:"owner"`
|
||||
OwnerName string `json:"owner_name"`
|
||||
EmbedModel string `json:"embed_model"`
|
||||
Dims int `json:"dims"`
|
||||
Writable bool `json:"writable"`
|
||||
}
|
||||
|
||||
// RemoteDoc is a document the collection holds.
|
||||
type RemoteDoc struct {
|
||||
Name string `json:"name"`
|
||||
Digest string `json:"digest"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Chunks int `json:"chunks"`
|
||||
AddedAt time.Time `json:"added_at"`
|
||||
}
|
||||
|
||||
// Manifest is a collection and its documents.
|
||||
type Manifest struct {
|
||||
Collection Collection `json:"collection"`
|
||||
Documents []RemoteDoc `json:"documents"`
|
||||
}
|
||||
|
||||
// ChunkOut is one chunk as sent.
|
||||
type ChunkOut struct {
|
||||
Text string `json:"text"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
}
|
||||
|
||||
// DocOut is one document as sent.
|
||||
type DocOut struct {
|
||||
Name string `json:"name"`
|
||||
Digest string `json:"digest"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Chunks []ChunkOut `json:"chunks"`
|
||||
}
|
||||
|
||||
// APIError is an answer that was not a success.
|
||||
type APIError struct {
|
||||
Status int
|
||||
Message string
|
||||
}
|
||||
|
||||
func (e *APIError) Error() string { return fmt.Sprintf("%d: %s", e.Status, e.Message) }
|
||||
|
||||
// do sends a request and decodes a JSON answer into out, waiting out answers
|
||||
// that say to retry.
|
||||
func (c *Client) do(ctx context.Context, method, path string, body any, out any) error {
|
||||
var payload []byte
|
||||
if body != nil {
|
||||
var err error
|
||||
if payload, err = json.Marshal(body); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
wait := c.Wait
|
||||
if wait <= 0 {
|
||||
wait = 15 * time.Minute
|
||||
}
|
||||
deadline := time.Now().Add(wait)
|
||||
for {
|
||||
req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(c.Base, "/")+path, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+c.Key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", "llmgw-kb")
|
||||
if c.Project != "" {
|
||||
req.Header.Set("X-Llmgw-Project", c.Project)
|
||||
}
|
||||
res, err := c.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(res.Body, 256<<20))
|
||||
res.Body.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if res.StatusCode < 300 {
|
||||
if out == nil {
|
||||
return nil
|
||||
}
|
||||
return json.Unmarshal(data, out)
|
||||
}
|
||||
retry := res.Header.Get("Retry-After")
|
||||
if retry != "" && (res.StatusCode == http.StatusTooManyRequests || res.StatusCode >= 500) && time.Now().Before(deadline) {
|
||||
secs, _ := strconv.Atoi(retry)
|
||||
d := time.Duration(max(secs, 1)) * time.Second
|
||||
if c.Waiting != nil {
|
||||
c.Waiting(errorMessage(data, res.Status), d)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(d):
|
||||
}
|
||||
continue
|
||||
}
|
||||
return &APIError{Status: res.StatusCode, Message: errorMessage(data, res.Status)}
|
||||
}
|
||||
}
|
||||
|
||||
// errorMessage is the message in an OpenAI-shaped error, or the body.
|
||||
func errorMessage(data []byte, status string) string {
|
||||
var e struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(data, &e) == nil && e.Error.Message != "" {
|
||||
return e.Error.Message
|
||||
}
|
||||
if s := strings.TrimSpace(string(data)); s != "" && len(s) < 500 {
|
||||
return s
|
||||
}
|
||||
return status
|
||||
}
|
||||
|
||||
// Collections lists what this key may read.
|
||||
func (c *Client) Collections(ctx context.Context) ([]Collection, error) {
|
||||
var out struct {
|
||||
Data []Collection `json:"data"`
|
||||
}
|
||||
err := c.do(ctx, http.MethodGet, "/v1/kb/collections", nil, &out)
|
||||
return out.Data, err
|
||||
}
|
||||
|
||||
// Manifest reads a collection's documents, by its name or id.
|
||||
func (c *Client) Manifest(ctx context.Context, collection string) (Manifest, error) {
|
||||
var m Manifest
|
||||
err := c.do(ctx, http.MethodGet, "/v1/kb/collections/"+url.PathEscape(collection)+"/manifest", nil, &m)
|
||||
return m, err
|
||||
}
|
||||
|
||||
// Put sends a document; changed is false when the collection already had it.
|
||||
func (c *Client) Put(ctx context.Context, collection string, d DocOut) (bool, error) {
|
||||
var out struct {
|
||||
Changed bool `json:"changed"`
|
||||
}
|
||||
err := c.do(ctx, http.MethodPut, "/v1/kb/collections/"+url.PathEscape(collection)+"/documents", d, &out)
|
||||
return out.Changed, err
|
||||
}
|
||||
|
||||
// Delete removes a document; one already gone is not an error.
|
||||
func (c *Client) Delete(ctx context.Context, collection, name string) error {
|
||||
err := c.do(ctx, http.MethodDelete, "/v1/kb/collections/"+url.PathEscape(collection)+"/documents?name="+url.QueryEscape(name), nil, nil)
|
||||
var ae *APIError
|
||||
if errors.As(err, &ae) && ae.Status == http.StatusNotFound && strings.HasPrefix(ae.Message, "no document") {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// embedBatch is how many chunks go in one embeddings request.
|
||||
const embedBatch = 32
|
||||
|
||||
// Embed has the gateway embed texts with a model, in batches, in order.
|
||||
func (c *Client) Embed(ctx context.Context, model string, texts []string) ([][]float32, error) {
|
||||
out := make([][]float32, 0, len(texts))
|
||||
for start := 0; start < len(texts); start += embedBatch {
|
||||
end := min(start+embedBatch, len(texts))
|
||||
var res struct {
|
||||
Data []struct {
|
||||
Index int `json:"index"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := c.do(ctx, http.MethodPost, "/v1/embeddings",
|
||||
map[string]any{"model": model, "input": texts[start:end], "encoding_format": "float"}, &res); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(res.Data) != end-start {
|
||||
return nil, fmt.Errorf("the gateway embedded %d of %d texts", len(res.Data), end-start)
|
||||
}
|
||||
batch := make([][]float32, end-start)
|
||||
for _, d := range res.Data {
|
||||
if d.Index < 0 || d.Index >= len(batch) {
|
||||
return nil, fmt.Errorf("an embedding came back with index %d", d.Index)
|
||||
}
|
||||
batch[d.Index] = d.Embedding
|
||||
}
|
||||
out = append(out, batch...)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
// Package kbsync is llmgw-kb's work (BACKLOG F2): it finds a person's
|
||||
// documents, keeps a ledger of what it has done with each, and sends a
|
||||
// knowledge base only what changed — text and vectors, never the files.
|
||||
package kbsync
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
_ "modernc.org/sqlite" // pure Go, so the client builds for every system
|
||||
)
|
||||
|
||||
// Ledger is the record, on this machine, of every document sent or skipped:
|
||||
// which file it came from, its size, modification time and digest, and what
|
||||
// happened. A file whose size and time have not changed is not read again to
|
||||
// learn that it has not changed, and --delete removes only what this
|
||||
// machine sent.
|
||||
type Ledger struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
// State is what happened to a document.
|
||||
type State string
|
||||
|
||||
// The states a document is left in.
|
||||
const (
|
||||
Sent State = "sent"
|
||||
Skipped State = "skipped"
|
||||
Failed State = "failed"
|
||||
)
|
||||
|
||||
// Entry is one document in the ledger.
|
||||
type Entry struct {
|
||||
Gateway string
|
||||
Collection string
|
||||
// Name is the document's name in the collection.
|
||||
Name string
|
||||
// File is the path on disk it came from, absolute.
|
||||
File string
|
||||
Size int64
|
||||
MTime time.Time
|
||||
Digest string
|
||||
State State
|
||||
Reason string
|
||||
Chunks int
|
||||
At time.Time
|
||||
}
|
||||
|
||||
// Run is one sync, for history.
|
||||
type Run struct {
|
||||
ID int64
|
||||
Gateway, Collection string
|
||||
Started, Finished time.Time
|
||||
Sent, Unchanged, Skipped, Failed, Delete int
|
||||
}
|
||||
|
||||
// DefaultLedgerPath is where the ledger is kept: the user's configuration
|
||||
// directory, which, unlike a cache directory, nothing cleans up.
|
||||
func DefaultLedgerPath() (string, error) {
|
||||
dir, err := os.UserConfigDir()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Join(dir, "llmgw-kb", "ledger.db"), nil
|
||||
}
|
||||
|
||||
const ledgerSchema = `
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
gateway TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
file TEXT NOT NULL,
|
||||
size INTEGER NOT NULL,
|
||||
mtime_ns INTEGER NOT NULL,
|
||||
digest TEXT NOT NULL,
|
||||
state TEXT NOT NULL,
|
||||
reason TEXT NOT NULL DEFAULT '',
|
||||
chunks INTEGER NOT NULL DEFAULT 0,
|
||||
at_ms INTEGER NOT NULL,
|
||||
PRIMARY KEY (gateway, collection, name)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS documents_by_file ON documents (file);
|
||||
CREATE TABLE IF NOT EXISTS runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
gateway TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
started_ms INTEGER NOT NULL,
|
||||
ended_ms INTEGER NOT NULL DEFAULT 0,
|
||||
sent INTEGER NOT NULL DEFAULT 0,
|
||||
unchanged INTEGER NOT NULL DEFAULT 0,
|
||||
skipped INTEGER NOT NULL DEFAULT 0,
|
||||
failed INTEGER NOT NULL DEFAULT 0,
|
||||
deleted INTEGER NOT NULL DEFAULT 0
|
||||
);`
|
||||
|
||||
// OpenLedger opens (creating) the ledger at a path.
|
||||
func OpenLedger(ctx context.Context, path string) (*Ledger, error) {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// WAL and a busy timeout: two syncs at once, from two terminals, wait
|
||||
// for each other rather than failing.
|
||||
db, err := sql.Open("sqlite", path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(10000)&_pragma=synchronous(NORMAL)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
db.SetMaxOpenConns(1)
|
||||
if _, err := db.ExecContext(ctx, ledgerSchema); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("opening the ledger %s: %w", path, err)
|
||||
}
|
||||
_ = os.Chmod(path, 0o600)
|
||||
return &Ledger{db: db}, nil
|
||||
}
|
||||
|
||||
// Close closes the ledger.
|
||||
func (l *Ledger) Close() error { return l.db.Close() }
|
||||
|
||||
const entryCols = `gateway, collection, name, file, size, mtime_ns, digest, state, reason, chunks, at_ms`
|
||||
|
||||
func scanEntry(sc interface{ Scan(...any) error }) (Entry, error) {
|
||||
var e Entry
|
||||
var mtime, at int64
|
||||
var state string
|
||||
err := sc.Scan(&e.Gateway, &e.Collection, &e.Name, &e.File, &e.Size, &mtime, &e.Digest, &state, &e.Reason, &e.Chunks, &at)
|
||||
e.MTime, e.At, e.State = time.Unix(0, mtime), time.UnixMilli(at), State(state)
|
||||
return e, err
|
||||
}
|
||||
|
||||
// Put records what happened to a document.
|
||||
func (l *Ledger) Put(ctx context.Context, e Entry) error {
|
||||
if e.At.IsZero() {
|
||||
e.At = time.Now()
|
||||
}
|
||||
_, err := l.db.ExecContext(ctx, `INSERT INTO documents (`+entryCols+`) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT (gateway, collection, name) DO UPDATE SET file = excluded.file, size = excluded.size,
|
||||
mtime_ns = excluded.mtime_ns, digest = excluded.digest, state = excluded.state,
|
||||
reason = excluded.reason, chunks = excluded.chunks, at_ms = excluded.at_ms`,
|
||||
e.Gateway, e.Collection, e.Name, e.File, e.Size, e.MTime.UnixNano(), e.Digest, string(e.State),
|
||||
e.Reason, e.Chunks, e.At.UnixMilli())
|
||||
return err
|
||||
}
|
||||
|
||||
// Remove forgets a document.
|
||||
func (l *Ledger) Remove(ctx context.Context, gateway, collection, name string) error {
|
||||
_, err := l.db.ExecContext(ctx, `DELETE FROM documents WHERE gateway = ? AND collection = ? AND name = ?`,
|
||||
gateway, collection, name)
|
||||
return err
|
||||
}
|
||||
|
||||
// Entries lists a collection's documents as this machine left them.
|
||||
func (l *Ledger) Entries(ctx context.Context, gateway, collection string) ([]Entry, error) {
|
||||
rows, err := l.db.QueryContext(ctx, `SELECT `+entryCols+` FROM documents WHERE gateway = ? AND collection = ? ORDER BY name`,
|
||||
gateway, collection)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Entry
|
||||
for rows.Next() {
|
||||
e, err := scanEntry(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, e)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// Known is the digest the ledger has for a file of this size and time, from
|
||||
// any collection: a file is hashed once however many collections it goes to.
|
||||
func (l *Ledger) Known(ctx context.Context, file string, size int64, mtime time.Time) (string, bool) {
|
||||
var d string
|
||||
err := l.db.QueryRowContext(ctx, `SELECT digest FROM documents WHERE file = ? AND size = ? AND mtime_ns = ? LIMIT 1`,
|
||||
file, size, mtime.UnixNano()).Scan(&d)
|
||||
return d, err == nil
|
||||
}
|
||||
|
||||
// StartRun records that a sync began.
|
||||
func (l *Ledger) StartRun(ctx context.Context, gateway, collection string) (int64, error) {
|
||||
res, err := l.db.ExecContext(ctx, `INSERT INTO runs (gateway, collection, started_ms) VALUES (?, ?, ?)`,
|
||||
gateway, collection, time.Now().UnixMilli())
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.LastInsertId()
|
||||
}
|
||||
|
||||
// EndRun records how a sync ended.
|
||||
func (l *Ledger) EndRun(ctx context.Context, id int64, r Run) error {
|
||||
_, err := l.db.ExecContext(ctx, `UPDATE runs SET ended_ms = ?, sent = ?, unchanged = ?, skipped = ?, failed = ?, deleted = ?
|
||||
WHERE id = ?`, time.Now().UnixMilli(), r.Sent, r.Unchanged, r.Skipped, r.Failed, r.Delete, id)
|
||||
return err
|
||||
}
|
||||
|
||||
// Runs lists the latest syncs, newest first.
|
||||
func (l *Ledger) Runs(ctx context.Context, limit int) ([]Run, error) {
|
||||
rows, err := l.db.QueryContext(ctx, `SELECT id, gateway, collection, started_ms, ended_ms, sent, unchanged, skipped, failed, deleted
|
||||
FROM runs ORDER BY id DESC LIMIT ?`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Run
|
||||
for rows.Next() {
|
||||
var r Run
|
||||
var s, e int64
|
||||
if err := rows.Scan(&r.ID, &r.Gateway, &r.Collection, &s, &e, &r.Sent, &r.Unchanged, &r.Skipped, &r.Failed, &r.Delete); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r.Started = time.UnixMilli(s)
|
||||
if e > 0 {
|
||||
r.Finished = time.UnixMilli(e)
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,574 @@
|
||||
package kbsync
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/text/unicode/norm"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/chunk"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/extract"
|
||||
)
|
||||
|
||||
// Options say what to sync and how.
|
||||
type Options struct {
|
||||
// Collection is the collection's name or id.
|
||||
Collection string
|
||||
// Roots are the directories and files to sync.
|
||||
Roots []string
|
||||
// Include, when set, are the only files read (globs; one without a "/"
|
||||
// matches a file's name, one with it the path under its root, "**"
|
||||
// crossing directories). Exclude are never read.
|
||||
Include, Exclude []string
|
||||
// Hidden includes files and directories whose names start with a dot.
|
||||
Hidden bool
|
||||
// DryRun reads and cuts everything that would be sent, and sends
|
||||
// nothing.
|
||||
DryRun bool
|
||||
// Delete removes from the collection documents this machine sent whose
|
||||
// files under these roots are gone.
|
||||
Delete bool
|
||||
// Retry reads again files skipped before, as when a helper has been
|
||||
// installed since.
|
||||
Retry bool
|
||||
// Workers is how many files are read and sent at once; 0 is 4.
|
||||
Workers int
|
||||
Extract extract.Options
|
||||
Chunk chunk.Options
|
||||
}
|
||||
|
||||
// Event is something a sync did to one document, for showing progress.
|
||||
type Event struct {
|
||||
// Kind is sent, unchanged, skipped, failed, deleted, would-send,
|
||||
// would-delete, or warning.
|
||||
Kind string
|
||||
Name string
|
||||
Detail string
|
||||
Chunks int
|
||||
Bytes int64
|
||||
}
|
||||
|
||||
// Report is what a sync did in all.
|
||||
type Report struct {
|
||||
Collection Collection
|
||||
Files int
|
||||
Sent int
|
||||
Unchanged int
|
||||
Skipped int
|
||||
Failed int
|
||||
Deleted int
|
||||
// Chunks and Bytes are what was sent, or with DryRun what would be.
|
||||
Chunks int
|
||||
Bytes int64
|
||||
}
|
||||
|
||||
// maxDocJSON is the most one document's request is, under the gateway's
|
||||
// limit with room to spare.
|
||||
const maxDocJSON = 60 << 20
|
||||
|
||||
// defaultExcludes are directories nobody means to put in a knowledge base.
|
||||
var defaultExcludes = map[string]bool{
|
||||
".git": true, ".svn": true, ".hg": true, "node_modules": true, ".venv": true, "venv": true,
|
||||
"__pycache__": true, ".idea": true, ".vscode": true, "$RECYCLE.BIN": true,
|
||||
"System Volume Information": true, ".Trash": true, ".Trashes": true, ".cache": true,
|
||||
}
|
||||
|
||||
type file struct {
|
||||
abs string
|
||||
name string
|
||||
size int64
|
||||
mtime time.Time
|
||||
}
|
||||
|
||||
// Sync brings a collection up to date with the files under the roots.
|
||||
func Sync(ctx context.Context, c *Client, l *Ledger, o Options, progress func(Event)) (Report, error) {
|
||||
var mu sync.Mutex
|
||||
var rep Report
|
||||
emit := func(e Event) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
switch e.Kind {
|
||||
case "sent", "would-send":
|
||||
rep.Sent++
|
||||
rep.Chunks += e.Chunks
|
||||
rep.Bytes += e.Bytes
|
||||
case "unchanged":
|
||||
rep.Unchanged++
|
||||
case "skipped":
|
||||
rep.Skipped++
|
||||
case "failed":
|
||||
rep.Failed++
|
||||
case "deleted", "would-delete":
|
||||
rep.Deleted++
|
||||
}
|
||||
if progress != nil {
|
||||
progress(e)
|
||||
}
|
||||
}
|
||||
|
||||
m, err := c.Manifest(ctx, o.Collection)
|
||||
if err != nil {
|
||||
return rep, fmt.Errorf("reading collection %s: %w", o.Collection, err)
|
||||
}
|
||||
rep.Collection = m.Collection
|
||||
if !m.Collection.Writable && !o.DryRun {
|
||||
return rep, fmt.Errorf("this key may read collection %s but not add to it", m.Collection.Name)
|
||||
}
|
||||
remote, chunks := map[string]string{}, map[string]int{}
|
||||
for _, d := range m.Documents {
|
||||
remote[d.Name], chunks[d.Name] = d.Digest, d.Chunks
|
||||
}
|
||||
gw := strings.TrimRight(c.Base, "/")
|
||||
coll := m.Collection.ID
|
||||
|
||||
files, roots, err := walk(o, func(name, why string) { emit(Event{Kind: "warning", Name: name, Detail: why}) })
|
||||
if err != nil {
|
||||
return rep, err
|
||||
}
|
||||
rep.Files = len(files)
|
||||
if len(files) == 0 {
|
||||
return rep, ErrNothing
|
||||
}
|
||||
|
||||
runID, err := l.StartRun(ctx, gw, coll)
|
||||
if err != nil {
|
||||
return rep, err
|
||||
}
|
||||
previous, err := l.Entries(ctx, gw, coll)
|
||||
if err != nil {
|
||||
return rep, err
|
||||
}
|
||||
byFile := map[string][]Entry{}
|
||||
for _, e := range previous {
|
||||
byFile[e.File] = append(byFile[e.File], e)
|
||||
}
|
||||
|
||||
s := &syncer{c: c, l: l, o: o, gw: gw, coll: m.Collection, remote: remote, chunks: chunks, byFile: byFile, emit: emit}
|
||||
workers := o.Workers
|
||||
if workers <= 0 {
|
||||
workers = 4
|
||||
}
|
||||
jobs := make(chan file)
|
||||
var wg sync.WaitGroup
|
||||
for range workers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for f := range jobs {
|
||||
s.one(ctx, f)
|
||||
}
|
||||
}()
|
||||
}
|
||||
for _, f := range files {
|
||||
select {
|
||||
case jobs <- f:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
close(jobs)
|
||||
wg.Wait()
|
||||
|
||||
if o.Delete && ctx.Err() == nil {
|
||||
seen := map[string]bool{}
|
||||
for _, f := range files {
|
||||
seen[f.abs] = true
|
||||
}
|
||||
for _, e := range previous {
|
||||
if e.State != Sent || seen[e.File] || !under(e.File, roots) {
|
||||
continue
|
||||
}
|
||||
if o.DryRun {
|
||||
emit(Event{Kind: "would-delete", Name: e.Name})
|
||||
continue
|
||||
}
|
||||
if err := c.Delete(ctx, coll, e.Name); err != nil {
|
||||
emit(Event{Kind: "failed", Name: e.Name, Detail: "deleting: " + err.Error()})
|
||||
continue
|
||||
}
|
||||
_ = l.Remove(ctx, gw, coll, e.Name)
|
||||
emit(Event{Kind: "deleted", Name: e.Name})
|
||||
}
|
||||
}
|
||||
_ = l.EndRun(context.WithoutCancel(ctx), runID, Run{Sent: rep.Sent, Unchanged: rep.Unchanged, Skipped: rep.Skipped,
|
||||
Failed: rep.Failed, Delete: rep.Deleted})
|
||||
return rep, ctx.Err()
|
||||
}
|
||||
|
||||
type syncer struct {
|
||||
c *Client
|
||||
l *Ledger
|
||||
o Options
|
||||
gw string
|
||||
coll Collection
|
||||
remote map[string]string
|
||||
chunks map[string]int
|
||||
byFile map[string][]Entry
|
||||
emit func(Event)
|
||||
}
|
||||
|
||||
func (s *syncer) entry(f file, name, digest string, st State, reason string, chunks int) Entry {
|
||||
return Entry{Gateway: s.gw, Collection: s.coll.ID, Name: name, File: f.abs, Size: f.size, MTime: f.mtime,
|
||||
Digest: digest, State: st, Reason: reason, Chunks: chunks}
|
||||
}
|
||||
|
||||
// one syncs one file: every document in it.
|
||||
func (s *syncer) one(ctx context.Context, f file) {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
digest, ok := s.l.Known(ctx, f.abs, f.size, f.mtime)
|
||||
if !ok {
|
||||
d, err := hashFile(f.abs)
|
||||
if err != nil {
|
||||
s.emit(Event{Kind: "skipped", Name: f.name, Detail: err.Error()})
|
||||
return
|
||||
}
|
||||
digest = d
|
||||
}
|
||||
prev := s.byFile[f.abs]
|
||||
|
||||
// The collection has this file as it is: nothing to read.
|
||||
if s.remote[f.name] == digest {
|
||||
s.record(ctx, s.entry(f, f.name, digest, Sent, "", s.chunks[f.name]))
|
||||
s.emit(Event{Kind: "unchanged", Name: f.name})
|
||||
return
|
||||
}
|
||||
// An archive or a mail is several documents; all of them as they are
|
||||
// is nothing to read either.
|
||||
if len(prev) > 0 && allSame(prev, digest, s.remote) && !(s.o.Retry && anySkipped(prev)) {
|
||||
for _, e := range prev {
|
||||
if e.State == Skipped {
|
||||
s.emit(Event{Kind: "skipped", Name: e.Name, Detail: e.Reason + " (as before; --retry reads it again)"})
|
||||
} else {
|
||||
s.emit(Event{Kind: "unchanged", Name: e.Name})
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
parts, err := extract.File(ctx, f.abs, s.o.Extract)
|
||||
if err != nil {
|
||||
s.skip(ctx, f, f.name, digest, err)
|
||||
return
|
||||
}
|
||||
current := map[string]bool{}
|
||||
for _, p := range parts {
|
||||
name := f.name
|
||||
if p.Name != "" {
|
||||
name = f.name + "!/" + p.Name
|
||||
}
|
||||
current[name] = true
|
||||
if p.Err != nil {
|
||||
s.skip(ctx, f, name, digest, p.Err)
|
||||
continue
|
||||
}
|
||||
s.send(ctx, f, name, digest, p.Text)
|
||||
}
|
||||
// Members an archive no longer has go with it.
|
||||
for _, e := range prev {
|
||||
if current[e.Name] || e.Name == f.name {
|
||||
continue
|
||||
}
|
||||
if s.o.DryRun {
|
||||
s.emit(Event{Kind: "would-delete", Name: e.Name, Detail: "no longer in " + f.name})
|
||||
continue
|
||||
}
|
||||
if e.State == Sent {
|
||||
if err := s.c.Delete(ctx, s.coll.ID, e.Name); err != nil {
|
||||
s.emit(Event{Kind: "failed", Name: e.Name, Detail: "deleting: " + err.Error()})
|
||||
continue
|
||||
}
|
||||
s.emit(Event{Kind: "deleted", Name: e.Name, Detail: "no longer in " + f.name})
|
||||
}
|
||||
_ = s.l.Remove(ctx, s.gw, s.coll.ID, e.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// allSame reports whether every document from a file was left with this
|
||||
// digest, and the collection still holds those that were sent.
|
||||
func allSame(prev []Entry, digest string, remote map[string]string) bool {
|
||||
for _, e := range prev {
|
||||
if e.Digest != digest {
|
||||
return false
|
||||
}
|
||||
switch e.State {
|
||||
case Sent:
|
||||
if remote[e.Name] != digest {
|
||||
return false
|
||||
}
|
||||
case Failed:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func anySkipped(es []Entry) bool {
|
||||
for _, e := range es {
|
||||
if e.State == Skipped {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *syncer) skip(ctx context.Context, f file, name, digest string, err error) {
|
||||
s.record(ctx, s.entry(f, name, digest, Skipped, err.Error(), 0))
|
||||
s.emit(Event{Kind: "skipped", Name: name, Detail: err.Error()})
|
||||
}
|
||||
|
||||
// record writes the ledger; a dry run leaves it as it was.
|
||||
func (s *syncer) record(ctx context.Context, e Entry) {
|
||||
if s.o.DryRun {
|
||||
return
|
||||
}
|
||||
_ = s.l.Put(ctx, e)
|
||||
}
|
||||
|
||||
// send cuts a document's text, embeds it and sends it.
|
||||
func (s *syncer) send(ctx context.Context, f file, name, digest, text string) {
|
||||
if s.remote[name] == digest {
|
||||
s.record(ctx, s.entry(f, name, digest, Sent, "", s.chunks[name]))
|
||||
s.emit(Event{Kind: "unchanged", Name: name})
|
||||
return
|
||||
}
|
||||
cs := chunk.Split(text, s.o.Chunk)
|
||||
if len(cs) == 0 {
|
||||
s.skip(ctx, f, name, digest, extract.ErrNoText)
|
||||
return
|
||||
}
|
||||
texts := make([]string, len(cs))
|
||||
var size int64
|
||||
for i, c := range cs {
|
||||
texts[i] = c.Text()
|
||||
size += int64(len(texts[i]))
|
||||
}
|
||||
// About ten bytes of JSON a dimension.
|
||||
size += int64(len(cs) * s.coll.Dims * 10)
|
||||
if size > maxDocJSON {
|
||||
s.fail(ctx, f, name, digest, fmt.Errorf("%d MiB to send is over the gateway's limit for one document; split the file", size>>20))
|
||||
return
|
||||
}
|
||||
if s.o.DryRun {
|
||||
s.emit(Event{Kind: "would-send", Name: name, Chunks: len(cs), Bytes: size})
|
||||
return
|
||||
}
|
||||
vecs, err := s.c.Embed(ctx, s.coll.EmbedModel, texts)
|
||||
if err != nil {
|
||||
s.fail(ctx, f, name, digest, fmt.Errorf("embedding with %s: %w", s.coll.EmbedModel, err))
|
||||
return
|
||||
}
|
||||
out := DocOut{Name: name, Digest: digest, Bytes: f.size, Chunks: make([]ChunkOut, len(cs))}
|
||||
for i := range cs {
|
||||
if len(vecs[i]) != s.coll.Dims {
|
||||
s.fail(ctx, f, name, digest, fmt.Errorf("%s gave %d dimensions and the collection takes %d",
|
||||
s.coll.EmbedModel, len(vecs[i]), s.coll.Dims))
|
||||
return
|
||||
}
|
||||
out.Chunks[i] = ChunkOut{Text: texts[i], Embedding: vecs[i]}
|
||||
}
|
||||
if _, err := s.c.Put(ctx, s.coll.ID, out); err != nil {
|
||||
s.fail(ctx, f, name, digest, err)
|
||||
return
|
||||
}
|
||||
s.record(ctx, s.entry(f, name, digest, Sent, "", len(cs)))
|
||||
s.emit(Event{Kind: "sent", Name: name, Chunks: len(cs), Bytes: size})
|
||||
}
|
||||
|
||||
func (s *syncer) fail(ctx context.Context, f file, name, digest string, err error) {
|
||||
s.record(ctx, s.entry(f, name, digest, Failed, err.Error(), 0))
|
||||
s.emit(Event{Kind: "failed", Name: name, Detail: err.Error()})
|
||||
}
|
||||
|
||||
func hashFile(p string) (string, error) {
|
||||
fh, err := os.Open(p)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer fh.Close()
|
||||
h := sha256.New()
|
||||
if _, err := io.Copy(h, fh); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "sha256:" + hex.EncodeToString(h.Sum(nil)), nil
|
||||
}
|
||||
|
||||
// under reports whether a path is one of the roots or inside one.
|
||||
func under(p string, roots []string) bool {
|
||||
for _, r := range roots {
|
||||
if p == r || strings.HasPrefix(p, strings.TrimRight(r, string(filepath.Separator))+string(filepath.Separator)) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// walk lists the files to sync, named as documents: the path under its root
|
||||
// with forward slashes, in Unicode NFC, prefixed by the root's own name when
|
||||
// there are several roots so that two roots' "README.md" stay apart.
|
||||
func walk(o Options, warn func(name, why string)) ([]file, []string, error) {
|
||||
inc, err := globs(o.Include)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
exc, err := globs(o.Exclude)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var out []file
|
||||
var roots []string
|
||||
seen := map[string]string{}
|
||||
for _, root := range o.Roots {
|
||||
abs, err := filepath.Abs(root)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
roots = append(roots, abs)
|
||||
prefix := ""
|
||||
if len(o.Roots) > 1 && st.IsDir() {
|
||||
prefix = norm.NFC.String(filepath.Base(abs)) + "/"
|
||||
}
|
||||
add := func(p string, info fs.FileInfo) {
|
||||
rel := filepath.Base(p)
|
||||
if st.IsDir() {
|
||||
rel, _ = filepath.Rel(abs, p)
|
||||
}
|
||||
name := prefix + norm.NFC.String(filepath.ToSlash(rel))
|
||||
if !wanted(name, inc, exc, o.Hidden) {
|
||||
return
|
||||
}
|
||||
if prev, ok := seen[strings.ToLower(name)]; ok && prev != name {
|
||||
warn(name, "differs from "+prev+" only in case; on Windows and macOS they are one file")
|
||||
}
|
||||
seen[strings.ToLower(name)] = name
|
||||
out = append(out, file{abs: p, name: name, size: info.Size(), mtime: info.ModTime()})
|
||||
}
|
||||
if !st.IsDir() {
|
||||
add(abs, st)
|
||||
continue
|
||||
}
|
||||
err = filepath.WalkDir(abs, func(p string, d fs.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
// A directory this user may not read is skipped, not fatal.
|
||||
warn(p, err.Error())
|
||||
if d != nil && d.IsDir() {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
base := d.Name()
|
||||
if d.IsDir() {
|
||||
if p != abs && (defaultExcludes[base] || (!o.Hidden && strings.HasPrefix(base, "."))) {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if !d.Type().IsRegular() {
|
||||
return nil
|
||||
}
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
add(p, info)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].name < out[j].name })
|
||||
return out, roots, nil
|
||||
}
|
||||
|
||||
// wanted applies the filters to a document name.
|
||||
func wanted(name string, inc, exc []*regexp.Regexp, hidden bool) bool {
|
||||
base := path.Base(name)
|
||||
if !hidden && strings.HasPrefix(base, ".") {
|
||||
return false
|
||||
}
|
||||
// Office's lock files, LibreOffice's, editors' backups.
|
||||
if strings.HasPrefix(base, "~$") || strings.HasPrefix(base, ".~lock.") || strings.HasSuffix(base, "~") ||
|
||||
base == "Thumbs.db" || base == "desktop.ini" || base == ".DS_Store" {
|
||||
return false
|
||||
}
|
||||
for _, re := range exc {
|
||||
if re.MatchString(name) || re.MatchString(base) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
if len(inc) > 0 {
|
||||
for _, re := range inc {
|
||||
if re.MatchString(name) || re.MatchString(base) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
return extract.Known(base)
|
||||
}
|
||||
|
||||
// globs compiles glob patterns: "*" within a name, "**" across directories,
|
||||
// "?" one character; case-insensitive, as names are on Windows and macOS.
|
||||
func globs(patterns []string) ([]*regexp.Regexp, error) {
|
||||
var out []*regexp.Regexp
|
||||
for _, p := range patterns {
|
||||
for _, one := range strings.Split(p, ",") {
|
||||
one = strings.TrimSpace(filepath.ToSlash(one))
|
||||
if one == "" {
|
||||
continue
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("(?i)^")
|
||||
for i := 0; i < len(one); i++ {
|
||||
switch c := one[i]; c {
|
||||
case '*':
|
||||
if i+1 < len(one) && one[i+1] == '*' {
|
||||
b.WriteString(".*")
|
||||
i++
|
||||
if i+1 < len(one) && one[i+1] == '/' {
|
||||
b.WriteString("/?")
|
||||
i++
|
||||
}
|
||||
} else {
|
||||
b.WriteString("[^/]*")
|
||||
}
|
||||
case '?':
|
||||
b.WriteString("[^/]")
|
||||
default:
|
||||
b.WriteString(regexp.QuoteMeta(string(c)))
|
||||
}
|
||||
}
|
||||
b.WriteString("$")
|
||||
re, err := regexp.Compile(b.String())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pattern %q: %w", one, err)
|
||||
}
|
||||
out = append(out, re)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ErrNothing is a sync that found no file to read.
|
||||
var ErrNothing = errors.New("no files to read under these paths")
|
||||
@@ -0,0 +1,272 @@
|
||||
package kbsync
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// fakeGateway holds one collection in memory and embeds by counting letters.
|
||||
type fakeGateway struct {
|
||||
mu sync.Mutex
|
||||
writable bool
|
||||
docs map[string]RemoteDoc
|
||||
puts int
|
||||
embedded int
|
||||
busyFirst bool
|
||||
}
|
||||
|
||||
func (g *fakeGateway) serve(t *testing.T) *Client {
|
||||
t.Helper()
|
||||
mux := http.NewServeMux()
|
||||
coll := Collection{ID: "kb_x", Name: "ik", EmbedModel: "emb", Dims: 3}
|
||||
mux.HandleFunc("GET /v1/kb/collections/{id}/manifest", func(w http.ResponseWriter, r *http.Request) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if r.PathValue("id") != "ik" && r.PathValue("id") != "kb_x" {
|
||||
http.Error(w, `{"error":{"message":"no collection"}}`, http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
m := Manifest{Collection: coll}
|
||||
m.Collection.Writable = g.writable
|
||||
for _, d := range g.docs {
|
||||
m.Documents = append(m.Documents, d)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(m)
|
||||
})
|
||||
mux.HandleFunc("PUT /v1/kb/collections/{id}/documents", func(w http.ResponseWriter, r *http.Request) {
|
||||
var d DocOut
|
||||
_ = json.NewDecoder(r.Body).Decode(&d)
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
g.puts++
|
||||
g.docs[d.Name] = RemoteDoc{Name: d.Name, Digest: d.Digest, Chunks: len(d.Chunks)}
|
||||
_, _ = w.Write([]byte(`{"changed":true}`))
|
||||
})
|
||||
mux.HandleFunc("DELETE /v1/kb/collections/{id}/documents", func(w http.ResponseWriter, r *http.Request) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
delete(g.docs, r.URL.Query().Get("name"))
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
})
|
||||
mux.HandleFunc("POST /v1/embeddings", func(w http.ResponseWriter, r *http.Request) {
|
||||
g.mu.Lock()
|
||||
busy := g.busyFirst
|
||||
g.busyFirst = false
|
||||
g.mu.Unlock()
|
||||
if busy {
|
||||
w.Header().Set("Retry-After", "1")
|
||||
http.Error(w, `{"error":{"message":"the model is loading"}}`, http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
var in struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&in)
|
||||
type row struct {
|
||||
Index int `json:"index"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
}
|
||||
var out struct {
|
||||
Data []row `json:"data"`
|
||||
}
|
||||
for i, s := range in.Input {
|
||||
out.Data = append(out.Data, row{i, []float32{float32(len(s)), 1, 0}})
|
||||
}
|
||||
g.mu.Lock()
|
||||
g.embedded += len(in.Input)
|
||||
g.mu.Unlock()
|
||||
_ = json.NewEncoder(w).Encode(out)
|
||||
})
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
h, _ := NewHTTP("")
|
||||
return &Client{Base: srv.URL, Key: "k", HTTP: h}
|
||||
}
|
||||
|
||||
func setup(t *testing.T) (*fakeGateway, *Client, *Ledger, string) {
|
||||
t.Helper()
|
||||
g := &fakeGateway{writable: true, docs: map[string]RemoteDoc{}}
|
||||
c := g.serve(t)
|
||||
l, err := OpenLedger(context.Background(), filepath.Join(t.TempDir(), "ledger.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = l.Close() })
|
||||
dir := t.TempDir()
|
||||
write(t, dir, "politika/izin.md", "# İzin\n\nYirmi gün.\n")
|
||||
write(t, dir, "politika/masraf.txt", "Masraflar finansa.\n")
|
||||
write(t, dir, "eski.doc", "\xd0\xcf\x11\xe0\x00binary")
|
||||
write(t, dir, "resim.bin", "\x00\x01")
|
||||
write(t, dir, ".git/config", "[core]")
|
||||
write(t, dir, "politika/~$izin.docx", "lock")
|
||||
return g, c, l, dir
|
||||
}
|
||||
|
||||
func write(t *testing.T, dir, name, s string) {
|
||||
t.Helper()
|
||||
p := filepath.Join(dir, filepath.FromSlash(name))
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(p, []byte(s), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func kinds(events []Event) map[string][]string {
|
||||
out := map[string][]string{}
|
||||
for _, e := range events {
|
||||
out[e.Kind] = append(out[e.Kind], e.Name)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func run(t *testing.T, c *Client, l *Ledger, o Options) (Report, map[string][]string) {
|
||||
t.Helper()
|
||||
var mu sync.Mutex
|
||||
var ev []Event
|
||||
rep, err := Sync(context.Background(), c, l, o, func(e Event) {
|
||||
mu.Lock()
|
||||
ev = append(ev, e)
|
||||
mu.Unlock()
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return rep, kinds(ev)
|
||||
}
|
||||
|
||||
func TestASecondRunSendsNothingAndReadsNothing(t *testing.T) {
|
||||
g, c, l, dir := setup(t)
|
||||
rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}})
|
||||
if rep.Sent != 2 || len(g.docs) != 2 {
|
||||
t.Fatalf("first run %+v %v", rep, k)
|
||||
}
|
||||
if _, ok := g.docs["politika/izin.md"]; !ok {
|
||||
t.Errorf("documents %v", g.docs)
|
||||
}
|
||||
// The legacy Word file is skipped with the reason, and no hidden,
|
||||
// lock or unknown file is touched.
|
||||
if len(k["skipped"]) != 1 || k["skipped"][0] != "eski.doc" {
|
||||
t.Errorf("skipped %v", k["skipped"])
|
||||
}
|
||||
|
||||
puts, embedded := g.puts, g.embedded
|
||||
rep, k = run(t, c, l, Options{Collection: "ik", Roots: []string{dir}})
|
||||
if rep.Sent != 0 || rep.Unchanged != 2 || g.puts != puts || g.embedded != embedded {
|
||||
t.Errorf("second run %+v %v", rep, k)
|
||||
}
|
||||
if rep.Skipped != 1 {
|
||||
t.Errorf("the skipped file was not reported as before: %+v", rep)
|
||||
}
|
||||
es, _ := l.Entries(context.Background(), c.Base, "kb_x")
|
||||
if len(es) != 3 {
|
||||
t.Errorf("ledger %+v", es)
|
||||
}
|
||||
runs, _ := l.Runs(context.Background(), 10)
|
||||
if len(runs) != 2 || runs[0].Unchanged != 2 {
|
||||
t.Errorf("runs %+v", runs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAChangedFileIsSentAgain(t *testing.T) {
|
||||
g, c, l, dir := setup(t)
|
||||
run(t, c, l, Options{Collection: "ik", Roots: []string{dir}})
|
||||
old := g.docs["politika/izin.md"].Digest
|
||||
write(t, dir, "politika/izin.md", "# İzin\n\nOtuz gün.\n")
|
||||
// A changed time is what makes it look again; some filesystems keep
|
||||
// seconds only.
|
||||
later := time.Now().Add(2 * time.Second)
|
||||
_ = os.Chtimes(filepath.Join(dir, "politika", "izin.md"), later, later)
|
||||
rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}})
|
||||
if rep.Sent != 1 || k["sent"][0] != "politika/izin.md" || g.docs["politika/izin.md"].Digest == old {
|
||||
t.Errorf("after a change %+v %v", rep, k)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteRemovesOnlyWhatThisMachineSentUnderTheseRoots(t *testing.T) {
|
||||
g, c, l, dir := setup(t)
|
||||
run(t, c, l, Options{Collection: "ik", Roots: []string{dir}})
|
||||
// Somebody else's document, sent from another machine.
|
||||
g.docs["baska/makine.md"] = RemoteDoc{Name: "baska/makine.md", Digest: "x"}
|
||||
// A file of this machine from another folder.
|
||||
other := t.TempDir()
|
||||
write(t, other, "not.md", "başka klasör")
|
||||
run(t, c, l, Options{Collection: "ik", Roots: []string{other}})
|
||||
|
||||
if err := os.Remove(filepath.Join(dir, "politika", "masraf.txt")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, Delete: true})
|
||||
if rep.Deleted != 1 || k["deleted"][0] != "politika/masraf.txt" {
|
||||
t.Errorf("deleted %+v %v", rep, k)
|
||||
}
|
||||
for _, keep := range []string{"baska/makine.md", "not.md", "politika/izin.md"} {
|
||||
if _, ok := g.docs[keep]; !ok {
|
||||
t.Errorf("%s was deleted", keep)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestADryRunSendsAndRecordsNothing(t *testing.T) {
|
||||
g, c, l, dir := setup(t)
|
||||
rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, DryRun: true})
|
||||
if rep.Sent != 2 || len(k["would-send"]) != 2 || rep.Bytes == 0 || rep.Chunks != 2 {
|
||||
t.Errorf("dry run %+v %v", rep, k)
|
||||
}
|
||||
if g.puts != 0 || g.embedded != 0 {
|
||||
t.Errorf("a dry run sent %d documents, embedded %d chunks", g.puts, g.embedded)
|
||||
}
|
||||
if es, _ := l.Entries(context.Background(), c.Base, "kb_x"); len(es) != 0 {
|
||||
t.Errorf("a dry run wrote the ledger: %+v", es)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAReadOnlyCollectionIsRefusedAndALoadingModelWaitedFor(t *testing.T) {
|
||||
g, c, l, dir := setup(t)
|
||||
g.writable = false
|
||||
if _, err := Sync(context.Background(), c, l, Options{Collection: "ik", Roots: []string{dir}}, nil); err == nil ||
|
||||
!strings.Contains(err.Error(), "not add") {
|
||||
t.Errorf("read-only gave %v", err)
|
||||
}
|
||||
g.writable, g.busyFirst = true, true
|
||||
var waited []string
|
||||
c.Waiting = func(why string, _ time.Duration) { waited = append(waited, why) }
|
||||
rep, _ := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, Workers: 1})
|
||||
if rep.Sent != 2 || len(waited) != 1 || !strings.Contains(waited[0], "loading") {
|
||||
t.Errorf("after a 503 %+v, waited %v", rep, waited)
|
||||
}
|
||||
if _, err := Sync(context.Background(), c, l, Options{Collection: "yok", Roots: []string{dir}}, nil); err == nil {
|
||||
t.Error("a collection that is not there was synced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFiltersAndNames(t *testing.T) {
|
||||
_, c, l, dir := setup(t)
|
||||
write(t, dir, "rapor/2026/ocak.md", "ocak")
|
||||
write(t, dir, "rapor/taslak.md", "taslak")
|
||||
rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, DryRun: true,
|
||||
Include: []string{"rapor/**/*.md"}, Exclude: []string{"taslak*"}})
|
||||
if rep.Sent != 1 || k["would-send"][0] != "rapor/2026/ocak.md" {
|
||||
t.Errorf("filtered %+v %v", rep, k)
|
||||
}
|
||||
// Two roots keep their names apart.
|
||||
other := t.TempDir()
|
||||
write(t, other, "politika/izin.md", "aynı ad")
|
||||
_, k = run(t, c, l, Options{Collection: "ik", Roots: []string{dir, other}, DryRun: true, Include: []string{"izin.md"}})
|
||||
if len(k["would-send"]) != 2 || k["would-send"][0] == k["would-send"][1] {
|
||||
t.Errorf("two roots: %v", k["would-send"])
|
||||
}
|
||||
if _, err := Sync(context.Background(), c, l, Options{Collection: "ik", Roots: []string{t.TempDir()}}, nil); !errors.Is(err, ErrNothing) {
|
||||
t.Errorf("an empty folder gave %v", err)
|
||||
}
|
||||
}
|
||||
@@ -159,6 +159,11 @@ func TestARunningClusterIsLeftAsItIs(t *testing.T) {
|
||||
if !a.ranAny("CREATE ROLE") {
|
||||
t.Error("the role was not made sure of on the leader")
|
||||
}
|
||||
// Enabled where the package is there, which on a running cluster is
|
||||
// all a rerun does about it: it installs nothing.
|
||||
if !a.ranAny("CREATE EXTENSION IF NOT EXISTS vector") {
|
||||
t.Error("pgvector was not enabled in llmgw's database on the leader")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTwoDatabaseMachinesWithoutAThirdVoteAreRefused(t *testing.T) {
|
||||
|
||||
+17
-1
@@ -502,6 +502,14 @@ if [ -n "$need" ]; then
|
||||
sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y $need >/dev/null
|
||||
fi
|
||||
v=$(ls /usr/lib/postgresql | sort -n | tail -1)
|
||||
# pgvector, for knowledge bases: on every member, since any of them may lead.
|
||||
# Not a reason to stop; the store says knowledge bases are off without it.
|
||||
if [ ! -f /usr/share/postgresql/$v/extension/vector.control ]; then
|
||||
sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y postgresql-$v-pgvector >/dev/null 2>&1 ||
|
||||
{ sudo DEBIAN_FRONTEND=noninteractive apt-get -qq update >/dev/null 2>&1 &&
|
||||
sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y postgresql-$v-pgvector >/dev/null 2>&1; } ||
|
||||
echo "no-pgvector"
|
||||
fi
|
||||
if [ -d /var/lib/postgresql/$v/main ]; then
|
||||
if sudo systemctl is-active --quiet postgresql@$v-main; then
|
||||
n=$(sudo -u postgres psql -Atc "select count(*) from pg_database where datname not in ('postgres','template0','template1')")
|
||||
@@ -553,13 +561,21 @@ sudo -u postgres psql -q -v ON_ERROR_STOP=1 -v pw="$pw" <<'SQL'
|
||||
SELECT format('CREATE ROLE %I LOGIN', 'llmgw') WHERE NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'llmgw') \gexec
|
||||
ALTER ROLE llmgw PASSWORD :'pw';
|
||||
SELECT format('CREATE DATABASE %I OWNER %I', 'llmgw', 'llmgw') WHERE NOT EXISTS (SELECT 1 FROM pg_database WHERE datname = 'llmgw') \gexec
|
||||
SQL`, nil)
|
||||
SQL
|
||||
sudo -u postgres psql -q -d llmgw -v ON_ERROR_STOP=1 <<'SQL' || echo "pgvector-failed"
|
||||
SELECT 'CREATE EXTENSION IF NOT EXISTS vector' WHERE EXISTS (SELECT 1 FROM pg_available_extensions WHERE name = 'vector') \gexec
|
||||
SQL
|
||||
sudo -u postgres psql -d llmgw -Atc "select 'vector=' || count(*) from pg_extension where extname = 'vector'"`, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("on %s: %v: %s", leader.Name, err, strings.TrimSpace(string(out)))
|
||||
}
|
||||
// Said as a change only when something was made: a rerun that finds
|
||||
// both is not a reason to restart the gateway.
|
||||
s.emit("database", !strings.Contains(string(out), "had=2"), "role and database %s on %s", Database, leader.Name)
|
||||
if !strings.Contains(string(out), "vector=1") {
|
||||
s.emit("pgvector", false, "pgvector is not enabled in %s, so knowledge bases stay off "+
|
||||
"(install postgresql-<version>-pgvector on every member and run setup again)", Database)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -251,6 +251,11 @@ func reusable(req *parsedRequest, anyAnswer bool) bool {
|
||||
if req.noStore || req.ThreadID != "" || req.ContentType != "" || !cacheEndpoints[req.Path] {
|
||||
return false
|
||||
}
|
||||
// An answer from knowledge bases is the caller's own: two keys of one
|
||||
// project can read different collections.
|
||||
if req.knowledge || req.Header.Get(knowledgeHeader) != "" {
|
||||
return false
|
||||
}
|
||||
if v := strings.ToLower(strings.TrimSpace(req.Header.Get(cacheHeader))); v == "no" || v == "off" || v == "bypass" {
|
||||
return false
|
||||
}
|
||||
|
||||
+41
-1
@@ -37,7 +37,7 @@ func (r *Router) registerChat(mux *http.ServeMux, pages *web.AccountPages) {
|
||||
if current == "" && !r.personalOff {
|
||||
current = personalContext
|
||||
}
|
||||
web.Chat(user, r.modelsFor(req, user), contexts, current, version.Version)(w, req)
|
||||
web.Chat(user, r.modelsFor(req, user), contexts, current, version.Version, r.chatKnowledge(req, user, current))(w, req)
|
||||
}))
|
||||
mux.HandleFunc("POST /chat/send", pages.Guard(r.chatSend))
|
||||
// Where a model stands while it loads, and a way to ask for it without
|
||||
@@ -75,6 +75,9 @@ func (r *Router) chatSend(w http.ResponseWriter, req *http.Request, user store.U
|
||||
// Temperature is the page's one sampling setting, held to the range
|
||||
// every engine takes.
|
||||
Temperature *float64 `json:"temperature"`
|
||||
// Knowledge is a collection to answer from, "none", or "" for the
|
||||
// default; whether it may be read is decided as for a key.
|
||||
Knowledge string `json:"knowledge"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &in); err != nil || in.Model == "" || len(in.Messages) == 0 {
|
||||
http.Error(w, "a model and at least one message are required", http.StatusBadRequest)
|
||||
@@ -112,6 +115,10 @@ func (r *Router) chatSend(w http.ResponseWriter, req *http.Request, user store.U
|
||||
next.Body = io.NopCloser(bytes.NewReader(upstream))
|
||||
next.ContentLength = int64(len(upstream))
|
||||
next.Header.Set("Content-Type", "application/json")
|
||||
next.Header.Del(knowledgeHeader)
|
||||
if k := strings.TrimSpace(in.Knowledge); k != "" {
|
||||
next.Header.Set(knowledgeHeader, k)
|
||||
}
|
||||
|
||||
// The same handler the API serves, so there is exactly one request path and
|
||||
// no second place for its rules to drift.
|
||||
@@ -309,3 +316,36 @@ func (r *Router) modelsFor(req *http.Request, user store.User) []web.ChatModel {
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// chatKnowledge is the collections the page offers in a context, and the
|
||||
// default that answers when none is chosen.
|
||||
func (r *Router) chatKnowledge(req *http.Request, user store.User, context string) web.ChatKnowledge {
|
||||
var out web.ChatKnowledge
|
||||
if r.store == nil || r.store.KnowledgeAvailable() != nil {
|
||||
return out
|
||||
}
|
||||
p, err := r.principalFor(req, user, context)
|
||||
if err != nil {
|
||||
return out
|
||||
}
|
||||
e := store.Effective{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID, OrgID: p.OrgID}
|
||||
cs, err := r.store.CollectionsFor(req.Context(), e, false)
|
||||
if err != nil {
|
||||
return out
|
||||
}
|
||||
names := map[string]string{}
|
||||
for _, c := range cs {
|
||||
out.Collections = append(out.Collections, web.ChatCollection{ID: c.ID, Label: c.Name})
|
||||
names[c.ID] = c.Name
|
||||
}
|
||||
if def, err := r.store.KBDefaultFor(req.Context(), e); err == nil {
|
||||
var shown []string
|
||||
for _, id := range def {
|
||||
if n, ok := names[id]; ok {
|
||||
shown = append(shown, n)
|
||||
}
|
||||
}
|
||||
out.Default = strings.Join(shown, ", ")
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package router
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -207,3 +208,31 @@ func TestAChatTurnCarriesItsTemperatureAndAsksForUsage(t *testing.T) {
|
||||
t.Errorf("sent %s", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheChatPageAsksForTheDocumentsChosenAndOnlyThose(t *testing.T) {
|
||||
// The picker's choice becomes the header; a header the page sent itself
|
||||
// is not trusted. On SQLite, asking is enough to tell: it is refused as
|
||||
// knowledge bases being off, before the model is asked.
|
||||
e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
|
||||
rt, user := chatRouter(t, e.srv.URL)
|
||||
waitForEndpoint(t, rt)
|
||||
send := func(knowledge, spoofed string) *httptest.ResponseRecorder {
|
||||
body := `{"model":"qwen2.5-coder-7b-instruct","knowledge":"` + knowledge + `","messages":[{"role":"user","content":"hi"}]}`
|
||||
req := httptest.NewRequest("POST", "/chat/send", strings.NewReader(body))
|
||||
if spoofed != "" {
|
||||
req.Header.Set(knowledgeHeader, spoofed)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
rt.chatSend(rec, req, user)
|
||||
return rec
|
||||
}
|
||||
if rec := send("politikalar", ""); rec.Code != http.StatusNotImplemented || !strings.Contains(rec.Body.String(), "SQLite") {
|
||||
t.Errorf("a collection chosen: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if rec := send("", "politikalar"); rec.Code != http.StatusOK {
|
||||
t.Errorf("a header the page set itself was used: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if rec := send("none", ""); rec.Code != http.StatusOK {
|
||||
t.Errorf("none: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
@@ -121,6 +121,8 @@ func (r *Router) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /v1/batches", r.handleListBatches)
|
||||
mux.HandleFunc("GET /v1/batches/{id}", r.handleGetBatch)
|
||||
mux.HandleFunc("POST /v1/batches/{id}/cancel", r.handleCancelBatch)
|
||||
// Knowledge bases: documents a person's llmgw-kb read and embedded.
|
||||
r.registerKB(mux)
|
||||
// Liveness: this process is up. It says nothing about whether it can
|
||||
// serve, which is why it is not the thing to put a load balancer on.
|
||||
mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
|
||||
@@ -245,6 +247,9 @@ func (r *Router) Handler() http.Handler {
|
||||
if r.store != nil && r.guardLibrary != nil {
|
||||
pages.RegisterAdminGuard(mux, r.store, r.guardLibrary)
|
||||
}
|
||||
if r.store != nil && r.settingsCheck != nil {
|
||||
pages.RegisterAdminSettings(mux, r.store, r.settingsCheck, r.settingsLocal)
|
||||
}
|
||||
pages.Register(mux)
|
||||
// The docs portal, with the menu of whoever is signed in, and only to
|
||||
// them when the site asked for that.
|
||||
@@ -405,6 +410,11 @@ type parsedRequest struct {
|
||||
reservation accounting.Reservation
|
||||
// guardIn is what the guard decided about the request.
|
||||
guardIn guard.Verdict
|
||||
// knowledge is set when passages from knowledge bases were added.
|
||||
knowledge bool
|
||||
// kbOrg is the organization whose passages were added, for the prefix
|
||||
// cache's salt.
|
||||
kbOrg string
|
||||
// noStore is a policy that keeps nothing of this request.
|
||||
noStore bool
|
||||
// inCluster is a project whose requests may not leave the cluster.
|
||||
@@ -573,6 +583,13 @@ func (r *Router) handleInference(path string) http.HandlerFunc {
|
||||
w.Header().Set(trimmedHeader, strconv.Itoa(dropped))
|
||||
}
|
||||
|
||||
// The passages a question asked knowledge bases for, before it is
|
||||
// judged: they are text the model reads, and a document can carry
|
||||
// instructions aimed at its reader.
|
||||
if !r.retrieve(req.Context(), w, parsed, start) {
|
||||
return
|
||||
}
|
||||
|
||||
// Judged as it will be sent: authorised, within its limits, and with
|
||||
// its stored conversation in it, which is text the model reads too.
|
||||
if r.guard != nil && !r.guardInput(req.Context(), w, parsed, start) {
|
||||
|
||||
@@ -0,0 +1,355 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
"golang.org/x/text/unicode/norm"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// Knowledge bases on the client listener (BACKLOG F2): what llmgw-kb, run by
|
||||
// a person on their own machine with their own key, reads and writes. The
|
||||
// documents are read, cut and embedded there; what arrives here is a
|
||||
// document's name, digest and chunks with their vectors.
|
||||
//
|
||||
// Which collections a caller may see is decided from the key's place in the
|
||||
// tree and nothing the request says, beyond naming a project for a personal
|
||||
// key. A collection the caller may not read is answered exactly as one that
|
||||
// does not exist.
|
||||
|
||||
// projectHeader names the project a personal key acts in; the person must be
|
||||
// a member of it.
|
||||
const projectHeader = "X-Llmgw-Project"
|
||||
|
||||
// Bounds on one document, so one request cannot hold the control plane's
|
||||
// memory or a transaction for long.
|
||||
const (
|
||||
kbMaxBody = 64 << 20
|
||||
kbMaxChunks = 20000
|
||||
kbMaxChunkRunes = 8000
|
||||
kbMaxNameRunes = 1024
|
||||
)
|
||||
|
||||
// KBCollectionView is a collection as a caller sees it.
|
||||
type KBCollectionView struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
// Owner is the owning scope as scope:ID; OwnerName how people know it,
|
||||
// "Acme / hr".
|
||||
Owner string `json:"owner"`
|
||||
OwnerName string `json:"owner_name,omitempty"`
|
||||
EmbedModel string `json:"embed_model"`
|
||||
Dims int `json:"dims"`
|
||||
// Writable is whether this caller may add and remove documents.
|
||||
Writable bool `json:"writable"`
|
||||
}
|
||||
|
||||
// KBDocumentView is a document as the manifest lists it.
|
||||
type KBDocumentView struct {
|
||||
Name string `json:"name"`
|
||||
Digest string `json:"digest"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Chunks int `json:"chunks"`
|
||||
AddedAt time.Time `json:"added_at"`
|
||||
}
|
||||
|
||||
// KBManifest is a collection and its documents: what a sync compares the
|
||||
// files on disk with.
|
||||
type KBManifest struct {
|
||||
Collection KBCollectionView `json:"collection"`
|
||||
Documents []KBDocumentView `json:"documents"`
|
||||
}
|
||||
|
||||
// KBPutDocument is one document, read and embedded by the client.
|
||||
type KBPutDocument struct {
|
||||
Name string `json:"name"`
|
||||
Digest string `json:"digest"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
Chunks []struct {
|
||||
Text string `json:"text"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"chunks"`
|
||||
}
|
||||
|
||||
// KBPutResult is what adding a document did.
|
||||
type KBPutResult struct {
|
||||
Name string `json:"name"`
|
||||
Chunks int `json:"chunks"`
|
||||
Changed bool `json:"changed"`
|
||||
}
|
||||
|
||||
func (r *Router) registerKB(mux *http.ServeMux) {
|
||||
mux.HandleFunc("GET /v1/kb/collections", r.handleKBCollections)
|
||||
mux.HandleFunc("GET /v1/kb/collections/{id}/manifest", r.handleKBManifest)
|
||||
mux.HandleFunc("PUT /v1/kb/collections/{id}/documents", r.handleKBPut)
|
||||
mux.HandleFunc("DELETE /v1/kb/collections/{id}/documents", r.handleKBDelete)
|
||||
}
|
||||
|
||||
// kbCaller is the caller's place in the tree, for deciding what it may read.
|
||||
// A personal key has none unless it names a project it is a member of.
|
||||
func (r *Router) kbCaller(w http.ResponseWriter, req *http.Request) (auth.Principal, storepkg.Effective, bool) {
|
||||
p, ok := auth.FromContext(req.Context())
|
||||
if !ok || r.store == nil {
|
||||
writeError(w, http.StatusNotImplemented, "invalid_request_error",
|
||||
"knowledge bases need authentication and the store, which this gateway runs without")
|
||||
return p, storepkg.Effective{}, false
|
||||
}
|
||||
if err := r.store.KnowledgeAvailable(); err != nil {
|
||||
writeError(w, http.StatusNotImplemented, "invalid_request_error", err.Error())
|
||||
return p, storepkg.Effective{}, false
|
||||
}
|
||||
e, err := r.kbEffective(req.Context(), p, req.Header.Get(projectHeader))
|
||||
if err != nil {
|
||||
var ke *kbError
|
||||
if errors.As(err, &ke) {
|
||||
writeError(w, ke.status, ke.kind, ke.msg)
|
||||
} else {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
}
|
||||
return p, e, false
|
||||
}
|
||||
return p, e, true
|
||||
}
|
||||
|
||||
func (r *Router) kbView(ctx context.Context, c storepkg.KBCollection, writable bool) KBCollectionView {
|
||||
return KBCollectionView{ID: c.ID, Name: c.Name, Owner: string(c.OwnerScope) + ":" + c.OwnerID,
|
||||
OwnerName: r.store.OwnerName(ctx, c.OwnerScope, c.OwnerID), EmbedModel: c.EmbedModel, Dims: c.Dims, Writable: writable}
|
||||
}
|
||||
|
||||
// kbCollection finds a collection by id or name among those the caller may
|
||||
// read (or write), answering 404 for one it may not.
|
||||
func (r *Router) kbCollection(w http.ResponseWriter, req *http.Request, e storepkg.Effective, write bool) (storepkg.KBCollection, bool) {
|
||||
want := req.PathValue("id")
|
||||
cs, err := r.store.CollectionsFor(req.Context(), e, write)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return storepkg.KBCollection{}, false
|
||||
}
|
||||
for _, c := range cs {
|
||||
if c.ID == want {
|
||||
return c, true
|
||||
}
|
||||
}
|
||||
for _, c := range cs {
|
||||
if c.Name == want {
|
||||
return c, true
|
||||
}
|
||||
}
|
||||
msg := "no collection " + want
|
||||
if e.OrgID == "" {
|
||||
msg += " (a personal key names its project in " + projectHeader + ")"
|
||||
}
|
||||
writeError(w, http.StatusNotFound, "not_found_error", msg)
|
||||
return storepkg.KBCollection{}, false
|
||||
}
|
||||
|
||||
func (r *Router) handleKBCollections(w http.ResponseWriter, req *http.Request) {
|
||||
_, e, ok := r.kbCaller(w, req)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
read, err := r.store.CollectionsFor(req.Context(), e, false)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return
|
||||
}
|
||||
write, err := r.store.CollectionsFor(req.Context(), e, true)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return
|
||||
}
|
||||
writable := map[string]bool{}
|
||||
for _, c := range write {
|
||||
writable[c.ID] = true
|
||||
}
|
||||
out := []KBCollectionView{}
|
||||
for _, c := range read {
|
||||
out = append(out, r.kbView(req.Context(), c, writable[c.ID]))
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": out})
|
||||
}
|
||||
|
||||
func (r *Router) handleKBManifest(w http.ResponseWriter, req *http.Request) {
|
||||
_, e, ok := r.kbCaller(w, req)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
c, ok := r.kbCollection(w, req, e, false)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
docs, err := r.store.Documents(req.Context(), c.ID)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return
|
||||
}
|
||||
writable, _ := r.store.CollectionsFor(req.Context(), e, true)
|
||||
m := KBManifest{Collection: r.kbView(req.Context(), c, containsCollection(writable, c.ID)), Documents: []KBDocumentView{}}
|
||||
for _, d := range docs {
|
||||
m.Documents = append(m.Documents, KBDocumentView{Name: d.Name, Digest: d.Digest, Bytes: d.Bytes,
|
||||
Chunks: d.Chunks, AddedAt: d.AddedAt})
|
||||
}
|
||||
writeJSON(w, http.StatusOK, m)
|
||||
}
|
||||
|
||||
func containsCollection(cs []storepkg.KBCollection, id string) bool {
|
||||
for _, c := range cs {
|
||||
if c.ID == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// kbWriter is a caller about to change a collection: one it may write, and a
|
||||
// caller whose policy does not forbid keeping what it sends.
|
||||
func (r *Router) kbWriter(w http.ResponseWriter, req *http.Request) (auth.Principal, storepkg.KBCollection, bool) {
|
||||
p, e, ok := r.kbCaller(w, req)
|
||||
if !ok {
|
||||
return p, storepkg.KBCollection{}, false
|
||||
}
|
||||
if r.noStore(guard.Subject{KeyID: p.KeyID, UserID: p.UserID, ProjectID: e.ProjectID,
|
||||
TeamID: e.TeamID, OrgID: e.OrgID, Path: req.URL.Path}) {
|
||||
writeError(w, http.StatusForbidden, "permission_error",
|
||||
"the gateway's policy for this caller keeps nothing, and a knowledge base is documents kept")
|
||||
return p, storepkg.KBCollection{}, false
|
||||
}
|
||||
c, ok := r.kbCollection(w, req, e, true)
|
||||
return p, c, ok
|
||||
}
|
||||
|
||||
// validDocName checks a document's name: what the client sends is a path
|
||||
// relative to what it synced, with forward slashes.
|
||||
func validDocName(n string) (string, error) {
|
||||
if !utf8.ValidString(n) {
|
||||
return "", errors.New("the name is not UTF-8")
|
||||
}
|
||||
n = norm.NFC.String(strings.TrimSpace(n))
|
||||
switch {
|
||||
case n == "":
|
||||
return "", errors.New("a document needs a name")
|
||||
case utf8.RuneCountInString(n) > kbMaxNameRunes:
|
||||
return "", fmt.Errorf("a name has at most %d characters", kbMaxNameRunes)
|
||||
case strings.ContainsFunc(n, unicode.IsControl):
|
||||
return "", errors.New("the name holds a control character")
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (r *Router) handleKBPut(w http.ResponseWriter, req *http.Request) {
|
||||
p, c, ok := r.kbWriter(w, req)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var in KBPutDocument
|
||||
dec := json.NewDecoder(http.MaxBytesReader(w, req.Body, kbMaxBody))
|
||||
if err := dec.Decode(&in); err != nil {
|
||||
var tooBig *http.MaxBytesError
|
||||
if errors.As(err, &tooBig) {
|
||||
writeError(w, http.StatusRequestEntityTooLarge, "invalid_request_error",
|
||||
fmt.Sprintf("a document is at most %d MiB of JSON; split the file", kbMaxBody>>20))
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", "the request is not valid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
name, err := validDocName(in.Name)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case in.Digest == "":
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", "a document needs the digest of its file")
|
||||
return
|
||||
case len(in.Chunks) == 0:
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", "a document needs at least one chunk")
|
||||
return
|
||||
case len(in.Chunks) > kbMaxChunks:
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error",
|
||||
fmt.Sprintf("%d chunks: a document has at most %d", len(in.Chunks), kbMaxChunks))
|
||||
return
|
||||
}
|
||||
chunks := make([]storepkg.KBChunk, len(in.Chunks))
|
||||
for i, ch := range in.Chunks {
|
||||
if !utf8.ValidString(ch.Text) || strings.TrimSpace(ch.Text) == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("chunk %d has no UTF-8 text", i))
|
||||
return
|
||||
}
|
||||
if utf8.RuneCountInString(ch.Text) > kbMaxChunkRunes {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error",
|
||||
fmt.Sprintf("chunk %d is over %d characters", i, kbMaxChunkRunes))
|
||||
return
|
||||
}
|
||||
if len(ch.Embedding) != c.Dims {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error",
|
||||
fmt.Sprintf("chunk %d has %d dimensions; collection %s is embedded by %s in %d", i,
|
||||
len(ch.Embedding), c.Name, c.EmbedModel, c.Dims))
|
||||
return
|
||||
}
|
||||
chunks[i] = storepkg.KBChunk{Body: strings.ToValidUTF8(ch.Text, ""), Embedding: ch.Embedding}
|
||||
}
|
||||
d, changed, err := r.store.PutDocument(req.Context(), c, storepkg.KBDocument{Name: name, Digest: in.Digest,
|
||||
Bytes: in.Bytes, AddedBy: kbActor(p)}, chunks)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return
|
||||
}
|
||||
if changed {
|
||||
r.kbAudit(req.Context(), p, "PUT /v1/kb/collections/{id}/documents", c, name, http.StatusOK)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, KBPutResult{Name: name, Chunks: d.Chunks, Changed: changed})
|
||||
}
|
||||
|
||||
func (r *Router) handleKBDelete(w http.ResponseWriter, req *http.Request) {
|
||||
p, c, ok := r.kbWriter(w, req)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
name, err := validDocName(req.URL.Query().Get("name"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
gone, err := r.store.DeleteDocument(req.Context(), c.ID, name)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "server_error", err.Error())
|
||||
return
|
||||
}
|
||||
if !gone {
|
||||
writeError(w, http.StatusNotFound, "not_found_error", "no document "+name+" in "+c.Name)
|
||||
return
|
||||
}
|
||||
r.kbAudit(req.Context(), p, "DELETE /v1/kb/collections/{id}/documents", c, name, http.StatusOK)
|
||||
writeJSON(w, http.StatusOK, map[string]any{"name": name, "deleted": true})
|
||||
}
|
||||
|
||||
// kbActor is who changed a collection, as the audit log names people: the
|
||||
// person, or the key when it belongs to no one.
|
||||
func kbActor(p auth.Principal) string {
|
||||
if p.UserID != "" {
|
||||
return "user:" + p.UserID
|
||||
}
|
||||
return "key:" + p.KeyID
|
||||
}
|
||||
|
||||
// kbAudit records a change to a collection. The document's name is kept; its
|
||||
// text never is.
|
||||
func (r *Router) kbAudit(ctx context.Context, p auth.Principal, action string, c storepkg.KBCollection, doc string, status int) {
|
||||
err := r.store.WriteAudit(context.WithoutCancel(ctx), storepkg.AuditEvent{Actor: kbActor(p), Source: "gateway",
|
||||
Action: action, Target: "collection:" + c.ID, Detail: c.Name + ": " + doc, Status: status})
|
||||
if err != nil {
|
||||
r.log.Warn("could not record a knowledge base change in the audit log", "collection", c.ID, "err", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// The whole way: llmgw-kb's sync against this gateway, embeddings served by
|
||||
// an engine through the ordinary path, chunks in PostgreSQL, found again by
|
||||
// a search.
|
||||
func TestASyncFromAMachineIsSearchable(t *testing.T) {
|
||||
rt := kbRouter(t)
|
||||
ctx := t.Context()
|
||||
s := rt.store
|
||||
sfx := strings.ToLower(store.NewID(""))
|
||||
org, _ := s.CreateOrg(ctx, "e2e-"+sfx)
|
||||
team, _ := s.CreateTeam(ctx, org, "t")
|
||||
proj, _ := s.CreateProject(ctx, team, "p")
|
||||
coll, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj,
|
||||
Name: "belgeler", EmbedModel: "qwen2.5-14b-instruct", Dims: 3})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv := httptest.NewServer(rt.Handler())
|
||||
t.Cleanup(srv.Close)
|
||||
h, _ := kbsync.NewHTTP("")
|
||||
c := &kbsync.Client{Base: srv.URL, Key: kbKey(t, s, proj, ""), HTTP: h}
|
||||
|
||||
dir := t.TempDir()
|
||||
for name, body := range map[string]string{
|
||||
"izin.md": "# Yıllık izin\n\nÇalışanlar yılda yirmi gün ücretli izin kullanır.\n",
|
||||
"masraf.md": "# Masraflar\n\nMasraflar finans ekibine iletilir.\n",
|
||||
} {
|
||||
if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
l, err := kbsync.OpenLedger(ctx, filepath.Join(t.TempDir(), "ledger.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer l.Close()
|
||||
rep, err := kbsync.Sync(ctx, c, l, kbsync.Options{Collection: "belgeler", Roots: []string{dir}}, nil)
|
||||
if err != nil || rep.Sent != 2 {
|
||||
t.Fatalf("sync: %+v %v", rep, err)
|
||||
}
|
||||
docs, _ := s.Documents(ctx, coll.ID)
|
||||
if len(docs) != 2 {
|
||||
t.Fatalf("documents %+v", docs)
|
||||
}
|
||||
hits, err := s.SearchChunks(context.Background(), store.KBSearch{OrgID: org, Collections: []store.KBCollection{coll},
|
||||
Text: "kaç gün izin", K: 1})
|
||||
if err != nil || len(hits) != 1 || hits[0].Document != "izin.md" || !strings.Contains(hits[0].Body, "Yıllık izin\n\n") {
|
||||
t.Errorf("search: %+v %v", hits, err)
|
||||
}
|
||||
// Again: nothing sent.
|
||||
rep, err = kbsync.Sync(ctx, c, l, kbsync.Options{Collection: "belgeler", Roots: []string{dir}}, nil)
|
||||
if err != nil || rep.Sent != 0 || rep.Unchanged != 2 {
|
||||
t.Errorf("second sync: %+v %v", rep, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func TestKnowledgeBasesSayWhyTheyAreOffOnSQLite(t *testing.T) {
|
||||
e := newFakeEngine(t, "qwen2.5-14b-instruct")
|
||||
rt, secret := pricedRouter(t, e.srv.URL)
|
||||
code, _, body := call(t, rt.Handler(), secret, http.MethodGet, "/v1/kb/collections", nil)
|
||||
if code != http.StatusNotImplemented || !strings.Contains(body, "SQLite") {
|
||||
t.Errorf("got %d %s", code, body)
|
||||
}
|
||||
}
|
||||
|
||||
// kbRouter is a router on the PostgreSQL test database, with pgvector.
|
||||
func kbRouter(t *testing.T) *Router {
|
||||
t.Helper()
|
||||
return kbRouterWith(t, newFakeEngine(t, "qwen2.5-14b-instruct").srv.URL)
|
||||
}
|
||||
|
||||
// kbRouterWith is a router on the PostgreSQL test database, with pgvector,
|
||||
// in front of an engine at engineURL.
|
||||
func kbRouterWith(t *testing.T, engineURL string) *Router {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||||
if dsn == "" {
|
||||
t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run these")
|
||||
}
|
||||
if u, err := url.Parse(dsn); err != nil || !strings.Contains(u.Path, "test") {
|
||||
t.Fatalf("refusing a database whose name does not say test: %s", dsn)
|
||||
}
|
||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: engineURL, Models: []string{"qwen2.5-14b-instruct"}})
|
||||
cfg.Auth.Enabled = true
|
||||
cfg.Auth.Store = dsn
|
||||
rt := newTestRouter(t, cfg)
|
||||
if err := rt.store.KnowledgeAvailable(); err != nil {
|
||||
t.Skipf("no pgvector in the test database: %v", err)
|
||||
}
|
||||
return rt
|
||||
}
|
||||
|
||||
// kbKey makes a key in a project, or a person's key with no project.
|
||||
func kbKey(t *testing.T, s *store.DB, project, user string) string {
|
||||
t.Helper()
|
||||
secret := store.KeyPrefix + strings.ToLower(store.NewID("")) + strings.Repeat("k", 27)
|
||||
sum := sha256.Sum256([]byte(secret))
|
||||
k := store.Key{ID: store.NewID("key-"), Hash: sum[:], Prefix: secret[:8], Name: "t", ProjectID: project, CreatedAt: time.Now()}
|
||||
var err error
|
||||
if user != "" {
|
||||
err = s.CreateUserKey(t.Context(), user, k)
|
||||
} else {
|
||||
err = s.CreateKey(t.Context(), k)
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return secret
|
||||
}
|
||||
|
||||
func vec(x ...float32) []float32 { return x }
|
||||
|
||||
func doc(name, digest string, embeddings ...[]float32) map[string]any {
|
||||
var chunks []map[string]any
|
||||
for i, e := range embeddings {
|
||||
chunks = append(chunks, map[string]any{"text": name + " parça " + string(rune('a'+i)), "embedding": e})
|
||||
}
|
||||
return map[string]any{"name": name, "digest": digest, "bytes": 10, "chunks": chunks}
|
||||
}
|
||||
|
||||
func TestKnowledgeBasesKeepTenantsApart(t *testing.T) {
|
||||
rt := kbRouter(t)
|
||||
h := rt.Handler()
|
||||
ctx := t.Context()
|
||||
s := rt.store
|
||||
// Unique names: the test database is shared and not emptied here.
|
||||
sfx := strings.ToLower(store.NewID(""))
|
||||
orgA, _ := s.CreateOrg(ctx, "a-"+sfx)
|
||||
teamA, _ := s.CreateTeam(ctx, orgA, "t")
|
||||
hr, _ := s.CreateProject(ctx, teamA, "hr")
|
||||
sales, _ := s.CreateProject(ctx, teamA, "sales")
|
||||
orgB, _ := s.CreateOrg(ctx, "b-"+sfx)
|
||||
teamB, _ := s.CreateTeam(ctx, orgB, "t")
|
||||
other, _ := s.CreateProject(ctx, teamB, "x")
|
||||
c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: hr,
|
||||
Name: "politikalar", EmbedModel: "emb", Dims: 3})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hrKey, salesKey, otherKey := kbKey(t, s, hr, ""), kbKey(t, s, sales, ""), kbKey(t, s, other, "")
|
||||
base := "/v1/kb/collections/"
|
||||
|
||||
// The owner's project writes, by name or by id.
|
||||
code, out, body := call(t, h, hrKey, http.MethodPut, base+"politikalar/documents", doc("izin.pdf", "d1", vec(1, 0, 0), vec(0, 1, 0)))
|
||||
if code != http.StatusOK || out["changed"] != true || out["chunks"] != float64(2) {
|
||||
t.Fatalf("owner's put: %d %s", code, body)
|
||||
}
|
||||
code, out, _ = call(t, h, hrKey, http.MethodPut, base+c.ID+"/documents", doc("izin.pdf", "d1", vec(1, 0, 0)))
|
||||
if code != http.StatusOK || out["changed"] != false {
|
||||
t.Errorf("the same digest again: %d %v", code, out)
|
||||
}
|
||||
code, _, body = call(t, h, hrKey, http.MethodGet, base+"politikalar/manifest", nil)
|
||||
if code != http.StatusOK || !strings.Contains(body, `"digest":"d1"`) || !strings.Contains(body, `"writable":true`) {
|
||||
t.Errorf("manifest: %d %s", code, body)
|
||||
}
|
||||
|
||||
// Another organization, and another project without a grant, see no
|
||||
// such collection — not a refusal that says it exists.
|
||||
for name, key := range map[string]string{"other org": otherKey, "other project": salesKey} {
|
||||
for _, m := range []string{http.MethodGet + " manifest", http.MethodPut + " documents"} {
|
||||
method, path, _ := strings.Cut(m, " ")
|
||||
code, _, body := call(t, h, key, method, base+c.ID+"/"+path, doc("x", "y", vec(1, 0, 0)))
|
||||
if code != http.StatusNotFound {
|
||||
t.Errorf("%s %s: %d %s", name, m, code, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A read grant lets sales read, not write.
|
||||
if err := s.Grant(ctx, store.KBGrant{Scope: store.ScopeProject, ScopeID: sales, Collection: c.ID, Perm: store.KBRead}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if code, _, body := call(t, h, salesKey, http.MethodGet, base+"politikalar/manifest", nil); code != http.StatusOK ||
|
||||
!strings.Contains(body, `"writable":false`) {
|
||||
t.Errorf("read grant, manifest: %d %s", code, body)
|
||||
}
|
||||
if code, _, _ := call(t, h, salesKey, http.MethodPut, base+"politikalar/documents", doc("x", "y", vec(1, 0, 0))); code != http.StatusNotFound {
|
||||
t.Errorf("read grant, put: %d", code)
|
||||
}
|
||||
if code, _, body := call(t, h, salesKey, http.MethodGet, "/v1/kb/collections", nil); code != http.StatusOK ||
|
||||
!strings.Contains(body, "politikalar") || !strings.Contains(body, `"owner_name":"a-`+sfx+` / t / hr"`) {
|
||||
t.Errorf("read grant, list: %d %s", code, body)
|
||||
}
|
||||
if _, _, body := call(t, h, otherKey, http.MethodGet, "/v1/kb/collections", nil); strings.Contains(body, "politikalar") {
|
||||
t.Errorf("another organization lists it: %s", body)
|
||||
}
|
||||
|
||||
// Vectors of the wrong size are refused before anything is written.
|
||||
if code, _, body := call(t, h, hrKey, http.MethodPut, base+"politikalar/documents", doc("b.md", "d", vec(1, 0))); code != http.StatusBadRequest ||
|
||||
!strings.Contains(body, "dimensions") {
|
||||
t.Errorf("wrong dimensions: %d %s", code, body)
|
||||
}
|
||||
|
||||
if code, _, body := call(t, h, hrKey, http.MethodDelete, base+"politikalar/documents?name=izin.pdf", nil); code != http.StatusOK {
|
||||
t.Errorf("delete: %d %s", code, body)
|
||||
}
|
||||
if code, _, _ := call(t, h, hrKey, http.MethodDelete, base+"politikalar/documents?name=izin.pdf", nil); code != http.StatusNotFound {
|
||||
t.Errorf("deleting again: %d", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPersonalKeyReadsOnlyInTheProjectItNames(t *testing.T) {
|
||||
rt := kbRouter(t)
|
||||
h := rt.Handler()
|
||||
ctx := t.Context()
|
||||
s := rt.store
|
||||
sfx := strings.ToLower(store.NewID(""))
|
||||
org, _ := s.CreateOrg(ctx, "p-"+sfx)
|
||||
team, _ := s.CreateTeam(ctx, org, "t")
|
||||
hr, _ := s.CreateProject(ctx, team, "hr")
|
||||
fin, _ := s.CreateProject(ctx, team, "fin")
|
||||
uid := store.NewID("u-")
|
||||
if err := s.CreateUser(ctx, store.User{ID: uid, Email: uid + "@example.com", Name: "u", PasswordHash: "x",
|
||||
State: store.UserActive, CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.AddMembership(ctx, uid, store.ScopeProject, hr, "member", "t", time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: hr,
|
||||
Name: "ik", EmbedModel: "emb", Dims: 3}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: fin,
|
||||
Name: "finans", EmbedModel: "emb", Dims: 3}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
key := kbKey(t, s, "", uid)
|
||||
get := func(project string) (int, string) {
|
||||
req, _ := http.NewRequest(http.MethodGet, "/v1/kb/collections", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if project != "" {
|
||||
req.Header.Set(projectHeader, project)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
return rec.Code, rec.Body.String()
|
||||
}
|
||||
if code, body := get(""); code != http.StatusOK || strings.Contains(body, "ik") {
|
||||
t.Errorf("no project named: %d %s", code, body)
|
||||
}
|
||||
if code, body := get(hr); code != http.StatusOK || !strings.Contains(body, `"ik"`) || strings.Contains(body, "finans") {
|
||||
t.Errorf("in hr: %d %s", code, body)
|
||||
}
|
||||
if code, _ := get(fin); code != http.StatusForbidden {
|
||||
t.Errorf("a project the person is not in: %d", code)
|
||||
}
|
||||
}
|
||||
@@ -45,8 +45,12 @@ type metrics struct {
|
||||
// cacheRequests counts what the answer cache did, cacheBytes and
|
||||
// cacheEntries what it holds.
|
||||
cacheRequests *prometheus.CounterVec
|
||||
cacheBytes prometheus.Gauge
|
||||
cacheEntries prometheus.Gauge
|
||||
// kbRetrievals counts requests answered from knowledge bases by what
|
||||
// the search found; kbRetrieval is how long finding it took.
|
||||
kbRetrievals *prometheus.CounterVec
|
||||
kbRetrieval prometheus.Histogram
|
||||
cacheBytes prometheus.Gauge
|
||||
cacheEntries prometheus.Gauge
|
||||
|
||||
tenant *tenantMetrics
|
||||
guard *guardMetrics
|
||||
@@ -138,6 +142,16 @@ func newMetrics() *metrics {
|
||||
"hit (answered from it), miss (asked the model, and kept the answer if it could), " +
|
||||
"bypass (not one whose answer is reused).",
|
||||
}, []string{"model", "result"}),
|
||||
kbRetrievals: prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Name: "llmgw_kb_retrievals_total",
|
||||
Help: "Chat requests that named knowledge-base collections, by what came of it: " +
|
||||
"found (passages given to the model), nothing (none matched), refused (not answered).",
|
||||
}, []string{"result"}),
|
||||
kbRetrieval: prometheus.NewHistogram(prometheus.HistogramOpts{
|
||||
Name: "llmgw_kb_retrieval_seconds",
|
||||
Help: "Time to find a request's passages: the question embedded, and the collections searched.",
|
||||
Buckets: []float64{0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10},
|
||||
}),
|
||||
cacheBytes: prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "llmgw_cache_bytes",
|
||||
Help: "Memory the answers kept for reuse take in this process, roughly.",
|
||||
@@ -161,6 +175,7 @@ func (m *metrics) collectors() []prometheus.Collector {
|
||||
m.adapterAffinity, m.adapterLoads, m.adapterEvictions,
|
||||
m.kvDecisions, m.kvTokenize,
|
||||
m.cacheRequests, m.cacheBytes, m.cacheEntries,
|
||||
m.kbRetrievals, m.kbRetrieval,
|
||||
}, slices.Concat(m.tenant.collectors(), m.guard.collectors(), m.token.collectors())...)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
|
||||
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// Answering from knowledge bases (BACKLOG F3): a chat request that names
|
||||
// collections in X-Llmgw-Knowledge has the passages its last question needs
|
||||
// found in them and put in front of the model, with their sources, before it
|
||||
// is judged and sent. Which collections it may name is decided by the key's
|
||||
// place in the tree, exactly as llmgw-kb's own requests are; one it may not
|
||||
// read is answered as one that does not exist.
|
||||
|
||||
const (
|
||||
// knowledgeHeader names the collections, by name or id, comma separated.
|
||||
knowledgeHeader = "X-Llmgw-Knowledge"
|
||||
// knowledgeKHeader is how many passages; kbDefaultK when absent.
|
||||
knowledgeKHeader = "X-Llmgw-Knowledge-K"
|
||||
// sourcesHeader answers which passages were given: document#chunk, each
|
||||
// path-escaped, in the order they were numbered.
|
||||
sourcesHeader = "x-llmgw-sources"
|
||||
|
||||
kbDefaultK = 5
|
||||
kbMaxK = 20
|
||||
)
|
||||
|
||||
// kbContextIntro is what the model is told about the passages. In English,
|
||||
// which every model reads; it is asked to answer in the question's language.
|
||||
const kbContextIntro = "Below are passages from the organisation's own documents, found for the user's last " +
|
||||
"question. Answer from them, and cite the ones you use as [1], [2] and so on. If they do not hold the " +
|
||||
"answer, say so rather than guessing. Answer in the language of the question."
|
||||
|
||||
const kbNothingFound = "The user asked a question meant to be answered from the organisation's documents, " +
|
||||
"and nothing in them matched it. Say so, rather than answering from general knowledge as if it were " +
|
||||
"the organisation's."
|
||||
|
||||
// kbEffective is a caller's place in the tree for reading collections: a
|
||||
// personal key has one only in a project it names and belongs to.
|
||||
func (r *Router) kbEffective(ctx context.Context, p auth.Principal, project string) (storepkg.Effective, error) {
|
||||
e := storepkg.Effective{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID, OrgID: p.OrgID}
|
||||
project = strings.TrimSpace(project)
|
||||
if project == "" || p.ProjectID != "" || p.UserID == "" {
|
||||
return e, nil
|
||||
}
|
||||
may, err := r.store.CanUseProject(ctx, p.UserID, project)
|
||||
if err != nil {
|
||||
return e, err
|
||||
}
|
||||
if !may {
|
||||
return e, &kbError{status: http.StatusForbidden, kind: "permission_error", msg: "you are not a member of project " + project}
|
||||
}
|
||||
eff, err := r.store.ResolveFor(ctx, p.KeyID, p.UserID, project)
|
||||
if err != nil {
|
||||
return e, err
|
||||
}
|
||||
e.ProjectID, e.TeamID, e.OrgID = eff.ProjectID, eff.TeamID, eff.OrgID
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// kbError is a refusal with the status and kind a client is answered with.
|
||||
type kbError struct {
|
||||
status int
|
||||
kind, msg string
|
||||
retryAfter string
|
||||
}
|
||||
|
||||
func (e *kbError) Error() string { return e.msg }
|
||||
|
||||
// retrieve finds and adds the passages a request asked for. It answers false
|
||||
// when it has answered the client itself.
|
||||
func (r *Router) retrieve(ctx context.Context, w http.ResponseWriter, req *parsedRequest, start time.Time) bool {
|
||||
names := splitList(req.Header.Get(knowledgeHeader))
|
||||
if len(names) == 1 && strings.EqualFold(names[0], "none") {
|
||||
// Asked for no passages, whatever the default.
|
||||
return true
|
||||
}
|
||||
// Without the header, the default of the caller's part of the tree, if
|
||||
// it has one. A default is the administrator's choice, not the
|
||||
// caller's, so it never refuses a request: a collection the caller
|
||||
// cannot read is left out, and a search that fails is answered without
|
||||
// passages.
|
||||
explicit := len(names) > 0
|
||||
if !explicit {
|
||||
if req.Path != "/v1/chat/completions" || req.internal || r.store == nil || r.store.KnowledgeAvailable() != nil {
|
||||
return true
|
||||
}
|
||||
p, ok := auth.FromContext(ctx)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
if names = r.kbDefaults(ctx, p, req.Header.Get(projectHeader)); len(names) == 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
fail := func(err error) bool {
|
||||
if !explicit {
|
||||
r.log.Warn("a knowledge-base default could not be searched; answering without it", "model", req.Model, "err", err)
|
||||
r.m.kbRetrievals.WithLabelValues("skipped").Inc()
|
||||
return true
|
||||
}
|
||||
var ke *kbError
|
||||
if !errors.As(err, &ke) {
|
||||
ke = &kbError{status: http.StatusInternalServerError, kind: "server_error", msg: err.Error()}
|
||||
}
|
||||
if ke.retryAfter != "" {
|
||||
w.Header().Set("Retry-After", ke.retryAfter)
|
||||
}
|
||||
writeError(w, ke.status, ke.kind, ke.msg)
|
||||
r.m.kbRetrievals.WithLabelValues("refused").Inc()
|
||||
r.m.observeRequest(r.metricModel(req.Model), "", ke.status, time.Since(start), nil, 0)
|
||||
return false
|
||||
}
|
||||
if req.Path != "/v1/chat/completions" {
|
||||
return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error",
|
||||
msg: knowledgeHeader + " is for chat completions"})
|
||||
}
|
||||
p, ok := auth.FromContext(ctx)
|
||||
if !ok || r.store == nil {
|
||||
return fail(&kbError{status: http.StatusNotImplemented, kind: "invalid_request_error",
|
||||
msg: "knowledge bases need authentication and the store, which this gateway runs without"})
|
||||
}
|
||||
if err := r.store.KnowledgeAvailable(); err != nil {
|
||||
return fail(&kbError{status: http.StatusNotImplemented, kind: "invalid_request_error", msg: err.Error()})
|
||||
}
|
||||
k := kbDefaultK
|
||||
if v := req.Header.Get(knowledgeKHeader); v != "" {
|
||||
n, err := strconv.Atoi(v)
|
||||
if err != nil || n < 1 || n > kbMaxK {
|
||||
return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error",
|
||||
msg: fmt.Sprintf("%s is 1 to %d", knowledgeKHeader, kbMaxK)})
|
||||
}
|
||||
k = n
|
||||
}
|
||||
began := time.Now()
|
||||
e, err := r.kbEffective(ctx, p, req.Header.Get(projectHeader))
|
||||
if err != nil {
|
||||
return fail(err)
|
||||
}
|
||||
allowed, err := r.store.CollectionsFor(ctx, e, false)
|
||||
if err != nil {
|
||||
return fail(err)
|
||||
}
|
||||
var colls []storepkg.KBCollection
|
||||
for _, n := range names {
|
||||
c, ok := pickCollection(allowed, n)
|
||||
if !ok && !explicit {
|
||||
continue
|
||||
}
|
||||
if !ok {
|
||||
msg := "no collection " + n
|
||||
if e.OrgID == "" {
|
||||
msg += " (a personal key names its project in " + projectHeader + ")"
|
||||
}
|
||||
return fail(&kbError{status: http.StatusNotFound, kind: "not_found_error", msg: msg})
|
||||
}
|
||||
colls = append(colls, c)
|
||||
}
|
||||
if len(colls) == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
question := lastUserText(req.UpstreamBody)
|
||||
if strings.TrimSpace(question) == "" && !explicit {
|
||||
return true
|
||||
}
|
||||
if strings.TrimSpace(question) == "" {
|
||||
return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error",
|
||||
msg: knowledgeHeader + " needs a last user message with text to search for"})
|
||||
}
|
||||
embeddings := map[string][]float32{}
|
||||
for _, c := range colls {
|
||||
if _, done := embeddings[c.EmbedModel]; done {
|
||||
continue
|
||||
}
|
||||
v, err := r.embedFor(ctx, p, c.EmbedModel, question)
|
||||
if err != nil {
|
||||
return fail(err)
|
||||
}
|
||||
embeddings[c.EmbedModel] = v
|
||||
}
|
||||
hits, err := r.store.SearchChunks(ctx, storepkg.KBSearch{OrgID: e.OrgID, Collections: colls,
|
||||
Embeddings: embeddings, Text: question, K: k})
|
||||
if err != nil {
|
||||
return fail(err)
|
||||
}
|
||||
|
||||
// What fits: the context the model has, less what the request already
|
||||
// holds and the answer it asked room for — the last passages go first.
|
||||
room := -1
|
||||
if _, entry := r.resolve(req.Model); entry.ContextLength > 0 {
|
||||
room = (entry.ContextLength-req.MaxTokens)*4 - len(req.UpstreamBody) - len(kbContextIntro) - 64
|
||||
}
|
||||
var b strings.Builder
|
||||
var sources []string
|
||||
for i, h := range hits {
|
||||
block := fmt.Sprintf("\n\n[%d] %s\n%s", i+1, h.Document, strings.TrimSpace(h.Body))
|
||||
if room >= 0 && b.Len()+len(block) > room {
|
||||
break
|
||||
}
|
||||
b.WriteString(block)
|
||||
sources = append(sources, url.PathEscape(h.Document)+"#"+strconv.Itoa(h.Ord))
|
||||
}
|
||||
text := kbNothingFound
|
||||
if len(sources) > 0 {
|
||||
text = kbContextIntro + b.String()
|
||||
}
|
||||
body, err := withSystemText(req.UpstreamBody, text)
|
||||
if err != nil {
|
||||
return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error", msg: err.Error()})
|
||||
}
|
||||
req.UpstreamBody = body
|
||||
req.PromptTokens = len(body) / 4
|
||||
req.knowledge = true
|
||||
req.kbOrg = e.OrgID
|
||||
// Set now: the proxy writes the headers when the answer starts.
|
||||
w.Header().Set(sourcesHeader, strings.Join(sources, ","))
|
||||
result := "found"
|
||||
if len(sources) == 0 {
|
||||
result = "nothing"
|
||||
}
|
||||
r.m.kbRetrievals.WithLabelValues(result).Inc()
|
||||
r.m.kbRetrieval.Observe(time.Since(began).Seconds())
|
||||
return true
|
||||
}
|
||||
|
||||
// pickCollection finds a collection by id, then by name.
|
||||
func pickCollection(cs []storepkg.KBCollection, want string) (storepkg.KBCollection, bool) {
|
||||
for _, c := range cs {
|
||||
if c.ID == want {
|
||||
return c, true
|
||||
}
|
||||
}
|
||||
for _, c := range cs {
|
||||
if strings.EqualFold(c.Name, want) {
|
||||
return c, true
|
||||
}
|
||||
}
|
||||
return storepkg.KBCollection{}, false
|
||||
}
|
||||
|
||||
func splitList(v string) []string {
|
||||
var out []string
|
||||
seen := map[string]bool{}
|
||||
for _, s := range strings.Split(v, ",") {
|
||||
if s = strings.TrimSpace(s); s != "" && !seen[s] {
|
||||
seen[s] = true
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// embedFor has a model embed the question as the caller: through the
|
||||
// ordinary path, so the caller's limits, budget and guardrails apply to it
|
||||
// and it is counted to them.
|
||||
func (r *Router) embedFor(ctx context.Context, p auth.Principal, model, text string) ([]float32, error) {
|
||||
body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}})
|
||||
hreq, err := http.NewRequestWithContext(auth.WithPrincipal(ctx, p), http.MethodPost, "/v1/embeddings", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hreq.Header.Set("Content-Type", "application/json")
|
||||
w := &bufferWriter{header: http.Header{}}
|
||||
r.handleInference("/v1/embeddings")(w, hreq)
|
||||
status := w.status
|
||||
if status == 0 {
|
||||
status = http.StatusOK
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
msg := errorMessageOf(w.buf.Bytes())
|
||||
ke := &kbError{status: status, kind: "server_error",
|
||||
msg: fmt.Sprintf("the knowledge base's model %s could not embed the question: %s", model, msg),
|
||||
retryAfter: w.header.Get("Retry-After")}
|
||||
if status < 500 {
|
||||
ke.kind = "invalid_request_error"
|
||||
}
|
||||
return nil, ke
|
||||
}
|
||||
var res struct {
|
||||
Data []struct {
|
||||
Embedding []float32 `json:"embedding"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(w.buf.Bytes(), &res); err != nil || len(res.Data) != 1 {
|
||||
return nil, fmt.Errorf("%s answered no embedding", model)
|
||||
}
|
||||
return res.Data[0].Embedding, nil
|
||||
}
|
||||
|
||||
func errorMessageOf(b []byte) string {
|
||||
var e struct {
|
||||
Error struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(b, &e) == nil && e.Error.Message != "" {
|
||||
return e.Error.Message
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
|
||||
// lastUserText is the text of a chat request's last user message: its
|
||||
// content, or the text parts of it.
|
||||
func lastUserText(body []byte) string {
|
||||
var in struct {
|
||||
Messages []struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
if json.Unmarshal(body, &in) != nil {
|
||||
return ""
|
||||
}
|
||||
for i := len(in.Messages) - 1; i >= 0; i-- {
|
||||
m := in.Messages[i]
|
||||
if m.Role != "user" {
|
||||
continue
|
||||
}
|
||||
var s string
|
||||
if json.Unmarshal(m.Content, &s) == nil {
|
||||
return s
|
||||
}
|
||||
var parts []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
if json.Unmarshal(m.Content, &parts) == nil {
|
||||
var b strings.Builder
|
||||
for _, p := range parts {
|
||||
if p.Type == "text" {
|
||||
b.WriteString(p.Text + "\n")
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
return ""
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// withSystemText adds text to a chat request's system message: appended to
|
||||
// the first one when it leads, or as a new one at the start. One system
|
||||
// message, first, because some models' templates refuse anything else.
|
||||
func withSystemText(body []byte, text string) ([]byte, error) {
|
||||
var top map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &top); err != nil {
|
||||
return nil, fmt.Errorf("the body is not a JSON object: %w", err)
|
||||
}
|
||||
var msgs []map[string]json.RawMessage
|
||||
if err := json.Unmarshal(top["messages"], &msgs); err != nil {
|
||||
return nil, errors.New("the request has no messages")
|
||||
}
|
||||
added := false
|
||||
if len(msgs) > 0 {
|
||||
var role string
|
||||
_ = json.Unmarshal(msgs[0]["role"], &role)
|
||||
if role == "system" || role == "developer" {
|
||||
var s string
|
||||
if json.Unmarshal(msgs[0]["content"], &s) == nil {
|
||||
msgs[0]["content"], _ = json.Marshal(strings.TrimRight(s, "\n") + "\n\n" + text)
|
||||
added = true
|
||||
} else {
|
||||
var parts []json.RawMessage
|
||||
if json.Unmarshal(msgs[0]["content"], &parts) == nil {
|
||||
part, _ := json.Marshal(map[string]string{"type": "text", "text": "\n\n" + text})
|
||||
msgs[0]["content"], _ = json.Marshal(append(parts, part))
|
||||
added = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if !added {
|
||||
content, _ := json.Marshal(text)
|
||||
role, _ := json.Marshal("system")
|
||||
msgs = append([]map[string]json.RawMessage{{"role": role, "content": content}}, msgs...)
|
||||
}
|
||||
var err error
|
||||
if top["messages"], err = json.Marshal(msgs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return json.Marshal(top)
|
||||
}
|
||||
|
||||
// kbDefaultTTL is how long a caller's default is remembered: a change to one
|
||||
// is in force within it, and a chat request does not cost a query.
|
||||
const kbDefaultTTL = 30 * time.Second
|
||||
|
||||
type kbDefaultEntry struct {
|
||||
at time.Time
|
||||
names []string
|
||||
}
|
||||
|
||||
// kbDefaults is the default collections in force for a caller.
|
||||
func (r *Router) kbDefaults(ctx context.Context, p auth.Principal, project string) []string {
|
||||
key := p.KeyID + "|" + p.ProjectID + "|" + project
|
||||
r.kbDefMu.Lock()
|
||||
e, ok := r.kbDefCache[key]
|
||||
r.kbDefMu.Unlock()
|
||||
if ok && time.Since(e.at) < kbDefaultTTL {
|
||||
return e.names
|
||||
}
|
||||
var names []string
|
||||
if eff, err := r.kbEffective(ctx, p, project); err == nil {
|
||||
names, err = r.store.KBDefaultFor(ctx, eff)
|
||||
if err != nil {
|
||||
r.log.Warn("could not read the knowledge-base defaults", "err", err)
|
||||
}
|
||||
}
|
||||
r.kbDefMu.Lock()
|
||||
if r.kbDefCache == nil || len(r.kbDefCache) > 10000 {
|
||||
r.kbDefCache = map[string]kbDefaultEntry{}
|
||||
}
|
||||
r.kbDefCache[key] = kbDefaultEntry{at: time.Now(), names: names}
|
||||
r.kbDefMu.Unlock()
|
||||
return names
|
||||
}
|
||||
|
||||
// withCacheSalt sets vLLM's cache_salt to one derived from an organization:
|
||||
// the same for its requests, so they still share a cache among themselves.
|
||||
func withCacheSalt(body []byte, org string) []byte {
|
||||
var top map[string]json.RawMessage
|
||||
if json.Unmarshal(body, &top) != nil {
|
||||
return body
|
||||
}
|
||||
sum := sha256.Sum256([]byte("llmgw-kb\x00" + org))
|
||||
top["cache_salt"], _ = json.Marshal(hex.EncodeToString(sum[:16]))
|
||||
out, err := json.Marshal(top)
|
||||
if err != nil {
|
||||
return body
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kb/kbsync"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
func TestPassagesJoinTheLeadingSystemMessage(t *testing.T) {
|
||||
type msg struct {
|
||||
Role string `json:"role"`
|
||||
Content any `json:"content"`
|
||||
}
|
||||
read := func(b []byte) []msg {
|
||||
var in struct {
|
||||
Messages []msg `json:"messages"`
|
||||
Model string
|
||||
}
|
||||
if err := json.Unmarshal(b, &in); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return in.Messages
|
||||
}
|
||||
// None: one is put first.
|
||||
out, err := withSystemText([]byte(`{"model":"m","messages":[{"role":"user","content":"q"}],"temperature":0.1}`), "CTX")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if m := read(out); len(m) != 2 || m[0].Role != "system" || m[0].Content != "CTX" || !strings.Contains(string(out), `"temperature":0.1`) {
|
||||
t.Errorf("no system message: %s", out)
|
||||
}
|
||||
// One leading: added to it, never a second.
|
||||
out, _ = withSystemText([]byte(`{"messages":[{"role":"system","content":"Be brief."},{"role":"user","content":"q"}]}`), "CTX")
|
||||
if m := read(out); len(m) != 2 || m[0].Content != "Be brief.\n\nCTX" {
|
||||
t.Errorf("leading system message: %s", out)
|
||||
}
|
||||
// Parts: a text part added.
|
||||
out, _ = withSystemText([]byte(`{"messages":[{"role":"system","content":[{"type":"text","text":"Be brief."}]},{"role":"user","content":"q"}]}`), "CTX")
|
||||
if !strings.Contains(string(out), `{"text":"\n\nCTX","type":"text"}`) {
|
||||
t.Errorf("parts: %s", out)
|
||||
}
|
||||
if _, err := withSystemText([]byte(`{"prompt":"x"}`), "CTX"); err == nil {
|
||||
t.Error("a request without messages was changed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheCacheSaltIsAnOrganizations(t *testing.T) {
|
||||
body := []byte(`{"model":"m","messages":[],"temperature":0}`)
|
||||
a1, a2, b := withCacheSalt(body, "org-a"), withCacheSalt(body, "org-a"), withCacheSalt(body, "org-b")
|
||||
var x, y struct {
|
||||
Salt string `json:"cache_salt"`
|
||||
T float64 `json:"temperature"`
|
||||
}
|
||||
_ = json.Unmarshal(a1, &x)
|
||||
_ = json.Unmarshal(b, &y)
|
||||
if string(a1) != string(a2) || x.Salt == "" || x.Salt == y.Salt || len(x.Salt) != 32 || !strings.Contains(string(a1), `"temperature":0`) {
|
||||
t.Errorf("salts %s / %s", a1, b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheLastUserMessageIsTheQuestion(t *testing.T) {
|
||||
body := `{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"a"},
|
||||
{"role":"user","content":[{"type":"text","text":"yıllık izin"},{"type":"image_url","image_url":{"url":"x"}}]}]}`
|
||||
if got := lastUserText([]byte(body)); strings.TrimSpace(got) != "yıllık izin" {
|
||||
t.Errorf("question %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledgeOnSQLiteSaysWhy(t *testing.T) {
|
||||
e := newFakeEngine(t, "qwen2.5-14b-instruct")
|
||||
rt, secret := pricedRouter(t, e.srv.URL)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
|
||||
strings.NewReader(`{"model":"q14","messages":[{"role":"user","content":"hi"}]}`))
|
||||
req.Header.Set("Authorization", "Bearer "+secret)
|
||||
req.Header.Set(knowledgeHeader, "docs")
|
||||
rec := httptest.NewRecorder()
|
||||
rt.Handler().ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusNotImplemented || !strings.Contains(rec.Body.String(), "SQLite") || e.calls.Load() != 0 {
|
||||
t.Errorf("got %d %s, engine asked %d times", rec.Code, rec.Body.String(), e.calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAChatIsAnsweredFromItsCollections(t *testing.T) {
|
||||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||||
if dsn == "" {
|
||||
t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
|
||||
}
|
||||
e := newFakeEngine(t, "qwen2.5-14b-instruct")
|
||||
rt := kbRouterWith(t, e.srv.URL)
|
||||
ctx := t.Context()
|
||||
s := rt.store
|
||||
sfx := strings.ToLower(store.NewID(""))
|
||||
org, _ := s.CreateOrg(ctx, "rag-"+sfx)
|
||||
team, _ := s.CreateTeam(ctx, org, "t")
|
||||
proj, _ := s.CreateProject(ctx, team, "p")
|
||||
other, _ := s.CreateOrg(ctx, "other-"+sfx)
|
||||
oteam, _ := s.CreateTeam(ctx, other, "t")
|
||||
oproj, _ := s.CreateProject(ctx, oteam, "p")
|
||||
if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj,
|
||||
Name: "politika-" + sfx, EmbedModel: "qwen2.5-14b-instruct", Dims: 3}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
key, otherKey := kbKey(t, s, proj, ""), kbKey(t, s, oproj, "")
|
||||
srv := httptest.NewServer(rt.Handler())
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
dir := t.TempDir()
|
||||
_ = os.WriteFile(filepath.Join(dir, "izin.md"), []byte("# Yıllık izin\n\nÇalışanlar yılda yirmi gün izin kullanır.\n"), 0o644)
|
||||
_ = os.WriteFile(filepath.Join(dir, "masraf.md"), []byte("# Masraf\n\nFişler ayın beşine kadar verilir.\n"), 0o644)
|
||||
h, _ := kbsync.NewHTTP("")
|
||||
l, _ := kbsync.OpenLedger(ctx, filepath.Join(t.TempDir(), "l.db"))
|
||||
defer l.Close()
|
||||
if rep, err := kbsync.Sync(ctx, &kbsync.Client{Base: srv.URL, Key: key, HTTP: h}, l,
|
||||
kbsync.Options{Collection: "politika-" + sfx, Roots: []string{dir}}, nil); err != nil || rep.Sent != 2 {
|
||||
t.Fatalf("sync %+v %v", rep, err)
|
||||
}
|
||||
|
||||
chat := func(key, collections string) *http.Response {
|
||||
t.Helper()
|
||||
req, _ := http.NewRequest(http.MethodPost, srv.URL+"/v1/chat/completions", strings.NewReader(
|
||||
`{"model":"qwen2.5-14b-instruct","messages":[{"role":"system","content":"Kısa yaz."},{"role":"user","content":"kaç gün yıllık izin var"}]}`))
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if collections != "" {
|
||||
req.Header.Set(knowledgeHeader, collections)
|
||||
}
|
||||
req.Header.Set(knowledgeKHeader, "1")
|
||||
res, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
res.Body.Close()
|
||||
return res
|
||||
}
|
||||
res := chat(key, "politika-"+sfx)
|
||||
if res.StatusCode != http.StatusOK || res.Header.Get(sourcesHeader) != "izin.md#0" {
|
||||
t.Fatalf("answered %d, sources %q", res.StatusCode, res.Header.Get(sourcesHeader))
|
||||
}
|
||||
sent, _ := e.last.Load().([]byte)
|
||||
var body struct {
|
||||
Messages []struct {
|
||||
Role, Content string
|
||||
}
|
||||
}
|
||||
_ = json.Unmarshal(sent, &body)
|
||||
if len(body.Messages) != 2 || body.Messages[0].Role != "system" ||
|
||||
!strings.HasPrefix(body.Messages[0].Content, "Kısa yaz.\n\nBelow are passages") ||
|
||||
!strings.Contains(body.Messages[0].Content, "[1] izin.md\nYıllık izin\n\nÇalışanlar yılda yirmi gün") ||
|
||||
strings.Contains(body.Messages[0].Content, "Fişler") {
|
||||
t.Errorf("the model was sent:\n%s", sent)
|
||||
}
|
||||
|
||||
// The engine here is no vLLM: no salt is sent to it.
|
||||
if strings.Contains(string(sent), "cache_salt") {
|
||||
t.Errorf("a salt was sent to an engine that is not vLLM:\n%s", sent)
|
||||
}
|
||||
|
||||
// A project default answers a request that names nothing, and none
|
||||
// turns it off.
|
||||
if err := s.SetKBDefault(ctx, store.KBDefault{Scope: store.ScopeProject, ScopeID: proj,
|
||||
Collections: []string{mustID(t, s, "politika-"+sfx)}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if res := chat(key, ""); res.Header.Get(sourcesHeader) != "izin.md#0" {
|
||||
t.Errorf("by default: sources %q", res.Header.Get(sourcesHeader))
|
||||
}
|
||||
if res := chat(key, "none"); res.Header.Get(sourcesHeader) != "" {
|
||||
t.Errorf("with none: sources %q", res.Header.Get(sourcesHeader))
|
||||
}
|
||||
sent, _ = e.last.Load().([]byte)
|
||||
if strings.Contains(string(sent), "Below are passages") {
|
||||
t.Errorf("none still added passages:\n%s", sent)
|
||||
}
|
||||
|
||||
// Another organization's key names it: as if it did not exist, and the
|
||||
// model is not asked.
|
||||
calls := e.calls.Load()
|
||||
if res := chat(otherKey, "politika-"+sfx); res.StatusCode != http.StatusNotFound || e.calls.Load() != calls {
|
||||
t.Errorf("another organization: %d, engine asked %d more times", res.StatusCode, e.calls.Load()-calls)
|
||||
}
|
||||
}
|
||||
|
||||
func mustID(t *testing.T, s *store.DB, name string) string {
|
||||
t.Helper()
|
||||
cs, err := s.Collections(t.Context())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, c := range cs {
|
||||
if c.Name == name {
|
||||
return c.ID
|
||||
}
|
||||
}
|
||||
t.Fatalf("no collection %s", name)
|
||||
return ""
|
||||
}
|
||||
|
||||
func TestAVLLMReplicaGetsTheOrganizationsSaltOnlyWithPassages(t *testing.T) {
|
||||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||||
if dsn == "" {
|
||||
t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
|
||||
}
|
||||
e := newFakeEngine(t, "qwen2.5-14b-instruct")
|
||||
cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Engine: "vllm", Models: []string{"qwen2.5-14b-instruct"}})
|
||||
cfg.Auth.Enabled = true
|
||||
cfg.Auth.Store = dsn
|
||||
rt := newTestRouter(t, cfg)
|
||||
if err := rt.store.KnowledgeAvailable(); err != nil {
|
||||
t.Skip(err)
|
||||
}
|
||||
ctx := t.Context()
|
||||
s := rt.store
|
||||
sfx := strings.ToLower(store.NewID(""))
|
||||
org, _ := s.CreateOrg(ctx, "salt-"+sfx)
|
||||
team, _ := s.CreateTeam(ctx, org, "t")
|
||||
proj, _ := s.CreateProject(ctx, team, "p")
|
||||
c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: proj,
|
||||
Name: "c-" + sfx, EmbedModel: "qwen2.5-14b-instruct", Dims: 3})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := s.PutDocument(ctx, c, store.KBDocument{Name: "a.md", Digest: "d"},
|
||||
[]store.KBChunk{{Body: "yıllık izin yirmi gün", Embedding: []float32{1, 1, 1}}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
key := kbKey(t, s, proj, "")
|
||||
ask := func(knowledge string) string {
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(
|
||||
`{"model":"qwen2.5-14b-instruct","messages":[{"role":"user","content":"yıllık izin"}]}`))
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if knowledge != "" {
|
||||
req.Header.Set(knowledgeHeader, knowledge)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
rt.Handler().ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
b, _ := e.last.Load().([]byte)
|
||||
return string(b)
|
||||
}
|
||||
if got := ask(c.Name); !strings.Contains(got, `"cache_salt":"`) {
|
||||
t.Errorf("no salt with passages:\n%s", got)
|
||||
}
|
||||
if got := ask(""); strings.Contains(got, "cache_salt") {
|
||||
t.Errorf("a salt without passages:\n%s", got)
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"net/netip"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
@@ -171,6 +172,15 @@ type Router struct {
|
||||
// none is configured, which leaves the charts off the page.
|
||||
history web.History
|
||||
histWindow time.Duration
|
||||
|
||||
// kbDefCache remembers each caller's default knowledge-base
|
||||
// collections for kbDefaultTTL.
|
||||
kbDefMu sync.Mutex
|
||||
kbDefCache map[string]kbDefaultEntry
|
||||
|
||||
// settingsCheck and settingsLocal mount the portal's settings page.
|
||||
settingsCheck func([]byte) error
|
||||
settingsLocal []string
|
||||
}
|
||||
|
||||
// SetClusterView gives the router the data the /admin page renders.
|
||||
@@ -194,6 +204,13 @@ func (r *Router) SetGuard(g guard.Guard) { r.guard = g }
|
||||
// SetGuardLibrary is the policies the guard was built from, for the pages.
|
||||
func (r *Router) SetGuardLibrary(c func() guard.Config) { r.guardLibrary = c }
|
||||
|
||||
// SetSettingsCheck enables the portal's settings page: check says whether a
|
||||
// document would load as the cluster's settings beside this machine's own
|
||||
// server.yaml, and local are the keys that file keeps.
|
||||
func (r *Router) SetSettingsCheck(check func([]byte) error, local []string) {
|
||||
r.settingsCheck, r.settingsLocal = check, local
|
||||
}
|
||||
|
||||
// SetStopped refuses, or again serves, every request for a model.
|
||||
func (r *Router) SetStopped(stopped bool) { r.stopped.Store(stopped) }
|
||||
|
||||
@@ -1073,6 +1090,14 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
|
||||
if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" {
|
||||
body = withAdapter(body, model, ad.index, ad.scale)
|
||||
}
|
||||
// Passages of one organization's documents never share cached
|
||||
// blocks with another's: a prefix found warm would say what another
|
||||
// tenant asked about, by how fast its answer started. vLLM keys its
|
||||
// prefix cache by this salt; other engines are not sent a field
|
||||
// they may refuse.
|
||||
if req.kbOrg != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" {
|
||||
body = withCacheSalt(body, req.kbOrg)
|
||||
}
|
||||
// What computes the answer, said before the proxy writes the headers.
|
||||
// A model can have replicas on cards and on CPUs at once, and the two
|
||||
// differ tenfold in speed: a reply that came slowly is otherwise
|
||||
|
||||
@@ -66,6 +66,20 @@ func newFakeEngine(t *testing.T, model string) *fakeEngine {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`, model)
|
||||
})
|
||||
// Three dimensions: the input's length in bytes, its number of words,
|
||||
// and 1.
|
||||
mux.HandleFunc("POST /v1/embeddings", func(w http.ResponseWriter, r *http.Request) {
|
||||
var in struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&in)
|
||||
var rows []string
|
||||
for i, s := range in.Input {
|
||||
rows = append(rows, fmt.Sprintf(`{"index":%d,"embedding":[%d,%d,1]}`, i, len(s), len(strings.Fields(s))))
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprintf(w, `{"model":%q,"data":[%s],"usage":{"prompt_tokens":%d}}`, model, strings.Join(rows, ","), len(in.Input))
|
||||
})
|
||||
f.srv = httptest.NewServer(mux)
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f
|
||||
|
||||
@@ -214,6 +214,12 @@ func (s *DB) deleteEmpty(ctx context.Context, scope Scope, id, del string, check
|
||||
busy = append(busy, plural(n, ch.noun))
|
||||
}
|
||||
}
|
||||
// Not an emptyCheck: the table is there only on PostgreSQL.
|
||||
if n, err := s.kbOwnedBy(ctx, c, scope, id); err != nil {
|
||||
return err
|
||||
} else if n > 0 {
|
||||
busy = append(busy, plural(n, "collection"))
|
||||
}
|
||||
if len(busy) > 0 {
|
||||
return fmt.Errorf("%w: it still has %s", ErrNotEmpty, strings.Join(busy, " and "))
|
||||
}
|
||||
@@ -222,6 +228,9 @@ func (s *DB) deleteEmpty(ctx context.Context, scope Scope, id, del string, check
|
||||
string(scope), string(MemberScope(scope)), id); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.kbForgetScope(ctx, c, scope, id); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := c.ExecContext(ctx, del, id)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -229,7 +238,7 @@ func (s *DB) deleteEmpty(ctx context.Context, scope Scope, id, del string, check
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return ErrNotFound
|
||||
}
|
||||
return nil
|
||||
return s.kbDropPartition(ctx, c, scope, id)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -248,6 +257,10 @@ func (s *DB) DeleteUser(ctx context.Context, userID string, at time.Time) error
|
||||
millis(at), userID); err != nil {
|
||||
return err
|
||||
}
|
||||
// A person's own collections go with the account.
|
||||
if err := s.kbForgetScope(ctx, c, ScopeUser, userID); err != nil {
|
||||
return err
|
||||
}
|
||||
res, err := c.ExecContext(ctx, `DELETE FROM users WHERE id = ?`, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,446 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"math"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestKnowledgeIsOffOnSQLite(t *testing.T) {
|
||||
s := openDB(t)
|
||||
tree(t, s)
|
||||
err := s.KnowledgeAvailable()
|
||||
if !errors.Is(err, ErrKnowledgeOff) || !strings.Contains(err.Error(), "SQLite") {
|
||||
t.Fatalf("SQLite store says %v", err)
|
||||
}
|
||||
if _, err := s.CreateCollection(context.Background(), KBCollection{
|
||||
OwnerScope: ScopeProject, OwnerID: "proj1", Name: "x", EmbedModel: "m", Dims: 3,
|
||||
}); !errors.Is(err, ErrKnowledgeOff) {
|
||||
t.Errorf("creating a collection on SQLite gave %v", err)
|
||||
}
|
||||
// Deleting from the tree still works without the tables.
|
||||
if err := s.DeleteProject(context.Background(), "proj1"); err != nil {
|
||||
t.Errorf("deleting a project on SQLite: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// kbTwoOrgs is org1 (team1/proj1, team2/proj2) and org2 (team3/proj3).
|
||||
func kbTwoOrgs(t *testing.T, s *DB) {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
tree(t, s)
|
||||
for _, err := range []error{
|
||||
s.PutTeam(ctx, Team{ID: "team2", OrgID: "org1", Name: "Finance"}),
|
||||
s.PutProject(ctx, Project{ID: "proj2", TeamID: "team2", Name: "Payroll"}),
|
||||
s.PutOrg(ctx, Org{ID: "org2", Name: "Other"}),
|
||||
s.PutTeam(ctx, Team{ID: "team3", OrgID: "org2", Name: "Research"}),
|
||||
s.PutProject(ctx, Project{ID: "proj3", TeamID: "team3", Name: "Chat"}),
|
||||
} {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// unit is a 3-dimensional unit vector pointing mostly along axis i.
|
||||
func unit(i int) []float32 {
|
||||
v := []float32{0.1, 0.1, 0.1}
|
||||
v[i] = 1
|
||||
n := float32(math.Sqrt(float64(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])))
|
||||
for j := range v {
|
||||
v[j] /= n
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func mustCollection(t *testing.T, s *DB, scope Scope, owner, org, name string) KBCollection {
|
||||
t.Helper()
|
||||
c, err := s.CreateCollection(context.Background(), KBCollection{
|
||||
OwnerScope: scope, OwnerID: owner, OrgID: org, Name: name, EmbedModel: "emb", Dims: 3,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("creating %s: %v", name, err)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func mustDoc(t *testing.T, s *DB, c KBCollection, name string, bodies ...string) {
|
||||
t.Helper()
|
||||
var chunks []KBChunk
|
||||
for i, b := range bodies {
|
||||
chunks = append(chunks, KBChunk{Body: b, Embedding: unit(i % 3)})
|
||||
}
|
||||
if _, _, err := s.PutDocument(context.Background(), c, KBDocument{Name: name, Digest: name + "-1", Bytes: 10}, chunks); err != nil {
|
||||
t.Fatalf("adding %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
|
||||
func openKB(t *testing.T) *DB {
|
||||
t.Helper()
|
||||
s := openPG(t)
|
||||
if err := s.KnowledgeAvailable(); err != nil {
|
||||
t.Skipf("the test database has no pgvector: %v", err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func names(cs []KBCollection) string {
|
||||
var out []string
|
||||
for _, c := range cs {
|
||||
out = append(out, c.Name)
|
||||
}
|
||||
return strings.Join(out, ",")
|
||||
}
|
||||
|
||||
func TestPostgresCollectionsReachOnlyTheirChain(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
mustCollection(t, s, ScopeOrg, "org1", "", "handbook")
|
||||
mustCollection(t, s, ScopeTeam, "team1", "", "research")
|
||||
pay := mustCollection(t, s, ScopeProject, "proj2", "", "payroll")
|
||||
mustCollection(t, s, ScopeProject, "proj3", "", "other-org")
|
||||
|
||||
read := func(e Effective) string {
|
||||
t.Helper()
|
||||
got, err := s.CollectionsFor(ctx, e, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return names(got)
|
||||
}
|
||||
inProj1 := Effective{KeyID: "k1", ProjectID: "proj1", TeamID: "team1", OrgID: "org1"}
|
||||
if got := read(inProj1); got != "handbook,research" {
|
||||
t.Errorf("proj1 reads %q", got)
|
||||
}
|
||||
if got := read(Effective{KeyID: "k3", ProjectID: "proj3", TeamID: "team3", OrgID: "org2"}); got != "other-org" {
|
||||
t.Errorf("org2 reads %q", got)
|
||||
}
|
||||
// A personal key with no project reads nothing, whatever its person
|
||||
// belongs to.
|
||||
if got := read(Effective{KeyID: "kp", UserID: "u1"}); got != "" {
|
||||
t.Errorf("a caller outside any organization reads %q", got)
|
||||
}
|
||||
|
||||
// A grant to a key adds that collection for that key alone.
|
||||
_, k := mustKey(t, s, "k1", "proj1")
|
||||
if err := s.Grant(ctx, KBGrant{Scope: ScopeKey, ScopeID: k.ID, Collection: pay.ID, Perm: KBRead}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := read(inProj1); got != "handbook,payroll,research" {
|
||||
t.Errorf("proj1 with a key grant reads %q", got)
|
||||
}
|
||||
if got := read(Effective{KeyID: "k-other", ProjectID: "proj1", TeamID: "team1", OrgID: "org1"}); got != "handbook,research" {
|
||||
t.Errorf("another key of proj1 reads %q", got)
|
||||
}
|
||||
// Read is not write.
|
||||
w, err := s.CollectionsFor(ctx, inProj1, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(names(w), "payroll") {
|
||||
t.Errorf("a read grant gave write: %s", names(w))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresGrantsNeverCrossOrganizations(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
c := mustCollection(t, s, ScopeProject, "proj1", "", "docs")
|
||||
for _, g := range []KBGrant{
|
||||
{Scope: ScopeOrg, ScopeID: "org2"},
|
||||
{Scope: ScopeTeam, ScopeID: "team3"},
|
||||
{Scope: ScopeProject, ScopeID: "proj3"},
|
||||
} {
|
||||
g.Collection, g.Perm = c.ID, KBRead
|
||||
if err := s.Grant(ctx, g); err == nil || !strings.Contains(err.Error(), "another organization") {
|
||||
t.Errorf("granting to %s %s gave %v", g.Scope, g.ScopeID, err)
|
||||
}
|
||||
}
|
||||
_, k3 := mustKey(t, s, "k3", "proj3")
|
||||
if err := s.Grant(ctx, KBGrant{Scope: ScopeKey, ScopeID: k3.ID, Collection: c.ID, Perm: KBRead}); err == nil {
|
||||
t.Error("a key of another organization was granted a collection")
|
||||
}
|
||||
// A person who belongs to nothing in the organization cannot own or be
|
||||
// granted one there.
|
||||
if err := s.CreateUser(ctx, User{ID: "u1", Email: "u1@example.com", Name: "u", PasswordHash: "x",
|
||||
State: UserActive, CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Grant(ctx, KBGrant{Scope: ScopeUser, ScopeID: "u1", Collection: c.ID, Perm: KBRead}); err == nil {
|
||||
t.Error("a person with no membership in the organization was granted a collection")
|
||||
}
|
||||
if err := s.AddMembership(ctx, "u1", ScopeProject, "proj2", "member", "t", time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Grant(ctx, KBGrant{Scope: ScopeUser, ScopeID: "u1", Collection: c.ID, Perm: KBRead}); err != nil {
|
||||
t.Errorf("a member of a project in the organization: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresSearchReturnsOnlyAllowedChunks(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
mine := mustCollection(t, s, ScopeProject, "proj1", "", "mine")
|
||||
secret := mustCollection(t, s, ScopeProject, "proj2", "", "secret")
|
||||
foreign := mustCollection(t, s, ScopeProject, "proj3", "", "foreign")
|
||||
mustDoc(t, s, mine, "a.md", "vacation policy is twenty days", "expenses go through finance")
|
||||
mustDoc(t, s, secret, "salaries.md", "vacation salaries are confidential")
|
||||
mustDoc(t, s, foreign, "theirs.md", "vacation at the other company")
|
||||
|
||||
hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{mine},
|
||||
Embeddings: map[string][]float32{"emb": unit(0)}, Text: "vacation policy", K: 5})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) == 0 || hits[0].Body != "vacation policy is twenty days" || hits[0].Document != "a.md" {
|
||||
t.Fatalf("best hit %+v", hits)
|
||||
}
|
||||
for _, h := range hits {
|
||||
if h.Collection != mine.ID {
|
||||
t.Errorf("a search of %s returned a chunk of %s: %q", mine.Name, h.Collection, h.Body)
|
||||
}
|
||||
}
|
||||
|
||||
// A collection of another organization is refused, not searched.
|
||||
if _, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{mine, foreign},
|
||||
Embeddings: map[string][]float32{"emb": unit(0)}, Text: "vacation", K: 5}); err == nil {
|
||||
t.Error("a search across organizations was run")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresRowSecurityHoldsWithoutTheFilter(t *testing.T) {
|
||||
// The second line: a query that forgets its WHERE gets nothing outside
|
||||
// the scope its transaction set, and nothing at all without one.
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
mine := mustCollection(t, s, ScopeProject, "proj1", "", "mine")
|
||||
other := mustCollection(t, s, ScopeProject, "proj2", "", "other")
|
||||
mustDoc(t, s, mine, "a.md", "one", "two")
|
||||
mustDoc(t, s, other, "b.md", "three")
|
||||
|
||||
count := func(org string, colls ...string) int {
|
||||
t.Helper()
|
||||
var n int
|
||||
err := s.inTx(ctx, func(c conn) error {
|
||||
if org != "" {
|
||||
if err := kbScope(ctx, c, org, colls); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return c.QueryRowContext(ctx, `SELECT COUNT(*) FROM kb_chunks`).Scan(&n)
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
if n := count(""); n != 0 {
|
||||
t.Errorf("with no scope set, %d chunks are visible", n)
|
||||
}
|
||||
if n := count("org1", mine.ID); n != 2 {
|
||||
t.Errorf("scoped to one collection, %d chunks are visible", n)
|
||||
}
|
||||
if n := count("org2", mine.ID); n != 0 {
|
||||
t.Errorf("a collection named under the wrong organization shows %d chunks", n)
|
||||
}
|
||||
// Nor can a write put a chunk outside its scope.
|
||||
err := s.inTx(ctx, func(c conn) error {
|
||||
if err := kbScope(ctx, c, "org1", []string{mine.ID}); err != nil {
|
||||
return err
|
||||
}
|
||||
docs, _ := s.Documents(ctx, other.ID)
|
||||
_, err := c.ExecContext(ctx, `INSERT INTO kb_chunks (org_id, collection, doc_id, ord, body, embedding)
|
||||
VALUES ('org1', ?, ?, 9, 'x', '[1,0,0]')`, other.ID, docs[0].ID)
|
||||
return err
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "row-level security") {
|
||||
t.Errorf("writing into another collection gave %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresDocumentsReplaceWholeAndSkipTheSame(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
c := mustCollection(t, s, ScopeProject, "proj1", "", "docs")
|
||||
put := func(digest string, bodies ...string) bool {
|
||||
t.Helper()
|
||||
var chunks []KBChunk
|
||||
for _, b := range bodies {
|
||||
chunks = append(chunks, KBChunk{Body: b, Embedding: unit(0)})
|
||||
}
|
||||
_, changed, err := s.PutDocument(ctx, c, KBDocument{Name: "a.md", Digest: digest}, chunks)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return changed
|
||||
}
|
||||
if !put("d1", "old one", "old two") {
|
||||
t.Error("a new document was not a change")
|
||||
}
|
||||
if put("d1", "ignored") {
|
||||
t.Error("the same digest again was a change")
|
||||
}
|
||||
if !put("d2", "new") {
|
||||
t.Error("a new digest was not a change")
|
||||
}
|
||||
hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{c},
|
||||
Embeddings: map[string][]float32{"emb": unit(0)}, K: 10})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) != 1 || hits[0].Body != "new" {
|
||||
t.Errorf("after replacing, the collection holds %+v", hits)
|
||||
}
|
||||
if _, _, err := s.PutDocument(ctx, c, KBDocument{Name: "b.md", Digest: "x"},
|
||||
[]KBChunk{{Body: "x", Embedding: []float32{1, 0}}}); err == nil {
|
||||
t.Error("a chunk of the wrong dimension was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresScopesWithCollectionsAreNotDeleted(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
c := mustCollection(t, s, ScopeProject, "proj2", "", "docs")
|
||||
mustDoc(t, s, c, "a.md", "x")
|
||||
if err := s.DeleteProject(ctx, "proj2"); !errors.Is(err, ErrNotEmpty) || !strings.Contains(err.Error(), "collection") {
|
||||
t.Errorf("deleting a project with a collection gave %v", err)
|
||||
}
|
||||
if err := s.DeleteCollection(ctx, c.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.DeleteProject(ctx, "proj2"); err != nil {
|
||||
t.Errorf("deleting the project once its collection is gone: %v", err)
|
||||
}
|
||||
var n int
|
||||
if err := s.raw.QueryRowContext(ctx, `SELECT COUNT(*) FROM pg_indexes WHERE indexname = $1`,
|
||||
kbVectorIndex(c.ID)).Scan(&n); err != nil || n != 0 {
|
||||
t.Errorf("the collection's index is left behind (%d, %v)", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresAPersonsCollectionsGoWithTheAccount(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
if err := s.CreateUser(ctx, User{ID: "u1", Email: "u1@example.com", Name: "u", PasswordHash: "x",
|
||||
State: UserActive, CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.CreateCollection(ctx, KBCollection{OwnerScope: ScopeUser, OwnerID: "u1", OrgID: "org1",
|
||||
Name: "notes", EmbedModel: "emb", Dims: 3}); err == nil {
|
||||
t.Error("a person outside the organization made a collection in it")
|
||||
}
|
||||
if err := s.AddMembership(ctx, "u1", ScopeTeam, "team1", "member", "t", time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c := mustCollection(t, s, ScopeUser, "u1", "org1", "notes")
|
||||
mustDoc(t, s, c, "n.md", "x")
|
||||
if err := s.DeleteUser(ctx, "u1", time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.Collection(ctx, c.ID); !errors.Is(err, ErrNotFound) {
|
||||
t.Errorf("the person's collection outlived the account: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresMeaningOutweighsSharedWords(t *testing.T) {
|
||||
// The PoC's case: the question shares its everyday words with one
|
||||
// passage and its meaning with another.
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
c := mustCollection(t, s, ScopeProject, "proj1", "", "docs")
|
||||
put := func(name, body string, v []float32) {
|
||||
t.Helper()
|
||||
if _, _, err := s.PutDocument(ctx, c, KBDocument{Name: name, Digest: name}, []KBChunk{{Body: body, Embedding: v}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
put("yillik.md", "yıllık izin kadar izin alır izin", unit(0))
|
||||
put("babalik.md", "babalık hakkı on gündür", unit(1))
|
||||
for _, f := range []string{"a", "b", "c"} {
|
||||
put(f+".md", "başka bir konu", unit(2))
|
||||
}
|
||||
hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{c},
|
||||
Embeddings: map[string][]float32{"emb": unit(1)}, Text: "babalar ne kadar izin alır", K: 1})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) != 1 || hits[0].Document != "babalik.md" {
|
||||
t.Errorf("best hit %+v", hits)
|
||||
}
|
||||
// And a question that holds a passage's exact words still finds it,
|
||||
// with nothing embedded at all (the model down).
|
||||
hits, err = s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{c}, Text: "babalık hakkı", K: 1})
|
||||
if err != nil || len(hits) != 1 || hits[0].Document != "babalik.md" || hits[0].Score != textWeight {
|
||||
t.Errorf("text alone: %+v %v", hits, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOrQueryMatchesAnyWord(t *testing.T) {
|
||||
if got := orQuery("Yıllık izin kaç gün? İzin & 'x' -- a"); got != "Yıllık | izin | gün" {
|
||||
t.Errorf("orQuery gave %q", got)
|
||||
}
|
||||
// Words every passage has are no evidence of anything.
|
||||
if got := orQuery("Babalar NE kadar izin alır, ve İÇİN"); got != "Babalar | izin | alır" {
|
||||
t.Errorf("stop words kept: %q", got)
|
||||
}
|
||||
if got := orQuery("?!"); got != "" {
|
||||
t.Errorf("no words gave %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresDefaultsAreTheNarrowestAndStayInTheirOrganization(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
hand := mustCollection(t, s, ScopeOrg, "org1", "", "handbook")
|
||||
team := mustCollection(t, s, ScopeTeam, "team1", "", "research")
|
||||
foreign := mustCollection(t, s, ScopeOrg, "org2", "", "theirs")
|
||||
if err := s.SetKBDefault(ctx, KBDefault{Scope: ScopeProject, ScopeID: "proj1", Collections: []string{foreign.ID}}); err == nil {
|
||||
t.Error("another organization's collection became a default")
|
||||
}
|
||||
for _, d := range []KBDefault{
|
||||
{Scope: ScopeOrg, ScopeID: "org1", Collections: []string{hand.ID}},
|
||||
{Scope: ScopeTeam, ScopeID: "team1", Collections: []string{team.ID, hand.ID}},
|
||||
} {
|
||||
if err := s.SetKBDefault(ctx, d); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
in := func(e Effective) string {
|
||||
got, err := s.KBDefaultFor(ctx, e)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return strings.Join(got, ",")
|
||||
}
|
||||
if got := in(Effective{KeyID: "k", ProjectID: "proj1", TeamID: "team1", OrgID: "org1"}); got != team.ID+","+hand.ID {
|
||||
t.Errorf("in team1: %q", got)
|
||||
}
|
||||
if got := in(Effective{KeyID: "k", ProjectID: "proj2", TeamID: "team2", OrgID: "org1"}); got != hand.ID {
|
||||
t.Errorf("in team2: %q", got)
|
||||
}
|
||||
if got := in(Effective{KeyID: "k", ProjectID: "proj3", TeamID: "team3", OrgID: "org2"}); got != "" {
|
||||
t.Errorf("in org2: %q", got)
|
||||
}
|
||||
// A collection removed leaves every default that named it.
|
||||
if err := s.DeleteCollection(ctx, hand.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := in(Effective{KeyID: "k", ProjectID: "proj1", TeamID: "team1", OrgID: "org1"}); got != team.ID {
|
||||
t.Errorf("after removing one: %q", got)
|
||||
}
|
||||
if got := in(Effective{KeyID: "k", ProjectID: "proj2", TeamID: "team2", OrgID: "org1"}); got != "" {
|
||||
t.Errorf("a default left empty stayed: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
-- Knowledge bases: collections of a company's documents that requests are
|
||||
-- answered from (BACKLOG F). PostgreSQL only: there is no SQLite file of this
|
||||
-- number, and on SQLite the feature is off.
|
||||
--
|
||||
-- The chunks themselves are not here. They need the pgvector extension, which
|
||||
-- only a superuser can create and which a cluster may not have installed; a
|
||||
-- migration that failed for want of it would stop the control plane from
|
||||
-- starting at all. The store creates kb_chunks when it opens, if it can
|
||||
-- (knowledge.go), and says why not when it cannot.
|
||||
|
||||
-- A collection has one owner in the tenancy tree and lives inside one
|
||||
-- organization, which no grant can cross.
|
||||
CREATE TABLE kb_collections (
|
||||
id TEXT PRIMARY KEY,
|
||||
org_id TEXT NOT NULL REFERENCES organizations(id),
|
||||
owner_scope TEXT NOT NULL,
|
||||
owner_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
-- Vectors from two models are not comparable, so a collection is embedded
|
||||
-- by one model and its dimension is fixed with it.
|
||||
embed_model TEXT NOT NULL,
|
||||
dims INTEGER NOT NULL,
|
||||
created_at BIGINT NOT NULL,
|
||||
created_by TEXT NOT NULL DEFAULT '',
|
||||
UNIQUE (org_id, name)
|
||||
);
|
||||
CREATE INDEX kb_collections_by_owner ON kb_collections(owner_scope, owner_id);
|
||||
|
||||
-- Who else may read or write a collection. Like guard_attachments, keyed by
|
||||
-- scope with everything below the scope included; unlike guard policies,
|
||||
-- grants add up rather than narrow.
|
||||
CREATE TABLE kb_grants (
|
||||
scope TEXT NOT NULL,
|
||||
scope_id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL REFERENCES kb_collections(id) ON DELETE CASCADE,
|
||||
perm TEXT NOT NULL,
|
||||
added_at BIGINT NOT NULL,
|
||||
added_by TEXT NOT NULL DEFAULT '',
|
||||
PRIMARY KEY (scope, scope_id, collection)
|
||||
);
|
||||
CREATE INDEX kb_grants_by_collection ON kb_grants(collection);
|
||||
|
||||
-- A document as added, by name within its collection. The digest is what
|
||||
-- decides whether adding it again changes anything.
|
||||
CREATE TABLE kb_documents (
|
||||
id TEXT PRIMARY KEY,
|
||||
collection TEXT NOT NULL REFERENCES kb_collections(id) ON DELETE CASCADE,
|
||||
name TEXT NOT NULL,
|
||||
digest TEXT NOT NULL,
|
||||
bytes BIGINT NOT NULL,
|
||||
chunks INTEGER NOT NULL,
|
||||
added_at BIGINT NOT NULL,
|
||||
added_by TEXT NOT NULL DEFAULT '',
|
||||
UNIQUE (collection, name)
|
||||
);
|
||||
@@ -0,0 +1,12 @@
|
||||
-- The collections a part of the tree answers from when a request names
|
||||
-- none (BACKLOG F3): on an organization, team, project or key, the narrowest
|
||||
-- one that has any winning. PostgreSQL only, as the collections are.
|
||||
CREATE TABLE kb_defaults (
|
||||
scope TEXT NOT NULL,
|
||||
scope_id TEXT NOT NULL,
|
||||
-- Collection IDs, comma separated, in the order given.
|
||||
collections TEXT NOT NULL,
|
||||
added_at BIGINT NOT NULL,
|
||||
added_by TEXT NOT NULL DEFAULT '',
|
||||
PRIMARY KEY (scope, scope_id)
|
||||
);
|
||||
@@ -55,9 +55,13 @@ func OpenPostgres(ctx context.Context, dsn string) (*DB, error) {
|
||||
_ = db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &DB{
|
||||
s := &DB{
|
||||
raw: db,
|
||||
dialect: PostgresDialect,
|
||||
db: bound{c: PostgresDialect, to: db},
|
||||
}, nil
|
||||
}
|
||||
// Not fatal: a store without pgvector serves everything but knowledge
|
||||
// bases, and says why when somebody asks for one.
|
||||
s.kbErr = s.ensureKnowledge(ctx)
|
||||
return s, nil
|
||||
}
|
||||
@@ -48,6 +48,10 @@ type DB struct {
|
||||
// personal is what a person's own use is held to when nobody wrote a
|
||||
// limit for them; see SetPersonalDefaults.
|
||||
personal *Limits
|
||||
|
||||
// kbErr is why this store cannot hold knowledge bases, nil when it can;
|
||||
// set once at open (knowledge.go).
|
||||
kbErr error
|
||||
}
|
||||
|
||||
// OpenSQLite opens (creating if needed) a database and brings it up to date.
|
||||
|
||||
@@ -90,7 +90,9 @@ type AccountPages struct {
|
||||
audit audit.Sink
|
||||
// hasGuard is set when /admin/guard is registered, and guardLibrary
|
||||
// with it: the policies in force, whose judges the mirror page serves.
|
||||
hasGuard bool
|
||||
hasGuard bool
|
||||
// hasSettings is set when /admin/settings is registered.
|
||||
hasSettings bool
|
||||
guardLibrary func() guard.Config
|
||||
// guardFor names the guardrail policies in force for someone, and
|
||||
// whether they keep nothing; nil when there are none.
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
// The cluster's settings, edited in the browser (BACKLOG E4): the document
|
||||
// `llmgw ctl config edit` opens, for whoever has no terminal on a control
|
||||
// plane. The same rules hold here as there — a version is checked as a
|
||||
// control plane would load it before it is kept, it is made from the newest
|
||||
// version or refused, and every change is in force on every control plane
|
||||
// within seconds and in the audit log.
|
||||
|
||||
// SettingsAdmin is where the cluster's settings are kept.
|
||||
type SettingsAdmin interface {
|
||||
ClusterSettings(ctx context.Context) (store.SettingsVersion, error)
|
||||
ClusterSettingsAt(ctx context.Context, version int64) (store.SettingsVersion, error)
|
||||
ClusterSettingsHistory(ctx context.Context, limit int) ([]store.SettingsVersion, error)
|
||||
PutClusterSettings(ctx context.Context, base int64, body, by, note string) (int64, error)
|
||||
}
|
||||
|
||||
type settingsPageData struct {
|
||||
Nav Nav
|
||||
Version string
|
||||
CSRF string
|
||||
// Current is the version in force; Shown the one in the editor, which
|
||||
// is an earlier one when ?version= asked for it.
|
||||
Current, Shown store.SettingsVersion
|
||||
Earlier bool
|
||||
// Body is what the editor holds: the version shown, or what was sent
|
||||
// and refused, so a refusal does not lose the edit.
|
||||
Body string
|
||||
Note string
|
||||
History []settingsRow
|
||||
LocalKeys []string
|
||||
Notice string
|
||||
Error string
|
||||
}
|
||||
|
||||
type settingsRow struct {
|
||||
Version int64
|
||||
At, By string
|
||||
Note string
|
||||
InForce bool
|
||||
Selected bool
|
||||
}
|
||||
|
||||
// RegisterAdminSettings mounts /admin/settings, for administrators only:
|
||||
// the settings hold webhook addresses and the like, so even reading them is
|
||||
// not every staff role's.
|
||||
func (p *AccountPages) RegisterAdminSettings(mux *http.ServeMux, s SettingsAdmin, check func([]byte) error, localKeys []string) {
|
||||
p.hasSettings = true
|
||||
page := func(w http.ResponseWriter, r *http.Request, admin store.User, d settingsPageData) {
|
||||
ctx, l := r.Context(), langOf(r)
|
||||
d.Nav, d.Version, d.CSRF, d.LocalKeys = p.navFor(r, admin, "settings"), p.version, csrfFor(r), localKeys
|
||||
cur, err := s.ClusterSettings(ctx)
|
||||
if err != nil {
|
||||
d.Error = l.Err(err)
|
||||
}
|
||||
d.Current = cur
|
||||
if d.Shown.Version == 0 {
|
||||
d.Shown = cur
|
||||
}
|
||||
d.Earlier = d.Shown.Version != cur.Version
|
||||
if d.Body == "" {
|
||||
d.Body = d.Shown.Body
|
||||
}
|
||||
if hist, err := s.ClusterSettingsHistory(ctx, 30); err == nil {
|
||||
for _, v := range hist {
|
||||
d.History = append(d.History, settingsRow{Version: v.Version, At: l.Date(v.At.UTC()) + " " + v.At.UTC().Format("15:04"),
|
||||
By: v.By, Note: v.Note, InForce: v.Version == cur.Version, Selected: v.Version == d.Shown.Version})
|
||||
}
|
||||
}
|
||||
noStore(w)
|
||||
render(w, r, "admin_settings", d)
|
||||
}
|
||||
adminsOnly := func(next func(w http.ResponseWriter, r *http.Request, admin store.User)) http.HandlerFunc {
|
||||
return p.adminOnly(func(w http.ResponseWriter, r *http.Request, admin store.User) {
|
||||
if !admin.IsAdmin {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if r.Method == http.MethodPost {
|
||||
sess, _, _ := p.session(r)
|
||||
if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
|
||||
http.Error(w, "bad request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
}
|
||||
next(w, r, admin)
|
||||
})
|
||||
}
|
||||
|
||||
mux.HandleFunc("GET /admin/settings", adminsOnly(func(w http.ResponseWriter, r *http.Request, admin store.User) {
|
||||
var d settingsPageData
|
||||
if q := r.URL.Query().Get("version"); q != "" {
|
||||
n, err := strconv.ParseInt(q, 10, 64)
|
||||
if v, verr := s.ClusterSettingsAt(r.Context(), n); err == nil && verr == nil && v.Version != 0 {
|
||||
d.Shown = v
|
||||
} else {
|
||||
d.Error = langOf(r).T("settings.err.no_version")
|
||||
}
|
||||
}
|
||||
page(w, r, admin, d)
|
||||
}))
|
||||
|
||||
mux.HandleFunc("POST /admin/settings", adminsOnly(func(w http.ResponseWriter, r *http.Request, admin store.User) {
|
||||
l := langOf(r)
|
||||
body := strings.ReplaceAll(r.FormValue("body"), "\r\n", "\n")
|
||||
note := strings.TrimSpace(r.FormValue("note"))
|
||||
base, _ := strconv.ParseInt(r.FormValue("base"), 10, 64)
|
||||
// The audit line records the form; the document can hold webhook
|
||||
// secrets, so it records the document's size, not the document.
|
||||
r.PostForm.Set("body", "("+strconv.Itoa(len(body))+" bytes)")
|
||||
refuse := func(msg string) {
|
||||
noteAudit(r, "cluster settings", "", msg)
|
||||
page(w, r, admin, settingsPageData{Body: body, Note: note, Error: msg})
|
||||
}
|
||||
if cur, err := s.ClusterSettings(r.Context()); err == nil && cur.Body == body {
|
||||
refuse(l.T("settings.err.unchanged"))
|
||||
return
|
||||
}
|
||||
if err := check([]byte(body)); err != nil {
|
||||
refuse(l.T("settings.err.invalid") + " " + err.Error())
|
||||
return
|
||||
}
|
||||
n, err := s.PutClusterSettings(r.Context(), base, body, admin.Email, note)
|
||||
switch {
|
||||
case errors.Is(err, store.ErrSettingsChanged):
|
||||
refuse(l.T("settings.err.changed"))
|
||||
return
|
||||
case err != nil:
|
||||
refuse(l.Err(err))
|
||||
return
|
||||
}
|
||||
noteAudit(r, "cluster settings", "version "+strconv.FormatInt(n, 10), "")
|
||||
page(w, r, admin, settingsPageData{Notice: l.T("settings.saved", n)})
|
||||
}))
|
||||
|
||||
mux.HandleFunc("POST /admin/settings/rollback", adminsOnly(func(w http.ResponseWriter, r *http.Request, admin store.User) {
|
||||
l := langOf(r)
|
||||
n, _ := strconv.ParseInt(r.FormValue("version"), 10, 64)
|
||||
old, err := s.ClusterSettingsAt(r.Context(), n)
|
||||
if err != nil || old.Version == 0 {
|
||||
page(w, r, admin, settingsPageData{Error: l.T("settings.err.no_version")})
|
||||
return
|
||||
}
|
||||
cur, err := s.ClusterSettings(r.Context())
|
||||
if err != nil {
|
||||
page(w, r, admin, settingsPageData{Error: l.Err(err)})
|
||||
return
|
||||
}
|
||||
if err := check([]byte(old.Body)); err != nil {
|
||||
// A version kept before a change to server.yaml may no longer
|
||||
// load beside it; refused here rather than breaking the cluster.
|
||||
noteAudit(r, "cluster settings", "", err.Error())
|
||||
page(w, r, admin, settingsPageData{Error: l.T("settings.err.invalid") + " " + err.Error()})
|
||||
return
|
||||
}
|
||||
v, err := s.PutClusterSettings(r.Context(), cur.Version, old.Body, admin.Email, "back to version "+strconv.FormatInt(n, 10))
|
||||
if err != nil {
|
||||
page(w, r, admin, settingsPageData{Error: l.Err(err)})
|
||||
return
|
||||
}
|
||||
noteAudit(r, "cluster settings", "version "+strconv.FormatInt(v, 10)+" from "+strconv.FormatInt(n, 10), "")
|
||||
page(w, r, admin, settingsPageData{Notice: l.T("settings.rolled_back", n, v)})
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||||
)
|
||||
|
||||
type fakeSettings struct{ versions []store.SettingsVersion }
|
||||
|
||||
func (f *fakeSettings) ClusterSettings(context.Context) (store.SettingsVersion, error) {
|
||||
if len(f.versions) == 0 {
|
||||
return store.SettingsVersion{}, nil
|
||||
}
|
||||
return f.versions[len(f.versions)-1], nil
|
||||
}
|
||||
|
||||
func (f *fakeSettings) ClusterSettingsAt(_ context.Context, v int64) (store.SettingsVersion, error) {
|
||||
for _, s := range f.versions {
|
||||
if s.Version == v {
|
||||
return s, nil
|
||||
}
|
||||
}
|
||||
return store.SettingsVersion{}, nil
|
||||
}
|
||||
|
||||
func (f *fakeSettings) ClusterSettingsHistory(context.Context, int) ([]store.SettingsVersion, error) {
|
||||
out := make([]store.SettingsVersion, len(f.versions))
|
||||
for i, v := range f.versions {
|
||||
out[len(f.versions)-1-i] = v
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (f *fakeSettings) PutClusterSettings(_ context.Context, base int64, body, by, note string) (int64, error) {
|
||||
cur, _ := f.ClusterSettings(context.Background())
|
||||
if base != cur.Version {
|
||||
return 0, store.ErrSettingsChanged
|
||||
}
|
||||
n := cur.Version + 1
|
||||
f.versions = append(f.versions, store.SettingsVersion{Version: n, Body: body, By: by, Note: note, At: time.Now()})
|
||||
return n, nil
|
||||
}
|
||||
|
||||
type auditSink struct{ events []store.AuditEvent }
|
||||
|
||||
func (a *auditSink) WriteAudit(_ context.Context, e store.AuditEvent) error {
|
||||
a.events = append(a.events, e)
|
||||
return nil
|
||||
}
|
||||
|
||||
func clusterSettingsPage(t *testing.T, user store.User) (*http.ServeMux, *fakeSettings, *auditSink) {
|
||||
t.Helper()
|
||||
s := &fakeSettings{versions: []store.SettingsVersion{{Version: 1, Body: "guard:\n enabled: true\n", By: "init", At: time.Now()}}}
|
||||
p := NewAccountPages(fakeAccounts{user: user}, nil, nil, "v0", slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
sink := &auditSink{}
|
||||
p.SetAudit(sink)
|
||||
mux := http.NewServeMux()
|
||||
p.RegisterAdminSettings(mux, s, func(b []byte) error {
|
||||
if strings.Contains(string(b), "bogus") {
|
||||
return errors.New("unknown key bogus")
|
||||
}
|
||||
return nil
|
||||
}, []string{"listen", "store"})
|
||||
return mux, s, sink
|
||||
}
|
||||
|
||||
func TestTheSettingsPageIsForAdministratorsOnly(t *testing.T) {
|
||||
for _, u := range []store.User{
|
||||
{ID: "u1", Email: "v@example.com", Role: store.RoleViewer},
|
||||
{ID: "u2", Email: "o@example.com", Role: store.RoleOperator},
|
||||
{ID: "u3", Email: "b@example.com", Role: store.RoleBilling},
|
||||
} {
|
||||
mux, _, _ := clusterSettingsPage(t, u)
|
||||
if w := mirrorGet(t, mux, "/admin/settings"); w.Code != http.StatusNotFound {
|
||||
t.Errorf("%s saw the settings: %d", u.Role, w.Code)
|
||||
}
|
||||
}
|
||||
mux, _, _ := clusterSettingsPage(t, store.User{ID: "a", Email: "a@example.com", IsAdmin: true})
|
||||
body := mirrorGet(t, mux, "/admin/settings").Body.String()
|
||||
for _, want := range []string{"guard:\n enabled: true", "Version 1, in force", "<code>listen</code>", `href="/admin/settings" aria-current="page"`} {
|
||||
if !strings.Contains(body, want) {
|
||||
t.Errorf("no %q in:\n%s", want, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsAreCheckedAppliedAndRolledBack(t *testing.T) {
|
||||
mux, s, sink := clusterSettingsPage(t, store.User{ID: "a", Email: "a@example.com", IsAdmin: true})
|
||||
post := func(form url.Values) string {
|
||||
form.Set("csrf", "tok")
|
||||
return mirrorPost(t, mux, "/admin/settings", form).Body.String()
|
||||
}
|
||||
// Refused: does not load. The edit stays in the editor; nothing kept.
|
||||
body := post(url.Values{"base": {"1"}, "body": {"bogus: 1\n"}})
|
||||
if !strings.Contains(body, "unknown key bogus") || !strings.Contains(body, "bogus: 1") || len(s.versions) != 1 {
|
||||
t.Errorf("an invalid document: %d versions\n%s", len(s.versions), body)
|
||||
}
|
||||
// Applied, with its note and who.
|
||||
body = post(url.Values{"base": {"1"}, "body": {"guard:\r\n enabled: false\r\n"}, "note": {"maintenance"}})
|
||||
if !strings.Contains(body, "Version 2 is in force") || len(s.versions) != 2 ||
|
||||
s.versions[1].Body != "guard:\n enabled: false\n" || s.versions[1].By != "a@example.com" || s.versions[1].Note != "maintenance" {
|
||||
t.Errorf("applying: %+v\n%s", s.versions, body)
|
||||
}
|
||||
// Made from a version no longer the newest: refused, not overwritten.
|
||||
body = post(url.Values{"base": {"1"}, "body": {"guard:\n enabled: true\n x: 1\n"}})
|
||||
if !strings.Contains(body, "someone changed the settings") || len(s.versions) != 2 {
|
||||
t.Errorf("a stale edit: %d versions\n%s", len(s.versions), body)
|
||||
}
|
||||
// An earlier version is shown, and put back as a new one.
|
||||
if body := mirrorGet(t, mux, "/admin/settings?version=1").Body.String(); !strings.Contains(body, "Put version 1 back in force") {
|
||||
t.Errorf("viewing version 1:\n%s", body)
|
||||
}
|
||||
w := mirrorPost(t, mux, "/admin/settings/rollback", url.Values{"csrf": {"tok"}, "version": {"1"}})
|
||||
if len(s.versions) != 3 || s.versions[2].Body != s.versions[0].Body || !strings.Contains(w.Body.String(), "as version 3") {
|
||||
t.Errorf("rollback: %+v\n%s", s.versions, w.Body.String())
|
||||
}
|
||||
// Without the CSRF token, nothing.
|
||||
if w := mirrorPost(t, mux, "/admin/settings", url.Values{"base": {"3"}, "body": {"a: 1\n"}}); w.Code != http.StatusBadRequest || len(s.versions) != 3 {
|
||||
t.Errorf("no token: %d", w.Code)
|
||||
}
|
||||
// Audited, without the document itself: it can hold secrets.
|
||||
var saw bool
|
||||
for _, e := range sink.events {
|
||||
if strings.Contains(e.Detail, "enabled") {
|
||||
t.Errorf("the document is in the audit log: %+v", e)
|
||||
}
|
||||
if e.Action == "POST /admin/settings" && strings.Contains(e.Detail, "version 2") {
|
||||
saw = true
|
||||
}
|
||||
}
|
||||
if !saw {
|
||||
t.Errorf("no audit line for version 2: %+v", sink.events)
|
||||
}
|
||||
}
|
||||
+23
-6
@@ -54,24 +54,41 @@ type ChatContext struct {
|
||||
Label string
|
||||
}
|
||||
|
||||
// ChatCollection is a knowledge-base collection the chat can answer from.
|
||||
type ChatCollection struct {
|
||||
ID, Label string
|
||||
}
|
||||
|
||||
type chatData struct {
|
||||
Nav Nav
|
||||
Contexts []ChatContext
|
||||
Context string
|
||||
User store.User
|
||||
Models []ChatModel
|
||||
CSRF string
|
||||
Version string
|
||||
Now string
|
||||
// Collections are what this person may read in the chosen context;
|
||||
// KnowledgeDefault names the ones its part of the tree answers from
|
||||
// when none is chosen, if any.
|
||||
Collections []ChatCollection
|
||||
KnowledgeDefault string
|
||||
User store.User
|
||||
Models []ChatModel
|
||||
CSRF string
|
||||
Version string
|
||||
Now string
|
||||
}
|
||||
|
||||
// ChatKnowledge is what the chat page offers to answer from.
|
||||
type ChatKnowledge struct {
|
||||
Collections []ChatCollection
|
||||
Default string
|
||||
}
|
||||
|
||||
// Chat renders the in-browser chat for one signed-in person.
|
||||
func Chat(user store.User, models []ChatModel, contexts []ChatContext, current, version string) http.HandlerFunc {
|
||||
func Chat(user store.User, models []ChatModel, contexts []ChatContext, current, version string, kb ChatKnowledge) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
noStore(w)
|
||||
render(w, r, "chat", chatData{
|
||||
Nav: navFrom(r, "chat"),
|
||||
Contexts: contexts, Context: current,
|
||||
Collections: kb.Collections, KnowledgeDefault: kb.Default,
|
||||
User: user, Models: models, CSRF: csrfFor(r),
|
||||
Version: version, Now: time.Now().UTC().Format(time.RFC3339),
|
||||
})
|
||||
|
||||
@@ -20,7 +20,7 @@ func chatPage(t *testing.T, names ...string) string {
|
||||
|
||||
func chatPageWith(t *testing.T, models []ChatModel) string {
|
||||
t.Helper()
|
||||
h := Chat(store.User{Email: "a@example.com"}, models, nil, "", "v1")
|
||||
h := Chat(store.User{Email: "a@example.com"}, models, nil, "", "v1", ChatKnowledge{})
|
||||
r := httptest.NewRequest(http.MethodGet, "http://x/chat", nil)
|
||||
w := httptest.NewRecorder()
|
||||
h(w, r)
|
||||
|
||||
@@ -12,7 +12,8 @@ import (
|
||||
|
||||
func TestTheChatPageSpeaksTurkishWhenAsked(t *testing.T) {
|
||||
h := Chat(store.User{Email: "a@example.com"},
|
||||
[]ChatModel{{Name: "qwen"}, {Name: "big-model", Simulated: true}}, nil, "", "v1")
|
||||
[]ChatModel{{Name: "qwen"}, {Name: "big-model", Simulated: true}}, nil, "", "v1",
|
||||
ChatKnowledge{Collections: []ChatCollection{{ID: "kb_x", Label: "politikalar"}}, Default: "politikalar"})
|
||||
r := httptest.NewRequest(http.MethodGet, "http://x/chat", nil)
|
||||
r.Header.Set("Accept-Language", "tr-TR,tr;q=0.9")
|
||||
w := httptest.NewRecorder()
|
||||
@@ -31,12 +32,18 @@ func TestTheChatPageSpeaksTurkishWhenAsked(t *testing.T) {
|
||||
// The script's own words come from the catalogue too.
|
||||
`"Uzunluk sınırında kesildi — devam etmesini isteyin."`,
|
||||
`decimal: ","`,
|
||||
// The documents picker, with the default named.
|
||||
"Belgeler",
|
||||
`<option value="">varsayılan (politikalar)</option>`,
|
||||
`<option value="none">yok</option>`,
|
||||
`<option value="kb_x">politikalar</option>`,
|
||||
`"Kaynaklar:"`,
|
||||
} {
|
||||
if !strings.Contains(page, want) {
|
||||
t.Errorf("the Turkish chat page does not have %q", want)
|
||||
}
|
||||
}
|
||||
for _, not := range []string{"New conversation", "Cut off at the length limit", "Ask something"} {
|
||||
for _, not := range []string{"New conversation", "Cut off at the length limit", "Ask something", "Documents", "Sources:"} {
|
||||
if strings.Contains(page, not) {
|
||||
t.Errorf("the Turkish chat page still says %q", not)
|
||||
}
|
||||
|
||||
@@ -19,7 +19,6 @@
|
||||
"chat.kept": "The conversation is kept in this browser, so a reload does not lose it; \"New conversation\" clears it. Nothing of it is stored on the gateway — what is recorded there is the same usage any request records: tokens, cost and latency.",
|
||||
"chat.empty": "No models are available to you right now.",
|
||||
"chat.empty_why": "Either nothing is placed yet, or your project's allowlist does not include anything that is. <code>/admin</code> shows the cluster if you have access to it.",
|
||||
|
||||
"chat.js.preparing": "Preparing…",
|
||||
"chat.js.download_failed": "download failed",
|
||||
"chat.js.stage.waiting": "Waiting for a node with room",
|
||||
@@ -41,7 +40,10 @@
|
||||
"chat.js.stopped_by_policy_empty": "(stopped by the gateway's policy{policy} before anything was shown)",
|
||||
"chat.js.stopped_empty": "(stopped before anything came back)",
|
||||
"chat.js.no_content": "(no content came back)",
|
||||
|
||||
"download.err.read": "could not read the request",
|
||||
"download.err.no_code": "there is no code in that reply to download"
|
||||
"download.err.no_code": "there is no code in that reply to download",
|
||||
"chat.knowledge": "Documents",
|
||||
"chat.knowledge.default": "default",
|
||||
"chat.knowledge.none": "none",
|
||||
"chat.js.sources": "Sources:"
|
||||
}
|
||||
@@ -10,5 +10,6 @@
|
||||
"nav.chat": "Chat",
|
||||
"nav.account": "Account",
|
||||
"nav.sign_out": "Sign out",
|
||||
"nav.docs": "Docs"
|
||||
"nav.docs": "Docs",
|
||||
"nav.settings": "Settings"
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"settings.title": "Settings",
|
||||
"settings.sub": "The cluster's settings: everything that is not one machine's own, the same on every control plane, in force within seconds of a change. <code>llmgw ctl config edit</code> edits the same document.",
|
||||
"settings.current": "Version %d, in force",
|
||||
"settings.earlier": "Version %d",
|
||||
"settings.earlier_sub": "An earlier version; version %d is in force. Rolling back keeps this one as a new version, so nothing is lost.",
|
||||
"settings.rollback": "Put version %d back in force",
|
||||
"settings.back_to_current": "back to the version in force",
|
||||
"settings.local_keys": "Not here, because each machine's server.yaml keeps them:",
|
||||
"settings.note": "Why",
|
||||
"settings.note_hint": "kept with the version, for whoever reads the history",
|
||||
"settings.save": "Check and apply",
|
||||
"settings.save_hint": "Checked as a control plane would load it before it is kept; a change made since this page was opened is not overwritten.",
|
||||
"settings.saved": "Version %d is in force; every control plane applies it within seconds.",
|
||||
"settings.rolled_back": "Version %d's settings are back in force, as version %d.",
|
||||
"settings.history": "History",
|
||||
"settings.col.version": "Version",
|
||||
"settings.col.at": "When",
|
||||
"settings.col.by": "By",
|
||||
"settings.col.note": "Why",
|
||||
"settings.in_force": "in force",
|
||||
"settings.view": "view",
|
||||
"settings.no_history": "No versions yet.",
|
||||
"settings.err.unchanged": "Nothing was changed.",
|
||||
"settings.err.invalid": "Not applied: a control plane could not load it.",
|
||||
"settings.err.changed": "Not applied: someone changed the settings since this page was opened. Your edit is below; compare it with the version now in force and apply it again.",
|
||||
"settings.err.no_version": "There is no such version."
|
||||
}
|
||||
@@ -19,7 +19,6 @@
|
||||
"chat.kept": "Sohbet bu tarayıcıda saklanır; sayfayı yenilemek onu silmez, \"Yeni sohbet\" temizler. Ağ geçidinde sohbetten hiçbir şey saklanmaz — orada kaydedilen, her isteğin kaydettiği kullanımla aynıdır: token, maliyet ve gecikme.",
|
||||
"chat.empty": "Şu anda kullanabileceğiniz bir model yok.",
|
||||
"chat.empty_why": "Ya henüz hiçbir model yerleştirilmedi ya da projenizin izin listesi yerleştirilmiş olanların hiçbirini içermiyor. Erişiminiz varsa <code>/admin</code> kümeyi gösterir.",
|
||||
|
||||
"chat.js.preparing": "Hazırlanıyor…",
|
||||
"chat.js.download_failed": "indirme başarısız oldu",
|
||||
"chat.js.stage.waiting": "Yeri olan bir düğüm bekleniyor",
|
||||
@@ -41,7 +40,10 @@
|
||||
"chat.js.stopped_by_policy_empty": "(hiçbir şey gösterilmeden ağ geçidinin politikası{policy} tarafından durduruldu)",
|
||||
"chat.js.stopped_empty": "(hiçbir yanıt gelmeden durduruldu)",
|
||||
"chat.js.no_content": "(hiçbir içerik gelmedi)",
|
||||
|
||||
"download.err.read": "istek okunamadı",
|
||||
"download.err.no_code": "bu yanıtta indirilecek kod yok"
|
||||
"download.err.no_code": "bu yanıtta indirilecek kod yok",
|
||||
"chat.knowledge": "Belgeler",
|
||||
"chat.knowledge.default": "varsayılan",
|
||||
"chat.knowledge.none": "yok",
|
||||
"chat.js.sources": "Kaynaklar:"
|
||||
}
|
||||
@@ -10,5 +10,6 @@
|
||||
"nav.chat": "Sohbet",
|
||||
"nav.account": "Hesap",
|
||||
"nav.sign_out": "Çıkış yap",
|
||||
"nav.docs": "Belgeler"
|
||||
"nav.docs": "Belgeler",
|
||||
"nav.settings": "Ayarlar"
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"settings.title": "Ayarlar",
|
||||
"settings.sub": "Kümenin ayarları: tek bir makineye ait olmayan her şey; her kontrol düzleminde aynı, değiştikten saniyeler sonra yürürlükte. <code>llmgw ctl config edit</code> aynı belgeyi düzenler.",
|
||||
"settings.current": "Sürüm %d, yürürlükte",
|
||||
"settings.earlier": "Sürüm %d",
|
||||
"settings.earlier_sub": "Önceki bir sürüm; yürürlükte olan sürüm %d. Geri almak bunu yeni bir sürüm olarak kaydeder, hiçbir şey kaybolmaz.",
|
||||
"settings.rollback": "%d numaralı sürümü yeniden yürürlüğe koy",
|
||||
"settings.back_to_current": "yürürlükteki sürüme dön",
|
||||
"settings.local_keys": "Burada değil, çünkü her makinenin kendi server.yaml dosyası tutar:",
|
||||
"settings.note": "Neden",
|
||||
"settings.note_hint": "sürümle birlikte saklanır, geçmişi okuyacak kişi için",
|
||||
"settings.save": "Denetle ve uygula",
|
||||
"settings.save_hint": "Saklanmadan önce bir kontrol düzleminin yükleyeceği gibi denetlenir; bu sayfa açıldıktan sonra yapılmış bir değişikliğin üzerine yazılmaz.",
|
||||
"settings.saved": "%d numaralı sürüm yürürlükte; her kontrol düzlemi saniyeler içinde uygular.",
|
||||
"settings.rolled_back": "%d numaralı sürümün ayarları yeniden yürürlükte; yeni sürüm numarası %d.",
|
||||
"settings.history": "Geçmiş",
|
||||
"settings.col.version": "Sürüm",
|
||||
"settings.col.at": "Ne zaman",
|
||||
"settings.col.by": "Kim",
|
||||
"settings.col.note": "Neden",
|
||||
"settings.in_force": "yürürlükte",
|
||||
"settings.view": "görüntüle",
|
||||
"settings.no_history": "Henüz sürüm yok.",
|
||||
"settings.err.unchanged": "Hiçbir şey değiştirilmedi.",
|
||||
"settings.err.invalid": "Uygulanmadı: bir kontrol düzlemi bunu yükleyemezdi.",
|
||||
"settings.err.changed": "Uygulanmadı: bu sayfa açıldıktan sonra biri ayarları değiştirdi. Düzenlemeniz aşağıda; yürürlükteki sürümle karşılaştırıp yeniden uygulayın.",
|
||||
"settings.err.no_version": "Böyle bir sürüm yok."
|
||||
}
|
||||
+11
-7
@@ -26,6 +26,9 @@ type Nav struct {
|
||||
Mirror, Quotas bool
|
||||
// Guard is whether the guardrails page was mounted.
|
||||
Guard bool
|
||||
// Settings offers the cluster's settings: mounted, and to an
|
||||
// administrator only.
|
||||
Settings bool
|
||||
// CSRF signs the sign-out form.
|
||||
CSRF string
|
||||
// Org offers the organization page: to an administrator, and to anyone
|
||||
@@ -38,13 +41,14 @@ type navKey struct{}
|
||||
// navFor builds the menu for a signed-in person.
|
||||
func (p *AccountPages) navFor(r *http.Request, user store.User, current string) Nav {
|
||||
n := Nav{
|
||||
Current: current,
|
||||
Admin: user.Staff(),
|
||||
Role: user.StaffRole(),
|
||||
Mirror: p.hasMirror,
|
||||
Quotas: p.hasQuotas,
|
||||
Guard: p.hasGuard,
|
||||
CSRF: csrfFor(r),
|
||||
Current: current,
|
||||
Admin: user.Staff(),
|
||||
Role: user.StaffRole(),
|
||||
Mirror: p.hasMirror,
|
||||
Quotas: p.hasQuotas,
|
||||
Guard: p.hasGuard,
|
||||
Settings: p.hasSettings && user.IsAdmin,
|
||||
CSRF: csrfFor(r),
|
||||
}
|
||||
// Only when there is an organization to show: a link that leads to "there
|
||||
// are none" on every page is noise.
|
||||
|
||||
@@ -905,3 +905,13 @@ a.sel-link.small { font-size: .78rem; color: var(--muted); margin-left: .4rem; }
|
||||
/* The language switcher at the foot of every page. */
|
||||
.langs a { color: var(--muted); }
|
||||
.langs b { font-weight: 600; }
|
||||
|
||||
/* The cluster's settings, edited as the YAML document they are. */
|
||||
.settings-form textarea, pre.settings-body {
|
||||
width: 100%; box-sizing: border-box; font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
|
||||
font-size: .85rem; line-height: 1.4; padding: .6rem .7rem; border: 1px solid var(--line);
|
||||
border-radius: 6px; background: var(--card); color: var(--fg); tab-size: 2; overflow-x: auto;
|
||||
}
|
||||
.settings-form label { display: block; margin: .8rem 0; }
|
||||
.settings-form input[type=text] { width: 100%; box-sizing: border-box; }
|
||||
.sr-only { position: absolute; width: 1px; height: 1px; overflow: hidden; clip: rect(0 0 0 0); white-space: nowrap; }
|
||||
@@ -0,0 +1,50 @@
|
||||
{{define "admin_settings"}}<!doctype html>
|
||||
<html lang="{{lang}}"><head>
|
||||
<meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
|
||||
<title>{{t "settings.title"}} · llmgw</title><link rel="stylesheet" href="{{stylesheet}}">
|
||||
</head><body><div class="wrap">
|
||||
<header>
|
||||
<h1>{{t "settings.title"}}</h1>
|
||||
<p class="sub">{{th "settings.sub"}}</p>
|
||||
{{template "nav" .Nav}}
|
||||
</header>
|
||||
{{if .Error}}<p class="error" role="alert">{{.Error}}</p>{{end}}
|
||||
{{if .Notice}}<p class="ok" role="status">{{.Notice}}</p>{{end}}
|
||||
|
||||
{{if .Earlier}}
|
||||
<h2>{{t "settings.earlier" .Shown.Version}}</h2>
|
||||
<p class="sub">{{t "settings.earlier_sub" .Current.Version}}</p>
|
||||
<pre class="settings-body">{{.Body}}</pre>
|
||||
<form method="post" action="/admin/settings/rollback">
|
||||
<input type="hidden" name="csrf" value="{{.CSRF}}">
|
||||
<input type="hidden" name="version" value="{{.Shown.Version}}">
|
||||
<button type="submit">{{t "settings.rollback" .Shown.Version}}</button>
|
||||
<a href="/admin/settings">{{t "settings.back_to_current"}}</a>
|
||||
</form>
|
||||
{{else}}
|
||||
<h2>{{t "settings.current" .Current.Version}}</h2>
|
||||
{{if .LocalKeys}}<p class="sub">{{t "settings.local_keys"}} {{range $i, $k := .LocalKeys}}{{if $i}}, {{end}}<code>{{$k}}</code>{{end}}.</p>{{end}}
|
||||
<form method="post" action="/admin/settings" class="settings-form">
|
||||
<input type="hidden" name="csrf" value="{{.CSRF}}">
|
||||
<input type="hidden" name="base" value="{{.Current.Version}}">
|
||||
<label for="settings-body" class="sr-only">{{t "settings.title"}}</label>
|
||||
<textarea id="settings-body" name="body" rows="30" spellcheck="false" class="mono">{{.Body}}</textarea>
|
||||
<label>{{t "settings.note"}} <span class="opt">{{t "settings.note_hint"}}</span>
|
||||
<input type="text" name="note" value="{{.Note}}" maxlength="200"></label>
|
||||
<button type="submit">{{t "settings.save"}}</button>
|
||||
<p class="sub">{{t "settings.save_hint"}}</p>
|
||||
</form>
|
||||
{{end}}
|
||||
|
||||
<h2>{{t "settings.history"}}</h2>
|
||||
{{if .History}}
|
||||
<table class="rows">
|
||||
<tr><th class="num">{{t "settings.col.version"}}</th><th>{{t "settings.col.at"}}</th><th>{{t "settings.col.by"}}</th><th>{{t "settings.col.note"}}</th><th></th></tr>
|
||||
{{range .History}}<tr{{if .Selected}} aria-current="true"{{end}}>
|
||||
<td class="num">{{.Version}}</td><td>{{.At}}</td><td>{{.By}}</td><td>{{.Note}}</td>
|
||||
<td>{{if .InForce}}<strong>{{t "settings.in_force"}}</strong>{{else}}<a href="/admin/settings?version={{.Version}}">{{t "settings.view"}}</a>{{end}}</td>
|
||||
</tr>{{end}}
|
||||
</table>
|
||||
{{else}}<p class="sub">{{t "settings.no_history"}}</p>{{end}}
|
||||
<footer>llmgw {{.Version}}</footer>
|
||||
</div></body></html>{{end}}
|
||||
@@ -34,6 +34,13 @@
|
||||
<label class="inline-label">{{t "chat.model"}}
|
||||
<select id="model">{{range .Models}}<option value="{{.Name}}" data-simulated="{{if .Simulated}}1{{else}}0{{end}}" data-ondemand="{{if .OnDemand}}1{{else}}0{{end}}">{{.Name}}{{with .Runs}} · {{.}}{{end}}{{if .Simulated}} {{t "chat.model.simulated"}}{{end}}{{if .OnDemand}} {{t "chat.model.not_loaded"}}{{end}}</option>{{end}}</select>
|
||||
</label>
|
||||
{{if .Collections}}<label class="inline-label">{{t "chat.knowledge"}}
|
||||
<select id="knowledge">
|
||||
{{if .KnowledgeDefault}}<option value="">{{t "chat.knowledge.default"}} ({{.KnowledgeDefault}})</option>
|
||||
<option value="none">{{t "chat.knowledge.none"}}</option>{{else}}<option value="none">{{t "chat.knowledge.none"}}</option>{{end}}
|
||||
{{range .Collections}}<option value="{{.ID}}">{{.Label}}</option>{{end}}
|
||||
</select>
|
||||
</label>{{end}}
|
||||
<button id="reset" type="button" class="link">{{t "chat.new"}}</button>
|
||||
</div>
|
||||
|
||||
@@ -89,6 +96,7 @@
|
||||
downloadZip: {{t "chat.js.download_zip"}},
|
||||
stoppedByPolicy: {{t "chat.js.stopped_by_policy"}},
|
||||
cutOff: {{t "chat.js.cut_off"}},
|
||||
sources: {{t "chat.js.sources"}},
|
||||
stoppedByPolicyEmpty: {{t "chat.js.stopped_by_policy_empty"}},
|
||||
stoppedEmpty: {{t "chat.js.stopped_empty"}},
|
||||
noContent: {{t "chat.js.no_content"}},
|
||||
@@ -381,6 +389,7 @@
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json", "X-CSRF-Token": csrf },
|
||||
body: JSON.stringify({ model: model.value, context: ctxValue(), temperature: +temp.value,
|
||||
knowledge: (document.getElementById("knowledge") || {}).value || "",
|
||||
messages: (system ? [{ role: "system", content: system }] : []).concat(messages) }),
|
||||
signal: busy.signal,
|
||||
});
|
||||
@@ -478,6 +487,16 @@
|
||||
note.textContent = say.cutOff;
|
||||
out.after(note);
|
||||
}
|
||||
// Which passages of the documents it was given, numbered as the
|
||||
// reply cites them. Ours, so a sibling placed with textContent.
|
||||
const src = resp.headers.get("x-llmgw-sources");
|
||||
if (src) {
|
||||
const note = document.createElement("p");
|
||||
note.className = "note sources";
|
||||
note.textContent = say.sources + " " + src.split(",").map((s, i) =>
|
||||
"[" + (i + 1) + "] " + decodeURIComponent(s.split("#")[0])).join(" ");
|
||||
out.after(note);
|
||||
}
|
||||
stats(out, tokens || pieces, t0, tFirst, performance.now(), stopped, resp.headers.get("x-llmgw-hardware"));
|
||||
// Once, at the end. Re-rendering every chunk would try to typeset half
|
||||
// an expression on each token and flicker its way through the reply.
|
||||
|
||||
@@ -20,6 +20,9 @@
|
||||
{{- if .Guard}}
|
||||
<a href="/admin/guard"{{if eq .Current "guard"}} aria-current="page"{{end}}>{{t "nav.guard"}}</a>
|
||||
{{- end}}
|
||||
{{- if .Settings}}
|
||||
<a href="/admin/settings"{{if eq .Current "settings"}} aria-current="page"{{end}}>{{t "nav.settings"}}</a>
|
||||
{{- end}}
|
||||
{{- /* An organization is a panel of the people page for an administrator;
|
||||
/org is the managers' view, and leads an administrator there. */}}
|
||||
<span class="menu-gap" aria-hidden="true"></span>
|
||||
|
||||
Reference in new issue
Block a user