Internal change

LiteRT-LM-PiperOrigin-RevId: 908791856

Author: hheydary

runtime/conversation/BUILD

@@ -181,7 +181,6 @@ cc_test(
         "//runtime/components/constrained_decoding:bitmap",
         "//runtime/components/constrained_decoding:constraint",
         "//runtime/components/constrained_decoding:external_constraint_config",
-        "//runtime/core:session_factory",  # buildcleaner: keep
         "//runtime/engine:engine_factory",
         "//runtime/engine:engine_interface",
         "//runtime/engine:engine_settings",

runtime/core/BUILD

@@ -43,7 +43,6 @@ package_group(
 )
 
 ENGINE_IMPL_COMMON_DEPS = [
-    ":session_factory",
     "@com_google_absl//absl/base:no_destructor",
     "@com_google_absl//absl/log",
     "@com_google_absl//absl/log:absl_check",
@@ -97,6 +96,7 @@ cc_library(
     visibility = [":engine_impl_users"],
     deps = ENGINE_IMPL_COMMON_DEPS + [
         "//runtime/components:default_static_gpu_samplers",
+        "//runtime/core:session_basic",
         "//runtime/executor:default_static_gpu_accelerator",
     ] + select({
         "//conditions:default": [],
@@ -107,7 +107,9 @@ cc_library(
 cc_library(
     name = "engine_impl_cpu_only",
     srcs = ["engine_impl.cc"],
-    deps = ENGINE_IMPL_COMMON_DEPS,
+    deps = ENGINE_IMPL_COMMON_DEPS + [
+        "//runtime/core:session_basic",
+    ],
     alwayslink = 1,
 )
 
@@ -307,44 +309,6 @@ cc_test(
     }),
 )
 
-cc_library(
-    name = "session_factory",
-    srcs = ["session_factory.cc"],
-    hdrs = ["session_factory.h"],
-    deps = [
-        ":session_basic",
-        "@com_google_absl//absl/base:nullability",
-        "@com_google_absl//absl/status:statusor",
-        "//runtime/components:tokenizer",
-        "//runtime/engine:engine_interface",
-        "//runtime/engine:engine_settings",
-        "//runtime/engine:io_types",
-        "//runtime/executor:audio_executor",
-        "//runtime/executor:llm_executor",
-        "//runtime/executor:vision_executor",
-        "//runtime/framework:threadpool",
-        "//runtime/proto:sampler_params_cc_proto",
-        "//runtime/util:litert_status_util",
-    ],
-)
-
-cc_test(
-    name = "session_factory_test",
-    srcs = ["session_factory_test.cc"],
-    deps = [
-        ":session_factory",
-        "@com_google_googletest//:gtest_main",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:string_view",
-        "//runtime/components:tokenizer",
-        "//runtime/engine:engine_settings",
-        "//runtime/executor:executor_settings_base",
-        "//runtime/executor:fake_llm_executor",
-        "//runtime/framework:threadpool",
-        "//runtime/util:test_utils",
-    ],
-)
-
 cc_library(
     name = "session_utils",
     srcs = ["session_utils.cc"],

runtime/core/engine_advanced_impl.cc

@@ -36,7 +36,7 @@
 #include "litert/cc/litert_macros.h"  // from @litert
 #include "runtime/components/model_resources.h"
 #include "runtime/components/tokenizer.h"
-#include "runtime/core/session_factory.h"
+#include "runtime/core/session_advanced.h"
 #include "runtime/engine/engine.h"
 #include "runtime/engine/engine_factory.h"
 #include "runtime/engine/engine_settings.h"
@@ -184,8 +184,8 @@ class EngineAdvancedImpl : public Engine {
 
     ASSIGN_OR_RETURN(
         auto session,
-        InitializeSessionAdvanced(execution_manager_, tokenizer_.get(), config,
-                                  std::move(session_benchmark_info)));
+        SessionAdvanced::Create(execution_manager_, tokenizer_.get(), config,
+                                std::move(session_benchmark_info)));
 
     if (benchmark_info_.has_value()) {
       auto session_benchmark_info_or = session->GetMutableBenchmarkInfo();

runtime/core/engine_impl.cc

@@ -36,7 +36,7 @@
 #include "litert/cc/litert_macros.h"  // from @litert
 #include "runtime/components/model_resources.h"
 #include "runtime/components/tokenizer.h"
-#include "runtime/core/session_factory.h"
+#include "runtime/core/session_basic.h"
 #include "runtime/engine/engine.h"
 #include "runtime/engine/engine_factory.h"
 #include "runtime/engine/engine_settings.h"
@@ -187,11 +187,11 @@ class EngineImpl : public Engine {
     }
     ASSIGN_OR_RETURN(
         auto session,
-        InitializeSessionBasic(executor_.get(), tokenizer_.get(),
-                               /*vision_executor=*/vision_executor_.get(),
-                               /*audio_executor=*/audio_executor_.get(), config,
-                               std::move(session_benchmark_info),
-                               worker_thread_pool_.get()));
+        SessionBasic::Create(executor_.get(), tokenizer_.get(),
+                             /*vision_executor=*/vision_executor_.get(),
+                             /*audio_executor=*/audio_executor_.get(), config,
+                             std::move(session_benchmark_info),
+                             worker_thread_pool_.get()));
     if (benchmark_info_.has_value()) {
       auto session_benchmark_info_or = session->GetMutableBenchmarkInfo();
       if (session_benchmark_info_or.ok()) {

runtime/core/session_factory.cc

@@ -15,8 +15,10 @@
 #ifndef THIRD_PARTY_ODML_LITERT_LM_RUNTIME_EXECUTOR_MOCK_LLM_EXECUTOR_H_
 #define THIRD_PARTY_ODML_LITERT_LM_RUNTIME_EXECUTOR_MOCK_LLM_EXECUTOR_H_
 
+#include <cstdint>
 #include <memory>
 #include <optional>
+#include <utility>
 #include <vector>
 
 #include "absl/status/status.h"  // from @com_google_absl
@@ -85,6 +87,23 @@ class FakeLlmExecutor : public LlmExecutor {
   };
   absl::StatusOr<int> GetCurrentStep() const override { return current_step_; }
 
+  absl::StatusOr<std::unique_ptr<LlmContext>> CreateNewContext(
+      std::optional<uint32_t> lora_id,
+      RuntimeConfig runtime_config) const override {
+    return std::make_unique<LlmContext>(
+        nullptr, std::make_unique<RuntimeConfig>(std::move(runtime_config)),
+        std::make_unique<RuntimeState>());
+  };
+
+  absl::Status RestoreContext(
+      std::unique_ptr<LlmContext> llm_context) override {
+    return absl::OkStatus();
+  };
+
+  absl::StatusOr<const ProcessedTokens*> GetProcessedTokens() const override {
+    return &processed_tokens_;
+  }
+
   absl::Status SetCurrentStep(int current_step) override {
     current_step_ = current_step;
     if (current_step >= prefill_tokens_total_) {
@@ -95,6 +114,11 @@ class FakeLlmExecutor : public LlmExecutor {
     return absl::OkStatus();
   }
 
+  absl::Status UpdateRuntimeConfig(
+      const RuntimeConfig& runtime_config) override {
+    return absl::OkStatus();
+  }
+
   // Sets the status to be returned by the Prefill function.
   void SetPrefillStatus(const absl::Status& status) {
     prefill_status_ = status;