langfun 0.1.2.dev202501080804__py3-none-any.whl → 0.1.2.dev202501240804__py3-none-any.whl

langfun/core/llms/deepseek.py

@@ -0,0 +1,117 @@
+# Copyright 2024 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Language models from DeepSeek."""
+
+import os
+from typing import Annotated, Any
+
+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-chat': pg.Dict(
+        in_service=True,
+        rpm=100,
+        tpm=1000000,
+        cost_per_1k_input_tokens=0.00014,
+        cost_per_1k_output_tokens=0.00028,
+    ),
+}
+
+
+# DeepSeek API uses an API format compatible with OpenAI.
+# Reference: https://api-docs.deepseek.com/
+@lf.use_init_args(['model'])
+class DeepSeek(openai_compatible.OpenAICompatible):
+  """DeepSeek model."""
+
+  model: pg.typing.Annotated[
+      pg.typing.Enum(
+          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
+      ),
+      'The name of the model to use.',
+  ]
+
+  api_endpoint: str = 'https://api.deepseek.com/chat/completions'
+
+  api_key: Annotated[
+      str | None,
+      (
+          'API key. If None, the key will be read from environment variable '
+          "'DEEPSEEK_API_KEY'."
+      ),
+  ] = None
+
+  @property
+  def headers(self) -> dict[str, Any]:
+    api_key = self.api_key or os.environ.get('DEEPSEEK_API_KEY', None)
+    if not api_key:
+      raise ValueError(
+          'Please specify `api_key` during `__init__` or set environment '
+          'variable `DEEPSEEK_API_KEY` with your DeepSeek API key.'
+      )
+    headers = super().headers
+    headers.update({
+        'Authorization': f'Bearer {api_key}',
+    })
+    return headers
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return f'DeepSeek({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
+
+  @classmethod
+  def dir(cls):
+    return [k for k, v in SUPPORTED_MODELS_AND_SETTINGS.items() if v.in_service]
+
+
+class DeepSeekChat(DeepSeek):
+  """DeepSeek Chat model.
+
+  Currently, it is powered by DeepSeek-V3 model, 64K input contenxt window and
+  8k max output tokens.
+  """
+
+  model = 'deepseek-chat'

langfun/core/llms/deepseek_test.py

@@ -0,0 +1,61 @@
+# Copyright 2023 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import unittest
+from langfun.core.llms import deepseek
+
+
+class DeepSeekTest(unittest.TestCase):
+  """Tests for DeepSeek language model."""
+
+  def test_dir(self):
+    self.assertIn('deepseek-chat', deepseek.DeepSeek.dir())
+
+  def test_key(self):
+    with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
+      _ = deepseek.DeepSeekChat().headers
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer test_key',
+        }
+    )
+
+  def test_model_id(self):
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').model_id,
+        'DeepSeek(deepseek-chat)',
+    )
+
+  def test_resource_id(self):
+    self.assertEqual(
+        deepseek.DeepSeekChat(api_key='test_key').resource_id,
+        'DeepSeek(deepseek-chat)',
+    )
+
+  def test_max_concurrency(self):
+    self.assertGreater(
+        deepseek.DeepSeekChat(api_key='test_key').max_concurrency, 0
+    )
+
+  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
+    )
+
+if __name__ == '__main__':
+  unittest.main()

langfun/core/llms/gemini.py

@@ -380,7 +380,7 @@ class Gemini(rest.REST):
     return (
         cost_per_1m_input_tokens * num_input_tokens
         + cost_per_1m_output_tokens * num_output_tokens
-    ) / 1000_1000
+    ) / 1000_000
 
   @property
   def model_id(self) -> str:

langfun/core/llms/groq.py

@@ -17,8 +17,7 @@ import os
 from typing import Annotated, Any
 
 import langfun.core as lf
-from langfun.core import modalities as lf_modalities
-from langfun.core.llms import rest
+from langfun.core.llms import openai_compatible
 import pyglove as pg
 
 
@@ -95,7 +94,7 @@ SUPPORTED_MODELS_AND_SETTINGS = {
 
 
 @lf.use_init_args(['model'])
-class Groq(rest.REST):
+class Groq(openai_compatible.OpenAICompatible):
   """Groq LLMs through REST APIs (OpenAI compatible).
 
   See https://platform.openai.com/docs/api-reference/chat
@@ -108,10 +107,6 @@ class Groq(rest.REST):
       'The name of the model to use.',
   ]
 
-  multimodal: Annotated[bool, 'Whether this model has multimodal support.'] = (
-      False
-  )
-
   api_key: Annotated[
       str | None,
       (
@@ -122,25 +117,19 @@ class Groq(rest.REST):
 
   api_endpoint: str = 'https://api.groq.com/openai/v1/chat/completions'
 
-  def _on_bound(self):
-    super()._on_bound()
-    self._api_key = None
-
-  def _initialize(self):
+  @property
+  def headers(self) -> dict[str, Any]:
     api_key = self.api_key or os.environ.get('GROQ_API_KEY', None)
     if not api_key:
       raise ValueError(
           'Please specify `api_key` during `__init__` or set environment '
           'variable `GROQ_API_KEY` with your Groq API key.'
       )
-    self._api_key = api_key
-
-  @property
-  def headers(self) -> dict[str, Any]:
-    return {
-        'Authorization': f'Bearer {self._api_key}',
-        'Content-Type': 'application/json',
-    }
+    headers = super().headers
+    headers.update({
+        'Authorization': f'Bearer {api_key}',
+    })
+    return headers
 
   @property
   def model_id(self) -> str:
@@ -170,90 +159,14 @@ class Groq(rest.REST):
         + cost_per_1k_output_tokens * num_output_tokens
     ) / 1000
 
-  def request(
-      self,
-      prompt: lf.Message,
-      sampling_options: lf.LMSamplingOptions
-  ) -> dict[str, Any]:
-    """Returns the JSON input for a message."""
-    request = dict()
-    request.update(self._request_args(sampling_options))
-    request.update(
-        dict(
-            messages=[
-                dict(role='user', content=self._content_from_message(prompt))
-            ]
-        )
-    )
-    return request
-
   def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
     """Returns a dict as request arguments."""
     # `logprobs` and `top_logprobs` flags are not supported on Groq yet.
-    args = dict(
-        model=self.model,
-        n=options.n,
-        stream=False,
-    )
-
-    if options.temperature is not None:
-      args['temperature'] = options.temperature
-    if options.max_tokens is not None:
-      args['max_tokens'] = options.max_tokens
-    if options.top_p is not None:
-      args['top_p'] = options.top_p
-    if options.stop:
-      args['stop'] = options.stop
+    args = super()._request_args(options)
+    args.pop('logprobs', None)
+    args.pop('top_logprobs', None)
     return args
 
-  def _content_from_message(self, prompt: lf.Message) -> list[dict[str, Any]]:
-    """Converts an message to Groq's content protocol (list of dicts)."""
-    # Refer: https://platform.openai.com/docs/api-reference/chat/create
-    content = []
-    for chunk in prompt.chunk():
-      if isinstance(chunk, str):
-        item = dict(type='text', text=chunk)
-      elif (
-          self.multimodal
-          and isinstance(chunk, lf_modalities.Image)
-          and chunk.uri
-      ):
-        # NOTE(daiyip): Groq only support image URL.
-        item = dict(type='image_url', image_url=chunk.uri)
-      else:
-        raise ValueError(f'Unsupported modality object: {chunk!r}.')
-      content.append(item)
-    return content
-
-  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
-    samples = [
-        lf.LMSample(self._message_from_choice(choice), score=0.0)
-        for choice in json['choices']
-    ]
-    usage = json['usage']
-    return lf.LMSamplingResult(
-        samples,
-        usage=lf.LMSamplingUsage(
-            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 _message_from_choice(self, choice: dict[str, Any]) -> lf.Message:
-    """Converts Groq's content protocol to message."""
-    # Refer: https://platform.openai.com/docs/api-reference/chat/create
-    content = choice['message']['content']
-    if isinstance(content, str):
-      return lf.AIMessage(content)
-    return lf.AIMessage.from_chunks(
-        [x['text'] for x in content if x['type'] == 'text']
-    )
-
 
 class GroqLlama3_2_3B(Groq):  # pylint: disable=invalid-name
   """Llama3.2-3B with 8K context window.

langfun/core/llms/groq_test.py

@@ -11,89 +11,10 @@
 # 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 Groq models."""
-
 import os
-from typing import Any
 import unittest
-from unittest import mock
-from langfun.core import modalities as lf_modalities
+import langfun.core as lf
 from langfun.core.llms import groq
-import pyglove as pg
-import requests
-
-
-def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
-  del url, kwargs
-
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str({
-      'choices': [{
-          'message': {
-              'content': [{
-                  'type': 'text',
-                  'text': (
-                      f'hello with temperature={json.get("temperature")}, '
-                      f'top_p={json.get("top_p")}, '
-                      f'max_tokens={json.get("max_tokens")}, '
-                      f'stop={json.get("stop")}.'
-                  ),
-              }],
-          }
-      }],
-      'usage': {
-          'prompt_tokens': 2,
-          'completion_tokens': 1,
-          'total_tokens': 3,
-      },
-  }).encode()
-  return response
-
-
-def mock_mm_requests_post(url: str, json: dict[str, Any], **kwargs):
-  del url, kwargs
-  v = json['messages'][0]['content'][0]
-  image = lf_modalities.Image.from_uri(v['image_url'])
-
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str({
-      'choices': [
-          {
-              'message': {
-                  'content': [{
-                      'type': 'text',
-                      'text': image.uri,
-                  }],
-              }
-          }
-      ],
-      'usage': {
-          'prompt_tokens': 2,
-          'completion_tokens': 1,
-          'total_tokens': 3,
-      },
-  }).encode()
-  return response
-
-
-def mock_requests_post_error(status_code, error_type, error_message):
-  def _mock_requests(url: str, json: dict[str, Any], **kwargs):
-    del url, json, kwargs
-    response = requests.Response()
-    response.status_code = status_code
-    response._content = pg.to_json_str(
-        {
-            'error': {
-                'type': error_type,
-                'message': error_message,
-            }
-        }
-    ).encode()
-    return response
-
-  return _mock_requests
 
 
 class AuthropicTest(unittest.TestCase):
@@ -101,69 +22,42 @@ class AuthropicTest(unittest.TestCase):
   def test_basics(self):
     self.assertEqual(groq.GroqMistral_8x7B().model_id, 'mixtral-8x7b-32768')
     self.assertEqual(groq.GroqMistral_8x7B().max_concurrency, 16)
+    self.assertEqual(groq.GroqMistral_8x7B().estimate_cost(100, 100), 4.8e-5)
+
+  def test_request_args(self):
+    args = groq.GroqMistral_8x7B()._request_args(
+        lf.LMSamplingOptions(
+            temperature=1.0, stop=['\n'], n=1, random_seed=123,
+            logprobs=True, top_logprobs=True
+        )
+    )
+    self.assertNotIn('logprobs', args)
+    self.assertNotIn('top_logprobs', args)
 
   def test_api_key(self):
     lm = groq.GroqMistral_8x7B()
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      lm('hi')
-
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_requests_post
-
-      lm = groq.GroqMistral_8x7B(api_key='fake key')
-      self.assertRegex(lm('hi').text, 'hello.*')
-
-      os.environ['GROQ_API_KEY'] = 'abc'
-      lm = groq.GroqMistral_8x7B()
-      self.assertRegex(lm('hi').text, 'hello.*')
-      del os.environ['GROQ_API_KEY']
-
-  def test_call(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_requests_post
-      lm = groq.GroqLlama3_70B(api_key='fake_key')
-      response = lm(
-          'hello',
-          temperature=0.0,
-          max_tokens=1024,
-          top_k=0.1,
-          top_p=0.2,
-          stop=['\n'],
-      )
-      self.assertEqual(
-          response.text,
-          (
-              'hello with temperature=0.0, top_p=0.2, '
-              "max_tokens=1024, stop=['\\n']."
-          ),
-      )
-      self.assertIsNotNone(response.usage)
-      self.assertIsNotNone(response.usage.prompt_tokens, 2)
-      self.assertIsNotNone(response.usage.completion_tokens, 1)
-      self.assertIsNotNone(response.usage.total_tokens, 3)
+      _ = lm.headers
 
-  def test_mm_call(self):
-    with mock.patch('requests.Session.post') as mock_mm_request:
-      mock_mm_request.side_effect = mock_mm_requests_post
-      lm = groq.GroqLlama3_70B(multimodal=True, api_key='fake_key')
-      response = lm(lf_modalities.Image.from_uri('https://fake/image.jpg'))
-      self.assertEqual(response.text, 'https://fake/image.jpg')
+    lm = groq.GroqMistral_8x7B(api_key='fake key')
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer fake key',
+        }
+    )
 
-  def test_call_errors(self):
-    for status_code, error_type, error_message in [
-        (429, 'rate_limit', 'Rate limit exceeded.'),
-        (503, 'service_unavailable', 'Service unavailable.'),
-        (500, 'bad_request', 'Bad request.'),
-    ]:
-      with mock.patch('requests.Session.post') as mock_mm_request:
-        mock_mm_request.side_effect = mock_requests_post_error(
-            status_code, error_type, error_message
-        )
-        lm = groq.GroqLlama3_70B(api_key='fake_key')
-        with self.assertRaisesRegex(
-            Exception, f'{status_code}:.*{error_type}'
-        ):
-          lm('hello', max_attempts=1)
+    os.environ['GROQ_API_KEY'] = 'abc'
+    lm = groq.GroqMistral_8x7B()
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer abc',
+        }
+    )
+    del os.environ['GROQ_API_KEY']
 
 
 if __name__ == '__main__':

langfun/core/llms/llama_cpp.py

@@ -13,72 +13,35 @@
 # limitations under the License.
 """Language models from llama.cpp."""
 
-from typing import Any
-
-import langfun.core as lf
-from langfun.core.llms import rest
+from typing import Annotated
+from langfun.core.llms import openai_compatible
 import pyglove as pg
 
 
-class LlamaCppRemote(rest.REST):
+@pg.use_init_args(['url', 'model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class LlamaCppRemote(openai_compatible.OpenAICompatible):
   """The remote LLaMA C++ model.
 
   The Remote LLaMA C++ models can be launched via
   https://github.com/ggerganov/llama.cpp/tree/master/examples/server
   """
+  url: Annotated[
+      str,
+      'The URL of the LLaMA C++ server.',
+  ]
+
+  model: Annotated[
+      str,
+      'The name of the model to use.',
+  ] = ''
 
-  @pg.explicit_method_override
-  def __init__(self, url: str, model: str | None = None, **kwargs):
-    super().__init__(api_endpoint=f'{url}/completion', model=model, **kwargs)
+  @property
+  def api_endpoint(self) -> str:
+    return self.url + '/completion'
 
   @property
   def model_id(self) -> str:
     """Returns a string to identify the model."""
     return f'LLaMAC++({self.model or ""})'
 
-  def request(
-      self, prompt: lf.Message, sampling_options: lf.LMSamplingOptions
-  ) -> dict[str, Any]:
-    """Returns the JSON input for a message."""
-    request = dict()
-    request.update(self._request_args(sampling_options))
-    # NOTE(daiyip): multi-modal is current not supported.
-    request['prompt'] = prompt.text
-    return request
-
-  def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
-    """Returns a dict as request arguments."""
-    args = dict(
-        n_predict=options.max_tokens or 1024,
-        top_k=options.top_k or 50,
-        top_p=options.top_p or 0.95,
-    )
-    if options.temperature is not None:
-      args['temperature'] = options.temperature
-    return args
-
-  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
-    return lf.LMSamplingResult(
-        [lf.LMSample(item['content'], score=0.0) for item in json['items']]
-    )
-
-  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    request = self.request(prompt, self.sampling_options)
-
-    def _sample_one_example(request):
-      response = self._session.post(
-          self.api_endpoint,
-          json=request,
-          timeout=self.timeout,
-      )
-      if response.status_code == 200:
-        return response.json()
-      else:
-        error_cls = self._error_cls_from_status(response.status_code)
-        raise error_cls(f'{response.status_code}: {response.content}')
-
-    items = self._parallel_execute_with_currency_control(
-        _sample_one_example,
-        [request] * (self.sampling_options.n or 1),
-    )
-    return self.result(dict(items=items))

langfun/core/llms/llama_cpp_test.py

@@ -11,48 +11,16 @@
 # 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 llama cpp models."""
-
-import typing
 import unittest
-from unittest import mock
-
 from langfun.core.llms import llama_cpp
 
 
-def mock_requests_post(url: str, json: typing.Dict[str, typing.Any], **kwargs):
-  del kwargs
-
-  class TEMP:
-    @property
-    def status_code(self):
-      return 200
-
-    def json(self):
-      return {"content": json["prompt"] + "\n" + url}
-
-  return TEMP()
-
-
 class LlamaCppRemoteTest(unittest.TestCase):
   """Tests for the LlamaCppRemote model."""
 
-  def test_call_completion(self):
-    with mock.patch("requests.Session.post") as mock_request:
-      mock_request.side_effect = mock_requests_post
-      lm = llama_cpp.LlamaCppRemote("http://127.0.0.1:8080")
-      [result] = lm.sample(["hello"], n=2)
-      self.assertEqual(
-          len(result.samples),
-          2
-      )
-      self.assertEqual(
-          str(result.samples[0].response),
-          "hello\nhttp://127.0.0.1:8080/completion",
-      )
-
-  def test_model_id(self):
+  def test_basics(self):
     lm = llama_cpp.LlamaCppRemote("http://127.0.0.1:8080")
+    self.assertEqual(lm.api_endpoint, "http://127.0.0.1:8080/completion")
     self.assertEqual(lm.model_id, "LLaMAC++()")
     lm = llama_cpp.LlamaCppRemote("xxx", model="x")
     self.assertEqual(lm.model_id, "LLaMAC++(x)")