Add image-gen module: Spring AI ImageModel options on the wire, a moderation gate, and safe storage of generated images in MinIO (sniffed content type, byte and pixel caps, server-chosen keys, capped URL fetcher); per-provider cost table dated 2026-10-28. ImageModel and ModerationModel are hand-written fakes; MinIO is real via Testcontainers; no image provider called.
Co-Authored-By: Claude Sonnet 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_017ic4b7rnC7jppGGdHGESvy
This commit is contained in:
Claude Sonnet 5.5
@@ -30,5 +30,6 @@ Runnable companion code for the Spring AI articles on [ankurm.com](https://ankur
|
||||
| [`rerank/`](rerank) | A bi-encoder (all-MiniLM-L6-v2) retrieves candidates, a cross-encoder (ms-marco-MiniLM-L-6-v2) rescores them, and a `DocumentPostProcessor` keeps the top 5 for `RetrievalAugmentationAdvisor`. Both models run in-process through ONNX Runtime 1.31.0 with DJL tokenizers 0.39.0 (no database, no API key). 60 questions over 94 passages: hit@1 41 to 49 and MRR 0.78 to 0.87, almost all of it on exact error codes, with paraphrases slightly worse; pool size, batch size, threads and int8 against latency on 2 vCPUs; and five traps reproduced (missing `token_type_ids`, swapped pair order, logits are not probabilities, truncation, a pool that is too small). Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Reranking RAG Results in Java: Cross-Encoder Reranker with ONNX Runtime Before the LLM](https://ankurm.com/reranking-rag-results-java-cross-encoder-onnx-runtime-spring-ai/) |
|
||||
| [`audio/`](audio) | A meeting-minutes generator: `OpenAiAudioTranscriptionModel` -> `ChatClient.entity(Minutes.class)` -> `OpenAiAudioSpeechModel`. Splitting a 30 MB recording at a silence under the 25 MB upload limit (and what a blind split does), SDK retries that re-upload the file, the 4096-character speech limit, three replies that fail to parse, and a cost model with formulas. The audio endpoints are a local mock and the chat model a stub; no real model was called. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Speech-to-Text and Text-to-Speech with Spring AI: Build a Meeting-Minutes Generator](https://ankurm.com/spring-ai-speech-to-text-text-to-speech-meeting-minutes/) |
|
||||
| [`graphrag-neo4j/`](graphrag-neo4j) | Entities and relations extracted with `ChatClient.entity(...)` into a typed Neo4j graph with a vector index on entity names (Testcontainers, `neo4j:2026.10.1-community`; also run on 5.26), then multi-hop retrieval by span linking, variable-length Cypher paths and path ranking, against plain vector RAG over the same 35 chunks and 9 questions: complete evidence 7/9 for the graph, 5/9 for vector k=8, 2/9 for k=4, with the traps reproduced (duplicate entity under an alias, an embedding merge that joins the wrong incidents, a reply with a sentence before the JSON fence, a hostile relation type, hop explosion). **The extraction step is a rule-based stand-in, not a language model**; Neo4j, embeddings and Cypher are real. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [GraphRAG in Java with Neo4j and Spring AI](https://ankurm.com/graphrag-java-neo4j-spring-ai-knowledge-graph/) |
|
||||
| [`image-gen/`](image-gen) | `ImageModel` and `ImagePrompt` with `OpenAiImageOptions` (what the real `OpenAiImageModel` puts on the wire, a per-call options object that silently resets the model, a `List<ImageMessage>` prompt that throws), a prompt gate with `ModerationModel` that fails closed, and safe storage in MinIO: sniffed content type, byte and pixel caps, server-chosen keys, a fetcher for provider URLs, presigned reads. A per-image price table from the vendors' pages dated 2026-10-28. No image provider called: the `ImageModel` and `ModerationModel` are hand-written fakes and a mock server; the MinIO tests use a real server through Testcontainers. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Image Generation and Vision in Spring AI: ImageModel, Prompt Options and Safe Storage](https://ankurm.com/spring-ai-image-generation-imagemodel-s3-minio-safe-storage/) |
|
||||
|
||||
Upgrading from Spring AI 1.x: [migration guide](https://ankurm.com/spring-ai-1-to-2-migration-guide/).
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
target/
|
||||
@@ -0,0 +1,45 @@
|
||||
# image-gen
|
||||
|
||||
Companion code for [Image Generation and Vision in Spring AI: ImageModel, Prompt Options and Safe Storage](https://ankurm.com/spring-ai-image-generation-imagemodel-s3-minio-safe-storage/), part of the [Spring AI series](../README.md) on ankurm.com. It is the generation half of [Multimodal Spring AI: Extract Structured Data from Images](https://ankurm.com/multimodal-spring-ai-extract-structured-data-from-images-receipts-java-records/).
|
||||
|
||||
A prompt goes through our own checks, into `ImageModel`, and the bytes that come back are verified and written to an S3-compatible bucket under a key we choose.
|
||||
|
||||
**No image provider was called.** The tests use two hand-written fakes -- `ScriptedImageModel` (an `ImageModel` that returns what it was told to) and `ScriptedModerationModel` -- and, for the wire tests, the real `OpenAiImageModel` against a local mock of `/v1/images/generations` (`FakeImagesServer`) that returns a fixed 64x64 PNG whatever the prompt says. Nothing here measures image quality, latency, or whether a real model would refuse a prompt. **The storage is real:** the `ImageStore` tests run against a MinIO server in Docker through Testcontainers (the image is `chainguard/minio`, an upstream MinIO build, because `minio/minio` and `quay.io/minio/minio` could not be pulled from the machine these were captured on; the exact version is in `output/05-store-minio.txt`). Prices in `CostModel` were copied from the providers' pages on 2026-10-28 and are not checked by any test.
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| Spring Boot | 4.1.2 (parent; GA on Maven Central 2026-10-22) |
|
||||
| Spring AI | 2.0.2 (`spring-ai-model`, `spring-ai-openai`, `spring-ai-stability-ai`, `spring-ai-google-genai-image`) |
|
||||
| AWS SDK for Java | 2.57.5 (`s3`) |
|
||||
| Testcontainers | 2.0.5 |
|
||||
| Java | 25 |
|
||||
| Run | `scripts/run-all.sh` (about 20 s, needs Docker for the MinIO tests) |
|
||||
|
||||
## Source
|
||||
|
||||
| File | What it is |
|
||||
|---|---|
|
||||
| [`ImageService`](src/main/java/com/ankurm/imagegen/ImageService.java) | prompt gate -> `ImageModel` -> bytes -> verify -> store; one bad image in a batch stores nothing |
|
||||
| [`PromptGate`](src/main/java/com/ankurm/imagegen/PromptGate.java) | length and control-character checks, then `ModerationModel`; fails closed when moderation is down |
|
||||
| [`OpenAiOptionsFactory`](src/main/java/com/ankurm/imagegen/OpenAiOptionsFactory.java) | size / quality / n checked against allow-lists, then `OpenAiImageOptions` with the model set explicitly |
|
||||
| [`ImageSniffer`](src/main/java/com/ankurm/imagegen/ImageSniffer.java), [`ImageVerifier`](src/main/java/com/ankurm/imagegen/ImageVerifier.java), [`ImagePolicy`](src/main/java/com/ankurm/imagegen/ImagePolicy.java) | magic-number sniffing, size cap, pixel cap read from the header, a real decode, allow-list of types |
|
||||
| [`RemoteFetcher`](src/main/java/com/ankurm/imagegen/RemoteFetcher.java) | downloads a provider-returned URL: https and listed hosts only, no redirects, hard byte cap |
|
||||
| [`ImageStore`](src/main/java/com/ankurm/imagegen/ImageStore.java) | S3/MinIO writer: server-chosen key, pinned `Content-Type`, `Content-Disposition: attachment`, presigned GET |
|
||||
| [`CostModel`](src/main/java/com/ankurm/imagegen/CostModel.java) | per-image prices with the date they were read and the page they came from |
|
||||
|
||||
## Captured output (`output/`)
|
||||
|
||||
| File | Shows |
|
||||
|---|---|
|
||||
| `01-api-shape.txt` | class names, portable vs provider options, what is *not* there (moderation option, image moderation) |
|
||||
| `02-wire-requests.txt` | JSON body sent by the real `OpenAiImageModel` for six option shapes, including the model reset and the `ImagePrompt(List)` failure |
|
||||
| `03-wire-response.txt` | what `ImageResponse` carries back (b64 or URL, revised prompt, no usage) |
|
||||
| `04-wire-refusal.txt` | a provider 400 `content_policy_violation` arrives as an SDK exception |
|
||||
| `05-store-minio.txt`, `05b-presigned.txt` | the object as MinIO holds it; presigned URL behaviour |
|
||||
| `06-verifier.txt`, `07-trailing-bytes.txt`, `08-base64-cap.txt` | twelve inputs against the verifier; the limit of sniffing; base64 refused before decoding |
|
||||
| `09-fetcher.txt` | redirect, scheme, host list, endless body |
|
||||
| `10-keys.txt`, `11-naive-vs-pinned.txt` | owner and id validation; trusting the sender's `Content-Type` |
|
||||
| `12-service-happy.txt` .. `15-url-results.txt` | the service end to end; refusals before any paid call; a batch with one bad image; URL results |
|
||||
| `16-cost.txt` | the per-image price table |
|
||||
|
||||
Keys, ports and timestamps differ between runs; everything else is stable.
|
||||
@@ -0,0 +1,25 @@
|
||||
# Spring AI 2.0.2 image classes (read from the jars)
|
||||
|
||||
interface : org.springframework.ai.image.ImageModel
|
||||
single method : [call]
|
||||
request : org.springframework.ai.image.ImagePrompt
|
||||
OpenAI : org.springframework.ai.openai.OpenAiImageModel
|
||||
Stability AI : org.springframework.ai.stabilityai.StabilityAiImageModel
|
||||
moderation interface : org.springframework.ai.moderation.ModerationModel
|
||||
moderation OpenAI : org.springframework.ai.openai.OpenAiModerationModel
|
||||
|
||||
portable ImageOptions getters : [getHeight, getModel, getN, getResponseFormat, getStyle, getWidth]
|
||||
OpenAiImageOptions getters : [getApiKey, getBaseUrl, getCredential, getCustomHeaders, getDeploymentName, getHeight, getMaxRetries, getMicrosoftDeploymentName, getMicrosoftFoundryServiceVersion, getModel, getN, getOrganizationId, getProxy, getQuality, getResponseFormat, getSize, getStyle, getTimeout, getUser, getWidth, isGitHubModels, isMicrosoftFoundry]
|
||||
StabilityAiImageOptions : [getCfgScale, getClipGuidancePreset, getHeight, getModel, getN, getResponseFormat, getSampler, getSeed, getSteps, getStyle, getStylePreset, getWidth]
|
||||
GoogleGenAiImageOptions : [getAspectRatio, getHeight, getImageSize, getLabels, getMaxOutputTokens, getModel, getN, getOutputCompressionQuality, getOutputMimeType, getPersonGeneration, getResponseFormat, getSafetyFilterLevel, getSeed, getStyle, getTemperature, getTopK, getTopP, getWidth]
|
||||
|
||||
OpenAiImageOptions.DEFAULT_IMAGE_MODEL = gpt-image-1-mini
|
||||
OpenAiModerationOptions.DEFAULT_MODERATION_MODEL = omni-moderation-latest
|
||||
OpenAiImageOptions has a getModeration() : false
|
||||
ModerationMessage getters : [getText] (text only, no media)
|
||||
Stability default engine / base URL : stable-diffusion-v1-6 @ https://api.stability.ai/v1
|
||||
GoogleGenAiImageModelName : [gemini-2.5-flash-image, gemini-3-pro-image, gemini-3.1-flash-image]
|
||||
Google SafetyFilterLevel : [BLOCK_LOW_AND_ABOVE, BLOCK_MEDIUM_AND_ABOVE, BLOCK_ONLY_HIGH, BLOCK_NONE, SAFETY_FILTER_LEVEL_UNSPECIFIED]
|
||||
Google PersonGeneration : [DONT_ALLOW, ALLOW_ADULT, ALLOW_ALL, PERSON_GENERATION_UNSPECIFIED]
|
||||
AbstractOpenAiOptions.DEFAULT_MAX_RETRIES (image options extend it) = 3
|
||||
portable builder gives : DefaultImageModelOptions
|
||||
@@ -0,0 +1,21 @@
|
||||
# What OpenAiImageModel sends (mock server; bodies recorded as received)
|
||||
|
||||
1. ImagePrompt("a red bicycle"), model default only
|
||||
POST /v1/images/generations
|
||||
{"prompt":"a red bicycle","model":"gpt-image-test"}
|
||||
|
||||
2. OpenAiImageOptions size=1536x1024 quality=high n=2
|
||||
{"prompt":"a red bicycle","model":"gpt-image-1-mini","n":2,"quality":"high","size":"1536x1024"}
|
||||
|
||||
3. portable ImageOptionsBuilder width=1024 height=1536
|
||||
{"prompt":"a red bicycle","model":"gpt-image-test","size":"1024x1536"}
|
||||
|
||||
4. OpenAiImageOptions user=user-8f2c
|
||||
{"prompt":"a red bicycle","model":"gpt-image-1-mini","user":"user-8f2c"}
|
||||
|
||||
5. OpenAiImageOptions responseFormat=b64_json
|
||||
{"prompt":"a red bicycle","model":"gpt-image-1-mini","response_format":"b64_json"}
|
||||
|
||||
6a. ImagePrompt(List<ImageMessage>) with no options: IllegalArgumentException: image options cannot be null
|
||||
6b. same two messages (weights 1.0 and -1.0) with empty OpenAiImageOptions
|
||||
{"prompt":"a red bicycle","model":"gpt-image-1-mini"}
|
||||
@@ -0,0 +1,12 @@
|
||||
# What comes back from OpenAiImageModel (mock server)
|
||||
|
||||
results : 1
|
||||
Image.getUrl() : null
|
||||
Image.getB64Json() len : 10500 chars
|
||||
metadata type : OpenAiImageGenerationMetadata
|
||||
revised prompt : a rewritten prompt
|
||||
response created : 1760000000
|
||||
response usage : EmptyUsage, total tokens 0 (the mock body carried usage.total_tokens=1068)
|
||||
|
||||
URL mode: Image.getUrl() = http://127.0.0.1:43039/files/ok.png
|
||||
URL mode: Image.getB64Json() = null
|
||||
@@ -0,0 +1,5 @@
|
||||
# A provider-side refusal (400 content_policy_violation from the mock)
|
||||
|
||||
exception class : com.openai.errors.BadRequestException
|
||||
message (first 200): 400: Your request was rejected as a result of our safety system.
|
||||
requests received : 1 (maxRetries=0)
|
||||
@@ -0,0 +1,13 @@
|
||||
# ImageStore on a real MinIO server
|
||||
|
||||
image : chainguard/minio:latest
|
||||
minio --version : minio version RELEASE.2026-10-23T20-06-51Z (commit-id=76c06ff1d2cc844050fdeb1d54bf8fa70db72278)
|
||||
Server header : MinIO
|
||||
stored key : generated/user-42/4cf346e4b7008b1874935d7969e8ccb1.png
|
||||
content type : image/png, 6363 bytes, 64x48
|
||||
sha256 : 7254899daa04ebf862c59d32d2d3280bb179e6c61899505c43aed169c6197bcb
|
||||
HEAD Content-Type : image/png
|
||||
HEAD Content-Disposition : attachment
|
||||
HEAD Cache-Control : private, max-age=3600
|
||||
HEAD metadata : {owner=user-42, sha256=7254899daa04ebf862c59d32d2d3280bb179e6c61899505c43aed169c6197bcb}
|
||||
read back equal : true
|
||||
@@ -0,0 +1,7 @@
|
||||
# Presigned GET (5 minutes) against MinIO
|
||||
|
||||
signed URL, parameters : X-Amz-Algorithm=... & X-Amz-Date=... & X-Amz-SignedHeaders=host & X-Amz-Credential=... & X-Amz-Expires=300 & X-Amz-Signature=...
|
||||
GET signed URL : 200, 742 bytes, Content-Type image/jpeg
|
||||
Content-Disposition : attachment
|
||||
same URL, other owner : 403
|
||||
same path, no signature : 403
|
||||
@@ -0,0 +1,15 @@
|
||||
# ImageVerifier decisions (cap 8 MiB, 16,000,000 pixels, png/jpeg/webp)
|
||||
|
||||
input verdict detail
|
||||
real PNG 64x48 ACCEPT image/png 64x48
|
||||
real JPEG 64x48 ACCEPT image/jpeg 64x48
|
||||
WebP header only (sniffed, not decoded) ACCEPT image/webp
|
||||
SVG with a script REJECT not-an-image: no known image signature in the first bytes
|
||||
HTML page REJECT not-an-image: no known image signature in the first bytes
|
||||
PDF REJECT not-an-image: no known image signature in the first bytes
|
||||
PNG signature + garbage REJECT undecodable: IIOException
|
||||
valid PNG + 40 bytes of HTML appended ACCEPT image/png 16x16
|
||||
PNG header claiming 40000x40000 REJECT too-many-pixels: 40000x40000 = 1600000000 > cap 16000000
|
||||
GIF (sniffed but not on the allow list) REJECT type-not-allowed: image/gif
|
||||
empty REJECT not-an-image: no known image signature in the first bytes
|
||||
9 MB of PNG-looking zeros REJECT too-large: 9437184 bytes > cap 8388608
|
||||
@@ -0,0 +1,7 @@
|
||||
# Sniff + decode does not catch bytes after the image
|
||||
|
||||
input : valid 16x16 PNG followed by 25 bytes of <script>
|
||||
ImageVerifier : ACCEPT (image/png)
|
||||
why it is still OK: the object is stored as Content-Type image/png with
|
||||
Content-Disposition: attachment, so a browser will not run it as HTML.
|
||||
The check is the header, not the sniffing.
|
||||
@@ -0,0 +1,4 @@
|
||||
# Oversized base64 is refused from its length, before decoding
|
||||
|
||||
payload chars : 12582912 (decodes to about 9437184 bytes; cap is 8388608)
|
||||
result : too-large: base64 payload decodes to about 9437184 bytes > cap 8388608
|
||||
@@ -0,0 +1,10 @@
|
||||
# RemoteFetcher against a local image host
|
||||
|
||||
case result
|
||||
allowed host, 200 OK 7875 bytes
|
||||
302 redirect to the same host fetch-status: 302
|
||||
endless body (cap 1 MiB) too-large: download passed the 1048576-byte cap
|
||||
host not on the list host-not-allowed: localhost
|
||||
file:// URL bad-scheme: file
|
||||
http:// when insecure is off bad-scheme: http
|
||||
cloud metadata address host-not-allowed: 169.254.169.254
|
||||
@@ -0,0 +1,15 @@
|
||||
# Key construction
|
||||
|
||||
200 puts of identical bytes -> 200 distinct ids (a content hash would give 1)
|
||||
owner '../admin' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
owner 'user-42/../../x' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
owner 'USER' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
owner '' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
owner 'a b' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
owner 'user-42\0' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
owner 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx' -> bad-owner: owner must match [a-z0-9][a-z0-9-]{0,31}
|
||||
id '../../etc/passwd' -> bad-id: an image id is 32 lowercase hex characters
|
||||
id 'ABCDEF0123456789ABCDEF0123456789' -> bad-id: an image id is 32 lowercase hex characters
|
||||
id '0123' -> bad-id: an image id is 32 lowercase hex characters
|
||||
id '0123456789abcdef0123456789abcde/' -> bad-id: an image id is 32 lowercase hex characters
|
||||
user-43 asks for user-42's id -> NoSuchKeyException (key is rebuilt under user-43/)
|
||||
@@ -0,0 +1,6 @@
|
||||
# Trusting the sender's Content-Type vs pinning it
|
||||
|
||||
naive writer : key naive/upload.png Content-Type = text/html (sender said so; the .png suffix did not matter)
|
||||
ImageStore : the verifier refuses these bytes first:
|
||||
not-an-image: no known image signature in the first bytes
|
||||
objects under naive/ : 1
|
||||
@@ -0,0 +1,9 @@
|
||||
# ImageService: scripted ImageModel, scripted moderation, real MinIO
|
||||
|
||||
moderation saw : [a red bicycle]
|
||||
ImageModel calls : 1
|
||||
options sent : OpenAiImageOptions model=gpt-image-test size=1536x1024 quality=high n=1
|
||||
prompt sent : a red bicycle
|
||||
stored : generated/happy/65ed822c50b77fecbcedfd866b567c54.png image/png 7875 bytes 64x64
|
||||
revised prompt : a red bicycle on a quiet street, photographic
|
||||
objects in bucket under generated/happy/ : 1
|
||||
@@ -0,0 +1,11 @@
|
||||
# What is refused before any image is paid for
|
||||
|
||||
flagged by moderation -> prompt-flagged: violence
|
||||
blank prompt -> empty-prompt: prompt is blank
|
||||
prompt over 4000 chars -> prompt-too-long: 4001 > 4000
|
||||
control character -> control-characters: prompt contains control characters
|
||||
size not in list -> bad-size: 9999x9999 not in [1024x1024, 1024x1536, 1536x1024]
|
||||
quality not in list -> bad-quality: ultra not in [high, low, medium]
|
||||
moderation endpoint down -> moderation-unavailable: IllegalStateException
|
||||
ImageModel calls : 0
|
||||
objects stored : 0
|
||||
@@ -0,0 +1,4 @@
|
||||
# A three-image response where the second is HTML
|
||||
|
||||
result : not-an-image: no known image signature in the first bytes
|
||||
objects under batch/ : 0
|
||||
@@ -0,0 +1,5 @@
|
||||
# Provider returns a URL instead of bytes
|
||||
|
||||
allowed host : stored generated/urls/f009e9267986ffc76e6b2aacec391b12.png (image/png, 7875 bytes)
|
||||
metadata-service URL : host-not-allowed: 169.254.169.254
|
||||
garbage base64 : bad-base64: payload is not valid base64
|
||||
@@ -0,0 +1,38 @@
|
||||
# Per-image prices as read from provider pages on 2026-10-28 (USD)
|
||||
|
||||
OpenAI, output only, text prompt tokens extra
|
||||
model quality 1024x1024 1024x1536 1536x1024
|
||||
gpt-image-2 low $0.006 $0.005 $0.005
|
||||
gpt-image-2 medium $0.053 $0.041 $0.041
|
||||
gpt-image-2 high $0.211 $0.165 $0.165
|
||||
gpt-image-1.5 low $0.009 $0.013 $0.013
|
||||
gpt-image-1.5 medium $0.034 $0.05 $0.05
|
||||
gpt-image-1.5 high $0.133 $0.2 $0.2
|
||||
gpt-image-1-mini low $0.005 $0.006 $0.006
|
||||
gpt-image-1-mini medium $0.011 $0.015 $0.015
|
||||
gpt-image-1-mini high $0.036 $0.052 $0.052
|
||||
|
||||
Google Gemini image models, standard tier
|
||||
gemini-3.1-flash-lite-image 1K $0.0336
|
||||
gemini-3.1-flash-image 0.5K $0.045
|
||||
gemini-3.1-flash-image 1K $0.067
|
||||
gemini-3.1-flash-image 2K $0.101
|
||||
gemini-3.1-flash-image 4K $0.151
|
||||
gemini-3-pro-image 1K-2K $0.134
|
||||
gemini-3-pro-image 4K $0.24
|
||||
gemini-2.5-flash-image 1K $0.039
|
||||
|
||||
Stability AI, 1 credit = $0.01
|
||||
Stable Image Ultra 8.0 credits = $0.080
|
||||
Stable Diffusion 3.5 Large 6.5 credits = $0.065
|
||||
Stable Diffusion 3.5 Large Turbo 4.0 credits = $0.040
|
||||
Stable Diffusion 3.5 Medium 3.5 credits = $0.035
|
||||
Stable Image Core 3.0 credits = $0.030
|
||||
Stable Diffusion 3.5 Flash 2.5 credits = $0.025
|
||||
|
||||
1,000 images a day for 30 days = 30,000 images
|
||||
gpt-image-2 medium 1024x1024 $1,590
|
||||
gpt-image-1-mini medium 1024x1024 $330
|
||||
gemini-3.1-flash-image 1K $2,010
|
||||
Stable Image Core $900
|
||||
Retrying every refusal once doubles the bill for those requests; the moderation call is separate.
|
||||
@@ -0,0 +1,83 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<parent>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-parent</artifactId>
|
||||
<version>4.1.2</version>
|
||||
<relativePath/>
|
||||
</parent>
|
||||
|
||||
<groupId>com.ankurm</groupId>
|
||||
<artifactId>image-gen</artifactId>
|
||||
<version>1.0.0</version>
|
||||
<name>image-gen</name>
|
||||
<description>Image generation with Spring AI ImageModel: options per provider, a scripted ImageModel for tests, safe storage of results in S3/MinIO (content sniffing, size cap, random keys), and a moderation gate of our own.</description>
|
||||
|
||||
<properties>
|
||||
<java.version>25</java.version>
|
||||
<spring-ai.version>2.0.2</spring-ai.version>
|
||||
<awssdk.version>2.57.5</awssdk.version>
|
||||
<testcontainers.version>2.0.5</testcontainers.version>
|
||||
</properties>
|
||||
|
||||
<dependencyManagement>
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bom</artifactId>
|
||||
<version>${spring-ai.version}</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>software.amazon.awssdk</groupId>
|
||||
<artifactId>bom</artifactId>
|
||||
<version>${awssdk.version}</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>testcontainers-bom</artifactId>
|
||||
<version>${testcontainers.version}</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
</dependencyManagement>
|
||||
|
||||
<dependencies>
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-model</artifactId></dependency>
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-client-chat</artifactId></dependency>
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-openai</artifactId></dependency>
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-stability-ai</artifactId></dependency>
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-google-genai-image</artifactId></dependency>
|
||||
<dependency><groupId>software.amazon.awssdk</groupId><artifactId>s3</artifactId></dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-test</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>testcontainers-junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<configuration>
|
||||
<argLine>-Xmx512m -Duser.timezone=UTC -Dstdout.encoding=UTF-8 -Dfile.encoding=UTF-8</argLine>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
Executable
+9
@@ -0,0 +1,9 @@
|
||||
#!/usr/bin/env bash
|
||||
# Regenerates every file under output/ from the test suite. No API key. Needs Docker for the
|
||||
# MinIO tests (Testcontainers pulls chainguard/minio, an upstream MinIO build; override with
|
||||
# -Dminio.image=...). The image models and the moderation model are hand-written fakes.
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
rm -rf target output
|
||||
mvn -q -B test 2>&1 | grep -E "Tests run:|BUILD|FAIL|ERROR" || true
|
||||
ls output
|
||||
@@ -0,0 +1,71 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
|
||||
/**
|
||||
* Per-image prices copied from each provider's own pricing pages, read on {@link #READ_ON}.
|
||||
* Prices change; the date is the whole claim. A provider or size missing here is missing because
|
||||
* the official page did not state a per-image figure we could read.
|
||||
*
|
||||
* <p>Sources:
|
||||
* <ul>
|
||||
* <li>OpenAI: https://developers.openai.com/api/docs/guides/image-generation ("Calculating costs" table)
|
||||
* and https://developers.openai.com/api/docs/pricing (token rates)</li>
|
||||
* <li>Google: https://ai.google.dev/gemini-api/docs/pricing (Gemini image models)</li>
|
||||
* <li>Stability AI: https://docs.stability.ai/pricing (credits; 1 credit = $0.01)</li>
|
||||
* </ul>
|
||||
*/
|
||||
public final class CostModel {
|
||||
|
||||
public static final String READ_ON = "2026-10-28";
|
||||
public static final double STABILITY_USD_PER_CREDIT = 0.01;
|
||||
|
||||
private CostModel() {}
|
||||
|
||||
/** OpenAI per-image output prices, 1024x1024 / 1024x1536 / 1536x1024, as printed on the guide. Text prompt tokens are extra. */
|
||||
public static Map<String, double[]> openAiPerImage() {
|
||||
Map<String, double[]> m = new LinkedHashMap<>();
|
||||
m.put("gpt-image-2 low", new double[] {0.006, 0.005, 0.005});
|
||||
m.put("gpt-image-2 medium", new double[] {0.053, 0.041, 0.041});
|
||||
m.put("gpt-image-2 high", new double[] {0.211, 0.165, 0.165});
|
||||
m.put("gpt-image-1.5 low", new double[] {0.009, 0.013, 0.013});
|
||||
m.put("gpt-image-1.5 medium", new double[] {0.034, 0.05, 0.05});
|
||||
m.put("gpt-image-1.5 high", new double[] {0.133, 0.2, 0.2});
|
||||
m.put("gpt-image-1-mini low", new double[] {0.005, 0.006, 0.006});
|
||||
m.put("gpt-image-1-mini medium", new double[] {0.011, 0.015, 0.015});
|
||||
m.put("gpt-image-1-mini high", new double[] {0.036, 0.052, 0.052});
|
||||
return m;
|
||||
}
|
||||
|
||||
/** Google: dollars per output image at the stated resolution, standard tier. */
|
||||
public static Map<String, Double> googlePerImage() {
|
||||
Map<String, Double> m = new LinkedHashMap<>();
|
||||
m.put("gemini-3.1-flash-lite-image 1K", 0.0336);
|
||||
m.put("gemini-3.1-flash-image 0.5K", 0.045);
|
||||
m.put("gemini-3.1-flash-image 1K", 0.067);
|
||||
m.put("gemini-3.1-flash-image 2K", 0.101);
|
||||
m.put("gemini-3.1-flash-image 4K", 0.151);
|
||||
m.put("gemini-3-pro-image 1K-2K", 0.134);
|
||||
m.put("gemini-3-pro-image 4K", 0.24);
|
||||
m.put("gemini-2.5-flash-image 1K", 0.039);
|
||||
return m;
|
||||
}
|
||||
|
||||
/** Stability: credits per call. */
|
||||
public static Map<String, Double> stabilityCredits() {
|
||||
Map<String, Double> m = new LinkedHashMap<>();
|
||||
m.put("Stable Image Ultra", 8.0);
|
||||
m.put("Stable Diffusion 3.5 Large", 6.5);
|
||||
m.put("Stable Diffusion 3.5 Large Turbo", 4.0);
|
||||
m.put("Stable Diffusion 3.5 Medium", 3.5);
|
||||
m.put("Stable Image Core", 3.0);
|
||||
m.put("Stable Diffusion 3.5 Flash", 2.5);
|
||||
return m;
|
||||
}
|
||||
|
||||
/** Dollars for {@code images} images with a flat per-image price. */
|
||||
public static double total(double perImage, int images) {
|
||||
return perImage * images;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* Every limit in one place so a test (and a reader) can see them.
|
||||
*
|
||||
* @param maxBytes largest image we will accept, decoded
|
||||
* @param maxPixels largest width x height we will accept, read from the header before decoding
|
||||
* @param maxPromptChars longest prompt we will send
|
||||
* @param allowedTypes sniffed MIME types we will store
|
||||
* @param fetchHosts hosts a provider-returned URL may point at (exact match, https only unless allowInsecure)
|
||||
* @param allowInsecure allow http:// fetches; only for tests against a local server
|
||||
*/
|
||||
public record ImagePolicy(long maxBytes, long maxPixels, int maxPromptChars, Set<String> allowedTypes,
|
||||
Set<String> fetchHosts, boolean allowInsecure) {
|
||||
|
||||
public static ImagePolicy defaults() {
|
||||
return new ImagePolicy(8L * 1024 * 1024, 16_000_000L, 4000,
|
||||
Set.of("image/png", "image/jpeg", "image/webp"), Set.of(), false);
|
||||
}
|
||||
|
||||
public ImagePolicy withMaxBytes(long v) {
|
||||
return new ImagePolicy(v, maxPixels, maxPromptChars, allowedTypes, fetchHosts, allowInsecure);
|
||||
}
|
||||
|
||||
public ImagePolicy withFetch(Set<String> hosts, boolean insecure) {
|
||||
return new ImagePolicy(maxBytes, maxPixels, maxPromptChars, allowedTypes, hosts, insecure);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
/**
|
||||
* What the HTTP layer hands the service. {@code owner} comes from the authenticated session, never
|
||||
* from the request body. {@code size} and {@code quality} are validated against the provider's list.
|
||||
*/
|
||||
public record ImageRequest(String owner, String prompt, String size, String quality, int n) {
|
||||
|
||||
public ImageRequest(String owner, String prompt) {
|
||||
this(owner, prompt, null, null, 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Base64;
|
||||
import java.util.List;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.springframework.ai.image.Image;
|
||||
import org.springframework.ai.image.ImageGeneration;
|
||||
import org.springframework.ai.image.ImageModel;
|
||||
import org.springframework.ai.image.ImageOptions;
|
||||
import org.springframework.ai.image.ImagePrompt;
|
||||
import org.springframework.ai.image.ImageResponse;
|
||||
import org.springframework.ai.openai.metadata.OpenAiImageGenerationMetadata;
|
||||
|
||||
/**
|
||||
* prompt check -> ImageModel -> bytes -> verify -> store. Each arrow is a place something can be
|
||||
* refused, and every refusal is a {@link RejectedException} with a stable reason code.
|
||||
*
|
||||
* <p>If one image of several fails verification nothing is stored: a half-saved batch is harder to
|
||||
* reason about than a retry.
|
||||
*/
|
||||
public final class ImageService {
|
||||
|
||||
private final ImageModel model;
|
||||
private final Function<ImageRequest, ImageOptions> options;
|
||||
private final PromptGate gate;
|
||||
private final ImagePolicy policy;
|
||||
private final ImageVerifier verifier;
|
||||
private final RemoteFetcher fetcher;
|
||||
private final ImageStore store;
|
||||
|
||||
public ImageService(ImageModel model, Function<ImageRequest, ImageOptions> options, PromptGate gate,
|
||||
ImagePolicy policy, RemoteFetcher fetcher, ImageStore store) {
|
||||
this.model = model;
|
||||
this.options = options;
|
||||
this.gate = gate;
|
||||
this.policy = policy;
|
||||
this.verifier = new ImageVerifier(policy);
|
||||
this.fetcher = fetcher;
|
||||
this.store = store;
|
||||
}
|
||||
|
||||
public List<StoredImage> generate(ImageRequest request) {
|
||||
gate.check(request.prompt());
|
||||
ImageOptions o = options.apply(request);
|
||||
ImageResponse response = model.call(new ImagePrompt(request.prompt(), o));
|
||||
|
||||
record Pending(byte[] bytes, Sniffed sniffed, String revised) {}
|
||||
List<Pending> verified = new ArrayList<>();
|
||||
for (ImageGeneration g : response.getResults()) {
|
||||
byte[] bytes = bytesOf(g.getOutput());
|
||||
Sniffed sniffed = verifier.verify(bytes);
|
||||
String revised = g.getMetadata() instanceof OpenAiImageGenerationMetadata m ? m.getRevisedPrompt() : null;
|
||||
verified.add(new Pending(bytes, sniffed, revised));
|
||||
}
|
||||
List<StoredImage> stored = new ArrayList<>();
|
||||
for (Pending p : verified) {
|
||||
stored.add(store.put(request.owner(), p.bytes(), p.sniffed(), p.revised()));
|
||||
}
|
||||
return stored;
|
||||
}
|
||||
|
||||
byte[] bytesOf(Image image) {
|
||||
if (image.getB64Json() != null) {
|
||||
String b64 = image.getB64Json();
|
||||
// Check the size from the string length BEFORE decoding, so a 2 GB string never becomes a 1.5 GB array.
|
||||
long approx = (long) b64.length() * 3 / 4;
|
||||
if (approx > policy.maxBytes() + 3) {
|
||||
throw new RejectedException("too-large", "base64 payload decodes to about " + approx + " bytes > cap " + policy.maxBytes());
|
||||
}
|
||||
try {
|
||||
return Base64.getDecoder().decode(b64);
|
||||
} catch (IllegalArgumentException e) {
|
||||
throw new RejectedException("bad-base64", "payload is not valid base64");
|
||||
}
|
||||
}
|
||||
if (image.getUrl() != null) {
|
||||
return fetcher.fetch(image.getUrl());
|
||||
}
|
||||
throw new RejectedException("empty-result", "provider returned neither b64_json nor url");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.util.Optional;
|
||||
|
||||
/**
|
||||
* Decides what a byte array is from its first bytes, ignoring every label the sender supplied
|
||||
* (file name, MIME type, URL suffix). Only four raster formats are recognised; SVG, HTML, PDF and
|
||||
* anything else come back empty and are refused by the caller.
|
||||
*
|
||||
* <p>Sniffing the magic number is necessary, not sufficient: a file can start like a PNG and
|
||||
* carry anything after it. That is why {@link ImageVerifier} also decodes it and why the store
|
||||
* pins the Content-Type itself. See the article's "safe storage" section.
|
||||
*/
|
||||
public final class ImageSniffer {
|
||||
|
||||
private ImageSniffer() {}
|
||||
|
||||
public static Optional<Sniffed> sniff(byte[] b) {
|
||||
if (b == null || b.length < 12) {
|
||||
return Optional.empty();
|
||||
}
|
||||
if (u(b, 0) == 0x89 && b[1] == 'P' && b[2] == 'N' && b[3] == 'G'
|
||||
&& u(b, 4) == 0x0D && u(b, 5) == 0x0A && u(b, 6) == 0x1A && u(b, 7) == 0x0A) {
|
||||
// IHDR: length(4) "IHDR"(4) width(4) height(4)
|
||||
int w = b.length >= 24 ? be32(b, 16) : 0;
|
||||
int h = b.length >= 24 ? be32(b, 20) : 0;
|
||||
return Optional.of(new Sniffed("image/png", "png", w, h));
|
||||
}
|
||||
if (u(b, 0) == 0xFF && u(b, 1) == 0xD8 && u(b, 2) == 0xFF) {
|
||||
int[] wh = jpegSize(b);
|
||||
return Optional.of(new Sniffed("image/jpeg", "jpg", wh[0], wh[1]));
|
||||
}
|
||||
if (b[0] == 'G' && b[1] == 'I' && b[2] == 'F' && b[3] == '8' && (b[4] == '7' || b[4] == '9') && b[5] == 'a') {
|
||||
int w = u(b, 6) | (u(b, 7) << 8);
|
||||
int h = u(b, 8) | (u(b, 9) << 8);
|
||||
return Optional.of(new Sniffed("image/gif", "gif", w, h));
|
||||
}
|
||||
if (b[0] == 'R' && b[1] == 'I' && b[2] == 'F' && b[3] == 'F' && b[8] == 'W' && b[9] == 'E' && b[10] == 'B' && b[11] == 'P') {
|
||||
return Optional.of(new Sniffed("image/webp", "webp", 0, 0));
|
||||
}
|
||||
return Optional.empty();
|
||||
}
|
||||
|
||||
/** Walks JPEG segments to the first start-of-frame marker. Returns {0, 0} if none is found. */
|
||||
private static int[] jpegSize(byte[] b) {
|
||||
int i = 2;
|
||||
while (i + 9 < b.length) {
|
||||
if (u(b, i) != 0xFF) {
|
||||
i++;
|
||||
continue;
|
||||
}
|
||||
int marker = u(b, i + 1);
|
||||
if (marker == 0xFF) {
|
||||
i++;
|
||||
continue;
|
||||
}
|
||||
if (marker == 0xD8 || marker == 0x01 || (marker >= 0xD0 && marker <= 0xD7)) {
|
||||
i += 2;
|
||||
continue;
|
||||
}
|
||||
int len = (u(b, i + 2) << 8) | u(b, i + 3);
|
||||
boolean sof = marker >= 0xC0 && marker <= 0xCF && marker != 0xC4 && marker != 0xC8 && marker != 0xCC;
|
||||
if (sof) {
|
||||
int h = (u(b, i + 5) << 8) | u(b, i + 6);
|
||||
int w = (u(b, i + 7) << 8) | u(b, i + 8);
|
||||
return new int[] {w, h};
|
||||
}
|
||||
i += 2 + len;
|
||||
}
|
||||
return new int[] {0, 0};
|
||||
}
|
||||
|
||||
private static int u(byte[] b, int i) {
|
||||
return b[i] & 0xFF;
|
||||
}
|
||||
|
||||
private static int be32(byte[] b, int i) {
|
||||
return (u(b, i) << 24) | (u(b, i + 1) << 16) | (u(b, i + 2) << 8) | u(b, i + 3);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.security.MessageDigest;
|
||||
import java.security.SecureRandom;
|
||||
import java.time.Duration;
|
||||
import java.util.HexFormat;
|
||||
import java.util.Map;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
import software.amazon.awssdk.core.sync.RequestBody;
|
||||
import software.amazon.awssdk.services.s3.S3Client;
|
||||
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
|
||||
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
|
||||
import software.amazon.awssdk.services.s3.presigner.S3Presigner;
|
||||
import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest;
|
||||
|
||||
/**
|
||||
* Writes generated images to an S3-compatible bucket (AWS S3 or MinIO).
|
||||
*
|
||||
* <p>No part of the object key is ever taken from the caller. The owner is a server-side id that
|
||||
* must match {@link #OWNER}; the name is 128 random bits from {@link SecureRandom}; the extension
|
||||
* comes from the sniffed type. A client that wants an image back sends the 32-hex id and we rebuild
|
||||
* the key, so {@code ../} and absolute paths have nowhere to go.
|
||||
*/
|
||||
public final class ImageStore {
|
||||
|
||||
public static final Pattern OWNER = Pattern.compile("[a-z0-9][a-z0-9-]{0,31}");
|
||||
public static final Pattern ID = Pattern.compile("[0-9a-f]{32}");
|
||||
private static final SecureRandom RANDOM = new SecureRandom();
|
||||
|
||||
private final S3Client s3;
|
||||
private final S3Presigner presigner;
|
||||
private final String bucket;
|
||||
|
||||
public ImageStore(S3Client s3, S3Presigner presigner, String bucket) {
|
||||
this.s3 = s3;
|
||||
this.presigner = presigner;
|
||||
this.bucket = bucket;
|
||||
}
|
||||
|
||||
public StoredImage put(String owner, byte[] bytes, Sniffed sniffed, String revisedPrompt) {
|
||||
requireOwner(owner);
|
||||
byte[] raw = new byte[16];
|
||||
RANDOM.nextBytes(raw);
|
||||
String id = HexFormat.of().formatHex(raw);
|
||||
String key = keyFor(owner, id, sniffed.extension());
|
||||
String sha = sha256(bytes);
|
||||
s3.putObject(PutObjectRequest.builder()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.contentType(sniffed.contentType()) // from the bytes, never from the sender
|
||||
.contentDisposition("attachment") // browsers download instead of rendering
|
||||
.cacheControl("private, max-age=3600")
|
||||
.metadata(Map.of("sha256", sha, "owner", owner))
|
||||
.build(), RequestBody.fromBytes(bytes));
|
||||
return new StoredImage(key, id, sniffed.contentType(), bytes.length, sniffed.width(), sniffed.height(), sha, revisedPrompt);
|
||||
}
|
||||
|
||||
/** Rebuilds the key from a validated owner and id; the extension is looked up, not supplied. */
|
||||
public String keyFor(String owner, String id, String extension) {
|
||||
requireOwner(owner);
|
||||
if (!ID.matcher(id).matches()) {
|
||||
throw new RejectedException("bad-id", "an image id is 32 lowercase hex characters");
|
||||
}
|
||||
return "generated/" + owner + "/" + id + "." + extension;
|
||||
}
|
||||
|
||||
public byte[] read(String owner, String id, String extension) {
|
||||
return s3.getObjectAsBytes(GetObjectRequest.builder().bucket(bucket).key(keyFor(owner, id, extension)).build()).asByteArray();
|
||||
}
|
||||
|
||||
public String presignedGet(String owner, String id, String extension, Duration ttl) {
|
||||
String key = keyFor(owner, id, extension);
|
||||
return presigner.presignGetObject(GetObjectPresignRequest.builder()
|
||||
.signatureDuration(ttl)
|
||||
.getObjectRequest(GetObjectRequest.builder().bucket(bucket).key(key).build())
|
||||
.build()).url().toString();
|
||||
}
|
||||
|
||||
private static void requireOwner(String owner) {
|
||||
if (owner == null || !OWNER.matcher(owner).matches()) {
|
||||
throw new RejectedException("bad-owner", "owner must match " + OWNER.pattern());
|
||||
}
|
||||
}
|
||||
|
||||
static String sha256(byte[] b) {
|
||||
try {
|
||||
return HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(b));
|
||||
} catch (java.security.NoSuchAlgorithmException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.IOException;
|
||||
import javax.imageio.ImageIO;
|
||||
|
||||
/**
|
||||
* Applies the policy to bytes: size cap, then sniff, then allowed type, then pixel cap read from
|
||||
* the header (before any decoder allocates a bitmap), then a real decode for the formats the JDK
|
||||
* can read. WebP is sniffed only, because the JDK has no WebP reader.
|
||||
*/
|
||||
public final class ImageVerifier {
|
||||
|
||||
private final ImagePolicy policy;
|
||||
|
||||
public ImageVerifier(ImagePolicy policy) {
|
||||
this.policy = policy;
|
||||
}
|
||||
|
||||
public Sniffed verify(byte[] bytes) {
|
||||
if (bytes.length > policy.maxBytes()) {
|
||||
throw new RejectedException("too-large", bytes.length + " bytes > cap " + policy.maxBytes());
|
||||
}
|
||||
Sniffed s = ImageSniffer.sniff(bytes)
|
||||
.orElseThrow(() -> new RejectedException("not-an-image", "no known image signature in the first bytes"));
|
||||
if (!policy.allowedTypes().contains(s.contentType())) {
|
||||
throw new RejectedException("type-not-allowed", s.contentType());
|
||||
}
|
||||
if (s.pixels() > policy.maxPixels()) {
|
||||
throw new RejectedException("too-many-pixels", s.width() + "x" + s.height() + " = " + s.pixels() + " > cap " + policy.maxPixels());
|
||||
}
|
||||
if (!s.contentType().equals("image/webp")) {
|
||||
try {
|
||||
BufferedImage img = ImageIO.read(new ByteArrayInputStream(bytes));
|
||||
if (img == null) {
|
||||
throw new RejectedException("undecodable", "signature says " + s.contentType() + " but no decoder accepts it");
|
||||
}
|
||||
} catch (IOException | RuntimeException e) {
|
||||
if (e instanceof RejectedException r) throw r;
|
||||
throw new RejectedException("undecodable", e.getClass().getSimpleName());
|
||||
}
|
||||
}
|
||||
return s;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.util.Set;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.springframework.ai.image.ImageOptions;
|
||||
import org.springframework.ai.openai.OpenAiImageOptions;
|
||||
|
||||
/**
|
||||
* Turns the caller's size and quality choices into {@link OpenAiImageOptions} - after checking
|
||||
* them against an allow-list, because the provider's own error for a bad value arrives late
|
||||
* (after the network round trip) and as an SDK exception.
|
||||
*/
|
||||
public final class OpenAiOptionsFactory implements Function<ImageRequest, ImageOptions> {
|
||||
|
||||
public static final Set<String> SIZES = Set.of("1024x1024", "1536x1024", "1024x1536");
|
||||
public static final Set<String> QUALITIES = Set.of("low", "medium", "high");
|
||||
|
||||
private final String model;
|
||||
|
||||
public OpenAiOptionsFactory(String model) {
|
||||
this.model = model;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ImageOptions apply(ImageRequest r) {
|
||||
String size = r.size() == null ? "1024x1024" : r.size();
|
||||
String quality = r.quality() == null ? "medium" : r.quality();
|
||||
if (!SIZES.contains(size)) {
|
||||
throw new RejectedException("bad-size", size + " not in " + new java.util.TreeSet<>(SIZES));
|
||||
}
|
||||
if (!QUALITIES.contains(quality)) {
|
||||
throw new RejectedException("bad-quality", quality + " not in " + new java.util.TreeSet<>(QUALITIES));
|
||||
}
|
||||
if (r.n() < 1 || r.n() > 4) {
|
||||
throw new RejectedException("bad-n", "n must be 1..4");
|
||||
}
|
||||
return OpenAiImageOptions.builder().model(model).size(size).quality(quality).n(r.n()).build();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
import org.springframework.ai.moderation.Categories;
|
||||
import org.springframework.ai.moderation.ModerationModel;
|
||||
import org.springframework.ai.moderation.ModerationPrompt;
|
||||
import org.springframework.ai.moderation.ModerationResult;
|
||||
|
||||
/**
|
||||
* The moderation step Spring AI does not run for you. It asks a {@link ModerationModel} about the
|
||||
* prompt before any image is requested and fails closed: if the moderation call itself throws, the
|
||||
* prompt is refused rather than waved through.
|
||||
*
|
||||
* <p>This inspects the TEXT. It says nothing about the picture that comes back.
|
||||
*/
|
||||
public final class PromptGate {
|
||||
|
||||
private final ModerationModel moderation;
|
||||
private final ImagePolicy policy;
|
||||
|
||||
public PromptGate(ModerationModel moderation, ImagePolicy policy) {
|
||||
this.moderation = moderation;
|
||||
this.policy = policy;
|
||||
}
|
||||
|
||||
public void check(String prompt) {
|
||||
if (prompt == null || prompt.isBlank()) {
|
||||
throw new RejectedException("empty-prompt", "prompt is blank");
|
||||
}
|
||||
if (prompt.length() > policy.maxPromptChars()) {
|
||||
throw new RejectedException("prompt-too-long", prompt.length() + " > " + policy.maxPromptChars());
|
||||
}
|
||||
if (prompt.chars().anyMatch(c -> Character.isISOControl(c) && c != '\n' && c != '\t')) {
|
||||
throw new RejectedException("control-characters", "prompt contains control characters");
|
||||
}
|
||||
List<String> flagged = new ArrayList<>();
|
||||
try {
|
||||
for (ModerationResult r : moderation.call(new ModerationPrompt(prompt)).getResult().getOutput().getResults()) {
|
||||
if (r.isFlagged()) {
|
||||
flagged.addAll(flaggedCategories(r.getCategories()));
|
||||
}
|
||||
}
|
||||
} catch (RuntimeException e) {
|
||||
throw new RejectedException("moderation-unavailable", e.getClass().getSimpleName());
|
||||
}
|
||||
if (!flagged.isEmpty()) {
|
||||
throw new RejectedException("prompt-flagged", String.join(",", flagged));
|
||||
}
|
||||
}
|
||||
|
||||
private static List<String> flaggedCategories(Categories c) {
|
||||
List<String> out = new ArrayList<>();
|
||||
if (c == null) return List.of("unspecified");
|
||||
if (c.isSexual()) out.add("sexual");
|
||||
if (c.isSexualMinors()) out.add("sexual/minors");
|
||||
if (c.isHate()) out.add("hate");
|
||||
if (c.isHarassment()) out.add("harassment");
|
||||
if (c.isSelfHarm()) out.add("self-harm");
|
||||
if (c.isViolence()) out.add("violence");
|
||||
if (c.isViolenceGraphic()) out.add("violence/graphic");
|
||||
return out.isEmpty() ? List.of("unspecified") : out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
/** A request or a provider result was refused by one of our own checks. {@code reason} is a stable code. */
|
||||
public class RejectedException extends RuntimeException {
|
||||
|
||||
private final String reason;
|
||||
|
||||
public RejectedException(String reason, String detail) {
|
||||
super(reason + ": " + detail);
|
||||
this.reason = reason;
|
||||
}
|
||||
|
||||
public String reason() {
|
||||
return reason;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.time.Duration;
|
||||
|
||||
/**
|
||||
* Downloads an image a provider returned as a URL (DALL-E 3 style {@code response_format=url}).
|
||||
* The URL comes from a third party, so it is treated as untrusted input: https only, the host must
|
||||
* be on a list, redirects are not followed, the read stops at the byte cap, and there is a timeout.
|
||||
*/
|
||||
public final class RemoteFetcher {
|
||||
|
||||
private final ImagePolicy policy;
|
||||
private final HttpClient http = HttpClient.newBuilder()
|
||||
.followRedirects(HttpClient.Redirect.NEVER)
|
||||
.connectTimeout(Duration.ofSeconds(5))
|
||||
.build();
|
||||
|
||||
public RemoteFetcher(ImagePolicy policy) {
|
||||
this.policy = policy;
|
||||
}
|
||||
|
||||
public byte[] fetch(String url) {
|
||||
URI uri;
|
||||
try {
|
||||
uri = URI.create(url);
|
||||
} catch (IllegalArgumentException e) {
|
||||
throw new RejectedException("bad-url", "not a URI");
|
||||
}
|
||||
boolean httpsOk = "https".equalsIgnoreCase(uri.getScheme());
|
||||
boolean httpOk = policy.allowInsecure() && "http".equalsIgnoreCase(uri.getScheme());
|
||||
if (!httpsOk && !httpOk) {
|
||||
throw new RejectedException("bad-scheme", String.valueOf(uri.getScheme()));
|
||||
}
|
||||
if (uri.getHost() == null || !policy.fetchHosts().contains(uri.getHost().toLowerCase())) {
|
||||
throw new RejectedException("host-not-allowed", String.valueOf(uri.getHost()));
|
||||
}
|
||||
try {
|
||||
HttpResponse<InputStream> r = http.send(
|
||||
HttpRequest.newBuilder(uri).timeout(Duration.ofSeconds(15)).GET().build(),
|
||||
HttpResponse.BodyHandlers.ofInputStream());
|
||||
if (r.statusCode() != 200) {
|
||||
r.body().close();
|
||||
throw new RejectedException("fetch-status", String.valueOf(r.statusCode()));
|
||||
}
|
||||
return readCapped(r.body());
|
||||
} catch (IOException e) {
|
||||
throw new RejectedException("fetch-failed", e.getClass().getSimpleName());
|
||||
} catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
throw new RejectedException("fetch-failed", "interrupted");
|
||||
}
|
||||
}
|
||||
|
||||
/** Reads at most maxBytes + 1 bytes so an endless body costs us one byte over the cap, not the heap. */
|
||||
private byte[] readCapped(InputStream in) throws IOException {
|
||||
try (in) {
|
||||
ByteArrayOutputStream out = new ByteArrayOutputStream();
|
||||
byte[] buf = new byte[8192];
|
||||
long total = 0;
|
||||
int n;
|
||||
while ((n = in.read(buf)) > 0) {
|
||||
total += n;
|
||||
if (total > policy.maxBytes()) {
|
||||
throw new RejectedException("too-large", "download passed the " + policy.maxBytes() + "-byte cap");
|
||||
}
|
||||
out.write(buf, 0, n);
|
||||
}
|
||||
return out.toByteArray();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
/**
|
||||
* What the bytes themselves say they are. {@code width} and {@code height} are 0 when the
|
||||
* format's header was not parsed (WebP).
|
||||
*/
|
||||
public record Sniffed(String contentType, String extension, int width, int height) {
|
||||
|
||||
public long pixels() {
|
||||
return (long) width * height;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
/** What we kept: where, what it really is, how big, and the provider's rewrite of the prompt if it sent one. */
|
||||
public record StoredImage(String key, String id, String contentType, int bytes, int width, int height,
|
||||
String sha256, String revisedPrompt) {}
|
||||
@@ -0,0 +1,72 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.lang.reflect.Method;
|
||||
import java.util.Arrays;
|
||||
import java.util.TreeSet;
|
||||
|
||||
import com.ankurm.imagegen.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.google.genai.image.GoogleGenAiImageOptions;
|
||||
import org.springframework.ai.image.ImageModel;
|
||||
import org.springframework.ai.image.ImageOptions;
|
||||
import org.springframework.ai.image.ImageOptionsBuilder;
|
||||
import org.springframework.ai.image.ImagePrompt;
|
||||
import org.springframework.ai.moderation.ModerationModel;
|
||||
import org.springframework.ai.openai.OpenAiImageModel;
|
||||
import org.springframework.ai.openai.OpenAiImageOptions;
|
||||
import org.springframework.ai.openai.OpenAiModerationModel;
|
||||
import org.springframework.ai.openai.OpenAiModerationOptions;
|
||||
import org.springframework.ai.stabilityai.StabilityAiImageModel;
|
||||
import org.springframework.ai.stabilityai.api.StabilityAiImageOptions;
|
||||
|
||||
/** The image API as shipped in the 2.0.2 jars: names, the portable options, each provider's extras. Nothing is sent. */
|
||||
class ApiShapeTest {
|
||||
|
||||
@Test
|
||||
void namesAndOptions() {
|
||||
try (Transcript t = new Transcript("01-api-shape.txt", "Spring AI 2.0.2 image classes (read from the jars)")) {
|
||||
t.line("interface : %s", ImageModel.class.getName());
|
||||
t.line("single method : %s", Arrays.stream(ImageModel.class.getDeclaredMethods())
|
||||
.filter(m -> !m.isSynthetic() && !m.isBridge() && java.lang.reflect.Modifier.isAbstract(m.getModifiers()))
|
||||
.map(Method::getName).toList());
|
||||
t.line("request : %s", ImagePrompt.class.getName());
|
||||
t.line("OpenAI : %s", OpenAiImageModel.class.getName());
|
||||
t.line("Stability AI : %s", StabilityAiImageModel.class.getName());
|
||||
t.line("moderation interface : %s", ModerationModel.class.getName());
|
||||
t.line("moderation OpenAI : %s", OpenAiModerationModel.class.getName());
|
||||
t.blank();
|
||||
t.line("portable ImageOptions getters : %s", getters(ImageOptions.class));
|
||||
t.line("OpenAiImageOptions getters : %s", getters(OpenAiImageOptions.class));
|
||||
t.line("StabilityAiImageOptions : %s", getters(StabilityAiImageOptions.class));
|
||||
t.line("GoogleGenAiImageOptions : %s", getters(GoogleGenAiImageOptions.class));
|
||||
t.blank();
|
||||
t.line("OpenAiImageOptions.DEFAULT_IMAGE_MODEL = %s", OpenAiImageOptions.DEFAULT_IMAGE_MODEL);
|
||||
t.line("OpenAiModerationOptions.DEFAULT_MODERATION_MODEL = %s", OpenAiModerationOptions.DEFAULT_MODERATION_MODEL);
|
||||
t.line("OpenAiImageOptions has a getModeration() : %s", getters(OpenAiImageOptions.class).contains("getModeration"));
|
||||
t.line("ModerationMessage getters : %s (text only, no media)", getters(org.springframework.ai.moderation.ModerationMessage.class));
|
||||
t.line("Stability default engine / base URL : %s @ %s", org.springframework.ai.stabilityai.api.StabilityAiApi.DEFAULT_IMAGE_MODEL, org.springframework.ai.stabilityai.api.StabilityAiApi.DEFAULT_BASE_URL);
|
||||
t.line("GoogleGenAiImageModelName : %s", java.util.Arrays.stream(org.springframework.ai.google.genai.image.GoogleGenAiImageModelName.values()).map(org.springframework.ai.google.genai.image.GoogleGenAiImageModelName::getName).toList());
|
||||
t.line("Google SafetyFilterLevel : %s", java.util.Arrays.toString(GoogleGenAiImageOptions.SafetyFilterLevel.values()));
|
||||
t.line("Google PersonGeneration : %s", java.util.Arrays.toString(GoogleGenAiImageOptions.PersonGeneration.values()));
|
||||
t.line("AbstractOpenAiOptions.DEFAULT_MAX_RETRIES (image options extend it) = %d", org.springframework.ai.openai.AbstractOpenAiOptions.DEFAULT_MAX_RETRIES);
|
||||
ImageOptions portable = ImageOptionsBuilder.builder().model("m").width(1024).height(1536).n(2).build();
|
||||
t.line("portable builder gives : %s", portable.getClass().getSimpleName());
|
||||
}
|
||||
assertThat(getters(OpenAiImageOptions.class)).doesNotContain("getModeration");
|
||||
assertThat(getters(ImageOptions.class)).doesNotContain("getQuality", "getSize");
|
||||
assertThat(getters(OpenAiImageOptions.class)).contains("getQuality", "getSize", "getUser");
|
||||
}
|
||||
|
||||
private static java.util.Set<String> getters(Class<?> c) {
|
||||
TreeSet<String> s = new TreeSet<>();
|
||||
for (Method m : c.getMethods()) {
|
||||
if (m.getParameterCount() == 0 && (m.getName().startsWith("get") || m.getName().startsWith("is"))
|
||||
&& m.getDeclaringClass() != Object.class && !m.getName().equals("getClass")) {
|
||||
s.add(m.getName());
|
||||
}
|
||||
}
|
||||
return s;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import com.ankurm.imagegen.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** Arithmetic only. The prices are in CostModel with the date they were read; this test does not check them against the web. */
|
||||
class CostTest {
|
||||
|
||||
@Test
|
||||
void perProviderTable() {
|
||||
try (Transcript t = new Transcript("16-cost.txt", "Per-image prices as read from provider pages on " + CostModel.READ_ON + " (USD)")) {
|
||||
t.line("OpenAI, output only, text prompt tokens extra");
|
||||
t.line("%-26s %10s %10s %10s", "model quality", "1024x1024", "1024x1536", "1536x1024");
|
||||
CostModel.openAiPerImage().forEach((k, v) -> t.line("%-26s %10s %10s %10s", k, "$" + v[0], "$" + v[1], "$" + v[2]));
|
||||
t.blank();
|
||||
t.line("Google Gemini image models, standard tier");
|
||||
CostModel.googlePerImage().forEach((k, v) -> t.line("%-34s $%s", k, v));
|
||||
t.blank();
|
||||
t.line("Stability AI, 1 credit = $%s", CostModel.STABILITY_USD_PER_CREDIT);
|
||||
CostModel.stabilityCredits().forEach((k, v) -> t.line("%-34s %4s credits = $%.3f", k, v, v * CostModel.STABILITY_USD_PER_CREDIT));
|
||||
t.blank();
|
||||
t.line("1,000 images a day for 30 days = 30,000 images");
|
||||
t.line("%-34s $%,.0f", "gpt-image-2 medium 1024x1024", CostModel.total(0.053, 30000));
|
||||
t.line("%-34s $%,.0f", "gpt-image-1-mini medium 1024x1024", CostModel.total(0.011, 30000));
|
||||
t.line("%-34s $%,.0f", "gemini-3.1-flash-image 1K", CostModel.total(0.067, 30000));
|
||||
t.line("%-34s $%,.0f", "Stable Image Core", CostModel.total(3.0 * 0.01, 30000));
|
||||
t.line("Retrying every refusal once doubles the bill for those requests; the moderation call is separate.");
|
||||
}
|
||||
assertThat(CostModel.total(0.053, 30000)).isEqualTo(1590.0, org.assertj.core.data.Offset.offset(0.001));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
import com.ankurm.imagegen.support.FakeImagesServer;
|
||||
import com.ankurm.imagegen.support.Transcript;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.image.ImageOptionsBuilder;
|
||||
import org.springframework.ai.image.ImageMessage;
|
||||
import org.springframework.ai.image.ImagePrompt;
|
||||
import org.springframework.ai.image.ImageResponse;
|
||||
import org.springframework.ai.openai.OpenAiImageModel;
|
||||
import org.springframework.ai.openai.OpenAiImageOptions;
|
||||
import org.springframework.ai.openai.metadata.OpenAiImageGenerationMetadata;
|
||||
|
||||
/**
|
||||
* The REAL OpenAiImageModel against a local mock of /v1/images/generations: what Spring AI puts on
|
||||
* the wire for each options object, and what comes back. The mock does not judge the values.
|
||||
*/
|
||||
class OpenAiWireTest {
|
||||
|
||||
static FakeImagesServer server;
|
||||
|
||||
@BeforeAll
|
||||
static void start() throws Exception {
|
||||
server = new FakeImagesServer();
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void stop() {
|
||||
server.close();
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void reset() {
|
||||
server.clear();
|
||||
server.mode(FakeImagesServer.Mode.B64);
|
||||
}
|
||||
|
||||
private static OpenAiImageModel model(String defaultModel) {
|
||||
return OpenAiImageModel.builder().options(OpenAiImageOptions.builder()
|
||||
.baseUrl(server.baseUrl()).apiKey("test-key-not-real").model(defaultModel).maxRetries(0).build()).build();
|
||||
}
|
||||
|
||||
@Test
|
||||
void requestsOnTheWire() {
|
||||
try (Transcript t = new Transcript("02-wire-requests.txt", "What OpenAiImageModel sends (mock server; bodies recorded as received)")) {
|
||||
OpenAiImageModel m = model("gpt-image-test");
|
||||
|
||||
m.call(new ImagePrompt("a red bicycle"));
|
||||
t.line("1. ImagePrompt(\"a red bicycle\"), model default only");
|
||||
t.line(" %s %s", server.seen().getLast().method(), server.seen().getLast().path());
|
||||
t.line(" %s", server.seen().getLast().body());
|
||||
t.blank();
|
||||
|
||||
m.call(new ImagePrompt("a red bicycle", OpenAiImageOptions.builder().size("1536x1024").quality("high").n(2).build()));
|
||||
t.line("2. OpenAiImageOptions size=1536x1024 quality=high n=2");
|
||||
t.line(" %s", server.seen().getLast().body());
|
||||
t.blank();
|
||||
|
||||
m.call(new ImagePrompt("a red bicycle", ImageOptionsBuilder.builder().width(1024).height(1536).build()));
|
||||
t.line("3. portable ImageOptionsBuilder width=1024 height=1536");
|
||||
t.line(" %s", server.seen().getLast().body());
|
||||
t.blank();
|
||||
|
||||
m.call(new ImagePrompt("a red bicycle", OpenAiImageOptions.builder().user("user-8f2c").build()));
|
||||
t.line("4. OpenAiImageOptions user=user-8f2c");
|
||||
t.line(" %s", server.seen().getLast().body());
|
||||
t.blank();
|
||||
|
||||
m.call(new ImagePrompt("a red bicycle", OpenAiImageOptions.builder().responseFormat("b64_json").build()));
|
||||
t.line("5. OpenAiImageOptions responseFormat=b64_json");
|
||||
t.line(" %s", server.seen().getLast().body());
|
||||
t.blank();
|
||||
|
||||
var twoMessages = java.util.List.of(new ImageMessage("a red bicycle", 1.0f), new ImageMessage("no people", -1.0f));
|
||||
try {
|
||||
m.call(new ImagePrompt(twoMessages));
|
||||
t.line("6a. ImagePrompt(List<ImageMessage>) with no options: accepted");
|
||||
} catch (RuntimeException e) {
|
||||
t.line("6a. ImagePrompt(List<ImageMessage>) with no options: %s: %s", e.getClass().getSimpleName(), e.getMessage());
|
||||
}
|
||||
m.call(new ImagePrompt(twoMessages, OpenAiImageOptions.builder().build()));
|
||||
t.line("6b. same two messages (weights 1.0 and -1.0) with empty OpenAiImageOptions");
|
||||
t.line(" %s", server.seen().getLast().body());
|
||||
}
|
||||
assertThat(server.seen()).hasSize(6);
|
||||
// the constructor default model is NOT inherited by a per-call OpenAiImageOptions that leaves model unset
|
||||
assertThat(server.seen().get(0).body()).contains("\"model\":\"gpt-image-test\"");
|
||||
assertThat(server.seen().get(1).body()).contains("\"model\":\"" + OpenAiImageOptions.DEFAULT_IMAGE_MODEL + "\"");
|
||||
assertThat(server.seen().get(1).body()).contains("1536x1024").contains("high");
|
||||
}
|
||||
|
||||
@Test
|
||||
void responseShape() {
|
||||
try (Transcript t = new Transcript("03-wire-response.txt", "What comes back from OpenAiImageModel (mock server)")) {
|
||||
ImageResponse r = model("gpt-image-test").call(new ImagePrompt("a red bicycle"));
|
||||
t.line("results : %d", r.getResults().size());
|
||||
var img = r.getResult().getOutput();
|
||||
t.line("Image.getUrl() : %s", img.getUrl());
|
||||
t.line("Image.getB64Json() len : %d chars", img.getB64Json().length());
|
||||
t.line("metadata type : %s", r.getResult().getMetadata().getClass().getSimpleName());
|
||||
t.line("revised prompt : %s", ((OpenAiImageGenerationMetadata) r.getResult().getMetadata()).getRevisedPrompt());
|
||||
t.line("response created : %s", r.getMetadata().getCreated());
|
||||
t.line("response usage : %s, total tokens %d (the mock body carried usage.total_tokens=1068)", r.getMetadata().getUsage().getClass().getSimpleName(), r.getMetadata().getUsage().getTotalTokens());
|
||||
t.blank();
|
||||
server.mode(FakeImagesServer.Mode.URL);
|
||||
ImageResponse u = model("gpt-image-test").call(new ImagePrompt("a red bicycle"));
|
||||
t.line("URL mode: Image.getUrl() = %s", u.getResult().getOutput().getUrl());
|
||||
t.line("URL mode: Image.getB64Json() = %s", u.getResult().getOutput().getB64Json());
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void providerRefusalIsAnException() {
|
||||
server.mode(FakeImagesServer.Mode.REFUSE_400);
|
||||
try (Transcript t = new Transcript("04-wire-refusal.txt", "A provider-side refusal (400 content_policy_violation from the mock)")) {
|
||||
Throwable thrown = null;
|
||||
try {
|
||||
model("gpt-image-test").call(new ImagePrompt("anything"));
|
||||
} catch (RuntimeException e) {
|
||||
thrown = e;
|
||||
}
|
||||
assertThat(thrown).isNotNull();
|
||||
t.line("exception class : %s", thrown.getClass().getName());
|
||||
t.line("message (first 200): %s", thrown.getMessage().substring(0, Math.min(200, thrown.getMessage().length())));
|
||||
t.line("requests received : %d (maxRetries=0)", server.seen().size());
|
||||
}
|
||||
assertThatThrownBy(() -> model("gpt-image-test").call(new ImagePrompt("anything"))).isInstanceOf(RuntimeException.class);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.imagegen.support.FakeImagesServer;
|
||||
import com.ankurm.imagegen.support.Fixtures;
|
||||
import com.ankurm.imagegen.support.Images;
|
||||
import com.ankurm.imagegen.support.Minio;
|
||||
import com.ankurm.imagegen.support.ScriptedImageModel;
|
||||
import com.ankurm.imagegen.support.ScriptedModerationModel;
|
||||
import com.ankurm.imagegen.support.Transcript;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.image.ImageOptions;
|
||||
import org.springframework.ai.openai.OpenAiImageOptions;
|
||||
import software.amazon.awssdk.services.s3.S3Client;
|
||||
import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
|
||||
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
|
||||
|
||||
/** ImageService end to end: scripted ImageModel + scripted ModerationModel + real MinIO. */
|
||||
class ServiceTest {
|
||||
|
||||
static Minio minio;
|
||||
static S3Client s3;
|
||||
static ImageStore store;
|
||||
static final String BUCKET = "service-test";
|
||||
final ImagePolicy policy = ImagePolicy.defaults();
|
||||
|
||||
@BeforeAll
|
||||
static void start() {
|
||||
minio = new Minio().start();
|
||||
s3 = minio.client();
|
||||
s3.createBucket(CreateBucketRequest.builder().bucket(BUCKET).build());
|
||||
store = new ImageStore(s3, minio.presigner(), BUCKET);
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void stop() {
|
||||
if (minio != null) minio.close();
|
||||
}
|
||||
|
||||
private int objects(String prefix) {
|
||||
return s3.listObjectsV2(ListObjectsV2Request.builder().bucket(BUCKET).prefix(prefix).build()).keyCount();
|
||||
}
|
||||
|
||||
@Test
|
||||
void happyPathAndOptionsReachTheModel() {
|
||||
ScriptedImageModel model = new ScriptedImageModel().thenB64Revised("a red bicycle on a quiet street, photographic", Images.png(64, 64));
|
||||
ScriptedModerationModel mod = ScriptedModerationModel.flagging("blood", "violence");
|
||||
ImageService svc = Fixtures.service(model, mod, policy, store);
|
||||
try (Transcript t = new Transcript("12-service-happy.txt", "ImageService: scripted ImageModel, scripted moderation, real MinIO")) {
|
||||
List<StoredImage> out = svc.generate(new ImageRequest("happy", "a red bicycle", "1536x1024", "high", 1));
|
||||
ImageOptions sent = model.seen().getFirst().getOptions();
|
||||
t.line("moderation saw : %s", mod.seen());
|
||||
t.line("ImageModel calls : %d", model.calls());
|
||||
t.line("options sent : %s model=%s size=%s quality=%s n=%d", sent.getClass().getSimpleName(), sent.getModel(),
|
||||
((OpenAiImageOptions) sent).getSize(), ((OpenAiImageOptions) sent).getQuality(), sent.getN());
|
||||
t.line("prompt sent : %s", model.seen().getFirst().getInstructions().getFirst().getText());
|
||||
StoredImage s = out.getFirst();
|
||||
t.line("stored : %s %s %d bytes %dx%d", s.key(), s.contentType(), s.bytes(), s.width(), s.height());
|
||||
t.line("revised prompt : %s", s.revisedPrompt());
|
||||
t.line("objects in bucket under generated/happy/ : %d", objects("generated/happy/"));
|
||||
assertThat(out).hasSize(1);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void refusalsStopBeforeTheModel() {
|
||||
ScriptedImageModel model = new ScriptedImageModel().thenB64(Images.png(8, 8));
|
||||
ScriptedModerationModel mod = ScriptedModerationModel.flagging("blood", "violence");
|
||||
ImageService svc = Fixtures.service(model, mod, policy, store);
|
||||
ScriptedModerationModel down = new ScriptedModerationModel(p -> null).failingWith(new IllegalStateException("moderation endpoint 503"));
|
||||
ImageService svcDown = Fixtures.service(model, down, policy, store);
|
||||
try (Transcript t = new Transcript("13-refusals.txt", "What is refused before any image is paid for")) {
|
||||
String[][] cases = {
|
||||
{"flagged by moderation", "a pool of blood", null, null},
|
||||
{"blank prompt", " ", null, null},
|
||||
{"prompt over 4000 chars", "x".repeat(4001), null, null},
|
||||
{"control character", "a cat\u0007", null, null},
|
||||
{"size not in list", "a cat", "9999x9999", null},
|
||||
{"quality not in list", "a cat", null, "ultra"}};
|
||||
for (String[] c : cases) {
|
||||
String r;
|
||||
try { svc.generate(new ImageRequest("refuse", c[1], c[2], c[3], 1)); r = "GENERATED?!"; } catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("%-24s -> %s", c[0], r);
|
||||
}
|
||||
String r;
|
||||
try { svcDown.generate(new ImageRequest("refuse", "a cat")); r = "GENERATED?!"; } catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("%-24s -> %s", "moderation endpoint down", r);
|
||||
t.line("ImageModel calls : %d", model.calls());
|
||||
t.line("objects stored : %d", objects("generated/refuse/"));
|
||||
}
|
||||
assertThat(model.calls()).isZero();
|
||||
assertThat(objects("generated/refuse/")).isZero();
|
||||
}
|
||||
|
||||
@Test
|
||||
void oneBadImageInABatchStoresNothing() {
|
||||
ScriptedImageModel model = new ScriptedImageModel().thenB64(Images.png(16, 16), "<html></html>".getBytes(), Images.png(16, 16));
|
||||
ImageService svc = Fixtures.service(model, ScriptedModerationModel.flagging("zzz", "x"), policy, store);
|
||||
try (Transcript t = new Transcript("14-batch.txt", "A three-image response where the second is HTML")) {
|
||||
String r;
|
||||
try { svc.generate(new ImageRequest("batch", "three cats", null, null, 3)); r = "STORED?!"; } catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("result : %s", r);
|
||||
t.line("objects under batch/ : %d", objects("generated/batch/"));
|
||||
}
|
||||
assertThat(objects("generated/batch/")).isZero();
|
||||
}
|
||||
|
||||
@Test
|
||||
void urlResultsGoThroughTheFetcherAndTheSameChecks() throws Exception {
|
||||
try (FakeImagesServer host = new FakeImagesServer(); Transcript t = new Transcript("15-url-results.txt", "Provider returns a URL instead of bytes")) {
|
||||
ImagePolicy local = Fixtures.withHost(policy, host.host());
|
||||
ScriptedImageModel okModel = new ScriptedImageModel().thenUrl(host.fileUrl("ok.png"));
|
||||
StoredImage s = Fixtures.service(okModel, ScriptedModerationModel.flagging("zzz", "x"), local, store)
|
||||
.generate(new ImageRequest("urls", "a cat")).getFirst();
|
||||
t.line("allowed host : stored %s (%s, %d bytes)", s.key(), s.contentType(), s.bytes());
|
||||
ScriptedImageModel evil = new ScriptedImageModel().thenUrl("http://169.254.169.254/latest/meta-data/");
|
||||
String r;
|
||||
try { Fixtures.service(evil, ScriptedModerationModel.flagging("zzz", "x"), local, store).generate(new ImageRequest("urls", "a cat")); r = "FETCHED?!"; }
|
||||
catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("metadata-service URL : %s", r);
|
||||
ScriptedImageModel nothing = new ScriptedImageModel().thenRawB64("not base64 !!!");
|
||||
try { Fixtures.service(nothing, ScriptedModerationModel.flagging("zzz", "x"), local, store).generate(new ImageRequest("urls", "a cat")); r = "STORED?!"; }
|
||||
catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("garbage base64 : %s", r);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void providerExceptionPropagates() {
|
||||
ScriptedImageModel model = new ScriptedImageModel().thenThrow(new IllegalStateException("provider 400 content_policy_violation"));
|
||||
ImageService svc = Fixtures.service(model, ScriptedModerationModel.flagging("zzz", "x"), policy, store);
|
||||
assertThatThrownBy(() -> svc.generate(new ImageRequest("prov", "a cat"))).isInstanceOf(IllegalStateException.class);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
import com.ankurm.imagegen.support.FakeImagesServer;
|
||||
import com.ankurm.imagegen.support.Images;
|
||||
import com.ankurm.imagegen.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** Our own byte checks and the fetcher, with no S3 involved: what is refused, and under which code. */
|
||||
class StorageChecksTest {
|
||||
|
||||
private final ImagePolicy policy = ImagePolicy.defaults();
|
||||
private final ImageVerifier verifier = new ImageVerifier(policy);
|
||||
|
||||
@Test
|
||||
void verifierDecisions() {
|
||||
Map<String, byte[]> cases = new LinkedHashMap<>();
|
||||
cases.put("real PNG 64x48", Images.png(64, 48));
|
||||
cases.put("real JPEG 64x48", Images.jpeg(64, 48));
|
||||
cases.put("WebP header only (sniffed, not decoded)", Images.webpHeaderOnly());
|
||||
cases.put("SVG with a script", "<svg xmlns='http://www.w3.org/2000/svg'><script>alert(1)</script></svg>".getBytes(StandardCharsets.UTF_8));
|
||||
cases.put("HTML page", "<!doctype html><html><body><script>alert(1)</script></body></html>".getBytes(StandardCharsets.UTF_8));
|
||||
cases.put("PDF", "%PDF-1.7\n1 0 obj<<>>endobj\n".getBytes(StandardCharsets.UTF_8));
|
||||
cases.put("PNG signature + garbage", Images.concat(new byte[] {(byte) 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A}, new byte[40]));
|
||||
cases.put("valid PNG + 40 bytes of HTML appended", Images.concat(Images.png(16, 16), "<script>alert(1)</script>".getBytes(StandardCharsets.UTF_8)));
|
||||
cases.put("PNG header claiming 40000x40000", Images.pngClaiming(40000, 40000));
|
||||
cases.put("GIF (sniffed but not on the allow list)", "GIF89a\u0001\u0000\u0001\u0000\u0000\u0000\u0000;".getBytes(StandardCharsets.ISO_8859_1));
|
||||
cases.put("empty", new byte[0]);
|
||||
cases.put("9 MB of PNG-looking zeros", withPngSignature(9 * 1024 * 1024));
|
||||
|
||||
try (Transcript t = new Transcript("06-verifier.txt", "ImageVerifier decisions (cap 8 MiB, 16,000,000 pixels, png/jpeg/webp)")) {
|
||||
t.line("%-45s %-9s %s", "input", "verdict", "detail");
|
||||
for (var e : cases.entrySet()) {
|
||||
String verdict;
|
||||
String detail;
|
||||
try {
|
||||
Sniffed s = verifier.verify(e.getValue());
|
||||
verdict = "ACCEPT";
|
||||
detail = s.contentType() + (s.width() > 0 ? " " + s.width() + "x" + s.height() : "");
|
||||
} catch (RejectedException r) {
|
||||
verdict = "REJECT";
|
||||
detail = r.getMessage();
|
||||
}
|
||||
t.line("%-45s %-9s %s", e.getKey(), verdict, detail);
|
||||
}
|
||||
}
|
||||
assertThat(verifier.verify(Images.png(64, 48)).contentType()).isEqualTo("image/png");
|
||||
assertReason(() -> verifier.verify(cases.get("HTML page")), "not-an-image");
|
||||
assertReason(() -> verifier.verify(cases.get("PNG header claiming 40000x40000")), "too-many-pixels");
|
||||
assertReason(() -> verifier.verify(cases.get("9 MB of PNG-looking zeros")), "too-large");
|
||||
assertReason(() -> verifier.verify(cases.get("PNG signature + garbage")), "undecodable");
|
||||
}
|
||||
|
||||
/** Honest limit, pinned by a test: a valid image with junk appended is still accepted. */
|
||||
@Test
|
||||
void trailingBytesSurvive() {
|
||||
byte[] polyglot = Images.concat(Images.png(16, 16), "<script>alert(1)</script>".getBytes(StandardCharsets.UTF_8));
|
||||
assertThat(verifier.verify(polyglot).contentType()).isEqualTo("image/png");
|
||||
try (Transcript t = new Transcript("07-trailing-bytes.txt", "Sniff + decode does not catch bytes after the image")) {
|
||||
t.line("input : valid 16x16 PNG followed by 25 bytes of <script>");
|
||||
t.line("ImageVerifier : ACCEPT (%s)", verifier.verify(polyglot).contentType());
|
||||
t.line("why it is still OK: the object is stored as Content-Type image/png with");
|
||||
t.line(" Content-Disposition: attachment, so a browser will not run it as HTML.");
|
||||
t.line(" The check is the header, not the sniffing.");
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void base64BeforeDecode() {
|
||||
// 12 MiB of 'A' is valid base64 for 9 MiB of zero bytes; it must be refused from the length alone.
|
||||
String huge = "A".repeat(12 * 1024 * 1024);
|
||||
var service = new ImageService(null, null, null, policy, null, null);
|
||||
try (Transcript t = new Transcript("08-base64-cap.txt", "Oversized base64 is refused from its length, before decoding")) {
|
||||
t.line("payload chars : %d (decodes to about %d bytes; cap is %d)", huge.length(), (long) huge.length() * 3 / 4, policy.maxBytes());
|
||||
try {
|
||||
service.bytesOf(new org.springframework.ai.image.Image(null, huge));
|
||||
} catch (RejectedException e) {
|
||||
t.line("result : %s", e.getMessage());
|
||||
assertThat(e.reason()).isEqualTo("too-large");
|
||||
return;
|
||||
}
|
||||
throw new AssertionError("expected a rejection");
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void fetcherRules() throws Exception {
|
||||
try (FakeImagesServer server = new FakeImagesServer(); Transcript t = new Transcript("09-fetcher.txt", "RemoteFetcher against a local image host")) {
|
||||
ImagePolicy allowLocal = Fixtures_.withHost(policy, server.host()).withMaxBytes(1024 * 1024);
|
||||
RemoteFetcher f = new RemoteFetcher(allowLocal);
|
||||
Map<String, Supplier<byte[]>> cases = new LinkedHashMap<>();
|
||||
cases.put("allowed host, 200", () -> f.fetch(server.fileUrl("ok.png")));
|
||||
cases.put("302 redirect to the same host", () -> f.fetch(server.fileUrl("redirect")));
|
||||
cases.put("endless body (cap 1 MiB)", () -> f.fetch(server.fileUrl("endless")));
|
||||
cases.put("host not on the list", () -> f.fetch("http://localhost:" + server.fileUrl("ok.png").split(":")[2]));
|
||||
cases.put("file:// URL", () -> f.fetch("file:///etc/passwd"));
|
||||
cases.put("http:// when insecure is off", () -> new RemoteFetcher(policy.withFetch(java.util.Set.of(server.host()), false)).fetch(server.fileUrl("ok.png")));
|
||||
cases.put("cloud metadata address", () -> f.fetch("http://169.254.169.254/latest/meta-data/"));
|
||||
t.line("%-34s %s", "case", "result");
|
||||
for (var e : cases.entrySet()) {
|
||||
String res;
|
||||
try {
|
||||
res = "OK " + e.getValue().get().length + " bytes";
|
||||
} catch (RejectedException r) {
|
||||
res = r.getMessage();
|
||||
}
|
||||
t.line("%-34s %s", e.getKey(), res);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static void assertReason(Runnable r, String reason) {
|
||||
assertThatThrownBy(r::run).isInstanceOfSatisfying(RejectedException.class, e -> assertThat(e.reason()).isEqualTo(reason));
|
||||
}
|
||||
|
||||
private static byte[] withPngSignature(int size) {
|
||||
byte[] b = new byte[size];
|
||||
byte[] sig = {(byte) 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A};
|
||||
System.arraycopy(sig, 0, b, 0, sig.length);
|
||||
return b;
|
||||
}
|
||||
|
||||
private static final class Fixtures_ {
|
||||
static ImagePolicy withHost(ImagePolicy p, String host) {
|
||||
return com.ankurm.imagegen.support.Fixtures.withHost(p, host);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package com.ankurm.imagegen;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpClient;
|
||||
import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.time.Duration;
|
||||
import java.util.HashSet;
|
||||
import java.util.Set;
|
||||
|
||||
import com.ankurm.imagegen.support.Images;
|
||||
import com.ankurm.imagegen.support.Minio;
|
||||
import com.ankurm.imagegen.support.Transcript;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import software.amazon.awssdk.services.s3.S3Client;
|
||||
import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
|
||||
import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
|
||||
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
|
||||
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
|
||||
import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
|
||||
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
|
||||
import software.amazon.awssdk.core.sync.RequestBody;
|
||||
|
||||
/** ImageStore against a REAL MinIO server (Docker, Testcontainers). */
|
||||
class StoreMinioTest {
|
||||
|
||||
static Minio minio;
|
||||
static S3Client s3;
|
||||
static ImageStore store;
|
||||
static final String BUCKET = "generated-images";
|
||||
|
||||
@BeforeAll
|
||||
static void start() {
|
||||
minio = new Minio().start();
|
||||
s3 = minio.client();
|
||||
s3.createBucket(CreateBucketRequest.builder().bucket(BUCKET).build());
|
||||
store = new ImageStore(s3, minio.presigner(), BUCKET);
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void stop() {
|
||||
if (minio != null) minio.close();
|
||||
}
|
||||
|
||||
@Test
|
||||
void putAndReadBack() {
|
||||
byte[] png = Images.png(64, 48);
|
||||
Sniffed sniffed = ImageSniffer.sniff(png).orElseThrow();
|
||||
try (Transcript t = new Transcript("05-store-minio.txt", "ImageStore on a real MinIO server")) {
|
||||
t.line("image : %s", minio.imageName());
|
||||
t.line("minio --version : %s", minio.version());
|
||||
t.line("Server header : %s", s3.headBucket(b -> b.bucket(BUCKET)).sdkHttpResponse().firstMatchingHeader("Server").orElse("(no Server header)"));
|
||||
StoredImage a = store.put("user-42", png, sniffed, "a revised prompt");
|
||||
t.line("stored key : %s", a.key());
|
||||
t.line("content type : %s, %d bytes, %dx%d", a.contentType(), a.bytes(), a.width(), a.height());
|
||||
t.line("sha256 : %s", a.sha256());
|
||||
HeadObjectResponse head = s3.headObject(HeadObjectRequest.builder().bucket(BUCKET).key(a.key()).build());
|
||||
t.line("HEAD Content-Type : %s", head.contentType());
|
||||
t.line("HEAD Content-Disposition : %s", head.contentDisposition());
|
||||
t.line("HEAD Cache-Control : %s", head.cacheControl());
|
||||
t.line("HEAD metadata : %s", new java.util.TreeMap<>(head.metadata()));
|
||||
byte[] back = store.read("user-42", a.id(), "png");
|
||||
t.line("read back equal : %s", java.util.Arrays.equals(back, png));
|
||||
assertThat(back).isEqualTo(png);
|
||||
assertThat(head.contentType()).isEqualTo("image/png");
|
||||
assertThat(head.contentDisposition()).isEqualTo("attachment");
|
||||
assertThat(a.key()).matches("generated/user-42/[0-9a-f]{32}\\.png");
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void presignedUrlIsReadableAndSigned() throws Exception {
|
||||
byte[] jpeg = Images.jpeg(32, 32);
|
||||
StoredImage s = store.put("user-42", jpeg, ImageSniffer.sniff(jpeg).orElseThrow(), null);
|
||||
String url = store.presignedGet("user-42", s.id(), "jpg", Duration.ofMinutes(5));
|
||||
HttpClient http = HttpClient.newHttpClient();
|
||||
HttpResponse<byte[]> ok = http.send(HttpRequest.newBuilder(URI.create(url)).build(), HttpResponse.BodyHandlers.ofByteArray());
|
||||
String tampered = url.replace("/user-42/", "/user-43/");
|
||||
HttpResponse<byte[]> bad = http.send(HttpRequest.newBuilder(URI.create(tampered)).build(), HttpResponse.BodyHandlers.ofByteArray());
|
||||
String bare = url.substring(0, url.indexOf('?'));
|
||||
HttpResponse<byte[]> unsigned = http.send(HttpRequest.newBuilder(URI.create(bare)).build(), HttpResponse.BodyHandlers.ofByteArray());
|
||||
try (Transcript t = new Transcript("05b-presigned.txt", "Presigned GET (5 minutes) against MinIO")) {
|
||||
t.line("signed URL, parameters : %s", url.substring(url.indexOf('?') + 1).replaceAll("=[^&]{12,}", "=...").replace("&", " & "));
|
||||
t.line("GET signed URL : %d, %d bytes, Content-Type %s", ok.statusCode(), ok.body().length, ok.headers().firstValue("Content-Type").orElse("?"));
|
||||
t.line("Content-Disposition : %s", ok.headers().firstValue("Content-Disposition").orElse("?"));
|
||||
t.line("same URL, other owner : %d", bad.statusCode());
|
||||
t.line("same path, no signature : %d", unsigned.statusCode());
|
||||
}
|
||||
assertThat(ok.statusCode()).isEqualTo(200);
|
||||
assertThat(bad.statusCode()).isIn(403, 404);
|
||||
assertThat(unsigned.statusCode()).isEqualTo(403);
|
||||
}
|
||||
|
||||
@Test
|
||||
void keysAreRandomAndCallerCannotChooseThem() {
|
||||
byte[] png = Images.png(8, 8);
|
||||
Sniffed s = ImageSniffer.sniff(png).orElseThrow();
|
||||
Set<String> ids = new HashSet<>();
|
||||
for (int i = 0; i < 200; i++) ids.add(store.put("user-42", png, s, null).id());
|
||||
try (Transcript t = new Transcript("10-keys.txt", "Key construction")) {
|
||||
t.line("200 puts of identical bytes -> %d distinct ids (a content hash would give 1)", ids.size());
|
||||
String[] owners = {"../admin", "user-42/../../x", "USER", "", "a b", "user-42\u0000", "x".repeat(40)};
|
||||
for (String o : owners) {
|
||||
String r;
|
||||
try { store.put(o, png, s, null); r = "STORED?!"; } catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("owner %-22s -> %s", "'" + o.replace("\u0000", "\\0") + "'", r);
|
||||
}
|
||||
String[] ids2 = {"../../etc/passwd", "ABCDEF0123456789ABCDEF0123456789", "0123", "0123456789abcdef0123456789abcde/"};
|
||||
for (String id : ids2) {
|
||||
String r;
|
||||
try { store.keyFor("user-42", id, "png"); r = "ACCEPTED?!"; } catch (RejectedException e) { r = e.getMessage(); }
|
||||
t.line("id %-34s -> %s", "'" + id + "'", r);
|
||||
}
|
||||
StoredImage mine = store.put("user-42", png, s, null);
|
||||
String r;
|
||||
try { store.read("user-43", mine.id(), "png"); r = "READ?!"; } catch (NoSuchKeyException e) { r = "NoSuchKeyException (key is rebuilt under user-43/)"; }
|
||||
t.line("user-43 asks for user-42's id -> %s", r);
|
||||
}
|
||||
assertThat(ids).hasSize(200);
|
||||
assertThatThrownBy(() -> store.put("../admin", png, s, null)).isInstanceOf(RejectedException.class);
|
||||
}
|
||||
|
||||
@Test
|
||||
void whatBucketPolicyDoesNotCover() {
|
||||
// The store pins Content-Type from the sniffed bytes. Show the contrast: a naive writer that trusts the sender.
|
||||
byte[] html = "<html><script>alert(1)</script></html>".getBytes();
|
||||
s3.putObject(PutObjectRequest.builder().bucket(BUCKET).key("naive/upload.png").contentType("text/html").build(), RequestBody.fromBytes(html));
|
||||
HeadObjectResponse naive = s3.headObject(HeadObjectRequest.builder().bucket(BUCKET).key("naive/upload.png").build());
|
||||
try (Transcript t = new Transcript("11-naive-vs-pinned.txt", "Trusting the sender's Content-Type vs pinning it")) {
|
||||
t.line("naive writer : key naive/upload.png Content-Type = %s (sender said so; the .png suffix did not matter)", naive.contentType());
|
||||
t.line("ImageStore : the verifier refuses these bytes first: ");
|
||||
try { new ImageVerifier(ImagePolicy.defaults()).verify(html); } catch (RejectedException e) { t.line(" %s", e.getMessage()); }
|
||||
t.line("objects under naive/ : %d", s3.listObjectsV2(ListObjectsV2Request.builder().bucket(BUCKET).prefix("naive/").build()).keyCount());
|
||||
}
|
||||
assertThat(naive.contentType()).isEqualTo("text/html");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
|
||||
import com.sun.net.httpserver.HttpExchange;
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
|
||||
/**
|
||||
* A MOCK of OpenAI's POST /v1/images/generations and /v1/images/edits on a local port, plus a tiny
|
||||
* "image host" for the URL-returning case. The real {@code OpenAiImageModel} (and the OpenAI Java
|
||||
* SDK under it) talk to it over HTTP, so the recorded request bodies are what Spring AI really sends.
|
||||
*
|
||||
* <p>NOT real: the pictures. The server returns a fixed 64x64 gradient PNG whatever the prompt says.
|
||||
* It checks nothing about size, quality or model names, so it cannot tell you whether a provider
|
||||
* would accept the values. The 400 "content_policy_violation" body copies the shape OpenAI's error
|
||||
* format uses; whether and when a real model refuses a prompt is not tested here.
|
||||
*/
|
||||
public final class FakeImagesServer implements AutoCloseable {
|
||||
|
||||
public record Seen(String method, String path, String contentType, String body) {}
|
||||
|
||||
public enum Mode { B64, URL, REFUSE_400 }
|
||||
|
||||
private final HttpServer server;
|
||||
private final List<Seen> seen = new CopyOnWriteArrayList<>();
|
||||
private volatile Mode mode = Mode.B64;
|
||||
private final byte[] png = Images.png(64, 64);
|
||||
|
||||
public FakeImagesServer() throws IOException {
|
||||
server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/v1/images/generations", this::generate);
|
||||
server.createContext("/v1/images/edits", this::generate);
|
||||
server.createContext("/files/ok.png", ex -> bytes(ex, 200, "image/png", png));
|
||||
server.createContext("/files/redirect", ex -> {
|
||||
ex.getResponseHeaders().add("Location", "/files/ok.png");
|
||||
ex.sendResponseHeaders(302, -1);
|
||||
ex.close();
|
||||
});
|
||||
server.createContext("/files/endless", this::endless);
|
||||
server.start();
|
||||
}
|
||||
|
||||
public String baseUrl() { return "http://127.0.0.1:" + server.getAddress().getPort() + "/v1"; }
|
||||
|
||||
public String host() { return "127.0.0.1"; }
|
||||
|
||||
public String fileUrl(String name) { return "http://127.0.0.1:" + server.getAddress().getPort() + "/files/" + name; }
|
||||
|
||||
public List<Seen> seen() { return seen; }
|
||||
|
||||
public void clear() { seen.clear(); }
|
||||
|
||||
public void mode(Mode m) { this.mode = m; }
|
||||
|
||||
@Override
|
||||
public void close() { server.stop(0); }
|
||||
|
||||
private void generate(HttpExchange ex) throws IOException {
|
||||
String body = new String(ex.getRequestBody().readAllBytes(), StandardCharsets.ISO_8859_1);
|
||||
String ct = ex.getRequestHeaders().getFirst("Content-Type");
|
||||
boolean multipart = ct != null && ct.startsWith("multipart/");
|
||||
seen.add(new Seen(ex.getRequestMethod(), ex.getRequestURI().getPath(), ct, multipart ? summarise(body) : body));
|
||||
String json;
|
||||
int status = 200;
|
||||
switch (mode) {
|
||||
case REFUSE_400 -> {
|
||||
status = 400;
|
||||
json = "{\"error\":{\"message\":\"Your request was rejected as a result of our safety system.\","
|
||||
+ "\"type\":\"invalid_request_error\",\"param\":null,\"code\":\"content_policy_violation\"}}";
|
||||
}
|
||||
case URL -> json = "{\"created\":1760000000,\"data\":[{\"url\":\"" + fileUrl("ok.png") + "\",\"revised_prompt\":\"a rewritten prompt\"}]}";
|
||||
default -> json = "{\"created\":1760000000,\"data\":[{\"b64_json\":\"" + Images.b64(png)
|
||||
+ "\",\"revised_prompt\":\"a rewritten prompt\"}],"
|
||||
+ "\"usage\":{\"input_tokens\":12,\"output_tokens\":1056,\"total_tokens\":1068,"
|
||||
+ "\"input_tokens_details\":{\"text_tokens\":12,\"image_tokens\":0}}}";
|
||||
}
|
||||
bytes(ex, status, "application/json", json.getBytes(StandardCharsets.UTF_8));
|
||||
}
|
||||
|
||||
/** Multipart bodies carry binary parts; keep only the part names and text values. */
|
||||
private static String summarise(String body) {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
for (String part : body.split("--[A-Za-z0-9'()+_,./:=?-]{10,}")) {
|
||||
int nameAt = part.indexOf("name=\"");
|
||||
if (nameAt < 0) continue;
|
||||
String name = part.substring(nameAt + 6, part.indexOf('"', nameAt + 6));
|
||||
boolean file = part.contains("filename=");
|
||||
int bodyAt = part.indexOf("\r\n\r\n");
|
||||
String value = bodyAt < 0 ? "" : part.substring(bodyAt + 4).trim();
|
||||
sb.append(name).append('=').append(file ? "<file, " + value.length() + " bytes>" : value).append("; ");
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
|
||||
private void endless(HttpExchange ex) throws IOException {
|
||||
ex.getResponseHeaders().add("Content-Type", "image/png");
|
||||
ex.sendResponseHeaders(200, 0);
|
||||
byte[] chunk = new byte[64 * 1024];
|
||||
try (OutputStream out = ex.getResponseBody()) {
|
||||
for (int i = 0; i < 4096; i++) { // up to 256 MB; the client should hang up long before
|
||||
out.write(chunk);
|
||||
}
|
||||
} catch (IOException ignored) {
|
||||
// the client stopped reading, which is the point
|
||||
}
|
||||
}
|
||||
|
||||
private static void bytes(HttpExchange ex, int status, String type, byte[] data) throws IOException {
|
||||
ex.getResponseHeaders().add("Content-Type", type);
|
||||
ex.sendResponseHeaders(status, data.length);
|
||||
try (OutputStream out = ex.getResponseBody()) {
|
||||
out.write(data);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.Set;
|
||||
|
||||
import com.ankurm.imagegen.ImagePolicy;
|
||||
import com.ankurm.imagegen.ImageService;
|
||||
import com.ankurm.imagegen.ImageStore;
|
||||
import com.ankurm.imagegen.OpenAiOptionsFactory;
|
||||
import com.ankurm.imagegen.PromptGate;
|
||||
import com.ankurm.imagegen.RemoteFetcher;
|
||||
import org.springframework.ai.image.ImageModel;
|
||||
import org.springframework.ai.moderation.ModerationModel;
|
||||
|
||||
public final class Fixtures {
|
||||
|
||||
private Fixtures() {}
|
||||
|
||||
public static byte[] utf8(String s) {
|
||||
return s.getBytes(StandardCharsets.UTF_8);
|
||||
}
|
||||
|
||||
public static ImageService service(ImageModel model, ModerationModel moderation, ImagePolicy policy, ImageStore store) {
|
||||
return new ImageService(model, new OpenAiOptionsFactory("gpt-image-test"), new PromptGate(moderation, policy),
|
||||
policy, new RemoteFetcher(policy), store);
|
||||
}
|
||||
|
||||
public static ImagePolicy withHost(ImagePolicy p, String host) {
|
||||
return p.withFetch(Set.of(host), true);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.awt.Color;
|
||||
import java.awt.image.BufferedImage;
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.Base64;
|
||||
import javax.imageio.ImageIO;
|
||||
|
||||
/** Real, decodable image bytes made in memory - the "generated" pictures in every test. */
|
||||
public final class Images {
|
||||
|
||||
private Images() {}
|
||||
|
||||
public static byte[] png(int w, int h) {
|
||||
return encode(w, h, "png");
|
||||
}
|
||||
|
||||
public static byte[] jpeg(int w, int h) {
|
||||
return encode(w, h, "jpg");
|
||||
}
|
||||
|
||||
private static byte[] encode(int w, int h, String fmt) {
|
||||
BufferedImage img = new BufferedImage(w, h, BufferedImage.TYPE_INT_RGB);
|
||||
for (int y = 0; y < h; y++) {
|
||||
for (int x = 0; x < w; x++) {
|
||||
img.setRGB(x, y, new Color((x * 255) / Math.max(1, w - 1), (y * 255) / Math.max(1, h - 1), 128).getRGB());
|
||||
}
|
||||
}
|
||||
try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
|
||||
ImageIO.write(img, fmt, out);
|
||||
return out.toByteArray();
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
|
||||
/** Smallest RIFF/WEBP header the sniffer recognises (not a decodable WebP; the verifier does not decode WebP). */
|
||||
public static byte[] webpHeaderOnly() {
|
||||
byte[] b = new byte[32];
|
||||
System.arraycopy("RIFF".getBytes(StandardCharsets.US_ASCII), 0, b, 0, 4);
|
||||
System.arraycopy("WEBP".getBytes(StandardCharsets.US_ASCII), 0, b, 8, 4);
|
||||
return b;
|
||||
}
|
||||
|
||||
/** A valid PNG whose IHDR has been rewritten to claim w x h. The pixel data no longer matches. */
|
||||
public static byte[] pngClaiming(int w, int h) {
|
||||
byte[] b = png(8, 8);
|
||||
put32(b, 16, w);
|
||||
put32(b, 20, h);
|
||||
return b;
|
||||
}
|
||||
|
||||
public static byte[] concat(byte[] a, byte[] b) {
|
||||
byte[] out = new byte[a.length + b.length];
|
||||
System.arraycopy(a, 0, out, 0, a.length);
|
||||
System.arraycopy(b, 0, out, a.length, b.length);
|
||||
return out;
|
||||
}
|
||||
|
||||
public static String b64(byte[] b) {
|
||||
return Base64.getEncoder().encodeToString(b);
|
||||
}
|
||||
|
||||
private static void put32(byte[] b, int i, int v) {
|
||||
b[i] = (byte) (v >>> 24);
|
||||
b[i + 1] = (byte) (v >>> 16);
|
||||
b[i + 2] = (byte) (v >>> 8);
|
||||
b[i + 3] = (byte) v;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.net.URI;
|
||||
import java.time.Duration;
|
||||
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import org.testcontainers.utility.DockerImageName;
|
||||
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
import software.amazon.awssdk.services.s3.S3Client;
|
||||
import software.amazon.awssdk.services.s3.S3Configuration;
|
||||
import software.amazon.awssdk.services.s3.presigner.S3Presigner;
|
||||
|
||||
/**
|
||||
* A real MinIO server in Docker (Testcontainers). The image is Chainguard's build of upstream MinIO
|
||||
* because minio/minio and quay.io/minio/minio were not pullable from the build sandbox; the version
|
||||
* is printed in output/05-store-minio.txt. Override with -Dminio.image=... to use another.
|
||||
*/
|
||||
public final class Minio implements AutoCloseable {
|
||||
|
||||
public static final String USER = "minio-test-user";
|
||||
public static final String PASSWORD = "minio-test-password-not-real";
|
||||
|
||||
private final GenericContainer<?> container = new GenericContainer<>(
|
||||
DockerImageName.parse(System.getProperty("minio.image", "chainguard/minio:latest")))
|
||||
.withEnv("MINIO_ROOT_USER", USER)
|
||||
.withEnv("MINIO_ROOT_PASSWORD", PASSWORD)
|
||||
.withCommand("server", "/tmp/data")
|
||||
.withExposedPorts(9000)
|
||||
.waitingFor(Wait.forHttp("/minio/health/live").forPort(9000).withStartupTimeout(Duration.ofSeconds(90)));
|
||||
|
||||
public Minio start() {
|
||||
container.start();
|
||||
return this;
|
||||
}
|
||||
|
||||
public String endpoint() {
|
||||
return "http://" + container.getHost() + ":" + container.getMappedPort(9000);
|
||||
}
|
||||
|
||||
public String imageName() {
|
||||
return container.getDockerImageName();
|
||||
}
|
||||
|
||||
public String version() {
|
||||
try {
|
||||
return container.execInContainer("/usr/bin/minio", "--version").getStdout().lines().findFirst().orElse("?");
|
||||
} catch (Exception e) {
|
||||
return "unknown (" + e.getClass().getSimpleName() + ")";
|
||||
}
|
||||
}
|
||||
|
||||
public S3Client client() {
|
||||
return S3Client.builder()
|
||||
.endpointOverride(URI.create(endpoint()))
|
||||
.region(Region.US_EAST_1)
|
||||
.credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create(USER, PASSWORD)))
|
||||
.serviceConfiguration(S3Configuration.builder().pathStyleAccessEnabled(true).build())
|
||||
.build();
|
||||
}
|
||||
|
||||
public S3Presigner presigner() {
|
||||
return S3Presigner.builder()
|
||||
.endpointOverride(URI.create(endpoint()))
|
||||
.region(Region.US_EAST_1)
|
||||
.credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create(USER, PASSWORD)))
|
||||
.serviceConfiguration(S3Configuration.builder().pathStyleAccessEnabled(true).build())
|
||||
.build();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
container.stop();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.util.ArrayDeque;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Queue;
|
||||
|
||||
import org.springframework.ai.image.Image;
|
||||
import org.springframework.ai.image.ImageGeneration;
|
||||
import org.springframework.ai.image.ImageModel;
|
||||
import org.springframework.ai.image.ImagePrompt;
|
||||
import org.springframework.ai.image.ImageResponse;
|
||||
import org.springframework.ai.openai.metadata.OpenAiImageGenerationMetadata;
|
||||
|
||||
/**
|
||||
* A hand-written {@link ImageModel}: it returns the responses it was given, in order, and records
|
||||
* every {@link ImagePrompt} it receives. No network, no provider, no randomness.
|
||||
*/
|
||||
public final class ScriptedImageModel implements ImageModel {
|
||||
|
||||
private final java.util.LinkedList<ImageResponse> script = new java.util.LinkedList<>();
|
||||
private final List<ImagePrompt> seen = new ArrayList<>();
|
||||
|
||||
public ScriptedImageModel thenB64(byte[]... images) {
|
||||
return thenB64Revised(null, images);
|
||||
}
|
||||
|
||||
public ScriptedImageModel thenB64Revised(String revisedPrompt, byte[]... images) {
|
||||
List<ImageGeneration> gens = new ArrayList<>();
|
||||
for (byte[] b : images) {
|
||||
gens.add(new ImageGeneration(new Image(null, Images.b64(b)),
|
||||
revisedPrompt == null ? null : new OpenAiImageGenerationMetadata(revisedPrompt)));
|
||||
}
|
||||
script.add(new ImageResponse(gens));
|
||||
return this;
|
||||
}
|
||||
|
||||
public ScriptedImageModel thenRawB64(String b64) {
|
||||
script.add(new ImageResponse(List.of(new ImageGeneration(new Image(null, b64)))));
|
||||
return this;
|
||||
}
|
||||
|
||||
public ScriptedImageModel thenUrl(String url) {
|
||||
script.add(new ImageResponse(List.of(new ImageGeneration(new Image(url, null)))));
|
||||
return this;
|
||||
}
|
||||
|
||||
public ScriptedImageModel thenThrow(RuntimeException e) {
|
||||
script.add(null);
|
||||
throwables.add(e);
|
||||
return this;
|
||||
}
|
||||
|
||||
private final Queue<RuntimeException> throwables = new java.util.LinkedList<>();
|
||||
|
||||
@Override
|
||||
public ImageResponse call(ImagePrompt prompt) {
|
||||
seen.add(prompt);
|
||||
if (script.isEmpty()) {
|
||||
throw new IllegalStateException("ScriptedImageModel has no response left for call #" + seen.size());
|
||||
}
|
||||
ImageResponse next = script.poll();
|
||||
if (next == null) {
|
||||
throw throwables.poll();
|
||||
}
|
||||
return next;
|
||||
}
|
||||
|
||||
public List<ImagePrompt> seen() {
|
||||
return seen;
|
||||
}
|
||||
|
||||
public int calls() {
|
||||
return seen.size();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.springframework.ai.moderation.Categories;
|
||||
import org.springframework.ai.moderation.Generation;
|
||||
import org.springframework.ai.moderation.Moderation;
|
||||
import org.springframework.ai.moderation.ModerationModel;
|
||||
import org.springframework.ai.moderation.ModerationPrompt;
|
||||
import org.springframework.ai.moderation.ModerationResponse;
|
||||
import org.springframework.ai.moderation.ModerationResult;
|
||||
|
||||
/**
|
||||
* A hand-written {@link ModerationModel}. A rule maps the prompt text to a flagged category name
|
||||
* (or null for "clean"). It is a stand-in for OpenAiModerationModel; it judges nothing.
|
||||
*/
|
||||
public final class ScriptedModerationModel implements ModerationModel {
|
||||
|
||||
private final Function<String, String> rule;
|
||||
private final List<String> seen = new ArrayList<>();
|
||||
private RuntimeException failure;
|
||||
|
||||
public ScriptedModerationModel(Function<String, String> rule) {
|
||||
this.rule = rule;
|
||||
}
|
||||
|
||||
public static ScriptedModerationModel flagging(String word, String category) {
|
||||
return new ScriptedModerationModel(p -> p.toLowerCase().contains(word) ? category : null);
|
||||
}
|
||||
|
||||
public ScriptedModerationModel failingWith(RuntimeException e) {
|
||||
this.failure = e;
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ModerationResponse call(ModerationPrompt prompt) {
|
||||
String text = prompt.getInstructions().getText();
|
||||
seen.add(text);
|
||||
if (failure != null) throw failure;
|
||||
String category = rule.apply(text);
|
||||
Categories.Builder c = Categories.builder();
|
||||
if ("violence".equals(category)) c.violence(true);
|
||||
if ("sexual".equals(category)) c.sexual(true);
|
||||
if ("hate".equals(category)) c.hate(true);
|
||||
ModerationResult r = ModerationResult.builder().flagged(category != null).categories(c.build()).build();
|
||||
Moderation m = Moderation.builder().id("scripted").model("scripted").results(List.of(r)).build();
|
||||
return new ModerationResponse(new Generation(m));
|
||||
}
|
||||
|
||||
public List<String> seen() {
|
||||
return seen;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package com.ankurm.imagegen.support;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.PrintWriter;
|
||||
import java.io.StringWriter;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
/**
|
||||
* Writes a numbered transcript under {@code output/} (module root, not {@code docs/}) and echoes
|
||||
* it to the console. Every console block quoted in the article comes out of one of these files.
|
||||
*/
|
||||
public final class Transcript implements AutoCloseable {
|
||||
|
||||
private final Path path;
|
||||
private final StringWriter buffer = new StringWriter();
|
||||
private final PrintWriter out = new PrintWriter(buffer);
|
||||
|
||||
public Transcript(String fileName, String title) {
|
||||
this.path = Path.of("output", fileName);
|
||||
out.println("# " + title);
|
||||
out.println();
|
||||
}
|
||||
|
||||
public Transcript line(String format, Object... args) {
|
||||
out.println(args.length == 0 ? format : String.format(format, args));
|
||||
return this;
|
||||
}
|
||||
|
||||
public Transcript blank() {
|
||||
out.println();
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
out.flush();
|
||||
try {
|
||||
Files.createDirectories(path.getParent());
|
||||
Files.writeString(path, buffer.toString());
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException("could not write " + path, e);
|
||||
}
|
||||
System.out.print(buffer);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user