langfun 0.1.1.dev20240826__py3-none-any.whl → 0.1.1.dev202408282153__py3-none-any.whl

langfun/core/llms/google_genai.py

@@ -18,12 +18,29 @@ import functools
 import os
 from typing import Annotated, Any, Literal
 
-import google.generativeai as genai
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 import pyglove as pg
 
 
+try:
+  import google.generativeai as genai   # pylint: disable=g-import-not-at-top
+  BlobDict = genai.types.BlobDict
+  GenerativeModel = genai.GenerativeModel
+  Completion = genai.types.Completion
+  GenerationConfig = genai.GenerationConfig
+  GenerateContentResponse = genai.types.GenerateContentResponse
+  ChatResponse = genai.types.ChatResponse
+except ImportError:
+  genai = None
+  BlobDict = Any
+  GenerativeModel = Any
+  Completion = Any
+  GenerationConfig = Any
+  GenerateContentResponse = Any
+  ChatResponse = Any
+
+
 @lf.use_init_args(['model'])
 class GenAI(lf.LanguageModel):
   """Language models provided by Google GenAI."""
@@ -59,10 +76,16 @@ class GenAI(lf.LanguageModel):
 
   def _on_bound(self):
     super()._on_bound()
+    if genai is None:
+      raise RuntimeError(
+          'Please install "langfun[llm-google-genai]" to use '
+          'Google Generative AI models.'
+      )
     self.__dict__.pop('_api_initialized', None)
 
   @functools.cached_property
   def _api_initialized(self):
+    assert genai is not None
     api_key = self.api_key or os.environ.get('GOOGLE_API_KEY', None)
     if not api_key:
       raise ValueError(
@@ -78,6 +101,7 @@ class GenAI(lf.LanguageModel):
   @classmethod
   def dir(cls) -> list[str]:
     """Lists generative models."""
+    assert genai is not None
     return [
         m.name.lstrip('models/')
         for m in genai.list_models()
@@ -100,7 +124,7 @@ class GenAI(lf.LanguageModel):
 
   def _generation_config(self, options: lf.LMSamplingOptions) -> dict[str, Any]:
     """Creates generation config from langfun sampling options."""
-    return genai.GenerationConfig(
+    return GenerationConfig(
         candidate_count=options.n,
         temperature=options.temperature,
         top_p=options.top_p,
@@ -111,7 +135,7 @@ class GenAI(lf.LanguageModel):
 
   def _content_from_message(
       self, prompt: lf.Message
-  ) -> list[str | genai.types.BlobDict]:
+  ) -> list[str | BlobDict]:
     """Gets Evergreen formatted content from langfun message."""
     formatted = lf.UserMessage(prompt.text)
     formatted.source = prompt
@@ -131,7 +155,7 @@ class GenAI(lf.LanguageModel):
             if modality.is_text:
               chunk = modality.to_text()
             else:
-              chunk = genai.types.BlobDict(
+              chunk = BlobDict(
                   data=modality.to_bytes(),
                   mime_type=modality.mime_type
               )
@@ -143,7 +167,7 @@ class GenAI(lf.LanguageModel):
     return chunks
 
   def _response_to_result(
-      self, response: genai.types.GenerateContentResponse | pg.Dict
+      self, response: GenerateContentResponse | pg.Dict
   ) -> lf.LMSamplingResult:
     """Parses generative response into message."""
     samples = []
@@ -182,8 +206,8 @@ class _LegacyGenerativeModel(pg.Object):
 
   def generate_content(
       self,
-      input_content: list[str | genai.types.BlobDict],
-      generation_config: genai.GenerationConfig,
+      input_content: list[str | BlobDict],
+      generation_config: GenerationConfig,
   ) -> pg.Dict:
     """Generate content."""
     segments = []
@@ -195,7 +219,7 @@ class _LegacyGenerativeModel(pg.Object):
 
   @abc.abstractmethod
   def generate(
-      self, prompt: str, generation_config: genai.GenerationConfig) -> pg.Dict:
+      self, prompt: str, generation_config: GenerationConfig) -> pg.Dict:
     """Generate response based on prompt."""
 
 
@@ -203,9 +227,10 @@ class _LegacyCompletionModel(_LegacyGenerativeModel):
   """Legacy GenAI completion model."""
 
   def generate(
-      self, prompt: str, generation_config: genai.GenerationConfig
+      self, prompt: str, generation_config: GenerationConfig
   ) -> pg.Dict:
-    completion: genai.types.Completion = genai.generate_text(
+    assert genai is not None
+    completion: Completion = genai.generate_text(
         model=f'models/{self.model}',
         prompt=prompt,
         temperature=generation_config.temperature,
@@ -227,9 +252,10 @@ class _LegacyChatModel(_LegacyGenerativeModel):
   """Legacy GenAI chat model."""
 
   def generate(
-      self, prompt: str, generation_config: genai.GenerationConfig
+      self, prompt: str, generation_config: GenerationConfig
   ) -> pg.Dict:
-    response: genai.types.ChatResponse = genai.chat(
+    assert genai is not None
+    response: ChatResponse = genai.chat(
         model=f'models/{self.model}',
         messages=prompt,
         temperature=generation_config.temperature,
@@ -253,8 +279,9 @@ class _ModelHub:
 
   def get(
       self, model_name: str
-  ) -> genai.GenerativeModel | _LegacyGenerativeModel:
+  ) -> GenerativeModel | _LegacyGenerativeModel:
     """Gets a generative model by model id."""
+    assert genai is not None
     model = self._model_cache.get(model_name, None)
     if model is None:
       model_info = genai.get_model(f'models/{model_name}')

langfun/core/llms/openai.py

@@ -16,15 +16,31 @@
 import collections
 import functools
 import os
-from typing import Annotated, Any, cast
+from typing import Annotated, Any
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
-import openai
-from openai import error as openai_error
-from openai import openai_object
 import pyglove as pg
 
+try:
+  import openai  # pylint: disable=g-import-not-at-top
+
+  if hasattr(openai, 'error'):
+    # For lower versions.
+    ServiceUnavailableError = openai.error.ServiceUnavailableError
+    RateLimitError = openai.error.RateLimitError
+    APITimeoutError = (
+        openai.error.APIError,
+        '.*The server had an error processing your request'
+    )
+  else:
+    # For higher versions.
+    ServiceUnavailableError = getattr(openai, 'InternalServerError')
+    RateLimitError = getattr(openai, 'RateLimitError')
+    APITimeoutError = getattr(openai, 'APITimeoutError')
+except ImportError:
+  openai = None
+
 
 # From https://platform.openai.com/settings/organization/limits
 _DEFAULT_TPM = 250000
@@ -119,6 +135,10 @@ class OpenAI(lf.LanguageModel):
   def _on_bound(self):
     super()._on_bound()
     self.__dict__.pop('_api_initialized', None)
+    if openai is None:
+      raise RuntimeError(
+          'Please install "langfun[llm-openai]" to use OpenAI models.'
+      )
 
   @functools.cached_property
   def _api_initialized(self):
@@ -149,6 +169,7 @@ class OpenAI(lf.LanguageModel):
 
   @classmethod
   def dir(cls):
+    assert openai is not None
     return openai.Model.list()
 
   @property
@@ -195,11 +216,11 @@ class OpenAI(lf.LanguageModel):
   ) -> list[lf.LMSamplingResult]:
 
     def _open_ai_completion(prompts):
+      assert openai is not None
       response = openai.Completion.create(
           prompt=[p.text for p in prompts],
           **self._get_request_args(self.sampling_options),
       )
-      response = cast(openai_object.OpenAIObject, response)
       # Parse response.
       samples_by_index = collections.defaultdict(list)
       for choice in response.choices:
@@ -222,12 +243,9 @@ class OpenAI(lf.LanguageModel):
         _open_ai_completion,
         [prompts],
         retry_on_errors=(
-            openai_error.ServiceUnavailableError,
-            openai_error.RateLimitError,
-            # Handling transient OpenAI server error (code 500). Check out
-            # https://platform.openai.com/docs/guides/error-codes/error-codes
-            (openai_error.APIError,
-             '.*The server had an error processing your request'),
+            ServiceUnavailableError,
+            RateLimitError,
+            APITimeoutError,
         ),
     )[0]
 
@@ -292,10 +310,8 @@ class OpenAI(lf.LanguageModel):
         )
       messages.append(dict(role='user', content=_content_from_message(prompt)))
 
-      response = cast(
-          openai_object.OpenAIObject,
-          openai.ChatCompletion.create(messages=messages, **request_args)
-      )
+      assert openai is not None
+      response = openai.ChatCompletion.create(messages=messages, **request_args)
 
       samples = []
       for choice in response.choices:
@@ -330,8 +346,9 @@ class OpenAI(lf.LanguageModel):
         _open_ai_chat_completion,
         prompts,
         retry_on_errors=(
-            openai_error.ServiceUnavailableError,
-            openai_error.RateLimitError,
+            ServiceUnavailableError,
+            RateLimitError,
+            APITimeoutError
         ),
     )
 

langfun/core/llms/vertexai.py

@@ -17,11 +17,28 @@ import functools
 import os
 from typing import Annotated, Any
 
-from google.auth import credentials as credentials_lib
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
 import pyglove as pg
 
+try:
+  # pylint: disable=g-import-not-at-top
+  from google.auth import credentials as credentials_lib
+  import vertexai
+  from google.cloud.aiplatform import models as aiplatform_models
+  from vertexai import generative_models
+  from vertexai import language_models
+  # pylint: enable=g-import-not-at-top
+
+  Credentials = credentials_lib.Credentials
+except ImportError:
+  credentials_lib = None   # pylint: disable=invalid-name
+  vertexai = None
+  generative_models = None
+  language_models = None
+  aiplatform_models = None
+  Credentials = Any
+
 
 SUPPORTED_MODELS_AND_SETTINGS = {
     'gemini-1.5-pro-001': pg.Dict(api='gemini', rpm=50),
@@ -78,7 +95,7 @@ class VertexAI(lf.LanguageModel):
   ] = None
 
   credentials: Annotated[
-      credentials_lib.Credentials | None,
+      Credentials | None,
       (
           'Credentials to use. If None, the default credentials to the '
           'environment will be used.'
@@ -93,6 +110,10 @@ class VertexAI(lf.LanguageModel):
   def _on_bound(self):
     super()._on_bound()
     self.__dict__.pop('_api_initialized', None)
+    if generative_models is None:
+      raise RuntimeError(
+          'Please install "langfun[llm-google-vertex]" to use Vertex AI models.'
+      )
 
   @functools.cached_property
   def _api_initialized(self):
@@ -112,7 +133,7 @@ class VertexAI(lf.LanguageModel):
 
     credentials = self.credentials
     # Placeholder for Google-internal credentials.
-    import vertexai
+    assert vertexai is not None
     vertexai.init(project=project, location=location, credentials=credentials)
     return True
 
@@ -138,7 +159,7 @@ class VertexAI(lf.LanguageModel):
       self, prompt: lf.Message, options: lf.LMSamplingOptions
   ) -> Any:  # generative_models.GenerationConfig
     """Creates generation config from langfun sampling options."""
-    from vertexai import generative_models
+    assert generative_models is not None
     # Users could use `metadata_json_schema` to pass additional
     # request arguments.
     json_schema = prompt.metadata.get('json_schema')
@@ -169,7 +190,7 @@ class VertexAI(lf.LanguageModel):
       self, prompt: lf.Message
   ) -> list[str | Any]:
     """Gets generation input from langfun message."""
-    from vertexai import generative_models
+    assert generative_models is not None
     chunks = []
 
     for lf_chunk in prompt.chunk():
@@ -296,8 +317,8 @@ class VertexAI(lf.LanguageModel):
 
   def _sample_endpoint_model(self, prompt: lf.Message) -> lf.LMSamplingResult:
     """Samples a text generation model."""
-    from google.cloud.aiplatform import models
-    model = models.Endpoint(self.endpoint_name)
+    assert aiplatform_models is not None
+    model = aiplatform_models.Endpoint(self.endpoint_name)
     # TODO(chengrun): Add support for stop_sequences.
     predict_options = dict(
         temperature=self.sampling_options.temperature
@@ -337,7 +358,7 @@ class _ModelHub:
     """Gets a generative model by model id."""
     model = self._generative_model_cache.get(model_id, None)
     if model is None:
-      from vertexai import generative_models
+      assert generative_models is not None
       model = generative_models.GenerativeModel(model_id)
       self._generative_model_cache[model_id] = model
     return model
@@ -348,7 +369,7 @@ class _ModelHub:
     """Gets a text generation model by model id."""
     model = self._text_generation_model_cache.get(model_id, None)
     if model is None:
-      from vertexai import language_models
+      assert language_models is not None
       model = language_models.TextGenerationModel.from_pretrained(model_id)
       self._text_generation_model_cache[model_id] = model
     return model

langfun/core/message.py

@@ -278,6 +278,16 @@ class Message(natural_language.NaturalLanguageFormattable, pg.Object):
   # API for supporting modalities.
   #
 
+  @property
+  def text_with_modality_hash(self) -> str:
+    """Returns text with modality object placeheld by their 8-byte MD5 hash."""
+    parts = [self.text]
+    for name, modality_obj in self.referred_modalities().items():
+      parts.append(
+          f'<{name}>{modality_obj.hash}</{name}>'
+      )
+    return ''.join(parts)
+
   def get_modality(
       self, var_name: str, default: Any = None, from_message_chain: bool = True
   ) -> modality.Modality | None:

langfun/core/message_test.py

@@ -282,6 +282,20 @@ class MessageTest(unittest.TestCase):
         },
     )
 
+  def test_text_with_modality_hash(self):
+    m = message.UserMessage(
+        'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>',
+        img1=CustomModality('foo'),
+        x=dict(img2=CustomModality('bar')),
+    )
+    self.assertEqual(
+        m.text_with_modality_hash,
+        (
+            'hi, this is a <<[[img1]]>> and <<[[x.img2]]>>'
+            '<img1>acbd18db</img1><x.img2>37b51d19</x.img2>'
+        )
+    )
+
   def test_chunking(self):
     m = message.UserMessage(
         inspect.cleandoc("""

langfun/core/modalities/image.py

@@ -15,9 +15,21 @@
 
 import functools
 import io
+from typing import Any
 
 from langfun.core.modalities import mime
-from PIL import Image as pil_image
+
+try:
+  from PIL import Image as pil_image    # pylint: disable=g-import-not-at-top
+  PILImage = pil_image.Image
+  pil_open = pil_image.open
+except ImportError:
+  PILImage = Any
+
+  def pil_open(*unused_args, **unused_kwargs):
+    raise RuntimeError(
+        'Please install "langfun[mime-pil]" to enable PIL image support.'
+    )
 
 
 class Image(mime.Mime):
@@ -34,14 +46,14 @@ class Image(mime.Mime):
 
   @functools.cached_property
   def size(self) -> tuple[int, int]:
-    img = pil_image.open(io.BytesIO(self.to_bytes()))
+    img = pil_open(io.BytesIO(self.to_bytes()))
     return img.size
 
-  def to_pil_image(self) -> pil_image.Image:
-    return pil_image.open(io.BytesIO(self.to_bytes()))
+  def to_pil_image(self) -> PILImage:   # pytype: disable=invalid-annotation
+    return pil_open(io.BytesIO(self.to_bytes()))
 
   @classmethod
-  def from_pil_image(cls, img: pil_image.Image) -> 'Image':
+  def from_pil_image(cls, img: PILImage) -> 'Image':  # pytype: disable=invalid-annotation
     buf = io.BytesIO()
     img.save(buf, format='PNG')
     return cls.from_bytes(buf.getvalue())

langfun/core/modalities/mime.py

@@ -17,11 +17,20 @@ import base64
 import functools
 from typing import Annotated, Iterable, Type, Union
 import langfun.core as lf
-import magic
 import pyglove as pg
 import requests
 
 
+try:
+  import magic    # pylint: disable=g-import-not-at-top
+  from_buffer = magic.from_buffer
+except ImportError:
+  def from_buffer(*unused_args, **unused_kwargs):
+    raise RuntimeError(
+        'Please install "langfun[mime-auto]" to enable automatic MIME support.'
+    )
+
+
 class Mime(lf.Modality):
   """Base for MIME data."""
 
@@ -38,7 +47,7 @@ class Mime(lf.Modality):
   @functools.cached_property
   def mime_type(self) -> str:
     """Returns the MIME type."""
-    mime = magic.from_buffer((self.to_bytes()), mime=True)
+    mime = from_buffer((self.to_bytes()), mime=True)
     if (
         self.MIME_PREFIX
         and not mime.lower().startswith(self.MIME_PREFIX)
@@ -136,14 +145,14 @@ class Mime(lf.Modality):
   def from_uri(cls, uri: str, **kwargs) -> 'Mime':
     if cls is Mime:
       content = cls.download(uri)
-      mime = magic.from_buffer(content, mime=True).lower()
+      mime = from_buffer(content, mime=True).lower()
       return cls.class_from_mime_type(mime)(uri=uri, content=content, **kwargs)
     return cls(uri=uri, content=None, **kwargs)
 
   @classmethod
   def from_bytes(cls, content: bytes | str, **kwargs) -> 'Mime':
     if cls is Mime:
-      mime = magic.from_buffer(content, mime=True).lower()
+      mime = from_buffer(content, mime=True).lower()
       return cls.class_from_mime_type(mime)(content=content, **kwargs)
     return cls(content=content, **kwargs)
 

langfun/core/modalities/ms_office.py

@@ -30,10 +30,15 @@ class Xlsx(mime.Mime):
   )
 
   def to_html(self) -> str:
-    import pandas as pd  # pylint: disable=g-import-not-at-top
-
-    df = pd.read_excel(io.BytesIO(self.to_bytes()))
-    return df.to_html()
+    try:
+      import pandas as pd  # pylint: disable=g-import-not-at-top
+      import openpyxl      # pylint: disable=g-import-not-at-top, unused-import
+      df = pd.read_excel(io.BytesIO(self.to_bytes()))
+      return df.to_html()
+    except ImportError as e:
+      raise RuntimeError(
+          'Please install "langfun[mime-xlsx]" to enable XLSX support.'
+      ) from e
 
   def _repr_html_(self) -> str:
     return self.to_html()
@@ -58,10 +63,14 @@ class Docx(mime.Mime):
   )
 
   def to_xml(self) -> str:
-    import docx  # pylint: disable=g-import-not-at-top
-
-    doc = docx.Document(io.BytesIO(self.to_bytes()))
-    return str(doc.element.xml)
+    try:
+      import docx  # pylint: disable=g-import-not-at-top
+      doc = docx.Document(io.BytesIO(self.to_bytes()))
+      return str(doc.element.xml)
+    except ImportError as e:
+      raise RuntimeError(
+          'Please install "langfun[mime-docx]" to enable Docx support.'
+      ) from e
 
   def _repr_html_(self) -> str:
     return self.to_xml()

langfun/core/modality.py

@@ -14,6 +14,8 @@
 """Interface for modality (e.g. Image, Video, etc.)."""
 
 import abc
+import functools
+import hashlib
 from typing import Any, ContextManager
 from langfun.core import component
 import pyglove as pg
@@ -35,6 +37,11 @@ class Modality(component.Component):
   REF_START = '<<[['
   REF_END = ']]>>'
 
+  def _on_bound(self):
+    super()._on_bound()
+    # Invalidate cached hash if modality member is changed.
+    self.__dict__.pop('hash', None)
+
   def format(self, *args, **kwargs) -> str:
     if self.referred_name is None or not pg.object_utils.thread_local_get(
         _TLS_MODALITY_AS_REF, False
@@ -46,6 +53,11 @@ class Modality(component.Component):
   def to_bytes(self) -> bytes:
     """Returns content in bytes."""
 
+  @functools.cached_property
+  def hash(self) -> str:
+    """Returns a 8-byte MD5 hash as the identifier for this modality object."""
+    return hashlib.md5(self.to_bytes()).hexdigest()[:8]
+
   @classmethod
   def text_marker(cls, var_name: str) -> str:
     """Returns a marker in the text for this object."""

langfun/core/modality_test.py

@@ -32,6 +32,7 @@ class ModalityTest(unittest.TestCase):
     v = CustomModality('a')
     self.assertIsNone(v.referred_name)
     self.assertEqual(str(v), "CustomModality(\n  content = 'a'\n)")
+    self.assertEqual(v.hash, '0cc175b9')
 
     _ = pg.Dict(metadata=pg.Dict(x=pg.Dict(metadata=pg.Dict(y=v))))
     self.assertEqual(v.referred_name, 'x.metadata.y')

langfun/core/structured/mapping.py

@@ -352,6 +352,12 @@ class Mapping(lf.LangFunc):
       lm_output = self.postprocess_response(lm_output)
       lm_output.result = self.postprocess_result(self.parse_result(lm_output))
     except Exception as e:  # pylint: disable=broad-exception-caught
+      if (self.lm.cache is not None
+          and lm_output.lm_input.cache_seed is not None):
+        success = self.lm.cache.delete(
+            self.lm, lm_output.lm_input, lm_output.lm_input.cache_seed
+        )
+        assert success
       if self.default == lf.RAISE_IF_HAS_ERROR:
         raise MappingError(lm_output, e) from e
       lm_output.result = self.default

langfun/core/structured/prompting_test.py

@@ -19,6 +19,7 @@ import unittest
 import langfun.core as lf
 from langfun.core import modalities
 from langfun.core.llms import fake
+from langfun.core.llms.cache import in_memory
 from langfun.core.structured import mapping
 from langfun.core.structured import prompting
 import pyglove as pg
@@ -799,15 +800,18 @@ class QueryStructureJsonTest(unittest.TestCase):
       self.assertIsNone(r.result[0].hotel)
 
   def test_bad_transform(self):
-    with lf.context(
-        lm=fake.StaticSequence(['3']),
-        override_attrs=True,
-    ):
-      with self.assertRaisesRegex(
-          mapping.MappingError,
-          'No JSON dict in the output',
+    with in_memory.lm_cache() as cache:
+      with lf.context(
+          lm=fake.StaticSequence(['3']),
+          override_attrs=True,
       ):
-        prompting.query('Compute 1 + 2', int, protocol='json')
+        with self.assertRaisesRegex(
+            mapping.MappingError,
+            'No JSON dict in the output',
+        ):
+          prompting.query('Compute 1 + 2', int, protocol='json', cache_seed=1)
+      # Make sure bad mapping does not impact cache.
+      self.assertEqual(len(cache), 0)
 
   def test_query(self):
     lm = fake.StaticSequence(['{"result": 1}'])

langfun/core/text_formatting.py

@@ -16,7 +16,11 @@
 import io
 import re
 from typing import Any
-import termcolor
+
+try:
+  import termcolor      # pylint: disable=g-import-not-at-top
+except ImportError:
+  termcolor = None
 
 
 # Regular expression for ANSI color characters.
@@ -49,6 +53,8 @@ def colored(
   Returns:
     A string with ANSI color characters embracing the entire text.
   """
+  if not termcolor:
+    return text
   return termcolor.colored(
       text,
       color=color,

langfun/core/text_formatting_test.py

@@ -42,6 +42,24 @@ class TextFormattingTest(unittest.TestCase):
     )
     self.assertEqual(text_formatting.decolored(colored_text), original_text)
 
+  def test_colored_without_termcolor(self):
+    termcolor = text_formatting.termcolor
+    text_formatting.termcolor = None
+    original_text = inspect.cleandoc("""
+        Hi {{ foo }}
+        {# print x if x is present #}
+        {% if x %}
+        {{ x }}
+        {% endif %}
+        """)
+
+    colored_text = text_formatting.colored_template(
+        text_formatting.colored(original_text, color='blue')
+    )
+    self.assertEqual(colored_text, original_text)
+    self.assertEqual(text_formatting.decolored(colored_text), original_text)
+    text_formatting.termcolor = termcolor
+
 
 if __name__ == '__main__':
   unittest.main()