langfun 0.1.2.dev202501060804__py3-none-any.whl → 0.1.2.dev202501100804__py3-none-any.whl

langfun/core/llms/openai_test.py

@@ -13,102 +13,9 @@
 # limitations under the License.
 """Tests for OpenAI models."""
 
-from typing import Any
 import unittest
-from unittest import mock
-
 import langfun.core as lf
-from langfun.core import modalities as lf_modalities
 from langfun.core.llms import openai
-import pyglove as pg
-import requests
-
-
-def mock_chat_completion_request(url: str, json: dict[str, Any], **kwargs):
-  del url, kwargs
-  messages = json['messages']
-  if len(messages) > 1:
-    system_message = f' system={messages[0]["content"]}'
-  else:
-    system_message = ''
-
-  if 'response_format' in json:
-    response_format = f' format={json["response_format"]["type"]}'
-  else:
-    response_format = ''
-
-  choices = []
-  for k in range(json['n']):
-    if json.get('logprobs'):
-      logprobs = dict(
-          content=[
-              dict(
-                  token='chosen_token',
-                  logprob=0.5,
-                  top_logprobs=[
-                      dict(
-                          token=f'alternative_token_{i + 1}',
-                          logprob=0.1
-                      ) for i in range(3)
-                  ]
-              )
-          ]
-      )
-    else:
-      logprobs = None
-
-    choices.append(dict(
-        message=dict(
-            content=(
-                f'Sample {k} for message.{system_message}{response_format}'
-            )
-        ),
-        logprobs=logprobs,
-    ))
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str(
-      dict(
-          choices=choices,
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=100,
-              completion_tokens=100,
-              total_tokens=200,
-          ),
-      )
-  ).encode()
-  return response
-
-
-def mock_chat_completion_request_vision(
-    url: str, json: dict[str, Any], **kwargs
-):
-  del url, kwargs
-  choices = []
-  urls = [
-      c['image_url']['url']
-      for c in json['messages'][0]['content'] if c['type'] == 'image_url'
-  ]
-  for k in range(json['n']):
-    choices.append(pg.Dict(
-        message=pg.Dict(
-            content=f'Sample {k} for message: {"".join(urls)}'
-        ),
-        logprobs=None,
-    ))
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str(
-      dict(
-          choices=choices,
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=100,
-              completion_tokens=100,
-              total_tokens=200,
-          ),
-      )
-  ).encode()
-  return response
 
 
 class OpenAITest(unittest.TestCase):
@@ -130,6 +37,15 @@ class OpenAITest(unittest.TestCase):
         openai.Gpt35(api_key='test_key').resource_id, 'OpenAI(text-davinci-003)'
     )
 
+  def test_headers(self):
+    self.assertEqual(
+        openai.Gpt35(api_key='test_key').headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer test_key',
+        },
+    )
+
   def test_max_concurrency(self):
     self.assertGreater(openai.Gpt35(api_key='test_key').max_concurrency, 0)
 
@@ -156,340 +72,14 @@ class OpenAITest(unittest.TestCase):
           )
       )
 
-  def test_call_chat_completion(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = openai.OpenAI(
-          model='gpt-4',
-          api_key='test_key',
-          organization='my_org',
-          project='my_project'
-      )
-      self.assertEqual(
-          lm('hello', sampling_options=lf.LMSamplingOptions(n=2)),
-          'Sample 0 for message.',
-      )
-
-  def test_call_chat_completion_with_logprobs(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = openai.OpenAI(
-          model='gpt-4',
-          api_key='test_key',
-          organization='my_org',
-          project='my_project'
-      )
-      results = lm.sample(['hello'], logprobs=True)
-      self.assertEqual(len(results), 1)
-      self.assertEqual(
-          results[0],
-          lf.LMSamplingResult(
-              [
-                  lf.LMSample(
-                      response=lf.AIMessage(
-                          text='Sample 0 for message.',
-                          metadata={
-                              'score': 0.0,
-                              'logprobs': [(
-                                  'chosen_token',
-                                  0.5,
-                                  [
-                                      ('alternative_token_1', 0.1),
-                                      ('alternative_token_2', 0.1),
-                                      ('alternative_token_3', 0.1),
-                                  ],
-                              )],
-                              'is_cached': False,
-                              'usage': lf.LMSamplingUsage(
-                                  prompt_tokens=100,
-                                  completion_tokens=100,
-                                  total_tokens=200,
-                                  estimated_cost=0.009,
-                              ),
-                          },
-                          tags=['lm-response'],
-                      ),
-                      logprobs=[(
-                          'chosen_token',
-                          0.5,
-                          [
-                              ('alternative_token_1', 0.1),
-                              ('alternative_token_2', 0.1),
-                              ('alternative_token_3', 0.1),
-                          ],
-                      )],
-                  )
-              ],
-              usage=lf.LMSamplingUsage(
-                  prompt_tokens=100,
-                  completion_tokens=100,
-                  total_tokens=200,
-                  estimated_cost=0.009,
-              ),
-          ),
-      )
-
-  def test_call_chat_completion_vision(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request_vision
-      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',
-        )
-    lm_3 = openai.Gpt35Turbo(api_key='test_key')
-    with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
-      lm_3(
-          lf.UserMessage(
-              'hello <<[[image]]>>',
-              image=lf_modalities.Image.from_uri('https://fake/image')
-          ),
-      )
-
-  def test_sample_chat_completion(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      openai.SUPPORTED_MODELS_AND_SETTINGS['gpt-4'].update({
-          'cost_per_1k_input_tokens': 1.0,
-          'cost_per_1k_output_tokens': 1.0,
-      })
-      lm = openai.OpenAI(api_key='test_key', model='gpt-4')
-      results = lm.sample(
-          ['hello', 'bye'], sampling_options=lf.LMSamplingOptions(n=3)
-      )
-
-    self.assertEqual(len(results), 2)
-    print(results[0])
-    self.assertEqual(
-        results[0],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200,
-                estimated_cost=0.2,
-            ),
-        ),
-    )
+  def test_estimate_cost(self):
     self.assertEqual(
-        results[1],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200,
-                estimated_cost=0.2,
-            ),
+        openai.Gpt4(api_key='test_key').estimate_cost(
+            num_input_tokens=100, num_output_tokens=100
         ),
+        0.009
     )
 
-  def test_sample_with_contextual_options(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = openai.OpenAI(api_key='test_key', model='text-davinci-003')
-      with lf.use_settings(sampling_options=lf.LMSamplingOptions(n=2)):
-        results = lm.sample(['hello'])
-
-    self.assertEqual(len(results), 1)
-    self.assertEqual(
-        results[0],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        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 message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=50,
-                            completion_tokens=50,
-                            total_tokens=100,
-                        ),
-                        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_call_with_system_message(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = openai.OpenAI(api_key='test_key', model='gpt-4')
-      self.assertEqual(
-          lm(
-              lf.UserMessage(
-                  'hello',
-                  system_message='hi',
-              ),
-              sampling_options=lf.LMSamplingOptions(n=2)
-          ),
-          '''Sample 0 for message. system=[{'type': 'text', 'text': 'hi'}]''',
-      )
-
-  def test_call_with_json_schema(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = openai.OpenAI(api_key='test_key', model='gpt-4')
-      self.assertEqual(
-          lm(
-              lf.UserMessage(
-                  'hello',
-                  json_schema={
-                      'type': 'object',
-                      'properties': {
-                          'name': {'type': 'string'},
-                      },
-                      'required': ['name'],
-                      'title': 'Person',
-                  }
-              ),
-              sampling_options=lf.LMSamplingOptions(n=2)
-          ),
-          'Sample 0 for message. format=json_schema',
-      )
-
-    # Test bad json schema.
-    with self.assertRaisesRegex(ValueError, '`json_schema` must be a dict'):
-      lm(lf.UserMessage('hello', json_schema='foo'))
-
-    with self.assertRaisesRegex(
-        ValueError, 'The root of `json_schema` must have a `title` field'
-    ):
-      lm(lf.UserMessage('hello', json_schema={}))
-
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/vertexai.py

@@ -90,6 +90,8 @@ class VertexAI(gemini.Gemini):
       )
 
     self._project = project
+    self._location = location
+
     credentials = self.credentials
     if credentials is None:
       # Use default credentials.
@@ -114,9 +116,10 @@ class VertexAI(gemini.Gemini):
 
   @property
   def api_endpoint(self) -> str:
+    assert self._api_initialized
     return (
-        f'https://{self.location}-aiplatform.googleapis.com/v1/projects/'
-        f'{self.project}/locations/{self.location}/publishers/google/'
+        f'https://{self._location}-aiplatform.googleapis.com/v1/projects/'
+        f'{self._project}/locations/{self._location}/publishers/google/'
         f'models/{self.model}:generateContent'
     )
 
@@ -126,6 +129,7 @@ class VertexAIGeminiFlash2_0ThinkingExp_20241219(VertexAI):  # pylint: disable=i
 
   api_version = 'v1alpha'
   model = 'gemini-2.0-flash-thinking-exp-1219'
+  timeout = None
 
 
 class VertexAIGeminiFlash2_0Exp(VertexAI):  # pylint: disable=invalid-name

langfun/core/llms/vertexai_test.py

@@ -41,7 +41,7 @@ class VertexAITest(unittest.TestCase):
     os.environ['VERTEXAI_LOCATION'] = 'us-central1'
     model = vertexai.VertexAIGeminiPro1()
     self.assertTrue(model.model_id.startswith('VertexAI('))
-    self.assertIsNotNone(model.api_endpoint)
+    self.assertIn('us-central1', model.api_endpoint)
     self.assertTrue(model._api_initialized)
     self.assertIsNotNone(model._session)
     del os.environ['VERTEXAI_PROJECT']

langfun/core/modalities/mime.py

@@ -138,9 +138,17 @@ class Mime(lf.Modality):
 
   @property
   def content_uri(self) -> str:
+    """Returns the URI with encoded content."""
     base64_content = base64.b64encode(self.to_bytes()).decode()
     return f'data:{self.mime_type};base64,{base64_content}'
 
+  @property
+  def embeddable_uri(self) -> str:
+    """Returns the URI that can be embedded in HTML."""
+    if self.uri and self.uri.lower().startswith(('http:', 'https:', 'ftp:')):
+      return self.uri
+    return self.content_uri
+
   @classmethod
   def from_uri(cls, uri: str, **kwargs) -> 'Mime':
     if cls is Mime:

langfun/core/modalities/mime_test.py

@@ -23,12 +23,12 @@ import pyglove as pg
 
 def mock_request(*args, **kwargs):
   del args, kwargs
-  return pg.Dict(content='foo')
+  return pg.Dict(content=b'foo')
 
 
 def mock_readfile(*args, **kwargs):
   del args, kwargs
-  return 'bar'
+  return b'bar'
 
 
 class CustomMimeTest(unittest.TestCase):
@@ -65,17 +65,32 @@ class CustomMimeTest(unittest.TestCase):
     ):
       mime.Custom('text/plain')
 
+  def test_uri(self):
+    content = mime.Custom.from_uri('http://mock/web/a.txt', mime='text/plain')
+    with mock.patch('requests.get') as mock_requests_stub:
+      mock_requests_stub.side_effect = mock_request
+      self.assertEqual(content.uri, 'http://mock/web/a.txt')
+      self.assertEqual(content.content_uri, 'data:text/plain;base64,Zm9v')
+      self.assertEqual(content.embeddable_uri, 'http://mock/web/a.txt')
+
+    content = mime.Custom.from_uri('a.txt', mime='text/plain')
+    with mock.patch('pyglove.io.readfile') as mock_readfile_stub:
+      mock_readfile_stub.side_effect = mock_readfile
+      self.assertEqual(content.uri, 'a.txt')
+      self.assertEqual(content.content_uri, 'data:text/plain;base64,YmFy')
+      self.assertEqual(content.embeddable_uri, 'data:text/plain;base64,YmFy')
+
   def test_from_uri(self):
     content = mime.Custom.from_uri('http://mock/web/a.txt', mime='text/plain')
     with mock.patch('requests.get') as mock_requests_stub:
       mock_requests_stub.side_effect = mock_request
-      self.assertEqual(content.to_bytes(), 'foo')
+      self.assertEqual(content.to_bytes(), b'foo')
       self.assertEqual(content.mime_type, 'text/plain')
 
     content = mime.Custom.from_uri('a.txt', mime='text/plain')
     with mock.patch('pyglove.io.readfile') as mock_readfile_stub:
       mock_readfile_stub.side_effect = mock_readfile
-      self.assertEqual(content.to_bytes(), 'bar')
+      self.assertEqual(content.to_bytes(), b'bar')
       self.assertEqual(content.mime_type, 'text/plain')
 
   def assert_html_content(self, html, expected):

langfun/core/modality_test.py

@@ -11,7 +11,6 @@
 # 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 modality."""
 from typing import Any
 import unittest
 

langfun/core/structured/mapping.py

@@ -46,15 +46,15 @@ class MappingError(Exception):  # pylint: disable=g-bad-exception-name
     r = io.StringIO()
     error_message = str(self.cause).rstrip()
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {error_message}', 'magenta'
         )
     )
     if include_lm_response:
       r.write('\n\n')
-      r.write(lf.colored('[LM Response]', 'blue', styles=['bold']))
+      r.write(pg.colored('[LM Response]', 'blue', styles=['bold']))
       r.write('\n')
-      r.write(lf.colored(self.lm_response.text, 'blue'))
+      r.write(pg.colored(self.lm_response.text, 'blue'))
     return r.getvalue()
 
 
@@ -163,27 +163,27 @@ class MappingExample(lf.NaturalLanguageFormattable,
   def natural_language_format(self) -> str:
     result = io.StringIO()
     if self.context:
-      result.write(lf.colored('[CONTEXT]\n', styles=['bold']))
-      result.write(lf.colored(self.context, color='magenta'))
+      result.write(pg.colored('[CONTEXT]\n', styles=['bold']))
+      result.write(pg.colored(self.context, color='magenta'))
       result.write('\n\n')
 
-    result.write(lf.colored('[INPUT]\n', styles=['bold']))
-    result.write(lf.colored(self.input_repr(), color='green'))
+    result.write(pg.colored('[INPUT]\n', styles=['bold']))
+    result.write(pg.colored(self.input_repr(), color='green'))
 
     if self.schema is not None:
       result.write('\n\n')
-      result.write(lf.colored('[SCHEMA]\n', styles=['bold']))
-      result.write(lf.colored(self.schema_repr(), color='red'))
+      result.write(pg.colored('[SCHEMA]\n', styles=['bold']))
+      result.write(pg.colored(self.schema_repr(), color='red'))
 
     if schema_lib.MISSING != self.output:
       result.write('\n\n')
-      result.write(lf.colored('[OUTPUT]\n', styles=['bold']))
-      result.write(lf.colored(self.output_repr(), color='blue'))
+      result.write(pg.colored('[OUTPUT]\n', styles=['bold']))
+      result.write(pg.colored(self.output_repr(), color='blue'))
 
     if self.metadata:
       result.write('\n\n')
-      result.write(lf.colored('[METADATA]\n', styles=['bold']))
-      result.write(lf.colored(str(self.metadata), color='cyan'))
+      result.write(pg.colored('[METADATA]\n', styles=['bold']))
+      result.write(pg.colored(str(self.metadata), color='cyan'))
     return result.getvalue().strip()
 
   @classmethod

langfun/core/structured/mapping_test.py

@@ -29,11 +29,11 @@ class MappingErrorTest(unittest.TestCase):
         lf.AIMessage('hi'), ValueError('Cannot parse message.')
     )
     self.assertEqual(
-        lf.text_formatting.decolored(str(error)),
+        pg.decolor(str(error)),
         'ValueError: Cannot parse message.\n\n[LM Response]\nhi',
     )
     self.assertEqual(
-        lf.text_formatting.decolored(error.format(include_lm_response=False)),
+        pg.decolor(error.format(include_lm_response=False)),
         'ValueError: Cannot parse message.',
     )