diff --git a/README.md b/README.md index 5032a73..f9588c3 100644 --- a/README.md +++ b/README.md @@ -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/) | | [`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/) | +| [`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/). diff --git a/advisors/.gitignore b/advisors/.gitignore new file mode 100644 index 0000000..2f7896d --- /dev/null +++ b/advisors/.gitignore @@ -0,0 +1 @@ +target/ diff --git a/advisors/README.md b/advisors/README.md new file mode 100644 index 0000000..91e2a40 --- /dev/null +++ b/advisors/README.md @@ -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; `` 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. diff --git a/advisors/output/01-chain-order.txt b/advisors/output/01-chain-order.txt new file mode 100644 index 0000000..9a8624d --- /dev/null +++ b/advisors/output/01-chain-order.txt @@ -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< diff --git a/advisors/output/02-chain-order-ties.txt b/advisors/output/02-chain-order-ties.txt new file mode 100644 index 0000000..ee80b68 --- /dev/null +++ b/advisors/output/02-chain-order-ties.txt @@ -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). diff --git a/advisors/output/03-chain-contents.txt b/advisors/output/03-chain-contents.txt new file mode 100644 index 0000000..4bc1e0c --- /dev/null +++ b/advisors/output/03-chain-contents.txt @@ -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 diff --git a/advisors/output/04-base-advisor-stream.txt b/advisors/output/04-base-advisor-stream.txt new file mode 100644 index 0000000..4f55048 --- /dev/null +++ b/advisors/output/04-base-advisor-stream.txt @@ -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 diff --git a/advisors/output/05-logging-advisor.txt b/advisors/output/05-logging-advisor.txt new file mode 100644 index 0000000..0132725 --- /dev/null +++ b/advisors/output/05-logging-advisor.txt @@ -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 diff --git a/advisors/output/06-logging-failure.txt b/advisors/output/06-logging-failure.txt new file mode 100644 index 0000000..320939b --- /dev/null +++ b/advisors/output/06-logging-failure.txt @@ -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 diff --git a/advisors/output/07-pii-redaction.txt b/advisors/output/07-pii-redaction.txt new file mode 100644 index 0000000..7d27c19 --- /dev/null +++ b/advisors/output/07-pii-redaction.txt @@ -0,0 +1,8 @@ +# PiiRedactionAdvisor on a call + +caller sends: Hi, I'm Priya. Email priya@example.com or call +91 98765 43210. Card 4111 1111 1111 1111, order 1234 5678 9012 3456, and again priya@example.com. +model was sent: Hi, I'm Priya. Email or call . Card , order 1234 5678 9012 3456, and again . +caller receives: You said: Hi, I'm Priya. Email priya@example.com or call +91 98765 43210. Card 4111 1111 1111 1111, order 1234 5678 9012 3456, and again priya@example.com. + +with restore switched off, the caller receives: + You said: Hi, I'm Priya. Email or call . Card , order 1234 5678 9012 3456, and again . diff --git a/advisors/output/08-pii-limits.txt b/advisors/output/08-pii-limits.txt new file mode 100644 index 0000000..e977f9e --- /dev/null +++ b/advisors/output/08-pii-limits.txt @@ -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 + diff --git a/advisors/output/09-pii-stream-boundary.txt b/advisors/output/09-pii-stream-boundary.txt new file mode 100644 index 0000000..def9563 --- /dev/null +++ b/advisors/output/09-pii-stream-boundary.txt @@ -0,0 +1,11 @@ +# Restoring placeholders in a stream + +the model streams 5-character chunks: Sure,| I wi|ll wr|ite t|o now|. + +restore each chunk on its own: + chunks: Sure,| I wi|ll wr|ite t|o now|. + joined: Sure, I will write to now. + +PiiRedactionAdvisor (holds back from an unfinished "<"): + chunks: Sure,| I wi|ll wr|ite t|o |priya@example.com now|. + joined: Sure, I will write to priya@example.com now. diff --git a/advisors/output/10-pii-order.txt b/advisors/output/10-pii-order.txt new file mode 100644 index 0000000..fb9bd59 --- /dev/null +++ b/advisors/output/10-pii-order.txt @@ -0,0 +1,19 @@ +# Redaction order versus logging and memory + +memory advisor is fixed at HIGHEST_PRECEDENCE + 200 +user says: My email is priya@example.com + +A redaction +100, logging +400 (redaction outside both): + log line saw: "My email is " + memory stored: "My email is " + model was sent: "My email is " + +B redaction +100, logging +50 (logging outside redaction): + log line saw: "My email is priya@example.com" + memory stored: "My email is " + model was sent: "My email is " + +C redaction +300 (inside the memory advisor), logging +400: + log line saw: "My email is " + memory stored: "My email is priya@example.com" + model was sent: "My email is " diff --git a/advisors/output/11-pii-multi-turn.txt b/advisors/output/11-pii-multi-turn.txt new file mode 100644 index 0000000..42d8cfc --- /dev/null +++ b/advisors/output/11-pii-multi-turn.txt @@ -0,0 +1,18 @@ +# Placeholders across two turns with memory + +turn 1: My email is priya@example.com +turn 2: Also cc bob@example.com + +numbering restarts on every request: + model was sent on turn 2: U:My email is | A:You said: My email is | U:Also cc + caller receives: You said: Also cc bob@example.com + +numbering kept per conversation (the default): + model was sent on turn 2: U:My email is | A:You said: My email is | U:Also cc + caller receives: You said: Also cc bob@example.com + +what the memory stores (placeholders, never the addresses): + USER My email is + ASSISTANT You said: My email is + USER Also cc + ASSISTANT You said: Also cc diff --git a/advisors/output/12-token-budget.txt b/advisors/output/12-token-budget.txt new file mode 100644 index 0000000..af28430 --- /dev/null +++ b/advisors/output/12-token-budget.txt @@ -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 diff --git a/advisors/output/13-token-budget-stream.txt b/advisors/output/13-token-budget-stream.txt new file mode 100644 index 0000000..9cd6e39 --- /dev/null +++ b/advisors/output/13-token-budget-stream.txt @@ -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. diff --git a/advisors/output/14-token-budget-stream-refusal.txt b/advisors/output/14-token-budget-stream-refusal.txt new file mode 100644 index 0000000..199a9a5 --- /dev/null +++ b/advisors/output/14-token-budget-stream-refusal.txt @@ -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 diff --git a/advisors/output/15-tool-loop-order.txt b/advisors/output/15-tool-loop-order.txt new file mode 100644 index 0000000..7815384 --- /dev/null +++ b/advisors/output/15-tool-loop-order.txt @@ -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 diff --git a/advisors/output/16-error-propagation.txt b/advisors/output/16-error-propagation.txt new file mode 100644 index 0000000..bd52214 --- /dev/null +++ b/advisors/output/16-error-propagation.txt @@ -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 diff --git a/advisors/output/17-context.txt b/advisors/output/17-context.txt new file mode 100644 index 0000000..27d5239 --- /dev/null +++ b/advisors/output/17-context.txt @@ -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 diff --git a/advisors/output/18-unit-test-no-model.txt b/advisors/output/18-unit-test-no-model.txt new file mode 100644 index 0000000..d5fcafd --- /dev/null +++ b/advisors/output/18-unit-test-no-model.txt @@ -0,0 +1,6 @@ +# Testing an advisor with a stub chain + +what the caller sent: Mail priya@example.com please +what reached the rest: Mail please +what the caller got back: Noted: Mail priya@example.com please +caller's request unchanged: true diff --git a/advisors/output/19-web-429.txt b/advisors/output/19-web-429.txt new file mode 100644 index 0000000..9154686 --- /dev/null +++ b/advisors/output/19-web-429.txt @@ -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) diff --git a/advisors/output/20-built-in-orders.txt b/advisors/output/20-built-in-orders.txt new file mode 100644 index 0000000..deca649 --- /dev/null +++ b/advisors/output/20-built-in-orders.txt @@ -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 diff --git a/advisors/output/21-tool-loop-budget.txt b/advisors/output/21-tool-loop-budget.txt new file mode 100644 index 0000000..c02e837 --- /dev/null +++ b/advisors/output/21-tool-loop-budget.txt @@ -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 diff --git a/advisors/pom.xml b/advisors/pom.xml new file mode 100644 index 0000000..4246ffc --- /dev/null +++ b/advisors/pom.xml @@ -0,0 +1,69 @@ + + + 4.0.0 + + + org.springframework.boot + spring-boot-starter-parent + 4.1.1 + + + + com.ankurm + advisors + 1.0.0 + advisors + Custom Spring AI 2.0 advisors: audit logging, PII redaction and token budgets, with the ordering and streaming traps reproduced offline. + + + 25 + 2.0.1 + + + + + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + + + + + + + org.springframework.boot + spring-boot-starter-webmvc + + + org.springframework.ai + spring-ai-starter-model-openai + + + + org.springframework.boot + spring-boot-starter-test + test + + + + + + + org.springframework.boot + spring-boot-maven-plugin + + + org.apache.maven.plugins + maven-surefire-plugin + + -Duser.timezone=UTC -Dstdout.encoding=UTF-8 -Dfile.encoding=UTF-8 + + + + + diff --git a/advisors/scripts/run-all.sh b/advisors/scripts/run-all.sh new file mode 100755 index 0000000..d4100e4 --- /dev/null +++ b/advisors/scripts/run-all.sh @@ -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 diff --git a/advisors/src/main/java/com/ankurm/advisors/AdvisorsApplication.java b/advisors/src/main/java/com/ankurm/advisors/AdvisorsApplication.java new file mode 100644 index 0000000..2d194a0 --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/AdvisorsApplication.java @@ -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); + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/LoggingAdvisor.java b/advisors/src/main/java/com/ankurm/advisors/advisor/LoggingAdvisor.java new file mode 100644 index 0000000..3b4a6d5 --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/LoggingAdvisor.java @@ -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 inside this + * advisor in the chain: the model, plus any advisor with a higher order number. + * + *

Message text is off by default. When it is on, what you read is whatever reaches this + * 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 sink; + + private final LongSupplier nanoClock; + + public LoggingAdvisor(String name, int order, boolean logContent, Consumer 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 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 = 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 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; + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/Orders.java b/advisors/src/main/java/com/ankurm/advisors/advisor/Orders.java new file mode 100644 index 0000000..1795fcf --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/Orders.java @@ -0,0 +1,28 @@ +package com.ankurm.advisors.advisor; + +import org.springframework.core.Ordered; + +/** + * Where the three custom advisors sit, written down once. A lower number runs earlier + * on the way in (and later on the way out), so it sits further out. + * + * 

+ *   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
+ *
+ */ +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() { + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/PiiRedactionAdvisor.java b/advisors/src/main/java/com/ankurm/advisors/advisor/PiiRedactionAdvisor.java new file mode 100644 index 0000000..3e65a9a --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/PiiRedactionAdvisor.java @@ -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 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 adviseStream(ChatClientRequest request, StreamAdvisorChain chain) { + return Flux.defer(() -> { + PiiRedactor.Session session = sessionFor(request); + Flux 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 } 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 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(); + } + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/PiiRedactor.java b/advisors/src/main/java/com/ankurm/advisors/advisor/PiiRedactor.java new file mode 100644 index 0000000..851ab3c --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/PiiRedactor.java @@ -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 }. 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. + * + *

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("(? placeholderByValue = new LinkedHashMap<>(); + + private final Map valueByPlaceholder = new LinkedHashMap<>(); + + private final Map 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 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 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; + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/Texts.java b/advisors/src/main/java/com/ankurm/advisors/advisor/Texts.java new file mode 100644 index 0000000..7ff08b3 --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/Texts.java @@ -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 f) { + List 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 f) { + ChatResponse original = response.chatResponse(); + if (original == null) { + return response; + } + List 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(); + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/TokenBudgetAdvisor.java b/advisors/src/main/java/com/ankurm/advisors/advisor/TokenBudgetAdvisor.java new file mode 100644 index 0000000..66ef483 --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/TokenBudgetAdvisor.java @@ -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: + *

    + *
  • a per-request cap on the estimated prompt size, so one pasted log file cannot cost + * a fortune;
    • + *
  • 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.
    • + *
+ * 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 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 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; + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/advisor/TokenBudgetExceededException.java b/advisors/src/main/java/com/ankurm/advisors/advisor/TokenBudgetExceededException.java new file mode 100644 index 0000000..edc6be3 --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/advisor/TokenBudgetExceededException.java @@ -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); + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/config/AdvisorConfig.java b/advisors/src/main/java/com/ankurm/advisors/config/AdvisorConfig.java new file mode 100644 index 0000000..74c876d --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/config/AdvisorConfig.java @@ -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(); + } +} diff --git a/advisors/src/main/java/com/ankurm/advisors/web/ChatController.java b/advisors/src/main/java/com/ankurm/advisors/web/ChatController.java new file mode 100644 index 0000000..81fda01 --- /dev/null +++ b/advisors/src/main/java/com/ankurm/advisors/web/ChatController.java @@ -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 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); + } +} diff --git a/advisors/src/main/resources/application.yml b/advisors/src/main/resources/application.yml new file mode 100644 index 0000000..a02ecff --- /dev/null +++ b/advisors/src/main/resources/application.yml @@ -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 diff --git a/advisors/src/test/java/com/ankurm/advisors/AdvisorWebTest.java b/advisors/src/test/java/com/ankurm/advisors/AdvisorWebTest.java new file mode 100644 index 0000000..1f21459 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/AdvisorWebTest.java @@ -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); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/BaseAdvisorStreamTest.java b/advisors/src/test/java/com/ankurm/advisors/BaseAdvisorStreamTest.java new file mode 100644 index 0000000..8df1532 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/BaseAdvisorStreamTest.java @@ -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 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 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()); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/BuiltInOrdersTest.java b/advisors/src/test/java/com/ankurm/advisors/BuiltInOrdersTest.java new file mode 100644 index 0000000..0da102e --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/BuiltInOrdersTest.java @@ -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 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(); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/ChainOrderTest.java b/advisors/src/test/java/com/ankurm/advisors/ChainOrderTest.java new file mode 100644 index 0000000..07fae71 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/ChainOrderTest.java @@ -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 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 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 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 callChain = new ArrayList<>(); + + final List 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 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"); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/ContextAndImmutabilityTest.java b/advisors/src/test/java/com/ankurm/advisors/ContextAndImmutabilityTest.java new file mode 100644 index 0000000..7acaa58 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/ContextAndImmutabilityTest.java @@ -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 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 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")); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/ErrorPropagationTest.java b/advisors/src/test/java/com/ankurm/advisors/ErrorPropagationTest.java new file mode 100644 index 0000000..e6c9768 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/ErrorPropagationTest.java @@ -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 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 build, Supplier 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"); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/LoggingAdvisorTest.java b/advisors/src/test/java/com/ankurm/advisors/LoggingAdvisorTest.java new file mode 100644 index 0000000..f14c51a --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/LoggingAdvisorTest.java @@ -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 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 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 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"); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/PiiMultiTurnTest.java b/advisors/src/test/java/com/ankurm/advisors/PiiMultiTurnTest.java new file mode 100644 index 0000000..9668169 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/PiiMultiTurnTest.java @@ -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 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 priya@example.com").advisors(a -> a.param(ChatMemory.CONVERSATION_ID, "c")).call() + .content(); + String reply = client.prompt().user("Also cc bob@example.com").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 priya@example.com"); + t.line("turn 2: Also cc bob@example.com"); + + 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 ").contains("U:Also cc "); + assertThat(scoped.modelSawTurn2()).contains("U:My email is ").contains("U:Also cc "); + assertThat(scoped.stored()).noneMatch(s -> s.contains("@")); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/PiiOrderTest.java b/advisors/src/test/java/com/ankurm/advisors/PiiOrderTest.java new file mode 100644 index 0000000..28b22f1 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/PiiOrderTest.java @@ -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 priya@example.com"; + + private record Outcome(String logged, String stored, String modelSaw) { + } + + private static Outcome run(int piiOffset, int loggingOffset) { + List 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("priya@example.com"); + assertThat(a.stored()).doesNotContain("priya@example.com"); + assertThat(b.logged()).contains("priya@example.com"); + assertThat(b.stored()).doesNotContain("priya@example.com"); + assertThat(c.stored()).contains("priya@example.com"); + assertThat(c.modelSaw()).doesNotContain("priya@example.com"); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/PiiRedactionTest.java b/advisors/src/test/java/com/ankurm/advisors/PiiRedactionTest.java new file mode 100644 index 0000000..852739d --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/PiiRedactionTest.java @@ -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 priya@example.com or call +91 98765 43210. " + + "Card 4111 1111 1111 1111, order 1234 5678 9012 3456, and again priya@example.com."; + + @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("", "", "").doesNotContain("priya@example.com", "98765", "4111"); + assertThat(sent).contains("1234 5678 9012 3456").contains("Priya"); + assertThat(sent.split("", -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 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 priya@example.com"; + + List naive = ChatClient.builder(model).defaultAdvisors(new NaiveStreamingPii()).build().prompt() + .user(input).stream().content().collectList().block(); + List raw = model.stream(new org.springframework.ai.chat.prompt.Prompt("")) + .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 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 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(); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/ToolLoopBudgetTest.java b/advisors/src/test/java/com/ankurm/advisors/ToolLoopBudgetTest.java new file mode 100644 index 0000000..d80e669 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/ToolLoopBudgetTest.java @@ -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); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/ToolLoopOrderTest.java b/advisors/src/test/java/com/ankurm/advisors/ToolLoopOrderTest.java new file mode 100644 index 0000000..29b1ca0 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/ToolLoopOrderTest.java @@ -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 run(int loggingOffset) { + List 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 outside = run(250); + List 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); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/UnitTestingWithoutModelTest.java b/advisors/src/test/java/com/ankurm/advisors/UnitTestingWithoutModelTest.java new file mode 100644 index 0000000..2c3e101 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/UnitTestingWithoutModelTest.java @@ -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 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 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 priya@example.com 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("priya@example.com")); + + assertThat(forwarded).isEqualTo("Mail please"); + assertThat(answer).isEqualTo("Noted: Mail priya@example.com please"); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/Http.java b/advisors/src/test/java/com/ankurm/advisors/support/Http.java new file mode 100644 index 0000000..c315e66 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/Http.java @@ -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 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; + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/NaiveStreamingPii.java b/advisors/src/test/java/com/ankurm/advisors/support/NaiveStreamingPii.java new file mode 100644 index 0000000..55be8ab --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/NaiveStreamingPii.java @@ -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 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)); + }); + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/Probe.java b/advisors/src/test/java/com/ankurm/advisors/support/Probe.java new file mode 100644 index 0000000..0e602d4 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/Probe.java @@ -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 events; + + public Probe(String name, int offsetFromHighest, List 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 adviseStream(ChatClientRequest request, StreamAdvisorChain chain) { + return Flux.defer(() -> { + events.add(name + ">"); + return chain.nextStream(request).doOnComplete(() -> events.add(name + "<")); + }); + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/RecordingModel.java b/advisors/src/test/java/com/ankurm/advisors/support/RecordingModel.java new file mode 100644 index 0000000..c8b3fe2 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/RecordingModel.java @@ -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. + * + *

{@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 prompts = new CopyOnWriteArrayList<>(); + + private Function replier = p -> "You said: " + p.getInstructions().getLast().getText(); + + private Consumer trace = s -> { + }; + + private int chunkSize = 8; + + private boolean streamUsage = true; + + private boolean finishReason = true; + + public RecordingModel replier(Function replier) { + this.replier = replier; + return this; + } + + public RecordingModel trace(Consumer 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 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 stream(Prompt prompt) { + return Flux.defer(() -> { + prompts.add(prompt); + trace.accept("model"); + String reply = replier.apply(prompt); + List 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(); + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/ScriptedToolModel.java b/advisors/src/test/java/com/ankurm/advisors/support/ScriptedToolModel.java new file mode 100644 index 0000000..e13891b --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/ScriptedToolModel.java @@ -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. + * + *

{@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. + * + *

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 script; + private final List capturedPrompts = new CopyOnWriteArrayList<>(); + + private ScriptedToolModel(Deque 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 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 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)); + } + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/Show.java b/advisors/src/test/java/com/ankurm/advisors/support/Show.java new file mode 100644 index 0000000..c11f4ff --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/Show.java @@ -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 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(); + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/TestModelConfig.java b/advisors/src/test/java/com/ankurm/advisors/support/TestModelConfig.java new file mode 100644 index 0000000..fcd8111 --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/TestModelConfig.java @@ -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(); + } +} diff --git a/advisors/src/test/java/com/ankurm/advisors/support/Transcript.java b/advisors/src/test/java/com/ankurm/advisors/support/Transcript.java new file mode 100644 index 0000000..c47b3ae --- /dev/null +++ b/advisors/src/test/java/com/ankurm/advisors/support/Transcript.java @@ -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); + } +}