langfun 0.1.2.dev202510270805__py3-none-any.whl → 0.1.2.dev202510290805__py3-none-any.whl

langfun/env/interface.py

@@ -170,6 +170,9 @@ class SessionTeardownError(SandboxError):
 class Environment(pg.Object):
   """Base class for an environment."""
 
+  # Disable symbolic comparison and hashing for environment objects.
+  use_symbolic_comparison = False
+
   @dataclasses.dataclass(frozen=True)
   class Id:
     """Identifier for an environment."""
@@ -220,6 +223,20 @@ class Environment(pg.Object):
   def id(self) -> Id:
     """Returns the identifier for the environment."""
 
+  @property
+  @abc.abstractmethod
+  def image_ids(self) -> list[str]:
+    """Returns the non-dynamic image IDs served by the environment."""
+
+  def image_id_for(self, feature: 'Feature') -> str:
+    """Returns the default image ID for the environment."""
+    for image_id in self.image_ids:
+      if feature.is_applicable(image_id):
+        return image_id
+    raise ValueError(
+        f'No image ID found for feature {feature.name} in {self.image_ids}.'
+    )
+
   @property
   @abc.abstractmethod
   def status(self) -> Status:
@@ -245,9 +262,16 @@ class Environment(pg.Object):
     """
 
   @abc.abstractmethod
-  def acquire(self) -> 'Sandbox':
+  def acquire(
+      self,
+      image_id: str | None = None,
+  ) -> 'Sandbox':
     """Acquires a free sandbox from the environment.
 
+    Args:
+      image_id: The image ID to use for the sandbox. If None, it will be
+        automatically determined by the environment.
+
     Returns:
       A free sandbox from the environment.
 
@@ -257,7 +281,7 @@ class Environment(pg.Object):
     """
 
   @abc.abstractmethod
-  def new_session_id(self) -> str:
+  def new_session_id(self, feature_hint: str | None = None) -> str:
     """Generates a new session ID."""
 
   #
@@ -298,33 +322,63 @@ class Environment(pg.Object):
 
   def sandbox(
       self,
+      image_id: str | None = None,
       session_id: str | None = None,
   ) -> ContextManager['Sandbox']:
     """Gets a sandbox from the environment and starts a new user session."""
-    return self.acquire().new_session(session_id)
+    return self.acquire(image_id=image_id).new_session(session_id)
 
   def __getattr__(self, name: str) -> Any:
-    """Gets a feature from a free sandbox from the environment.
+    """Gets a feature session from a free sandbox from the environment.
 
     Example:
       ```
       with XboxEnvironment(
           features={'selenium': SeleniumFeature()}
       ) as env:
-        driver = env.selenium.get_driver()
+        with env.selenium() as selenium:
+          driver = selenium.get_driver()
       ```
 
     Args:
       name: The name of the feature.
 
     Returns:
-      A feature from a free sandbox from the environment.
+      A callable `(image_id, *, session_id) -> ContextManager[Feature]` that
+      creates a context manager for the requested feature under a new client
+      session.
     """
     if name in self.features:
-      return self.acquire().features[name]
+      return _feature_session_creator(self, self.features[name])
     raise AttributeError(name)
 
 
+@contextlib.contextmanager
+def _session_for_feature(
+    environment: Environment,
+    feature: 'Feature',
+    image_id: str | None = None,
+    session_id: str | None = None,
+) -> Iterator['Feature']:
+  """Returns a context manager for a session for a feature."""
+  if image_id is None:
+    image_id = environment.image_id_for(feature)
+  elif not feature.is_applicable(image_id):
+    raise ValueError(
+        f'Feature {feature.name!r} is not applicable to image {image_id!r}.'
+    )
+  sandbox = environment.acquire(image_id=image_id)
+  with sandbox.new_session(session_id=session_id, feature_hint=feature.name):
+    yield sandbox.features[feature.name]
+
+
+def _feature_session_creator(environment: Environment, feature: 'Feature'):
+  """Returns a callable that returns a context manager for a feature session."""
+  def fn(image_id: str | None = None, session_id: str | None = None):
+    return _session_for_feature(environment, feature, image_id, session_id)
+  return fn
+
+
 # Enable automatic conversion from str to Environment.Id.
 pg.typing.register_converter(str, Environment.Id, Environment.Id)
 
@@ -332,14 +386,18 @@ pg.typing.register_converter(str, Environment.Id, Environment.Id)
 class Sandbox(pg.Object):
   """Interface for sandboxes."""
 
+  # Disable symbolic comparison and hashing for sandbox objects.
+  use_symbolic_comparison = False
+
   @dataclasses.dataclass(frozen=True, slots=True)
   class Id:
     """Identifier for a sandbox."""
     environment_id: Environment.Id
+    image_id: str
     sandbox_id: str
 
     def __str__(self) -> str:
-      return f'{self.environment_id}/{self.sandbox_id}'
+      return f'{self.environment_id}/{self.image_id}:{self.sandbox_id}'
 
     def working_dir(self, root_dir: str | None) -> str | None:
       """Returns the download directory for the sandbox."""
@@ -347,6 +405,7 @@ class Sandbox(pg.Object):
         return None
       return os.path.join(
           self.environment_id.working_dir(root_dir),
+          _make_path_compatible(self.image_id),
           _make_path_compatible(self.sandbox_id)
       )
 
@@ -438,6 +497,11 @@ class Sandbox(pg.Object):
   def id(self) -> Id:
     """Returns the identifier for the sandbox."""
 
+  @property
+  @abc.abstractmethod
+  def image_id(self) -> str:
+    """Returns the image ID used for bootstrapping the sandbox."""
+
   @property
   @abc.abstractmethod
   def environment(self) -> Environment:
@@ -626,6 +690,8 @@ class Sandbox(pg.Object):
   def new_session(
       self,
       session_id: str | None = None,
+      *,
+      feature_hint: str | None = None,
   ) -> Iterator['Sandbox']:
     """Context manager for obtaining a sandbox for a user session.
 
@@ -636,6 +702,9 @@ class Sandbox(pg.Object):
     Args:
       session_id: The identifier for the user session. If not provided, a random
         ID will be generated.
+      feature_hint: A hint of which feature is the main user intent when
+        starting the session. This is used for generating the session id if
+        `session_id` is not provided.
 
     Yields:
       The sandbox for the user session.
@@ -646,7 +715,7 @@ class Sandbox(pg.Object):
         errors.
     """
     if session_id is None:
-      session_id = self.environment.new_session_id()
+      session_id = self.environment.new_session_id(feature_hint)
     self.start_session(session_id)
     try:
       yield self
@@ -683,6 +752,9 @@ class Sandbox(pg.Object):
 class Feature(pg.Object):
   """Interface for sandbox features."""
 
+  # Disable symbolic comparison and hashing for sandbox objects.
+  allow_symbolic_comparison = False
+
   @property
   @abc.abstractmethod
   def name(self) -> str:
@@ -700,6 +772,10 @@ class Feature(pg.Object):
       AssertError: If the feature is not set up with a sandbox yet.
     """
 
+  @abc.abstractmethod
+  def is_applicable(self, image_id: str) -> bool:
+    """Returns True if the feature is applicable to the given image."""
+
   @abc.abstractmethod
   def setup(self, sandbox: Sandbox) -> None:
     """Sets up the feature, which is called once when the sandbox is up.

langfun/env/interface_test.py

@@ -29,12 +29,13 @@ class IdTest(unittest.TestCase):
   def test_sandbox_id(self):
     sandbox_id = interface.Sandbox.Id(
         environment_id=interface.Environment.Id('env'),
+        image_id='image:2025_01_01_00_00_00',
         sandbox_id='sandbox'
     )
-    self.assertEqual(str(sandbox_id), 'env/sandbox')
+    self.assertEqual(str(sandbox_id), 'env/image:2025_01_01_00_00_00:sandbox')
     self.assertEqual(
         sandbox_id.working_dir(root_dir='/tmp'),
-        '/tmp/env/sandbox'
+        '/tmp/env/image_2025_01_01_00_00_00/sandbox'
     )
     self.assertIsNone(sandbox_id.working_dir(root_dir=None))
 

langfun/env/load_balancers_test.py

@@ -22,6 +22,7 @@ from langfun.env import load_balancers
 class TestingSandbox(interface.Sandbox):
   sandbox_id: str
   status: interface.Sandbox.Status = interface.Sandbox.Status.READY
+  image_id: str = 'test_image'
 
   def _on_bound(self) -> None:
     super()._on_bound()
@@ -31,6 +32,7 @@ class TestingSandbox(interface.Sandbox):
   def id(self) -> interface.Sandbox.Id:
     return interface.Sandbox.Id(
         environment_id=interface.Environment.Id('testing-env'),
+        image_id=self.image_id,
         sandbox_id=self.sandbox_id
     )
 

langfun/env/test_utils.py

@@ -27,7 +27,7 @@ import pyglove as pg
 
 class TestingEnvironment(base_environment.BaseEnvironment):
   """Testing environment for unit tests."""
-
+  image_ids: list[str] = ['test_image']
   housekeep_interval: float = 0.0
   simulate_start_error: Type[BaseException] | None = None
   simulate_shutdown_error: Type[BaseException] | None = None
@@ -45,6 +45,7 @@ class TestingEnvironment(base_environment.BaseEnvironment):
 
   def _create_sandbox(
       self,
+      image_id: str,
       sandbox_id: str,
       reusable: bool,
       proactive_session_setup: bool,
@@ -54,8 +55,10 @@ class TestingEnvironment(base_environment.BaseEnvironment):
         environment=self,
         id=interface.Sandbox.Id(
             environment_id=self.id,
+            image_id=image_id,
             sandbox_id=sandbox_id
         ),
+        image_id=image_id,
         reusable=reusable,
         proactive_session_setup=proactive_session_setup,
         keepalive_interval=keepalive_interval,

{langfun-0.1.2.dev202510270805.dist-info → langfun-0.1.2.dev202510290805.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202510270805
+Version: 0.1.2.dev202510290805
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202510270805.dist-info → langfun-0.1.2.dev202510290805.dist-info}/RECORD

@@ -6,7 +6,7 @@ langfun/assistant/capabilities/gui/drawing.py,sha256=8wgol61P7HovLg5EaevRmDPXTFu
 langfun/assistant/capabilities/gui/drawing_test.py,sha256=d6LQ1ctG78YRi58UVBdndwyEqTC_ITdk191oA3tASxM,3421
 langfun/assistant/capabilities/gui/location.py,sha256=QYJlY5kUNEwiZFiPYRyFAO7bD2ez4jL5hYn1_AK1ulc,8643
 langfun/assistant/capabilities/gui/location_test.py,sha256=pQUOH1sKuAwjTTYKu615RUnecc_lNrddfvkyxf9297w,9051
-langfun/core/__init__.py,sha256=GC4amBXybjSfYVrEQSLFb_KBhH75uD95GKkHNiJfwkY,5011
+langfun/core/__init__.py,sha256=-HrGZb0Gv4nH6R2I2dXOVfkWKrGQfEu9UHBmXsG-x-E,5056
 langfun/core/async_support.py,sha256=WF4sflm0Q-UHJ8lPtlEo9hSwqXqm1kfaAYVPTyVP3n0,3062
 langfun/core/async_support_test.py,sha256=fMz1ulGrfUCuHp7RtYktuBwpbRN3kCBKnB4LFvaXSAA,1754
 langfun/core/component.py,sha256=g1kQM0bryYYYWVDrSMnHfc74wIBbpfe5_B3s-UIP5GE,3028
@@ -22,7 +22,7 @@ langfun/core/language_model_test.py,sha256=UIQpovmVfJnJ2EDOi9OZG_1UNClWs9lkTaE22
 langfun/core/logging.py,sha256=7IGAhp7mGokZxxqtL-XZvFLKaZ5k3F5_Xp2NUtR4GwE,9136
 langfun/core/logging_test.py,sha256=vbVGOQxwMmVSiFfbt2897gUt-8nqDpV64jCAeUG_q5U,6924
 langfun/core/memory.py,sha256=vyXVvfvSdLLJAzdIupnbn3k26OgclCx-OJ7gddS5e1Y,2070
-langfun/core/message.py,sha256=m36GDLtuDVUURkYORTEj61Uho6eILwSkCjrksK0HLWs,33964
+langfun/core/message.py,sha256=DShgb0pHCfy8IvMzFccTO6mG0WS3O226UZNoxzNOGF4,34100
 langfun/core/message_test.py,sha256=hjnzxJbjLNau30Ex7h5B0qC8qUU4uscoUcEVayRiXFQ,38355
 langfun/core/modality.py,sha256=h56hfCyxX_rK_BU3pFxXpGD79TvFeRXVo71QvpzmZiw,6596
 langfun/core/modality_test.py,sha256=0MQz6e39XOd2UjAIc5Vq9Eems1IUgOxD7Skf0gTMGTg,3665
@@ -53,7 +53,7 @@ langfun/core/coding/python/sandboxing.py,sha256=ppeE4E9rwVpMO2wxZN6dA2OwSMyD2hQ9
 langfun/core/coding/python/sandboxing_test.py,sha256=H_0_pd-_uS-ci5yYhmDTR6-hyzosAFkExziAHndfdDo,2023
 langfun/core/data/__init__.py,sha256=qllw9ST1vveZv-1E0VM5hezn1YH-OcqGI-yFqQYnWgI,732
 langfun/core/data/conversion/__init__.py,sha256=ZcGntBruvvZSfESO-Tha1nzHfgWEK7I1u78Jw8gsoVU,883
-langfun/core/data/conversion/anthropic.py,sha256=0xf1sWt6WhmG8EsaeLY4Ee7S3FcDWKPOtn6caCBhtrQ,4431
+langfun/core/data/conversion/anthropic.py,sha256=8JceJL4Z7qxgprKaByT1uDiyxpUTeHl2rgVEWRHr1Ms,4511
 langfun/core/data/conversion/anthropic_test.py,sha256=BMIyLgFOtFuJYJAkSLR1nwG7UOwtoT61R5xAXtJfSMY,8624
 langfun/core/data/conversion/gemini.py,sha256=mLft9j30W8aIi5a7tOdjfO2v8O8a61FSDKf2k-co93M,5825
 langfun/core/data/conversion/gemini_test.py,sha256=LeSToAqEhBHekSCjHFkC_cLN7Xv_cn-nnCURS95511M,7366
@@ -123,9 +123,11 @@ langfun/core/llms/cache/in_memory.py,sha256=i58oiQL28RDsq37dwqgVpC2mBETJjIEFS20y
 langfun/core/llms/cache/in_memory_test.py,sha256=AG9OGdxfyW9LdteFRF-aOIhxiH5Kna5U-pQk91ljHCQ,10538
 langfun/core/mcp/__init__.py,sha256=cyVP_YTjOmbjhYg8BE7-RnE4Txt8wDukYC0anhHKpuo,286
 langfun/core/mcp/client.py,sha256=NBo4W2fdvYyywlUHynAbyh7LNohILccTkCiisj-UEEg,3642
-langfun/core/mcp/client_test.py,sha256=Vjbh7bXM8xXMMe_QOnFJ4rE0wITovAZbwk7Nu-cqf8g,2766
-langfun/core/mcp/session.py,sha256=dyuNXpLvM0TgLoUmWEwy5XCRf-IK7H_BbcNI2oFa_ek,5626
-langfun/core/mcp/tool.py,sha256=eap46OK4lrw5G5uCwPi4kUMNdaN1Bgy0cUi0bdLJI5Q,3607
+langfun/core/mcp/client_test.py,sha256=ytplbrnFDqnHkDQuf5UGJaFf-rXUtwONCJcmNZ_rNUY,2221
+langfun/core/mcp/session.py,sha256=yFjt0TzV0HqkanFmcTb_JvTkJor2KTEdMoAsCSu3KuE,5163
+langfun/core/mcp/session_test.py,sha256=V4UVH9uVDL10RUWUr4jldcZqn4HYB75WgkjAnNg6cVw,1952
+langfun/core/mcp/tool.py,sha256=PRh4zHRXaSpYmjSrhh0vvoOs-Gi-4plDucKcM7o0umY,6056
+langfun/core/mcp/tool_test.py,sha256=oKsIpkyb-2Pr7_Fyqn1WNX8B07_izfdUyo4fIAp7x3o,6372
 langfun/core/mcp/testing/simple_mcp_client.py,sha256=U5pUXa0SnB2-DdRI2vPhrdVUv14dX_OeS4mPLFpllMc,924
 langfun/core/mcp/testing/simple_mcp_server.py,sha256=dw4t6ERWMdmi6kDE38RU5oYu5MQbEX-GJ6CMxGcV-WE,994
 langfun/core/memories/__init__.py,sha256=HpghfZ-w1NQqzJXBx8Lz0daRhB2rcy2r9Xm491SBhC4,773
@@ -173,23 +175,23 @@ langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3AB
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
 langfun/env/__init__.py,sha256=m6Y8y16ms9lytvO_r-Br8HmTp2rNDhb3R6JJaH5dEEk,1491
-langfun/env/base_environment.py,sha256=m7IJNHqt4dlgBRwv0E8BxKtARvcNJis9FrrALguzX1E,19472
-langfun/env/base_feature.py,sha256=gGUCWLcSaQZ7MoXNKogXEdcpfoqJxcfnoXhWQjqi8jk,6462
-langfun/env/base_sandbox.py,sha256=yG_4qG5MVIEP5uyXO2KssZFSFDewwNvF7Hr_qE7yGfQ,37301
-langfun/env/base_test.py,sha256=DrBtyweO4v8Fz3Oz-gnpJ4W_9hRhuVXA1CTLvXJDa9s,61099
-langfun/env/interface.py,sha256=2Amf-_op7dGRF8c4-wYxcFxs1UCBYz1AB20Lk7__V4E,25724
-langfun/env/interface_test.py,sha256=hfQn4RRTEo1YfVHXTPzH1puzD14BTo8R_5v1IpXVZ90,1398
+langfun/env/base_environment.py,sha256=8Cpwb4D37zHj5AK4qgmGIPyxIMEDxTIPHXqxTSNkE_M,25030
+langfun/env/base_feature.py,sha256=JDEhL9LkbBHB0c603guchry7cy_zaIReg5vqExyQQgg,6902
+langfun/env/base_sandbox.py,sha256=aBzPWJGszWuwv8BObvBoIVtl8aiOp41YjTEfVbYYDKQ,37425
+langfun/env/base_test.py,sha256=PqFLqT-IwOe3zH4JWeCtUdDjSqG9KQ4RBQebo67AIqc,74983
+langfun/env/interface.py,sha256=_XguM6xjPQuiTUUrs-4_3ejm8rnolO5cfxpuYE6yoX8,28520
+langfun/env/interface_test.py,sha256=1ZeLPsB9QUyervNrg2LdkJeYx-rQYQ1Zq_IzIKJtS5M,1496
 langfun/env/load_balancers.py,sha256=qRhCthqzjZIQBwta8qC1C0s0J-VQAVomJQqI7Nqv-r4,1948
-langfun/env/load_balancers_test.py,sha256=73QQG-Lufy5DBhQuzWNUaaH24YiZFdZEHe67TkJn7a0,3708
-langfun/env/test_utils.py,sha256=rnIQZ2ZpiFuB7yfl5ckGtacBoySAwguzBzXzqhyo8jw,14042
+langfun/env/load_balancers_test.py,sha256=E1ksJ5Xsoj5aJTM4i9EKDiWKDFljs75FpoeJpbtCKL8,3771
+langfun/env/test_utils.py,sha256=5_lCkTvXC1J9unPdgmqn5PxaucU7GCcro0iFtjehJo0,14160
 langfun/env/event_handlers/__init__.py,sha256=H34n-TbKSgtxqBhE-yAti8fY6weF2_v3yw59M9_zmGM,443
 langfun/env/event_handlers/base.py,sha256=eGdJ6N5em9kX-c9wzm1TdnRP5_5IAltX5JTHILdjzLM,10124
 langfun/env/event_handlers/event_logger.py,sha256=3dbPjBe53dBgntYHlyLlj_77hVecPSXkmKeiGXMxlO0,12699
-langfun/env/event_handlers/event_logger_test.py,sha256=PGof3rPllNnyzs3Yp8kaOHLeTkVrzUgCJwlODTrVRKI,9111
-langfun/env/event_handlers/metric_writer.py,sha256=NgJKsd6xWOtEd0IjYi7coGEaqGYkkPcDjXN9CQ3vxPU,18043
-langfun/env/event_handlers/metric_writer_test.py,sha256=deqCIdo-Q89rRCF2lerU3GpWmFG_SvDx162UQIZxPgE,5338
-langfun-0.1.2.dev202510270805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202510270805.dist-info/METADATA,sha256=cQSg_RGaOWibKP72oV20cfWjqzwXA9oMRhDjRILfApQ,7522
-langfun-0.1.2.dev202510270805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-langfun-0.1.2.dev202510270805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202510270805.dist-info/RECORD,,
+langfun/env/event_handlers/event_logger_test.py,sha256=ryupxaEP9D8wdtSsSwZRSZwqFaHCaSD-bFSea_T7QJo,9133
+langfun/env/event_handlers/metric_writer.py,sha256=F_Gk1lpJX5SZ6-Hyrf_-utf4gvSKvMmcov8VkKogZCU,19618
+langfun/env/event_handlers/metric_writer_test.py,sha256=sntUifTPHGixUshIgVBX4Q9vJL-xbeS0Cpd5X5hSQyQ,5955
+langfun-0.1.2.dev202510290805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202510290805.dist-info/METADATA,sha256=0QBqKpkm1e7B23GJpMUX5XQWfJEWwn_6nl-Qw_m2tos,7522
+langfun-0.1.2.dev202510290805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+langfun-0.1.2.dev202510290805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202510290805.dist-info/RECORD,,