langfun 0.1.2.dev202503040804__py3-none-any.whl → 0.1.2.dev202503060804__py3-none-any.whl

langfun/core/concurrent_test.py

@@ -560,6 +560,7 @@ class ConcurrentMapTest(unittest.TestCase):
               fun, [1, 2, 3], timeout=1.5, max_workers=1, show_progress=True
           )
       ], key=lambda x: x[0])
+      string_io.flush()
     self.assertEqual(   # pylint: disable=g-generic-assert
         output,
         [
@@ -570,7 +571,6 @@ class ConcurrentMapTest(unittest.TestCase):
     )
     output = string_io.getvalue()
     self.assertIn('100%', output)
-    self.assertIn('TimeIt=foo (', output)
 
   def test_concurrent_map_with_showing_progress_and_status_fn(self):
     def fun(x):

langfun/core/eval/v2/checkpointing.py

@@ -361,7 +361,7 @@ class SequenceWriter:
         example,
         hide_default_values=True,
         save_ref_value=True,
-        exclude_input=True
+        exclude_input=False,
     )
     with self._lock:
       if self._sequence_writer is None:

langfun/core/llms/__init__.py

@@ -36,6 +36,7 @@ from langfun.core.llms.anthropic import Anthropic
 # Base models by serving platforms.
 from langfun.core.llms.vertexai import VertexAI
 from langfun.core.llms.groq import Groq
+from langfun.core.llms.azure_openai import AzureOpenAI
 
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI

langfun/core/llms/azure_openai.py

@@ -0,0 +1,90 @@
+# Copyright 2025 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.
+"""OpenAI models hosted on Azure."""
+import os
+from typing import Annotated, Any
+
+import langfun.core as lf
+from langfun.core.llms import openai
+import pyglove as pg
+
+
+@lf.use_init_args(['model', 'deployment_name'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class AzureOpenAI(openai.OpenAI):
+  """Azure OpenAI model service.
+
+  This service interacts with the Azure OpenAI API to generate chat completions.
+  It uses the deployment_name and API version to construct the endpoint, and
+  authenticates using an API key provided via parameter or the
+  AZURE_OPENAI_API_KEY environment variable.
+
+  Example:
+      lm = AzureOpenAI(
+          model='gpt-4o',
+          deployment_name='gpt-4o',
+          api_version='2024-08-01-preview',
+          azure_endpoint='https://trackname.openai.azure.com/',
+          api_key='token'
+      )
+      response = lf.query(prompt="what the capital of France", lm=lm)
+      print(response)
+  """
+
+  deployment_name: Annotated[
+      str,
+      'The name of the Azure OpenAI deployment.'
+  ] = 'gpt-4'
+
+  api_version: Annotated[
+      str,
+      'The API version for Azure OpenAI.'
+  ] = '2023-05-15'
+
+  azure_endpoint: Annotated[
+      str,
+      (
+          'The base URL for Azure OpenAI '
+          '(e.g. "https://<your-resource>.openai.azure.com/")'
+      )
+  ] = 'https://api.openai.azure.com/'
+
+  def _initialize(self):
+    # Authentication
+    self._api_key = self.api_key or os.environ.get('AZURE_OPENAI_API_KEY')
+    if not self._api_key:
+      raise ValueError(
+          'Azure OpenAI requires an API key. Please provide '
+          'via `api_key` argument or AZURE_OPENAI_API_KEY '
+          'environment variable.'
+      )
+
+    # Endpoint construction
+    self._api_endpoint = (
+        f"{self.azure_endpoint.rstrip('/')}/openai/deployments/"
+        f"{self.deployment_name}/chat/completions"
+        f"?api-version={self.api_version}"
+    )
+
+  @property
+  def api_endpoint(self) -> str:
+    return self._api_endpoint
+
+  @property
+  def headers(self) -> dict[str, Any]:
+    headers = super().headers
+    headers.update({
+        'api-key': self._api_key,
+    })
+    return headers

langfun/core/llms/azure_openai_test.py

@@ -0,0 +1,47 @@
+# Copyright 2025 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.
+import os
+import unittest
+from langfun.core.llms import azure_openai
+
+
+class AzureOpenAITest(unittest.TestCase):
+  def test_api_key_missing(self):
+    # Ensure that ValueError is raised when API key is not provided
+    os.environ.pop('AZURE_OPENAI_API_KEY', None)
+    with self.assertRaisesRegex(ValueError, 'Azure OpenAI requires an API key'):
+      azure = azure_openai.AzureOpenAI(model='gpt-4')
+      azure._initialize()
+
+  def test_api_endpoint(self):
+    # Test that api_endpoint is properly constructed
+    azure = azure_openai.AzureOpenAI(model='gpt-4', api_key='test_key')
+    azure._initialize()
+    expected = (
+        'https://api.openai.azure.com/openai/deployments/gpt-4/chat/'
+        'completions?api-version=2023-05-15'
+    )
+    self.assertEqual(azure.api_endpoint, expected)
+
+  def test_headers(self):
+    # Test that headers contain the api-key with correct value
+    azure = azure_openai.AzureOpenAI(model='gpt-4', api_key='test_key')
+    azure._initialize()
+    headers = azure.headers
+    self.assertIn('api-key', headers)
+    self.assertEqual(headers.get('api-key'), 'test_key')
+
+
+if __name__ == '__main__':
+  unittest.main()

{langfun-0.1.2.dev202503040804.dist-info → langfun-0.1.2.dev202503060804.dist-info}/METADATA

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

{langfun-0.1.2.dev202503040804.dist-info → langfun-0.1.2.dev202503060804.dist-info}/RECORD

@@ -3,7 +3,7 @@ langfun/core/__init__.py,sha256=S6YS15fLy5M3L3qaSa3XhDgku-klUm1LJxFPB-Wl-tY,4596
 langfun/core/component.py,sha256=g1kQM0bryYYYWVDrSMnHfc74wIBbpfe5_B3s-UIP5GE,3028
 langfun/core/component_test.py,sha256=0CxTgjAud3aj8wBauFhG2FHDqrxCTl4OI4gzQTad-40,9254
 langfun/core/concurrent.py,sha256=zY-pXqlGqss_GI20tM1gXvyW8QepVPUuFNmutcIdhbI,32760
-langfun/core/concurrent_test.py,sha256=rc5T-2giWgtbwNuN6gmei7Uwo66HsJeeRtXZCpya_QU,17590
+langfun/core/concurrent_test.py,sha256=roMFze0EKuyPbmG6DZzz8K8VGsZwWzc0F1uJZTFROC4,17572
 langfun/core/console.py,sha256=V_mOiFi9oGh8gLsUeR56pdFDkuvYOpvQt7DY1KUTWTA,2535
 langfun/core/console_test.py,sha256=pBOcuNMJdVELywvroptfcRtJMsegMm3wSlHAL2TdxVk,1679
 langfun/core/langfunc.py,sha256=G50YgoVZ0y1GFw2ev41MlOqr6qa8YakbvNC0h_E0PiA,11140
@@ -52,7 +52,7 @@ langfun/core/eval/patching_test.py,sha256=8kCd54Egjju22FMgtJuxEsrXkW8ifs-UUBHtrC
 langfun/core/eval/scoring.py,sha256=_DvnlgI1SdRVaOojao_AkV3pnenfCPOqyhvlg-Sw-5M,6322
 langfun/core/eval/scoring_test.py,sha256=adQEeuDux11dt9lkJIzLYNmqYSi9h-Mi2Cr0lUUTx9I,4546
 langfun/core/eval/v2/__init__.py,sha256=qoa6zKdFXOFyCX6vay6OdgPf1eUhYGoHYAxe35qECGk,1628
-langfun/core/eval/v2/checkpointing.py,sha256=u-MrnwQbm0T-BDcn9pPXs_FeKPzMYm1-pVos0DiTqgM,11769
+langfun/core/eval/v2/checkpointing.py,sha256=esGX-fB5qr_DEFoYOFapl-Ve7VEtX0_xTK45TYKC4xo,11771
 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=PgX1DoG3z-En7KVZe2hgSI7sMrLodvtT_UcUGMkQ8FQ,22562
@@ -73,9 +73,11 @@ langfun/core/eval/v2/reporting.py,sha256=7rL9LLmGYnQ5HIjqRqsOMkUlBl4BmFPEL6Vlofq
 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=0hJgKhVvwAw8Cgm4_yb26o_1Aw_evzxONcdlKgOmJqg,7983
+langfun/core/llms/__init__.py,sha256=F2nqEv9Cc-1QALrlPnIEp-TZiSVTi6U8WWb1YwDt1w0,8038
 langfun/core/llms/anthropic.py,sha256=7kY1-QFpzh2Oye8H1Mtko7QJWrnrhnhe1Wj16nOZXvk,23473
 langfun/core/llms/anthropic_test.py,sha256=SSK7OTx3gMYE1NMAi_PqQqeNsCkZAcVJvl_OCEOhyzk,7145
+langfun/core/llms/azure_openai.py,sha256=-KkSLaR54MlsIqz_XIwv0TnsBnvNTAxnjA2Q2O2u5KM,2733
+langfun/core/llms/azure_openai_test.py,sha256=lkMZkQdJBV97fTM4C4z8qNfvr6spgiN5G4hvVUIVr0M,1735
 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/deepseek.py,sha256=jvTxdXPr-vH6HNakn_Ootx1heDg8Fen2FUkUW36bpCs,5247
@@ -148,8 +150,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.dev202503040804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202503040804.dist-info/METADATA,sha256=1e8sjr4qC3ro4yDNCB36YSYkDwnLJeJ6twkUJev2KjQ,8172
-langfun-0.1.2.dev202503040804.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-langfun-0.1.2.dev202503040804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202503040804.dist-info/RECORD,,
+langfun-0.1.2.dev202503060804.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202503060804.dist-info/METADATA,sha256=EjrGcGHR07d5MlmqA2tOO1pGzISqvGXUHpaTCWHjNzM,8172
+langfun-0.1.2.dev202503060804.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+langfun-0.1.2.dev202503060804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202503060804.dist-info/RECORD,,