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

langfun/core/__init__.py

@@ -77,11 +77,6 @@ from langfun.core.concurrent import concurrent_map
 from langfun.core.concurrent import with_context_access
 from langfun.core.concurrent import with_retry
 
-# Utility libraries for text formatting.
-from langfun.core.text_formatting import colored
-from langfun.core.text_formatting import colored_print as print  # pylint: disable=redefined-builtin
-from langfun.core.text_formatting import colored_template
-
 # Interface for natural language formattable.
 from langfun.core.natural_language import NaturalLanguageFormattable
 

langfun/core/coding/python/correction.py

@@ -40,6 +40,7 @@ def run_with_correction(
     sandbox: bool | None = None,
     timeout: int | None = 5,
     returns_code: bool = False,
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
 ) -> Any | tuple[Any, str]:
   """Correct code with a language model via self-play.
@@ -62,6 +63,7 @@ def run_with_correction(
       timeout. Applicable only when sandbox is set to True.
     returns_code: If True, the return value is a tuple of (result, final code).
       Otherwise the return value is the result only.
+    returns_stdout: If True, the stdout (a str) will be returned.
     outputs_intermediate: If True, intermediate output will be outputted as a
       dict, with the last line's value accessible by key '__result__'. Otherwise
       the value of the last line will be returned.
@@ -87,6 +89,7 @@ def run_with_correction(
             global_vars=global_vars,
             sandbox=sandbox,
             timeout=timeout,
+            returns_stdout=returns_stdout,
             outputs_intermediate=outputs_intermediate,
         )
     )
@@ -189,9 +192,7 @@ def correct(
 def _error_feedback_str(error: Exception) -> str:
   """Returns the error str for feedback."""
   if isinstance(error, errors.CodeError):
-    return lf.text_formatting.decolored(
-        error.format(include_complete_code=False)
-    )
+    return pg.decolor(error.format(include_complete_code=False))
   else:
     return f"Encountered {error.__class__.__name__}: {error}"
 

langfun/core/coding/python/errors.py

@@ -17,7 +17,8 @@ import io
 import sys
 import textwrap
 import traceback
-import langfun.core as lf
+
+import pyglove as pg
 
 
 class CodeError(RuntimeError):
@@ -62,13 +63,13 @@ class CodeError(RuntimeError):
     if 'line' not in error_message and self.lineno is not None:
       error_message += f' (<unknown>, line {self.lineno})'
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {error_message}', 'magenta'))
 
     if self.lineno is not None:
       r.write('\n\n')
       r.write(textwrap.indent(
-          lf.colored(
+          pg.colored(
               self.code_lines(self.lineno - 1, self.end_lineno), 'magenta'),
           ' ' * 2
       ))
@@ -76,14 +77,14 @@ class CodeError(RuntimeError):
 
     if include_complete_code:
       r.write('\n')
-      r.write(lf.colored('[Generated Code]', 'green', styles=['bold']))
+      r.write(pg.colored('[Generated Code]', 'green', styles=['bold']))
       r.write('\n\n')
-      r.write(lf.colored('  ```python\n', 'green'))
+      r.write(pg.colored('  ```python\n', 'green'))
       r.write(textwrap.indent(
-          lf.colored(self.code, 'green'),
+          pg.colored(self.code, 'green'),
           ' ' * 2
       ))
-      r.write(lf.colored('\n  ```\n', 'green'))
+      r.write(pg.colored('\n  ```\n', 'green'))
     return r.getvalue()
 
 
@@ -98,10 +99,10 @@ class SerializationError(RuntimeError):
     r = io.StringIO()
     cause_message = str(self.cause).rstrip()
     if self.message:
-      r.write(lf.colored(self.message, 'magenta'))
+      r.write(pg.colored(self.message, 'magenta'))
       r.write('\n\n')
     r.write(
-        lf.colored(
+        pg.colored(
             f'{self.cause.__class__.__name__}: {cause_message}', 'magenta'
         )
     )

langfun/core/coding/python/execution.py

@@ -57,6 +57,7 @@ def evaluate(
     *,
     global_vars: dict[str, Any] | None = None,
     permission: permissions.CodePermission | None = None,
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
 ) -> Any | dict[str, Any]:
   """Executes Python code.
@@ -71,14 +72,17 @@ def evaluate(
     global_vars: An optional dict as the globals that could be referenced by the
       code.
     permission: Permission for the Python code to run.
-    outputs_intermediate: If True, intermediate output will be outputted as a
-      dict, with the last line's value accessible by key '__result__'. Otherwise
-      the value of the last line will be returned.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
 
   Returns:
-    The value of the last line of the code. Or a dict of variable name to
-    their values if `outputs_intermediate` is set to True, with the final result
-    accessible by key '__result__'.
+    The value of the last line of the code block. Or a dict of variable
+    names of all locals to their evaluated values as the output of the code to
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
   """
   # Set up the permission and context.
   permission = permission or permissions.get_permission()
@@ -136,6 +140,8 @@ def evaluate(
         raise errors.CodeError(code, e) from e
       global_vars[RESULT_KEY] = list(global_vars.values())[-1]
 
+  if returns_stdout:
+    return stdout.getvalue()
   if outputs_intermediate:
     outputs = {}
     for k, v in global_vars.items():
@@ -258,6 +264,7 @@ def run(
     *,
     global_vars: dict[str, Any] | None = None,
     permission: permissions.CodePermission | None = None,
+    returns_stdout: bool = False,
     outputs_intermediate: bool = False,
     sandbox: bool | None = None,
     timeout: float | None = None,
@@ -273,9 +280,11 @@ def run(
     code: Python code to run.
     global_vars: An optional dict of
     permission: Permission for the Python code to run.
-    outputs_intermediate: If True, all variables created as locals will be
-      returned, with the final result accessible by key '__result__'. Otherwise
-      only the final result will be returned.
+    returns_stdout: If True, the stdout (a str) will be returned.
+    outputs_intermediate: Applicable when returns_stdout is False. If True,
+      intermediate output will be outputted as a dict, with the last line's
+      value accessible by key '__result__' and the std output accessible by
+      key '__stdout__'. Otherwise the value of the last line will be returned.
     sandbox: If True, run code in sandbox; If False, run code in current
       process. If None, run in sandbox first, if the output could not be
       serialized and pass to current process, run the code again in current
@@ -285,7 +294,8 @@ def run(
   Returns:
     The value of the last line of the code block. Or a dict of variable
     names of all locals to their evaluated values as the output of the code to
-    run. The value for the last line can be accessed by key '__result__'.
+    run. The value for the last line can be accessed by key '__result__'. Or the
+    stdout as a str.
 
   Raises:
     TimeoutError: If the execution time exceeds the timeout.
@@ -293,5 +303,6 @@ def run(
   """
   return call(
       evaluate, code=code, global_vars=global_vars, permission=permission,
-      outputs_intermediate=outputs_intermediate,
-      sandbox=sandbox, timeout=timeout)
+      returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate,
+      sandbox=sandbox, timeout=timeout
+  )

langfun/core/coding/python/execution_test.py

@@ -63,6 +63,15 @@ class EvaluateTest(unittest.TestCase):
         ),
         3,
     )
+    with self.assertRaisesRegex(errors.CodeError, 'ValueError'):
+      execution.evaluate(
+          """
+          def foo():
+            raise ValueError("intentional error")
+          foo()
+          """,
+          permission=permissions.CodePermission.ALL
+      )
 
   def test_class_def(self):
     ret = execution.evaluate(
@@ -102,16 +111,20 @@ class EvaluateTest(unittest.TestCase):
     self.assertIs(ret['__result__'], ret['bar'])
 
   def test_function_def_and_call(self):
-    ret = execution.evaluate(
+    code = (
         """
         def foo(x, y):
           return x + y
 
         def bar(z):
+          print(f'z is {z}')
           return z + foo(z, z)
 
         bar(1)
-        """,
+        """
+    )
+    ret = execution.evaluate(
+        code,
         permission=permissions.CodePermission.ALL,
         outputs_intermediate=True,
     )
@@ -119,6 +132,12 @@ class EvaluateTest(unittest.TestCase):
         list(ret.keys()), ['foo', 'bar', '__result__', '__stdout__']
     )
     self.assertEqual(ret['__result__'], 3)
+    ret = execution.evaluate(
+        code,
+        permission=permissions.CodePermission.ALL,
+        returns_stdout=True,
+    )
+    self.assertEqual(ret, 'z is 1\n')
 
   def test_complex(self):
     ret = execution.evaluate(

langfun/core/coding/python/generation.py

@@ -88,6 +88,8 @@ class PythonCode(pg.Object):
       sandbox: bool | None = None,
       timeout: int | None = 5,
       global_vars: dict[str, Any] | None = None,
+      returns_stdout: bool = False,
+      outputs_intermediate: bool = False,
       autofix: int = 3,
       autofix_lm: lf.LanguageModel | None = None,
   ) -> Any:
@@ -101,13 +103,22 @@ class PythonCode(pg.Object):
       timeout: Timeout in seconds. If None, there is no timeout. Applicable when
         sandbox is set to True.
       global_vars: Global variables that could be accessed from the source code.
+      returns_stdout: If True, the stdout (a str) will be returned.
+      outputs_intermediate: Applicable when returns_stdout is False. If True,
+        intermediate output will be outputted as a dict, with the last line's
+        value accessible by key '__result__' and the std output accessible by 
+        key '__stdout__'. Otherwise the value of the last line will be returned.
       autofix: Number of attempts to auto fix the generated code. If 0, autofix
         is disabled.
       autofix_lm: Language model to be used. If not specified, it will try to
         use the `lm` under `lf.context`.
 
     Returns:
-      The value of the last expression in the source code.
+      The value of the last expression in the source code. Or a dict of local
+      variable names defined in the source code to their values if
+      `outputs_intermediate` is set to True. The value for the last line can be
+      accessed by key '__result__'. Or the stdout as a str if `returns_stdout`
+      is set to True.
 
     Raises:
       TimeoutError: If `sandbox` is True and timeout has reached.
@@ -121,6 +132,8 @@ class PythonCode(pg.Object):
         max_attempts=autofix,
         lm=autofix_lm,
         returns_code=True,
+        returns_stdout=returns_stdout,
+        outputs_intermediate=outputs_intermediate,
     )
     self.rebind(source=updated_code)
     return result
@@ -158,18 +171,14 @@ class PythonCode(pg.Object):
       TimeoutError: If `sandbox` is True and timeout has reached.
       Exception: Any errors that the source code has raised.
     """
-    result, updated_code = correction.run_with_correction(
-        self.source,
-        global_vars=global_vars,
+    return self(
         sandbox=sandbox,
         timeout=timeout,
+        global_vars=global_vars,
+        autofix=autofix,
+        autofix_lm=autofix_lm,
         outputs_intermediate=True,
-        max_attempts=autofix,
-        lm=autofix_lm,
-        returns_code=True,
     )
-    self.rebind(source=updated_code)
-    return result
 
 
 class PythonFunction(pg.Object):

langfun/core/concurrent.py

@@ -25,7 +25,6 @@ import time
 from typing import Any, Callable, Iterable, Iterator, Literal, Sequence, Tuple, Type, Union
 
 from langfun.core import component
-from langfun.core import text_formatting
 import pyglove as pg
 
 
@@ -844,10 +843,10 @@ class _ConsoleProgressControl(_ProgressControl):
   def refresh(self):
     s = io.StringIO()
     if self.label is not None:
-      s.write(text_formatting.colored(self.label, 'red', styles=['bold']))
+      s.write(pg.colored(self.label, 'red', styles=['bold']))
       s.write(': ')
     s.write(
-        text_formatting.colored(
+        pg.colored(
             '%d%% (%d/%d)' %
             (
                 self._progress * 100 // self.total,

langfun/core/console.py

@@ -15,7 +15,7 @@
 
 import sys
 from typing import Any
-from langfun.core.text_formatting import colored
+import pyglove as pg
 
 
 def write(
@@ -42,10 +42,15 @@ def write(
   """
   # Print title if present.
   if title is not None:
-    print(colored(title, styles=['bold']))
+    print(pg.colored(title, styles=['bold']))
 
   # Print body.
-  print(colored(str(value), color=color, background=background, styles=styles))
+  print(dir(pg.utils))
+  print(
+      pg.colored(
+          str(value), color=color, background=background, styles=styles
+      )
+  )
 
 
 try:

langfun/core/eval/base.py

@@ -1298,7 +1298,7 @@ class Evaluation(Evaluable):
             id=self.id,
             dir=self.dir,
             model=self.lm.model_id,
-            prompt_template=lf.text_formatting.decolored(str(self.prompt)),
+            prompt_template=pg.decolor(str(self.prompt)),
             method=self.method,
             schema_fn=str(self.schema_fn),
         ),
@@ -2110,8 +2110,7 @@ class Summary(pg.Object):
 
 def _format_error(error: Exception):
   """Formats an error into a string."""
-  return (f'({error.__class__.__name__}) '
-          + lf.text_formatting.decolored(str(error)))
+  return (f'({error.__class__.__name__}) ' + pg.decolor(str(error)))
 
 
 def _error_key(error: Exception) -> str:

langfun/core/eval/v2/reporting.py

@@ -51,6 +51,7 @@ class HtmlReporter(experiment_lib.Plugin):
     self._update_thread = None
     self._stop_update = False
     self._stop_update_experiment_ids = set()
+    self._summary_lock = None
     self._experiment_index_lock = None
 
   def on_run_start(
@@ -62,6 +63,7 @@ class HtmlReporter(experiment_lib.Plugin):
     self._last_experiment_report_time = {leaf.id: 0 for leaf in root.leaf_nodes}
     self._stop_update = False
     self._stop_update_experiment_ids = set()
+    self._summary_lock = threading.Lock()
     self._experiment_index_lock = {
         leaf.id: threading.Lock() for leaf in root.leaf_nodes
     }
@@ -137,21 +139,23 @@ class HtmlReporter(experiment_lib.Plugin):
     """Maybe update the summary of current run."""
     run = runner.current_run
     def _summary():
-      run.experiment.to_html(
+      html = run.experiment.to_html(
           collapse_level=None,
           extra_flags=dict(
               current_run=run, interactive=False, card_view=True,
           )
-      ).save(
-          run.output_path_for(run.experiment, _SUMMARY_FILE)
       )
+      with self._summary_lock:
+        html.save(
+            run.output_path_for(run.experiment, _SUMMARY_FILE)
+        )
 
     if force or (time.time() - self._last_summary_time > self.summary_interval):
+      self._last_summary_time = time.time()
       if background:
         runner.background_run(_summary)
       else:
         _summary()
-      self._last_summary_time = time.time()
 
   def _maybe_update_experiment_html(
       self,

langfun/core/language_model.py

@@ -434,7 +434,10 @@ class LanguageModel(component.Component):
   def __init__(self, *args, **kwargs) -> None:
     """Overrides __init__ to pass through **kwargs to sampling options."""
 
-    sampling_options = kwargs.pop('sampling_options', LMSamplingOptions())
+    sampling_options = kwargs.pop(
+        'sampling_options',
+        pg.clone(self.__schema__.fields['sampling_options'].default_value)
+    )
     sampling_options_delta = {}
 
     for k, v in kwargs.items():
@@ -650,7 +653,7 @@ class LanguageModel(component.Component):
     """Outputs debugging information about the model."""
     title_suffix = ''
     if usage.total_tokens != 0:
-      title_suffix = console.colored(
+      title_suffix = pg.colored(
           f' (total {usage.total_tokens} tokens)', 'red'
       )
 
@@ -669,7 +672,7 @@ class LanguageModel(component.Component):
     """Outputs debugging information about the prompt."""
     title_suffix = ''
     if usage.prompt_tokens != 0:
-      title_suffix = console.colored(f' ({usage.prompt_tokens} tokens)', 'red')
+      title_suffix = pg.colored(f' ({usage.prompt_tokens} tokens)', 'red')
 
     console.write(
         # We use metadata 'formatted_text' for scenarios where the prompt text
@@ -700,7 +703,7 @@ class LanguageModel(component.Component):
     if usage.completion_tokens != 0:
       title_suffix += f'{usage.completion_tokens} tokens '
     title_suffix += f'in {elapse:.2f} seconds)'
-    title_suffix = console.colored(title_suffix, 'red')
+    title_suffix = pg.colored(title_suffix, 'red')
 
     console.write(
         str(response) + '\n',

langfun/core/language_model_test.py

@@ -117,6 +117,21 @@ class LanguageModelTest(unittest.TestCase):
     self.assertEqual(lm.sampling_options.top_k, 2)
     self.assertEqual(lm.max_attempts, 2)
 
+  def test_subclassing(self):
+
+    class ChildModel(lm_lib.LanguageModel):
+
+      sampling_options = lm_lib.LMSamplingOptions(
+          temperature=0.5, top_k=20
+      )
+
+      def _sample(self, *args, **kwargs):
+        pass
+
+    lm = ChildModel(top_k=10)
+    self.assertEqual(lm.sampling_options.temperature, 0.5)
+    self.assertEqual(lm.sampling_options.top_k, 10)
+
   def test_sample(self):
     lm = MockModel(top_k=1)
     self.assertEqual(

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
 
@@ -141,6 +144,10 @@ from langfun.core.llms.groq import GroqWhisper_Large_v3Turbo
 # LLaMA C++ models.
 from langfun.core.llms.llama_cpp import LlamaCppRemote
 
+# DeepSeek models.
+from langfun.core.llms.deepseek import DeepSeek
+from langfun.core.llms.deepseek import DeepSeekChat
+
 # Placeholder for Google-internal imports.
 
 # Include cache as sub-module.

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()