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

langfun/core/llms/anthropic_test.py

@@ -19,6 +19,7 @@ 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 anthropic
 import pyglove as pg
@@ -119,6 +120,17 @@ class AnthropicTest(unittest.TestCase):
     )
     self.assertGreater(anthropic.Claude3Haiku().max_concurrency, 0)
 
+  def test_model_alias(self):
+    # Alias will be normalized to the official version.
+    self.assertEqual(
+        anthropic.Anthropic('claude-3-5-sonnet-20241022').model_id,
+        'claude-3-5-sonnet-20241022'
+    )
+    self.assertEqual(
+        anthropic.Anthropic('claude-3-5-sonnet-v2@20241022').model_id,
+        'claude-3-5-sonnet-20241022'
+    )
+
   def test_api_key(self):
     lm = anthropic.Claude3Haiku()
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
@@ -151,6 +163,7 @@ class AnthropicTest(unittest.TestCase):
       self.assertIsNotNone(response.usage.prompt_tokens, 2)
       self.assertIsNotNone(response.usage.completion_tokens, 1)
       self.assertIsNotNone(response.usage.total_tokens, 3)
+      self.assertGreater(response.usage.estimated_cost, 0)
 
   def test_mm_call(self):
     with mock.patch('requests.Session.post') as mock_mm_request:
@@ -162,7 +175,7 @@ class AnthropicTest(unittest.TestCase):
   def test_pdf_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')
+      lm = anthropic.Claude35Sonnet(api_key='fake_key')
       response = lm(lf_modalities.PDF.from_bytes(pdf_content), lm=lm)
       self.assertEqual(response.text, 'document: application/pdf')
 
@@ -182,6 +195,12 @@ class AnthropicTest(unittest.TestCase):
         ):
           lm('hello', max_attempts=1)
 
+  def test_lm_get(self):
+    self.assertIsInstance(
+        lf.LanguageModel.get('claude-3-5-sonnet-latest'),
+        anthropic.Anthropic,
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/deepseek.py

@@ -13,33 +13,81 @@
 # limitations under the License.
 """Language models from DeepSeek."""
 
+import datetime
+import functools
 import os
-from typing import Annotated, Any
+from typing import Annotated, Any, Final
 
 import langfun.core as lf
 from langfun.core.llms import openai_compatible
 import pyglove as pg
 
-SUPPORTED_MODELS_AND_SETTINGS = {
-    # pylint: disable=g-line-too-long
-    # TODO(yifenglu): The RPM and TPM are arbitrary numbers. Update them once DeepSeek provides concrete guidelines.
-    # DeepSeek doesn't control the rate limit at the moment: https://api-docs.deepseek.com/quick_start/rate_limit
-    # The cost is based on: https://api-docs.deepseek.com/quick_start/pricing
-    'deepseek-reasoner': pg.Dict(
+
+class DeepSeekModelInfo(lf.ModelInfo):
+  """DeepSeek model info."""
+
+  LINKS = dict(
+      models='https://api-docs.deepseek.com/quick_start/pricing',
+      pricing='https://api-docs.deepseek.com/quick_start/pricing',
+      rate_limits='https://api-docs.deepseek.com/quick_start/rate_limit',
+      error_codes='https://api-docs.deepseek.com/quick_start/error_codes',
+  )
+
+  provider: Final[str] = 'DeepSeek'  # pylint: disable=invalid-name
+
+  api_model_name: Annotated[
+      str,
+      'The model name used in the DeepSeek API.'
+  ]
+
+
+SUPPORTED_MODELS = [
+    DeepSeekModelInfo(
+        model_id='deepseek-r1',
         in_service=True,
-        rpm=100,
-        tpm=1000000,
-        cost_per_1k_input_tokens=0.00055,
-        cost_per_1k_output_tokens=0.00219,
+        model_type='thinking',
+        api_model_name='deepseek-reasoner',
+        description='DeepSeek Reasoner model (01/20/2025).',
+        url='https://api-docs.deepseek.com/news/news250120',
+        release_date=datetime.datetime(2025, 1, 20),
+        input_modalities=lf.ModelInfo.TEXT_INPUT_ONLY,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=64_000,
+            max_output_tokens=8_000,
+            max_cot_tokens=32_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.14,
+            cost_per_1m_input_tokens=0.55,
+            cost_per_1m_output_tokens=2.19,
+        ),
+        # No rate limits is enforced by DeepSeek for now.
+        rate_limits=None
     ),
-    'deepseek-chat': pg.Dict(
+    DeepSeekModelInfo(
+        model_id='deepseek-v3',
         in_service=True,
-        rpm=100,
-        tpm=1000000,
-        cost_per_1k_input_tokens=0.00014,
-        cost_per_1k_output_tokens=0.00028,
+        model_type='instruction-tuned',
+        api_model_name='deepseek-chat',
+        description='DeepSeek V3 model (12/26/2024).',
+        url='https://api-docs.deepseek.com/news/news1226',
+        release_date=datetime.datetime(2024, 12, 26),
+        input_modalities=lf.ModelInfo.TEXT_INPUT_ONLY,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=64_000,
+            max_output_tokens=8_000,
+        ),
+        pricing=lf.ModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.07,
+            cost_per_1m_input_tokens=0.27,
+            cost_per_1m_output_tokens=1.1,
+        ),
+        # No rate limits is enforced by DeepSeek for now.
+        rate_limits=None
     ),
-}
+]
+
+_SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 
 # DeepSeek API uses an API format compatible with OpenAI.
@@ -50,7 +98,7 @@ class DeepSeek(openai_compatible.OpenAICompatible):
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
-          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
+          pg.MISSING_VALUE, [m.model_id for m in SUPPORTED_MODELS]
       ),
       'The name of the model to use.',
   ]
@@ -79,56 +127,47 @@ class DeepSeek(openai_compatible.OpenAICompatible):
     })
     return headers
 
-  @property
-  def model_id(self) -> str:
-    """Returns a string to identify the model."""
-    return f'DeepSeek({self.model})'
+  @functools.cached_property
+  def model_info(self) -> DeepSeekModelInfo:
+    return _SUPPORTED_MODELS_BY_ID[self.model]
 
-  @property
-  def max_concurrency(self) -> int:
-    rpm = SUPPORTED_MODELS_AND_SETTINGS[self.model].get('rpm', 0)
-    tpm = SUPPORTED_MODELS_AND_SETTINGS[self.model].get('tpm', 0)
-    return self.rate_to_max_concurrency(
-        requests_per_min=rpm, tokens_per_min=tpm
-    )
-
-  def estimate_cost(
-      self, num_input_tokens: int, num_output_tokens: int
-  ) -> float | None:
-    """Estimate the cost based on usage."""
-    cost_per_1k_input_tokens = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
-        'cost_per_1k_input_tokens', None
-    )
-    cost_per_1k_output_tokens = SUPPORTED_MODELS_AND_SETTINGS[self.model].get(
-        'cost_per_1k_output_tokens', None
-    )
-    if cost_per_1k_output_tokens is None or cost_per_1k_input_tokens is None:
-      return None
-    return (
-        cost_per_1k_input_tokens * num_input_tokens
-        + cost_per_1k_output_tokens * num_output_tokens
-    ) / 1000
+  def _request_args(
+      self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+    """Returns a dict as request arguments."""
+    # NOTE(daiyip): Replace model name with the API model name instead of the
+    # model ID.
+    args = super()._request_args(options)
+    args['model'] = self.model_info.api_model_name
+    return args
 
   @classmethod
   def dir(cls):
-    return [k for k, v in SUPPORTED_MODELS_AND_SETTINGS.items() if v.in_service]
+    return [m.model_id for m in SUPPORTED_MODELS if m.in_service]
 
 
-class DeepSeekReasoner(DeepSeek):
+class DeepSeekR1(DeepSeek):
   """DeepSeek Reasoner model.
 
   Currently it is powered by DeepSeek-R1 model, 64k input context, 8k max
   output, 32k max CoT output.
   """
 
-  model = 'deepseek-reasoner'
+  model = 'deepseek-r1'
 
 
-class DeepSeekChat(DeepSeek):
+class DeepSeekV3(DeepSeek):
   """DeepSeek Chat model.
 
   Currently, it is powered by DeepSeek-V3 model, 64K input contenxt window and
   8k max output tokens.
   """
 
-  model = 'deepseek-chat'
+  model = 'deepseek-v3'
+
+
+def _register_deepseek_models():
+  """Registers DeepSeek models."""
+  for m in SUPPORTED_MODELS:
+    lf.LanguageModel.register(m.model_id, DeepSeek)
+
+_register_deepseek_models()

langfun/core/llms/deepseek_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
+import langfun.core as lf
 from langfun.core.llms import deepseek
 
 
@@ -19,13 +20,13 @@ class DeepSeekTest(unittest.TestCase):
   """Tests for DeepSeek language model."""
 
   def test_dir(self):
-    self.assertIn('deepseek-chat', deepseek.DeepSeek.dir())
+    self.assertIn('deepseek-v3', deepseek.DeepSeek.dir())
 
   def test_key(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      _ = deepseek.DeepSeekChat().headers
+      _ = deepseek.DeepSeekV3().headers
     self.assertEqual(
-        deepseek.DeepSeekChat(api_key='test_key').headers,
+        deepseek.DeepSeekV3(api_key='test_key').headers,
         {
             'Content-Type': 'application/json',
             'Authorization': 'Bearer test_key',
@@ -34,27 +35,25 @@ class DeepSeekTest(unittest.TestCase):
 
   def test_model_id(self):
     self.assertEqual(
-        deepseek.DeepSeekChat(api_key='test_key').model_id,
-        'DeepSeek(deepseek-chat)',
+        deepseek.DeepSeekV3(api_key='test_key').model_id,
+        'deepseek-v3',
     )
 
   def test_resource_id(self):
     self.assertEqual(
-        deepseek.DeepSeekChat(api_key='test_key').resource_id,
-        'DeepSeek(deepseek-chat)',
+        deepseek.DeepSeekV3(api_key='test_key').resource_id,
+        'deepseek://deepseek-v3',
     )
 
-  def test_max_concurrency(self):
-    self.assertGreater(
-        deepseek.DeepSeekChat(api_key='test_key').max_concurrency, 0
+  def test_request(self):
+    request = deepseek.DeepSeekV3(api_key='test_key').request(
+        lf.UserMessage('hi'), lf.LMSamplingOptions(temperature=0.0),
     )
+    self.assertEqual(request['model'], 'deepseek-chat')
 
-  def test_estimate_cost(self):
-    self.assertEqual(
-        deepseek.DeepSeekChat(api_key='test_key').estimate_cost(
-            num_input_tokens=100, num_output_tokens=100
-        ),
-        4.2e-5
+  def test_lm_get(self):
+    self.assertIsInstance(
+        lf.LanguageModel.get('deepseek-v3'), deepseek.DeepSeek
     )
 
 if __name__ == '__main__':

langfun/core/llms/fake.py

@@ -14,6 +14,7 @@
 """Fake LMs for testing."""
 
 import abc
+import functools
 from typing import Annotated
 import langfun.core as lf
 
@@ -44,6 +45,11 @@ class Fake(lf.LanguageModel):
       )
     return results
 
+  @functools.cached_property
+  def model_info(self) -> lf.ModelInfo:
+    """Returns the specification of the model."""
+    return lf.ModelInfo(model_id=self.__class__.__name__)
+
   @abc.abstractmethod
   def _response_from(self, prompt: lf.Message) -> lf.Message:
     """Returns the response for the given prompt."""