langfun 0.0.2.dev20240603__py3-none-any.whl → 0.0.2.dev20240604__py3-none-any.whl

langfun/__init__.py

@@ -63,7 +63,7 @@ Image = modalities.Image
 Video = modalities.Video
 PDF = modalities.PDF
 
-# Error types.
+# Additional error types.
 MappingError = structured.MappingError
 SchemaError = structured.SchemaError
 JsonError = structured.JsonError

langfun/core/__init__.py

@@ -106,6 +106,11 @@ from langfun.core.language_model import LMScoringResult
 from langfun.core.language_model import LMCache
 from langfun.core.language_model import LMDebugMode
 
+from langfun.core.language_model import LMError
+from langfun.core.language_model import RetryableLMError
+from langfun.core.language_model import RateLimitError
+from langfun.core.language_model import TemporaryLMError
+
 # Components for building agents.
 from langfun.core.memory import Memory
 

langfun/core/language_model.py

@@ -29,6 +29,32 @@ TOKENS_PER_REQUEST = 250  # Estimated num tokens for a single request
 DEFAULT_MAX_CONCURRENCY = 1  # Use this as max concurrency if no RPM or TPM data
 
 
+#
+# Common errors during calling language models.
+#
+
+
+class LMError(RuntimeError):
+  """Base class for language model errors."""
+
+
+class RetryableLMError(LMError):
+  """Base class for LLM errors that can be solved by retrying."""
+
+
+class RateLimitError(RetryableLMError):
+  """Error for rate limit reached."""
+
+
+class TemporaryLMError(RetryableLMError):
+  """Error for temporary service issues that can be retried."""
+
+
+#
+# Language model input/output interfaces.
+#
+
+
 class LMSample(pg.Object):
   """Response candidate."""
 
@@ -445,7 +471,7 @@ class LanguageModel(component.Component):
           None,
           Union[Type[Exception], Tuple[Type[Exception], str]],
           Sequence[Union[Type[Exception], Tuple[Type[Exception], str]]],
-      ] = None,
+      ] = RetryableLMError,
   ) -> Any:
     """Helper method for subclasses for implementing _sample."""
     return concurrent.concurrent_execute(

langfun/core/llms/__init__.py

@@ -24,6 +24,9 @@ from langfun.core.llms.fake import StaticMapping
 from langfun.core.llms.fake import StaticResponse
 from langfun.core.llms.fake import StaticSequence
 
+# REST-based models.
+from langfun.core.llms.rest import REST
+
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
 from langfun.core.llms.google_genai import GeminiPro

langfun/core/llms/anthropic.py

@@ -14,14 +14,13 @@
 """Language models from Anthropic."""
 
 import base64
-import functools
 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
 import pyglove as pg
-import requests
 
 
 SUPPORTED_MODELS_AND_SETTINGS = {
@@ -38,24 +37,8 @@ SUPPORTED_MODELS_AND_SETTINGS = {
 }
 
 
-class AnthropicError(Exception):  # pylint: disable=g-bad-exception-name
-  """Base class for Anthropic errors."""
-
-
-class RateLimitError(AnthropicError):
-  """Error for rate limit reached."""
-
-
-class OverloadedError(AnthropicError):
-  """Anthropic's server is temporarily overloaded."""
-
-
-_ANTHROPIC_MESSAGE_API_ENDPOINT = 'https://api.anthropic.com/v1/messages'
-_ANTHROPIC_API_VERSION = '2023-06-01'
-
-
 @lf.use_init_args(['model'])
-class Anthropic(lf.LanguageModel):
+class Anthropic(rest.REST):
   """Anthropic LLMs (Claude) through REST APIs.
 
   See https://docs.anthropic.com/claude/reference/messages_post
@@ -80,14 +63,18 @@ class Anthropic(lf.LanguageModel):
       ),
   ] = None
 
+  api_endpoint: str = 'https://api.anthropic.com/v1/messages'
+
+  api_version: Annotated[
+      str,
+      'Anthropic API version.'
+  ] = '2023-06-01'
+
   def _on_bound(self):
     super()._on_bound()
     self._api_key = None
-    self.__dict__.pop('_api_initialized', None)
-    self.__dict__.pop('_session', None)
 
-  @functools.cached_property
-  def _api_initialized(self):
+  def _initialize(self):
     api_key = self.api_key or os.environ.get('ANTHROPIC_API_KEY', None)
     if not api_key:
       raise ValueError(
@@ -95,18 +82,14 @@ class Anthropic(lf.LanguageModel):
           'variable `ANTHROPIC_API_KEY` with your Anthropic API key.'
       )
     self._api_key = api_key
-    return True
 
-  @functools.cached_property
-  def _session(self) -> requests.Session:
-    assert self._api_initialized
-    s = requests.Session()
-    s.headers.update({
+  @property
+  def headers(self) -> dict[str, Any]:
+    return {
         'x-api-key': self._api_key,
-        'anthropic-version': _ANTHROPIC_API_VERSION,
+        'anthropic-version': self.api_version,
         'content-type': 'application/json',
-    })
-    return s
+    }
 
   @property
   def model_id(self) -> str:
@@ -121,13 +104,24 @@ class Anthropic(lf.LanguageModel):
         requests_per_min=rpm, tokens_per_min=tpm
     )
 
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    assert self._api_initialized
-    return self._parallel_execute_with_currency_control(
-        self._sample_single, prompts, retry_on_errors=(RateLimitError)
+  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 _get_request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+  def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
     """Returns a dict as request arguments."""
     # Authropic requires `max_tokens` to be specified.
     max_tokens = (
@@ -174,6 +168,19 @@ class Anthropic(lf.LanguageModel):
     else:
       return [dict(type='text', text=prompt.text)]
 
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    message = self._message_from_content(json['content'])
+    input_tokens = json['usage']['input_tokens']
+    output_tokens = json['usage']['output_tokens']
+    return lf.LMSamplingResult(
+        [lf.LMSample(message)],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=input_tokens,
+            completion_tokens=output_tokens,
+            total_tokens=input_tokens + output_tokens,
+        ),
+    )
+
   def _message_from_content(self, content: list[dict[str, Any]]) -> lf.Message:
     """Converts Anthropic's content protocol to message."""
     # Refer: https://docs.anthropic.com/claude/reference/messages-examples
@@ -181,49 +188,6 @@ class Anthropic(lf.LanguageModel):
         [x['text'] for x in content if x['type'] == 'text']
     )
 
-  def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
-    """Parses Anthropic's response."""
-    # NOTE(daiyip): Refer https://docs.anthropic.com/claude/reference/errors
-    if response.status_code == 200:
-      output = response.json()
-      message = self._message_from_content(output['content'])
-      input_tokens = output['usage']['input_tokens']
-      output_tokens = output['usage']['output_tokens']
-      return lf.LMSamplingResult(
-          [lf.LMSample(message)],
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=input_tokens,
-              completion_tokens=output_tokens,
-              total_tokens=input_tokens + output_tokens,
-          ),
-      )
-    else:
-      if response.status_code == 429:
-        error_cls = RateLimitError
-      elif response.status_code in (502, 529):
-        error_cls = OverloadedError
-      else:
-        error_cls = AnthropicError
-      raise error_cls(f'{response.status_code}: {response.content}')
-
-  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    request = dict()
-    request.update(self._get_request_args(self.sampling_options))
-    request.update(
-        dict(
-            messages=[
-                dict(role='user', content=self._content_from_message(prompt))
-            ]
-        )
-    )
-    try:
-      response = self._session.post(
-          _ANTHROPIC_MESSAGE_API_ENDPOINT, json=request, timeout=self.timeout,
-      )
-      return self._parse_response(response)
-    except ConnectionError as e:
-      raise OverloadedError(str(e)) from e
-
 
 class Claude3(Anthropic):
   """Base class for Claude 3 models. 200K input tokens and 4K output tokens."""

langfun/core/llms/anthropic_test.py

@@ -160,7 +160,7 @@ class AnthropicTest(unittest.TestCase):
         with self.assertRaisesRegex(
             Exception, f'.*{status_code}: .*{error_message}'
         ):
-          lm('hello', lm=lm, max_attempts=1)
+          lm('hello', max_attempts=1)
 
 
 if __name__ == '__main__':

langfun/core/llms/groq.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 """Language models from Groq."""
 
-import functools
 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
 import pyglove as pg
-import requests
 
 
 SUPPORTED_MODELS_AND_SETTINGS = {
@@ -33,23 +32,8 @@ SUPPORTED_MODELS_AND_SETTINGS = {
 }
 
 
-class GroqError(Exception):  # pylint: disable=g-bad-exception-name
-  """Base class for Groq errors."""
-
-
-class RateLimitError(GroqError):
-  """Error for rate limit reached."""
-
-
-class OverloadedError(GroqError):
-  """Groq's server is temporarily overloaded."""
-
-
-_CHAT_COMPLETE_API_ENDPOINT = 'https://api.groq.com/openai/v1/chat/completions'
-
-
 @lf.use_init_args(['model'])
-class Groq(lf.LanguageModel):
+class Groq(rest.REST):
   """Groq LLMs through REST APIs (OpenAI compatible).
 
   See https://platform.openai.com/docs/api-reference/chat
@@ -74,14 +58,13 @@ class Groq(lf.LanguageModel):
       ),
   ] = None
 
+  api_endpoint: str = 'https://api.groq.com/openai/v1/chat/completions'
+
   def _on_bound(self):
     super()._on_bound()
     self._api_key = None
-    self.__dict__.pop('_api_initialized', None)
-    self.__dict__.pop('_session', None)
 
-  @functools.cached_property
-  def _api_initialized(self):
+  def _initialize(self):
     api_key = self.api_key or os.environ.get('GROQ_API_KEY', None)
     if not api_key:
       raise ValueError(
@@ -89,17 +72,13 @@ class Groq(lf.LanguageModel):
           'variable `GROQ_API_KEY` with your Groq API key.'
       )
     self._api_key = api_key
-    return True
 
-  @functools.cached_property
-  def _session(self) -> requests.Session:
-    assert self._api_initialized
-    s = requests.Session()
-    s.headers.update({
+  @property
+  def headers(self) -> dict[str, Any]:
+    return {
         'Authorization': f'Bearer {self._api_key}',
         'Content-Type': 'application/json',
-    })
-    return s
+    }
 
   @property
   def model_id(self) -> str:
@@ -110,7 +89,24 @@ class Groq(lf.LanguageModel):
   def max_concurrency(self) -> int:
     return SUPPORTED_MODELS_AND_SETTINGS[self.model].max_concurrency
 
-  def _get_request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+  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(
@@ -148,6 +144,21 @@ class Groq(lf.LanguageModel):
       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'],
+        ),
+    )
+
   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
@@ -158,62 +169,6 @@ class Groq(lf.LanguageModel):
         [x['text'] for x in content if x['type'] == 'text']
     )
 
-  def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
-    """Parses Groq's response."""
-    # Refer: https://platform.openai.com/docs/api-reference/chat/object
-    if response.status_code == 200:
-      output = response.json()
-      samples = [
-          lf.LMSample(self._message_from_choice(choice), score=0.0)
-          for choice in output['choices']
-      ]
-      usage = output['usage']
-      return lf.LMSamplingResult(
-          samples,
-          usage=lf.LMSamplingUsage(
-              prompt_tokens=usage['prompt_tokens'],
-              completion_tokens=usage['completion_tokens'],
-              total_tokens=usage['total_tokens'],
-          ),
-      )
-    else:
-      # https://platform.openai.com/docs/guides/error-codes/api-errors
-      if response.status_code == 429:
-        error_cls = RateLimitError
-      elif response.status_code in (500, 502, 503):
-        error_cls = OverloadedError
-      else:
-        error_cls = GroqError
-      raise error_cls(f'{response.status_code}: {response.content}')
-
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    assert self._api_initialized
-    return self._parallel_execute_with_currency_control(
-        self._sample_single,
-        prompts,
-        retry_on_errors=(RateLimitError, OverloadedError),
-    )
-
-  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
-    request = dict()
-    request.update(self._get_request_args(self.sampling_options))
-    request.update(
-        dict(
-            messages=[
-                dict(role='user', content=self._content_from_message(prompt))
-            ]
-        )
-    )
-    try:
-      response = self._session.post(
-          _CHAT_COMPLETE_API_ENDPOINT,
-          json=request,
-          timeout=self.timeout,
-      )
-      return self._parse_response(response)
-    except ConnectionError as e:
-      raise OverloadedError(str(e)) from e
-
 
 class GroqLlama3_8B(Groq):  # pylint: disable=invalid-name
   """Llama3-8B with 8K context window.

langfun/core/llms/groq_test.py

@@ -163,7 +163,7 @@ class AuthropicTest(unittest.TestCase):
         with self.assertRaisesRegex(
             Exception, f'{status_code}:.*{error_type}'
         ):
-          lm('hello', lm=lm, max_attempts=1)
+          lm('hello', max_attempts=1)
 
 
 if __name__ == '__main__':

langfun/core/llms/llama_cpp.py

@@ -13,62 +13,72 @@
 # limitations under the License.
 """Language models from llama.cpp."""
 
-from typing import Annotated
+from typing import Any
 
 import langfun.core as lf
-import requests
+from langfun.core.llms import rest
+import pyglove as pg
 
 
-@lf.use_init_args(["url"])
-class LlamaCppRemote(lf.LanguageModel):
+class LlamaCppRemote(rest.REST):
   """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 name of the model to use.",
-  ] = ""
-
-  name: Annotated[
-      str,
-      "The abbreviation for the LLaMA CPP-based model name.",
-  ] = ""
+  @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 model_id(self) -> str:
     """Returns a string to identify the model."""
-    return f"LLaMAC++({self.name})"
+    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(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    def _complete_fn(cur_prompts):
-      results = []
-      for prompt in cur_prompts:
-        result = lf.LMSamplingResult()
-        for _ in range(self.sampling_options.n or 1):
-          data = {
-              "prompt": prompt.text,
-              "n_predict": self.sampling_options.max_tokens,
-              "top_k": self.sampling_options.top_k or 50,
-              "top_p": self.sampling_options.top_p or 0.95,
-          }
-          if self.sampling_options.temperature is not None:
-            data["temperature"] = self.sampling_options.temperature
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    request = self.request(prompt, self.sampling_options)
 
-          response = requests.post(
-              f"{self.url}/completion",
-              json=data,
-              headers={"Content-Type": "application/json"},
-              timeout=self.timeout,
-          )
-          decoded_response = response.json()
-          response = decoded_response["content"]
-          result.samples.append(lf.LMSample(response, score=0.0))
-        results.append(result)
-      return results
+    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}')
 
-    return self._parallel_execute_with_currency_control(
-        _complete_fn, [prompts]
-    )[0]
+    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

@@ -17,7 +17,6 @@ import typing
 import unittest
 from unittest import mock
 
-import langfun.core as lf
 from langfun.core.llms import llama_cpp
 
 
@@ -25,6 +24,9 @@ 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}
@@ -36,19 +38,23 @@ class LlamaCppRemoteTest(unittest.TestCase):
   """Tests for the LlamaCppRemote model."""
 
   def test_call_completion(self):
-    with mock.patch("requests.post") as mock_request:
+    with mock.patch("requests.Session.post") as mock_request:
       mock_request.side_effect = mock_requests_post
-      lm = llama_cpp.LlamaCppRemote(url="http://127.0.0.1:8080")
-      response = lm("hello", sampling_options=lf.LMSamplingOptions(n=1))
+      lm = llama_cpp.LlamaCppRemote("http://127.0.0.1:8080")
+      [result] = lm.sample(["hello"], n=2)
       self.assertEqual(
-          response.text,
+          len(result.samples),
+          2
+      )
+      self.assertEqual(
+          str(result.samples[0].response),
           "hello\nhttp://127.0.0.1:8080/completion",
       )
 
-  def test_name(self):
-    lm = llama_cpp.LlamaCppRemote()
+  def test_model_id(self):
+    lm = llama_cpp.LlamaCppRemote("http://127.0.0.1:8080")
     self.assertEqual(lm.model_id, "LLaMAC++()")
-    lm = llama_cpp.LlamaCppRemote(url="xxx", name="x")
+    lm = llama_cpp.LlamaCppRemote("xxx", model="x")
     self.assertEqual(lm.model_id, "LLaMAC++(x)")
 
 

langfun/core/llms/rest.py

@@ -0,0 +1,112 @@
+# 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.
+"""Base class for language models through REST APIs."""
+
+import functools
+from typing import Annotated, Any, Callable
+
+import langfun.core as lf
+import requests
+
+
+class REST(lf.LanguageModel):
+  """REST-based language model."""
+
+  api_endpoint: Annotated[
+      str,
+      'The endpoint of the REST API.'
+  ]
+
+  request: Annotated[
+      Callable[[lf.Message, lf.LMSamplingOptions], dict[str, Any]],
+      'A function to convert a Langfun message to a JSON request.'
+  ]
+
+  result: Annotated[
+      Callable[[dict[str, Any]], lf.LMSamplingResult],
+      'A function to convert a JSON response to an LMSamplingResult.'
+  ]
+
+  model: Annotated[
+      str | None,
+      'Model ID.'
+  ] = None
+
+  headers: Annotated[
+      dict[str, Any] | None,
+      'The headers for the REST API.'
+  ] = None
+
+  @property
+  def model_id(self) -> str:
+    """Returns a string to identify the model."""
+    return self.model or 'unknown'
+
+  @functools.cached_property
+  def _api_initialized(self) -> bool:
+    """Returns whether the API is initialized."""
+    self._initialize()
+    return True
+
+  def _initialize(self) -> None:
+    """Initializes the API. Subclasses can override."""
+
+  @functools.cached_property
+  def _session(self) -> requests.Session:
+    assert self._api_initialized
+    s = requests.Session()
+    s.headers.update(self.headers or {})
+    return s
+
+  def _on_bound(self):
+    super()._on_bound()
+    self.__dict__.pop('_session', None)
+    self.__dict__.pop('_api_initialized', None)
+
+  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
+    assert self._api_initialized
+    return self._parallel_execute_with_currency_control(
+        self._sample_single, prompts
+    )
+
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    try:
+      response = self._session.post(
+          self.api_endpoint,
+          json=self.request(prompt, self.sampling_options),
+          timeout=self.timeout,
+      )
+      return self._parse_response(response)
+    except ConnectionError as e:
+      raise lf.LMError(str(e)) from e
+
+  def _error(self, status_code: int, content: str) -> lf.LMError:
+    if status_code == 429:
+      error_cls = lf.RateLimitError
+    elif status_code in (
+        500,  # Server side issue (might be bug).
+        502,  # Bad gateway (upstream issue, might retry).
+        503,  # Servers currently under load, retry after a brief wait.
+    ):
+      error_cls = lf.TemporaryLMError
+    else:
+      error_cls = lf.LMError
+    return error_cls(f'{status_code}: {content}')
+
+  def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
+    """Parses Anthropic's response."""
+    if response.status_code == 200:
+      return self.result(response.json())
+    else:
+      raise self._error(response.status_code, response.content)

langfun/core/llms/rest_test.py

@@ -0,0 +1,111 @@
+# 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.
+"""Tests for REST models."""
+
+from typing import Any
+import unittest
+from unittest import mock
+import langfun.core as lf
+from langfun.core.llms import rest
+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({
+      'content': [(
+          f'hello with temperature={json.get("temperature")}, '
+          f'top_k={json.get("top_k")}, '
+          f'top_p={json.get("top_p")}, '
+          f'max_tokens={json.get("max_tokens")}, '
+          f'stop={json.get("stop_sequences")}.'
+      )],
+  }).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 RestTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._lm = rest.REST(
+        api_endpoint='https://fake-api.com',
+        request=lambda x, o: dict(
+            model='test-model',
+            prompt=x.text,
+            temperature=0.0,
+            top_k=0.1,
+            top_p=0.2,
+            stop_sequences=['\n'],
+            max_tokens=4096,
+        ),
+        result=lambda x: lf.LMSamplingResult(
+            [lf.LMSample(c) for c in x['content']]),
+        headers=dict(api_key='fake_key'),
+    )
+
+  def test_call(self):
+    with mock.patch('requests.Session.post') as mock_request:
+      mock_request.side_effect = mock_requests_post
+      self.assertEqual(self._lm.model_id, 'unknown')
+      response = self._lm(
+          'hello', temperature=0.0, top_k=0.1, top_p=0.2, stop=['\n'])
+      self.assertEqual(
+          response.text,
+          (
+              'hello with temperature=0.0, top_k=0.1, top_p=0.2, '
+              "max_tokens=4096, stop=['\\n']."
+          ),
+      )
+      self.assertIsNone(response.usage)
+
+  def test_call_errors(self):
+    for status_code, error_type, error_message in [
+        (429, 'rate_limit', 'Rate limit exceeded.'),
+        (529, '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
+        )
+        with self.assertRaisesRegex(
+            Exception, f'.*{status_code}: .*{error_message}'
+        ):
+          self._lm('hello', max_attempts=1)
+
+
+if __name__ == '__main__':
+  unittest.main()

{langfun-0.0.2.dev20240603.dist-info → langfun-0.0.2.dev20240604.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langfun
-Version: 0.0.2.dev20240603
+Version: 0.0.2.dev20240604
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.0.2.dev20240603.dist-info → langfun-0.0.2.dev20240604.dist-info}/RECORD

@@ -1,5 +1,5 @@
-langfun/__init__.py,sha256=LFsDp22pTeJHmzzKEg2OLmSVOPAym00DyF38LmrL2n4,2263
-langfun/core/__init__.py,sha256=nFJx6X7oB7IIWsAQqjbgZ_ScH-gsKg53YgAkuDvY0cw,4296
+langfun/__init__.py,sha256=P62MnqA6-f0h8iYfQ3MT6Yg7a4qRnQeb4GrIn6dcSnY,2274
+langfun/core/__init__.py,sha256=ZheiCpop_GAZbVpnSS-uPBJaEEM15Td5xFGGizSGqko,4514
 langfun/core/component.py,sha256=oxesbC0BoE_TbtxwW5x-BAZWxZyyJbuPiX5S38RqCv0,9909
 langfun/core/component_test.py,sha256=uR-_Sz_42Jxc5qzLIB-f5_pXmNwnC01Xlbv5NOQSeSU,8021
 langfun/core/concurrent.py,sha256=TRc49pJ3HQro2kb5FtcWkHjhBm8UcgE8RJybU5cU3-0,24537
@@ -8,7 +8,7 @@ langfun/core/console.py,sha256=bk5rNPNm9rMGW5YT2HixxU04p2umnoabn5SDz6Dqe88,2317
 langfun/core/console_test.py,sha256=5SYJdxpJGLgdSSQqqMPoA1X6jpsLD8rgcyk-EgI65oE,1077
 langfun/core/langfunc.py,sha256=RvIcRjIq0jWYRu1xim-FYe4HSrt97r3GMBO_PuagUmw,11060
 langfun/core/langfunc_test.py,sha256=_mfARnakX3oji5HDigFSLMd6yQ2wma-2Mgbztwqn73g,8501
-langfun/core/language_model.py,sha256=owNCgefGoPeRCHrxBhMtNdOj3orbeVml4eqLf1n211o,20760
+langfun/core/language_model.py,sha256=PocBg1t3uB0a_bJntLW5aagHhNbZsVdp2iduSBEW6ro,21240
 langfun/core/language_model_test.py,sha256=NZaSUls6cZdtxiqkqumWbtkx9zgNiJlsviYZOWkuHig,20137
 langfun/core/memory.py,sha256=f-asN1F7Vehgdn_fK84v73GrEUOxRtaW934keutTKjk,2416
 langfun/core/message.py,sha256=Rw3yC9HyGRjMhfDgyNjGlSCALEyDDbJ0_o6qTXeeDiQ,15738
@@ -48,19 +48,21 @@ langfun/core/eval/patching.py,sha256=R0s2eAd1m97exQt06dmUL0V_MBG0W2Hxg7fhNB7cXW0
 langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrCG1L6w,4775
 langfun/core/eval/scoring.py,sha256=1J7IATo-8FXUR0SBqk9icztHiM0lWkBFcWUo-vUURgQ,6376
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
-langfun/core/llms/__init__.py,sha256=XHK_ZpfEppCF-ixfpIvmrOvH2P6XgkjMhS7zBa8yYk4,4302
-langfun/core/llms/anthropic.py,sha256=7W9YdPN3SlAFhAIQlihMkrpo7tTY_4NvD0KIlCrqcsk,8505
-langfun/core/llms/anthropic_test.py,sha256=TMM30myyEhwF99Le4RvJEXOn8RYl0q1FRkt9Q9nl1jk,5540
+langfun/core/llms/__init__.py,sha256=3G7pJISeClgHGV34Gy2t_Nih4N08UhGbWe6uAff8TnA,4364
+langfun/core/llms/anthropic.py,sha256=pBYe8dVwswxKaqhNjA_jtZbyfvOaXtEo399Zty242iA,7097
+langfun/core/llms/anthropic_test.py,sha256=T-swuMkfnlgs8Fpif4rtXs579exGk0TsbLMirXDZCkg,5533
 langfun/core/llms/fake.py,sha256=Dd7-6ka9pFf3fcWZyczamjOqQ91MOI-m7We3Oc9Ffmo,2927
 langfun/core/llms/fake_test.py,sha256=ipKfdOcuqVcJ8lDXVpnBVb9HHG0hAVkFkMoHpWjC2cI,7212
 langfun/core/llms/google_genai.py,sha256=Rl5a5CyF_6Y0BYYArKk8yMaenv1rH3MUQLy6b3dfMRI,10202
 langfun/core/llms/google_genai_test.py,sha256=iTISk3tJ4-3gjWmzcKQhEbH3ke4AkEiCu8rAGtB7SvU,7535
-langfun/core/llms/groq.py,sha256=NaGItVL_pkOpqPpI4bPGU27xLFRoaeizZ49v2s-4ERs,7844
-langfun/core/llms/groq_test.py,sha256=M6GtlrsOvDun_j-sR8cPh4W_moHWZNSTiThu3kuwbbc,5281
-langfun/core/llms/llama_cpp.py,sha256=Y_KkMUf3Xfac49koMUtUslKl3h-HWp3-ntq7Jaa3bdo,2385
-langfun/core/llms/llama_cpp_test.py,sha256=ZxC6defGd_HX9SFRU9U4cJiQnBKundbOrchbXuC1Z2M,1683
+langfun/core/llms/groq.py,sha256=pqtyOZ_1_OJMOg8xATWT_B_SVbuT9nMRf4VkH9GzW8g,6308
+langfun/core/llms/groq_test.py,sha256=GYF_Qtq5S1H1TrKH38t6_lkdroqT7v-joYLDKnmS9e0,5274
+langfun/core/llms/llama_cpp.py,sha256=9tXQntSCDtjTF3bnyJrAPCr4N6wycy5nXYvp9uduygE,2843
+langfun/core/llms/llama_cpp_test.py,sha256=MWO_qaOeKjRniGjcaWPDScd7HPaIJemqUZoslrt4FPs,1806
 langfun/core/llms/openai.py,sha256=IN46gIqfY6aEEfxCPNmyH1hrep3oWBhJDwVFilfqNkM,13657
 langfun/core/llms/openai_test.py,sha256=QWDzTgi8F2Z9u9ip6alK4rDEp_YraVTxWlDX5XOsKJk,14858
+langfun/core/llms/rest.py,sha256=laopuq-zD8V-3Y6eFDngftHEbE66VlUkCD2-rvvRaLU,3388
+langfun/core/llms/rest_test.py,sha256=Zw08Xbl_O2OQuUglLXHsPsY5KW2VEcPGl1gR8PRHuFY,3449
 langfun/core/llms/vertexai.py,sha256=eILbXoMSza5r4FLGlIdH6-eD8Ggy9Z4PdjLaBDxy29A,11162
 langfun/core/llms/vertexai_test.py,sha256=G18BG36h5KvmX2zutDTLjtYCRjTuP_nWIFm4FMnLnyY,7651
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
@@ -111,8 +113,8 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=DYVrkk7uNKCqJGEHH31HssU2BPuMItU1vJLzfcXIlYg,2156
-langfun-0.0.2.dev20240603.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.0.2.dev20240603.dist-info/METADATA,sha256=G7sGJIdQQ5xbDndwUtXcIw1m-xHA1taiAHngx0mamkk,3550
-langfun-0.0.2.dev20240603.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-langfun-0.0.2.dev20240603.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.0.2.dev20240603.dist-info/RECORD,,
+langfun-0.0.2.dev20240604.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.0.2.dev20240604.dist-info/METADATA,sha256=WGM_N4Nizh0eT4OvlSgopXgcnHy7uset6O7iAjzqUdk,3550
+langfun-0.0.2.dev20240604.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+langfun-0.0.2.dev20240604.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.0.2.dev20240604.dist-info/RECORD,,