diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b2f30c..3bd94ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,18 @@ # Changelog +## Unreleased - P-f + +- Add OpenAI-compatible embeddings (including Azure), explicit CHAT/EMBEDDING + model kinds, editable discovery hints, and budgeted EMBED usage. +- Add admin Knowledge sources with streamed Review/Result Index now and Clear + index actions. Indexing uses the clicking admin's UI API session, chunks text, + skips unchanged hashes and reconciles disappeared records. +- Add content-free in-memory and optional JDBC pgvector stores, per-instance + schemas, HNSW cosine search and transactional record replacement. +- Add permission-filtered `search_knowledge`, fresh UI API reads, existing + masking/redaction, PHI embedding guards, classification propagation and citations. +- Add embedding HTTP contracts, real UI API RAG security/golden tests and + always-on JDBC contracts; the pgvector container test skips without Docker. ## Unreleased - Answer presentation - Add `present_results`: server-formatted metrics, tables, computed percentages, diff --git a/CLAUDE.md b/CLAUDE.md index 403ab3c..8945998 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -46,6 +46,8 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def - `qbit-assistant-provider-bedrock`: Claude through the Anthropic Bedrock backend, AWS default credentials and SigV4; IRSA/web-identity hosts supply STS. - `qbit-assistant-provider-vertex`: Claude through the Anthropic Vertex backend and Gemini through OpenAI-compatible Chat Completions, with refreshed Google ADC tokens. +- `qbit-assistant-rag-pgvector`: optional host-supplied JDBC vector store, isolated by instance schema; no QQQ RDBMS internals. + ## Invariants (do not break) - **Sessions:** @@ -53,6 +55,7 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def - Every route uses `ExecutorSessionUtils.setupSession` (the user's own session). - `AssistantSecurity` only ever adds the user's OWN id key, or the all-access key for `assistant.admin` holders. - **Data access:** + - The vector index stores no record content: only source/table/key/field/chunk, content hash and embedding. Retrieval re-reads every hit through the UI API as the asking user, then applies the existing get_record exposure/masking path. - Only `persistence/` calls QueryAction, InsertAction and the other data actions (ArchUnit). - Data tools run only through `UiApiToolbox` → `UiApiGateway` (spec'd `/qqq/v1` operations with the user's `sessionUUID`); never core data actions for user data. Only `navigate` and `search_help` use `AiToolKernel.execute`. - Data access (`DataAccessService`) and the policy only ever narrow what the UI API allows; a table whose AI exposure is NONE stays NONE. @@ -99,4 +102,4 @@ QQQ style: - `develop` publishes SNAPSHOTs. - `release/*` publishes RCs; the orb bumps `` in the root POM. - `v*` tags publish releases. -- All five modules and the parent POM publish. The flatten plugin resolves `${revision}`. +- All six modules and the parent POM publish. The flatten plugin resolves `${revision}`. diff --git a/README.md b/README.md index a980696..ad5a707 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,68 @@ Design spec: `QRun-IO/qqq` `docs/specs/2026-10-11-qqq-assistant-qbit-design.md`. | `com.kingsrook.qbits:qbit-assistant` | The QBit: conversation tables, turn engine, `/qqq/v1/assistant` routes (JSON + Server-Sent Events), budgets, kill switch, the provider SPI, and a scripted provider for tests and demos. | | `com.kingsrook.qbits:qbit-assistant-provider-anthropic` | The Claude provider (`com.anthropic:anthropic-java`). Add it to use Claude. The core QBit has no LLM SDK dependency. | +Data tools use the QQQ UI API (`/qqq/v1`) with the signed-in user's session. The UI API enforces permissions and record locks; assistant data-access policy and field exposure can only narrow access. Only record-free navigation and help use the core AI tool kernel. Every tool result reaches the model wrapped as untrusted data. + +## RAG setup + +Knowledge retrieval is optional. The vector index stores **no record content**: +only source/table/key/field/chunk identifiers, SHA-256 hashes and embeddings. +Indexing reads through the UI API as the clicking administrator. Retrieval +re-reads each hit through that same API as the asking user, silently filters +403/404 results, applies `get_record` field exposure/masking/redaction and returns +fresh snippets with `/app/...` citations. `search_knowledge` follows `allowedTools` +and is offered only with an enabled source the user may read in the UI. + +1. Register the OpenAI-compatible provider module and configure the embedding + endpoint through the existing provider settings/URL allowlist. Discover or add + a model, set **Model kind = EMBEDDING** and its dimensions. Known discovery + defaults are `bge-m3` (1024) and `nomic-embed-text` (768); these remain editable. + CHAT models cannot be selected for embedding, and EMBEDDING models cannot answer + chat turns. Each embedding batch contains at most 64 inputs. +2. Use the default `InMemoryVectorStore` for tests/small hosts, or add + `com.kingsrook.qbits:qbit-assistant-rag-pgvector` at the QBit's version for + persistence. The host supplies its JDBC driver and DataSource; this module + does not depend on QQQ RDBMS internals. +3. For pgvector, have the DBA install the extension once in the chosen database: + `CREATE EXTENSION IF NOT EXISTS vector`. Give the application a distinct schema + per QQQ instance. Configure the store before producing the QBit: + + ```java + PgVectorStore vectors = new PgVectorStore(dataSource, 1024) + .withSchema("assistant_rag_my_instance"); + vectors.ensureSchema(); + config.withVectorStore(vectors); + config.withUiApiBaseUrl("http://127.0.0.1:8000"); + ``` + + The UI API URL is this server's own listening address. Configure it explicitly + so Index now works before any chat request has taught the runtime its local + port. The store creates a content-free `assistant_rag_chunk` table and an HNSW + cosine index. Its fixed dimension must match every source using the store; + different models are searched separately. No credentials go in source rows. + See the [pgvector installation/index documentation](https://github.com/pgvector/pgvector) + and the reference DDL in [schema-postgres.sql](docs/schema-postgres.sql). +4. Open **Assistant → Knowledge sources** as an `assistant.admin`. Add the table, + comma-separated text fields, optional title field and QQQ query-filter JSON, + embedding model, and chunk size/overlap (defaults 1500/200). Tables with AI data + access NONE, `assistant_*` tables, and fields with NONE/METADATA_ONLY exposure + are refused. MASKED fields are masked before embedding and in snippets. +5. Select **Index now**. Review makes no embedding calls; Result reports records + seen, chunks embedded, unchanged, removed and errors. Re-running unchanged + content sends no embeddings. A complete successful scan removes disappeared + records; failed scans preserve unvisited entries. **Clear index** removes only + that source's vectors. Changes after indexing are read fresh, but semantic + ranking becomes current only after re-indexing. + +Embedding usage has purpose `EMBED`, counts toward budgets, and observes the kill +switch and provider concurrency limits. PHI sources and queries in PHI +conversations require an eligible embedding provider; the eligibility flag is an +operator assertion, as for chat. Automatic budget downgrades cannot change an +embedding model's vector space, so an incompatible downgrade stops the call. + +Scheduling is deferred to P-g. Its designated-user runner can call +`RagService.index` within that user's authenticated context; P-f creates no system +session, stored admin credentials, scheduler or background content crawler. Data tools call the QQQ UI API through `UiApiToolbox` and `UiApiGateway` with the signed-in user's session. The UI enforces permissions and record security; the assistant further narrows access and caps and redacts returned values. Data tool results reach the model wrapped as untrusted data. Only record-free navigation and help use the core AI tool kernel. ## Requirements diff --git a/UPGRADE.md b/UPGRADE.md index 646546a..fd719a8 100644 --- a/UPGRADE.md +++ b/UPGRADE.md @@ -1,5 +1,36 @@ # Upgrading qbit-assistant +## P-f: knowledge retrieval + +Apply the **P-f** assistant-table DDL in +[docs/schema-postgres.sql](docs/schema-postgres.sql) before deployment. It adds +`assistant_model.model_kind` (existing rows default to CHAT), +`embedding_dimensions`, and the admin-only `assistant_rag_source` table. +No sources are seeded. Existing chat selections remain unchanged; explicitly +EMBEDDING models are excluded from chat, classifier, fallback and downgrade +selection. Correct discovery's kind/dimensions if a model name is ambiguous. + +For persistent vectors, add `qbit-assistant-rag-pgvector` at the same version as +the QBit, supply the host's JDBC DataSource/driver, and install pgvector in the +database. A DBA runs `CREATE EXTENSION IF NOT EXISTS vector`; the application +role needs schema creation rights for initial setup (or a DBA-created schema), +then normal DML rights. `ensureSchema()` checks for an existing extension first +and does not require extension installation privileges when it is already present. +The optional pgvector DDL is a **commented reference block**, separate from the +assistant-table migration. See [RAG setup](README.md#rag-setup). + +Use a distinct validated schema for each QQQ instance. `vector(N)` fixes dimensions +for that store; changing dimensions requires a new schema and re-indexing. +Sources can use different models with that dimension; ranking is restricted to +sources using the query's model. If the weights behind an unchanged model id +change, clear and rebuild its sources. Editing source chunk/filter/model settings +or registry model identity/dimensions automatically clears affected vectors. + +Every indexing batch produces EMBED usage and updates existing budget rollups. +PHI source content and PHI conversation queries require a HIPAA-eligible embedding +provider. No background runner or designated credentials are installed in P-f. +The default in-memory store is lost on restart and must be indexed again. + ## P-e: profile routing and provider fallback Apply the additive **P-e** ALTER TABLE statements at the end of diff --git a/docs/DESIGN.md b/docs/DESIGN.md index 5aa86f2..2081531 100644 --- a/docs/DESIGN.md +++ b/docs/DESIGN.md @@ -422,6 +422,53 @@ AI exposure NONE is not enough, because the assistant acts as the user and an ad - Admin record screens, additive Postgres DDL and the flow-layout home expose these settings and counters without granting AI access to assistant configuration. +### P-f as built + +- `AssistantModelProvider.embed` is optional and unsupported by default. The + OpenAI-compatible implementation shares chat's JDK HTTP authentication, TLS, + Azure deployment paths, timeout and sanitized error mapping. A request accepts + 1–64 inputs; `EmbeddingService` batches larger jobs and records each call as + EMBED, including daily rollups and provider/model budget accounting. +- Registry rows have CHAT/EMBEDDING kinds and embedding dimensions. Legacy rows + remain CHAT. Discovery suggests EMBEDDING for ids containing embed, bge, e5, + gte or nomic; administrators can correct it. Chat routing/fallback/downgrades + cannot select an embedding model. Embedding resolution is explicit and checks + fresh provider enablement and HIPAA eligibility before dispatch. +- `assistant_rag_source` is an admin-only navigation table. Its text fields must + exist and allow VALUES/MASKED exposure; application data access must allow READ. + Index now and Clear index are standard streamed ETL record actions with Review + and Result, never navigation entries. Indexing pages the UI API as the clicking + admin, reuses get_record's masking path, chunks deterministically and compares + hashes. Complete scans alone prune disappeared records; per-record replacement + also removes obsolete field chunks. Configuration/model identity changes clear + stale vectors. No source text is saved on a source or vector row. +- `VectorStore` supplies upsert/delete/search/count plus content-free chunk + reconciliation and atomic record replacement. The in-memory implementation uses + exact cosine; the optional pgvector module uses host JDBC, a validated instance + schema, bound values and an HNSW cosine index. Sources from other embedding + models are excluded before ranking. The DBA installs pgvector; existing + installations do not require extension-installation privileges on startup. +- `KnowledgeTool` is separate from the unchanged `UiApiToolbox`. A small + ContextBuilder hook adds it only for an enabled source, policy permission and + UI-readable source table. Search embeds the question, requests 4× the result + limit, then re-reads each hit through get_record with the asking user's session. + 403/404 hits contribute only to the filtered count. Safe current values produce + snippets and `/app` citations; record refs propagate touched tables into P-c + classification and the normal tool-call log. The citation instruction is in + the static rules block. The result limit is capped at eight. +- PHI source content and PHI conversation queries only reach eligible embedding + providers. Embedding budgets, kill switch and provider concurrency checks apply + to each batch. A budget downgrade into another vector space stops embedding. +- Tests exercise real fixture UI API reads, denials, record locks, field masking, + PHI refusal, classification, reconciliation, embedding wire contracts and a + golden cited answer. JDBC shape/transaction tests always run; a pgvector:pg16 + Testcontainers test skips on hosts without Docker. +- P-g may invoke the manual service within a designated user's authenticated + context. P-f supplies no scheduled runner or stored/elevated session. Multiple + replicas should serialize manual re-indexing of a source at the host level; + per-record pgvector replacement is transactional, while scan reconciliation is + not a snapshot of concurrently changing application data. + ### Phases (each a green PR) | | Scope | Depends on | |---|---|---| diff --git a/docs/schema-postgres.sql b/docs/schema-postgres.sql index 45a895b..4eeae53 100644 --- a/docs/schema-postgres.sql +++ b/docs/schema-postgres.sql @@ -477,3 +477,50 @@ ALTER TABLE assistant_turn ADD COLUMN route_decision VARCHAR(16); ALTER TABLE assistant_turn ADD COLUMN route_source VARCHAR(16); ALTER TABLE assistant_turn ADD COLUMN fallback_count INTEGER DEFAULT 0; ALTER TABLE assistant_turn ADD COLUMN provider_id INTEGER REFERENCES assistant_provider(id); + + +-- P-f: embedding registry and admin-only knowledge-source configuration. +ALTER TABLE assistant_model ADD COLUMN model_kind VARCHAR(16) DEFAULT 'CHAT' + CHECK (model_kind IN ('CHAT', 'EMBEDDING')); +ALTER TABLE assistant_model ADD COLUMN embedding_dimensions INTEGER + CHECK (embedding_dimensions > 0); + +CREATE TABLE assistant_rag_source ( + id SERIAL PRIMARY KEY, + name VARCHAR(200) NOT NULL, + enabled BOOLEAN DEFAULT TRUE, + table_name VARCHAR(200) NOT NULL, + field_names VARCHAR(500) NOT NULL, + title_field_name VARCHAR(200), + filter_json TEXT, + embedding_model_id INTEGER NOT NULL REFERENCES assistant_model(id), + chunk_chars INTEGER DEFAULT 1500 CHECK (chunk_chars BETWEEN 100 AND 16000), + chunk_overlap_chars INTEGER DEFAULT 200 CHECK (chunk_overlap_chars >= 0 AND chunk_overlap_chars < chunk_chars), + last_indexed_at TIMESTAMPTZ, + last_index_status VARCHAR(200), + last_index_error VARCHAR(500), + indexed_record_count INTEGER, + indexed_chunk_count INTEGER, + create_date TIMESTAMPTZ, + modify_date TIMESTAMPTZ +); + +-- Optional pgvector part: reference only. Run the extension step as a DBA. +-- PgVectorStore.ensureSchema() performs the remainder for the host's schema/N. +-- Use a distinct schema per QQQ instance; match N to the embedding model. +-- CREATE EXTENSION IF NOT EXISTS vector; +-- CREATE SCHEMA IF NOT EXISTS assistant_rag; +-- CREATE TABLE IF NOT EXISTS assistant_rag.assistant_rag_chunk ( +-- id BIGSERIAL PRIMARY KEY, +-- source_id INT NOT NULL, +-- table_name TEXT NOT NULL, +-- record_pk TEXT NOT NULL, +-- field_name TEXT NOT NULL, +-- chunk_index INT NOT NULL, +-- content_hash CHAR(64) NOT NULL, +-- embedding vector(1024) NOT NULL, +-- indexed_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, +-- UNIQUE(source_id, table_name, record_pk, field_name, chunk_index) +-- ); +-- CREATE INDEX IF NOT EXISTS assistant_rag_chunk_cosine +-- ON assistant_rag.assistant_rag_chunk USING hnsw (embedding vector_cosine_ops); diff --git a/pom.xml b/pom.xml index b3d2499..173f755 100644 --- a/pom.xml +++ b/pom.xml @@ -55,6 +55,7 @@ qbit-assistant-provider-openai qbit-assistant-provider-bedrock qbit-assistant-provider-vertex + qbit-assistant-rag-pgvector diff --git a/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiCompatibleModelProvider.java b/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiCompatibleModelProvider.java index 073d39f..9a54e97 100644 --- a/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiCompatibleModelProvider.java +++ b/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiCompatibleModelProvider.java @@ -38,6 +38,8 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; import com.kingsrook.qbits.assistant.provider.AssistantModelProvider; import com.kingsrook.qbits.assistant.provider.CompletableStreamHandle; +import com.kingsrook.qbits.assistant.provider.EmbeddingRequest; +import com.kingsrook.qbits.assistant.provider.EmbeddingResponse; import com.kingsrook.qbits.assistant.provider.ModelCapabilities; import com.kingsrook.qbits.assistant.provider.ModelEventListener; import com.kingsrook.qbits.assistant.provider.ModelPart; @@ -82,6 +84,84 @@ public class OpenAiCompatibleModelProvider implements AssistantModelProvider + /******************************************************************************* + ** One bounded batch, with the deadline covering headers and the entire body. + ** Restore input order from the response indexes; reject malformed vectors. + *******************************************************************************/ + @Override + public EmbeddingResponse embed(EmbeddingRequest request) throws ModelProviderException + { + ObjectNode body = OpenAiHttp.JSON.createObjectNode(); + body.put("model", request.model()); + body.set("input", OpenAiHttp.JSON.valueToTree(request.input())); + var pending = http.client.sendAsync(http.request(request.model(), http.settings.getRequestTimeoutSeconds(), "embeddings") + .POST(java.net.http.HttpRequest.BodyPublishers.ofString(body.toString())).build(), HttpResponse.BodyHandlers.ofString()); + try + { + HttpResponse response = pending.get(http.settings.getRequestTimeoutSeconds(), TimeUnit.SECONDS); + if(response.statusCode() != 200) + { + throw (http.error(response.statusCode(), response.body(), OpenAiHttp.requestId(response), false)); + } + JsonNode json = OpenAiHttp.JSON.readTree(response.body()); + if(json.has("error")) + { + throw (http.error(500, response.body(), OpenAiHttp.requestId(response), false)); + } + float[][] vectors = new float[request.input().size()][]; + Integer dimensions = null; + for(JsonNode item : json.path("data")) + { + Integer index = item.path("index").asInt(-1); + JsonNode values = item.path("embedding"); + if(index < 0 || index >= vectors.length || vectors[index] != null || !values.isArray() || values.isEmpty() + || (dimensions != null && dimensions != values.size())) + { + throw (new IOException("Invalid embedding response.")); + } + dimensions = values.size(); + float[] vector = new float[dimensions]; + for(Integer i = 0; i < vector.length; i++) + { + if(!values.get(i).isNumber()) + { + throw (new IOException("Invalid embedding coordinate.")); + } + vector[i] = (float) values.get(i).asDouble(); + if(!Float.isFinite(vector[i])) + { + throw (new IOException("Non-finite embedding coordinate.")); + } + } + vectors[index] = vector; + } + if(java.util.Arrays.stream(vectors).anyMatch(java.util.Objects::isNull)) + { + throw (new IOException("Missing embedding.")); + } + return (new EmbeddingResponse(List.of(vectors), new ModelUsage().withInputTokens(json.path("usage").path("prompt_tokens").asLong()), OpenAiHttp.requestId(response))); + } + catch(java.util.concurrent.TimeoutException e) + { + throw (new ModelProviderException(ModelProviderException.Kind.TIMEOUT, null, "The provider did not answer in time.")); + } + catch(InterruptedException e) + { + Thread.currentThread().interrupt(); + throw (new ModelProviderException(ModelProviderException.Kind.NETWORK, null, "The embedding request was interrupted.")); + } + catch(java.util.concurrent.ExecutionException | IOException e) + { + throw (OpenAiHttp.transportError(e)); + } + finally + { + pending.cancel(true); + } + } + + + /******************************************************************************* ** Conservative defaults; registry rows can override model-specific facts. *******************************************************************************/ diff --git a/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiHttp.java b/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiHttp.java index 0313a79..f0891c1 100644 --- a/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiHttp.java +++ b/qbit-assistant-provider-openai/src/main/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiHttp.java @@ -84,13 +84,23 @@ final class OpenAiHttp ** Azure uses deployment paths and a version query parameter. *******************************************************************************/ HttpRequest.Builder request(String model, Integer seconds) throws ModelProviderException + { + return (request(model, seconds, "chat/completions")); + } + + + + /******************************************************************************* + ** Chat and embeddings share authentication, TLS, deadlines and Azure paths. + *******************************************************************************/ + HttpRequest.Builder request(String model, Integer seconds, String operation) throws ModelProviderException { try { - String path = model == null ? "/models" : "/chat/completions"; + String path = model == null ? "/models" : "/" + operation; if("AZURE_OPENAI".equals(settings.getKind())) { - path = model == null ? "/openai/models" : "/openai/deployments/" + encode(model) + "/chat/completions"; + path = model == null ? "/openai/models" : "/openai/deployments/" + encode(model) + "/" + operation; path += "?api-version=" + encode(settings.getApiVersion() == null || settings.getApiVersion().isBlank() ? "2024-10-21" : settings.getApiVersion()); } HttpRequest.Builder builder = HttpRequest.newBuilder(URI.create(base + path)).timeout(Duration.ofSeconds(seconds)).header("Content-Type", "application/json"); diff --git a/qbit-assistant-provider-openai/src/test/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiEmbeddingTest.java b/qbit-assistant-provider-openai/src/test/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiEmbeddingTest.java new file mode 100644 index 0000000..bb398b4 --- /dev/null +++ b/qbit-assistant-provider-openai/src/test/java/com/kingsrook/qbits/assistant/provider/openai/OpenAiEmbeddingTest.java @@ -0,0 +1,133 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.provider.openai; + + +import java.util.List; +import com.kingsrook.qbits.assistant.provider.EmbeddingRequest; +import com.kingsrook.qbits.assistant.provider.ModelProviderException; +import com.kingsrook.qbits.assistant.provider.ProviderSettings; +import org.junit.jupiter.api.Test; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + + +/******************************************************************************* + ** Embedding wire contracts through the same local HTTP harness as chat. + *******************************************************************************/ +class OpenAiEmbeddingTest +{ + /******************************************************************************* + ** Input arrays, response indexes, usage and authentication are preserved. + *******************************************************************************/ + @Test + void shapeOrderAndAzure() throws Exception + { + for(String kind : List.of("OPENAI", "AZURE_OPENAI", "OLLAMA")) + { + try(MockOpenAiServer server = new MockOpenAiServer()) + { + server.enqueue(200, "{\"data\":[{\"index\":1,\"embedding\":[0,1]},{\"index\":0,\"embedding\":[1,0]}],\"usage\":{\"prompt_tokens\":7}}"); + var provider = new OpenAiCompatibleModelProvider(new ProviderSettings().withKind(kind).withBaseUrl(server.url()).withApiKey("synthetic-key").withApiVersion("2024-10-21")); + var response = provider.embed(new EmbeddingRequest("bge-m3", List.of("first", "second"))); + assertThat(response.vectors()).hasSize(2); + assertThat(response.vectors().getFirst()).containsExactly(1, 0); + assertThat(response.usage().getInputTokens()).isEqualTo(7L); + var captured = server.requests.getFirst(); + assertThat(captured.path()).isEqualTo(kind.equals("AZURE_OPENAI") ? "/v1/openai/deployments/bge-m3/embeddings?api-version=2024-10-21" : "/v1/embeddings"); + var body = OpenAiHttp.JSON.readTree(captured.body()); + assertThat(body.path("model").asText()).isEqualTo("bge-m3"); + assertThat(body.path("input").size()).isEqualTo(2); + assertThat(body.has("messages")).isFalse(); + assertThat(body.has("stream")).isFalse(); + assertThat(kind.equals("AZURE_OPENAI") ? captured.apiKey() : captured.authorization()).contains("synthetic-key"); + } + } + } + + + + /******************************************************************************* + ** Reuse chat's safe error categories, without echoing credentials. + *******************************************************************************/ + @Test + void errorsAndMalformedVectors() throws Exception + { + try(MockOpenAiServer server = new MockOpenAiServer()) + { + var provider = new OpenAiCompatibleModelProvider(new ProviderSettings().withKind("OPENAI").withBaseUrl(server.url()).withApiKey("synthetic-key")); + server.enqueue(429, "{\"error\":{\"message\":\"synthetic-key rejected\"}}"); + assertThatThrownBy(() -> provider.embed(new EmbeddingRequest("embed", List.of("text")))) + .isInstanceOfSatisfying(ModelProviderException.class, e -> + { + assertThat(e.getKind()).isEqualTo(ModelProviderException.Kind.RATE_LIMITED); + assertThat(e.getMessage()).doesNotContain("synthetic-key"); + }); + for(String body : List.of("{}", "{\"data\":[{\"index\":0,\"embedding\":[\"not numeric\"]}]}", + "{\"data\":[{\"index\":2,\"embedding\":[1]}]}", "{\"data\":[{\"index\":0,\"embedding\":[1]},{\"index\":0,\"embedding\":[1]}]}")) + { + server.enqueue(200, body); + assertThatThrownBy(() -> provider.embed(new EmbeddingRequest("embed", List.of("text")))).isInstanceOf(ModelProviderException.class); + } + assertThatThrownBy(() -> new EmbeddingRequest("embed", java.util.Collections.nCopies(65, "text"))).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new EmbeddingRequest("embed", List.of())).isInstanceOf(IllegalArgumentException.class); + } + } + + + + /******************************************************************************* + ** Deadline covers a stalled response body, not only HTTP headers. + *******************************************************************************/ + @Test + void bodyTimeout() throws Exception + { + try(MockOpenAiServer server = new MockOpenAiServer()) + { + server.stall(true); + var provider = new OpenAiCompatibleModelProvider(new ProviderSettings().withKind("OLLAMA").withBaseUrl(server.url()).withRequestTimeoutSeconds(1)); + assertThatThrownBy(() -> provider.embed(new EmbeddingRequest("embed", List.of("text")))) + .isInstanceOfSatisfying(ModelProviderException.class, e -> assertThat(e.getKind()).isEqualTo(ModelProviderException.Kind.TIMEOUT)); + } + } + + + + /******************************************************************************* + ** The transport accepts a complete 64-input batch and rejects overflow + ** before sending. EmbeddingService owns splitting larger indexing jobs. + *******************************************************************************/ + @Test + void maximumBatch() throws Exception + { + try(MockOpenAiServer server = new MockOpenAiServer()) + { + var body = OpenAiHttp.JSON.createObjectNode(); + var data = body.putArray("data"); + for(Integer i = 0; i < 64; i++) + { + data.addObject().put("index", i).putArray("embedding").add(1).add(0); + } + server.enqueue(200, body.toString()); + var provider = new OpenAiCompatibleModelProvider(new ProviderSettings().withKind("OLLAMA").withBaseUrl(server.url())); + assertThat(provider.embed(new EmbeddingRequest("embed", java.util.Collections.nCopies(64, "text"))).vectors()).hasSize(64); + assertThat(OpenAiHttp.JSON.readTree(server.requests.getFirst().body()).path("input").size()).isEqualTo(64); + assertThatThrownBy(() -> provider.embed(new EmbeddingRequest("embed", java.util.Collections.nCopies(65, "text")))).isInstanceOf(IllegalArgumentException.class); + assertThat(server.requests).hasSize(1); + } + } +} diff --git a/qbit-assistant-rag-pgvector/pom.xml b/qbit-assistant-rag-pgvector/pom.xml new file mode 100644 index 0000000..999a459 --- /dev/null +++ b/qbit-assistant-rag-pgvector/pom.xml @@ -0,0 +1,130 @@ + + + + + 4.0.0 + + + com.kingsrook.qbits + qbit-assistant-parent + ${revision} + + + qbit-assistant-rag-pgvector + jar + + QBit Assistant - pgvector + Content-free per-instance vector storage through JDBC + https://github.com/QRun-IO/qbit-assistant + + + + com.kingsrook.qbits + qbit-assistant + + + com.kingsrook.qqq + qqq-backend-core + + + org.testcontainers + junit-jupiter + 1.21.4 + test + + + org.testcontainers + postgresql + 1.21.4 + test + + + org.postgresql + postgresql + 42.7.14 + test + + + org.junit.jupiter + junit-jupiter + test + + + org.assertj + assertj-core + test + + + + + + + org.apache.maven.plugins + maven-surefire-plugin + + + ${test.excludedGroups} + true + + + + org.jacoco + jacoco-maven-plugin + + + prepare-agent + + prepare-agent + + + + report + verify + + report + + + + check + + check + + + + + BUNDLE + + + INSTRUCTION + COVEREDRATIO + 0.70 + + + CLASS + COVEREDRATIO + 0.90 + + + + + + + + + + + diff --git a/qbit-assistant-rag-pgvector/src/main/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorStore.java b/qbit-assistant-rag-pgvector/src/main/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorStore.java new file mode 100644 index 0000000..ca81dfd --- /dev/null +++ b/qbit-assistant-rag-pgvector/src/main/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorStore.java @@ -0,0 +1,452 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag.pgvector; + + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.StringJoiner; +import javax.sql.DataSource; +import com.kingsrook.qbits.assistant.rag.ChunkVector; +import com.kingsrook.qbits.assistant.rag.VectorStore; +import com.kingsrook.qqq.backend.core.exceptions.QException; + + +/******************************************************************************* + ** Host-supplied JDBC only, without QQQ RDBMS internals. No text is ever stored. + *******************************************************************************/ +public class PgVectorStore implements VectorStore +{ + private final ConnectionSupplier connections; + private final Integer dimensions; + private String schema = "assistant_rag"; + + + + /******************************************************************************* + ** The host owns credentials and connection pooling; each call closes its lease. + *******************************************************************************/ + @FunctionalInterface + public interface ConnectionSupplier + { + /******************************************************************************* + ** Return a connection owned by this caller until it closes it. + *******************************************************************************/ + Connection open() throws SQLException; + } + + + + /******************************************************************************* + ** Use any host DataSource and a fixed vector dimension for this instance. + *******************************************************************************/ + public PgVectorStore(DataSource dataSource, Integer dimensions) + { + this(dataSource::getConnection, dimensions); + } + + + + /******************************************************************************* + ** HNSW vector cosine indexes support up to 2000 dimensions. + *******************************************************************************/ + public PgVectorStore(ConnectionSupplier connections, Integer dimensions) + { + if(connections == null || dimensions == null || dimensions < 1 || dimensions > 2000) + { + throw (new IllegalArgumentException("A connection supplier and 1 to 2000 dimensions are required.")); + } + this.connections = connections; + this.dimensions = dimensions; + } + + + + /******************************************************************************* + ** Configure before use. Identifiers cannot be parameters, so validate strictly. + *******************************************************************************/ + public PgVectorStore withSchema(String schema) + { + if(schema == null || !schema.matches("[a-z_][a-z0-9_]{0,62}")) + { + throw (new IllegalArgumentException("The vector schema must be a lower-case SQL identifier of at most 63 characters.")); + } + this.schema = schema; + return (this); + } + + + + /******************************************************************************* + ** The DBA normally installs vector. An existing extension needs no elevated + ** privileges; if absent, attempt installation and report an actionable error. + *******************************************************************************/ + public void ensureSchema() throws QException + { + try(Connection connection = connections.open()) + { + try(PreparedStatement check = connection.prepareStatement("SELECT 1 FROM pg_extension WHERE extname = ?")) + { + check.setString(1, "vector"); + try(ResultSet rows = check.executeQuery()) + { + if(!rows.next()) + { + try(PreparedStatement extension = connection.prepareStatement("CREATE EXTENSION IF NOT EXISTS vector")) + { + extension.execute(); + } + catch(SQLException e) + { + throw (new QException("Ask the DBA to run CREATE EXTENSION IF NOT EXISTS vector in this database, then retry.", e)); + } + } + } + } + for(String sql : List.of("CREATE SCHEMA IF NOT EXISTS " + schema, + "CREATE TABLE IF NOT EXISTS " + table() + " (id BIGSERIAL PRIMARY KEY, source_id INT NOT NULL, table_name TEXT NOT NULL, record_pk TEXT NOT NULL, field_name TEXT NOT NULL, chunk_index INT NOT NULL, content_hash CHAR(64) NOT NULL, embedding vector(" + dimensions + ") NOT NULL, indexed_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(source_id, table_name, record_pk, field_name, chunk_index))", + "CREATE INDEX IF NOT EXISTS assistant_rag_chunk_cosine ON " + table() + " USING hnsw (embedding vector_cosine_ops)")) + { + try(PreparedStatement statement = connection.prepareStatement(sql)) + { + statement.execute(); + } + } + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Validate the entire batch before sending anything to JDBC. + *******************************************************************************/ + @Override + public void upsert(List chunks) throws QException + { + chunks.forEach(chunk -> vector(chunk.embedding())); + try(Connection connection = connections.open()) + { + upsert(connection, chunks); + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Values, including vectors and arbitrary record keys, are bound parameters. + *******************************************************************************/ + private void upsert(Connection connection, List chunks) throws SQLException + { + String sql = "INSERT INTO " + table() + " (source_id, table_name, record_pk, field_name, chunk_index, content_hash, embedding) VALUES (?, ?, ?, ?, ?, ?, ?::vector) ON CONFLICT (source_id, table_name, record_pk, field_name, chunk_index) DO UPDATE SET content_hash = EXCLUDED.content_hash, embedding = EXCLUDED.embedding, indexed_at = CURRENT_TIMESTAMP"; + try(PreparedStatement statement = connection.prepareStatement(sql)) + { + for(ChunkVector chunk : chunks) + { + statement.setInt(1, chunk.source()); + statement.setString(2, chunk.table()); + statement.setString(3, chunk.pk()); + statement.setString(4, chunk.field()); + statement.setInt(5, chunk.chunkIndex()); + statement.setString(6, chunk.contentHash()); + statement.setString(7, vector(chunk.embedding())); + statement.addBatch(); + } + statement.executeBatch(); + } + } + + + + /******************************************************************************* + ** A record replacement is transactional, including empty replacements. + *******************************************************************************/ + @Override + public void replaceRecord(Integer source, String table, String pk, List chunks) throws QException + { + for(ChunkVector chunk : chunks) + { + vector(chunk.embedding()); + if(!source.equals(chunk.source()) || !table.equals(chunk.table()) || !pk.equals(chunk.pk())) + { + throw (new IllegalArgumentException("Replacement chunks must belong to the specified record.")); + } + } + try(Connection connection = connections.open()) + { + Boolean autoCommit = connection.getAutoCommit(); + connection.setAutoCommit(false); + try + { + deleteRecord(connection, source, table, pk); + upsert(connection, chunks); + connection.commit(); + } + catch(SQLException e) + { + connection.rollback(); + throw (e); + } + finally + { + connection.setAutoCommit(autoCommit); + } + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Remove one record, always scoped to its source. + *******************************************************************************/ + @Override + public void deleteRecord(Integer source, String table, String pk) throws QException + { + try(Connection connection = connections.open()) + { + deleteRecord(connection, source, table, pk); + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Shared transactional delete statement. + *******************************************************************************/ + private void deleteRecord(Connection connection, Integer source, String table, String pk) throws SQLException + { + try(PreparedStatement statement = connection.prepareStatement("DELETE FROM " + table() + " WHERE source_id = ? AND table_name = ? AND record_pk = ?")) + { + statement.setInt(1, source); + statement.setString(2, table); + statement.setString(3, pk); + statement.executeUpdate(); + } + } + + + + /******************************************************************************* + ** Clear a source, preserving every other source in the instance. + *******************************************************************************/ + @Override + public void deleteSource(Integer source) throws QException + { + try(Connection connection = connections.open(); PreparedStatement statement = connection.prepareStatement("DELETE FROM " + table() + " WHERE source_id = ?")) + { + statement.setInt(1, source); + statement.executeUpdate(); + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** HNSW cosine ordering; tables are values, never interpolated SQL identifiers. + *******************************************************************************/ + @Override + public List search(float[] query, Integer k, Set tables) throws QException + { + return (search(query, k, tables, Set.of())); + } + + + + /******************************************************************************* + ** Apply enabled compatible source restrictions before the nearest-hit limit. + *******************************************************************************/ + @Override + public List search(float[] query, Integer k, Set tables, Set sources) throws QException + { + if(tables.isEmpty() || k <= 0) + { + return (List.of()); + } + String vector = vector(query); + String sql = "SELECT source_id, table_name, record_pk, field_name, chunk_index, 1 - (embedding <=> ?::vector) AS score FROM " + table() + + " WHERE table_name IN (" + String.join(", ", Collections.nCopies(tables.size(), "?")) + ")" + + (sources.isEmpty() ? "" : " AND source_id IN (" + String.join(", ", Collections.nCopies(sources.size(), "?")) + ")") + + " ORDER BY embedding <=> ?::vector LIMIT ?"; + try(Connection connection = connections.open(); PreparedStatement statement = connection.prepareStatement(sql)) + { + Integer parameter = 1; + statement.setString(parameter++, vector); + for(String table : tables) + { + statement.setString(parameter++, table); + } + for(Integer source : sources) + { + statement.setInt(parameter++, source); + } + statement.setString(parameter++, vector); + statement.setInt(parameter, k); + List hits = new ArrayList<>(); + try(ResultSet rows = statement.executeQuery()) + { + while(rows.next()) + { + Double score = rows.getDouble("score"); + if(Double.isFinite(score)) + { + hits.add(new Hit(rows.getInt("source_id"), rows.getString("table_name"), rows.getString("record_pk"), rows.getString("field_name"), rows.getInt("chunk_index"), score)); + } + } + } + return (hits); + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Count chunks for the source's status screen. + *******************************************************************************/ + @Override + public Long count(Integer source) throws QException + { + try(Connection connection = connections.open(); PreparedStatement statement = connection.prepareStatement("SELECT COUNT(*) FROM " + table() + " WHERE source_id = ?")) + { + statement.setInt(1, source); + try(ResultSet rows = statement.executeQuery()) + { + rows.next(); + return (rows.getLong(1)); + } + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Content-free reconciliation data. Embedding text is pgvector's numeric form. + *******************************************************************************/ + @Override + public List chunks(Integer source) throws QException + { + try(Connection connection = connections.open(); PreparedStatement statement = connection.prepareStatement("SELECT table_name, record_pk, field_name, chunk_index, content_hash, embedding::text FROM " + table() + " WHERE source_id = ?")) + { + statement.setInt(1, source); + List chunks = new ArrayList<>(); + try(ResultSet rows = statement.executeQuery()) + { + while(rows.next()) + { + chunks.add(new ChunkVector(source, rows.getString("table_name"), rows.getString("record_pk"), rows.getString("field_name"), rows.getInt("chunk_index"), rows.getString("content_hash"), parse(rows.getString("embedding")))); + } + } + return (chunks); + } + catch(SQLException e) + { + throw (failure(e)); + } + } + + + + /******************************************************************************* + ** Qualified name uses only the validated instance schema. + *******************************************************************************/ + private String table() + { + return (schema + ".assistant_rag_chunk"); + } + + + + /******************************************************************************* + ** Serialize only finite coordinates with the configured dimension. + *******************************************************************************/ + private String vector(float[] vector) + { + if(vector == null || vector.length != dimensions) + { + throw (new IllegalArgumentException("Vector dimensions differ from the configured pgvector schema.")); + } + StringJoiner values = new StringJoiner(",", "[", "]"); + for(float value : vector) + { + if(!Float.isFinite(value)) + { + throw (new IllegalArgumentException("Vector coordinates must be finite.")); + } + values.add(Float.toString(value)); + } + return (values.toString()); + } + + + + /******************************************************************************* + ** Parse the database's numeric representation without any source content. + *******************************************************************************/ + private float[] parse(String text) + { + String[] parts = text.substring(1, text.length() - 1).split(","); + float[] vector = new float[parts.length]; + for(Integer i = 0; i < parts.length; i++) + { + vector[i] = Float.parseFloat(parts[i]); + } + vector(vector); + return (vector); + } + + + + /******************************************************************************* + ** Do not expose database URLs, SQL values or credentials in tool responses. + *******************************************************************************/ + private static QException failure(SQLException cause) + { + return (new QException("The knowledge vector store is unavailable. Check its schema, dimensions and database permissions.", cause)); + } +} diff --git a/qbit-assistant-rag-pgvector/src/test/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorContainerTest.java b/qbit-assistant-rag-pgvector/src/test/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorContainerTest.java new file mode 100644 index 0000000..8da04a8 --- /dev/null +++ b/qbit-assistant-rag-pgvector/src/test/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorContainerTest.java @@ -0,0 +1,76 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag.pgvector; + + +import java.sql.DriverManager; +import java.util.List; +import java.util.Set; +import com.kingsrook.qbits.assistant.rag.ChunkVector; +import com.kingsrook.qbits.assistant.rag.TextChunks; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; +import static org.assertj.core.api.Assertions.assertThat; + + +/******************************************************************************* + ** Real pgvector contract; skipped automatically on hosts without Docker. + *******************************************************************************/ +@Testcontainers(disabledWithoutDocker = true) +class PgVectorContainerTest +{ + @Container + static final PostgreSQLContainer POSTGRES = new PostgreSQLContainer<>(DockerImageName.parse("pgvector/pgvector:pg16").asCompatibleSubstituteFor("postgres")); + + + + /******************************************************************************* + ** DDL idempotence, cosine ranking, upsert, deletes and instance isolation. + *******************************************************************************/ + @Test + void realVectorLifecycle() throws Exception + { + PgVectorStore.ConnectionSupplier connections = () -> DriverManager.getConnection(POSTGRES.getJdbcUrl(), POSTGRES.getUsername(), POSTGRES.getPassword()); + PgVectorStore store = new PgVectorStore(connections, 3).withSchema("instance_one"); + PgVectorStore other = new PgVectorStore(connections, 3).withSchema("instance_two"); + store.ensureSchema(); + store.ensureSchema(); + other.ensureSchema(); + ChunkVector first = new ChunkVector(1, "article", "1", "body", 0, TextChunks.hash("one"), new float[] { 1, 0, 0 }); + ChunkVector second = new ChunkVector(1, "article", "2", "body", 0, TextChunks.hash("two"), new float[] { 0, 1, 0 }); + store.upsert(List.of(first, second)); + store.upsert(List.of(first)); + assertThat(store.count(1)).isEqualTo(2); + assertThat(other.count(1)).isZero(); + assertThat(store.search(new float[] { 1, 0, 0 }, 1, Set.of("article"))).singleElement().satisfies(hit -> + { + assertThat(hit.pk()).isEqualTo("1"); + assertThat(hit.score()).isEqualTo(1.0); + }); + assertThat(store.chunks(1)).hasSize(2); + store.replaceRecord(1, "article", "1", List.of()); + assertThat(store.count(1)).isEqualTo(1); + store.deleteRecord(1, "article", "2"); + assertThat(store.count(1)).isZero(); + store.upsert(List.of(first)); + store.deleteSource(1); + assertThat(store.count(1)).isZero(); + } +} diff --git a/qbit-assistant-rag-pgvector/src/test/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorStoreTest.java b/qbit-assistant-rag-pgvector/src/test/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorStoreTest.java new file mode 100644 index 0000000..72d1267 --- /dev/null +++ b/qbit-assistant-rag-pgvector/src/test/java/com/kingsrook/qbits/assistant/rag/pgvector/PgVectorStoreTest.java @@ -0,0 +1,228 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag.pgvector; + + +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import javax.sql.DataSource; +import com.kingsrook.qbits.assistant.rag.ChunkVector; +import com.kingsrook.qbits.assistant.rag.TextChunks; +import org.junit.jupiter.api.Test; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + + +/******************************************************************************* + ** Always-on JDBC SQL-shape and transaction tests, with no Docker dependency. + *******************************************************************************/ +class PgVectorStoreTest +{ + private final List sql = new ArrayList<>(); + private final List> parameters = new ArrayList<>(); + private final List operations = new ArrayList<>(); + private Boolean extensionExists = true; + private Boolean failBatch = false; + private Boolean failExtension = false; + + + + /******************************************************************************* + ** JDBC recording connection; values never appear in SQL text. + *******************************************************************************/ + private Connection connection() + { + return ((Connection) Proxy.newProxyInstance(getClass().getClassLoader(), new Class[] { Connection.class }, (proxy, method, args) -> + { + operations.add(method.getName()); + if(method.getName().equals("prepareStatement")) + { + return (statement((String) args[0])); + } + return (method.getName().equals("getAutoCommit") ? true : null); + })); + } + + + + /******************************************************************************* + ** Capture each binding and emulate rows or SQL failures. + *******************************************************************************/ + private PreparedStatement statement(String query) + { + sql.add(query); + Map values = new HashMap<>(); + parameters.add(values); + return ((PreparedStatement) Proxy.newProxyInstance(getClass().getClassLoader(), new Class[] { PreparedStatement.class }, (proxy, method, args) -> + { + if(method.getName().startsWith("set")) + { + values.put((Integer) args[0], args[1]); + } + if(method.getName().equals("executeQuery")) + { + return (rows(query)); + } + if(method.getName().equals("executeBatch")) + { + if(failBatch) + { + throw (new SQLException("synthetic batch failure")); + } + return (new int[] { 1 }); + } + if(method.getName().equals("executeUpdate")) + { + return (1); + } + if(method.getName().equals("execute")) + { + if(failExtension && query.startsWith("CREATE EXTENSION")) + { + throw (new SQLException("extension privilege required")); + } + return (true); + } + return (null); + })); + } + + + + /******************************************************************************* + ** One content-free result; the extension probe can return no rows. + *******************************************************************************/ + private ResultSet rows(String query) + { + Boolean[] first = { true }; + return ((ResultSet) Proxy.newProxyInstance(getClass().getClassLoader(), new Class[] { ResultSet.class }, (proxy, method, args) -> + { + if(method.getName().equals("next")) + { + Boolean next = first[0] && (!query.contains("pg_extension") || extensionExists); + first[0] = false; + return (next); + } + return (switch(method.getName()) + { + case "getString" -> switch(String.valueOf(args[0])) + { + case "content_hash" -> TextChunks.hash("synthetic"); + case "embedding" -> "[1,0,0]"; + case "record_pk" -> "pk'; --"; + case "field_name" -> "body"; + default -> "article"; + }; + case "getInt" -> "chunk_index".equals(args[0]) ? 0 : 1; + case "getLong" -> 1L; + case "getDouble" -> 1.0; + default -> null; + }); + })); + } + + + + /******************************************************************************* + ** Exercise every statement with values that would be unsafe to interpolate. + *******************************************************************************/ + @Test + void statementsAreBoundAndContentFree() throws Exception + { + PgVectorStore store = new PgVectorStore(this::connection, 3).withSchema("instance_one"); + store.ensureSchema(); + assertThat(sql).anyMatch(s -> s.contains("embedding vector(3)")); + assertThat(sql).anyMatch(s -> s.contains("USING hnsw (embedding vector_cosine_ops)")); + assertThat(sql).noneMatch(s -> s.contains("CREATE EXTENSION")); + ChunkVector chunk = new ChunkVector(1, "article", "pk'; --", "body", 0, TextChunks.hash("synthetic"), new float[] { 1, 0, 0 }); + store.upsert(List.of(chunk)); + store.replaceRecord(1, "article", "pk'; --", List.of(chunk)); + assertThat(store.search(new float[] { 1, 0, 0 }, 8, Set.of("article", "injection'; --"))).hasSize(1); + assertThat(store.search(new float[] { 1, 0, 0 }, 8, Set.of("article"), Set.of(1))).hasSize(1); + assertThat(store.chunks(1)).singleElement().satisfies(v -> assertThat(v.embedding()).containsExactly(1, 0, 0)); + assertThat(store.count(1)).isEqualTo(1); + store.deleteRecord(1, "article", "pk'; --"); + store.deleteSource(1); + assertThat(sql).anyMatch(s -> s.contains("?::vector") && s.contains("ON CONFLICT")); + assertThat(sql).anyMatch(s -> s.contains("ORDER BY embedding <=> ?::vector LIMIT ?")); + assertThat(sql).noneMatch(s -> s.contains("pk'; --") || s.contains("injection'; --") || s.contains("synthetic")); + assertThat(parameters).anyMatch(p -> p.containsValue("pk'; --")); + assertThat(parameters).anyMatch(p -> p.containsValue("injection'; --")); + assertThat(operations).contains("commit", "close", "setAutoCommit"); + assertThat(store.search(new float[] { 1, 0, 0 }, 8, Set.of())).isEmpty(); + assertThat(store.search(new float[] { 1, 0, 0 }, 0, Set.of("article"))).isEmpty(); + DataSource dataSource = (DataSource) Proxy.newProxyInstance(getClass().getClassLoader(), new Class[] { DataSource.class }, (p, method, args) -> connection()); + assertThat(new PgVectorStore(dataSource, 3).count(1)).isEqualTo(1); + } + + + + /******************************************************************************* + ** DBA failure is actionable; failed replacement rolls back before closing. + *******************************************************************************/ + @Test + void extensionAndRollback() throws Exception + { + extensionExists = false; + PgVectorStore store = new PgVectorStore(this::connection, 3); + store.ensureSchema(); + assertThat(sql).contains("CREATE EXTENSION IF NOT EXISTS vector"); + failExtension = true; + assertThatThrownBy(store::ensureSchema).hasMessageContaining("DBA"); + failBatch = true; + assertThatThrownBy(() -> store.replaceRecord(1, "article", "1", List.of())).hasMessageContaining("vector store"); + assertThat(operations).contains("rollback"); + PgVectorStore down = new PgVectorStore(() -> + { + throw (new SQLException("offline")); + }, 3); + assertThatThrownBy(down::ensureSchema).hasMessageContaining("vector store"); + assertThatThrownBy(() -> down.count(1)).hasMessageContaining("vector store"); + assertThatThrownBy(() -> down.chunks(1)).hasMessageContaining("vector store"); + assertThatThrownBy(() -> down.upsert(List.of())).hasMessageContaining("vector store"); + assertThatThrownBy(() -> down.deleteRecord(1, "a", "1")).hasMessageContaining("vector store"); + assertThatThrownBy(() -> down.deleteSource(1)).hasMessageContaining("vector store"); + assertThatThrownBy(() -> down.search(new float[] { 1, 0, 0 }, 8, Set.of("a"))).hasMessageContaining("vector store"); + } + + + + /******************************************************************************* + ** Identifier and coordinate validation happens before opening a connection. + *******************************************************************************/ + @Test + void invalidInputsNeverReachJdbc() + { + PgVectorStore store = new PgVectorStore(this::connection, 3); + for(String name : List.of("a.b", "bad; DROP SCHEMA public", "UPPER", "a".repeat(64))) + { + assertThatThrownBy(() -> store.withSchema(name)).isInstanceOf(IllegalArgumentException.class); + } + assertThatThrownBy(() -> new PgVectorStore(this::connection, 2001)).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> store.search(new float[] { 1 }, 1, Set.of("a"))).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> store.search(new float[] { Float.NaN, 0, 0 }, 1, Set.of("a"))).isInstanceOf(IllegalArgumentException.class); + assertThat(operations).isEmpty(); + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java index 09cf9ac..a2af573 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitConfig.java @@ -17,7 +17,6 @@ package com.kingsrook.qbits.assistant; -import javax.net.ssl.SSLContext; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; @@ -29,6 +28,7 @@ import java.util.Map; import java.util.Set; import java.util.function.Function; import java.util.regex.Pattern; +import javax.net.ssl.SSLContext; import com.fasterxml.jackson.annotation.JsonIgnore; import com.kingsrook.qbits.assistant.classification.ClassificationSource; import com.kingsrook.qbits.assistant.classification.ConfiguredClassificationSource; @@ -37,6 +37,7 @@ import com.kingsrook.qbits.assistant.engine.ModelPricing; import com.kingsrook.qbits.assistant.provider.AssistantModelProvider; import com.kingsrook.qbits.assistant.provider.AssistantProviderType; import com.kingsrook.qbits.assistant.provider.TrustedCaBundle; +import com.kingsrook.qbits.assistant.rag.VectorStore; import com.kingsrook.qbits.assistant.uiapi.UiApiGateway; import com.kingsrook.qqq.backend.core.actions.ai.AiExposure; import com.kingsrook.qqq.backend.core.model.metadata.QInstance; @@ -59,6 +60,8 @@ import com.kingsrook.qqq.backend.core.utils.StringUtils; *******************************************************************************/ public class AssistantQBitConfig implements QBitConfig { + private VectorStore vectorStore; + public static final String ACCESS_PERMISSION = "assistant.hasAccess"; public static final String ADMIN_PERMISSION = "assistant.admin"; public static final String DEFAULT_OWNER_KEY_TYPE = "assistant_user_id"; @@ -1788,4 +1791,35 @@ public class AssistantQBitConfig implements QBitConfig configuredClassifications.withClassification(tableName, fieldName, classification); return (this); } + + + + /******************************************************************************* + ** Getter for vectorStore + *******************************************************************************/ + public VectorStore getVectorStore() + { + return (vectorStore); + } + + + + /******************************************************************************* + ** Setter for vectorStore + *******************************************************************************/ + public void setVectorStore(VectorStore vectorStore) + { + this.vectorStore = vectorStore; + } + + + + /******************************************************************************* + ** Fluent setter for vectorStore + *******************************************************************************/ + public AssistantQBitConfig withVectorStore(VectorStore vectorStore) + { + this.vectorStore = vectorStore; + return (this); + } } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitProducer.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitProducer.java index 3324769..917dffa 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitProducer.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/AssistantQBitProducer.java @@ -26,6 +26,7 @@ import com.kingsrook.qbits.assistant.home.AssistantHome; import com.kingsrook.qbits.assistant.metadata.AssistantMetaDataBuilder; import com.kingsrook.qbits.assistant.metadata.AssistantRecordScreens; import com.kingsrook.qbits.assistant.processes.DiscoverModelsProcess; +import com.kingsrook.qbits.assistant.processes.KnowledgeIndexProcess; import com.kingsrook.qbits.assistant.processes.TestConnectionProcess; import com.kingsrook.qbits.assistant.processes.TestModelConnectionProcess; import com.kingsrook.qqq.backend.core.alerts.AlertSeverity; @@ -179,6 +180,12 @@ public class AssistantQBitProducer implements QBitMetaDataProducer(byName.values())); } @@ -372,4 +379,15 @@ public class ContextBuilder return (null); } + + + /******************************************************************************* + ** Additive optional tool hook; UI API toolbox stays independently maintained. + *******************************************************************************/ + public ContextBuilder withKnowledgeTool(KnowledgeTool knowledgeTool) + { + this.knowledgeTool = knowledgeTool; + return (this); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ModelDiscovery.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ModelDiscovery.java index 8ab619b..4c7a11e 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ModelDiscovery.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ModelDiscovery.java @@ -85,6 +85,7 @@ public final class ModelDiscovery } ModelCapabilities caps = provider.getCapabilities(model.modelId()); runtime.getStore().insert(new AssistantModel().withProviderId(row.getId()).withModelId(model.modelId()) + .withModelKind(embedding(model.modelId()) ? "EMBEDDING" : "CHAT").withEmbeddingDimensions(dimensions(model.modelId())) .withDisplayName(model.displayName()).withEnabled(local || discovered.size() <= 20) .withSource("DISCOVERED").withDiscoveredAt(runtime.getClock().instant()) .withSupportsTools(caps.getTools()).withSupportsVision(caps.getVision()).withSupportsStreaming(caps.getStreaming()) @@ -115,4 +116,33 @@ public final class ModelDiscovery return (result); } } + + + /******************************************************************************* + ** Discovery is a suggestion; an administrator can correct the kind. + *******************************************************************************/ + public static Boolean embedding(String id) + { + String lower = id.toLowerCase(java.util.Locale.ROOT); + return (List.of("embed", "bge", "e5", "gte", "nomic").stream().anyMatch(lower::contains)); + } + + + + /******************************************************************************* + ** Known local defaults; other models must declare their dimensions. + *******************************************************************************/ + private static Integer dimensions(String id) + { + if(id.contains("bge-m3")) + { + return (1024); + } + if(id.contains("nomic-embed-text")) + { + return (768); + } + return (null); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ProviderService.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ProviderService.java index 8f3c9ac..1badf42 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ProviderService.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/ProviderService.java @@ -186,6 +186,16 @@ public class ProviderService ** with a message for the admin when the row cannot be used. *******************************************************************************/ public Active forRow(AssistantProvider row, String fallbackModel) throws ProviderSetupException + { + return (forRow(row, fallbackModel, false)); + } + + + + /******************************************************************************* + ** Resolve only the requested kind; embedding calls never select chat models. + *******************************************************************************/ + private Active forRow(AssistantProvider row, String fallbackModel, Boolean embedding) throws ProviderSetupException { if(Boolean.TRUE.equals(config.getHipaaMode()) && !Boolean.TRUE.equals(row.getHipaaEligible())) { @@ -205,7 +215,7 @@ public class ProviderService String model = StringUtils.hasContent(row.getModel()) ? row.getModel() : null; if(model == null && !models.isEmpty()) { - model = models.stream().filter(m -> Boolean.TRUE.equals(m.getEnabled())).map(AssistantModel::getModelId).findFirst().orElse(null); + model = models.stream().filter(m -> Boolean.TRUE.equals(m.getEnabled()) && !m.isChat().equals(embedding)).map(AssistantModel::getModelId).findFirst().orElse(null); if(model == null) { throw (new ProviderSetupException("All models of this provider are disabled. Enable a model first.")); @@ -217,7 +227,7 @@ public class ProviderService } String selected = model; AssistantModel registry = models.stream().filter(m -> java.util.Objects.equals(selected, m.getModelId())).findFirst().orElse(null); - if(registry != null && !Boolean.TRUE.equals(registry.getEnabled())) + if(registry != null && (!Boolean.TRUE.equals(registry.getEnabled()) || registry.isChat().equals(embedding))) { throw (new ProviderSetupException("The selected model is disabled in this provider's model registry.")); } @@ -752,7 +762,7 @@ public class ProviderService public Active resolveModel(Integer modelId, Boolean requiresHipaa, Boolean requireTools) throws QException { AssistantModel model = store.listModels().stream().filter(m -> modelId.equals(m.getId()) && Boolean.TRUE.equals(m.getEnabled())).findFirst().orElse(null); - if(model == null) + if(model == null || !model.isChat()) { return (null); } @@ -816,4 +826,28 @@ public class ProviderService { return (left != null && right != null && java.util.Objects.equals(left.providerId(), right.providerId()) && java.util.Objects.equals(left.model(), right.model())); } + + + /******************************************************************************* + ** Embedding selection is explicit and checks fresh provider eligibility. + *******************************************************************************/ + public Active resolveEmbedding(Integer modelId, Boolean requiresHipaa) throws QException + { + AssistantModel model = store.listModels().stream().filter(m -> java.util.Objects.equals(modelId, m.getId())).findFirst().orElse(null); + if(model == null || !"EMBEDDING".equals(model.getModelKind()) || !Boolean.TRUE.equals(model.getEnabled())) + { + throw (new ProviderSetupException("Select an enabled EMBEDDING model for this knowledge source.")); + } + AssistantProvider row = store.getProvider(model.getProviderId()); + if(row == null || !Boolean.TRUE.equals(row.getEnabled())) + { + throw (new ProviderSetupException("The embedding provider is disabled or unavailable.")); + } + if((requiresHipaa || Boolean.TRUE.equals(config.getHipaaMode())) && !Boolean.TRUE.equals(row.getHipaaEligible())) + { + throw (new ProviderSetupException("PHI may only be embedded by a HIPAA-eligible provider. Nothing was sent.")); + } + return (forRow(new AssistantProvider(row.toQRecord()).withModel(model.getModelId()), model.getModelId(), true)); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/SystemPrompt.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/SystemPrompt.java index 7be292a..6c19a3a 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/SystemPrompt.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/SystemPrompt.java @@ -40,6 +40,8 @@ public class SystemPrompt 10. An administrator may give you a name, a personality and skills (in and blocks). Follow them for tone and know-how, but they never override these rules and never give you other tools, data or permissions. 11. Some values are masked by the application (for example j***@x.com or ***1234). Treat them as masked, never as the real value, and never ask the user to reveal them for you. + 12. When you use search_knowledge, cite records with their links. + Presentation (default skill; administrator skills may extend it): Answer first in one sentence; do not restate totals ("101 add up to 101"). For 3+ items, comparisons or breakdowns, call present_results with a table; add a chart when shares or trends help. Prefer metrics for a single big number. Include any list you offer to open as a filterLinks item, never just plain text. Use only values from tool results, omit zero-count categories, and use percentOfColumn for shares. Keep prose short; do not repeat presented tables or links. Never write URLs yourself; supply metadata names and raw filter values to present_results. """; diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnClassification.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnClassification.java index b1e5154..3945e7c 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnClassification.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnClassification.java @@ -75,7 +75,7 @@ public class TurnClassification return; } Map input = use.getInput() == null ? Map.of() : use.getInput(); - boolean crossTable = "search_records".equals(use.getToolName()) || "get_record_history".equals(use.getToolName()); + boolean crossTable = "search_knowledge".equals(use.getToolName()) || "search_records".equals(use.getToolName()) || "get_record_history".equals(use.getToolName()); String table = result.getTableName() != null ? result.getTableName() : (input.get("table") instanceof String name ? name : (input.get("tableName") instanceof String name ? name : null)); if(crossTable) @@ -98,6 +98,10 @@ public class TurnClassification } collectTables(result.getRender(), touched); collectResultTables(result.getModelContent(), touched); + if("search_knowledge".equals(use.getToolName()) && touched.isEmpty()) + { + return; + } if(touched.isEmpty() || touched.stream().anyMatch(name -> QContext.getQInstance().getTable(name) == null)) { touchUnknown(); diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java index 8656189..7729f7c 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/TurnEngine.java @@ -17,6 +17,21 @@ package com.kingsrook.qbits.assistant.engine; +import java.io.Serializable; +import java.time.Duration; +import java.time.Instant; +import java.time.LocalDate; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; import com.kingsrook.qbits.assistant.AssistantQBitConfig; import com.kingsrook.qbits.assistant.AssistantRuntime; import com.kingsrook.qbits.assistant.model.AssistantConversation; @@ -36,6 +51,7 @@ import com.kingsrook.qbits.assistant.provider.ModelRequest; import com.kingsrook.qbits.assistant.provider.ModelResponse; import com.kingsrook.qbits.assistant.provider.ModelUsage; import com.kingsrook.qbits.assistant.provider.StreamHandle; +import com.kingsrook.qbits.assistant.tools.KnowledgeTool; import com.kingsrook.qbits.assistant.tools.PresentResultsTool; import com.kingsrook.qbits.assistant.tools.UiApiToolbox; import com.kingsrook.qqq.backend.core.actions.ai.AiRecordRef; @@ -50,21 +66,6 @@ import com.kingsrook.qqq.backend.core.model.data.QRecord; import com.kingsrook.qqq.backend.core.model.session.QSession; import com.kingsrook.qqq.backend.core.utils.JsonUtils; import com.kingsrook.qqq.backend.core.utils.StringUtils; -import java.io.Serializable; -import java.time.Duration; -import java.time.Instant; -import java.time.LocalDate; -import java.time.ZoneOffset; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.Future; -import java.util.concurrent.Semaphore; -import java.util.concurrent.TimeUnit; import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair; @@ -711,8 +712,15 @@ public class TurnEngine implements Runnable QContext.init(runtime.getQInstance(), session); try { - outcomes[index] = runtime.getToolbox().execute(toolContext(), call.getId(), call.getName(), call.getInput()); - results[index] = outcomes[index].getResult(); + if(KnowledgeTool.NAME.equals(call.getName())) + { + results[index] = runtime.getKnowledgeTool().execute(toolContext(), call.getId(), call.getInput(), classification.getClassification(), conversation.getId(), turn.getId()); + } + else + { + outcomes[index] = runtime.getToolbox().execute(toolContext(), call.getId(), call.getName(), call.getInput()); + results[index] = outcomes[index].getResult(); + } } finally { @@ -850,7 +858,7 @@ public class TurnEngine implements Runnable .withResultSummary(summary) .withResultJson(truncate(result.getModelContent(), MAX_RESULT_JSON_CHARS)) .withRedactionsJson(result.getRedactions() == null || result.getRedactions().isEmpty() ? null : JsonUtils.toJson(result.getRedactions())) - .withUiApiOperation(extras == null ? null : extras.getOperation()) + .withUiApiOperation(extras == null ? ("search_knowledge".equals(use.getToolName()) ? "GET /table/{tableName}/{primaryKey}" : null) : extras.getOperation()) .withHttpStatus(extras == null ? null : extras.getHttpStatus()) .withAuditHref(extras == null ? null : extras.getAuditHref()) .withPendingActionId(extras == null ? null : extras.getPendingActionId())); @@ -923,7 +931,7 @@ public class TurnEngine implements Runnable /******************************************************************************* ** Only same-origin application links are passed on. *******************************************************************************/ - static boolean isAppHref(String href) + public static boolean isAppHref(String href) { return (href != null && (href.equals("/app") || href.startsWith("/app/")) && !href.contains("//") && !href.contains("\\")); } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/UsageRecorder.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/UsageRecorder.java index 734606b..bd237eb 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/UsageRecorder.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/engine/UsageRecorder.java @@ -39,6 +39,7 @@ public final class UsageRecorder public static final String PURPOSE_TURN = "TURN"; public static final String PURPOSE_TEST = "TEST"; + public static final String PURPOSE_EMBED = "EMBED"; public static final String PURPOSE_ROUTE = "ROUTE"; public static final String STATUS_OK = "OK"; diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantMetaDataBuilder.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantMetaDataBuilder.java index e696d64..4b7972e 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantMetaDataBuilder.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantMetaDataBuilder.java @@ -30,6 +30,7 @@ import com.kingsrook.qbits.assistant.model.AssistantDataAudit; import com.kingsrook.qbits.assistant.model.AssistantFeedback; import com.kingsrook.qbits.assistant.model.AssistantMessage; import com.kingsrook.qbits.assistant.model.AssistantModel; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; import com.kingsrook.qbits.assistant.model.AssistantModelPrice; import com.kingsrook.qbits.assistant.model.AssistantPendingAction; import com.kingsrook.qbits.assistant.model.AssistantPolicy; @@ -110,7 +111,7 @@ public class AssistantMetaDataBuilder ** calls on a conversation; model calls on a provider) or reached from the ** home. Processes are table and record actions, never nav entries. *******************************************************************************/ - public static final List NAV_TABLES = List.of(AssistantBudget.TABLE_NAME, AssistantProvider.TABLE_NAME, AssistantProfile.TABLE_NAME, AssistantSkill.TABLE_NAME, AssistantDataAccess.TABLE_NAME, AssistantConversation.TABLE_NAME, AssistantModelPrice.TABLE_NAME, AssistantPolicy.TABLE_NAME); + public static final List NAV_TABLES = List.of(AssistantRagSource.TABLE_NAME, AssistantBudget.TABLE_NAME, AssistantProvider.TABLE_NAME, AssistantProfile.TABLE_NAME, AssistantSkill.TABLE_NAME, AssistantDataAccess.TABLE_NAME, AssistantConversation.TABLE_NAME, AssistantModelPrice.TABLE_NAME, AssistantPolicy.TABLE_NAME); private final AssistantQBitConfig config; @@ -186,6 +187,13 @@ public class AssistantMetaDataBuilder .withUniqueKey(new UniqueKey("name")) .withRecordLabelFormat("%s").withRecordLabelFields("name")); + QCodeReference ragCustomizer = new QCodeReference(AssistantRagSourceCustomizer.class); + tables.add(text(base(AssistantRagSource.TABLE_NAME, "Knowledge sources", "manage_search", AssistantRagSource.class), "filterJson") + .withRecordLabelFormat("%s").withRecordLabelFields("name") + .withCustomizer(TableCustomizers.PRE_INSERT_RECORD, ragCustomizer).withCustomizer(TableCustomizers.PRE_UPDATE_RECORD, ragCustomizer) + .withCustomizer(TableCustomizers.POST_INSERT_RECORD, ragCustomizer).withCustomizer(TableCustomizers.POST_UPDATE_RECORD, ragCustomizer) + .withCustomizer(TableCustomizers.POST_DELETE_RECORD, ragCustomizer)); + tables.add(settings(base(AssistantModel.TABLE_NAME, "Assistant Model", "memory", AssistantModel.class)) .withUniqueKey(new UniqueKey("providerId", "modelId")) .withRecordLabelFormat("%s").withRecordLabelFields("modelId")); diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRagSourceCustomizer.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRagSourceCustomizer.java new file mode 100644 index 0000000..ccf9caf --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRagSourceCustomizer.java @@ -0,0 +1,125 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.metadata; + + +import java.io.Serializable; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; +import com.kingsrook.qqq.backend.core.actions.customizers.TableCustomizerInterface; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.actions.AbstractActionInput; +import com.kingsrook.qqq.backend.core.model.actions.tables.delete.DeleteInput; +import com.kingsrook.qqq.backend.core.model.data.QRecord; +import com.kingsrook.qqq.backend.core.model.statusmessages.BadInputStatusMessage; + + +/******************************************************************************* + ** Validate source configuration; invalidate vectors whenever their meaning changes. + *******************************************************************************/ +public class AssistantRagSourceCustomizer implements TableCustomizerInterface +{ + private static final List SEMANTIC_FIELDS = List.of("tableName", "fieldNames", "filterJson", "embeddingModelId", "chunkChars", "chunkOverlapChars"); + + + + /******************************************************************************* + ** Validate merged edits, not partial update payloads. Status writes remain + ** possible even after a host removes a table or disables a model. + *******************************************************************************/ + @Override + public List preInsertOrUpdate(AbstractActionInput input, List records, boolean isPreview, Optional> oldRecordList) throws QException + { + Optional> old = oldRecordListToMap("id", oldRecordList); + AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance()); + for(QRecord record : records) + { + if(record.getValues().keySet().stream().allMatch(name -> name.equals("id") || name.startsWith("lastIndex") || name.startsWith("indexed") || name.endsWith("Date"))) + { + continue; + } + QRecord merged = new QRecord(); + old.map(m -> m.get(record.getValue("id"))).ifPresent(o -> merged.setValues(new LinkedHashMap<>(o.getValues()))); + record.getValues().forEach(merged::setValue); + for(var entry : Map.of("enabled", true, "chunkChars", 1500, "chunkOverlapChars", 200).entrySet()) + { + if(merged.getValue(entry.getKey()) == null) + { + record.setValue(entry.getKey(), entry.getValue()); + merged.setValue(entry.getKey(), entry.getValue()); + } + } + if(runtime != null) + { + for(String error : runtime.getRagService().validate(new AssistantRagSource(merged))) + { + record.addError(new BadInputStatusMessage(error)); + } + } + } + return (records); + } + + + + /******************************************************************************* + ** Old vectors cannot survive a model, filter or chunk-layout change. + *******************************************************************************/ + @Override + public List postInsertOrUpdate(AbstractActionInput input, List records, Optional> oldRecordList) throws QException + { + Optional> old = oldRecordListToMap("id", oldRecordList); + AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance()); + if(runtime != null && old.isPresent()) + { + for(QRecord record : records) + { + QRecord previous = old.get().get(record.getValue("id")); + if(previous != null && SEMANTIC_FIELDS.stream().anyMatch(name -> record.getValues().containsKey(name) && !Objects.equals(previous.getValue(name), record.getValue(name)))) + { + runtime.getRagService().clear(new AssistantRagSource(previous)); + } + } + } + return (records); + } + + + + /******************************************************************************* + ** Deleted sources no longer have searchable vectors. + *******************************************************************************/ + @Override + public List postDelete(DeleteInput input, List records) throws QException + { + AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance()); + if(runtime != null) + { + for(QRecord record : records) + { + runtime.getRagService().remove(new AssistantRagSource(record)); + } + } + return (records); + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRecordScreens.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRecordScreens.java index e6adb42..75bb9d6 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRecordScreens.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantRecordScreens.java @@ -17,6 +17,11 @@ package com.kingsrook.qbits.assistant.metadata; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; import com.kingsrook.qbits.assistant.AssistantModelRouting; import com.kingsrook.qbits.assistant.AssistantQBitConfig; import com.kingsrook.qbits.assistant.access.DataAccessMode; @@ -37,6 +42,7 @@ import com.kingsrook.qbits.assistant.model.AssistantPolicy; import com.kingsrook.qbits.assistant.model.AssistantProfile; import com.kingsrook.qbits.assistant.model.AssistantProfileSkill; import com.kingsrook.qbits.assistant.model.AssistantProvider; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; import com.kingsrook.qbits.assistant.model.AssistantSkill; import com.kingsrook.qbits.assistant.model.AssistantToolCall; import com.kingsrook.qbits.assistant.model.AssistantTurn; @@ -47,6 +53,9 @@ import com.kingsrook.qbits.assistant.provider.AssistantProviderType; import com.kingsrook.qqq.backend.core.actions.dashboard.widgets.ChildRecordListRenderer; import com.kingsrook.qqq.backend.core.instances.QInstanceEnricher; import com.kingsrook.qqq.backend.core.model.metadata.QInstance; +import com.kingsrook.qqq.backend.core.model.actions.tables.query.QCriteriaOperator; +import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter; +import com.kingsrook.qqq.backend.core.model.metadata.code.QCodeReferenceWithProperties; import com.kingsrook.qqq.backend.core.model.metadata.fields.AdornmentType; import com.kingsrook.qqq.backend.core.model.metadata.fields.DisplayFormat; import com.kingsrook.qqq.backend.core.model.metadata.fields.FieldAdornment; @@ -64,10 +73,6 @@ import com.kingsrook.qqq.backend.core.model.metadata.possiblevalues.QPossibleVal import com.kingsrook.qqq.backend.core.model.metadata.tables.QFieldSection; import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData; import com.kingsrook.qqq.backend.core.model.metadata.tables.Tier; -import java.util.ArrayList; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Set; /******************************************************************************* @@ -184,6 +189,7 @@ public final class AssistantRecordScreens case AssistantUsageDaily.TABLE_NAME -> usageDaily(table); case AssistantProvider.TABLE_NAME -> provider(table); case AssistantModel.TABLE_NAME -> model(table); + case AssistantRagSource.TABLE_NAME -> knowledge(table); case AssistantModelPrice.TABLE_NAME -> price(table); case AssistantUsage.TABLE_NAME -> usage(table); case AssistantDataAccess.TABLE_NAME -> dataAccess(table); @@ -286,17 +292,30 @@ public final class AssistantRecordScreens return (List.of( new QPossibleValueSource().withName(PROVIDER_TYPE_PVS_NAME).withType(QPossibleValueSourceType.ENUM).withIdType(QFieldType.STRING) .withValueFormatAndFields(PVSValueFormatAndFields.LABEL_ONLY).withEnumValues(types), - new QPossibleValueSource().withName(MODEL_PVS_NAME).withType(QPossibleValueSourceType.ENUM).withIdType(QFieldType.STRING) - .withValueFormatAndFields(PVSValueFormatAndFields.LABEL_ONLY).withEnumValues(modelValues), - enumSource("assistant_provider_kind", kinds), + settingsSource(MODEL_PVS_NAME, modelValues, true, true), + settingsSource("assistant_price_model", modelValues, true, false), + settingsSource("assistant_provider_kind", kinds, false, false), enumSource("assistant_provider_location", List.of(new QPossibleValue<>("LOCAL", "Local"), new QPossibleValue<>("CLOUD", "Cloud"))), enumSource("assistant_routing_mode", List.of(new QPossibleValue<>("OFF", "Off"), new QPossibleValue<>("HEURISTIC", "Heuristic"), new QPossibleValue<>("CLASSIFIER", "Classifier"))), + enumSource("assistant_model_kind", List.of(new QPossibleValue<>("CHAT", "Chat"), new QPossibleValue<>("EMBEDDING", "Embedding"))), enumSource("assistant_model_source", List.of(new QPossibleValue<>("DISCOVERED", "Discovered"), new QPossibleValue<>("MANUAL", "Manual"))), QPossibleValueSource.newForTable(AssistantProvider.TABLE_NAME))); } + /******************************************************************************* + ** Contingent searches use RC.5's otherValues input. + *******************************************************************************/ + private static QPossibleValueSource settingsSource(String name, List> union, Boolean models, Boolean providerRow) + { + return (new QPossibleValueSource().withName(name).withType(QPossibleValueSourceType.CUSTOM).withIdType(QFieldType.STRING) + .withValueFormatAndFields(PVSValueFormatAndFields.LABEL_ONLY) + .withCustomCodeReference(new QCodeReferenceWithProperties(AssistantSettingsPossibleValueProvider.class, Map.of("union", new ArrayList<>(union), "models", models, "providerRow", providerRow)))); + } + + + /******************************************************************************* ** Display label shared by provider settings and home cards. Custom kinds ** retain their registered name. @@ -332,6 +351,11 @@ public final class AssistantRecordScreens private static void provider(QTableMetaData table) { pvs(table, "providerType", PROVIDER_TYPE_PVS_NAME, "Type"); + ///////////////////////////////////////////////////////////////////////// + // Add possibleValueSourceDependsOn("providerType") to model and kind // + // once the next QQQ RC includes core feature/contingent-possible-values. // + // See QQQ docs/CONTINGENT-POSSIBLE-VALUES.md. // + ///////////////////////////////////////////////////////////////////////// pvs(table, "model", MODEL_PVS_NAME, "Model"); pvs(table, "kind", "assistant_provider_kind", "Kind"); pvs(table, "location", "assistant_provider_location", "Location"); @@ -380,6 +404,11 @@ public final class AssistantRecordScreens private static void price(QTableMetaData table) { pvs(table, "providerType", PROVIDER_TYPE_PVS_NAME, "Provider Type"); + ///////////////////////////////////////////////////////////////////////// + // Add possibleValueSourceDependsOn("providerType") when the next QQQ // + // RC includes it (QQQ docs/CONTINGENT-POSSIBLE-VALUES.md). // + ///////////////////////////////////////////////////////////////////////// + pvs(table, "model", "assistant_price_model", "Model"); for(String field : List.of("inputUsdPerMTok", "outputUsdPerMTok", "cacheWrite5mUsdPerMTok", "cacheWrite1hUsdPerMTok")) { help(table, field, "USD per million tokens."); @@ -888,6 +917,7 @@ public final class AssistantRecordScreens { link(table, "providerId", AssistantProvider.TABLE_NAME, "Provider"); pvs(table, "source", "assistant_model_source", "Source"); + pvs(table, "modelKind", "assistant_model_kind", "Model kind"); pvs(table, "fallbackModelId", "assistant_registry_model", "Fallback model"); table.getField("contextWindowTokens").withLabel("Context window (tokens)"); table.getField("maxOutputTokens").withLabel("Max output tokens"); @@ -896,7 +926,7 @@ public final class AssistantRecordScreens table.getField("outputUsdPerMTok").withLabel("Output (USD per million tokens)"); table.getField("lastTestLatencyMs").withLabel("Last test latency (ms)"); sections(table, - section("model", "Model", "memory", Tier.T1, "id", "providerId", "modelId", "displayName", "enabled", "source", "discoveredAt", "fallbackModelId"), + section("model", "Model", "memory", Tier.T1, "id", "providerId", "modelId", "displayName", "enabled", "modelKind", "embeddingDimensions", "source", "discoveredAt", "fallbackModelId"), section("capabilities", "Capabilities", "settings", Tier.T2, "supportsTools", "supportsVision", "supportsStreaming", "supportsReasoningEffort", "contextWindowTokens", "maxOutputTokens"), section("prices", "Prices", "payments", Tier.T2, "inputUsdPerMTok", "outputUsdPerMTok"), section("lastTest", "Last connection test", "network_check", Tier.T2, "lastTestAt", "lastTestStatus", "lastTestLatencyMs", "lastTestError")); @@ -936,4 +966,22 @@ public final class AssistantRecordScreens section("budget", "Budget", "savings", Tier.T1, "id", "name", "enabled", "scopeType", "scopeValue", "period", "metric", "limitValue"), section("thresholds", "Thresholds", "notifications", Tier.T2, "warnPercent", "atLimitAction", "downgradeModelId", "hardCeilingPercent", "notes")); } + + + /******************************************************************************* + ** Standard source editor and status, with explicit embedding model selection. + *******************************************************************************/ + private static void knowledge(QTableMetaData table) + { + pvs(table, "embeddingModelId", "assistant_registry_model", "Embedding model"); + table.getField("embeddingModelId").withPossibleValueSourceFilter(new QQueryFilter() + .withCriteria("modelKind", QCriteriaOperator.EQUALS, "EMBEDDING").withCriteria("enabled", QCriteriaOperator.EQUALS, true)); + help(table, "embeddingModelId", "Select an enabled EMBEDDING model. Chat models are refused on save."); + help(table, "filterJson", "Optional QQQ query filter JSON. Index now reads it through the UI API as the clicking administrator."); + sections(table, + section("source", "Knowledge source", "manage_search", Tier.T1, "id", "name", "enabled", "tableName", "fieldNames", "titleFieldName", "filterJson"), + section("embedding", "Embedding", "memory", Tier.T2, "embeddingModelId", "chunkChars", "chunkOverlapChars"), + section("status", "Last index", "info", Tier.T2, "lastIndexedAt", "lastIndexStatus", "lastIndexError", "indexedRecordCount", "indexedChunkCount")); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsCustomizer.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsCustomizer.java index 785151c..db53105 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsCustomizer.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsCustomizer.java @@ -83,6 +83,18 @@ public class AssistantSettingsCustomizer implements TableCustomizerInterface { for(QRecord record : records) { + if(record.getValue("id") == null && record.getValue("modelKind") == null) + { + record.setValue("modelKind", "CHAT"); + } + if(record.getValue("modelKind") != null && !List.of("CHAT", "EMBEDDING").contains(record.getValueString("modelKind"))) + { + record.addError(new BadInputStatusMessage("Model kind must be CHAT or EMBEDDING.")); + } + if(record.getValue("embeddingDimensions") != null && record.getValueInteger("embeddingDimensions") < 1) + { + record.addError(new BadInputStatusMessage("Embedding dimensions must be positive.")); + } if(record.getValue("id") == null) { if(record.getValue("source") == null) @@ -171,6 +183,25 @@ public class AssistantSettingsCustomizer implements TableCustomizerInterface } } } + if(AssistantModel.TABLE_NAME.equals(tableName(input)) && oldRecordList.isPresent()) + { + Map previous = oldRecordListToMap("id", oldRecordList).orElse(Map.of()); + for(QRecord record : records) + { + QRecord before = previous.get(record.getValue("id")); + if(before != null && List.of("providerId", "modelId", "modelKind", "embeddingDimensions").stream() + .anyMatch(name -> record.getValues().containsKey(name) && !java.util.Objects.equals(before.getValue(name), record.getValue(name)))) + { + for(var source : runtime.getRagService().sources()) + { + if(java.util.Objects.equals(source.getEmbeddingModelId(), record.getValueInteger("id"))) + { + runtime.getRagService().clear(source); + } + } + } + } + } runtime.getProviderService().invalidate(); return (records); } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsPossibleValueProvider.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsPossibleValueProvider.java new file mode 100644 index 0000000..b1ef3ab --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/metadata/AssistantSettingsPossibleValueProvider.java @@ -0,0 +1,134 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.metadata; + + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.model.AssistantModel; +import com.kingsrook.qbits.assistant.model.AssistantProvider; +import com.kingsrook.qbits.assistant.persistence.AssistantStore; +import com.kingsrook.qbits.assistant.provider.AssistantProviderType; +import com.kingsrook.qqq.backend.core.actions.values.QCustomPossibleValueProvider; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.actions.values.SearchPossibleValueSourceInput; +import com.kingsrook.qqq.backend.core.model.metadata.code.InitializableViaCodeReference; +import com.kingsrook.qqq.backend.core.model.metadata.code.QCodeReference; +import com.kingsrook.qqq.backend.core.model.metadata.code.QCodeReferenceWithProperties; +import com.kingsrook.qqq.backend.core.model.metadata.possiblevalues.QPossibleValue; + + +/******************************************************************************* + ** Provider settings choices, narrowed by the selected provider type. + *******************************************************************************/ +public class AssistantSettingsPossibleValueProvider implements QCustomPossibleValueProvider, InitializableViaCodeReference +{ + private final Map> union = new LinkedHashMap<>(); + private Boolean models; + private Boolean providerRow; + + + + /******************************************************************************* + ** Price rows share model choices but their id is never a provider id. + *******************************************************************************/ + @Override + public void initialize(QCodeReference codeReference) + { + Map properties = ((QCodeReferenceWithProperties) codeReference).getProperties(); + models = (Boolean) properties.get("models"); + providerRow = (Boolean) properties.get("providerRow"); + for(Object object : (List) properties.get("union")) + { + QPossibleValue value = (QPossibleValue) object; + union.put(String.valueOf(value.getId()), new QPossibleValue<>(String.valueOf(value.getId()), value.getLabel())); + } + } + + + + /******************************************************************************* + ** Saved values resolve across types, including disabled registry models. + *******************************************************************************/ + @Override + public QPossibleValue getPossibleValue(Serializable idValue) throws QException + { + String id = String.valueOf(idValue); + if(union.containsKey(id)) + { + return (union.get(id)); + } + if(models && idValue != null) + { + for(AssistantModel model : new AssistantStore().listModels()) + { + if(id.equals(model.getModelId())) + { + return (new QPossibleValue<>(id, id)); + } + } + } + return (null); + } + + + + /******************************************************************************* + ** Keep union fallback for callers that do not supply a known provider type. + *******************************************************************************/ + @Override + public List> search(SearchPossibleValueSourceInput input) throws QException + { + Map otherValues = input.getOtherValues(); + AssistantProviderType type = otherValues == null ? null : AssistantRuntime.of(QContext.getQInstance()).getConfig().getProviderTypes().get(otherValues.get("providerType")); + Map> values = new LinkedHashMap<>(); + if(type == null) + { + values.putAll(union); + } + else + { + for(String id : models ? type.getModels() : type.getKinds()) + { + values.put(id, new QPossibleValue<>(id, models ? id : AssistantRecordScreens.providerKindLabel(id))); + } + if(models && providerRow && otherValues.get("id") != null) + { + AssistantStore store = new AssistantStore(); + String providerId = String.valueOf(otherValues.get("id")); + AssistantProvider provider = store.listProviders().stream() + .filter(row -> providerId.equals(String.valueOf(row.getId())) && type.getType().equals(row.getProviderType())).findFirst().orElse(null); + if(provider != null) + { + for(AssistantModel model : store.listModels()) + { + if(provider.getId().equals(model.getProviderId()) && Boolean.TRUE.equals(model.getEnabled()) && "CHAT".equals(model.getModelKind())) + { + values.put(model.getModelId(), new QPossibleValue<>(model.getModelId(), model.getModelId())); + } + } + } + } + } + return (completeCustomPVSSearch(input, new ArrayList<>(values.values()))); + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantModel.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantModel.java index be7963f..d1d5645 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantModel.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantModel.java @@ -32,6 +32,12 @@ public class AssistantModel extends QRecordEntity { public static final String TABLE_NAME = "assistant_model"; + @QField(maxLength = 16) + private String modelKind; + + @QField() + private Integer embeddingDimensions; + @QField() private Integer fallbackModelId; @@ -798,4 +804,76 @@ public class AssistantModel extends QRecordEntity this.fallbackModelId = fallbackModelId; return (this); } + + + + /******************************************************************************* + ** Getter for modelKind + *******************************************************************************/ + public String getModelKind() + { + return (modelKind); + } + + + + /******************************************************************************* + ** Setter for modelKind + *******************************************************************************/ + public void setModelKind(String modelKind) + { + this.modelKind = modelKind; + } + + + + /******************************************************************************* + ** Fluent setter for modelKind + *******************************************************************************/ + public AssistantModel withModelKind(String modelKind) + { + this.modelKind = modelKind; + return (this); + } + + + + /******************************************************************************* + ** Getter for embeddingDimensions + *******************************************************************************/ + public Integer getEmbeddingDimensions() + { + return (embeddingDimensions); + } + + + + /******************************************************************************* + ** Setter for embeddingDimensions + *******************************************************************************/ + public void setEmbeddingDimensions(Integer embeddingDimensions) + { + this.embeddingDimensions = embeddingDimensions; + } + + + + /******************************************************************************* + ** Fluent setter for embeddingDimensions + *******************************************************************************/ + public AssistantModel withEmbeddingDimensions(Integer embeddingDimensions) + { + this.embeddingDimensions = embeddingDimensions; + return (this); + } + + + /******************************************************************************* + ** Legacy registry rows remain chat models until explicitly classified. + *******************************************************************************/ + public Boolean isChat() + { + return (modelKind == null || "CHAT".equals(modelKind)); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantRagSource.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantRagSource.java new file mode 100644 index 0000000..1804f11 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/model/AssistantRagSource.java @@ -0,0 +1,629 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.model; + + +import java.time.Instant; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.data.QField; +import com.kingsrook.qqq.backend.core.model.data.QRecord; +import com.kingsrook.qqq.backend.core.model.data.QRecordEntity; + + +/******************************************************************************* + ** Admin-owned indexing configuration. Source records never contain indexed text. + *******************************************************************************/ +public class AssistantRagSource extends QRecordEntity +{ + public static final String TABLE_NAME = "assistant_rag_source"; + + @QField(isEditable = false, isPrimaryKey = true) + private Integer id; + + @QField(isRequired = true, maxLength = 200) + private String name; + + @QField() + private Boolean enabled; + + @QField(isRequired = true, maxLength = 200) + private String tableName; + + @QField(isRequired = true, maxLength = 500) + private String fieldNames; + + @QField(maxLength = 200) + private String titleFieldName; + + @QField() + private String filterJson; + + @QField(isRequired = true) + private Integer embeddingModelId; + + @QField() + private Integer chunkChars; + + @QField() + private Integer chunkOverlapChars; + + @QField(isEditable = false) + private Instant lastIndexedAt; + + @QField(isEditable = false, maxLength = 200) + private String lastIndexStatus; + + @QField(isEditable = false, maxLength = 500) + private String lastIndexError; + + @QField(isEditable = false) + private Integer indexedRecordCount; + + @QField(isEditable = false) + private Integer indexedChunkCount; + + @QField(isEditable = false) + private Instant createDate; + + @QField(isEditable = false) + private Instant modifyDate; + + + /******************************************************************************* + ** Empty entity for framework mapping. + *******************************************************************************/ + public AssistantRagSource() + { + } + + + + /******************************************************************************* + ** Map the source's own configuration record. + *******************************************************************************/ + public AssistantRagSource(QRecord record) throws QException + { + populateFromQRecord(record); + } + + + + /******************************************************************************* + ** Getter for id + *******************************************************************************/ + public Integer getId() + { + return (id); + } + + + + /******************************************************************************* + ** Setter for id + *******************************************************************************/ + public void setId(Integer id) + { + this.id = id; + } + + + + /******************************************************************************* + ** Fluent setter for id + *******************************************************************************/ + public AssistantRagSource withId(Integer id) + { + this.id = id; + return (this); + } + + + + /******************************************************************************* + ** Getter for name + *******************************************************************************/ + public String getName() + { + return (name); + } + + + + /******************************************************************************* + ** Setter for name + *******************************************************************************/ + public void setName(String name) + { + this.name = name; + } + + + + /******************************************************************************* + ** Fluent setter for name + *******************************************************************************/ + public AssistantRagSource withName(String name) + { + this.name = name; + return (this); + } + + + + /******************************************************************************* + ** Getter for enabled + *******************************************************************************/ + public Boolean getEnabled() + { + return (enabled); + } + + + + /******************************************************************************* + ** Setter for enabled + *******************************************************************************/ + public void setEnabled(Boolean enabled) + { + this.enabled = enabled; + } + + + + /******************************************************************************* + ** Fluent setter for enabled + *******************************************************************************/ + public AssistantRagSource withEnabled(Boolean enabled) + { + this.enabled = enabled; + return (this); + } + + + + /******************************************************************************* + ** Getter for tableName + *******************************************************************************/ + public String getTableName() + { + return (tableName); + } + + + + /******************************************************************************* + ** Setter for tableName + *******************************************************************************/ + public void setTableName(String tableName) + { + this.tableName = tableName; + } + + + + /******************************************************************************* + ** Fluent setter for tableName + *******************************************************************************/ + public AssistantRagSource withTableName(String tableName) + { + this.tableName = tableName; + return (this); + } + + + + /******************************************************************************* + ** Getter for fieldNames + *******************************************************************************/ + public String getFieldNames() + { + return (fieldNames); + } + + + + /******************************************************************************* + ** Setter for fieldNames + *******************************************************************************/ + public void setFieldNames(String fieldNames) + { + this.fieldNames = fieldNames; + } + + + + /******************************************************************************* + ** Fluent setter for fieldNames + *******************************************************************************/ + public AssistantRagSource withFieldNames(String fieldNames) + { + this.fieldNames = fieldNames; + return (this); + } + + + + /******************************************************************************* + ** Getter for titleFieldName + *******************************************************************************/ + public String getTitleFieldName() + { + return (titleFieldName); + } + + + + /******************************************************************************* + ** Setter for titleFieldName + *******************************************************************************/ + public void setTitleFieldName(String titleFieldName) + { + this.titleFieldName = titleFieldName; + } + + + + /******************************************************************************* + ** Fluent setter for titleFieldName + *******************************************************************************/ + public AssistantRagSource withTitleFieldName(String titleFieldName) + { + this.titleFieldName = titleFieldName; + return (this); + } + + + + /******************************************************************************* + ** Getter for filterJson + *******************************************************************************/ + public String getFilterJson() + { + return (filterJson); + } + + + + /******************************************************************************* + ** Setter for filterJson + *******************************************************************************/ + public void setFilterJson(String filterJson) + { + this.filterJson = filterJson; + } + + + + /******************************************************************************* + ** Fluent setter for filterJson + *******************************************************************************/ + public AssistantRagSource withFilterJson(String filterJson) + { + this.filterJson = filterJson; + return (this); + } + + + + /******************************************************************************* + ** Getter for embeddingModelId + *******************************************************************************/ + public Integer getEmbeddingModelId() + { + return (embeddingModelId); + } + + + + /******************************************************************************* + ** Setter for embeddingModelId + *******************************************************************************/ + public void setEmbeddingModelId(Integer embeddingModelId) + { + this.embeddingModelId = embeddingModelId; + } + + + + /******************************************************************************* + ** Fluent setter for embeddingModelId + *******************************************************************************/ + public AssistantRagSource withEmbeddingModelId(Integer embeddingModelId) + { + this.embeddingModelId = embeddingModelId; + return (this); + } + + + + /******************************************************************************* + ** Getter for chunkChars + *******************************************************************************/ + public Integer getChunkChars() + { + return (chunkChars); + } + + + + /******************************************************************************* + ** Setter for chunkChars + *******************************************************************************/ + public void setChunkChars(Integer chunkChars) + { + this.chunkChars = chunkChars; + } + + + + /******************************************************************************* + ** Fluent setter for chunkChars + *******************************************************************************/ + public AssistantRagSource withChunkChars(Integer chunkChars) + { + this.chunkChars = chunkChars; + return (this); + } + + + + /******************************************************************************* + ** Getter for chunkOverlapChars + *******************************************************************************/ + public Integer getChunkOverlapChars() + { + return (chunkOverlapChars); + } + + + + /******************************************************************************* + ** Setter for chunkOverlapChars + *******************************************************************************/ + public void setChunkOverlapChars(Integer chunkOverlapChars) + { + this.chunkOverlapChars = chunkOverlapChars; + } + + + + /******************************************************************************* + ** Fluent setter for chunkOverlapChars + *******************************************************************************/ + public AssistantRagSource withChunkOverlapChars(Integer chunkOverlapChars) + { + this.chunkOverlapChars = chunkOverlapChars; + return (this); + } + + + + /******************************************************************************* + ** Getter for lastIndexedAt + *******************************************************************************/ + public Instant getLastIndexedAt() + { + return (lastIndexedAt); + } + + + + /******************************************************************************* + ** Setter for lastIndexedAt + *******************************************************************************/ + public void setLastIndexedAt(Instant lastIndexedAt) + { + this.lastIndexedAt = lastIndexedAt; + } + + + + /******************************************************************************* + ** Fluent setter for lastIndexedAt + *******************************************************************************/ + public AssistantRagSource withLastIndexedAt(Instant lastIndexedAt) + { + this.lastIndexedAt = lastIndexedAt; + return (this); + } + + + + /******************************************************************************* + ** Getter for lastIndexStatus + *******************************************************************************/ + public String getLastIndexStatus() + { + return (lastIndexStatus); + } + + + + /******************************************************************************* + ** Setter for lastIndexStatus + *******************************************************************************/ + public void setLastIndexStatus(String lastIndexStatus) + { + this.lastIndexStatus = lastIndexStatus; + } + + + + /******************************************************************************* + ** Fluent setter for lastIndexStatus + *******************************************************************************/ + public AssistantRagSource withLastIndexStatus(String lastIndexStatus) + { + this.lastIndexStatus = lastIndexStatus; + return (this); + } + + + + /******************************************************************************* + ** Getter for lastIndexError + *******************************************************************************/ + public String getLastIndexError() + { + return (lastIndexError); + } + + + + /******************************************************************************* + ** Setter for lastIndexError + *******************************************************************************/ + public void setLastIndexError(String lastIndexError) + { + this.lastIndexError = lastIndexError; + } + + + + /******************************************************************************* + ** Fluent setter for lastIndexError + *******************************************************************************/ + public AssistantRagSource withLastIndexError(String lastIndexError) + { + this.lastIndexError = lastIndexError; + return (this); + } + + + + /******************************************************************************* + ** Getter for indexedRecordCount + *******************************************************************************/ + public Integer getIndexedRecordCount() + { + return (indexedRecordCount); + } + + + + /******************************************************************************* + ** Setter for indexedRecordCount + *******************************************************************************/ + public void setIndexedRecordCount(Integer indexedRecordCount) + { + this.indexedRecordCount = indexedRecordCount; + } + + + + /******************************************************************************* + ** Fluent setter for indexedRecordCount + *******************************************************************************/ + public AssistantRagSource withIndexedRecordCount(Integer indexedRecordCount) + { + this.indexedRecordCount = indexedRecordCount; + return (this); + } + + + + /******************************************************************************* + ** Getter for indexedChunkCount + *******************************************************************************/ + public Integer getIndexedChunkCount() + { + return (indexedChunkCount); + } + + + + /******************************************************************************* + ** Setter for indexedChunkCount + *******************************************************************************/ + public void setIndexedChunkCount(Integer indexedChunkCount) + { + this.indexedChunkCount = indexedChunkCount; + } + + + + /******************************************************************************* + ** Fluent setter for indexedChunkCount + *******************************************************************************/ + public AssistantRagSource withIndexedChunkCount(Integer indexedChunkCount) + { + this.indexedChunkCount = indexedChunkCount; + return (this); + } + + + + /******************************************************************************* + ** Getter for createDate + *******************************************************************************/ + public Instant getCreateDate() + { + return (createDate); + } + + + + /******************************************************************************* + ** Setter for createDate + *******************************************************************************/ + public void setCreateDate(Instant createDate) + { + this.createDate = createDate; + } + + + + /******************************************************************************* + ** Fluent setter for createDate + *******************************************************************************/ + public AssistantRagSource withCreateDate(Instant createDate) + { + this.createDate = createDate; + return (this); + } + + + + /******************************************************************************* + ** Getter for modifyDate + *******************************************************************************/ + public Instant getModifyDate() + { + return (modifyDate); + } + + + + /******************************************************************************* + ** Setter for modifyDate + *******************************************************************************/ + public void setModifyDate(Instant modifyDate) + { + this.modifyDate = modifyDate; + } + + + + /******************************************************************************* + ** Fluent setter for modifyDate + *******************************************************************************/ + public AssistantRagSource withModifyDate(Instant modifyDate) + { + this.modifyDate = modifyDate; + return (this); + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/processes/KnowledgeIndexProcess.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/processes/KnowledgeIndexProcess.java new file mode 100644 index 0000000..71555b4 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/processes/KnowledgeIndexProcess.java @@ -0,0 +1,231 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.processes; + + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.classification.DataClassification; +import com.kingsrook.qbits.assistant.metadata.AssistantMetaDataBuilder; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; +import com.kingsrook.qbits.assistant.rag.RagService; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.actions.processes.ProcessSummaryLineInterface; +import com.kingsrook.qqq.backend.core.model.actions.processes.RunBackendStepInput; +import com.kingsrook.qqq.backend.core.model.actions.processes.RunBackendStepOutput; +import com.kingsrook.qqq.backend.core.model.actions.processes.Status; +import com.kingsrook.qqq.backend.core.model.data.QRecord; +import com.kingsrook.qqq.backend.core.model.metadata.layout.QIcon; +import com.kingsrook.qqq.backend.core.model.metadata.processes.QProcessMetaData; +import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.AbstractLoadStep; +import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.AbstractTransformStep; +import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.ExtractViaQueryStep; +import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.ProcessSummaryProviderInterface; +import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.StreamedETLWithFrontendProcess; + + +/******************************************************************************* + ** Standard Review/Result record actions. No scheduler or service-user runner. + *******************************************************************************/ +public final class KnowledgeIndexProcess +{ + public static final String INDEX = "assistantKnowledgeIndex"; + public static final String CLEAR = "assistantKnowledgeClear"; + + + + /******************************************************************************* + ** Utility class. + *******************************************************************************/ + private KnowledgeIndexProcess() + { + } + + + + /******************************************************************************* + ** Both processes are table actions, never navigation entries. + *******************************************************************************/ + public static List processes() + { + List processes = new ArrayList<>(); + for(Boolean clear : List.of(false, true)) + { + Map values = new HashMap<>(); + values.put(StreamedETLWithFrontendProcess.FIELD_SOURCE_TABLE, AssistantRagSource.TABLE_NAME); + values.put(StreamedETLWithFrontendProcess.FIELD_DESTINATION_TABLE, AssistantRagSource.TABLE_NAME); + values.put(StreamedETLWithFrontendProcess.FIELD_SUPPORTS_FULL_VALIDATION, true); + values.put(StreamedETLWithFrontendProcess.FIELD_DO_FULL_VALIDATION, true); + values.put(StreamedETLWithFrontendProcess.FIELD_PREVIEW_MESSAGE, clear ? "Clear vectors for the selected knowledge sources." : "Read content through the UI API as you and embed changed chunks. Embedding usage counts toward budgets."); + QProcessMetaData process = StreamedETLWithFrontendProcess.defineProcessMetaData(ExtractViaQueryStep.class, clear ? ClearReview.class : Review.class, clear ? Clear.class : Index.class, values); + process.setName(clear ? CLEAR : INDEX); + process.setLabel(clear ? "Clear index" : "Index now"); + process.setIcon(new QIcon(clear ? "delete_sweep" : "manage_search")); + process.setTableName(AssistantRagSource.TABLE_NAME); + process.setPermissionRules(AssistantMetaDataBuilder.adminRules()); + processes.add(process); + } + return (processes); + } + + + + /******************************************************************************* + ** Review validates configuration without sending content to any provider. + *******************************************************************************/ + public static class Review extends AbstractTransformStep implements ProcessSummaryProviderInterface + { + private final ArrayList lines = new ArrayList<>(); + + + + /******************************************************************************* + ** Source records only; application records are read at Result time. + *******************************************************************************/ + @Override + public void runOnePage(RunBackendStepInput input, RunBackendStepOutput output) throws QException + { + AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance()); + for(QRecord record : input.getRecords()) + { + AssistantRagSource source = new AssistantRagSource(record); + List errors = problems(runtime, source); + lines.add(TestConnectionProcess.single(errors.isEmpty() ? Status.INFO : Status.ERROR, + source.getName() + (errors.isEmpty() ? ": source selected; no content has been sent." : ": " + String.join("; ", errors)), source.getId())); + output.getRecords().add(record); + } + } + + + + /******************************************************************************* + ** Standard streamed ETL Review lines. + *******************************************************************************/ + @Override + public ArrayList getProcessSummary(RunBackendStepOutput output, boolean isForResultScreen) + { + return (lines); + } + + + + /******************************************************************************* + ** Review reports PHI/model setup blocks before any indexing dispatch. + *******************************************************************************/ + protected List problems(AssistantRuntime runtime, AssistantRagSource source) throws QException + { + List errors = runtime.getRagService().validate(source); + if(errors.isEmpty()) + { + try + { + runtime.getProviderService().resolveEmbedding(source.getEmbeddingModelId(), + runtime.getConfig().getClassificationSource().highest(source.getTableName()) == DataClassification.PHI); + } + catch(QException e) + { + errors.add(e.getMessage()); + } + } + return (errors); + } + } + + + + /******************************************************************************* + ** Broken or disabled source configuration must never prevent clearing it. + *******************************************************************************/ + public static class ClearReview extends Review + { + /******************************************************************************* + ** Clearing has no model or source-content prerequisites. + *******************************************************************************/ + @Override + protected List problems(AssistantRuntime runtime, AssistantRagSource source) + { + return (List.of()); + } + } + + + + /******************************************************************************* + ** Result performs indexing under the clicking administrator's QContext. + *******************************************************************************/ + public static class Index extends AbstractLoadStep implements ProcessSummaryProviderInterface + { + protected final ArrayList lines = new ArrayList<>(); + + + + /******************************************************************************* + ** Page of source configuration records, not application content. + *******************************************************************************/ + @Override + public void runOnePage(RunBackendStepInput input, RunBackendStepOutput output) throws QException + { + AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance()); + for(QRecord record : input.getRecords()) + { + AssistantRagSource source = new AssistantRagSource(record); + RagService.Result result = runtime.getRagService().index(source); + lines.add(TestConnectionProcess.single(result.errors() == 0 ? Status.OK : Status.ERROR, result.summary(), source.getId())); + output.getRecords().add(record); + } + } + + + + /******************************************************************************* + ** Counts from this run, suitable for the standard Result screen. + *******************************************************************************/ + @Override + public ArrayList getProcessSummary(RunBackendStepOutput output, boolean isForResultScreen) + { + return (lines); + } + } + + + + /******************************************************************************* + ** Clear shares the process presentation but makes no embedding calls. + *******************************************************************************/ + public static class Clear extends Index + { + /******************************************************************************* + ** Remove only the selected sources' vectors. + *******************************************************************************/ + @Override + public void runOnePage(RunBackendStepInput input, RunBackendStepOutput output) throws QException + { + for(QRecord record : input.getRecords()) + { + AssistantRagSource source = new AssistantRagSource(record); + AssistantRuntime.of(QContext.getQInstance()).getRagService().clear(source); + lines.add(TestConnectionProcess.single(Status.OK, source.getName() + ": index cleared.", source.getId())); + output.getRecords().add(record); + } + } + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/AssistantModelProvider.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/AssistantModelProvider.java index c67bfce..5bc10a5 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/AssistantModelProvider.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/AssistantModelProvider.java @@ -58,4 +58,14 @@ public interface AssistantModelProvider return (null); } + + + /******************************************************************************* + ** Optional embedding capability. Unsupported adapters fail before transport. + *******************************************************************************/ + default EmbeddingResponse embed(EmbeddingRequest request) throws ModelProviderException + { + throw (new ModelProviderException(ModelProviderException.Kind.INVALID_REQUEST, null, "Embeddings are unsupported by this provider.")); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/EmbeddingRequest.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/EmbeddingRequest.java new file mode 100644 index 0000000..2a903fb --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/EmbeddingRequest.java @@ -0,0 +1,39 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.provider; + + +import java.util.List; + + +/******************************************************************************* + ** One embedding batch. Providers accept at most 64 inputs per request. + *******************************************************************************/ +public record EmbeddingRequest(String model, List input) +{ + /******************************************************************************* + ** Snapshot the inputs and reject empty or oversized batches before transport. + *******************************************************************************/ + public EmbeddingRequest + { + input = List.copyOf(input); + if(model == null || model.isBlank() || input.isEmpty() || input.size() > 64) + { + throw (new IllegalArgumentException("An embedding request needs a model and 1 to 64 inputs.")); + } + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/EmbeddingResponse.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/EmbeddingResponse.java new file mode 100644 index 0000000..9101c56 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/EmbeddingResponse.java @@ -0,0 +1,28 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.provider; + + +import java.util.List; + + +/******************************************************************************* + ** Ordered vectors and billable usage. No input text is retained in the response. + *******************************************************************************/ +public record EmbeddingResponse(List vectors, ModelUsage usage, String requestId) +{ +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/RegistryModelProvider.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/RegistryModelProvider.java index 024947c..b83ad44 100644 --- a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/RegistryModelProvider.java +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/provider/RegistryModelProvider.java @@ -94,4 +94,15 @@ public class RegistryModelProvider implements AssistantModelProvider { return (delegate.countTokens(request)); } + + + /******************************************************************************* + ** Preserve the delegate's optional embedding capability. + *******************************************************************************/ + @Override + public EmbeddingResponse embed(EmbeddingRequest request) throws ModelProviderException + { + return (delegate.embed(request)); + } + } diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/ChunkVector.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/ChunkVector.java new file mode 100644 index 0000000..1cf5a16 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/ChunkVector.java @@ -0,0 +1,55 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +/******************************************************************************* + ** Content-free index entry. The embedding is copied on entry and on read. + *******************************************************************************/ +public record ChunkVector(Integer source, String table, String pk, String field, Integer chunkIndex, String contentHash, float[] embedding) +{ + /******************************************************************************* + ** No mutable vector can alter an indexed entry after insertion. + *******************************************************************************/ + public ChunkVector + { + if(source == null || table == null || pk == null || field == null || chunkIndex == null || chunkIndex < 0 + || contentHash == null || !contentHash.matches("[a-f0-9]{64}") || embedding == null || embedding.length == 0) + { + throw (new IllegalArgumentException("Incomplete chunk vector metadata.")); + } + embedding = embedding.clone(); + for(float value : embedding) + { + if(!Float.isFinite(value)) + { + throw (new IllegalArgumentException("Embedding coordinates must be finite.")); + } + } + } + + + + /******************************************************************************* + ** Return a defensive copy. + *******************************************************************************/ + @Override + public float[] embedding() + { + return (embedding.clone()); + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/EmbeddingService.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/EmbeddingService.java new file mode 100644 index 0000000..8aed0d7 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/EmbeddingService.java @@ -0,0 +1,170 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +import java.time.LocalDate; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.Semaphore; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.AssistantSecurity; +import com.kingsrook.qbits.assistant.classification.DataClassification; +import com.kingsrook.qbits.assistant.engine.BudgetService; +import com.kingsrook.qbits.assistant.engine.PolicyService; +import com.kingsrook.qbits.assistant.engine.ProviderService; +import com.kingsrook.qbits.assistant.engine.UsageRecorder; +import com.kingsrook.qbits.assistant.model.AssistantDataAudit; +import com.kingsrook.qbits.assistant.provider.EmbeddingRequest; +import com.kingsrook.qbits.assistant.provider.EmbeddingResponse; +import com.kingsrook.qbits.assistant.provider.ModelProviderException; +import com.kingsrook.qbits.assistant.provider.ModelResponse; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.exceptions.QException; + + +/******************************************************************************* + ** Embedding dispatch uses fresh eligibility, kill-switch and budget checks per batch. + *******************************************************************************/ +public class EmbeddingService +{ + private final AssistantRuntime runtime; + + + + /******************************************************************************* + ** Bind the per-instance accounting and provider services. + *******************************************************************************/ + public EmbeddingService(AssistantRuntime runtime) + { + this.runtime = runtime; + } + + + + /******************************************************************************* + ** Batch at most 64 inputs. Every dispatched batch records EMBED usage, + ** including failures, and updates the existing daily budget rollups. + *******************************************************************************/ + public List embed(Integer modelId, List texts, DataClassification classification, String tables, Integer conversationId, Integer turnId) throws QException + { + List vectors = new ArrayList<>(); + for(Integer offset = 0; offset < texts.size(); offset += 64) + { + vectors.addAll(batch(modelId, texts.subList(offset, Math.min(offset + 64, texts.size())), classification, tables, conversationId, turnId)); + } + return (vectors); + } + + + + /******************************************************************************* + ** Audit protected dispatch refusals without any text, before transport. + *******************************************************************************/ + public ProviderService.Active resolve(Integer modelId, DataClassification classification, String tables, Integer conversationId, Integer turnId) throws QException + { + try + { + return (runtime.getProviderService().resolveEmbedding(modelId, classification == DataClassification.PHI)); + } + catch(ProviderService.ProviderSetupException e) + { + if(classification == DataClassification.PHI || Boolean.TRUE.equals(runtime.getConfig().getHipaaMode())) + { + var model = runtime.getStore().listModels().stream().filter(m -> Objects.equals(m.getId(), modelId)).findFirst().orElse(null); + var provider = model == null ? null : runtime.getStore().getProvider(model.getProviderId()); + runtime.getStore().insert(new AssistantDataAudit().withUserId(AssistantSecurity.userId(QContext.getQSession())).withConversationId(conversationId).withTurnId(turnId) + .withProviderId(provider == null ? null : provider.getId()).withProviderName(provider == null ? null : provider.getName()) + .withModel(model == null ? null : model.getModelId()).withHipaaEligible(provider != null && Boolean.TRUE.equals(provider.getHipaaEligible())) + .withClassification(classification.name()).withTables(tables).withDecision("BLOCKED").withReason(e.getMessage()).withCreateDate(runtime.getClock().instant())); + } + throw (e); + } + } + + + + /******************************************************************************* + ** Never downgrade an embedding call into a different vector space. + *******************************************************************************/ + private List batch(Integer modelId, List texts, DataClassification classification, String tables, Integer conversationId, Integer turnId) throws QException + { + String userId = AssistantSecurity.userId(QContext.getQSession()); + ProviderService.Active active = resolve(modelId, classification, tables, conversationId, turnId); + PolicyService.EffectivePolicy policy = runtime.getPolicyService().current(true); + BudgetService.Decision budget = runtime.getBudgetService().evaluate(userId, active, classification == DataClassification.PHI, false); + Long cap = ProviderService.capMicros(active.row()); + if(policy.killSwitch()) + { + throw (new QException("The assistant kill switch is on; embeddings were not sent.")); + } + if(budget.stopped() || !Objects.equals(active.model(), budget.active().model()) || !Objects.equals(active.providerId(), budget.active().providerId()) + || runtime.getPolicyService().checkConfiguredBudget(policy, userId, 0, 0).hardExceeded() + || (cap != null && runtime.getProviderService().monthToDateCostMicros(active.providerId()) >= cap)) + { + throw (new QException("The embedding budget has been reached; embeddings were not sent.")); + } + if(classification.isAtLeast(DataClassification.PII)) + { + runtime.getStore().insert(new AssistantDataAudit().withUserId(userId).withConversationId(conversationId).withTurnId(turnId) + .withProviderId(active.providerId()).withProviderName(active.row().getName()).withModel(active.model()) + .withHipaaEligible(Boolean.TRUE.equals(active.row().getHipaaEligible())).withClassification(classification.name()) + .withTables(tables).withDecision("ALLOWED").withReason("Embedding provider permitted for this classification.").withCreateDate(runtime.getClock().instant())); + } + Semaphore gate = runtime.getProviderService().gate(active.row()); + Boolean acquired = gate == null || gate.tryAcquire(); + if(!acquired) + { + throw (new QException("Embedding provider concurrency limit reached.")); + } + Long started = System.nanoTime(); + ModelResponse response = null; + ModelProviderException error = null; + try + { + EmbeddingResponse embedded = active.provider().embed(new EmbeddingRequest(active.model(), texts)); + response = new ModelResponse().withModel(active.model()).withRequestId(embedded.requestId()).withUsage(embedded.usage()); + if(embedded.vectors().size() != texts.size()) + { + throw (new ModelProviderException(ModelProviderException.Kind.OTHER, null, "The embedding provider returned the wrong number of vectors.")); + } + return (embedded.vectors()); + } + catch(ModelProviderException e) + { + error = e; + throw (new QException("Embedding request failed (" + e.getKind() + "). Check provider settings and retry.")); + } + finally + { + if(gate != null) + { + gate.release(); + } + var pricing = runtime.getProviderService().pricing(active, active.model()); + runtime.getStore().insert(UsageRecorder.row(new UsageRecorder.Call(userId, conversationId, turnId, active, UsageRecorder.PURPOSE_EMBED), + response, error, (System.nanoTime() - started) / 1_000_000, pricing)); + var usage = response == null ? (error == null ? null : error.getPartialUsage()) : response.getUsage(); + if(usage != null) + { + runtime.getStore().addUsage(userId, LocalDate.now(runtime.getClock().withZone(ZoneOffset.UTC)), usage, pricing.costMicros(usage), false); + } + } + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/InMemoryVectorStore.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/InMemoryVectorStore.java new file mode 100644 index 0000000..d8ccdb3 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/InMemoryVectorStore.java @@ -0,0 +1,157 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + + +/******************************************************************************* + ** Exact cosine search for tests and small hosts. Never stores source text. + *******************************************************************************/ +public class InMemoryVectorStore implements VectorStore +{ + private final Map, ChunkVector> vectors = new LinkedHashMap<>(); + + + + /******************************************************************************* + ** Upsert immutable entries. + *******************************************************************************/ + @Override + public synchronized void upsert(List chunks) + { + for(ChunkVector chunk : chunks) + { + vectors.put(List.of(chunk.source(), chunk.table(), chunk.pk(), chunk.field(), chunk.chunkIndex()), chunk); + } + } + + + + /******************************************************************************* + ** Remove a record within its source only. + *******************************************************************************/ + @Override + public synchronized void deleteRecord(Integer source, String table, String pk) + { + vectors.values().removeIf(v -> v.source().equals(source) && v.table().equals(table) && v.pk().equals(pk)); + } + + + + /******************************************************************************* + ** A source can be cleared independently of other sources. + *******************************************************************************/ + @Override + public synchronized void deleteSource(Integer source) + { + vectors.values().removeIf(v -> v.source().equals(source)); + } + + + + /******************************************************************************* + ** Search only the named tables and vectors with compatible dimensions. + *******************************************************************************/ + @Override + public synchronized List search(float[] query, Integer k, Set tables) + { + return (search(query, k, tables, Set.of())); + } + + + + /******************************************************************************* + ** Different embedding models are different spaces even at equal dimensions. + *******************************************************************************/ + @Override + public synchronized List search(float[] query, Integer k, Set tables, Set sources) + { + List hits = new ArrayList<>(); + for(ChunkVector vector : vectors.values()) + { + if(tables.contains(vector.table()) && (sources.isEmpty() || sources.contains(vector.source())) && query.length == vector.embedding().length) + { + Double score = cosine(query, vector.embedding()); + if(Double.isFinite(score)) + { + hits.add(new Hit(vector.source(), vector.table(), vector.pk(), vector.field(), vector.chunkIndex(), score)); + } + } + } + return (hits.stream().sorted(Comparator.comparing(Hit::score).reversed()).limit(Math.max(0, k)).toList()); + } + + + + /******************************************************************************* + ** Zero vectors have no direction and never match. + *******************************************************************************/ + private static Double cosine(float[] a, float[] b) + { + Double dot = 0.0; + Double left = 0.0; + Double right = 0.0; + for(Integer i = 0; i < a.length; i++) + { + dot += (double) a[i] * b[i]; + left += (double) a[i] * a[i]; + right += (double) b[i] * b[i]; + } + return (left == 0 || right == 0 ? Double.NaN : dot / Math.sqrt(left * right)); + } + + + + /******************************************************************************* + ** Count by source. + *******************************************************************************/ + @Override + public synchronized Long count(Integer source) + { + return (vectors.values().stream().filter(v -> v.source().equals(source)).count()); + } + + + + /******************************************************************************* + ** Snapshot the content-free entries for reconciliation. + *******************************************************************************/ + @Override + public synchronized List chunks(Integer source) + { + return (vectors.values().stream().filter(v -> v.source().equals(source)).toList()); + } + + + + /******************************************************************************* + ** Readers see either complete version of this record. + *******************************************************************************/ + @Override + public synchronized void replaceRecord(Integer source, String table, String pk, List chunks) + { + deleteRecord(source, table, pk); + upsert(chunks); + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/RagService.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/RagService.java new file mode 100644 index 0000000..e0a376d --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/RagService.java @@ -0,0 +1,446 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import com.kingsrook.qbits.assistant.AssistantQBitConfig; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.engine.PolicyService; +import com.kingsrook.qbits.assistant.model.AssistantModel; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; +import com.kingsrook.qbits.assistant.tools.UiApiToolbox; +import com.kingsrook.qbits.assistant.uiapi.UiApiCall; +import com.kingsrook.qbits.assistant.uiapi.UiApiResult; +import com.kingsrook.qqq.backend.core.actions.ai.AiExposure; +import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome; +import com.kingsrook.qqq.backend.core.actions.ai.exposure.AiExposurePolicy; +import com.kingsrook.qqq.backend.core.actions.ai.output.AiRedactor; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.actions.tables.query.QFilterOrderBy; +import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter; +import com.kingsrook.qqq.backend.core.model.data.QRecord; +import com.kingsrook.qqq.backend.core.model.metadata.fields.QFieldMetaData; +import com.kingsrook.qqq.backend.core.model.metadata.fields.QFieldType; +import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData; +import com.kingsrook.qqq.backend.core.utils.JsonUtils; + + +/******************************************************************************* + ** Manual indexing as the clicking administrator, and shared safe UI API reads. + *******************************************************************************/ +public class RagService +{ + private final AssistantRuntime runtime; + private final VectorStore vectors; + private final AiExposurePolicy exposure; + private final UiApiToolbox reader; + private final EmbeddingService embeddings; + + + + /******************************************************************************* + ** Reuse get_record's helpers through its public execution path, without + ** imposing the chat cell-length cap on indexing or snippet boundaries. + *******************************************************************************/ + public RagService(AssistantRuntime runtime, VectorStore vectors) + { + this.runtime = runtime; + this.vectors = vectors; + var settings = AssistantRuntime.kernelSettings(runtime.getConfig()); + exposure = new AiExposurePolicy(settings); + reader = new UiApiToolbox(runtime.getQInstance(), runtime.getDataAccessService(), exposure, new AiRedactor(settings.getCustomRedactionPatterns()), Integer.MAX_VALUE); + embeddings = new EmbeddingService(runtime); + } + + + + /******************************************************************************* + ** Configuration only; never query application data through AssistantStore. + *******************************************************************************/ + public List sources() throws QException + { + return (runtime.getStore().query(AssistantRagSource.TABLE_NAME, new QQueryFilter(), AssistantRagSource.class)); + } + + + + /******************************************************************************* + ** Source definitions only narrow existing AI table/field exposure. + *******************************************************************************/ + public List validate(AssistantRagSource source) throws QException + { + List errors = new ArrayList<>(); + QTableMetaData table = runtime.getQInstance().getTable(source.getTableName()); + if(table == null || source.getTableName().startsWith("assistant_") || !runtime.getDataAccessService().forTable(source.getTableName()).mode().allowsRead()) + { + errors.add("Knowledge sources require an application table with AI data access READ or READ_WRITE, never NONE or assistant_*."); + return (errors); + } + if(fields(source).isEmpty()) + { + errors.add("Select at least one text field."); + } + Set selected = new LinkedHashSet<>(fields(source)); + if(source.getTitleFieldName() != null && !source.getTitleFieldName().isBlank()) + { + selected.add(source.getTitleFieldName()); + } + for(String name : selected) + { + QFieldMetaData field = table.getFields().get(name); + AiExposure level = field == null ? AiExposure.NONE : exposure.getFieldExposure(table, field); + if(field == null || !Set.of(QFieldType.STRING, QFieldType.TEXT).contains(field.getType()) || !level.allowsValues()) + { + errors.add("Knowledge fields must be text fields with VALUES or MASKED exposure: " + name); + } + } + if(size(source) < 100 || size(source) > 16000 || overlap(source) < 0 || overlap(source) >= size(source)) + { + errors.add("Chunk size must be 100 to 16000; overlap must be nonnegative and smaller than chunk size."); + } + AssistantModel model = runtime.getStore().listModels().stream().filter(m -> Objects.equals(m.getId(), source.getEmbeddingModelId())).findFirst().orElse(null); + if(model == null || !"EMBEDDING".equals(model.getModelKind()) || !Boolean.TRUE.equals(model.getEnabled()) || model.getEmbeddingDimensions() == null || model.getEmbeddingDimensions() < 1) + { + errors.add("Select an enabled EMBEDDING model with positive embedding dimensions."); + } + try + { + filter(source); + } + catch(QException e) + { + errors.add("The query filter must be valid QQQ query filter JSON."); + } + return (errors); + } + + + + /******************************************************************************* + ** Transient safe values from get_record, with the asking user's session. + *******************************************************************************/ + public UiApiToolbox.Outcome read(UiApiToolbox.Context context, String table, String pk) + { + PolicyService.ToolPolicy policy = context.policy(); + UiApiToolbox.Context internal = new UiApiToolbox.Context(context.sessionUUID(), context.userTimezone(), context.gateway(), + new PolicyService.ToolPolicy(true, Set.of(), policy == null ? 50 : policy.maxRowsPerCall(), policy == null || policy.redactPii()), null); + return (reader.execute(internal, "knowledge-read", UiApiToolbox.GET_RECORD, Map.of("table", table, "id", pk))); + } + + + + /******************************************************************************* + ** Read already masked/redacted fields from the existing record render. + *******************************************************************************/ + public static Map values(UiApiToolbox.Outcome outcome) + { + Map values = new LinkedHashMap<>(); + if(outcome.getResult().getRender() instanceof Map render && render.get("fields") instanceof List fields) + { + for(Object item : fields) + { + if(item instanceof Map field && field.get("name") instanceof String name && field.get("value") != null) + { + values.put(name, String.valueOf(field.get("value"))); + } + } + } + return (values); + } + + + + /******************************************************************************* + ** Index result counters contain no record values or provider diagnostics. + *******************************************************************************/ + public record Result(Integer recordsSeen, Integer embedded, Integer unchanged, Integer removed, Integer errors, String reason) + { + /******************************************************************************* + ** Standard process summary wording. + *******************************************************************************/ + public String summary() + { + return ("Records seen: " + recordsSeen + "; chunks embedded: " + embedded + "; unchanged: " + unchanged + "; removed: " + removed + "; errors: " + errors + + (reason == null ? "" : ". " + reason)); + } + } + + + + /******************************************************************************* + ** Complete scans alone remove disappeared records. Failed or denied pages + ** never erase previously indexed records. Re-indexing is deliberately manual. + *******************************************************************************/ + public synchronized Result index(AssistantRagSource source) throws QException + { + requireAdmin(); + Integer seen = 0; + Integer embedded = 0; + Integer unchanged = 0; + Integer removed = 0; + String reason = null; + try + { + List errors = validate(source); + if(!errors.isEmpty()) + { + throw (new QException(errors.getFirst())); + } + var classification = runtime.getConfig().getClassificationSource().highest(source.getTableName()); + embeddings.resolve(source.getEmbeddingModelId(), classification, source.getTableName(), null, null); + Map> old = new LinkedHashMap<>(); + vectors.chunks(source.getId()).forEach(v -> old.computeIfAbsent(v.pk(), key -> new ArrayList<>()).add(v)); + Set seenIds = new LinkedHashSet<>(); + QTableMetaData table = runtime.getQInstance().getTable(source.getTableName()); + QQueryFilter filter = filter(source); + filter.setOrderBys(List.of(new QFilterOrderBy(table.getPrimaryKeyField(), true))); + filter.setLimit(100); + UiApiToolbox.Context context = context(); + for(Integer skip = 0; ; skip += 100) + { + checkKillSwitch(); + filter.setSkip(skip); + UiApiResult page = runtime.getUiApiGateway().invoke(new UiApiCall("POST", "/table/{tableName}/query").withPathParam("tableName", source.getTableName()) + .withSessionUUID(context.sessionUUID()).withUserTimezone(context.userTimezone()).withJsonBody("{\"filter\":" + JsonUtils.toJson(filter) + "}")); + if(!page.isSuccess()) + { + throw (new QException("Index query answered HTTP " + page.status() + "; existing chunks were preserved.")); + } + if(!(page.json() instanceof Map json) || !(json.get("records") instanceof List records)) + { + throw (new QException("The UI API returned an invalid query page; existing chunks were preserved.")); + } + for(Object item : records) + { + if(!(item instanceof Map record) || !(record.get("values") instanceof Map raw) || raw.get(table.getPrimaryKeyField()) == null) + { + throw (new QException("The UI API returned a record without its primary key.")); + } + String pk = String.valueOf(raw.get(table.getPrimaryKeyField())); + if(!seenIds.add(pk)) + { + throw (new QException("Records moved during paging. Retry indexing; disappeared-record cleanup was skipped.")); + } + seen++; + checkKillSwitch(); + UiApiToolbox.Outcome read = read(context, source.getTableName(), pk); + if(read.getResult().getOutcome() != AiToolOutcome.OK) + { + throw (new QException("Index record read was refused; existing chunks were preserved.")); + } + List previous = old.getOrDefault(pk, List.of()); + List next = new ArrayList<>(); + List texts = new ArrayList<>(); + List names = new ArrayList<>(); + List indexes = new ArrayList<>(); + Map values = values(read); + for(String field : fields(source)) + { + List chunks = TextChunks.split(values.getOrDefault(field, ""), size(source), overlap(source)); + for(Integer i = 0; i < chunks.size(); i++) + { + Integer index = i; + String hash = TextChunks.hash(chunks.get(i)); + ChunkVector same = previous.stream().filter(v -> v.field().equals(field) && v.chunkIndex().equals(index) && v.contentHash().equals(hash)).findFirst().orElse(null); + if(same != null) + { + next.add(same); + unchanged++; + } + else + { + texts.add(chunks.get(i)); + names.add(field); + indexes.add(i); + } + } + } + List generated = embeddings.embed(source.getEmbeddingModelId(), texts, classification, source.getTableName(), null, null); + Integer dimensions = runtime.getStore().listModels().stream().filter(m -> m.getId().equals(source.getEmbeddingModelId())).findFirst().orElseThrow().getEmbeddingDimensions(); + for(Integer i = 0; i < generated.size(); i++) + { + if(generated.get(i).length != dimensions) + { + throw (new QException("Embedding dimensions differ from the model registry; no record replacement was made.")); + } + next.add(new ChunkVector(source.getId(), source.getTableName(), pk, names.get(i), indexes.get(i), TextChunks.hash(texts.get(i)), generated.get(i))); + } + embedded += generated.size(); + removed += (int) previous.stream().filter(v -> next.stream().noneMatch(n -> n.field().equals(v.field()) && n.chunkIndex().equals(v.chunkIndex()))).count(); + if(!texts.isEmpty() || next.size() != previous.size()) + { + vectors.replaceRecord(source.getId(), source.getTableName(), pk, next); + } + } + if(records.size() < 100) + { + break; + } + } + checkKillSwitch(); + for(var entry : old.entrySet()) + { + if(!seenIds.contains(entry.getKey())) + { + vectors.deleteRecord(source.getId(), source.getTableName(), entry.getKey()); + removed += entry.getValue().size(); + } + } + } + catch(QException e) + { + reason = e.getMessage(); + } + catch(RuntimeException e) + { + reason = "Indexing failed. Check the source and vector store configuration."; + } + Integer chunkCount = source.getIndexedChunkCount(); + try + { + chunkCount = vectors.count(source.getId()).intValue(); + } + catch(QException e) + { + reason = "The vector store is unavailable; the indexed chunk count could not be refreshed."; + } + Result result = new Result(seen, embedded, unchanged, removed, reason == null ? 0 : 1, reason); + runtime.getStore().updateFields(AssistantRagSource.TABLE_NAME, new QRecord().withValue("id", source.getId()) + .withValue("lastIndexedAt", runtime.getClock().instant()).withValue("lastIndexStatus", reason == null ? "OK" : "ERROR") + .withValue("lastIndexError", reason == null ? null : reason.substring(0, Math.min(500, reason.length()))) + .withValue("indexedRecordCount", seen).withValue("indexedChunkCount", chunkCount)); + return (result); + } + + + + /******************************************************************************* + ** Clear is an admin record action, independent of the current model setup. + *******************************************************************************/ + public synchronized void clear(AssistantRagSource source) throws QException + { + requireAdmin(); + vectors.deleteSource(source.getId()); + runtime.getStore().updateFields(AssistantRagSource.TABLE_NAME, new QRecord().withValue("id", source.getId()).withValue("indexedRecordCount", 0) + .withValue("indexedChunkCount", 0).withValue("lastIndexStatus", "CLEARED").withValue("lastIndexError", null)); + } + + + + /******************************************************************************* + ** Only a real clicking administrator may launch manual indexing. + *******************************************************************************/ + private static void requireAdmin() throws QException + { + if(QContext.getQSession() == null || !QContext.getQSession().hasPermission(AssistantQBitConfig.ADMIN_PERMISSION)) + { + throw (new QException("Knowledge indexing requires assistant.admin.")); + } + } + + + + /******************************************************************************* + ** Stop between UI API reads as well as before every embedding dispatch. + *******************************************************************************/ + private void checkKillSwitch() throws QException + { + if(runtime.getPolicyService().current(true).killSwitch()) + { + throw (new QException("The assistant kill switch is on.")); + } + } + + + + /******************************************************************************* + ** Current caller context, never a stored or service session. + *******************************************************************************/ + public UiApiToolbox.Context context() throws QException + { + return (new UiApiToolbox.Context(QContext.getQSession().getUuid(), QContext.getQSession().getValue("userTimezone"), runtime.getUiApiGateway(), runtime.getPolicyService().current(true).tools(), null)); + } + + + + /******************************************************************************* + ** Configured text field names in stable order. + *******************************************************************************/ + public static List fields(AssistantRagSource source) + { + return (source.getFieldNames() == null ? List.of() : java.util.Arrays.stream(source.getFieldNames().split(",")).map(String::trim).filter(s -> !s.isEmpty()).distinct().toList()); + } + + + + /******************************************************************************* + ** Default chunk size. + *******************************************************************************/ + public static Integer size(AssistantRagSource source) + { + return (source.getChunkChars() == null ? 1500 : source.getChunkChars()); + } + + + + /******************************************************************************* + ** Default overlap. + *******************************************************************************/ + public static Integer overlap(AssistantRagSource source) + { + return (source.getChunkOverlapChars() == null ? 200 : source.getChunkOverlapChars()); + } + + + + /******************************************************************************* + ** Source filter is QQQ filter JSON, passed unchanged apart from paging/order. + *******************************************************************************/ + private static QQueryFilter filter(AssistantRagSource source) throws QException + { + try + { + QQueryFilter filter = source.getFilterJson() == null || source.getFilterJson().isBlank() ? new QQueryFilter() : JsonUtils.toObject(source.getFilterJson(), QQueryFilter.class); + if(filter == null) + { + throw (new IllegalArgumentException("null filter")); + } + return (filter); + } + catch(Exception e) + { + throw (new QException("The query filter must be valid QQQ query filter JSON.")); + } + } + /******************************************************************************* + ** Remove a deleted source without attempting to update its missing row. + *******************************************************************************/ + public synchronized void remove(AssistantRagSource source) throws QException + { + requireAdmin(); + vectors.deleteSource(source.getId()); + } + +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/TextChunks.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/TextChunks.java new file mode 100644 index 0000000..92d16a9 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/TextChunks.java @@ -0,0 +1,80 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.ArrayList; +import java.util.HexFormat; +import java.util.List; + + +/******************************************************************************* + ** Deterministic overlapping chunks. Text exists only during indexing/retrieval. + *******************************************************************************/ +public final class TextChunks +{ + /******************************************************************************* + ** Utility class. + *******************************************************************************/ + private TextChunks() + { + } + + + + /******************************************************************************* + ** Fixed character boundaries are reproduced on the fresh UI API value. + *******************************************************************************/ + public static List split(String text, Integer size, Integer overlap) + { + if(size == null || overlap == null || size < 1 || overlap < 0 || overlap >= size) + { + throw (new IllegalArgumentException("Chunk overlap must be nonnegative and smaller than chunk size.")); + } + List chunks = new ArrayList<>(); + for(Integer start = 0; start < text.length(); start += size - overlap) + { + Integer end = Math.min(text.length(), start + size); + chunks.add(text.substring(start, end)); + if(end == text.length()) + { + break; + } + } + return (chunks); + } + + + + /******************************************************************************* + ** SHA-256 permits unchanged detection without retaining any text. + *******************************************************************************/ + public static String hash(String text) + { + try + { + return (HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(text.getBytes(StandardCharsets.UTF_8)))); + } + catch(NoSuchAlgorithmException e) + { + throw (new IllegalStateException("SHA-256 is required by the Java runtime.", e)); + } + } +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/VectorStore.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/VectorStore.java new file mode 100644 index 0000000..a124449 --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/rag/VectorStore.java @@ -0,0 +1,76 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +import java.util.List; +import java.util.Set; +import com.kingsrook.qqq.backend.core.exceptions.QException; + + +/******************************************************************************* + ** Per-instance vectors only. Every hit must be re-read through the UI API. + *******************************************************************************/ +public interface VectorStore +{ + /******************************************************************************* + ** Source identity disambiguates field/chunk layouts on the same table. + *******************************************************************************/ + record Hit(Integer source, String table, String pk, String field, Integer chunkIndex, Double score) + { + } + + /******************************************************************************* + ** Upsert by source/table/key/field/chunk, without storing content. + *******************************************************************************/ + void upsert(List chunks) throws QException; + + /******************************************************************************* + ** Delete all chunks for a record within a source. + *******************************************************************************/ + void deleteRecord(Integer source, String table, String pk) throws QException; + + /******************************************************************************* + ** Remove a source's entire index. + *******************************************************************************/ + void deleteSource(Integer source) throws QException; + + /******************************************************************************* + ** Cosine similarity, highest first; an empty table set returns no hits. + *******************************************************************************/ + List search(float[] query, Integer k, Set tables) throws QException; + + /******************************************************************************* + ** Restrict ranking to compatible enabled sources, before applying the limit. + *******************************************************************************/ + List search(float[] query, Integer k, Set tables, Set sources) throws QException; + + /******************************************************************************* + ** Count indexed chunks. + *******************************************************************************/ + Long count(Integer source) throws QException; + + /******************************************************************************* + ** Content-free reconciliation seam for unchanged and disappeared chunks. + *******************************************************************************/ + List chunks(Integer source) throws QException; + + /******************************************************************************* + ** Atomically replace one record, removing shortened or removed fields. + *******************************************************************************/ + void replaceRecord(Integer source, String table, String pk, List chunks) throws QException; +} diff --git a/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/tools/KnowledgeTool.java b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/tools/KnowledgeTool.java new file mode 100644 index 0000000..fab725e --- /dev/null +++ b/qbit-assistant/src/main/java/com/kingsrook/qbits/assistant/tools/KnowledgeTool.java @@ -0,0 +1,227 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.tools; + + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import com.kingsrook.qbits.assistant.AssistantRuntime; +import com.kingsrook.qbits.assistant.classification.DataClassification; +import com.kingsrook.qbits.assistant.engine.PolicyService; +import com.kingsrook.qbits.assistant.engine.TurnEngine; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; +import com.kingsrook.qbits.assistant.provider.ModelToolDefinition; +import com.kingsrook.qbits.assistant.rag.EmbeddingService; +import com.kingsrook.qbits.assistant.rag.RagService; +import com.kingsrook.qbits.assistant.rag.TextChunks; +import com.kingsrook.qbits.assistant.rag.VectorStore; +import com.kingsrook.qqq.backend.core.actions.ai.AiRecordRef; +import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome; +import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult; +import com.kingsrook.qqq.backend.core.actions.ai.output.AiDataEnvelope; +import com.kingsrook.qqq.backend.core.exceptions.QException; +import com.kingsrook.qqq.backend.core.model.session.QSession; +import com.kingsrook.qqq.backend.core.utils.JsonUtils; + + +/******************************************************************************* + ** Permission-filtered semantic retrieval. Every candidate is re-read as the caller. + *******************************************************************************/ +public class KnowledgeTool +{ + public static final String NAME = "search_knowledge"; + private final AssistantRuntime runtime; + private final RagService rag; + private final VectorStore vectors; + + + + /******************************************************************************* + ** Bind the existing services; the vector store contains no record content. + *******************************************************************************/ + public KnowledgeTool(AssistantRuntime runtime, RagService rag, VectorStore vectors) + { + this.runtime = runtime; + this.rag = rag; + this.vectors = vectors; + } + + + + /******************************************************************************* + ** A source must be enabled and readable in the UI before offering the tool. + *******************************************************************************/ + public Boolean offered(PolicyService.ToolPolicy policy, QSession session) + { + if(session == null || (policy != null && !policy.allows(NAME))) + { + return (false); + } + try + { + Set tables = new LinkedHashSet<>(); + for(AssistantRagSource source : rag.sources()) + { + if(Boolean.TRUE.equals(source.getEnabled()) && runtime.getDataAccessService().forTable(source.getTableName()).mode().allowsRead() && rag.validate(source).isEmpty()) + { + tables.add(source.getTableName()); + } + } + if(tables.isEmpty()) + { + return (false); + } + return (runtime.getAccessExplainer().explain(runtime.getUiApiGateway(), session.getUuid(), policy, null, null).stream() + .anyMatch(t -> tables.contains(t.table()) && t.uiRead() && t.mode().allowsRead())); + } + catch(QException e) + { + return (false); + } + } + + + + /******************************************************************************* + ** Stable, small schema in the same sorted tool list as other tools. + *******************************************************************************/ + public ModelToolDefinition definition() + { + return (new ModelToolDefinition().withName(NAME).withDescription("Search indexed application knowledge. Cite returned records using their links.").withStrict(false) + .withInputSchema(Map.of("type", "object", "properties", Map.of("query", Map.of("type", "string"), + "tables", Map.of("type", "array", "items", Map.of("type", "string")), "limit", Map.of("type", "integer", "minimum", 1, "maximum", 8)), "required", List.of("query")))); + } + + + + /******************************************************************************* + ** Expected denials disclose only a filtered count, never hit identifiers. + *******************************************************************************/ + public AiToolResult execute(UiApiToolbox.Context context, String callId, Map input, DataClassification classification, Integer conversationId, Integer turnId) + { + AiToolResult result = new AiToolResult().withToolName(NAME).withToolCallId(callId); + Long started = System.nanoTime(); + try + { + if(context.policy() != null && !context.policy().allows(NAME)) + { + return (result.withOutcome(AiToolOutcome.DENIED).withSummary("Not allowed").withModelContent("The policy does not allow search_knowledge.")); + } + if(input == null || !(input.get("query") instanceof String query) || query.isBlank() || query.length() > 8000 + || (input.containsKey("tables") && (!(input.get("tables") instanceof List requested) || requested.stream().anyMatch(t -> !(t instanceof String)))) + || (input.containsKey("limit") && (!(input.get("limit") instanceof Number n) || n.doubleValue() != n.intValue() || n.intValue() < 1))) + { + return (result.withOutcome(AiToolOutcome.INVALID_INPUT).withSummary("Invalid search").withModelContent("Supply a nonempty query (at most 8000 characters), optional table names and a positive integer limit.")); + } + Integer limit = input.get("limit") instanceof Number n ? Math.min(8, n.intValue()) : 5; + Set requested = input.get("tables") instanceof List list ? new LinkedHashSet<>(list) : Set.of(); + Map sources = new LinkedHashMap<>(); + Map> modelTables = new LinkedHashMap<>(); + for(AssistantRagSource source : rag.sources()) + { + if(Boolean.TRUE.equals(source.getEnabled()) && (requested.isEmpty() || requested.contains(source.getTableName())) + && runtime.getDataAccessService().forTable(source.getTableName()).mode().allowsRead() && rag.validate(source).isEmpty()) + { + sources.put(source.getId(), source); + modelTables.computeIfAbsent(source.getEmbeddingModelId(), key -> new LinkedHashSet<>()).add(source.getTableName()); + } + } + List hits = new ArrayList<>(); + EmbeddingService embeddings = new EmbeddingService(runtime); + for(var model : modelTables.entrySet()) + { + float[] vector = embeddings.embed(model.getKey(), List.of(query), classification, String.join(", ", model.getValue()), conversationId, turnId).getFirst(); + Set compatible = new LinkedHashSet<>(); + sources.values().stream().filter(s -> model.getKey().equals(s.getEmbeddingModelId())).forEach(s -> compatible.add(s.getId())); + for(VectorStore.Hit hit : vectors.search(vector, limit * 4, model.getValue(), compatible)) + { + AssistantRagSource source = sources.get(hit.source()); + if(source != null && model.getKey().equals(source.getEmbeddingModelId()) && source.getTableName().equals(hit.table()) && RagService.fields(source).contains(hit.field())) + { + hits.add(hit); + } + } + } + hits.sort(Comparator.comparing(VectorStore.Hit::score).reversed()); + List> records = new ArrayList<>(); + Set> returned = new LinkedHashSet<>(); + Integer filtered = 0; + for(VectorStore.Hit hit : hits) + { + if(records.size() >= limit) + { + break; + } + if(runtime.getPolicyService().current(true).killSwitch()) + { + throw (new QException("The assistant kill switch is on.")); + } + UiApiToolbox.Outcome read = rag.read(context, hit.table(), hit.pk()); + if(read.getResult().getOutcome() == AiToolOutcome.DENIED || read.getResult().getOutcome() == AiToolOutcome.NOT_FOUND) + { + filtered++; + continue; + } + if(read.getResult().getOutcome() != AiToolOutcome.OK) + { + throw (new QException("A knowledge record could not be read through the UI API. Retry the search.")); + } + Map values = RagService.values(read); + AssistantRagSource source = sources.get(hit.source()); + List chunks = TextChunks.split(values.getOrDefault(hit.field(), ""), RagService.size(source), RagService.overlap(source)); + if(hit.chunkIndex() >= chunks.size() || !returned.add(List.of(hit.table(), hit.pk(), hit.field(), String.valueOf(hit.chunkIndex())))) + { + filtered++; + continue; + } + String href = String.valueOf(((Map) read.getResult().getRender()).get("href")); + if(!TurnEngine.isAppHref(href)) + { + filtered++; + continue; + } + String title = values.getOrDefault(source.getTitleFieldName(), hit.pk()); + title = title.substring(0, Math.min(300, title.length())); + String snippet = chunks.get(hit.chunkIndex()); + snippet = snippet.substring(0, Math.min(1500, snippet.length())); + records.add(Map.of("table", hit.table(), "recordId", hit.pk(), "title", title, "snippet", snippet, "score", hit.score(), "href", href)); + result.getRecordRefs().add(new AiRecordRef().withTableName(hit.table()).withId(hit.pk()).withLabel(title).withHref(href)); + result.getRedactions().addAll(read.getResult().getRedactions()); + } + Map content = Map.of("results", records, "filtered", filtered); + return (result.withOutcome(AiToolOutcome.OK).withRowCount(records.size()).withSummary("Found " + records.size() + " knowledge results; filtered " + filtered) + .withPermissionCheck("uiApi:GET /table/{tableName}/{primaryKey}").withModelContent(AiDataEnvelope.wrap(NAME, null, JsonUtils.toJson(content)))); + } + catch(QException e) + { + return (result.withOutcome(AiToolOutcome.ERROR).withSummary("Knowledge search failed").withModelContent(e.getMessage())); + } + catch(RuntimeException e) + { + return (result.withOutcome(AiToolOutcome.ERROR).withSummary("Knowledge search failed").withModelContent("Knowledge search could not complete. Check the source and embedding configuration.")); + } + finally + { + result.withDurationMs((System.nanoTime() - started) / 1_000_000); + } + } +} diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantQBitProducerTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantQBitProducerTest.java index 719e8a1..38a5c19 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantQBitProducerTest.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/AssistantQBitProducerTest.java @@ -88,7 +88,7 @@ class AssistantQBitProducerTest extends BaseAssistantTest // the nav: the home plus the main tables; processes are actions // ////////////////////////////////////////////////////////////////////// assertThat(qInstance.getApp(AssistantMetaDataBuilder.APP_NAME).getChildren()).extracting(c -> c.getName()) - .containsExactly("assistant_budget", "assistant_provider", "assistant_profile", "assistant_skill", "assistant_data_access", "assistant_conversation", "assistant_model_price", "assistant_policy"); + .containsExactly("assistant_rag_source", "assistant_budget", "assistant_provider", "assistant_profile", "assistant_skill", "assistant_data_access", "assistant_conversation", "assistant_model_price", "assistant_policy"); assertThat(qInstance.getApp(AssistantMetaDataBuilder.APP_NAME).getWidgets()).containsExactly("assistantHome", "assistantAccess"); assertThat(qInstance.getProcess(AssistantMetaDataBuilder.RETENTION_PROCESS_NAME).getTableName()).isEqualTo("assistant_conversation"); assertThat(qInstance.getProcess(AssistantMetaDataBuilder.RETENTION_PROCESS_NAME)).isNotNull(); diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/rag/RagTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/rag/RagTest.java new file mode 100644 index 0000000..74136df --- /dev/null +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/rag/RagTest.java @@ -0,0 +1,578 @@ +/* + * Copyright 2026 Kingsrook, LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.kingsrook.qbits.assistant.rag; + + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import com.kingsrook.qbits.assistant.AssistantQBitConfig; +import com.kingsrook.qbits.assistant.classification.DataClassification; +import com.kingsrook.qbits.assistant.engine.PolicyService; +import com.kingsrook.qbits.assistant.engine.TurnClassification; +import com.kingsrook.qbits.assistant.model.AssistantConversation; +import com.kingsrook.qbits.assistant.model.AssistantDataAccess; +import com.kingsrook.qbits.assistant.model.AssistantModel; +import com.kingsrook.qbits.assistant.model.AssistantPolicy; +import com.kingsrook.qbits.assistant.model.AssistantProvider; +import com.kingsrook.qbits.assistant.model.AssistantRagSource; +import com.kingsrook.qbits.assistant.model.AssistantUsage; +import com.kingsrook.qbits.assistant.processes.KnowledgeIndexProcess; +import com.kingsrook.qbits.assistant.provider.AssistantModelProvider; +import com.kingsrook.qbits.assistant.provider.AssistantProviderType; +import com.kingsrook.qbits.assistant.provider.EmbeddingRequest; +import com.kingsrook.qbits.assistant.provider.EmbeddingResponse; +import com.kingsrook.qbits.assistant.provider.ModelPart; +import com.kingsrook.qbits.assistant.provider.ModelUsage; +import com.kingsrook.qbits.assistant.provider.fake.ScriptedModelProvider; +import com.kingsrook.qbits.assistant.testing.AssistantTestFixture; +import com.kingsrook.qbits.assistant.testing.BaseAssistantTest; +import com.kingsrook.qbits.assistant.tools.KnowledgeTool; +import com.kingsrook.qbits.assistant.tools.UiApiToolbox; +import com.kingsrook.qbits.assistant.uiapi.UiApiCall; +import com.kingsrook.qqq.backend.core.actions.ai.AiExposure; +import com.kingsrook.qqq.backend.core.actions.ai.AiFieldMetaData; +import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome; +import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult; +import com.kingsrook.qqq.backend.core.context.QContext; +import com.kingsrook.qqq.backend.core.model.actions.processes.RunBackendStepInput; +import com.kingsrook.qqq.backend.core.model.actions.processes.RunBackendStepOutput; +import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter; +import com.kingsrook.qqq.backend.core.model.data.QRecord; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + + +/******************************************************************************* + ** Real UI API indexing/retrieval, with deterministic local embeddings and no credentials. + *******************************************************************************/ +class RagTest extends BaseAssistantTest +{ + private final InMemoryVectorStore vectors = new InMemoryVectorStore(); + private final Embeddings embeddings = new Embeddings(); + private AssistantRagSource source; + private AssistantModel model; + private AssistantProvider embeddingProvider; + + + + /******************************************************************************* + ** A hashed bag of words; capture batches solely to assert routing and masking. + *******************************************************************************/ + static class Embeddings extends ScriptedModelProvider + { + final List batches = new ArrayList<>(); + + + + /******************************************************************************* + ** Same text always yields the same 16-dimensional vector. + *******************************************************************************/ + @Override + public EmbeddingResponse embed(EmbeddingRequest request) + { + batches.add(request); + List vectors = new ArrayList<>(); + for(String text : request.input()) + { + float[] vector = new float[16]; + for(String word : text.toLowerCase(java.util.Locale.ROOT).split("[^a-z0-9]+")) + { + vector[Math.floorMod(word.hashCode(), vector.length)]++; + } + vectors.add(vector); + } + return (new EmbeddingResponse(vectors, new ModelUsage().withInputTokens((long) request.input().size()), "embed-test")); + } + } + + + + /******************************************************************************* + ** Keyless local type with the same provider instance for chat and embeddings. + *******************************************************************************/ + class Type implements AssistantProviderType + { + /******************************************************************************* + ** Test type identity. + *******************************************************************************/ + @Override + public String getType() + { + return ("rag-test"); + } + + + + /******************************************************************************* + ** Test label. + *******************************************************************************/ + @Override + public String getLabel() + { + return ("RAG test"); + } + + + + /******************************************************************************* + ** Chat default deliberately differs from the embedding model. + *******************************************************************************/ + @Override + public List getModels() + { + return (List.of("chat", "bge-m3", "nomic-embed-text")); + } + + + + /******************************************************************************* + ** No credentials for this in-process fake. + *******************************************************************************/ + @Override + public String getDefaultApiKeyEnvName() + { + return (null); + } + + + + /******************************************************************************* + ** Never contacted; satisfies provider configuration validation. + *******************************************************************************/ + @Override + public String getDefaultBaseUrl() + { + return ("http://localhost:11434/v1"); + } + + + + /******************************************************************************* + ** Local test location. + *******************************************************************************/ + @Override + public String getDefaultLocation(String kind) + { + return ("LOCAL"); + } + + + + /******************************************************************************* + ** Return the test recording provider. + *******************************************************************************/ + @Override + public AssistantModelProvider create(String key, String base) + { + return (embeddings); + } + } + + + + /******************************************************************************* + ** Configure only local fakes and the optional PHI fixture table. + *******************************************************************************/ + @Override + protected void customizeConfig(AssistantQBitConfig config) + { + config.withVectorStore(vectors).withProviderType(new Type()).withDataClassification("patient", DataClassification.PHI); + } + + + + /******************************************************************************* + ** Registered EMBEDDING model and a source in the real metadata-backed store. + *******************************************************************************/ + @BeforeEach + void sources() throws Exception + { + as(AssistantTestFixture.ADMIN); + embeddingProvider = config.getRuntime().getStore().insert(new AssistantProvider().withName("Local embeddings").withProviderType("rag-test") + .withModel("chat").withEnabled(true).withHipaaEligible(false)); + model = config.getRuntime().getStore().insert(new AssistantModel().withProviderId(embeddingProvider.getId()).withModelId("bge-m3") + .withModelKind("EMBEDDING").withEmbeddingDimensions(16).withEnabled(true)); + config.getRuntime().getStore().insert(new AssistantModel().withProviderId(embeddingProvider.getId()).withModelId("chat").withEnabled(true)); + source = config.getRuntime().getStore().insert(new AssistantRagSource().withName("Blog knowledge").withTableName("blog_post") + .withFieldNames("title,content").withTitleFieldName("title").withEmbeddingModelId(model.getId())); + assertThat(source.getChunkChars()).isEqualTo(1500); + } + + + + /******************************************************************************* + ** The session used both for configuration reads and UI API calls. + *******************************************************************************/ + private void as(String user) + { + QContext.init(qInstance, AssistantTestFixture.ownerSession(user, config)); + } + + + + /******************************************************************************* + ** Retrieval helper, retaining the caller's exact session UUID. + *******************************************************************************/ + private AiToolResult search(String user, DataClassification classification) throws Exception + { + as(user); + return (config.getRuntime().getKnowledgeTool().execute(config.getRuntime().getRagService().context(), "search-1", Map.of("query", "Drain upgrade verify", "limit", 8), classification, null, null)); + } + + + + /******************************************************************************* + ** Unchanged content performs no embedding; disappeared records are removed. + *******************************************************************************/ + @Test + void indexesReconcilesAndReports() throws Exception + { + RagService rag = config.getRuntime().getRagService(); + RunBackendStepInput input = new RunBackendStepInput().withRecords(List.of(source.toQRecord())); + RunBackendStepOutput output = new RunBackendStepOutput(); + KnowledgeIndexProcess.Review review = new KnowledgeIndexProcess.Review(); + review.runOnePage(input, output); + assertThat(review.getProcessSummary(output, false)).hasSize(1); + assertThat(embeddings.batches).isEmpty(); + KnowledgeIndexProcess.Index index = new KnowledgeIndexProcess.Index(); + index.runOnePage(input, new RunBackendStepOutput()); + assertThat(index.getProcessSummary(output, true)).hasSize(1); + assertThat(vectors.count(source.getId())).isPositive(); + Integer calls = embeddings.batches.size(); + RagService.Result unchanged = rag.index(source); + assertThat(unchanged.errors()).as(unchanged.reason()).isZero(); + assertThat(unchanged.embedded()).isZero(); + assertThat(unchanged.unchanged()).isEqualTo(vectors.count(source.getId()).intValue()); + assertThat(embeddings.batches).hasSize(calls); + assertThat(vectors.chunks(source.getId())).allSatisfy(v -> assertThat(v.contentHash()).matches("[a-f0-9]{64}")); + assertThat(ChunkVector.class.getRecordComponents()).extracting(java.lang.reflect.RecordComponent::getName) + .containsExactly("source", "table", "pk", "field", "chunkIndex", "contentHash", "embedding"); + assertThat(config.getRuntime().getUiApiGateway().invoke(new UiApiCall("DELETE", "/table/{tableName}/{primaryKey}") + .withPathParam("tableName", "blog_post").withPathParam("primaryKey", "1").withSessionUUID(AssistantTestFixture.ADMIN)).status()).isEqualTo(200); + RagService.Result changed = rag.index(source); + assertThat(changed.errors()).as(changed.reason()).isZero(); + assertThat(changed.removed()).isEqualTo(2); + assertThat(vectors.chunks(source.getId())).noneMatch(v -> v.pk().equals("1")); + new KnowledgeIndexProcess.Clear().runOnePage(input, new RunBackendStepOutput()); + assertThat(vectors.count(source.getId())).isZero(); + } + + + + /******************************************************************************* + ** UI read permission is checked again for every vector hit, with no identifiers + ** or snippets leaking from denied records. Exposure is reused from get_record. + *******************************************************************************/ + @Test + void retrievesAsUserAndFiltersUiDenials() throws Exception + { + assertThat(config.getRuntime().getRagService().index(source).errors()).isZero(); + AiToolResult allowed = search(AssistantTestFixture.EDITOR, DataClassification.PUBLIC); + assertThat(allowed.getOutcome()).as(allowed.getModelContent()).isEqualTo(AiToolOutcome.OK); + assertThat(allowed.getRowCount()).isPositive(); + assertThat(allowed.getModelContent()).contains("Drain, upgrade, verify.", "/app/blog_post/").doesNotContain(AssistantTestFixture.HIDDEN_SENTINEL, AssistantTestFixture.PASSWORD_SENTINEL); + assertThat(config.getRuntime().getKnowledgeTool().offered(PolicyService.ToolPolicy.defaults(), QContext.getQSession())).isTrue(); + AiToolResult denied = search(AssistantTestFixture.ANALYST, DataClassification.PUBLIC); + assertThat(denied.getOutcome()).isEqualTo(AiToolOutcome.OK); + assertThat(denied.getRowCount()).isZero(); + assertThat(denied.getSummary()).doesNotEndWith("filtered 0"); + assertThat(denied.getRecordRefs()).isEmpty(); + assertThat(denied.getModelContent()).doesNotContain("blog_post", "Drain", "/app/"); + assertThat(config.getRuntime().getKnowledgeTool().offered(PolicyService.ToolPolicy.defaults(), QContext.getQSession())).isFalse(); + assertThat(config.getRuntime().getUiApiGateway().invoke(new UiApiCall("GET", "/table/{tableName}/{primaryKey}") + .withPathParam("tableName", "blog_post").withPathParam("primaryKey", "1").withSessionUUID(AssistantTestFixture.ANALYST)).status()).isEqualTo(403); + } + + + + /******************************************************************************* + ** Fresh field/table restrictions narrow both index validation and tool offering. + *******************************************************************************/ + @Test + void validatesExposureAndMasksBeforeEmbedding() throws Exception + { + RagService rag = config.getRuntime().getRagService(); + assertThat(rag.validate(new AssistantRagSource(source.toQRecord()).withFieldNames("internalNotes"))).isNotEmpty(); + assertThat(rag.validate(new AssistantRagSource(source.toQRecord()).withFieldNames("id"))).isNotEmpty(); + assertThat(rag.validate(new AssistantRagSource(source.toQRecord()).withTableName("assistant_model"))).isNotEmpty(); + assertThat(rag.validate(new AssistantRagSource(source.toQRecord()).withChunkChars(100).withChunkOverlapChars(100))).isNotEmpty(); + qInstance.getTable("blog_post").getField("content").withSupplementalMetaData(new AiFieldMetaData().withExposure(AiExposure.MASKED)); + assertThat(rag.index(source).errors()).isZero(); + assertThat(embeddings.batches.stream().flatMap(b -> b.input().stream())).noneMatch(s -> s.contains("Drain, upgrade")); + assertThat(search(AssistantTestFixture.EDITOR, DataClassification.PUBLIC).getModelContent()).doesNotContain("Drain, upgrade"); + as(AssistantTestFixture.ADMIN); + config.getRuntime().getStore().insert(new AssistantDataAccess().withObjectType("TABLE").withObjectName("blog_post").withMode("NONE")); + assertThat(rag.validate(source)).anyMatch(s -> s.contains("NONE")); + assertThat(config.getRuntime().getKnowledgeTool().offered(PolicyService.ToolPolicy.defaults(), QContext.getQSession())).isFalse(); + assertThat(search(AssistantTestFixture.EDITOR, DataClassification.PUBLIC).getRowCount()).isZero(); + } + + + + /******************************************************************************* + ** PHI dispatch fails before text leaves the instance; eligible retrieval + ** contributes table classification through the existing record-ref mechanism. + *******************************************************************************/ + @Test + void phiEmbeddingGuardAndClassification() throws Exception + { + AssistantRagSource patient = config.getRuntime().getStore().insert(new AssistantRagSource().withName("Patients").withTableName("patient") + .withFieldNames("diagnosis").withEmbeddingModelId(model.getId())); + RagService.Result blocked = config.getRuntime().getRagService().index(patient); + assertThat(blocked.errors()).isEqualTo(1); + assertThat(blocked.reason()).contains("HIPAA-eligible"); + assertThat(embeddings.batches).isEmpty(); + assertThat(search(AssistantTestFixture.EDITOR, DataClassification.PHI).getModelContent()).contains("HIPAA-eligible"); + assertThat(embeddings.batches).isEmpty(); + as(AssistantTestFixture.ADMIN); + config.getRuntime().getStore().updateFields(AssistantProvider.TABLE_NAME, new QRecord().withValue("id", embeddingProvider.getId()).withValue("hipaaEligible", true)); + assertThat(config.getRuntime().getRagService().index(patient).errors()).isZero(); + as(AssistantTestFixture.EDITOR); + AiToolResult result = config.getRuntime().getKnowledgeTool().execute(config.getRuntime().getRagService().context(), "phi", Map.of("query", "diagnosis", "tables", List.of("patient")), DataClassification.PUBLIC, null, null); + assertThat(result.getRowCount()).isPositive(); + TurnClassification classification = new TurnClassification(config.getClassificationSource(), new AssistantConversation().withMaxClassification("PUBLIC"), null); + classification.observe(ModelPart.toolUse("phi", KnowledgeTool.NAME, Map.of()), result); + assertThat(classification.getClassification()).isEqualTo(DataClassification.PHI); + } + + + + /******************************************************************************* + ** Batches are capped, accounted as EMBED, and blocked by the kill switch. + *******************************************************************************/ + @Test + void batchesUsageAndSelection() throws Exception + { + new EmbeddingService(config.getRuntime()).embed(model.getId(), java.util.Collections.nCopies(130, "word"), DataClassification.PUBLIC, "blog_post", null, null); + assertThat(embeddings.batches).extracting(b -> b.input().size()).containsExactly(64, 64, 2); + assertThat(config.getRuntime().getStore().query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class)) + .hasSize(3).allSatisfy(row -> assertThat(row.getPurpose()).isEqualTo("EMBED")); + assertThat(config.getRuntime().getProviderService().resolveModel(model.getId(), false)).isNull(); + assertThat(config.getRuntime().getProviderService().forRow(new AssistantProvider(embeddingProvider.toQRecord()).withModel(null), "chat").model()).isEqualTo("chat"); + assertThatThrownBy(() -> config.getRuntime().getProviderService().resolveEmbedding(-1, false)).hasMessageContaining("EMBEDDING"); + assertThat(com.kingsrook.qbits.assistant.engine.ModelDiscovery.run(config.getRuntime(), embeddingProvider).ok()).isTrue(); + assertThat(config.getRuntime().getStore().listModels()).filteredOn(m -> m.getModelId().equals("nomic-embed-text")) + .singleElement().satisfies(m -> + { + assertThat(m.getModelKind()).isEqualTo("EMBEDDING"); + assertThat(m.getEmbeddingDimensions()).isEqualTo(768); + }); + config.getRuntime().getStore().insert(new AssistantPolicy().withName("Stop").withKillSwitch(true)); + assertThatThrownBy(() -> new EmbeddingService(config.getRuntime()).embed(model.getId(), List.of("blocked"), DataClassification.PUBLIC, "blog_post", null, null)).hasMessageContaining("kill switch"); + assertThat(embeddings.batches).hasSize(3); + assertThat(com.kingsrook.qbits.assistant.engine.ConnectionTest.runModel(config.getRuntime(), model, AssistantTestFixture.ADMIN).status()).isEqualTo("ERROR"); + } + + + + /******************************************************************************* + ** A full golden turn cites fresh knowledge and writes the standard tool log. + *******************************************************************************/ + @Test + @Tag("golden") + void goldenAnswerCitesKnowledge() throws Exception + { + assertThat(config.getRuntime().getRagService().index(source).errors()).isZero(); + embeddings.thenToolUse(KnowledgeTool.NAME, Map.of("query", "Drain upgrade verify")) + .thenRespond(com.kingsrook.qbits.assistant.testing.EvalSupport.answer(request -> + { + Map data = com.kingsrook.qbits.assistant.testing.EvalSupport.firstResult(request); + Map record = (Map) ((List) data.get("results")).getFirst(); + return (record.get("snippet") + " [Source](" + record.get("href") + ")"); + })); + var outcome = runner.run(AssistantTestFixture.EDITOR, "Find the knowledge about safely upgrading"); + assertThat(outcome.errorCode()).isNull(); + assertThat(outcome.text()).contains("[Source](/app/blog_post/"); + assertThat(outcome.toolCalls()).singleElement().satisfies(call -> + { + assertThat(call.getToolName()).isEqualTo(KnowledgeTool.NAME); + assertThat(call.getOutcome()).isEqualTo("OK"); + assertThat(call.getUiApiOperation()).startsWith("GET /table/"); + }); + } + + + + /******************************************************************************* + ** Defensive vectors and chunk boundary behavior, including invalid inputs. + *******************************************************************************/ + @Test + void vectorAndChunkContracts() throws Exception + { + float[] raw = { 1, 0 }; + ChunkVector chunk = new ChunkVector(1, "table", "key", "body", 0, TextChunks.hash("text"), raw); + raw[0] = 9; + assertThat(chunk.embedding()).containsExactly(1, 0); + chunk.embedding()[0] = 5; + assertThat(chunk.embedding()).containsExactly(1, 0); + InMemoryVectorStore store = new InMemoryVectorStore(); + store.upsert(List.of(chunk)); + assertThat(store.search(new float[] { 1, 0 }, 1, Set.of("table"))).singleElement().satisfies(hit -> assertThat(hit.score()).isEqualTo(1.0)); + assertThat(store.search(new float[] { 0, 0 }, 1, Set.of("table"))).isEmpty(); + assertThat(store.search(new float[] { 1 }, 1, Set.of("table"))).isEmpty(); + assertThat(store.search(raw, 1, Set.of())).isEmpty(); + store.replaceRecord(1, "table", "key", List.of()); + assertThat(store.count(1)).isZero(); + assertThat(TextChunks.split("abcdefgh", 4, 1)).containsExactly("abcd", "defg", "gh"); + assertThat(TextChunks.split("", 4, 0)).isEmpty(); + assertThatThrownBy(() -> TextChunks.split("a", 1, 1)).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new ChunkVector(1, "t", "p", "f", 0, "invalid", raw)).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new ChunkVector(1, "t", "p", "f", 0, TextChunks.hash("x"), new float[] { Float.NaN })).isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> new ScriptedModelProvider().embed(new EmbeddingRequest("chat", List.of("text")))).hasMessageContaining("unsupported"); + } + + + + /******************************************************************************* + ** Test auth supplies a record-lock key for each real user's UI API session. + *******************************************************************************/ + public static class LockedSessions extends AssistantTestFixture.TestSessionCustomizer + { + /******************************************************************************* + ** The admin can index all records; each editor can see only its own key. + *******************************************************************************/ + @Override + public void customizeSession(com.kingsrook.qqq.backend.core.model.metadata.QInstance instance, com.kingsrook.qqq.backend.core.model.session.QSession session, Map context) + { + super.customizeSession(instance, session, context); + session.withSecurityKeyValue("rag_owner", session.getUuid()); + if(session.getUuid().startsWith("admin")) + { + session.withSecurityKeyValue("rag_all", true); + } + } + } + + + + /******************************************************************************* + ** The owner and non-owner hit identical UI record-lock enforcement. + *******************************************************************************/ + @Test + void recordLocksFilterNonOwner() throws Exception + { + qInstance.addSecurityKeyType(new com.kingsrook.qqq.backend.core.model.metadata.security.QSecurityKeyType().withName("rag_owner").withAllAccessKeyName("rag_all")); + qInstance.getAuthentication().withCustomizer(new com.kingsrook.qqq.backend.core.model.metadata.code.QCodeReference(LockedSessions.class)); + var table = qInstance.getTable("blog_post"); + table.withRecordSecurityLock(new com.kingsrook.qqq.backend.core.model.metadata.security.RecordSecurityLock().withFieldName("slug").withSecurityKeyType("rag_owner")); + assertThat(config.getRuntime().getUiApiGateway().invoke(new UiApiCall("PATCH", "/table/{tableName}/{primaryKey}") + .withPathParam("tableName", "blog_post").withPathParam("primaryKey", "1").withSessionUUID(AssistantTestFixture.ADMIN) + .withJsonBody(Map.of("slug", AssistantTestFixture.EDITOR))).status()).isEqualTo(200); + RagService.Result indexed = config.getRuntime().getRagService().index(source); + assertThat(indexed.errors()).as(indexed.reason()).isZero(); + AiToolResult owner = search(AssistantTestFixture.EDITOR, DataClassification.PUBLIC); + assertThat(owner.getRowCount()).as(owner.getModelContent()).isPositive(); + assertThat(owner.getRecordRefs()).allSatisfy(ref -> assertThat(String.valueOf(ref.getId())).isEqualTo("1")); + AiToolResult other = search(AssistantTestFixture.OTHER_EDITOR, DataClassification.PUBLIC); + assertThat(other.getRowCount()).isZero(); + assertThat(other.getSummary()).doesNotEndWith("filtered 0"); + Integer status = config.getRuntime().getUiApiGateway().invoke(new UiApiCall("GET", "/table/{tableName}/{primaryKey}") + .withPathParam("tableName", "blog_post").withPathParam("primaryKey", "1").withSessionUUID(AssistantTestFixture.OTHER_EDITOR)).status(); + assertThat(status).isIn(403, 404); + } + + + + /******************************************************************************* + ** Filter paging and chunk shrinkage reconcile without holding record content. + *******************************************************************************/ + @Test + void sourceFilterAndChangedChunks() throws Exception + { + source.setFilterJson("{\"criteria\":[{\"fieldName\":\"id\",\"operator\":\"EQUALS\",\"values\":[1]}]}"); + source.setFieldNames("content"); + source.setChunkChars(100); + source.setChunkOverlapChars(20); + config.getRuntime().getStore().update(source); + var gateway = config.getRuntime().getUiApiGateway(); + var update = new UiApiCall("PATCH", "/table/{tableName}/{primaryKey}").withPathParam("tableName", "blog_post").withPathParam("primaryKey", "1").withSessionUUID(AssistantTestFixture.ADMIN); + assertThat(gateway.invoke(update.withJsonBody(Map.of("content", "Knowledge ".repeat(40)))).status()).isEqualTo(200); + RagService.Result first = config.getRuntime().getRagService().index(source); + assertThat(first.errors()).as(first.reason()).isZero(); + assertThat(first.recordsSeen()).isEqualTo(1); + assertThat(first.embedded()).isGreaterThan(1); + assertThat(gateway.invoke(update.withJsonBody(Map.of("content", "Short updated knowledge"))).status()).isEqualTo(200); + assertThat(search(AssistantTestFixture.EDITOR, DataClassification.PUBLIC).getModelContent()).contains("Short updated knowledge").doesNotContain("Knowledge Knowledge"); + as(AssistantTestFixture.ADMIN); + RagService.Result shortened = config.getRuntime().getRagService().index(source); + assertThat(shortened.errors()).as(shortened.reason()).isZero(); + assertThat(shortened.embedded()).isEqualTo(1); + assertThat(shortened.removed()).isEqualTo(first.embedded() - 1); + source.setFieldNames("title"); + config.getRuntime().getStore().update(source); + assertThat(vectors.count(source.getId())).isZero(); + assertThat(config.getRuntime().getRagService().validate(new AssistantRagSource(source.toQRecord()).withFilterJson("{"))).isNotEmpty(); + as(AssistantTestFixture.EDITOR); + assertThatThrownBy(() -> config.getRuntime().getRagService().index(source)).hasMessageContaining("assistant.admin"); + } + + + + /******************************************************************************* + ** Standard process execution reaches both Review and Result backend steps. + *******************************************************************************/ + @Test + void standardProcessAndPolicy() throws Exception + { + var input = new com.kingsrook.qqq.backend.core.model.actions.processes.RunProcessInput(); + input.setProcessName(KnowledgeIndexProcess.INDEX); + input.setCallback(com.kingsrook.qqq.backend.core.actions.processes.QProcessCallbackFactory.forPrimaryKey("id", source.getId())); + input.setFrontendStepBehavior(com.kingsrook.qqq.backend.core.model.actions.processes.RunProcessInput.FrontendStepBehavior.SKIP); + new com.kingsrook.qqq.backend.core.actions.processes.RunProcessAction().execute(input); + assertThat(vectors.count(source.getId())).isPositive(); + assertThat(com.kingsrook.qbits.assistant.engine.ConnectionTest.runModel(config.getRuntime(), model, AssistantTestFixture.ADMIN).status()).isEqualTo("OK"); + RunBackendStepOutput clearOutput = new RunBackendStepOutput(); + KnowledgeIndexProcess.ClearReview clearReview = new KnowledgeIndexProcess.ClearReview(); + clearReview.runOnePage(new RunBackendStepInput().withRecords(List.of(source.toQRecord())), clearOutput); + assertThat(clearReview.getProcessSummary(clearOutput, false)).hasSize(1); + var policy = new PolicyService.ToolPolicy(true, Set.of("search_help"), 50, true); + assertThat(config.getRuntime().getKnowledgeTool().offered(policy, QContext.getQSession())).isFalse(); + var context = new UiApiToolbox.Context(AssistantTestFixture.ADMIN, null, config.getRuntime().getUiApiGateway(), policy, null); + assertThat(config.getRuntime().getKnowledgeTool().execute(context, "denied", Map.of("query", "x"), DataClassification.PUBLIC, null, null).getOutcome()).isEqualTo(AiToolOutcome.DENIED); + assertThat(config.getRuntime().getKnowledgeTool().execute(config.getRuntime().getRagService().context(), "invalid", Map.of(), DataClassification.PUBLIC, null, null).getOutcome()).isEqualTo(AiToolOutcome.INVALID_INPUT); + config.getBudgets().withUserDailyTokensHard(1L); + assertThatThrownBy(() -> new EmbeddingService(config.getRuntime()).embed(model.getId(), List.of("blocked"), DataClassification.PUBLIC, "blog_post", null, null)).hasMessageContaining("budget"); + } + + + + /******************************************************************************* + ** A multi-page scan does not stop at the first UI API query page. + *******************************************************************************/ + @Test + void indexesMultiplePages() throws Exception + { + List records = new ArrayList<>(); + for(Integer i = 0; i < 105; i++) + { + records.add(new QRecord().withValue("title", "Extra page " + i).withValue("content", "Paged knowledge " + i)); + } + new com.kingsrook.qqq.backend.core.actions.tables.InsertAction().execute(new com.kingsrook.qqq.backend.core.model.actions.tables.insert.InsertInput("page").withRecords(records)); + source.setTableName("page"); + source.setFieldNames("content"); + config.getRuntime().getStore().update(source); + RagService.Result result = config.getRuntime().getRagService().index(source); + assertThat(result.errors()).as(result.reason()).isZero(); + assertThat(result.recordsSeen()).isEqualTo(117); + assertThat(vectors.count(source.getId())).isEqualTo(117); + } +} diff --git a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/settings/AssistantSettingsTest.java b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/settings/AssistantSettingsTest.java index d18cefc..9d14b52 100644 --- a/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/settings/AssistantSettingsTest.java +++ b/qbit-assistant/src/test/java/com/kingsrook/qbits/assistant/settings/AssistantSettingsTest.java @@ -66,6 +66,13 @@ import com.kingsrook.qqq.backend.core.model.data.QRecord; import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData; import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.StreamedETLWithFrontendProcess; import com.kingsrook.qqq.backend.core.utils.JsonUtils; +import com.kingsrook.qbits.assistant.model.AssistantModel; +import com.kingsrook.qbits.assistant.persistence.AssistantStore; +import com.kingsrook.qqq.backend.core.actions.values.SearchPossibleValueSourceAction; +import com.kingsrook.qqq.backend.core.actions.values.QCustomPossibleValueProvider; +import com.kingsrook.qqq.backend.core.model.actions.values.SearchPossibleValueSourceInput; +import com.kingsrook.qqq.backend.core.actions.customizers.QCodeLoader; +import com.kingsrook.qqq.backend.core.model.metadata.possiblevalues.QPossibleValue; import java.io.Serializable; import java.math.BigDecimal; import java.util.HashMap; @@ -222,8 +229,9 @@ class AssistantSettingsTest extends BaseAssistantTest ** Tables, nav, actions, child lists and exposure. *******************************************************************************/ @Test - void testMetaData() + void testMetaData() throws Exception { + asAdmin(); for(String name : List.of(AssistantProvider.TABLE_NAME, AssistantModelPrice.TABLE_NAME, AssistantUsage.TABLE_NAME)) { QTableMetaData table = qInstance.getTable(name); @@ -237,7 +245,7 @@ class AssistantSettingsTest extends BaseAssistantTest assertThat(qInstance.getProcess(TestConnectionProcess.NAME).getLabel()).isEqualTo("Test connection"); assertThat(qInstance.getWidget(AssistantHome.WIDGET_NAME)).isNotNull(); assertThat(qInstance.getPossibleValueSource(AssistantRecordScreens.PROVIDER_TYPE_PVS_NAME).getEnumValues()).extracting(v -> String.valueOf(v.getId())).containsExactly(FAKE_TYPE); - assertThat(qInstance.getPossibleValueSource(AssistantRecordScreens.MODEL_PVS_NAME).getEnumValues()).extracting(v -> String.valueOf(v.getId())).contains("fake-model-1", "claude-sonnet-5-5"); + assertThat(new SearchPossibleValueSourceAction().execute(new SearchPossibleValueSourceInput().withPossibleValueSourceName(AssistantRecordScreens.MODEL_PVS_NAME)).getResults()).extracting(v -> String.valueOf(v.getId())).contains("fake-model-1", "claude-sonnet-5-5"); assertThat(qInstance.getTable(AssistantProvider.TABLE_NAME).getSections()).anyMatch(s -> s.getWidgetName() != null); assertThat(qInstance.getTable(AssistantProvider.TABLE_NAME).getSections()).anyMatch(s -> s.getFieldNames() != null && s.getFieldNames().contains("apiKeyStatus")); } @@ -570,4 +578,112 @@ class AssistantSettingsTest extends BaseAssistantTest return (((RawHTML) output.getWidgetData()).getHtml()); } + + + /******************************************************************************* + ** Provider models are scoped, with only this row's enabled chat models added. + *******************************************************************************/ + @Test + void testContingentProviderModels() throws Exception + { + asAdmin(); + Integer id = providerId(); + AssistantStore store = new AssistantStore(); + store.insert(new AssistantModel().withProviderId(id).withModelId("registered-chat").withEnabled(true).withModelKind("CHAT")); + store.insert(new AssistantModel().withProviderId(id).withModelId("disabled-chat").withEnabled(false).withModelKind("CHAT")); + store.insert(new AssistantModel().withProviderId(id).withModelId("embedding").withEnabled(true).withModelKind("EMBEDDING")); + Integer otherId = insertProvider(Map.of("name", "Other")).getRecords().get(0).getValueInteger("id"); + store.insert(new AssistantModel().withProviderId(otherId).withModelId("other-chat").withEnabled(true).withModelKind("CHAT")); + assertThat(settingsValues(AssistantRecordScreens.MODEL_PVS_NAME, Map.of("providerType", FAKE_TYPE, "id", id.toString()))) + .extracting(value -> String.valueOf(value.getId())).containsExactly("fake-model-1", "fake-model-2", "registered-chat"); + assertThat(settingsValues(AssistantRecordScreens.MODEL_PVS_NAME, Map.of("providerType", FAKE_TYPE))) + .extracting(value -> String.valueOf(value.getId())).containsExactly("fake-model-1", "fake-model-2"); + assertThat(new SearchPossibleValueSourceAction().execute(new SearchPossibleValueSourceInput() + .withPossibleValueSourceName(AssistantRecordScreens.MODEL_PVS_NAME).withOtherValues(Map.of("providerType", FAKE_TYPE)).withSearchTerm("fake-model-2")).getResults()) + .extracting(value -> String.valueOf(value.getId())).containsExactly("fake-model-2"); + } + + + + /******************************************************************************* + ** Kinds use the selected type and keep their existing labels. + *******************************************************************************/ + @Test + void testContingentProviderKinds() throws Exception + { + assertThat(settingsValues("assistant_provider_kind", Map.of("providerType", FAKE_TYPE))) + .extracting(value -> String.valueOf(value.getId())).containsExactly("FAKE"); + assertThat(settingsValues("assistant_provider_kind", Map.of())).anySatisfy(value -> + { + assertThat(value.getId()).isEqualTo("AZURE_OPENAI"); + assertThat(value.getLabel()).isEqualTo("Azure OpenAI"); + }); + } + + + + /******************************************************************************* + ** A price id must not be interpreted as a provider id. + *******************************************************************************/ + @Test + void testContingentPriceModels() throws Exception + { + asAdmin(); + Integer id = providerId(); + new AssistantStore().insert(new AssistantModel().withProviderId(id).withModelId("registered-chat").withEnabled(true).withModelKind("CHAT")); + String source = qInstance.getTable(AssistantModelPrice.TABLE_NAME).getField("model").getPossibleValueSourceName(); + assertThat(settingsValues(source, Map.of("providerType", FAKE_TYPE, "id", id))) + .extracting(value -> String.valueOf(value.getId())).containsExactly("fake-model-1", "fake-model-2"); + } + + + + /******************************************************************************* + ** Missing and unknown types retain the original union, including pricing. + *******************************************************************************/ + @Test + void testContingentUnionFallback() throws Exception + { + for(String source : List.of(AssistantRecordScreens.MODEL_PVS_NAME, "assistant_price_model", "assistant_provider_kind")) + { + List> union = settingsValues(source, null); + assertThat(settingsValues(source, Map.of())).usingRecursiveComparison().isEqualTo(union); + assertThat(settingsValues(source, Map.of("providerType", "unknown"))).usingRecursiveComparison().isEqualTo(union); + } + assertThat(settingsValues(AssistantRecordScreens.MODEL_PVS_NAME, null)).extracting(value -> String.valueOf(value.getId())) + .contains("fake-model-1", "claude-sonnet-5-5"); + } + + + + /******************************************************************************* + ** Saved values resolve without a selected type, even disabled registry ids. + *******************************************************************************/ + @Test + void testSavedPossibleValue() throws Exception + { + asAdmin(); + Integer id = providerId(); + new AssistantStore().insert(new AssistantModel().withProviderId(id).withModelId("saved-model").withEnabled(false).withModelKind("CHAT")); + for(String source : List.of(AssistantRecordScreens.MODEL_PVS_NAME, "assistant_price_model", "assistant_provider_kind")) + { + QCustomPossibleValueProvider custom = QCodeLoader.getAdHoc(QCustomPossibleValueProvider.class, qInstance.getPossibleValueSource(source).getCustomCodeReference()); + String saved = source.equals("assistant_provider_kind") ? "OPENAI" : "saved-model"; + assertThat(custom.getPossibleValue(saved).getId()).isEqualTo(saved); + assertThat(custom.getPossibleValue("not-known")).isNull(); + } + } + + + + /******************************************************************************* + ** Exercise the registered source through QQQ's search action. + *******************************************************************************/ + private List> settingsValues(String source, Map otherValues) throws Exception + { + asAdmin(); + return (new SearchPossibleValueSourceAction().execute(new SearchPossibleValueSourceInput() + .withPossibleValueSourceName(source).withOtherValues(otherValues)).getResults()); + } + }