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

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/core/llms/deepseek.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
 
 SUPPORTED_MODELS_AND_SETTINGS = {
@@ -39,7 +38,7 @@ SUPPORTED_MODELS_AND_SETTINGS = {
 # DeepSeek API uses an API format compatible with OpenAI.
 # Reference: https://api-docs.deepseek.com/
 @lf.use_init_args(['model'])
-class DeepSeek(rest.REST):
+class DeepSeek(openai_compatible.OpenAICompatible):
   """DeepSeek model."""
 
   model: pg.typing.Annotated[
@@ -51,10 +50,6 @@ class DeepSeek(rest.REST):
 
   api_endpoint: str = 'https://api.deepseek.com/chat/completions'
 
-  multimodal: Annotated[bool, 'Whether this model has multimodal support.'] = (
-      False
-  )
-
   api_key: Annotated[
       str | None,
       (
@@ -63,25 +58,18 @@ class DeepSeek(rest.REST):
       ),
   ] = None
 
-  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('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.'
       )
-    self._api_key = api_key
-
-  @property
-  def headers(self) -> dict[str, Any]:
-    headers = {
-        'Content-Type': 'application/json',
-        'Authorization': f'Bearer {self._api_key}',
-    }
+    headers = super().headers
+    headers.update({
+        'Authorization': f'Bearer {api_key}',
+    })
     return headers
 
   @property
@@ -118,138 +106,6 @@ class DeepSeek(rest.REST):
   def dir(cls):
     return [k for k, v in SUPPORTED_MODELS_AND_SETTINGS.items() if v.in_service]
 
-  def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
-    # Reference:
-    # https://platform.openai.com/docs/api-reference/completions/create
-    # NOTE(daiyip): options.top_k is not applicable.
-    args = dict(
-        model=self.model,
-        n=options.n,
-        top_logprobs=options.top_logprobs,
-    )
-    if options.logprobs:
-      args['logprobs'] = options.logprobs
-
-    if options.temperature is not None:
-      args['temperature'] = options.temperature
-    if options.max_tokens is not None:
-      args['max_completion_tokens'] = options.max_tokens
-    if options.top_p is not None:
-      args['top_p'] = options.top_p
-    if options.stop:
-      args['stop'] = options.stop
-    if options.random_seed is not None:
-      args['seed'] = options.random_seed
-    return args
-
-  def _content_from_message(self, message: lf.Message):
-    """Returns a OpenAI content object from a Langfun message."""
-
-    def _uri_from(chunk: lf.Modality) -> str:
-      if chunk.uri and chunk.uri.lower().startswith(
-          ('http:', 'https:', 'ftp:')
-      ):
-        return chunk.uri
-      return chunk.content_uri
-
-    content = []
-    for chunk in message.chunk():
-      if isinstance(chunk, str):
-        item = dict(type='text', text=chunk)
-      elif isinstance(chunk, lf_modalities.Image) and self.multimodal:
-        item = dict(type='image_url', image_url=dict(url=_uri_from(chunk)))
-      else:
-        raise ValueError(f'Unsupported modality: {chunk!r}.')
-      content.append(item)
-    return content
-
-  def request(
-      self, prompt: lf.Message, sampling_options: lf.LMSamplingOptions
-  ) -> dict[str, Any]:
-    """Returns the JSON input for a message."""
-    request_args = self._request_args(sampling_options)
-
-    # Users could use `metadata_json_schema` to pass additional
-    # request arguments.
-    json_schema = prompt.metadata.get('json_schema')
-    if json_schema is not None:
-      if not isinstance(json_schema, dict):
-        raise ValueError(f'`json_schema` must be a dict, got {json_schema!r}.')
-      if 'title' not in json_schema:
-        raise ValueError(
-            'The root of `json_schema` must have a `title` field, '
-            f'got {json_schema!r}.'
-        )
-      request_args.update(
-          response_format=dict(
-              type='json_schema',
-              json_schema=dict(
-                  schema=json_schema,
-                  name=json_schema['title'],
-                  strict=True,
-              ),
-          )
-      )
-      prompt.metadata.formatted_text = (
-          prompt.text
-          + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
-          + pg.to_json_str(request_args['response_format'], json_indent=2)
-      )
-
-    # Prepare messages.
-    messages = []
-    # Users could use `metadata_system_message` to pass system message.
-    system_message = prompt.metadata.get('system_message')
-    if system_message:
-      system_message = lf.SystemMessage.from_value(system_message)
-      messages.append(
-          dict(
-              role='system', content=self._content_from_message(system_message)
-          )
-      )
-    messages.append(
-        dict(role='user', content=self._content_from_message(prompt))
-    )
-    request = dict()
-    request.update(request_args)
-    request['messages'] = messages
-    return request
-
-  def _parse_choice(self, choice: dict[str, Any]) -> lf.LMSample:
-    # Reference:
-    # https://platform.openai.com/docs/api-reference/chat/object
-    logprobs = None
-    choice_logprobs = choice.get('logprobs')
-    if choice_logprobs:
-      logprobs = [
-          (
-              t['token'],
-              t['logprob'],
-              [(tt['token'], tt['logprob']) for tt in t['top_logprobs']],
-          )
-          for t in choice_logprobs['content']
-      ]
-    return lf.LMSample(
-        choice['message']['content'],
-        score=0.0,
-        logprobs=logprobs,
-    )
-
-  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
-    usage = json['usage']
-    return lf.LMSamplingResult(
-        samples=[self._parse_choice(choice) for choice in json['choices']],
-        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'],
-            ),
-        ),
-    )
-
 
 class DeepSeekChat(DeepSeek):
   """DeepSeek Chat model.

langfun/core/llms/deepseek_test.py

@@ -11,72 +11,8 @@
 # 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 OpenAI models."""
-
-from typing import Any
 import unittest
-from unittest import mock
-
-import langfun.core as lf
 from langfun.core.llms import deepseek
-import pyglove as pg
-import requests
-
-
-def mock_chat_completion_request(url: str, json: dict[str, Any], **kwargs):
-  del url, kwargs
-  messages = json['messages']
-  if len(messages) > 1:
-    system_message = f' system={messages[0]["content"]}'
-  else:
-    system_message = ''
-
-  if 'response_format' in json:
-    response_format = f' format={json["response_format"]["type"]}'
-  else:
-    response_format = ''
-
-  choices = []
-  for k in range(json['n']):
-    if json.get('logprobs'):
-      logprobs = dict(
-          content=[
-              dict(
-                  token='chosen_token',
-                  logprob=0.5,
-                  top_logprobs=[
-                      dict(
-                          token=f'alternative_token_{i + 1}',
-                          logprob=0.1
-                      ) for i in range(3)
-                  ]
-              )
-          ]
-      )
-    else:
-      logprobs = None
-
-    choices.append(dict(
-        message=dict(
-            content=(
-                f'Sample {k} for message.{system_message}{response_format}'
-            )
-        ),
-        logprobs=logprobs,
-    ))
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str(
-      dict(
-          choices=choices,
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=100,
-              completion_tokens=100,
-              total_tokens=200,
-          ),
-      )
-  ).encode()
-  return response
 
 
 class DeepSeekTest(unittest.TestCase):
@@ -87,7 +23,14 @@ class DeepSeekTest(unittest.TestCase):
 
   def test_key(self):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      deepseek.DeepSeekChat()('hi')
+      _ = 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(
@@ -106,333 +49,13 @@ class DeepSeekTest(unittest.TestCase):
         deepseek.DeepSeekChat(api_key='test_key').max_concurrency, 0
     )
 
-  def test_request_args(self):
-    self.assertEqual(
-        deepseek.DeepSeekChat(api_key='test_key')._request_args(
-            lf.LMSamplingOptions(
-                temperature=1.0, stop=['\n'], n=1, random_seed=123
-            )
-        ),
-        dict(
-            model='deepseek-chat',
-            top_logprobs=None,
-            n=1,
-            temperature=1.0,
-            stop=['\n'],
-            seed=123,
-        ),
-    )
-
-  def test_call_chat_completion(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = deepseek.DeepSeek(model='deepseek-chat', api_key='test_key')
-      self.assertEqual(
-          lm('hello', sampling_options=lf.LMSamplingOptions(n=2)),
-          'Sample 0 for message.',
-      )
-
-  def test_call_chat_completion_with_logprobs(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = deepseek.DeepSeek(model='deepseek-chat', api_key='test_key')
-      results = lm.sample(['hello'], logprobs=True)
-      self.assertEqual(len(results), 1)
-      expected = lf.LMSamplingResult(
-          [
-              lf.LMSample(
-                  response=lf.AIMessage(
-                      text='Sample 0 for message.',
-                      metadata={
-                          'score': 0.0,
-                          'logprobs': [(
-                              'chosen_token',
-                              0.5,
-                              [
-                                  ('alternative_token_1', 0.1),
-                                  ('alternative_token_2', 0.1),
-                                  ('alternative_token_3', 0.1),
-                              ],
-                          )],
-                          'is_cached': False,
-                          'usage': lf.LMSamplingUsage(
-                              prompt_tokens=100,
-                              completion_tokens=100,
-                              total_tokens=200,
-                              estimated_cost=4.2e-05,
-                          ),
-                      },
-                      tags=['lm-response'],
-                  ),
-                  logprobs=[(
-                      'chosen_token',
-                      0.5,
-                      [
-                          ('alternative_token_1', 0.1),
-                          ('alternative_token_2', 0.1),
-                          ('alternative_token_3', 0.1),
-                      ],
-                  )],
-              )
-          ],
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=100,
-              completion_tokens=100,
-              total_tokens=200,
-              estimated_cost=4.2e-05,
-          ),
-      )
-      self.assertTrue(pg.eq(results[0], expected))
-
-  def test_sample_chat_completion(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      deepseek.SUPPORTED_MODELS_AND_SETTINGS['deepseek-chat'].update({
-          'cost_per_1k_input_tokens': 1.0,
-          'cost_per_1k_output_tokens': 1.0,
-      })
-      lm = deepseek.DeepSeek(api_key='test_key', model='deepseek-chat')
-      results = lm.sample(
-          ['hello', 'bye'], sampling_options=lf.LMSamplingOptions(n=3)
-      )
-
-    self.assertEqual(len(results), 2)
-    print(results[0])
-    self.assertEqual(
-        results[0],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200,
-                estimated_cost=0.2,
-            ),
-        ),
-    )
+  def test_estimate_cost(self):
     self.assertEqual(
-        results[1],
-        lf.LMSamplingResult(
-            [
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 0 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 1 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-                lf.LMSample(
-                    lf.AIMessage(
-                        'Sample 2 for message.',
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=33,
-                            completion_tokens=33,
-                            total_tokens=66,
-                            estimated_cost=0.2 / 3,
-                        ),
-                        tags=[lf.Message.TAG_LM_RESPONSE],
-                    ),
-                    score=0.0,
-                    logprobs=None,
-                ),
-            ],
-            usage=lf.LMSamplingUsage(
-                prompt_tokens=100, completion_tokens=100, total_tokens=200,
-                estimated_cost=0.2,
-            ),
+        deepseek.DeepSeekChat(api_key='test_key').estimate_cost(
+            num_input_tokens=100, num_output_tokens=100
         ),
+        4.2e-5
     )
 
-  def test_sample_with_contextual_options(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = deepseek.DeepSeek(api_key='test_key', model='deepseek-chat')
-      with lf.use_settings(sampling_options=lf.LMSamplingOptions(n=2)):
-        results = lm.sample(['hello'])
-
-    self.assertEqual(len(results), 1)
-    expected = lf.LMSamplingResult(
-        samples=[
-            lf.LMSample(
-                response=lf.AIMessage(
-                    text='Sample 0 for message.',
-                    sender='AI',
-                    metadata=pg.Dict(
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=50,
-                            completion_tokens=50,
-                            total_tokens=100,
-                            num_requests=1,
-                            estimated_cost=0.1,
-                        ),
-                    ),
-                    tags=['lm-response'],
-                ),
-                score=0.0,
-                logprobs=None,
-            ),
-            lf.LMSample(
-                response=lf.AIMessage(
-                    text='Sample 1 for message.',
-                    sender='AI',
-                    metadata=pg.Dict(
-                        score=0.0,
-                        logprobs=None,
-                        is_cached=False,
-                        usage=lf.LMSamplingUsage(
-                            prompt_tokens=50,
-                            completion_tokens=50,
-                            total_tokens=100,
-                            num_requests=1,
-                            estimated_cost=0.1,
-                        ),
-                    ),
-                    tags=['lm-response'],
-                ),
-                score=0.0,
-                logprobs=None,
-            ),
-        ],
-        usage=lf.LMSamplingUsage(
-            prompt_tokens=100,
-            completion_tokens=100,
-            total_tokens=200,
-            num_requests=1,
-            estimated_cost=0.2,
-        ),
-        is_cached=False,
-    )
-    self.assertTrue(pg.eq(results[0], expected))
-
-  def test_call_with_system_message(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = deepseek.DeepSeek(api_key='test_key', model='deepseek-chat')
-      self.assertEqual(
-          lm(
-              lf.UserMessage(
-                  'hello',
-                  system_message='hi',
-              ),
-              sampling_options=lf.LMSamplingOptions(n=2)
-          ),
-          '''Sample 0 for message. system=[{'type': 'text', 'text': 'hi'}]''',
-      )
-
-  def test_call_with_json_schema(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_chat_completion_request
-      lm = deepseek.DeepSeek(api_key='test_key', model='deepseek-chat')
-      self.assertEqual(
-          lm(
-              lf.UserMessage(
-                  'hello',
-                  json_schema={
-                      'type': 'object',
-                      'properties': {
-                          'name': {'type': 'string'},
-                      },
-                      'required': ['name'],
-                      'title': 'Person',
-                  }
-              ),
-              sampling_options=lf.LMSamplingOptions(n=2)
-          ),
-          'Sample 0 for message. format=json_schema',
-      )
-
-    # Test bad json schema.
-    with self.assertRaisesRegex(ValueError, '`json_schema` must be a dict'):
-      lm(lf.UserMessage('hello', json_schema='foo'))
-
-    with self.assertRaisesRegex(
-        ValueError, 'The root of `json_schema` must have a `title` field'
-    ):
-      lm(lf.UserMessage('hello', json_schema={}))
-
-
 if __name__ == '__main__':
   unittest.main()