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

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/{gemini.py → google_genai.py}

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Gemini models exposed through Google Generative AI APIs."""
 
+import abc
 import functools
 import os
 from typing import Annotated, Any, Literal
@@ -20,14 +21,20 @@ from typing import Annotated, Any, Literal
 import google.generativeai as genai
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
+import pyglove as pg
 
 
 @lf.use_init_args(['model'])
-class Gemini(lf.LanguageModel):
-  """Language model served on VertexAI."""
+class GenAI(lf.LanguageModel):
+  """Language models provided by Google GenAI."""
 
   model: Annotated[
-      Literal['gemini-pro', 'gemini-pro-vision', ''],
+      Literal[
+          'gemini-pro',
+          'gemini-pro-vision',
+          'text-bison-001',
+          'chat-bison-001',
+      ],
       'Model name.',
   ]
 
@@ -35,7 +42,8 @@ class Gemini(lf.LanguageModel):
       str | None,
       (
           'API key. If None, the key will be read from environment variable '
-          "'GOOGLE_API_KEY'."
+          "'GOOGLE_API_KEY'. "
+          'Get an API key at https://ai.google.dev/tutorials/setup'
       ),
   ] = None
 
@@ -43,6 +51,9 @@ class Gemini(lf.LanguageModel):
       False
   )
 
+  # Set the default max concurrency to 8 workers.
+  max_concurrency = 8
+
   def _on_bound(self):
     super()._on_bound()
     self.__dict__.pop('_api_initialized', None)
@@ -67,7 +78,11 @@ class Gemini(lf.LanguageModel):
     return [
         m.name.lstrip('models/')
         for m in genai.list_models()
-        if 'generateContent' in m.supported_generation_methods
+        if (
+            'generateContent' in m.supported_generation_methods
+            or 'generateText' in m.supported_generation_methods
+            or 'generateMessage' in m.supported_generation_methods
+        )
     ]
 
   @property
@@ -80,11 +95,6 @@ class Gemini(lf.LanguageModel):
     """Returns a string to identify the resource for rate control."""
     return self.model_id
 
-  @property
-  def max_concurrency(self) -> int:
-    """Max concurrent requests."""
-    return 8
-
   def _generation_config(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
     """Creates generation config from langfun sampling options."""
     return genai.GenerationConfig(
@@ -117,7 +127,7 @@ class Gemini(lf.LanguageModel):
     return chunks
 
   def _response_to_result(
-      self, response: genai.types.GenerateContentResponse
+      self, response: genai.types.GenerateContentResponse | pg.Dict
   ) -> lf.LMSamplingResult:
     """Parses generative response into message."""
     samples = []
@@ -149,17 +159,97 @@ class Gemini(lf.LanguageModel):
     return self._response_to_result(response)
 
 
+class _LegacyGenerativeModel(pg.Object):
+  """Base for legacy GenAI generative model."""
+
+  model: str
+
+  def generate_content(
+      self,
+      input_content: list[str | genai.types.BlobDict],
+      generation_config: genai.GenerationConfig,
+  ) -> pg.Dict:
+    """Generate content."""
+    segments = []
+    for s in input_content:
+      if not isinstance(s, str):
+        raise ValueError(f'Unsupported modality: {s!r}')
+      segments.append(s)
+    return self.generate(' '.join(segments), generation_config)
+
+  @abc.abstractmethod
+  def generate(
+      self, prompt: str, generation_config: genai.GenerationConfig) -> pg.Dict:
+    """Generate response based on prompt."""
+
+
+class _LegacyCompletionModel(_LegacyGenerativeModel):
+  """Legacy GenAI completion model."""
+
+  def generate(
+      self, prompt: str, generation_config: genai.GenerationConfig
+  ) -> pg.Dict:
+    completion: genai.types.Completion = genai.generate_text(
+        model=f'models/{self.model}',
+        prompt=prompt,
+        temperature=generation_config.temperature,
+        top_k=generation_config.top_k,
+        top_p=generation_config.top_p,
+        candidate_count=generation_config.candidate_count,
+        max_output_tokens=generation_config.max_output_tokens,
+        stop_sequences=generation_config.stop_sequences,
+    )
+    return pg.Dict(
+        candidates=[
+            pg.Dict(content=pg.Dict(parts=[pg.Dict(text=c['output'])]))
+            for c in completion.candidates
+        ]
+    )
+
+
+class _LegacyChatModel(_LegacyGenerativeModel):
+  """Legacy GenAI chat model."""
+
+  def generate(
+      self, prompt: str, generation_config: genai.GenerationConfig
+  ) -> pg.Dict:
+    response: genai.types.ChatResponse = genai.chat(
+        model=f'models/{self.model}',
+        messages=prompt,
+        temperature=generation_config.temperature,
+        top_k=generation_config.top_k,
+        top_p=generation_config.top_p,
+        candidate_count=generation_config.candidate_count,
+    )
+    return pg.Dict(
+        candidates=[
+            pg.Dict(content=pg.Dict(parts=[pg.Dict(text=c['content'])]))
+            for c in response.candidates
+        ]
+    )
+
+
 class _ModelHub:
   """Google Generative AI model hub."""
 
   def __init__(self):
     self._model_cache = {}
 
-  def get(self, model_name: str) -> genai.GenerativeModel:
+  def get(
+      self, model_name: str
+  ) -> genai.GenerativeModel | _LegacyGenerativeModel:
     """Gets a generative model by model id."""
     model = self._model_cache.get(model_name, None)
     if model is None:
-      model = genai.GenerativeModel(model_name)
+      model_info = genai.get_model(f'models/{model_name}')
+      if 'generateContent' in model_info.supported_generation_methods:
+        model = genai.GenerativeModel(model_name)
+      elif 'generateText' in model_info.supported_generation_methods:
+        model = _LegacyCompletionModel(model_name)
+      elif 'generateMessage' in model_info.supported_generation_methods:
+        model = _LegacyChatModel(model_name)
+      else:
+        raise ValueError(f'Unsupported model: {model_name!r}')
       self._model_cache[model_name] = model
     return model
 
@@ -172,14 +262,26 @@ _GOOGLE_GENAI_MODEL_HUB = _ModelHub()
 #
 
 
-class GeminiPro(Gemini):
+class GeminiPro(GenAI):
   """Gemini Pro model."""
 
   model = 'gemini-pro'
 
 
-class GeminiProVision(Gemini):
+class GeminiProVision(GenAI):
   """Gemini Pro vision model."""
 
   model = 'gemini-pro-vision'
   multimodal = True
+
+
+class Palm2(GenAI):
+  """PaLM2 model."""
+
+  model = 'text-bison-001'
+
+
+class Palm2_IT(GenAI):  # pylint: disable=invalid-name
+  """PaLM2 instruction-tuned model."""
+
+  model = 'chat-bison-001'

langfun/core/llms/{gemini_test.py → google_genai_test.py}

@@ -20,7 +20,7 @@ from unittest import mock
 from google import generativeai as genai
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
-from langfun.core.llms import gemini
+from langfun.core.llms import google_genai
 import pyglove as pg
 
 
@@ -36,6 +36,29 @@ example_image = (
 )
 
 
+def mock_get_model(model_name, *args, **kwargs):
+  del args, kwargs
+  if 'gemini' in model_name:
+    method = 'generateContent'
+  elif 'chat' in model_name:
+    method = 'generateMessage'
+  else:
+    method = 'generateText'
+  return pg.Dict(supported_generation_methods=[method])
+
+
+def mock_generate_text(*, model, prompt, **kwargs):
+  return pg.Dict(
+      candidates=[pg.Dict(output=f'{prompt} to {model} with {kwargs}')]
+  )
+
+
+def mock_chat(*, model, messages, **kwargs):
+  return pg.Dict(
+      candidates=[pg.Dict(content=f'{messages} to {model} with {kwargs}')]
+  )
+
+
 def mock_generate_content(content, generation_config, **kwargs):
   del kwargs
   c = generation_config
@@ -68,12 +91,12 @@ def mock_generate_content(content, generation_config, **kwargs):
   )
 
 
-class GeminiTest(unittest.TestCase):
-  """Tests for Evergreen language model."""
+class GenAITest(unittest.TestCase):
+  """Tests for Google GenAI model."""
 
   def test_content_from_message_text_only(self):
     text = 'This is a beautiful day'
-    model = gemini.GeminiPro()
+    model = google_genai.GeminiPro()
     chunks = model._content_from_message(lf.UserMessage(text))
     self.assertEqual(chunks, [text])
 
@@ -85,9 +108,9 @@ class GeminiTest(unittest.TestCase):
 
     # Non-multimodal model.
     with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
-      gemini.GeminiPro()._content_from_message(message)
+      google_genai.GeminiPro()._content_from_message(message)
 
-    model = gemini.GeminiProVision()
+    model = google_genai.GeminiProVision()
     chunks = model._content_from_message(message)
     self.maxDiff = None
     self.assertEqual(
@@ -118,7 +141,7 @@ class GeminiTest(unittest.TestCase):
             ],
         ),
     )
-    model = gemini.GeminiProVision()
+    model = google_genai.GeminiProVision()
     result = model._response_to_result(response)
     self.assertEqual(
         result,
@@ -129,34 +152,79 @@ class GeminiTest(unittest.TestCase):
     )
 
   def test_model_hub(self):
-    model = gemini._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
+    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(gemini._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), 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`'):
-      _ = gemini.GeminiPro()._api_initialized
+      _ = google_genai.GeminiPro()._api_initialized
 
-    self.assertTrue(gemini.GeminiPro(api_key='abc')._api_initialized)
+    self.assertTrue(google_genai.GeminiPro(api_key='abc')._api_initialized)
     os.environ['GOOGLE_API_KEY'] = 'abc'
-    self.assertTrue(gemini.GeminiPro()._api_initialized)
+    self.assertTrue(google_genai.GeminiPro()._api_initialized)
     del os.environ['GOOGLE_API_KEY']
 
   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
       mock_generate.side_effect = mock_generate_content
 
-      lm = gemini.GeminiPro(api_key='test_key')
+      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.'
           ),
       )
+      genai.get_model = orig_get_model
+
+  def test_call_with_legacy_completion_model(self):
+    orig_get_model = genai.get_model
+    genai.get_model = mock_get_model
+    orig_generate_text = genai.generate_text
+    genai.generate_text = mock_generate_text
+
+    lm = google_genai.Palm2(api_key='test_key')
+    self.maxDiff = None
+    self.assertEqual(
+        lm('hello', temperature=2.0, top_k=20).text,
+        (
+            "hello to models/text-bison-001 with {'temperature': 2.0, "
+            "'top_k': 20, 'top_p': None, 'candidate_count': 1, "
+            "'max_output_tokens': None, 'stop_sequences': None}"
+        ),
+    )
+    genai.get_model = orig_get_model
+    genai.generate_text = orig_generate_text
+
+  def test_call_with_legacy_chat_model(self):
+    orig_get_model = genai.get_model
+    genai.get_model = mock_get_model
+    orig_chat = genai.chat
+    genai.chat = mock_chat
+
+    lm = google_genai.Palm2_IT(api_key='test_key')
+    self.maxDiff = None
+    self.assertEqual(
+        lm('hello', temperature=2.0, top_k=20).text,
+        (
+            "hello to models/chat-bison-001 with {'temperature': 2.0, "
+            "'top_k': 20, 'top_p': None, 'candidate_count': 1}"
+        ),
+    )
+    genai.get_model = orig_get_model
+    genai.chat = orig_chat
 
 
 if __name__ == '__main__':