langfun 0.0.2.dev20240422__py3-none-any.whl → 0.0.2.dev20240425__py3-none-any.whl

langfun/core/llms/anthropic.py

@@ -26,12 +26,15 @@ import requests
 
 SUPPORTED_MODELS_AND_SETTINGS = {
     # See https://docs.anthropic.com/claude/docs/models-overview
-    'claude-3-opus-20240229': pg.Dict(max_tokens=4096, max_concurrency=16),
-    'claude-3-sonnet-20240229': pg.Dict(max_tokens=4096, max_concurrency=16),
-    'claude-3-haiku-20240307': pg.Dict(max_tokens=4096, max_concurrency=16),
-    'claude-2.1': pg.Dict(max_tokens=4096, max_concurrency=16),
-    'claude-2.0': pg.Dict(max_tokens=4096, max_concurrency=16),
-    'claude-instant-1.2': pg.Dict(max_tokens=4096, max_concurrency=16),
+    # Rate limits from https://docs.anthropic.com/claude/reference/rate-limits
+    #     RPM/TPM for Claude-2.1, Claude-2.0, and Claude-Instant-1.2 estimated
+    #     as RPM/TPM of the largest-available model (Claude-3-Opus).
+    'claude-3-opus-20240229': pg.Dict(max_tokens=4096, rpm=4000, tpm=400000),
+    'claude-3-sonnet-20240229': pg.Dict(max_tokens=4096, rpm=4000, tpm=400000),
+    'claude-3-haiku-20240307': pg.Dict(max_tokens=4096, rpm=4000, tpm=400000),
+    'claude-2.1': pg.Dict(max_tokens=4096, rpm=4000, tpm=400000),
+    'claude-2.0': pg.Dict(max_tokens=4096, rpm=4000, tpm=400000),
+    'claude-instant-1.2': pg.Dict(max_tokens=4096, rpm=4000, tpm=400000),
 }
 
 
@@ -81,6 +84,7 @@ class Anthropic(lf.LanguageModel):
     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):
@@ -93,6 +97,17 @@ class Anthropic(lf.LanguageModel):
     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({
+        'x-api-key': self._api_key,
+        'anthropic-version': _ANTHROPIC_API_VERSION,
+        'content-type': 'application/json',
+    })
+    return s
+
   @property
   def model_id(self) -> str:
     """Returns a string to identify the model."""
@@ -100,7 +115,11 @@ class Anthropic(lf.LanguageModel):
 
   @property
   def max_concurrency(self) -> int:
-    return SUPPORTED_MODELS_AND_SETTINGS[self.model].max_concurrency
+    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 _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     assert self._api_initialized
@@ -165,8 +184,8 @@ class Anthropic(lf.LanguageModel):
   def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
     """Parses Anthropic's response."""
     # NOTE(daiyip): Refer https://docs.anthropic.com/claude/reference/errors
-    output = response.json()
     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']
@@ -181,12 +200,11 @@ class Anthropic(lf.LanguageModel):
     else:
       if response.status_code == 429:
         error_cls = RateLimitError
-      elif response.status_code == 529:
+      elif response.status_code in (502, 529):
         error_cls = OverloadedError
       else:
         error_cls = AnthropicError
-      error = output['error']
-      raise error_cls(f'{error["type"]}: {error["message"]}')
+      raise error_cls(f'{response.status_code}: {response.content}')
 
   def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult:
     request = dict()
@@ -198,17 +216,13 @@ class Anthropic(lf.LanguageModel):
             ]
         )
     )
-    response = requests.post(
-        _ANTHROPIC_MESSAGE_API_ENDPOINT,
-        json=request,
-        headers={
-            'x-api-key': self._api_key,
-            'anthropic-version': _ANTHROPIC_API_VERSION,
-            'content-type': 'application/json',
-        },
-        timeout=self.timeout,
-    )
-    return self._parse_response(response)
+    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):

langfun/core/llms/anthropic_test.py

@@ -98,20 +98,20 @@ def mock_requests_post_error(status_code, error_type, error_message):
   return _mock_requests
 
 
-class AuthropicTest(unittest.TestCase):
+class AnthropicTest(unittest.TestCase):
 
   def test_basics(self):
     self.assertEqual(
         anthropic.Claude3Haiku().model_id, 'claude-3-haiku-20240307'
     )
-    self.assertEqual(anthropic.Claude3Haiku().max_concurrency, 16)
+    self.assertGreater(anthropic.Claude3Haiku().max_concurrency, 0)
 
   def test_api_key(self):
     lm = anthropic.Claude3Haiku()
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
       lm('hi')
 
-    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 = anthropic.Claude3Haiku(api_key='fake key')
@@ -123,7 +123,7 @@ class AuthropicTest(unittest.TestCase):
       del os.environ['ANTHROPIC_API_KEY']
 
   def test_call(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 = anthropic.Claude3Haiku(api_key='fake_key')
       response = lm('hello', temperature=0.0, top_k=0.1, top_p=0.2, stop=['\n'])
@@ -140,7 +140,7 @@ class AuthropicTest(unittest.TestCase):
       self.assertIsNotNone(response.usage.total_tokens, 3)
 
   def test_mm_call(self):
-    with mock.patch('requests.post') as mock_mm_request:
+    with mock.patch('requests.Session.post') as mock_mm_request:
       mock_mm_request.side_effect = mock_mm_requests_post
       lm = anthropic.Claude3Haiku(api_key='fake_key')
       response = lm(lf_modalities.Image.from_bytes(image_content), lm=lm)
@@ -152,13 +152,13 @@ class AuthropicTest(unittest.TestCase):
         (529, 'service_unavailable', 'Service unavailable.'),
         (500, 'bad_request', 'Bad request.'),
     ]:
-      with mock.patch('requests.post') as mock_mm_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
         )
         lm = anthropic.Claude3Haiku(api_key='fake_key')
         with self.assertRaisesRegex(
-            Exception, f'{error_type}: {error_message}'
+            Exception, f'.*{status_code}: .*{error_message}'
         ):
           lm('hello', lm=lm, max_attempts=1)
 

langfun/core/llms/groq.py

@@ -78,6 +78,7 @@ class Groq(lf.LanguageModel):
     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):
@@ -85,11 +86,21 @@ class Groq(lf.LanguageModel):
     if not api_key:
       raise ValueError(
           'Please specify `api_key` during `__init__` or set environment '
-          'variable `GROQ_API_KEY` with your Anthropic API key.'
+          '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({
+        'Authorization': f'Bearer {self._api_key}',
+        'Content-Type': 'application/json',
+    })
+    return s
+
   @property
   def model_id(self) -> str:
     """Returns a string to identify the model."""
@@ -119,7 +130,7 @@ class Groq(lf.LanguageModel):
     return args
 
   def _content_from_message(self, prompt: lf.Message) -> list[dict[str, Any]]:
-    """Converts an message to Anthropic's content protocol (list of dicts)."""
+    """Converts an message to Groq's content protocol (list of dicts)."""
     # Refer: https://platform.openai.com/docs/api-reference/chat/create
     content = []
     for chunk in prompt.chunk():
@@ -138,7 +149,7 @@ class Groq(lf.LanguageModel):
     return content
 
   def _message_from_choice(self, choice: dict[str, Any]) -> lf.Message:
-    """Converts Anthropic's content protocol to message."""
+    """Converts Groq's content protocol to message."""
     # Refer: https://platform.openai.com/docs/api-reference/chat/create
     content = choice['message']['content']
     if isinstance(content, str):
@@ -148,10 +159,10 @@ class Groq(lf.LanguageModel):
     )
 
   def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult:
-    """Parses Anthropic's response."""
+    """Parses Groq's response."""
     # Refer: https://platform.openai.com/docs/api-reference/chat/object
-    output = response.json()
     if response.status_code == 200:
+      output = response.json()
       samples = [
           lf.LMSample(self._message_from_choice(choice), score=0.0)
           for choice in output['choices']
@@ -169,12 +180,11 @@ class Groq(lf.LanguageModel):
       # https://platform.openai.com/docs/guides/error-codes/api-errors
       if response.status_code == 429:
         error_cls = RateLimitError
-      elif response.status_code in (500, 503):
+      elif response.status_code in (500, 502, 503):
         error_cls = OverloadedError
       else:
         error_cls = GroqError
-      error = output['error']
-      raise error_cls(f'{error["type"]}: {error["message"]}')
+      raise error_cls(f'{response.status_code}: {response.content}')
 
   def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
     assert self._api_initialized
@@ -194,16 +204,15 @@ class Groq(lf.LanguageModel):
             ]
         )
     )
-    response = requests.post(
-        _CHAT_COMPLETE_API_ENDPOINT,
-        json=request,
-        headers={
-            'Authorization': f'Bearer {self._api_key}',
-            'Content-Type': 'application/json',
-        },
-        timeout=self.timeout,
-    )
-    return self._parse_response(response)
+    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

langfun/core/llms/groq_test.py

@@ -107,7 +107,7 @@ class AuthropicTest(unittest.TestCase):
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
       lm('hi')
 
-    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 = groq.GroqMistral_8x7B(api_key='fake key')
@@ -119,7 +119,7 @@ class AuthropicTest(unittest.TestCase):
       del os.environ['GROQ_API_KEY']
 
   def test_call(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 = groq.GroqLlama3_70B(api_key='fake_key')
       response = lm(
@@ -143,7 +143,7 @@ class AuthropicTest(unittest.TestCase):
       self.assertIsNotNone(response.usage.total_tokens, 3)
 
   def test_mm_call(self):
-    with mock.patch('requests.post') as mock_mm_request:
+    with mock.patch('requests.Session.post') as mock_mm_request:
       mock_mm_request.side_effect = mock_mm_requests_post
       lm = groq.GroqLlama3_70B(multimodal=True, api_key='fake_key')
       response = lm(lf_modalities.Image.from_uri('https://fake/image.jpg'))
@@ -155,13 +155,13 @@ class AuthropicTest(unittest.TestCase):
         (503, 'service_unavailable', 'Service unavailable.'),
         (500, 'bad_request', 'Bad request.'),
     ]:
-      with mock.patch('requests.post') as mock_mm_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
         )
         lm = groq.GroqLlama3_70B(api_key='fake_key')
         with self.assertRaisesRegex(
-            Exception, f'{error_type}: {error_message}'
+            Exception, f'{status_code}:.*{error_type}'
         ):
           lm('hello', lm=lm, max_attempts=1)
 

langfun/core/llms/openai.py

@@ -26,54 +26,55 @@ from openai import openai_object
 import pyglove as pg
 
 
-SUPPORTED_MODELS_AND_SETTINGS = [
-    # Model name, max concurrent requests.
-    # The concurrent requests is estimated by TPM/RPM from
-    # https://platform.openai.com/account/limits
-    # GPT-4 Turbo models.
-    ('gpt-4-turbo', 8),  # GPT-4 Turbo with Vision
-    ('gpt-4-turbo-2024-04-09', 8),  # GPT-4-Turbo with Vision, 04/09/2024
-    ('gpt-4-turbo-preview', 8),  # GPT-4 Turbo Preview
-    ('gpt-4-0125-preview', 8),  # GPT-4 Turbo Preview, 01/25/2024
-    ('gpt-4-1106-preview', 8),  # GPT-4 Turbo Preview, 11/06/2023
-    ('gpt-4-vision-preview', 8),  # GPT-4 Turbo Vision Preview.
-    ('gpt-4-1106-vision-preview', 8),  # GPT-4 Turbo Vision Preview, 11/06/2023
-    # GPT-4 models.
-    ('gpt-4', 4),
-    ('gpt-4-0613', 4),
-    ('gpt-4-0314', 4),
-    ('gpt-4-32k', 4),
-    ('gpt-4-32k-0613', 4),
-    ('gpt-4-32k-0314', 4),
-    # GPT-3.5 Turbo models.
-    ('gpt-3.5-turbo', 16),
-    ('gpt-3.5-turbo-0125', 16),
-    ('gpt-3.5-turbo-1106', 16),
-    ('gpt-3.5-turbo-0613', 16),
-    ('gpt-3.5-turbo-0301', 16),
-    ('gpt-3.5-turbo-16k', 16),
-    ('gpt-3.5-turbo-16k-0613', 16),
-    ('gpt-3.5-turbo-16k-0301', 16),
-    # GPT-3.5 models.
-    ('text-davinci-003', 8),  # GPT-3.5, trained with RHLF.
-    ('text-davinci-002', 4),  # Trained with SFT but no RHLF.
-    ('code-davinci-002', 4),
-    # GPT-3 instruction-tuned models.
-    ('text-curie-001', 4),
-    ('text-babbage-001', 4),
-    ('text-ada-001', 4),
-    ('davinci', 4),
-    ('curie', 4),
-    ('babbage', 4),
-    ('ada', 4),
-    # GPT-3 base models without instruction tuning.
-    ('babbage-002', 4),
-    ('davinci-002', 4),
-]
-
-
-# Model concurreny setting.
-_MODEL_CONCURRENCY = {m[0]: m[1] for m in SUPPORTED_MODELS_AND_SETTINGS}
+# From https://platform.openai.com/settings/organization/limits
+_DEFAULT_TPM = 250000
+_DEFAULT_RPM = 3000
+
+SUPPORTED_MODELS_AND_SETTINGS = {
+    # Models from https://platform.openai.com/docs/models
+    # RPM is from https://platform.openai.com/docs/guides/rate-limits
+    # GPT-4-Turbo models
+    'gpt-4-turbo': pg.Dict(rpm=10000, tpm=1500000),
+    'gpt-4-turbo-2024-04-09': pg.Dict(rpm=10000, tpm=1500000),
+    'gpt-4-turbo-preview': pg.Dict(rpm=10000, tpm=1500000),
+    'gpt-4-0125-preview': pg.Dict(rpm=10000, tpm=1500000),
+    'gpt-4-1106-preview': pg.Dict(rpm=10000, tpm=1500000),
+    'gpt-4-vision-preview': pg.Dict(rpm=10000, tpm=1500000),
+    'gpt-4-1106-vision-preview': pg.Dict(
+        rpm=10000, tpm=1500000
+    ),
+    # GPT-4 models
+    'gpt-4': pg.Dict(rpm=10000, tpm=300000),
+    'gpt-4-0613': pg.Dict(rpm=10000, tpm=300000),
+    'gpt-4-0314': pg.Dict(rpm=10000, tpm=300000),
+    'gpt-4-32k': pg.Dict(rpm=10000, tpm=300000),
+    'gpt-4-32k-0613': pg.Dict(rpm=10000, tpm=300000),
+    'gpt-4-32k-0314': pg.Dict(rpm=10000, tpm=300000),
+    # GPT-3.5-Turbo models
+    'gpt-3.5-turbo': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-0125': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-1106': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-0613': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-0301': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-16k': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-16k-0613': pg.Dict(rpm=10000, tpm=2000000),
+    'gpt-3.5-turbo-16k-0301': pg.Dict(rpm=10000, tpm=2000000),
+    # GPT-3.5 models
+    'text-davinci-003': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'text-davinci-002': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'code-davinci-002': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    # GPT-3 instruction-tuned models
+    'text-curie-001': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'text-babbage-001': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'text-ada-001': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'davinci': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'curie': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'babbage': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'ada': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    # GPT-3 base models
+    'babbage-002': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'davinci-002': pg.Dict(rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+}
 
 
 @lf.use_init_args(['model'])
@@ -82,7 +83,7 @@ class OpenAI(lf.LanguageModel):
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
-          pg.MISSING_VALUE, [m[0] for m in SUPPORTED_MODELS_AND_SETTINGS]
+          pg.MISSING_VALUE, list(SUPPORTED_MODELS_AND_SETTINGS.keys())
       ),
       'The name of the model to use.',
   ] = 'gpt-3.5-turbo'
@@ -134,7 +135,11 @@ class OpenAI(lf.LanguageModel):
 
   @property
   def max_concurrency(self) -> int:
-    return _MODEL_CONCURRENCY[self.model]
+    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
+    )
 
   @classmethod
   def dir(cls):

langfun/core/llms/openai_test.py

@@ -11,7 +11,7 @@
 # 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."""
+"""Tests for OpenAI models."""
 
 import unittest
 from unittest import mock
@@ -85,7 +85,7 @@ def mock_chat_completion_query_vision(messages, *, n=1, **kwargs):
   )
 
 
-class OpenaiTest(unittest.TestCase):
+class OpenAITest(unittest.TestCase):
   """Tests for OpenAI language model."""
 
   def test_model_id(self):
@@ -98,7 +98,7 @@ class OpenaiTest(unittest.TestCase):
     )
 
   def test_max_concurrency(self):
-    self.assertEqual(openai.Gpt35(api_key='test_key').max_concurrency, 8)
+    self.assertGreater(openai.Gpt35(api_key='test_key').max_concurrency, 0)
 
   def test_get_request_args(self):
     self.assertEqual(

langfun/core/structured/__init__.py

@@ -51,6 +51,7 @@ from langfun.core.structured.schema_generation import default_classgen_examples
 from langfun.core.structured.function_generation import function_gen
 
 from langfun.core.structured.mapping import Mapping
+from langfun.core.structured.mapping import MappingError
 from langfun.core.structured.mapping import MappingExample
 
 from langfun.core.structured.parsing import ParseStructure

langfun/core/structured/completion_test.py

@@ -17,7 +17,6 @@ import inspect
 import unittest
 
 import langfun.core as lf
-from langfun.core import coding
 from langfun.core import modalities
 from langfun.core.llms import fake
 from langfun.core.structured import completion
@@ -608,7 +607,7 @@ class CompleteStructureTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          coding.CodeError,
+          mapping.MappingError,
           'Expect .* but encountered .*',
       ):
         completion.complete(Activity.partial(), autofix=0)

langfun/core/structured/mapping.py

@@ -20,6 +20,43 @@ from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
+class MappingError(Exception):  # pylint: disable=g-bad-exception-name
+  """Mapping error."""
+
+  def __init__(self, lm_response: lf.Message, cause: Exception):
+    self._lm_response = lm_response
+    self._cause = cause
+
+  @property
+  def lm_response(self) -> lf.Message:
+    """Returns the LM response that failed to be mapped."""
+    return self._lm_response
+
+  @property
+  def cause(self) -> Exception:
+    """Returns the cause of the error."""
+    return self._cause
+
+  def __str__(self) -> str:
+    return self.format(include_lm_response=True)
+
+  def format(self, include_lm_response: bool = True) -> str:
+    """Formats the mapping error."""
+    r = io.StringIO()
+    error_message = str(self.cause).rstrip()
+    r.write(
+        lf.colored(
+            f'{self.cause.__class__.__name__}: {error_message}', 'magenta'
+        )
+    )
+    if include_lm_response:
+      r.write('\n\n')
+      r.write(lf.colored('[LM Response]', 'blue', styles=['bold']))
+      r.write('\n')
+      r.write(lf.colored(self.lm_response.text, 'blue'))
+    return r.getvalue()
+
+
 @pg.use_init_args(['input', 'output', 'schema', 'context'])
 class MappingExample(lf.NaturalLanguageFormattable, lf.Component):
   """Mapping example between text, schema and structured value."""
@@ -308,7 +345,7 @@ class Mapping(lf.LangFunc):
       lm_output.result = self.postprocess_result(self.parse_result(lm_output))
     except Exception as e:  # pylint: disable=broad-exception-caught
       if self.default == lf.RAISE_IF_HAS_ERROR:
-        raise e
+        raise MappingError(lm_output, e) from e
       lm_output.result = self.default
     return lm_output
 

langfun/core/structured/mapping_test.py

@@ -16,10 +16,27 @@
 import inspect
 import unittest
 
+import langfun.core as lf
 from langfun.core.structured import mapping
 import pyglove as pg
 
 
+class MappingErrorTest(unittest.TestCase):
+
+  def test_format(self):
+    error = mapping.MappingError(
+        lf.AIMessage('hi'), ValueError('Cannot parse message.')
+    )
+    self.assertEqual(
+        lf.text_formatting.decolored(str(error)),
+        'ValueError: Cannot parse message.\n\n[LM Response]\nhi',
+    )
+    self.assertEqual(
+        lf.text_formatting.decolored(error.format(include_lm_response=False)),
+        'ValueError: Cannot parse message.',
+    )
+
+
 class MappingExampleTest(unittest.TestCase):
 
   def test_basics(self):

langfun/core/structured/parsing_test.py

@@ -17,11 +17,9 @@ import inspect
 import unittest
 
 import langfun.core as lf
-from langfun.core import coding
 from langfun.core.llms import fake
 from langfun.core.structured import mapping
 from langfun.core.structured import parsing
-from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
@@ -255,7 +253,7 @@ class ParseStructurePythonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          coding.CodeError,
+          mapping.MappingError,
           'name .* is not defined',
       ):
         parsing.parse('three', int)
@@ -546,7 +544,7 @@ class ParseStructureJsonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          schema_lib.JsonError,
+          mapping.MappingError,
           'No JSON dict in the output',
       ):
         parsing.parse('three', int, protocol='json')

langfun/core/structured/prompting_test.py

@@ -17,12 +17,10 @@ import inspect
 import unittest
 
 import langfun.core as lf
-from langfun.core import coding
 from langfun.core import modalities
 from langfun.core.llms import fake
 from langfun.core.structured import mapping
 from langfun.core.structured import prompting
-from langfun.core.structured import schema as schema_lib
 import pyglove as pg
 
 
@@ -439,7 +437,7 @@ class QueryStructurePythonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          coding.CodeError,
+          mapping.MappingError,
           'name .* is not defined',
       ):
         prompting.query('Compute 1 + 2', int)
@@ -677,7 +675,7 @@ class QueryStructureJsonTest(unittest.TestCase):
         override_attrs=True,
     ):
       with self.assertRaisesRegex(
-          schema_lib.JsonError,
+          mapping.MappingError,
           'No JSON dict in the output',
       ):
         prompting.query('Compute 1 + 2', int, protocol='json')

langfun/core/structured/schema_generation_test.py

@@ -14,8 +14,8 @@
 import inspect
 import unittest
 
-import langfun.core.coding as lf_coding
 from langfun.core.llms import fake
+from langfun.core.structured import mapping
 from langfun.core.structured import schema_generation
 
 
@@ -92,7 +92,7 @@ class GenerateClassTest(unittest.TestCase):
     )
     self.assertIs(cls.__name__, 'B')
 
-    with self.assertRaises(lf_coding.CodeError):
+    with self.assertRaises(mapping.MappingError):
       schema_generation.generate_class(
           'Foo',
           'Generate a Foo class with a field pointing to another class A',