Internal change

LiteRT-LM-PiperOrigin-RevId: 859298975

Author: ai-edge-bot

c/BUILD

@@ -51,7 +51,8 @@ cc_library(
         "@litert//litert/c/internal:litert_logging",
         "//runtime/conversation",
         "//runtime/conversation:io_types",
-        "//runtime/engine:engine_impl_selected",
+        "//runtime/core:engine_impl",
+        "//runtime/engine:engine_factory",
         "//runtime/engine:engine_interface",
         "//runtime/engine:engine_settings",
         "//runtime/engine:io_types",

c/engine.cc

@@ -32,6 +32,7 @@
 #include "runtime/conversation/conversation.h"
 #include "runtime/conversation/io_types.h"
 #include "runtime/engine/engine.h"
+#include "runtime/engine/engine_factory.h"
 #include "runtime/engine/engine_settings.h"
 #include "runtime/engine/io_types.h"
 #include "runtime/executor/executor_settings_base.h"
@@ -92,6 +93,7 @@ CreateConversationCallback(LiteRtLmStreamCallback callback, void* user_data) {
 using ::litert::lm::Conversation;
 using ::litert::lm::ConversationConfig;
 using ::litert::lm::Engine;
+using ::litert::lm::EngineFactory;
 using ::litert::lm::EngineSettings;
 using ::litert::lm::InputText;
 using ::litert::lm::JsonMessage;
@@ -329,7 +331,10 @@ LiteRtLmEngine* litert_lm_engine_create(
     return nullptr;
   }
 
-  auto engine = Engine::CreateEngine(*settings->settings);
+  absl::StatusOr<std::unique_ptr<Engine>> engine;
+    engine = EngineFactory::CreateDefault(*settings->settings);
+
+
   if (!engine.ok()) {
     ABSL_LOG(ERROR) << "Failed to create engine: " << engine.status();
     return nullptr;
@@ -339,6 +344,7 @@ LiteRtLmEngine* litert_lm_engine_create(
   c_engine->engine = *std::move(engine);
   return c_engine;
 }
+
 void litert_lm_engine_delete(LiteRtLmEngine* engine) { delete engine; }
 
 LiteRtLmSession* litert_lm_engine_create_session(

runtime/core/BUILD

@@ -35,6 +35,8 @@ package_group(
         "//third_party/odml/litert_lm/js/packages/core/src/cpp",
         # for e2e tests.
         "//runtime/e2e_tests/...",
+        "//third_party/odml/litert_lm/c/...",
+        "//c/...",
     ],
 )
 

runtime/core/engine_advanced_legacy_impl.cc

@@ -54,11 +54,11 @@
 
 namespace litert::lm {
 
-absl::StatusOr<std::unique_ptr<Engine>> CreateEngineAdvancedLegacy(
-    EngineSettings engine_settings);
-
 namespace {
 
+absl::StatusOr<std::unique_ptr<Engine>> CreateEngineAdvancedLegacy(
+    EngineSettings engine_settings, absl::string_view input_prompt_as_hint);
+
 namespace oi = ::odml::infra;
 
 class EngineAdvancedLegacyImpl : public Engine {
@@ -74,7 +74,8 @@ class EngineAdvancedLegacyImpl : public Engine {
 
  private:
   friend absl::StatusOr<std::unique_ptr<Engine>>
-  litert::lm::CreateEngineAdvancedLegacy(EngineSettings engine_settings);
+  CreateEngineAdvancedLegacy(EngineSettings engine_settings,
+                               absl::string_view input_prompt_as_hint);
 
   explicit EngineAdvancedLegacyImpl(
       EngineSettings engine_settings,
@@ -177,18 +178,10 @@ const EngineSettings& EngineAdvancedLegacyImpl::GetEngineSettings() const {
   return engine_settings_;
 }
 
-}  // namespace
-
-// Method to create Engine.
-absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
-    EngineSettings settings, absl::string_view input_prompt_as_hint) {
-  return CreateEngineAdvancedLegacy(std::move(settings));
-}
-
-// Method to create Engine.
+// The factory for Advanced Legacy TfLite.
 absl::StatusOr<std::unique_ptr<Engine>> CreateEngineAdvancedLegacy(
-    EngineSettings engine_settings) {
-  ABSL_LOG(INFO) << "Constructing legacy EngineImpl...";
+    EngineSettings engine_settings, absl::string_view input_prompt_as_hint) {
+  ABSL_LOG(INFO) << "Constructing advanced legacy EngineImpl...";
   std::optional<BenchmarkInfo> benchmark_info;
   if (engine_settings.IsBenchmarkEnabled()) {
     benchmark_info = std::make_optional<BenchmarkInfo>(
@@ -237,8 +230,8 @@ absl::StatusOr<std::unique_ptr<Engine>> CreateEngineAdvancedLegacy(
   }
   // Update and load the parameters from the model file and convert the tokens
   // to ids.
-  RETURN_IF_ERROR(
-      engine_settings.MaybeUpdateAndValidate(*tokenizer, &llm_metadata));
+  RETURN_IF_ERROR(engine_settings.MaybeUpdateAndValidate(
+      *tokenizer, &llm_metadata, input_prompt_as_hint));
 
   ASSIGN_OR_RETURN(auto executor,
                    BuildExecutor(*model_resources, engine_settings));
@@ -301,11 +294,15 @@ absl::StatusOr<std::unique_ptr<Engine>> CreateEngineAdvancedLegacy(
   return llm_impl;
 };
 
+}  // namespace
+
 LITERT_LM_REGISTER_ENGINE(EngineFactory::EngineType::kAdvancedLegacyTfLite,
-                          [](EngineSettings settings,
-                             absl::string_view input_prompt_as_hint) {
-                            return Engine::CreateEngine(std::move(settings),
-                                                        input_prompt_as_hint);
-                          });
+                          CreateEngineAdvancedLegacy);
+
+// Method to create Engine.
+absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
+    EngineSettings settings, absl::string_view input_prompt_as_hint) {
+  return CreateEngineAdvancedLegacy(std::move(settings), input_prompt_as_hint);
+}
 
 }  // namespace litert::lm

runtime/core/engine_impl.cc

@@ -136,8 +136,6 @@ absl::StatusOr<Environment&> GetEnvironment(
   return **kEnvironment;
 }
 
-}  // namespace
-
 class EngineImpl : public Engine {
  public:
   ~EngineImpl() override {
@@ -224,8 +222,7 @@ class EngineImpl : public Engine {
   std::unique_ptr<ThreadPool> worker_thread_pool_;
 };
 
-// Method to create Engine.
-absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
+absl::StatusOr<std::unique_ptr<Engine>> CreateEngineCompiled(
     EngineSettings engine_settings, absl::string_view input_prompt_as_hint) {
   std::optional<BenchmarkInfo> benchmark_info =
       engine_settings.IsBenchmarkEnabled()
@@ -329,12 +326,15 @@ absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
 
   return llm_impl;
 };
+}  // namespace
 
 LITERT_LM_REGISTER_ENGINE(EngineFactory::EngineType::kLiteRTCompiledModel,
-                          [](EngineSettings settings,
-                             absl::string_view input_prompt_as_hint) {
-                            return Engine::CreateEngine(std::move(settings),
-                                                        input_prompt_as_hint);
-                          });
+                          CreateEngineCompiled);
+
+// Method to create Engine.
+absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
+    EngineSettings engine_settings, absl::string_view input_prompt_as_hint) {
+  return CreateEngineCompiled(std::move(engine_settings), input_prompt_as_hint);
+}
 
 }  // namespace litert::lm

runtime/core/engine_legacy_impl.cc

@@ -132,6 +132,7 @@ absl::StatusOr<Environment&> GetEnvironment() {
 }
 }  // namespace
 
+namespace {
 class EngineImpl : public Engine {
  public:
   ~EngineImpl() override {
@@ -231,8 +232,7 @@ class EngineImpl : public Engine {
   std::unique_ptr<ThreadPool> worker_thread_pool_;
 };
 
-// Method to create Engine.
-absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
+absl::StatusOr<std::unique_ptr<Engine>> CreateEngineLegacy(
     EngineSettings engine_settings, absl::string_view input_prompt_as_hint) {
   ABSL_LOG(INFO) << "Constructing legacy EngineImpl...";
   std::optional<BenchmarkInfo> benchmark_info;
@@ -369,12 +369,15 @@ absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
       std::move(benchmark_info), std::move(worker_thread_pool));
   return llm_impl;
 };
+}  // namespace
 
 LITERT_LM_REGISTER_ENGINE(EngineFactory::EngineType::kLegacyTfLite,
-                          [](EngineSettings settings,
-                             absl::string_view input_prompt_as_hint) {
-                            return Engine::CreateEngine(std::move(settings),
-                                                        input_prompt_as_hint);
-                          });
+                          CreateEngineLegacy);
+
+// Method to create Engine.
+absl::StatusOr<std::unique_ptr<Engine>> Engine::CreateEngine(
+    EngineSettings engine_settings, absl::string_view input_prompt_as_hint) {
+  return CreateEngineLegacy(std::move(engine_settings), input_prompt_as_hint);
+}
 
 }  // namespace litert::lm