langfun 0.1.2.dev202412310804__py3-none-any.whl → 0.1.2.dev202501050804__py3-none-any.whl

langfun/core/llms/openai.py

@@ -32,6 +32,13 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     # o1 (preview) models.
     # Pricing in US dollars, from https://openai.com/api/pricing/
     # as of 2024-10-10.
+    'o1': pg.Dict(
+        in_service=True,
+        rpm=10000,
+        tpm=5000000,
+        cost_per_1k_input_tokens=0.015,
+        cost_per_1k_output_tokens=0.06,
+    ),
     'o1-preview': pg.Dict(
         in_service=True,
         rpm=10000,
@@ -255,25 +262,17 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     ),
     # GPT-3.5 models
     'text-davinci-003': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     'text-davinci-002': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     'code-davinci-002': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     # GPT-3 instruction-tuned models (Deprecated)
     'text-curie-001': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
+        in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM
     ),
     'text-babbage-001': pg.Dict(
         in_service=False,
@@ -290,32 +289,12 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         rpm=_DEFAULT_RPM,
         tpm=_DEFAULT_TPM,
     ),
-    'curie': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
-    'babbage': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
-    'ada': pg.Dict(
-        in_service=False,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
+    'curie': pg.Dict(in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'babbage': pg.Dict(in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'ada': pg.Dict(in_service=False, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
     # GPT-3 base models that are still in service.
-    'babbage-002': pg.Dict(
-        in_service=True,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
-    'davinci-002': pg.Dict(
-        in_service=True,
-        rpm=_DEFAULT_RPM,
-        tpm=_DEFAULT_TPM
-    ),
+    'babbage-002': pg.Dict(in_service=True, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
+    'davinci-002': pg.Dict(in_service=True, rpm=_DEFAULT_RPM, tpm=_DEFAULT_TPM),
 }
 
 
@@ -569,6 +548,13 @@ class OpenAI(rest.REST):
     )
 
 
+class GptO1(OpenAI):
+  """GPT-O1."""
+
+  model = 'o1'
+  multimodal = True
+
+
 class GptO1Preview(OpenAI):
   """GPT-O1."""
   model = 'o1-preview'

langfun/core/llms/vertexai.py

@@ -110,7 +110,13 @@ SUPPORTED_MODELS_AND_SETTINGS = {
     ),
     # TODO(sharatsharat): Update costs when published
     'gemini-2.0-flash-exp': pg.Dict(
-        rpm=20,
+        rpm=10,
+        cost_per_1k_input_chars=0.000,
+        cost_per_1k_output_chars=0.000,
+    ),
+    # TODO(yifenglu): Update costs when published
+    'gemini-2.0-flash-thinking-exp-1219': pg.Dict(
+        rpm=10,
         cost_per_1k_input_chars=0.000,
         cost_per_1k_output_chars=0.000,
     ),
@@ -415,6 +421,12 @@ class VertexAIGeminiFlash2_0Exp(VertexAIGemini2_0):  # pylint: disable=invalid-n
   model = 'gemini-2.0-flash-exp'
 
 
+class VertexAIGeminiFlash2_0ThinkingExp(VertexAIGemini2_0):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 2.0 Flash model."""
+
+  model = 'gemini-2.0-flash-thinking-exp-1219'
+
+
 class VertexAIGemini1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 model."""
 

{langfun-0.1.2.dev202412310804.dist-info → langfun-0.1.2.dev202501050804.dist-info}/METADATA

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

{langfun-0.1.2.dev202412310804.dist-info → langfun-0.1.2.dev202501050804.dist-info}/RECORD

@@ -1,5 +1,5 @@
 langfun/__init__.py,sha256=fhfPXpHN7GoGqixpFfqhQkYxFs_siP_LhbjZhd3lhio,2497
-langfun/core/__init__.py,sha256=xlvFTXc7IKUTs8aCFRFhzOLTmmeuhXgk9yx2InBLNiA,4937
+langfun/core/__init__.py,sha256=oboXhlpA93wk4p1K52-tx0pbBt49r6Y1aCmiavBV03Y,4830
 langfun/core/component.py,sha256=HVrEoTL1Y01iqOHC3FYdbAOnffqfHHtGJXoK1vkdEwo,11583
 langfun/core/component_test.py,sha256=sG-T2wpvBfHqWGZE7sc4NayJj2aj5QFBzSwFiwrGEIc,10376
 langfun/core/concurrent.py,sha256=QMNYhB_PyjvVJtabMokpzotZRYvyE9iYu2QsgwDk7M4,29552
@@ -19,8 +19,6 @@ langfun/core/modality.py,sha256=K8pUGuMpfWcOtVcXC_OqVjro1-RhHF6ddQni61DuYzM,4166
 langfun/core/modality_test.py,sha256=7SwhixFME2Q1sIXRgJx97EZFiIyC31A9NVr6_nDtFv4,2441
 langfun/core/natural_language.py,sha256=3ynSnaYQnjE60LIPK5fyMgdIjubnPYZwzGq4rWPeloE,1177
 langfun/core/natural_language_test.py,sha256=LHGU_1ytbkGuSZQFIFP7vP3dBlcY4-A12fT6dbjUA0E,1424
-langfun/core/repr_utils.py,sha256=yyc2xeK8sKQOOLcxwIV9E50Wf-B0TjXrnS07tSx_ea8,6491
-langfun/core/repr_utils_test.py,sha256=ULG7gvgoyqQFWi0m6g2-E0GorNEr1nnZ0J_sZVQKz80,3036
 langfun/core/sampling.py,sha256=SCnS5PFJWNVxSKvSkSCNRUmruvScun8UcNN4gafuXcw,5866
 langfun/core/sampling_test.py,sha256=U7PANpMsl9E_pa4_Y4FzesSjcwg-u-LKHGCWSgv-8FY,3663
 langfun/core/subscription.py,sha256=euawEuSZP-BHydaT-AQpfYFL0m5pWPGcW0upFhrojqc,10930
@@ -51,22 +49,22 @@ langfun/core/coding/python/permissions_test.py,sha256=w5EDb8QxpxgJyZkojyzVWQvDfg
 langfun/core/eval/__init__.py,sha256=OEXr1ZRuvLuhJJfuQ1ZWQ-SvYzjyrtiAAEogYaB7E6o,1933
 langfun/core/eval/base.py,sha256=ajHUS_BdzBIDKEwAmMzne0lJi8HcDhPmyw_plO0p7G0,75814
 langfun/core/eval/base_test.py,sha256=-LsIV9DXlDal0EnOlaWpibJvfef0NbxtZAm0OH_abAE,27189
-langfun/core/eval/matching.py,sha256=UnjdM_ebPqXKJamY4lvL3AYxrMIz3LqkjRTnHJ5xsYc,9349
+langfun/core/eval/matching.py,sha256=AVKkGoc-BaHEzgSBamaAk3194TgqckDe_dinpS6LrXI,9323
 langfun/core/eval/matching_test.py,sha256=QCoYEuf4b_1bkHqUCuRzKMbXHrV3AB2FCOBivo1stC4,5249
 langfun/core/eval/patching.py,sha256=R0s2eAd1m97exQt06dmUL0V_MBG0W2Hxg7fhNB7cXW0,3866
 langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrCG1L6w,4775
-langfun/core/eval/scoring.py,sha256=B69IsIxiPs1xZcOBFIhZF70YmDue2Siik-CPL2bh33s,6254
+langfun/core/eval/scoring.py,sha256=_DvnlgI1SdRVaOojao_AkV3pnenfCPOqyhvlg-Sw-5M,6322
 langfun/core/eval/scoring_test.py,sha256=O8olHbrUEg60gMxwOkWzKBJZpZoUlmVnBANX5Se2SXM,4546
 langfun/core/eval/v2/__init__.py,sha256=qoa6zKdFXOFyCX6vay6OdgPf1eUhYGoHYAxe35qECGk,1628
-langfun/core/eval/v2/checkpointing.py,sha256=5Koc9IM5OFAhLaO7sAZC979vEFxR8a_ssXcaqG_iVw0,10325
-langfun/core/eval/v2/checkpointing_test.py,sha256=i9qmEJk90kz6SB7OBGAHCogKVyazUZxAZlKXZt9hccI,4435
-langfun/core/eval/v2/eval_test_helper.py,sha256=pDpZTBnWRR5xjJv3Uy3NWEzArqlL8FTMOgeR4C53F5M,2348
-langfun/core/eval/v2/evaluation.py,sha256=NFBGAWw2BtW7H0zcoZhfWtz59Psra84eshJm73uAFwg,21807
-langfun/core/eval/v2/evaluation_test.py,sha256=GmV1TiqX1V15st2qpcGWooM5hudomQVjW5kajovGDvE,6231
-langfun/core/eval/v2/example.py,sha256=fURrvdNmMsVMqoEErcsmLmC6Xq3ny16dYsnLH8HVlcY,9626
-langfun/core/eval/v2/example_test.py,sha256=WcJmU7IQQXvjFia63mokySC4CqxzVL9Wso1sC5F0YK8,3032
-langfun/core/eval/v2/experiment.py,sha256=DuuqTJA0jmep-CM098MYutKkarLIcd1Do9P0W_HhdBg,29800
-langfun/core/eval/v2/experiment_test.py,sha256=zSMHYqC9cA0k61U71pCSYTAJ6yK2_b6Dml5btc-bKzQ,9133
+langfun/core/eval/v2/checkpointing.py,sha256=u-MrnwQbm0T-BDcn9pPXs_FeKPzMYm1-pVos0DiTqgM,11769
+langfun/core/eval/v2/checkpointing_test.py,sha256=R-R8SFzworuYnMmGGcvE9f4oiYkb8HMoZ0m4euF3pus,8466
+langfun/core/eval/v2/eval_test_helper.py,sha256=A9w0FbRKieB3yym8x34kY24FfptHrp7m6_z3i8HXSFI,3922
+langfun/core/eval/v2/evaluation.py,sha256=kARf0pG6SrpN__IMeFsw8DV_5mT_tl52pJrr6wIZdsw,22482
+langfun/core/eval/v2/evaluation_test.py,sha256=0l0DqJTF8PZGA2Q1OlaF4YIax4ZhSk0ewOCvuVW1XAk,6658
+langfun/core/eval/v2/example.py,sha256=4-LNr8Ke-fhaF6gyeXX4JMyw0s8YkVTC63pXZ-CXKrE,10144
+langfun/core/eval/v2/example_test.py,sha256=1DNm6EuyZOq827DKvf3oTRVFkMNM_qTnLUpvOjpgz5I,3419
+langfun/core/eval/v2/experiment.py,sha256=qYWx22KMfoUa4ieSq1bt7NE8L9dgoiJpuNOQGT1IBQw,32723
+langfun/core/eval/v2/experiment_test.py,sha256=CqpDsDai2DiIU-SzpVmqFzM_ZxxkVYKd0Gr1Uvcvkuw,13546
 langfun/core/eval/v2/metric_values.py,sha256=_B905bC-jxrYPLSEcP2M8MaHZOVMz_bVrUw8YC4arCE,4660
 langfun/core/eval/v2/metric_values_test.py,sha256=ab2oF_HsIwrSy459108ggyjgefHSPn8UVILR4dRwx14,2634
 langfun/core/eval/v2/metrics.py,sha256=bl8i6u-ZHRBz4hAc3LzsZ2Dc7ZRQcuTYeUhhH-GxfF0,10628
@@ -75,28 +73,28 @@ langfun/core/eval/v2/progress.py,sha256=azZgssQgNdv3IgjKEaQBuGI5ucFDNbdi02P4z_nQ
 langfun/core/eval/v2/progress_test.py,sha256=YU7VHzmy5knPZwj9vpBN3rQQH2tukj9eKHkuBCI62h8,2540
 langfun/core/eval/v2/progress_tracking.py,sha256=l9fEkz4oP5McpZzf72Ua7PYm3lAWtRru7gRWNf8H0ms,6083
 langfun/core/eval/v2/progress_tracking_test.py,sha256=fouMVJkFJqHjbhQJngGLGCmA9x3n0dU4USI2dY163mg,2291
-langfun/core/eval/v2/reporting.py,sha256=qFxoxqQhNGxh1APg6-JVDlIfniGdty2Rh8aXQMc1aak,6612
-langfun/core/eval/v2/reporting_test.py,sha256=4nobW6pcaatiZh8u4xciexciaiZNDlDoJci157Wp_RI,1492
-langfun/core/eval/v2/runners.py,sha256=t6_yHAJ4HWufK4wvh_OntKcok2KquA5ARIHIk1vvEwc,15870
+langfun/core/eval/v2/reporting.py,sha256=KF4pE2H1qj3mJcgkv_c5YYFjrU-uJCk_-fuu891Olzs,8061
+langfun/core/eval/v2/reporting_test.py,sha256=UmYSAQvD3AIXsSyWQ-WD2uLtEISYpmBeoKY5u5Qwc8E,5696
+langfun/core/eval/v2/runners.py,sha256=DKEmSlGXjOXKWFdBhTpLy7tMsBHZHd1Brl3hWIngsSQ,15931
 langfun/core/eval/v2/runners_test.py,sha256=A37fKK2MvAVTiShsg_laluJzJ9AuAQn52k7HPbfD0Ks,11666
-langfun/core/llms/__init__.py,sha256=lWXKjGHv66ShG7AE_Bc4QM7SDTxJdfoQMn3PF0lr0sU,6461
-langfun/core/llms/anthropic.py,sha256=afKZmdiLcosS_UEBlB8WKyf1K-zeXgwtPAx6ofg2Gww,13989
+langfun/core/llms/__init__.py,sha256=6mi0IKTNfq6kymZPlPGA2V7YF1xDLrBCPytojeFMMeA,6716
+langfun/core/llms/anthropic.py,sha256=a5MmnFsBA0CbfvwzXT1v_0fqLRMrhUNdh1tx6469PQ4,14357
 langfun/core/llms/anthropic_test.py,sha256=-2U4kc_pgBM7wqxu8RuxzyHPGww1EAWqKUvN4PW8Btw,8058
 langfun/core/llms/compositional.py,sha256=csW_FLlgL-tpeyCOTVvfUQkMa_zCN5Y2I-YbSNuK27U,2872
 langfun/core/llms/compositional_test.py,sha256=4eTnOer-DncRKGaIJW2ZQQMLnt5r2R0UIx_DYOvGAQo,2027
 langfun/core/llms/fake.py,sha256=gCHBYBLvBCsC78HI1hpoqXCS-p1FMTgY1P1qh_sGBPk,3070
 langfun/core/llms/fake_test.py,sha256=2h13qkwEz_JR0mtUDPxdAhQo7MueXaFSwsD2DIRDW9g,7653
-langfun/core/llms/google_genai.py,sha256=AAYOsSyeNIfHduIL9ZBzLhA8_acZUDMzHhS7AwUbOlM,11603
+langfun/core/llms/google_genai.py,sha256=3iAmLMcBXxkfiiI8BN0S6trKCfyfuajCIHIGpnCrtTg,11973
 langfun/core/llms/google_genai_test.py,sha256=zw14sgWmk0P_irHyb7vpPy1WAuLEE0PmyfiFElu03sA,7686
 langfun/core/llms/groq.py,sha256=dCnR3eAECEKuKKAAj-PDTs8NRHl6CQPdf57m1f6a79U,10312
 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=l49v6RubfInvV0iG114AymTKNogTX4u4N-UFCeSgIxw,20963
+langfun/core/llms/openai.py,sha256=dLDVBB47nJ30XCwjJpAZMc55ZlZXB__PcfcICCRNuXQ,20995
 langfun/core/llms/openai_test.py,sha256=kOWa1nf-nJvtYY10REUw5wojh3ZgfU8tRaCZ8wUgJbA,16623
 langfun/core/llms/rest.py,sha256=sWbYUV8S3SuOg9giq7xwD-xDRfaF7NP_ig7bI52-Rj4,3442
 langfun/core/llms/rest_test.py,sha256=NZ3Nf0XQVpT9kLP5cBVo_yBHLI7vWTYhWQxYEJVMGs4,3472
-langfun/core/llms/vertexai.py,sha256=oEd665IBwzCTlHuLEMrCdwgQzrFB5ERcnxw6nrYNSyk,14990
+langfun/core/llms/vertexai.py,sha256=EPPswgaTfPZQ_GGa_dWsqWPV9uRjCmIH2Iwgm1YXOqM,15377
 langfun/core/llms/vertexai_test.py,sha256=ffcA5yPecnQy_rhkuYAw_6o1iLW8AR8FgswmHt6aAys,6725
 langfun/core/llms/cache/__init__.py,sha256=QAo3InUMDM_YpteNnVCSejI4zOsnjSMWKJKzkb3VY64,993
 langfun/core/llms/cache/base.py,sha256=rt3zwmyw0y9jsSGW-ZbV1vAfLxQ7_3AVk0l2EySlse4,3918
@@ -148,8 +146,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=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202412310804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202412310804.dist-info/METADATA,sha256=NXUwqSAFDqcqo2k0EyFp4e4WFxfAyUyYL--hw9BBK10,8281
-langfun-0.1.2.dev202412310804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-langfun-0.1.2.dev202412310804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202412310804.dist-info/RECORD,,
+langfun-0.1.2.dev202501050804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202501050804.dist-info/METADATA,sha256=kNngwBqJlzR5sCW7VWtWMwZLA2uItkoRB8q1mXZe5ys,8281
+langfun-0.1.2.dev202501050804.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+langfun-0.1.2.dev202501050804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202501050804.dist-info/RECORD,,

langfun/core/repr_utils.py

@@ -1,204 +0,0 @@
-# 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.
-"""Helpers for implementing _repr_xxx_ methods."""
-
-import collections
-import contextlib
-import html
-import io
-from typing import Any, Callable, Iterator
-
-from langfun.core import component
-import pyglove as pg
-
-
-class Html(pg.Object):
-  """A HTML adapter for rendering."""
-  content: str
-
-  def _repr_html_(self) -> str:
-    return self.content
-
-
-@contextlib.contextmanager
-def share_parts() -> Iterator[dict[str, int]]:
-  """Context manager for defining the context (scope) of shared content.
-
-  Under the context manager, call to `lf.write_shared` with the same content
-  will be written only once. This is useful for writing shared content such as
-  shared style and script sections in the HTML.
-
-  Example:
-    ```
-    class Foo(pg.Object):
-      def _repr_html_(self) -> str:
-        s = io.StringIO()
-        lf.repr_utils.write_shared_part(s, '<style>..</style>')
-        lf.repr_utils.write_shared_part(s, '<script>..</script>')
-        return s.getvalue()
-
-    with lf.repr_utils.share_parts() as share_parts:
-      # The <style> and <script> section will be written only once.
-      lf.console.display(Foo())
-      lf.console.display(Foo())
-
-    # Assert that the shared content is attempted to be written twice.
-    assert share_parts['<style>..</style>'] == 2
-    ```
-
-  Yields:
-    A dictionary mapping the shared content to the number of times it is
-    attempted to be written.
-  """
-  context = component.context_value(
-      '__shared_parts__', collections.defaultdict(int)
-  )
-  with component.context(__shared_parts__=context):
-    try:
-      yield context
-    finally:
-      pass
-
-
-def write_maybe_shared(s: io.StringIO, content: str) -> bool:
-  """Writes a maybe shared part to an string stream.
-
-  Args:
-    s: The string stream to write to.
-    content: A maybe shared content to write.
-
-  Returns:
-    True if the content is written to the string. False if the content is
-    already written under the same share context.
-  """
-  context = component.context_value('__shared_parts__', None)
-  if context is None:
-    s.write(content)
-    return True
-  written = content in context
-  if not written:
-    s.write(content)
-  context[content] += 1
-  return not written
-
-
-def html_repr(
-    value: dict[str, Any],
-    item_color: Callable[
-        [str, str],
-        tuple[
-            str | None,   # Label text color
-            str | None,   # Label background color
-            str | None,   # Value text color
-            str | None,   # Value background color
-        ]
-    ] | None = None    # pylint: disable=bad-whitespace
-) -> str:
-  """Writes a list of key-value pairs to an string stream.
-
-  Args:
-    value: A value to be rendered in HTML.
-    item_color: A function that takes the key and value and returns a tuple
-      of four strings, the text color and background color of the label and
-      value respectively. If None, a default color scheme will be used.
-
-  Returns:
-    The HTML representation of the value.
-  """
-  s = io.StringIO()
-  s.write('<div style="padding-left: 20px; margin-top: 10px">')
-  s.write('<table style="border-top: 1px solid #EEEEEE;">')
-  item_color = item_color or (lambda k, v: (None, '#F1C40F', None, None))
-
-  def maybe_html_format(v: Any, root_indent: int) -> str | None:
-    del root_indent
-    if hasattr(v, '_repr_html_'):
-      return v._repr_html_()  # pylint: disable=protected-access
-    # Fall back to the default format.
-    return None
-
-  with (pg.str_format(custom_format=maybe_html_format),
-        pg.repr_format(custom_format=maybe_html_format)):
-    for k, v in pg.object_utils.flatten(value).items():
-      if isinstance(v, pg.Ref):
-        v = v.value
-      if hasattr(v, '_repr_html_'):
-        cs = v._repr_html_()  # pylint: disable=protected-access
-      else:
-        cs = html.escape(v) if isinstance(v, str) else escape_quoted(str(v))
-        cs = f'<span style="white-space: pre-wrap">{cs}</span>'
-
-      key_color, key_bg_color, value_color, value_bg_color = item_color(k, v)
-      key_span = html_round_text(
-          k,
-          text_color=key_color,
-          background_color=key_bg_color,
-          margin_bottom='0px'
-      )
-      value_color_style = f'color: {value_color};' if value_color else ''
-      value_bg_color_style = (
-          f'background-color: {value_bg_color};' if value_bg_color else ''
-      )
-      s.write(
-          '<tr>'
-          '<td style="padding: 5px; vertical-align: top; '
-          f'border-bottom: 1px solid #EEEEEE">{key_span}</td>'
-          '<td style="padding: 15px 5px 5px 5px; vertical-align: top; '
-          'border-bottom: 1px solid #EEEEEE;'
-          f'{value_color_style}{value_bg_color_style}">{cs}</td></tr>'
-      )
-    s.write('</table></div>')
-    return s.getvalue()
-
-
-def html_round_text(
-    text: str,
-    *,
-    text_color: str = 'black',
-    background_color: str = '#EEEEEE',
-    display: str = 'inline-block',
-    margin_top: str = '5px',
-    margin_bottom: str = '5px',
-    whitespace: str = 'pre-wrap') -> str:
-  """Renders a HTML span with rounded corners."""
-  color_style = f'color: {text_color};' if text_color else ''
-  bg_color_style = (
-      f'background-color: {background_color};' if background_color else ''
-  )
-  return (
-      f'<span style="{color_style}{bg_color_style}'
-      f'display:{display}; border-radius:10px; padding:5px; '
-      f'margin-top: {margin_top}; margin-bottom: {margin_bottom}; '
-      f'white-space: {whitespace}">{text}</span>'
-  )
-
-
-def escape_quoted(s: str):
-  """Escape quoted parts within a string."""
-  r = io.StringIO()
-  quote_char = None
-  quote_start = -1
-  for i, c in enumerate(s):
-    if c in ('\'', '"'):
-      if quote_char is None:
-        quote_char = c
-        quote_start = i
-      elif quote_char == c:
-        r.write(c)
-        r.write(html.escape(s[quote_start + 1:i]))
-        r.write(c)
-        quote_char = None
-    elif quote_char is None:
-      r.write(c)
-  return r.getvalue()

langfun/core/repr_utils_test.py

@@ -1,90 +0,0 @@
-# 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 langfun.core.repr_utils."""
-
-import io
-import unittest
-
-from langfun.core import repr_utils
-import pyglove as pg
-
-
-class SharingContentTest(unittest.TestCase):
-
-  def test_sharing(self):
-    s = io.StringIO()
-
-    self.assertTrue(repr_utils.write_maybe_shared(s, '<hr>'))
-    self.assertTrue(repr_utils.write_maybe_shared(s, '<hr>'))
-
-    with repr_utils.share_parts() as ctx1:
-      self.assertTrue(repr_utils.write_maybe_shared(s, '<style></style>'))
-      self.assertFalse(repr_utils.write_maybe_shared(s, '<style></style>'))
-
-      with repr_utils.share_parts() as ctx2:
-        self.assertIs(ctx2, ctx1)
-        self.assertFalse(repr_utils.write_maybe_shared(s, '<style></style>'))
-        self.assertTrue(repr_utils.write_maybe_shared(s, '<style>a</style>'))
-        self.assertFalse(repr_utils.write_maybe_shared(s, '<style>a</style>'))
-        self.assertTrue(repr_utils.write_maybe_shared(s, '<style>b</style>'))
-
-      with repr_utils.share_parts() as ctx3:
-        self.assertIs(ctx3, ctx1)
-        self.assertFalse(repr_utils.write_maybe_shared(s, '<style></style>'))
-        self.assertFalse(repr_utils.write_maybe_shared(s, '<style>a</style>'))
-        self.assertFalse(repr_utils.write_maybe_shared(s, '<style>a</style>'))
-        self.assertFalse(repr_utils.write_maybe_shared(s, '<style>b</style>'))
-
-    self.assertEqual(
-        s.getvalue(),
-        '<hr><hr><style></style><style>a</style><style>b</style>'
-    )
-    self.assertEqual(ctx1['<style></style>'], 4)
-    self.assertEqual(ctx1['<style>b</style>'], 2)
-    self.assertEqual(ctx1['<style>a</style>'], 4)
-
-  def test_escape_quoted(self):
-    self.assertEqual(
-        repr_utils.escape_quoted(str('<a>')), '<a>'
-    )
-    self.assertEqual(
-        repr_utils.escape_quoted('x=<a>, b="<a>"'),
-        'x=<a>, b="&lt;a&gt;"'
-    )
-
-  def test_html(self):
-    html = repr_utils.Html('<div>foo</div>')
-    self.assertEqual(html.content, '<div>foo</div>')
-    self.assertEqual(html._repr_html_(), '<div>foo</div>')
-
-  def test_html_repr(self):
-    class Foo(pg.Object):
-      x: int
-
-    class Bar(pg.Object):
-
-      def _repr_html_(self):
-        return '<bar>'
-
-    html = repr_utils.html_repr(
-        {'foo': pg.Ref(Foo(1)), 'bar': Bar(), 'baz': '<lf_image>'}
-    )
-    self.assertIn('foo</span>', html)
-    self.assertIn('<bar>', html)
-    self.assertIn('&lt;lf_image&gt;', html)
-    self.assertNotIn('Ref', html)
-
-
-if __name__ == '__main__':
-  unittest.main()