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

langfun/core/llms/anthropic_test.py

@@ -0,0 +1,167 @@
+# Copyright 2023 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 Anthropic models."""
+
+import base64
+import os
+from typing import Any
+import unittest
+from unittest import mock
+from langfun.core import modalities as lf_modalities
+from langfun.core.llms import anthropic
+import pyglove as pg
+import requests
+
+
+def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'content': [{
+          'type': 'text',
+          'text': (
+              f'hello with temperature={json.get("temperature")}, '
+              f'top_k={json.get("top_k")}, '
+              f'top_p={json.get("top_p")}, '
+              f'max_tokens={json.get("max_tokens")}, '
+              f'stop={json.get("stop_sequences")}.'
+          ),
+      }],
+      'usage': {
+          'input_tokens': 2,
+          'output_tokens': 1,
+      },
+  }).encode()
+  return response
+
+
+image_content = (
+    b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x18\x00\x00\x00\x18\x04'
+    b'\x03\x00\x00\x00\x12Y \xcb\x00\x00\x00\x18PLTE\x00\x00'
+    b'\x00fff_chaag_cg_ch^ci_ciC\xedb\x94\x00\x00\x00\x08tRNS'
+    b'\x00\n\x9f*\xd4\xff_\xf4\xe4\x8b\xf3a\x00\x00\x00>IDATx'
+    b'\x01c \x05\x08)"\xd8\xcc\xae!\x06pNz\x88k\x19\\Q\xa8"\x10'
+    b'\xc1\x14\x95\x01%\xc1\n\xa143Ta\xa8"D-\x84\x03QM\x98\xc3'
+    b'\x1a\x1a\x1a@5\x0e\x04\xa0q\x88\x05\x00\x07\xf8\x18\xf9'
+    b'\xdao\xd0|\x00\x00\x00\x00IEND\xaeB`\x82'
+)
+
+
+def mock_mm_requests_post(url: str, json: dict[str, Any], **kwargs):
+  del url, kwargs
+  v = json['messages'][0]['content'][0]
+  image = lf_modalities.Image.from_bytes(base64.b64decode(v['source']['data']))
+
+  response = requests.Response()
+  response.status_code = 200
+  response._content = pg.to_json_str({
+      'content': [{
+          'type': 'text',
+          'text': f'{v["type"]}: {image.mime_type}',
+      }],
+      'usage': {
+          'input_tokens': 2,
+          'output_tokens': 1,
+      },
+  }).encode()
+  return response
+
+
+def mock_requests_post_error(status_code, error_type, error_message):
+  def _mock_requests(url: str, json: dict[str, Any], **kwargs):
+    del url, json, kwargs
+    response = requests.Response()
+    response.status_code = status_code
+    response._content = pg.to_json_str(
+        {
+            'error': {
+                'type': error_type,
+                'message': error_message,
+            }
+        }
+    ).encode()
+    return response
+
+  return _mock_requests
+
+
+class AnthropicTest(unittest.TestCase):
+
+  def test_basics(self):
+    self.assertEqual(
+        anthropic.Claude3Haiku().model_id, 'claude-3-haiku-20240307'
+    )
+    self.assertGreater(anthropic.Claude3Haiku().max_concurrency, 0)
+
+  def test_api_key(self):
+    lm = anthropic.Claude3Haiku()
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      lm('hi')
+
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_requests_post
+
+      lm = anthropic.Claude3Haiku(api_key='fake key')
+      self.assertRegex(lm('hi').text, 'hello.*')
+
+      os.environ['ANTHROPIC_API_KEY'] = 'abc'
+      lm = anthropic.Claude3Haiku()
+      self.assertRegex(lm('hi').text, 'hello.*')
+      del os.environ['ANTHROPIC_API_KEY']
+
+  def test_call(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_requests_post
+      lm = anthropic.Claude3Haiku(api_key='fake_key')
+      response = lm('hello', temperature=0.0, top_k=0.1, top_p=0.2, stop=['\n'])
+      self.assertEqual(
+          response.text,
+          (
+              'hello with temperature=0.0, top_k=0.1, top_p=0.2, '
+              "max_tokens=4096, stop=['\\n']."
+          ),
+      )
+      self.assertIsNotNone(response.usage)
+      self.assertIsNotNone(response.usage.prompt_tokens, 2)
+      self.assertIsNotNone(response.usage.completion_tokens, 1)
+      self.assertIsNotNone(response.usage.total_tokens, 3)
+
+  def test_mm_call(self):
+    with mock.patch('requests.Session.post') as mock_mm_request:
+      mock_mm_request.side_effect = mock_mm_requests_post
+      lm = anthropic.Claude3Haiku(api_key='fake_key')
+      response = lm(lf_modalities.Image.from_bytes(image_content), lm=lm)
+      self.assertEqual(response.text, 'image: image/png')
+
+  def test_call_errors(self):
+    for status_code, error_type, error_message in [
+        (429, 'rate_limit', 'Rate limit exceeded.'),
+        (529, 'service_unavailable', 'Service unavailable.'),
+        (500, 'bad_request', 'Bad request.'),
+    ]:
+      with mock.patch('requests.Session.post') as mock_mm_request:
+        mock_mm_request.side_effect = mock_requests_post_error(
+            status_code, error_type, error_message
+        )
+        lm = anthropic.Claude3Haiku(api_key='fake_key')
+        with self.assertRaisesRegex(
+            Exception, f'.*{status_code}: .*{error_message}'
+        ):
+          lm('hello', lm=lm, max_attempts=1)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/cache/in_memory_test.py

@@ -44,28 +44,37 @@ class InMemoryLMCacheTest(unittest.TestCase):
     self.assertEqual(
         list(cache.keys()),
         [
-            ('a', (0.0, 1024, 1, 40, None, None), 0),
-            ('a', (0.0, 1024, 1, 40, None, None), 1),
-            ('b', (0.0, 1024, 1, 40, None, None), 0),
-            ('c', (0.0, 1024, 1, 40, None, None), 0),
+            ('a', (None, None, 1, 40, None, None), 0),
+            ('a', (None, None, 1, 40, None, None), 1),
+            ('b', (None, None, 1, 40, None, None), 0),
+            ('c', (None, None, 1, 40, None, None), 0),
         ],
     )
     self.assertEqual(
         list(cache.keys('StaticSequence')),
         [
-            ('a', (0.0, 1024, 1, 40, None, None), 0),
-            ('a', (0.0, 1024, 1, 40, None, None), 1),
-            ('b', (0.0, 1024, 1, 40, None, None), 0),
-            ('c', (0.0, 1024, 1, 40, None, None), 0),
+            ('a', (None, None, 1, 40, None, None), 0),
+            ('a', (None, None, 1, 40, None, None), 1),
+            ('b', (None, None, 1, 40, None, None), 0),
+            ('c', (None, None, 1, 40, None, None), 0),
         ],
     )
 
     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(
@@ -90,19 +99,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
         list(cache.items()),
         [
             (
-                ('a', (0.0, 1024, 1, 40, None, None), 0),
+                ('a', (None, None, 1, 40, None, None), 0),
                 cache_entry('1'),
             ),
             (
-                ('a', (0.0, 1024, 1, 40, None, None), 1),
+                ('a', (None, None, 1, 40, None, None), 1),
                 cache_entry('2', 1),
             ),
             (
-                ('b', (0.0, 1024, 1, 40, None, None), 0),
+                ('b', (None, None, 1, 40, None, None), 0),
                 cache_entry('3'),
             ),
             (
-                ('c', (0.0, 1024, 1, 40, None, None), 0),
+                ('c', (None, None, 1, 40, None, None), 0),
                 cache_entry('4'),
             ),
         ],
@@ -111,19 +120,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
         list(cache.items('StaticSequence')),
         [
             (
-                ('a', (0.0, 1024, 1, 40, None, None), 0),
+                ('a', (None, None, 1, 40, None, None), 0),
                 cache_entry('1'),
             ),
             (
-                ('a', (0.0, 1024, 1, 40, None, None), 1),
+                ('a', (None, None, 1, 40, None, None), 1),
                 cache_entry('2', 1),
             ),
             (
-                ('b', (0.0, 1024, 1, 40, None, None), 0),
+                ('b', (None, None, 1, 40, None, None), 0),
                 cache_entry('3'),
             ),
             (
-                ('c', (0.0, 1024, 1, 40, None, None), 0),
+                ('c', (None, None, 1, 40, None, None), 0),
                 cache_entry('4'),
             ),
         ],
@@ -161,15 +170,15 @@ class InMemoryLMCacheTest(unittest.TestCase):
     self.assertEqual(
         list(cache.keys()),
         [
-            ('a', (0.0, 1024, 1, 40, None, None), 0),
-            ('a', (1.0, 1024, 1, 40, None, None), 0),
+            ('a', (None, None, 1, 40, None, None), 0),
+            ('a', (1.0, None, 1, 40, None, None), 0),
         ],
     )
 
   def test_different_model(self):
     cache = in_memory.InMemory()
-    lm1 = fake.StaticSequence(['1', '2', '3'], cache=cache)
-    lm2 = fake.Echo(cache=cache)
+    lm1 = fake.StaticSequence(['1', '2', '3'], cache=cache, temperature=0.0)
+    lm2 = fake.Echo(cache=cache, temperature=0.0)
 
     self.assertEqual(lm1('a'), '1')
     self.assertEqual(lm2('a'), 'a')
@@ -180,15 +189,15 @@ class InMemoryLMCacheTest(unittest.TestCase):
     self.assertEqual(
         list(cache.keys('StaticSequence')),
         [
-            ('a', (0.0, 1024, 1, 40, None, None), 0),
-            ('b', (0.0, 1024, 1, 40, None, None), 0),
+            ('a', (0.0, None, 1, 40, None, None), 0),
+            ('b', (0.0, None, 1, 40, None, None), 0),
         ],
     )
     self.assertEqual(
         list(cache.keys('Echo')),
         [
-            ('a', (0.0, 1024, 1, 40, None, None), 0),
-            ('b', (0.0, 1024, 1, 40, None, None), 0),
+            ('a', (0.0, None, 1, 40, None, None), 0),
+            ('b', (0.0, None, 1, 40, None, None), 0),
         ],
     )
     self.assertEqual(len(cache), 4)

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,24 @@ 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(
+                        lf.AIMessage(
+                            'hi',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 2, 4),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                lf.LMSamplingUsage(2, 2, 4))
+        ]
     )
 
   def test_call(self):
@@ -34,8 +51,8 @@ class EchoTest(unittest.TestCase):
     with contextlib.redirect_stdout(string_io):
       self.assertEqual(lm('hi'), 'hi')
     debug_info = string_io.getvalue()
-    self.assertIn('[0] LM INFO:', debug_info)
-    self.assertIn('[0] PROMPT SENT TO LM:', debug_info)
+    self.assertIn('[0] LM INFO', debug_info)
+    self.assertIn('[0] PROMPT SENT TO LM', debug_info)
     self.assertIn('[0] LM RESPONSE', debug_info)
 
   def test_score(self):
@@ -53,11 +70,45 @@ 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(
+                        lf.AIMessage(
+                            canned_response,
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 38, 40),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                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(
+                        lf.AIMessage(
+                            canned_response,
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(15, 38, 53),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(15, 38, 53)
+            )
+        ],
     )
 
   def test_call(self):
@@ -69,8 +120,8 @@ class StaticResponseTest(unittest.TestCase):
       self.assertEqual(lm('hi'), canned_response)
 
     debug_info = string_io.getvalue()
-    self.assertIn('[0] LM INFO:', debug_info)
-    self.assertIn('[0] PROMPT SENT TO LM:', debug_info)
+    self.assertIn('[0] LM INFO', debug_info)
+    self.assertIn('[0] PROMPT SENT TO LM', debug_info)
     self.assertIn('[0] LM RESPONSE', debug_info)
 
 
@@ -85,8 +136,38 @@ 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(
+                        lf.AIMessage(
+                            'Hello',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 5, 7),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'I am fine, how about you?',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(12, 25, 37),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(KeyError):
@@ -104,8 +185,38 @@ 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(
+                        lf.AIMessage(
+                            'Hello',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(2, 5, 7),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(2, 5, 7)
+            ),
+            lf.LMSamplingResult(
+                [
+                    lf.LMSample(
+                        lf.AIMessage(
+                            'I am fine, how about you?',
+                            score=1.0,
+                            logprobs=None,
+                            usage=lf.LMSamplingUsage(12, 25, 37),
+                            tags=[lf.Message.TAG_LM_RESPONSE],
+                        ),
+                        score=1.0,
+                        logprobs=None,
+                    )
+                ],
+                usage=lf.LMSamplingUsage(12, 25, 37)
+            )
         ]
     )
     with self.assertRaises(IndexError):

langfun/core/llms/google_genai_test.py

@@ -152,10 +152,15 @@ class GenAITest(unittest.TestCase):
     )
 
   def test_model_hub(self):
+    orig_get_model = genai.get_model
+    genai.get_model = mock_get_model
+
     model = google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
     self.assertIsNotNone(model)
     self.assertIs(google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), model)
 
+    genai.get_model = orig_get_model
+
   def test_api_key_check(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
       _ = google_genai.GeminiPro()._api_initialized
@@ -167,7 +172,7 @@ class GenAITest(unittest.TestCase):
 
   def test_call(self):
     with mock.patch(
-        'google.generativeai.generative_models.GenerativeModel.generate_content'
+        'google.generativeai.GenerativeModel.generate_content',
     ) as mock_generate:
       orig_get_model = genai.get_model
       genai.get_model = mock_get_model
@@ -176,7 +181,7 @@ class GenAITest(unittest.TestCase):
       lm = google_genai.GeminiPro(api_key='test_key')
       self.maxDiff = None
       self.assertEqual(
-          lm('hello', temperature=2.0, top_k=20).text,
+          lm('hello', temperature=2.0, top_k=20, max_tokens=1024).text,
           (
               'This is a response to hello with n=1, temperature=2.0, '
               'top_p=None, top_k=20, max_tokens=1024, stop=None.'
@@ -197,7 +202,7 @@ class GenAITest(unittest.TestCase):
         (
             "hello to models/text-bison-001 with {'temperature': 2.0, "
             "'top_k': 20, 'top_p': None, 'candidate_count': 1, "
-            "'max_output_tokens': 1024, 'stop_sequences': None}"
+            "'max_output_tokens': None, 'stop_sequences': None}"
         ),
     )
     genai.get_model = orig_get_model