mockserver_client/
llm.rs

1//! Fluent LLM-mocking builders for the MockServer Rust client.
2//!
3//! This module is an idiomatic Rust port of the Java/Node/Python client LLM
4//! builders (`org.mockserver.client.Llm` / `LlmMockBuilder` /
5//! `LlmConversationBuilder` / `LlmFailoverBuilder` / `TurnBuilder`) and the
6//! underlying model classes (`Completion`, `ToolUse`, `Usage`,
7//! `StreamingPhysics`, `EmbeddingResponse`). The builders produce exactly the
8//! same expectation wire JSON the other clients emit, so a mock scripted from
9//! Rust is byte-for-byte equivalent (modulo JSON key order, which the server
10//! ignores) to one scripted from Java.
11//!
12//! The expectation action is carried in the `httpLlmResponse` field of an
13//! expectation (a sibling of `httpRequest`, `scenarioName`, `scenarioState`,
14//! `newScenarioState`, `times`, `timeToLive`, `httpResponse`). Null/unset
15//! fields are omitted from the JSON, matching the server's `NON_NULL`
16//! serialization.
17//!
18//! # Example
19//!
20//! ```
21//! use mockserver_client::llm::{llm_mock, completion, Provider};
22//!
23//! let expectation = llm_mock("/v1/chat/completions")
24//!     .with_provider(Provider::OPENAI)
25//!     .with_model("gpt-4o")
26//!     .responding_with(completion().with_text("Hello!"))
27//!     .build();
28//!
29//! assert_eq!(expectation["httpLlmResponse"]["provider"], "OPENAI");
30//! ```
31
32use serde_json::{json, Map, Value};
33
34use crate::error::Result;
35use crate::MockServerClient;
36
37// ---------------------------------------------------------------------------
38// Provider / Role string constants (serialized as the UPPERCASE enum name)
39// ---------------------------------------------------------------------------
40
41/// LLM provider names. Serialized on the wire as the upper-case enum name
42/// (mirrors `org.mockserver.model.Provider`).
43pub struct Provider;
44
45impl Provider {
46    pub const ANTHROPIC: &'static str = "ANTHROPIC";
47    pub const OPENAI: &'static str = "OPENAI";
48    pub const OPENAI_RESPONSES: &'static str = "OPENAI_RESPONSES";
49    pub const GEMINI: &'static str = "GEMINI";
50    pub const BEDROCK: &'static str = "BEDROCK";
51    pub const AZURE_OPENAI: &'static str = "AZURE_OPENAI";
52    pub const OLLAMA: &'static str = "OLLAMA";
53}
54
55/// Parsed-message roles (mirrors `org.mockserver.llm.ParsedMessage.Role`).
56pub struct Role;
57
58impl Role {
59    pub const USER: &'static str = "USER";
60    pub const ASSISTANT: &'static str = "ASSISTANT";
61    pub const TOOL: &'static str = "TOOL";
62    pub const SYSTEM: &'static str = "SYSTEM";
63}
64
65// ---------------------------------------------------------------------------
66// Small helpers
67// ---------------------------------------------------------------------------
68
69/// Build a `{method: POST, path}` request matcher object.
70fn post_matcher(path: &str) -> Value {
71    json!({ "method": "POST", "path": path })
72}
73
74/// Insert a key only when the value is present (mirrors NON_NULL omission).
75fn insert_some(map: &mut Map<String, Value>, key: &str, value: Option<Value>) {
76    if let Some(v) = value {
77        map.insert(key.to_string(), v);
78    }
79}
80
81// ---------------------------------------------------------------------------
82// ToolUse
83// ---------------------------------------------------------------------------
84
85/// A single tool/function call emitted by the assistant
86/// (mirrors `org.mockserver.model.ToolUse`).
87#[derive(Debug, Clone, Default, PartialEq)]
88pub struct ToolUse {
89    name: String,
90    id: Option<String>,
91    arguments: Option<String>,
92}
93
94impl ToolUse {
95    /// Create a tool call with the given tool name.
96    pub fn new(name: impl Into<String>) -> Self {
97        Self {
98            name: name.into(),
99            id: None,
100            arguments: None,
101        }
102    }
103
104    /// Set the tool-call id.
105    pub fn with_id(mut self, id: impl Into<String>) -> Self {
106        self.id = Some(id.into());
107        self
108    }
109
110    /// Set the tool name.
111    pub fn with_name(mut self, name: impl Into<String>) -> Self {
112        self.name = name.into();
113        self
114    }
115
116    /// Set the call arguments as a JSON string (matching the Java API).
117    pub fn with_arguments(mut self, arguments: impl Into<String>) -> Self {
118        self.arguments = Some(arguments.into());
119        self
120    }
121
122    fn to_value(&self) -> Value {
123        let mut map = Map::new();
124        insert_some(&mut map, "id", self.id.clone().map(Value::from));
125        map.insert("name".to_string(), Value::from(self.name.clone()));
126        insert_some(&mut map, "arguments", self.arguments.clone().map(Value::from));
127        Value::Object(map)
128    }
129}
130
131/// Factory mirroring `ToolUse.toolUse(name)`.
132pub fn tool_use(name: impl Into<String>) -> ToolUse {
133    ToolUse::new(name)
134}
135
136// ---------------------------------------------------------------------------
137// Usage
138// ---------------------------------------------------------------------------
139
140/// Token usage counts for a completion (mirrors `org.mockserver.model.Usage`).
141#[derive(Debug, Clone, Default, PartialEq)]
142pub struct Usage {
143    input_tokens: Option<i64>,
144    output_tokens: Option<i64>,
145}
146
147impl Usage {
148    /// Create an empty usage.
149    pub fn new() -> Self {
150        Self::default()
151    }
152
153    /// Set the input-token count (must be `>= 0`).
154    ///
155    /// # Panics
156    /// Panics if `input_tokens < 0`.
157    pub fn with_input_tokens(mut self, input_tokens: i64) -> Self {
158        assert!(input_tokens >= 0, "input_tokens must be >= 0");
159        self.input_tokens = Some(input_tokens);
160        self
161    }
162
163    /// Set the output-token count (must be `>= 0`).
164    ///
165    /// # Panics
166    /// Panics if `output_tokens < 0`.
167    pub fn with_output_tokens(mut self, output_tokens: i64) -> Self {
168        assert!(output_tokens >= 0, "output_tokens must be >= 0");
169        self.output_tokens = Some(output_tokens);
170        self
171    }
172
173    fn to_value(&self) -> Value {
174        let mut map = Map::new();
175        insert_some(&mut map, "inputTokens", self.input_tokens.map(Value::from));
176        insert_some(&mut map, "outputTokens", self.output_tokens.map(Value::from));
177        Value::Object(map)
178    }
179}
180
181/// Factory mirroring `Usage.usage()`.
182pub fn usage() -> Usage {
183    Usage::new()
184}
185
186/// Factory mirroring `Usage.inputTokens(n)`.
187pub fn input_tokens(n: i64) -> Usage {
188    Usage::new().with_input_tokens(n)
189}
190
191/// Factory mirroring `Usage.outputTokens(n)`.
192pub fn output_tokens(n: i64) -> Usage {
193    Usage::new().with_output_tokens(n)
194}
195
196// ---------------------------------------------------------------------------
197// StreamingPhysics
198// ---------------------------------------------------------------------------
199
200/// Controls the timing physics of a streamed (SSE) completion
201/// (mirrors `org.mockserver.model.StreamingPhysics`).
202///
203/// `timeToFirstToken` serialises as a `Delay`: `{ timeUnit, value }`.
204#[derive(Debug, Clone, Default, PartialEq)]
205pub struct StreamingPhysics {
206    time_to_first_token: Option<Value>,
207    tokens_per_second: Option<i64>,
208    jitter: Option<f64>,
209    seed: Option<i64>,
210}
211
212impl StreamingPhysics {
213    /// Create empty streaming physics.
214    pub fn new() -> Self {
215        Self::default()
216    }
217
218    /// Set time-to-first-token as a delay value with the given time unit
219    /// (e.g. `"MILLISECONDS"`).
220    pub fn with_time_to_first_token(mut self, value: i64, time_unit: impl Into<String>) -> Self {
221        self.time_to_first_token = Some(json!({ "timeUnit": time_unit.into(), "value": value }));
222        self
223    }
224
225    /// Set tokens-per-second (must be in `[1, 10000]`).
226    ///
227    /// # Panics
228    /// Panics if `tokens_per_second` is outside `1..=10000`.
229    pub fn with_tokens_per_second(mut self, tokens_per_second: i64) -> Self {
230        assert!(
231            (1..=10000).contains(&tokens_per_second),
232            "tokens_per_second must be between 1 and 10000"
233        );
234        self.tokens_per_second = Some(tokens_per_second);
235        self
236    }
237
238    /// Set the inter-token jitter (must be in `[0.0, 1.0]`).
239    ///
240    /// # Panics
241    /// Panics if `jitter` is outside `0.0..=1.0`.
242    pub fn with_jitter(mut self, jitter: f64) -> Self {
243        assert!(
244            (0.0..=1.0).contains(&jitter),
245            "jitter must be between 0.0 and 1.0"
246        );
247        self.jitter = Some(jitter);
248        self
249    }
250
251    /// Set the deterministic seed.
252    pub fn with_seed(mut self, seed: i64) -> Self {
253        self.seed = Some(seed);
254        self
255    }
256
257    fn to_value(&self) -> Value {
258        let mut map = Map::new();
259        insert_some(&mut map, "timeToFirstToken", self.time_to_first_token.clone());
260        insert_some(&mut map, "tokensPerSecond", self.tokens_per_second.map(Value::from));
261        insert_some(&mut map, "jitter", self.jitter.map(Value::from));
262        insert_some(&mut map, "seed", self.seed.map(Value::from));
263        Value::Object(map)
264    }
265}
266
267/// Factory mirroring `StreamingPhysics.streamingPhysics()`.
268pub fn streaming_physics() -> StreamingPhysics {
269    StreamingPhysics::new()
270}
271
272/// Factory mirroring `StreamingPhysics.tokensPerSecond(n)`.
273pub fn tokens_per_second(n: i64) -> StreamingPhysics {
274    StreamingPhysics::new().with_tokens_per_second(n)
275}
276
277/// Factory mirroring `StreamingPhysics.jitter(j)`.
278pub fn jitter(j: f64) -> StreamingPhysics {
279    StreamingPhysics::new().with_jitter(j)
280}
281
282// ---------------------------------------------------------------------------
283// Completion
284// ---------------------------------------------------------------------------
285
286/// A mocked LLM chat/completion response, provider-agnostic
287/// (mirrors `org.mockserver.model.Completion`).
288///
289/// MockServer re-encodes this into the wire shape of the configured provider
290/// when a request is served.
291#[derive(Debug, Clone, Default, PartialEq)]
292pub struct Completion {
293    text: Option<String>,
294    tool_calls: Option<Vec<ToolUse>>,
295    stop_reason: Option<String>,
296    usage: Option<Usage>,
297    streaming: Option<bool>,
298    streaming_physics: Option<StreamingPhysics>,
299    output_schema: Option<String>,
300    model: Option<String>,
301}
302
303impl Completion {
304    /// Create an empty completion.
305    pub fn new() -> Self {
306        Self::default()
307    }
308
309    /// Set the response text.
310    pub fn with_text(mut self, text: impl Into<String>) -> Self {
311        self.text = Some(text.into());
312        self
313    }
314
315    /// Append a single tool call.
316    pub fn with_tool_call(mut self, tool_call: ToolUse) -> Self {
317        self.tool_calls.get_or_insert_with(Vec::new).push(tool_call);
318        self
319    }
320
321    /// Replace all tool calls.
322    pub fn with_tool_calls(mut self, tool_calls: Vec<ToolUse>) -> Self {
323        self.tool_calls = Some(tool_calls);
324        self
325    }
326
327    /// Set the stop reason.
328    pub fn with_stop_reason(mut self, stop_reason: impl Into<String>) -> Self {
329        self.stop_reason = Some(stop_reason.into());
330        self
331    }
332
333    /// Set token usage.
334    pub fn with_usage(mut self, usage: Usage) -> Self {
335        self.usage = Some(usage);
336        self
337    }
338
339    /// Set the streaming flag explicitly.
340    pub fn with_streaming(mut self, streaming: bool) -> Self {
341        self.streaming = Some(streaming);
342        self
343    }
344
345    /// Enable streaming (mirrors Java `completion().streaming()`).
346    pub fn streaming(self) -> Self {
347        self.with_streaming(true)
348    }
349
350    /// Set streaming physics. Does NOT touch the `streaming` flag — enable it
351    /// explicitly via [`with_streaming`](Self::with_streaming) / [`streaming`](Self::streaming).
352    pub fn with_streaming_physics(mut self, physics: StreamingPhysics) -> Self {
353        self.streaming_physics = Some(physics);
354        self
355    }
356
357    /// Set the structured-output JSON schema (a JSON string, matching Java).
358    pub fn with_output_schema(mut self, output_schema: impl Into<String>) -> Self {
359        self.output_schema = Some(output_schema.into());
360        self
361    }
362
363    /// Set the model name carried by the completion.
364    pub fn with_model(mut self, model: impl Into<String>) -> Self {
365        self.model = Some(model.into());
366        self
367    }
368
369    fn to_value(&self) -> Value {
370        let mut map = Map::new();
371        insert_some(&mut map, "text", self.text.clone().map(Value::from));
372        insert_some(
373            &mut map,
374            "toolCalls",
375            self.tool_calls
376                .as_ref()
377                .map(|calls| Value::Array(calls.iter().map(ToolUse::to_value).collect())),
378        );
379        insert_some(&mut map, "stopReason", self.stop_reason.clone().map(Value::from));
380        insert_some(&mut map, "usage", self.usage.as_ref().map(Usage::to_value));
381        insert_some(&mut map, "streaming", self.streaming.map(Value::from));
382        insert_some(
383            &mut map,
384            "streamingPhysics",
385            self.streaming_physics.as_ref().map(StreamingPhysics::to_value),
386        );
387        insert_some(&mut map, "outputSchema", self.output_schema.clone().map(Value::from));
388        insert_some(&mut map, "model", self.model.clone().map(Value::from));
389        Value::Object(map)
390    }
391}
392
393/// Factory mirroring `Completion.completion()`.
394pub fn completion() -> Completion {
395    Completion::new()
396}
397
398// ---------------------------------------------------------------------------
399// EmbeddingResponse
400// ---------------------------------------------------------------------------
401
402/// A mocked embedding response — vector shape and determinism
403/// (mirrors `org.mockserver.model.EmbeddingResponse`).
404#[derive(Debug, Clone, Default, PartialEq)]
405pub struct EmbeddingResponse {
406    dimensions: Option<i64>,
407    deterministic_from_input: Option<bool>,
408    seed: Option<i64>,
409}
410
411impl EmbeddingResponse {
412    /// Create an empty embedding response.
413    pub fn new() -> Self {
414        Self::default()
415    }
416
417    /// Set the embedding vector dimensionality.
418    pub fn with_dimensions(mut self, dimensions: i64) -> Self {
419        self.dimensions = Some(dimensions);
420        self
421    }
422
423    /// Whether the embedding should be deterministically derived from input.
424    pub fn with_deterministic_from_input(mut self, deterministic: bool) -> Self {
425        self.deterministic_from_input = Some(deterministic);
426        self
427    }
428
429    /// Set the deterministic seed.
430    pub fn with_seed(mut self, seed: i64) -> Self {
431        self.seed = Some(seed);
432        self
433    }
434
435    fn to_value(&self) -> Value {
436        let mut map = Map::new();
437        insert_some(&mut map, "dimensions", self.dimensions.map(Value::from));
438        insert_some(
439            &mut map,
440            "deterministicFromInput",
441            self.deterministic_from_input.map(Value::from),
442        );
443        insert_some(&mut map, "seed", self.seed.map(Value::from));
444        Value::Object(map)
445    }
446}
447
448/// Factory mirroring `EmbeddingResponse.embedding()`.
449pub fn embedding() -> EmbeddingResponse {
450    EmbeddingResponse::new()
451}
452
453// ---------------------------------------------------------------------------
454// A completion-or-embedding response payload
455// ---------------------------------------------------------------------------
456
457/// Either a [`Completion`] or an [`EmbeddingResponse`], accepted by
458/// `responding_with` builder methods.
459#[derive(Debug, Clone, PartialEq)]
460pub enum LlmResponseBody {
461    Completion(Completion),
462    Embedding(EmbeddingResponse),
463}
464
465impl From<Completion> for LlmResponseBody {
466    fn from(c: Completion) -> Self {
467        LlmResponseBody::Completion(c)
468    }
469}
470
471impl From<EmbeddingResponse> for LlmResponseBody {
472    fn from(e: EmbeddingResponse) -> Self {
473        LlmResponseBody::Embedding(e)
474    }
475}
476
477/// Build the `httpLlmResponse` action object.
478fn build_llm_response(
479    provider: Option<&str>,
480    model: Option<&str>,
481    completion: Option<&Completion>,
482    embedding: Option<&EmbeddingResponse>,
483    conversation_predicates: Option<Value>,
484    chaos: Option<Value>,
485) -> Value {
486    let mut map = Map::new();
487    insert_some(&mut map, "provider", provider.map(Value::from));
488    insert_some(&mut map, "model", model.map(Value::from));
489    insert_some(&mut map, "completion", completion.map(Completion::to_value));
490    insert_some(&mut map, "embedding", embedding.map(EmbeddingResponse::to_value));
491    insert_some(&mut map, "conversationPredicates", conversation_predicates);
492    insert_some(&mut map, "chaos", chaos);
493    Value::Object(map)
494}
495
496// ---------------------------------------------------------------------------
497// IsolationSource
498// ---------------------------------------------------------------------------
499
500const ISOLATION_MARKER: &str = "__iso=";
501
502/// Where to read the per-session isolation key from an inbound request
503/// (mirrors `org.mockserver.llm.IsolationSource`). Encodes as `"<kind>:<name>"`.
504#[derive(Debug, Clone, PartialEq)]
505pub struct IsolationSource {
506    kind: String,
507    name: String,
508}
509
510impl IsolationSource {
511    /// `kind` is one of `"header"`, `"query_parameter"`, `"cookie"`.
512    fn new(kind: &str, name: impl Into<String>) -> Self {
513        Self {
514            kind: kind.to_string(),
515            name: name.into(),
516        }
517    }
518
519    /// Encode as `"<kind>:<name>"`.
520    pub fn encode(&self) -> String {
521        format!("{}:{}", self.kind, self.name)
522    }
523}
524
525/// Isolate conversations by a request header.
526pub fn header(name: impl Into<String>) -> IsolationSource {
527    IsolationSource::new("header", name)
528}
529
530/// Isolate conversations by a query parameter.
531pub fn query_parameter(name: impl Into<String>) -> IsolationSource {
532    IsolationSource::new("query_parameter", name)
533}
534
535/// Isolate conversations by a cookie.
536pub fn cookie(name: impl Into<String>) -> IsolationSource {
537    IsolationSource::new("cookie", name)
538}
539
540// ---------------------------------------------------------------------------
541// LlmMockBuilder — a single completion or embedding mock
542// ---------------------------------------------------------------------------
543
544/// Fluent builder for a single LLM mock expectation
545/// (mirrors `org.mockserver.client.LlmMockBuilder`).
546#[derive(Debug, Clone)]
547pub struct LlmMockBuilder {
548    path: String,
549    provider: Option<String>,
550    model: Option<String>,
551    completion: Option<Completion>,
552    embedding: Option<EmbeddingResponse>,
553}
554
555impl LlmMockBuilder {
556    /// Set the LLM provider (use a [`Provider`] constant).
557    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
558        self.provider = Some(provider.into());
559        self
560    }
561
562    /// Set the model name.
563    pub fn with_model(mut self, model: impl Into<String>) -> Self {
564        self.model = Some(model.into());
565        self
566    }
567
568    /// Set the response body. Passing a [`Completion`] clears any embedding and
569    /// vice-versa.
570    pub fn responding_with(mut self, response: impl Into<LlmResponseBody>) -> Self {
571        match response.into() {
572            LlmResponseBody::Completion(c) => {
573                self.completion = Some(c);
574                self.embedding = None;
575            }
576            LlmResponseBody::Embedding(e) => {
577                self.embedding = Some(e);
578                self.completion = None;
579            }
580        }
581        self
582    }
583
584    /// Build the single expectation JSON.
585    pub fn build(&self) -> Value {
586        let mut map = Map::new();
587        map.insert("httpRequest".to_string(), post_matcher(&self.path));
588        map.insert(
589            "httpLlmResponse".to_string(),
590            build_llm_response(
591                self.provider.as_deref(),
592                self.model.as_deref(),
593                self.completion.as_ref(),
594                self.embedding.as_ref(),
595                None,
596                None,
597            ),
598        );
599        Value::Object(map)
600    }
601
602    /// Build and register the expectation on the given client.
603    pub fn apply_to(&self, client: &MockServerClient) -> Result<Value> {
604        client.upsert_raw(self.build())
605    }
606}
607
608/// Entry point mirroring `LlmMockBuilder.llmMock(path)`.
609pub fn llm_mock(path: impl Into<String>) -> LlmMockBuilder {
610    LlmMockBuilder {
611        path: path.into(),
612        provider: None,
613        model: None,
614        completion: None,
615        embedding: None,
616    }
617}
618
619// ---------------------------------------------------------------------------
620// LlmConversationBuilder + TurnBuilder
621// ---------------------------------------------------------------------------
622
623const SCENARIO_PREFIX: &str = "__llm_conv_";
624const DONE_STATE: &str = "__done";
625
626/// Sub-builder configuring one turn of a conversation mock
627/// (mirrors `org.mockserver.client.TurnBuilder`).
628#[derive(Debug, Clone, Default)]
629pub struct TurnBuilder {
630    turn_index: Option<i64>,
631    latest_message_contains: Option<String>,
632    latest_message_matches: Option<String>,
633    latest_message_role: Option<String>,
634    contains_tool_result_for: Option<String>,
635    semantic_match_against: Option<String>,
636    normalization: Option<Value>,
637    chaos: Option<Value>,
638    completion: Option<Completion>,
639}
640
641impl TurnBuilder {
642    /// Match when the conversation is at the given (0-based) turn index.
643    pub fn when_turn_index(mut self, n: i64) -> Self {
644        self.turn_index = Some(n);
645        self
646    }
647
648    /// Match when the latest message contains the given substring.
649    pub fn when_latest_message_contains(mut self, text: impl Into<String>) -> Self {
650        self.latest_message_contains = Some(text.into());
651        self
652    }
653
654    /// Match when the latest message matches the given regex.
655    pub fn when_latest_message_matches(mut self, regex: impl Into<String>) -> Self {
656        self.latest_message_matches = Some(regex.into());
657        self
658    }
659
660    /// Match when the latest message has the given role (use a [`Role`] constant).
661    pub fn when_latest_message_role(mut self, role: impl Into<String>) -> Self {
662        self.latest_message_role = Some(role.into());
663        self
664    }
665
666    /// Match when the latest message contains a tool result for the given tool.
667    pub fn when_contains_tool_result_for(mut self, tool_name: impl Into<String>) -> Self {
668        self.contains_tool_result_for = Some(tool_name.into());
669        self
670    }
671
672    /// Match semantically against the given expected meaning.
673    pub fn when_semantic_match(mut self, expected_meaning: impl Into<String>) -> Self {
674        self.semantic_match_against = Some(expected_meaning.into());
675        self
676    }
677
678    /// Set normalization options (a JSON object). Not counted as a predicate.
679    pub fn with_normalization(mut self, normalization: Value) -> Self {
680        self.normalization = Some(normalization);
681        self
682    }
683
684    /// Set a chaos profile (a JSON object) for this turn.
685    pub fn with_chaos(mut self, chaos: Value) -> Self {
686        self.chaos = Some(chaos);
687        self
688    }
689
690    /// Set the completion this turn responds with.
691    pub fn responding_with(mut self, completion: Completion) -> Self {
692        self.completion = Some(completion);
693        self
694    }
695
696    /// Build the `conversationPredicates` object (always, even when empty).
697    fn predicates_value(&self) -> Value {
698        let mut map = Map::new();
699        insert_some(&mut map, "turnIndex", self.turn_index.map(Value::from));
700        insert_some(
701            &mut map,
702            "latestMessageContains",
703            self.latest_message_contains.clone().map(Value::from),
704        );
705        insert_some(
706            &mut map,
707            "latestMessageMatches",
708            self.latest_message_matches.clone().map(Value::from),
709        );
710        insert_some(
711            &mut map,
712            "latestMessageRole",
713            self.latest_message_role.clone().map(Value::from),
714        );
715        insert_some(
716            &mut map,
717            "containsToolResultFor",
718            self.contains_tool_result_for.clone().map(Value::from),
719        );
720        insert_some(
721            &mut map,
722            "semanticMatchAgainst",
723            self.semantic_match_against.clone().map(Value::from),
724        );
725        insert_some(&mut map, "normalization", self.normalization.clone());
726        Value::Object(map)
727    }
728
729    /// True if at least one predicate (excluding `normalization`) is set.
730    fn has_any_predicate(&self) -> bool {
731        self.turn_index.is_some()
732            || self.latest_message_contains.is_some()
733            || self.latest_message_matches.is_some()
734            || self.latest_message_role.is_some()
735            || self.contains_tool_result_for.is_some()
736            || self.semantic_match_against.is_some()
737    }
738}
739
740/// Builder for multi-turn LLM conversation mocks with scenario advancement
741/// (mirrors `org.mockserver.client.LlmConversationBuilder`).
742#[derive(Debug, Clone, Default)]
743pub struct LlmConversationBuilder {
744    path: Option<String>,
745    provider: Option<String>,
746    model: Option<String>,
747    isolation_source: Option<IsolationSource>,
748    turns: Vec<TurnBuilder>,
749}
750
751impl LlmConversationBuilder {
752    /// Set the request path.
753    pub fn with_path(mut self, path: impl Into<String>) -> Self {
754        self.path = Some(path.into());
755        self
756    }
757
758    /// Set the LLM provider (use a [`Provider`] constant).
759    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
760        self.provider = Some(provider.into());
761        self
762    }
763
764    /// Set the model name.
765    pub fn with_model(mut self, model: impl Into<String>) -> Self {
766        self.model = Some(model.into());
767        self
768    }
769
770    /// Isolate per-session conversations by the given source.
771    pub fn isolate_by(mut self, source: IsolationSource) -> Self {
772        self.isolation_source = Some(source);
773        self
774    }
775
776    /// Add a turn to the conversation. Build the turn fluently and pass it back.
777    pub fn turn(mut self, turn: TurnBuilder) -> Self {
778        self.turns.push(turn);
779        self
780    }
781
782    /// Build the list of conversation expectations.
783    ///
784    /// The auto-generated scenario id makes the output non-deterministic between
785    /// calls; tests should assert structure, not the exact uuid.
786    pub fn build(&self) -> Vec<Value> {
787        self.build_with_id(&new_uuid())
788    }
789
790    /// Build with an explicit conversation id (used by tests for determinism).
791    pub fn build_with_id(&self, conversation_uuid: &str) -> Vec<Value> {
792        assert!(!self.turns.is_empty(), "At least one turn must be defined");
793        let path = self.path.as_deref().expect("Path must be set");
794        let provider = self.provider.as_deref().expect("Provider must be set");
795
796        let mut scenario_name = format!("{SCENARIO_PREFIX}{conversation_uuid}");
797        if let Some(source) = &self.isolation_source {
798            scenario_name = format!("{scenario_name}{ISOLATION_MARKER}{}", source.encode());
799        }
800
801        let n = self.turns.len();
802        let mut expectations = Vec::with_capacity(n);
803        for (i, turn) in self.turns.iter().enumerate() {
804            let scenario_state = if i == 0 {
805                "Started".to_string()
806            } else {
807                format!("turn_{i}")
808            };
809            let new_scenario_state = if i < n - 1 {
810                format!("turn_{}", i + 1)
811            } else {
812                DONE_STATE.to_string()
813            };
814
815            let predicates = if turn.has_any_predicate() {
816                Some(turn.predicates_value())
817            } else {
818                None
819            };
820
821            let llm_response = build_llm_response(
822                Some(provider),
823                self.model.as_deref(),
824                turn.completion.as_ref(),
825                None,
826                predicates,
827                turn.chaos.clone(),
828            );
829
830            let mut map = Map::new();
831            map.insert("httpRequest".to_string(), post_matcher(path));
832            map.insert("scenarioName".to_string(), Value::from(scenario_name.clone()));
833            map.insert("scenarioState".to_string(), Value::from(scenario_state));
834            map.insert("newScenarioState".to_string(), Value::from(new_scenario_state));
835            map.insert("httpLlmResponse".to_string(), llm_response);
836            expectations.push(Value::Object(map));
837        }
838        expectations
839    }
840
841    /// Build and register the expectations on the given client.
842    pub fn apply_to(&self, client: &MockServerClient) -> Result<Value> {
843        client.upsert_raw(Value::Array(self.build()))
844    }
845}
846
847/// Entry point mirroring `LlmConversationBuilder.conversation()`.
848pub fn conversation() -> LlmConversationBuilder {
849    LlmConversationBuilder::default()
850}
851
852/// Entry point mirroring `TurnBuilder` construction (used with
853/// [`LlmConversationBuilder::turn`]).
854pub fn turn() -> TurnBuilder {
855    TurnBuilder::default()
856}
857
858// ---------------------------------------------------------------------------
859// LlmFailoverBuilder — N failures then a success completion
860// ---------------------------------------------------------------------------
861
862/// The default error body for the given status code (mirrors the Java/Node/Python
863/// `defaultErrorBody`). Exposed for parity/testing.
864pub fn default_error_body(status_code: u16) -> String {
865    let (type_, message) = match status_code {
866        429 => (
867            "rate_limit_error",
868            "Rate limit exceeded. Please retry after a brief wait.".to_string(),
869        ),
870        500 => (
871            "internal_server_error",
872            "An internal error occurred. Please retry your request.".to_string(),
873        ),
874        502 => (
875            "bad_gateway",
876            "Bad gateway. The upstream server returned an invalid response.".to_string(),
877        ),
878        503 => (
879            "service_unavailable",
880            "The service is temporarily overloaded. Please retry later.".to_string(),
881        ),
882        _ => ("error", format!("Request failed with status {status_code}")),
883    };
884    // Compact, fixed-order JSON exactly matching the other clients.
885    format!(r#"{{"error":{{"type":"{type_}","message":"{message}"}}}}"#)
886}
887
888#[derive(Debug, Clone, PartialEq)]
889struct FailureSpec {
890    status_code: u16,
891    error_body: Option<String>,
892}
893
894/// Builder for provider failover/retry scenarios
895/// (mirrors `org.mockserver.client.LlmFailoverBuilder`).
896///
897/// Produces an ordered list of expectations: failure expectations (limited
898/// `times`) first, then a single success expectation with unlimited `times`.
899/// Consecutive identical failures are coalesced into one expectation with
900/// `times.remainingTimes = count`.
901#[derive(Debug, Clone, Default)]
902pub struct LlmFailoverBuilder {
903    path: Option<String>,
904    provider: Option<String>,
905    model: Option<String>,
906    failures: Vec<FailureSpec>,
907    success_completion: Option<Completion>,
908}
909
910impl LlmFailoverBuilder {
911    /// Set the request path.
912    pub fn with_path(mut self, path: impl Into<String>) -> Self {
913        self.path = Some(path.into());
914        self
915    }
916
917    /// Set the LLM provider (use a [`Provider`] constant).
918    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
919        self.provider = Some(provider.into());
920        self
921    }
922
923    /// Set the model name.
924    pub fn with_model(mut self, model: impl Into<String>) -> Self {
925        self.model = Some(model.into());
926        self
927    }
928
929    /// Add a single failure with the default error body for `status_code`.
930    ///
931    /// # Panics
932    /// Panics if `status_code` is outside `100..=599`.
933    pub fn fail_with(self, status_code: u16) -> Self {
934        self.fail_with_count(status_code, 1)
935    }
936
937    /// Add a single failure with a custom error body.
938    ///
939    /// # Panics
940    /// Panics if `status_code` is outside `100..=599`.
941    pub fn fail_with_body(mut self, status_code: u16, error_body: impl Into<String>) -> Self {
942        validate_status_code(status_code);
943        self.failures.push(FailureSpec {
944            status_code,
945            error_body: Some(error_body.into()),
946        });
947        self
948    }
949
950    /// Add `count` consecutive failures with the default error body.
951    ///
952    /// # Panics
953    /// Panics if `status_code` is outside `100..=599` or `count < 1`.
954    pub fn fail_with_count(mut self, status_code: u16, count: u32) -> Self {
955        validate_status_code(status_code);
956        assert!(count >= 1, "count must be >= 1");
957        for _ in 0..count {
958            self.failures.push(FailureSpec {
959                status_code,
960                error_body: None,
961            });
962        }
963        self
964    }
965
966    /// Set the success completion served after all failures are consumed.
967    pub fn then_respond_with(mut self, completion: Completion) -> Self {
968        self.success_completion = Some(completion);
969        self
970    }
971
972    /// Number of failure attempts configured.
973    pub fn failure_count(&self) -> usize {
974        self.failures.len()
975    }
976
977    fn coalesce(&self) -> Vec<(u16, Option<String>, u32)> {
978        let mut result: Vec<(u16, Option<String>, u32)> = Vec::new();
979        for spec in &self.failures {
980            if let Some(last) = result.last_mut() {
981                if last.0 == spec.status_code && last.1 == spec.error_body {
982                    last.2 += 1;
983                    continue;
984                }
985            }
986            result.push((spec.status_code, spec.error_body.clone(), 1));
987        }
988        result
989    }
990
991    /// Build the list of failover expectations.
992    pub fn build(&self) -> Vec<Value> {
993        let path = self.path.as_deref().expect("Path must be set");
994        let provider = self.provider.as_deref().expect("Provider must be set");
995        assert!(!self.failures.is_empty(), "At least one failure must be defined");
996        let success = self
997            .success_completion
998            .as_ref()
999            .expect("Success completion must be set via then_respond_with()");
1000
1001        let mut expectations = Vec::new();
1002        for (status_code, error_body, count) in self.coalesce() {
1003            let body = error_body.unwrap_or_else(|| default_error_body(status_code));
1004            expectations.push(json!({
1005                "httpRequest": post_matcher(path),
1006                "times": { "remainingTimes": count, "unlimited": false },
1007                "timeToLive": { "unlimited": true },
1008                "httpResponse": {
1009                    "statusCode": status_code,
1010                    "headers": [{ "name": "Content-Type", "values": ["application/json"] }],
1011                    "body": body
1012                }
1013            }));
1014        }
1015
1016        let mut success_map = Map::new();
1017        success_map.insert("httpRequest".to_string(), post_matcher(path));
1018        success_map.insert(
1019            "times".to_string(),
1020            json!({ "remainingTimes": 0, "unlimited": true }),
1021        );
1022        success_map.insert("timeToLive".to_string(), json!({ "unlimited": true }));
1023        success_map.insert(
1024            "httpLlmResponse".to_string(),
1025            build_llm_response(
1026                Some(provider),
1027                self.model.as_deref(),
1028                Some(success),
1029                None,
1030                None,
1031                None,
1032            ),
1033        );
1034        expectations.push(Value::Object(success_map));
1035        expectations
1036    }
1037
1038    /// Build and register the expectations on the given client.
1039    pub fn apply_to(&self, client: &MockServerClient) -> Result<Value> {
1040        client.upsert_raw(Value::Array(self.build()))
1041    }
1042}
1043
1044/// Entry point mirroring `LlmFailoverBuilder.llmFailover()`.
1045pub fn llm_failover() -> LlmFailoverBuilder {
1046    LlmFailoverBuilder::default()
1047}
1048
1049fn validate_status_code(status_code: u16) {
1050    assert!(
1051        (100..=599).contains(&status_code),
1052        "statusCode must be between 100 and 599"
1053    );
1054}
1055
1056// ---------------------------------------------------------------------------
1057// Minimal RFC4122 v4-ish UUID (dependency-free) for scenario names.
1058// ---------------------------------------------------------------------------
1059
1060fn new_uuid() -> String {
1061    // Derive 16 pseudo-random bytes from a few entropy sources. This does not
1062    // need to be cryptographically strong — it only needs to be unique enough
1063    // to keep concurrent conversation scenarios from colliding.
1064    use std::collections::hash_map::RandomState;
1065    use std::hash::{BuildHasher, Hasher};
1066    use std::sync::atomic::{AtomicU64, Ordering};
1067    use std::time::{SystemTime, UNIX_EPOCH};
1068
1069    // A process-wide monotonic counter guarantees that two calls never share the
1070    // same input even if every other entropy source (clock, stack address, the
1071    // per-call RandomState seeds) happened to collide.
1072    static COUNTER: AtomicU64 = AtomicU64::new(0);
1073    let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
1074
1075    let nanos = SystemTime::now()
1076        .duration_since(UNIX_EPOCH)
1077        .map(|d| d.as_nanos())
1078        .unwrap_or(0);
1079    let stack_marker = &nanos as *const _ as usize;
1080
1081    let mut hasher = RandomState::new().build_hasher();
1082    hasher.write_u128(nanos);
1083    hasher.write_usize(stack_marker);
1084    hasher.write_u64(seq);
1085    let h1 = hasher.finish();
1086    let mut hasher2 = RandomState::new().build_hasher();
1087    hasher2.write_u64(h1);
1088    hasher2.write_u128(nanos.rotate_left(17));
1089    hasher2.write_u64(seq.rotate_left(32));
1090    let h2 = hasher2.finish();
1091
1092    let mut bytes = [0u8; 16];
1093    bytes[..8].copy_from_slice(&h1.to_le_bytes());
1094    bytes[8..].copy_from_slice(&h2.to_le_bytes());
1095    // version 4 + variant bits
1096    bytes[6] = (bytes[6] & 0x0f) | 0x40;
1097    bytes[8] = (bytes[8] & 0x3f) | 0x80;
1098
1099    let h = |b: &[u8]| b.iter().map(|x| format!("{x:02x}")).collect::<String>();
1100    format!(
1101        "{}-{}-{}-{}-{}",
1102        h(&bytes[0..4]),
1103        h(&bytes[4..6]),
1104        h(&bytes[6..8]),
1105        h(&bytes[8..10]),
1106        h(&bytes[10..16]),
1107    )
1108}
mockserver_client/llm.rs

mockserver_client/
llm.rs
