langfun 0.0.2.dev20240413__py3-none-any.whl → 0.0.2.dev20240415__py3-none-any.whl

langfun/core/__init__.py

@@ -99,6 +99,7 @@ from langfun.core.modality import ModalityRef
 from langfun.core.language_model import LanguageModel
 from langfun.core.language_model import LMSample
 from langfun.core.language_model import LMSamplingOptions
+from langfun.core.language_model import LMSamplingUsage
 from langfun.core.language_model import LMSamplingResult
 from langfun.core.language_model import LMScoringResult
 from langfun.core.language_model import LMCache

langfun/core/eval/base.py

@@ -1565,6 +1565,7 @@ class Summary(pg.Object):
         results.append(
             pg.Dict(
                 experiment=entry,
+                dir=entry.dir,
                 metrics=entry.result.metrics if entry.result else None,
             )
         )

langfun/core/eval/base_test.py

@@ -194,6 +194,7 @@ class EvaluationTest(unittest.TestCase):
             cache_seed=0,
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(387, 24, 411),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )

langfun/core/langfunc_test.py

@@ -82,7 +82,9 @@ class LangFuncCallTest(unittest.TestCase):
     self.assertEqual(i.tags, ['rendered'])
 
     r = l()
-    self.assertEqual(r, message.AIMessage('Hello!!!', score=0.0, logprobs=None))
+    self.assertEqual(
+        r, message.AIMessage('Hello!!!', score=0.0, logprobs=None, usage=None)
+    )
     self.assertEqual(r.tags, ['lm-response', 'lm-output'])
     self.assertEqual(r.source, message.UserMessage('Hello'))
     self.assertEqual(r.source.tags, ['rendered', 'lm-input'])
@@ -106,7 +108,7 @@ class LangFuncCallTest(unittest.TestCase):
       self.assertEqual(l.render(), 'Hello')
       r = l()
       self.assertEqual(
-          r, message.AIMessage('Hello!!!', score=0.0, logprobs=None)
+          r, message.AIMessage('Hello!!!', score=0.0, logprobs=None, usage=None)
       )
       self.assertEqual(r.tags, ['lm-response', 'lm-output'])
 

langfun/core/language_model.py

@@ -47,6 +47,14 @@ class LMSample(pg.Object):
   ] = None
 
 
+class LMSamplingUsage(pg.Object):
+  """Usage information per completion."""
+
+  prompt_tokens: int
+  completion_tokens: int
+  total_tokens: int
+
+
 class LMSamplingResult(pg.Object):
   """Language model response."""
 
@@ -58,6 +66,11 @@ class LMSamplingResult(pg.Object):
       ),
   ] = []
 
+  usage: Annotated[
+      LMSamplingUsage | None,
+      'Usage information. Currently only OpenAI models are supported.',
+  ] = None
+
 
 class LMSamplingOptions(component.Component):
   """Language model sampling options."""
@@ -424,6 +437,8 @@ class LanguageModel(component.Component):
       logprobs = result.samples[0].logprobs
       response.set('score', result.samples[0].score)
       response.metadata.logprobs = logprobs
+      response.metadata.usage = result.usage
+
       elapse = time.time() - request_start
       self._debug(prompt, response, call_counter, elapse)
       return response

langfun/core/language_model_test.py

@@ -38,9 +38,19 @@ class MockModel(lm_lib.LanguageModel):
     def fake_sample(prompts):
       if context.attempt >= self.failures_before_attempt:
         return [
-            lm_lib.LMSamplingResult([lm_lib.LMSample(  # pylint: disable=g-complex-comprehension
-                response=prompt.text * self.sampling_options.top_k,
-                score=self.sampling_options.temperature or -1.0)])
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(  # pylint: disable=g-complex-comprehension
+                        response=prompt.text * self.sampling_options.top_k,
+                        score=self.sampling_options.temperature or -1.0,
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(
+                    prompt_tokens=100,
+                    completion_tokens=100,
+                    total_tokens=200,
+                ),
+            )
             for prompt in prompts
         ]
       context.attempt += 1
@@ -100,8 +110,14 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar']),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample('foo', score=-1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample('bar', score=-1.0)]),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('foo', score=-1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('bar', score=-1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
         ],
     )
     # Test override sampling_options.
@@ -112,10 +128,12 @@ class LanguageModelTest(unittest.TestCase):
         ),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo' * 2, score=0.5)]
+                [lm_lib.LMSample('foo' * 2, score=0.5)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar' * 2, score=0.5)]
+                [lm_lib.LMSample('bar' * 2, score=0.5)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
         ],
     )
@@ -123,18 +141,26 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar'], temperature=1.0),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample('foo', score=1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample('bar', score=1.0)]),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('foo', score=1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [lm_lib.LMSample('bar', score=1.0)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
         ],
     )
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar'], top_k=2, temperature=0.7),
         [
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('foo' * 2, score=0.7)]
+                [lm_lib.LMSample('foo' * 2, score=0.7)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
             lm_lib.LMSamplingResult(
-                [lm_lib.LMSample('bar' * 2, score=0.7)]
+                [lm_lib.LMSample('bar' * 2, score=0.7)],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
             ),
         ],
     )
@@ -144,6 +170,8 @@ class LanguageModelTest(unittest.TestCase):
     response = lm(prompt='foo')
     self.assertEqual(response.text, 'foo')
     self.assertEqual(response.score, -1.0)
+    self.assertIsNone(response.logprobs)
+    self.assertEqual(response.usage, lm_lib.LMSamplingUsage(100, 100, 200))
 
     # Test override sampling_options.
     self.assertEqual(
@@ -158,11 +186,24 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'bar']),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('foo', cache_seed=0), score=-1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('bar', cache_seed=0), score=-1.0)]),
-        ])
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('foo', cache_seed=0), score=-1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('bar', cache_seed=0), score=-1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+        ],
+    )
     self.assertEqual(cache.stats.num_queries, 2)
     self.assertEqual(cache.stats.num_hits, 0)
     self.assertEqual(cache.stats.num_updates, 2)
@@ -181,10 +222,22 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(prompts=['foo', 'baz'], temperature=1.0),
         [
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('foo', cache_seed=0), score=1.0)]),
-            lm_lib.LMSamplingResult([lm_lib.LMSample(
-                message_lib.AIMessage('baz', cache_seed=0), score=1.0)]),
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('foo', cache_seed=0), score=1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
+            lm_lib.LMSamplingResult(
+                [
+                    lm_lib.LMSample(
+                        message_lib.AIMessage('baz', cache_seed=0), score=1.0
+                    )
+                ],
+                usage=lm_lib.LMSamplingUsage(100, 100, 200),
+            ),
         ],
     )
     self.assertEqual(cache.stats.num_queries, 6)

langfun/core/llms/cache/in_memory_test.py

@@ -62,10 +62,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
 
     def cache_entry(response_text, cache_seed=0):
       return base.LMCacheEntry(
-          lf.LMSamplingResult([
-              lf.LMSample(
-                  lf.AIMessage(response_text, cache_seed=cache_seed), score=1.0)
-          ])
+          lf.LMSamplingResult(
+              [
+                  lf.LMSample(
+                      lf.AIMessage(response_text, cache_seed=cache_seed),
+                      score=1.0
+                  )
+              ],
+              usage=lf.LMSamplingUsage(
+                  1,
+                  len(response_text),
+                  len(response_text) + 1,
+              )
+          )
       )
 
     self.assertEqual(

langfun/core/llms/fake.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Fake LMs for testing."""
 
+import abc
 from typing import Annotated
 import langfun.core as lf
 
@@ -23,15 +24,32 @@ class Fake(lf.LanguageModel):
   def _score(self, prompt: lf.Message, completions: list[lf.Message]):
     return [lf.LMScoringResult(score=-i * 1.0) for i in range(len(completions))]
 
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
+    results = []
+    for prompt in prompts:
+      response = self._response_from(prompt)
+      results.append(
+          lf.LMSamplingResult(
+              [lf.LMSample(response, 1.0)],
+              usage=lf.LMSamplingUsage(
+                  prompt_tokens=len(prompt.text),
+                  completion_tokens=len(response.text),
+                  total_tokens=len(prompt.text) + len(response.text),
+              )
+          )
+      )
+    return results
+
+  @abc.abstractmethod
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    """Returns the response for the given prompt."""
+
 
 class Echo(Fake):
   """A simple echo language model for testing."""
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(prompt.text, 1.0)])
-        for prompt in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(prompt.text)
 
 
 @lf.use_init_args(['response'])
@@ -43,11 +61,8 @@ class StaticResponse(Fake):
       'A canned response that will be returned regardless of the prompt.'
   ]
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(self.response, 1.0)])
-        for _ in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(self.response)
 
 
 @lf.use_init_args(['mapping'])
@@ -59,11 +74,8 @@ class StaticMapping(Fake):
       'A mapping from prompt to response.'
   ]
 
-  def _sample(self, prompts: list[str]) -> list[lf.LMSamplingResult]:
-    return [
-        lf.LMSamplingResult([lf.LMSample(self.mapping[prompt], 1.0)])
-        for prompt in prompts
-    ]
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    return lf.AIMessage(self.mapping[prompt])
 
 
 @lf.use_init_args(['sequence'])
@@ -79,10 +91,7 @@ class StaticSequence(Fake):
     super()._on_bound()
     self._pos = 0
 
-  def _sample(self, prompts: list[str]) -> list[lf.LMSamplingResult]:
-    results = []
-    for _ in prompts:
-      results.append(lf.LMSamplingResult(
-          [lf.LMSample(self.sequence[self._pos], 1.0)]))
-      self._pos += 1
-    return results
+  def _response_from(self, prompt: lf.Message) -> lf.Message:
+    r = lf.AIMessage(self.sequence[self._pos])
+    self._pos += 1
+    return r

langfun/core/llms/fake_test.py

@@ -25,7 +25,12 @@ class EchoTest(unittest.TestCase):
   def test_sample(self):
     lm = fakelm.Echo()
     self.assertEqual(
-        lm.sample(['hi']), [lf.LMSamplingResult([lf.LMSample('hi', 1.0)])]
+        lm.sample(['hi']),
+        [
+            lf.LMSamplingResult(
+                [lf.LMSample('hi', 1.0)],
+                lf.LMSamplingUsage(2, 2, 4))
+        ]
     )
 
   def test_call(self):
@@ -53,11 +58,21 @@ class StaticResponseTest(unittest.TestCase):
     lm = fakelm.StaticResponse(canned_response)
     self.assertEqual(
         lm.sample(['hi']),
-        [lf.LMSamplingResult([lf.LMSample(canned_response, 1.0)])],
+        [
+            lf.LMSamplingResult(
+                [lf.LMSample(canned_response, 1.0)],
+                usage=lf.LMSamplingUsage(2, 38, 40)
+            )
+        ],
     )
     self.assertEqual(
         lm.sample(['Tell me a joke.']),
-        [lf.LMSamplingResult([lf.LMSample(canned_response, 1.0)])],
+        [
+            lf.LMSamplingResult(
+                [lf.LMSample(canned_response, 1.0)],
+                usage=lf.LMSamplingUsage(15, 38, 53)
+            )
+        ],
     )
 
   def test_call(self):
@@ -85,8 +100,14 @@ class StaticMappingTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(['Hi', 'How are you?']),
         [
-            lf.LMSamplingResult([lf.LMSample('Hello', 1.0)]),
-            lf.LMSamplingResult([lf.LMSample('I am fine, how about you?', 1.0)])
+            lf.LMSamplingResult(
+                [lf.LMSample('Hello', 1.0)],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [lf.LMSample('I am fine, how about you?', 1.0)],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(KeyError):
@@ -104,8 +125,14 @@ class StaticSequenceTest(unittest.TestCase):
     self.assertEqual(
         lm.sample(['Hi', 'How are you?']),
         [
-            lf.LMSamplingResult([lf.LMSample('Hello', 1.0)]),
-            lf.LMSamplingResult([lf.LMSample('I am fine, how about you?', 1.0)])
+            lf.LMSamplingResult(
+                [lf.LMSample('Hello', 1.0)],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [lf.LMSample('I am fine, how about you?', 1.0)],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(IndexError):

langfun/core/llms/openai.py

@@ -26,20 +26,6 @@ from openai import openai_object
 import pyglove as pg
 
 
-class Usage(pg.Object):
-  """Usage information per completion."""
-
-  prompt_tokens: int
-  completion_tokens: int
-  total_tokens: int
-
-
-class LMSamplingResult(lf.LMSamplingResult):
-  """LMSamplingResult with usage information."""
-
-  usage: Usage | None = None
-
-
 SUPPORTED_MODELS_AND_SETTINGS = [
     # Model name, max concurrent requests.
     # The concurrent requests is estimated by TPM/RPM from
@@ -181,7 +167,7 @@ class OpenAI(lf.LanguageModel):
       args['stop'] = options.stop
     return args
 
-  def _sample(self, prompts: list[lf.Message]) -> list[LMSamplingResult]:
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     assert self._api_initialized
     if self.is_chat_model:
       return self._chat_complete_batch(prompts)
@@ -189,7 +175,8 @@ class OpenAI(lf.LanguageModel):
       return self._complete_batch(prompts)
 
   def _complete_batch(
-      self, prompts: list[lf.Message]) -> list[LMSamplingResult]:
+      self, prompts: list[lf.Message]
+  ) -> list[lf.LMSamplingResult]:
 
     def _open_ai_completion(prompts):
       response = openai.Completion.create(
@@ -204,13 +191,13 @@ class OpenAI(lf.LanguageModel):
             lf.LMSample(choice.text.strip(), score=choice.logprobs or 0.0)
         )
 
-      usage = Usage(
+      usage = lf.LMSamplingUsage(
           prompt_tokens=response.usage.prompt_tokens,
           completion_tokens=response.usage.completion_tokens,
           total_tokens=response.usage.total_tokens,
       )
       return [
-          LMSamplingResult(
+          lf.LMSamplingResult(
               samples_by_index[index], usage=usage if index == 0 else None
           )
           for index in sorted(samples_by_index.keys())
@@ -231,7 +218,7 @@ class OpenAI(lf.LanguageModel):
 
   def _chat_complete_batch(
       self, prompts: list[lf.Message]
-  ) -> list[LMSamplingResult]:
+  ) -> list[lf.LMSamplingResult]:
     def _open_ai_chat_completion(prompt: lf.Message):
       if self.multimodal:
         content = []
@@ -272,9 +259,9 @@ class OpenAI(lf.LanguageModel):
             )
         )
 
-      return LMSamplingResult(
+      return lf.LMSamplingResult(
           samples=samples,
-          usage=Usage(
+          usage=lf.LMSamplingUsage(
               prompt_tokens=response.usage.prompt_tokens,
               completion_tokens=response.usage.completion_tokens,
               total_tokens=response.usage.total_tokens,

langfun/core/llms/openai_test.py

@@ -32,11 +32,14 @@ def mock_completion_query(prompt, *, n=1, **kwargs):
           text=f'Sample {k} for prompt {i}.',
           logprobs=k / 10,
       ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+  return pg.Dict(
+      choices=choices,
+      usage=lf.LMSamplingUsage(
+          prompt_tokens=100,
+          completion_tokens=100,
+          total_tokens=200,
+      ),
+  )
 
 
 def mock_chat_completion_query(messages, *, n=1, **kwargs):
@@ -49,11 +52,14 @@ def mock_chat_completion_query(messages, *, n=1, **kwargs):
         ),
         logprobs=None,
     ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+  return pg.Dict(
+      choices=choices,
+      usage=lf.LMSamplingUsage(
+          prompt_tokens=100,
+          completion_tokens=100,
+          total_tokens=200,
+      ),
+  )
 
 
 def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
@@ -69,11 +75,14 @@ def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
         ),
         logprobs=None,
     ))
-  return pg.Dict(choices=choices, usage=openai.Usage(
-      prompt_tokens=100,
-      completion_tokens=100,
-      total_tokens=200,
-  ))
+  return pg.Dict(
+      choices=choices,
+      usage=lf.LMSamplingUsage(
+          prompt_tokens=100,
+          completion_tokens=100,
+          total_tokens=200,
+      ),
+  )
 
 
 class OpenaiTest(unittest.TestCase):
@@ -169,18 +178,28 @@ class OpenaiTest(unittest.TestCase):
       )
 
     self.assertEqual(len(results), 2)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-        lf.LMSample('Sample 2 for prompt 0.', score=0.2),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-
-    self.assertEqual(results[1], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 1.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 1.', score=0.1),
-        lf.LMSample('Sample 2 for prompt 1.', score=0.2),
-    ]))
+    self.assertEqual(
+        results[0],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for prompt 0.', score=0.0),
+                lf.LMSample('Sample 1 for prompt 0.', score=0.1),
+                lf.LMSample('Sample 2 for prompt 0.', score=0.2),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
+
+    self.assertEqual(
+        results[1],
+        lf.LMSamplingResult([
+            lf.LMSample('Sample 0 for prompt 1.', score=0.0),
+            lf.LMSample('Sample 1 for prompt 1.', score=0.1),
+            lf.LMSample('Sample 2 for prompt 1.', score=0.2),
+        ]),
+    )
 
   def test_sample_chat_completion(self):
     with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
@@ -191,18 +210,32 @@ class OpenaiTest(unittest.TestCase):
       )
 
     self.assertEqual(len(results), 2)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for message.', score=0.0),
-        lf.LMSample('Sample 1 for message.', score=0.0),
-        lf.LMSample('Sample 2 for message.', score=0.0),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
-    self.assertEqual(results[1], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for message.', score=0.0),
-        lf.LMSample('Sample 1 for message.', score=0.0),
-        lf.LMSample('Sample 2 for message.', score=0.0),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
+    self.assertEqual(
+        results[0],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for message.', score=0.0),
+                lf.LMSample('Sample 1 for message.', score=0.0),
+                lf.LMSample('Sample 2 for message.', score=0.0),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
+    self.assertEqual(
+        results[1],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for message.', score=0.0),
+                lf.LMSample('Sample 1 for message.', score=0.0),
+                lf.LMSample('Sample 2 for message.', score=0.0),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
 
   def test_sample_with_contextual_options(self):
     with mock.patch('openai.Completion.create') as mock_completion:
@@ -212,11 +245,18 @@ class OpenaiTest(unittest.TestCase):
         results = lm.sample(['hello'])
 
     self.assertEqual(len(results), 1)
-    self.assertEqual(results[0], openai.LMSamplingResult([
-        lf.LMSample('Sample 0 for prompt 0.', score=0.0),
-        lf.LMSample('Sample 1 for prompt 0.', score=0.1),
-    ], usage=openai.Usage(
-        prompt_tokens=100, completion_tokens=100, total_tokens=200)))
+    self.assertEqual(
+        results[0],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample('Sample 0 for prompt 0.', score=0.0),
+                lf.LMSample('Sample 1 for prompt 0.', score=0.1),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
 
 
 if __name__ == '__main__':

langfun/core/structured/completion_test.py

@@ -583,6 +583,7 @@ class CompleteStructureTest(unittest.TestCase):
             result=Activity(description='foo'),
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(553, 27, 580),
             tags=['lm-response', 'lm-output', 'transformed']
         )
     )

langfun/core/structured/parsing_test.py

@@ -280,13 +280,15 @@ class ParseStructurePythonTest(unittest.TestCase):
         ),
         1,
     )
+    r = parsing.parse(
+        'the answer is 1', int, user_prompt='what is 0 + 1?', lm=lm,
+        returns_message=True
+    )
     self.assertEqual(
-        parsing.parse(
-            'the answer is 1', int, user_prompt='what is 0 + 1?', lm=lm,
-            returns_message=True
-        ),
+        r,
         lf.AIMessage(
             '1', score=1.0, result=1, logprobs=None,
+            usage=lf.LMSamplingUsage(652, 1, 653),
             tags=['lm-response', 'lm-output', 'transformed']
         ),
     )
@@ -634,13 +636,18 @@ class CallTest(unittest.TestCase):
       )
 
   def test_call_with_returning_message(self):
+    r = parsing.call(
+        'Compute 1 + 2', int, lm=fake.StaticSequence(['three', '3']),
+        returns_message=True
+    )
     self.assertEqual(
-        parsing.call(
-            'Compute 1 + 2', int, lm=fake.StaticSequence(['three', '3']),
-            returns_message=True
-        ),
+        r,
         lf.AIMessage(
-            '3', result=3, score=1.0, logprobs=None,
+            '3',
+            result=3,
+            score=1.0,
+            logprobs=None,
+            usage=lf.LMSamplingUsage(315, 1, 316),
             tags=['lm-response', 'lm-output', 'transformed']
         ),
     )

langfun/core/structured/prompting_test.py

@@ -77,6 +77,7 @@ class QueryTest(unittest.TestCase):
             result=1,
             score=1.0,
             logprobs=None,
+            usage=lf.LMSamplingUsage(323, 1, 324),
             tags=['lm-response', 'lm-output', 'transformed'],
         ),
     )

langfun/core/templates/selfplay_test.py

@@ -56,7 +56,9 @@ class SelfPlayTest(unittest.TestCase):
     g = NumberGuess(target_num=10)
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 10])):
-      self.assertEqual(g(), lf.AIMessage('10', score=0.0, logprobs=None))
+      self.assertEqual(
+          g(), lf.AIMessage('10', score=0.0, logprobs=None, usage=None)
+      )
 
     self.assertEqual(g.num_turns, 4)
 
@@ -64,7 +66,9 @@ class SelfPlayTest(unittest.TestCase):
     g = NumberGuess(target_num=10, max_turns=10)
 
     with lf.context(lm=NumberGuesser(guesses=[50, 20, 5, 2, 5, 4])):
-      self.assertEqual(g(), lf.AIMessage('2', score=0.0, logprobs=None))
+      self.assertEqual(
+          g(), lf.AIMessage('2', score=0.0, logprobs=None, usage=None)
+      )
 
     self.assertEqual(g.num_turns, 10)
 

{langfun-0.0.2.dev20240413.dist-info → langfun-0.0.2.dev20240415.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240413
+Version: 0.0.2.dev20240415
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.0.2.dev20240413.dist-info → langfun-0.0.2.dev20240415.dist-info}/RECORD

@@ -1,5 +1,5 @@
 langfun/__init__.py,sha256=PqX3u18BC0szYIMu00j-RKxvwkNPwXtAFZ-96oxrQ0M,1841
-langfun/core/__init__.py,sha256=sVcPl89lWYHQ1cUoaLaM8dErCovugJo5e2F3A_94Q3Y,4192
+langfun/core/__init__.py,sha256=6QEuXOZ9BXxm6TjpaMXuLwUBTYO3pkFDqn9QVBXyyPQ,4248
 langfun/core/component.py,sha256=VRPfDB_2jEnxcB3-HoiVjG4ID-SMenNPIsytb0uXMPg,9674
 langfun/core/component_test.py,sha256=VAPd6V_-odAe8rBvesW3ogYDd6OSqRq4FaPhfgOM4Zg,7949
 langfun/core/concurrent.py,sha256=TRc49pJ3HQro2kb5FtcWkHjhBm8UcgE8RJybU5cU3-0,24537
@@ -7,9 +7,9 @@ langfun/core/concurrent_test.py,sha256=mwFMZhDUdppnDr7vDSTwcbMHwrdsIoKJwRYNtl4ZW
 langfun/core/console.py,sha256=bk5rNPNm9rMGW5YT2HixxU04p2umnoabn5SDz6Dqe88,2317
 langfun/core/console_test.py,sha256=5SYJdxpJGLgdSSQqqMPoA1X6jpsLD8rgcyk-EgI65oE,1077
 langfun/core/langfunc.py,sha256=WXdTc3QsmGD_n80KD9dFRr5MHpGZ9E_y_Rhtk4t9-3w,11852
-langfun/core/langfunc_test.py,sha256=rRxz2OOka5qagTSS1IcJ1Ij3mjjWawPFe1n9zYtGST8,8340
-langfun/core/language_model.py,sha256=D3aU7ep1MFnyMWYCfvbA3ZK9DgP_wk0PogXo1Kmvk4Q,17185
-langfun/core/language_model_test.py,sha256=bTyQVsH5JAxEzzzuq8VO8bVa9kiAMeiahzrxLxnOuQs,11380
+langfun/core/langfunc_test.py,sha256=sQaKuZpGGmG80GRifhbxkj7nfzQLJKj4Vuw5y1s1K3U,8378
+langfun/core/language_model.py,sha256=Tzswu0hyXOQOZ3fZ_Mz_Cc0ei7tVj8rTay9jJEgM6mI,17510
+langfun/core/language_model_test.py,sha256=KvXXOr64TsSs3WkEALCLLZSlz09i7hBiHDOZ_8Eq8_o,13047
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
 langfun/core/message.py,sha256=QhvV9t5qaryPcruyxxcXi3gm9QDInkSldwTtK6sVJ3c,15734
 langfun/core/message_test.py,sha256=Z23pUM5vPnDrYkIIibe2KL73D5HKur_awI0ut_EQFQA,9501
@@ -40,25 +40,25 @@ langfun/core/coding/python/parsing_test.py,sha256=9vAWF484kWIm6JZq8NFiMgKUDhXV-d
 langfun/core/coding/python/permissions.py,sha256=1QWGHvzL8MM0Ok_auQ9tURqZHtdOfJaDpBzZ29GUE-c,2544
 langfun/core/coding/python/permissions_test.py,sha256=w5EDb8QxpxgJyZkojyzVWQvDfg366zn99-g__6TbPQ0,2699
 langfun/core/eval/__init__.py,sha256=iDA2OcJ3kR6ixZizXIY3N9LsjkaVrfTbSClTiSP8ekY,1291
-langfun/core/eval/base.py,sha256=Op-DO-YV8sL8mQvCfbzLfDDL6bDMuTtNYeyp5_QCBsQ,55328
-langfun/core/eval/base_test.py,sha256=mjdQ3ukxc7BhsVJkFJvqtz9EVhSR0OGL9j1zf_AfXR4,21540
+langfun/core/eval/base.py,sha256=TZAmcdRBtzwMG1V3e_NgyJXg7J6dWMdMBrHvBnFuFho,55359
+langfun/core/eval/base_test.py,sha256=OuuXFW_lX9bGhyd__kvlDSNJVne-5cSlnm-qDhyvOcc,21592
 langfun/core/eval/matching.py,sha256=aqNlYrlav7YmsB7rUlsdfoi1RLA5CYqn2RGPxRlPc78,9599
 langfun/core/eval/matching_test.py,sha256=FFHYD7IDuKe5RMjkx74ksukiwUhO5a_SS340JaIPMws,4898
 langfun/core/eval/scoring.py,sha256=aKeanBJf1yO3Q9JEtgPWoiZk_3M_GiqwXVXX7x_g22w,6172
 langfun/core/eval/scoring_test.py,sha256=YH1cIxBWtfdKcAV9Fh10vLkV5J-gxk8b6nxW4Z2u5pk,4024
 langfun/core/llms/__init__.py,sha256=gROJ8AjMq_ebXFcEfsyzYGCS6NsGfzf9d43nLu_TIdw,2504
-langfun/core/llms/fake.py,sha256=dVzOrW27RZ1p3DdQoRCRZs_vfoQcTcNrlWxia7oqmvw,2499
-langfun/core/llms/fake_test.py,sha256=Qk_Yoi4Z7P9o6f8Q_BZkaSlvxH89ZVsDxnVIbSBRBXk,3555
+langfun/core/llms/fake.py,sha256=b-Xk5IPTbUt-elsyzd_i3n1tqzc_kgETXrEvgJruSMk,2824
+langfun/core/llms/fake_test.py,sha256=AThvNyhZbkpsn-YO798uLgqB6TSw5XP2SKpKvcXEytw,4188
 langfun/core/llms/google_genai.py,sha256=n8zyJwh9UCTgb6-8LyvmjVNFGZQ4-zfzZ0ulkhHAnR8,8624
 langfun/core/llms/google_genai_test.py,sha256=_UcGTfl16-aDUlEWFC2W2F8y9jPUs53RBYA6MOCpGXw,7525
 langfun/core/llms/llama_cpp.py,sha256=Y_KkMUf3Xfac49koMUtUslKl3h-HWp3-ntq7Jaa3bdo,2385
 langfun/core/llms/llama_cpp_test.py,sha256=ZxC6defGd_HX9SFRU9U4cJiQnBKundbOrchbXuC1Z2M,1683
-langfun/core/llms/openai.py,sha256=uOJDflucpKZv3TPZwaeDSp9QMs2oDFuzh5Jm5j4dlm4,11680
-langfun/core/llms/openai_test.py,sha256=ulzp5uzEmEvnqZ21D0FP6eaiH1xMQ59FaLHoqA0lTgc,7570
+langfun/core/llms/openai.py,sha256=1EUd8WTI6EpcU_fzD90-4M11RdL9Mj4S9zfrzUZIyGM,11463
+langfun/core/llms/openai_test.py,sha256=hiByS95g3pXtjB2XfIdVCKiAZDb_-Qirb2_LsSyskpY,8166
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=cFfYvOIUae842pncqCAsRvqXCk2AnAsRYVx0mcIoAeY,3338
 langfun/core/llms/cache/in_memory.py,sha256=YfFyJEhLs73cUiB0ZfhMxYpdE8Iuxxw-dvMFwGHTSHw,4742
-langfun/core/llms/cache/in_memory_test.py,sha256=guHHjislh1Mj3-GBARICMh-qq5gh4fwZQ7SI5kQEAeQ,8510
+langfun/core/llms/cache/in_memory_test.py,sha256=D-n26h__rVXQO51WRFhRfq5sw1oifRLx2SvCQWuNEm8,8747
 langfun/core/memories/__init__.py,sha256=HpghfZ-w1NQqzJXBx8Lz0daRhB2rcy2r9Xm491SBhC4,773
 langfun/core/memories/conversation_history.py,sha256=c9amD8hCxGFiZuVAzkP0dOMWSp8L90uvwkOejjuBqO0,1835
 langfun/core/memories/conversation_history_test.py,sha256=AaW8aNoFjxNusanwJDV0r3384Mg0eAweGmPx5DIkM0Y,2052
@@ -71,15 +71,15 @@ langfun/core/modalities/video.py,sha256=25M4XsNG5XEWRy57LYT_a6_aMURMPAgC41B3weEX
 langfun/core/modalities/video_test.py,sha256=jYuI2m8S8zDCAVBPEUbbpP205dXAht90A2_PHWo4-r8,2039
 langfun/core/structured/__init__.py,sha256=SpObW-HKpyKvkLlX8FV5ixz7CRm098j2aGfOguM3AUI,3462
 langfun/core/structured/completion.py,sha256=skBxt6V_fv2TBUKnzFgnPMbVY8HSYn8sY04MLok2yvs,7299
-langfun/core/structured/completion_test.py,sha256=98UCgA4gzfp6H6HgP2s2kcKs25YH3k4Nxj1rgAvmVBw,19249
+langfun/core/structured/completion_test.py,sha256=0FJreSmz0Umsj47dIlOyCjBXUa7janIplXhg1CbLT4U,19301
 langfun/core/structured/description.py,sha256=SXW4MJvshFjbR-0gw6rE21o6WXq12UlRXawvDBXMZFA,5211
 langfun/core/structured/description_test.py,sha256=UtZGjSFUaQ6130t1E5tcL7ODu0xIefkapb53TbnqsK8,7362
 langfun/core/structured/mapping.py,sha256=7JInwZLmQdu7asHhC0vFLJNOCBnY-hrD6v5RQgf-xKk,11020
 langfun/core/structured/mapping_test.py,sha256=07DDCGbwytQHSMm7fCi5-Ly-JNgdV4ubHZq0wthX4A4,3338
 langfun/core/structured/parsing.py,sha256=keoVqEfzAbdULh6GawWFsTQzU91MzJXYFZjXGXLaD8g,11492
-langfun/core/structured/parsing_test.py,sha256=2_Uf3LYNRON1-5ysEr75xiG_cAxR3ZiixSfvUQu6mOQ,20846
+langfun/core/structured/parsing_test.py,sha256=9rUe7ipRhltQv7y8NXgR98lBXhSVKnfRM9TSAyVdxbs,20980
 langfun/core/structured/prompting.py,sha256=mOmCWNVMnBk4rI7KBlEm5kmusPXoAKiWcohhzaw-s2o,7427
-langfun/core/structured/prompting_test.py,sha256=luJoJ16h0CkKmZv0-elOD2xLhqa7exZwHUTa9J15wqs,19894
+langfun/core/structured/prompting_test.py,sha256=csOzqHRp6T3KGp7Dsm0vS-BkZdQ4ALRt09iiFNz_YmA,19945
 langfun/core/structured/schema.py,sha256=mJXirgqx3N7SA9zBO_ISHrzcV-ZRshLhnMJyCcSjGjY,25057
 langfun/core/structured/schema_generation.py,sha256=U3nRQsqmMZg_qIVDh2fiY3K4JLfsAL1LcKzIFP1iXFg,5316
 langfun/core/structured/schema_generation_test.py,sha256=cfZyP0gHno2fXy_c9vsVdvHmqKQSfuyUsCtfO3JFmYQ,2945
@@ -94,9 +94,9 @@ langfun/core/templates/conversation_test.py,sha256=RryYyIhfc34dLWOs6GfPQ8HU8mXpK
 langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fikKhwhzwhpKI,1460
 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=IB5rWbjK_9CTkqEo1BclQPzFAKcIiusJckH8J19HFgI,2096
-langfun-0.0.2.dev20240413.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240413.dist-info/METADATA,sha256=O7nH000mRIqk6_yW-mchTjLq3iB6MFVbebb3sE_QMek,3405
-langfun-0.0.2.dev20240413.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240413.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240413.dist-info/RECORD,,
+langfun/core/templates/selfplay_test.py,sha256=DYVrkk7uNKCqJGEHH31HssU2BPuMItU1vJLzfcXIlYg,2156
+langfun-0.0.2.dev20240415.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240415.dist-info/METADATA,sha256=V_zKk0hFMrBR6jMyr0C0v71Y4RJ9GL9b0uAkBerHIIw,3405
+langfun-0.0.2.dev20240415.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240415.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240415.dist-info/RECORD,,