langfun 0.0.2.dev20240601__tar.gz → 0.0.2.dev20240604__tar.gz

{langfun-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/PKG-INFO

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

{langfun-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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-0.0.2.dev20240604/langfun/core/llms/llama_cpp.py

@@ -0,0 +1,84 @@
+# 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 Any
+
+import langfun.core as lf
+from langfun.core.llms import rest
+import pyglove as pg
+
+
+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
+  """
+
+  @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.model or ""})'
+
+  def request(
+      self, prompt: lf.Message, sampling_options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    """Returns the JSON input for a message."""
+    request = dict()
+    request.update(self._request_args(sampling_options))
+    # NOTE(daiyip): multi-modal is current not supported.
+    request['prompt'] = prompt.text
+    return request
+
+  def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+    """Returns a dict as request arguments."""
+    args = dict(
+        n_predict=options.max_tokens or 1024,
+        top_k=options.top_k or 50,
+        top_p=options.top_p or 0.95,
+    )
+    if options.temperature is not None:
+      args['temperature'] = options.temperature
+    return args
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    return lf.LMSamplingResult(
+        [lf.LMSample(item['content'], score=0.0) for item in json['items']]
+    )
+
+  def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
+    request = self.request(prompt, self.sampling_options)
+
+    def _sample_one_example(request):
+      response = self._session.post(
+          self.api_endpoint,
+          json=request,
+          timeout=self.timeout,
+      )
+      if response.status_code == 200:
+        return response.json()
+      else:
+        error_cls = self._error_cls_from_status(response.status_code)
+        raise error_cls(f'{response.status_code}: {response.content}')
+
+    items = self._parallel_execute_with_currency_control(
+        _sample_one_example,
+        [request] * (self.sampling_options.n or 1),
+    )
+    return self.result(dict(items=items))

{langfun-0.0.2.dev20240601 → langfun-0.0.2.dev20240604}/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)")