diff --git a/sdk_v2/cpp/CMakeLists.txt b/sdk_v2/cpp/CMakeLists.txt index cc07cd0d9..5d5082181 100644 --- a/sdk_v2/cpp/CMakeLists.txt +++ b/sdk_v2/cpp/CMakeLists.txt @@ -214,12 +214,14 @@ set(FOUNDRY_LOCAL_SOURCES src/inferencing/generative/embeddings/embeddings_session.cc src/inferencing/generative/chat/chat_generator.cc src/inferencing/session/session.cc + src/inferencing/session/request.cc src/inferencing/session/session_manager.cc src/inferencing/session/tool_registry.cc src/inferencing/generative/chat/chat_session.cc src/inferencing/generative/chat/chat_template.cc src/inferencing/generative/chat/chat_transcript.cc src/inferencing/generative/chat/media_input.cc + src/inferencing/generative/chat/prepared_chat_prompt.cc src/configuration.cc src/download/blob_download_state.cc src/download/blob_downloader.cc diff --git a/sdk_v2/cpp/include/foundry_local/foundry_local_c.h b/sdk_v2/cpp/include/foundry_local/foundry_local_c.h index 4bce94459..f2f378904 100644 --- a/sdk_v2/cpp/include/foundry_local/foundry_local_c.h +++ b/sdk_v2/cpp/include/foundry_local/foundry_local_c.h @@ -126,6 +126,7 @@ FL_TYPE(ModelList); // Request accumulates parameters and input items // Response accumulates output items FL_TYPE(Request); +FL_TYPE(RequestPreflight); FL_TYPE(Response); // Opaque type for a session. Create with loaded Model so Model:Session is 1:M @@ -377,6 +378,18 @@ typedef struct flUsage { /* V3 fields go here. Read only when version >= 3. */ } flUsage; +/// Exact token-budget preflight result for a request in the captured session state. +typedef struct flRequestPreflightResult { + uint32_t version; ///< Set to FOUNDRY_LOCAL_API_VERSION. + int64_t prompt_tokens; ///< Exact number of prompt tokens after request preparation. + int64_t output_reserve_tokens; ///< Tokens reserved for generated output. + int64_t required_tokens; ///< Total tokens required: prompt plus output reserve. + int64_t context_limit_tokens; ///< Model context-window limit. + bool fits; ///< Whether required_tokens fits within the context limit. + int64_t deficit_tokens; ///< Tokens over budget, or 0 when fits is true. + /* V3 fields go here. Read only when version >= 3. */ +} flRequestPreflightResult; + /// Information about a discoverable execution provider. /// Returned by Manager_GetDiscoverableEps. Storage is owned by the Manager; the /// returned pointers and `name` strings are stable for the Manager's lifetime. @@ -953,6 +966,21 @@ struct flInferenceApi { FL_API_STATUS(Session_UndoTurns, _In_ flSession* session, size_t count); // End V1 + + /// Capture the request and current chat-session state for an exact token-budget preflight. + /// The returned one-shot handle remains valid after the source request and session are released. The Manager and + /// its model runtime must remain alive until the handle is executed and released. Execute may run on another + /// thread; callers must not execute and release the same handle concurrently. + FL_API_STATUS(Session_CreateRequestPreflight, _In_ const flSession* session, _In_ const flRequest* request, + _Outptr_ flRequestPreflight** out_preflight); + + /// Synchronously execute a preflight. This operation may be expensive. + /// The caller must initialize out_result->version to FOUNDRY_LOCAL_API_VERSION. + FL_API_STATUS(RequestPreflight_Execute, _In_ flRequestPreflight* preflight, + _Out_ flRequestPreflightResult* out_result); + + /// Release a preflight handle. Passing nullptr is allowed. + FL_TYPE_RELEASE(RequestPreflight); }; /* --- Configuration API ------------------------------------------------- */ diff --git a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h index 3a0e0c16c..93fe54448 100644 --- a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h +++ b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.h @@ -1037,6 +1037,29 @@ class Request { detail::Base handle_; }; +/// Captured exact token-budget preflight operation. +/// +/// Capture is separate from execution so callers can capture request and session state on one +/// thread, then move the operation to a worker for the potentially expensive execution. The +/// source Request and ChatSession do not need to remain alive after capture. The owning Manager +/// and model runtime must remain alive until this operation is executed and destroyed. +class RequestPreflight { + public: + RequestPreflight(const RequestPreflight&) = delete; + RequestPreflight& operator=(const RequestPreflight&) = delete; + RequestPreflight(RequestPreflight&&) noexcept = default; + RequestPreflight& operator=(RequestPreflight&&) noexcept = default; + + /// Execute the captured preflight operation. + flRequestPreflightResult Execute(); + + private: + friend class ChatSession; + explicit RequestPreflight(flRequestPreflight* preflight); + + detail::Base handle_; +}; + /// Wrapper for an opaque flResponse. class Response { public: @@ -1104,6 +1127,16 @@ class ChatSession : public Session { /// Undo the last `count` turns and remove their input messages and assistant replies from history. /// Retained inference state is reused when it can be restored safely; otherwise it is rebuilt on the next request. void UndoTurns(size_t count); + + /// Capture request and session state for a later exact token-budget preflight. + /// The returned operation can be moved to a worker and executed after the source request and + /// session are destroyed. The owning Manager and model runtime must remain alive until the + /// returned operation is executed and destroyed. + RequestPreflight CaptureRequestPreflight(const Request& request) const; + + /// Synchronously capture and execute an exact token-budget preflight without mutating the request or session. + /// This operation may be expensive. The owning Manager and model runtime must remain alive for the call. + flRequestPreflightResult PreflightRequest(const Request& request) const; }; /// Session for automatic-speech-recognition (transcription) models. diff --git a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h index ded9498f7..36ef680e5 100644 --- a/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h +++ b/sdk_v2/cpp/include/foundry_local/foundry_local_cpp.inline.h @@ -1160,6 +1160,20 @@ inline flRequest* detail::CreateRequest() { return req; } +// =========================================================================== +// RequestPreflight +// =========================================================================== + +inline RequestPreflight::RequestPreflight(flRequestPreflight* preflight) + : handle_(preflight, detail::inference_api()->RequestPreflight_Release) {} + +inline flRequestPreflightResult RequestPreflight::Execute() { + flRequestPreflightResult result{}; + result.version = FOUNDRY_LOCAL_API_VERSION; + Check(detail::inference_api()->RequestPreflight_Execute(handle_.get_mutable(), &result)); + return result; +} + // =========================================================================== // Response // =========================================================================== @@ -1270,6 +1284,16 @@ inline void ChatSession::UndoTurns(size_t count) { Check(detail::inference_api()->Session_UndoTurns(handle_.get_mutable(), count)); } +inline RequestPreflight ChatSession::CaptureRequestPreflight(const Request& request) const { + flRequestPreflight* preflight = nullptr; + Check(detail::inference_api()->Session_CreateRequestPreflight(handle_.get(), request.native_handle(), &preflight)); + return RequestPreflight(preflight); +} + +inline flRequestPreflightResult ChatSession::PreflightRequest(const Request& request) const { + return CaptureRequestPreflight(request).Execute(); +} + // =========================================================================== // AudioSession // =========================================================================== diff --git a/sdk_v2/cpp/src/c_api.cc b/sdk_v2/cpp/src/c_api.cc index 92ebc0556..929f699f3 100644 --- a/sdk_v2/cpp/src/c_api.cc +++ b/sdk_v2/cpp/src/c_api.cc @@ -1903,6 +1903,51 @@ FL_API_STATUS_IMPL(Session_UndoTurnsImpl, flSession* session, size_t count) { API_IMPL_END } +FL_API_STATUS_IMPL(Session_CreateRequestPreflightImpl, const flSession* session, const flRequest* request, + flRequestPreflight** out_preflight) { + API_IMPL_BEGIN + if (!out_preflight) { + return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "null out_preflight"); + } + + *out_preflight = nullptr; + if (!session || !request) { + return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "null argument"); + } + + auto operation = AsImpl(session)->CreateRequestPreflight(*AsImpl(request)); + *out_preflight = reinterpret_cast(operation.release()); + return nullptr; + API_IMPL_END +} + +FL_API_STATUS_IMPL(RequestPreflight_ExecuteImpl, flRequestPreflight* preflight, + flRequestPreflightResult* out_result) { + API_IMPL_BEGIN + if (!preflight || !out_result) { + return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "null argument"); + } + if (out_result->version != FOUNDRY_LOCAL_API_VERSION) { + return MakeStatus(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "out_result->version is not supported"); + } + + auto* operation = reinterpret_cast(preflight); + const auto result = operation->Execute(); + out_result->version = FOUNDRY_LOCAL_API_VERSION; + out_result->prompt_tokens = result.prompt_tokens; + out_result->output_reserve_tokens = result.output_reserve_tokens; + out_result->required_tokens = result.required_tokens; + out_result->context_limit_tokens = result.context_limit_tokens; + out_result->fits = result.fits; + out_result->deficit_tokens = result.deficit_tokens; + return nullptr; + API_IMPL_END +} + +static void FL_API_CALL RequestPreflight_ReleaseImpl(flRequestPreflight* preflight) FL_NO_EXCEPTION { + delete reinterpret_cast(preflight); +} + static const flInferenceApi g_inference_api = { Request_CreateImpl, Request_ReleaseImpl, @@ -1926,10 +1971,18 @@ static const flInferenceApi g_inference_api = { Session_RemoveToolDefinitionImpl, Session_GetTurnCountImpl, Session_UndoTurnsImpl, + Session_CreateRequestPreflightImpl, + RequestPreflight_ExecuteImpl, + RequestPreflight_ReleaseImpl, }; static_assert(offsetof(flInferenceApi, Session_UndoTurns) / sizeof(void*) == 21, - "Size of version 1 Inference API cannot change"); + "Version 1 Inference API prefix cannot change"); +static_assert(offsetof(flInferenceApi, Session_CreateRequestPreflight) / sizeof(void*) == 22, + "Version 2 Inference API must append after the version 1 prefix"); +static_assert(sizeof(flInferenceApi) == + offsetof(flInferenceApi, Session_CreateRequestPreflight) + 3 * sizeof(void*), + "Version 2 Inference API must append exactly three preflight functions"); // ======================================================================== // Sub-API accessors diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.cc b/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.cc index c3fac9869..3985a1a53 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.cc @@ -3,11 +3,20 @@ #include "inferencing/generative/chat/chat_generator.h" #include "exception.h" +#include "inferencing/generative/chat/prepared_chat_prompt.h" namespace fl { void ChatGenerator::Close() {} +int ChatGenerator::AppendPreparedPrompt(const std::vector&, + const PreparedChatPrompt&, + GenAIModelInstance&, + const ToolCallContext&, + const SearchOptions&) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "generator does not support prepared prompt append"); +} + void ChatGenerator::RewindTo(int /*token_count*/) { FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "This generator does not support rewinding retained model state"); } diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h b/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h index 012559e54..cbfc284a4 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h @@ -13,6 +13,7 @@ namespace fl { class GenAIModelInstance; +struct PreparedChatPrompt; struct TranscriptMessage; struct SearchOptions; struct ToolCallContext; @@ -96,6 +97,13 @@ class ChatGenerator { const ToolCallContext& tool_ctx, const SearchOptions& options) = 0; + /// Append a prompt already rendered and tokenized by the authoritative request preparation path. + virtual int AppendPreparedPrompt(const std::vector& new_messages, + const PreparedChatPrompt& prompt, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx, + const SearchOptions& options); + /// Whether the prompt for the active turn ends inside a reasoning block opened by the chat template. virtual bool PromptOpensReasoning() const { return false; } diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc index 7450fa289..5934c22a1 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc @@ -10,6 +10,7 @@ #include "inferencing/generative/chat/onnx_chat_engine.h" #include "inferencing/generative/chat/onnx_chat_generator.h" #include "inferencing/generative/chat/onnx_engine_chat_stream.h" +#include "inferencing/generative/chat/prepared_chat_prompt.h" #include "inferencing/generative/chat/reasoning_stream_splitter.h" #include "inferencing/generative/chat/stop_strings.h" #include "inferencing/generative/genai_model_instance.h" @@ -74,6 +75,25 @@ std::unique_ptr CreateTextChatGenerator(const chat_internal::Prep return OnnxChatGenerator::Create(messages, options, model, tool_ctx, use_full_context); } +std::unique_ptr CreatePreparedTextChatGenerator( + const chat_internal::PreparedChatMessages& messages, + PreparedChatPrompt prepared, + const SearchOptions& options, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx, + bool use_full_context, + const TextChatGeneratorFactory& factory) { + if (factory) { + return factory(messages, options, model, tool_ctx, use_full_context); + } + + if (model.GetGenAIConfig().GetChatBackendKind() != ChatBackendKind::kGenerator) { + return OnnxEngineChatStream::CreatePrepared(std::move(prepared), options, model, tool_ctx); + } + + return OnnxChatGenerator::CreatePrepared(std::move(prepared), options, model, tool_ctx, use_full_context); +} + using TextSegment = ReasoningStreamSplitter::Segment; /// Value of `key` in `options`, or an empty string when it is absent. @@ -866,15 +886,21 @@ void ChatSession::SetSessionOptionsImpl(const KeyValuePairs& options) { session_options_ = SearchOptions::FromParameters(options); } -ToolCallContext ChatSession::BuildToolCallContext(const Request& request, - const std::vector& definitions) const { +namespace { + +ToolCallContext BuildToolCallContextForRequest(const Request& request, + const std::vector& definitions, + const SearchOptions& session_options, + const ModelInfo& model_info, + GenAIModelInstance& model, + ILogger& logger) { ToolCallContext tool_ctx; tool_ctx.tool_call_start = GetOptionOrEmpty(request.options, FOUNDRY_LOCAL_MODEL_PROP_TOOL_CALL_START_STR); tool_ctx.tool_call_end = GetOptionOrEmpty(request.options, FOUNDRY_LOCAL_MODEL_PROP_TOOL_CALL_END_STR); // Fall back to model info properties if not specified in the request - const auto& info = CatalogModel().Info(); + const auto& info = model_info; // Check if the model supports tool calling const auto* tool_calling_val = info.GetPropertyInt(FOUNDRY_LOCAL_MODEL_PROP_SUPPORTS_TOOL_CALLING_INT); @@ -899,7 +925,7 @@ ToolCallContext ChatSession::BuildToolCallContext(const Request& request, // Catalog metadata is immutable and may not contain markers for models whose // tokenizer defines them dynamically. Read those markers from the loaded GenAI // model without mutating the published ModelInfo. - const auto& tag_info = model_.GetTagInfo(); + const auto& tag_info = model.GetTagInfo(); if (tool_ctx.tool_call_start.empty()) { tool_ctx.tool_call_start = tag_info.bot_str; } @@ -971,7 +997,7 @@ ToolCallContext ChatSession::BuildToolCallContext(const Request& request, // ParseToolChoice rejects unknown values with FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT. auto tool_choice = SearchOptions::ParseToolChoice(request.options); if (!tool_choice.has_value()) { - tool_choice = session_options_.tool_choice; + tool_choice = session_options.tool_choice; } // User guidance is independent of generated tool guidance and must remain authoritative even when a malformed @@ -1000,11 +1026,238 @@ ToolCallContext ChatSession::BuildToolCallContext(const Request& request, } } - chat_session_internal::ApplyRawEnvelopeGuidance(tool_ctx, logger_); + chat_session_internal::ApplyRawEnvelopeGuidance(tool_ctx, logger); return tool_ctx; } +} // namespace + +struct PreparedChatRequest { + TranscriptIngest ingest; + MediaInput media; + ToolCallContext tool_context; + SearchOptions options; + ChatBackendKind backend_kind = ChatBackendKind::kGenerator; + std::string system_prompt; + std::vector reply_inputs; + std::optional messages; + PreparedChatPrompt prompt; + std::optional host_max_output_tokens; + std::string json_model_name; + bool json_passthrough = false; +}; + +namespace { + +void ResolvePreparedGenerationLimit(PreparedChatRequest& prepared, bool media_turn) { + if (chat_session_internal::ShouldEnforceHostOutputLimit(prepared.backend_kind, media_turn)) { + prepared.host_max_output_tokens = + ResolveMaxOutputTokens(prepared.options, GetDefaultMaxOutputTokens(media_turn)); + } +} + +std::unique_ptr PrepareChatRequest( + const Request& request, + const ChatTranscript& transcript, + const KeyValuePairs& base_session_options, + const SearchOptions& chat_session_options, + const std::vector& tool_definitions, + const ModelInfo& model_info, + GenAIModelInstance& model, + ILogger& logger, + const ChatMessagePreparer& message_preparer) { + auto prepared = std::make_unique(); + + for (const auto* item : request.items) { + if (item->type != FOUNDRY_LOCAL_ITEM_TEXT) { + continue; + } + + const auto& text_item = static_cast(*item); + if (text_item.text_type != FOUNDRY_LOCAL_TEXT_ITEM_TYPE_OPENAI_JSON) { + continue; + } + + auto request_json = nlohmann::json::parse(text_item.text); + auto chat_request = request_json.get(); + chat_completions::ApplyCatalogDefaults(chat_request, model_info.model_settings); + prepared->json_model_name = chat_request.model; + + Request internal_request; + internal_request.forced_tool_choice = request.forced_tool_choice; + internal_request.raw_envelope_descriptor = request.raw_envelope_descriptor; + chat_completions::BuildRequestItems(chat_request, internal_request); + if (internal_request.items.empty()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, + "the request has nothing to generate from: `messages` carried no content"); + } + + auto definitions = request.prepared_tool_definitions.has_value() + ? *request.prepared_tool_definitions + : chat_completions::ExtractToolDefinitions(chat_request, internal_request); + chat_completions::MapRequestParameters(chat_request, internal_request); + chat_completions::MapGuidance(chat_request, internal_request); + chat_completions::MapStopSequences(chat_request, internal_request); + + internal_request.options = MergeKeyValuePairs(request.options, internal_request.options); + + std::vector request_definitions; + if (request.prepared_tool_definitions.has_value()) { + if (!tool_definitions.empty()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, + "Tool definitions cannot be used with OpenAI JSON input; the JSON payload must be fully " + "self-contained"); + } + + request_definitions = std::move(definitions); + } else { + request_definitions = + chat_session_internal::BuildJsonRequestToolDefinitions(std::move(definitions), tool_definitions); + } + + if (!internal_request.raw_envelope_descriptor.has_value() && chat_request.metadata.has_value()) { + const auto descriptor = chat_request.metadata->find(tools::kRawEnvelopeMetadataKey); + if (descriptor != chat_request.metadata->end()) { + internal_request.raw_envelope_descriptor = tools::ParseRawEnvelopeDescriptor(descriptor->second); + } + } + if (internal_request.raw_envelope_descriptor.has_value()) { + tools::ValidateRawEnvelopeTool(*internal_request.raw_envelope_descriptor, request_definitions); + } + + prepared->tool_context = BuildToolCallContextForRequest(internal_request, request_definitions, + chat_session_options, model_info, model, logger); + prepared->options = + SearchOptions::FromParameters(MergeKeyValuePairs(base_session_options, internal_request.options)); + prepared->backend_kind = model.GetGenAIConfig().GetChatBackendKind(); + + auto messages = BuildTranscriptMessages(internal_request.items, prepared->tool_context.tool_kinds); + const ChatTranscript payload_transcript; + payload_transcript.ValidateInputs(messages); + prepared->reply_inputs = messages; + prepared->messages.emplace(message_preparer(std::move(messages), model.HasPositionalToolResults())); + prepared->prompt = PrepareTextChatPrompt(*prepared->messages, model, prepared->tool_context); + prepared->json_passthrough = true; + ResolvePreparedGenerationLimit(*prepared, /*media_turn=*/false); + return prepared; + } + + prepared->tool_context = BuildToolCallContextForRequest(request, tool_definitions, + chat_session_options, model_info, model, logger); + prepared->ingest = IngestRequestItems(request.items, request.item_segment_starts, + prepared->tool_context.tool_kinds); + transcript.ValidateInputs(prepared->ingest.messages); + prepared->media = CollectMediaInput(request); + const bool media_turn = !prepared->media.Empty(); + ValidateMediaTurn(prepared->media, prepared->ingest.messages, + {.session_has_history = !transcript.Empty(), + .tools_declared = prepared->tool_context.HasTools()}); + + const auto effective_kvp = MergeKeyValuePairs(base_session_options, request.options); + prepared->options = SearchOptions::FromParameters(effective_kvp); + prepared->backend_kind = model.GetGenAIConfig().GetChatBackendKind(); + prepared->system_prompt = GetOptionOrEmpty(effective_kvp, kSystemPromptOption); + + if (!TurnCanGenerate(prepared->ingest.messages, + {.media = media_turn, + .history = !transcript.Empty(), + .system_prefix = !prepared->system_prompt.empty()})) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, + "the request has nothing to generate from: no message content, no image or audio, no instructions, " + "and no conversation to continue"); + } + + std::vector all_messages; + const auto& committed = transcript.Messages(); + all_messages.reserve(committed.size() + prepared->ingest.messages.size() + + (prepared->system_prompt.empty() ? 0u : 1u)); + all_messages.insert(all_messages.end(), committed.begin(), committed.end()); + all_messages.insert(all_messages.end(), prepared->ingest.messages.begin(), prepared->ingest.messages.end()); + all_messages = WithSystemPrompt(prepared->system_prompt, std::move(all_messages)); + prepared->messages.emplace(message_preparer(std::move(all_messages), model.HasPositionalToolResults())); + + if (media_turn) { + auto media_messages = prepared->media.messages; + if (!prepared->system_prompt.empty()) { + media_messages.insert(media_messages.begin(), + MessageItem(FOUNDRY_LOCAL_ROLE_SYSTEM, prepared->system_prompt)); + } + + prepared->prompt = + PrepareMediaChatPrompt(media_messages, model, prepared->media.images, prepared->media.audios, + prepared->tool_context); + } else { + prepared->prompt = PrepareTextChatPrompt(*prepared->messages, model, prepared->tool_context); + } + + ResolvePreparedGenerationLimit(*prepared, media_turn); + return prepared; +} + +class ChatRequestPreflightOperation final : public Session::RequestPreflightOperation { + public: + ChatRequestPreflightOperation(Request request, + ChatTranscript transcript, + KeyValuePairs base_session_options, + SearchOptions chat_session_options, + std::vector tool_definitions, + ModelInfo model_info, + GenAIModelInstance& model, + ILogger& logger, + ChatMessagePreparer message_preparer) + : request_(std::move(request)), + transcript_(std::move(transcript)), + base_session_options_(std::move(base_session_options)), + chat_session_options_(std::move(chat_session_options)), + tool_definitions_(std::move(tool_definitions)), + model_info_(std::move(model_info)), + model_(model), + logger_(logger), + message_preparer_(std::move(message_preparer)) { + model_.AcquireSession(); + } + + ~ChatRequestPreflightOperation() override { + model_.ReleaseSession(); + } + + RequestBudget Execute() override { + if (!request_.has_value()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "request preflight operation has already been executed"); + } + + auto request = std::move(*request_); + request_.reset(); + auto prepared = PrepareChatRequest(request, transcript_, base_session_options_, chat_session_options_, + tool_definitions_, model_info_, model_, logger_, message_preparer_); + const auto context_limit = static_cast(GetModelMaxContextLength(model_.GetGenAIConfig())); + const auto output_reserve = + ResolveOutputReserve(prepared->options, prepared->backend_kind, !prepared->media.Empty(), + prepared->prompt.prompt_token_count, context_limit); + return ComputeRequestBudget(prepared->prompt.prompt_token_count, output_reserve, context_limit); + } + + private: + std::optional request_; + ChatTranscript transcript_; + KeyValuePairs base_session_options_; + SearchOptions chat_session_options_; + std::vector tool_definitions_; + ModelInfo model_info_; + GenAIModelInstance& model_; + ILogger& logger_; + ChatMessagePreparer message_preparer_; +}; + +} // namespace + +std::unique_ptr ChatSession::CreateRequestPreflightImpl(Request request) const { + return std::make_unique( + std::move(request), transcript_, SessionOptions(), session_options_, ToolDefinitions(), CatalogModel().Info(), + model_, logger_, message_preparer_); +} + void ChatSession::ProcessGeneratedOutput(std::vector events, const ToolCallContext& tool_ctx, const SearchOptions& effective_options, @@ -1079,88 +1332,29 @@ void ChatSession::ProcessGeneratedOutput(std::vector event } void ChatSession::ProcessRequestImpl(const Request& request, Response& response) { - // OpenAI chat completions JSON pass-through: a TEXT item tagged OPENAI_JSON. Routes to a separate handler that - // never uses the cached generator or the transcript (the JSON payload is self-contained). - for (const auto* item : request.items) { - if (item->type == FOUNDRY_LOCAL_ITEM_TEXT) { - const auto& text_item = static_cast(*item); + auto prepared = PrepareChatRequest(request, transcript_, SessionOptions(), session_options_, ToolDefinitions(), + CatalogModel().Info(), Model(), logger_, message_preparer_); - if (text_item.text_type == FOUNDRY_LOCAL_TEXT_ITEM_TYPE_OPENAI_JSON) { - ProcessChatCompletionsJson(text_item.text, request, response); - return; - } - } + if (prepared->json_passthrough) { + ProcessChatCompletionsJson(*prepared, request, response); + return; } - // One snapshot of the session's tools for this whole turn, taken before anything reads them. The registry is safe - // to mutate from another thread while a request generates, so taking it once is what makes a turn - // self-consistent: the calls it replays, the prompt it builds, and the calls it produces are all resolved against - // the same tool set. - // - // A turn appended to a cached generator does not rebuild the prompt and so keeps the kinds from the turn that - // did. That stays consistent because a turn carrying tool activity always invalidates the cached generator - // below, so any turn that actually replays a call is also the turn that rebuilds the prompt from this snapshot. - auto turn_tool_ctx = BuildToolCallContext(request, ToolDefinitions()); - - // Collect this turn's input messages locally — nothing reaches the transcript until the turn commits. Replay - // segment boundaries travel with the items so a reconstructed conversation regroups exactly as it was committed. - // - // Replayed tool calls are normalized with this turn's kinds: a custom tool's call comes back as the text payload - // this session handed out, so it must be rewrapped rather than rejected as malformed JSON. - auto ingest = IngestRequestItems(request.items, request.item_segment_starts, turn_tool_ctx.tool_kinds); + auto turn_tool_ctx = prepared->tool_context; + auto& ingest = prepared->ingest; auto inputs = std::move(ingest.messages); - - // Reject bad tool-call correlation before any generator work: a rejected turn must cost nothing and leave no - // state. This runs before the has-anything-to-say check below so a tool result that answers nothing is reported - // as the correlation error it is — an empty result for an unknown call ID is a broken conversation, not an - // empty request. - transcript_.ValidateInputs(inputs); - - // Media is single-shot. One conversation-scoped rule covers a live session and a conversation replayed into this - // request after the session cache dropped it, so a continuation is rejected identically either way. - auto media = CollectMediaInput(request); + auto& media = prepared->media; const bool media_turn = !media.Empty(); - ValidateMediaTurn(media, inputs, - {.session_has_history = !transcript_.Empty(), .tools_declared = turn_tool_ctx.HasTools()}); - - // Merge session-level and per-request options once for this turn. - auto effective_kvp = MergedOptions(request.options); - SearchOptions effective_options = SearchOptions::FromParameters(effective_kvp); - const ChatBackendKind backend_kind = Model().GetGenAIConfig().GetChatBackendKind(); - - // Request-scoped system prefix. It is not conversation history: it never enters the transcript, so it cannot - // accumulate a copy per turn, and the value this request carries is the only one used. It is baked into a - // generator's prompt, so a changed prefix has to rebuild while an unchanged one keeps the KV cache. - const std::string turn_system_prompt = GetOptionOrEmpty(effective_kvp, kSystemPromptOption); - - // One gate for every way a turn can carry meaning: its own messages, media bytes, the conversation behind it, or - // the instructions in front of it. A turn with none of them is the caller's mistake, so it is a client error — - // and a turn with any of them generates, whether the conversation is held in this session or was replayed into - // the request after the cache dropped it. - if (!TurnCanGenerate(inputs, {.media = media_turn, - .history = !transcript_.Empty(), - .system_prefix = !turn_system_prompt.empty()})) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, - "the request has nothing to generate from: no message content, no image or audio, no instructions, " - "and no conversation to continue"); - } + const auto& effective_options = prepared->options; + const auto backend_kind = prepared->backend_kind; + const auto& turn_system_prompt = prepared->system_prompt; + auto& prepared_messages = *prepared->messages; int prompt_tokens = 0; // Empty until this turn's input is appended to an existing generator. A rebuilt generator bakes the input into its // prompt, so there is no pre-turn boundary that undo could rewind back to. std::optional pre_turn_token_count; - // The complete authoritative prompt for this turn. Engine uses this to verify that its resident raw tokens are an - // exact prefix before reusing them; otherwise it replaces the conversation and submits the full prompt. - std::vector all_messages; - const auto& committed = transcript_.Messages(); - all_messages.reserve(committed.size() + inputs.size() + (turn_system_prompt.empty() ? 0u : 1u)); - all_messages.insert(all_messages.end(), committed.begin(), committed.end()); - all_messages.insert(all_messages.end(), inputs.begin(), inputs.end()); - all_messages = WithSystemPrompt(turn_system_prompt, std::move(all_messages)); - auto prepared_messages = - message_preparer_(std::move(all_messages), Model().HasPositionalToolResults()); - // Classic Generator cannot append tool exchanges or a changed prefix in isolation. Engine always renders the full // transcript and performs token-prefix reconciliation, so it can decide safely whether to reuse or replace state. const bool retained_tool_context_changed = @@ -1215,7 +1409,11 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response) if (cached_generator_) { pre_turn_token_count = cached_generator_->TokenCount(); try { - cached_generator_->AppendMessages(inputs, prepared_messages, Model(), turn_tool_ctx, effective_options); + if (text_generator_factory_) { + cached_generator_->AppendMessages(inputs, prepared_messages, Model(), turn_tool_ctx, effective_options); + } else { + cached_generator_->AppendPreparedPrompt(inputs, prepared->prompt, Model(), turn_tool_ctx, effective_options); + } prompt_tokens = cached_generator_->TokenCount(); cached_tool_ctx_ = turn_tool_ctx; } catch (const RetainedPromptMismatchError&) { @@ -1232,18 +1430,12 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response) std::unique_ptr generator; if (media_turn) { - // Media is single-shot: the generator is dropped after the turn because retained text state cannot reconstruct - // media bytes. The system prefix is projected into the media prompt but never enters the transcript. - auto media_messages = std::move(media.messages); - if (!turn_system_prompt.empty()) { - media_messages.insert(media_messages.begin(), MessageItem(FOUNDRY_LOCAL_ROLE_SYSTEM, turn_system_prompt)); - } - - generator = OnnxChatGenerator::CreateWithMedia(media_messages, effective_options, Model(), media.images, - media.audios, tool_ctx, /*use_full_context*/ false); + generator = OnnxChatGenerator::CreatePrepared(std::move(prepared->prompt), effective_options, Model(), tool_ctx, + /*use_full_context=*/false); } else { - generator = CreateTextChatGenerator(prepared_messages, effective_options, Model(), tool_ctx, - /*use_full_context=*/true, text_generator_factory_); + generator = CreatePreparedTextChatGenerator(prepared_messages, std::move(prepared->prompt), effective_options, + Model(), tool_ctx, /*use_full_context=*/true, + text_generator_factory_); } prompt_tokens = generator->PromptTokenCount(); @@ -1253,10 +1445,7 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response) system_prompt_ = turn_system_prompt; } - std::optional host_max_output; - if (chat_session_internal::ShouldEnforceHostOutputLimit(backend_kind, media_turn)) { - host_max_output = ResolveMaxOutputTokens(effective_options, GetDefaultMaxOutputTokens(media_turn)); - } + const auto host_max_output = prepared->host_max_output_tokens; // Generate token-by-token with optional streaming. // Check request.canceled each iteration — a streaming callback returning @@ -1551,99 +1740,19 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response) } } -void ChatSession::ProcessChatCompletionsJson(const std::string& request_json, const Request& original_request, +void ChatSession::ProcessChatCompletionsJson(PreparedChatRequest& prepared, const Request& original_request, Response& response) { - // Consult mutable session state exactly once. JSON requests otherwise derive their complete tool - // context from their own payload, so later registration cannot alter this request's interpretation. - const auto session_tool_definitions = ToolDefinitions(); - - // Parse the OpenAI chat completions request - auto req_json = nlohmann::json::parse(request_json); - auto req = req_json.get(); - - // Apply catalog defaults passed via request options - chat_completions::ApplyCatalogDefaults(req, CatalogModel().Info().model_settings); - - std::string model_name = req.model; + const auto& model_name = prepared.json_model_name; std::string completion_id = chat_completions::GenerateCompletionId(); auto now = std::chrono::system_clock::now(); int64_t created = std::chrono::duration_cast(now.time_since_epoch()).count(); - // Build the internal request from the chat completions request - Request internal_request; - internal_request.forced_tool_choice = original_request.forced_tool_choice; - internal_request.raw_envelope_descriptor = original_request.raw_envelope_descriptor; - - // We don't use history_ for this request as it's for backwards compat and all messages come from the input. - chat_completions::BuildRequestItems(req, internal_request); - if (internal_request.items.empty()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, - "the request has nothing to generate from: `messages` carried no content"); - } - - auto tool_definitions = original_request.prepared_tool_definitions.has_value() - ? *original_request.prepared_tool_definitions - : chat_completions::ExtractToolDefinitions(req, internal_request); - chat_completions::MapRequestParameters(req, internal_request); - chat_completions::MapGuidance(req, internal_request); - chat_completions::MapStopSequences(req, internal_request); - - // Merge options from the original request (e.g. tool_call_start/end from model properties) - for (const auto& [key, value] : original_request.options) { - if (internal_request.options.find(key) == internal_request.options.end()) { - internal_request.options[key] = value; - } - } - - // Build request-local definitions. They never enter the session registry: this payload is - // self-contained and concurrent registration must not alter either its prompt or call parsing. - std::vector request_tool_definitions; - if (original_request.prepared_tool_definitions.has_value()) { - if (!session_tool_definitions.empty()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, - "Tool definitions cannot be used with OpenAI JSON input; " - "the JSON payload must be fully self-contained"); - } - - request_tool_definitions = std::move(tool_definitions); - } else { - request_tool_definitions = chat_session_internal::BuildJsonRequestToolDefinitions( - std::move(tool_definitions), session_tool_definitions); - } - - if (!internal_request.raw_envelope_descriptor.has_value() && req.metadata.has_value()) { - const auto descriptor = req.metadata->find(tools::kRawEnvelopeMetadataKey); - if (descriptor != req.metadata->end()) { - internal_request.raw_envelope_descriptor = - tools::ParseRawEnvelopeDescriptor(descriptor->second); - } - } - if (internal_request.raw_envelope_descriptor.has_value()) { - tools::ValidateRawEnvelopeTool(*internal_request.raw_envelope_descriptor, request_tool_definitions); - } - - const auto tool_ctx = BuildToolCallContext(internal_request, request_tool_definitions); - - // Merge session-level and per-request options once. - auto effective_kvp = MergedOptions(internal_request.options); - SearchOptions options = SearchOptions::FromParameters(effective_kvp); - - // Collect transcript messages from the internal request for the generator. - // The session transcript is not used here — everything comes from the parsed JSON input. - // The context above already snapshotted the kinds that shape this prompt, so replayed calls in the payload are - // normalized with exactly the kinds the produced calls are read back with. - auto messages = BuildTranscriptMessages(internal_request.items, tool_ctx.tool_kinds); - - // The payload is self-contained, so correlate its tool calls and results against an empty transcript. This gives - // the same stable errors a session turn would produce for a history the model cannot interpret. - const ChatTranscript payload_transcript; - payload_transcript.ValidateInputs(messages); - - // Create generator - const auto prepared_messages = - chat_internal::PrepareChatMessages(messages, Model().HasPositionalToolResults()); - auto generator = CreateTextChatGenerator(prepared_messages, options, Model(), tool_ctx, - /*use_full_context=*/false, text_generator_factory_); + const auto& tool_ctx = prepared.tool_context; + const auto& options = prepared.options; + const auto& prepared_messages = *prepared.messages; + auto generator = + CreatePreparedTextChatGenerator(prepared_messages, std::move(prepared.prompt), options, Model(), tool_ctx, + /*use_full_context=*/false, text_generator_factory_); int prompt_tokens = generator->PromptTokenCount(); auto streaming_callback = CreateCallbackHandler(original_request); @@ -1686,7 +1795,7 @@ void ChatSession::ProcessChatCompletionsJson(const std::string& request_json, co // A trailing assistant message is a prefill. If it already carries a tool call, visible text generated after it // would be merged into that same turn when the client replays the response, so seed the same ordering guard used // by the stateful path. - auto turn_guard = AssistantTurnGuard::ForReplyTo(messages, 0); + auto turn_guard = AssistantTurnGuard::ForReplyTo(prepared.reply_inputs, 0); // Use the same typed segments for streaming and final response construction. auto splitter = CreateReasoningSplitter(tool_ctx, Model(), generator->PromptOpensReasoning()); diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h index 5b2050b61..165e02b78 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h @@ -27,6 +27,7 @@ namespace fl { class GenAIModelInstance; class ChatGenerator; +struct PreparedChatRequest; enum class BackendTerminationCause; using TextChatGeneratorFactory = @@ -203,12 +204,7 @@ class ChatSession : public Session { /// Process a request: extracts items and parameters from the generic request, generates a response, and on /// success commits the turn to the transcript. void ProcessRequestImpl(const Request& request, Response& response) override; - - /// Build tool calling context from request parameters and a snapshot of the session's tool definitions. - /// - /// The snapshot is supplied by the caller rather than read here so that one turn resolves its replayed calls, its - /// prompt, and its produced calls against the same tool set. - ToolCallContext BuildToolCallContext(const Request& request, const std::vector& definitions) const; + std::unique_ptr CreateRequestPreflightImpl(Request request) const override; /// Build final response items from the typed segments and tool calls produced during generation. void ProcessGeneratedOutput(std::vector events, @@ -227,8 +223,7 @@ class ChatSession : public Session { /// request. Parses the JSON, converts to internal items, runs generation, and produces an OPENAI_JSON-tagged /// TextItem response with the OpenAI ChatCompletionResponse. /// Does not use or update the transcript or the cached generator. - void ProcessChatCompletionsJson(const std::string& request_json, const Request& original_request, - Response& response); + void ProcessChatCompletionsJson(PreparedChatRequest& prepared, const Request& original_request, Response& response); /// Drop the cached generator and its tool context. Called whenever the generator's KV cache can no longer be /// trusted to match the committed transcript — the next turn then rebuilds from full committed history. diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.cc b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.cc index 3fe519b78..c2af626ed 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.cc @@ -16,24 +16,6 @@ namespace fl { -namespace { - -/// Probe whether the rendered prompt leaves a reasoning block open. Uses the encoded prompt token IDs when they are -/// available (the text path) and falls back to the rendered text for the media path, which has no encoded sequence -/// of its own. -bool DetectPromptOpensReasoning(const std::string& prompt, - const OgaSequences* sequences, - const ReasoningMarkers& markers) { - std::span prompt_token_ids; - if (sequences != nullptr && sequences->Count() > 0) { - prompt_token_ids = {sequences->SequenceData(0), sequences->SequenceCount(0)}; - } - - return PromptOpensReasoning(prompt_token_ids, markers, prompt); -} - -} // namespace - namespace onnx_chat_generator_internal { TurnTermination ClassifyTurnTermination(bool cancelled, @@ -260,17 +242,22 @@ int OnnxChatGenerator::AppendMessages(const std::vector& new_ GenAIModelInstance& model, const ToolCallContext& tool_ctx, const SearchOptions& options) { - if (new_messages.empty() || full_messages.Empty()) { + auto prepared = PrepareTextChatPrompt(full_messages, model, tool_ctx); + return AppendPreparedPrompt(new_messages, prepared, model, tool_ctx, options); +} + +int OnnxChatGenerator::AppendPreparedPrompt(const std::vector& new_messages, + const PreparedChatPrompt& prepared, + GenAIModelInstance&, + const ToolCallContext&, + const SearchOptions& options) { + if (new_messages.empty() || prepared.token_ids.empty()) { FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "new_messages and full_messages must not be empty"); } // Render and tokenize the authoritative full transcript once. Generated text is not guaranteed to round-trip // through decode/encode to the same token IDs, so resident state is reusable only when it is an exact prefix. - std::string prompt = BuildChatPrompt(full_messages, model, tool_ctx.tools_json); - auto full_sequences = EncodePrompt(prompt, model); - const auto full_count = full_sequences->SequenceCount(0); - const auto* full_data = full_sequences->SequenceData(0); - const std::span full_prompt(full_data, full_count); + const std::span full_prompt(prepared.token_ids); const std::span resident(generator_->GetSequenceData(0), generator_->GetSequenceCount(0)); const auto suffix_start = chat_internal::FindUnmatchedPromptSuffix(resident, full_prompt); if (!suffix_start.has_value()) { @@ -294,8 +281,8 @@ int OnnxChatGenerator::AppendMessages(const std::vector& new_ // Re-probe: the appended segment ends with this turn's assistant generation prefix, so it — not the original // prompt — determines whether generation resumes inside a template-opened reasoning block. - prompt_opens_reasoning_ = DetectPromptOpensReasoning(prompt, full_sequences.get(), reasoning_markers_); - prompt_token_count_ = static_cast(full_count); + prompt_opens_reasoning_ = fl::PromptOpensReasoning(full_prompt, reasoning_markers_, prepared.prompt); + prompt_token_count_ = static_cast(prepared.prompt_token_count); turn_start_token_count_ = TokenCount(); max_output_tokens_ = ResolveMaxOutputTokens(options); ResetTurnState(); @@ -391,12 +378,8 @@ std::unique_ptr OnnxChatGenerator::Create( GenAIModelInstance& model, const ToolCallContext& tool_ctx, bool use_full_context) { - if (messages.Empty()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "messages must not be empty"); - } - - std::string prompt = BuildChatPrompt(messages, model, tool_ctx.tools_json); - return CreateImpl(prompt, options, model, tool_ctx, use_full_context, /*images=*/{}, /*audios=*/{}); + return CreatePrepared(PrepareTextChatPrompt(messages, model, tool_ctx), options, model, tool_ctx, + use_full_context); } std::unique_ptr OnnxChatGenerator::CreateWithMedia( @@ -407,152 +390,48 @@ std::unique_ptr OnnxChatGenerator::CreateWithMedia( const std::vector& audios, const ToolCallContext& tool_ctx, bool use_full_context) { - if (messages.empty()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "messages must not be empty"); - } - - if (images.empty() && audios.empty()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, - "CreateWithMedia requires at least one image or audio input"); - } - - if (!model.IsMultiModal()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "image or audio input requires a multimodal model"); - } - - if (!model.GetPreprocessor().HasMultiModalProcessor()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "model has no multimodal processor available for media input"); - } - - std::string messages_json = TransformMessagesForMedia(messages); - const char* tools_ptr = tool_ctx.tools_json.empty() ? nullptr : tool_ctx.tools_json.c_str(); - std::string prompt = - model.GetPreprocessor().ApplyChatTemplate(messages_json.c_str(), tools_ptr, /*add_generation_prompt=*/true); - - return CreateImpl(prompt, options, model, tool_ctx, use_full_context, images, audios); + return CreatePrepared(PrepareMediaChatPrompt(messages, model, images, audios, tool_ctx), options, model, tool_ctx, + use_full_context); } -std::unique_ptr OnnxChatGenerator::CreateImpl(const std::string& prompt, - const SearchOptions& options, - GenAIModelInstance& model, - const ToolCallContext& tool_ctx, - bool use_full_context, - const std::vector& images, - const std::vector& audios) { - const bool media_branch = !images.empty() || !audios.empty(); - - // 1. Token budgeting. - // Text path: encode the prompt up front so we know its token count. - // Media path: process inputs first and read the expanded input_ids shape. - std::unique_ptr sequences; - int input_token_count = 0; - - if (!media_branch) { - sequences = EncodePrompt(prompt, model); - input_token_count = static_cast(sequences->SequenceCount(0)); - } - - // Process media before sizing the generator so max_length includes the - // exact token expansion produced by the multimodal processor. - std::unique_ptr named_tensors; - if (media_branch) { - std::unique_ptr oga_images; - if (!images.empty()) { - std::vector> image_bytes; - image_bytes.reserve(images.size()); - std::vector buffers; - buffers.reserve(images.size()); - std::vector sizes; - sizes.reserve(images.size()); - - for (const auto* img : images) { - if (img == nullptr) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "image entry must not be null"); - } - - image_bytes.push_back(img->ReadBytes()); - buffers.push_back(image_bytes.back().data()); - sizes.push_back(image_bytes.back().size()); - } - - oga_images = OgaImages::Load(buffers.data(), sizes.data(), buffers.size()); - } - - std::unique_ptr oga_audios; - if (!audios.empty()) { - std::vector audio_buffers; - audio_buffers.reserve(audios.size()); - std::vector audio_sizes; - audio_sizes.reserve(audios.size()); - for (const auto* audio : audios) { - if (audio == nullptr || audio->data == nullptr || audio->data_size == 0) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "audio entry must contain bytes"); - } - audio_buffers.push_back(audio->data); - audio_sizes.push_back(audio->data_size); - } - - oga_audios = OgaAudios::Load(audio_buffers.data(), audio_sizes.data(), audio_buffers.size()); - } - - named_tensors = model.GetPreprocessor().ProcessMedia(prompt.c_str(), oga_images.get(), oga_audios.get()); - auto input_ids = named_tensors->Get("input_ids"); - auto input_shape = input_ids->Shape(); - if (input_shape.empty() || input_shape.back() <= 0) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "multimodal processor returned invalid input_ids"); - } - input_token_count = static_cast(input_shape.back()); - } - - // 2. Create GeneratorParams from the model +std::unique_ptr OnnxChatGenerator::CreatePrepared( + PreparedChatPrompt prepared, + const SearchOptions& options, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx, + bool use_full_context) { + const int input_token_count = static_cast(prepared.prompt_token_count); auto gen_params = OgaGeneratorParams::Create(model.GetOgaModel()); - - // 3. Apply search options (temperature, top_p, max_length, etc.) and validate token budget. - // Media inputs use a larger default because preprocessing expands them into tokens. - const int default_max_output_tokens = GetDefaultMaxOutputTokens(media_branch); + const int default_max_output_tokens = GetDefaultMaxOutputTokens(prepared.HasMedia()); const int max_length = ApplySearchOptions(options, input_token_count, model.GetGenAIConfig(), *gen_params, model.EP(), use_full_context, default_max_output_tokens); const int max_output_tokens = ResolveMaxOutputTokens(options, default_max_output_tokens); - // 4. Build guidance from the actual rendered prompt state, then reuse that state to seed stream reasoning. auto reasoning_markers = ResolveReasoningMarkers(tool_ctx, model); - const bool prompt_opens_reasoning = DetectPromptOpensReasoning(prompt, sequences.get(), reasoning_markers); + const bool prompt_opens_reasoning = + fl::PromptOpensReasoning(prepared.token_ids, reasoning_markers, prepared.prompt); ApplyGuidanceOptions(tool_ctx, prompt_opens_reasoning, *gen_params); - // 5. Create the Generator and feed it the prompt. - // Text path: append the encoded token sequences. - // Media path: process inputs via OgaMultiModalProcessor and feed the - // resulting named tensors via SetInputs (which extracts input_ids and - // appends them internally — do NOT also call AppendTokenSequences). std::unique_ptr generator; try { generator = OgaGenerator::Create(model.GetOgaModel(), *gen_params); - - if (media_branch) { - generator->SetInputs(*named_tensors); + if (prepared.HasMedia()) { + generator->SetInputs(*prepared.media_tensors); } else { + auto sequences = OgaSequences::Create(); + sequences->Append(prepared.token_ids.data(), prepared.token_ids.size()); generator->AppendTokenSequences(*sequences); } } catch (const std::runtime_error& e) { FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, std::string("failed to create generator: ") + e.what()); } - // 6. Create tokenizer stream (single-decode path). auto stream = model.GetPreprocessor().CreateTokenizerStream(); - - // `std::make_unique` constructs inside the library helper, which does not have - // access to this class's private constructor. - return std::unique_ptr(new OnnxChatGenerator(std::move(gen_params), - std::move(generator), - std::move(stream), - model, - input_token_count, - max_length, - max_output_tokens, - std::move(reasoning_markers), - prompt_opens_reasoning, - std::move(named_tensors))); + return std::unique_ptr( + new OnnxChatGenerator(std::move(gen_params), std::move(generator), std::move(stream), model, + input_token_count, max_length, max_output_tokens, std::move(reasoning_markers), + prompt_opens_reasoning, std::move(prepared.media_tensors))); } } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.h b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.h index 70723ffb9..933afd195 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.h @@ -4,6 +4,7 @@ #include "inferencing/generative/chat/chat_generator.h" #include "inferencing/generative/chat/chat_template.h" +#include "inferencing/generative/chat/prepared_chat_prompt.h" #include "inferencing/generative/chat/reasoning_stream_splitter.h" #include "inferencing/generative/chat/search_options.h" #include "inferencing/generative/toolcalling/tool_call_context.h" @@ -79,6 +80,11 @@ class OnnxChatGenerator : public ChatGenerator { GenAIModelInstance& model, const ToolCallContext& tool_ctx, const SearchOptions& options) override; + int AppendPreparedPrompt(const std::vector& new_messages, + const PreparedChatPrompt& prompt, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx, + const SearchOptions& options) override; int AppendMessages(const std::vector& new_messages, const std::vector& full_messages, GenAIModelInstance& model, @@ -126,6 +132,12 @@ class OnnxChatGenerator : public ChatGenerator { const ToolCallContext& tool_ctx = {}, bool use_full_context = false); + static std::unique_ptr CreatePrepared(PreparedChatPrompt prepared, + const SearchOptions& options, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx = {}, + bool use_full_context = false); + // ---- Static helpers exposed for unit testing ---- /// Build the JSON messages array fed to OgaTokenizer::ApplyChatTemplate when @@ -149,18 +161,6 @@ class OnnxChatGenerator : public ChatGenerator { bool prompt_opens_reasoning, std::unique_ptr named_tensors = nullptr); - // Shared implementation for text and media creation paths. The caller supplies the fully rendered prompt because - // the two paths project their messages differently: text builds from the transcript, media rewrites MessageItems - // so the template inserts media sentinels. Both share search-options validation, guidance setup, media tensor - // preparation, and generator construction. - static std::unique_ptr CreateImpl(const std::string& prompt, - const SearchOptions& options, - GenAIModelInstance& model, - const ToolCallContext& tool_ctx, - bool use_full_context, - const std::vector& images, - const std::vector& audios); - void ResetTurnState(); std::unique_ptr gen_params_; diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.cc b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.cc index 28a9a44e4..bf75531e9 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.cc @@ -57,18 +57,6 @@ std::optional MapTerminationCause(uint32_t reason) { } // namespace onnx_engine_chat_stream_internal -namespace { - -bool DetectPromptOpensReasoning(const std::string& prompt, - const OgaSequences& sequences, - const ToolCallContext& tool_ctx, - GenAIModelInstance& model) { - const std::span token_ids(sequences.SequenceData(0), sequences.SequenceCount(0)); - return PromptOpensReasoning(token_ids, ResolveReasoningMarkers(tool_ctx, model), prompt); -} - -} // namespace - OnnxEngineChatStream::OnnxEngineChatStream( OnnxChatEngine& engine, std::shared_ptr conversation, @@ -179,18 +167,25 @@ int OnnxEngineChatStream::AppendMessages(const std::vector& n GenAIModelInstance& model, const ToolCallContext& tool_ctx, const SearchOptions& options) { - if (new_messages.empty() || full_messages.Empty()) { + const auto prepared = PrepareTextChatPrompt(full_messages, model, tool_ctx); + return AppendPreparedPrompt(new_messages, prepared, model, tool_ctx, options); +} + +int OnnxEngineChatStream::AppendPreparedPrompt(const std::vector& new_messages, + const PreparedChatPrompt& prepared, + GenAIModelInstance&, + const ToolCallContext& tool_ctx, + const SearchOptions& options) { + if (new_messages.empty() || prepared.token_ids.empty()) { FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "new_messages and full_messages must not be empty"); } - auto prompt = BuildChatPrompt(full_messages, model, tool_ctx.tools_json); - auto sequences = EncodePrompt(prompt, model); - const int count = static_cast(sequences->SequenceCount(0)); - const auto* data = sequences->SequenceData(0); - const std::span full_prompt(data, static_cast(count)); + const int count = static_cast(prepared.prompt_token_count); + const std::span full_prompt(prepared.token_ids); const auto resident_tokens = engine_.ResidentTokens(conversation_); const auto suffix_start = chat_internal::FindUnmatchedPromptSuffix(resident_tokens, full_prompt); - const bool prompt_opens_reasoning = DetectPromptOpensReasoning(prompt, *sequences, tool_ctx, model); + const bool prompt_opens_reasoning = + fl::PromptOpensReasoning(full_prompt, ResolveReasoningMarkers(tool_ctx, model_), prepared.prompt); // Keep the previous decoder intact if admission fails. Once admitted, start a fresh stream so partial UTF-8/BPE // state from the prior turn cannot affect generated tokens; prompt tokens are never decoded. @@ -258,8 +253,16 @@ std::unique_ptr OnnxEngineChatStream::Create( const SearchOptions& options, GenAIModelInstance& model, const ToolCallContext& tool_ctx) { - if (messages.Empty()) { - FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "messages must not be empty"); + return CreatePrepared(PrepareTextChatPrompt(messages, model, tool_ctx), options, model, tool_ctx); +} + +std::unique_ptr OnnxEngineChatStream::CreatePrepared( + PreparedChatPrompt prepared, + const SearchOptions& options, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx) { + if (prepared.token_ids.empty()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "prepared text prompt must not be empty"); } auto* engine = model.GetChatEngine(); @@ -267,20 +270,16 @@ std::unique_ptr OnnxEngineChatStream::Create( FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "model does not own a chat Engine"); } - auto prompt = BuildChatPrompt(messages, model, tool_ctx.tools_json); - auto sequences = EncodePrompt(prompt, model); - const int prompt_token_count = static_cast(sequences->SequenceCount(0)); - const bool prompt_opens_reasoning = DetectPromptOpensReasoning(prompt, *sequences, tool_ctx, model); + const int prompt_token_count = static_cast(prepared.prompt_token_count); + const bool prompt_opens_reasoning = + fl::PromptOpensReasoning(prepared.token_ids, ResolveReasoningMarkers(tool_ctx, model), prepared.prompt); auto stream = model.GetPreprocessor().CreateTokenizerStream(); auto conversation = engine->CreateConversation(options, tool_ctx, prompt_token_count); try { - const auto* data = sequences->SequenceData(0); - engine->BeginTurn(conversation, std::span(data, static_cast(prompt_token_count)), options, - tool_ctx, prompt_opens_reasoning); + engine->BeginTurn(conversation, prepared.token_ids, options, tool_ctx, prompt_opens_reasoning); auto result = std::unique_ptr( - new OnnxEngineChatStream(*engine, std::move(conversation), std::move(stream), model, - prompt_token_count)); + new OnnxEngineChatStream(*engine, std::move(conversation), std::move(stream), model, prompt_token_count)); result->prompt_opens_reasoning_ = prompt_opens_reasoning; return result; } catch (...) { diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.h b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.h index e2774a9ea..883b88bd4 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_stream.h @@ -4,6 +4,7 @@ #include "inferencing/generative/chat/chat_generator.h" #include "inferencing/generative/chat/onnx_chat_engine.h" +#include "inferencing/generative/chat/prepared_chat_prompt.h" #include "inferencing/generative/chat/search_options.h" #include "inferencing/generative/toolcalling/tool_call_context.h" @@ -41,6 +42,11 @@ class OnnxEngineChatStream final : public ChatGenerator { GenAIModelInstance& model, const ToolCallContext& tool_ctx, const SearchOptions& options) override; + int AppendPreparedPrompt(const std::vector& new_messages, + const PreparedChatPrompt& prompt, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx, + const SearchOptions& options) override; int AppendMessages(const std::vector& new_messages, const std::vector& full_messages, GenAIModelInstance& model, @@ -54,6 +60,10 @@ class OnnxEngineChatStream final : public ChatGenerator { const SearchOptions& options, GenAIModelInstance& model, const ToolCallContext& tool_ctx); + static std::unique_ptr CreatePrepared(PreparedChatPrompt prepared, + const SearchOptions& options, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx); static std::unique_ptr Create( const chat_internal::PreparedChatMessages& messages, const SearchOptions& options, diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/prepared_chat_prompt.cc b/sdk_v2/cpp/src/inferencing/generative/chat/prepared_chat_prompt.cc new file mode 100644 index 000000000..26f331678 --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/generative/chat/prepared_chat_prompt.cc @@ -0,0 +1,115 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "inferencing/generative/chat/prepared_chat_prompt.h" + +#include "exception.h" +#include "inferencing/generative/chat/onnx_chat_generator.h" +#include "inferencing/generative/genai_model_instance.h" +#include "inferencing/generative/toolcalling/tool_call_context.h" +#include "items/audio_item.h" +#include "items/image_item.h" + +#include + +namespace fl { + +PreparedChatPrompt::PreparedChatPrompt() = default; +PreparedChatPrompt::~PreparedChatPrompt() = default; +PreparedChatPrompt::PreparedChatPrompt(PreparedChatPrompt&&) noexcept = default; +PreparedChatPrompt& PreparedChatPrompt::operator=(PreparedChatPrompt&&) noexcept = default; + +PreparedChatPrompt PrepareTextChatPrompt(const chat_internal::PreparedChatMessages& messages, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx) { + if (messages.Empty()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "messages must not be empty"); + } + + PreparedChatPrompt prepared; + prepared.prompt = BuildChatPrompt(messages, model, tool_ctx.tools_json); + auto sequences = EncodePrompt(prepared.prompt, model); + const auto count = sequences->SequenceCount(0); + const auto* data = sequences->SequenceData(0); + prepared.token_ids.assign(data, data + count); + prepared.prompt_token_count = static_cast(count); + return prepared; +} + +PreparedChatPrompt PrepareMediaChatPrompt(const std::vector& messages, + GenAIModelInstance& model, + const std::vector& images, + const std::vector& audios, + const ToolCallContext& tool_ctx) { + if (messages.empty()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "messages must not be empty"); + } + if (images.empty() && audios.empty()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "media preparation requires image or audio input"); + } + if (!model.IsMultiModal()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "image or audio input requires a multimodal model"); + } + if (!model.GetPreprocessor().HasMultiModalProcessor()) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "model has no multimodal processor available for media input"); + } + + PreparedChatPrompt prepared; + const auto messages_json = OnnxChatGenerator::TransformMessagesForMedia(messages); + const char* tools = tool_ctx.tools_json.empty() ? nullptr : tool_ctx.tools_json.c_str(); + prepared.prompt = model.GetPreprocessor().ApplyChatTemplate(messages_json.c_str(), tools, true); + + std::unique_ptr oga_images; + std::vector> image_bytes; + if (!images.empty()) { + image_bytes.reserve(images.size()); + std::vector buffers; + std::vector sizes; + buffers.reserve(images.size()); + sizes.reserve(images.size()); + + for (const auto* image : images) { + if (image == nullptr) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "image entry must not be null"); + } + + image_bytes.push_back(image->ReadBytes()); + buffers.push_back(image_bytes.back().data()); + sizes.push_back(image_bytes.back().size()); + } + + oga_images = OgaImages::Load(buffers.data(), sizes.data(), buffers.size()); + } + + std::unique_ptr oga_audios; + if (!audios.empty()) { + std::vector buffers; + std::vector sizes; + buffers.reserve(audios.size()); + sizes.reserve(audios.size()); + + for (const auto* audio : audios) { + if (audio == nullptr || audio->data == nullptr || audio->data_size == 0) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "audio entry must contain bytes"); + } + + buffers.push_back(audio->data); + sizes.push_back(audio->data_size); + } + + oga_audios = OgaAudios::Load(buffers.data(), sizes.data(), buffers.size()); + } + + prepared.media_tensors = + model.GetPreprocessor().ProcessMedia(prepared.prompt.c_str(), oga_images.get(), oga_audios.get()); + const auto input_ids = prepared.media_tensors->Get("input_ids"); + const auto shape = input_ids->Shape(); + if (shape.empty() || shape.back() <= 0) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "multimodal processor returned invalid input_ids"); + } + + prepared.prompt_token_count = shape.back(); + return prepared; +} + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/prepared_chat_prompt.h b/sdk_v2/cpp/src/inferencing/generative/chat/prepared_chat_prompt.h new file mode 100644 index 000000000..1660a44de --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/generative/chat/prepared_chat_prompt.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include "inferencing/generative/chat/chat_template.h" + +#include +#include +#include +#include + +struct OgaNamedTensors; + +namespace fl { + +class AudioItem; +class GenAIModelInstance; +class ImageItem; +class MessageItem; +struct ToolCallContext; + +/// Fully rendered and tokenized model input. Constructing this value performs no generation or conversation mutation. +struct PreparedChatPrompt { + PreparedChatPrompt(); + ~PreparedChatPrompt(); + PreparedChatPrompt(PreparedChatPrompt&&) noexcept; + PreparedChatPrompt& operator=(PreparedChatPrompt&&) noexcept; + + PreparedChatPrompt(const PreparedChatPrompt&) = delete; + PreparedChatPrompt& operator=(const PreparedChatPrompt&) = delete; + + std::string prompt; + std::vector token_ids; + std::unique_ptr media_tensors; + int64_t prompt_token_count = 0; + + bool HasMedia() const noexcept { + return media_tensors != nullptr; + } +}; + +PreparedChatPrompt PrepareTextChatPrompt(const chat_internal::PreparedChatMessages& messages, + GenAIModelInstance& model, + const ToolCallContext& tool_ctx); + +PreparedChatPrompt PrepareMediaChatPrompt(const std::vector& messages, + GenAIModelInstance& model, + const std::vector& images, + const std::vector& audios, + const ToolCallContext& tool_ctx); + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc index fbdfee0c5..5a9445a82 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc @@ -11,6 +11,7 @@ #include #include +#include namespace fl { @@ -70,6 +71,53 @@ int GetModelMaxContextLength(const GenAIConfig& config) { return model_max_length; } +RequestBudget ComputeRequestBudget(int64_t prompt_tokens, + int64_t output_reserve_tokens, + int64_t context_limit_tokens) { + if (prompt_tokens < 0) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "prompt token count must not be negative"); + } + + if (output_reserve_tokens < 0) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "output token reserve must not be negative"); + } + + if (context_limit_tokens < 1) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "model context length must be a positive integer"); + } + + if (prompt_tokens > (std::numeric_limits::max)() - output_reserve_tokens) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "request token budget overflow"); + } + + const auto required_tokens = prompt_tokens + output_reserve_tokens; + const bool fits = required_tokens <= context_limit_tokens; + return { + .prompt_tokens = prompt_tokens, + .output_reserve_tokens = output_reserve_tokens, + .required_tokens = required_tokens, + .context_limit_tokens = context_limit_tokens, + .fits = fits, + .deficit_tokens = fits ? 0 : required_tokens - context_limit_tokens, + }; +} + +int64_t ResolveOutputReserve(const SearchOptions& options, + ChatBackendKind backend_kind, + bool has_media, + int64_t prompt_tokens, + int64_t context_limit_tokens) { + if (options.max_output_tokens.has_value()) { + return ResolveMaxOutputTokens(options, GetDefaultMaxOutputTokens(has_media)); + } + + if (backend_kind == ChatBackendKind::kEngine && !has_media) { + return std::max(0, context_limit_tokens - prompt_tokens); + } + + return GetDefaultMaxOutputTokens(has_media); +} + std::optional ResolveTurnGuidanceOptions(const ToolCallContext& tool_ctx, bool prompt_opens_reasoning) { if (tool_ctx.guidance_disabled) { diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h index 366adfecb..4732e2d4b 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/search_options.h @@ -91,6 +91,15 @@ struct SearchOptions { bool HasSameRetainedGenerationSettings(const SearchOptions& other, ChatBackendKind backend_kind) const; }; +struct RequestBudget { + int64_t prompt_tokens = 0; + int64_t output_reserve_tokens = 0; + int64_t required_tokens = 0; + int64_t context_limit_tokens = 0; + bool fits = false; + int64_t deficit_tokens = 0; +}; + inline constexpr int kDefaultChatTextMaxOutputTokens = 2048; inline constexpr int kDefaultChatMediaMaxOutputTokens = 3072; @@ -103,9 +112,21 @@ constexpr int GetDefaultMaxOutputTokens(bool has_media) noexcept { int ResolveMaxOutputTokens(const SearchOptions& options, int default_max_output_tokens = kDefaultChatTextMaxOutputTokens); -/// Return the model's total context window from genai_config.json. +/// Return the model's total context window from search.max_length. int GetModelMaxContextLength(const GenAIConfig& config); +/// Compute the checked prompt plus output token budget. +RequestBudget ComputeRequestBudget(int64_t prompt_tokens, + int64_t output_reserve_tokens, + int64_t context_limit_tokens); + +/// Resolve the output reserve using the same explicit limits and backend defaults as generation. +int64_t ResolveOutputReserve(const SearchOptions& options, + ChatBackendKind backend_kind, + bool has_media, + int64_t prompt_tokens, + int64_t context_limit_tokens); + /// Resolve the guidance configuration that should apply to a single tool-only turn. std::optional ResolveTurnGuidanceOptions(const ToolCallContext& tool_ctx, bool prompt_opens_reasoning); diff --git a/sdk_v2/cpp/src/inferencing/session/request.cc b/sdk_v2/cpp/src/inferencing/session/request.cc new file mode 100644 index 000000000..4686ce63a --- /dev/null +++ b/sdk_v2/cpp/src/inferencing/session/request.cc @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "inferencing/session/request.h" + +#include "exception.h" +#include "items/audio_item.h" +#include "items/image_item.h" +#include "items/message_item.h" +#include "items/text_item.h" +#include "items/tool_call_item.h" +#include "items/tool_result_item.h" + +#include + +namespace fl { +namespace { + +std::unique_ptr CloneChatItem(const Item& item) { + switch (item.type) { + case FOUNDRY_LOCAL_ITEM_TEXT: { + const auto& text = static_cast(item); + return std::make_unique(text.text, text.text_type); + } + case FOUNDRY_LOCAL_ITEM_MESSAGE: + return std::make_unique(static_cast(item)); + case FOUNDRY_LOCAL_ITEM_TOOL_CALL: { + const auto& call = static_cast(item); + auto clone = std::make_unique(call.call_id, call.name, call.arguments, call.replayed_from_store, + call.kind, call.declared_kind, call.generated_encoding); + clone->replayed_arguments = call.replayed_arguments; + clone->replayed_kind = call.replayed_kind; + return clone; + } + case FOUNDRY_LOCAL_ITEM_TOOL_RESULT: { + const auto& result = static_cast(item); + return std::make_unique(result.call_id, result.result); + } + case FOUNDRY_LOCAL_ITEM_IMAGE: { + const auto& image = static_cast(item); + if (image.data != nullptr && image.data_size > 0) { + const auto* begin = static_cast(image.data); + auto clone = std::make_unique(std::vector(begin, begin + image.data_size), image.format); + clone->uri = image.uri; + return clone; + } + + return std::make_unique(image.uri, image.format); + } + case FOUNDRY_LOCAL_ITEM_AUDIO: { + const auto& audio = static_cast(item); + if (audio.data == nullptr || audio.data_size == 0) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, + "request preflight audio input must contain bytes"); + } + + const auto* begin = static_cast(audio.data); + auto clone = std::make_unique( + std::vector(begin, begin + audio.data_size), audio.format); + clone->sample_rate = audio.sample_rate; + clone->channels = audio.channels; + return clone; + } + default: + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, + "request preflight does not support " + std::string(Item::TypeName(item.type)) + " items"); + } +} + +} // namespace + +Request Request::CaptureChatSnapshot() const { + Request snapshot; + snapshot.options = options; + snapshot.prepared_tool_definitions = prepared_tool_definitions; + snapshot.forced_tool_choice = forced_tool_choice; + snapshot.raw_envelope_descriptor = raw_envelope_descriptor; + snapshot.item_segment_starts = item_segment_starts; + snapshot.items.reserve(items.size()); + + for (const auto* item : items) { + if (item == nullptr) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_ARGUMENT, "request preflight snapshot cannot contain a null item"); + } + + snapshot.AddOwnedItem(CloneChatItem(*item)); + } + + return snapshot; +} + +} // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/request.h b/sdk_v2/cpp/src/inferencing/session/request.h index 0ae4f3e07..d12e71b6a 100644 --- a/sdk_v2/cpp/src/inferencing/session/request.h +++ b/sdk_v2/cpp/src/inferencing/session/request.h @@ -82,6 +82,9 @@ struct Request { Request(const Request&) = delete; Request& operator=(const Request&) = delete; + /// Deep-copy every field read during chat preparation, including inline media bytes. + Request CaptureChatSnapshot() const; + /// Add a pre-allocated owned item. void AddOwnedItem(std::unique_ptr item) { items.push_back(item.get()); diff --git a/sdk_v2/cpp/src/inferencing/session/session.cc b/sdk_v2/cpp/src/inferencing/session/session.cc index d2f34bb51..562f75cfb 100644 --- a/sdk_v2/cpp/src/inferencing/session/session.cc +++ b/sdk_v2/cpp/src/inferencing/session/session.cc @@ -86,6 +86,21 @@ void Session::UndoTurns(size_t /*count*/) { FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "UndoTurns is not supported for this session type"); } +std::unique_ptr Session::CreateRequestPreflight(const Request& request) const { + auto lock = LockRequestMutex(); + if (Type() != SessionType::kChat) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "request preflight is only supported for chat sessions"); + } + + auto snapshot = request.CaptureChatSnapshot(); + ValidateRequestItems(snapshot); + return CreateRequestPreflightImpl(std::move(snapshot)); +} + +std::unique_ptr Session::CreateRequestPreflightImpl(Request) const { + FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "request preflight is only supported for chat sessions"); +} + void Session::AddToolDefinition(ToolDefinition tool_def) { tool_registry_.Add(std::move(tool_def)); } diff --git a/sdk_v2/cpp/src/inferencing/session/session.h b/sdk_v2/cpp/src/inferencing/session/session.h index f2db93c5c..5ffbfb34c 100644 --- a/sdk_v2/cpp/src/inferencing/session/session.h +++ b/sdk_v2/cpp/src/inferencing/session/session.h @@ -17,6 +17,7 @@ #include "inferencing/session/response.h" #include "inferencing/session/tool_registry.h" #include "inferencing/session/types.h" +#include "inferencing/generative/chat/search_options.h" #include "util/key_value_pairs.h" namespace fl { @@ -34,6 +35,21 @@ class Model; // forward declaration /// - Future: predictive inference, realtime audio, multi-modal class Session { public: + class RequestPreflightOperation { + public: + virtual ~RequestPreflightOperation() = default; + + RequestPreflightOperation(const RequestPreflightOperation&) = delete; + RequestPreflightOperation& operator=(const RequestPreflightOperation&) = delete; + RequestPreflightOperation(RequestPreflightOperation&&) = delete; + RequestPreflightOperation& operator=(RequestPreflightOperation&&) = delete; + + virtual RequestBudget Execute() = 0; + + protected: + RequestPreflightOperation() = default; + }; + virtual ~Session(); Session(Session&&) = default; @@ -54,6 +70,9 @@ class Session { /// in-flight callbacks and ensures the Response is fully populated on return. void ProcessRequest(const Request& request, Response& response); + /// Capture a chat request and session state under the same serialization boundary as generation and undo. + std::unique_ptr CreateRequestPreflight(const Request& request) const; + /// Signal every in-flight request on this session to cancel. Only sets each request's atomic /// cancel flag — never blocks and never joins — so it is safe to call from a shutdown path while /// another thread holds a manager lock. Generation loops poll the flag and stop within ~one token. @@ -96,6 +115,7 @@ class Session { /// Session-level parameters overlaid onto each request. void SetSessionOptions(const KeyValuePairs& options) { + auto lock = LockRequestMutex(); session_options_ = options; SetSessionOptionsImpl(session_options_); } @@ -120,16 +140,7 @@ class Session { /// Returns a copy of session options with request options overlaid (request wins on conflict). /// Derived classes call this when they want a single resolved option set. KeyValuePairs MergedOptions(const KeyValuePairs& request_options) const { - if (request_options.empty()) { - return session_options_; - } - - KeyValuePairs merged = session_options_; - for (const auto& [key, value] : request_options) { - merged.Add(key, value); - } - - return merged; + return MergeKeyValuePairs(session_options_, request_options); } /// Derived classes implement the actual generation logic. @@ -137,6 +148,8 @@ class Session { /// Requests are serialized if the derived class does not opt into concurrency via allow_concurrent_requests_. virtual void ProcessRequestImpl(const Request& request, Response& response) = 0; + virtual std::unique_ptr CreateRequestPreflightImpl(Request request) const; + /// Create a per-request callback handler. Returns nullptr if no callback is set. /// The handler is owned by the caller (unique_ptr) and drains+joins on destruction. std::unique_ptr CreateCallbackHandler(const Request& request) { diff --git a/sdk_v2/cpp/src/util/key_value_pairs.h b/sdk_v2/cpp/src/util/key_value_pairs.h index 3f281e650..6160e1759 100644 --- a/sdk_v2/cpp/src/util/key_value_pairs.h +++ b/sdk_v2/cpp/src/util/key_value_pairs.h @@ -152,4 +152,18 @@ class KeyValuePairs { mutable bool dirty_ = false; }; +/// Return base options with request options overlaid; request values win on collision. +inline KeyValuePairs MergeKeyValuePairs(const KeyValuePairs& base, const KeyValuePairs& request) { + if (request.empty()) { + return base; + } + + auto merged = base; + for (const auto& [key, value] : request) { + merged.Add(key, value); + } + + return merged; +} + } // namespace fl diff --git a/sdk_v2/cpp/test/internal_api/chat/dynamic_engine_chat_test.cc b/sdk_v2/cpp/test/internal_api/chat/dynamic_engine_chat_test.cc index 81839d3c8..a9a590b60 100644 --- a/sdk_v2/cpp/test/internal_api/chat/dynamic_engine_chat_test.cc +++ b/sdk_v2/cpp/test/internal_api/chat/dynamic_engine_chat_test.cc @@ -6,6 +6,7 @@ #include "inferencing/generative/chat/onnx_chat_engine.h" #include "inferencing/generative/chat/onnx_engine_chat_stream.h" #include "inferencing/model_load_manager.h" +#include "c_api_types.h" #include "internal_api/test_helpers.h" #include "internal_api/test_model_cache.h" #include "items/text_item.h" @@ -193,6 +194,44 @@ class DynamicEngineChatTest : public ::testing::Test { TelemetryLogger telemetry_{"dynamic-engine-test", test::NullLog()}; }; +TEST_F(DynamicEngineChatTest, NativeCAbiPreflightCapturesAndExecutesExactlyOnce) { + const auto* api = FoundryLocalGetApi(FOUNDRY_LOCAL_API_VERSION); + ASSERT_NE(api, nullptr); + const auto* inference_api = api->GetInferenceApi(); + + auto session = std::make_unique(CatalogModel(), ModelInstance(), *logger_, telemetry_); + auto request = std::make_unique(MakeRequest("Count this exact prompt.", 17)); + request->canceled.store(true, std::memory_order_relaxed); + EXPECT_FALSE(request->CaptureChatSnapshot().canceled.load(std::memory_order_relaxed)); + const auto expected_prompt_tokens = + static_cast(EncodeUserPrompt("Count this exact prompt.", ModelInstance()).size()); + + flRequestPreflight* preflight = nullptr; + ASSERT_EQ(inference_api->Session_CreateRequestPreflight( + AsHandle(session.get()), AsHandle(request.get()), &preflight), + nullptr); + ASSERT_NE(preflight, nullptr); + + request.reset(); + session.reset(); + + flRequestPreflightResult result{}; + result.version = FOUNDRY_LOCAL_API_VERSION; + ASSERT_EQ(inference_api->RequestPreflight_Execute(preflight, &result), nullptr); + EXPECT_EQ(result.prompt_tokens, expected_prompt_tokens); + EXPECT_EQ(result.output_reserve_tokens, 17); + EXPECT_EQ(result.required_tokens, result.prompt_tokens + result.output_reserve_tokens); + EXPECT_TRUE(result.fits); + EXPECT_EQ(result.deficit_tokens, 0); + + auto* repeated = inference_api->RequestPreflight_Execute(preflight, &result); + ASSERT_NE(repeated, nullptr); + EXPECT_EQ(api->Status_GetErrorCode(repeated), FOUNDRY_LOCAL_ERROR_INVALID_USAGE); + api->Status_Release(repeated); + + inference_api->RequestPreflight_Release(preflight); +} + TEST_F(DynamicEngineChatTest, RetainedContinuationReportsFreshPromptUsageParity) { ChatSession session(CatalogModel(), ModelInstance(), *logger_, telemetry_); @@ -213,6 +252,11 @@ TEST_F(DynamicEngineChatTest, RetainedContinuationReportsFreshPromptUsageParity) const auto expected_second_prompt_tokens = EncodeMessages(full_history, ModelInstance()).size(); auto second = MakeRequest(kSecondPrompt); + const auto second_budget = session.CreateRequestPreflight(second)->Execute(); + EXPECT_EQ(second_budget.prompt_tokens, expected_second_prompt_tokens); + EXPECT_EQ(second_budget.output_reserve_tokens, 32); + EXPECT_EQ(session.TurnCount(), 1u); + Response second_response; session.ProcessRequest(second, second_response); diff --git a/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc b/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc index 8e3a517ac..5c11d3f11 100644 --- a/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc +++ b/sdk_v2/cpp/test/internal_api/chat/search_options_test.cc @@ -50,6 +50,29 @@ TEST(SearchOptionsParsingTest, ExplicitOutputLimitOverridesTurnDefault) { EXPECT_EQ(ResolveMaxOutputTokens(explicit_limit, GetDefaultMaxOutputTokens(/*has_media=*/true)), 64); } +TEST(SearchOptionsParsingTest, RequestBudgetReportsExactFitAndDeficit) { + const auto exact = ComputeRequestBudget(90, 10, 100); + EXPECT_TRUE(exact.fits); + EXPECT_EQ(exact.required_tokens, 100); + EXPECT_EQ(exact.deficit_tokens, 0); + + const auto over = ComputeRequestBudget(91, 10, 100); + EXPECT_FALSE(over.fits); + EXPECT_EQ(over.required_tokens, 101); + EXPECT_EQ(over.deficit_tokens, 1); +} + +TEST(SearchOptionsParsingTest, OutputReserveMatchesBackendGenerationPolicy) { + SearchOptions explicit_limit; + explicit_limit.max_output_tokens = 64; + EXPECT_EQ(ResolveOutputReserve(explicit_limit, ChatBackendKind::kEngine, false, 25, 100), 64); + + const SearchOptions defaults; + EXPECT_EQ(ResolveOutputReserve(defaults, ChatBackendKind::kGenerator, false, 25, 100), 2048); + EXPECT_EQ(ResolveOutputReserve(defaults, ChatBackendKind::kGenerator, true, 25, 100), 3072); + EXPECT_EQ(ResolveOutputReserve(defaults, ChatBackendKind::kEngine, false, 25, 100), 75); +} + TEST(SearchOptionsParsingTest, RetainedGenerationSettingsAreBackendAware) { SearchOptions first; first.temperature = 0.5f; diff --git a/sdk_v2/cs/README.md b/sdk_v2/cs/README.md index 3fbecefb9..1b6391c7d 100644 --- a/sdk_v2/cs/README.md +++ b/sdk_v2/cs/README.md @@ -236,6 +236,30 @@ chatClient.Settings.TopP = 0.9f; chatClient.Settings.FrequencyPenalty = 0.0f; // Nonzero OpenAI penalties are not currently supported. ``` +#### Exact request preflight + +Use `ChatSession.PreflightRequestAsync` to ask the native runtime for the exact token budget of +a request in the session's current conversation state. The request and session state are captured +synchronously; the potentially expensive calculation then runs asynchronously. The source request +and session may be disposed after the method returns. The caller must not dispose +`FoundryLocalManager` until the returned task completes. + +```csharp +using var session = new ChatSession(model); +using var request = new Request() + .AddItem(MessageItem.User("Explain async/await in C#.")); + +var preflight = await session.PreflightRequestAsync(request); + +Console.WriteLine( + $"Prompt: {preflight.PromptTokens}, reserve: {preflight.OutputReserveTokens}, " + + $"required: {preflight.RequiredTokens}, limit: {preflight.ContextLimitTokens}, " + + $"fits: {preflight.Fits}, deficit: {preflight.DeficitTokens}"); +``` + +The values are returned directly by the native preflight ABI. The C# SDK does not apply defaults, +recalculate the budget, rewrite the request, or impose a local fit policy. + ### Audio Transcription ```csharp diff --git a/sdk_v2/cs/src/ChatSession.cs b/sdk_v2/cs/src/ChatSession.cs index cc7de3338..b67b7d114 100644 --- a/sdk_v2/cs/src/ChatSession.cs +++ b/sdk_v2/cs/src/ChatSession.cs @@ -7,6 +7,7 @@ namespace Microsoft.AI.Foundry.Local; using Microsoft.AI.Foundry.Local.Detail.Interop; +using Microsoft.AI.Foundry.Local.Detail.Native; /// /// A chat session for chat-completion models. @@ -126,4 +127,38 @@ public void UndoTurns(ulong count) ThrowIfDisposed(); GetNativeSession().UndoTurns(count); } + + /// + /// Capture this session's current state and synchronously, then + /// execute an exact token-budget preflight asynchronously. The captured operation remains valid + /// after the source session and request are disposed. The + /// must not be disposed until the returned task completes. + /// + public Task PreflightRequestAsync( + Request request, + CancellationToken ct = default) + { + ThrowIfDisposed(); + Detail.Throw.IfNull(request); + + var nativeSession = GetNativeSession(); + var status = Api.Inference.SessionCreateRequestPreflight( + nativeSession.Ptr, request.Ptr, out var preflightPtr); + Api.CheckStatus(status); + + var operation = new RequestPreflightOperation(preflightPtr); + return ExecutePreflightAsync(operation, ct); + } + + private static async Task ExecutePreflightAsync( + RequestPreflightOperation operation, + CancellationToken ct) + { +#pragma warning disable IDISP007 // Ownership is transferred by PreflightRequestAsync to this async helper. + using (operation) +#pragma warning restore IDISP007 + { + return await Task.Run(operation.Execute, ct).ConfigureAwait(false); + } + } } diff --git a/sdk_v2/cs/src/Detail/NativeMethods.cs b/sdk_v2/cs/src/Detail/NativeMethods.cs index 0248490e2..35d185594 100644 --- a/sdk_v2/cs/src/Detail/NativeMethods.cs +++ b/sdk_v2/cs/src/Detail/NativeMethods.cs @@ -208,6 +208,21 @@ public struct FlUsage public long TotalTokens; } +[StructLayout(LayoutKind.Sequential)] +public struct FlRequestPreflightResult +{ + public uint Version; + // 4 bytes implicit padding + public long PromptTokens; + public long OutputReserveTokens; + public long RequiredTokens; + public long ContextLimitTokens; + [MarshalAs(UnmanagedType.U1)] + public bool Fits; + // 7 bytes implicit padding + public long DeficitTokens; +} + [StructLayout(LayoutKind.Sequential)] public struct FlBytesData { @@ -690,6 +705,17 @@ public delegate IntPtr FlInference_SessionSetStreamingCallbackDelegate(IntPtr se [UnmanagedFunctionPointer(CallingConvention.Winapi)] public delegate IntPtr FlInference_SessionUndoTurnsDelegate(IntPtr session, UIntPtr count); +[UnmanagedFunctionPointer(CallingConvention.Winapi)] +public delegate IntPtr FlInference_SessionCreateRequestPreflightDelegate( + IntPtr session, IntPtr request, out IntPtr outPreflight); + +[UnmanagedFunctionPointer(CallingConvention.Winapi)] +public delegate IntPtr FlInference_RequestPreflightExecuteDelegate( + IntPtr preflight, ref FlRequestPreflightResult outResult); + +[UnmanagedFunctionPointer(CallingConvention.Winapi)] +public delegate void FlInference_RequestPreflightReleaseDelegate(IntPtr preflight); + // --- Configuration API (flConfigurationApi) delegates --- [UnmanagedFunctionPointer(CallingConvention.Winapi)] @@ -942,6 +968,11 @@ public struct FlInferenceApi public FlInference_SessionRemoveToolDefinitionDelegate SessionRemoveToolDefinition; public FlInference_SessionGetTurnCountDelegate SessionGetTurnCount; public FlInference_SessionUndoTurnsDelegate SessionUndoTurns; + + // Request preflight (appended in V2) + public FlInference_SessionCreateRequestPreflightDelegate SessionCreateRequestPreflight; + public FlInference_RequestPreflightExecuteDelegate RequestPreflightExecute; + public FlInference_RequestPreflightReleaseDelegate RequestPreflightRelease; } /// Configuration API table. diff --git a/sdk_v2/cs/src/Detail/RequestPreflightOperation.cs b/sdk_v2/cs/src/Detail/RequestPreflightOperation.cs new file mode 100644 index 000000000..35d3729ae --- /dev/null +++ b/sdk_v2/cs/src/Detail/RequestPreflightOperation.cs @@ -0,0 +1,49 @@ +// -------------------------------------------------------------------------------------------------------------------- +// +// Copyright (c) Microsoft. All rights reserved. +// +// -------------------------------------------------------------------------------------------------------------------- + +namespace Microsoft.AI.Foundry.Local.Detail.Native; + +using Microsoft.AI.Foundry.Local.Detail.Interop; + +/// +/// Owns a one-shot native request-preflight operation. +/// +internal sealed class RequestPreflightOperation : IDisposable +{ + private IntPtr _ptr; + + internal RequestPreflightOperation(IntPtr ptr) + { + _ptr = ptr; + } + + internal RequestPreflightResult Execute() + { + var nativeResult = new FlRequestPreflightResult + { + Version = NativeMethods.ApiVersion, + }; + + Api.CheckStatus(Api.Inference.RequestPreflightExecute(_ptr, ref nativeResult)); + + return new RequestPreflightResult( + nativeResult.PromptTokens, + nativeResult.OutputReserveTokens, + nativeResult.RequiredTokens, + nativeResult.ContextLimitTokens, + nativeResult.Fits, + nativeResult.DeficitTokens); + } + + public void Dispose() + { + var ptr = Interlocked.Exchange(ref _ptr, IntPtr.Zero); + if (ptr != IntPtr.Zero) + { + Api.Inference.RequestPreflightRelease(ptr); + } + } +} diff --git a/sdk_v2/cs/src/RequestPreflightResult.cs b/sdk_v2/cs/src/RequestPreflightResult.cs new file mode 100644 index 000000000..7a78284ab --- /dev/null +++ b/sdk_v2/cs/src/RequestPreflightResult.cs @@ -0,0 +1,41 @@ +// -------------------------------------------------------------------------------------------------------------------- +// +// Copyright (c) Microsoft. All rights reserved. +// +// -------------------------------------------------------------------------------------------------------------------- + +namespace Microsoft.AI.Foundry.Local; + +/// +/// Exact token-budget preflight result for a request captured in a chat session's current state. +/// +public sealed class RequestPreflightResult +{ + internal RequestPreflightResult( + long promptTokens, + long outputReserveTokens, + long requiredTokens, + long contextLimitTokens, + bool fits, + long deficitTokens) + { + PromptTokens = promptTokens; + OutputReserveTokens = outputReserveTokens; + RequiredTokens = requiredTokens; + ContextLimitTokens = contextLimitTokens; + Fits = fits; + DeficitTokens = deficitTokens; + } + + public long PromptTokens { get; } + + public long OutputReserveTokens { get; } + + public long RequiredTokens { get; } + + public long ContextLimitTokens { get; } + + public bool Fits { get; } + + public long DeficitTokens { get; } +} diff --git a/sdk_v2/cs/test/FoundryLocal.Tests/ChatSessionTests.cs b/sdk_v2/cs/test/FoundryLocal.Tests/ChatSessionTests.cs index 7634f1620..04507501c 100644 --- a/sdk_v2/cs/test/FoundryLocal.Tests/ChatSessionTests.cs +++ b/sdk_v2/cs/test/FoundryLocal.Tests/ChatSessionTests.cs @@ -69,6 +69,30 @@ public async Task Chat_NoStreaming_Succeeds() Console.WriteLine($"Response: {content}"); } + [Test] + public async Task RequestPreflight_CaptureSurvivesSourceDisposal() + { + var session = new ChatSession(model!); + session.SetOptions(new RequestOptions { Search = new SearchOptions { MaxOutputTokens = 32 } }); + var request = new Request(); + request.AddItem(MessageItem.User("Count the tokens in this request.")); + + var preflightTask = session.PreflightRequestAsync(request); + + request.Dispose(); + session.Dispose(); + + var result = await preflightTask.ConfigureAwait(false); + + await Assert.That(result.PromptTokens).IsGreaterThan(0L); + await Assert.That(result.OutputReserveTokens).IsEqualTo(32L); + await Assert.That(result.RequiredTokens) + .IsEqualTo(result.PromptTokens + result.OutputReserveTokens); + await Assert.That(result.ContextLimitTokens).IsGreaterThan(0L); + await Assert.That(result.Fits).IsTrue(); + await Assert.That(result.DeficitTokens).IsEqualTo(0L); + } + [Test] public async Task Chat_Streaming_Succeeds() { diff --git a/sdk_v2/cs/test/FoundryLocal.Tests/RequestPreflightAbiTests.cs b/sdk_v2/cs/test/FoundryLocal.Tests/RequestPreflightAbiTests.cs new file mode 100644 index 000000000..8e21722a2 --- /dev/null +++ b/sdk_v2/cs/test/FoundryLocal.Tests/RequestPreflightAbiTests.cs @@ -0,0 +1,58 @@ +// -------------------------------------------------------------------------------------------------------------------- +// +// Copyright (c) Microsoft. All rights reserved. +// +// -------------------------------------------------------------------------------------------------------------------- + +namespace Microsoft.AI.Foundry.Local.Tests; + +using System.Runtime.InteropServices; +using System.Threading.Tasks; + +using Microsoft.AI.Foundry.Local.Detail.Interop; + +#pragma warning disable TUnitAssertions0005 + +internal sealed class RequestPreflightAbiTests +{ + [Test] + public async Task ResultLayout_MatchesNativeNaturalAlignment() + { + await Assert.That(NativeMethods.ApiVersion).IsEqualTo(2u); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.Version)).ToInt64()).IsEqualTo(0L); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.PromptTokens)).ToInt64()).IsEqualTo(8L); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.OutputReserveTokens)).ToInt64()).IsEqualTo(16L); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.RequiredTokens)).ToInt64()).IsEqualTo(24L); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.ContextLimitTokens)).ToInt64()).IsEqualTo(32L); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.Fits)).ToInt64()).IsEqualTo(40L); + await Assert.That(Marshal.OffsetOf( + nameof(FlRequestPreflightResult.DeficitTokens)).ToInt64()).IsEqualTo(48L); + await Assert.That(Marshal.SizeOf()).IsEqualTo(56); + } + + [Test] + public async Task InferenceVtable_PreflightSlotsAreAppendedInHeaderOrder() + { + var pointerSize = IntPtr.Size; + + await Assert.That(Marshal.OffsetOf( + nameof(FlInferenceApi.SessionCreateRequestPreflight)).ToInt64()) + .IsEqualTo(22L * pointerSize); + await Assert.That(Marshal.OffsetOf( + nameof(FlInferenceApi.RequestPreflightExecute)).ToInt64()) + .IsEqualTo(23L * pointerSize); + await Assert.That(Marshal.OffsetOf( + nameof(FlInferenceApi.RequestPreflightRelease)).ToInt64()) + .IsEqualTo(24L * pointerSize); + await Assert.That(Marshal.SizeOf()).IsEqualTo(25 * pointerSize); + } + +} + +#pragma warning restore TUnitAssertions0005 diff --git a/sdk_v2/js/README.md b/sdk_v2/js/README.md index f0a4e5d2a..8ba2ecd04 100644 --- a/sdk_v2/js/README.md +++ b/sdk_v2/js/README.md @@ -36,6 +36,9 @@ your TS/JS code The addon talks to the **C++ wrapper**, never directly to the C ABI. If the wrapper is missing something, fix the wrapper rather than reaching past it. +Request preflight captures state synchronously, executes token counting on a worker, and keeps the +manager-owned runtime alive until completion. + --- ## 2. Prerequisites diff --git a/sdk_v2/js/native/src/session.cc b/sdk_v2/js/native/src/session.cc index 5e92dbc9a..75d4d5e7e 100644 --- a/sdk_v2/js/native/src/session.cc +++ b/sdk_v2/js/native/src/session.cc @@ -59,6 +59,17 @@ Napi::Value ResponseToJs(Napi::Env env, foundry_local::Response& resp) { return out; } +Napi::Value RequestPreflightResultToJs(Napi::Env env, const flRequestPreflightResult& result) { + Napi::Object out = Napi::Object::New(env); + out.Set("promptTokens", Napi::Number::New(env, static_cast(result.prompt_tokens))); + out.Set("outputReserveTokens", Napi::Number::New(env, static_cast(result.output_reserve_tokens))); + out.Set("requiredTokens", Napi::Number::New(env, static_cast(result.required_tokens))); + out.Set("contextLimitTokens", Napi::Number::New(env, static_cast(result.context_limit_tokens))); + out.Set("fits", Napi::Boolean::New(env, result.fits)); + out.Set("deficitTokens", Napi::Number::New(env, static_cast(result.deficit_tokens))); + return out; +} + void ThrowFoundryLocalError(Napi::Env env, int code, const std::string& msg) { Napi::Error err = Napi::Error::New(env, msg); Napi::Object value = err.Value(); @@ -69,7 +80,7 @@ void ThrowFoundryLocalError(Napi::Env env, int code, const std::string& msg) { foundry_local::Request* UnwrapRequest(Napi::Env env, const Napi::Value& v) { if (!v.IsObject()) { - Napi::TypeError::New(env, "processRequest(request): expected a Request instance") + Napi::TypeError::New(env, "request argument must be a Request instance") .ThrowAsJavaScriptException(); return nullptr; } @@ -80,13 +91,13 @@ foundry_local::Request* UnwrapRequest(Napi::Env env, const Napi::Value& v) { return nullptr; } if (!obj.InstanceOf(data->request_ctor.Value())) { - Napi::TypeError::New(env, "processRequest(request): argument is not a Request instance") + Napi::TypeError::New(env, "request argument must be a Request instance") .ThrowAsJavaScriptException(); return nullptr; } Request* req = Napi::ObjectWrap::Unwrap(obj); if (req == nullptr || req->native() == nullptr) { - Napi::TypeError::New(env, "processRequest(request): Request is not initialized") + Napi::TypeError::New(env, "request argument is not initialized") .ThrowAsJavaScriptException(); return nullptr; } @@ -292,6 +303,7 @@ Napi::Function ChatSession::Init(Napi::Env env) { return DefineClass(env, "ChatSession", { InstanceMethod("processRequest", &ChatSession::ProcessRequest), + InstanceMethod("preflightRequest", &ChatSession::PreflightRequest), InstanceMethod("processStreamingRequest", &ChatSession::ProcessStreamingRequest), InstanceMethod("setOptions", &ChatSession::SetOptions), InstanceMethod("addToolDefinition", &ChatSession::AddToolDefinition), @@ -353,6 +365,37 @@ Napi::Value ChatSession::ProcessRequest(const Napi::CallbackInfo& info) { return ProcessRequestOn(env, impl_.get(), info[0], std::move(owner)); } +Napi::Value ChatSession::PreflightRequest(const Napi::CallbackInfo& info) { + Napi::Env env = info.Env(); + if (ThrowIfDisposed(env)) return env.Undefined(); + if (info.Length() < 1) { + Napi::TypeError::New(env, "preflightRequest(request: Request)").ThrowAsJavaScriptException(); + return env.Undefined(); + } + foundry_local::Request* request = UnwrapRequest(env, info[0]); + if (request == nullptr) return env.Undefined(); + + std::shared_ptr preflight; + CallCheckedVoid(env, [&]() { + preflight = std::make_shared( + impl_->CaptureRequestPreflight(*request)); + }); + if (env.IsExceptionPending() || preflight == nullptr) return env.Undefined(); + + // Capture owns the request/session snapshot; only the Manager runtime must remain alive. + Napi::ObjectReference manager_ref = Napi::Reference::New(manager_.Value(), 1); + return PromiseWorker::Run( + env, + [preflight = std::move(preflight)]() mutable { + struct ResetOperation { + std::shared_ptr& operation; + ~ResetOperation() { operation.reset(); } + } reset{preflight}; + return preflight->Execute(); + }, + RequestPreflightResultToJs, std::move(manager_ref)); +} + Napi::Value ChatSession::ProcessStreamingRequest(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (ThrowIfDisposed(env)) return env.Undefined(); diff --git a/sdk_v2/js/native/src/session.h b/sdk_v2/js/native/src/session.h index 2b1db7b86..c005c7e33 100644 --- a/sdk_v2/js/native/src/session.h +++ b/sdk_v2/js/native/src/session.h @@ -7,6 +7,8 @@ // * new ChatSession(model) — sync construction; underlying // flSession_Create is fast. // * session.processRequest(request) -> Promise (PromiseWorker) +// * session.preflightRequest(request) -> Promise — +// captures in-memory state synchronously, then executes it on a worker. // * session.processStreamingRequest(request, onItem) -> Promise — streaming bridge via // Napi::ThreadSafeFunction; resolves with the terminal Response after every item callback drains. // The JS layer wraps this in an AsyncIterable whose `.response` promise carries the resolved value. @@ -40,6 +42,7 @@ class ChatSession : public Napi::ObjectWrap { private: Napi::Value ProcessRequest(const Napi::CallbackInfo& info); + Napi::Value PreflightRequest(const Napi::CallbackInfo& info); Napi::Value ProcessStreamingRequest(const Napi::CallbackInfo& info); Napi::Value SetOptions(const Napi::CallbackInfo& info); Napi::Value AddToolDefinition(const Napi::CallbackInfo& info); diff --git a/sdk_v2/js/src/detail/native.ts b/sdk_v2/js/src/detail/native.ts index 4167acdd3..a151d824d 100644 --- a/sdk_v2/js/src/detail/native.ts +++ b/sdk_v2/js/src/detail/native.ts @@ -120,6 +120,15 @@ export interface NativeResponse { usage: { promptTokens: number; completionTokens: number; totalTokens: number }; } +export interface NativeRequestPreflightResult { + promptTokens: number; + outputReserveTokens: number; + requiredTokens: number; + contextLimitTokens: number; + fits: boolean; + deficitTokens: number; +} + export interface NativeRequestCtor { new (): NativeRequest; } @@ -175,6 +184,7 @@ export interface NativeSession { } export interface NativeChatSession extends NativeSession { + preflightRequest(request: NativeRequest): Promise; addToolDefinition( definition: | { diff --git a/sdk_v2/js/src/index.ts b/sdk_v2/js/src/index.ts index 712eced03..823a1b834 100644 --- a/sdk_v2/js/src/index.ts +++ b/sdk_v2/js/src/index.ts @@ -58,6 +58,7 @@ export { type ToolKind, type StreamOptions, type StreamingResponse, + type RequestPreflightResult, } from "./session.js"; export { Request, type RequestOptions, type RequestToolChoice, type SearchOptions } from "./request.js"; export type { Response, FinishReason, TokenUsage } from "./response.js"; diff --git a/sdk_v2/js/src/request.ts b/sdk_v2/js/src/request.ts index 801f0399c..198116726 100644 --- a/sdk_v2/js/src/request.ts +++ b/sdk_v2/js/src/request.ts @@ -88,7 +88,7 @@ export class Request { export function unwrapNativeRequest(request: Request): NativeRequest { const n = nativeByRequest.get(request); if (n === undefined) { - throw new TypeError("Session.processRequest: argument is not a valid Request"); + throw new TypeError("Session request argument is not a valid Request"); } return n; } diff --git a/sdk_v2/js/src/session.ts b/sdk_v2/js/src/session.ts index 932c6638d..b556daab2 100644 --- a/sdk_v2/js/src/session.ts +++ b/sdk_v2/js/src/session.ts @@ -54,6 +54,16 @@ export interface StreamingResponse extends AsyncIterable { readonly response: Promise; } +/** Exact token budget for a request captured against the current chat-session state. */ +export interface RequestPreflightResult { + readonly promptTokens: number; + readonly outputReserveTokens: number; + readonly requiredTokens: number; + readonly contextLimitTokens: number; + readonly fits: boolean; + readonly deficitTokens: number; +} + function rejectEmbeddedNul(value: unknown, argumentName: string): void { if (typeof value === "string" && value.includes("\0")) { throw new TypeError(`${argumentName} must not contain an embedded NUL character`); @@ -366,6 +376,14 @@ export class ChatSession extends Session { return this.native as NativeChatSession; } + /** + * Compute the exact token budget without generation or state changes. + * Request and session state are captured before worker execution is queued. + */ + preflightRequest(request: Request): Promise { + return this.#nativeChat.preflightRequest(unwrapNativeRequest(request)); + } + /** * Register a tool definition available to the model for the rest of the * session. Mirrors `foundry_local::ChatSession::AddToolDefinition`. diff --git a/sdk_v2/js/test/chat-session.test.ts b/sdk_v2/js/test/chat-session.test.ts index c4073b51d..3bc95243f 100644 --- a/sdk_v2/js/test/chat-session.test.ts +++ b/sdk_v2/js/test/chat-session.test.ts @@ -95,6 +95,52 @@ describe.skipIf(!haveTestModelCache)("ChatSession (real model, non-streaming)", 2 * 60_000, ); + it( + "preflightRequest returns an exact budget matching inference without changing history", + async () => { + if (session === undefined) throw new Error("fixture missing"); + const request = new Request() + .addItem(Item.systemMessage("You are concise.")) + .addItem(Item.userMessage("Explain why the sky is blue.")) + .setOptions({ search: { maxOutputTokens: 64, temperature: 0 } }); + const before = session.turnCount; + + const result = await session.preflightRequest(request); + + expect(result.promptTokens).toBeGreaterThan(0); + expect(result.outputReserveTokens).toBe(64); + expect(result.requiredTokens).toBe(result.promptTokens + result.outputReserveTokens); + expect(result.contextLimitTokens).toBeGreaterThan(0); + expect(result.fits).toBe(result.requiredTokens <= result.contextLimitTokens); + expect(result.deficitTokens).toBe(Math.max(0, result.requiredTokens - result.contextLimitTokens)); + expect(session.turnCount).toBe(before); + + const response = await session.processRequest(request); + expect(result.promptTokens).toBe(response.usage.promptTokens); + }, + 2 * 60_000, + ); + + it( + "preflightRequest owns captured state after source session disposal", + async () => { + if (fixture === undefined) throw new Error("fixture missing"); + const source = new ChatSession(fixture.model); + const request = new Request() + .addItem(Item.userMessage("Count the prompt tokens.")) + .setOptions({ search: { maxOutputTokens: 32 } }); + + const pending = source.preflightRequest(request); + request.setOptions({ search: { maxOutputTokens: 96 } }); + source.dispose(); + + await expect(pending).resolves.toMatchObject({ + outputReserveTokens: 32, + }); + }, + 2 * 60_000, + ); + it( "undoTurns rewinds the conversation", async () => { diff --git a/sdk_v2/python/README.md b/sdk_v2/python/README.md index 69990b531..9c61476de 100644 --- a/sdk_v2/python/README.md +++ b/sdk_v2/python/README.md @@ -244,6 +244,22 @@ model.unload() `ChatSession` is stateful across turns. `session.turn_count` reports how many requests have been processed; `session.undo_turns(n)` rewinds history. +Before processing a request, use `preflight_request()` to synchronously ask the native runtime for +the request's token budget: + +```python +model = manager.catalog.get_model("qwen2.5-0.5b") +model.load() +with ChatSession(model) as session: + with Request().add_item(MessageItem.user("Summarize this conversation.")) as req: + budget = session.preflight_request(req) + print(f"Required tokens: {budget.required_tokens}; fits: {budget.fits}") +model.unload() +``` + +The SDK reports the native result without estimating tokens, changing options, truncating context, +or rewriting the request. + ### Multi-turn conversations Each call to `process_request` extends the session's turn history. Build a new `Request` per turn: diff --git a/sdk_v2/python/src/foundry_local_sdk/__init__.py b/sdk_v2/python/src/foundry_local_sdk/__init__.py index c5b79a498..727fbc79e 100644 --- a/sdk_v2/python/src/foundry_local_sdk/__init__.py +++ b/sdk_v2/python/src/foundry_local_sdk/__init__.py @@ -43,6 +43,7 @@ from foundry_local_sdk.session_types import ( FinishReason, RequestOptions, + RequestPreflightResult, SearchOptions, TokenUsage, ToolChoice, @@ -105,6 +106,7 @@ "TokenUsage", "SearchOptions", "RequestOptions", + "RequestPreflightResult", "ToolChoice", "Request", "Response", diff --git a/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py b/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py index 8ca717576..a5df69652 100644 --- a/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py +++ b/sdk_v2/python/src/foundry_local_sdk/_native/build_cffi.py @@ -172,6 +172,7 @@ typedef struct flModelInfo flModelInfo; typedef struct flModelList flModelList; typedef struct flRequest flRequest; +typedef struct flRequestPreflight flRequestPreflight; typedef struct flResponse flResponse; typedef struct flSession flSession; typedef struct flStatus flStatus; @@ -194,6 +195,16 @@ int64_t total_tokens; } flUsage; +typedef struct flRequestPreflightResult { + uint32_t version; + int64_t prompt_tokens; + int64_t output_reserve_tokens; + int64_t required_tokens; + int64_t context_limit_tokens; + _Bool fits; + int64_t deficit_tokens; +} flRequestPreflightResult; + typedef struct flEpInfo { uint32_t version; const char* name; @@ -467,6 +478,9 @@ flStatusPtr (*Session_RemoveToolDefinition)(flSession* session, const char* tool_name, bool* out_removed); size_t (*Session_GetTurnCount)(const flSession* session); flStatusPtr (*Session_UndoTurns)(flSession* session, size_t count); + flStatusPtr (*Session_CreateRequestPreflight)(const flSession* session, const flRequest* request, flRequestPreflight** out_preflight); + flStatusPtr (*RequestPreflight_Execute)(flRequestPreflight* preflight, flRequestPreflightResult* out_result); + void (*RequestPreflight_Release)(flRequestPreflight* instance); } flInferenceApi; /* ----------------------------------------------------------------------- diff --git a/sdk_v2/python/src/foundry_local_sdk/session.py b/sdk_v2/python/src/foundry_local_sdk/session.py index afd0de319..671e86c5a 100644 --- a/sdk_v2/python/src/foundry_local_sdk/session.py +++ b/sdk_v2/python/src/foundry_local_sdk/session.py @@ -15,7 +15,7 @@ from foundry_local_sdk.items import Item from foundry_local_sdk.request import Request from foundry_local_sdk.response import Response - from foundry_local_sdk.session_types import RequestOptions + from foundry_local_sdk.session_types import RequestOptions, RequestPreflightResult # Stamped on every versioned struct this module builds. Must match the version requested from # FoundryLocalGetApi (see _native/api.py): a tool definition carrying `kind` is only read as such @@ -487,6 +487,36 @@ def __init__(self, model: "IModel") -> None: ) super().__init__(model) + def preflight_request(self, request: "Request") -> "RequestPreflightResult": + """Synchronously return the native token budget for a request.""" + self._check_open() + + from foundry_local_sdk._native import ffi + from foundry_local_sdk._native.api import api + from foundry_local_sdk.session_types import RequestPreflightResult + + out_preflight = ffi.new("flRequestPreflight**") + api.check_status( + api.inference.Session_CreateRequestPreflight( + self._ptr, request._ptr, out_preflight + ) + ) + preflight = out_preflight[0] + try: + result = ffi.new("flRequestPreflightResult*") + result.version = _API_VERSION + api.check_status(api.inference.RequestPreflight_Execute(preflight, result)) + return RequestPreflightResult( + prompt_tokens=int(result.prompt_tokens), + output_reserve_tokens=int(result.output_reserve_tokens), + required_tokens=int(result.required_tokens), + context_limit_tokens=int(result.context_limit_tokens), + fits=bool(result.fits), + deficit_tokens=int(result.deficit_tokens), + ) + finally: + api.inference.RequestPreflight_Release(preflight) + def add_tool_definition(self, name: str, description: str, json_schema: str) -> "ChatSession": """Register a function tool so the model can request tool calls. Returns self (fluent). diff --git a/sdk_v2/python/src/foundry_local_sdk/session_types.py b/sdk_v2/python/src/foundry_local_sdk/session_types.py index b558fd570..6c98e024f 100644 --- a/sdk_v2/python/src/foundry_local_sdk/session_types.py +++ b/sdk_v2/python/src/foundry_local_sdk/session_types.py @@ -23,6 +23,18 @@ class TokenUsage: total_tokens: int +@dataclass(frozen=True) +class RequestPreflightResult: + """Exact native token-budget result from a one-shot captured request/session state.""" + + prompt_tokens: int + output_reserve_tokens: int + required_tokens: int + context_limit_tokens: int + fits: bool + deficit_tokens: int + + class _SessionParam: """Internal — well-known parameter key strings for the native KVP wire format. These mirror the ``FOUNDRY_LOCAL_PARAM_*`` macros in ``foundry_local_c.h``. Use the typed ``RequestOptions`` / ``SearchOptions`` API diff --git a/sdk_v2/python/test/unit/test_imports.py b/sdk_v2/python/test/unit/test_imports.py index da0efbe1d..d6cfa95f9 100644 --- a/sdk_v2/python/test/unit/test_imports.py +++ b/sdk_v2/python/test/unit/test_imports.py @@ -48,6 +48,7 @@ "TokenUsage", "SearchOptions", "RequestOptions", + "RequestPreflightResult", "ToolChoice", "Request", "Response", diff --git a/sdk_v2/python/test/unit/test_request_preflight.py b/sdk_v2/python/test/unit/test_request_preflight.py new file mode 100644 index 000000000..fac0b6886 --- /dev/null +++ b/sdk_v2/python/test/unit/test_request_preflight.py @@ -0,0 +1,106 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Focused tests for synchronous request preflight.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import pytest + +from foundry_local_sdk import RequestPreflightResult +from foundry_local_sdk._native import ffi +from foundry_local_sdk.exception import FoundryLocalException +from foundry_local_sdk.session import ChatSession, _API_VERSION + +_SESSION_PTR = ffi.cast("flSession *", 1) +_REQUEST_PTR = ffi.cast("flRequest *", 2) +_PREFLIGHT_PTR = ffi.cast("flRequestPreflight *", 3) + + +def _fake_session() -> ChatSession: + session = ChatSession.__new__(ChatSession) + session._closed = False + session._ptr = _SESSION_PTR + session._stream_thread = None + session._stream_request = None + return session + + +def _install_native(monkeypatch: pytest.MonkeyPatch, api: object) -> None: + native_api_module = importlib.import_module("foundry_local_sdk._native.api") + monkeypatch.setattr(native_api_module, "api", api) + + +def test_preflight_request_maps_result_and_releases_operation(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + def create(session_ptr, request_ptr, out): + assert session_ptr == _SESSION_PTR + assert request_ptr == _REQUEST_PTR + out[0] = _PREFLIGHT_PTR + calls.append("create") + return ffi.NULL + + def execute(operation, result): + assert operation == _PREFLIGHT_PTR + assert result.version == _API_VERSION + result.prompt_tokens = 11 + result.output_reserve_tokens = 7 + result.required_tokens = 18 + result.context_limit_tokens = 16 + result.fits = False + result.deficit_tokens = 2 + calls.append("execute") + return ffi.NULL + + fake_api = SimpleNamespace( + inference=SimpleNamespace( + Session_CreateRequestPreflight=create, + RequestPreflight_Execute=execute, + RequestPreflight_Release=lambda _operation: calls.append("release"), + ), + check_status=lambda status: assert_null(status), + ) + _install_native(monkeypatch, fake_api) + + result = _fake_session().preflight_request(SimpleNamespace(_ptr=_REQUEST_PTR)) + + assert result == RequestPreflightResult(11, 7, 18, 16, False, 2) + assert calls == ["create", "execute", "release"] + + +def test_preflight_request_execute_error_releases_operation(monkeypatch: pytest.MonkeyPatch) -> None: + execute_status = object() + releases = [] + + def create(_session_ptr, _request_ptr, out): + out[0] = _PREFLIGHT_PTR + return ffi.NULL + + def check_status(status): + if status is execute_status: + raise FoundryLocalException("preflight failed") + assert_null(status) + + fake_api = SimpleNamespace( + inference=SimpleNamespace( + Session_CreateRequestPreflight=create, + RequestPreflight_Execute=lambda _operation, _result: execute_status, + RequestPreflight_Release=releases.append, + ), + check_status=check_status, + ) + _install_native(monkeypatch, fake_api) + + with pytest.raises(FoundryLocalException, match="preflight failed"): + _fake_session().preflight_request(SimpleNamespace(_ptr=_REQUEST_PTR)) + + assert releases == [_PREFLIGHT_PTR] + + +def assert_null(status) -> None: + assert status == ffi.NULL diff --git a/sdk_v2/python/test/unit/test_session_types.py b/sdk_v2/python/test/unit/test_session_types.py index 828eea423..05d9113e1 100644 --- a/sdk_v2/python/test/unit/test_session_types.py +++ b/sdk_v2/python/test/unit/test_session_types.py @@ -12,6 +12,7 @@ from foundry_local_sdk import ( FinishReason, RequestOptions, + RequestPreflightResult, SearchOptions, TokenUsage, ToolChoice, @@ -53,6 +54,13 @@ def test_equality_by_value(self): assert a != c +class TestRequestPreflightResult: + def test_is_frozen_dataclass(self): + result = RequestPreflightResult(10, 20, 30, 40, True, 0) + with pytest.raises(dataclasses.FrozenInstanceError): + result.fits = False # type: ignore[misc] + + class TestToolChoice: def test_wire_values(self): # Must match the FOUNDRY_LOCAL_TOOL_CHOICE_* C++ enum's serialised form. diff --git a/sdk_v2/rust/README.md b/sdk_v2/rust/README.md index 4a8e9eeac..b6ee4766b 100644 --- a/sdk_v2/rust/README.md +++ b/sdk_v2/rust/README.md @@ -21,6 +21,7 @@ The Foundry Local Rust SDK provides an async Rust interface for running AI model - **Embedded web service** — Start a local HTTP server for OpenAI-compatible API access - **WinML support** — Automatic execution provider download on Windows for NPU/GPU acceleration - **Configurable inference** — Control temperature, max tokens, top-k, top-p, frequency penalty, random seed, and more +- **Exact request preflight** — Capture a chat request and session state, then asynchronously query its native token budget before generation - **Async-first** — Every operation is `async`; designed for use with the `tokio` runtime - **Safe FFI** — Dynamically loads the native Foundry Local engine (`foundry_local`) with a safe Rust wrapper @@ -611,8 +612,8 @@ engine's dependencies regardless of rpath/search-path setup. On platforms where ### Runtime Loading At runtime, the SDK uses `libloading` to dynamically load the `foundry_local` library, resolve the -API function table via `FoundryLocalGetApi`, and cache the sub-API tables. No static linking or -system-wide installation is required. +API version 2 function table via `FoundryLocalGetApi`, and cache the sub-API tables. No static +linking or system-wide installation is required. ## Platform Support diff --git a/sdk_v2/rust/docs/api.md b/sdk_v2/rust/docs/api.md index ae7a66a94..9ab69bec2 100644 --- a/sdk_v2/rust/docs/api.md +++ b/sdk_v2/rust/docs/api.md @@ -27,6 +27,7 @@ - [Inference API](#inference-api) - [Session](#session) - [ChatSession](#chatsession) + - [RequestPreflightResult](#requestpreflightresult) - [EmbeddingsSession](#embeddingssession) - [AudioSession](#audiosession) - [ItemQueue](#itemqueue) @@ -435,8 +436,20 @@ impl Deref for ChatSession { type Target = Session; } | `remove_tool_definition` | `async fn remove_tool_definition(&self, name: impl Into) -> Result` | Remove a tool by name; returns whether one was removed. | | `turn_count` | `fn turn_count(&self) -> usize` | The number of completed conversation turns. | | `undo_turns` | `async fn undo_turns(&self, count: usize) -> Result<(), FoundryLocalError>` | Rewind the last `count` turns. | +| `preflight_request` | `fn preflight_request(&self, request: Request) -> impl Future> + Send + 'static` | Synchronously capture the request and current conversation state before returning an awaitable future; only exact native token-budget execution is deferred to a blocking worker. | | `into_session` | `fn into_session(self) -> Session` | Consume this handle, yielding the base session. | +### RequestPreflightResult + +Exact native token-budget result returned by +[`ChatSession::preflight_request`](#chatsession). Its public value fields are +`prompt_tokens`, `output_reserve_tokens`, `required_tokens`, +`context_limit_tokens`, `fits`, and `deficit_tokens`. + +Request and session state is captured synchronously before the future is returned. +URI-backed media is resolved during execution. The captured operation executes +once on the SDK's blocking worker. + ### EmbeddingsSession An embeddings-oriented session producing dense vectors for text input. diff --git a/sdk_v2/rust/src/detail/ffi.rs b/sdk_v2/rust/src/detail/ffi.rs index 3b2291685..5db171a73 100644 --- a/sdk_v2/rust/src/detail/ffi.rs +++ b/sdk_v2/rust/src/detail/ffi.rs @@ -14,10 +14,7 @@ use core::ffi::c_void; use std::os::raw::{c_char, c_int}; -/// The library is built against this API version (`FOUNDRY_LOCAL_API_VERSION`). -/// -/// This is both the version requested from `FoundryLocalGetApi` and the version stamped on every -/// versioned struct built here, so the two can never disagree. +/// API version requested from `FoundryLocalGetApi` and stamped on versioned C structs. pub const FOUNDRY_LOCAL_API_VERSION: u32 = 2; // ── Opaque handle types ────────────────────────────────────────────────────── @@ -42,6 +39,7 @@ opaque_type!(flModel); opaque_type!(flModelInfo); opaque_type!(flModelList); opaque_type!(flRequest); +opaque_type!(flRequestPreflight); opaque_type!(flResponse); opaque_type!(flSession); opaque_type!(flStatus); @@ -182,6 +180,17 @@ pub struct flUsage { pub total_tokens: i64, } +#[repr(C)] +pub struct flRequestPreflightResult { + pub version: u32, + pub prompt_tokens: i64, + pub output_reserve_tokens: i64, + pub required_tokens: i64, + pub context_limit_tokens: i64, + pub fits: bool, + pub deficit_tokens: i64, +} + #[repr(C)] pub struct flEpInfo { pub version: u32, @@ -620,6 +629,16 @@ pub struct flInferenceApiVtable { pub Session_GetTurnCount: unsafe extern "system" fn(session: *const flSession) -> usize, pub Session_UndoTurns: unsafe extern "system" fn(session: *mut flSession, count: usize) -> flStatusPtr, + pub Session_CreateRequestPreflight: unsafe extern "system" fn( + session: *const flSession, + request: *const flRequest, + out_preflight: *mut *mut flRequestPreflight, + ) -> flStatusPtr, + pub RequestPreflight_Execute: unsafe extern "system" fn( + preflight: *mut flRequestPreflight, + out_result: *mut flRequestPreflightResult, + ) -> flStatusPtr, + pub RequestPreflight_Release: unsafe extern "system" fn(instance: *mut flRequestPreflight), } /// Configuration API table (`flConfigurationApi`). diff --git a/sdk_v2/rust/src/detail/session.rs b/sdk_v2/rust/src/detail/session.rs index d6d2829ac..089ed2166 100644 --- a/sdk_v2/rust/src/detail/session.rs +++ b/sdk_v2/rust/src/detail/session.rs @@ -100,6 +100,46 @@ impl Drop for NativeRequest { } } +// ── Request preflight ──────────────────────────────────────────────────────── + +pub(crate) struct NativeRequestPreflight { + api: Arc, + ptr: *mut flRequestPreflight, + _manager: Arc, +} + +// SAFETY: the native preflight is a self-contained, exclusively owned snapshot. +// Moving that handle to one blocking worker does not permit concurrent access, +// and `_manager` keeps the runtime alive until the native operation is released. +unsafe impl Send for NativeRequestPreflight {} + +impl NativeRequestPreflight { + pub(crate) fn execute(self) -> Result { + let mut result = flRequestPreflightResult { + version: FOUNDRY_LOCAL_API_VERSION, + prompt_tokens: 0, + output_reserve_tokens: 0, + required_tokens: 0, + context_limit_tokens: 0, + fits: false, + deficit_tokens: 0, + }; + let status = + unsafe { (self.api.inference_api().RequestPreflight_Execute)(self.ptr, &mut result) }; + self.api.check(status)?; + Ok(result) + } +} + +impl Drop for NativeRequestPreflight { + fn drop(&mut self) { + if !self.ptr.is_null() { + unsafe { (self.api.inference_api().RequestPreflight_Release)(self.ptr) }; + self.ptr = ptr::null_mut(); + } + } +} + // ── Response ───────────────────────────────────────────────────────────────── pub(crate) struct NativeResponse { @@ -408,6 +448,27 @@ impl NativeSession { self.api.check(status) } + pub(crate) fn create_request_preflight( + &self, + request: &NativeRequest, + ) -> Result { + let _guard = self.lock_ops(); + let mut ptr: *mut flRequestPreflight = ptr::null_mut(); + let status = unsafe { + (self.api.inference_api().Session_CreateRequestPreflight)( + self.ptr, + request.ptr, + &mut ptr, + ) + }; + self.api.check(status)?; + Ok(NativeRequestPreflight { + api: Arc::clone(&self.api), + ptr, + _manager: Arc::clone(&self._manager), + }) + } + pub(crate) fn process_request(&self, request: &NativeRequest) -> Result { let mut resp: *mut flResponse = ptr::null_mut(); let status = unsafe { diff --git a/sdk_v2/rust/src/lib.rs b/sdk_v2/rust/src/lib.rs index 93d7605d9..f30f7388c 100644 --- a/sdk_v2/rust/src/lib.rs +++ b/sdk_v2/rust/src/lib.rs @@ -29,8 +29,8 @@ pub use self::item_queue::ItemQueue; pub use self::request::{Request, RequestOptions, SearchOptions, ToolChoice}; pub use self::response::{FinishReason, Response, Usage}; pub use self::session::{ - AudioSession, ChatSession, CustomToolDefinition, EmbeddingsSession, ItemStream, Session, - ToolDefinition, + AudioSession, ChatSession, CustomToolDefinition, EmbeddingsSession, ItemStream, + RequestPreflightResult, Session, ToolDefinition, }; pub use self::types::{ ChatResponseFormat, ChatToolChoice, DeviceType, EpDownloadResult, EpInfo, ModelInfo, diff --git a/sdk_v2/rust/src/session.rs b/sdk_v2/rust/src/session.rs index fe56efa7b..39700d651 100644 --- a/sdk_v2/rust/src/session.rs +++ b/sdk_v2/rust/src/session.rs @@ -17,7 +17,9 @@ use std::task::{Context, Poll}; use tokio::sync::{mpsc::UnboundedReceiver, oneshot}; use crate::detail::api::{Api, Kvps}; -use crate::detail::ffi::{FOUNDRY_LOCAL_TOOL_KIND_CUSTOM, FOUNDRY_LOCAL_TOOL_KIND_FUNCTION}; +use crate::detail::ffi::{ + flRequestPreflightResult, FOUNDRY_LOCAL_TOOL_KIND_CUSTOM, FOUNDRY_LOCAL_TOOL_KIND_FUNCTION, +}; use crate::detail::model::Model; use crate::detail::session::{run_item_streaming, NativeItemQueue, NativeRequest, NativeSession}; use crate::detail::task::spawn_blocking; @@ -43,6 +45,36 @@ pub struct Session { inner: Arc, } +/// Exact token-budget preflight for a request captured against a chat session. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RequestPreflightResult { + /// Exact prompt tokens after native request preparation. + pub prompt_tokens: i64, + /// Tokens reserved for generated output. + pub output_reserve_tokens: i64, + /// Total tokens required by the request. + pub required_tokens: i64, + /// Model context-window limit. + pub context_limit_tokens: i64, + /// Whether the required tokens fit within the context limit. + pub fits: bool, + /// Tokens over budget, or zero when the request fits. + pub deficit_tokens: i64, +} + +impl From for RequestPreflightResult { + fn from(value: flRequestPreflightResult) -> Self { + Self { + prompt_tokens: value.prompt_tokens, + output_reserve_tokens: value.output_reserve_tokens, + required_tokens: value.required_tokens, + context_limit_tokens: value.context_limit_tokens, + fits: value.fits, + deficit_tokens: value.deficit_tokens, + } + } +} + impl Session { /// Open a session on a loaded model. /// @@ -456,6 +488,31 @@ impl ChatSession { spawn_blocking(move || inner.undo_turns(count)).await } + /// Capture this request and the current conversation state, then execute an exact token-budget + /// preflight on a blocking worker. + /// + /// Capture is completed synchronously before this method returns its future. The native + /// operation owns its in-memory state independently of later source-handle mutation or + /// destruction. URI-backed media is resolved during execution and is not frozen at capture. + /// Execution is one-shot. The captured operation may execute and be released on a worker thread, + /// but exclusive Rust ownership prevents execution from racing with release. + pub fn preflight_request( + &self, + request: Request, + ) -> impl std::future::Future> + Send + 'static { + let inner = &self.session.inner; + let captured = (|| { + let native = NativeRequest::new(Arc::clone(&inner.api))?; + populate_native_request(&inner.api, &native, &request)?; + inner.create_request_preflight(&native) + })(); + + async move { + let preflight = captured?; + spawn_blocking(move || preflight.execute().map(RequestPreflightResult::from)).await + } + } + /// Consume this handle, yielding the underlying base [`Session`]. pub fn into_session(self) -> Session { self.session