langfun 0.0.2.dev20240330__py3-none-any.whl → 0.0.2.dev20240429__py3-none-any.whl

langfun/core/llms/openai_test.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for openai models."""
+"""Tests for OpenAI models."""
 
 import unittest
 from unittest import mock
@@ -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,14 +75,17 @@ 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):
+class OpenAITest(unittest.TestCase):
   """Tests for OpenAI language model."""
 
   def test_model_id(self):
@@ -89,7 +98,7 @@ class OpenaiTest(unittest.TestCase):
     )
 
   def test_max_concurrency(self):
-    self.assertEqual(openai.Gpt35(api_key='test_key').max_concurrency, 8)
+    self.assertGreater(openai.Gpt35(api_key='test_key').max_concurrency, 0)
 
   def test_get_request_args(self):
     self.assertEqual(
@@ -121,7 +130,6 @@ class OpenaiTest(unittest.TestCase):
             top_logprobs=None,
             n=1,
             temperature=1.0,
-            max_tokens=1024,
             stream=False,
             timeout=120.0,
             stop=['\n'],
@@ -149,17 +157,19 @@ class OpenaiTest(unittest.TestCase):
   def test_call_chat_completion_vision(self):
     with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
       mock_chat_completion.side_effect = mock_chat_completion_query_vision
-      lm = openai.Gpt4TurboVision(api_key='test_key')
-      self.assertEqual(
-          lm(
-              lf.UserMessage(
-                  'hello {{image}}',
-                  image=lf_modalities.Image.from_uri('https://fake/image')
-              ),
-              sampling_options=lf.LMSamplingOptions(n=2)
-          ),
-          'Sample 0 for message: https://fake/image',
-      )
+      lm_1 = openai.Gpt4Turbo(api_key='test_key')
+      lm_2 = openai.Gpt4VisionPreview(api_key='test_key')
+      for lm in (lm_1, lm_2):
+        self.assertEqual(
+            lm(
+                lf.UserMessage(
+                    'hello {{image}}',
+                    image=lf_modalities.Image.from_uri('https://fake/image')
+                ),
+                sampling_options=lf.LMSamplingOptions(n=2)
+            ),
+            'Sample 0 for message: https://fake/image',
+        )
 
   def test_sample_completion(self):
     with mock.patch('openai.Completion.create') as mock_completion:
@@ -170,18 +180,101 @@ 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(
+                    lf.AIMessage(
+                        'Sample 0 for prompt 0.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for prompt 0.',
+                        score=0.1,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.1,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for prompt 0.',
+                        score=0.2,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.2,
+                    logprobs=None,
+                ),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
+    self.assertEqual(
+        results[1],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for prompt 1.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=None,
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for prompt 1.',
+                        score=0.1,
+                        logprobs=None,
+                        usage=None,
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.1,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for prompt 1.',
+                        score=0.2,
+                        logprobs=None,
+                        usage=None,
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.2,
+                    logprobs=None,
+                ),
+            ],
+        ),
+    )
 
   def test_sample_chat_completion(self):
     with mock.patch('openai.ChatCompletion.create') as mock_chat_completion:
@@ -192,18 +285,116 @@ 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(
+                    lf.AIMessage(
+                        'Sample 0 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
+    self.assertEqual(
+        results[1],
+        lf.LMSamplingResult(
+            [
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 0 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 2 for message.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=33,
+                            completion_tokens=33,
+                            total_tokens=66
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+            ],
+            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:
@@ -213,11 +404,46 @@ 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(
+                    lf.AIMessage(
+                        'Sample 0 for prompt 0.',
+                        score=0.0,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=50,
+                            completion_tokens=50,
+                            total_tokens=100,
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.0,
+                    logprobs=None,
+                ),
+                lf.LMSample(
+                    lf.AIMessage(
+                        'Sample 1 for prompt 0.',
+                        score=0.1,
+                        logprobs=None,
+                        usage=lf.LMSamplingUsage(
+                            prompt_tokens=50,
+                            completion_tokens=50,
+                            total_tokens=100,
+                        ),
+                        tags=[lf.Message.TAG_LM_RESPONSE],
+                    ),
+                    score=0.1,
+                    logprobs=None,
+                ),
+            ],
+            usage=lf.LMSamplingUsage(
+                prompt_tokens=100, completion_tokens=100, total_tokens=200
+            ),
+        ),
+    )
 
 
 if __name__ == '__main__':

langfun/core/modalities/video.py

@@ -15,9 +15,7 @@
 
 import base64
 from typing import cast
-
 from langfun.core.modalities import mime
-import magic
 
 
 class Video(mime.MimeType):
@@ -29,6 +27,11 @@ class Video(mime.MimeType):
 
   @property
   def mime_type(self) -> str:
+    # TODO(daiyip): after cl/619658455, LaunchPad binaries cannot import `magic`
+    # correctly. This is to mitigate the issue for major Langfun users who do
+    # not use Video. We shall move this import out once the issue is fixed.
+    import magic  # pylint: disable=g-import-not-at-top
+
     video_mime_type = magic.from_buffer(self.to_bytes(), mime=True)
     if 'video/' not in video_mime_type:
       raise ValueError(f'Not a video: {video_mime_type!r}.')

langfun/core/structured/__init__.py

@@ -48,7 +48,10 @@ from langfun.core.structured.schema_generation import generate_class
 from langfun.core.structured.schema_generation import classgen_example
 from langfun.core.structured.schema_generation import default_classgen_examples
 
+from langfun.core.structured.function_generation import function_gen
+
 from langfun.core.structured.mapping import Mapping
+from langfun.core.structured.mapping import MappingError
 from langfun.core.structured.mapping import MappingExample
 
 from langfun.core.structured.parsing import ParseStructure

langfun/core/structured/completion_test.py

@@ -17,7 +17,6 @@ import inspect
 import unittest
 
 import langfun.core as lf
-from langfun.core import coding
 from langfun.core import modalities
 from langfun.core.llms import fake
 from langfun.core.structured import completion
@@ -583,6 +582,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']
         )
     )
@@ -607,7 +607,7 @@ class CompleteStructureTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          coding.CodeError,
+          mapping.MappingError,
           'Expect .* but encountered .*',
       ):
         completion.complete(Activity.partial(), autofix=0)