diff --git a/model_gateway/src/routers/grpc/context.rs b/model_gateway/src/routers/grpc/context.rs index 9d6a8438d9..36a816fa88 100644 --- a/model_gateway/src/routers/grpc/context.rs +++ b/model_gateway/src/routers/grpc/context.rs @@ -490,6 +490,10 @@ pub(crate) struct ResponseState { /// response_processing runs. pub skip_special_tokens: Option, + /// Response-phase view of the request, set by request building so + /// response processing never reads the payload after dispatch. + pub request_view: Option, + /// Execution result (streams from workers) pub execution_result: Option, @@ -637,18 +641,6 @@ impl RequestContext { ) } - /// Get chat request (panics if not chat) - #[expect( - clippy::panic, - reason = "typed accessor: caller guarantees variant via RequestType construction" - )] - pub fn chat_request(&self) -> &ChatCompletionRequest { - match &self.input.request_type { - RequestType::Chat(req) => req.as_ref(), - _ => panic!("Expected chat request"), - } - } - /// Get Arc clone of chat request (panics if not chat) #[expect( clippy::panic, @@ -661,18 +653,6 @@ impl RequestContext { } } - /// Get generate request (panics if not generate) - #[expect( - clippy::panic, - reason = "typed accessor: caller guarantees variant via RequestType construction" - )] - pub fn generate_request(&self) -> &GenerateRequest { - match &self.input.request_type { - RequestType::Generate(req) => req.as_ref(), - _ => panic!("Expected generate request"), - } - } - /// Get Arc clone of generate request (panics if not generate) #[expect( clippy::panic, @@ -685,22 +665,6 @@ impl RequestContext { } } - /// Get completion request (panics if not completion) - #[expect( - dead_code, - reason = "ref accessor provided for API completeness alongside Arc accessor" - )] - #[expect( - clippy::panic, - reason = "typed accessor: caller guarantees variant via RequestType construction" - )] - pub fn completion_request(&self) -> &CompletionRequest { - match &self.input.request_type { - RequestType::Completion(req) => req.as_ref(), - _ => panic!("Expected completion request"), - } - } - /// Get Arc clone of completion request (panics if not completion) #[expect( clippy::panic, @@ -725,22 +689,6 @@ impl RequestContext { } } - /// Get messages request (panics if not messages) - #[expect( - dead_code, - reason = "scaffolding for Messages API pipeline, wired in follow-up PR" - )] - #[expect( - clippy::panic, - reason = "typed accessor: caller guarantees variant via RequestType construction" - )] - pub fn messages_request(&self) -> &CreateMessageRequest { - match &self.input.request_type { - RequestType::Messages(req) => req.as_ref(), - _ => panic!("Expected messages request"), - } - } - /// Get Arc clone of messages request (panics if not messages) #[expect( clippy::panic, diff --git a/model_gateway/src/routers/grpc/pipeline.rs b/model_gateway/src/routers/grpc/pipeline.rs index 1a7c63da26..842b3fc5ff 100644 --- a/model_gateway/src/routers/grpc/pipeline.rs +++ b/model_gateway/src/routers/grpc/pipeline.rs @@ -1588,7 +1588,7 @@ mod alias_pipeline_tests { ); assert_eq!(ctx.input.model_id, CANONICAL_MODEL); - assert_eq!(ctx.generate_request().model, CANONICAL_MODEL); + assert_eq!(ctx.generate_request_arc().model, CANONICAL_MODEL); for stage in pipeline.stages.iter() { assert!(stage.execute(&mut ctx).await.unwrap().is_none()); @@ -1598,7 +1598,7 @@ mod alias_pipeline_tests { } assert_eq!(ctx.input.model_id, CANONICAL_MODEL); - assert_eq!(ctx.generate_request().model, CANONICAL_MODEL); + assert_eq!(ctx.generate_request_arc().model, CANONICAL_MODEL); assert!(ctx.state.tokenizer.is_some()); match ctx.state.workers.as_ref().unwrap() { WorkerSelection::Disaggregated { @@ -1612,6 +1612,377 @@ mod alias_pipeline_tests { } } +#[cfg(test)] +mod request_release_tests { + use std::{ + pin::Pin, + sync::{ + atomic::{AtomicBool, Ordering}, + Weak, + }, + time::Duration, + }; + + use futures::Stream; + use llm_tokenizer::{traits::Tokenizer, MockTokenizer, TokenizerRegistry}; + use openai_protocol::{ + completion::CompletionRequest, model_card::ModelCard, worker::HealthCheckConfig, + }; + use portpicker::pick_unused_port; + use smg_grpc_client::{common_proto as common, tokenspeed_proto as ts}; + use tokio_stream::wrappers::ReceiverStream; + use tonic::{transport::Server, Request as TonicRequest, Response as TonicResponse, Status}; + use ts::token_speed_scheduler_server::{TokenSpeedScheduler, TokenSpeedSchedulerServer}; + + use super::*; + use crate::{ + config::types::PolicyConfig, + worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, WorkerType}, + }; + + const MODEL: &str = "request-release-test-model"; + + type GenStream = Pin> + Send>>; + type KvEventStream = Pin> + Send>>; + type TokenizerStream = + Pin> + Send>>; + + /// TokenSpeed stub gated on the parsed request's drop probe: it withholds + /// its tokens until the probe reaches zero strong references or a + /// deadline passes, recording the outcome in `released`. An ungated stub + /// (no probe) answers immediately -- used for the PD prefill leg. + #[derive(Clone, Default)] + struct GatedScheduler { + probe: Option>, + released: Arc, + } + + impl GatedScheduler { + async fn await_probe(probe: &Weak, released: &AtomicBool) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while probe.strong_count() > 0 && tokio::time::Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(2)).await; + } + released.store(probe.strong_count() == 0, Ordering::SeqCst); + } + } + + fn generate_frames(request_id: &str) -> Vec> { + use ts::generate_response::Response as GenResp; + vec![ + Ok(ts::GenerateResponse { + request_id: request_id.to_string(), + response: Some(GenResp::Chunk(ts::GenerateStreamChunk { + token_ids: vec![100], + prompt_tokens: 2, + completion_tokens: 1, + cached_tokens: 0, + output_logprobs: None, + index: 0, + })), + }), + Ok(ts::GenerateResponse { + request_id: request_id.to_string(), + response: Some(GenResp::Complete(ts::GenerateComplete { + output_ids: vec![100], + finish_reason: "stop".to_string(), + prompt_tokens: 2, + completion_tokens: 1, + cached_tokens: 0, + output_logprobs: None, + matched_stop: None, + index: 0, + })), + }), + ] + } + + #[tonic::async_trait] + impl TokenSpeedScheduler for GatedScheduler { + type GenerateStream = GenStream; + type SubscribeKvEventsStream = KvEventStream; + type GetTokenizerStream = TokenizerStream; + + #[expect( + clippy::disallowed_methods, + reason = "test stub; the gate task ends at its deadline" + )] + async fn generate( + &self, + request: TonicRequest, + ) -> Result, Status> { + let request_id = request.into_inner().request_id; + let (tx, rx) = tokio::sync::mpsc::channel(8); + let probe = self.probe.clone(); + let released = Arc::clone(&self.released); + tokio::spawn(async move { + if let Some(probe) = probe { + Self::await_probe(&probe, &released).await; + } + for frame in generate_frames(&request_id) { + if tx.send(frame).await.is_err() { + return; + } + } + }); + Ok(TonicResponse::new(Box::pin(ReceiverStream::new(rx)))) + } + + async fn health_check( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Ok(TonicResponse::new(ts::HealthCheckResponse { + healthy: true, + message: "ok".to_string(), + })) + } + + async fn abort( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Ok(TonicResponse::new(ts::AbortResponse { + success: true, + message: String::new(), + })) + } + + async fn get_model_info( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Ok(TonicResponse::new(ts::GetModelInfoResponse { + model_path: MODEL.to_string(), + tokenizer_path: MODEL.to_string(), + served_model_name: MODEL.to_string(), + model_type: "mock".to_string(), + architectures: vec!["MockForCausalLM".to_string()], + max_context_length: 32768, + max_req_input_len: 32768, + vocab_size: 32000, + eos_token_ids: vec![2], + pad_token_id: 0, + bos_token_id: 1, + weight_version: "mock".to_string(), + default_sampling_params_json: String::new(), + supports_vision: false, + ..Default::default() + })) + } + + async fn get_server_info( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Ok(TonicResponse::new(ts::GetServerInfoResponse { + max_total_num_tokens: 1_000_000, + tokenspeed_version: "mock".to_string(), + ..Default::default() + })) + } + + async fn get_loads( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Err(Status::unimplemented("release-test stub")) + } + + async fn subscribe_kv_events( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Err(Status::unimplemented("release-test stub")) + } + + async fn flush_cache( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Err(Status::unimplemented("release-test stub")) + } + + async fn start_profile( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Err(Status::unimplemented("release-test stub")) + } + + async fn stop_profile( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Err(Status::unimplemented("release-test stub")) + } + + async fn get_tokenizer( + &self, + _request: TonicRequest, + ) -> Result, Status> { + Err(Status::unimplemented("release-test stub")) + } + } + + #[expect( + clippy::disallowed_methods, + reason = "test helper; the stub server lives for the test process" + )] + async fn spawn_stub(scheduler: GatedScheduler) -> u16 { + let port = pick_unused_port().expect("no free port for release-test stub"); + let addr = format!("127.0.0.1:{port}").parse().expect("stub addr"); + tokio::spawn(async move { + Server::builder() + .add_service(TokenSpeedSchedulerServer::new(scheduler)) + .serve(addr) + .await + .expect("release-test stub server"); + }); + for _ in 0..100 { + if tokio::net::TcpStream::connect(("127.0.0.1", port)) + .await + .is_ok() + { + return port; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("release-test stub on port {port} never came up"); + } + + fn register_worker(registry: &WorkerRegistry, port: u16, worker_type: WorkerType) { + let worker = BasicWorkerBuilder::new(format!("grpc://127.0.0.1:{port}")) + .worker_type(worker_type) + .connection_mode(ConnectionMode::Grpc) + .runtime_type(RuntimeType::TokenSpeed) + .model(ModelCard::new(MODEL)) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + registry + .register(Arc::new(worker)) + .expect("register release-test worker"); + } + + async fn components(worker_registry: Arc) -> Arc { + let tokenizer_registry = Arc::new(TokenizerRegistry::new()); + let tokenizer = Arc::new(MockTokenizer::new()) as Arc; + tokenizer_registry + .load( + "tokenizer-id", + MODEL, + "test", + || async move { Ok(tokenizer) }, + ) + .await + .expect("load mock tokenizer"); + Arc::new(SharedComponents { + tokenizer_registry, + worker_registry, + tool_parser_factory: ToolParserFactory::default(), + reasoning_parser_factory: ReasoningParserFactory::default(), + parser_resolver: utils::ParserResolver::disabled(), + multimodal: None, + }) + } + + fn completion_request(stream: bool) -> Arc { + Arc::new( + serde_json::from_value(serde_json::json!({ + "model": MODEL, + "prompt": "Hello world", + "stream": stream, + })) + .expect("completion request"), + ) + } + + fn completion_pipeline(worker_registry: &Arc, mode: Mode) -> RequestPipeline { + let deps = PipelineDeps::pair( + worker_registry.clone(), + Arc::new(PolicyRegistry::new(PolicyConfig::Random)), + None, + ); + RequestPipeline::build(Endpoint::Completion, mode, &deps).expect("completion pipeline") + } + + async fn run_and_drain( + pipeline: RequestPipeline, + components: Arc, + request: Arc, + ) -> bytes::Bytes { + let response = pipeline + .execute_completion(request, None, MODEL.to_string(), components, None, None) + .await; + assert_eq!(response.status(), http::StatusCode::OK); + axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("drain SSE body") + } + + /// The stream task must run off an extracted view: the stub refuses to + /// emit tokens until the parsed request has been freed, so a stream that + /// still pinned it would stall past the stub's deadline. + #[tokio::test] + async fn streaming_releases_parsed_request_before_first_token() { + let request = completion_request(true); + let released = Arc::new(AtomicBool::new(false)); + let port = spawn_stub(GatedScheduler { + probe: Some(Arc::downgrade(&request)), + released: Arc::clone(&released), + ..Default::default() + }) + .await; + + let worker_registry = Arc::new(WorkerRegistry::new()); + register_worker(&worker_registry, port, WorkerType::Regular); + let pipeline = completion_pipeline(&worker_registry, Mode::Regular); + let components = components(worker_registry).await; + + let body = run_and_drain(pipeline, components, request).await; + + assert!( + released.load(Ordering::SeqCst), + "the parsed request must be freed before the first token" + ); + let body = String::from_utf8_lossy(&body); + assert!(body.contains("data: [DONE]"), "stream must finish: {body}"); + } + + /// grpc_pd twin: the decode leg's stream task must not pin the parsed + /// request either (the prefill stub answers immediately). + #[tokio::test] + async fn pd_streaming_releases_parsed_request_before_first_token() { + let request = completion_request(true); + let released = Arc::new(AtomicBool::new(false)); + let prefill_port = spawn_stub(GatedScheduler::default()).await; + let decode_port = spawn_stub(GatedScheduler { + probe: Some(Arc::downgrade(&request)), + released: Arc::clone(&released), + ..Default::default() + }) + .await; + + let worker_registry = Arc::new(WorkerRegistry::new()); + register_worker(&worker_registry, prefill_port, WorkerType::Prefill); + register_worker(&worker_registry, decode_port, WorkerType::Decode); + let pipeline = completion_pipeline(&worker_registry, Mode::PrefillDecode); + let components = components(worker_registry).await; + + let body = run_and_drain(pipeline, components, request).await; + + assert!( + released.load(Ordering::SeqCst), + "the parsed request must be freed before the first decode token" + ); + let body = String::from_utf8_lossy(&body); + assert!(body.contains("data: [DONE]"), "stream must finish: {body}"); + } +} + #[cfg(test)] mod rate_limit_reserve_tests { use llm_tokenizer::{traits::Tokenizer, MockTokenizer, TokenizerRegistry}; diff --git a/model_gateway/src/routers/grpc/regular/mod.rs b/model_gateway/src/routers/grpc/regular/mod.rs index 38311b10a1..22c47fd379 100644 --- a/model_gateway/src/routers/grpc/regular/mod.rs +++ b/model_gateway/src/routers/grpc/regular/mod.rs @@ -7,3 +7,4 @@ pub(crate) mod processor; pub(crate) mod responses; pub(crate) mod stages; pub(crate) mod streaming; +pub(crate) mod views; diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 2513d1e6cd..ba8a7e5bfb 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -11,18 +11,17 @@ use llm_tokenizer::{ traits::Tokenizer, }; use openai_protocol::{ - chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse}, - common::{ - FunctionCallResponse, StringOrArray, Tool, ToolCall, ToolChoice, ToolChoiceValue, Usage, - }, - completion::{CompletionChoice, CompletionRequest, CompletionResponse}, - generate::{GenerateMetaInfo, GenerateRequest, GenerateResponse}, - messages::{self, CreateMessageRequest, Message}, + chat::{ChatChoice, ChatCompletionMessage, ChatCompletionResponse}, + common::{FunctionCallResponse, Tool, ToolCall, ToolChoice, ToolChoiceValue, Usage}, + completion::{CompletionChoice, CompletionResponse}, + generate::{GenerateMetaInfo, GenerateResponse}, + messages::{self, Message}, }; use reasoning_parser::ParserFactory as ReasoningParserFactory; use tool_parser::ParserFactory as ToolParserFactory; use tracing::{error, warn}; +use super::views::{ChatRequestView, CompletionRequestView, MessagesRequestView}; use crate::routers::{ error, grpc::{ @@ -61,7 +60,8 @@ impl ResponseProcessor { &self, complete: &ProtoGenerateComplete, index: usize, - original_request: &ChatCompletionRequest, + original_request: &ChatRequestView, + model: &str, tokenizer: &Arc, stop_decoder: &mut StopSequenceDecoder, history_tool_calls_count: usize, @@ -113,7 +113,7 @@ impl ResponseProcessor { if let Some(mut parser) = utils::create_reasoning_parser( &self.reasoning_parser_factory, reasoning_parser_name, - &original_request.model, + model, ) { // If the template injected `` in the prefill (thinking toggle // is supported and effectively ON), start in reasoning mode. @@ -170,14 +170,14 @@ impl ResponseProcessor { (tool_calls, processed_text) = utils::parse_json_schema_response( &processed_text, original_request.tool_choice.as_ref(), - &original_request.model, + model, history_tool_calls_count, ); } else if tool_parser_available { (tool_calls, processed_text) = self .parse_tool_calls( &processed_text, - &original_request.model, + model, tool_parser_name, original_request.tools.as_deref().unwrap_or(&[]), history_tool_calls_count, @@ -239,26 +239,27 @@ impl ResponseProcessor { pub async fn process_non_streaming_chat_response( &self, execution_result: ExecutionResult, - chat_request: Arc, + chat_request: ChatRequestView, dispatch: DispatchMetadata, tokenizer: Arc, stop_decoder: &mut StopSequenceDecoder, request_logprobs: bool, ) -> Result { - let reasoning_parser_name = self.parser_resolver.reasoning_parser(&chat_request.model); - let tool_parser_name = self.parser_resolver.tool_parser(&chat_request.model); + let model = &dispatch.model; + let reasoning_parser_name = self.parser_resolver.reasoning_parser(model); + let tool_parser_name = self.parser_resolver.tool_parser(model); // Collect all responses from the execution result let all_responses = response_collection::collect_responses(execution_result, request_logprobs).await?; - let history_tool_calls_count = utils::get_history_tool_calls_count(&chat_request); + let history_tool_calls_count = chat_request.history_tool_calls_count; // Check parser availability once upfront (not per choice) let reasoning_parser_available = chat_request.separate_reasoning && utils::check_reasoning_parser_availability( &self.reasoning_parser_factory, reasoning_parser_name.as_deref(), - &chat_request.model, + model, ); let tool_choice_enabled = !matches!( @@ -271,22 +272,18 @@ impl ResponseProcessor { && utils::check_tool_parser_availability( &self.tool_parser_factory, tool_parser_name.as_deref(), - &chat_request.model, + model, ); // Log once per request (not per choice) if chat_request.separate_reasoning && !reasoning_parser_available { tracing::debug!( - "No reasoning parser found for model '{}', skipping reasoning parsing", - chat_request.model + "No reasoning parser found for model '{model}', skipping reasoning parsing" ); } if chat_request.tools.is_some() && tool_choice_enabled && !tool_parser_available { - tracing::debug!( - "No tool parser found for model '{}', skipping tool call parsing", - chat_request.model - ); + tracing::debug!("No tool parser found for model '{model}', skipping tool call parsing"); } // Process all choices @@ -297,6 +294,7 @@ impl ResponseProcessor { complete, index, &chat_request, + model, &tokenizer, stop_decoder, history_tool_calls_count, @@ -396,7 +394,6 @@ impl ResponseProcessor { pub async fn process_non_streaming_generate_response( &self, execution_result: ExecutionResult, - _generate_request: Arc, dispatch: DispatchMetadata, stop_decoder: &mut StopSequenceDecoder, request_logprobs: bool, @@ -504,15 +501,14 @@ impl ResponseProcessor { pub async fn process_non_streaming_messages_response( &self, execution_result: ExecutionResult, - messages_request: Arc, + messages_request: MessagesRequestView, dispatch: DispatchMetadata, tokenizer: Arc, stop_decoder: &mut StopSequenceDecoder, ) -> Result { - let reasoning_parser_name = self - .parser_resolver - .reasoning_parser(&messages_request.model); - let tool_parser_name = self.parser_resolver.tool_parser(&messages_request.model); + let model = &dispatch.model; + let reasoning_parser_name = self.parser_resolver.reasoning_parser(model); + let tool_parser_name = self.parser_resolver.tool_parser(model); // Collect all responses (no logprobs for Messages API) let all_responses = response_collection::collect_responses(execution_result, false).await?; @@ -549,7 +545,7 @@ impl ResponseProcessor { let reasoning_requires_special_tokens = utils::reasoning_parser_requires_special_tokens( &self.reasoning_parser_factory, reasoning_parser_name.as_deref(), - &messages_request.model, + model, ); let separate_reasoning = reasoning_requires_special_tokens || matches!( @@ -563,7 +559,7 @@ impl ResponseProcessor { && utils::check_reasoning_parser_availability( &self.reasoning_parser_factory, reasoning_parser_name.as_deref(), - &messages_request.model, + model, ); let tool_choice_enabled = !matches!( @@ -572,25 +568,21 @@ impl ResponseProcessor { ); let tool_parser_available = tool_choice_enabled - && messages_request.tools.is_some() + && messages_request.has_tools && utils::check_tool_parser_availability( &self.tool_parser_factory, tool_parser_name.as_deref(), - &messages_request.model, + model, ); if separate_reasoning && !reasoning_parser_available { tracing::debug!( - "No reasoning parser found for model '{}', reasoning content will not be separated", - messages_request.model + "No reasoning parser found for model '{model}', reasoning content will not be separated" ); } - if messages_request.tools.is_some() && tool_choice_enabled && !tool_parser_available { - tracing::debug!( - "No tool parser found for model '{}', skipping tool call parsing", - messages_request.model - ); + if messages_request.has_tools && tool_choice_enabled && !tool_parser_available { + tracing::debug!("No tool parser found for model '{model}', skipping tool call parsing"); } // Decode tokens through stop decoder @@ -636,7 +628,7 @@ impl ResponseProcessor { if let Some(mut parser) = utils::create_reasoning_parser( &self.reasoning_parser_factory, reasoning_parser_name.as_deref(), - &messages_request.model, + model, ) { // If thinking is effectively ON and template has a toggle, start in reasoning mode. { @@ -670,7 +662,7 @@ impl ResponseProcessor { // Step 2: Parse tool calls let mut tool_calls: Option> = None; - if tool_choice_enabled && messages_request.tools.is_some() { + if tool_choice_enabled && messages_request.has_tools { // Check if JSON schema constraint was used (specific tool or any/required mode) let has_structural_tag = self .tool_parser_factory @@ -692,24 +684,17 @@ impl ResponseProcessor { (tool_calls, processed_text) = utils::parse_json_schema_response( &processed_text, chat_tool_choice.as_ref(), - &messages_request.model, - utils::message_utils::get_history_tool_calls_count_messages(&messages_request), + model, + messages_request.history_tool_calls_count, ); } else if tool_parser_available { - let chat_tools = messages_request - .tools - .as_deref() - .map(utils::message_utils::extract_chat_tools) - .unwrap_or_default(); (tool_calls, processed_text) = self .parse_tool_calls( &processed_text, - &messages_request.model, + model, tool_parser_name.as_deref(), - &chat_tools, - utils::message_utils::get_history_tool_calls_count_messages( - &messages_request, - ), + &messages_request.chat_tools, + messages_request.history_tool_calls_count, ) .await; } @@ -824,21 +809,18 @@ impl ResponseProcessor { pub async fn process_non_streaming_completion_response( &self, execution_result: ExecutionResult, - completion_req: Arc, + completion_req: CompletionRequestView, dispatch: DispatchMetadata, _tokenizer: Arc, stop_decoder: &mut StopSequenceDecoder, ) -> Result { - let request_logprobs = completion_req.logprobs.is_some(); + let request_logprobs = completion_req.logprobs; let per_prompt_results = match execution_result { ExecutionResult::Batch { results } => results, other => vec![other], }; - let choices_per_prompt = completion_req.n.unwrap_or(1).max(1); - let prompt_texts: Vec<&str> = match &completion_req.prompt { - StringOrArray::String(text) => vec![text.as_str()], - StringOrArray::Array(texts) => texts.iter().map(String::as_str).collect(), - }; + let choices_per_prompt = completion_req.choices_per_prompt; + let prompt_texts = &completion_req.prompt_texts; // Drain all sub-streams concurrently; decoding below stays sequential // (shared stop decoder). @@ -854,7 +836,10 @@ impl ResponseProcessor { let mut choices = Vec::new(); for (prompt_index, all_responses) in collected.into_iter().enumerate() { - let prompt_text = prompt_texts.get(prompt_index).copied().unwrap_or_default(); + let prompt_text = prompt_texts + .get(prompt_index) + .map(String::as_str) + .unwrap_or_default(); let index_offset = prompt_index as u32 * choices_per_prompt; // n>1 choices share one prompt: max within a prompt, summed across prompts. let mut prompt_tokens = 0u32; diff --git a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs index aeac125acb..841e74275c 100644 --- a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs @@ -13,6 +13,7 @@ use crate::routers::{ ClientSelection, ExecutionPlan, ExecutionPlanKind, PreparationOutput, RequestContext, }, multimodal::{assemble_multimodal_data, assemble_multimodal_data_after_encode}, + regular::views, utils, }, }; @@ -173,6 +174,12 @@ impl PipelineStage for ChatRequestBuildingStage { helpers::maybe_inject_pd_rendezvous(&mut proto_request, workers); } + // Last request reader before dispatch: extract the response-phase view + // so response processing never needs the (possibly released) payload. + ctx.state.response.request_view = Some(views::RequestView::Chat( + views::ChatRequestView::from(chat_request.as_ref()), + )); + ctx.state.execution_plan = Some(ExecutionPlan::generate(self.plan_kind, proto_request)); Ok(None) } diff --git a/model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs b/model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs index d27eba5e51..710423025f 100644 --- a/model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs +++ b/model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs @@ -14,7 +14,7 @@ use crate::routers::{ grpc::{ common::stages::{helpers, PipelineStage, RateLimitCell}, context::{FinalResponse, RequestContext}, - regular::{processor, streaming}, + regular::{processor, streaming, views}, }, }; @@ -89,13 +89,27 @@ impl ChatResponseProcessingStage { ) })?; + // Response-phase view set by request building; the payload itself may + // already be released. + let Some(views::RequestView::Chat(chat_request)) = ctx.state.response.request_view.take() + else { + error!( + function = "ChatResponseProcessingStage::execute", + "Request view not set" + ); + return Err(error::internal_error( + "request_view_not_set", + "Request view not set", + )); + }; + if is_streaming { // Read derived skip_special_tokens (set in preparation, survives request_building .take()) let skip_special_tokens = ctx .state .response .skip_special_tokens - .unwrap_or_else(|| ctx.chat_request().skip_special_tokens); + .unwrap_or(chat_request.skip_special_tokens); // Reserved (if tenant rate limiting is enabled): settled with real // usage inside the streaming processor on success, or abandoned @@ -107,13 +121,14 @@ impl ChatResponseProcessingStage { .as_deref() .and_then(RateLimitCell::take_for_streaming_handoff); - // Streaming: Use StreamingProcessor and return SSE response + // Streaming: Use StreamingProcessor and return SSE response. The + // stream task consumes the view, never the parsed request. let response = self .streaming_processor .clone() .process_streaming_response( execution_result, - ctx.chat_request_arc(), // Cheap Arc clone (8 bytes) + chat_request, dispatch, tokenizer, skip_special_tokens, @@ -133,9 +148,7 @@ impl ChatResponseProcessingStage { } // Non-streaming: Delegate to ResponseProcessor - let request_logprobs = ctx.chat_request().logprobs; - - let chat_request = ctx.chat_request_arc(); + let request_logprobs = chat_request.logprobs; let stop_decoder = ctx.state.response.stop_decoder.as_mut().ok_or_else(|| { error!( diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs index 99807062c7..f1f0b3369a 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs @@ -24,6 +24,7 @@ use crate::routers::{ RequestContext, RequestType, WorkerSelection, }, proto_wrapper::ProtoGenerateRequest, + regular::views, }, }; @@ -214,6 +215,12 @@ impl PipelineStage for CompletionRequestBuildingStage { } }; + // Last request reader before dispatch: extract the response-phase view + // so response processing never needs the (possibly released) payload. + ctx.state.response.request_view = Some(views::RequestView::Completion( + views::CompletionRequestView::from(completion_request.as_ref()), + )); + ctx.state.execution_plan = Some(plan); Ok(None) } diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs b/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs index 5942eee7e3..25d0fecb5b 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs @@ -17,7 +17,7 @@ use crate::routers::{ grpc::{ common::stages::{helpers, PipelineStage, RateLimitCell}, context::{FinalResponse, RequestContext}, - regular::{processor, streaming}, + regular::{processor, streaming, views}, }, }; @@ -76,6 +76,21 @@ impl PipelineStage for CompletionResponseProcessingStage { ) })?; + // Response-phase view set by request building; the payload itself may + // already be released. + let Some(views::RequestView::Completion(completion_request)) = + ctx.state.response.request_view.take() + else { + error!( + function = "CompletionResponseProcessingStage::execute", + "Request view not set" + ); + return Err(error::internal_error( + "request_view_not_set", + "Request view not set", + )); + }; + if is_streaming { // Reserved (if tenant rate limiting is enabled): settled with real // usage inside the streaming processor on success, or abandoned @@ -87,12 +102,14 @@ impl PipelineStage for CompletionResponseProcessingStage { .as_deref() .and_then(RateLimitCell::take_for_streaming_handoff); + // Streaming: the stream tasks consume the view, never the parsed + // request. let response = self .streaming_processor .clone() .process_completion_streaming_response( execution_result, - ctx.completion_request_arc(), + completion_request, dispatch, tokenizer, reservation.clone(), @@ -111,8 +128,6 @@ impl PipelineStage for CompletionResponseProcessingStage { } // Non-streaming path - let completion_request = ctx.completion_request_arc(); - let stop_decoder = ctx.state.response.stop_decoder.as_mut().ok_or_else(|| { error!( function = "CompletionResponseProcessingStage::execute", diff --git a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs index b18dc1d7ab..7f2eecf543 100644 --- a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs @@ -9,6 +9,7 @@ use crate::routers::{ grpc::{ common::stages::{helpers, PipelineStage}, context::{ClientSelection, ExecutionPlan, ExecutionPlanKind, RequestContext}, + regular::views, }, }; @@ -104,6 +105,12 @@ impl PipelineStage for GenerateRequestBuildingStage { helpers::maybe_inject_pd_rendezvous(&mut proto_request, workers); } + // Last request reader before dispatch: extract the response-phase view + // so response processing never needs the (possibly released) payload. + ctx.state.response.request_view = Some(views::RequestView::Generate( + views::GenerateRequestView::from(generate_request.as_ref()), + )); + ctx.state.execution_plan = Some(ExecutionPlan::generate(self.plan_kind, proto_request)); Ok(None) } diff --git a/model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs b/model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs index 59834d957e..5f40566a11 100644 --- a/model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs +++ b/model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs @@ -11,7 +11,7 @@ use crate::routers::{ grpc::{ common::stages::{helpers, PipelineStage, RateLimitCell}, context::{FinalResponse, RequestContext}, - regular::{processor, streaming}, + regular::{processor, streaming, views}, }, }; @@ -89,6 +89,21 @@ impl GenerateResponseProcessingStage { ) })?; + // Response-phase view set by request building; the payload itself may + // already be released. + let Some(views::RequestView::Generate(generate_request)) = + ctx.state.response.request_view.take() + else { + error!( + function = "GenerateResponseProcessingStage::execute", + "Request view not set" + ); + return Err(error::internal_error( + "request_view_not_set", + "Request view not set", + )); + }; + if is_streaming { // Reserved (if tenant rate limiting is enabled): settled with real // usage inside the streaming processor on success, or abandoned @@ -106,7 +121,7 @@ impl GenerateResponseProcessingStage { .clone() .process_streaming_generate( execution_result, - ctx.generate_request_arc(), // Cheap Arc clone (8 bytes) + generate_request, dispatch, tokenizer, reservation.clone(), @@ -125,8 +140,7 @@ impl GenerateResponseProcessingStage { } // Non-streaming: Delegate to ResponseProcessor - let request_logprobs = ctx.generate_request().return_logprob.unwrap_or(false); - let generate_request = ctx.generate_request_arc(); + let request_logprobs = generate_request.return_logprob; let stop_decoder = ctx.state.response.stop_decoder.as_mut().ok_or_else(|| { error!( @@ -143,7 +157,6 @@ impl GenerateResponseProcessingStage { .processor .process_non_streaming_generate_response( execution_result, - generate_request, dispatch, stop_decoder, request_logprobs, diff --git a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs index 78d4f477f8..d20d82d3e9 100644 --- a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs @@ -14,6 +14,7 @@ use crate::routers::{ ClientSelection, ExecutionPlan, ExecutionPlanKind, PreparationOutput, RequestContext, }, multimodal::{assemble_multimodal_data, assemble_multimodal_data_after_encode}, + regular::views, utils, }, }; @@ -173,6 +174,12 @@ impl PipelineStage for MessageRequestBuildingStage { helpers::maybe_inject_pd_rendezvous(&mut proto_request, workers); } + // Last request reader before dispatch: extract the response-phase view + // so response processing never needs the (possibly released) payload. + ctx.state.response.request_view = Some(views::RequestView::Messages( + views::MessagesRequestView::from(messages_request.as_ref()), + )); + ctx.state.execution_plan = Some(ExecutionPlan::generate(self.plan_kind, proto_request)); Ok(None) } diff --git a/model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs b/model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs index cf12204cea..cb4d72662b 100644 --- a/model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs +++ b/model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs @@ -15,7 +15,7 @@ use crate::routers::{ grpc::{ common::stages::{helpers, PipelineStage, RateLimitCell}, context::{FinalResponse, RequestContext}, - regular::{processor, streaming}, + regular::{processor, streaming, views}, }, }; @@ -77,6 +77,21 @@ impl PipelineStage for MessageResponseProcessingStage { ) })?; + // Response-phase view set by request building; the payload itself may + // already be released. + let Some(views::RequestView::Messages(messages_request)) = + ctx.state.response.request_view.take() + else { + error!( + function = "MessageResponseProcessingStage::execute", + "Request view not set" + ); + return Err(error::internal_error( + "request_view_not_set", + "Request view not set", + )); + }; + if is_streaming { // Read derived skip_special_tokens (set in preparation, survives request_building .take()) let skip_special_tokens = ctx.state.response.skip_special_tokens.unwrap_or(true); @@ -91,13 +106,14 @@ impl PipelineStage for MessageResponseProcessingStage { .as_deref() .and_then(RateLimitCell::take_for_streaming_handoff); - // Streaming: use StreamingProcessor and return SSE response + // Streaming: use StreamingProcessor and return SSE response. The + // stream task consumes the view, never the parsed request. let response = self .streaming_processor .clone() .process_messages_streaming_response( execution_result, - ctx.messages_request_arc(), + messages_request, dispatch, tokenizer, skip_special_tokens, @@ -117,8 +133,6 @@ impl PipelineStage for MessageResponseProcessingStage { } // Non-streaming: delegate to ResponseProcessor - let messages_request = ctx.messages_request_arc(); - let stop_decoder = ctx.state.response.stop_decoder.as_mut().ok_or_else(|| { error!( function = "MessageResponseProcessingStage::execute", diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index c08d78a171..1d3cf8d3a0 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -16,16 +16,15 @@ use llm_tokenizer::{ traits::Tokenizer, }; use openai_protocol::{ - chat::{ChatCompletionRequest, ChatCompletionStreamResponse}, + chat::ChatCompletionStreamResponse, common::{ ChatLogProbs, FunctionCallDelta, StringOrArray, Tool, ToolCallDelta, ToolChoice, ToolChoiceValue, Usage, }, - completion::{CompletionRequest, CompletionStreamChoice, CompletionStreamResponse}, - generate::GenerateRequest, + completion::{CompletionStreamChoice, CompletionStreamResponse}, messages::{ - self, ContentBlock, ContentBlockDelta, CreateMessageRequest, Message, MessageDelta, - MessageDeltaUsage, MessageStreamEvent, + self, ContentBlock, ContentBlockDelta, Message, MessageDelta, MessageDeltaUsage, + MessageStreamEvent, }, }; use reasoning_parser::{ParserFactory as ReasoningParserFactory, ParserResult, ReasoningParser}; @@ -33,6 +32,9 @@ use serde_json::{json, Value}; use tool_parser::{ParserFactory as ToolParserFactory, StreamingParseResult, ToolParser}; use tracing::{debug, error, warn}; +use super::views::{ + ChatRequestView, CompletionRequestView, GenerateRequestView, MessagesRequestView, +}; use crate::{ observability::metrics::{metrics_labels, Metrics, StreamingMetricsParams}, rate_limit::{SharedReservationHandle, UsageSettlement}, @@ -123,7 +125,7 @@ impl StreamingProcessor { pub async fn process_streaming_response( self: Arc, execution_result: context::ExecutionResult, - chat_request: Arc, + chat_request: ChatRequestView, dispatch: context::DispatchMetadata, tokenizer: Arc, skip_special_tokens: bool, @@ -238,7 +240,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - original_request: Arc, + original_request: ChatRequestView, tx: &SseSender, reservation: Option>, ) -> Result<(), String> { @@ -265,7 +267,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - original_request: Arc, + original_request: ChatRequestView, tx: &SseSender, pd_timing: Option, reservation: Option>, @@ -278,7 +280,7 @@ impl StreamingProcessor { let separate_reasoning = original_request.separate_reasoning; let tool_choice = &original_request.tool_choice; let tools = &original_request.tools; - let history_tool_calls_count = utils::get_history_tool_calls_count(&original_request); + let history_tool_calls_count = original_request.history_tool_calls_count; let stream_options = &original_request.stream_options; // Phase 1: Initialize state tracking (per-index for n>1 support) @@ -730,7 +732,7 @@ impl StreamingProcessor { // produce one for every expected `n>1` choice has only partial // usage -- settling with that would understate the real cost. // Keep the reserved amount as final instead. - let expected_choices = original_request.n.unwrap_or(1).max(1); + let expected_choices = original_request.expected_choices; if (prompt_tokens.len() as u32) < expected_choices { handle.close_reserved_only().await; } else { @@ -765,7 +767,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - original_request: Arc, + original_request: ChatRequestView, tx: &SseSender, pd_timing: context::PdTiming, reservation: Option>, @@ -820,7 +822,7 @@ impl StreamingProcessor { pub async fn process_streaming_generate( self: Arc, execution_result: context::ExecutionResult, - generate_request: Arc, + generate_request: GenerateRequestView, dispatch: context::DispatchMetadata, tokenizer: Arc, reservation: Option>, @@ -835,15 +837,10 @@ impl StreamingProcessor { .weight_version .clone() .unwrap_or_else(|| "default".to_string()), - return_logprob: generate_request.return_logprob.unwrap_or(false), + return_logprob: generate_request.return_logprob, backend_type: self.backend_type, model: dispatch.model.clone(), - expected_choices: generate_request - .sampling_params - .as_ref() - .and_then(|p| p.n) - .unwrap_or(1) - .max(1), + expected_choices: generate_request.expected_choices, }; // Spawn background task based on execution mode @@ -1716,7 +1713,7 @@ impl StreamingProcessor { pub async fn process_messages_streaming_response( self: Arc, execution_result: context::ExecutionResult, - messages_request: Arc, + messages_request: MessagesRequestView, dispatch: context::DispatchMetadata, tokenizer: Arc, skip_special_tokens: bool, @@ -1845,7 +1842,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - original_request: Arc, + original_request: MessagesRequestView, tx: &SseSender, reservation: Option>, ) -> Result<(), String> { @@ -1858,7 +1855,7 @@ impl StreamingProcessor { let request_id = &dispatch.request_id; let model = &dispatch.model; - let has_tools = original_request.tools.is_some(); + let has_tools = original_request.has_tools; // Content block state machine let mut current_block_index: u32 = 0; @@ -1968,15 +1965,10 @@ impl StreamingProcessor { Some(messages::ToolChoice::Tool { .. }) ); - let history_tool_calls_count = - message_utils::get_history_tool_calls_count_messages(&original_request); + let history_tool_calls_count = original_request.history_tool_calls_count; - // Pre-convert Messages tools to Chat tools for parser reuse (done once upfront) - let chat_tools: Vec = original_request - .tools - .as_deref() - .map(message_utils::extract_chat_tools) - .unwrap_or_default(); + // Messages tools pre-converted to Chat tools for parser reuse + let chat_tools: &[Tool] = &original_request.chat_tools; // Create fresh streaming tool parser (not pooled — streaming parsers maintain state) let mut streaming_tool_parser: Option> = @@ -2208,7 +2200,7 @@ impl StreamingProcessor { } } else if let Some(ref mut parser) = streaming_tool_parser { // Regular/required tool choice: use incremental parser - match parser.parse_incremental(&normal_text, &chat_tools).await { + match parser.parse_incremental(&normal_text, chat_tools).await { Ok(StreamingParseResult { normal_text: text, calls, @@ -2537,7 +2529,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - original_request: Arc, + original_request: MessagesRequestView, tx: &SseSender, reservation: Option>, ) -> Result<(), String> { @@ -2583,7 +2575,7 @@ impl StreamingProcessor { pub async fn process_completion_streaming_response( self: Arc, execution_result: context::ExecutionResult, - completion_request: Arc, + completion_request: CompletionRequestView, dispatch: context::DispatchMetadata, tokenizer: Arc, reservation: Option>, @@ -2606,20 +2598,13 @@ impl StreamingProcessor { )] tokio::spawn(async move { let start_time = Instant::now(); - let choices_per_prompt = completion_request.n.unwrap_or(1).max(1); + let choices_per_prompt = completion_request.choices_per_prompt; let echo = completion_request.echo; - let include_usage = completion_request - .stream_options - .as_ref() - .and_then(|opts| opts.include_usage) - .unwrap_or(false); - let prompt_texts: Vec<&str> = match &completion_request.prompt { - StringOrArray::String(text) => vec![text.as_str()], - StringOrArray::Array(texts) => texts.iter().map(String::as_str).collect(), - }; + let include_usage = completion_request.include_usage; // Fail-fast: the first stream error cancels the remaining units // (their streams abort on drop) and fails the whole request. + let completion_request = &completion_request; let outcomes = try_join_all(units.into_iter().enumerate().map(|(prompt_index, unit)| { let stop_params = ( @@ -2631,9 +2616,12 @@ impl StreamingProcessor { ); let dispatch = dispatch.clone(); let tokenizer = tokenizer.clone(); - let completion_request = completion_request.clone(); let prompt_text = if echo { - prompt_texts.get(prompt_index).copied().unwrap_or_default() + completion_request + .prompt_texts + .get(prompt_index) + .map(String::as_str) + .unwrap_or_default() } else { "" }; @@ -2805,7 +2793,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - completion_request: Arc, + completion_request: &CompletionRequestView, prompt_text: &str, index_offset: u32, tx: &SseSender, @@ -2819,8 +2807,6 @@ impl StreamingProcessor { let echo = completion_request.echo; let suffix = completion_request.suffix.as_deref(); - // TODO: wire per-token logprob streaming when backend support is available - let _request_logprobs = completion_request.logprobs.is_some(); let mut stop_decoders: HashMap = HashMap::new(); let mut is_firsts: HashMap = HashMap::new(); @@ -3096,7 +3082,7 @@ impl StreamingProcessor { // via a `Complete` message. A clean EOF partway through this unit's // `n>1` choices (some completed, others didn't) must not be treated // as full usage. - let expected_choices = completion_request.n.unwrap_or(1).max(1); + let expected_choices = completion_request.choices_per_prompt; let saw_complete = completed_indices.len() as u32 >= expected_choices; Ok(CompletionStreamOutcome { @@ -3119,7 +3105,7 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), - original_request: Arc, + original_request: &CompletionRequestView, prompt_text: &str, index_offset: u32, tx: &SseSender, diff --git a/model_gateway/src/routers/grpc/regular/views.rs b/model_gateway/src/routers/grpc/regular/views.rs new file mode 100644 index 0000000000..dd515fd6d3 --- /dev/null +++ b/model_gateway/src/routers/grpc/regular/views.rs @@ -0,0 +1,162 @@ +//! Response-phase views of the parsed request. +//! +//! Response processing and stream tasks consume these instead of the parsed +//! request, so the request (and its multimodal payloads) is never pinned for +//! the lifetime of the response. + +use std::collections::HashMap; + +use openai_protocol::{ + chat::ChatCompletionRequest, + common::{StreamOptions, StringOrArray, Tool, ToolChoice}, + completion::CompletionRequest, + generate::GenerateRequest, + messages::{self, CreateMessageRequest}, +}; +use serde_json::Value; + +use crate::routers::grpc::utils; + +/// Response-phase view of one request, set by request building (the last +/// request reader before dispatch) and taken by response processing. +pub(crate) enum RequestView { + Chat(ChatRequestView), + Generate(GenerateRequestView), + Completion(CompletionRequestView), + Messages(MessagesRequestView), +} + +pub(crate) struct ChatRequestView { + pub separate_reasoning: bool, + pub tool_choice: Option, + pub tools: Option>, + pub history_tool_calls_count: usize, + pub stream_options: Option, + pub chat_template_kwargs: Option>, + pub reasoning_effort: Option, + /// `sampling_params.n`, normalized. + pub expected_choices: u32, + pub logprobs: bool, + pub stop: Option, + pub stop_token_ids: Option>, + pub no_stop_trim: bool, + pub ignore_eos: bool, + /// Fallback when preparation derived no override. + pub skip_special_tokens: bool, +} + +impl From<&ChatCompletionRequest> for ChatRequestView { + fn from(request: &ChatCompletionRequest) -> Self { + Self { + separate_reasoning: request.separate_reasoning, + tool_choice: request.tool_choice.clone(), + tools: request.tools.clone(), + history_tool_calls_count: utils::get_history_tool_calls_count(request), + stream_options: request.stream_options.clone(), + chat_template_kwargs: request.chat_template_kwargs.clone(), + reasoning_effort: request.reasoning_effort.clone(), + expected_choices: request.n.unwrap_or(1).max(1), + logprobs: request.logprobs, + stop: request.stop.clone(), + stop_token_ids: request.stop_token_ids.clone(), + no_stop_trim: request.no_stop_trim, + ignore_eos: request.ignore_eos, + skip_special_tokens: request.skip_special_tokens, + } + } +} + +pub(crate) struct GenerateRequestView { + pub return_logprob: bool, + /// `sampling_params.n`, normalized. + pub expected_choices: u32, +} + +impl From<&GenerateRequest> for GenerateRequestView { + fn from(request: &GenerateRequest) -> Self { + Self { + return_logprob: request.return_logprob.unwrap_or(false), + expected_choices: request + .sampling_params + .as_ref() + .and_then(|p| p.n) + .unwrap_or(1) + .max(1), + } + } +} + +pub(crate) struct MessagesRequestView { + pub thinking: Option, + pub tool_choice: Option, + pub has_tools: bool, + pub history_tool_calls_count: usize, + /// Messages tools pre-converted to Chat tools for parser reuse. + pub chat_tools: Vec, + pub stop_sequences: Option>, +} + +impl From<&CreateMessageRequest> for MessagesRequestView { + fn from(request: &CreateMessageRequest) -> Self { + Self { + thinking: request.thinking.clone(), + tool_choice: request.tool_choice.clone(), + has_tools: request.tools.is_some(), + history_tool_calls_count: utils::message_utils::get_history_tool_calls_count_messages( + request, + ), + chat_tools: request + .tools + .as_deref() + .map(utils::message_utils::extract_chat_tools) + .unwrap_or_default(), + stop_sequences: request.stop_sequences.clone(), + } + } +} + +pub(crate) struct CompletionRequestView { + /// `n`, normalized. + pub choices_per_prompt: u32, + pub echo: bool, + pub suffix: Option, + pub logprobs: bool, + pub include_usage: bool, + /// Populated only when `echo` (choices prepend their prompt text). + pub prompt_texts: Vec, + pub stop: Option, + pub stop_token_ids: Option>, + pub skip_special_tokens: bool, + pub no_stop_trim: bool, + pub ignore_eos: bool, +} + +impl From<&CompletionRequest> for CompletionRequestView { + fn from(request: &CompletionRequest) -> Self { + let prompt_texts = if request.echo { + match &request.prompt { + StringOrArray::String(text) => vec![text.clone()], + StringOrArray::Array(texts) => texts.clone(), + } + } else { + Vec::new() + }; + Self { + choices_per_prompt: request.n.unwrap_or(1).max(1), + echo: request.echo, + suffix: request.suffix.clone(), + logprobs: request.logprobs.is_some(), + include_usage: request + .stream_options + .as_ref() + .and_then(|opts| opts.include_usage) + .unwrap_or(false), + prompt_texts, + stop: request.stop.clone(), + stop_token_ids: request.stop_token_ids.clone(), + skip_special_tokens: request.skip_special_tokens, + no_stop_trim: request.no_stop_trim, + ignore_eos: request.ignore_eos, + } + } +}