diff --git a/.gitignore b/.gitignore index 72316932b..2e70b5fbb 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,6 @@ # OpenNLP models (downloaded by tools/fetch-opennlp-models.sh) /java/gateway-core/src/main/resources/opennlp/ + +# genMetadata model cache/archives (tools/fetch-gen-metadata-model.sh) +.build/ diff --git a/docs/configuration/remote-resources.md b/docs/configuration/remote-resources.md index bd8c4c8de..000d54d94 100644 --- a/docs/configuration/remote-resources.md +++ b/docs/configuration/remote-resources.md @@ -28,30 +28,39 @@ mounted locally. ## Terraform Configuration -By default, the host modules in this repository (`aws-host` and `gcp-host`) will configure the -`REMOTE_RESOURCE_BUCKET` for you if you set the `enable_remote_resources` variable to `true`. This -automatically wires the **artifacts bucket** (used for deployment bundles) as the remote resource bucket. +Host modules (`aws-host` and `gcp-host`) enable remote resource loading **per connector** via `enable_remote_resources` on each `api_connectors`, `bulk_connectors`, or `webhook_collectors` entry. This wires the **artifacts bucket** (used for deployment bundles) as the remote resource bucket for only those instances. > [!IMPORTANT] > - If you configure an existing bucket (e.g., by providing `artifacts_bucket_name`), the bucket must already exist. > - The Terraform runner (the credentials running the `terraform` command) must have sufficient IAM permissions on that bucket to apply permissions (since it will grant read access to the proxy's service account or Lambda execution role). +### Connector specs and custom connectors + +Prebuilt connectors in `worklytics-connector-specs` may set `enable_remote_resources` or `enable_gen_metadata` per connector (e.g. `msft-copilot` enables gen metadata). Those flags flow through to `aws-host` / `gcp-host` when the connector is enabled. + +For ad-hoc connectors, set the flags on `custom_api_connectors` or `custom_bulk_connectors` in your root module (see `infra/examples-dev/*/variables.tf`). + ### AWS (`aws-host`) ```hcl module "psoxy" { source = "../../modules/aws-host" - # ... existing configuration ... - - # Enable remote resource loading from the artifacts S3 bucket - enable_remote_resources = true + api_connectors = { + "my-api" = { + source_kind = "..." + source_auth_strategy = "..." + target_host = "..." + enable_remote_resources = true # OpenNLP, rules.yaml in bucket, etc. + enable_gen_metadata = false + } + } } ``` -This will: -- Set `REMOTE_RESOURCE_BUCKET` on every Lambda function to the artifacts bucket name -- Grant `s3:GetObject` permission on the configured path prefixes in the bucket to each Lambda's execution role +This will, for that connector only: +- Set `REMOTE_RESOURCE_BUCKET` on the Lambda to the artifacts bucket name +- Grant `s3:GetObject` permission on the configured path prefixes in the bucket to that Lambda's execution role ### GCP (`gcp-host`) @@ -59,22 +68,28 @@ This will: module "psoxy" { source = "../../modules/gcp-host" - # ... existing configuration ... - - # Enable remote resource loading from the artifacts GCS bucket - enable_remote_resources = true + api_connectors = { + "gcal" = { + source_kind = "..." + source_auth_strategy = "..." + target_host = "..." + enable_remote_resources = true + } + } } ``` -This will: -- Set `REMOTE_RESOURCE_BUCKET` on every Cloud Function to the artifacts bucket name -- Grant `roles/storage.objectViewer` on the bucket, scoped to the configured path prefixes using IAM conditions, to each function's service account +This will, for that connector only: +- Set `REMOTE_RESOURCE_BUCKET` on the Cloud Function to the artifacts bucket name +- Grant `roles/storage.objectViewer` on the bucket, scoped to the configured path prefixes using IAM conditions, to that function's service account + +Remote resource paths use `/` as a hierarchy separator within the bucket (e.g. `psoxy-dev-erik/GCAL/rules.yaml` for shared prefix `psoxy-dev-erik/` and connector `gcal`). They are distinct from secret / parameter prefixes, which use a trailing `_` to separate names (e.g. `psoxy-dev-erik_GCAL_SOURCE`). When `INSTANCE_RESOURCE_PATH` / `SHARED_RESOURCE_PATH` are not set, psoxy falls back to the config paths and normalizes trailing `_` to `/` and strips any leading `/`. ## Environment Variables | Variable | Description | Required | |---------------------------|----------------------------------------------------------------------------------------------------|----------| -| `REMOTE_RESOURCE_BUCKET` | Name of the S3/GCS bucket containing remote resources. | No | +| `REMOTE_RESOURCE_BUCKET` | Name of the S3/GCS bucket from which to load remote resources. | Yes | | `INSTANCE_RESOURCE_PATH` | Path prefix for instance-specific resources within the bucket. Defaults to `PATH_TO_INSTANCE_CONFIG`. | No | | `SHARED_RESOURCE_PATH` | Path prefix for shared resources (NLP models, etc.) within the bucket. Defaults to `PATH_TO_SHARED_CONFIG`. | No | @@ -83,12 +98,12 @@ This will: The Terraform modules automatically grant minimal read permissions following the Principle of Least Privilege. Access is limited to the configured resource path prefixes within the bucket: -- **AWS**: `s3:GetObject` only for objects under `{INSTANCE_RESOURCE_PATH}/` and `{SHARED_RESOURCE_PATH}/` -- **GCP**: object read access only for objects under `{INSTANCE_RESOURCE_PATH}/` and `{SHARED_RESOURCE_PATH}/`, enforced with IAM Conditions - -No write, delete, or list permissions are granted. When an object is missing or inaccessible, S3 may return 403 (citing `s3:ListBucket`) rather than 404 if the caller lacks bucket list permission; psoxy treats that as unavailable (non-fatal) and continues its resource lookup chain (e.g. falling back to prebuilt rules). +- **AWS**: `s3:GetObject` only for objects under `{INSTANCE_RESOURCE_PATH}/` and + `{SHARED_RESOURCE_PATH}/` +- **GCP**: object read access only for objects under `{INSTANCE_RESOURCE_PATH}/` and + `{SHARED_RESOURCE_PATH}/`, enforced with IAM Conditions -Remote resource paths use `/` as a hierarchy separator within the bucket (e.g. `psoxy-dev-erik/GCAL/rules.yaml` for shared prefix `psoxy-dev-erik/` and connector `gcal`). They are distinct from secret / parameter prefixes, which use a trailing `_` to separate names (e.g. `psoxy-dev-erik_GCAL_SOURCE`). When `INSTANCE_RESOURCE_PATH` / `SHARED_RESOURCE_PATH` are not set, psoxy falls back to the config paths and normalizes trailing `_` to `/` and strips any leading `/`. +No write, delete, or list permissions are granted. ## Use Cases @@ -98,8 +113,7 @@ if no `RULES` config property (env var, parameter store entry, etc.) is found. ### NLP Models (alpha) OpenNLP model files (`en-sent.bin`, `en-pos-maxent.bin`, `en-chunker.bin`) are **not** bundled in -deployment JARs. If your connector rules use `sentenceMetadata` augments, you must upload these -models to the remote resources bucket yourself (requires `enable_remote_resources = true`). +deployment JARs. If your connector rules use `sentenceMetadata` augments, set `enable_remote_resources = true` on that API connector and upload these models to the remote resources bucket. Place them under `{SHARED_RESOURCE_PATH}/opennlp/` (e.g. `{SHARED_RESOURCE_PATH}/opennlp/en-sent.bin`). `{SHARED_RESOURCE_PATH}` defaults to @@ -130,9 +144,32 @@ aws s3 cp en-sent.bin s3://{REMOTE_RESOURCE_BUCKET}/{SHARED_RESOURCE_PATH}/openn gsutil cp en-sent.bin gs://{REMOTE_RESOURCE_BUCKET}/{SHARED_RESOURCE_PATH}/opennlp/en-sent.bin ``` -### LLM Weights (future) -Smaller language models that fit in memory can be placed in the shared resource path for -on-the-fly inference within the proxy. +### LLM model archives (genMetadata BETA) +Upload a **zip** of a Jlama SafeTensors model directory (must include `config.json`) for the **genMetadata** augment to `{SHARED_RESOURCE_PATH}/llm/` in the bucket. The archive name is derived from `PSOXY_GEN_MODEL` (default `tjake/Llama-3.2-1B-Instruct-JQ4`): slashes become `__`, with a `.zip` suffix — e.g. `llm/tjake__Llama-3.2-1B-Instruct-JQ4.zip`. Set per-connector `enable_gen_metadata = true` on `api_connectors` (which also enables remote resource loading for that connector). See [gen-metadata-augment.md](../development/gen-metadata-augment.md) for HuggingFace ids, cloud backends, and ops detail (cloud inference does not use `llm/` archives). + +**Helper script** (download from HuggingFace, zip, upload): + +```bash +# AWS — PREFIX is your SHARED_RESOURCE_PATH within the bucket (trailing slash optional) +./tools/fetch-gen-metadata-model.sh s3://REMOTE_RESOURCE_BUCKET/PREFIX/ + +# GCP — optional MODEL_ID overrides PSOXY_GEN_MODEL / default +./tools/fetch-gen-metadata-model.sh gs://REMOTE_RESOURCE_BUCKET/PREFIX/ tjake/Llama-3.2-1B-Instruct-JQ4 + +# Use an existing local SafeTensors directory instead of downloading +./tools/fetch-gen-metadata-model.sh --from-dir /path/to/model-dir s3://REMOTE_RESOURCE_BUCKET/PREFIX/ +``` + +**Manual upload:** + +```bash +cd /path/to/model-dir && zip -r ../tjake__Llama-3.2-1B-Instruct-JQ4.zip . +aws s3 cp ../tjake__Llama-3.2-1B-Instruct-JQ4.zip \ + s3://{REMOTE_RESOURCE_BUCKET}/{SHARED_RESOURCE_PATH}/llm/tjake__Llama-3.2-1B-Instruct-JQ4.zip + +gsutil cp ../tjake__Llama-3.2-1B-Instruct-JQ4.zip \ + gs://{REMOTE_RESOURCE_BUCKET}/{SHARED_RESOURCE_PATH}/llm/tjake__Llama-3.2-1B-Instruct-JQ4.zip +``` ## Uploading Resources diff --git a/docs/development/augments.md b/docs/development/augments.md index f55f1f490..c77732d4d 100644 --- a/docs/development/augments.md +++ b/docs/development/augments.md @@ -40,14 +40,16 @@ Inspired by OData's `@`-annotation pattern (where metadata about a property `Foo +{sourceProperty}:{augmentFunction} ``` -When `innerJsonPath` is set and matches nested JSON inside a string field, each inner match is -augmented in place within the parsed embedded JSON (mirroring the legacy `textDigest` transform -with `isJsonEscaped`), then the modified structure is re-serialized as a string: +When `innerJsonPath` is set and matches nested JSON inside a string field, all inner matches are +grouped under a single augment property, keyed by normalized inner path suffix: ```jsonc { "content": "{ ... escaped AdaptiveCard JSON ... }", - "+content:textDigest": "{\"type\":\"AdaptiveCard\",\"version\":\"1.0\",\"body\":[{\"type\":\"TextBlock\",\"text\":\"{\\\"length\\\":2572,\\\"word_count\\\":154}\",\"wrap\":true},...]}" + "+content:textDigest": { + "body[0].text": { "length": 2572, "word_count": 154 }, + "body[1].text": { "length": 980, "word_count": 154 } + } } ``` @@ -58,8 +60,8 @@ with `isJsonEscaped`), then the modified structure is re-serialized as a string: | `:` | Separator. | | `{augmentFunction}` | The name of the augment function (e.g. `textDigest`). | -For embedded JSON, inner matched fields (e.g. `text`) are replaced with serialized augment output -and the whole embedded JSON is re-serialized as a string value. +Inner map keys use `{innerPathSuffix}` — the normalized concrete path within parsed inner JSON +(e.g. `body[0].text`). **Example.** For a response containing `body.content`, the augment property produced by the `textDigest` function would be: @@ -193,7 +195,7 @@ public abstract class Augment { ### `outputSchema` — Output Validation -Each augment rule carries an optional `outputSchema` property of type `JsonSchemaFilter`. This schema is applied as a **predicate** (not a filter) to the value produced by the augment's `compute()` method: +Each augment rule may carry an `outputSchema` property of type `JsonSchemaFilter`. This schema is applied as a **predicate** (not a filter) to the value produced by the augment's `compute()` method. **`genMetadata` requires `outputSchema`.** Validation is implemented in `AugmentProcessor` (v0.6.x BETA); failures emit `X-Psoxy-Warning: augment-output-schema-mismatch`. | Outcome | Action | |---|---| diff --git a/docs/development/gen-metadata-augment.md b/docs/development/gen-metadata-augment.md new file mode 100644 index 000000000..5d1b4f26b --- /dev/null +++ b/docs/development/gen-metadata-augment.md @@ -0,0 +1,126 @@ +# genMetadata Augment (BETA) + +> **Status:** BETA · PoC on branch `s225-gen-metadata-poc` +> **Since:** v0.6.x +> **Relates to:** [augments.md](augments.md), [sentence-metadata-augment.md](sentence-metadata-augment.md), [remote-resources.md](../configuration/remote-resources.md) + +## Overview + +The **genMetadata** augment generates structured JSON metadata alongside a source field using a pluggable generative backend. All backends share a [LangChain4j](https://github.com/langchain4j/langchain4j) `ChatModel` integration inside `psoxy-core`. **Local** inference uses embedded [Jlama](https://github.com/tjake/Jlama) (pure Java). **Bedrock** and **Vertex** cloud backends are planned. + +Output appears as a sibling property: `+{sourceProperty}:genMetadata`. + +## Rule configuration (required) + +| Field | Required | Description | +|-------|----------|-------------| +| `jsonPaths` | yes | Source values to process | +| `prompt` | yes | Task instruction (included in rules SHA) | +| `outputSchema` | yes | JSON Schema predicate; invalid output is suppressed | + +No `model`, `backend`, `quality`, or `maxTokens` in rules. Generation limits and model selection are deployment settings (see env vars below). The `outputSchema` gate ensures only structured, expected fields reach the response; it does not cap generation length — use `PSOXY_GEN_MAX_TOKENS` for that. + +## Deployment configuration (env) + +Read via `ConfigService` / `ProxyConfigProperty`: + +| Variable | Default | Purpose | +|----------|---------|---------| +| `PSOXY_GEN_BACKEND` | `local` | `local` (Jlama), `bedrock` and `vertex` (future) | +| `PSOXY_GEN_MODEL` | `tjake/Llama-3.2-1B-Instruct-JQ4` | Jlama HuggingFace id (`owner/name`) or logical id for a cached local archive | +| `PSOXY_GEN_TIMEOUT_SECONDS` | `15` | Max seconds per genMetadata `chat()` call (and model-load wait) | +| `PSOXY_GEN_MAX_INPUT_CHARS` | `4096` | Truncate source text before prompting | +| `PSOXY_GEN_MAX_TOKENS` | `256` | Max tokens to generate per inference | +| `PSOXY_GEN_META_RETRIES` | `2` | Total inference attempts per augment when output is unparseable or fails `outputSchema` (minimum 1; `2` = one retry) | +| `ENABLE_GEN_METADATA` | unset | Set to `true` when that connector has Terraform `enable_gen_metadata = true` | + +When `enable_gen_metadata` is set on an API connector, Terraform **appends** Jlama JVM flags to any existing `JAVA_TOOL_OPTIONS` from `general_environment_variables` or that connector's `environment_variables` (`--add-modules=jdk.incubator.vector --enable-preview --enable-native-access=ALL-UNNAMED`). + +## Infrastructure + +Set **`enable_gen_metadata = true`** on individual `api_connectors` entries (API connectors only). The `msft-copilot` entry in `worklytics-connector-specs` sets this by default when that connector is enabled; custom API connectors can set it in `custom_api_connectors`: + +- Sets `ENABLE_GEN_METADATA=true` on that function +- Sets `JAVA_TOOL_OPTIONS` for Jlama on that function's JVM +- Floors memory at **4096 MB** on that connector unless a higher `memory_size_mb` / `available_memory_mb` is set +- Enables remote resource loading for that connector (for `llm/*.zip` model archives) + +You do not need a separate runtime failure if memory is too low — enable the flag on connectors whose rules use `genMetadata`. Without the flag, augments still run but return `augment-gen-unavailable` if the model cannot load. + +Manual setup (without Terraform flags): set env vars yourself, `memory_size_mb = 4096`, enable remote bucket access, and upload a model archive (see below). + +## Java / LangChain4j + +BETA uses **LangChain4j** with the **Jlama** provider for local embedded inference (`langchain4j` **1.15.0**, `langchain4j-jlama` **1.15.0-beta25**, `jlama-native` in `psoxy-core`). Maven Central does not publish a non-beta `langchain4j-jlama` artifact for the 1.x `ChatModel` API (older **0.36.2** uses a different API). The beta suffix is LangChain4j’s release channel for this integration module, not a separate fork. This avoids JNI bindings to llama.cpp that can take down the JVM on native faults. The same `ChatModel` abstraction will back **Bedrock** (`langchain4j-bedrock`, AWS bundle only) and **Vertex** (`langchain4j-vertex-ai-gemini`, GCP bundle only) when implemented. + +Jlama loads models in **SafeTensors** layout (HuggingFace-style directory with `config.json`), not single-file GGUF. + +### Local model deployment + +1. **HuggingFace id (dev / first run):** set `PSOXY_GEN_MODEL` to a Jlama-compatible id such as `tjake/Llama-3.2-1B-Instruct-JQ4`. Jlama may download weights on first use (requires outbound network from the function). +2. **Production (recommended):** zip a SafeTensors model directory and upload to `{SHARED_RESOURCE_PATH}/llm/{cache-dir-name}.zip`, where `{cache-dir-name}` is `PSOXY_GEN_MODEL` with `/` replaced by `__` (e.g. `tjake__Llama-3.2-1B-Instruct-JQ4.zip` for `tjake/Llama-3.2-1B-Instruct-JQ4`). The runtime extracts the archive into a temp cache before loading. + +**Example `PSOXY_GEN_MODEL` values** (Jlama / HuggingFace; prefer small instruct models on the 4096 MB floor): + +| `PSOXY_GEN_MODEL` | Notes | +|-------------------|-------| +| `tjake/Llama-3.2-1B-Instruct-JQ4` | Default; pre-quantized Jlama build | +| `tjake/Llama-3.2-3B-Instruct-JQ4` | Higher quality; verify memory | +| `tjake/Qwen2.5-1.5B-Instruct-JQ4` | Strong small instruct | +| `tjake/Phi-3.5-mini-instruct-JQ4` | Compact | +| `tjake/gemma-2-2b-it-JQ4` | Google small instruct | + +See [tjake on Hugging Face](https://huggingface.co/tjake) for other pre-quantized `-JQ4` builds. Logical ids without `/` (e.g. `llama-3.2-1b-instruct`) work when the matching `llm/llama-3.2-1b-instruct.zip` archive is present. + +## MS Copilot PoC: prompt classification + +Classifies `body.content` into one of 11 categories (`category` only). See commented example in `docs/sources/microsoft-365/msft-copilot/msft-copilot.yaml` and `MS_COPILOT_GEN_METADATA_AUGMENT` in `PrebuiltSanitizerRules.java`. + +## Local integration tests + +Opt-in Jlama tests live in `LangChain4jGenMetadataBackendIntegrationTest` (token-budget regressions from live Copilot PoC). Run from the `java/` directory: + +```bash +cd java +mvn test -pl core -am -Pgen-metadata-integration +``` + +Requires a Jlama model under `~/.jlama` (downloaded on first run) or set `PSOXY_GEN_MODEL_CACHE`. Optional: `PSOXY_GEN_MODEL` to override the default HuggingFace id. + +## Error handling + +Non-fatal: failures throw `AugmentProcessingException` (caught in `AugmentProcessor`) and surface as `X-Psoxy-Warning` headers on the response. + +| Code | Meaning | +|------|---------| +| `augment-gen-unavailable` | Model missing / not loaded / unsupported backend | +| `augment-gen-inference-failed` | Inference or JSON parse failed | +| `augment-output-schema-mismatch` | Output failed `outputSchema` predicate | +| `augment-conflict-skipped` | Upstream `+` properties present | + +## Roadmap + +### Cloud backends (`PSOXY_GEN_BACKEND=bedrock` | `vertex`) — planned + +Implement via LangChain4j in `psoxy-core`, behind the existing `GenMetadataBackend` SPI: + +| Backend | Module | Bundle | Credentials | +|---------|--------|--------|-------------| +| `bedrock` | `langchain4j-bedrock` | **AWS** Lambda / `psoxy-aws` only | Lambda execution role | +| `vertex` | `langchain4j-vertex-ai-gemini` | **GCP** Cloud Functions / `psoxy-gcp` only | Function service account | + +`PSOXY_GEN_MODEL` becomes the cloud model id. Use native JSON schema (`ResponseFormat`) where supported; no `llm/` remote weights. Wrong backend on the wrong platform should fail clearly. IAM for Bedrock / Vertex AI will ship with per-connector `enable_gen_metadata` extensions. + +### Alternative local engine: java-llama.cpp (`de.kherud`) — conditional + +[java-llama.cpp](https://github.com/kherud/java-llama.cpp) (GGUF via JNI) is **not** in use. Revisit only if **both** are true: + +1. Embedded local models prove valuable enough vs cloud inference, and +2. Jlama performance is a significant bottleneck for production workloads. + +Rationale for deferring: JNI/native faults in llama.cpp can crash the JVM; LangChain4j + Jlama keeps inference in managed Java with a single `ChatModel` surface for local and cloud. + +### Deferred + +- **Rule-level `backend`, `model`, or `quality`** — env-only for BETA. +- **ONNX Runtime GenAI, DJL, Ollama sidecars** — out of scope for in-process serverless. diff --git a/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized/response_beta.json b/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized/response_beta.json index 15b4d0447..730d24def 100644 --- a/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized/response_beta.json +++ b/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized/response_beta.json @@ -31,7 +31,16 @@ { "attachmentId":"4062fb240a03490b98ccd7c86bb2bcbf", "contentType":"application/vnd.microsoft.card.adaptive", - "+content:textDigest":"{\"type\":\"AdaptiveCard\",\"version\":\"1.0\",\"body\":[{\"type\":\"TextBlock\",\"text\":\"{\\\"length\\\":2572,\\\"word_count\\\":154}\",\"wrap\":true},{\"type\":\"TextBlock\",\"id\":\"MessageTextField\",\"text\":\"{\\\"length\\\":980,\\\"word_count\\\":154}\",\"wrap\":true}]}" + "+content:textDigest":{ + "body[0].text":{ + "length":2572, + "word_count":154 + }, + "body[1].text":{ + "length":980, + "word_count":154 + } + } } ], "links":[ diff --git a/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized_no-userIds/response_beta.json b/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized_no-userIds/response_beta.json index 3d199b754..fe26a2a84 100644 --- a/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized_no-userIds/response_beta.json +++ b/docs/sources/microsoft-365/msft-copilot/example-api-responses/sanitized_no-userIds/response_beta.json @@ -31,7 +31,16 @@ { "attachmentId":"p~b7-iqzej17OJkneCXlO1ouZozsG_aj0atv5hHVpKgxM0kzjFc7S6sHiTYXoRAC45pipdQSn_BCn3vJl8dMpaHjBr-dCj3A5A3i2D4QhjLWA", "contentType":"application/vnd.microsoft.card.adaptive", - "+content:textDigest":"{\"type\":\"AdaptiveCard\",\"version\":\"1.0\",\"body\":[{\"type\":\"TextBlock\",\"text\":\"{\\\"length\\\":2572,\\\"word_count\\\":154}\",\"wrap\":true},{\"type\":\"TextBlock\",\"id\":\"MessageTextField\",\"text\":\"{\\\"length\\\":980,\\\"word_count\\\":154}\",\"wrap\":true}]}" + "+content:textDigest":{ + "body[0].text":{ + "length":2572, + "word_count":154 + }, + "body[1].text":{ + "length":980, + "word_count":154 + } + } } ], "links":[ diff --git a/docs/sources/microsoft-365/msft-copilot/msft-copilot.yaml b/docs/sources/microsoft-365/msft-copilot/msft-copilot.yaml index 5921567c9..d1ee8e0f9 100644 --- a/docs/sources/microsoft-365/msft-copilot/msft-copilot.yaml +++ b/docs/sources/microsoft-365/msft-copilot/msft-copilot.yaml @@ -11,6 +11,41 @@ endpoints: - ! jsonPaths: - "$..body.content" + # BETA genMetadata PoC — enable when PSOXY_GEN_* env and 4096MB memory are configured: + # - ! + # jsonPaths: + # - "$..body.content" + # prompt: > + # Return ONLY a JSON object with a single property "category" (no other keys). + # Do NOT return a JSON Schema (never use type, properties, required, or enum as keys). + # Examples: {"category":"Research and Ideation"} or {"category":"Excluded"} + # + # Classify the input text into exactly one category. Use the category name verbatim: + # Email Drafting, General Content Drafting, Email Editing and Refinement, + # General Content Editing and Refinement, Summarization and Synthesis, + # Research and Ideation, Analysis and Data Work, Code Generation and Development, + # Workflow and Task Management, Uncategorized, Excluded. + # Use "Uncategorized" when the prompt is substantive but does not clearly fit a specific category. + # Use "Excluded" for social niceties (greetings, thanks, pleasantries only) or prompts + # that are too short or non-substantive to classify. + # outputSchema: + # type: object + # properties: + # category: + # type: string + # enum: + # - Email Drafting + # - General Content Drafting + # - Email Editing and Refinement + # - General Content Editing and Refinement + # - Summarization and Synthesis + # - Research and Ideation + # - Analysis and Data Work + # - Code Generation and Development + # - Workflow and Task Management + # - Uncategorized + # - Excluded + # required: [category] transforms: - ! jsonPaths: diff --git a/infra/examples-dev/aws/main.tf b/infra/examples-dev/aws/main.tf index ca1a37e6b..7cdc70b51 100644 --- a/infra/examples-dev/aws/main.tf +++ b/infra/examples-dev/aws/main.tf @@ -55,13 +55,15 @@ module "worklytics_connectors" { # the naming convention of `{source-identifier}.tf`, eg `msft-365.tf` # lines below merge results of those files back into single maps of sources locals { - api_connectors = merge( - module.worklytics_connectors.enabled_api_connectors, - module.worklytics_connectors_google_workspace.enabled_api_connectors, - local.msft_api_connectors_with_auth, - var.custom_api_connectors, - {} - ) + api_connectors = { + for k, v in merge( + module.worklytics_connectors.enabled_api_connectors, + module.worklytics_connectors_google_workspace.enabled_api_connectors, + local.msft_api_connectors_with_auth, + var.custom_api_connectors, + {} + ) : k => merge(v, { enable_remote_resources = try(v.enable_remote_resources, true) }) + } custom_bulk_connectors_with_defaults = { for k, v in var.custom_bulk_connectors : k => merge({ @@ -71,10 +73,12 @@ locals { }, v) } - bulk_connectors = merge( - module.worklytics_connectors.enabled_bulk_connectors, - local.custom_bulk_connectors_with_defaults, - ) + bulk_connectors = { + for k, v in merge( + module.worklytics_connectors.enabled_bulk_connectors, + local.custom_bulk_connectors_with_defaults, + ) : k => merge(v, { enable_remote_resources = try(v.enable_remote_resources, true) }) + } source_authorization_todos = concat( module.worklytics_connectors.todos, @@ -158,7 +162,8 @@ module "psoxy" { webhook_collectors = { for k, v in var.webhook_collectors : k => merge( v, { - example_payload = try(file(v.example_payload_file), null) + example_payload = try(file(v.example_payload_file), null) + enable_remote_resources = try(v.enable_remote_resources, true) } ) } provision_bucket_public_access_block = var.provision_bucket_public_access_block @@ -167,8 +172,6 @@ module "psoxy" { custom_side_outputs = var.custom_side_outputs todo_step = local.max_auth_todo_step todos_as_local_files = var.todos_as_local_files - enable_remote_resources = true - # vpc_config = { # vpc_id = aws_default_vpc.default.id @@ -268,6 +271,28 @@ output "todos_3" { value = var.todos_as_outputs ? join("\n", values(module.connection_in_worklytics)[*].todo) : null } +output "todos_4" { + description = "Remote resource uploads (OpenNLP, genMetadata LLM) after deploy, in markdown format." + value = var.todos_as_outputs ? join("\n\n", compact([ + module.psoxy.remote_resource_opennlp_todo, + module.psoxy.remote_resource_gen_metadata_todo, + ])) : null +} + +resource "local_file" "todo_4_upload_opennlp_models" { + count = var.todos_as_local_files && module.psoxy.remote_resource_opennlp_todo != null ? 1 : 0 + + filename = "TODO 4 - upload OpenNLP models.md" + content = module.psoxy.remote_resource_opennlp_todo +} + +resource "local_file" "todo_4_upload_gen_metadata_model" { + count = var.todos_as_local_files && module.psoxy.remote_resource_gen_metadata_todo != null ? 1 : 0 + + filename = "TODO 4 - upload genMetadata LLM model.md" + content = module.psoxy.remote_resource_gen_metadata_todo +} + # although should be sensitive such that Terraform won't echo it to command line or expose it, leave # commented out in example until needed # if you uncomment it, you will then be able to obtain the value through `terraform output --raw pseudonym_salt` diff --git a/infra/examples-dev/aws/variables.tf b/infra/examples-dev/aws/variables.tf index d46673fe4..8f680c200 100644 --- a/infra/examples-dev/aws/variables.tf +++ b/infra/examples-dev/aws/variables.tf @@ -259,6 +259,9 @@ variable "custom_api_connectors" { oauth_scopes_needed = optional(list(string), []) environment_variables = optional(map(string), {}) enable_async_processing = optional(bool, false) + enable_remote_resources = optional(bool, false) + enable_gen_metadata = optional(bool, false) + memory_size_mb = optional(number) example_api_calls = optional(list(string), []) example_api_requests = optional(list(object({ method = optional(string, "GET") @@ -285,10 +288,12 @@ variable "custom_api_connectors" { description = "map of custom API connectors to provision" default = { # "custom-api" = { - # source_kind = "my-custom-api" - # source_auth_strategy = "bearer" - # target_host = "api.example.com" - # example_api_calls = ["/v1/users"] + # source_kind = "my-custom-api" + # source_auth_strategy = "bearer" + # target_host = "api.example.com" + # enable_remote_resources = true # if rules use sentenceMetadata or remote rules.yaml + # enable_gen_metadata = false + # example_api_calls = ["/v1/users"] # secured_variables = [ # { name = "API_KEY" } # ] @@ -340,10 +345,11 @@ variable "custom_bulk_connectors" { transforms = optional(list(map(string)), []) })), {}) })) - memory_size_mb = optional(number, null) - settings_to_provide = optional(map(string), {}) - example_file = optional(string) - example_files = optional(list(string), []) + memory_size_mb = optional(number, null) + enable_remote_resources = optional(bool, false) + settings_to_provide = optional(map(string), {}) + example_file = optional(string) + example_files = optional(list(string), []) })) description = "specs of custom bulk connectors to create" @@ -425,12 +431,13 @@ variable "webhook_collectors" { rotation_days = optional(number, null) # null means no rotation; if > 0, will rotate every N days key_spec = optional(string, "RSA_2048") # RSA_2048, RSA_3072, or RSA_4096; defaults to RSA_2048, which should be sufficient this use-case }), null) - auth_public_keys = optional(list(string), []) # list of public keys to use for verifying webhook signatures; if empty AND no auth keys provision, no app-level auth will be done - allow_origins = optional(list(string), ["*"]) # list of origins to allow for CORS, eg 'https://my-app.com'; if you want to allow all origins, use ['*'] (the default) - output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket (e.g., 'events_', 'webhooks/') - keep_warm_instances = optional(number, null) # if set to 1+, keeps that many Lambda instances warm to eliminate cold starts; adds cost (~$11/month per instance) but improves reliability - example_payload_file = optional(string, null) # path to example payload file to use for testing; if provided, will be used in the test script - example_identity = optional(string, null) # example identity to use for testing; if provided, will be used to test the collector + auth_public_keys = optional(list(string), []) # list of public keys to use for verifying webhook signatures; if empty AND no auth keys provision, no app-level auth will be done + allow_origins = optional(list(string), ["*"]) # list of origins to allow for CORS, eg 'https://my-app.com'; if you want to allow all origins, use ['*'] (the default) + output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket (e.g., 'events_', 'webhooks/') + keep_warm_instances = optional(number, null) # if set to 1+, keeps that many Lambda instances warm to eliminate cold starts; adds cost (~$11/month per instance) but improves reliability + example_payload_file = optional(string, null) # path to example payload file to use for testing; if provided, will be used in the test script + example_identity = optional(string, null) # example identity to use for testing; if provided, will be used to test the collector + enable_remote_resources = optional(bool, false) })) default = {} diff --git a/infra/examples-dev/gcp/main.tf b/infra/examples-dev/gcp/main.tf index d3c0b2db5..cba564efd 100644 --- a/infra/examples-dev/gcp/main.tf +++ b/infra/examples-dev/gcp/main.tf @@ -64,13 +64,15 @@ module "worklytics_connectors" { # the naming convention of `{source-identifier}.tf`, eg `msft-365.tf` # lines below merge results of those files back into single maps of sources locals { - api_connectors = merge( - module.worklytics_connectors.enabled_api_connectors, - module.worklytics_connectors_google_workspace.enabled_api_connectors, - local.msft_api_connectors_with_auth, - var.custom_api_connectors, - {} - ) + api_connectors = { + for k, v in merge( + module.worklytics_connectors.enabled_api_connectors, + module.worklytics_connectors_google_workspace.enabled_api_connectors, + local.msft_api_connectors_with_auth, + var.custom_api_connectors, + {} + ) : k => merge(v, { enable_remote_resources = try(v.enable_remote_resources, true) }) + } custom_bulk_connectors_with_defaults = { for k, v in var.custom_bulk_connectors : k => merge({ @@ -80,10 +82,12 @@ locals { }, v) } - bulk_connectors = merge( - module.worklytics_connectors.enabled_bulk_connectors, - local.custom_bulk_connectors_with_defaults, - ) + bulk_connectors = { + for k, v in merge( + module.worklytics_connectors.enabled_bulk_connectors, + local.custom_bulk_connectors_with_defaults, + ) : k => merge(v, { enable_remote_resources = try(v.enable_remote_resources, true) }) + } source_authorization_todos = concat( @@ -123,7 +127,8 @@ module "psoxy" { webhook_collectors = { for k, v in var.webhook_collectors : k => merge( v, { - example_payload = try(file(v.example_payload_file), null) + example_payload = try(file(v.example_payload_file), null) + enable_remote_resources = try(v.enable_remote_resources, true) } ) } non_production_connectors = var.non_production_connectors @@ -146,7 +151,6 @@ module "psoxy" { tf_gcp_principal_email = var.gcp_terraform_sa_account_email provision_project_level_iam = var.provision_project_level_iam bucket_access_logs_destination = var.bucket_access_logs_destination - enable_remote_resources = true } locals { @@ -253,6 +257,28 @@ output "todos_3" { value = var.todos_as_outputs ? join("\n", values(module.connection_in_worklytics)[*].todo) : null } +output "todos_4" { + description = "Remote resource uploads (OpenNLP, genMetadata LLM) after deploy, in markdown format." + value = var.todos_as_outputs ? join("\n\n", compact([ + module.psoxy.remote_resource_opennlp_todo, + module.psoxy.remote_resource_gen_metadata_todo, + ])) : null +} + +resource "local_file" "todo_4_upload_opennlp_models" { + count = var.todos_as_local_files && module.psoxy.remote_resource_opennlp_todo != null ? 1 : 0 + + filename = "TODO 4 - upload OpenNLP models.md" + content = module.psoxy.remote_resource_opennlp_todo +} + +resource "local_file" "todo_4_upload_gen_metadata_model" { + count = var.todos_as_local_files && module.psoxy.remote_resource_gen_metadata_todo != null ? 1 : 0 + + filename = "TODO 4 - upload genMetadata LLM model.md" + content = module.psoxy.remote_resource_gen_metadata_todo +} + # although should be sensitive such that Terraform won't echo it to command line or expose it, leave # commented out in example until needed # if you uncomment it, you will then be able to obtain the value through `terraform output --raw pseudonym_salt` diff --git a/infra/examples-dev/gcp/variables.tf b/infra/examples-dev/gcp/variables.tf index 99dc9d1ee..11e109b2a 100644 --- a/infra/examples-dev/gcp/variables.tf +++ b/infra/examples-dev/gcp/variables.tf @@ -254,6 +254,9 @@ variable "custom_api_connectors" { oauth_scopes_needed = optional(list(string), []) environment_variables = optional(map(string), {}) enable_async_processing = optional(bool, false) + enable_remote_resources = optional(bool, false) + enable_gen_metadata = optional(bool, false) + available_memory_mb = optional(number) example_api_calls = optional(list(string), []) example_api_requests = optional(list(object({ method = optional(string, "GET") @@ -280,10 +283,12 @@ variable "custom_api_connectors" { description = "map of custom API connectors to provision" default = { # "custom-api" = { - # source_kind = "my-custom-api" - # source_auth_strategy = "bearer" - # target_host = "api.example.com" - # example_api_calls = ["/v1/users"] + # source_kind = "my-custom-api" + # source_auth_strategy = "bearer" + # target_host = "api.example.com" + # enable_remote_resources = true + # enable_gen_metadata = false + # example_api_calls = ["/v1/users"] # secured_variables = [ # { name = "API_KEY" } # ] @@ -308,11 +313,12 @@ variable "webhook_collectors" { rotation_days = optional(number, null) # null means no rotation; if > 0, will rotate every N days key_spec = optional(string, "RSA_SIGN_PKCS1_2048_SHA256") # see https://cloud.google.com/kms/docs/reference/rest/v1/CryptoKeyVersionAlgorithm }), null) - auth_public_keys = optional(list(string), []) # list of public keys to use for verifying webhook signatures; if empty AND no auth keys provision, no app-level auth will be done - allow_origins = optional(list(string), ["*"]) # list of origins to allow for CORS, eg 'https://my-app.com'; if you want to allow all origins, use ['*'] (the default) - output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket (e.g., 'events_', 'webhooks/') - example_payload_file = optional(string, null) # path to example payload file to use for testing; if provided, will be used in the test script - example_identity = optional(string, null) # example identity to use for testing; if provided, will be used to test the collector + auth_public_keys = optional(list(string), []) # list of public keys to use for verifying webhook signatures; if empty AND no auth keys provision, no app-level auth will be done + allow_origins = optional(list(string), ["*"]) # list of origins to allow for CORS, eg 'https://my-app.com'; if you want to allow all origins, use ['*'] (the default) + output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket (e.g., 'events_', 'webhooks/') + example_payload_file = optional(string, null) # path to example payload file to use for testing; if provided, will be used in the test script + example_identity = optional(string, null) # example identity to use for testing; if provided, will be used to test the collector + enable_remote_resources = optional(bool, false) })) default = {} @@ -341,12 +347,13 @@ variable "custom_bulk_connectors" { transforms = optional(list(map(string)), []) }))) })) - available_memory_mb = optional(number) - timeout_seconds = optional(number) - rules_file = optional(string) - settings_to_provide = optional(map(string), {}) - example_file = optional(string) - example_files = optional(list(string), []) + available_memory_mb = optional(number) + timeout_seconds = optional(number) + enable_remote_resources = optional(bool, false) + rules_file = optional(string) + settings_to_provide = optional(map(string), {}) + example_file = optional(string) + example_files = optional(list(string), []) })) description = "specs of custom bulk connectors to create" diff --git a/infra/modules/aws-host/main.tf b/infra/modules/aws-host/main.tf index a9470b59c..cbc7df4cb 100644 --- a/infra/modules/aws-host/main.tf +++ b/infra/modules/aws-host/main.tf @@ -86,6 +86,24 @@ locals { # proxy caller role requires direct lambda access if API Gateway v2 is not used and there are API connectors caller_requires_direct_lambda_access = !local.use_api_gateway_v2 && length(module.api_connector) > 0 + + needs_opennlp_model_upload = length([ + for k, v in merge(var.api_connectors, var.bulk_connectors, var.webhook_collectors) : k + if try(v.enable_remote_resources, false) + ]) > 0 + + needs_gen_metadata_model_upload = length([ + for k, v in var.api_connectors : k if try(v.enable_gen_metadata, false) + ]) > 0 + + gen_metadata_model_id = coalesce( + try(var.general_environment_variables["PSOXY_GEN_MODEL"], null), + "tjake/Llama-3.2-1B-Instruct-JQ4" + ) + gen_metadata_archive_name = "${replace(local.gen_metadata_model_id, "/", "__")}.zip" + + # Appended to JAVA_TOOL_OPTIONS when enable_gen_metadata (Jlama / Vector API). See gen-metadata-augment.md. + jlama_java_tool_options = "--add-modules=jdk.incubator.vector --enable-preview --enable-native-access=ALL-UNNAMED" } module "psoxy" { @@ -241,7 +259,10 @@ module "api_connector" { side_output_original = try(local.custom_original_side_outputs[each.key], null) side_output_sanitized = try(local.sanitized_side_outputs[each.key], null) enable_async_processing = each.value.enable_async_processing - memory_size_mb = each.value.enable_async_processing ? 1024 : 512 # default is 512; double it for async case, to give additional margin + memory_size_mb = max( + try(each.value.enable_gen_metadata, false) ? coalesce(each.value.memory_size_mb, 4096) : (each.value.enable_async_processing ? 1024 : 512), + try(each.value.enable_gen_metadata, false) ? 4096 : 0 + ) todos_as_local_files = var.todos_as_local_files todo_step = var.todo_step @@ -259,11 +280,15 @@ module "api_connector" { }, try(each.value.environment_variables, {}), var.general_environment_variables, + try(each.value.enable_gen_metadata, false) ? { + ENABLE_GEN_METADATA = "true" + JAVA_TOOL_OPTIONS = trimspace("${lookup(merge(try(each.value.environment_variables, {}), var.general_environment_variables), "JAVA_TOOL_OPTIONS", "")} ${local.jlama_java_tool_options}") + } : {}, ) - remote_resource_bucket = var.enable_remote_resources ? module.psoxy.artifacts_bucket_name : null - remote_resource_instance_path = var.enable_remote_resources ? local.connector_instance_resource_path[each.key] : null - remote_resource_shared_path = var.enable_remote_resources ? local.shared_resource_path : null + remote_resource_bucket = (try(each.value.enable_remote_resources, false) || try(each.value.enable_gen_metadata, false)) ? module.psoxy.artifacts_bucket_name : null + remote_resource_instance_path = (try(each.value.enable_remote_resources, false) || try(each.value.enable_gen_metadata, false)) ? local.connector_instance_resource_path[each.key] : null + remote_resource_shared_path = (try(each.value.enable_remote_resources, false) || try(each.value.enable_gen_metadata, false)) ? local.shared_resource_path : null } @@ -349,9 +374,9 @@ module "bulk_connector" { var.general_environment_variables ) - remote_resource_bucket = var.enable_remote_resources ? module.psoxy.artifacts_bucket_name : null - remote_resource_instance_path = var.enable_remote_resources ? local.connector_instance_resource_path[each.key] : null - remote_resource_shared_path = var.enable_remote_resources ? local.shared_resource_path : null + remote_resource_bucket = try(each.value.enable_remote_resources, false) ? module.psoxy.artifacts_bucket_name : null + remote_resource_instance_path = try(each.value.enable_remote_resources, false) ? local.connector_instance_resource_path[each.key] : null + remote_resource_shared_path = try(each.value.enable_remote_resources, false) ? local.shared_resource_path : null } @@ -401,9 +426,9 @@ module "webhook_collectors" { var.general_environment_variables, ) - remote_resource_bucket = var.enable_remote_resources ? module.psoxy.artifacts_bucket_name : null - remote_resource_instance_path = var.enable_remote_resources ? local.connector_instance_resource_path[each.key] : null - remote_resource_shared_path = var.enable_remote_resources ? local.shared_resource_path : null + remote_resource_bucket = try(each.value.enable_remote_resources, false) ? module.psoxy.artifacts_bucket_name : null + remote_resource_instance_path = try(each.value.enable_remote_resources, false) ? local.connector_instance_resource_path[each.key] : null + remote_resource_shared_path = try(each.value.enable_remote_resources, false) ? local.shared_resource_path : null } # Policy to allow test caller to invoke webhook collector urls and sign webhook requests diff --git a/infra/modules/aws-host/remote_resource_todos.tf b/infra/modules/aws-host/remote_resource_todos.tf new file mode 100644 index 000000000..e8d354f48 --- /dev/null +++ b/infra/modules/aws-host/remote_resource_todos.tf @@ -0,0 +1,63 @@ +# Markdown TODOs for uploading remote resource assets (OpenNLP, genMetadata LLM archives). +# Rendered as outputs only; upload is performed outside Terraform via tools/*.sh scripts. + +locals { + remote_resource_s3_prefix = "s3://${module.psoxy.artifacts_bucket_name}/${local.shared_resource_path}" + + opennlp_connector_ids = join(", ", [ + for k, v in merge(var.api_connectors, var.bulk_connectors, var.webhook_collectors) : k + if try(v.enable_remote_resources, false) + ]) + + gen_metadata_connector_ids = join(", ", [ + for k, v in var.api_connectors : k if try(v.enable_gen_metadata, false) + ]) +} + +output "remote_resource_opennlp_todo" { + description = "TODO (markdown) for uploading OpenNLP models when any connector has enable_remote_resources." + value = local.needs_opennlp_model_upload ? trimspace(<<-EOT + ## Upload OpenNLP models (sentenceMetadata augment) + + Connectors with `enable_remote_resources`: ${local.opennlp_connector_ids} + + OpenNLP binaries are not bundled in deployment JARs. From the **psoxy repository root**, download models and upload to the artifacts / remote-resources bucket in one step: + + ```bash + ./tools/fetch-opennlp-models.sh ${local.remote_resource_s3_prefix} + ``` + + That command uses this deployment's artifacts bucket and shared prefix (`${local.shared_resource_path}`). It requires `curl` and the AWS CLI (`aws`) with permission to write objects under that prefix. + + See [remote-resources.md](https://github.com/worklytics/psoxy/blob/main/docs/configuration/remote-resources.md). + EOT + ) : null +} + +output "remote_resource_gen_metadata_todo" { + description = "TODO (markdown) for uploading genMetadata LLM model archive when any API connector has enable_gen_metadata." + value = local.needs_gen_metadata_model_upload ? trimspace(<<-EOT + ## Upload genMetadata LLM model archive (BETA) + + API connectors with `enable_gen_metadata`: ${local.gen_metadata_connector_ids} + + Production deployments should load Jlama weights from the remote resources bucket instead of downloading from HuggingFace at runtime. + + Expected object key: `${local.shared_resource_path}llm/${local.gen_metadata_archive_name}` (from `PSOXY_GEN_MODEL=${local.gen_metadata_model_id}`). + + From the **psoxy repository root**, download the HuggingFace model, zip it, and upload to the artifacts bucket in one step: + + ```bash + ./tools/fetch-gen-metadata-model.sh ${local.remote_resource_s3_prefix} ${local.gen_metadata_model_id} + ``` + + That command uses: + - **Destination:** `${local.remote_resource_s3_prefix}` (this deployment's artifacts bucket + shared resource prefix) + - **Model:** `${local.gen_metadata_model_id}` (override with a second argument or `PSOXY_GEN_MODEL`) + + Requires `zip`, [HuggingFace CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli) (`pip install huggingface_hub`), and the AWS CLI (`aws`) with upload access. To use a local SafeTensors directory instead of downloading: add `--from-dir /path/to/model-dir` before the `s3://` URI. + + See [gen-metadata-augment.md](https://github.com/worklytics/psoxy/blob/main/docs/development/gen-metadata-augment.md). + EOT + ) : null +} diff --git a/infra/modules/aws-host/variables.tf b/infra/modules/aws-host/variables.tf index d9c9e1d5b..36a72afff 100644 --- a/infra/modules/aws-host/variables.tf +++ b/infra/modules/aws-host/variables.tf @@ -197,6 +197,9 @@ variable "api_connectors" { oauth_scopes_needed = optional(list(string), []) environment_variables = optional(map(string), {}) enable_async_processing = optional(bool, false) + enable_remote_resources = optional(bool, false) + enable_gen_metadata = optional(bool, false) + memory_size_mb = optional(number) example_api_calls = optional(list(string), []) example_api_requests = optional(list(object({ method = optional(string, "GET") @@ -289,13 +292,14 @@ variable "bulk_connectors" { transforms = optional(list(map(string)), []) }))) })) - rules_file = optional(string) - rules_raw = optional(string, null) - example_file = optional(string) - example_files = optional(list(string), []) - instructions_template = optional(string) - memory_size_mb = optional(number) - settings_to_provide = optional(map(string), {}) + rules_file = optional(string) + rules_raw = optional(string, null) + example_file = optional(string) + example_files = optional(list(string), []) + instructions_template = optional(string) + memory_size_mb = optional(number) + settings_to_provide = optional(map(string), {}) + enable_remote_resources = optional(bool, false) })) description = "map of connector id => bulk connectors to provision" @@ -373,12 +377,13 @@ variable "webhook_collectors" { rotation_days = optional(number, null) # null means no rotation; if > 0, will rotate every N days key_spec = optional(string, "RSA_2048") # RSA_2048, RSA_3072, or RSA_4096; defaults to RSA_2048, which should be sufficient this use-case }), null) - auth_public_keys = optional(list(string), []) # list of public keys to use for verifying webhook signatures; if empty AND no auth keys provision, no app-level auth will be done - allow_origins = optional(list(string), ["*"]) # list of origins to allow for CORS, eg 'https://my-app.com'; if you want to allow all origins, use ['*'] (the default) - output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket - keep_warm_instances = optional(number, null) # if set, keeps N Lambda instances warm to eliminate cold starts; adds cost (~$11/month per instance) but improves reliability - example_identity = optional(string, null) # example identity to use in test payloads - example_payload = optional(string, null) # example payload content to use in test scripts + auth_public_keys = optional(list(string), []) # list of public keys to use for verifying webhook signatures; if empty AND no auth keys provision, no app-level auth will be done + allow_origins = optional(list(string), ["*"]) # list of origins to allow for CORS, eg 'https://my-app.com'; if you want to allow all origins, use ['*'] (the default) + output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket + keep_warm_instances = optional(number, null) # if set, keeps N Lambda instances warm to eliminate cold starts; adds cost (~$11/month per instance) but improves reliability + example_identity = optional(string, null) # example identity to use in test payloads + example_payload = optional(string, null) # example payload content to use in test scripts + enable_remote_resources = optional(bool, false) })) default = {} @@ -473,8 +478,3 @@ variable "artifacts_bucket_name" { } -variable "enable_remote_resources" { - type = bool - description = "**beta** Whether to enable remote resource loading from the artifacts S3 bucket (rules, NLP models, etc.). When true, sets REMOTE_RESOURCE_BUCKET env var and grants s3:GetObject to each Lambda. Default will change to `true` in next major version." - default = false # will change to true in 0.7.x -} diff --git a/infra/modules/gcp-host/main.tf b/infra/modules/gcp-host/main.tf index 78e83ce13..c8f04607b 100644 --- a/infra/modules/gcp-host/main.tf +++ b/infra/modules/gcp-host/main.tf @@ -74,6 +74,24 @@ locals { try(local.api_connector_rules_raw[k], null) != null ? local.api_connector_rules_raw[k] : null ) } + + needs_opennlp_model_upload = length([ + for k, v in merge(var.api_connectors, var.bulk_connectors, var.webhook_collectors) : k + if try(v.enable_remote_resources, false) + ]) > 0 + + needs_gen_metadata_model_upload = length([ + for k, v in var.api_connectors : k if try(v.enable_gen_metadata, false) + ]) > 0 + + gen_metadata_model_id = coalesce( + try(var.general_environment_variables["PSOXY_GEN_MODEL"], null), + "tjake/Llama-3.2-1B-Instruct-JQ4" + ) + gen_metadata_archive_name = "${replace(local.gen_metadata_model_id, "/", "__")}.zip" + + # Appended to JAVA_TOOL_OPTIONS when enable_gen_metadata (Jlama / Vector API). See gen-metadata-augment.md. + jlama_java_tool_options = "--add-modules=jdk.incubator.vector --enable-preview --enable-native-access=ALL-UNNAMED" } # TODO: probably pull all the way to the top level bc 1) proper tf style, 2) simplifies customization if it doesn't work for a particular environment @@ -264,7 +282,11 @@ module "api_connector" { allowed_data_access_ip_blocks = var.allowed_data_access_ip_blocks instance_concurrency = var.api_connector_instance_concurrency max_instance_count = var.max_instances_per_api_connector - timeout_seconds = coalesce(try(each.value.timeout_seconds, null), 180) + available_memory_mb = max( + try(each.value.enable_gen_metadata, false) ? coalesce(each.value.available_memory_mb, 4096) : 1024, + try(each.value.enable_gen_metadata, false) ? 4096 : 0 + ) + timeout_seconds = coalesce(try(each.value.timeout_seconds, null), 180) environment_variables = merge( @@ -280,11 +302,15 @@ module "api_connector" { }, try(each.value.environment_variables, {}), var.general_environment_variables, + try(each.value.enable_gen_metadata, false) ? { + ENABLE_GEN_METADATA = "true" + JAVA_TOOL_OPTIONS = trimspace("${lookup(merge(try(each.value.environment_variables, {}), var.general_environment_variables), "JAVA_TOOL_OPTIONS", "")} ${local.jlama_java_tool_options}") + } : {}, ) - remote_resource_bucket = var.enable_remote_resources ? module.psoxy.artifacts_bucket_name : null - remote_resource_instance_path = var.enable_remote_resources ? local.connector_instance_resource_path[each.key] : null - remote_resource_shared_path = var.enable_remote_resources ? local.shared_resource_path : null + remote_resource_bucket = (try(each.value.enable_remote_resources, false) || try(each.value.enable_gen_metadata, false)) ? module.psoxy.artifacts_bucket_name : null + remote_resource_instance_path = (try(each.value.enable_remote_resources, false) || try(each.value.enable_gen_metadata, false)) ? local.connector_instance_resource_path[each.key] : null + remote_resource_shared_path = (try(each.value.enable_remote_resources, false) || try(each.value.enable_gen_metadata, false)) ? local.shared_resource_path : null secret_bindings = merge( local.secrets_bound_as_env_vars[each.key], @@ -378,9 +404,9 @@ module "webhook_collector" { var.general_environment_variables, ) - remote_resource_bucket = var.enable_remote_resources ? module.psoxy.artifacts_bucket_name : null - remote_resource_instance_path = var.enable_remote_resources ? local.connector_instance_resource_path[each.key] : null - remote_resource_shared_path = var.enable_remote_resources ? local.shared_resource_path : null + remote_resource_bucket = try(each.value.enable_remote_resources, false) ? module.psoxy.artifacts_bucket_name : null + remote_resource_instance_path = try(each.value.enable_remote_resources, false) ? local.connector_instance_resource_path[each.key] : null + remote_resource_shared_path = try(each.value.enable_remote_resources, false) ? local.shared_resource_path : null secret_bindings = module.psoxy.secrets @@ -446,9 +472,9 @@ module "bulk_connector" { var.general_environment_variables, ) - remote_resource_bucket = var.enable_remote_resources ? module.psoxy.artifacts_bucket_name : null - remote_resource_instance_path = var.enable_remote_resources ? local.connector_instance_resource_path[each.key] : null - remote_resource_shared_path = var.enable_remote_resources ? local.shared_resource_path : null + remote_resource_bucket = try(each.value.enable_remote_resources, false) ? module.psoxy.artifacts_bucket_name : null + remote_resource_instance_path = try(each.value.enable_remote_resources, false) ? local.connector_instance_resource_path[each.key] : null + remote_resource_shared_path = try(each.value.enable_remote_resources, false) ? local.shared_resource_path : null depends_on = [ module.psoxy # some of the set-up IAM grants done there, but not EXPLICITLY passed out as outputs and into above as inputs, are required; so make this explicit diff --git a/infra/modules/gcp-host/output.tf b/infra/modules/gcp-host/output.tf index 26b30e5cd..6b78ea838 100644 --- a/infra/modules/gcp-host/output.tf +++ b/infra/modules/gcp-host/output.tf @@ -19,6 +19,11 @@ output "artifacts_bucket_id" { value = module.psoxy.artifacts_bucket_id } +output "artifacts_bucket_name" { + description = "Name of the artifacts bucket (remote resources / deployment bundle)." + value = module.psoxy.artifacts_bucket_name +} + output "api_connector_instances" { value = local.api_instances } diff --git a/infra/modules/gcp-host/remote_resource_todos.tf b/infra/modules/gcp-host/remote_resource_todos.tf new file mode 100644 index 000000000..d5ed41459 --- /dev/null +++ b/infra/modules/gcp-host/remote_resource_todos.tf @@ -0,0 +1,63 @@ +# Markdown TODOs for uploading remote resource assets (OpenNLP, genMetadata LLM archives). +# Rendered as outputs only; upload is performed outside Terraform via tools/*.sh scripts. + +locals { + remote_resource_gcs_prefix = "gs://${module.psoxy.artifacts_bucket_name}/${local.shared_resource_path}" + + opennlp_connector_ids = join(", ", [ + for k, v in merge(var.api_connectors, var.bulk_connectors, var.webhook_collectors) : k + if try(v.enable_remote_resources, false) + ]) + + gen_metadata_connector_ids = join(", ", [ + for k, v in var.api_connectors : k if try(v.enable_gen_metadata, false) + ]) +} + +output "remote_resource_opennlp_todo" { + description = "TODO (markdown) for uploading OpenNLP models when any connector has enable_remote_resources." + value = local.needs_opennlp_model_upload ? trimspace(<<-EOT + ## Upload OpenNLP models (sentenceMetadata augment) + + Connectors with `enable_remote_resources`: ${local.opennlp_connector_ids} + + OpenNLP binaries are not bundled in deployment JARs. From the **psoxy repository root**, download models and upload to the artifacts / remote-resources bucket in one step: + + ```bash + ./tools/fetch-opennlp-models.sh ${local.remote_resource_gcs_prefix} + ``` + + That command uses this deployment's artifacts bucket and shared prefix (`${local.shared_resource_path}`). It requires `curl` and `gsutil` with permission to write objects under that prefix. + + See [remote-resources.md](https://github.com/worklytics/psoxy/blob/main/docs/configuration/remote-resources.md). + EOT + ) : null +} + +output "remote_resource_gen_metadata_todo" { + description = "TODO (markdown) for uploading genMetadata LLM model archive when any API connector has enable_gen_metadata." + value = local.needs_gen_metadata_model_upload ? trimspace(<<-EOT + ## Upload genMetadata LLM model archive (BETA) + + API connectors with `enable_gen_metadata`: ${local.gen_metadata_connector_ids} + + Production deployments should load Jlama weights from the remote resources bucket instead of downloading from HuggingFace at runtime. + + Expected object key: `${local.shared_resource_path}llm/${local.gen_metadata_archive_name}` (from `PSOXY_GEN_MODEL=${local.gen_metadata_model_id}`). + + From the **psoxy repository root**, download the HuggingFace model, zip it, and upload to the artifacts bucket in one step: + + ```bash + ./tools/fetch-gen-metadata-model.sh ${local.remote_resource_gcs_prefix} ${local.gen_metadata_model_id} + ``` + + That command uses: + - **Destination:** `${local.remote_resource_gcs_prefix}` (this deployment's artifacts bucket + shared resource prefix) + - **Model:** `${local.gen_metadata_model_id}` (override with a second argument or `PSOXY_GEN_MODEL`) + + Requires `zip`, [HuggingFace CLI](https://huggingface.co/docs/huggingface_hub/en/guides/cli) (`pip install huggingface_hub`), and `gsutil` with upload access. To use a local SafeTensors directory instead of downloading: add `--from-dir /path/to/model-dir` before the `gs://` URI. + + See [gen-metadata-augment.md](https://github.com/worklytics/psoxy/blob/main/docs/development/gen-metadata-augment.md). + EOT + ) : null +} diff --git a/infra/modules/gcp-host/variables.tf b/infra/modules/gcp-host/variables.tf index a5a65cff0..ec9cfced6 100644 --- a/infra/modules/gcp-host/variables.tf +++ b/infra/modules/gcp-host/variables.tf @@ -193,6 +193,9 @@ variable "api_connectors" { oauth_scopes_needed = optional(list(string), []) environment_variables = optional(map(string), {}) enable_async_processing = optional(bool, false) + enable_remote_resources = optional(bool, false) + enable_gen_metadata = optional(bool, false) + available_memory_mb = optional(number) example_api_calls = optional(list(string), []) example_api_requests = optional(list(object({ method = optional(string, "GET") @@ -240,8 +243,9 @@ variable "webhook_collectors" { batch_processing_frequency_minutes = optional(number, 5) # frequency (in minutes) at which to batch process webhooks output_path_prefix = optional(string, "") # optional path prefix to prepend to webhook output files in bucket - example_identity = optional(string, null) # example identity to use in test payloads - example_payload = optional(string, null) # example payload content to use in test scripts + example_identity = optional(string, null) # example identity to use in test payloads + example_payload = optional(string, null) # example payload content to use in test scripts + enable_remote_resources = optional(bool, false) })) default = {} @@ -266,14 +270,15 @@ variable "bulk_connectors" { transforms = optional(list(map(string)), []) }))) })) - rules_file = optional(string) - rules_raw = optional(string, null) - example_file = optional(string) - example_files = optional(list(string), []) - instructions_template = optional(string) - settings_to_provide = optional(map(string), {}) - available_memory_mb = optional(number) - timeout_seconds = optional(number) + rules_file = optional(string) + rules_raw = optional(string, null) + example_file = optional(string) + example_files = optional(list(string), []) + instructions_template = optional(string) + settings_to_provide = optional(map(string), {}) + available_memory_mb = optional(number) + timeout_seconds = optional(number) + enable_remote_resources = optional(bool, false) })) description = "map of connector id => bulk connectors to provision" @@ -461,8 +466,3 @@ variable "max_instances_per_api_connector" { } } -variable "enable_remote_resources" { - type = bool - description = "**beta** Whether to enable remote resource loading from the artifacts GCS bucket (rules, NLP models, etc.). When true, sets REMOTE_RESOURCE_BUCKET env var and grants roles/storage.objectViewer to each Cloud Function. Default will change to `true` in next major version." - default = false # will change to true in 0.7.x -} diff --git a/infra/modules/worklytics-connector-specs/main.tf b/infra/modules/worklytics-connector-specs/main.tf index 33e69a107..9a58ae829 100644 --- a/infra/modules/worklytics-connector-specs/main.tf +++ b/infra/modules/worklytics-connector-specs/main.tf @@ -1512,6 +1512,10 @@ locals { google_workspace_sources_backwards = { for k, v in local.google_workspace_sources : k => merge(v, { example_calls : try(v.example_api_calls, []) }, + { + enable_remote_resources = try(v.enable_remote_resources, false) + enable_gen_metadata = try(v.enable_gen_metadata, false) + }, try(local._resolve_rules_raw[k], null) != null ? { rules_raw : local._resolve_rules_raw[k], rules_file : null } : {} ) } @@ -1520,6 +1524,10 @@ locals { msft_365_connectors_backwards = { for k, v in local.msft_365_connectors : k => merge(v, { example_calls : try(v.example_api_calls, []) }, + { + enable_remote_resources = try(v.enable_remote_resources, false) + enable_gen_metadata = try(v.enable_gen_metadata, false) + }, try(local._resolve_rules_raw[k], null) != null ? { rules_raw : local._resolve_rules_raw[k], rules_file : null } : {} ) } @@ -1533,12 +1541,17 @@ locals { oauth_long_access_connectors_backwards_with_rules_raw = { for k, v in local.oauth_long_access_connectors : k => merge(v, { example_calls : try(v.example_api_calls, []) }, + { + enable_remote_resources = try(v.enable_remote_resources, false) + enable_gen_metadata = try(v.enable_gen_metadata, false) + }, try(local._resolve_rules_raw[k], null) != null ? { rules_raw : local._resolve_rules_raw[k], rules_file : null } : {} ) } bulk_connectors_with_rules_raw = { for k, v in local.bulk_connectors : k => merge(v, + { enable_remote_resources = try(v.enable_remote_resources, false) }, try(local._resolve_rules_raw[k], null) != null ? { rules_raw : local._resolve_rules_raw[k], rules_file : null } : {} ) } diff --git a/infra/modules/worklytics-connector-specs/msft-365.tf b/infra/modules/worklytics-connector-specs/msft-365.tf index e827f5218..b2ee67653 100644 --- a/infra/modules/worklytics-connector-specs/msft-365.tf +++ b/infra/modules/worklytics-connector-specs/msft-365.tf @@ -183,6 +183,7 @@ EOT environment_variables : local.msft_365_environment_variables external_todo : null enable_side_output : false + enable_gen_metadata : true example_api_calls : [ "/v1.0/users", "/beta/copilot/users/${local.example_msft_user_guid}/interactionHistory/getAllEnterpriseInteractions" diff --git a/java/core/pom.xml b/java/core/pom.xml index 542218e99..8ed3c549c 100644 --- a/java/core/pom.xml +++ b/java/core/pom.xml @@ -184,6 +184,24 @@ 10.5 + + + dev.langchain4j + langchain4j + ${dependency.langchain4j.version} + + + dev.langchain4j + langchain4j-jlama + ${dependency.langchain4j-jlama.version} + + + com.github.tjake + jlama-native + 0.8.4 + ${os.detected.classifier} + + org.mockito mockito-junit-jupiter @@ -222,6 +240,13 @@ + + + kr.motd.maven + os-maven-plugin + 1.7.1 + + org.apache.maven.plugins @@ -230,7 +255,7 @@ - -javaagent:"${settings.localRepository}/org/mockito/mockito-core/${dependency.mockito-junit-jupiter.version}/mockito-core-${dependency.mockito-junit-jupiter.version}.jar" + -javaagent:"${settings.localRepository}/org/mockito/mockito-core/${dependency.mockito-junit-jupiter.version}/mockito-core-${dependency.mockito-junit-jupiter.version}.jar" --add-modules=jdk.incubator.vector --enable-preview --enable-native-access=ALL-UNNAMED 2 @@ -269,4 +294,26 @@ + + + gen-metadata-integration + + + + org.apache.maven.plugins + maven-surefire-plugin + + + **/LangChain4jGenMetadataBackendIntegrationTest.java + + + true + + + + + + + + diff --git a/java/core/src/main/java/co/worklytics/psoxy/FunctionRuntimeModule.java b/java/core/src/main/java/co/worklytics/psoxy/FunctionRuntimeModule.java index 5e073bde4..c87c07b33 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/FunctionRuntimeModule.java +++ b/java/core/src/main/java/co/worklytics/psoxy/FunctionRuntimeModule.java @@ -9,9 +9,8 @@ import javax.inject.Singleton; import javax.net.ssl.SSLContext; import javax.net.ssl.SSLSocketFactory; -import com.avaulta.gateway.resources.ResourceService; import com.avaulta.gateway.rules.WebhookCollectionRules; -import com.avaulta.gateway.rules.augments.SentenceMetadataProcessor; +import com.avaulta.gateway.rules.augments.AugmentValidation; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.api.client.http.HttpRequestFactory; import com.google.api.client.http.javanet.NetHttpTransport; @@ -26,9 +25,7 @@ import co.worklytics.psoxy.gateway.auth.Base64KeyClient; import co.worklytics.psoxy.gateway.auth.PublicKeyStoreClient; import co.worklytics.psoxy.gateway.impl.CompositeConfigService; -import co.worklytics.psoxy.gateway.impl.CompositeResourceService; import co.worklytics.psoxy.gateway.impl.EnvVarsConfigService; -import co.worklytics.psoxy.gateway.impl.LocalFileResourceService; import co.worklytics.psoxy.gateway.impl.WebhookSanitizer; import co.worklytics.psoxy.gateway.impl.output.NoOutput; import co.worklytics.psoxy.gateway.impl.output.OutputUtils; @@ -52,6 +49,7 @@ @Module( includes = { FunctionRuntimeModule.Bindings.class, + ResourceServiceBindingsModule.class, } ) public class FunctionRuntimeModule { @@ -125,43 +123,6 @@ static ConfigService configService(EnvVarsConfigService envVarsConfigService, .build(); } - /** - * Provides the instance-scoped ResourceService, composing local FS with the - * platform-specific remote ResourceService (S3/GCS), using failover semantics. - * - *

Failover order: local filesystem ({@link ResourceService#DEFAULT_LOCAL_RESOURCE_PATH}) - * → remote cloud storage → no-op

- * - *

The local FS path is always checked first (no env var needed). It can be populated - * via deployment layers, Lambda layers, init scripts, etc.

- */ - @Provides @Singleton - static ResourceService instanceResourceService(@Named("Remote") ResourceService remoteResourceService) { - // always layer local FS on top of remote — local is a fast path / override - return CompositeResourceService.builder() - .preferred(new LocalFileResourceService(ResourceService.DEFAULT_LOCAL_RESOURCE_PATH)) - .fallback(remoteResourceService) - .build(); - } - - /** - * Resource resolution for OpenNLP models: instance-local (FS / instance bucket) first, - * then shared remote bucket (e.g. org-wide model artifacts). - */ - @Provides @Singleton @Named("OpenNlp") - static ResourceService openNlpResourceService(ResourceService instanceResourceService, - @Named("SharedRemote") ResourceService sharedRemoteResourceService) { - return CompositeResourceService.builder() - .preferred(instanceResourceService) - .fallback(sharedRemoteResourceService) - .build(); - } - - @Provides @Singleton - static SentenceMetadataProcessor sentenceMetadataProcessor(@Named("OpenNlp") ResourceService openNlpResourceService) { - return new SentenceMetadataProcessor(openNlpResourceService); - } - @Provides @Singleton @Named("async") static ApiSanitizedDataOutput apiSanitizedDataOutput(OutputUtils outputUtils) { return outputUtils.asyncOutput(); @@ -201,7 +162,10 @@ static NoOutput noOutput() { co.worklytics.psoxy.rules.RulesUtils rulesUtils) { String rulesStr = configService.getConfigPropertyOrError(ProxyConfigProperty.RULES); String yamlEncodedRules = rulesUtils.decodeToYaml(rulesStr); - return webhookSanitizerFactory.create(objectMapper.readerFor(WebhookCollectionRules.class).readValue(yamlEncodedRules)); + WebhookCollectionRules rules = + objectMapper.readerFor(WebhookCollectionRules.class).readValue(yamlEncodedRules); + AugmentValidation.validateWebhookEndpoints(rules.getEndpoints()); + return webhookSanitizerFactory.create(rules); } @Provides @Singleton diff --git a/java/core/src/main/java/co/worklytics/psoxy/PsoxyModule.java b/java/core/src/main/java/co/worklytics/psoxy/PsoxyModule.java index 316baff37..55334646b 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/PsoxyModule.java +++ b/java/core/src/main/java/co/worklytics/psoxy/PsoxyModule.java @@ -38,7 +38,15 @@ import co.worklytics.psoxy.gateway.SourceAuthStrategy; import co.worklytics.psoxy.gateway.auth.Base64KeyClient; import co.worklytics.psoxy.gateway.impl.EnvVarsConfigService; +import com.avaulta.gateway.resources.ResourceService; +import com.avaulta.gateway.rules.augments.GenMetadataBackend; +import com.avaulta.gateway.rules.augments.GenMetadataProcessor; +import com.avaulta.gateway.rules.augments.UnavailableGenMetadataBackend; import co.worklytics.psoxy.impl.AugmentProcessor; +import co.worklytics.psoxy.impl.gen.GenMetadataChatModelFactory; +import co.worklytics.psoxy.impl.gen.GenMetadataConfig; +import co.worklytics.psoxy.impl.gen.GenMetadataPromptBudget; +import co.worklytics.psoxy.impl.gen.LangChain4jGenMetadataBackend; import co.worklytics.psoxy.gateway.impl.oauth.OAuthRefreshTokenSourceAuthStrategy; import co.worklytics.psoxy.storage.BulkDataSanitizerFactory; import co.worklytics.psoxy.storage.impl.BulkDataSanitizerFactoryImpl; @@ -382,4 +390,33 @@ Base64KeyClient base64KeyClient() { Base64.Encoder provideBase64Encoder() { return Base64.getEncoder(); } + + /** + * Wires genMetadata backend. Model weights are resolved via {@code @Named("ForGenMetadata")} + * {@link ResourceService} from {@code ResourceServiceBindingsModule}. + */ + @Provides + @Singleton + static GenMetadataBackend genMetadataBackend( + ConfigService configService, + ObjectMapper objectMapper, + GenMetadataPromptBudget promptBudget, + GenMetadataChatModelFactory chatModelFactory) { + GenMetadataConfig config = GenMetadataConfig.from(configService); + if (GenMetadataConfig.BACKEND_LOCAL.equalsIgnoreCase(config.getBackend())) { + return new LangChain4jGenMetadataBackend(config, objectMapper, promptBudget, chatModelFactory); + } + return new UnavailableGenMetadataBackend(); + } + + @Provides + @Singleton + static GenMetadataProcessor genMetadataProcessor(GenMetadataBackend genMetadataBackend, + ObjectMapper objectMapper, + ConfigService configService, + JsonSchemaValidationUtils jsonSchemaValidationUtils) { + GenMetadataConfig config = GenMetadataConfig.from(configService); + return new GenMetadataProcessor(genMetadataBackend, objectMapper, config.getMaxInputChars(), + config.getMaxAttempts(), jsonSchemaValidationUtils); + } } diff --git a/java/core/src/main/java/co/worklytics/psoxy/RESTApiSanitizer.java b/java/core/src/main/java/co/worklytics/psoxy/RESTApiSanitizer.java index 6b920b121..afa1330b9 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/RESTApiSanitizer.java +++ b/java/core/src/main/java/co/worklytics/psoxy/RESTApiSanitizer.java @@ -5,6 +5,7 @@ import java.io.OutputStream; import java.net.URL; import java.util.Collection; +import java.util.List; import java.util.Optional; import co.worklytics.psoxy.rules.RESTRules; import lombok.NonNull; @@ -59,6 +60,13 @@ public interface RESTApiSanitizer { */ String sanitize(String httpMethod, URL url, String jsonResponse); + /** + * Warning codes from augment processing on the last {@link #sanitize} call for this thread. + * Suitable for {@link ProcessedDataMetadataFields#WARNING} response headers. + */ + default List getLastSanitizationWarnings() { + return List.of(); + } Pseudonymizer getPseudonymizer(); diff --git a/java/core/src/main/java/co/worklytics/psoxy/ResourceServiceBindingsModule.java b/java/core/src/main/java/co/worklytics/psoxy/ResourceServiceBindingsModule.java new file mode 100644 index 000000000..ee9dc4575 --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/ResourceServiceBindingsModule.java @@ -0,0 +1,62 @@ +package co.worklytics.psoxy; + +import javax.inject.Named; +import javax.inject.Singleton; + +import com.avaulta.gateway.resources.ResourceService; +import com.avaulta.gateway.rules.augments.SentenceMetadataProcessor; + +import co.worklytics.psoxy.gateway.impl.CompositeResourceService; +import co.worklytics.psoxy.gateway.impl.LocalFileResourceService; +import dagger.Module; +import dagger.Provides; + +/** + * Shared {@link ResourceService} composition for augment runtimes (local FS → instance remote → shared remote). + * + *

Platform modules must provide {@code @Named("Remote")} and {@code @Named("SharedRemote")}.

+ */ +@Module +public class ResourceServiceBindingsModule { + + @Provides + @Singleton + static ResourceService instanceResourceService(@Named("Remote") ResourceService remoteResourceService) { + return CompositeResourceService.builder() + .preferred(new LocalFileResourceService(ResourceService.DEFAULT_LOCAL_RESOURCE_PATH)) + .fallback(remoteResourceService) + .build(); + } + + @Provides + @Singleton + @Named("ForGenMetadata") + static ResourceService genMetadataResourceService(@Named("Remote") ResourceService remoteResourceService, + @Named("SharedRemote") ResourceService sharedRemoteResourceService) { + ResourceService instanceResourceService = CompositeResourceService.builder() + .preferred(new LocalFileResourceService(ResourceService.DEFAULT_LOCAL_RESOURCE_PATH)) + .fallback(remoteResourceService) + .build(); + return CompositeResourceService.builder() + .preferred(instanceResourceService) + .fallback(sharedRemoteResourceService) + .build(); + } + + @Provides + @Singleton + @Named("OpenNlp") + static ResourceService openNlpResourceService(ResourceService instanceResourceService, + @Named("SharedRemote") ResourceService sharedRemoteResourceService) { + return CompositeResourceService.builder() + .preferred(instanceResourceService) + .fallback(sharedRemoteResourceService) + .build(); + } + + @Provides + @Singleton + static SentenceMetadataProcessor sentenceMetadataProcessor(@Named("OpenNlp") ResourceService openNlpResourceService) { + return new SentenceMetadataProcessor(openNlpResourceService); + } +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/Warning.java b/java/core/src/main/java/co/worklytics/psoxy/Warning.java index a60cf5cfc..f50efc91f 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/Warning.java +++ b/java/core/src/main/java/co/worklytics/psoxy/Warning.java @@ -3,6 +3,11 @@ public enum Warning { COMPRESSION_NOT_REQUESTED, + + AUGMENT_OUTPUT_SCHEMA_MISMATCH, + AUGMENT_GEN_INFERENCE_FAILED, + AUGMENT_GEN_UNAVAILABLE, + AUGMENT_CONFLICT_SKIPPED, ; public String asHttpHeaderCode() { diff --git a/java/core/src/main/java/co/worklytics/psoxy/gateway/ProcessedContent.java b/java/core/src/main/java/co/worklytics/psoxy/gateway/ProcessedContent.java index 02d32b33f..95339d55a 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/gateway/ProcessedContent.java +++ b/java/core/src/main/java/co/worklytics/psoxy/gateway/ProcessedContent.java @@ -6,7 +6,9 @@ import java.io.Serializable; import java.nio.charset.Charset; import java.nio.charset.StandardCharsets; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; @@ -49,6 +51,12 @@ public class ProcessedContent implements Serializable { @Builder.Default Map metadata = new HashMap<>(); + /** + * Non-fatal warning codes from sanitization (e.g. augment failures), for {@code X-Psoxy-Warning} headers. + */ + @Builder.Default + List sanitizationWarnings = Collections.emptyList(); + /** * the actual content */ diff --git a/java/core/src/main/java/co/worklytics/psoxy/gateway/ProxyConfigProperty.java b/java/core/src/main/java/co/worklytics/psoxy/gateway/ProxyConfigProperty.java index bd46c2861..8c636e4c5 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/gateway/ProxyConfigProperty.java +++ b/java/core/src/main/java/co/worklytics/psoxy/gateway/ProxyConfigProperty.java @@ -139,6 +139,33 @@ public enum ProxyConfigProperty implements ConfigService.ConfigProperty { */ SOURCE, + /** BETA: generative backend for genMetadata augment (`local`, `bedrock`, `vertex`). */ + PSOXY_GEN_BACKEND, + + /** BETA: Jlama HuggingFace model id or logical id for `{SHARED_RESOURCE_PATH}/llm/{id}.zip`. */ + PSOXY_GEN_MODEL, + + /** BETA: max seconds for a single genMetadata {@code chat()} call (and model-load wait). */ + PSOXY_GEN_TIMEOUT_SECONDS, + + /** BETA: max source characters passed into genMetadata prompts. */ + PSOXY_GEN_MAX_INPUT_CHARS, + + /** BETA: max tokens to generate per genMetadata inference. */ + PSOXY_GEN_MAX_TOKENS, + + /** + * BETA: total genMetadata inference attempts per augment when output is unparseable or fails + * {@code outputSchema} validation (minimum 1). Default {@code 2} = one retry. + */ + PSOXY_GEN_META_RETRIES, + + /** + * BETA: set when Terraform {@code enable_gen_metadata} is used; documents that deployment + * is sized for local generative augments. + */ + ENABLE_GEN_METADATA, + /** * Whether the proxy should follow HTTP redirects (3xx responses) when calling source APIs. * diff --git a/java/core/src/main/java/co/worklytics/psoxy/gateway/impl/ApiDataRequestHandler.java b/java/core/src/main/java/co/worklytics/psoxy/gateway/impl/ApiDataRequestHandler.java index c1b638d29..544d91019 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/gateway/impl/ApiDataRequestHandler.java +++ b/java/core/src/main/java/co/worklytics/psoxy/gateway/impl/ApiDataRequestHandler.java @@ -574,6 +574,10 @@ && isSafeMethod(requestToSourceApi.getRequestMethod())) { proxyResponseContent = sanitizationResult.getContentAsString(); sanitizationResult.getMetadata().entrySet() .forEach(e -> builder.header(e.getKey(), e.getValue())); + sanitizationResult.getSanitizationWarnings().forEach(warningCode -> + builder.multivaluedHeader( + Pair.of(ProcessedDataMetadataFields.WARNING.getHttpHeader(), + warningCode))); } @@ -633,11 +637,11 @@ ProcessedContent sanitize(HttpEventRequest request, RequestUrls requestUrls, metadata.put(ProcessedDataMetadataFields.PII_SALT_SHA256.getMetadataKey(), healthCheckRequestHandler.piiSaltHash()); - // q: add instance id to the metadata?? return ProcessedContent.builder() .contentType(originalContent.getContentType()) .contentCharset(originalContent.getContentCharset()) .metadata(metadata) + .sanitizationWarnings(sanitizerForRequest.getLastSanitizationWarnings()) .content(sanitized.getBytes(originalContent.getContentCharset())) .build(); } diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessingException.java b/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessingException.java new file mode 100644 index 000000000..6c31d5c06 --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessingException.java @@ -0,0 +1,28 @@ +package co.worklytics.psoxy.impl; + +import co.worklytics.psoxy.Warning; +import lombok.Getter; + +/** + * Non-fatal augment failure to be caught by {@link AugmentProcessor} and surfaced as + * {@code X-Psoxy-Warning} response headers. + */ +@Getter +public class AugmentProcessingException extends Exception { + + private final Warning warning; + + public AugmentProcessingException(Warning warning, String message) { + super(message); + this.warning = warning; + } + + public AugmentProcessingException(Warning warning, String message, Throwable cause) { + super(message, cause); + this.warning = warning; + } + + public String getWarningCode() { + return warning.asHttpHeaderCode(); + } +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessor.java b/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessor.java index e9abbe3f2..5db65b50d 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessor.java +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/AugmentProcessor.java @@ -1,29 +1,35 @@ package co.worklytics.psoxy.impl; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.avaulta.gateway.rules.JsonSchemaValidationUtils; import com.avaulta.gateway.rules.augments.Augment; +import com.avaulta.gateway.rules.augments.GenMetadataAugmentException; +import com.avaulta.gateway.rules.augments.GenMetadataProcessor; import com.avaulta.gateway.rules.augments.SentenceMetadataProcessor; +import com.fasterxml.jackson.databind.ObjectMapper; import com.jayway.jsonpath.Configuration; -import com.jayway.jsonpath.DocumentContext; import com.jayway.jsonpath.JsonPath; import com.jayway.jsonpath.Option; import com.jayway.jsonpath.PathNotFoundException; +import co.worklytics.psoxy.Warning; +import lombok.NonNull; import lombok.extern.java.Log; import javax.inject.Inject; +import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.TreeMap; import java.util.concurrent.ConcurrentHashMap; import java.util.logging.Level; /** * Applies augments to a JSON document, adding synthetic sibling properties - * named {@code +{sourceProperty}:{augmentFunction}}. When {@code innerJsonPath} is used, each - * inner match is augmented in place within the parsed embedded JSON, then the modified structure - * is stored as the augment property value (mirroring {@code Transform.TextDigest} with - * {@code isJsonEscaped}). + * named {@code +{sourceProperty}:{augmentFunction}}. When {@code innerJsonPath} is used, multiple + * inner matches are grouped under that property as an object keyed by inner path suffix. * - *

Processing is intentionally non-fatal: any augment failure (exception, timeout, - * schema validation) results in the augment property being omitted with a warning logged. + *

Processing is intentionally non-fatal: {@link AugmentProcessingException} and other failures + * omit the augment property and add {@code X-Psoxy-Warning} codes via the returned list. * *

NOTE: this class must be thread-safe. A single instance may be shared across * concurrent requests. Compiled JsonPaths are cached in a ConcurrentHashMap. @@ -41,6 +47,8 @@ public class AugmentProcessor { /** Default config: reads return matched values (used for parent/source lookups). */ final Configuration jsonConfiguration; + final JsonSchemaValidationUtils jsonSchemaValidationUtils; + final ObjectMapper objectMapper; /** * Same underlying provider as {@link #jsonConfiguration}, but with {@link Option#AS_PATH_LIST}: @@ -52,20 +60,22 @@ public class AugmentProcessor { final Configuration pathListConfiguration; final SentenceMetadataProcessor sentenceMetadataProcessor; + final GenMetadataProcessor genMetadataProcessor; @Inject public AugmentProcessor(Configuration jsonConfiguration, - SentenceMetadataProcessor sentenceMetadataProcessor) { + JsonSchemaValidationUtils jsonSchemaValidationUtils, + ObjectMapper objectMapper, + SentenceMetadataProcessor sentenceMetadataProcessor, + GenMetadataProcessor genMetadataProcessor) { this.jsonConfiguration = jsonConfiguration; + this.jsonSchemaValidationUtils = jsonSchemaValidationUtils; + this.objectMapper = objectMapper; this.pathListConfiguration = jsonConfiguration.setOptions(Option.AS_PATH_LIST); this.sentenceMetadataProcessor = sentenceMetadataProcessor; + this.genMetadataProcessor = genMetadataProcessor; } - /** - * Cache of pre-compiled JsonPaths, keyed by Augment identity. - * Avoids recompiling paths on every request (same pattern as - * {@code RESTApiSanitizerImpl.compiledTransforms}). - */ private final Map> compiledAugmentPaths = new ConcurrentHashMap<>(); /** @@ -78,33 +88,35 @@ public AugmentProcessor(Configuration jsonConfiguration, * * @param augments the augment rules to apply * @param document the Jayway JSONPath document (will be mutated in place) + * @return warning header codes to add to the response (may be empty) */ - public void applyAugments(List augments, Object document) { + public List applyAugments(List augments, Object document) { if (augments == null || augments.isEmpty()) { - return; + return List.of(); } - // Conflict check: if any existing property starts with "+", skip all augments + List warnings = new ArrayList<>(); + if (hasConflictingProperties(document)) { log.warning("Response contains properties starting with '" + AUGMENT_PROPERTY_PREFIX + "'; skipping all augment processing to avoid conflicts."); - return; + warnings.add(Warning.AUGMENT_CONFLICT_SKIPPED.asHttpHeaderCode()); + return List.copyOf(warnings); } for (Augment augment : augments) { try { - applyAugment(augment, document); + warnings.addAll(applyAugment(augment, document)); } catch (Exception e) { log.log(Level.WARNING, "Augment '" + augment.getFunctionName() + "' failed; skipping.", e); } } + return List.copyOf(warnings); } - /** - * Apply a single augment to the document using its pre-compiled paths. - */ - private void applyAugment(Augment augment, Object document) { + private List applyAugment(Augment augment, Object document) { + List warnings = new ArrayList<>(); List paths = compiledAugmentPaths.computeIfAbsent(augment, a -> a.getJsonPaths().stream() .map(JsonPath::compile) @@ -112,7 +124,7 @@ private void applyAugment(Augment augment, Object document) { for (JsonPath compiledPath : paths) { try { - applyAugmentAtPath(augment, document, compiledPath); + warnings.addAll(applyAugmentAtPath(augment, document, compiledPath)); } catch (PathNotFoundException e) { // expected if path doesn't match this particular document — no-op } catch (Exception e) { @@ -121,47 +133,41 @@ private void applyAugment(Augment augment, Object document) { + compiledPath.getPath() + "'; skipping.", e); } } + return warnings; } - /** - * Apply an augment to all values matching a single compiled JSON path. - * - *

Uses {@code AS_PATH_LIST} to resolve the path expression to concrete paths - * (no wildcards), then derives the parent of each concrete path to insert the - * sibling augment property. - */ @SuppressWarnings("unchecked") - private void applyAugmentAtPath(Augment augment, Object document, JsonPath compiledPath) { + private List applyAugmentAtPath(Augment augment, Object document, JsonPath compiledPath) { + List warnings = new ArrayList<>(); List resolvedPaths; try { resolvedPaths = compiledPath.read(document, pathListConfiguration); } catch (PathNotFoundException e) { - return; + return warnings; } if (resolvedPaths == null || resolvedPaths.isEmpty()) { - return; + return warnings; } for (String concretePath : resolvedPaths) { try { applyAugmentAtConcretePath(augment, document, concretePath); + } catch (AugmentProcessingException e) { + log.log(Level.WARNING, e.getMessage(), e); + warnings.add(e.getWarningCode()); } catch (Exception e) { log.log(Level.WARNING, "Augment '" + augment.getFunctionName() + "' failed at concrete path '" + concretePath + "'; skipping.", e); } } + return warnings; } - /** - * Apply an augment at a single concrete (fully-resolved) JSON path. - * - *

Reads the source value, computes the augment, and inserts the result as a - * sibling property in the parent object. - */ @SuppressWarnings("unchecked") - private void applyAugmentAtConcretePath(Augment augment, Object document, String concretePath) { + private void applyAugmentAtConcretePath(Augment augment, Object document, String concretePath) + throws AugmentProcessingException { String leafFieldName = extractLeafFieldNameFromConcrete(concretePath); String parentPath = extractParentFromConcrete(concretePath); @@ -195,25 +201,74 @@ private void applyAugmentAtConcretePath(Augment augment, Object document, String String augmentPropertyName = buildAugmentPropertyName(leafFieldName, augment.getFunctionName()); - try { - if (hasInnerJsonPath(augment) && sourceValue instanceof String jsonStr && !jsonStr.isEmpty()) { - applyInnerJsonPathAugments(augment, (Map) parent, leafFieldName, jsonStr); - return; + if (hasInnerJsonPath(augment) && sourceValue instanceof String jsonStr && !jsonStr.isEmpty()) { + applyInnerJsonPathAugments(augment, (Map) parent, leafFieldName, jsonStr); + return; + } + + Object augmentValue = invokeCompute(augment, sourceValue); + + if (augmentValue == null) { + if (augment instanceof Augment.GenMetadata) { + throw new AugmentProcessingException(Warning.AUGMENT_GEN_UNAVAILABLE, + "genMetadata returned no value for property '" + augmentPropertyName + "'"); } + return; + } - Object augmentValue = computeAugmentValue(augment, sourceValue); - if (augmentValue == null) { - return; + if (!validateOutputSchema(augment, augmentValue, augmentPropertyName)) { + throw new AugmentProcessingException(Warning.AUGMENT_OUTPUT_SCHEMA_MISMATCH, + "Augment '" + augment.getFunctionName() + + "' output failed schema validation for property '" + augmentPropertyName + + "'"); + } + + ((Map) parent).put(augmentPropertyName, augmentValue); + } + + private Object invokeCompute(Augment augment, Object input) throws AugmentProcessingException { + try { + if (augment instanceof Augment.SentenceMetadata sentenceMetadata) { + return sentenceMetadataProcessor.compute(sentenceMetadata, input); + } + if (augment instanceof Augment.GenMetadata genMetadata) { + return genMetadataProcessor.compute(genMetadata, input); } + return augment.compute(input); + } catch (GenMetadataAugmentException e) { + throw toAugmentProcessingException(e); + } + } - // TODO: validate against outputSchema if present (predicate check) - // For PoC, outputSchema validation is deferred + private static AugmentProcessingException toAugmentProcessingException( + GenMetadataAugmentException e) { + Warning warning = switch (e.getCode()) { + case UNAVAILABLE -> Warning.AUGMENT_GEN_UNAVAILABLE; + case INFERENCE_FAILED -> Warning.AUGMENT_GEN_INFERENCE_FAILED; + }; + return e.getCause() != null + ? new AugmentProcessingException(warning, e.getMessage(), e.getCause()) + : new AugmentProcessingException(warning, e.getMessage()); + } - ((Map) parent).put(augmentPropertyName, augmentValue); + private boolean validateOutputSchema(@NonNull Augment augment, @NonNull Object augmentValue, + @NonNull String augmentPropertyName) { + JsonSchemaFilter outputSchema = augment.getOutputSchema(); + if (outputSchema == null) { + return true; + } + try { + String json = objectMapper.writeValueAsString(augmentValue); + boolean valid = jsonSchemaValidationUtils.validateJsonBySchema(json, outputSchema); + if (!valid) { + log.warning("Augment '" + augment.getFunctionName() + + "' output for '" + augmentPropertyName + "' failed schema validation"); + } + return valid; } catch (Exception e) { - log.log(Level.WARNING, - "Augment '" + augment.getFunctionName() + "' compute failed for property '" - + augmentPropertyName + "'; omitting.", e); + log.log(Level.WARNING, "Failed to validate augment output schema for '" + + augmentPropertyName + "'", e); + return false; } } @@ -222,20 +277,19 @@ private boolean hasInnerJsonPath(Augment augment) { } /** - * When {@link Augment#getInnerJsonPath()} is set, resolve each concrete inner path, replace - * matched inner field values with serialized augment output, then store the modified embedded - * JSON re-serialized as a string under {@code +{outerLeaf}:{fn}} (mirroring - * {@code Transform.TextDigest} with {@code isJsonEscaped}). + * When {@link Augment#getInnerJsonPath()} is set, resolve each concrete inner path and group + * results under {@code +{outerLeaf}:{fn}} keyed by inner path suffix. */ + @SuppressWarnings("unchecked") private void applyInnerJsonPathAugments(Augment augment, Map parent, - String leafFieldName, String jsonStr) { - DocumentContext innerContext = JsonPath.parse(jsonStr); - Object innerDocument = innerContext.json(); + String leafFieldName, String jsonStr) + throws AugmentProcessingException { + JsonPath innerCompiled = getCompiledPath(augment.getInnerJsonPath()); + Object innerDocument = JsonPath.parse(jsonStr).json(); List innerConcretePaths; try { - innerConcretePaths = getCompiledPath(augment.getInnerJsonPath()) - .read(innerDocument, pathListConfiguration); + innerConcretePaths = innerCompiled.read(innerDocument, pathListConfiguration); } catch (PathNotFoundException e) { return; } @@ -243,20 +297,27 @@ private void applyInnerJsonPathAugments(Augment augment, Map par return; } - boolean anyApplied = false; + Map innerAugments = new TreeMap<>(); for (String innerConcretePath : innerConcretePaths) { try { - Object innerValue = innerContext.read(innerConcretePath); - Object augmentValue = computeAugmentValue(augment, innerValue); + Object innerValue = getCompiledPath(innerConcretePath).read(innerDocument, jsonConfiguration); + Object augmentValue = invokeCompute(augment, innerValue); if (augmentValue == null) { continue; } + if (!validateOutputSchema(augment, augmentValue, + buildAugmentPropertyName(leafFieldName, augment.getFunctionName()))) { + log.warning("Augment '" + augment.getFunctionName() + + "' output failed schema validation at inner path '" + innerConcretePath + + "'; skipping."); + continue; + } - // TODO: validate against outputSchema if present (predicate check) - String serializedAugment = jsonConfiguration.jsonProvider().toJson(augmentValue); - innerContext.set(innerConcretePath, serializedAugment); - anyApplied = true; + String innerKey = toInnerPathSuffix(innerConcretePath); + innerAugments.put(innerKey, augmentValue); + } catch (AugmentProcessingException e) { + log.log(Level.WARNING, e.getMessage(), e); } catch (Exception e) { log.log(Level.WARNING, "Augment '" + augment.getFunctionName() + "' failed at inner path '" @@ -264,19 +325,11 @@ private void applyInnerJsonPathAugments(Augment augment, Map par } } - if (anyApplied) { - parent.put(buildAugmentPropertyName(leafFieldName, augment.getFunctionName()), - innerContext.jsonString()); + if (!innerAugments.isEmpty()) { + parent.put(buildAugmentPropertyName(leafFieldName, augment.getFunctionName()), innerAugments); } } - private Object computeAugmentValue(Augment augment, Object sourceValue) { - if (augment instanceof Augment.SentenceMetadata sentenceMetadata) { - return sentenceMetadataProcessor.compute(sentenceMetadata, sourceValue); - } - return augment.compute(sourceValue); - } - /** * Build augment property name: {@code +content:textDigest}. */ @@ -334,10 +387,6 @@ private boolean hasConflictingPropertiesInMap(Map map) { return false; } - /** - * Extract the leaf field name from a concrete (resolved) JSON path. - * Concrete paths use bracket notation: {@code $['body']['content']} → {@code "content"}. - */ static String extractLeafFieldNameFromConcrete(String concretePath) { int lastBracket = concretePath.lastIndexOf("['"); if (lastBracket >= 0) { @@ -349,10 +398,6 @@ static String extractLeafFieldNameFromConcrete(String concretePath) { return null; } - /** - * Extract the parent path from a concrete (resolved) JSON path. - * {@code $['body']['content']} → {@code $['body']}. - */ static String extractParentFromConcrete(String concretePath) { int lastBracket = concretePath.lastIndexOf("['"); if (lastBracket > 0) { @@ -365,7 +410,6 @@ static String extractParentFromConcrete(String concretePath) { return null; } - // keep legacy helpers for tests that use them directly static String extractLeafFieldName(String jsonPath) { String cleaned = jsonPath.replaceAll("\\[\\*]$", ""); int lastDot = cleaned.lastIndexOf('.'); diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/RESTApiSanitizerImpl.java b/java/core/src/main/java/co/worklytics/psoxy/impl/RESTApiSanitizerImpl.java index 15d02b821..70213dfe2 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/impl/RESTApiSanitizerImpl.java +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/RESTApiSanitizerImpl.java @@ -13,6 +13,7 @@ import java.net.URL; import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Locale; @@ -138,6 +139,13 @@ public RESTApiSanitizerImpl(@Assisted RESTRules rules, @Inject AugmentProcessor augmentProcessor; + private final ThreadLocal> lastAugmentWarnings = new ThreadLocal<>(); + + @Override + public List getLastSanitizationWarnings() { + List warnings = lastAugmentWarnings.get(); + return warnings == null ? List.of() : warnings; + } @Override public boolean isAllowed(@NonNull String httpMethod, @NonNull URL url) { @@ -253,6 +261,7 @@ public void sanitize(@NonNull String httpMethod, @NonNull URL url, InputStream originalStream, OutputStream outputStream) throws IOException { + lastAugmentWarnings.set(List.of()); if (!isAllowed(httpMethod, url)) { throw new IllegalStateException(String.format( "Sanitizer called to sanitize response that should not have been retrieved: %s", @@ -265,6 +274,7 @@ public void sanitize(@NonNull String httpMethod, } else { final Endpoint endpoint = matchingEndpoint.get().getValue(); final JsonFactory factory = objectMapper.getFactory(); + List aggregatedAugmentWarnings = new ArrayList<>(); try (JsonParser parser = factory.createParser(originalStream); OutputStreamWriter writer = new OutputStreamWriter(outputStream, StandardCharsets.UTF_8); BufferedWriter bufferedWriter = new BufferedWriter(writer)) { @@ -280,7 +290,7 @@ public void sanitize(@NonNull String httpMethod, Object node = jsonConfiguration.jsonProvider() .parse(objectMapper.writeValueAsString(objectMapper.readTree(parser))); - Object sanitized = sanitize(endpoint, node); + Object sanitized = sanitize(endpoint, node, aggregatedAugmentWarnings); if (!first) bufferedWriter.write("\n"); bufferedWriter.write(jsonConfiguration.jsonProvider().toJson(sanitized)); @@ -288,6 +298,7 @@ public void sanitize(@NonNull String httpMethod, } bufferedWriter.flush(); } + lastAugmentWarnings.set(List.copyOf(aggregatedAugmentWarnings)); } } @@ -296,12 +307,15 @@ public void sanitize(@NonNull String httpMethod, * * if endpoint allows for ndjson, we expect this to be just a single JSON object at a time, and * you call this once per row in the response + * + * TODO: return a result type (document + warnings) instead of mutating {@code augmentWarnings} + * for NDJSON aggregation; a {@code Pair} or small record would invert control here. */ - private Object sanitize(Endpoint endpoint, Object jsonResponse) { + private Object sanitize(Endpoint endpoint, Object jsonResponse, List augmentWarnings) { // 1. Augments: add synthetic sibling properties before any filtering/transforms if (ObjectUtils.isNotEmpty(endpoint.getAugments())) { - augmentProcessor.applyAugments(endpoint.getAugments(), jsonResponse); + augmentWarnings.addAll(augmentProcessor.applyAugments(endpoint.getAugments(), jsonResponse)); } // 2. Response schema filter (allow-list); auto-passes "+" properties when augments ran diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/WebhookSanitizerImpl.java b/java/core/src/main/java/co/worklytics/psoxy/impl/WebhookSanitizerImpl.java index 4a2fae9f7..0fb458505 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/impl/WebhookSanitizerImpl.java +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/WebhookSanitizerImpl.java @@ -11,6 +11,7 @@ import com.jayway.jsonpath.Configuration; import com.jayway.jsonpath.JsonPath; import com.jayway.jsonpath.PathNotFoundException; +import co.worklytics.psoxy.impl.AugmentProcessor; import dagger.assisted.Assisted; import dagger.assisted.AssistedInject; import lombok.extern.java.Log; @@ -36,6 +37,9 @@ public class WebhookSanitizerImpl implements WebhookSanitizer { @Inject Configuration jsonConfiguration; + @Inject + AugmentProcessor augmentProcessor; + final WebhookCollectionRules webhookRules; Map> compiledTransforms = new ConcurrentHashMap<>(); @@ -174,10 +178,16 @@ public ProcessedContent sanitize(HttpEventRequest request) { WebhookCollectionRules.WebhookEndpoint endpoint = webhookRules.getEndpoints().get(0); String sanitizedContent = new String(request.getBody(), StandardCharsets.UTF_8); //just assume UTF-8, always - if (ObjectUtils.isNotEmpty(endpoint.getTransforms())) { + if (ObjectUtils.isNotEmpty(endpoint.getAugments()) + || ObjectUtils.isNotEmpty(endpoint.getTransforms())) { Object document = jsonConfiguration.jsonProvider().parse(request.getBody()); - for (Transform transform : endpoint.getTransforms()) { - sanitizerUtils.applyTransform(pseudonymizer, transform, document, compiledTransforms); + if (ObjectUtils.isNotEmpty(endpoint.getAugments())) { + augmentProcessor.applyAugments(endpoint.getAugments(), document); + } + if (ObjectUtils.isNotEmpty(endpoint.getTransforms())) { + for (Transform transform : endpoint.getTransforms()) { + sanitizerUtils.applyTransform(pseudonymizer, transform, document, compiledTransforms); + } } sanitizedContent = jsonConfiguration.jsonProvider().toJson(document); } diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataChatModelFactory.java b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataChatModelFactory.java new file mode 100644 index 000000000..47c72328f --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataChatModelFactory.java @@ -0,0 +1,113 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.resources.ResourceService; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.jlama.JlamaChatModel; +import lombok.extern.java.Log; + +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.logging.Level; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; + +/** + * Creates LangChain4j {@link ChatModel} instances for genMetadata backends. + */ +@Log +@Singleton +public class GenMetadataChatModelFactory { + + private final ResourceService resourceService; + + @Inject + public GenMetadataChatModelFactory(@Named("ForGenMetadata") ResourceService resourceService) { + this.resourceService = resourceService; + } + + public boolean supports(GenMetadataConfig config) { + return GenMetadataConfig.BACKEND_LOCAL.equalsIgnoreCase(config.getBackend()); + } + + /** + * Do not set {@code maxTokens} on {@link JlamaChatModel}: Jlama treats that value as the + * absolute KV-cache position ceiling (prompt + completion), not max new tokens. When unset, + * LangChain4j defaults it to the model's {@code contextLength}. Output length is bounded by + * {@link GenMetadataConfig#getMaxTokens()} via prompt budget reservation in + * {@link GenMetadataPromptBudget}. + */ + public ChatModel buildLocal(GenMetadataConfig config, Path modelCacheDir) throws Exception { + return buildLocal(config, modelCacheDir, null); + } + + /** + * @param jlamaMaxTokens when non-null, passed to {@link JlamaChatModel#maxTokens(Integer)}. + * Jlama treats this as the absolute KV position ceiling (prompt + completion), not max new + * tokens — see class javadoc. Used by integration tests to reproduce the misconfiguration. + */ + public ChatModel buildLocal(GenMetadataConfig config, Path modelCacheDir, Integer jlamaMaxTokens) + throws Exception { + materializeModelArchiveIfPresent(config, modelCacheDir); + String jlamaModelName = resolveJlamaModelName(config, modelCacheDir); + JlamaChatModel.JlamaChatModelBuilder builder = JlamaChatModel.builder() + .modelCachePath(modelCacheDir) + .modelName(jlamaModelName) + .temperature(0f); + if (jlamaMaxTokens != null) { + builder.maxTokens(jlamaMaxTokens); + } + return builder.build(); + } + + String resolveJlamaModelName(GenMetadataConfig config, Path modelCacheDir) { + Path localDir = modelCacheDir.resolve(config.localModelCacheDirName()); + if (Files.exists(localDir.resolve("config.json"))) { + return config.localModelCacheDirName(); + } + return config.getModelId(); + } + + void materializeModelArchiveIfPresent(GenMetadataConfig config, Path modelCacheDir) { + String archivePath = config.localModelArchivePath(); + resourceService.getResource(archivePath).ifPresent(stream -> { + try (InputStream in = stream) { + Path targetDir = modelCacheDir.resolve(config.localModelCacheDirName()); + if (Files.exists(targetDir.resolve("config.json"))) { + return; + } + Files.createDirectories(targetDir); + unzip(in, targetDir); + log.info("Materialized genMetadata model archive to " + targetDir); + } catch (Exception e) { + log.log(Level.WARNING, + "Failed to materialize genMetadata model archive from " + archivePath, e); + throw new RuntimeException(e); + } + }); + } + + private void unzip(InputStream zipStream, Path targetDir) throws Exception { + Path normalizedTarget = targetDir.toAbsolutePath().normalize(); + try (ZipInputStream zis = new ZipInputStream(zipStream)) { + ZipEntry entry; + while ((entry = zis.getNextEntry()) != null) { + Path out = normalizedTarget.resolve(entry.getName()).normalize(); + if (!out.startsWith(normalizedTarget)) { + throw new IOException("Zip entry escapes target directory: " + entry.getName()); + } + if (entry.isDirectory()) { + Files.createDirectories(out); + } else { + Files.createDirectories(out.getParent()); + Files.copy(zis, out); + } + zis.closeEntry(); + } + } + } +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataConfig.java b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataConfig.java new file mode 100644 index 000000000..4d2a08914 --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataConfig.java @@ -0,0 +1,94 @@ +package co.worklytics.psoxy.impl.gen; + +import co.worklytics.psoxy.gateway.ConfigService; +import co.worklytics.psoxy.gateway.ProxyConfigProperty; +import lombok.Builder; +import lombok.Value; +import org.apache.commons.lang3.StringUtils; + +import java.util.Optional; + +/** + * Deployment configuration for {@link com.avaulta.gateway.rules.augments.GenMetadataBackend}. + */ +@Value +@Builder +public class GenMetadataConfig { + + public static final int DEFAULT_MAX_ATTEMPTS = 2; + + public static final String BACKEND_LOCAL = "local"; + public static final String BACKEND_BEDROCK = "bedrock"; + public static final String BACKEND_VERTEX = "vertex"; + + /** + * Default Jlama model (HuggingFace {@code owner/name}); pre-quantized builds from + * tjake on Hugging Face. + */ + public static final String DEFAULT_MODEL = "tjake/Llama-3.2-1B-Instruct-JQ4"; + + String backend; + String modelId; + int timeoutSeconds; + int maxInputChars; + int maxTokens; + /** Total inference attempts per augment (including the first try). */ + @Builder.Default + int maxAttempts = DEFAULT_MAX_ATTEMPTS; + + public static GenMetadataConfig from(ConfigService configService) { + String backend = configService.getConfigPropertyAsOptional(ProxyConfigProperty.PSOXY_GEN_BACKEND) + .filter(StringUtils::isNotBlank) + .orElse(BACKEND_LOCAL) + .trim() + .toLowerCase(); + String model = configService.getConfigPropertyAsOptional(ProxyConfigProperty.PSOXY_GEN_MODEL) + .filter(StringUtils::isNotBlank) + .orElse(DEFAULT_MODEL) + .trim(); + int timeout = configService.getConfigPropertyAsOptional(ProxyConfigProperty.PSOXY_GEN_TIMEOUT_SECONDS) + .flatMap(GenMetadataConfig::parsePositiveInt) + .orElse(15); + int maxInput = configService.getConfigPropertyAsOptional(ProxyConfigProperty.PSOXY_GEN_MAX_INPUT_CHARS) + .flatMap(GenMetadataConfig::parsePositiveInt) + .orElse(4096); + int maxTokens = configService.getConfigPropertyAsOptional(ProxyConfigProperty.PSOXY_GEN_MAX_TOKENS) + .flatMap(GenMetadataConfig::parsePositiveInt) + .orElse(256); + int maxAttempts = configService.getConfigPropertyAsOptional(ProxyConfigProperty.PSOXY_GEN_META_RETRIES) + .flatMap(GenMetadataConfig::parsePositiveInt) + .orElse(2); + return GenMetadataConfig.builder() + .backend(backend) + .modelId(model) + .timeoutSeconds(timeout) + .maxInputChars(maxInput) + .maxTokens(maxTokens) + .maxAttempts(maxAttempts) + .build(); + } + + private static Optional parsePositiveInt(String raw) { + try { + int v = Integer.parseInt(raw.trim()); + return v > 0 ? Optional.of(v) : Optional.empty(); + } catch (NumberFormatException e) { + return Optional.empty(); + } + } + + /** + * Optional zip archive in remote resources containing a Jlama SafeTensors model directory + * ({@code config.json} at the root of the archive or in a single top-level folder). + */ + public String localModelArchivePath() { + return "llm/" + localModelCacheDirName() + ".zip"; + } + + /** + * Directory name under the Jlama model cache for this {@link #modelId}. + */ + public String localModelCacheDirName() { + return modelId.replace("/", "__"); + } +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBudget.java b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBudget.java new file mode 100644 index 000000000..a11ddaf6d --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBudget.java @@ -0,0 +1,61 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.fasterxml.jackson.databind.ObjectMapper; + +import javax.inject.Inject; +import javax.inject.Singleton; + +/** + * Truncates genMetadata input so prompts fit local model context alongside reserved output tokens. + */ +@Singleton +public class GenMetadataPromptBudget { + + /** Rough chars-per-token for English text (conservative for JSON-escaped input). */ + private static final int CHARS_PER_TOKEN_ESTIMATE = 4; + + /** Template, labels, and chat formatting overhead beyond task + schema + input. */ + private static final int PROMPT_OVERHEAD_TOKENS = 64; + + private final ObjectMapper objectMapper; + + @Inject + public GenMetadataPromptBudget(ObjectMapper objectMapper) { + this.objectMapper = objectMapper; + } + + public String fitInputData(String inputData, String taskPrompt, JsonSchemaFilter outputSchema, + int maxInputChars, int contextLength, int maxOutputTokens) { + String schemaJson = schemaJson(outputSchema); + int contextChars = contextCharsBudget(contextLength, maxOutputTokens); + int overheadChars = charEstimate( + GenMetadataPromptBuilder.SYSTEM_PROMPT.length() + + taskPrompt.length() + + schemaJson.length() + + PROMPT_OVERHEAD_TOKENS * CHARS_PER_TOKEN_ESTIMATE); + int budgetChars = Math.min(maxInputChars, Math.max(0, contextChars - overheadChars)); + if (inputData.length() <= budgetChars) { + return inputData; + } + return inputData.substring(0, budgetChars); + } + + private int contextCharsBudget(int contextLength, int maxOutputTokens) { + int outputReserve = Math.max(1, maxOutputTokens); + int inputTokens = contextLength - outputReserve - PROMPT_OVERHEAD_TOKENS; + return Math.max(0, inputTokens * CHARS_PER_TOKEN_ESTIMATE); + } + + private int charEstimate(int chars) { + return Math.max(0, chars); + } + + private String schemaJson(JsonSchemaFilter outputSchema) { + try { + return objectMapper.writeValueAsString(outputSchema); + } catch (Exception e) { + return "{}"; + } + } +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBuilder.java b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBuilder.java new file mode 100644 index 000000000..596908933 --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBuilder.java @@ -0,0 +1,51 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.langchain4j.data.message.ChatMessage; +import dev.langchain4j.data.message.SystemMessage; +import dev.langchain4j.data.message.UserMessage; +import lombok.experimental.UtilityClass; + +import java.util.List; + +/** + * Builds chat messages for genMetadata inference (local and cloud backends). + */ +@UtilityClass +class GenMetadataPromptBuilder { + + static final String SYSTEM_PROMPT = + "You are a data-processing component in a privacy proxy. " + + "Respond with exactly one JSON object that is an INSTANCE of the task result, " + + "not a JSON Schema definition. " + + "Never include schema keywords such as type, properties, required, or enum. " + + "No markdown fences, no prose before or after the JSON."; + + static List toMessages(String taskPrompt, JsonSchemaFilter outputSchema, + String inputData, ObjectMapper objectMapper) { + return List.of( + SystemMessage.from(SYSTEM_PROMPT), + UserMessage.from(userContent(taskPrompt, outputSchema, inputData, objectMapper)) + ); + } + + static String userContent(String taskPrompt, JsonSchemaFilter outputSchema, String inputData, + ObjectMapper objectMapper) { + String schemaJson; + try { + schemaJson = objectMapper.writeValueAsString(outputSchema); + } catch (Exception e) { + schemaJson = "{}"; + } + return """ + Task: %s + + Your response must be a JSON object that validates against this schema (return an instance, NOT the schema itself): + %s + + Input data to process: + %s + """.formatted(taskPrompt.trim(), schemaJson, inputData); + } +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackend.java b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackend.java new file mode 100644 index 000000000..b4400de5e --- /dev/null +++ b/java/core/src/main/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackend.java @@ -0,0 +1,288 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.rules.augments.GenMetadataBackend; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.langchain4j.data.message.ChatMessage; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; +import lombok.AccessLevel; +import lombok.Getter; +import lombok.extern.java.Log; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantLock; +import java.util.logging.Level; + +/** + * genMetadata inference via LangChain4j {@link ChatModel} (Jlama for local embedded models). + * + *

Thread-safety: per-{@link GenMetadataConfig#getModelId()} {@link #loadLocks} give single-flight + * model load and zip extraction; {@link #inferenceLocks} serialize {@link ChatModel} use (Jlama is not + * safe for concurrent inference on one instance). {@link ObjectMapper} is shared read-only. + */ +@Log +public class LangChain4jGenMetadataBackend implements GenMetadataBackend { + + enum LoadState { + ABSENT, LOADING, READY, FAILED + } + + static final class ModelHandle { + final LoadState state; + final ChatModel chatModel; + final Exception failure; + + ModelHandle(LoadState state, ChatModel chatModel, Exception failure) { + this.state = state; + this.chatModel = chatModel; + this.failure = failure; + } + + static ModelHandle loading() { + return new ModelHandle(LoadState.LOADING, null, null); + } + + static ModelHandle ready(ChatModel chatModel) { + return new ModelHandle(LoadState.READY, chatModel, null); + } + + static ModelHandle failed(Exception e) { + return new ModelHandle(LoadState.FAILED, null, e); + } + } + + private final GenMetadataConfig config; + private final ObjectMapper objectMapper; + private final GenMetadataPromptBudget promptBudget; + private final GenMetadataChatModelFactory chatModelFactory; + /** + * Writable Jlama cache (extracted zips, HF downloads). Not the same as {@link ResourceService}, + * which only supplies read-only model archives from remote storage. + */ + @Getter(AccessLevel.PACKAGE) + private final Path modelCacheDir; + + private final ConcurrentHashMap models = new ConcurrentHashMap<>(); + private final ConcurrentHashMap loadLocks = new ConcurrentHashMap<>(); + private final ConcurrentHashMap inferenceLocks = new ConcurrentHashMap<>(); + private final ExecutorService chatExecutor = Executors.newCachedThreadPool(chatThreadFactory()); + + public LangChain4jGenMetadataBackend(GenMetadataConfig config, ObjectMapper objectMapper, + GenMetadataPromptBudget promptBudget, + GenMetadataChatModelFactory chatModelFactory) { + this(config, objectMapper, promptBudget, chatModelFactory, null); + } + + /** + * @param modelCacheDir when non-null, used instead of a fresh temp directory (integration tests) + */ + LangChain4jGenMetadataBackend(GenMetadataConfig config, ObjectMapper objectMapper, + GenMetadataPromptBudget promptBudget, + GenMetadataChatModelFactory chatModelFactory, + Path modelCacheDir) { + this.config = config; + this.objectMapper = objectMapper; + this.promptBudget = promptBudget; + this.chatModelFactory = chatModelFactory; + try { + this.modelCacheDir = modelCacheDir != null + ? modelCacheDir + : Files.createTempDirectory("psoxy-jlama-cache"); + if (modelCacheDir == null) { + this.modelCacheDir.toFile().deleteOnExit(); + } + } catch (Exception e) { + throw new IllegalStateException("Failed to create genMetadata model cache directory", e); + } + } + + @Override + public Object generate(String taskPrompt, JsonSchemaFilter outputSchema, String inputData) { + ModelHandle handle = resolveModel(); + if (handle.state != LoadState.READY || handle.chatModel == null) { + return null; + } + + String fittedInput = promptBudget.fitInputData( + inputData, + taskPrompt, + outputSchema, + config.getMaxInputChars(), + assumedContextLength(), + config.getMaxTokens()); + List messages = + GenMetadataPromptBuilder.toMessages(taskPrompt, outputSchema, fittedInput, objectMapper); + + ReentrantLock inferenceLock = + inferenceLocks.computeIfAbsent(config.getModelId(), k -> new ReentrantLock()); + boolean acquired = false; + try { + acquired = inferenceLock.tryLock(config.getTimeoutSeconds(), TimeUnit.SECONDS); + if (!acquired) { + log.warning("genMetadata inference lock timeout for model " + config.getModelId() + + " after " + config.getTimeoutSeconds() + "s"); + return null; + } + Instant inferenceStartedAt = Instant.now(); + long inferenceStartedNanos = System.nanoTime(); + log.info("genMetadata LLM inference started at " + inferenceStartedAt + + " modelId=" + config.getModelId()); + ChatResponse response; + try { + response = chatWithTimeout(handle.chatModel, messages); + } finally { + long inferenceMs = TimeUnit.NANOSECONDS.toMillis( + System.nanoTime() - inferenceStartedNanos); + log.info("genMetadata LLM inference completed in " + inferenceMs + "ms" + + " modelId=" + config.getModelId()); + } + if (response == null || response.aiMessage() == null) { + return null; + } + String text = response.aiMessage().text(); + if (text != null && !text.isBlank()) { + log.info("genMetadata raw model response: " + truncateForLog(text)); + } + return text; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return null; + } catch (Exception e) { + log.log(Level.WARNING, "genMetadata local inference failed", e); + return null; + } finally { + if (acquired) { + inferenceLock.unlock(); + } + } + } + + /** + * Conservative context length for pre-inference input budgeting before the model is loaded. + * Llama 3.2 1B supports much larger windows; Jlama defaults to model config when generating. + */ + int assumedContextLength() { + return 8192; + } + + ModelHandle resolveModel() { + String modelKey = config.getModelId(); + ModelHandle existing = models.get(modelKey); + if (existing != null) { + return switch (existing.state) { + case READY, FAILED -> existing; + case LOADING -> waitForLoad(modelKey); + case ABSENT -> loadModel(modelKey); + }; + } + return loadModel(modelKey); + } + + private ModelHandle loadModel(String modelKey) { + ReentrantLock loadLock = loadLocks.computeIfAbsent(modelKey, k -> new ReentrantLock()); + loadLock.lock(); + try { + ModelHandle existing = models.get(modelKey); + if (existing != null) { + if (existing.state == LoadState.LOADING) { + return waitForLoad(modelKey); + } + return existing; + } + models.put(modelKey, ModelHandle.loading()); + try { + ChatModel chatModel = + chatModelFactory.buildLocal(config, modelCacheDir); + ModelHandle ready = ModelHandle.ready(chatModel); + models.put(modelKey, ready); + log.info("Loaded genMetadata LangChain4j model: " + modelKey); + return ready; + } catch (Exception e) { + log.log(Level.WARNING, + "Failed to load genMetadata model '" + modelKey + "' (archive path: " + + config.localModelArchivePath() + ")", e); + ModelHandle failed = ModelHandle.failed(e); + models.put(modelKey, failed); + return failed; + } + } finally { + loadLock.unlock(); + } + } + + ChatResponse chatWithTimeout(ChatModel chatModel, List messages) throws Exception { + ChatRequest request = ChatRequest.builder().messages(messages).build(); + Future future = chatExecutor.submit(() -> chatModel.chat(request)); + try { + return future.get(config.getTimeoutSeconds(), TimeUnit.SECONDS); + } catch (TimeoutException e) { + future.cancel(true); + log.warning("genMetadata LLM inference timed out after " + config.getTimeoutSeconds() + + "s modelId=" + config.getModelId()); + return null; + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception ex) { + throw ex; + } + throw e; + } + } + + private static ThreadFactory chatThreadFactory() { + AtomicInteger threadNumber = new AtomicInteger(1); + return runnable -> { + Thread thread = new Thread(runnable, "genMetadata-chat-" + threadNumber.getAndIncrement()); + thread.setDaemon(true); + return thread; + }; + } + + private static final int MAX_LOG_OUTPUT_CHARS = 2000; + + private static String truncateForLog(String value) { + if (value == null) { + return "null"; + } + if (value.length() <= MAX_LOG_OUTPUT_CHARS) { + return value; + } + return value.substring(0, MAX_LOG_OUTPUT_CHARS) + "... (" + value.length() + " chars total)"; + } + + private ModelHandle waitForLoad(String modelKey) { + long deadline = System.nanoTime() + + TimeUnit.SECONDS.toNanos(config.getTimeoutSeconds()); + while (System.nanoTime() < deadline) { + ModelHandle handle = models.get(modelKey); + if (handle != null && handle.state != LoadState.LOADING) { + return handle; + } + try { + Thread.sleep(50); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return ModelHandle.failed(e); + } + } + log.warning("Timed out waiting for genMetadata model load: " + modelKey); + return models.getOrDefault(modelKey, ModelHandle.failed( + new IllegalStateException("model load timeout"))); + } + +} diff --git a/java/core/src/main/java/co/worklytics/psoxy/rules/Rules2.java b/java/core/src/main/java/co/worklytics/psoxy/rules/Rules2.java index 065aefbcd..1caa71931 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/rules/Rules2.java +++ b/java/core/src/main/java/co/worklytics/psoxy/rules/Rules2.java @@ -13,6 +13,7 @@ import java.util.stream.Collectors; import java.util.stream.Stream; import com.avaulta.gateway.rules.Endpoint; +import com.avaulta.gateway.rules.augments.AugmentValidation; import com.avaulta.gateway.rules.JsonSchemaFilter; import com.avaulta.gateway.rules.transforms.Transform; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; @@ -151,7 +152,9 @@ public static synchronized Rules2 load(String path) { if (is == null) { throw new IllegalArgumentException("Resource not found: " + path); } - return mapper.readerFor(Rules2.class).readValue(is); + Rules2 rules = mapper.readerFor(Rules2.class).readValue(is); + AugmentValidation.validateEndpoints(rules.getEndpoints()); + return rules; } catch (IOException e) { throw new RuntimeException("Failed to load rules from: " + path, e); } diff --git a/java/core/src/main/java/co/worklytics/psoxy/rules/Validator.java b/java/core/src/main/java/co/worklytics/psoxy/rules/Validator.java index 9e0ede1b0..4748efe09 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/rules/Validator.java +++ b/java/core/src/main/java/co/worklytics/psoxy/rules/Validator.java @@ -2,6 +2,7 @@ import com.avaulta.gateway.pseudonyms.PseudonymEncoder; import com.avaulta.gateway.rules.*; +import com.avaulta.gateway.rules.augments.AugmentValidation; import com.avaulta.gateway.rules.transforms.Transform; import com.google.common.base.Preconditions; import com.jayway.jsonpath.JsonPath; @@ -54,6 +55,7 @@ public void validate(@NonNull ColumnarRules rules) { public void validate(@NonNull RecordRules rules) { Preconditions.checkNotNull(rules.getFormat()); + AugmentValidation.validateRecordRules(rules); } public void validate(@NonNull MultiTypeBulkDataRules rules) { diff --git a/java/core/src/main/java/co/worklytics/psoxy/rules/msft/PrebuiltSanitizerRules.java b/java/core/src/main/java/co/worklytics/psoxy/rules/msft/PrebuiltSanitizerRules.java index 076e108da..817457307 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/rules/msft/PrebuiltSanitizerRules.java +++ b/java/core/src/main/java/co/worklytics/psoxy/rules/msft/PrebuiltSanitizerRules.java @@ -720,6 +720,49 @@ public class PrebuiltSanitizerRules { .jsonPath("$..body.content") .build(); + /** + * BETA genMetadata PoC — prompt classification into fixed categories. Not enabled on the + * default endpoint until deployment provides a Jlama model archive, {@code enable_remote_resources}, + * and {@code memory_size_mb >= 4096}. See {@code docs/development/gen-metadata-augment.md}. + */ + static final Augment.GenMetadata MS_COPILOT_GEN_METADATA_AUGMENT = Augment.GenMetadata.builder() + .jsonPath("$..body.content") + .prompt(""" + Return ONLY a JSON object with a single property "category" (no other keys). + Do NOT return a JSON Schema (never use type, properties, required, or enum as keys). + Examples: {"category":"Research and Ideation"} or {"category":"Excluded"} + + Classify the input text into exactly one category. Use the category name verbatim: + Email Drafting, General Content Drafting, Email Editing and Refinement, + General Content Editing and Refinement, Summarization and Synthesis, + Research and Ideation, Analysis and Data Work, Code Generation and Development, + Workflow and Task Management, Uncategorized, Excluded. + Use "Uncategorized" when the prompt is substantive but does not clearly fit a specific category. + Use "Excluded" for social niceties (greetings, thanks, pleasantries only) or prompts + that are too short or non-substantive to classify. + """) + .outputSchema(com.avaulta.gateway.rules.JsonSchemaFilter.builder() + .type("object") + .required(List.of("category")) + .properties(java.util.Map.of( + "category", com.avaulta.gateway.rules.JsonSchemaFilter.builder() + .type("string") + .enumValues(List.of( + "Email Drafting", + "General Content Drafting", + "Email Editing and Refinement", + "General Content Editing and Refinement", + "Summarization and Synthesis", + "Research and Ideation", + "Analysis and Data Work", + "Code Generation and Development", + "Workflow and Task Management", + "Uncategorized", + "Excluded")) + .build())) + .build()) + .build(); + static final Transform.Redact MS_COPILOT_CONTENT_REDACT = Transform.Redact.builder() .jsonPath("$..body.content") .jsonPath("$..attachments[*].content") diff --git a/java/core/src/main/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImpl.java b/java/core/src/main/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImpl.java index af1bbe682..ffaf2343d 100644 --- a/java/core/src/main/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImpl.java +++ b/java/core/src/main/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImpl.java @@ -33,6 +33,7 @@ import co.worklytics.psoxy.Pseudonymizer; import co.worklytics.psoxy.gateway.BulkModeConfig; import co.worklytics.psoxy.gateway.StorageEventRequest; +import co.worklytics.psoxy.impl.AugmentProcessor; import co.worklytics.psoxy.storage.BulkDataSanitizer; import dagger.assisted.Assisted; import dagger.assisted.AssistedInject; @@ -41,6 +42,7 @@ import lombok.NonNull; import lombok.Value; import lombok.extern.java.Log; +import org.apache.commons.lang3.ObjectUtils; @Log public class RecordBulkDataSanitizerImpl implements BulkDataSanitizer { @@ -59,6 +61,9 @@ public class RecordBulkDataSanitizerImpl implements BulkDataSanitizer { @Inject BulkModeConfig bulkModeConfig; + @Inject + AugmentProcessor augmentProcessor; + @Inject JsonSchemaFilterUtils jsonSchemaFilterUtils; @@ -108,7 +113,7 @@ public void sanitize(@NonNull StorageEventRequest request, Map record; while ((record = recordReader.readRecord()) != null) { try { - Map sanitized = applyTransforms(record, compiledTransforms); + Map sanitized = sanitizeRecord(record, compiledTransforms); recordWriter.writeRecord(sanitized); } catch (UnmatchedPseudonymization e) { log.warning("Skipped record due to UnmatchedPseudonymization: " + e.getPath()); @@ -208,7 +213,19 @@ private RecordRules.Format formatFromSourceObjectPath(String sourceObjectPath) { /** - * Apply transforms, then optional output schema filter. + * Apply augments then transforms to a single record. + */ + Map sanitizeRecord(Map document, + List> compiledTransforms) + throws UnmatchedPseudonymization { + if (ObjectUtils.isNotEmpty(rules.getAugments())) { + augmentProcessor.applyAugments(rules.getAugments(), document); + } + return applyTransforms(document, compiledTransforms); + } + + /** + * Apply the compiled transforms to the document, then optional output schema filter. * * @param document JSON "document object" * @param compiledTransforms ordered list of compiled transforms diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/AugmentProcessorTest.java b/java/core/src/test/java/co/worklytics/psoxy/impl/AugmentProcessorTest.java index aece216a6..c5dfeab79 100644 --- a/java/core/src/test/java/co/worklytics/psoxy/impl/AugmentProcessorTest.java +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/AugmentProcessorTest.java @@ -1,8 +1,14 @@ package co.worklytics.psoxy.impl; +import co.worklytics.psoxy.Warning; import com.avaulta.gateway.resources.ResourceService; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.avaulta.gateway.rules.JsonSchemaValidationUtils; import com.avaulta.gateway.rules.augments.Augment; +import com.avaulta.gateway.rules.augments.GenMetadataProcessor; import com.avaulta.gateway.rules.augments.SentenceMetadataProcessor; +import com.avaulta.gateway.rules.augments.UnavailableGenMetadataBackend; +import com.fasterxml.jackson.databind.ObjectMapper; import com.jayway.jsonpath.Configuration; import com.jayway.jsonpath.Option; import com.jayway.jsonpath.spi.json.JacksonJsonProvider; @@ -23,9 +29,11 @@ class AugmentProcessorTest { AugmentProcessor augmentProcessor; Configuration jsonConfiguration; + ObjectMapper objectMapper; @BeforeEach void setup() { + objectMapper = new ObjectMapper(); jsonConfiguration = Configuration.builder() .jsonProvider(new JacksonJsonProvider()) .mappingProvider(new JacksonMappingProvider()) @@ -33,7 +41,13 @@ void setup() { .build(); ResourceService noModels = path -> Optional.empty(); - augmentProcessor = new AugmentProcessor(jsonConfiguration, new SentenceMetadataProcessor(noModels)); + GenMetadataProcessor genMetadataProcessor = + new GenMetadataProcessor(new UnavailableGenMetadataBackend(), objectMapper); + augmentProcessor = new AugmentProcessor(jsonConfiguration, + new JsonSchemaValidationUtils(objectMapper), + objectMapper, + new SentenceMetadataProcessor(noModels), + genMetadataProcessor); } @Test @@ -295,10 +309,20 @@ void applyAugments_innerJsonPath_extractsFromEscapedJson() { Map resultAttachment = (Map) ((List) document.get("attachments")).get(0); assertFalse(resultAttachment.containsKey("+content.body[0].text:textDigest")); - String innerAugments = (String) resultAttachment.get("+content:textDigest"); + @SuppressWarnings("unchecked") + Map innerAugments = (Map) resultAttachment.get("+content:textDigest"); assertNotNull(innerAugments); - assertTrue(innerAugments.contains("\"text\":\"{\\\"length\\\":11,\\\"word_count\\\":2}\"")); - assertTrue(innerAugments.contains("\"text\":\"{\\\"length\\\":12,\\\"word_count\\\":2}\"")); + assertEquals(2, innerAugments.size()); + @SuppressWarnings("unchecked") + Map digest0 = (Map) innerAugments.get("body[0].text"); + @SuppressWarnings("unchecked") + Map digest1 = (Map) innerAugments.get("body[1].text"); + assertNotNull(digest0); + assertNotNull(digest1); + assertEquals(11, digest0.get("length")); + assertEquals(2, digest0.get("word_count")); + assertEquals(12, digest1.get("length")); + assertEquals(2, digest1.get("word_count")); } @Test @@ -345,4 +369,47 @@ void augmentPropertyName_usesConstants() { + AugmentProcessor.AUGMENT_SEPARATOR + "textDigest"; assertEquals("+content:textDigest", expected); } + + @SneakyThrows + @Test + void applyAugments_outputSchemaMismatch_omitsPropertyAndWarns() { + Augment.TextDigest augment = Augment.TextDigest.builder() + .jsonPath("$.body.content") + .outputSchema(JsonSchemaFilter.builder() + .type("object") + .required(List.of("category")) + .properties(Map.of( + "category", JsonSchemaFilter.builder().type("string").build())) + .build()) + .build(); + + Map body = new LinkedHashMap<>(); + body.put("content", "Hello world"); + Map document = new LinkedHashMap<>(); + document.put("body", body); + + List warnings = augmentProcessor.applyAugments(List.of(augment), document); + + @SuppressWarnings("unchecked") + Map resultBody = (Map) document.get("body"); + assertFalse(resultBody.containsKey("+content:textDigest")); + assertTrue(warnings.contains(Warning.AUGMENT_OUTPUT_SCHEMA_MISMATCH.asHttpHeaderCode())); + } + + @Test + void applyAugments_conflictDetection_reportsWarning() { + Augment.TextDigest augment = Augment.TextDigest.builder() + .jsonPath("$.body.content") + .build(); + + Map body = new LinkedHashMap<>(); + body.put("content", "Hello world"); + body.put("+content:someExisting", "conflict"); + Map document = new LinkedHashMap<>(); + document.put("body", body); + + List warnings = augmentProcessor.applyAugments(List.of(augment), document); + + assertTrue(warnings.contains(Warning.AUGMENT_CONFLICT_SKIPPED.asHttpHeaderCode())); + } } diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/WebhookSanitizerAugmentsTest.java b/java/core/src/test/java/co/worklytics/psoxy/impl/WebhookSanitizerAugmentsTest.java new file mode 100644 index 000000000..9a0439bb2 --- /dev/null +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/WebhookSanitizerAugmentsTest.java @@ -0,0 +1,91 @@ +package co.worklytics.psoxy.impl; + +import co.worklytics.psoxy.Pseudonymizer; +import co.worklytics.psoxy.gateway.HttpEventRequest; +import co.worklytics.psoxy.gateway.HttpEventRequestDto; +import co.worklytics.psoxy.gateway.ProcessedContent; +import co.worklytics.psoxy.utils.email.EmailAddressParser; +import com.avaulta.gateway.pseudonyms.impl.UrlSafeTokenPseudonymEncoder; +import com.avaulta.gateway.rules.JsonSchemaValidationUtils; +import com.avaulta.gateway.rules.WebhookCollectionRules; +import com.avaulta.gateway.rules.augments.Augment; +import com.avaulta.gateway.rules.transforms.Transform; +import com.avaulta.gateway.tokens.DeterministicTokenizationStrategy; +import com.avaulta.gateway.tokens.ReversibleTokenizationStrategy; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.jayway.jsonpath.Configuration; +import com.jayway.jsonpath.Option; +import com.jayway.jsonpath.spi.json.JacksonJsonProvider; +import com.jayway.jsonpath.spi.mapper.JacksonMappingProvider; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +class WebhookSanitizerAugmentsTest { + + WebhookSanitizerImpl sanitizer; + SanitizerUtils sanitizerUtils; + Configuration jsonConfiguration; + ObjectMapper objectMapper; + + @BeforeEach + void setUp() { + objectMapper = new ObjectMapper(); + jsonConfiguration = Configuration.builder() + .jsonProvider(new JacksonJsonProvider()) + .mappingProvider(new JacksonMappingProvider()) + .options(Option.SUPPRESS_EXCEPTIONS) + .build(); + sanitizerUtils = new SanitizerUtils( + jsonConfiguration, + mock(ReversibleTokenizationStrategy.class), + mock(UrlSafeTokenPseudonymEncoder.class), + mock(EmailAddressParser.class), + mock(ReversibleTokenizationStrategy.class), + mock(DeterministicTokenizationStrategy.class)); + + WebhookCollectionRules.WebhookEndpoint endpoint = WebhookCollectionRules.WebhookEndpoint.builder() + .augment(Augment.TextDigest.builder().jsonPath("$.content").build()) + .transform(Transform.Redact.ofPaths("$.content")) + .build(); + WebhookCollectionRules webhookRules = WebhookCollectionRules.builder() + .endpoint(endpoint) + .build(); + + sanitizer = new WebhookSanitizerImpl(webhookRules); + sanitizer.pseudonymizer = mock(Pseudonymizer.class); + sanitizer.emailAddressParser = mock(EmailAddressParser.class); + sanitizer.sanitizerUtils = sanitizerUtils; + sanitizer.jsonConfiguration = jsonConfiguration; + com.fasterxml.jackson.databind.ObjectMapper om = objectMapper; + sanitizer.augmentProcessor = new AugmentProcessor(jsonConfiguration, + new JsonSchemaValidationUtils(om), + om, + new com.avaulta.gateway.rules.augments.SentenceMetadataProcessor(path -> java.util.Optional.empty()), + new com.avaulta.gateway.rules.augments.GenMetadataProcessor( + new com.avaulta.gateway.rules.augments.UnavailableGenMetadataBackend(), om)); + } + + @Test + void sanitize_appliesAugmentsBeforeTransforms() { + String json = "{\"content\":\"hello world test\"}"; + HttpEventRequest request = HttpEventRequestDto.builder() + .headers(Map.of("Content-Type", List.of("application/json"))) + .body(json.getBytes(StandardCharsets.UTF_8)) + .build(); + + ProcessedContent result = sanitizer.sanitize(request); + String output = new String(result.getContent(), StandardCharsets.UTF_8); + + assertTrue(output.contains("\"+content:textDigest\"")); + assertTrue(output.contains("\"word_count\"")); + assertFalse(output.contains("\"content\":\"hello world test\"")); + } +} diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/gen/GenMetadataIntegrationSupport.java b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/GenMetadataIntegrationSupport.java new file mode 100644 index 000000000..140291030 --- /dev/null +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/GenMetadataIntegrationSupport.java @@ -0,0 +1,174 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.rules.JsonSchemaFilter; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * Opt-in helpers for local Jlama integration tests. + */ +final class GenMetadataIntegrationSupport { + + static final String ENABLE_PROPERTY = "psoxy.genMetadata.integration"; + static final String MODEL_CACHE_ENV = "PSOXY_GEN_MODEL_CACHE"; + static final String MODEL_ID_ENV = "PSOXY_GEN_MODEL"; + + private GenMetadataIntegrationSupport() { + } + + static boolean enabled() { + if ("true".equalsIgnoreCase(System.getProperty(ENABLE_PROPERTY))) { + return true; + } + String env = System.getenv("PSOXY_GEN_INTEGRATION_TEST"); + return "true".equalsIgnoreCase(env); + } + + static GenMetadataConfig defaultConfig() { + String modelId = Optional.ofNullable(System.getenv(MODEL_ID_ENV)) + .filter(s -> !s.isBlank()) + .orElse(GenMetadataConfig.DEFAULT_MODEL); + return GenMetadataConfig.builder() + .backend(GenMetadataConfig.BACKEND_LOCAL) + .modelId(modelId.trim()) + .timeoutSeconds(120) + .maxInputChars(4096) + .maxTokens(256) + .build(); + } + + static Path modelCacheDir(GenMetadataConfig config) { + String override = System.getenv(MODEL_CACHE_ENV); + if (override != null && !override.isBlank()) { + return Path.of(override.trim()); + } + return Path.of(System.getProperty("user.home"), ".jlama"); + } + + static boolean isModelMaterialized(GenMetadataConfig config, Path cacheDir) { + if (Files.isRegularFile(cacheDir.resolve(config.localModelCacheDirName()).resolve("config.json"))) { + return true; + } + if (Files.isRegularFile(cacheDir.resolve(config.getModelId().replace("/", "_")).resolve("config.json"))) { + return true; + } + if (Files.isRegularFile(cacheDir.resolve(config.getModelId()).resolve("config.json"))) { + return true; + } + try (var stream = Files.list(cacheDir)) { + return stream.anyMatch(path -> Files.isRegularFile(path.resolve("config.json"))); + } catch (Exception e) { + return false; + } + } + + /** MS Copilot PoC classification task (mirrors {@code PrebuiltSanitizerRules}). */ + static String copilotTaskPrompt() { + return """ + Return ONLY a JSON object with a single property "category" (no other keys). + Do NOT return a JSON Schema (never use type, properties, required, or enum as keys). + Examples: {"category":"Research and Ideation"} or {"category":"Excluded"} + + Classify the input text into exactly one category. Use the category name verbatim: + Email Drafting, General Content Drafting, Email Editing and Refinement, + General Content Editing and Refinement, Summarization and Synthesis, + Research and Ideation, Analysis and Data Work, Code Generation and Development, + Workflow and Task Management, Uncategorized, Excluded. + Use "Uncategorized" when the prompt is substantive but does not clearly fit a specific category. + Use "Excluded" for social niceties (greetings, thanks, pleasantries only) or prompts + that are too short or non-substantive to classify. + """; + } + + static JsonSchemaFilter copilotOutputSchema() { + return JsonSchemaFilter.builder() + .type("object") + .required(List.of("category")) + .properties(Map.of( + "category", JsonSchemaFilter.builder() + .type("string") + .enumValues(List.of( + "Email Drafting", + "General Content Drafting", + "Email Editing and Refinement", + "General Content Editing and Refinement", + "Summarization and Synthesis", + "Research and Ideation", + "Analysis and Data Work", + "Code Generation and Development", + "Workflow and Task Management", + "Uncategorized", + "Excluded")) + .build())) + .build(); + } + + static final List COPILOT_CATEGORIES = List.of( + "Email Drafting", + "General Content Drafting", + "Email Editing and Refinement", + "General Content Editing and Refinement", + "Summarization and Synthesis", + "Research and Ideation", + "Analysis and Data Work", + "Code Generation and Development", + "Workflow and Task Management", + "Uncategorized", + "Excluded"); + + /** + * Realistic short Copilot research prompt (typical interaction size). + */ + static String shortResearchIdeationPrompt() { + return """ + Help me brainstorm differentiated AI analytics workflows for enterprise buyers who care \ + about privacy and governance. Compare two competitors and suggest discovery questions \ + for customer interviews about research and ideation use cases. + """; + } + + /** + * Long Copilot user prompt (~3500 chars) that exceeded the misconfigured Jlama budget in live testing. + */ + static String longResearchIdeationPrompt() { + String paragraph = """ + Help me brainstorm a product strategy for our team analytics platform. I want to compare \ + how competitors position AI-assisted insights, identify three differentiated workflows we \ + could build this quarter, and outline risks if we rely too heavily on automated summaries \ + without human review. Include examples of enterprise buyers who care about privacy, \ + governance, and export controls. Suggest discovery questions I should ask customer interviews \ + next week about how they use Copilot-style assistants for research and ideation sessions. + """; + StringBuilder sb = new StringBuilder(); + while (sb.length() < 3500) { + sb.append(paragraph); + if (sb.length() < 3500) { + sb.append('\n'); + } + } + return sb.substring(0, 3500); + } + + static String promptOfLength(int chars) { + String seed = "Analyze quarterly revenue trends and recommend actions. "; + StringBuilder sb = new StringBuilder(); + while (sb.length() < chars) { + sb.append(seed); + } + return sb.substring(0, chars); + } + + static boolean messageChainContains(Throwable t, String fragment) { + while (t != null) { + if (t.getMessage() != null && t.getMessage().contains(fragment)) { + return true; + } + t = t.getCause(); + } + return false; + } +} diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBudgetTest.java b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBudgetTest.java new file mode 100644 index 000000000..77b75a879 --- /dev/null +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/GenMetadataPromptBudgetTest.java @@ -0,0 +1,60 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class GenMetadataPromptBudgetTest { + + private GenMetadataPromptBudget promptBudget; + + @BeforeEach + void setUp() { + promptBudget = new GenMetadataPromptBudget(new ObjectMapper()); + } + + @Test + void fitInputData_respectsMaxInputChars() { + JsonSchemaFilter schema = JsonSchemaFilter.builder() + .type("object") + .properties(Map.of("category", JsonSchemaFilter.builder().type("string").build())) + .build(); + String input = "x".repeat(5000); + String fitted = promptBudget.fitInputData(input, "Classify", schema, 1024, 8192, 256); + assertTrue(fitted.length() <= 1024); + } + + @Test + void fitInputData_reservesOutputTokensWithinContext() { + JsonSchemaFilter schema = JsonSchemaFilter.builder() + .type("object") + .properties(Map.of("category", JsonSchemaFilter.builder().type("string").build())) + .build(); + String taskPrompt = "Classify"; + String input = "x".repeat(10_000); + int contextLength = 512; + int maxOutputTokens = 256; + String fitted = promptBudget.fitInputData( + input, taskPrompt, schema, 10_000, contextLength, maxOutputTokens); + assertTrue(fitted.length() < input.length()); + int contextChars = (contextLength - maxOutputTokens - 64) * 4; + assertTrue(fitted.length() <= contextChars); + } + + @Test + void fitInputData_leavesShortInputUnchanged() { + JsonSchemaFilter schema = JsonSchemaFilter.builder() + .type("object") + .properties(Map.of("category", JsonSchemaFilter.builder().type("string").build())) + .build(); + String input = "short prompt"; + String fitted = promptBudget.fitInputData(input, "Classify", schema, 4096, 8192, 256); + assertEquals(input, fitted); + } +} diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendConcurrencyTest.java b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendConcurrencyTest.java new file mode 100644 index 000000000..cc0427d28 --- /dev/null +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendConcurrencyTest.java @@ -0,0 +1,70 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.resources.ResourceService; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies single-flight model load behavior (without requiring a real model file). + */ +class LangChain4jGenMetadataBackendConcurrencyTest { + + @Test + void resolveModel_singleFlightOnConcurrentMiss() throws InterruptedException { + AtomicInteger loadAttempts = new AtomicInteger(); + ResourceService resourceService = objectPath -> { + loadAttempts.incrementAndGet(); + try { + Thread.sleep(100); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return java.util.Optional.empty(); + }; + + GenMetadataConfig config = GenMetadataConfig.builder() + .backend(GenMetadataConfig.BACKEND_LOCAL) + .modelId("test-model") + .timeoutSeconds(5) + .maxInputChars(1024) + .maxTokens(64) + .build(); + + LangChain4jGenMetadataBackend backend = new LangChain4jGenMetadataBackend( + config, new ObjectMapper(), new GenMetadataPromptBudget(new ObjectMapper()), + new GenMetadataChatModelFactory(resourceService)); + + int threads = 8; + ExecutorService pool = Executors.newFixedThreadPool(threads); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + + for (int i = 0; i < threads; i++) { + pool.submit(() -> { + try { + start.await(); + backend.resolveModel(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + done.countDown(); + } + }); + } + + start.countDown(); + assertTrue(done.await(10, TimeUnit.SECONDS)); + pool.shutdown(); + + assertEquals(1, loadAttempts.get()); + } +} diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendIntegrationTest.java b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendIntegrationTest.java new file mode 100644 index 000000000..d10b482cd --- /dev/null +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendIntegrationTest.java @@ -0,0 +1,216 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.resources.ResourceService; +import com.avaulta.gateway.rules.augments.GenMetadataBackend; +import com.avaulta.gateway.rules.augments.GenMetadataAugmentException; +import com.avaulta.gateway.rules.augments.GenMetadataProcessor; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.langchain4j.data.message.ChatMessage; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +import java.nio.file.Path; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +/** + * Local Jlama integration tests for genMetadata token-budget regressions seen in live Copilot PoC runs. + * + *

Skipped unless {@code -Dpsoxy.genMetadata.integration=true} (or {@code PSOXY_GEN_INTEGRATION_TEST=true}) + * and a Jlama model is already materialized under {@code ~/.jlama} or {@code PSOXY_GEN_MODEL_CACHE}. + * + *

Run from {@code java/} (not the repo root): + * {@code mvn test -pl core -am -Pgen-metadata-integration} + * + *

Equivalent manual invocation: + * {@code mvn test -pl core -am -Dtest=LangChain4jGenMetadataBackendIntegrationTest -Dpsoxy.genMetadata.integration=true -Dsurefire.failIfNoSpecifiedTests=false} + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class LangChain4jGenMetadataBackendIntegrationTest { + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final GenMetadataPromptBudget PROMPT_BUDGET = new GenMetadataPromptBudget(OBJECT_MAPPER); + private static final ResourceService NO_REMOTE_RESOURCES = path -> java.util.Optional.empty(); + private static final GenMetadataChatModelFactory CHAT_MODEL_FACTORY = + new GenMetadataChatModelFactory(NO_REMOTE_RESOURCES); + + private GenMetadataConfig config; + private Path modelCacheDir; + + @BeforeAll + void requireLocalModel() throws Exception { + assumeTrue(GenMetadataIntegrationSupport.enabled(), + "Set -Dpsoxy.genMetadata.integration=true to run local Jlama tests"); + config = GenMetadataIntegrationSupport.defaultConfig(); + modelCacheDir = GenMetadataIntegrationSupport.modelCacheDir(config); + try { + CHAT_MODEL_FACTORY.buildLocal(config, modelCacheDir); + } catch (Exception e) { + assumeTrue(false, "Failed to load Jlama model '" + config.getModelId() + "' from " + + modelCacheDir + ": " + e.getMessage() + + " (set " + GenMetadataIntegrationSupport.MODEL_CACHE_ENV + " if cached elsewhere)"); + } + } + + @Test + void misconfiguredJlamaMaxTokens_rejectsLongCopilotPrompt() throws Exception { + ChatModel broken = CHAT_MODEL_FACTORY.buildLocal( + config, modelCacheDir, config.getMaxTokens()); + + String inputJson = OBJECT_MAPPER.writeValueAsString( + GenMetadataIntegrationSupport.longResearchIdeationPrompt()); + List messages = GenMetadataPromptBuilder.toMessages( + GenMetadataIntegrationSupport.copilotTaskPrompt(), + GenMetadataIntegrationSupport.copilotOutputSchema(), + inputJson, + OBJECT_MAPPER); + + Exception failure = null; + try { + broken.chat(ChatRequest.builder().messages(messages).build()); + } catch (Exception e) { + failure = e; + } + assertNotNull(failure, "expected misconfigured Jlama maxTokens to reject long Copilot prompt"); + assertTrue( + GenMetadataIntegrationSupport.messageChainContains(failure, "Prompt exceeds max tokens"), + "expected Jlama 'Prompt exceeds max tokens' failure, got: " + failure); + } + + @Test + void misconfiguredJlamaMaxTokens_yieldsUnparseableJsonForSomePromptLengths() throws Exception { + GenMetadataBackend brokenBackend = new MisconfiguredJlamaBackend( + config, modelCacheDir, OBJECT_MAPPER, CHAT_MODEL_FACTORY, config.getMaxTokens()); + GenMetadataProcessor processor = new GenMetadataProcessor(brokenBackend, OBJECT_MAPPER, 4096); + + String taskPrompt = GenMetadataIntegrationSupport.copilotTaskPrompt(); + JsonSchemaFilter schema = GenMetadataIntegrationSupport.copilotOutputSchema(); + + boolean sawUnparseable = false; + for (int inputChars : new int[] {200, 400, 600, 800, 1000, 1200, 1600, 2000}) { + String input = GenMetadataIntegrationSupport.promptOfLength(inputChars); + try { + processor.process(taskPrompt, schema, input); + } catch (GenMetadataAugmentException e) { + if (e.getCode() == GenMetadataAugmentException.Code.INFERENCE_FAILED) { + sawUnparseable = true; + break; + } + throw e; + } + } + assertTrue(sawUnparseable, + "expected misconfigured maxTokens to sometimes return truncated JSON that fails parsing " + + "(live failure: \"classification\": \"Research and Ide\")"); + } + + @Test + void productionBackend_inferenceSucceedsForLongCopilotPrompt() { + LangChain4jGenMetadataBackend backend = new LangChain4jGenMetadataBackend( + config, OBJECT_MAPPER, PROMPT_BUDGET, CHAT_MODEL_FACTORY, modelCacheDir); + String inputJson = serializeQuoted(GenMetadataIntegrationSupport.longResearchIdeationPrompt()); + Object raw = backend.generate( + GenMetadataIntegrationSupport.copilotTaskPrompt(), + GenMetadataIntegrationSupport.copilotOutputSchema(), + inputJson); + assertNotNull(raw, + "fixed backend must not fail with 'Prompt exceeds max tokens' on long Copilot prompts"); + } + + @Test + void productionBackend_classifiesShortResearchPrompt() { + LangChain4jGenMetadataBackend backend = new LangChain4jGenMetadataBackend( + config, OBJECT_MAPPER, PROMPT_BUDGET, CHAT_MODEL_FACTORY, modelCacheDir); + GenMetadataProcessor processor = new GenMetadataProcessor(backend, OBJECT_MAPPER, 4096); + + @SuppressWarnings("unchecked") + Map result = (Map) processor.process( + GenMetadataIntegrationSupport.copilotTaskPrompt(), + GenMetadataIntegrationSupport.copilotOutputSchema(), + GenMetadataIntegrationSupport.shortResearchIdeationPrompt()); + + assertFalse(result.containsKey("type") && "object".equals(result.get("type")), + "model echoed output schema instead of classification: " + result); + + String category = result.values().stream() + .filter(String.class::isInstance) + .map(String.class::cast) + .filter(GenMetadataIntegrationSupport.COPILOT_CATEGORIES::contains) + .findFirst() + .orElse(null); + assertNotNull(category, + "expected parseable JSON containing a valid Copilot category, got: " + result); + } + + private static String serializeQuoted(String text) { + try { + return OBJECT_MAPPER.writeValueAsString(text); + } catch (Exception e) { + throw new IllegalStateException(e); + } + } + + /** + * Simulates pre-fix backend: passes {@code PSOXY_GEN_MAX_TOKENS} through to Jlama as the KV ceiling + * and does not apply {@link GenMetadataPromptBudget}. + */ + static final class MisconfiguredJlamaBackend implements GenMetadataBackend { + + private final GenMetadataConfig config; + private final Path modelCacheDir; + private final ObjectMapper objectMapper; + private final int jlamaMaxTokens; + private final GenMetadataChatModelFactory chatModelFactory; + private volatile ChatModel chatModel; + + MisconfiguredJlamaBackend(GenMetadataConfig config, Path modelCacheDir, + ObjectMapper objectMapper, GenMetadataChatModelFactory chatModelFactory, + int jlamaMaxTokens) { + this.config = config; + this.modelCacheDir = modelCacheDir; + this.objectMapper = objectMapper; + this.chatModelFactory = chatModelFactory; + this.jlamaMaxTokens = jlamaMaxTokens; + } + + @Override + public Object generate(String taskPrompt, JsonSchemaFilter outputSchema, String inputData) { + try { + ChatModel model = chatModel(); + List messages = GenMetadataPromptBuilder.toMessages( + taskPrompt, outputSchema, inputData, objectMapper); + ChatResponse response = model.chat(ChatRequest.builder().messages(messages).build()); + if (response == null || response.aiMessage() == null) { + return null; + } + return response.aiMessage().text(); + } catch (Exception e) { + return null; + } + } + + private ChatModel chatModel() throws Exception { + ChatModel local = chatModel; + if (local == null) { + synchronized (this) { + local = chatModel; + if (local == null) { + local = chatModelFactory.buildLocal(config, modelCacheDir, jlamaMaxTokens); + chatModel = local; + } + } + } + return local; + } + } +} diff --git a/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendTimeoutTest.java b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendTimeoutTest.java new file mode 100644 index 000000000..51f5680dc --- /dev/null +++ b/java/core/src/test/java/co/worklytics/psoxy/impl/gen/LangChain4jGenMetadataBackendTimeoutTest.java @@ -0,0 +1,55 @@ +package co.worklytics.psoxy.impl.gen; + +import com.avaulta.gateway.resources.ResourceService; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.fasterxml.jackson.databind.ObjectMapper; +import dev.langchain4j.data.message.ChatMessage; +import dev.langchain4j.data.message.UserMessage; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; +import org.junit.jupiter.api.Test; + +import java.nio.file.Files; +import java.util.List; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.assertNull; + +class LangChain4jGenMetadataBackendTimeoutTest { + + @Test + void chatWithTimeout_returnsNullWhenModelExceedsLimit() throws Exception { + GenMetadataConfig config = GenMetadataConfig.builder() + .backend(GenMetadataConfig.BACKEND_LOCAL) + .modelId("test-model") + .timeoutSeconds(1) + .maxInputChars(4096) + .maxTokens(64) + .build(); + + ResourceService noModels = path -> Optional.empty(); + LangChain4jGenMetadataBackend backend = new LangChain4jGenMetadataBackend( + config, + new ObjectMapper(), + new GenMetadataPromptBudget(new ObjectMapper()), + new GenMetadataChatModelFactory(noModels), + Files.createTempDirectory("psoxy-jlama-timeout-test")); + + ChatModel slowModel = new ChatModel() { + @Override + public ChatResponse chat(ChatRequest request) { + try { + Thread.sleep(2_000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return ChatResponse.builder().build(); + } + }; + + List messages = List.of(UserMessage.from("hello")); + ChatResponse response = backend.chatWithTimeout(slowModel, messages); + assertNull(response); + } +} diff --git a/java/core/src/test/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImplTest.java b/java/core/src/test/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImplTest.java index a294ca243..4da4b9658 100644 --- a/java/core/src/test/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImplTest.java +++ b/java/core/src/test/java/co/worklytics/psoxy/storage/impl/RecordBulkDataSanitizerImplTest.java @@ -144,6 +144,31 @@ void base() { } + @SneakyThrows + @Test + void augmentsAppliedBeforeRedact() { + this.setUpWithRules("---\n" + + "format: \"NDJSON\"\n" + + "augments:\n" + + "- !\n" + + " jsonPaths:\n" + + " - \"$.content\"\n" + + "transforms:\n" + + "- redact: \"content\"\n"); + + String input = "{\"content\":\"hello world test\"}\n"; + + storageHandler.handle(BulkDataTestUtils.request("export/file.ndjson"), + BulkDataTestUtils.transform(rules), + () -> new ByteArrayInputStream(input.getBytes(StandardCharsets.UTF_8)), + outputStreamSupplier); + + String output = new String(outputStream.toByteArray(), StandardCharsets.UTF_8); + assertTrue(output.contains("\"+content:textDigest\"")); + assertTrue(output.contains("\"word_count\"")); + assertTrue(output.contains("\"content\":null")); + } + @Test void noTransforms() throws IOException { this.setUpWithRules("---\n" + diff --git a/java/core/src/test/java/co/worklytics/test/MockModules.java b/java/core/src/test/java/co/worklytics/test/MockModules.java index a51919841..03207c8f4 100644 --- a/java/core/src/test/java/co/worklytics/test/MockModules.java +++ b/java/core/src/test/java/co/worklytics/test/MockModules.java @@ -10,6 +10,7 @@ import co.worklytics.psoxy.gateway.output.OutputLocation; import com.avaulta.gateway.resources.ResourceService; import com.avaulta.gateway.rules.augments.SentenceMetadataProcessor; +import com.avaulta.gateway.rules.augments.SentenceMetadataProcessor; import co.worklytics.psoxy.rules.RESTRules; import co.worklytics.psoxy.utils.RandomNumberGenerator; import com.avaulta.gateway.rules.BulkDataRules; @@ -248,6 +249,13 @@ static SentenceMetadataProcessor sentenceMetadataProcessor( @Named("OpenNlp") ResourceService openNlpResourceService) { return new SentenceMetadataProcessor(openNlpResourceService); } + + @Provides + @Singleton + @Named("ForGenMetadata") + static ResourceService forGenMetadataResourceService() { + return new NoOpResourceService(); + } } @Module diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaFilter.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaFilter.java index 717d09d6e..cc3d0f243 100644 --- a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaFilter.java +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaFilter.java @@ -25,7 +25,6 @@ @JsonInclude(JsonInclude.Include.NON_NULL) @JsonIgnoreProperties({ "title", - "required", // not relevant to 'filter' use case "additionalProperties", // not currently supported "$schema", // not helpful in filter use-case, although maybe should include in future }) @@ -59,6 +58,17 @@ public Map getProperties() { //only applicable if type==array JsonSchemaFilter items; + /** + * Allowed values for string validation (JSON Schema {@code enum} keyword). + */ + @JsonProperty("enum") + List enumValues; + + /** + * Required property names when validating augment output (predicate use case). + */ + List required; + // @JsonProperties("$defs") on property, lombok getter/setter don't seem to do the right thing // get java.lang.IllegalArgumentException: Unrecognized field "definitions" (class com.avaulta.gateway.rules.SchemaRuleUtils$JsonSchema) // // when calling objectMapper.convertValue(((JsonNode) schema), JsonSchema.class);?? diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaValidationUtils.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaValidationUtils.java index 6881d0a6e..b16f2615c 100644 --- a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaValidationUtils.java +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/JsonSchemaValidationUtils.java @@ -29,6 +29,8 @@ @RequiredArgsConstructor public class JsonSchemaValidationUtils { + private static final int MAX_LOG_OUTPUT_CHARS = 2000; + final ObjectMapper objectMapper; private final static JsonMetaSchema metaSchema = JsonMetaSchema.builder(JsonMetaSchema.getV202012()) // .format(PatternFormat.of("pseudonym", @@ -45,6 +47,26 @@ public class JsonSchemaValidationUtils { .maximumSize(100) .build(CacheLoader.from(this::getJsonSchema)); + private final LoadingCache jsonSchemaFilterCache = + CacheBuilder.newBuilder() + .maximumSize(100) + .build(CacheLoader.from(this::getJsonSchemaFromFilter)); + + @SneakyThrows + public boolean validateJsonBySchema(String jsonString, JsonSchemaFilter schema) { + if (schema == null) { + return true; + } + JsonNode deserialized = objectMapper.readTree(jsonString); + JsonSchema jsonSchema = jsonSchemaFilterCache.get(schema); + Set validationMessages = jsonSchema.validate(deserialized); + if (!validationMessages.isEmpty()) { + log.warning("Validation failed for augment output: " + validationMessages + + "; output=" + truncateForLog(jsonString)); + } + return validationMessages.isEmpty(); + } + @SneakyThrows public boolean validateJsonBySchema(String jsonString, com.avaulta.gateway.rules.JsonSchema schema) { @@ -153,6 +175,21 @@ private JsonSchema getJsonSchema(com.avaulta.gateway.rules.JsonSchema schema) { return jsonSchemaFactory.getSchema(schemaNode); } + private JsonSchema getJsonSchemaFromFilter(JsonSchemaFilter schema) { + JsonNode schemaNode = objectMapper.valueToTree(schema); + return jsonSchemaFactory.getSchema(schemaNode); + } + + private static String truncateForLog(String value) { + if (value == null) { + return "null"; + } + if (value.length() <= MAX_LOG_OUTPUT_CHARS) { + return value; + } + return value.substring(0, MAX_LOG_OUTPUT_CHARS) + "... (" + value.length() + " chars total)"; + } + com.avaulta.gateway.rules.JsonSchema rewritePseudonymToPattern( com.avaulta.gateway.rules.JsonSchema schema) { if (schema.getFormat() != null diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/RecordRules.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/RecordRules.java index b8d2c2660..bafbe1988 100644 --- a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/RecordRules.java +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/RecordRules.java @@ -1,8 +1,10 @@ package com.avaulta.gateway.rules; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Optional; +import com.avaulta.gateway.rules.augments.Augment; import com.avaulta.gateway.rules.transforms.RecordTransform; import com.avaulta.gateway.rules.transforms.RecordTransformDeserializer; import com.fasterxml.jackson.annotation.JsonIgnore; @@ -62,6 +64,15 @@ public enum Format { @JsonDeserialize(contentUsing = RecordTransformDeserializer.class) List transforms; + /** + * Augments to compute and inject as synthetic sibling properties, run before transforms. + * + * @see Augments Design Doc + */ + @JsonInclude(JsonInclude.Include.NON_EMPTY) + @Singular + List augments; + /** * Optional JSON schema allow-list applied to each record after transforms. * Properties not declared in the schema (including unknown future event payload fields) are @@ -82,6 +93,7 @@ public Optional getOutputSchemaFilterOptional() { */ public RecordRules() { this.transforms = Collections.emptyList(); + this.augments = new ArrayList<>(); } //setter to ensure we get a List, even when coming through jackson @@ -93,5 +105,13 @@ public void setTransforms(List transforms) { } } + public void setAugments(List augments) { + if (augments == null) { + this.augments = Collections.emptyList(); + } else { + this.augments = augments; + } + } + } diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/WebhookCollectionRules.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/WebhookCollectionRules.java index df0514d1c..8ef65f3ee 100644 --- a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/WebhookCollectionRules.java +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/WebhookCollectionRules.java @@ -5,6 +5,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; +import com.avaulta.gateway.rules.augments.Augment; import com.avaulta.gateway.rules.transforms.Transform; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonPropertyOrder; @@ -73,6 +74,16 @@ public static class WebhookEndpoint implements Serializable { //TODO: schemaFilter for request payload? avoids risk of unexpected fields included in request payload + /** + * Augments to compute and inject as synthetic sibling properties, run before transforms. + * + * @see Augments Design Doc + */ + @Setter + @JsonInclude(value = JsonInclude.Include.NON_EMPTY) + @Singular + List augments; + /** * a list of transforms to apply to the request payload */ @@ -87,6 +98,7 @@ public static class WebhookEndpoint implements Serializable { * 2) Explicit instantiation of @Singular fields required to avoid Lombok warnings about ignored default values. */ public WebhookEndpoint() { + this.augments = new ArrayList<>(); this.transforms = new ArrayList<>(); } } diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/Augment.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/Augment.java index bb7f27d14..3aa4845f0 100644 --- a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/Augment.java +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/Augment.java @@ -29,6 +29,7 @@ @JsonSubTypes({ @JsonSubTypes.Type(value = Augment.TextDigest.class, name = "textDigest"), @JsonSubTypes.Type(value = Augment.SentenceMetadata.class, name = "sentenceMetadata"), + @JsonSubTypes.Type(value = Augment.GenMetadata.class, name = "genMetadata"), }) @SuperBuilder(toBuilder = true) @AllArgsConstructor @@ -236,4 +237,39 @@ static Set signalWords(List configured, List defaults) { return Set.copyOf(result); } } + + /** + * BETA: Generates structured metadata via LangChain4j (Jlama local; Bedrock/Vertex planned). + * Requires {@link #outputSchema} and {@link #prompt}; model/backend selection is deployment config. + */ + @SuperBuilder(toBuilder = true) + @NoArgsConstructor + @AllArgsConstructor + @Getter + @EqualsAndHashCode(callSuper = true) + public static class GenMetadata extends Augment { + + /** + * Task instruction passed to the generative backend. + */ + @JsonInclude(JsonInclude.Include.NON_NULL) + String prompt; + + @JsonIgnore + @Override + public String getFunctionName() { + return "genMetadata"; + } + + @Override + public Object compute(Object input) { + // Computed at runtime by AugmentProcessor via injected GenMetadataProcessor. + return null; + } + + @Override + protected boolean canEqual(Object other) { + return other instanceof GenMetadata; + } + } } diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/AugmentValidation.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/AugmentValidation.java new file mode 100644 index 000000000..138465a4f --- /dev/null +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/AugmentValidation.java @@ -0,0 +1,80 @@ +package com.avaulta.gateway.rules.augments; + +import com.avaulta.gateway.rules.Endpoint; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.avaulta.gateway.rules.RecordRules; +import com.avaulta.gateway.rules.WebhookCollectionRules; +import org.apache.commons.lang3.StringUtils; + +import java.util.ArrayList; +import java.util.List; + +/** + * Validates augment rules at load time. + */ +public final class AugmentValidation { + + private AugmentValidation() {} + + public static void validateAugments(Iterable augments) { + List errors = new ArrayList<>(); + collectAugmentErrors(augments, errors); + throwIfErrors(errors); + } + + public static void validateEndpoints(Iterable endpoints) { + List errors = new ArrayList<>(); + if (endpoints != null) { + for (Endpoint endpoint : endpoints) { + collectAugmentErrors(endpoint.getAugments(), errors); + } + } + throwIfErrors(errors); + } + + public static void validateRecordRules(RecordRules rules) { + if (rules != null) { + validateAugments(rules.getAugments()); + } + } + + public static void validateWebhookEndpoints( + Iterable endpoints) { + List errors = new ArrayList<>(); + if (endpoints != null) { + for (WebhookCollectionRules.WebhookEndpoint endpoint : endpoints) { + collectAugmentErrors(endpoint.getAugments(), errors); + } + } + throwIfErrors(errors); + } + + private static void collectAugmentErrors(Iterable augments, List errors) { + if (augments != null) { + for (Augment augment : augments) { + validateAugment(augment, errors); + } + } + } + + private static void throwIfErrors(List errors) { + if (!errors.isEmpty()) { + throw new IllegalArgumentException( + "Invalid augment configuration: " + String.join("; ", errors)); + } + } + + static void validateAugment(Augment augment, List errors) { + if (augment instanceof Augment.GenMetadata gen) { + if (StringUtils.isBlank(gen.getPrompt())) { + errors.add("genMetadata augment requires non-blank prompt"); + } + if (gen.getOutputSchema() == null) { + errors.add("genMetadata augment requires outputSchema"); + } + if (gen.getJsonPaths() == null || gen.getJsonPaths().isEmpty()) { + errors.add("genMetadata augment requires at least one jsonPath"); + } + } + } +} diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataAugmentException.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataAugmentException.java new file mode 100644 index 000000000..0408651e1 --- /dev/null +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataAugmentException.java @@ -0,0 +1,29 @@ +package com.avaulta.gateway.rules.augments; + +/** + * Runtime exception from genMetadata processing; translated to + * {@link co.worklytics.psoxy.impl.AugmentProcessingException} in the core module. + */ +public class GenMetadataAugmentException extends RuntimeException { + + public enum Code { + UNAVAILABLE, + INFERENCE_FAILED, + } + + private final Code code; + + public GenMetadataAugmentException(Code code, String message) { + super(message); + this.code = code; + } + + public GenMetadataAugmentException(Code code, String message, Throwable cause) { + super(message, cause); + this.code = code; + } + + public Code getCode() { + return code; + } +} diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataBackend.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataBackend.java new file mode 100644 index 000000000..c15a1ff01 --- /dev/null +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataBackend.java @@ -0,0 +1,21 @@ +package com.avaulta.gateway.rules.augments; + +import com.avaulta.gateway.rules.JsonSchemaFilter; + +/** + * Pluggable backend for {@link Augment.GenMetadata} inference. + * Implementations use LangChain4j {@code ChatModel} adapters (Jlama for local embedded models; + * Bedrock and Vertex planned for cloud). + */ +public interface GenMetadataBackend { + + /** + * Generate JSON metadata for the given input. + * + * @param taskPrompt augment rule task prompt + * @param outputSchema required output schema predicate + * @param inputData JSON-serialized source value + * @return parsed JSON object (typically a Map), or raw JSON string from the model + */ + Object generate(String taskPrompt, JsonSchemaFilter outputSchema, String inputData); +} diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataProcessor.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataProcessor.java new file mode 100644 index 000000000..5f5901208 --- /dev/null +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/GenMetadataProcessor.java @@ -0,0 +1,235 @@ +package com.avaulta.gateway.rules.augments; + +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.avaulta.gateway.rules.JsonSchemaValidationUtils; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.commons.lang3.StringUtils; + +import java.time.Instant; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; + +/** + * Facade for {@link Augment.GenMetadata} — delegates to a configured {@link GenMetadataBackend}. + */ +public class GenMetadataProcessor { + + private static final Logger log = Logger.getLogger(GenMetadataProcessor.class.getName()); + + private static final int DEFAULT_MAX_INPUT_CHARS = 4096; + private static final int DEFAULT_MAX_ATTEMPTS = 2; + private static final int MAX_LOG_OUTPUT_CHARS = 2000; + + private final GenMetadataBackend backend; + private final ObjectMapper objectMapper; + private final JsonSchemaValidationUtils jsonSchemaValidationUtils; + private final int maxInputChars; + private final int maxAttempts; + + public GenMetadataProcessor(GenMetadataBackend backend, ObjectMapper objectMapper, int maxInputChars, + int maxAttempts, JsonSchemaValidationUtils jsonSchemaValidationUtils) { + this.backend = backend != null ? backend : new UnavailableGenMetadataBackend(); + this.objectMapper = objectMapper != null ? objectMapper : new ObjectMapper(); + this.jsonSchemaValidationUtils = jsonSchemaValidationUtils != null + ? jsonSchemaValidationUtils + : new JsonSchemaValidationUtils(this.objectMapper); + this.maxInputChars = maxInputChars > 0 ? maxInputChars : DEFAULT_MAX_INPUT_CHARS; + this.maxAttempts = maxAttempts > 0 ? maxAttempts : DEFAULT_MAX_ATTEMPTS; + } + + public GenMetadataProcessor(GenMetadataBackend backend, ObjectMapper objectMapper, int maxInputChars) { + this(backend, objectMapper, maxInputChars, DEFAULT_MAX_ATTEMPTS, null); + } + + public GenMetadataProcessor(GenMetadataBackend backend, ObjectMapper objectMapper) { + this(backend, objectMapper, DEFAULT_MAX_INPUT_CHARS); + } + + /** + * Compute genMetadata output for a single augment invocation. + */ + public Object compute(Augment.GenMetadata augment, Object input) { + return process(augment.getPrompt(), augment.getOutputSchema(), input); + } + + public Object process(String taskPrompt, JsonSchemaFilter outputSchema, Object input) { + if (StringUtils.isBlank(taskPrompt) || outputSchema == null) { + throw new GenMetadataAugmentException(GenMetadataAugmentException.Code.UNAVAILABLE, + "genMetadata missing prompt or outputSchema"); + } + String inputJson = serializeInput(input); + if (inputJson == null) { + throw new GenMetadataAugmentException(GenMetadataAugmentException.Code.UNAVAILABLE, + "genMetadata input empty or not serializable"); + } + try { + for (int attempt = 1; attempt <= maxAttempts; attempt++) { + if (attempt > 1) { + log.info("genMetadata inference retry attempt " + attempt + " of " + maxAttempts); + } + Map parsed = inferOnce(taskPrompt, outputSchema, inputJson); + if (parsed != null && validatesOutputSchema(parsed, outputSchema)) { + return parsed; + } + } + throw new GenMetadataAugmentException(GenMetadataAugmentException.Code.INFERENCE_FAILED, + "genMetadata output failed schema validation after " + maxAttempts + " attempt(s)"); + } catch (GenMetadataAugmentException e) { + throw e; + } catch (Exception e) { + log.log(Level.WARNING, "genMetadata inference failed", e); + throw new GenMetadataAugmentException(GenMetadataAugmentException.Code.INFERENCE_FAILED, + "genMetadata inference failed", e); + } + } + + private Map inferOnce(String taskPrompt, JsonSchemaFilter outputSchema, String inputJson) + throws Exception { + Instant startedAt = Instant.now(); + long startedNanos = System.nanoTime(); + log.info("genMetadata augment inference call started at " + startedAt); + Object raw; + try { + raw = backend.generate(taskPrompt, outputSchema, inputJson); + } finally { + long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedNanos); + log.info("genMetadata augment inference call completed in " + elapsedMs + "ms"); + } + if (raw instanceof String rawText) { + log.info("genMetadata raw backend response: " + truncateForLog(rawText)); + } else if (raw != null) { + log.info("genMetadata backend returned non-string type: " + + raw.getClass().getSimpleName()); + } + Map parsed = parseModelJson(raw); + if (parsed == null) { + log.warning("genMetadata backend returned unparseable output"); + return null; + } + log.info("genMetadata parsed output keys: " + parsed.keySet() + + "; value=" + truncateForLog(serializeForLog(parsed))); + return parsed; + } + + private boolean validatesOutputSchema(Map parsed, JsonSchemaFilter outputSchema) { + try { + String json = objectMapper.writeValueAsString(parsed); + return jsonSchemaValidationUtils.validateJsonBySchema(json, outputSchema); + } catch (Exception e) { + log.log(Level.WARNING, "Failed to validate genMetadata output schema", e); + return false; + } + } + + String serializeInput(Object input) { + if (input == null) { + return null; + } + try { + if (input instanceof String text) { + if (text.isEmpty()) { + return null; + } + String truncated = truncate(text); + return objectMapper.writeValueAsString(truncated); + } + String serialized = objectMapper.writeValueAsString(input); + return truncateSerialized(serialized); + } catch (Exception e) { + log.log(Level.WARNING, "Failed to serialize genMetadata input", e); + return null; + } + } + + Map parseModelJson(Object raw) { + if (raw == null) { + return null; + } + if (raw instanceof Map map) { + return toSortedMap(map); + } + if (raw instanceof String response) { + String json = extractJsonObject(response); + if (json == null) { + return null; + } + try { + @SuppressWarnings("unchecked") + Map map = objectMapper.readValue(json, Map.class); + return new TreeMap<>(map); + } catch (Exception e) { + log.log(Level.WARNING, "Failed to parse genMetadata JSON response: " + json, e); + return null; + } + } + return null; + } + + private String serializeForLog(Map parsed) { + try { + return objectMapper.writeValueAsString(parsed); + } catch (Exception e) { + return String.valueOf(parsed); + } + } + + private static String truncateForLog(String value) { + if (value == null) { + return "null"; + } + if (value.length() <= MAX_LOG_OUTPUT_CHARS) { + return value; + } + return value.substring(0, MAX_LOG_OUTPUT_CHARS) + "... (" + value.length() + " chars total)"; + } + + static String extractJsonObject(String response) { + if (StringUtils.isBlank(response)) { + return null; + } + String trimmed = response.trim(); + int start = trimmed.indexOf('{'); + int end = trimmed.lastIndexOf('}'); + if (start >= 0 && end > start) { + return trimmed.substring(start, end + 1); + } + return trimmed.startsWith("{") ? trimmed : null; + } + + private TreeMap toSortedMap(Map raw) { + TreeMap sorted = new TreeMap<>(); + raw.forEach((k, v) -> { + if (k != null) { + sorted.put(String.valueOf(k), v); + } + }); + return sorted; + } + + private String truncate(String text) { + if (text.length() <= maxInputChars) { + return text; + } + return text.substring(0, maxInputChars); + } + + /** + * Truncate JSON-serialized non-string inputs (outer quotes included in length budget). + */ + private String truncateSerialized(String serialized) { + if (serialized.length() <= maxInputChars) { + return serialized; + } + if (serialized.startsWith("\"") && serialized.endsWith("\"")) { + int contentBudget = maxInputChars - 2; + if (contentBudget <= 0) { + return "\"\""; + } + return "\"" + serialized.substring(1, 1 + contentBudget) + "\""; + } + return serialized.substring(0, maxInputChars); + } +} diff --git a/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/UnavailableGenMetadataBackend.java b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/UnavailableGenMetadataBackend.java new file mode 100644 index 000000000..77c4188e8 --- /dev/null +++ b/java/gateway-core/src/main/java/com/avaulta/gateway/rules/augments/UnavailableGenMetadataBackend.java @@ -0,0 +1,15 @@ +package com.avaulta.gateway.rules.augments; + +import com.avaulta.gateway.rules.JsonSchemaFilter; + +/** + * Default backend when genMetadata is not configured or the selected backend is unsupported. + */ +public class UnavailableGenMetadataBackend implements GenMetadataBackend { + + @Override + public Object generate(String taskPrompt, JsonSchemaFilter outputSchema, String inputData) { + throw new GenMetadataAugmentException(GenMetadataAugmentException.Code.UNAVAILABLE, + "genMetadata backend is not available"); + } +} diff --git a/java/gateway-core/src/test/java/com/avaulta/gateway/rules/augments/AugmentValidationTest.java b/java/gateway-core/src/test/java/com/avaulta/gateway/rules/augments/AugmentValidationTest.java new file mode 100644 index 000000000..28d150425 --- /dev/null +++ b/java/gateway-core/src/test/java/com/avaulta/gateway/rules/augments/AugmentValidationTest.java @@ -0,0 +1,52 @@ +package com.avaulta.gateway.rules.augments; + +import com.avaulta.gateway.rules.Endpoint; +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.avaulta.gateway.rules.RecordRules; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; + +class AugmentValidationTest { + + @Test + void validateGenMetadata_requiresPromptAndSchema() { + Augment.GenMetadata invalid = Augment.GenMetadata.builder() + .jsonPath("$.body.content") + .build(); + + assertThrows(IllegalArgumentException.class, + () -> AugmentValidation.validateEndpoints(List.of( + Endpoint.builder().augment(invalid).pathTemplate("/test").build()))); + } + + @Test + void validateGenMetadata_valid() { + Augment.GenMetadata valid = Augment.GenMetadata.builder() + .jsonPath("$.body.content") + .prompt("Classify the prompt") + .outputSchema(JsonSchemaFilter.builder() + .type("object") + .properties(java.util.Map.of( + "category", JsonSchemaFilter.builder().type("string").build())) + .build()) + .build(); + + assertDoesNotThrow(() -> AugmentValidation.validateEndpoints(List.of( + Endpoint.builder().augment(valid).pathTemplate("/test").build()))); + } + + @Test + void validateRecordRules_genMetadataRequiresPromptAndSchema() { + Augment.GenMetadata invalid = Augment.GenMetadata.builder() + .jsonPath("$.content") + .build(); + + assertThrows(IllegalArgumentException.class, + () -> AugmentValidation.validateRecordRules( + RecordRules.builder().augment(invalid).build())); + } +} diff --git a/java/gateway-core/src/test/java/com/avaulta/gateway/rules/augments/GenMetadataProcessorTest.java b/java/gateway-core/src/test/java/com/avaulta/gateway/rules/augments/GenMetadataProcessorTest.java new file mode 100644 index 000000000..540258d19 --- /dev/null +++ b/java/gateway-core/src/test/java/com/avaulta/gateway/rules/augments/GenMetadataProcessorTest.java @@ -0,0 +1,120 @@ +package com.avaulta.gateway.rules.augments; + +import com.avaulta.gateway.rules.JsonSchemaFilter; +import com.avaulta.gateway.rules.JsonSchemaValidationUtils; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class GenMetadataProcessorTest { + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + private static JsonSchemaFilter categorySchema() { + return JsonSchemaFilter.builder() + .type("object") + .required(List.of("category")) + .properties(Map.of("category", JsonSchemaFilter.builder().type("string").build())) + .build(); + } + + @Test + void process_delegatesToBackend() { + JsonSchemaFilter schema = categorySchema(); + + GenMetadataProcessor processor = new GenMetadataProcessor( + (taskPrompt, outputSchema, inputData) -> { + TreeMap result = new TreeMap<>(); + result.put("category", "Excluded"); + return result; + }, + OBJECT_MAPPER, + 4096); + + @SuppressWarnings("unchecked") + Map out = (Map) processor.process( + "Classify the prompt", schema, "hello"); + assertEquals("Excluded", out.get("category")); + } + + @Test + void process_retriesOnSchemaMismatch() { + JsonSchemaFilter schema = categorySchema(); + AtomicInteger calls = new AtomicInteger(); + + GenMetadataProcessor processor = new GenMetadataProcessor( + (taskPrompt, outputSchema, inputData) -> { + if (calls.incrementAndGet() == 1) { + return """ + {"type":"object","properties":{"category":{"type":"string"}},"required":["category"]} + """; + } + return "{\"category\":\"Excluded\"}"; + }, + OBJECT_MAPPER, + 4096, + 2, + new JsonSchemaValidationUtils(OBJECT_MAPPER)); + + @SuppressWarnings("unchecked") + Map out = (Map) processor.process( + "Classify the prompt", schema, "hello"); + assertEquals("Excluded", out.get("category")); + assertEquals(2, calls.get()); + } + + @Test + void process_failsAfterExhaustingRetries() { + JsonSchemaFilter schema = categorySchema(); + + GenMetadataProcessor processor = new GenMetadataProcessor( + (taskPrompt, outputSchema, inputData) -> """ + {"type":"object","properties":{"category":{"type":"string"}},"required":["category"]} + """, + OBJECT_MAPPER, + 4096, + 1, + new JsonSchemaValidationUtils(OBJECT_MAPPER)); + + assertThrows(GenMetadataAugmentException.class, + () -> processor.process("Classify the prompt", schema, "hello")); + } + + @Test + void parseModelJson_extractsFromMarkdownFences() { + GenMetadataProcessor processor = new GenMetadataProcessor( + new UnavailableGenMetadataBackend(), OBJECT_MAPPER); + Object out = processor.parseModelJson(""" + ```json + {"category": "Email Drafting"} + ``` + """); + @SuppressWarnings("unchecked") + Map map = (Map) out; + assertEquals("Email Drafting", map.get("category")); + } + + @Test + void process_throwsWhenPromptMissing() { + GenMetadataProcessor processor = new GenMetadataProcessor( + new UnavailableGenMetadataBackend(), OBJECT_MAPPER); + assertThrows(GenMetadataAugmentException.class, + () -> processor.process(null, JsonSchemaFilter.builder().type("object").build(), "x")); + } + + @Test + void serializeInput_truncatesNonStringValues() throws Exception { + GenMetadataProcessor processor = new GenMetadataProcessor( + new UnavailableGenMetadataBackend(), OBJECT_MAPPER, 20); + String serialized = processor.serializeInput(Map.of("text", "a".repeat(100))); + assertTrue(serialized.length() <= 20); + } +} diff --git a/java/impl/cmd-line/src/main/java/co/worklytics/psoxy/CmdLineModule.java b/java/impl/cmd-line/src/main/java/co/worklytics/psoxy/CmdLineModule.java index 77def2001..af1bca30f 100644 --- a/java/impl/cmd-line/src/main/java/co/worklytics/psoxy/CmdLineModule.java +++ b/java/impl/cmd-line/src/main/java/co/worklytics/psoxy/CmdLineModule.java @@ -6,8 +6,9 @@ import co.worklytics.psoxy.gateway.ProxyConfigProperty; import co.worklytics.psoxy.gateway.SecretStore; import co.worklytics.psoxy.gateway.impl.BlindlyOptimisticLockService; - import co.worklytics.psoxy.gateway.impl.GoogleCloudPlatformServiceAccountKeyAuthStrategy; +import co.worklytics.psoxy.gateway.impl.NoOpResourceService; +import com.avaulta.gateway.resources.ResourceService; import com.avaulta.gateway.tokens.impl.Sha256DeterministicTokenizationStrategy; import com.google.cloud.secretmanager.v1.AccessSecretVersionResponse; import com.google.cloud.secretmanager.v1.SecretManagerServiceClient; @@ -16,13 +17,17 @@ import dagger.Provides; import lombok.AllArgsConstructor; import lombok.NoArgsConstructor; +import javax.inject.Named; import javax.inject.Singleton; import java.io.IOException; @NoArgsConstructor @AllArgsConstructor @Module( - includes = CmdLineModule.Bindings.class + includes = { + CmdLineModule.Bindings.class, + ResourceServiceBindingsModule.class, + } ) public class CmdLineModule { @@ -40,6 +45,17 @@ static ApiModeConfig apiModeConfig(ConfigService configService) { return ApiModeConfig.fromConfigService(configService); } + /** CLI has no cloud bucket; local FS under {@link ResourceService#DEFAULT_LOCAL_RESOURCE_PATH} only. */ + @Provides @Singleton @Named("Remote") + static ResourceService remoteResourceService() { + return new NoOpResourceService(); + } + + @Provides @Singleton @Named("SharedRemote") + static ResourceService sharedRemoteResourceService() { + return new NoOpResourceService(); + } + @Provides @Singleton SecretStore secretStore(CommandLineConfigServiceFactory factory) { return factory.create(args); diff --git a/java/pom.xml b/java/pom.xml index 987ed9ef3..e47c40661 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -29,6 +29,9 @@ 2.9.0 2.5.9 false + + 1.15.0 + 1.15.0-beta25 1.82 3.5.3 true @@ -113,4 +116,24 @@ + + + gen-metadata-integration + + + + + + org.apache.maven.plugins + maven-surefire-plugin + + false + + + + + + + + diff --git a/tools/fetch-gen-metadata-model.sh b/tools/fetch-gen-metadata-model.sh new file mode 100755 index 000000000..df6e07605 --- /dev/null +++ b/tools/fetch-gen-metadata-model.sh @@ -0,0 +1,184 @@ +#!/usr/bin/env bash +# Downloads a Jlama-compatible HuggingFace model (SafeTensors layout), zips it for remote loading, +# and optionally uploads the archive to a remote bucket for cloud genMetadata deployments. +# +# Usage: +# ./tools/fetch-gen-metadata-model.sh +# ./tools/fetch-gen-metadata-model.sh s3://BUCKET/SHARED_RESOURCE_PATH/ [MODEL_ID] +# ./tools/fetch-gen-metadata-model.sh gs://BUCKET/SHARED_RESOURCE_PATH/ [MODEL_ID] +# ./tools/fetch-gen-metadata-model.sh --from-dir /path/to/model-dir s3://BUCKET/PREFIX/ [MODEL_ID] +# +# MODEL_ID defaults to PSOXY_GEN_MODEL or tjake/Llama-3.2-1B-Instruct-JQ4. +# Upload path: PREFIX/llm/{MODEL_ID with / replaced by __}.zip +# +# Requires: zip. For download: huggingface-cli (pip install huggingface_hub). For upload: aws or gsutil. + +set -euo pipefail + +# Use semantic colors dynamically based on terminal capability +if [ -t 1 ] && command -v tput >/dev/null 2>&1; then + ERR=$(tput setaf 1) + SUCCESS=$(tput setaf 2) + WARN=$(tput setaf 3) + INFO=$(tput setaf 4) + NC=$(tput sgr0) +else + ERR='\033[0;31m' + SUCCESS='\033[0;32m' + WARN='\033[1;33m' + INFO='\033[0;34m' + NC='\033[0m' +fi + +DEFAULT_MODEL_ID="tjake/Llama-3.2-1B-Instruct-JQ4" +FROM_DIR="" +REMOTE_URI="" +MODEL_ID="${PSOXY_GEN_MODEL:-$DEFAULT_MODEL_ID}" + +usage() { + printf "Usage: %s [--from-dir DIR] [s3://BUCKET/PREFIX/ | gs://BUCKET/PREFIX/] [MODEL_ID]\n" "$(basename "$0")" + printf " Downloads (or uses --from-dir), zips to .build/gen-metadata-models/llm/, optionally uploads.\n" +} + +while [ $# -gt 0 ]; do + case "$1" in + -h|--help) + usage + exit 0 + ;; + --from-dir) + shift + FROM_DIR="${1:?--from-dir requires a path}" + shift + ;; + s3://*|gs://*) + REMOTE_URI="$1" + shift + ;; + *) + if [ -z "${REMOTE_URI}" ] && [[ "$1" == s3://* || "$1" == gs://* ]]; then + REMOTE_URI="$1" + else + MODEL_ID="$1" + fi + shift + ;; + esac +done + +if [ -n "${REMOTE_URI}" ] && [[ "${REMOTE_URI}" != s3://* && "${REMOTE_URI}" != gs://* ]]; then + printf "${ERR}Remote destination must be an s3:// or gs:// URI.${NC}\n" >&2 + usage >&2 + exit 1 +fi + +archive_stem="${MODEL_ID//\//__}" +archive_name="${archive_stem}.zip" + +cd "$(dirname "$0")/.." +DEST_ROOT=".build/gen-metadata-models" +MODEL_DIR="${DEST_ROOT}/${archive_stem}" +ARCHIVE_PATH="${DEST_ROOT}/llm/${archive_name}" + +mkdir -p "${DEST_ROOT}/llm" + +if [ -n "${FROM_DIR}" ]; then + if [ ! -f "${FROM_DIR}/config.json" ]; then + printf "${ERR}--from-dir must contain config.json at its root (Jlama SafeTensors layout).${NC}\n" >&2 + exit 1 + fi + rm -rf "${MODEL_DIR}" + cp -R "${FROM_DIR}" "${MODEL_DIR}" +else + if [ -f "${MODEL_DIR}/config.json" ]; then + printf "${WARN}Using existing model directory %s${NC}\n" "${MODEL_DIR}" + else + if ! command -v huggingface-cli >/dev/null 2>&1; then + printf "${ERR}huggingface-cli not found. Install with: pip install huggingface_hub${NC}\n" >&2 + printf "${ERR}Or pass --from-dir pointing at a local SafeTensors model directory.${NC}\n" >&2 + exit 1 + fi + rm -rf "${MODEL_DIR}" + mkdir -p "${MODEL_DIR}" + printf "${INFO}Downloading %s from HuggingFace to %s...${NC}\n" "${MODEL_ID}" "${MODEL_DIR}" + huggingface-cli download "${MODEL_ID}" --local-dir "${MODEL_DIR}" + fi +fi + +if [ ! -f "${MODEL_DIR}/config.json" ]; then + printf "${ERR}Model directory missing config.json: %s${NC}\n" "${MODEL_DIR}" >&2 + exit 1 +fi + +ARCHIVE_ABS="$(cd "${DEST_ROOT}/llm" && pwd)/${archive_name}" +printf "${INFO}Creating archive %s...${NC}\n" "${ARCHIVE_ABS}" +rm -f "${ARCHIVE_ABS}" +( + cd "${MODEL_DIR}" + zip -qr "${ARCHIVE_ABS}" . +) + +printf "${SUCCESS}Created %s${NC}\n" "${ARCHIVE_ABS}" +ARCHIVE_PATH="${ARCHIVE_ABS}" + +if [ -z "${REMOTE_URI}" ]; then + exit 0 +fi + +normalize_prefix() { + local prefix="$1" + if [ -n "$prefix" ] && [ "${prefix: -1}" != "/" ]; then + prefix="${prefix}/" + fi + printf '%s' "$prefix" +} + +upload_to_s3() { + local uri="$1" + local path="${uri#s3://}" + local bucket="${path%%/*}" + local prefix="" + + if [ "$path" != "$bucket" ]; then + prefix="${path#*/}" + fi + prefix="$(normalize_prefix "$prefix")" + + if ! command -v aws >/dev/null 2>&1; then + printf "${ERR}aws CLI is required to upload to S3.${NC}\n" >&2 + exit 1 + fi + + local object_key="${prefix}llm/${archive_name}" + printf "${INFO}Uploading to s3://%s/%s...${NC}\n" "$bucket" "$object_key" + aws s3 cp "${ARCHIVE_PATH}" "s3://${bucket}/${object_key}" +} + +upload_to_gcs() { + local uri="$1" + local path="${uri#gs://}" + local bucket="${path%%/*}" + local prefix="" + + if [ "$path" != "$bucket" ]; then + prefix="${path#*/}" + fi + prefix="$(normalize_prefix "$prefix")" + + if ! command -v gsutil >/dev/null 2>&1; then + printf "${ERR}gsutil is required to upload to GCS.${NC}\n" >&2 + exit 1 + fi + + local object_name="${prefix}llm/${archive_name}" + printf "${INFO}Uploading to gs://%s/%s...${NC}\n" "$bucket" "$object_name" + gsutil cp "${ARCHIVE_PATH}" "gs://${bucket}/${object_name}" +} + +if [[ "${REMOTE_URI}" == s3://* ]]; then + upload_to_s3 "${REMOTE_URI}" +elif [[ "${REMOTE_URI}" == gs://* ]]; then + upload_to_gcs "${REMOTE_URI}" +fi + +printf "${SUCCESS}genMetadata model archive uploaded to %s${NC}\n" "${REMOTE_URI}"