langfun 0.1.2.dev202501090804__tar.gz → 0.1.2.dev202501100804__tar.gz

{langfun-0.1.2.dev202501090804 → langfun-0.1.2.dev202501100804}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: langfun
-Version: 0.1.2.dev202501090804
+Version: 0.1.2.dev202501100804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202501090804 → langfun-0.1.2.dev202501100804}/langfun/core/llms/__init__.py

@@ -57,6 +57,9 @@ from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_002
 from langfun.core.llms.vertexai import VertexAIGeminiFlash1_5_001
 from langfun.core.llms.vertexai import VertexAIGeminiPro1
 
+# Base for OpenAI-compatible models.
+from langfun.core.llms.openai_compatible import OpenAICompatible
+
 # OpenAI models.
 from langfun.core.llms.openai import OpenAI
 

langfun-0.1.2.dev202501100804/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-0.1.2.dev202501100804/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-0.1.2.dev202501090804 → langfun-0.1.2.dev202501100804}/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-0.1.2.dev202501100804/langfun/core/llms/groq_test.py

@@ -0,0 +1,64 @@
+# 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 os
+import unittest
+import langfun.core as lf
+from langfun.core.llms import groq
+
+
+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.headers
+
+    lm = groq.GroqMistral_8x7B(api_key='fake key')
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer fake key',
+        }
+    )
+
+    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__':
+  unittest.main()

langfun-0.1.2.dev202501100804/langfun/core/llms/llama_cpp.py

@@ -0,0 +1,47 @@
+# 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.
+"""Language models from llama.cpp."""
+
+from typing import Annotated
+from langfun.core.llms import openai_compatible
+import pyglove as pg
+
+
+@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.',
+  ] = ''
+
+  @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 ""})'
+

{langfun-0.1.2.dev202501090804 → langfun-0.1.2.dev202501100804}/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)")