langfun 0.1.2.dev202502110804__py3-none-any.whl → 0.1.2.dev202502120804__py3-none-any.whl

langfun/core/llms/openai_compatible.py

@@ -29,10 +29,6 @@ class OpenAICompatible(rest.REST):
       str, 'The name of the model to use.',
   ] = ''
 
-  multimodal: Annotated[
-      bool, 'Whether this model has multimodal support.'
-  ] = False
-
   @property
   def headers(self) -> dict[str, Any]:
     return {
@@ -71,7 +67,8 @@ class OpenAICompatible(rest.REST):
     for chunk in message.chunk():
       if isinstance(chunk, str):
         item = dict(type='text', text=chunk)
-      elif isinstance(chunk, lf_modalities.Image) and self.multimodal:
+      elif (isinstance(chunk, lf_modalities.Image)
+            and self.supports_input(chunk.mime_type)):
         item = dict(type='image_url', image_url=dict(url=chunk.embeddable_uri))
       else:
         raise ValueError(f'Unsupported modality: {chunk!r}.')
@@ -162,18 +159,6 @@ class OpenAICompatible(rest.REST):
             prompt_tokens=usage['prompt_tokens'],
             completion_tokens=usage['completion_tokens'],
             total_tokens=usage['total_tokens'],
-            estimated_cost=self.estimate_cost(
-                num_input_tokens=usage['prompt_tokens'],
-                num_output_tokens=usage['completion_tokens'],
-            )
+
         ),
     )
-
-  def estimate_cost(
-      self,
-      num_input_tokens: int,
-      num_output_tokens: int
-  ) -> float | None:
-    """Estimate the cost based on usage."""
-    del num_input_tokens, num_output_tokens
-    return None

langfun/core/llms/openai_compatible_test.py

@@ -207,37 +207,52 @@ class OpenAIComptibleTest(unittest.TestCase):
       )
 
   def test_call_chat_completion_vision(self):
+
+    class FakeImage(lf_modalities.Image):
+
+      @property
+      def mime_type(self) -> str:
+        return 'image/png'
+
     with mock.patch('requests.Session.post') as mock_request:
       mock_request.side_effect = mock_chat_completion_request_vision
       lm_1 = openai_compatible.OpenAICompatible(
           api_endpoint='https://test-server',
           model='test-model1',
-          multimodal=True
       )
       lm_2 = openai_compatible.OpenAICompatible(
           api_endpoint='https://test-server',
           model='test-model2',
-          multimodal=True
       )
       for lm in (lm_1, lm_2):
         self.assertEqual(
             lm(
                 lf.UserMessage(
                     'hello <<[[image]]>>',
-                    image=lf_modalities.Image.from_uri('https://fake/image')
+                    image=FakeImage.from_uri('https://fake/image')
                 ),
                 sampling_options=lf.LMSamplingOptions(n=2)
             ),
             'Sample 0 for message: https://fake/image',
         )
-    lm_3 = openai_compatible.OpenAICompatible(
+
+    class TextOnlyModel(openai_compatible.OpenAICompatible):
+
+      class ModelInfo(lf.ModelInfo):
+        input_modalities: list[str] = lf.ModelInfo.TEXT_INPUT_ONLY
+
+      @property
+      def model_info(self) -> lf.ModelInfo:
+        return TextOnlyModel.ModelInfo('text-only-model')
+
+    lm_3 = TextOnlyModel(
         api_endpoint='https://test-server', model='test-model3'
     )
     with self.assertRaisesRegex(ValueError, 'Unsupported modality'):
       lm_3(
           lf.UserMessage(
               'hello <<[[image]]>>',
-              image=lf_modalities.Image.from_uri('https://fake/image')
+              image=FakeImage.from_uri('https://fake/image')
           ),
       )
 

langfun/core/llms/openai_test.py

@@ -30,11 +30,12 @@ class OpenAITest(unittest.TestCase):
 
   def test_model_id(self):
     self.assertEqual(
-        openai.Gpt35(api_key='test_key').model_id, 'OpenAI(text-davinci-003)')
+        openai.Gpt35(api_key='test_key').model_id, 'text-davinci-003')
 
   def test_resource_id(self):
     self.assertEqual(
-        openai.Gpt35(api_key='test_key').resource_id, 'OpenAI(text-davinci-003)'
+        openai.Gpt35(api_key='test_key').resource_id,
+        'openai://text-davinci-003'
     )
 
   def test_headers(self):
@@ -47,7 +48,9 @@ class OpenAITest(unittest.TestCase):
     )
 
   def test_max_concurrency(self):
-    self.assertGreater(openai.Gpt35(api_key='test_key').max_concurrency, 0)
+    self.assertGreater(
+        openai.Gpt4o(api_key='test_key').max_concurrency, 0
+    )
 
   def test_request_args(self):
     self.assertEqual(
@@ -75,11 +78,18 @@ class OpenAITest(unittest.TestCase):
   def test_estimate_cost(self):
     self.assertEqual(
         openai.Gpt4(api_key='test_key').estimate_cost(
-            num_input_tokens=100, num_output_tokens=100
+            lf.LMSamplingUsage(
+                total_tokens=200,
+                prompt_tokens=100,
+                completion_tokens=100,
+            )
         ),
         0.009
     )
 
+  def test_lm_get(self):
+    self.assertIsInstance(lf.LanguageModel.get('gpt-4o'), openai.OpenAI)
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/rest.py

@@ -49,11 +49,6 @@ class REST(lf.LanguageModel):
       'The headers for the REST API.'
   ] = None
 
-  @property
-  def model_id(self) -> str:
-    """Returns a string to identify the model."""
-    return self.model or 'unknown'
-
   @functools.cached_property
   def _api_initialized(self) -> bool:
     """Returns whether the API is initialized."""
@@ -76,7 +71,6 @@ class REST(lf.LanguageModel):
 
   def _on_bound(self):
     super()._on_bound()
-    self.__dict__.pop('_session', None)
     self.__dict__.pop('_api_initialized', None)
 
   def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
@@ -121,3 +115,14 @@ class REST(lf.LanguageModel):
       return self.result(response.json())
     else:
       raise self._error(response.status_code, response.content)
+
+  @property
+  def max_concurrency(self) -> int | None:
+    """Returns the max concurrency for this model."""
+    rate_limits = self.model_info.rate_limits
+    if rate_limits is not None:
+      return self.estimate_max_concurrency(
+          max_requests_per_minute=rate_limits.max_requests_per_minute,
+          max_tokens_per_minute=rate_limits.max_tokens_per_minute
+      )
+    return None