langfun 0.1.2.dev202505020804__py3-none-any.whl → 0.1.2.dev202505030803__py3-none-any.whl

langfun/core/agentic/action.py

@@ -27,7 +27,152 @@ import pyglove as pg
 
 
 class Action(pg.Object):
-  """Base class for agent actions."""
+  """Base class for Langfun's agentic actions.
+
+  # Developing Actions
+
+  In Langfun, an `Action` is a class representing a task an agent can execute.
+  To define custom actions, subclass `lf.agentic.Action` and implement the
+  `call` method, which contains the logic for the action's execution.
+
+  ```python
+  class Calculate(lf.agentic.Action):
+    expression: str
+
+    def call(self, session: Session, *, lm: lf.LanguageModel, **kwargs):
+      return session.query(expression, float, lm=lm)
+  ```
+
+  Key aspects of the `call` method:
+
+  - `session` (First Argument): An `lf.Session` object required to make queries,
+    perform logging, and add metadata to the action. It also tracks the
+    execution of the action and its sub-actions.
+
+    - Use `session.query(...)` to make calls to a Language Model.
+    - Use `session.debug(...)`, `session.info(...)`, `session.warning(...)`,
+      and `session.error(...)` for adding logs associated with the
+      current action.
+    - Use `session.add_metadata(...)` to associate custom metadata with
+      the current action.
+
+  - Keyword Arguments (e.g., lm): Arguments required for the action's execution
+    (like a language model) should be defined as keyword arguments.
+
+  - **kwargs: Include **kwargs to allow:
+
+    - Users to pass additional arguments to child actions.
+    - The action to gracefully handle extra arguments passed by parent actions.
+
+  # Using Actions
+
+  ## Creating Action objects
+  Action objects can be instantiated in two primary ways:
+
+  - Direct instantiation by Users:
+
+    ```
+    calculate_action = Calculate(expression='1 + 1')
+    ```
+
+  - Generation by Language Models (LLMs): LLMs can generate Action objects when
+    provided with an "action space" (a schema defining possible actions). The
+    LLM populates the action's attributes. User code can then invoke the
+    generated action.
+
+  ```python
+  import pyglove as pg
+  import langfun as lf
+
+  # Define possible actions for the LLM
+  class Search(lf.agentic.Action):
+    query: str
+    def call(self, session: lf.Session, *, lm: lf.LanguageModel, **kwargs):
+      # Placeholder for actual search logic
+      return f"Results for: {self.query}"
+
+  class DirectAnswer(lf.agentic.Action):
+    answer: str
+    def call(self, session: lf.Session, *, lm: lf.LanguageModel, **kwargs):
+      return self.answer
+
+  # Define the schema for the LLM's output
+  class NextStep(pg.Object):
+    step_by_step_thoughts: list[str]
+    next_action: Calculate | Search | DirectAnswer
+
+  # Query the LLM to determine the next step
+  next_step = lf.query(
+      'What is the next step for {{question}}?',
+      NextStep,
+      question='why is the sky blue?'
+  )
+  # Execute the action chosen by the LLM
+  result = next_step.next_action()
+  print(result)
+  ```
+
+  ## Invoking Actions and Managing Sessions:
+
+  When an action is called, the session argument (the first argument to call)
+  is handled as follows:
+
+  - Implicit Session Management: If no session is explicitly provided when
+    calling an action, Langfun automatically creates and passes one.
+
+    ```python
+    calc = Calculate(expression='1 + 1')
+
+    # A session is implicitly created and passed here.
+    result = calc()
+    print(result)
+
+    # Access the implicitly created session.
+    # print(calc.session)
+    ```
+
+  - Explicit Session Management: You can create and manage `lf.Session` objects
+    explicitly. This is useful for customizing session identifiers or managing
+    a shared context for multiple actions.
+
+  ```python
+  calc = Calculate(expression='1 + 1')
+
+  # Explicitly create and pass a session.
+  with lf.Session(id='my_agent_session') as session:
+    result = calc(session=session) # Pass the session explicitly
+    print(result)
+  ```
+
+  ## Accessing Execution Trajectory:
+
+  After an action is executed, the Session object holds a record of its
+  execution, known as the trajectory. This includes queries made and any
+  sub-actions performed.
+
+  - To access all queries issued directly by the root action:
+
+    ```python
+    print(session.root.execution.queries)
+    ```
+  - To access all actions issued by the root action and any of its
+    sub-actions (recursively):
+
+    ```python
+    print(session.root.execution.all_queries)
+    ```
+  - To access all child actions issued by the root action:
+
+    ```python
+    print(session.root.execution.actions)
+    ```
+
+  - To access all the actions in the sub-tree issued by the root action:
+
+    ```python
+    print(session.root.execution.all_actions)
+    ```
+  """
 
   def _on_bound(self):
     super()._on_bound()
@@ -60,6 +205,8 @@ class Action(pg.Object):
     """Executes the action."""
     if session is None:
       session = Session()
+      session.start()
+
       if show_progress:
         lf.console.display(pg.view(session, name='agent_session'))
 
@@ -107,8 +254,14 @@ class Action(pg.Object):
               action=self,
               error=error
           )
+        if self._session is not None:
+          self._session.end(result=None, error=error)
         raise
-      return result
+
+    if self._session is not None:
+      # Session is created by current action. Stop the session.
+      self._session.end(result)
+    return result
 
   @abc.abstractmethod
   def call(self, session: 'Session', **kwargs) -> Any:
@@ -229,9 +382,6 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           remove_class=['running'],
       )
 
-  def __len__(self) -> int:
-    return len(self.items)
-
   @property
   def has_started(self) -> bool:
     return self.start_time is not None
@@ -306,6 +456,22 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           for x in branch._iter_subtree(item_cls):  # pylint: disable=protected-access
             yield x
 
+  #
+  # Shortcut methods to operate on the execution trace.
+  #
+
+  def __len__(self) -> int:
+    return len(self.items)
+
+  def __iter__(self) -> Iterator[TracedItem]:
+    return iter(self.items)
+
+  def __bool__(self) -> bool:
+    return bool(self.items)
+
+  def __getitem__(self, index: int) -> TracedItem:
+    return self.items[index]
+
   def append(self, item: TracedItem) -> None:
     """Appends an item to the sequence."""
     with pg.notify_on_change(False):
@@ -935,6 +1101,44 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
           skip_notification=True
       )
 
+  def start(self) -> None:
+    """Starts the session."""
+    self.root.execution.start()
+
+  def end(
+      self,
+      result: Any,
+      error: pg.utils.ErrorInfo | None = None,
+      metadata: dict[str, Any] | None = None,
+  ) -> None:
+    """Ends the session."""
+    self.root.end(result, error, metadata)
+
+  def __enter__(self):
+    """Enters the session."""
+    self.start()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    """Exits the session."""
+    # We allow users to explicitly end the session with specified result
+    # and metadata.
+    if self.root.execution.has_stopped:
+      return
+
+    if exc_val is not None:
+      result, metadata = None, None
+      error = pg.utils.ErrorInfo.from_exception(exc_val)
+    else:
+      actions = self.root.actions
+      if actions:
+        result = actions[-1].result
+        error = actions[-1].error
+        metadata = actions[-1].metadata
+      else:
+        result, error, metadata = None, None, None
+    self.end(result, error, metadata)
+
   #
   # Context-manager for information tracking.
   #
@@ -942,8 +1146,12 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   @contextlib.contextmanager
   def track_action(self, action: Action) -> Iterator[ActionInvocation]:
     """Track the execution of an action."""
-    if not self._current_execution.has_started:
-      self._current_execution.start()
+    if not self.root.execution.has_started:
+      raise ValueError(
+          'Please call `Session.start() / Session.end()` explicitly, '
+          'or use `with Session(...) as session: ...` context manager to '
+          'signal the start and end of the session.'
+      )
 
     invocation = ActionInvocation(pg.maybe_ref(action))
     action._invocation = invocation  # pylint: disable=protected-access
@@ -960,12 +1168,6 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
     finally:
       self._current_execution = parent_execution
       self._current_action = parent_action
-      if parent_action is self.root:
-        parent_action.end(
-            result=invocation.result,
-            metadata=invocation.metadata,
-            error=invocation.error
-        )
 
   @contextlib.contextmanager
   def track_phase(self, name: str | None) -> Iterator[ExecutionTrace]:
@@ -1255,6 +1457,21 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
     """Returns the final result of the session."""
     return self.root.result
 
+  @property
+  def has_started(self) -> bool:
+    """Returns whether the session has started."""
+    return self.root.execution.has_started
+
+  @property
+  def has_stopped(self) -> bool:
+    """Returns whether the session has stopped."""
+    return self.root.execution.has_stopped
+
+  @property
+  def has_error(self) -> bool:
+    """Returns whether the session has an error."""
+    return self.root.has_error
+
   @property
   def current_action(self) -> ActionInvocation:
     """Returns the current invocation."""

langfun/core/agentic/action_eval.py

@@ -34,17 +34,20 @@ class ActionEval(lf.eval.v2.Evaluation):
   def process(self, example: lf.eval.v2.Example) -> tuple[str, dict[str, Any]]:
     example_input = example.input
     action = example_input.action
-    session = action_lib.Session(id=f'{self.id}#example-{example.id}')
 
-    # NOTE(daiyip): Setting session as metadata before action execution, so we
-    # could use `Evaluation.state.in_progress_examples` to access the session
-    # for status reporting from other threads.
-    example.metadata['session'] = session
+    # We explicitly create a session here to use a custom session ID.
+    with action_lib.Session(id=f'{self.id}#example-{example.id}') as session:
+
+      # NOTE(daiyip): Setting session as metadata before action execution, so we
+      # could use `Evaluation.state.in_progress_examples` to access the session
+      # for status reporting from other threads.
+      example.metadata['session'] = session
+
+      with lf.logging.use_log_level('fatal'):
+        kwargs = self.action_args.copy()
+        kwargs.update(verbose=True)
+        action(session=session, **kwargs)
 
-    with lf.logging.use_log_level('fatal'):
-      kwargs = self.action_args.copy()
-      kwargs.update(verbose=True)
-      action(session=session, **kwargs)
     return session.final_result, dict(session=session)
 
 #
@@ -76,8 +79,9 @@ class ActionEvalV1(lf_eval.Matching):
 
   def process(self, example: pg.Dict, **kwargs):
     action = example.action
-    session = action_lib.Session(id=str(getattr(example, 'id', '<empty>')))
-    action(session=session, lm=self.lm, **kwargs)
+    with action_lib.Session(
+        id=str(getattr(example, 'id', '<empty>'))) as session:
+      action(session=session, lm=self.lm, **kwargs)
     return session.as_message()
 
   def answer(self, output: Any, example: pg.Dict) -> Any:

langfun/core/agentic/action_eval_test.py

@@ -68,7 +68,6 @@ class ActionEvalV1Test(unittest.TestCase):
 
     s = FooEval()
     result = s.run(summary=False)
-    pg.print(result)
     self.assertEqual(
         result,
         dict(

langfun/core/agentic/action_test.py

@@ -98,7 +98,7 @@ class ExecutionTraceTest(unittest.TestCase):
     self.assertEqual(action_invocation.execution.id, '/a1')
 
     root.execution.reset()
-    self.assertEqual(len(root.execution), 0)
+    self.assertFalse(root.execution)
 
 
 class SessionTest(unittest.TestCase):
@@ -112,12 +112,18 @@ class SessionTest(unittest.TestCase):
 
     session = action_lib.Session(id='agent@1')
     self.assertEqual(session.id, 'agent@1')
+    self.assertFalse(session.has_started)
+    self.assertFalse(session.has_stopped)
 
     # Render HTML view to trigger dynamic update during execution.
     _ = session.to_html()
 
-    self.assertEqual(foo(session, lm=lm, verbose=True), 3)
+    with session:
+      result = foo(session, lm=lm, verbose=True)
 
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertEqual(result, 3)
     self.assertIsNone(foo.session)
     self.assertEqual(foo.result, 3)
     self.assertEqual(
@@ -135,8 +141,8 @@ class SessionTest(unittest.TestCase):
     self.assertIsNone(root.parent_action)
     self.assertEqual(root.id, 'agent@1:')
     self.assertEqual(root.execution.id, 'agent@1:')
-    self.assertEqual(len(root.execution.items), 1)
-    self.assertIs(root.execution.items[0].action, foo)
+    self.assertEqual(len(root.execution), 1)
+    self.assertIs(root.execution[0].action, foo)
 
     self.assertTrue(root.execution.has_started)
     self.assertTrue(root.execution.has_stopped)
@@ -160,14 +166,14 @@ class SessionTest(unittest.TestCase):
     self.assertEqual(root.usage_summary.total.num_requests, 6)
 
     # Inspecting the top-level action (Foo)
-    foo_invocation = root.execution.items[0]
+    foo_invocation = root.execution[0]
     self.assertIs(foo_invocation.parent_action, root)
     self.assertEqual(foo_invocation.id, 'agent@1:/a1')
     self.assertEqual(foo_invocation.execution.id, 'agent@1:/a1')
     self.assertEqual(len(foo_invocation.execution.items), 4)
 
     # Prepare phase.
-    prepare_phase = foo_invocation.execution.items[0]
+    prepare_phase = foo_invocation.execution[0]
     self.assertIsInstance(prepare_phase, action_lib.ExecutionTrace)
     self.assertEqual(prepare_phase.id, 'agent@1:/a1/prepare')
     self.assertEqual(len(prepare_phase.items), 2)
@@ -179,7 +185,7 @@ class SessionTest(unittest.TestCase):
     self.assertEqual(prepare_phase.items[1].id, 'agent@1:/a1/prepare/q1')
 
     # Tracked queries.
-    query_invocation = foo_invocation.execution.items[1]
+    query_invocation = foo_invocation.execution[1]
     self.assertIsInstance(query_invocation, lf_structured.QueryInvocation)
     self.assertEqual(query_invocation.id, 'agent@1:/a1/q2')
     self.assertIs(query_invocation.lm, lm)
@@ -197,7 +203,7 @@ class SessionTest(unittest.TestCase):
     )
 
     # Tracked parallel executions.
-    parallel_executions = foo_invocation.execution.items[2]
+    parallel_executions = foo_invocation.execution[2]
     self.assertEqual(parallel_executions.id, 'agent@1:/a1/p1')
     self.assertIsInstance(parallel_executions, action_lib.ParallelExecutions)
     self.assertEqual(len(parallel_executions), 3)
@@ -209,7 +215,7 @@ class SessionTest(unittest.TestCase):
     self.assertEqual(len(parallel_executions[2].queries), 1)
 
     # Invocation to Bar.
-    bar_invocation = foo_invocation.execution.items[3]
+    bar_invocation = foo_invocation.execution[3]
     self.assertIs(bar_invocation.parent_action, foo_invocation)
     self.assertEqual(bar_invocation.id, 'agent@1:/a1/a1')
     self.assertIsInstance(bar_invocation, action_lib.ActionInvocation)
@@ -240,10 +246,10 @@ class SessionTest(unittest.TestCase):
     root = session.root
     self.assertRegex(root.id, 'agent@.*:')
     self.assertTrue(root.has_error)
-    foo_invocation = root.execution.items[0]
+    foo_invocation = root.execution[0]
     self.assertIsInstance(foo_invocation, action_lib.ActionInvocation)
     self.assertTrue(foo_invocation.has_error)
-    bar_invocation = foo_invocation.execution.items[3]
+    bar_invocation = foo_invocation.execution[3]
     self.assertIsInstance(bar_invocation, action_lib.ActionInvocation)
     self.assertTrue(bar_invocation.has_error)
 
@@ -265,11 +271,146 @@ class SessionTest(unittest.TestCase):
     root = session.root
     self.assertRegex(root.id, 'agent@.*:')
     self.assertTrue(root.has_error)
-    foo_invocation = root.execution.items[0]
+    foo_invocation = root.execution[0]
     self.assertIsInstance(foo_invocation, action_lib.ActionInvocation)
     self.assertTrue(foo_invocation.has_error)
     self.assertEqual(len(foo_invocation.execution.items), 2)
 
+  def test_succeeded_with_implicit_session(self):
+    lm = fake.StaticResponse('lm response')
+    foo = Foo(1)
+    foo(lm=lm, verbose=True)
+    session = foo.session
+    self.assertIsNotNone(session)
+    self.assertIsInstance(session.root.action, action_lib.RootAction)
+    self.assertIs(session.current_action, session.root)
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertEqual(session.final_result, 3)
+    self.assertFalse(session.root.has_error)
+    self.assertEqual(session.root.metadata, {})
+
+  def test_failed_with_implicit_session(self):
+    lm = fake.StaticResponse('lm response')
+    foo = Foo(1, simulate_action_error=True)
+    with self.assertRaisesRegex(ValueError, 'Bar error'):
+      foo(lm=lm)
+    session = foo.session
+    self.assertIsNotNone(session)
+    self.assertIsInstance(session.root.action, action_lib.RootAction)
+    self.assertIs(session.current_action, session.root)
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertTrue(session.has_error)
+    self.assertIsInstance(session.root.error, pg.utils.ErrorInfo)
+    self.assertIn('Bar error', str(session.root.error))
+
+  def test_succeeded_with_explicit_session(self):
+    lm = fake.StaticResponse('lm response')
+    foo = Foo(1)
+    self.assertIsNone(foo.session)
+    self.assertIsNone(foo.result)
+    self.assertIsNone(foo.metadata)
+
+    session = action_lib.Session(id='agent@1')
+    self.assertEqual(session.id, 'agent@1')
+    self.assertFalse(session.has_started)
+    self.assertFalse(session.has_stopped)
+
+    with session:
+      result = foo(session, lm=lm, verbose=True)
+
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertEqual(result, 3)
+    self.assertIsNone(foo.session)
+    self.assertEqual(foo.result, 3)
+    self.assertEqual(
+        foo.metadata, dict(note='foo', subtask_0=0, subtask_1=1, subtask_2=2)
+    )
+    self.assertIs(session.final_result, foo.result)
+    self.assertFalse(session.has_error)
+
+  def test_succeeded_with_explicit_session_start_end(self):
+    lm = fake.StaticResponse('lm response')
+    foo = Foo(1)
+    self.assertIsNone(foo.session)
+    self.assertIsNone(foo.result)
+    self.assertIsNone(foo.metadata)
+
+    session = action_lib.Session(id='agent@1')
+    self.assertEqual(session.id, 'agent@1')
+    self.assertFalse(session.has_started)
+    self.assertFalse(session.has_stopped)
+
+    session.start()
+    result = foo(session, lm=lm, verbose=True)
+    session.end(result)
+
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertEqual(result, 3)
+    self.assertIsNone(foo.session)
+    self.assertEqual(foo.result, 3)
+    self.assertEqual(
+        foo.metadata, dict(note='foo', subtask_0=0, subtask_1=1, subtask_2=2)
+    )
+    self.assertIs(session.final_result, foo.result)
+    self.assertFalse(session.has_error)
+
+  def test_failed_with_explicit_session(self):
+    lm = fake.StaticResponse('lm response')
+    foo = Foo(1, simulate_action_error=True)
+    session = action_lib.Session(id='agent@1')
+    with self.assertRaisesRegex(ValueError, 'Bar error'):
+      with session:
+        foo(session, lm=lm, verbose=True)
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertTrue(session.has_error)
+    self.assertIsNone(session.final_result)
+    self.assertIsInstance(session.root.error, pg.utils.ErrorInfo)
+    self.assertIn('Bar error', str(session.root.error))
+
+  def test_failed_with_explicit_session_without_start(self):
+    lm = fake.StaticResponse('lm response')
+    foo = Foo(1, simulate_action_error=True)
+    session = action_lib.Session(id='agent@1')
+    with self.assertRaisesRegex(ValueError, 'Please call `Session.start'):
+      foo(session, lm=lm, verbose=True)
+
+  def test_succeed_with_multiple_actions(self):
+    lm = fake.StaticResponse('lm response')
+    with action_lib.Session() as session:
+      x = Bar()(session, lm=lm)
+      y = Bar()(session, lm=lm)
+      self.assertTrue(session.has_started)
+      self.assertFalse(session.has_stopped)
+      session.add_metadata(note='root metadata')
+      session.end(x + y)
+
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertEqual(session.final_result, 2 + 2)
+    self.assertEqual(len(session.root.execution), 2)
+    self.assertEqual(session.root.metadata, dict(note='root metadata'))
+
+  def test_failed_with_multiple_actions(self):
+    lm = fake.StaticResponse('lm response')
+    with self.assertRaisesRegex(ValueError, 'Bar error'):
+      with action_lib.Session() as session:
+        x = Bar()(session, lm=lm)
+        y = Bar(simulate_action_error=True)(session, lm=lm)
+        session.end(x + y)
+
+    self.assertTrue(session.has_started)
+    self.assertTrue(session.has_stopped)
+    self.assertTrue(session.has_error)
+    self.assertIsInstance(session.root.error, pg.utils.ErrorInfo)
+    self.assertEqual(len(session.root.execution), 2)
+    self.assertFalse(session.root.execution[0].has_error)
+    self.assertTrue(session.root.execution[1].has_error)
+
   def test_log(self):
     session = action_lib.Session()
     session.debug('hi', x=1, y=2)

langfun/core/llms/gemini.py

@@ -605,13 +605,13 @@ class Gemini(rest.REST):
           raise lf.ModalityError(f'Unsupported modality: {chunk!r}') from e
       return chunk
 
-    contents = []
     if system_message := prompt.get('system_message'):
       assert isinstance(system_message, lf.SystemMessage), type(system_message)
-      contents.append(
-          system_message.as_format(
-              'gemini', chunk_preprocessor=modality_conversion)
+      request['systemInstruction'] = system_message.as_format(
+          'gemini', chunk_preprocessor=modality_conversion
       )
+
+    contents = []
     contents.append(
         prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
     )
@@ -647,6 +647,11 @@ class Gemini(rest.REST):
           + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
           + pg.to_json_str(json_schema, json_indent=2)
       )
+    if options.max_thinking_tokens is not None:
+      config['thinkingConfig'] = {
+          'thinkingBudget': options.max_thinking_tokens
+      }
+
     return config
 
   def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
@@ -659,18 +664,25 @@ class Gemini(rest.REST):
     # NOTE(daiyip): We saw cases that `candidatesTokenCount` is not present.
     # Therefore, we use 0 as the default value.
     output_tokens = usage.get('candidatesTokenCount', 0)
+    thinking_tokens = usage.get('thoughtsTokenCount', 0)
+    total_tokens = usage.get('totalTokenCount', 0)
 
     return lf.LMSamplingResult(
         [lf.LMSample(message) for message in messages],
         usage=lf.LMSamplingUsage(
             prompt_tokens=input_tokens,
             completion_tokens=output_tokens,
-            total_tokens=input_tokens + output_tokens,
+            total_tokens=total_tokens,
+            completion_tokens_details={
+                'thinking_tokens': thinking_tokens,
+            },
         ),
     )
 
   def _error(self, status_code: int, content: str) -> lf.LMError:
-    if (status_code == 400
-        and b'exceeds the maximum number of tokens' in content):
+    if (
+        status_code == 400
+        and b'exceeds the maximum number of tokens' in content
+    ):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)

langfun/core/llms/gemini_test.py

@@ -38,14 +38,21 @@ example_image = (
 def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
   del url, kwargs
   c = pg.Dict(json['generationConfig'])
-  content = '\n'.join(c['parts'][0]['text'] for c in json['contents'])
+  parts = []
+  if system_instruction := json.get('systemInstruction'):
+    parts.extend([p['text'] for p in system_instruction.get('parts', [])])
+
+  # Add text from the main contents.
+  for c_item in json.get('contents', []):
+    for p in c_item.get('parts', []):
+      parts.append(p['text'])
+  content = '\n'.join(parts)
   response = requests.Response()
   response.status_code = 200
   response._content = pg.to_json_str({
       'candidates': [
           {
               'content': {
-                  'role': 'model',
                   'parts': [
                       {
                           'text': (
@@ -146,6 +153,30 @@ class GeminiTest(unittest.TestCase):
             }
         ),
     )
+
+    # Add test for thinkingConfig.
+    actual = model._generation_config(
+        lf.UserMessage('hi'),
+        lf.LMSamplingOptions(
+            max_thinking_tokens=100,
+        ),
+    )
+    self.assertEqual(
+        actual,
+        dict(
+            candidateCount=1,
+            temperature=None,
+            topP=None,
+            topK=40,
+            maxOutputTokens=None,
+            stopSequences=None,
+            responseLogprobs=False,
+            logprobs=None,
+            seed=None,
+            thinkingConfig={'thinkingBudget': 100},
+        ),
+    )
+
     with self.assertRaisesRegex(
         ValueError, '`json_schema` must be a dict, got'
     ):

langfun/core/llms/vertexai.py

@@ -28,6 +28,7 @@ import pyglove as pg
 try:
   # pylint: disable=g-import-not-at-top
   from google import auth as google_auth
+  from google.auth import exceptions as auth_exceptions
   from google.auth import credentials as credentials_lib
   from google.auth.transport import requests as auth_requests
   # pylint: enable=g-import-not-at-top
@@ -35,6 +36,7 @@ try:
   Credentials = credentials_lib.Credentials
 except ImportError:
   google_auth = None
+  auth_exceptions = None
   credentials_lib = None
   auth_requests = None
   Credentials = Any
@@ -134,6 +136,16 @@ class VertexAI(rest.REST):
     assert auth_requests is not None
     return auth_requests.AuthorizedSession(self._credentials)
 
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    assert auth_exceptions is not None
+    try:
+      return super()._sample_single(prompt)
+    except (
+        auth_exceptions.RefreshError,
+    ) as e:
+      raise lf.TemporaryLMError(
+          f'Failed to refresh Google authentication credentials: {e}'
+      ) from e
 
 #
 # Gemini models served by Vertex AI.

langfun/core/llms/vertexai_test.py

@@ -19,6 +19,7 @@ from unittest import mock
 
 from google.auth import exceptions
 import langfun.core as lf
+from langfun.core.llms import rest
 from langfun.core.llms import vertexai
 import pyglove as pg
 
@@ -51,6 +52,22 @@ class VertexAITest(unittest.TestCase):
     del os.environ['VERTEXAI_PROJECT']
     del os.environ['VERTEXAI_LOCATION']
 
+  def test_auth_refresh_error(self):
+    def _auth_refresh_error(*args, **kwargs):
+      del args, kwargs
+      raise exceptions.RefreshError('Cannot refresh token')
+
+    with self.assertRaisesRegex(
+        lf.concurrent.RetryError,
+        'Failed to refresh Google authentication credentials'
+    ):
+      with mock.patch.object(rest.REST, '_sample_single') as mock_sample_single:
+        mock_sample_single.side_effect = _auth_refresh_error
+        model = vertexai.VertexAIGemini15Pro(
+            project='abc', location='us-central1', max_attempts=1
+        )
+        model('hi')
+
 
 class VertexAIAnthropicTest(unittest.TestCase):
   """Tests for VertexAI Anthropic models."""

{langfun-0.1.2.dev202505020804.dist-info → langfun-0.1.2.dev202505030803.dist-info}/METADATA

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

{langfun-0.1.2.dev202505020804.dist-info → langfun-0.1.2.dev202505030803.dist-info}/RECORD

@@ -26,10 +26,10 @@ langfun/core/subscription_test.py,sha256=Y4ZdbZEwm83YNZBxHff0QR4QUa4rdaNXA3_jfIc
 langfun/core/template.py,sha256=jNhYSrbLIn9kZOa03w5QZbyjgfnzJzE_ZrrMvvWY4t4,24929
 langfun/core/template_test.py,sha256=AQv_m9qE93WxhEhSlm1xaBgB4hu0UVtA53dljngkUW0,17090
 langfun/core/agentic/__init__.py,sha256=qR3jlfUO4rhIoYdRDLz-d22YZf3FvU4FW88vsjiGDQQ,1224
-langfun/core/agentic/action.py,sha256=9P7xDiZVUV9MvJDfuAfLx-xa7qvS5F0EOGWDQnjAZBw,38931
-langfun/core/agentic/action_eval.py,sha256=NwjQ5hR-7YT6mo2q0mbDOgmNCKzTMpEzslYtR3fjXJY,4862
-langfun/core/agentic/action_eval_test.py,sha256=tRUkWmOE9p0rpNOq19xAY2oDEnYsEEykjg6sUpAwJk0,2832
-langfun/core/agentic/action_test.py,sha256=9EZKgLaBrqTErSRoxtrSlzmCz_cbnwWu0ZqpwKLst-s,10224
+langfun/core/agentic/action.py,sha256=3m2-k07Zz8qrzOdOa7xPl7fRH3I0c3VsMLR86_JBCcU,45359
+langfun/core/agentic/action_eval.py,sha256=JXhS5qEjWu9EZ0chDsjWxCqPAV26PUCBijtUYxiDeO4,4975
+langfun/core/agentic/action_eval_test.py,sha256=7AkOwNbUX-ZgR1R0a7bvUZ5abNTUV7blf_8Mnrwb-II,2811
+langfun/core/agentic/action_test.py,sha256=ezqg3tKlVwgLMnHKUmOdtxpnntuL8YIvhcTCSSdb8oc,15468
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
 langfun/core/coding/python/__init__.py,sha256=4ByknuoNU-mOIHwHKnTtmo6oD64oMFtlqPlYWmA5Wic,1736
 langfun/core/coding/python/correction.py,sha256=7zBedlhQKMPA4cfchUMxAOFl6Zl5RqCyllRHGWys40s,7092
@@ -92,8 +92,8 @@ langfun/core/llms/deepseek.py,sha256=jvTxdXPr-vH6HNakn_Ootx1heDg8Fen2FUkUW36bpCs
 langfun/core/llms/deepseek_test.py,sha256=DvROWPlDuow5E1lfoSkhyGt_ELA19JoQoDsTnRgDtTg,1847
 langfun/core/llms/fake.py,sha256=xmgCkk9y0I4x0IT32SZ9_OT27aLadXH8PRiYNo5VTd4,3265
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
-langfun/core/llms/gemini.py,sha256=k1uQy1aykkPtCpjnm56I74yMYIoOHN9j1ON7O8LDBJI,24111
-langfun/core/llms/gemini_test.py,sha256=d9Pvf3xmHgofv8AKXmbnfndsScxmgR5q_ctSIvEXYrU,6808
+langfun/core/llms/gemini.py,sha256=ZtUo2lQMSByYlzSALWae3KxFiKNtOOwGbkFwZTI1dO0,24472
+langfun/core/llms/gemini_test.py,sha256=Ve9X2Wvwu9wVFHpKZDP-qoM1_hzB4kgt6_HR9wxtNkg,7592
 langfun/core/llms/google_genai.py,sha256=j8W22WFvkT80Fw-r7Rg-e7MKhcSwljZkmtuufwSEn5s,5051
 langfun/core/llms/google_genai_test.py,sha256=NKNtpebArQ9ZR7Qsnhd2prFIpMjleojy6o6VMXkJ1zY,1502
 langfun/core/llms/groq.py,sha256=S9V10kFo3cgX89qPgt_umq-SpRnxEDLTt_hJmpERfbo,12066
@@ -106,8 +106,8 @@ langfun/core/llms/openai_compatible_test.py,sha256=KwOMA7tsmOxFBjezltkBDSU77AvOQ
 langfun/core/llms/openai_test.py,sha256=gwuO6aoa296iM2welWV9ua4KF8gEVGsEPakgbtkWkFQ,2687
 langfun/core/llms/rest.py,sha256=MCybcHApJcf49lubLnDzScN9Oc2IWY_JnMHIGdbDOuU,4474
 langfun/core/llms/rest_test.py,sha256=_zM7nV8DEVyoXNiQOnuwJ917mWjki0614H88rNmDboE,5020
-langfun/core/llms/vertexai.py,sha256=jCO1AjB3kBdBDNznygFpeXMZy-a7Lcap0NTe5Y7Wzx4,18205
-langfun/core/llms/vertexai_test.py,sha256=dOprP_uLNmXHYxMoX_hMPMsjKR-e_B5nKHjhlMCQoOQ,4252
+langfun/core/llms/vertexai.py,sha256=4t_Noj7cqzLNmESYCYzz9Ndodd_K4I4zxVLmljJ7r3E,18630
+langfun/core/llms/vertexai_test.py,sha256=0M4jsPOXGagdzPfEdJixmyLdhmmERePZWSFfTwnaYCQ,4875
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
 langfun/core/llms/cache/in_memory.py,sha256=i58oiQL28RDsq37dwqgVpC2mBETJjIEFS20yHiV5MKU,5185
@@ -156,8 +156,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202505020804.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202505020804.dist-info/METADATA,sha256=AWOuM3J68NYl8LFoEFVYgZyCDxyIDvcxwz79ce1GTSE,8178
-langfun-0.1.2.dev202505020804.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-langfun-0.1.2.dev202505020804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202505020804.dist-info/RECORD,,
+langfun-0.1.2.dev202505030803.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202505030803.dist-info/METADATA,sha256=g86bSMVJK0G9ZcJSYdMdBiHudgkbJ4sxiO5KvoHJwt8,8178
+langfun-0.1.2.dev202505030803.dist-info/WHEEL,sha256=7ciDxtlje1X8OhobNuGgi1t-ACdFSelPnSmDPrtlobY,91
+langfun-0.1.2.dev202505030803.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202505030803.dist-info/RECORD,,

{langfun-0.1.2.dev202505020804.dist-info → langfun-0.1.2.dev202505030803.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.1.0)
+Generator: setuptools (80.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any