Add advisors module: custom logging, PII redaction and token-budget advisors
Tests pin down chain ordering (including ties), BaseAdvisor stream behaviour, redaction order versus memory and logging, the tool loop, and how a refusal surfaces on calls, streams and over HTTP. Co-Authored-By: Claude Sonnet 5 <[email protected]> Claude-Session: https://claude.ai/code/session_01Ja4jkzrbQ4LQZBNrb5mkZE
This commit is contained in:
Claude
@@ -13,5 +13,6 @@ Runnable companion code for the Spring AI articles on [ankurm.com](https://ankur
|
|||||||
| [`structured-output/`](structured-output) | `ChatClient.entity()` mapping LLM responses to Java records, lists and maps; `StructuredOutputValidationAdvisor` retrying non-conforming JSON with a real enum-constrained schema, including a captured run that exhausts every retry without throwing. Spring Boot 4.1.1, Spring AI 2.0.1, Java 25. | [Structured Output in Spring AI 2.0](https://ankurm.com/spring-ai-2-0-structured-output/) |
|
| [`structured-output/`](structured-output) | `ChatClient.entity()` mapping LLM responses to Java records, lists and maps; `StructuredOutputValidationAdvisor` retrying non-conforming JSON with a real enum-constrained schema, including a captured run that exhausts every retry without throwing. Spring Boot 4.1.1, Spring AI 2.0.1, Java 25. | [Structured Output in Spring AI 2.0](https://ankurm.com/spring-ai-2-0-structured-output/) |
|
||||||
| [`ollama-local/`](ollama-local) | Chat and embeddings against a real local `qwen2.5:0.5b`/`all-minilm`, no API key, driven by a Testcontainers-managed Ollama container started from a baked image; a confirmed model unload via `keep_alive: 0` and `/api/ps`, not a scripted model anywhere. Spring Boot 4.1.1, Spring AI 2.0.1, Testcontainers 2.0.5, Java 25. | [Run LLMs Locally with Spring AI and Ollama](https://ankurm.com/spring-ai-2-0-ollama-local/) |
|
| [`ollama-local/`](ollama-local) | Chat and embeddings against a real local `qwen2.5:0.5b`/`all-minilm`, no API key, driven by a Testcontainers-managed Ollama container started from a baked image; a confirmed model unload via `keep_alive: 0` and `/api/ps`, not a scripted model anywhere. Spring Boot 4.1.1, Spring AI 2.0.1, Testcontainers 2.0.5, Java 25. | [Run LLMs Locally with Spring AI and Ollama](https://ankurm.com/spring-ai-2-0-ollama-local/) |
|
||||||
| [`chat-memory/`](chat-memory) | `MessageChatMemoryAdvisor`, `MessageWindowChatMemory`, the JDBC and Redis `ChatMemoryRepository`, per-user conversation IDs and a token-budget memory of our own, with the traps reproduced against a real PostgreSQL 16 and Redis Stack: a 36-character `conversation_id`, tool messages dropped on save, concurrent writers, a 1.x table under the 2.0 repository, and a Redis repository that silently steps aside for a custom `ChatMemory`. Spring Boot 4.1.1, Spring AI 2.0.1, Java 25. | [Chat Memory in Spring AI 2.0: JDBC, Redis and Windowed Conversations](https://ankurm.com/spring-ai-2-0-chat-memory-jdbc-redis-windowed-conversations/) |
|
| [`chat-memory/`](chat-memory) | `MessageChatMemoryAdvisor`, `MessageWindowChatMemory`, the JDBC and Redis `ChatMemoryRepository`, per-user conversation IDs and a token-budget memory of our own, with the traps reproduced against a real PostgreSQL 16 and Redis Stack: a 36-character `conversation_id`, tool messages dropped on save, concurrent writers, a 1.x table under the 2.0 repository, and a Redis repository that silently steps aside for a custom `ChatMemory`. Spring Boot 4.1.1, Spring AI 2.0.1, Java 25. | [Chat Memory in Spring AI 2.0: JDBC, Redis and Windowed Conversations](https://ankurm.com/spring-ai-2-0-chat-memory-jdbc-redis-windowed-conversations/) |
|
||||||
|
| [`advisors/`](advisors) | Three custom advisors -- a logger, a PII redactor (with a stream-safe restore) and a per-request / per-user token budget -- and tests for how the chain is ordered, what `BaseAdvisor` does on a stream, where an advisor sits relative to memory and the tool loop, and what a refusal looks like on a call, a stream and over HTTP (429). A recording stub model, no live model. Spring Boot 4.1.1, Spring AI 2.0.1, Java 25. | [Writing Custom Advisors in Spring AI 2.0: Logging, PII Redaction and Token Budgets](https://ankurm.com/spring-ai-2-0-custom-advisors-logging-pii-redaction-token-budgets/) |
|
||||||
|
|
||||||
Upgrading from Spring AI 1.x: [migration guide](https://ankurm.com/spring-ai-1-to-2-migration-guide/).
|
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,65 @@
|
|||||||
|
# advisors
|
||||||
|
|
||||||
|
Companion code for [Writing Custom Advisors in Spring AI 2.0: Logging, PII Redaction and Token Budgets](https://ankurm.com/spring-ai-2-0-custom-advisors-logging-pii-redaction-token-budgets/), part of the [Spring AI series](../README.md) on ankurm.com.
|
||||||
|
|
||||||
|
Three advisors of our own -- a logger, a PII redactor and a token budget -- plus the tests that pin down how the advisor chain orders them, what happens on a stream, and what a caller sees when one of them refuses a request.
|
||||||
|
|
||||||
|
There is no live model anywhere. [`RecordingModel`](src/test/java/com/ankurm/advisors/support/RecordingModel.java) records every prompt it is sent and answers with a fixed sentence, which is exactly what an advisor test needs: *what reached the model, and what came back*. No database, no Docker, no API key.
|
||||||
|
|
||||||
|
## Versions
|
||||||
|
|
||||||
|
| Component | Version |
|
||||||
|
|---|---|
|
||||||
|
| Spring Boot | 4.1.1 |
|
||||||
|
| Spring AI | 2.0.1 (`spring-ai-client-chat` 2.0.1) |
|
||||||
|
| Reactor Core | 3.8.7 |
|
||||||
|
| JTokkit (token estimates) | 1.1.0 |
|
||||||
|
| Java | 25 (LTS) |
|
||||||
|
|
||||||
|
## Quickstart
|
||||||
|
|
||||||
|
```bash
|
||||||
|
scripts/run-all.sh # runs all 21 tests and regenerates output/01 .. 21
|
||||||
|
```
|
||||||
|
|
||||||
|
Two consecutive runs produce byte-identical files. To run the app itself: `OPENAI_API_KEY=... mvn spring-boot:run`, then `POST /chat` with an `X-User` header and `{"text": "..."}`.
|
||||||
|
|
||||||
|
## What's here
|
||||||
|
|
||||||
|
| File | What it shows |
|
||||||
|
|---|---|
|
||||||
|
| [`advisor/Orders.java`](src/main/java/com/ankurm/advisors/advisor/Orders.java) | The order numbers, and the layers they sit between (memory +200, tool loop +300, model last) |
|
||||||
|
| [`advisor/LoggingAdvisor.java`](src/main/java/com/ankurm/advisors/advisor/LoggingAdvisor.java) | Request / response / stream / failure lines; content off by default |
|
||||||
|
| [`advisor/PiiRedactor.java`](src/main/java/com/ankurm/advisors/advisor/PiiRedactor.java) | Email, Luhn-checked card number and phone patterns; `<EMAIL_1>` placeholders |
|
||||||
|
| [`advisor/PiiRedactionAdvisor.java`](src/main/java/com/ankurm/advisors/advisor/PiiRedactionAdvisor.java) | Redact on the way in, restore on the way out, on calls and on streams (buffers a split placeholder) |
|
||||||
|
| [`advisor/TokenBudgetAdvisor.java`](src/main/java/com/ankurm/advisors/advisor/TokenBudgetAdvisor.java) | Per-request and per-user limits, checked before the model is called |
|
||||||
|
| [`advisor/Texts.java`](src/main/java/com/ankurm/advisors/advisor/Texts.java) | Rewriting request and response text with the `mutate()` copies |
|
||||||
|
| [`config/AdvisorConfig.java`](src/main/java/com/ankurm/advisors/config/AdvisorConfig.java) | The advisors as beans and one `ChatClient` that uses them with Spring AI's memory advisor |
|
||||||
|
| [`web/ChatController.java`](src/main/java/com/ankurm/advisors/web/ChatController.java) | `POST /chat`; the user comes from a header only so the tests need no login |
|
||||||
|
|
||||||
|
## Output files
|
||||||
|
|
||||||
|
Every file is written by the test named in the right column, and every console block in the article is quoted from one of them.
|
||||||
|
|
||||||
|
| File | Written by |
|
||||||
|
|---|---|
|
||||||
|
| `01-chain-order.txt`, `02-chain-order-ties.txt`, `03-chain-contents.txt` | `ChainOrderTest` |
|
||||||
|
| `04-base-advisor-stream.txt` | `BaseAdvisorStreamTest` |
|
||||||
|
| `05-logging-advisor.txt`, `06-logging-failure.txt` | `LoggingAdvisorTest` |
|
||||||
|
| `07-pii-redaction.txt`, `08-pii-limits.txt`, `09-pii-stream-boundary.txt` | `PiiRedactionTest` |
|
||||||
|
| `10-pii-order.txt` | `PiiOrderTest` |
|
||||||
|
| `11-pii-multi-turn.txt` | `PiiMultiTurnTest` |
|
||||||
|
| `12-token-budget.txt`, `13-token-budget-stream.txt`, `14-token-budget-stream-refusal.txt` | `TokenBudgetTest` |
|
||||||
|
| `15-tool-loop-order.txt` | `ToolLoopOrderTest` |
|
||||||
|
| `16-error-propagation.txt` | `ErrorPropagationTest` |
|
||||||
|
| `17-context.txt` | `ContextAndImmutabilityTest` |
|
||||||
|
| `18-unit-test-no-model.txt` | `UnitTestingWithoutModelTest` |
|
||||||
|
| `19-web-429.txt` | `AdvisorWebTest` |
|
||||||
|
| `20-built-in-orders.txt` | `BuiltInOrdersTest` |
|
||||||
|
| `21-tool-loop-budget.txt` | `ToolLoopBudgetTest` |
|
||||||
|
|
||||||
|
`TokenBudgetTest.aStreamIsRefusedAsAnErrorSignal` makes Spring AI's `MessageAggregator` log an `ERROR ... Aggregation Error` stack trace on the console. That is the refusal being reported, not a test failure.
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
JDK 25 and Maven.
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
# Advisor chain: what decides who runs first
|
||||||
|
|
||||||
|
registered as C300, A100, B200 (the number is HIGHEST_PRECEDENCE + n)
|
||||||
|
call: A100> B200> C300> model C300< B200< A100<
|
||||||
|
stream: A100> B200> C300> model C300< B200< A100<
|
||||||
|
|
||||||
|
plus one advisor added on the request with .advisors(...) at n=150:
|
||||||
|
call: A100> R150> B200> C300> model C300< B200< R150< A100<
|
||||||
@@ -0,0 +1,6 @@
|
|||||||
|
# Two advisors with the same order number
|
||||||
|
|
||||||
|
registered X, Y -> Y> X> model X< Y<
|
||||||
|
registered Y, X -> X> Y> model Y< X<
|
||||||
|
|
||||||
|
For equal numbers the advisor registered LAST runs first (outermost).
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
# What is actually in the chain
|
||||||
|
|
||||||
|
one custom advisor at HIGHEST_PRECEDENCE + 100; the chain a call runs through:
|
||||||
|
Dump HIGHEST_PRECEDENCE + 100
|
||||||
|
Tool Calling Advisor HIGHEST_PRECEDENCE + 300
|
||||||
|
call LOWEST_PRECEDENCE
|
||||||
|
|
||||||
|
and the chain a stream runs through:
|
||||||
|
Dump HIGHEST_PRECEDENCE + 100
|
||||||
|
Tool Calling Advisor HIGHEST_PRECEDENCE + 300
|
||||||
|
stream LOWEST_PRECEDENCE
|
||||||
|
|
||||||
|
the same, with ToolCallingAdvisor.builder().build() added by hand:
|
||||||
|
Dump HIGHEST_PRECEDENCE + 100
|
||||||
|
Tool Calling Advisor HIGHEST_PRECEDENCE + 300
|
||||||
|
call LOWEST_PRECEDENCE
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
# BaseAdvisor on a call and on a stream
|
||||||
|
|
||||||
|
call: before x1, after x1
|
||||||
|
before ran on thread: the caller's thread
|
||||||
|
|
||||||
|
stream: the model streamed 5 chunks: one |two |thre|e fo|ur
|
||||||
|
before x1, after x1
|
||||||
|
after saw only: "ur"
|
||||||
|
before ran on thread: a boundedElastic worker, not the caller's thread
|
||||||
|
|
||||||
|
stream whose last chunk has no finish reason: before x1, after x0
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# LoggingAdvisor: a call and a stream
|
||||||
|
|
||||||
|
call, content logging off (the default):
|
||||||
|
[LoggingAdvisor] request messages=2 roles=SU
|
||||||
|
[LoggingAdvisor] response chars=14 tokens=12+3 took 5 ms
|
||||||
|
|
||||||
|
stream:
|
||||||
|
[LoggingAdvisor] request messages=2 roles=SU
|
||||||
|
[LoggingAdvisor] complete 3 chunks, 14 chars tokens=12+3 took 5 ms
|
||||||
|
|
||||||
|
call, content logging on:
|
||||||
|
[LoggingAdvisor] request messages=1 roles=U last="Where is my refund?"
|
||||||
|
[LoggingAdvisor] response chars=14 tokens=5+3 text="Refund issued." took 5 ms
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
# LoggingAdvisor when the model call fails
|
||||||
|
|
||||||
|
caller got: IllegalStateException: provider returned 503
|
||||||
|
[LoggingAdvisor] request messages=1 roles=U
|
||||||
|
[LoggingAdvisor] failed IllegalStateException after 5 ms
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
# PiiRedactionAdvisor on a call
|
||||||
|
|
||||||
|
caller sends: Hi, I'm Priya. Email [email protected] or call +91 98765 43210. Card 4111 1111 1111 1111, order 1234 5678 9012 3456, and again [email protected].
|
||||||
|
model was sent: Hi, I'm Priya. Email <EMAIL_1> or call <PHONE_1>. Card <CARD_1>, order 1234 5678 9012 3456, and again <EMAIL_1>.
|
||||||
|
caller receives: You said: Hi, I'm Priya. Email [email protected] or call +91 98765 43210. Card 4111 1111 1111 1111, order 1234 5678 9012 3456, and again [email protected].
|
||||||
|
|
||||||
|
with restore switched off, the caller receives:
|
||||||
|
You said: Hi, I'm Priya. Email <EMAIL_1> or call <PHONE_1>. Card <CARD_1>, order 1234 5678 9012 3456, and again <EMAIL_1>.
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
# What pattern-based redaction misses
|
||||||
|
|
||||||
|
in: My name is Priya Sharma and I live at 14 Hill Road, Bandra, Mumbai 400050.
|
||||||
|
out: My name is Priya Sharma and I live at 14 Hill Road, Bandra, Mumbai 400050.
|
||||||
|
|
||||||
|
in: Passport N1234567, PAN ABCDE1234F.
|
||||||
|
out: Passport N1234567, PAN ABCDE1234F.
|
||||||
|
|
||||||
|
in: Write to priya (at) example (dot) com
|
||||||
|
out: Write to priya (at) example (dot) com
|
||||||
|
|
||||||
|
in: Card 4111-1111-1111-1112 and order 1234 5678 9012 3456
|
||||||
|
out: Card 4111-1111-1111-1112 and order 1234 5678 9012 3456
|
||||||
|
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
# Restoring placeholders in a stream
|
||||||
|
|
||||||
|
the model streams 5-character chunks: Sure,| I wi|ll wr|ite t|o <EM|AIL_1|> now|.
|
||||||
|
|
||||||
|
restore each chunk on its own:
|
||||||
|
chunks: Sure,| I wi|ll wr|ite t|o <EM|AIL_1|> now|.
|
||||||
|
joined: Sure, I will write to <EMAIL_1> now.
|
||||||
|
|
||||||
|
PiiRedactionAdvisor (holds back from an unfinished "<"):
|
||||||
|
chunks: Sure,| I wi|ll wr|ite t|o |[email protected] now|.
|
||||||
|
joined: Sure, I will write to [email protected] now.
|
||||||
@@ -0,0 +1,19 @@
|
|||||||
|
# Redaction order versus logging and memory
|
||||||
|
|
||||||
|
memory advisor is fixed at HIGHEST_PRECEDENCE + 200
|
||||||
|
user says: My email is [email protected]
|
||||||
|
|
||||||
|
A redaction +100, logging +400 (redaction outside both):
|
||||||
|
log line saw: "My email is <EMAIL_1>"
|
||||||
|
memory stored: "My email is <EMAIL_1>"
|
||||||
|
model was sent: "My email is <EMAIL_1>"
|
||||||
|
|
||||||
|
B redaction +100, logging +50 (logging outside redaction):
|
||||||
|
log line saw: "My email is [email protected]"
|
||||||
|
memory stored: "My email is <EMAIL_1>"
|
||||||
|
model was sent: "My email is <EMAIL_1>"
|
||||||
|
|
||||||
|
C redaction +300 (inside the memory advisor), logging +400:
|
||||||
|
log line saw: "My email is <EMAIL_1>"
|
||||||
|
memory stored: "My email is [email protected]"
|
||||||
|
model was sent: "My email is <EMAIL_1>"
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
# Placeholders across two turns with memory
|
||||||
|
|
||||||
|
turn 1: My email is [email protected]
|
||||||
|
turn 2: Also cc [email protected]
|
||||||
|
|
||||||
|
numbering restarts on every request:
|
||||||
|
model was sent on turn 2: U:My email is <EMAIL_1> | A:You said: My email is <EMAIL_1> | U:Also cc <EMAIL_1>
|
||||||
|
caller receives: You said: Also cc [email protected]
|
||||||
|
|
||||||
|
numbering kept per conversation (the default):
|
||||||
|
model was sent on turn 2: U:My email is <EMAIL_1> | A:You said: My email is <EMAIL_1> | U:Also cc <EMAIL_2>
|
||||||
|
caller receives: You said: Also cc [email protected]
|
||||||
|
|
||||||
|
what the memory stores (placeholders, never the addresses):
|
||||||
|
USER My email is <EMAIL_1>
|
||||||
|
ASSISTANT You said: My email is <EMAIL_1>
|
||||||
|
USER Also cc <EMAIL_2>
|
||||||
|
ASSISTANT You said: Also cc <EMAIL_2>
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
# TokenBudgetAdvisor: 40 tokens per request, 60 per user
|
||||||
|
|
||||||
|
alice asks short questions; "spent" is the usage the model reported:
|
||||||
|
call 1: answered, spent=13, model calls=1
|
||||||
|
call 2: answered, spent=26, model calls=2
|
||||||
|
call 3: answered, spent=39, model calls=3
|
||||||
|
call 4: answered, spent=52, model calls=4
|
||||||
|
call 5: answered, spent=65, model calls=5
|
||||||
|
call 6: refused (user alice has used 65 of 60 tokens), spent=65, model calls=5
|
||||||
|
|
||||||
|
bob pastes a stack trace of 103 estimated tokens:
|
||||||
|
refused: prompt is about 103 tokens, the limit per request is 40
|
||||||
|
model calls: 5 (was 5), bob's spent: 0
|
||||||
|
|
||||||
|
bob then asks a short question: answered, bob spent=13, alice spent=65
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
# Token accounting on a stream
|
||||||
|
|
||||||
|
same question, same answer:
|
||||||
|
call, usage from the response: 9 tokens
|
||||||
|
stream, usage on the last chunk: 9 tokens
|
||||||
|
stream, no usage reported (estimated): 9 tokens
|
||||||
|
|
||||||
|
A stream that reports no usage is billed by the provider all the same. The estimate above
|
||||||
|
matches only because this scripted model and the advisor use the same tokenizer.
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
# A refusal on the stream path
|
||||||
|
|
||||||
|
subscriber got: TokenBudgetExceededException
|
||||||
|
message: prompt is about 13 tokens, the limit per request is 5
|
||||||
|
model calls: 0
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# Where an advisor sits relative to the tool loop
|
||||||
|
|
||||||
|
one question, one tool call, so the model is called twice (ToolCallingAdvisor is at +300)
|
||||||
|
|
||||||
|
logging advisor at +250 (outside the tool loop):
|
||||||
|
[LoggingAdvisor] request messages=1 roles=U
|
||||||
|
[LoggingAdvisor] response chars=20 took 5 ms
|
||||||
|
|
||||||
|
logging advisor at +400 (inside the tool loop):
|
||||||
|
[LoggingAdvisor] request messages=1 roles=U
|
||||||
|
[LoggingAdvisor] response chars=0 took 5 ms
|
||||||
|
[LoggingAdvisor] request messages=3 roles=UAT
|
||||||
|
[LoggingAdvisor] response chars=20 took 5 ms
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
# Where a refusal surfaces
|
||||||
|
|
||||||
|
an advisor that throws IllegalStateException("refused"), model calls counted afterwards
|
||||||
|
|
||||||
|
plain advisor, call failed while running it with IllegalStateException: refused
|
||||||
|
plain advisor, stream, throws eagerly failed while running it with IllegalStateException: refused
|
||||||
|
plain advisor, stream, error inside Flux.defer failed while running it with IllegalStateException: refused
|
||||||
|
BaseAdvisor.before, call failed while running it with IllegalStateException: refused
|
||||||
|
BaseAdvisor.before, stream failed while running it with IllegalStateException: Stream processing failed (cause: IllegalStateException: refused)
|
||||||
|
|
||||||
|
model calls: 0
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
# Advisor context: request down, response up
|
||||||
|
|
||||||
|
outer: original request has tenant=null, the copy it forwards has tenant=acme
|
||||||
|
inner: request context has tenant=acme
|
||||||
|
outer: response context has verdict=clean, tenant=acme
|
||||||
|
|
||||||
|
req.context().put(...) on the incoming request: allowed
|
||||||
@@ -0,0 +1,6 @@
|
|||||||
|
# Testing an advisor with a stub chain
|
||||||
|
|
||||||
|
what the caller sent: Mail [email protected] please
|
||||||
|
what reached the rest: Mail <EMAIL_1> please
|
||||||
|
what the caller got back: Noted: Mail [email protected] please
|
||||||
|
caller's request unchanged: true
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
# Over HTTP: 200, then 429
|
||||||
|
|
||||||
|
limits for this run: 60 tokens per request, 30 per user
|
||||||
|
dana request 1 -> HTTP 200, model calls so far: 1
|
||||||
|
dana request 2 -> HTTP 200, model calls so far: 2
|
||||||
|
dana request 3 -> HTTP 429, model calls so far: 2
|
||||||
|
dana request 4 -> HTTP 429, model calls so far: 2
|
||||||
|
|
||||||
|
erin pastes a long stack trace -> HTTP 429, model calls: 2 (was 2)
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
# Built-in advisor orders in 2.0.1
|
||||||
|
|
||||||
|
SimpleLoggerAdvisor getOrder() = 0
|
||||||
|
SafeGuardAdvisor getOrder() = 0
|
||||||
|
MessageChatMemoryAdvisor getOrder() = -2147483448 (HIGHEST_PRECEDENCE + 200)
|
||||||
|
Tool Calling Advisor getOrder() = -2147483348 (HIGHEST_PRECEDENCE + 300)
|
||||||
|
|
||||||
|
Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER = HIGHEST_PRECEDENCE + 200
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
# Token budget and the tool loop
|
||||||
|
|
||||||
|
the model reports 100+10 tokens for round 1 (asks for the tool) and 130+20 for round 2 (answers)
|
||||||
|
so the provider would bill 260 tokens for this one question
|
||||||
|
|
||||||
|
budget advisor at +250 (outside the tool loop): recorded 260
|
||||||
|
budget advisor at +400 (inside the tool loop): recorded 260
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
<?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.1</version>
|
||||||
|
<relativePath/>
|
||||||
|
</parent>
|
||||||
|
|
||||||
|
<groupId>com.ankurm</groupId>
|
||||||
|
<artifactId>advisors</artifactId>
|
||||||
|
<version>1.0.0</version>
|
||||||
|
<name>advisors</name>
|
||||||
|
<description>Custom Spring AI 2.0 advisors: audit logging, PII redaction and token budgets, with the ordering and streaming traps reproduced offline.</description>
|
||||||
|
|
||||||
|
<properties>
|
||||||
|
<java.version>25</java.version>
|
||||||
|
<spring-ai.version>2.0.1</spring-ai.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>
|
||||||
|
</dependencies>
|
||||||
|
</dependencyManagement>
|
||||||
|
|
||||||
|
<dependencies>
|
||||||
|
<dependency>
|
||||||
|
<groupId>org.springframework.boot</groupId>
|
||||||
|
<artifactId>spring-boot-starter-webmvc</artifactId>
|
||||||
|
</dependency>
|
||||||
|
<dependency>
|
||||||
|
<groupId>org.springframework.ai</groupId>
|
||||||
|
<artifactId>spring-ai-starter-model-openai</artifactId>
|
||||||
|
</dependency>
|
||||||
|
|
||||||
|
<dependency>
|
||||||
|
<groupId>org.springframework.boot</groupId>
|
||||||
|
<artifactId>spring-boot-starter-test</artifactId>
|
||||||
|
<scope>test</scope>
|
||||||
|
</dependency>
|
||||||
|
</dependencies>
|
||||||
|
|
||||||
|
<build>
|
||||||
|
<plugins>
|
||||||
|
<plugin>
|
||||||
|
<groupId>org.springframework.boot</groupId>
|
||||||
|
<artifactId>spring-boot-maven-plugin</artifactId>
|
||||||
|
</plugin>
|
||||||
|
<plugin>
|
||||||
|
<groupId>org.apache.maven.plugins</groupId>
|
||||||
|
<artifactId>maven-surefire-plugin</artifactId>
|
||||||
|
<configuration>
|
||||||
|
<argLine>-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/ (01-21). The test suite writes all of them itself through
|
||||||
|
# the Transcript helper. No database, no Docker and no API key is needed: the "model" is a
|
||||||
|
# recording stub, so what the tests prove is what reached it and what came back.
|
||||||
|
set -euo pipefail
|
||||||
|
cd "$(dirname "$0")/.."
|
||||||
|
rm -rf target
|
||||||
|
mvn -q -B test 2>&1 | grep -E "Tests run:|BUILD|FAIL" || true
|
||||||
|
ls output
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import org.springframework.boot.SpringApplication;
|
||||||
|
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
||||||
|
|
||||||
|
@SpringBootApplication
|
||||||
|
public class AdvisorsApplication {
|
||||||
|
|
||||||
|
public static void main(String[] args) {
|
||||||
|
SpringApplication.run(AdvisorsApplication.class, args);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.function.Consumer;
|
||||||
|
import java.util.function.LongSupplier;
|
||||||
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
|
import java.util.concurrent.atomic.AtomicLong;
|
||||||
|
import java.util.concurrent.atomic.AtomicReference;
|
||||||
|
|
||||||
|
import org.slf4j.Logger;
|
||||||
|
import org.slf4j.LoggerFactory;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.messages.Message;
|
||||||
|
import org.springframework.ai.chat.metadata.Usage;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* One log line going in, one coming out, with latency and token usage. It wraps {@code
|
||||||
|
* chain.nextCall} / {@code nextStream}, so "took" is the time of everything <em>inside</em> this
|
||||||
|
* advisor in the chain: the model, plus any advisor with a higher order number.
|
||||||
|
*
|
||||||
|
* <p>Message text is off by default. When it is on, what you read is whatever reaches <em>this</em>
|
||||||
|
* position in the chain, which is the point of the ordering experiments in the article.
|
||||||
|
*/
|
||||||
|
public final class LoggingAdvisor implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
private static final Logger LOG = LoggerFactory.getLogger(LoggingAdvisor.class);
|
||||||
|
|
||||||
|
private final String name;
|
||||||
|
|
||||||
|
private final int order;
|
||||||
|
|
||||||
|
private final boolean logContent;
|
||||||
|
|
||||||
|
private final Consumer<String> sink;
|
||||||
|
|
||||||
|
private final LongSupplier nanoClock;
|
||||||
|
|
||||||
|
public LoggingAdvisor(String name, int order, boolean logContent, Consumer<String> sink, LongSupplier nanoClock) {
|
||||||
|
this.name = name;
|
||||||
|
this.order = order;
|
||||||
|
this.logContent = logContent;
|
||||||
|
this.sink = sink;
|
||||||
|
this.nanoClock = nanoClock;
|
||||||
|
}
|
||||||
|
|
||||||
|
public static LoggingAdvisor toSlf4j(int order, boolean logContent) {
|
||||||
|
return new LoggingAdvisor("LoggingAdvisor", order, logContent, LOG::info, System::nanoTime);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return name;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return order;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
long start = nanoClock.getAsLong();
|
||||||
|
sink.accept("[" + name + "] request " + describe(request));
|
||||||
|
try {
|
||||||
|
ChatClientResponse response = chain.nextCall(request);
|
||||||
|
sink.accept("[" + name + "] response " + describe(response) + " took " + millis(start) + " ms");
|
||||||
|
return response;
|
||||||
|
}
|
||||||
|
catch (RuntimeException e) {
|
||||||
|
sink.accept("[" + name + "] failed " + e.getClass().getSimpleName() + " after " + millis(start) + " ms");
|
||||||
|
throw e;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
return Flux.defer(() -> {
|
||||||
|
long start = nanoClock.getAsLong();
|
||||||
|
sink.accept("[" + name + "] request " + describe(request));
|
||||||
|
AtomicInteger chunks = new AtomicInteger();
|
||||||
|
AtomicInteger chars = new AtomicInteger();
|
||||||
|
AtomicReference<Usage> usage = new AtomicReference<>();
|
||||||
|
return chain.nextStream(request).doOnNext(chunk -> {
|
||||||
|
chunks.incrementAndGet();
|
||||||
|
chars.addAndGet(Texts.text(chunk).length());
|
||||||
|
Usage u = usageOf(chunk);
|
||||||
|
if (u != null) {
|
||||||
|
usage.set(u);
|
||||||
|
}
|
||||||
|
}).doOnComplete(() -> sink.accept("[" + name + "] complete " + chunks + " chunks, " + chars + " chars"
|
||||||
|
+ usageText(usage.get()) + " took " + millis(start) + " ms")).doOnError(e -> sink
|
||||||
|
.accept("[" + name + "] failed " + e.getClass().getSimpleName() + " after " + millis(start) + " ms"));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
private String describe(ChatClientRequest request) {
|
||||||
|
List<Message> messages = request.prompt().getInstructions();
|
||||||
|
StringBuilder roles = new StringBuilder();
|
||||||
|
for (Message m : messages) {
|
||||||
|
roles.append(m.getMessageType().name().charAt(0));
|
||||||
|
}
|
||||||
|
String text = "messages=" + messages.size() + " roles=" + roles;
|
||||||
|
if (logContent) {
|
||||||
|
text += " last=\"" + messages.getLast().getText() + "\"";
|
||||||
|
}
|
||||||
|
return text;
|
||||||
|
}
|
||||||
|
|
||||||
|
private String describe(ChatClientResponse response) {
|
||||||
|
String text = "chars=" + Texts.text(response).length() + usageText(usageOf(response));
|
||||||
|
if (logContent) {
|
||||||
|
text += " text=\"" + Texts.text(response) + "\"";
|
||||||
|
}
|
||||||
|
return text;
|
||||||
|
}
|
||||||
|
|
||||||
|
private static Usage usageOf(ChatClientResponse response) {
|
||||||
|
if (response.chatResponse() == null || response.chatResponse().getMetadata() == null) {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
Usage u = response.chatResponse().getMetadata().getUsage();
|
||||||
|
return u != null && u.getTotalTokens() != null && u.getTotalTokens() > 0 ? u : null;
|
||||||
|
}
|
||||||
|
|
||||||
|
private static String usageText(Usage u) {
|
||||||
|
return u == null ? "" : " tokens=" + u.getPromptTokens() + "+" + u.getCompletionTokens();
|
||||||
|
}
|
||||||
|
|
||||||
|
private long millis(long startNanos) {
|
||||||
|
return (nanoClock.getAsLong() - startNanos) / 1_000_000;
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Where the three custom advisors sit, written down once. A <em>lower</em> number runs earlier
|
||||||
|
* on the way in (and later on the way out), so it sits further out.
|
||||||
|
*
|
||||||
|
* <pre>
|
||||||
|
* HIGHEST + 100 PiiRedactionAdvisor (this module) nothing outside it sees raw text
|
||||||
|
* HIGHEST + 200 MessageChatMemoryAdvisor (Spring AI) memory stores the redacted text
|
||||||
|
* HIGHEST + 250 TokenBudgetAdvisor (this module) counts history + question
|
||||||
|
* HIGHEST + 300 ToolCallingAdvisor (Spring AI) the tool loop
|
||||||
|
* HIGHEST + 400 LoggingAdvisor (this module) logs every model round trip
|
||||||
|
* LOWEST ChatModelCallAdvisor / ChatModelStreamAdvisor (Spring AI) the model itself
|
||||||
|
* </pre>
|
||||||
|
*/
|
||||||
|
public final class Orders {
|
||||||
|
|
||||||
|
public static final int PII_REDACTION = Ordered.HIGHEST_PRECEDENCE + 100;
|
||||||
|
|
||||||
|
public static final int TOKEN_BUDGET = Ordered.HIGHEST_PRECEDENCE + 250;
|
||||||
|
|
||||||
|
public static final int LOGGING = Ordered.HIGHEST_PRECEDENCE + 400;
|
||||||
|
|
||||||
|
private Orders() {
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,144 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import java.util.concurrent.ConcurrentHashMap;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.memory.ChatMemory;
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.chat.model.ChatResponse;
|
||||||
|
import org.springframework.ai.chat.model.Generation;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
import reactor.core.publisher.Mono;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Redacts PII from everything sent to the model and, optionally, puts the originals back into the
|
||||||
|
* answer the caller sees. It implements {@link CallAdvisor} and {@link StreamAdvisor} directly
|
||||||
|
* rather than {@code BaseAdvisor}: {@code BaseAdvisor.after} is only invoked for the last chunk of
|
||||||
|
* a stream, which is no use for restoring text that arrives in pieces.
|
||||||
|
*/
|
||||||
|
public final class PiiRedactionAdvisor implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
private final int order;
|
||||||
|
|
||||||
|
private final boolean restore;
|
||||||
|
|
||||||
|
private final boolean conversationScoped;
|
||||||
|
|
||||||
|
private final ConcurrentHashMap<String, PiiRedactor.Session> vault = new ConcurrentHashMap<>();
|
||||||
|
|
||||||
|
/** Conversation-scoped placeholders (the default): the same value keeps the same placeholder across turns. */
|
||||||
|
public PiiRedactionAdvisor(int order, boolean restore) {
|
||||||
|
this(order, restore, true);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @param conversationScoped {@code false} starts a fresh numbering on every request, which is
|
||||||
|
* simpler and wrong as soon as memory replays an earlier turn ("<EMAIL_1>" then means two people)
|
||||||
|
*/
|
||||||
|
public PiiRedactionAdvisor(int order, boolean restore, boolean conversationScoped) {
|
||||||
|
this.order = order;
|
||||||
|
this.restore = restore;
|
||||||
|
this.conversationScoped = conversationScoped;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Drops the stored originals for a conversation. The vault holds real PII in memory: bound it. */
|
||||||
|
public void forget(String conversationId) {
|
||||||
|
vault.remove(conversationId);
|
||||||
|
}
|
||||||
|
|
||||||
|
private PiiRedactor.Session sessionFor(ChatClientRequest request) {
|
||||||
|
Object id = request.context().get(ChatMemory.CONVERSATION_ID);
|
||||||
|
if (!conversationScoped || id == null) {
|
||||||
|
return new PiiRedactor.Session();
|
||||||
|
}
|
||||||
|
return vault.computeIfAbsent(id.toString(), k -> new PiiRedactor.Session());
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return "PiiRedactionAdvisor";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return order;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
PiiRedactor.Session session = sessionFor(request);
|
||||||
|
ChatClientResponse response = chain.nextCall(Texts.mapRequestText(request, session::redact));
|
||||||
|
return restore ? Texts.mapResponseText(response, session::restore) : response;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
return Flux.defer(() -> {
|
||||||
|
PiiRedactor.Session session = sessionFor(request);
|
||||||
|
Flux<ChatClientResponse> upstream = chain.nextStream(Texts.mapRequestText(request, session::redact));
|
||||||
|
if (!restore) {
|
||||||
|
return upstream;
|
||||||
|
}
|
||||||
|
StreamRestorer restorer = new StreamRestorer(session);
|
||||||
|
return upstream.map(restorer::onChunk).concatWith(Mono.defer(restorer::flush).flux());
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* A placeholder such as {@code <EMAIL_1>} can be split across chunks ("<EMA" then "IL_1>"),
|
||||||
|
* and restoring each chunk on its own would miss it. So text from the last {@code <} that has
|
||||||
|
* no {@code >} yet is held back until the next chunk (or the end of the stream) completes it.
|
||||||
|
*/
|
||||||
|
static final class StreamRestorer {
|
||||||
|
|
||||||
|
private static final int LONGEST_PLACEHOLDER = 16;
|
||||||
|
|
||||||
|
private final PiiRedactor.Session session;
|
||||||
|
|
||||||
|
private final StringBuilder pending = new StringBuilder();
|
||||||
|
|
||||||
|
private ChatClientResponse last;
|
||||||
|
|
||||||
|
StreamRestorer(PiiRedactor.Session session) {
|
||||||
|
this.session = session;
|
||||||
|
}
|
||||||
|
|
||||||
|
ChatClientResponse onChunk(ChatClientResponse chunk) {
|
||||||
|
String text = Texts.text(chunk);
|
||||||
|
last = chunk;
|
||||||
|
if (text.isEmpty()) {
|
||||||
|
return chunk;
|
||||||
|
}
|
||||||
|
pending.append(text);
|
||||||
|
int cut = safeCut();
|
||||||
|
String emit = session.restore(pending.substring(0, cut));
|
||||||
|
pending.delete(0, cut);
|
||||||
|
return Texts.mapResponseText(chunk, t -> emit);
|
||||||
|
}
|
||||||
|
|
||||||
|
Mono<ChatClientResponse> flush() {
|
||||||
|
if (pending.isEmpty() || last == null) {
|
||||||
|
return Mono.empty();
|
||||||
|
}
|
||||||
|
String tail = session.restore(pending.toString());
|
||||||
|
pending.setLength(0);
|
||||||
|
ChatResponse tailResponse = ChatResponse.builder()
|
||||||
|
.generations(java.util.List.of(new Generation(new AssistantMessage(tail))))
|
||||||
|
.build();
|
||||||
|
return Mono.just(ChatClientResponse.builder().chatResponse(tailResponse).context(last.context()).build());
|
||||||
|
}
|
||||||
|
|
||||||
|
private int safeCut() {
|
||||||
|
int open = pending.lastIndexOf("<");
|
||||||
|
if (open >= 0 && pending.indexOf(">", open) < 0 && pending.length() - open <= LONGEST_PLACEHOLDER) {
|
||||||
|
return open;
|
||||||
|
}
|
||||||
|
return pending.length();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,113 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import java.util.HashMap;
|
||||||
|
import java.util.LinkedHashMap;
|
||||||
|
import java.util.Map;
|
||||||
|
import java.util.regex.Matcher;
|
||||||
|
import java.util.regex.Pattern;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Finds emails, card numbers and phone numbers in free text and swaps each for a numbered
|
||||||
|
* placeholder such as {@code <EMAIL_1>}. The same value always gets the same placeholder within
|
||||||
|
* one {@link Session}, so the model can still tell "the same address twice" from "two addresses",
|
||||||
|
* and {@link Session#restore} can put the originals back into the answer.
|
||||||
|
*
|
||||||
|
* <p>This is pattern matching, not understanding: it finds what looks like an email, a Luhn-valid
|
||||||
|
* card number and a 10-digit or NANP-style phone number. It does not find names, street addresses
|
||||||
|
* or an ID number in a shape it does not know. Treat it as a floor, not as a guarantee.
|
||||||
|
*/
|
||||||
|
public final class PiiRedactor {
|
||||||
|
|
||||||
|
private static final Pattern EMAIL = Pattern.compile("[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}");
|
||||||
|
|
||||||
|
private static final Pattern CARD = Pattern.compile("(?<!\\d)(?:\\d[ -]?){12,18}\\d(?!\\d)");
|
||||||
|
|
||||||
|
private static final Pattern PHONE = Pattern
|
||||||
|
.compile("(?<![\\w+])(?:\\+\\d{1,3}[ -]?)?(?:\\d{5}[ -]?\\d{5}|\\(?\\d{3}\\)?[ -]\\d{3}-\\d{4})(?!\\w)");
|
||||||
|
|
||||||
|
private PiiRedactor() {
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Which original value became which placeholder. Thread-safe. One session can span a whole
|
||||||
|
* conversation, so "<EMAIL_1>" means the same address in turn 1 and in turn 5.
|
||||||
|
*/
|
||||||
|
public static final class Session {
|
||||||
|
|
||||||
|
private final Map<String, String> placeholderByValue = new LinkedHashMap<>();
|
||||||
|
|
||||||
|
private final Map<String, String> valueByPlaceholder = new LinkedHashMap<>();
|
||||||
|
|
||||||
|
private final Map<String, Integer> counters = new HashMap<>();
|
||||||
|
|
||||||
|
public synchronized String redact(String text) {
|
||||||
|
if (text == null || text.isEmpty()) {
|
||||||
|
return text;
|
||||||
|
}
|
||||||
|
String out = replace(EMAIL, text, "EMAIL", false);
|
||||||
|
out = replace(CARD, out, "CARD", true);
|
||||||
|
return replace(PHONE, out, "PHONE", false);
|
||||||
|
}
|
||||||
|
|
||||||
|
public synchronized String restore(String text) {
|
||||||
|
if (text == null || valueByPlaceholder.isEmpty()) {
|
||||||
|
return text;
|
||||||
|
}
|
||||||
|
String out = text;
|
||||||
|
for (Map.Entry<String, String> e : valueByPlaceholder.entrySet()) {
|
||||||
|
out = out.replace(e.getKey(), e.getValue());
|
||||||
|
}
|
||||||
|
return out;
|
||||||
|
}
|
||||||
|
|
||||||
|
public synchronized boolean isEmpty() {
|
||||||
|
return valueByPlaceholder.isEmpty();
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Placeholder to original value. Never log this: it is the PII. */
|
||||||
|
public synchronized Map<String, String> mapping() {
|
||||||
|
return Map.copyOf(valueByPlaceholder);
|
||||||
|
}
|
||||||
|
|
||||||
|
private String replace(Pattern pattern, String text, String kind, boolean luhn) {
|
||||||
|
Matcher m = pattern.matcher(text);
|
||||||
|
StringBuilder sb = new StringBuilder();
|
||||||
|
while (m.find()) {
|
||||||
|
String value = m.group();
|
||||||
|
if (luhn && !luhnValid(value)) {
|
||||||
|
m.appendReplacement(sb, Matcher.quoteReplacement(value));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
String placeholder = placeholderByValue.computeIfAbsent(kind + ":" + value, k -> {
|
||||||
|
String p = "<" + kind + "_" + counters.merge(kind, 1, Integer::sum) + ">";
|
||||||
|
valueByPlaceholder.put(p, value);
|
||||||
|
return p;
|
||||||
|
});
|
||||||
|
m.appendReplacement(sb, Matcher.quoteReplacement(placeholder));
|
||||||
|
}
|
||||||
|
m.appendTail(sb);
|
||||||
|
return sb.toString();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
static boolean luhnValid(String candidate) {
|
||||||
|
String digits = candidate.replaceAll("\\D", "");
|
||||||
|
if (digits.length() < 13 || digits.length() > 19) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
int sum = 0;
|
||||||
|
boolean dbl = false;
|
||||||
|
for (int i = digits.length() - 1; i >= 0; i--) {
|
||||||
|
int d = digits.charAt(i) - '0';
|
||||||
|
if (dbl) {
|
||||||
|
d *= 2;
|
||||||
|
if (d > 9) {
|
||||||
|
d -= 9;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sum += d;
|
||||||
|
dbl = !dbl;
|
||||||
|
}
|
||||||
|
return sum % 10 == 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.function.UnaryOperator;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.chat.messages.Message;
|
||||||
|
import org.springframework.ai.chat.messages.SystemMessage;
|
||||||
|
import org.springframework.ai.chat.messages.UserMessage;
|
||||||
|
import org.springframework.ai.chat.model.ChatResponse;
|
||||||
|
import org.springframework.ai.chat.model.Generation;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Two rewriting helpers shared by the advisors. Everything in a {@code ChatClientRequest} and a
|
||||||
|
* {@code ChatClientResponse} is immutable, so "changing the prompt" always means building a new one;
|
||||||
|
* these keep the parts an advisor should not touch (options, metadata, media, tool calls).
|
||||||
|
*/
|
||||||
|
public final class Texts {
|
||||||
|
|
||||||
|
private Texts() {
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Applies {@code f} to the text of every system, user and assistant message in the prompt. */
|
||||||
|
public static ChatClientRequest mapRequestText(ChatClientRequest request, UnaryOperator<String> f) {
|
||||||
|
List<Message> rewritten = new ArrayList<>();
|
||||||
|
for (Message m : request.prompt().getInstructions()) {
|
||||||
|
rewritten.add(switch (m) {
|
||||||
|
case UserMessage u -> u.mutate().text(f.apply(u.getText())).build();
|
||||||
|
case SystemMessage s -> s.mutate().text(f.apply(s.getText())).build();
|
||||||
|
case AssistantMessage a when !a.hasToolCalls() -> a.mutate().content(f.apply(a.getText())).build();
|
||||||
|
default -> m;
|
||||||
|
});
|
||||||
|
}
|
||||||
|
return request.mutate().prompt(request.prompt().mutate().messages(rewritten).build()).build();
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Applies {@code f} to the text of every generation, keeping metadata and usage. */
|
||||||
|
public static ChatClientResponse mapResponseText(ChatClientResponse response, UnaryOperator<String> f) {
|
||||||
|
ChatResponse original = response.chatResponse();
|
||||||
|
if (original == null) {
|
||||||
|
return response;
|
||||||
|
}
|
||||||
|
List<Generation> generations = new ArrayList<>();
|
||||||
|
for (Generation g : original.getResults()) {
|
||||||
|
AssistantMessage out = g.getOutput();
|
||||||
|
String text = out.getText();
|
||||||
|
if (text == null || out.hasToolCalls()) {
|
||||||
|
generations.add(g);
|
||||||
|
}
|
||||||
|
else {
|
||||||
|
generations.add(new Generation(out.mutate().content(f.apply(text)).build(), g.getMetadata()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ChatResponse rewritten = ChatResponse.builder().from(original).generations(generations).build();
|
||||||
|
return response.mutate().chatResponse(rewritten).build();
|
||||||
|
}
|
||||||
|
|
||||||
|
/** The text of the first generation, or empty. */
|
||||||
|
public static String text(ChatClientResponse response) {
|
||||||
|
ChatResponse r = response.chatResponse();
|
||||||
|
if (r == null || r.getResult() == null || r.getResult().getOutput().getText() == null) {
|
||||||
|
return "";
|
||||||
|
}
|
||||||
|
return r.getResult().getOutput().getText();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,129 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import java.util.concurrent.ConcurrentHashMap;
|
||||||
|
import java.util.concurrent.atomic.AtomicLong;
|
||||||
|
import java.util.concurrent.atomic.LongAdder;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.messages.Message;
|
||||||
|
import org.springframework.ai.chat.metadata.Usage;
|
||||||
|
import org.springframework.ai.tokenizer.TokenCountEstimator;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Two limits, both checked before the model is called:
|
||||||
|
* <ul>
|
||||||
|
* <li>a per-request cap on the <em>estimated</em> prompt size, so one pasted log file cannot cost
|
||||||
|
* a fortune;</li>
|
||||||
|
* <li>a running per-user total built from the token usage the provider reports on each response,
|
||||||
|
* so a user who has used their allowance is refused without a model call.</li>
|
||||||
|
* </ul>
|
||||||
|
* The user comes from the request context (parameter {@link #USER_KEY}); the state is in memory,
|
||||||
|
* so it is per instance and resets on restart.
|
||||||
|
*/
|
||||||
|
public final class TokenBudgetAdvisor implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
public static final String USER_KEY = "budget_user";
|
||||||
|
|
||||||
|
private final TokenCountEstimator estimator;
|
||||||
|
|
||||||
|
private final int maxPromptTokens;
|
||||||
|
|
||||||
|
private final long maxTokensPerUser;
|
||||||
|
|
||||||
|
private final int order;
|
||||||
|
|
||||||
|
private final ConcurrentHashMap<String, LongAdder> spent = new ConcurrentHashMap<>();
|
||||||
|
|
||||||
|
public TokenBudgetAdvisor(TokenCountEstimator estimator, int maxPromptTokens, long maxTokensPerUser, int order) {
|
||||||
|
this.estimator = estimator;
|
||||||
|
this.maxPromptTokens = maxPromptTokens;
|
||||||
|
this.maxTokensPerUser = maxTokensPerUser;
|
||||||
|
this.order = order;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return "TokenBudgetAdvisor";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return order;
|
||||||
|
}
|
||||||
|
|
||||||
|
public long spent(String user) {
|
||||||
|
LongAdder a = spent.get(user);
|
||||||
|
return a == null ? 0 : a.sum();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
String user = userOf(request);
|
||||||
|
check(user, request);
|
||||||
|
ChatClientResponse response = chain.nextCall(request);
|
||||||
|
Usage usage = usageOf(response);
|
||||||
|
record(user, usage != null ? usage.getTotalTokens() : estimate(request) + estimator.estimate(Texts.text(response)));
|
||||||
|
return response;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
return Flux.defer(() -> {
|
||||||
|
String user = userOf(request);
|
||||||
|
check(user, request);
|
||||||
|
AtomicLong reported = new AtomicLong(-1);
|
||||||
|
StringBuilder streamed = new StringBuilder();
|
||||||
|
return chain.nextStream(request).doOnNext(chunk -> {
|
||||||
|
Usage usage = usageOf(chunk);
|
||||||
|
if (usage != null) {
|
||||||
|
reported.set(usage.getTotalTokens());
|
||||||
|
}
|
||||||
|
streamed.append(Texts.text(chunk));
|
||||||
|
}).doOnComplete(() -> record(user,
|
||||||
|
reported.get() >= 0 ? reported.get() : estimate(request) + estimator.estimate(streamed.toString())));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
private void check(String user, ChatClientRequest request) {
|
||||||
|
int prompt = estimate(request);
|
||||||
|
if (prompt > maxPromptTokens) {
|
||||||
|
throw new TokenBudgetExceededException(
|
||||||
|
"prompt is about " + prompt + " tokens, the limit per request is " + maxPromptTokens);
|
||||||
|
}
|
||||||
|
if (spent(user) >= maxTokensPerUser) {
|
||||||
|
throw new TokenBudgetExceededException(
|
||||||
|
"user " + user + " has used " + spent(user) + " of " + maxTokensPerUser + " tokens");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private int estimate(ChatClientRequest request) {
|
||||||
|
int total = 0;
|
||||||
|
for (Message m : request.prompt().getInstructions()) {
|
||||||
|
total += estimator.estimate(m.getText());
|
||||||
|
}
|
||||||
|
return total;
|
||||||
|
}
|
||||||
|
|
||||||
|
private void record(String user, long tokens) {
|
||||||
|
spent.computeIfAbsent(user, k -> new LongAdder()).add(tokens);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static String userOf(ChatClientRequest request) {
|
||||||
|
Object u = request.context().get(USER_KEY);
|
||||||
|
return u == null ? "anonymous" : u.toString();
|
||||||
|
}
|
||||||
|
|
||||||
|
private static Usage usageOf(ChatClientResponse response) {
|
||||||
|
if (response.chatResponse() == null || response.chatResponse().getMetadata() == null) {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
Usage u = response.chatResponse().getMetadata().getUsage();
|
||||||
|
return u != null && u.getTotalTokens() != null && u.getTotalTokens() > 0 ? u : null;
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
package com.ankurm.advisors.advisor;
|
||||||
|
|
||||||
|
import org.springframework.http.HttpStatus;
|
||||||
|
import org.springframework.web.bind.annotation.ResponseStatus;
|
||||||
|
|
||||||
|
/** Thrown before the model is called, so the request costs nothing. Maps to HTTP 429 in the web layer. */
|
||||||
|
@ResponseStatus(HttpStatus.TOO_MANY_REQUESTS)
|
||||||
|
public class TokenBudgetExceededException extends RuntimeException {
|
||||||
|
|
||||||
|
public TokenBudgetExceededException(String message) {
|
||||||
|
super(message);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
package com.ankurm.advisors.config;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.LoggingAdvisor;
|
||||||
|
import com.ankurm.advisors.advisor.Orders;
|
||||||
|
import com.ankurm.advisors.advisor.PiiRedactionAdvisor;
|
||||||
|
import com.ankurm.advisors.advisor.TokenBudgetAdvisor;
|
||||||
|
import com.knuddels.jtokkit.api.EncodingType;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
|
||||||
|
import org.springframework.ai.chat.memory.ChatMemory;
|
||||||
|
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||||
|
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||||
|
import org.springframework.ai.chat.model.ChatModel;
|
||||||
|
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
|
||||||
|
import org.springframework.ai.tokenizer.TokenCountEstimator;
|
||||||
|
import org.springframework.beans.factory.annotation.Value;
|
||||||
|
import org.springframework.context.annotation.Bean;
|
||||||
|
import org.springframework.context.annotation.Configuration;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The advisors as beans, and one {@link ChatClient} with all of them plus Spring AI's own memory
|
||||||
|
* advisor. The order numbers are what decide the chain, not the order they are listed in
|
||||||
|
* {@code defaultAdvisors(...)}; see {@link Orders}.
|
||||||
|
*/
|
||||||
|
@Configuration
|
||||||
|
public class AdvisorConfig {
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
TokenCountEstimator tokenCountEstimator() {
|
||||||
|
return new JTokkitTokenCountEstimator(EncodingType.O200K_BASE);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
PiiRedactionAdvisor piiRedactionAdvisor() {
|
||||||
|
return new PiiRedactionAdvisor(Orders.PII_REDACTION, true);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
TokenBudgetAdvisor tokenBudgetAdvisor(TokenCountEstimator estimator,
|
||||||
|
@Value("${app.budget.max-prompt-tokens:2000}") int maxPromptTokens,
|
||||||
|
@Value("${app.budget.max-tokens-per-user:20000}") long maxTokensPerUser) {
|
||||||
|
return new TokenBudgetAdvisor(estimator, maxPromptTokens, maxTokensPerUser, Orders.TOKEN_BUDGET);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
LoggingAdvisor loggingAdvisor(@Value("${app.logging.log-content:false}") boolean logContent) {
|
||||||
|
return LoggingAdvisor.toSlf4j(Orders.LOGGING, logContent);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
ChatMemory chatMemory() {
|
||||||
|
return MessageWindowChatMemory.builder().chatMemoryRepository(new InMemoryChatMemoryRepository()).build();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
ChatClient chatClient(ChatModel model, PiiRedactionAdvisor pii, TokenBudgetAdvisor budget, LoggingAdvisor logging,
|
||||||
|
ChatMemory memory) {
|
||||||
|
return ChatClient.builder(model)
|
||||||
|
.defaultSystem("You are a terse support assistant.")
|
||||||
|
.defaultAdvisors(logging, budget, MessageChatMemoryAdvisor.builder(memory).build(), pii)
|
||||||
|
.build();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
package com.ankurm.advisors.web;
|
||||||
|
|
||||||
|
import java.util.Map;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.TokenBudgetAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.memory.ChatMemory;
|
||||||
|
import org.springframework.web.bind.annotation.PostMapping;
|
||||||
|
import org.springframework.web.bind.annotation.RequestBody;
|
||||||
|
import org.springframework.web.bind.annotation.RequestHeader;
|
||||||
|
import org.springframework.web.bind.annotation.RestController;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* One conversation per user. The user comes from an {@code X-User} header only so the tests need
|
||||||
|
* no login; in a real service it is the authenticated principal, never a header the client picks.
|
||||||
|
*/
|
||||||
|
@RestController
|
||||||
|
public class ChatController {
|
||||||
|
|
||||||
|
public record Ask(String text) {
|
||||||
|
}
|
||||||
|
|
||||||
|
private final ChatClient chat;
|
||||||
|
|
||||||
|
public ChatController(ChatClient chat) {
|
||||||
|
this.chat = chat;
|
||||||
|
}
|
||||||
|
|
||||||
|
@PostMapping("/chat")
|
||||||
|
Map<String, String> ask(@RequestHeader("X-User") String user, @RequestBody Ask ask) {
|
||||||
|
String reply = chat.prompt()
|
||||||
|
.user(ask.text())
|
||||||
|
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "conv-" + user).param(TokenBudgetAdvisor.USER_KEY, user))
|
||||||
|
.call()
|
||||||
|
.content();
|
||||||
|
return Map.of("reply", reply);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,17 @@
|
|||||||
|
spring:
|
||||||
|
application:
|
||||||
|
name: advisors
|
||||||
|
ai:
|
||||||
|
model:
|
||||||
|
chat: openai
|
||||||
|
openai:
|
||||||
|
api-key: ${OPENAI_API_KEY:sk-not-set}
|
||||||
|
chat:
|
||||||
|
model: gpt-5-mini
|
||||||
|
|
||||||
|
app:
|
||||||
|
budget:
|
||||||
|
max-prompt-tokens: 2000
|
||||||
|
max-tokens-per-user: 20000
|
||||||
|
logging:
|
||||||
|
log-content: false
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.support.Http;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.TestModelConfig;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.beans.factory.annotation.Autowired;
|
||||||
|
import org.springframework.boot.test.context.SpringBootTest;
|
||||||
|
import org.springframework.boot.test.web.server.LocalServerPort;
|
||||||
|
import org.springframework.context.annotation.Import;
|
||||||
|
|
||||||
|
/** The whole app, over HTTP, with a scripted model: a refused request is a 429, and the model is never asked. */
|
||||||
|
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
|
||||||
|
properties = { "spring.ai.model.chat=none", "app.budget.max-prompt-tokens=60", "app.budget.max-tokens-per-user=30" })
|
||||||
|
@Import(TestModelConfig.class)
|
||||||
|
class AdvisorWebTest {
|
||||||
|
|
||||||
|
@LocalServerPort
|
||||||
|
int port;
|
||||||
|
|
||||||
|
@Autowired
|
||||||
|
RecordingModel model;
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void aRefusedRequestIsA429AndTheModelIsNeverCalled() {
|
||||||
|
try (Transcript t = new Transcript("19-web-429.txt", "Over HTTP: 200, then 429")) {
|
||||||
|
Http http = new Http(port);
|
||||||
|
t.line("limits for this run: 60 tokens per request, 30 per user");
|
||||||
|
int first = 0;
|
||||||
|
for (int i = 1; i <= 4; i++) {
|
||||||
|
Http.Reply r = http.post("dana", "/chat", "{\"text\":\"Where is my refund?\"}");
|
||||||
|
t.line("dana request %d -> HTTP %d, model calls so far: %d", i, r.status(), model.callCount());
|
||||||
|
if (i == 1) {
|
||||||
|
first = r.status();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
int callsBefore = model.callCount();
|
||||||
|
String big = "java.lang.NullPointerException at com.example.Checkout.pay ".repeat(10);
|
||||||
|
Http.Reply refused = http.post("erin", "/chat", "{\"text\":\"" + big + "\"}");
|
||||||
|
t.blank().line("erin pastes a long stack trace -> HTTP %d, model calls: %d (was %d)", refused.status(),
|
||||||
|
model.callCount(), callsBefore);
|
||||||
|
|
||||||
|
assertThat(first).isEqualTo(200);
|
||||||
|
assertThat(refused.status()).isEqualTo(429);
|
||||||
|
assertThat(model.callCount()).isEqualTo(callsBefore);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,94 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.concurrent.CopyOnWriteArrayList;
|
||||||
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@code BaseAdvisor} is the friendly base class: implement {@code before} and {@code after} and
|
||||||
|
* it wires both the call and the stream path. What it does with a stream is not what the name
|
||||||
|
* "after" suggests.
|
||||||
|
*/
|
||||||
|
class BaseAdvisorStreamTest {
|
||||||
|
|
||||||
|
/** Counts what BaseAdvisor calls, and on which thread. */
|
||||||
|
static final class Counting implements BaseAdvisor {
|
||||||
|
|
||||||
|
final AtomicInteger befores = new AtomicInteger();
|
||||||
|
|
||||||
|
final List<String> afterTexts = new CopyOnWriteArrayList<>();
|
||||||
|
|
||||||
|
volatile String beforeThread;
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return "Counting";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return Ordered.HIGHEST_PRECEDENCE + 100;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientRequest before(ChatClientRequest request, AdvisorChain chain) {
|
||||||
|
befores.incrementAndGet();
|
||||||
|
beforeThread = Thread.currentThread().getName();
|
||||||
|
return request;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse after(ChatClientResponse response, AdvisorChain chain) {
|
||||||
|
afterTexts.add(response.chatResponse().getResult().getOutput().getText());
|
||||||
|
return response;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void afterRunsOncePerCallButOnlyOnTheLastChunkOfAStream() {
|
||||||
|
try (Transcript t = new Transcript("04-base-advisor-stream.txt", "BaseAdvisor on a call and on a stream")) {
|
||||||
|
Counting advisor = new Counting();
|
||||||
|
RecordingModel model = new RecordingModel().replier(p -> "one two three four").chunkSize(4);
|
||||||
|
ChatClient client = ChatClient.builder(model).defaultAdvisors(advisor).build();
|
||||||
|
|
||||||
|
client.prompt().user("hi").call().content();
|
||||||
|
t.line("call: before x%d, after x%d", advisor.befores.get(), advisor.afterTexts.size());
|
||||||
|
t.line(" before ran on thread: %s", advisor.beforeThread.equals(Thread.currentThread().getName())
|
||||||
|
? "the caller's thread" : advisor.beforeThread);
|
||||||
|
String callThread = advisor.beforeThread;
|
||||||
|
|
||||||
|
Counting streamed = new Counting();
|
||||||
|
ChatClient streamClient = ChatClient.builder(model).defaultAdvisors(streamed).build();
|
||||||
|
List<String> chunks = streamClient.prompt().user("hi").stream().content().collectList().block();
|
||||||
|
t.blank().line("stream: the model streamed %d chunks: %s", chunks.size(), String.join("|", chunks));
|
||||||
|
t.line(" before x%d, after x%d", streamed.befores.get(), streamed.afterTexts.size());
|
||||||
|
t.line(" after saw only: \"%s\"", streamed.afterTexts.getFirst());
|
||||||
|
t.line(" before ran on thread: %s", streamed.beforeThread.startsWith("boundedElastic")
|
||||||
|
? "a boundedElastic worker, not the caller's thread" : streamed.beforeThread);
|
||||||
|
|
||||||
|
Counting noFinish = new Counting();
|
||||||
|
ChatClient noFinishClient = ChatClient.builder(new RecordingModel().replier(p -> "one two three four").chunkSize(4)
|
||||||
|
.finishReason(false)).defaultAdvisors(noFinish).build();
|
||||||
|
noFinishClient.prompt().user("hi").stream().content().blockLast();
|
||||||
|
t.blank().line("stream whose last chunk has no finish reason: before x%d, after x%d", noFinish.befores.get(),
|
||||||
|
noFinish.afterTexts.size());
|
||||||
|
assertThat(noFinish.afterTexts).isEmpty();
|
||||||
|
assertThat(advisor.afterTexts).hasSize(1);
|
||||||
|
assertThat(streamed.afterTexts).hasSize(1);
|
||||||
|
assertThat(chunks.size()).isGreaterThan(1);
|
||||||
|
assertThat(callThread).isEqualTo(Thread.currentThread().getName());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.ToolCallingAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||||
|
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||||
|
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/** Where the advisors that ship with Spring AI sit, read from the objects themselves. */
|
||||||
|
class BuiltInOrdersTest {
|
||||||
|
|
||||||
|
private static String row(Advisor a) {
|
||||||
|
long offset = (long) a.getOrder() - Ordered.HIGHEST_PRECEDENCE;
|
||||||
|
return "%-28s getOrder() = %d%s".formatted(a.getName(), a.getOrder(),
|
||||||
|
offset < 1000 ? " (HIGHEST_PRECEDENCE + " + offset + ")" : "");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void theShippedAdvisorsAndTheirOrders() {
|
||||||
|
try (Transcript t = new Transcript("20-built-in-orders.txt", "Built-in advisor orders in 2.0.1")) {
|
||||||
|
var memory = MessageWindowChatMemory.builder().chatMemoryRepository(new InMemoryChatMemoryRepository()).build();
|
||||||
|
List<Advisor> shipped = List.of(new SimpleLoggerAdvisor(), new SafeGuardAdvisor(List.of("secret")),
|
||||||
|
MessageChatMemoryAdvisor.builder(memory).build(), ToolCallingAdvisor.builder().build());
|
||||||
|
shipped.forEach(a -> t.line(row(a)));
|
||||||
|
t.blank().line("Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER = HIGHEST_PRECEDENCE + %d",
|
||||||
|
(long) Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER - Ordered.HIGHEST_PRECEDENCE);
|
||||||
|
assertThat(shipped).isNotEmpty();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,147 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.support.Probe;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.ToolCallingAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* How Spring AI turns a list of advisors into a chain. The chain is sorted by
|
||||||
|
* {@code getOrder()}; the order you pass them to {@code defaultAdvisors(...)} only matters for ties.
|
||||||
|
*/
|
||||||
|
class ChainOrderTest {
|
||||||
|
|
||||||
|
|
||||||
|
private static String joined(List<String> events) {
|
||||||
|
return String.join(" ", events);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void theChainIsSortedByOrderNotByRegistrationOrder() {
|
||||||
|
try (Transcript t = new Transcript("01-chain-order.txt", "Advisor chain: what decides who runs first")) {
|
||||||
|
List<String> events = new ArrayList<>();
|
||||||
|
RecordingModel model = new RecordingModel().trace(events::add);
|
||||||
|
ChatClient client = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new Probe("C300", 300, events), new Probe("A100", 100, events), new Probe("B200", 200, events))
|
||||||
|
.build();
|
||||||
|
|
||||||
|
client.prompt().user("hi").call().content();
|
||||||
|
String call = joined(events);
|
||||||
|
t.line("registered as C300, A100, B200 (the number is HIGHEST_PRECEDENCE + n)");
|
||||||
|
t.line("call: %s", call);
|
||||||
|
|
||||||
|
events.clear();
|
||||||
|
client.prompt().user("hi").stream().content().blockLast();
|
||||||
|
String stream = joined(events);
|
||||||
|
t.line("stream: %s", stream);
|
||||||
|
assertThat(call).isEqualTo("A100> B200> C300> model C300< B200< A100<");
|
||||||
|
assertThat(stream).isEqualTo(call);
|
||||||
|
|
||||||
|
events.clear();
|
||||||
|
ChatClient withRequestLevel = client.mutate().build();
|
||||||
|
withRequestLevel.prompt().user("hi").advisors(new Probe("R150", 150, events)).call().content();
|
||||||
|
t.blank().line("plus one advisor added on the request with .advisors(...) at n=150:");
|
||||||
|
t.line("call: %s", joined(events));
|
||||||
|
assertThat(joined(events)).isEqualTo("A100> R150> B200> C300> model C300< B200< R150< A100<");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void twoAdvisorsWithTheSameOrderKeepRegistrationOrder() {
|
||||||
|
try (Transcript t = new Transcript("02-chain-order-ties.txt", "Two advisors with the same order number")) {
|
||||||
|
List<String> events = new ArrayList<>();
|
||||||
|
RecordingModel model = new RecordingModel().trace(events::add);
|
||||||
|
|
||||||
|
ChatClient xy = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new Probe("X", 100, events), new Probe("Y", 100, events)).build();
|
||||||
|
xy.prompt().user("hi").call().content();
|
||||||
|
t.line("registered X, Y -> %s", joined(events));
|
||||||
|
String first = joined(events);
|
||||||
|
|
||||||
|
events.clear();
|
||||||
|
ChatClient yx = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new Probe("Y", 100, events), new Probe("X", 100, events)).build();
|
||||||
|
yx.prompt().user("hi").call().content();
|
||||||
|
t.line("registered Y, X -> %s", joined(events));
|
||||||
|
|
||||||
|
t.blank().line("For equal numbers the advisor registered LAST runs first (outermost).");
|
||||||
|
assertThat(first).isEqualTo("Y> X> model X< Y<");
|
||||||
|
assertThat(joined(events)).isEqualTo("X> Y> model Y< X<");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static String describe(int order) {
|
||||||
|
return order == Integer.MAX_VALUE ? "LOWEST_PRECEDENCE"
|
||||||
|
: "HIGHEST_PRECEDENCE + " + (order - (long) Integer.MIN_VALUE);
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Lists the chain the way an advisor inside it sees it. */
|
||||||
|
private static final class Dump implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
final List<String> callChain = new ArrayList<>();
|
||||||
|
|
||||||
|
final List<String> streamChain = new ArrayList<>();
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return "Dump";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return Ordered.HIGHEST_PRECEDENCE + 100;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
chain.getCallAdvisors().forEach(a -> callChain.add(String.format("%-24s %s", a.getName(), describe(a.getOrder()))));
|
||||||
|
return chain.nextCall(request);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
chain.getStreamAdvisors()
|
||||||
|
.forEach(a -> streamChain.add(String.format("%-24s %s", a.getName(), describe(a.getOrder()))));
|
||||||
|
return chain.nextStream(request);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void chainCanBeInspectedFromInsideAnAdvisor() {
|
||||||
|
try (Transcript t = new Transcript("03-chain-contents.txt", "What is actually in the chain")) {
|
||||||
|
Dump dump = new Dump();
|
||||||
|
ChatClient client = ChatClient.builder(new RecordingModel()).defaultAdvisors(dump).build();
|
||||||
|
client.prompt().user("hi").call().content();
|
||||||
|
t.line("one custom advisor at HIGHEST_PRECEDENCE + 100; the chain a call runs through:");
|
||||||
|
dump.callChain.forEach(s -> t.line(" %s", s));
|
||||||
|
client.prompt().user("hi").stream().content().blockLast();
|
||||||
|
t.blank().line("and the chain a stream runs through:");
|
||||||
|
dump.streamChain.forEach(s -> t.line(" %s", s));
|
||||||
|
|
||||||
|
Dump withExplicit = new Dump();
|
||||||
|
ChatClient explicit = ChatClient.builder(new RecordingModel())
|
||||||
|
.defaultAdvisors(withExplicit, ToolCallingAdvisor.builder().build()).build();
|
||||||
|
explicit.prompt().user("hi").call().content();
|
||||||
|
t.blank().line("the same, with ToolCallingAdvisor.builder().build() added by hand:");
|
||||||
|
withExplicit.callChain.forEach(s -> t.line(" %s", s));
|
||||||
|
|
||||||
|
assertThat(dump.callChain).hasSize(3);
|
||||||
|
assertThat(dump.callChain.get(1)).contains("Tool Calling Advisor");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/** Requests and responses are records. An advisor changes them by making a copy, and context is how advisors talk. */
|
||||||
|
class ContextAndImmutabilityTest {
|
||||||
|
|
||||||
|
private static CallAdvisor advisor(String name, int offset, java.util.function.BiFunction<ChatClientRequest, CallAdvisorChain, ChatClientResponse> body) {
|
||||||
|
return new CallAdvisor() {
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return name;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return Ordered.HIGHEST_PRECEDENCE + offset;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
return body.apply(request, chain);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void contextTravelsDownOnTheRequestAndUpOnTheResponse() {
|
||||||
|
try (Transcript t = new Transcript("17-context.txt", "Advisor context: request down, response up")) {
|
||||||
|
List<String> seen = new ArrayList<>();
|
||||||
|
CallAdvisor outer = advisor("outer", 100, (req, chain) -> {
|
||||||
|
ChatClientRequest copy = req.mutate().context("tenant", "acme").build();
|
||||||
|
seen.add("outer: original request has tenant=" + req.context().get("tenant")
|
||||||
|
+ ", the copy it forwards has tenant=" + copy.context().get("tenant"));
|
||||||
|
ChatClientResponse resp = chain.nextCall(copy);
|
||||||
|
seen.add("outer: response context has verdict=" + resp.context().get("verdict")
|
||||||
|
+ ", tenant=" + resp.context().get("tenant"));
|
||||||
|
return resp;
|
||||||
|
});
|
||||||
|
CallAdvisor inner = advisor("inner", 200, (req, chain) -> {
|
||||||
|
seen.add("inner: request context has tenant=" + req.context().get("tenant"));
|
||||||
|
ChatClientResponse resp = chain.nextCall(req);
|
||||||
|
return resp.mutate().context("verdict", "clean").build();
|
||||||
|
});
|
||||||
|
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
ChatClient.builder(model).defaultAdvisors(outer, inner).build().prompt().user("hi").call().content();
|
||||||
|
seen.forEach(t::line);
|
||||||
|
|
||||||
|
boolean mutable;
|
||||||
|
try {
|
||||||
|
ChatClient.builder(model).defaultAdvisors(advisor("poke", 100, (req, chain) -> {
|
||||||
|
req.context().put("sneaky", "yes");
|
||||||
|
return chain.nextCall(req);
|
||||||
|
})).build().prompt().user("hi").call().content();
|
||||||
|
mutable = true;
|
||||||
|
}
|
||||||
|
catch (UnsupportedOperationException e) {
|
||||||
|
mutable = false;
|
||||||
|
}
|
||||||
|
t.blank().line("req.context().put(...) on the incoming request: %s",
|
||||||
|
mutable ? "allowed" : "UnsupportedOperationException (the map is unmodifiable)");
|
||||||
|
|
||||||
|
assertThat(seen).anyMatch(s -> s.startsWith("inner:") && s.endsWith("tenant=acme"));
|
||||||
|
assertThat(seen).anyMatch(s -> s.contains("original request has tenant=null"));
|
||||||
|
assertThat(seen).anyMatch(s -> s.contains("verdict=clean"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,136 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.function.Supplier;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* What the caller sees when an advisor refuses a request: on a call, and on a stream. An advisor that throws
|
||||||
|
* from {@code adviseStream} before it returns a Flux and one that throws inside the Flux are different things.
|
||||||
|
*/
|
||||||
|
class ErrorPropagationTest {
|
||||||
|
|
||||||
|
private static final class Refuse implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
private final boolean insideFlux;
|
||||||
|
|
||||||
|
Refuse(boolean insideFlux) {
|
||||||
|
this.insideFlux = insideFlux;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return insideFlux ? "refuse-inside-flux" : "refuse-eagerly";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return Ordered.HIGHEST_PRECEDENCE + 100;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
throw new IllegalStateException("refused");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
if (insideFlux) {
|
||||||
|
return Flux.defer(() -> Flux.error(new IllegalStateException("refused")));
|
||||||
|
}
|
||||||
|
throw new IllegalStateException("refused");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static final class RefusingBase implements BaseAdvisor {
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return "refusing-base";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return Ordered.HIGHEST_PRECEDENCE + 100;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientRequest before(ChatClientRequest request, AdvisorChain chain) {
|
||||||
|
throw new IllegalStateException("refused");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse after(ChatClientResponse response, AdvisorChain chain) {
|
||||||
|
return response;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Runs the action and reports at which point it failed and with what. */
|
||||||
|
private static String outcome(Supplier<Object> build, Supplier<Object> consume) {
|
||||||
|
String phase = "building the Flux";
|
||||||
|
try {
|
||||||
|
build.get();
|
||||||
|
phase = "running it";
|
||||||
|
consume.get();
|
||||||
|
return "no error";
|
||||||
|
}
|
||||||
|
catch (Throwable e) {
|
||||||
|
String cause = e.getCause() == null ? ""
|
||||||
|
: " (cause: " + e.getCause().getClass().getSimpleName() + ": " + e.getCause().getMessage() + ")";
|
||||||
|
return "failed while " + phase + " with " + e.getClass().getSimpleName() + ": " + e.getMessage() + cause;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void whereARefusalSurfacesOnCallAndOnStream() {
|
||||||
|
try (Transcript t = new Transcript("16-error-propagation.txt", "Where a refusal surfaces")) {
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
t.line("an advisor that throws IllegalStateException(\"refused\"), model calls counted afterwards");
|
||||||
|
t.blank();
|
||||||
|
|
||||||
|
String[] labels = { "plain advisor, call", "plain advisor, stream, throws eagerly",
|
||||||
|
"plain advisor, stream, error inside Flux.defer", "BaseAdvisor.before, call",
|
||||||
|
"BaseAdvisor.before, stream" };
|
||||||
|
String[] results = new String[5];
|
||||||
|
|
||||||
|
ChatClient eager = ChatClient.builder(model).defaultAdvisors(new Refuse(false)).build();
|
||||||
|
ChatClient deferred = ChatClient.builder(model).defaultAdvisors(new Refuse(true)).build();
|
||||||
|
ChatClient base = ChatClient.builder(model).defaultAdvisors(new RefusingBase()).build();
|
||||||
|
|
||||||
|
results[0] = outcome(() -> null, () -> eager.prompt().user("hi").call().content());
|
||||||
|
results[1] = outcome(() -> eager.prompt().user("hi").stream().content(),
|
||||||
|
() -> eager.prompt().user("hi").stream().content().blockLast());
|
||||||
|
results[2] = outcome(() -> deferred.prompt().user("hi").stream().content(),
|
||||||
|
() -> deferred.prompt().user("hi").stream().content().blockLast());
|
||||||
|
results[3] = outcome(() -> null, () -> base.prompt().user("hi").call().content());
|
||||||
|
results[4] = outcome(() -> base.prompt().user("hi").stream().content(),
|
||||||
|
() -> base.prompt().user("hi").stream().content().blockLast());
|
||||||
|
|
||||||
|
for (int i = 0; i < labels.length; i++) {
|
||||||
|
t.line("%-48s %s", labels[i], results[i]);
|
||||||
|
}
|
||||||
|
t.blank().line("model calls: %d", model.callCount());
|
||||||
|
|
||||||
|
assertThat(model.callCount()).isZero();
|
||||||
|
for (String r : results) {
|
||||||
|
assertThat(r).contains("IllegalStateException").contains("refused");
|
||||||
|
}
|
||||||
|
assertThat(results[4]).contains("Stream processing failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,74 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.concurrent.atomic.AtomicLong;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.LoggingAdvisor;
|
||||||
|
import com.ankurm.advisors.advisor.Orders;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
|
||||||
|
/** The audit-logging advisor on a call and on a stream. A fake clock makes "took N ms" deterministic. */
|
||||||
|
class LoggingAdvisorTest {
|
||||||
|
|
||||||
|
/** Advances 5 ms every time it is read. */
|
||||||
|
private static AtomicLong clock() {
|
||||||
|
return new AtomicLong();
|
||||||
|
}
|
||||||
|
|
||||||
|
private static LoggingAdvisor logger(List<String> sink, boolean content) {
|
||||||
|
AtomicLong nanos = clock();
|
||||||
|
return new LoggingAdvisor("LoggingAdvisor", Orders.LOGGING, content, sink::add, () -> nanos.addAndGet(5_000_000));
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void oneLineInOneLineOutWithLatencyAndUsage() {
|
||||||
|
try (Transcript t = new Transcript("05-logging-advisor.txt", "LoggingAdvisor: a call and a stream")) {
|
||||||
|
List<String> log = new ArrayList<>();
|
||||||
|
RecordingModel model = new RecordingModel().replier(p -> "Refund issued.").chunkSize(6);
|
||||||
|
ChatClient client = ChatClient.builder(model).defaultSystem("You are a terse support assistant.")
|
||||||
|
.defaultAdvisors(logger(log, false)).build();
|
||||||
|
|
||||||
|
client.prompt().user("Where is my refund?").call().content();
|
||||||
|
t.line("call, content logging off (the default):");
|
||||||
|
log.forEach(l -> t.line(" %s", l));
|
||||||
|
assertThat(log).hasSize(2);
|
||||||
|
|
||||||
|
log.clear();
|
||||||
|
client.prompt().user("Where is my refund?").stream().content().blockLast();
|
||||||
|
t.blank().line("stream:");
|
||||||
|
log.forEach(l -> t.line(" %s", l));
|
||||||
|
|
||||||
|
log.clear();
|
||||||
|
ChatClient withText = ChatClient.builder(model).defaultAdvisors(logger(log, true)).build();
|
||||||
|
withText.prompt().user("Where is my refund?").call().content();
|
||||||
|
t.blank().line("call, content logging on:");
|
||||||
|
log.forEach(l -> t.line(" %s", l));
|
||||||
|
assertThat(log.getFirst()).contains("Where is my refund?");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void aFailureIsLoggedAndRethrown() {
|
||||||
|
try (Transcript t = new Transcript("06-logging-failure.txt", "LoggingAdvisor when the model call fails")) {
|
||||||
|
List<String> log = new ArrayList<>();
|
||||||
|
RecordingModel model = new RecordingModel().replier(p -> {
|
||||||
|
throw new IllegalStateException("provider returned 503");
|
||||||
|
});
|
||||||
|
ChatClient client = ChatClient.builder(model).defaultAdvisors(logger(log, false)).build();
|
||||||
|
try {
|
||||||
|
client.prompt().user("hi").call().content();
|
||||||
|
}
|
||||||
|
catch (IllegalStateException e) {
|
||||||
|
t.line("caller got: %s: %s", e.getClass().getSimpleName(), e.getMessage());
|
||||||
|
}
|
||||||
|
log.forEach(l -> t.line(" %s", l));
|
||||||
|
assertThat(log.getLast()).contains("failed").contains("IllegalStateException");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.Orders;
|
||||||
|
import com.ankurm.advisors.advisor.PiiRedactionAdvisor;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Show;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
|
||||||
|
import org.springframework.ai.chat.memory.ChatMemory;
|
||||||
|
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||||
|
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||||
|
|
||||||
|
/** Placeholders and memory: numbering that restarts every request gives one label to two people. */
|
||||||
|
class PiiMultiTurnTest {
|
||||||
|
|
||||||
|
private record Turns(String modelSawTurn2, String callerGotTurn2, List<String> stored) {
|
||||||
|
}
|
||||||
|
|
||||||
|
private static Turns twoTurns(boolean conversationScoped) {
|
||||||
|
ChatMemory memory = MessageWindowChatMemory.builder().chatMemoryRepository(new InMemoryChatMemoryRepository()).build();
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
ChatClient client = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new PiiRedactionAdvisor(Orders.PII_REDACTION, true, conversationScoped),
|
||||||
|
MessageChatMemoryAdvisor.builder(memory).build())
|
||||||
|
.build();
|
||||||
|
client.prompt().user("My email is [email protected]").advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "c")).call()
|
||||||
|
.content();
|
||||||
|
String reply = client.prompt().user("Also cc [email protected]").advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "c"))
|
||||||
|
.call().content();
|
||||||
|
return new Turns(Show.inline(model.lastPrompt().getInstructions()), reply,
|
||||||
|
memory.get("c").stream().map(Show::one).toList());
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void placeholdersMustSurviveAcrossTurnsWhenMemoryReplaysThem() {
|
||||||
|
try (Transcript t = new Transcript("11-pii-multi-turn.txt", "Placeholders across two turns with memory")) {
|
||||||
|
t.line("turn 1: My email is [email protected]");
|
||||||
|
t.line("turn 2: Also cc [email protected]");
|
||||||
|
|
||||||
|
Turns perRequest = twoTurns(false);
|
||||||
|
t.blank().line("numbering restarts on every request:");
|
||||||
|
t.line(" model was sent on turn 2: %s", perRequest.modelSawTurn2());
|
||||||
|
t.line(" caller receives: %s", perRequest.callerGotTurn2());
|
||||||
|
|
||||||
|
Turns scoped = twoTurns(true);
|
||||||
|
t.blank().line("numbering kept per conversation (the default):");
|
||||||
|
t.line(" model was sent on turn 2: %s", scoped.modelSawTurn2());
|
||||||
|
t.line(" caller receives: %s", scoped.callerGotTurn2());
|
||||||
|
t.blank().line("what the memory stores (placeholders, never the addresses):");
|
||||||
|
scoped.stored().forEach(s -> t.line(" %s", s));
|
||||||
|
|
||||||
|
assertThat(perRequest.modelSawTurn2()).contains("U:My email is <EMAIL_1>").contains("U:Also cc <EMAIL_1>");
|
||||||
|
assertThat(scoped.modelSawTurn2()).contains("U:My email is <EMAIL_1>").contains("U:Also cc <EMAIL_2>");
|
||||||
|
assertThat(scoped.stored()).noneMatch(s -> s.contains("@"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,83 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.concurrent.atomic.AtomicLong;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.LoggingAdvisor;
|
||||||
|
import com.ankurm.advisors.advisor.PiiRedactionAdvisor;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||||
|
import org.springframework.ai.chat.memory.ChatMemory;
|
||||||
|
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||||
|
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Three places to put the redaction advisor relative to the logging advisor and Spring AI's memory
|
||||||
|
* advisor (HIGHEST_PRECEDENCE + 200), and what each of the three ends up holding.
|
||||||
|
*/
|
||||||
|
class PiiOrderTest {
|
||||||
|
|
||||||
|
private static final String INPUT = "My email is [email protected]";
|
||||||
|
|
||||||
|
private record Outcome(String logged, String stored, String modelSaw) {
|
||||||
|
}
|
||||||
|
|
||||||
|
private static Outcome run(int piiOffset, int loggingOffset) {
|
||||||
|
List<String> log = new ArrayList<>();
|
||||||
|
AtomicLong nanos = new AtomicLong();
|
||||||
|
LoggingAdvisor logging = new LoggingAdvisor("LoggingAdvisor", Ordered.HIGHEST_PRECEDENCE + loggingOffset, true,
|
||||||
|
log::add, () -> nanos.addAndGet(5_000_000));
|
||||||
|
ChatMemory memory = MessageWindowChatMemory.builder().chatMemoryRepository(new InMemoryChatMemoryRepository()).build();
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
ChatClient client = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(logging, new PiiRedactionAdvisor(Ordered.HIGHEST_PRECEDENCE + piiOffset, true),
|
||||||
|
MessageChatMemoryAdvisor.builder(memory).build())
|
||||||
|
.build();
|
||||||
|
client.prompt().user(INPUT).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "c")).call().content();
|
||||||
|
String logged = log.getFirst().replaceFirst(".*last=", "");
|
||||||
|
return new Outcome(logged, "\"" + memory.get("c").getFirst().getText() + "\"",
|
||||||
|
"\"" + model.lastPrompt().getInstructions().getLast().getText() + "\"");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void whereTheRedactionAdvisorSitsDecidesWhatIsLoggedAndWhatIsStored() {
|
||||||
|
try (Transcript t = new Transcript("10-pii-order.txt", "Redaction order versus logging and memory")) {
|
||||||
|
t.line("memory advisor is fixed at HIGHEST_PRECEDENCE + %d",
|
||||||
|
Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER - (long) Ordered.HIGHEST_PRECEDENCE);
|
||||||
|
t.line("user says: %s", INPUT);
|
||||||
|
|
||||||
|
Outcome a = run(100, 400);
|
||||||
|
t.blank().line("A redaction +100, logging +400 (redaction outside both):");
|
||||||
|
t.line(" log line saw: %s", a.logged());
|
||||||
|
t.line(" memory stored: %s", a.stored());
|
||||||
|
t.line(" model was sent: %s", a.modelSaw());
|
||||||
|
|
||||||
|
Outcome b = run(100, 50);
|
||||||
|
t.blank().line("B redaction +100, logging +50 (logging outside redaction):");
|
||||||
|
t.line(" log line saw: %s", b.logged());
|
||||||
|
t.line(" memory stored: %s", b.stored());
|
||||||
|
t.line(" model was sent: %s", b.modelSaw());
|
||||||
|
|
||||||
|
Outcome c = run(300, 400);
|
||||||
|
t.blank().line("C redaction +300 (inside the memory advisor), logging +400:");
|
||||||
|
t.line(" log line saw: %s", c.logged());
|
||||||
|
t.line(" memory stored: %s", c.stored());
|
||||||
|
t.line(" model was sent: %s", c.modelSaw());
|
||||||
|
|
||||||
|
assertThat(a.logged()).doesNotContain("[email protected]");
|
||||||
|
assertThat(a.stored()).doesNotContain("[email protected]");
|
||||||
|
assertThat(b.logged()).contains("[email protected]");
|
||||||
|
assertThat(b.stored()).doesNotContain("[email protected]");
|
||||||
|
assertThat(c.stored()).contains("[email protected]");
|
||||||
|
assertThat(c.modelSaw()).doesNotContain("[email protected]");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,104 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.regex.Matcher;
|
||||||
|
import java.util.regex.Pattern;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.Orders;
|
||||||
|
import com.ankurm.advisors.advisor.PiiRedactionAdvisor;
|
||||||
|
import com.ankurm.advisors.support.NaiveStreamingPii;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Show;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
|
||||||
|
/** PII redaction before the model, and restoring the originals afterwards. */
|
||||||
|
class PiiRedactionTest {
|
||||||
|
|
||||||
|
private static final String INPUT = "Hi, I'm Priya. Email [email protected] or call +91 98765 43210. "
|
||||||
|
+ "Card 4111 1111 1111 1111, order 1234 5678 9012 3456, and again [email protected].";
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void theModelSeesPlaceholdersAndTheCallerSeesTheOriginals() {
|
||||||
|
try (Transcript t = new Transcript("07-pii-redaction.txt", "PiiRedactionAdvisor on a call")) {
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
ChatClient restoring = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new PiiRedactionAdvisor(Orders.PII_REDACTION, true)).build();
|
||||||
|
String reply = restoring.prompt().user(INPUT).call().content();
|
||||||
|
|
||||||
|
t.line("caller sends: %s", INPUT);
|
||||||
|
t.line("model was sent: %s", model.lastPrompt().getInstructions().getLast().getText());
|
||||||
|
t.line("caller receives: %s", reply);
|
||||||
|
|
||||||
|
ChatClient plain = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new PiiRedactionAdvisor(Orders.PII_REDACTION, false)).build();
|
||||||
|
t.blank().line("with restore switched off, the caller receives:");
|
||||||
|
t.line(" %s", plain.prompt().user(INPUT).call().content());
|
||||||
|
|
||||||
|
String sent = model.prompts().getFirst().getInstructions().getLast().getText();
|
||||||
|
assertThat(sent).contains("<EMAIL_1>", "<PHONE_1>", "<CARD_1>").doesNotContain("[email protected]", "98765", "4111");
|
||||||
|
assertThat(sent).contains("1234 5678 9012 3456").contains("Priya");
|
||||||
|
assertThat(sent.split("<EMAIL_1>", -1)).hasSize(3);
|
||||||
|
assertThat(reply).isEqualTo("You said: " + INPUT);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void whatItDoesNotCatch() {
|
||||||
|
try (Transcript t = new Transcript("08-pii-limits.txt", "What pattern-based redaction misses")) {
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
ChatClient client = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new PiiRedactionAdvisor(Orders.PII_REDACTION, true)).build();
|
||||||
|
List<String> inputs = List.of(
|
||||||
|
"My name is Priya Sharma and I live at 14 Hill Road, Bandra, Mumbai 400050.",
|
||||||
|
"Passport N1234567, PAN ABCDE1234F.",
|
||||||
|
"Write to priya (at) example (dot) com",
|
||||||
|
"Card 4111-1111-1111-1112 and order 1234 5678 9012 3456");
|
||||||
|
for (String in : inputs) {
|
||||||
|
client.prompt().user(in).call().content();
|
||||||
|
String sent = model.lastPrompt().getInstructions().getLast().getText();
|
||||||
|
t.line("in: %s", in);
|
||||||
|
t.line("out: %s", sent);
|
||||||
|
t.blank();
|
||||||
|
}
|
||||||
|
assertThat(model.prompts().get(0).getInstructions().getLast().getText()).contains("Priya Sharma", "14 Hill Road");
|
||||||
|
assertThat(model.prompts().get(1).getInstructions().getLast().getText()).contains("N1234567", "ABCDE1234F");
|
||||||
|
assertThat(model.prompts().get(3).getInstructions().getLast().getText()).contains("4111-1111-1111-1112");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void aPlaceholderSplitAcrossChunksSurvivesOnlyWithTheBufferedRestore() {
|
||||||
|
try (Transcript t = new Transcript("09-pii-stream-boundary.txt", "Restoring placeholders in a stream")) {
|
||||||
|
Pattern placeholder = Pattern.compile("<[A-Z]+_\\d+>");
|
||||||
|
RecordingModel model = new RecordingModel().chunkSize(5).replier(p -> {
|
||||||
|
Matcher m = placeholder.matcher(p.getInstructions().getLast().getText());
|
||||||
|
return "Sure, I will write to " + (m.find() ? m.group() : "nobody") + " now.";
|
||||||
|
});
|
||||||
|
String input = "Please email [email protected]";
|
||||||
|
|
||||||
|
List<String> naive = ChatClient.builder(model).defaultAdvisors(new NaiveStreamingPii()).build().prompt()
|
||||||
|
.user(input).stream().content().collectList().block();
|
||||||
|
List<String> raw = model.stream(new org.springframework.ai.chat.prompt.Prompt("<EMAIL_1>"))
|
||||||
|
.map(r -> r.getResult().getOutput().getText()).collectList().block();
|
||||||
|
t.line("the model streams 5-character chunks: %s", String.join("|", raw));
|
||||||
|
t.blank().line("restore each chunk on its own:");
|
||||||
|
t.line(" chunks: %s", String.join("|", naive));
|
||||||
|
t.line(" joined: %s", String.join("", naive));
|
||||||
|
|
||||||
|
List<String> buffered = ChatClient.builder(model)
|
||||||
|
.defaultAdvisors(new PiiRedactionAdvisor(Orders.PII_REDACTION, true)).build().prompt().user(input)
|
||||||
|
.stream().content().collectList().block();
|
||||||
|
t.blank().line("PiiRedactionAdvisor (holds back from an unfinished \"<\"):");
|
||||||
|
t.line(" chunks: %s", String.join("|", buffered));
|
||||||
|
t.line(" joined: %s", String.join("", buffered));
|
||||||
|
|
||||||
|
assertThat(raw).contains("o <EM", "AIL_1");
|
||||||
|
assertThat(String.join("", naive)).contains("<EMA").doesNotContain("[email protected]");
|
||||||
|
assertThat(String.join("", buffered)).isEqualTo("Sure, I will write to [email protected] now.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,112 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.Orders;
|
||||||
|
import com.ankurm.advisors.advisor.TokenBudgetAdvisor;
|
||||||
|
import com.ankurm.advisors.advisor.TokenBudgetExceededException;
|
||||||
|
import com.ankurm.advisors.support.RecordingModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import com.knuddels.jtokkit.api.EncodingType;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
|
||||||
|
import org.springframework.ai.tokenizer.TokenCountEstimator;
|
||||||
|
|
||||||
|
/** A per-request cap and a per-user running total, both enforced before the model is called. */
|
||||||
|
class TokenBudgetTest {
|
||||||
|
|
||||||
|
private static final TokenCountEstimator ESTIMATOR = new JTokkitTokenCountEstimator(EncodingType.O200K_BASE);
|
||||||
|
|
||||||
|
private static String ask(ChatClient client, String user, String text) {
|
||||||
|
return client.prompt().user(text).advisors(a -> a.param(TokenBudgetAdvisor.USER_KEY, user)).call().content();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void refusesBeforeTheModelIsCalled() {
|
||||||
|
try (Transcript t = new Transcript("12-token-budget.txt", "TokenBudgetAdvisor: 40 tokens per request, 60 per user")) {
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
TokenBudgetAdvisor budget = new TokenBudgetAdvisor(ESTIMATOR, 40, 60, Orders.TOKEN_BUDGET);
|
||||||
|
ChatClient client = ChatClient.builder(model).defaultAdvisors(budget).build();
|
||||||
|
|
||||||
|
t.line("alice asks short questions; \"spent\" is the usage the model reported:");
|
||||||
|
for (int i = 1; i <= 6; i++) {
|
||||||
|
try {
|
||||||
|
ask(client, "alice", "Where is my refund?");
|
||||||
|
t.line(" call %d: answered, spent=%d, model calls=%d", i, budget.spent("alice"), model.callCount());
|
||||||
|
}
|
||||||
|
catch (TokenBudgetExceededException e) {
|
||||||
|
t.line(" call %d: refused (%s), spent=%d, model calls=%d", i, e.getMessage(), budget.spent("alice"),
|
||||||
|
model.callCount());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
int before = model.callCount();
|
||||||
|
String paste = "java.lang.NullPointerException at com.example.Checkout.pay(Checkout.java:88) ".repeat(6);
|
||||||
|
t.blank().line("bob pastes a stack trace of %d estimated tokens:", ESTIMATOR.estimate(paste));
|
||||||
|
assertThatThrownBy(() -> ask(client, "bob", paste)).isInstanceOf(TokenBudgetExceededException.class)
|
||||||
|
.satisfies(e -> t.line(" refused: %s", e.getMessage()));
|
||||||
|
t.line(" model calls: %d (was %d), bob's spent: %d", model.callCount(), before, budget.spent("bob"));
|
||||||
|
|
||||||
|
ask(client, "bob", "Where is my refund?");
|
||||||
|
t.blank().line("bob then asks a short question: answered, bob spent=%d, alice spent=%d", budget.spent("bob"),
|
||||||
|
budget.spent("alice"));
|
||||||
|
|
||||||
|
assertThat(model.callCount()).isEqualTo(before + 1);
|
||||||
|
assertThat(budget.spent("alice")).isGreaterThanOrEqualTo(60);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void streamsAreCountedFromTheLastChunkOrEstimatedWhenThereIsNoUsage() {
|
||||||
|
try (Transcript t = new Transcript("13-token-budget-stream.txt", "Token accounting on a stream")) {
|
||||||
|
TokenBudgetAdvisor budget = new TokenBudgetAdvisor(ESTIMATOR, 1000, 100_000, Orders.TOKEN_BUDGET);
|
||||||
|
RecordingModel withUsage = new RecordingModel().replier(p -> "Refund issued today.").chunkSize(6);
|
||||||
|
ChatClient client = ChatClient.builder(withUsage).defaultAdvisors(budget).build();
|
||||||
|
|
||||||
|
client.prompt().user("Where is my refund?").advisors(a -> a.param(TokenBudgetAdvisor.USER_KEY, "u1")).call()
|
||||||
|
.content();
|
||||||
|
long callSpent = budget.spent("u1");
|
||||||
|
client.prompt().user("Where is my refund?").advisors(a -> a.param(TokenBudgetAdvisor.USER_KEY, "u2")).stream()
|
||||||
|
.content().blockLast();
|
||||||
|
long streamSpent = budget.spent("u2");
|
||||||
|
|
||||||
|
RecordingModel noUsage = new RecordingModel().replier(p -> "Refund issued today.").chunkSize(6)
|
||||||
|
.streamUsage(false);
|
||||||
|
ChatClient.builder(noUsage).defaultAdvisors(budget).build().prompt().user("Where is my refund?")
|
||||||
|
.advisors(a -> a.param(TokenBudgetAdvisor.USER_KEY, "u3")).stream().content().blockLast();
|
||||||
|
long estimated = budget.spent("u3");
|
||||||
|
|
||||||
|
t.line("same question, same answer:");
|
||||||
|
t.line(" call, usage from the response: %d tokens", callSpent);
|
||||||
|
t.line(" stream, usage on the last chunk: %d tokens", streamSpent);
|
||||||
|
t.line(" stream, no usage reported (estimated): %d tokens", estimated);
|
||||||
|
t.blank().line("A stream that reports no usage is billed by the provider all the same. The estimate above");
|
||||||
|
t.line("matches only because this scripted model and the advisor use the same tokenizer.");
|
||||||
|
|
||||||
|
assertThat(streamSpent).isEqualTo(callSpent);
|
||||||
|
assertThat(estimated).isEqualTo(callSpent);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void aStreamIsRefusedAsAnErrorSignal() {
|
||||||
|
try (Transcript t = new Transcript("14-token-budget-stream-refusal.txt", "A refusal on the stream path")) {
|
||||||
|
RecordingModel model = new RecordingModel();
|
||||||
|
TokenBudgetAdvisor budget = new TokenBudgetAdvisor(ESTIMATOR, 5, 100, Orders.TOKEN_BUDGET);
|
||||||
|
ChatClient client = ChatClient.builder(model).defaultAdvisors(budget).build();
|
||||||
|
try {
|
||||||
|
client.prompt().user("This question is longer than five tokens, so it is refused.").stream().content()
|
||||||
|
.blockLast();
|
||||||
|
}
|
||||||
|
catch (RuntimeException e) {
|
||||||
|
t.line("subscriber got: %s", e.getClass().getSimpleName());
|
||||||
|
t.line("message: %s", e.getMessage());
|
||||||
|
assertThat(e).isInstanceOf(TokenBudgetExceededException.class);
|
||||||
|
}
|
||||||
|
t.line("model calls: %d", model.callCount());
|
||||||
|
assertThat(model.callCount()).isZero();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.TokenBudgetAdvisor;
|
||||||
|
import com.ankurm.advisors.support.ScriptedToolModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import com.knuddels.jtokkit.api.EncodingType;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.support.ToolCallbacks;
|
||||||
|
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
|
||||||
|
import org.springframework.ai.tool.annotation.Tool;
|
||||||
|
import org.springframework.ai.tool.annotation.ToolParam;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/** Does a budget advisor count every model round trip of a tool loop, or only what comes out of it? */
|
||||||
|
class ToolLoopBudgetTest {
|
||||||
|
|
||||||
|
static class Weather {
|
||||||
|
|
||||||
|
@Tool(description = "Current temperature for a city")
|
||||||
|
String currentWeather(@ToolParam(description = "city name") String city) {
|
||||||
|
return "31C in " + city;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static long spentWith(int offset) {
|
||||||
|
TokenBudgetAdvisor budget = new TokenBudgetAdvisor(new JTokkitTokenCountEstimator(EncodingType.O200K_BASE), 10_000, 100_000,
|
||||||
|
Ordered.HIGHEST_PRECEDENCE + offset);
|
||||||
|
ScriptedToolModel model = ScriptedToolModel.builder().reportingUsage(100, 10)
|
||||||
|
.thenCallTools(new AssistantMessage.ToolCall("call-1", "function", "currentWeather", "{\"city\":\"Mumbai\"}"))
|
||||||
|
.reportingUsage(130, 20).thenRespond("It is 31C in Mumbai.").build();
|
||||||
|
ChatClient.builder(model).defaultToolCallbacks(ToolCallbacks.from(new Weather())).defaultAdvisors(budget).build().prompt()
|
||||||
|
.user("Weather in Mumbai?").advisors(a -> a.param(TokenBudgetAdvisor.USER_KEY, "u")).call().content();
|
||||||
|
return budget.spent("u");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void everyRoundIsCountedWhereverTheBudgetSits() {
|
||||||
|
try (Transcript t = new Transcript("21-tool-loop-budget.txt", "Token budget and the tool loop")) {
|
||||||
|
long outside = spentWith(250);
|
||||||
|
long inside = spentWith(400);
|
||||||
|
t.line("the model reports 100+10 tokens for round 1 (asks for the tool) and 130+20 for round 2 (answers)");
|
||||||
|
t.line("so the provider would bill 260 tokens for this one question");
|
||||||
|
t.blank().line("budget advisor at +250 (outside the tool loop): recorded %d", outside);
|
||||||
|
t.line("budget advisor at +400 (inside the tool loop): recorded %d", inside);
|
||||||
|
assertThat(inside).isEqualTo(260);
|
||||||
|
assertThat(outside).isEqualTo(260);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.concurrent.atomic.AtomicLong;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.LoggingAdvisor;
|
||||||
|
import com.ankurm.advisors.support.ScriptedToolModel;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClient;
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.support.ToolCallbacks;
|
||||||
|
import org.springframework.ai.tool.annotation.Tool;
|
||||||
|
import org.springframework.ai.tool.annotation.ToolParam;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The tool loop runs at HIGHEST_PRECEDENCE + 300. An advisor outside it (a lower number) is entered once
|
||||||
|
* per question; an advisor inside it (a higher number) is entered once per model round trip.
|
||||||
|
*/
|
||||||
|
class ToolLoopOrderTest {
|
||||||
|
|
||||||
|
static class Weather {
|
||||||
|
|
||||||
|
@Tool(description = "Current temperature for a city")
|
||||||
|
String currentWeather(@ToolParam(description = "city name") String city) {
|
||||||
|
return "31C in " + city;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static List<String> run(int loggingOffset) {
|
||||||
|
List<String> log = new ArrayList<>();
|
||||||
|
AtomicLong nanos = new AtomicLong();
|
||||||
|
LoggingAdvisor logging = new LoggingAdvisor("LoggingAdvisor", Ordered.HIGHEST_PRECEDENCE + loggingOffset, false,
|
||||||
|
log::add, () -> nanos.addAndGet(5_000_000));
|
||||||
|
ScriptedToolModel model = ScriptedToolModel.builder()
|
||||||
|
.thenCallTools(new AssistantMessage.ToolCall("call-1", "function", "currentWeather", "{\"city\":\"Mumbai\"}"))
|
||||||
|
.thenRespond("It is 31C in Mumbai.")
|
||||||
|
.build();
|
||||||
|
ChatClient.builder(model).defaultToolCallbacks(ToolCallbacks.from(new Weather())).defaultAdvisors(logging).build()
|
||||||
|
.prompt().user("Weather in Mumbai?").call().content();
|
||||||
|
return log;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void anAdvisorInsideTheToolLoopRunsOncePerModelRoundTrip() {
|
||||||
|
try (Transcript t = new Transcript("15-tool-loop-order.txt", "Where an advisor sits relative to the tool loop")) {
|
||||||
|
List<String> outside = run(250);
|
||||||
|
List<String> inside = run(400);
|
||||||
|
t.line("one question, one tool call, so the model is called twice (ToolCallingAdvisor is at +300)");
|
||||||
|
t.blank().line("logging advisor at +250 (outside the tool loop):");
|
||||||
|
outside.forEach(l -> t.line(" %s", l));
|
||||||
|
t.blank().line("logging advisor at +400 (inside the tool loop):");
|
||||||
|
inside.forEach(l -> t.line(" %s", l));
|
||||||
|
|
||||||
|
assertThat(outside).hasSize(2);
|
||||||
|
assertThat(inside).hasSize(4);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,73 @@
|
|||||||
|
package com.ankurm.advisors;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.concurrent.atomic.AtomicReference;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.Orders;
|
||||||
|
import com.ankurm.advisors.advisor.PiiRedactionAdvisor;
|
||||||
|
import com.ankurm.advisors.advisor.Texts;
|
||||||
|
import com.ankurm.advisors.support.Transcript;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.chat.model.ChatResponse;
|
||||||
|
import org.springframework.ai.chat.model.Generation;
|
||||||
|
import org.springframework.ai.chat.prompt.Prompt;
|
||||||
|
|
||||||
|
/** The smallest advisor test: no ChatClient, no ChatModel, no Spring. A stub chain stands in for "the rest". */
|
||||||
|
class UnitTestingWithoutModelTest {
|
||||||
|
|
||||||
|
/** The rest of the chain: records what reached it and answers with a fixed sentence built from the request. */
|
||||||
|
private static final class StubChain implements CallAdvisorChain {
|
||||||
|
|
||||||
|
final AtomicReference<ChatClientRequest> received = new AtomicReference<>();
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse nextCall(ChatClientRequest request) {
|
||||||
|
received.set(request);
|
||||||
|
String last = request.prompt().getUserMessage().getText();
|
||||||
|
return ChatClientResponse.builder()
|
||||||
|
.chatResponse(ChatResponse.builder()
|
||||||
|
.generations(List.of(new Generation(new AssistantMessage("Noted: " + last))))
|
||||||
|
.build())
|
||||||
|
.build();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public List<CallAdvisor> getCallAdvisors() {
|
||||||
|
return new ArrayList<>();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public CallAdvisorChain copy(CallAdvisor after) {
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void anAdvisorIsAnOrdinaryObjectYouCanCallDirectly() {
|
||||||
|
try (Transcript t = new Transcript("18-unit-test-no-model.txt", "Testing an advisor with a stub chain")) {
|
||||||
|
PiiRedactionAdvisor advisor = new PiiRedactionAdvisor(Orders.PII_REDACTION, true, false);
|
||||||
|
StubChain chain = new StubChain();
|
||||||
|
ChatClientRequest request = ChatClientRequest.builder().prompt(new Prompt("Mail [email protected] please")).build();
|
||||||
|
|
||||||
|
ChatClientResponse response = advisor.adviseCall(request, chain);
|
||||||
|
|
||||||
|
String forwarded = chain.received.get().prompt().getUserMessage().getText();
|
||||||
|
String answer = Texts.text(response);
|
||||||
|
t.line("what the caller sent: %s", request.prompt().getUserMessage().getText());
|
||||||
|
t.line("what reached the rest: %s", forwarded);
|
||||||
|
t.line("what the caller got back: %s", answer);
|
||||||
|
t.line("caller's request unchanged: %s", request.prompt().getUserMessage().getText().contains("[email protected]"));
|
||||||
|
|
||||||
|
assertThat(forwarded).isEqualTo("Mail <EMAIL_1> please");
|
||||||
|
assertThat(answer).isEqualTo("Noted: Mail [email protected] please");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import java.net.URI;
|
||||||
|
import java.net.http.HttpClient;
|
||||||
|
import java.net.http.HttpRequest;
|
||||||
|
import java.net.http.HttpResponse;
|
||||||
|
import java.util.regex.Matcher;
|
||||||
|
import java.util.regex.Pattern;
|
||||||
|
|
||||||
|
/** A synchronous JDK HttpClient -- one send() per request, so a transcript line is one request. */
|
||||||
|
public final class Http {
|
||||||
|
|
||||||
|
public record Reply(int status, String body) {
|
||||||
|
}
|
||||||
|
|
||||||
|
private final HttpClient client = HttpClient.newHttpClient();
|
||||||
|
private final String base;
|
||||||
|
|
||||||
|
public Http(int port) {
|
||||||
|
this.base = "http://127.0.0.1:" + port;
|
||||||
|
}
|
||||||
|
|
||||||
|
public Reply post(String user, String path, String json) {
|
||||||
|
return send(HttpRequest.newBuilder(URI.create(base + path)).header("X-User", user)
|
||||||
|
.header("Content-Type", "application/json").POST(HttpRequest.BodyPublishers.ofString(json)).build());
|
||||||
|
}
|
||||||
|
|
||||||
|
public Reply get(String user, String path) {
|
||||||
|
return send(HttpRequest.newBuilder(URI.create(base + path)).header("X-User", user).GET().build());
|
||||||
|
}
|
||||||
|
|
||||||
|
private Reply send(HttpRequest request) {
|
||||||
|
try {
|
||||||
|
HttpResponse<String> r = client.send(request, HttpResponse.BodyHandlers.ofString());
|
||||||
|
return new Reply(r.statusCode(), r.body());
|
||||||
|
}
|
||||||
|
catch (Exception e) {
|
||||||
|
throw new IllegalStateException(e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
public static String field(String json, String name) {
|
||||||
|
Matcher m = Pattern.compile("\"" + name + "\":\"([^\"]*)\"").matcher(json);
|
||||||
|
return m.find() ? m.group(1) : null;
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import com.ankurm.advisors.advisor.PiiRedactor;
|
||||||
|
import com.ankurm.advisors.advisor.Texts;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The obvious way to restore placeholders in a stream, and the wrong one: restore each chunk on its
|
||||||
|
* own. It works whenever a placeholder happens to arrive whole and fails whenever a chunk boundary
|
||||||
|
* falls inside one. Kept only so a test can show the failure next to the fix.
|
||||||
|
*/
|
||||||
|
public final class NaiveStreamingPii implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return "NaiveStreamingPii";
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return Ordered.HIGHEST_PRECEDENCE + 100;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
PiiRedactor.Session session = new PiiRedactor.Session();
|
||||||
|
return Texts.mapResponseText(chain.nextCall(Texts.mapRequestText(request, session::redact)), session::restore);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
return Flux.defer(() -> {
|
||||||
|
PiiRedactor.Session session = new PiiRedactor.Session();
|
||||||
|
return chain.nextStream(Texts.mapRequestText(request, session::redact))
|
||||||
|
.map(chunk -> Texts.mapResponseText(chunk, session::restore));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,54 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||||
|
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||||
|
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||||
|
import org.springframework.core.Ordered;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/** An advisor that only records when it is entered ("X>") and left ("X<"). */
|
||||||
|
public final class Probe implements CallAdvisor, StreamAdvisor {
|
||||||
|
|
||||||
|
private final String name;
|
||||||
|
|
||||||
|
private final int order;
|
||||||
|
|
||||||
|
private final List<String> events;
|
||||||
|
|
||||||
|
public Probe(String name, int offsetFromHighest, List<String> events) {
|
||||||
|
this.name = name;
|
||||||
|
this.order = Ordered.HIGHEST_PRECEDENCE + offsetFromHighest;
|
||||||
|
this.events = events;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public String getName() {
|
||||||
|
return name;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public int getOrder() {
|
||||||
|
return order;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
|
||||||
|
events.add(name + ">");
|
||||||
|
ChatClientResponse response = chain.nextCall(request);
|
||||||
|
events.add(name + "<");
|
||||||
|
return response;
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
|
||||||
|
return Flux.defer(() -> {
|
||||||
|
events.add(name + ">");
|
||||||
|
return chain.nextStream(request).doOnComplete(() -> events.add(name + "<"));
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,128 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.Collections;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.concurrent.CopyOnWriteArrayList;
|
||||||
|
import java.util.function.Consumer;
|
||||||
|
import java.util.function.Function;
|
||||||
|
|
||||||
|
import com.knuddels.jtokkit.api.EncodingType;
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.chat.messages.Message;
|
||||||
|
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
|
||||||
|
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
||||||
|
import org.springframework.ai.chat.metadata.DefaultUsage;
|
||||||
|
import org.springframework.ai.chat.model.ChatModel;
|
||||||
|
import org.springframework.ai.chat.model.ChatResponse;
|
||||||
|
import org.springframework.ai.chat.model.Generation;
|
||||||
|
import org.springframework.ai.chat.prompt.Prompt;
|
||||||
|
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
|
||||||
|
import org.springframework.ai.tokenizer.TokenCountEstimator;
|
||||||
|
import reactor.core.publisher.Flux;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* A scripted {@link ChatModel}: no network, no API key, no randomness. It records every {@link
|
||||||
|
* Prompt} it is sent, answers with whatever the {@code replier} function returns (by default it
|
||||||
|
* echoes the last message, which makes redaction visible: the model "repeats" what it was given),
|
||||||
|
* and reports token usage counted with the same JTokkit estimator the budget advisor uses, so the
|
||||||
|
* numbers are self-consistent. It says nothing about what a real model would reply.
|
||||||
|
*
|
||||||
|
* <p>{@link #stream} splits the reply into {@code chunkSize}-character pieces. The last piece
|
||||||
|
* carries a finish reason and, when {@code streamUsage} is on, the usage -- the shape OpenAI's
|
||||||
|
* streaming API has when usage reporting is switched on.
|
||||||
|
*/
|
||||||
|
public class RecordingModel implements ChatModel {
|
||||||
|
|
||||||
|
private static final TokenCountEstimator ESTIMATOR = new JTokkitTokenCountEstimator(EncodingType.O200K_BASE);
|
||||||
|
|
||||||
|
private final List<Prompt> prompts = new CopyOnWriteArrayList<>();
|
||||||
|
|
||||||
|
private Function<Prompt, String> replier = p -> "You said: " + p.getInstructions().getLast().getText();
|
||||||
|
|
||||||
|
private Consumer<String> trace = s -> {
|
||||||
|
};
|
||||||
|
|
||||||
|
private int chunkSize = 8;
|
||||||
|
|
||||||
|
private boolean streamUsage = true;
|
||||||
|
|
||||||
|
private boolean finishReason = true;
|
||||||
|
|
||||||
|
public RecordingModel replier(Function<Prompt, String> replier) {
|
||||||
|
this.replier = replier;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
public RecordingModel trace(Consumer<String> trace) {
|
||||||
|
this.trace = trace;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
public RecordingModel chunkSize(int chunkSize) {
|
||||||
|
this.chunkSize = chunkSize;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
public RecordingModel streamUsage(boolean streamUsage) {
|
||||||
|
this.streamUsage = streamUsage;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Whether the last streamed chunk carries a finish reason (real providers do; this lets a test see what happens if not). */
|
||||||
|
public RecordingModel finishReason(boolean finishReason) {
|
||||||
|
this.finishReason = finishReason;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
public List<Prompt> prompts() {
|
||||||
|
return Collections.unmodifiableList(prompts);
|
||||||
|
}
|
||||||
|
|
||||||
|
public Prompt lastPrompt() {
|
||||||
|
return prompts.getLast();
|
||||||
|
}
|
||||||
|
|
||||||
|
public int callCount() {
|
||||||
|
return prompts.size();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatResponse call(Prompt prompt) {
|
||||||
|
prompts.add(prompt);
|
||||||
|
trace.accept("model");
|
||||||
|
String reply = replier.apply(prompt);
|
||||||
|
return new ChatResponse(List.of(new Generation(new AssistantMessage(reply), finish())), metadata(prompt, reply));
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public Flux<ChatResponse> stream(Prompt prompt) {
|
||||||
|
return Flux.defer(() -> {
|
||||||
|
prompts.add(prompt);
|
||||||
|
trace.accept("model");
|
||||||
|
String reply = replier.apply(prompt);
|
||||||
|
List<ChatResponse> chunks = new ArrayList<>();
|
||||||
|
for (int i = 0; i < reply.length(); i += chunkSize) {
|
||||||
|
String piece = reply.substring(i, Math.min(reply.length(), i + chunkSize));
|
||||||
|
boolean last = i + chunkSize >= reply.length();
|
||||||
|
chunks.add(last
|
||||||
|
? new ChatResponse(List.of(finishReason ? new Generation(new AssistantMessage(piece), finish()) : new Generation(new AssistantMessage(piece))),
|
||||||
|
streamUsage ? metadata(prompt, reply) : ChatResponseMetadata.builder().build())
|
||||||
|
: new ChatResponse(List.of(new Generation(new AssistantMessage(piece)))));
|
||||||
|
}
|
||||||
|
return Flux.fromIterable(chunks);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
private static ChatGenerationMetadata finish() {
|
||||||
|
return ChatGenerationMetadata.builder().finishReason("STOP").build();
|
||||||
|
}
|
||||||
|
|
||||||
|
private static ChatResponseMetadata metadata(Prompt prompt, String reply) {
|
||||||
|
int in = 0;
|
||||||
|
for (Message m : prompt.getInstructions()) {
|
||||||
|
in += ESTIMATOR.estimate(m.getText());
|
||||||
|
}
|
||||||
|
return ChatResponseMetadata.builder().usage(new DefaultUsage(in, ESTIMATOR.estimate(reply))).build();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||||
|
import org.springframework.ai.chat.model.ChatModel;
|
||||||
|
import org.springframework.ai.chat.model.ChatResponse;
|
||||||
|
import org.springframework.ai.chat.model.Generation;
|
||||||
|
import org.springframework.ai.chat.prompt.ChatOptions;
|
||||||
|
import org.springframework.ai.chat.prompt.Prompt;
|
||||||
|
import org.springframework.ai.model.tool.ToolCallingChatOptions;
|
||||||
|
|
||||||
|
import java.util.ArrayDeque;
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.Deque;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.Queue;
|
||||||
|
import java.util.concurrent.CopyOnWriteArrayList;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* A hand-written {@link ChatModel} that returns a pre-programmed queue of {@link ChatResponse}s,
|
||||||
|
* one per {@link #call(Prompt)} invocation, instead of calling a real LLM API. {@code ChatModel}
|
||||||
|
* has exactly one abstract method -- confirmed with {@code javap} against
|
||||||
|
* {@code spring-ai-model-2.0.1.jar}, everything else on the interface has a default implementation
|
||||||
|
* -- so this is the entire surface a deterministic test double needs to implement.
|
||||||
|
*
|
||||||
|
* <p>{@link org.springframework.ai.chat.client.advisor.ToolCallingAdvisor ToolCallingAdvisor} (and
|
||||||
|
* its subclass, {@code ToolSearchToolCallingAdvisor}) drive the tool-calling loop by calling the
|
||||||
|
* underlying {@link ChatModel} once per round: once to get the model's first response (which may
|
||||||
|
* contain tool calls), then once more per round of tool results fed back in, until a response
|
||||||
|
* comes back with no tool calls. Queuing responses here lets a test assert the exact shape of that
|
||||||
|
* loop -- how many rounds it took, what tool calls appeared, what the final answer was -- without
|
||||||
|
* an API key, network access, or the nondeterminism of an actual model.
|
||||||
|
*
|
||||||
|
* <p>Every {@link Prompt} the advisor sends is recorded in {@link #capturedPrompts()} so a test can
|
||||||
|
* also assert on what the advisor sent back on the next round -- in particular, that a
|
||||||
|
* {@code ToolResponseMessage} was appended after tool execution.
|
||||||
|
*/
|
||||||
|
public class ScriptedToolModel implements ChatModel {
|
||||||
|
|
||||||
|
private final Queue<ChatResponse> script;
|
||||||
|
private final List<Prompt> capturedPrompts = new CopyOnWriteArrayList<>();
|
||||||
|
|
||||||
|
private ScriptedToolModel(Deque<ChatResponse> script) {
|
||||||
|
this.script = script;
|
||||||
|
}
|
||||||
|
|
||||||
|
public static Builder builder() {
|
||||||
|
return new Builder();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ChatResponse call(Prompt prompt) {
|
||||||
|
this.capturedPrompts.add(prompt);
|
||||||
|
ChatResponse next = this.script.poll();
|
||||||
|
if (next == null) {
|
||||||
|
throw new IllegalStateException(
|
||||||
|
"ScriptedChatModel ran out of queued responses after " + this.capturedPrompts.size()
|
||||||
|
+ " calls. Prompts so far: " + this.capturedPrompts);
|
||||||
|
}
|
||||||
|
return next;
|
||||||
|
}
|
||||||
|
|
||||||
|
public List<Prompt> capturedPrompts() {
|
||||||
|
return List.copyOf(this.capturedPrompts);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* {@code DefaultChatClientUtils} builds every outgoing {@link Prompt}'s options from
|
||||||
|
* {@code chatModel.getOptions().mutate()} -- not {@code getDefaultOptions()}, which is a
|
||||||
|
* separate default method nobody in the request-building path actually calls. Confirmed by
|
||||||
|
* disassembling both: {@link ChatModel#getOptions()}'s default body is a bare
|
||||||
|
* {@code ChatOptions.builder().build()}, a plain {@link ChatOptions} that is not a
|
||||||
|
* {@link ToolCallingChatOptions}. Since {@code ToolCallingAdvisor.adviseCall} starts with an
|
||||||
|
* {@code instanceof ToolCallingChatOptions} check on that exact object and falls straight
|
||||||
|
* through to the underlying model with no tool loop at all when it fails, leaving this method's
|
||||||
|
* default in place silently turns every tool call in this repository into a no-op -- confirmed
|
||||||
|
* the hard way, by a first version of this class that overrode {@code getDefaultOptions()}
|
||||||
|
* instead and watched every test below get back an empty answer after exactly one model call.
|
||||||
|
* Real providers (OpenAI, Anthropic) return their own {@code ToolCallingChatOptions}
|
||||||
|
* implementation from {@code getOptions()} for the same reason.
|
||||||
|
*/
|
||||||
|
@Override
|
||||||
|
public ChatOptions getOptions() {
|
||||||
|
return ToolCallingChatOptions.builder().build();
|
||||||
|
}
|
||||||
|
|
||||||
|
public int callCount() {
|
||||||
|
return this.capturedPrompts.size();
|
||||||
|
}
|
||||||
|
|
||||||
|
public static final class Builder {
|
||||||
|
|
||||||
|
private final Deque<ChatResponse> script = new ArrayDeque<>();
|
||||||
|
|
||||||
|
private int promptTokens = -1;
|
||||||
|
|
||||||
|
private int completionTokens;
|
||||||
|
|
||||||
|
private Builder() {
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Responses queued after this call report this usage, the way a provider does. */
|
||||||
|
public Builder reportingUsage(int prompt, int completion) {
|
||||||
|
this.promptTokens = prompt;
|
||||||
|
this.completionTokens = completion;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
private ChatResponse response(Generation generation) {
|
||||||
|
if (promptTokens < 0) {
|
||||||
|
return new ChatResponse(List.of(generation));
|
||||||
|
}
|
||||||
|
return ChatResponse.builder().generations(List.of(generation)).metadata(
|
||||||
|
org.springframework.ai.chat.metadata.ChatResponseMetadata.builder()
|
||||||
|
.usage(new org.springframework.ai.chat.metadata.DefaultUsage(promptTokens, completionTokens)).build())
|
||||||
|
.build();
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Queues a plain-text final answer with no tool calls -- ends the tool-calling loop. */
|
||||||
|
public Builder thenRespond(String text) {
|
||||||
|
this.script.add(response(new Generation(new AssistantMessage(text))));
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Queues an assistant turn that calls one or more tools, continuing the loop. */
|
||||||
|
public Builder thenCallTools(AssistantMessage.ToolCall... toolCalls) {
|
||||||
|
AssistantMessage message = AssistantMessage.builder()
|
||||||
|
.content("")
|
||||||
|
.toolCalls(List.of(toolCalls))
|
||||||
|
.build();
|
||||||
|
this.script.add(response(new Generation(message)));
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
public ScriptedToolModel build() {
|
||||||
|
return new ScriptedToolModel(new ArrayDeque<>(this.script));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,31 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import java.util.List;
|
||||||
|
|
||||||
|
import org.springframework.ai.chat.messages.Message;
|
||||||
|
|
||||||
|
/** One-line renderings of messages so transcripts stay short and diff cleanly. */
|
||||||
|
public final class Show {
|
||||||
|
|
||||||
|
private Show() {
|
||||||
|
}
|
||||||
|
|
||||||
|
public static String one(Message m) {
|
||||||
|
return String.format("%-9s %s", m.getMessageType(), m.getText());
|
||||||
|
}
|
||||||
|
|
||||||
|
public static String inline(List<Message> messages) {
|
||||||
|
return messages.stream().map(Show::brief).reduce((a, b) -> a + " | " + b).orElse("(empty)");
|
||||||
|
}
|
||||||
|
|
||||||
|
private static String brief(Message m) {
|
||||||
|
char kind = m.getMessageType().name().charAt(0);
|
||||||
|
if (m instanceof org.springframework.ai.chat.messages.AssistantMessage a && a.hasToolCalls()) {
|
||||||
|
return kind + ":[tool call " + a.getToolCalls().getFirst().name() + "]";
|
||||||
|
}
|
||||||
|
if (m instanceof org.springframework.ai.chat.messages.ToolResponseMessage r) {
|
||||||
|
return kind + ":[tool result " + r.getResponses().getFirst().responseData() + "]";
|
||||||
|
}
|
||||||
|
return kind + ":" + m.getText();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
package com.ankurm.advisors.support;
|
||||||
|
|
||||||
|
import org.springframework.boot.test.context.TestConfiguration;
|
||||||
|
import org.springframework.context.annotation.Bean;
|
||||||
|
|
||||||
|
/** Puts one shared {@link RecordingModel} in the context so a test can inspect what it was sent. */
|
||||||
|
@TestConfiguration(proxyBeanMethods = false)
|
||||||
|
public class TestModelConfig {
|
||||||
|
|
||||||
|
@Bean
|
||||||
|
RecordingModel recordingModel() {
|
||||||
|
return new RecordingModel();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,47 @@
|
|||||||
|
package com.ankurm.advisors.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/} (repository root, not {@code docs/}) and
|
||||||
|
* echoes it to the console. Every console block quoted in the article comes out of one of these
|
||||||
|
* files verbatim.
|
||||||
|
*/
|
||||||
|
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