langfun 0.1.2.dev202505010804__tar.gz → 0.1.2.dev202505030803__tar.gz

{langfun-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/PKG-INFO

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

{langfun-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/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-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/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-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/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-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/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-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/langfun/core/eval/v2/evaluation.py

@@ -167,7 +167,6 @@ class Evaluation(experiment_lib.Experiment):
       example.input = self.example_input_by_id(example.id)
 
     checkpointed = self._state.ckpt_example(example.id)
-    self._state.update(example, in_progress=True)
 
     with pg.timeit('evaluate') as timeit, lf.track_usages() as usage_summary:
       if checkpointed is None or checkpointed.has_error:
@@ -182,6 +181,7 @@ class Evaluation(experiment_lib.Experiment):
               f'contains error: {checkpointed.error}'
           )
         example.start_time = time.time()
+        self._state.update(example, in_progress=True)
         self._process(example, raise_if_has_error=raise_if_has_error)
       else:
         self.info(
@@ -189,6 +189,7 @@ class Evaluation(experiment_lib.Experiment):
             'is available and error free.'
         )
         example.start_time = checkpointed.start_time
+        self._state.update(example, in_progress=True)
 
         # Use the output and metadata obtained from the previous processing.
         example.output = checkpointed.output

{langfun-0.1.2.dev202505010804 → langfun-0.1.2.dev202505030803}/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)