langfun 0.1.2.dev202412020805__py3-none-any.whl → 0.1.2.dev202412050804__py3-none-any.whl

langfun/core/llms/vertexai_test.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 """Tests for Gemini models."""
 
+import base64
 import os
 from typing import Any
 import unittest
 from unittest import mock
 
-from google.cloud.aiplatform import models as aiplatform_models
-from vertexai import generative_models
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 from langfun.core.llms import vertexai
@@ -39,33 +38,6 @@ example_image = (
 )
 
 
-def mock_generate_content(content, generation_config, **kwargs):
-  del kwargs
-  c = pg.Dict(generation_config.to_dict())
-  return generative_models.GenerationResponse.from_dict({
-      'candidates': [
-          {
-              'index': 0,
-              'content': {
-                  'role': 'model',
-                  'parts': [
-                      {
-                          'text': (
-                              f'This is a response to {content[0]} with '
-                              f'temperature={c.temperature}, '
-                              f'top_p={c.top_p}, '
-                              f'top_k={c.top_k}, '
-                              f'max_tokens={c.max_output_tokens}, '
-                              f'stop={"".join(c.stop_sequences)}.'
-                          )
-                      },
-                  ],
-              },
-          },
-      ]
-  })
-
-
 def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
   del url, kwargs
   c = pg.Dict(json['generationConfig'])
@@ -100,273 +72,7 @@ def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
   return response
 
 
-def mock_endpoint_predict(instances, **kwargs):
-  del kwargs
-  assert len(instances) == 1
-  return aiplatform_models.Prediction(
-      predictions=[
-          f"This is a response to {instances[0]['prompt']} with"
-          f" temperature={instances[0]['temperature']},"
-          f" top_p={instances[0]['top_p']}, top_k={instances[0]['top_k']},"
-          f" max_tokens={instances[0]['max_tokens']}."
-      ],
-      deployed_model_id='',
-  )
-
-
 class VertexAITest(unittest.TestCase):
-  """Tests for Vertex model."""
-
-  def test_content_from_message_text_only(self):
-    text = 'This is a beautiful day'
-    model = vertexai.VertexAIGeminiPro1Vision()
-    chunks = model._content_from_message(lf.UserMessage(text))
-    self.assertEqual(chunks, [text])
-
-  def test_content_from_message_mm(self):
-    message = lf.UserMessage(
-        'This is an <<[[image]]>>, what is it?',
-        image=lf_modalities.Image.from_bytes(example_image),
-    )
-
-    # Non-multimodal model.
-    with self.assertRaisesRegex(lf.ModalityError, 'Unsupported modality'):
-      vertexai.VertexAIPalm2()._content_from_message(message)
-
-    model = vertexai.VertexAIGeminiPro1Vision()
-    chunks = model._content_from_message(message)
-    self.maxDiff = None
-    self.assertEqual([chunks[0], chunks[2]], ['This is an', ', what is it?'])
-    self.assertIsInstance(chunks[1], generative_models.Part)
-
-  def test_generation_response_to_message_text_only(self):
-    response = generative_models.GenerationResponse.from_dict({
-        'candidates': [
-            {
-                'index': 0,
-                'content': {
-                    'role': 'model',
-                    'parts': [
-                        {
-                            'text': 'hello world',
-                        },
-                    ],
-                },
-            },
-        ],
-    })
-    model = vertexai.VertexAIGeminiPro1Vision()
-    message = model._generation_response_to_message(response)
-    self.assertEqual(message, lf.AIMessage('hello world'))
-
-  def test_model_hub(self):
-    with mock.patch(
-        'vertexai.generative_models.'
-        'GenerativeModel.__init__'
-    ) as mock_model_init:
-      mock_model_init.side_effect = lambda *args, **kwargs: None
-      model = vertexai._VERTEXAI_MODEL_HUB.get_generative_model(
-          'gemini-1.0-pro'
-      )
-      self.assertIsNotNone(model)
-      self.assertIs(
-          vertexai._VERTEXAI_MODEL_HUB.get_generative_model('gemini-1.0-pro'),
-          model,
-      )
-
-    with mock.patch(
-        'vertexai.language_models.'
-        'TextGenerationModel.from_pretrained'
-    ) as mock_model_init:
-
-      class TextGenerationModel:
-        pass
-
-      mock_model_init.side_effect = lambda *args, **kw: TextGenerationModel()
-      model = vertexai._VERTEXAI_MODEL_HUB.get_text_generation_model(
-          'text-bison'
-      )
-      self.assertIsNotNone(model)
-      self.assertIs(
-          vertexai._VERTEXAI_MODEL_HUB.get_text_generation_model('text-bison'),
-          model,
-      )
-
-  def test_project_and_location_check(self):
-    with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
-      _ = vertexai.VertexAIGeminiPro1Vision()._api_initialized
-
-    with self.assertRaisesRegex(ValueError, 'Please specify `location`'):
-      _ = vertexai.VertexAIGeminiPro1Vision(project='abc')._api_initialized
-
-    self.assertTrue(
-        vertexai.VertexAIGeminiPro1Vision(
-            project='abc', location='us-central1'
-        )._api_initialized
-    )
-
-    os.environ['VERTEXAI_PROJECT'] = 'abc'
-    os.environ['VERTEXAI_LOCATION'] = 'us-central1'
-    self.assertTrue(vertexai.VertexAIGeminiPro1Vision()._api_initialized)
-    del os.environ['VERTEXAI_PROJECT']
-    del os.environ['VERTEXAI_LOCATION']
-
-  def test_generation_config(self):
-    model = vertexai.VertexAIGeminiPro1Vision()
-    json_schema = {
-        'type': 'object',
-        'properties': {
-            'name': {'type': 'string'},
-        },
-        'required': ['name'],
-        'title': 'Person',
-    }
-    config = model._generation_config(
-        lf.UserMessage('hi', json_schema=json_schema),
-        lf.LMSamplingOptions(
-            temperature=2.0,
-            top_p=1.0,
-            top_k=20,
-            max_tokens=1024,
-            stop=['\n'],
-        ),
-    )
-    actual = config.to_dict()
-    # There is a discrepancy between the `property_ordering` in the
-    # Google-internal version and the open-source version.
-    actual['response_schema'].pop('property_ordering', None)
-    if pg.KeyPath.parse('response_schema.type_').get(actual):
-      actual['response_schema']['type'] = actual['response_schema'].pop('type_')
-    if pg.KeyPath.parse('response_schema.properties.name.type_').get(actual):
-      actual['response_schema']['properties']['name']['type'] = actual[
-          'response_schema']['properties']['name'].pop('type_')
-
-    self.assertEqual(
-        actual,
-        dict(
-            temperature=2.0,
-            top_p=1.0,
-            top_k=20.0,
-            max_output_tokens=1024,
-            stop_sequences=['\n'],
-            response_mime_type='application/json',
-            response_schema={
-                'type': 'OBJECT',
-                'properties': {
-                    'name': {'type': 'STRING'}
-                },
-                'required': ['name'],
-                'title': 'Person',
-            }
-        ),
-    )
-    with self.assertRaisesRegex(
-        ValueError, '`json_schema` must be a dict, got'
-    ):
-      model._generation_config(
-          lf.UserMessage('hi', json_schema='not a dict'),
-          lf.LMSamplingOptions(),
-      )
-
-  def test_call_generative_model(self):
-    with mock.patch(
-        'vertexai.generative_models.'
-        'GenerativeModel.__init__'
-    ) as mock_model_init:
-      mock_model_init.side_effect = lambda *args, **kwargs: None
-
-      with mock.patch(
-          'vertexai.generative_models.'
-          'GenerativeModel.generate_content'
-      ) as mock_generate:
-        mock_generate.side_effect = mock_generate_content
-
-        lm = vertexai.VertexAIGeminiPro1Vision(
-            project='abc', location='us-central1'
-        )
-        self.assertEqual(
-            lm(
-                'hello',
-                temperature=2.0,
-                top_p=1.0,
-                top_k=20,
-                max_tokens=1024,
-                stop='\n',
-            ).text,
-            (
-                'This is a response to hello with temperature=2.0, '
-                'top_p=1.0, top_k=20.0, max_tokens=1024, stop=\n.'
-            ),
-        )
-
-  def test_call_text_generation_model(self):
-    with mock.patch(
-        'vertexai.language_models.'
-        'TextGenerationModel.from_pretrained'
-    ) as mock_model_init:
-
-      class TextGenerationModel:
-
-        def predict(self, prompt, **kwargs):
-          c = pg.Dict(kwargs)
-          return pg.Dict(
-              text=(
-                  f'This is a response to {prompt} with '
-                  f'temperature={c.temperature}, '
-                  f'top_p={c.top_p}, '
-                  f'top_k={c.top_k}, '
-                  f'max_tokens={c.max_output_tokens}, '
-                  f'stop={"".join(c.stop_sequences)}.'
-              )
-          )
-
-      mock_model_init.side_effect = lambda *args, **kw: TextGenerationModel()
-      lm = vertexai.VertexAIPalm2(project='abc', location='us-central1')
-      self.assertEqual(
-          lm(
-              'hello',
-              temperature=2.0,
-              top_p=1.0,
-              top_k=20,
-              max_tokens=1024,
-              stop='\n',
-          ).text,
-          (
-              'This is a response to hello with temperature=2.0, '
-              'top_p=1.0, top_k=20, max_tokens=1024, stop=\n.'
-          ),
-      )
-
-  def test_call_endpoint_model(self):
-    with mock.patch(
-        'google.cloud.aiplatform.models.Endpoint.__init__'
-    ) as mock_model_init:
-      mock_model_init.side_effect = lambda *args, **kwargs: None
-      with mock.patch(
-          'google.cloud.aiplatform.models.Endpoint.predict'
-      ) as mock_model_predict:
-
-        mock_model_predict.side_effect = mock_endpoint_predict
-        lm = vertexai.VertexAI(
-            'custom',
-            endpoint_name='123',
-            project='abc',
-            location='us-central1',
-        )
-        self.assertEqual(
-            lm(
-                'hello',
-                temperature=2.0,
-                top_p=1.0,
-                top_k=20,
-                max_tokens=50,
-            ),
-            'This is a response to hello with temperature=2.0, top_p=1.0,'
-            ' top_k=20, max_tokens=50.',
-        )
-
-
-class VertexRestfulAITest(unittest.TestCase):
   """Tests for Vertex model with REST API."""
 
   def test_content_from_message_text_only(self):
@@ -376,9 +82,9 @@ class VertexRestfulAITest(unittest.TestCase):
     self.assertEqual(chunks, {'role': 'user', 'parts': [{'text': text}]})
 
   def test_content_from_message_mm(self):
+    image = lf_modalities.Image.from_bytes(example_image)
     message = lf.UserMessage(
-        'This is an <<[[image]]>>, what is it?',
-        image=lf_modalities.Image.from_bytes(example_image),
+        'This is an <<[[image]]>>, what is it?', image=image
     )
 
     # Non-multimodal model.
@@ -386,46 +92,25 @@ class VertexRestfulAITest(unittest.TestCase):
       vertexai.VertexAIGeminiPro1()._content_from_message(message)
 
     model = vertexai.VertexAIGeminiPro1Vision()
-    chunks = model._content_from_message(message)
-    self.maxDiff = None
-    self.assertEqual([chunks[0], chunks[2]], ['This is an', ', what is it?'])
-    self.assertIsInstance(chunks[1], generative_models.Part)
-
-  def test_generation_response_to_message_text_only(self):
-    response = generative_models.GenerationResponse.from_dict({
-        'candidates': [
-            {
-                'index': 0,
-                'content': {
-                    'role': 'model',
-                    'parts': [
-                        {
-                            'text': 'hello world',
-                        },
-                    ],
+    content = model._content_from_message(message)
+    self.assertEqual(
+        content,
+        {
+            'role': 'user',
+            'parts': [
+                {'text': 'This is an'},
+                {
+                    'inlineData': {
+                        'data': base64.b64encode(example_image).decode(),
+                        'mimeType': 'image/png',
+                    }
                 },
-            },
-        ],
-    })
-    model = vertexai.VertexAIGeminiPro1Vision()
-    message = model._generation_response_to_message(response)
-    self.assertEqual(message, lf.AIMessage('hello world'))
-
-  def test_model_hub(self):
-    with mock.patch(
-        'vertexai.generative_models.'
-        'GenerativeModel.__init__'
-    ) as mock_model_init:
-      mock_model_init.side_effect = lambda *args, **kwargs: None
-      model = vertexai._VERTEXAI_MODEL_HUB.get_generative_model(
-          'gemini-1.0-pro'
-      )
-      self.assertIsNotNone(model)
-      self.assertIs(
-          vertexai._VERTEXAI_MODEL_HUB.get_generative_model('gemini-1.0-pro'),
-          model,
-      )
+                {'text': ', what is it?'},
+            ],
+        },
+    )
 
+  @mock.patch.object(vertexai.VertexAI, 'credentials', new=True)
   def test_project_and_location_check(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `project`'):
       _ = vertexai.VertexAIGeminiPro1()._api_initialized
@@ -496,6 +181,7 @@ class VertexRestfulAITest(unittest.TestCase):
           lf.LMSamplingOptions(),
       )
 
+  @mock.patch.object(vertexai.VertexAI, 'credentials', new=True)
   def test_call_model(self):
     with mock.patch('requests.Session.post') as mock_generate:
       mock_generate.side_effect = mock_requests_post

langfun/core/structured/__init__.py

@@ -69,6 +69,8 @@ from langfun.core.structured.prompting import query
 from langfun.core.structured.prompting import query_prompt
 from langfun.core.structured.prompting import query_output
 from langfun.core.structured.prompting import query_reward
+from langfun.core.structured.prompting import QueryInvocation
+from langfun.core.structured.prompting import track_queries
 
 from langfun.core.structured.description import DescribeStructure
 from langfun.core.structured.description import describe

langfun/core/structured/prompting.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 """Symbolic query."""
 
+import contextlib
 import functools
-from typing import Any, Callable, Type, Union
+from typing import Annotated, Any, Callable, Iterator, Type, Union
 
 import langfun.core as lf
 from langfun.core.llms import fake
@@ -102,7 +103,7 @@ def _query_structure_cls(
 
 
 def query(
-    prompt: Union[str, pg.Symbolic],
+    prompt: Union[str, lf.Template, Any],
     schema: Union[
         schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
     ] = None,
@@ -119,7 +120,7 @@ def query(
     skip_lm: bool = False,
     **kwargs,
 ) -> Any:
-  """Parse a natural langugage message based on schema.
+  """Queries an language model for a (maybe) structured output.
 
   Examples:
 
@@ -189,59 +190,93 @@ def query(
   """
     # Internal usage logging.
 
+  # Normalize query schema.
   # When `lf.query` is used for symbolic completion, schema is automatically
   # inferred when it is None.
   if isinstance(prompt, pg.Symbolic) and prompt.sym_partial and schema is None:
     schema = prompt.__class__
 
-  # Create a copy of the prompt if it has a parent object, so all child modality
-  # objects could be referred by path relative to the prompt.
-  if isinstance(prompt, lf.Template) and prompt.sym_parent:
-    prompt = prompt.clone()
-
-  if schema in (None, str):
-    # Query with natural language output.
-    output = lf.LangFunc.from_value(prompt, **kwargs)(
-        lm=lm, cache_seed=cache_seed, skip_lm=skip_lm
+  # Normalize query input.
+  if isinstance(prompt, (lf.Message, str)):
+    # Query with structured output.
+    prompt_kwargs = kwargs.copy()
+    prompt_kwargs.pop('template_str', None)
+    query_input = lf.Template.from_value(prompt, **prompt_kwargs)
+  elif isinstance(prompt, lf.Template):
+    # Create a copy of the prompt if it has a parent object, so all child
+    # modality objects could be referred by path relative to the prompt.
+    query_input = prompt.clone() if prompt.sym_parent is not None else prompt
+
+    # Attach template metadata from kwargs. This is used to pass through fields
+    # from kwargs to the rendered message.
+    template_metadata = {
+        k: v for k, v in kwargs.items() if k.startswith('metadata_')
+    }
+    query_input.rebind(
+        template_metadata, skip_notification=True, raise_on_no_change=False
     )
-    if response_postprocess:
-      processed_text = response_postprocess(output.text)
-      if processed_text != output.text:
-        output = lf.AIMessage(processed_text, source=output)
-    return output if returns_message else output.text
-
-  # Query with structured output.
-  prompt_kwargs = kwargs.copy()
-
-  # NOTE(daiyip): when `template_str` is passed in, it's intended to modify the
-  # QueryStructure template string. Therefore, we pop out the argument for
-  # prompt rendering.
-  prompt_kwargs.pop('template_str', None)
-
-  if isinstance(prompt, (str, lf.Message, lf.Template)):
-    prompt = lf.Template.from_value(prompt, **prompt_kwargs).render(lm=lm)
+  elif pg.MISSING_VALUE == prompt:
+    query_input = lf.UserMessage('')
   else:
-    prompt = schema_lib.mark_missing(prompt)
-
-  output = _query_structure_cls(protocol)(
-      input=prompt,
-      schema=schema,
-      default=default,
-      examples=examples,
-      response_postprocess=response_postprocess,
-      autofix=autofix if protocol == 'python' else 0,
-      **kwargs,
-  )(
-      lm=lm,
-      autofix_lm=autofix_lm or lm,
-      cache_seed=cache_seed,
-      skip_lm=skip_lm,
-  )
-  return output if returns_message else output.result
+    query_input = schema_lib.mark_missing(prompt)
+
+  with lf.track_usages() as usage_summary:
+    if schema in (None, str):
+      # Query with natural language output.
+      output_message = lf.LangFunc.from_value(query_input, **kwargs)(
+          lm=lm, cache_seed=cache_seed, skip_lm=skip_lm
+      )
+      if response_postprocess:
+        processed_text = response_postprocess(output_message.text)
+        if processed_text != output_message.text:
+          output_message = lf.AIMessage(processed_text, source=output_message)
+    else:
+      # Query with structured output.
+      output_message = _query_structure_cls(protocol)(
+          input=(
+              query_input.render(lm=lm)
+              if isinstance(query_input, lf.Template)
+              else query_input
+          ),
+          schema=schema,
+          default=default,
+          examples=examples,
+          response_postprocess=response_postprocess,
+          autofix=autofix if protocol == 'python' else 0,
+          **kwargs,
+      )(
+          lm=lm,
+          autofix_lm=autofix_lm or lm,
+          cache_seed=cache_seed,
+          skip_lm=skip_lm,
+      )
+
+  def _result(message: lf.Message):
+    return message.text if schema in (None, str) else message.result
+
+  # Track the query invocations.
+  if pg.MISSING_VALUE != prompt and not skip_lm:
+    trackers = lf.context_value('__query_trackers__', [])
+    if trackers:
+      invocation = QueryInvocation(
+          input=pg.Ref(query_input),
+          schema=(
+              schema_lib.Schema.from_value(schema)
+              if schema not in (None, str) else None
+          ),
+          output=pg.Ref(_result(output_message)),
+          lm=pg.Ref(lm),
+          examples=pg.Ref(examples) if examples else [],
+          usage_summary=usage_summary,
+      )
+      for i, (tracker, include_child_scopes) in enumerate(trackers):
+        if i == 0 or include_child_scopes:
+          tracker.append(invocation)
+  return output_message if returns_message else _result(output_message)
 
 
 def query_prompt(
-    prompt: Union[str, pg.Symbolic],
+    prompt: Union[str, lf.Template, Any],
     schema: Union[
         schema_lib.Schema, Type[Any], list[Type[Any]], dict[str, Any], None
     ] = None,
@@ -264,7 +299,7 @@ def query_output(
   kwargs.pop('prompt', None)
   kwargs.pop('lm', None)
   return query(
-      'Unused prompt', schema, lm=fake.StaticResponse(response), **kwargs
+      pg.MISSING_VALUE, schema, lm=fake.StaticResponse(response), **kwargs
   )
 
 
@@ -320,3 +355,69 @@ def _reward_fn(cls) -> Callable[
     args = [self, input, expected_output, metadata]
     return cls.__reward__(*args[:num_args])
   return _reward
+
+
+class QueryInvocation(pg.Object):
+  """A class to represent the invocation of `lf.query`."""
+
+  input: Annotated[
+      Union[lf.Template, pg.Symbolic],
+      'Mapping input of `lf.query`.'
+  ]
+  schema: pg.typing.Annotated[
+      schema_lib.schema_spec(noneable=True),
+      'Schema of `lf.query`.'
+  ]
+  output: Annotated[
+      Any,
+      'Mapping output of `lf.query`.'
+  ]
+  lm: Annotated[
+      lf.LanguageModel,
+      'Language model used for `lf.query`.'
+  ]
+  examples: Annotated[
+      list[mapping.MappingExample],
+      'Fewshot exemplars for `lf.query`.'
+  ]
+  usage_summary: Annotated[
+      lf.UsageSummary,
+      'Usage summary for `lf.query`.'
+  ]
+
+
+@contextlib.contextmanager
+def track_queries(
+    include_child_scopes: bool = True
+) -> Iterator[list[QueryInvocation]]:
+  """Track all queries made during the context.
+
+  Example:
+
+    ```
+    with lf.track_queries() as queries:
+      lf.query('hi', lm=lm)
+      lf.query('What is this {{image}}?', lm=lm, image=image)
+
+    print(queries)
+    ```
+
+  Args:
+    include_child_scopes: If True, the queries made in child scopes will be
+      included in the returned list. Otherwise, only the queries made in the
+      current scope will be included.
+
+  Yields:
+    A list of `QueryInvocation` objects representing the queries made during
+    the context.
+  """
+  trackers = lf.context_value('__query_trackers__', [])
+  tracker = []
+
+  with lf.context(
+      __query_trackers__=[(tracker, include_child_scopes)] + trackers
+  ):
+    try:
+      yield tracker
+    finally:
+      pass