Internal change

LiteRT-LM-PiperOrigin-RevId: 890131978

Author: matthewchan-g

python/litert_lm/BUILD

@@ -59,22 +59,28 @@ cc_binary(
         "nowindows",
     ],
     deps = [
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:log_severity",
         "@com_google_absl//absl/log:globals",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
+        "@nlohmann_json//:json",
         "@nanobind",
         "@nanobind_json",
         "@litert//litert/c/internal:litert_logging",
         "//runtime/conversation",
+        "//runtime/conversation:io_types",
         "//runtime/core:engine_impl",
         "//runtime/engine:engine_factory",
         "//runtime/engine:engine_interface",
         "//runtime/engine:engine_settings",
-        "@rules_python//python/cc:current_py_cc_headers",
+        "//runtime/engine:io_types",
+        "//runtime/executor:executor_settings_base",
+        "@rules_python//python/cc:current_py_cc_headers",  # buildcleaner: keep
         "@litert//tflite:minimal_logging",
-        "@litert//tflite/core/c:private_c_api_types",
     ],
 )
 

python/litert_lm/engine_test.py

@@ -148,6 +148,14 @@ def test_create_conversation_with_messages(self):
     ):
       self.assertEqual(conversation.messages, messages)
 
+  def test_create_conversation_with_extra_context(self):
+    extra_context = {"key": "value"}
+    with (
+        self._create_engine() as engine,
+        engine.create_conversation(extra_context=extra_context) as conversation,
+    ):
+      self.assertEqual(conversation.extra_context, extra_context)
+
   def test_str_input_support(self):
     with (
         self._create_engine() as engine,

python/litert_lm/interfaces.py

@@ -20,7 +20,6 @@
 import collections.abc
 import dataclasses
 import enum
-import pathlib
 from typing import Any
 
 
@@ -103,6 +102,7 @@ def create_conversation(
           collections.abc.Sequence[collections.abc.Callable[..., Any]] | None
       ) = None,
       tool_event_handler: ToolEventHandler | None = None,
+      extra_context: collections.abc.Mapping[str, Any] | None = None,
   ) -> AbstractConversation:
     """Creates a new conversation for this engine.
 
@@ -111,6 +111,7 @@ def create_conversation(
           message is a mapping that should contain 'role' and 'content' keys.
         tools: A list of Python functions to be used as tools.
         tool_event_handler: A handler for tool call and tool response events.
+        extra_context: Extra context for the conversation.
     """
 
   @abc.abstractmethod
@@ -129,6 +130,7 @@ class AbstractConversation(abc.ABC):
       messages: A sequence of messages for the conversation preface.
       tools: A list of Python functions to be used as tools.
       tool_event_handler: A handler for tool call and tool response events.
+      extra_context: Extra context for the chat template.
   """
 
   def __init__(
@@ -141,6 +143,7 @@ def __init__(
           collections.abc.Sequence[collections.abc.Callable[..., Any]] | None
       ) = None,
       tool_event_handler: ToolEventHandler | None = None,
+      extra_context: collections.abc.Mapping[str, Any] | None = None,
   ):
     """Initializes the instance.
 
@@ -149,10 +152,12 @@ def __init__(
           message is a mapping that should contain 'role' and 'content' keys.
         tools: A list of Python functions to be used as tools.
         tool_event_handler: A handler for tool call and tool response events.
+        extra_context: Extra context for the chat template.
     """
     self.messages = messages or []
     self.tools = tools or []
     self.tool_event_handler = tool_event_handler
+    self.extra_context = extra_context or {}
 
   def __enter__(self) -> AbstractConversation:
     """Initializes the conversation."""

python/litert_lm/litert_lm.cc

@@ -13,30 +13,43 @@
 // limitations under the License.
 
 #include <deque>
+#include <exception>
 #include <map>
+#include <memory>
+#include <optional>
 #include <sstream>
 #include <stdexcept>
+#include <string>
 #include <string_view>
 #include <utility>
+#include <variant>
+#include <vector>
 
 #include "nanobind/nanobind.h"
 #include "nanobind/stl/optional.h"  // IWYU pragma: keep
-#include "nanobind/stl/shared_ptr.h"
+#include "nanobind/stl/shared_ptr.h"  // IWYU pragma: keep
 #include "nanobind/stl/string_view.h"  // IWYU pragma: keep
 #include "nanobind/stl/unique_ptr.h"   // IWYU pragma: keep
 #include "nanobind/stl/variant.h"      // IWYU pragma: keep
 #include "nanobind/stl/vector.h"       // IWYU pragma: keep
 #include "absl/base/log_severity.h"  // from @com_google_absl
+#include "absl/base/thread_annotations.h"  // from @com_google_absl
 #include "absl/log/globals.h"  // from @com_google_absl
+#include "absl/status/status.h"  // from @com_google_absl
 #include "absl/status/statusor.h"  // from @com_google_absl
+#include "absl/strings/string_view.h"  // from @com_google_absl
 #include "absl/synchronization/mutex.h"  // from @com_google_absl
 #include "absl/time/time.h"  // from @com_google_absl
+#include "nlohmann/json_fwd.hpp"  // from @nlohmann_json
 #include "nanobind_json/nanobind_json.hpp"  // from @nanobind_json  // IWYU pragma: keep
 #include "litert/c/internal/litert_logging.h"  // from @litert
 #include "runtime/conversation/conversation.h"
+#include "runtime/conversation/io_types.h"
 #include "runtime/engine/engine.h"
 #include "runtime/engine/engine_factory.h"
-#include "tflite/core/c/c_api_types.h"  // from @litert
+#include "runtime/engine/engine_settings.h"
+#include "runtime/engine/io_types.h"
+#include "runtime/executor/executor_settings_base.h"
 #include "tflite/logger.h"  // from @litert
 #include "tflite/minimal_logging.h"  // from @litert
 
@@ -468,7 +481,8 @@ NB_MODULE(litert_lm_ext, module) {
       .def(
           "create_conversation",
           [](const nb::object& self, const nb::handle& messages,
-             const nb::handle& tools, const nb::handle& tool_event_handler) {
+             const nb::handle& tools, const nb::handle& tool_event_handler,
+             const nb::handle& extra_context) {
             Engine& engine = nb::cast<Engine&>(self);
 
             auto builder = ConversationConfig::Builder();
@@ -503,6 +517,12 @@ NB_MODULE(litert_lm_ext, module) {
               has_preface = true;
             }
 
+            if (!extra_context.is_none()) {
+              json_preface.extra_context =
+                  nb::cast<nlohmann::json>(extra_context);
+              has_preface = true;
+            }
+
             if (has_preface) {
               builder.SetPreface(json_preface);
             }
@@ -515,6 +535,7 @@ NB_MODULE(litert_lm_ext, module) {
             nb::object py_conversation = nb::cast(std::move(conversation));
             py_conversation.attr("_tool_map") = py_tool_map;
             py_conversation.attr("tool_event_handler") = tool_event_handler;
+            py_conversation.attr("extra_context") = extra_context;
             if (messages.is_none()) {
               py_conversation.attr("messages") = nb::list();
             } else {
@@ -529,7 +550,8 @@ NB_MODULE(litert_lm_ext, module) {
           },
           nb::kw_only(), nb::arg("messages") = nb::none(),
           nb::arg("tools") = nb::none(),
-          nb::arg("tool_event_handler") = nb::none())
+          nb::arg("tool_event_handler") = nb::none(),
+          nb::arg("extra_context") = nb::none())
       .def(
           "create_session",
           [](Engine& self) {
@@ -557,6 +579,7 @@ NB_MODULE(litert_lm_ext, module) {
           "run_prefill",
           [](Engine::Session& self, const std::vector<std::string>& contents) {
             std::vector<InputData> input_data;
+            input_data.reserve(contents.size());
             for (const auto& text : contents) {
               input_data.emplace_back(InputText(text));
             }
@@ -580,6 +603,7 @@ NB_MODULE(litert_lm_ext, module) {
           [](Engine::Session& self, const std::vector<std::string>& target_text,
              bool store_token_lengths) {
             std::vector<absl::string_view> target_text_views;
+            target_text_views.reserve(target_text.size());
             for (const auto& text : target_text) {
               target_text_views.push_back(text);
             }

python/litert_lm_cli/main.py

@@ -226,9 +226,10 @@ def add_numbers(a: float, b: float) -> float:
         """Adds two numbers."""
         return a + b
 
-      # Provides the "system instruction" and "tools"
+      # Provides the "system instruction", "tools", and "extra_context"
       system_instruction = "You are a helpful assistant."
       tools = [add_numbers]
+      extra_context = {"key": "value"}
       ```
 
     Args:

python/litert_lm_cli/model.py

@@ -36,16 +36,16 @@
 
 
 def load_preset(preset: str):
-  """Loads a preset file and returns the tools and messages."""
+  """Loads a preset file and returns the tools, messages and extra_context."""
   print(f"Loading preset from {preset}:")
   if not os.path.exists(preset):
     print(f"Preset file not found: {preset}")
-    return None, None
+    return None, None, None
 
   spec = importlib.util.spec_from_file_location("user_tools", preset)
   if not spec or not spec.loader:
     print(f"Failed to load tools from {preset}")
-    return None, None
+    return None, None, None
 
   user_tools = importlib.util.module_from_spec(spec)
   spec.loader.exec_module(user_tools)
@@ -71,7 +71,11 @@ def load_preset(preset: str):
   for tool in tools:
     print(f"  - {getattr(tool, "__name__", str(tool))}")
 
-  return tools, messages
+  extra_context = getattr(user_tools, "extra_context", None)
+  if extra_context:
+    print(f"- Extra context: {extra_context}")
+
+  return tools, messages, extra_context
 
 
 _GREEN = "\033[92m"
@@ -147,9 +151,10 @@ def run_interactive(
 
       tools = None
       messages = None
+      extra_context = None
       if preset:
-        tools, messages = load_preset(preset)
-        if tools is None:
+        tools, messages, extra_context = load_preset(preset)
+        if tools is None and messages is None and extra_context is None:
           return
 
       handler = LoggingToolEventHandler() if tools else None
@@ -164,7 +169,10 @@ def run_interactive(
       with (
           engine_cm as engine,
           engine.create_conversation(
-              tools=tools, messages=messages, tool_event_handler=handler
+              tools=tools,
+              messages=messages,
+              tool_event_handler=handler,
+              extra_context=extra_context,
           ) as conversation,
       ):
         if prompt: