genesis-flow 1.0.1__py3-none-any.whl → 1.0.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. {genesis_flow-1.0.1.dist-info → genesis_flow-1.0.4.dist-info}/METADATA +32 -28
  2. {genesis_flow-1.0.1.dist-info → genesis_flow-1.0.4.dist-info}/RECORD +29 -26
  3. mlflow/azure/config.py +13 -7
  4. mlflow/azure/connection_factory.py +15 -2
  5. mlflow/data/dataset_source_registry.py +8 -0
  6. mlflow/gateway/providers/bedrock.py +298 -0
  7. mlflow/genai/datasets/databricks_evaluation_dataset_source.py +77 -0
  8. mlflow/genai/datasets/evaluation_dataset.py +8 -5
  9. mlflow/genai/scorers/base.py +22 -14
  10. mlflow/langchain/utils/chat.py +10 -0
  11. mlflow/models/container/__init__.py +2 -2
  12. mlflow/spark/__init__.py +1286 -0
  13. mlflow/store/artifact/azure_blob_artifact_repo.py +1 -1
  14. mlflow/store/artifact/azure_data_lake_artifact_repo.py +1 -1
  15. mlflow/store/artifact/gcs_artifact_repo.py +1 -1
  16. mlflow/store/artifact/local_artifact_repo.py +2 -1
  17. mlflow/store/artifact/s3_artifact_repo.py +173 -3
  18. mlflow/tracing/client.py +139 -49
  19. mlflow/tracing/export/mlflow_v3.py +8 -11
  20. mlflow/tracing/provider.py +5 -1
  21. mlflow/tracking/_model_registry/client.py +5 -1
  22. mlflow/tracking/_tracking_service/utils.py +17 -5
  23. mlflow/utils/file_utils.py +2 -1
  24. mlflow/utils/rest_utils.py +4 -0
  25. mlflow/version.py +2 -2
  26. {genesis_flow-1.0.1.dist-info → genesis_flow-1.0.4.dist-info}/WHEEL +0 -0
  27. {genesis_flow-1.0.1.dist-info → genesis_flow-1.0.4.dist-info}/entry_points.txt +0 -0
  28. {genesis_flow-1.0.1.dist-info → genesis_flow-1.0.4.dist-info}/licenses/LICENSE.txt +0 -0
  29. {genesis_flow-1.0.1.dist-info → genesis_flow-1.0.4.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: genesis-flow
3
- Version: 1.0.1
4
- Summary: Genesis-Flow: Lightweight, secure MLflow fork for Genesis platform
3
+ Version: 1.0.4
4
+ Summary: Genesis-Flow: MLflow v3.1.4 compatible fork for Genesis platform
5
5
  Maintainer-email: Databricks <mlflow-oss-maintainers@googlegroups.com>
6
6
  License: Copyright 2018 Databricks, Inc. All rights reserved.
7
7
 
@@ -268,6 +268,36 @@ Requires-Dist: databricks-sdk<1,>=0.20.0; extra == "databricks"
268
268
  Requires-Dist: azure-storage-file-datalake>12; extra == "databricks"
269
269
  Requires-Dist: boto3>1; extra == "databricks"
270
270
  Requires-Dist: botocore; extra == "databricks"
271
+ Requires-Dist: databricks-agents<2.0,>=1.0.0; extra == "databricks"
272
+ Provides-Extra: mlserver
273
+ Requires-Dist: mlserver!=1.3.1,<2.0.0,>=1.2.0; extra == "mlserver"
274
+ Requires-Dist: mlserver-mlflow!=1.3.1,<2.0.0,>=1.2.0; extra == "mlserver"
275
+ Provides-Extra: gateway
276
+ Requires-Dist: fastapi<1; extra == "gateway"
277
+ Requires-Dist: uvicorn[standard]<1; extra == "gateway"
278
+ Requires-Dist: watchfiles<2; extra == "gateway"
279
+ Requires-Dist: aiohttp<4; extra == "gateway"
280
+ Requires-Dist: boto3<2,>=1.28.56; extra == "gateway"
281
+ Requires-Dist: tiktoken<1; extra == "gateway"
282
+ Requires-Dist: slowapi<1,>=0.1.9; extra == "gateway"
283
+ Provides-Extra: genai
284
+ Requires-Dist: fastapi<1; extra == "genai"
285
+ Requires-Dist: uvicorn[standard]<1; extra == "genai"
286
+ Requires-Dist: watchfiles<2; extra == "genai"
287
+ Requires-Dist: aiohttp<4; extra == "genai"
288
+ Requires-Dist: boto3<2,>=1.28.56; extra == "genai"
289
+ Requires-Dist: tiktoken<1; extra == "genai"
290
+ Requires-Dist: slowapi<1,>=0.1.9; extra == "genai"
291
+ Provides-Extra: sqlserver
292
+ Requires-Dist: mlflow-dbstore; extra == "sqlserver"
293
+ Provides-Extra: aliyun-oss
294
+ Requires-Dist: aliyunstoreplugin; extra == "aliyun-oss"
295
+ Provides-Extra: xethub
296
+ Requires-Dist: mlflow-xethub; extra == "xethub"
297
+ Provides-Extra: jfrog
298
+ Requires-Dist: mlflow-jfrog-plugin; extra == "jfrog"
299
+ Provides-Extra: langchain
300
+ Requires-Dist: langchain<=0.3.25,>=0.1.0; extra == "langchain"
271
301
  Provides-Extra: auth
272
302
  Requires-Dist: Flask-WTF<2; extra == "auth"
273
303
  Provides-Extra: extras
@@ -761,32 +791,6 @@ pre-commit install
761
791
 
762
792
  # Run tests
763
793
  pytest tests/
764
-
765
- # 1. Install build tools
766
- pip install build twine
767
-
768
- # 2. Build the package
769
- python -m build
770
-
771
- # 3. Upload to PyPI
772
- python -m twine upload dist/*
773
-
774
- # Or upload to TestPyPI first
775
- python -m twine upload --repository testpypi dist/*
776
-
777
- Before uploading, make sure you have:
778
- - A PyPI account and API token
779
- - Configure your credentials in ~/.pypirc:
780
-
781
- [pypi]
782
- username = __token__
783
- password = pypi-<your-token>
784
-
785
- [testpypi]
786
- repository = https://test.pypi.org/legacy/
787
- username = __token__
788
- password = pypi-<your-token>
789
-
790
794
  ```
791
795
 
792
796
  ### Code Quality
@@ -1,4 +1,4 @@
1
- genesis_flow-1.0.1.dist-info/licenses/LICENSE.txt,sha256=Y5U1Xebzka__NZlqMPtBsYm0mRpMtUmTrONatpoL-ig,11382
1
+ genesis_flow-1.0.4.dist-info/licenses/LICENSE.txt,sha256=Y5U1Xebzka__NZlqMPtBsYm0mRpMtUmTrONatpoL-ig,11382
2
2
  mlflow/__init__.py,sha256=-_r__N5Afed81pLVtr2wKbHQIA0aj9u9n_7kWGxLWi4,11194
3
3
  mlflow/__main__.py,sha256=_PcdoxKehR_a2MI6GqBfzYzRCXZhVyDCSdbxDWVlWd4,39
4
4
  mlflow/cli.py,sha256=f1ObrWZ03HgRiRoVEE1Gffe-dGcSY7CxJyEFgb5VUMM,26137
@@ -11,7 +11,7 @@ mlflow/mismatch.py,sha256=G8YMN05L2x3k0NZpRtVNygEcGUXxc3WBzR-VdJw2iao,1050
11
11
  mlflow/ml_package_versions.py,sha256=jCYVUZQ82hMWeskEjGDBk-EX-EnThuMM_qBhrI7SsuI,10533
12
12
  mlflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
13
13
  mlflow/runs.py,sha256=qd8EtlSKj9499tSpzhtXp9Vm1iO5KK8P85KwZShyGH0,2519
14
- mlflow/version.py,sha256=ZRrRr7fEH9NnTqdz8RHUdIATR7mSDWNjh-8RWVilGDs,744
14
+ mlflow/version.py,sha256=Wrx1SBDS4H0U3Z7Yc728mt1X9vzv3FeCnpiA0ABdGu8,755
15
15
  mlflow/ag2/__init__.py,sha256=M298Ynewc2rbH6qNq_eGzuameavonOBLPy6ypgNdbQY,2225
16
16
  mlflow/ag2/ag2_logger.py,sha256=vnKExbqJerKBrcHyAatjPYe9buIjndu2530J4bhSo_8,11357
17
17
  mlflow/anthropic/__init__.py,sha256=qCo8IhTJHUomooKQrLKbYv-I8hHHS1iziTDanwS4wE4,1306
@@ -23,8 +23,8 @@ mlflow/autogen/chat.py,sha256=UwuhppwnlSs4fuv5D2uaugrl5lPpyvTkaHzWMnqR5E0,4497
23
23
  mlflow/azure/__init__.py,sha256=3B0PXyQj_3CAhvvMU-siUSDr9Qpf5q6s2PMIlyAwzVQ,737
24
24
  mlflow/azure/auth_handler.py,sha256=oIlzPdhuZcyg-sKowGLpinP0KJOFfsKrUnEKe0Mj5j8,9171
25
25
  mlflow/azure/client.py,sha256=dUaMFj04OLrlyvYSRTHXF4CaLHCGUdSzhjM5u8ucv8g,11509
26
- mlflow/azure/config.py,sha256=R86GuPP-yFT2UfKJr7VwhBNtx_o2QA-kvdlWQZwqzwM,4259
27
- mlflow/azure/connection_factory.py,sha256=VGrrtw3TUfpqGCiJGOclG4RcbXUsjBhE9_ybaA2eUjE,11676
26
+ mlflow/azure/config.py,sha256=0DwiEYN3TefmdaGO9KogNNflX_-0wOeHwKt2hMxKEdU,4491
27
+ mlflow/azure/connection_factory.py,sha256=1hcwPmc5xpP_MS1t3CIlzYfJ_mzEn2GBxnz5TJ1sHOY,12391
28
28
  mlflow/azure/exceptions.py,sha256=3jQo_kimK4g94LmSCtzfRQqEOl6bDBtwAQ6utXbgw8k,494
29
29
  mlflow/azure/stores.py,sha256=dlmdDVZRINgJw2wxzLhEsnAp1W76bR3NhQ0MftH9T74,12732
30
30
  mlflow/azure/utils.py,sha256=uU4OWU9ffz3KhLf0IgGqabZTTiAskLrNGVH5bIcwKUU,5563
@@ -43,7 +43,7 @@ mlflow/data/code_dataset_source.py,sha256=YtetLm5O-U-8AeRzt_CUcqfDvmyNOQZ0FlDG-k
43
43
  mlflow/data/dataset.py,sha256=ZcFl9TQVVzv0hJH_xn44qVC2H6WlcmoWD7_bDr6WUho,4119
44
44
  mlflow/data/dataset_registry.py,sha256=jm4L9P6GnKJ728O_cp6vbNewv5Kl-rp-Z8dXyKR0bw8,6491
45
45
  mlflow/data/dataset_source.py,sha256=5ta8hL8Wnalsmf0-qU65SB0apyFq8TItviOCfs5cwKY,3550
46
- mlflow/data/dataset_source_registry.py,sha256=QxQ1PiMvHqvrkDsccTydBYFCslYFuGSURLH0U_uitbQ,8342
46
+ mlflow/data/dataset_source_registry.py,sha256=GsWoxTbQ13YdPgMkd642mpoqAwKNuf6qGZ7fWA4ysHg,8576
47
47
  mlflow/data/delta_dataset_source.py,sha256=1chBnweHmtRkWPzi1_2itPFzDguYYaGVOnTiGww5Wu8,6027
48
48
  mlflow/data/digest_utils.py,sha256=O4NP6z2m604uPqbTuk-1ksCUhTJCNMOQb419FM4mzeY,2937
49
49
  mlflow/data/evaluation_dataset.py,sha256=FygCtfd_Ccfq2ohHRlq1lghdRGiy-YzF-0bMea3Smuk,20707
@@ -135,13 +135,15 @@ mlflow/evaluation/evaluation.py,sha256=YAO2LRKEt74W29cIZ2tSaEs_qymOpt1MkjrNCHuth
135
135
  mlflow/evaluation/evaluation_tag.py,sha256=Ynl0fxz49lFPTv7kIvi4CePKlg3yXG2Ci6nT9PmtUNI,1682
136
136
  mlflow/evaluation/fluent.py,sha256=Qj67ni9dupg40XsPDb9q3FduxlFWfITYayPBt5CExaI,1849
137
137
  mlflow/evaluation/utils.py,sha256=wrb1CtJhjVJB1ITRTlav2WisfqO8gPU8_nzx0CTQ8zo,6375
138
+ mlflow/gateway/providers/bedrock.py,sha256=UVsTKNkSaJsCTjPYQmsSNwG1CPkM4trTwZIsF-AINZw,10475
138
139
  mlflow/gemini/__init__.py,sha256=-mPzQ9hFBLsNMkP3RjOp7cQAzyhHl2hcqVNYle46kKM,2481
139
140
  mlflow/gemini/autolog.py,sha256=DPFsW22c8yb8Z_Iv5wuTleyOrnfZQbmMWR34FN6PzcE,6585
140
141
  mlflow/gemini/chat.py,sha256=BqPjA2BAplw5JSqKJ87pGZC0Rasij28Ae57uGU6YnGI,9475
141
142
  mlflow/genai/__init__.py,sha256=Lc_eDeUQpIlGGGZbCs_Kj3dpn9NHm6QydFV3XsFJwE4,1612
142
143
  mlflow/genai/scheduled_scorers.py,sha256=lLd9LVTeC3AfUqVo_52t_8FIXHlI4ubd1q-qjTtLE00,18097
143
144
  mlflow/genai/datasets/__init__.py,sha256=UXgb_Ces-hYVvv7WwWw42Kbz3_RN9cW8GUUMiTwpq9A,2095
144
- mlflow/genai/datasets/evaluation_dataset.py,sha256=w1Bfp0VxMZpAHZoh_RsQA79PI79oeY9OznU7Dfgqh2I,4705
145
+ mlflow/genai/datasets/databricks_evaluation_dataset_source.py,sha256=27P-CGE1lEtsRAJisaywliGWAzq9BBGjUhN3QpXe-xI,2405
146
+ mlflow/genai/datasets/evaluation_dataset.py,sha256=1EzhmMupmlRkyjVuOPIE_axJrFAAVhPT_U2WEfGfnd0,4772
145
147
  mlflow/genai/evaluation/__init__.py,sha256=aHMqUtH-4Bvgae2wlQMXL5cLYjwtI3U7KhTZ_v_dGHE,106
146
148
  mlflow/genai/evaluation/base.py,sha256=Eu8yqoEapy3qxAe3AzCBREMKsH7RU9B0UeSMEXpjChA,15974
147
149
  mlflow/genai/evaluation/constant.py,sha256=CKhzmcf67O9ZJjrEIIKUp74-Lc3IlXNnBcorGcCotCc,736
@@ -164,7 +166,7 @@ mlflow/genai/optimize/optimizers/utils/dspy_mipro_callback.py,sha256=RsRnAuEK5_5
164
166
  mlflow/genai/optimize/optimizers/utils/dspy_mipro_utils.py,sha256=VHzDqQU5VUndGBCh1kA9cZpZpwLmkkqrnna5QCeUqD0,500
165
167
  mlflow/genai/prompts/__init__.py,sha256=EGzQX6JVnxX9Ie8gqJnI9CFjQJYHJp1ltviQhM9vO44,7165
166
168
  mlflow/genai/scorers/__init__.py,sha256=lU2WYlye6rYi3xMV_5V204QeZv7htvJmAt_Khcq2EGI,553
167
- mlflow/genai/scorers/base.py,sha256=FkzT7wLZzQOZMDCCGgtX__DH8B_zaQKhiHhIL6eisfk,19861
169
+ mlflow/genai/scorers/base.py,sha256=NxsOIigau84kUORAQ7ebv9WnIK5kARsSa4WKdW5TyxI,20208
168
170
  mlflow/genai/scorers/builtin_scorers.py,sha256=WB3P58wbPk8XM2UHwUxw2tXKWllv5Nr5KdUxiKENn-o,27727
169
171
  mlflow/genai/scorers/scorer_utils.py,sha256=NTUUVkikAIQO52wOLwwSCsM8JBiuz1ah7K7sbJmjUEE,4849
170
172
  mlflow/genai/scorers/validation.py,sha256=28tIEy2j2NxUx-ooWR2Go4_gk3akGVB3hq1RTaOcaf0,6390
@@ -186,7 +188,7 @@ mlflow/langchain/model.py,sha256=RF8Vy61hVem4-ULsHRVfEzzLV6BFgcsNnEw1uMaO6G8,391
186
188
  mlflow/langchain/output_parsers.py,sha256=U7sYzic0ZBCoU-NCrMZWAePprn5PG_ZldfQEvRScyrQ,4957
187
189
  mlflow/langchain/retriever_chain.py,sha256=jjLhPLNn7kSEJNTCGKNDyF7HbiVAQ3TY6fbunjJYc_Q,5272
188
190
  mlflow/langchain/runnables.py,sha256=Hn2dcDtFAykLdXB6Q-6wRoqLIg1rO3ztBwhu2_iu3wQ,19770
189
- mlflow/langchain/utils/chat.py,sha256=HUzYQ5maPaARIVNqhDoWhR8vNDwCFpg-H407GCGGESA,14195
191
+ mlflow/langchain/utils/chat.py,sha256=1UfQ0JazRwxFQUYbzsliCWshIO8P-09_R7i95BuE8Ss,14563
190
192
  mlflow/langchain/utils/logging.py,sha256=nzGjLVr-0JgRnVNTpmwz0oNGS4E5_GtF2sULxndOX3U,23411
191
193
  mlflow/langchain/utils/serialization.py,sha256=VZbuOS2USnzBIR3Mvm4xsL7v-g1L9W9PftEBn6pBja4,1233
192
194
  mlflow/legacy_databricks_cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -232,7 +234,7 @@ mlflow/models/resources.py,sha256=blqbEthnPNqFwUdtRYdyDOZ2nR6Q-3R6Jwf9eotegQw,11
232
234
  mlflow/models/signature.py,sha256=A-odLBU1TI1u-1GQrru4pJ0HFIGiCdmad3_Y9vCdz7g,25978
233
235
  mlflow/models/utils.py,sha256=7pr3WX8uNvaX36GXhM32DJjk-2-q6_7vDbE5jzNNYpQ,82897
234
236
  mlflow/models/wheeled_model.py,sha256=U8zqEZtnFjwNuji_6DBCaL3sOF3Hczm_PJSUGwGOGI0,11886
235
- mlflow/models/container/__init__.py,sha256=P-uQFU5TQb3PZONbr_gGuweGwcXFBbav2j2Kllem56Y,10636
237
+ mlflow/models/container/__init__.py,sha256=Y1HlcvGwOkYTQbjIMIJLIUd9gtR6Q5jSbhoiYMgoWLw,10636
236
238
  mlflow/models/container/scoring_server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
237
239
  mlflow/models/container/scoring_server/nginx.conf,sha256=XwBf_NWOylw3MulHAZgyVMjorgoLt0EkRtf4lPJYj80,710
238
240
  mlflow/models/evaluation/__init__.py,sha256=MbagaX5WlRiWeysM1RO0q0SEe0nRhBZG2tyc7vo1VP8,528
@@ -367,6 +369,7 @@ mlflow/sklearn/utils.py,sha256=CM9z-UD65RBN4sIdEoJWhazkdmE9Wa9InjsU61y1XfE,38912
367
369
  mlflow/smolagents/__init__.py,sha256=V2sfkuWxQ090CaIjwBh7icD0bSgYP3jKD5MRzy__uW0,2202
368
370
  mlflow/smolagents/autolog.py,sha256=5GJpqmXs8dd2tf4jkL6fk7Fxb0TsxAMykB865M-BpK8,4770
369
371
  mlflow/smolagents/chat.py,sha256=tsrS9dm-iIigSbZQ3A6fA_vUXlnr77hMAI65PNY9Ono,980
372
+ mlflow/spark/__init__.py,sha256=-h5YdE-BVQPW9ke9Lfu33W9qlNqrn_FXPHkNsVVL8ic,54655
370
373
  mlflow/store/__init__.py,sha256=9uQ8sjjcbpTkeNTV1RjPFLlW8b6oDaeae1nDwbLVtNs,281
371
374
  mlflow/store/_unity_catalog/__init__.py,sha256=idiiLNYc7lzfOaPwY_QwTfvH1A6XMM0pkPzi-voVUsw,63
372
375
  mlflow/store/_unity_catalog/lineage/__init__.py,sha256=X-C_OnRoaueY0AoIty-WokUpnoRgoDGxhZBkDzYkquo,62
@@ -379,8 +382,8 @@ mlflow/store/_unity_catalog/registry/utils.py,sha256=aOfKcdh8liGODGzRTMs6R_ddItU
379
382
  mlflow/store/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
380
383
  mlflow/store/artifact/artifact_repo.py,sha256=2FBg7jEIpXLYF8vGDwyuV6VyyHJsDZHD4P2Hzi8KH8Y,18654
381
384
  mlflow/store/artifact/artifact_repository_registry.py,sha256=dMUX-uZLlNDFRaASHnMnEcHQpk5UvtSWSBNRQId6Ldk,7099
382
- mlflow/store/artifact/azure_blob_artifact_repo.py,sha256=Ud-ccqA1wTUdn-q5y-mGg_Xd2lYu9hRafKLpfPOJCJU,12417
383
- mlflow/store/artifact/azure_data_lake_artifact_repo.py,sha256=pwrm5iLeIe1uA6jjB-6ZKsob4Oj07PpHoMTezK5a4Ag,12138
385
+ mlflow/store/artifact/azure_blob_artifact_repo.py,sha256=kB-GgvyaKIHKL4CirmdRzw3tX-uc5OYoNbsxUtSeyO4,12417
386
+ mlflow/store/artifact/azure_data_lake_artifact_repo.py,sha256=9jG3rYUR4FSOqpLlo7xcBMKha3pTlTsvt_lDNUUREg0,12138
384
387
  mlflow/store/artifact/cli.py,sha256=TsUlxBTxx4Iyudyd-zZmhbsHL9zVyjAyJMnqqYzss-U,5375
385
388
  mlflow/store/artifact/cloud_artifact_repo.py,sha256=qQRuWgrKU7USTqoG7Urf6g4OEaJaNraEuwFrt_tjfPM,13332
386
389
  mlflow/store/artifact/databricks_artifact_repo.py,sha256=tGlaVlGcBvppDIGr-no1X7kSdzgcwtmiOJcO8jFP9t0,30377
@@ -391,17 +394,17 @@ mlflow/store/artifact/databricks_sdk_artifact_repo.py,sha256=UytO9rYFuQk1c_nWBFE
391
394
  mlflow/store/artifact/databricks_sdk_models_artifact_repo.py,sha256=izMSMpimTMFUWzVlmrWxWWs74B0kfeCYurT_fwNhP_4,3458
392
395
  mlflow/store/artifact/dbfs_artifact_repo.py,sha256=Z449LgfNJ1UverObdceuZF6FIBUOIG1Lb9545u6W1W4,10758
393
396
  mlflow/store/artifact/ftp_artifact_repo.py,sha256=A5aCDhRpdpbbxXGkFTHJgYpCMS7R7KDJ6CWxRX0jYhI,5291
394
- mlflow/store/artifact/gcs_artifact_repo.py,sha256=F77lALp9SDh3X2NwcKhEWYoGyQ2o9m1bVnMySTmu-QM,12098
397
+ mlflow/store/artifact/gcs_artifact_repo.py,sha256=v_msi2tLs76raxOGoBnEUsU0uf5BbgFEXr9ifY3Z1Mc,12098
395
398
  mlflow/store/artifact/hdfs_artifact_repo.py,sha256=1YFjdF51av_Zm0B_unfSlOme15hsULWjL9bvtou5nGg,7956
396
399
  mlflow/store/artifact/http_artifact_repo.py,sha256=FSQ-IAoiMVc0K3mgHMz6ALVgqFRXfXXU43IBKxAdps4,9088
397
- mlflow/store/artifact/local_artifact_repo.py,sha256=LPVNoCyS-QHdCY4F0RsuGPsCD7hClv2ZYNJ-azaGoIY,6083
400
+ mlflow/store/artifact/local_artifact_repo.py,sha256=u8GNoZZhtau2khzgjNKOLKRtdnDLl7dsvguZPTK8Qiw,6126
398
401
  mlflow/store/artifact/mlflow_artifacts_repo.py,sha256=vHF4f-IB4xB7arcGHZXy5C4SqUrjzNNX30OLH0bQwvs,3647
399
402
  mlflow/store/artifact/models_artifact_repo.py,sha256=TXyj7uMBk_e8IMlu79xaQo3jttU_FME8kXbeGW2jfvc,11093
400
403
  mlflow/store/artifact/optimized_s3_artifact_repo.py,sha256=QMNvnvcPmYwSM_PzM1KWmvuElQRwqg-nV1QFaErGs2g,15510
401
404
  mlflow/store/artifact/presigned_url_artifact_repo.py,sha256=p9I3V3hGn73Em6hRI6PflN9DWVDWAeL-X7CDa-XIfYg,7378
402
405
  mlflow/store/artifact/r2_artifact_repo.py,sha256=IQNMA5MUffsMOSxYyiF16lGnxWviU7br__utXjCmuJ8,2713
403
406
  mlflow/store/artifact/runs_artifact_repo.py,sha256=GPhn92ODs-POctsY5jxh5A18V5FXbldMEHxzWxN2NbE,11450
404
- mlflow/store/artifact/s3_artifact_repo.py,sha256=F7pWPTu-cm0c1v0-Id7WMLhnuYLZdata6tOR9soPwWE,12878
407
+ mlflow/store/artifact/s3_artifact_repo.py,sha256=_4cyQdvY15LQWkcWDiaLZNO2thJK3QJWyoXhACoch6c,21253
405
408
  mlflow/store/artifact/sftp_artifact_repo.py,sha256=1Sn0w_VNEepHka1YpRao7Cq4u-SRKxfu7AZ3jv8PK-U,5508
406
409
  mlflow/store/artifact/uc_volume_artifact_repo.py,sha256=6gCjr_usIAXBi-Cd6-TF1XqMJEEnW9mnu4IXJXGeXWA,3457
407
410
  mlflow/store/artifact/unity_catalog_models_artifact_repo.py,sha256=BGk5QEYwoNWSSGa3dXNnu7lqKz9LBy1LOEzeopDKL5Y,7534
@@ -479,18 +482,18 @@ mlflow/system_metrics/metrics/network_monitor.py,sha256=przQFWg7P8iY7pZHpihj97Ip
479
482
  mlflow/system_metrics/metrics/rocm_monitor.py,sha256=n-UvfDCLjfXKrwBLqVr2WR-BCbqaLsspZAKjpPIm6Pw,4426
480
483
  mlflow/tracing/__init__.py,sha256=b_ZzQtaP5ctjnIeIFUG7OJ1iK7_ULePJOgFK5A4aKTk,499
481
484
  mlflow/tracing/assessment.py,sha256=USb21LblbH4p2bqe6KnNWHtaBxINXOgQuBNDrI2hf74,12236
482
- mlflow/tracing/client.py,sha256=6-dSnFNePnEAjvwN1meciku8e6tAJnalX3MWpBo3-Ns,21610
485
+ mlflow/tracing/client.py,sha256=vrwje1RuXifyM8fbFwqloA_3_zpSlLPLY-pLK9CBMEw,25075
483
486
  mlflow/tracing/config.py,sha256=JNuIFBdetSctdN0zLRqDosAR-g86Ux5P2MIQYjSkx8g,4507
484
487
  mlflow/tracing/constant.py,sha256=_RnKWkv0RBr-ik3GHJpX7GxN0pPcz-HC1iuednnVPeo,4120
485
488
  mlflow/tracing/destination.py,sha256=ZtHYxGx9oLzSdpLqDpMYmpL9tClcSqVAD8fgYII5kto,2757
486
489
  mlflow/tracing/fluent.py,sha256=HNA-cmm7WeDmLUxOBNXxb9Bcqt7jwoq5W48YanU785k,53397
487
- mlflow/tracing/provider.py,sha256=db0_8yu3YoBTMZiyFbUbVBdQrTX2pGtuvA3tAWmNcGc,16677
490
+ mlflow/tracing/provider.py,sha256=L-rLpRB4YZVrjr_u8Y57iIk_g1b-RWR-Y6KHWAy0Zuk,16977
488
491
  mlflow/tracing/trace_manager.py,sha256=sjMg-V4kaON-xz5jl5cPF53H4ZNMqW1QSMz0_1WI2XM,7281
489
492
  mlflow/tracing/display/__init__.py,sha256=XoiwHbWNbQVuDklpcm-X3cG8pgdg8vRZYtjYC2os4m0,1265
490
493
  mlflow/tracing/display/display_handler.py,sha256=ui39ZbJ5PthiEYfDTYlPP44UpWaMmL5GhcBsqMcjuDA,6282
491
494
  mlflow/tracing/export/async_export_queue.py,sha256=2Slz5Ab_6FM8GcJRhZ1G7jqRGnRlIGYziGzWDT0NOcU,6766
492
495
  mlflow/tracing/export/inference_table.py,sha256=nI2MjgaJQbKhwFKtMGnpcupPLWLAuCJWilY5j-_Asq0,6044
493
- mlflow/tracing/export/mlflow_v3.py,sha256=PaRtWX5BqCoxsQTcFK2lIwYbwFr4J_ynzSnBHSCmUm4,5898
496
+ mlflow/tracing/export/mlflow_v3.py,sha256=8LaWwy2vc8Hxu1ttECJJvFhAB4TceBOg7XFsWiXTYvQ,5797
494
497
  mlflow/tracing/export/utils.py,sha256=C9xB1wAqopblJ1MSoUZVHHYnS2o_MhACM_j3Qg01ZZs,2201
495
498
  mlflow/tracing/processor/base_mlflow.py,sha256=vPpcFaQQOsA_9FUKqexPDKRCQluGsv9sQHSKXZLdkiQ,8389
496
499
  mlflow/tracing/processor/inference_table.py,sha256=VmzCWAOQyFHrWmZVfxO--CsYg7czaAsW205mWD5wuCw,7505
@@ -517,14 +520,14 @@ mlflow/tracking/metric_value_conversion_utils.py,sha256=AXxTd9Cot55d9bs4LUQvoeiF
517
520
  mlflow/tracking/multimedia.py,sha256=vm11i1lxw-gXh8j5kxmdXOFSwQrj8JmdV0AjPiw79Y4,6279
518
521
  mlflow/tracking/registry.py,sha256=I1Dzf3mXyhN0CDcOD3J_Y8bHuWv9fXmLsjABWPQhuNk,3524
519
522
  mlflow/tracking/_model_registry/__init__.py,sha256=_HtGwD9WY1x2-02TgEftcWrKKqBnJNjB9YdRc6sqLe4,41
520
- mlflow/tracking/_model_registry/client.py,sha256=MeyyyxCndQNRdRaAj-AKsp5mny6kE10vkB2OT9Z9r2M,27122
523
+ mlflow/tracking/_model_registry/client.py,sha256=6URbRwOCiNSKl6_LO1so6tStMTOHCTVhvYEX77xfiOA,27371
521
524
  mlflow/tracking/_model_registry/fluent.py,sha256=DYXqDICZvjpiSLxcb2LqOVp165phGHr__nOHQ3kCA6M,32343
522
525
  mlflow/tracking/_model_registry/registry.py,sha256=szFyA0NMxtsD-aAGVDmag0yxXIhMUeUttf_DcEGmuu4,3251
523
526
  mlflow/tracking/_model_registry/utils.py,sha256=0oyWU_iKleLVeEaWLb-glhkJlArKX0U3kj5iVWv3zdI,9377
524
527
  mlflow/tracking/_tracking_service/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
525
528
  mlflow/tracking/_tracking_service/client.py,sha256=wkFfP8r8vDA0YPGdE8q8DHn0LY-BeeWgxqWBnAqi_ug,36067
526
529
  mlflow/tracking/_tracking_service/registry.py,sha256=o_mi4_MN8Cp3WSnqVThZwABAsJJVBneZmxyqqQEQAeI,2419
527
- mlflow/tracking/_tracking_service/utils.py,sha256=ryY_fChkUrrBLMjGsjZc1IJ7SrKNWYPBDlm3VWHAZQs,9549
530
+ mlflow/tracking/_tracking_service/utils.py,sha256=yHUYcsPoVz2tz0OYeFJ2aKaGw66PTnLh48Yz4voXF0o,10256
528
531
  mlflow/tracking/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
529
532
  mlflow/tracking/context/abstract_context.py,sha256=OrlpO7gd_y2UP4pEpvwMKtfHI8XjtmPxO_2SujthUnI,1060
530
533
  mlflow/tracking/context/databricks_cluster_context.py,sha256=EORKG7JPBCnFv2Dupqd7FYgsTtwh4ZHB05T9W704IFI,520
@@ -587,7 +590,7 @@ mlflow/utils/env_manager.py,sha256=iLX1KxlTmoslXixTTKO896_SOrGEpiDIwMYcbthnb0E,4
587
590
  mlflow/utils/env_pack.py,sha256=QnuDhCjcqsXsYRZFPn1l0S8zHMs2wxbSSlkbzg-P7WY,4972
588
591
  mlflow/utils/environment.py,sha256=yODdh3keEvsl0l8cgm5r6pLcnDKOpWqJJXK7RhRGSeo,37390
589
592
  mlflow/utils/exception_utils.py,sha256=BHedTMmi3NvVbHzppKXmPubxLQq6OS7Q44fnpnT4bCw,365
590
- mlflow/utils/file_utils.py,sha256=1mzHEmVx7tWH0TVgyTuFWzfWpfzrevEnW4M1ZgVPDoo,31465
593
+ mlflow/utils/file_utils.py,sha256=joq-v2D5ffSV762mPbI__fABwQrIy0ELQ3z7AzkZTVo,31518
591
594
  mlflow/utils/git_utils.py,sha256=_uKXn2h58Xg2KvnanOf0eoZd23muWY__Lnl3nYmyblQ,2361
592
595
  mlflow/utils/gorilla.py,sha256=tAAbZgNVoCQwdGfPli-8yuK-fL1liNjjuMp-QuW6SwI,24049
593
596
  mlflow/utils/lazy_load.py,sha256=VJF0HV1fDEjenu8BFHOWfmDgJWME9GI8Fc9NG1VsN2o,1726
@@ -607,7 +610,7 @@ mlflow/utils/proto_json_utils.py,sha256=mku-o9v-uk_KVOAY6HFTVLDuAWtpWAg5NY_E1lMB
607
610
  mlflow/utils/pydantic_utils.py,sha256=-wsOgoARRTCYJEaTzEH0LVMiJX6ichzMHE2s4DiEJgw,1798
608
611
  mlflow/utils/request_utils.py,sha256=ZMctmkm0SuucaMwu_AO36v09VjY3Ql97WXPM5o4gl4w,10127
609
612
  mlflow/utils/requirements_utils.py,sha256=79BL9iFqg9H4VRYiqk-gTLlrfLyUY7KVG4BDV8NRZyo,26669
610
- mlflow/utils/rest_utils.py,sha256=yLeFOP135qV2OFcvtb_fFtj18MATRoxV4wJ5UshnEyk,26540
613
+ mlflow/utils/rest_utils.py,sha256=R5MZk0Ck358VeW-DKqpfSyDHbWjg_slQTpRfr4ZLDLM,26664
611
614
  mlflow/utils/search_logged_model_utils.py,sha256=JWczQxogo4u9Gr3gfoSVf7tVny1e-rxfAchh3Vqy6RU,4289
612
615
  mlflow/utils/search_utils.py,sha256=PgXtwlViHvmahY63otykh5iJVW0PR9DIwetrvCCBw2U,86142
613
616
  mlflow/utils/secure_loading.py,sha256=yJpKwrhzINmo0yVdp6fXvbNkZu_HEdwyB8deljGDXxk,10095
@@ -638,8 +641,8 @@ mlflow/utils/autologging_utils/metrics_queue.py,sha256=bwpMX7Go6xFxrpYROi6rDBdVt
638
641
  mlflow/utils/autologging_utils/safety.py,sha256=IwbTbusyE87Hc4qkhhvMikoaZqX9kpr972FWS2B8goc,51465
639
642
  mlflow/utils/autologging_utils/versioning.py,sha256=2hSN4KXFWEJCcopDdLG6BiPeqSoqjETeNMuUBsCCwlI,3762
640
643
  mlflow/utils/import_hooks/__init__.py,sha256=werje98Woelkbwrhtlb8wmRdt3RtiL--LqGru7Xh3YU,13589
641
- genesis_flow-1.0.1.dist-info/METADATA,sha256=QOs7W_wCIn4bYsjQVwLX8JE9zoW1dBO4rEJsf_qQmzU,32586
642
- genesis_flow-1.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
643
- genesis_flow-1.0.1.dist-info/entry_points.txt,sha256=fwlE6mYWea9uCNIsmXdT5qMIOZ8rYbTvO6XcZR8fBbE,495
644
- genesis_flow-1.0.1.dist-info/top_level.txt,sha256=wm8UqYyUHI21EvrTDHb3eYICy0dOVDLBhAL-jp5zbuI,7
645
- genesis_flow-1.0.1.dist-info/RECORD,,
644
+ genesis_flow-1.0.4.dist-info/METADATA,sha256=0ZYCp0XfX9xral_05JfSsI8gh0jXigpbeHEANiP2EIM,33401
645
+ genesis_flow-1.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
646
+ genesis_flow-1.0.4.dist-info/entry_points.txt,sha256=fwlE6mYWea9uCNIsmXdT5qMIOZ8rYbTvO6XcZR8fBbE,495
647
+ genesis_flow-1.0.4.dist-info/top_level.txt,sha256=wm8UqYyUHI21EvrTDHb3eYICy0dOVDLBhAL-jp5zbuI,7
648
+ genesis_flow-1.0.4.dist-info/RECORD,,
mlflow/azure/config.py CHANGED
@@ -84,13 +84,19 @@ class AzureAuthConfig:
84
84
  or os.getenv("MLFLOW_AZURE_DEBUG", "false").lower() == "true"
85
85
  )
86
86
 
87
- # Additional environment variable checks
88
- # Only override auth method if not explicitly set
89
- if (os.getenv("MLFLOW_POSTGRES_USE_MANAGED_IDENTITY", "false").lower() == "true" and
90
- auth_method is None and
91
- os.getenv("MLFLOW_AZURE_AUTH_METHOD") is None):
92
- self.auth_enabled = True
93
- self.auth_method = AuthMethod.MANAGED_IDENTITY
87
+ # IMPORTANT: Do not automatically enable Managed Identity based on other env vars
88
+ # The auth_enabled flag should be the single source of truth
89
+ # This prevents unintended authentication attempts when Helm hasn't enabled it
90
+
91
+ # Log configuration decision for debugging
92
+ import logging
93
+ logger = logging.getLogger(__name__)
94
+ logger.debug(
95
+ "AzureAuthConfig initialized: auth_enabled=%s, auth_method=%s, source=%s",
96
+ self.auth_enabled,
97
+ self.auth_method.value,
98
+ "explicit" if auth_enabled is not None else "env_var"
99
+ )
94
100
 
95
101
  # Validate configuration
96
102
  self._validate()
@@ -33,9 +33,10 @@ class ConnectionFactory:
33
33
  )
34
34
 
35
35
  logger.info(
36
- "Initialized connection factory: azure_auth_enabled=%s, auth_method=%s",
36
+ "Initialized connection factory: azure_auth_enabled=%s, auth_method=%s, auth_flag=%s",
37
37
  self.config.should_use_azure_auth,
38
38
  self.config.auth_method.value if self.config.should_use_azure_auth else "sql_auth",
39
+ self.config.auth_enabled,
39
40
  )
40
41
 
41
42
  def create_engine(
@@ -56,11 +57,23 @@ class ConnectionFactory:
56
57
  ConnectionError: If connection configuration fails
57
58
  ConfigurationError: If the URI is invalid
58
59
  """
60
+ # Check if auth_method is in URI
61
+ has_auth_method_param = "auth_method=" in database_uri
62
+
59
63
  logger.info(
60
- "Creating database engine: connection_info=%s, azure_auth=%s",
64
+ "Creating database engine: connection_info=%s, azure_auth=%s, has_auth_param=%s, auth_enabled_env=%s",
61
65
  sanitize_connection_string_for_logging(database_uri),
62
66
  self.config.should_use_azure_auth,
67
+ has_auth_method_param,
68
+ self.config.auth_enabled,
63
69
  )
70
+
71
+ # Warn if there's a mismatch between URI parameter and configuration
72
+ if has_auth_method_param and not self.config.should_use_azure_auth:
73
+ logger.warning(
74
+ "Database URI contains auth_method parameter but Azure authentication is not enabled. "
75
+ "The auth_method parameter will be ignored. Enable Azure authentication in Helm values to use Managed Identity."
76
+ )
64
77
 
65
78
  try:
66
79
  # Parse and validate the URI
@@ -217,3 +217,11 @@ try:
217
217
  _dataset_source_registry.register(UCVolumeDatasetSource)
218
218
  except ImportError:
219
219
  pass
220
+ try:
221
+ from mlflow.genai.datasets.databricks_evaluation_dataset_source import (
222
+ DatabricksEvaluationDatasetSource,
223
+ )
224
+
225
+ _dataset_source_registry.register(DatabricksEvaluationDatasetSource)
226
+ except ImportError:
227
+ pass
@@ -0,0 +1,298 @@
1
+ import json
2
+ import time
3
+ from enum import Enum
4
+
5
+ from mlflow.gateway.config import AmazonBedrockConfig, AWSIdAndKey, AWSRole, RouteConfig
6
+ from mlflow.gateway.constants import (
7
+ MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS,
8
+ )
9
+ from mlflow.gateway.exceptions import AIGatewayConfigException, AIGatewayException
10
+ from mlflow.gateway.providers.anthropic import AnthropicAdapter
11
+ from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
12
+ from mlflow.gateway.providers.cohere import CohereAdapter
13
+ from mlflow.gateway.providers.utils import rename_payload_keys
14
+ from mlflow.gateway.schemas import completions
15
+
16
+ AWS_BEDROCK_ANTHROPIC_MAXIMUM_MAX_TOKENS = 8191
17
+
18
+
19
+ class AmazonBedrockAnthropicAdapter(AnthropicAdapter):
20
+ @classmethod
21
+ def chat_to_model(cls, payload, config):
22
+ payload = super().chat_to_model(payload, config)
23
+ # "model" keys are not supported in Bedrock"
24
+ payload.pop("model", None)
25
+ return payload
26
+
27
+ @classmethod
28
+ def completions_to_model(cls, payload, config):
29
+ payload = super().completions_to_model(payload, config)
30
+
31
+ if "\n\nHuman:" not in payload.get("stop_sequences", []):
32
+ payload.setdefault("stop_sequences", []).append("\n\nHuman:")
33
+
34
+ payload["max_tokens_to_sample"] = min(
35
+ payload.get("max_tokens_to_sample", MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS),
36
+ AWS_BEDROCK_ANTHROPIC_MAXIMUM_MAX_TOKENS,
37
+ )
38
+
39
+ # "model" keys are not supported in Bedrock"
40
+ payload.pop("model", None)
41
+ return payload
42
+
43
+ @classmethod
44
+ def model_to_completions(cls, payload, config):
45
+ payload["model"] = config.model.name
46
+ return super().model_to_completions(payload, config)
47
+
48
+
49
+ class AWSTitanAdapter(ProviderAdapter):
50
+ # TODO handle top_p, top_k, etc.
51
+ @classmethod
52
+ def completions_to_model(cls, payload, config):
53
+ n = payload.pop("n", 1)
54
+ if n != 1:
55
+ raise AIGatewayException(
56
+ status_code=422,
57
+ detail=f"'n' must be '1' for AWS Titan models. Received value: '{n}'.",
58
+ )
59
+
60
+ # The range of Titan's temperature is 0-1, but ours is 0-2, so we halve it
61
+ if "temperature" in payload:
62
+ payload["temperature"] = 0.5 * payload["temperature"]
63
+ return {
64
+ "inputText": payload.pop("prompt"),
65
+ "textGenerationConfig": rename_payload_keys(
66
+ payload, {"max_tokens": "maxTokenCount", "stop": "stopSequences"}
67
+ ),
68
+ }
69
+
70
+ @classmethod
71
+ def model_to_completions(cls, resp, config):
72
+ return completions.ResponsePayload(
73
+ created=int(time.time()),
74
+ object="text_completion",
75
+ model=config.model.name,
76
+ choices=[
77
+ completions.Choice(
78
+ index=idx,
79
+ text=candidate.get("outputText"),
80
+ finish_reason=None,
81
+ )
82
+ for idx, candidate in enumerate(resp.get("results", []))
83
+ ],
84
+ usage=completions.CompletionsUsage(
85
+ prompt_tokens=None,
86
+ completion_tokens=None,
87
+ total_tokens=None,
88
+ ),
89
+ )
90
+
91
+ @classmethod
92
+ def embeddings_to_model(cls, payload, config):
93
+ raise NotImplementedError
94
+
95
+ @classmethod
96
+ def model_to_embeddings(cls, resp, config):
97
+ raise NotImplementedError
98
+
99
+
100
+ class AI21Adapter(ProviderAdapter):
101
+ # TODO handle top_p, top_k, etc.
102
+ @classmethod
103
+ def completions_to_model(cls, payload, config):
104
+ return rename_payload_keys(
105
+ payload,
106
+ {
107
+ "stop": "stopSequences",
108
+ "n": "numResults",
109
+ "max_tokens": "maxTokens",
110
+ },
111
+ )
112
+
113
+ @classmethod
114
+ def model_to_completions(cls, resp, config):
115
+ return completions.ResponsePayload(
116
+ created=int(time.time()),
117
+ object="text_completion",
118
+ model=config.model.name,
119
+ choices=[
120
+ completions.Choice(
121
+ index=idx,
122
+ text=candidate.get("data", {}).get("text"),
123
+ finish_reason=None,
124
+ )
125
+ for idx, candidate in enumerate(resp.get("completions", []))
126
+ ],
127
+ usage=completions.CompletionsUsage(
128
+ prompt_tokens=None,
129
+ completion_tokens=None,
130
+ total_tokens=None,
131
+ ),
132
+ )
133
+
134
+ @classmethod
135
+ def embeddings_to_model(cls, payload, config):
136
+ raise NotImplementedError
137
+
138
+ @classmethod
139
+ def model_to_embeddings(cls, resp, config):
140
+ raise NotImplementedError
141
+
142
+
143
+ class AmazonBedrockModelProvider(Enum):
144
+ AMAZON = "amazon"
145
+ COHERE = "cohere"
146
+ AI21 = "ai21"
147
+ ANTHROPIC = "anthropic"
148
+
149
+ @property
150
+ def adapter_class(self) -> type[ProviderAdapter]:
151
+ return AWS_MODEL_PROVIDER_TO_ADAPTER.get(self)
152
+
153
+ @classmethod
154
+ def of_str(cls, name: str):
155
+ name = name.lower()
156
+
157
+ for opt in cls:
158
+ if opt.name.lower() in name or opt.value.lower() in name:
159
+ return opt
160
+
161
+
162
+ AWS_MODEL_PROVIDER_TO_ADAPTER = {
163
+ AmazonBedrockModelProvider.COHERE: CohereAdapter,
164
+ AmazonBedrockModelProvider.ANTHROPIC: AmazonBedrockAnthropicAdapter,
165
+ AmazonBedrockModelProvider.AMAZON: AWSTitanAdapter,
166
+ AmazonBedrockModelProvider.AI21: AI21Adapter,
167
+ }
168
+
169
+
170
+ class AmazonBedrockProvider(BaseProvider):
171
+ NAME = "Amazon Bedrock"
172
+ CONFIG_TYPE = AmazonBedrockConfig
173
+
174
+ def __init__(self, config: RouteConfig):
175
+ super().__init__(config)
176
+
177
+ if config.model.config is None or not isinstance(config.model.config, AmazonBedrockConfig):
178
+ raise TypeError(f"Invalid config type {config.model.config}")
179
+ self.bedrock_config: AmazonBedrockConfig = config.model.config
180
+ self._client = None
181
+ self._client_created = 0
182
+
183
+ def _client_expired(self):
184
+ if not isinstance(self.bedrock_config.aws_config, AWSRole):
185
+ return False
186
+
187
+ return (
188
+ (time.monotonic_ns() - self._client_created)
189
+ >= (self.bedrock_config.aws_config.session_length_seconds) * 1_000_000_000,
190
+ )
191
+
192
+ def get_bedrock_client(self):
193
+ import boto3
194
+ import botocore.exceptions
195
+
196
+ if self._client is not None and not self._client_expired():
197
+ return self._client
198
+
199
+ session = boto3.Session(**self._construct_session_args())
200
+
201
+ try:
202
+ self._client, self._client_created = (
203
+ session.client(
204
+ service_name="bedrock-runtime",
205
+ **self._construct_client_args(session),
206
+ ),
207
+ time.monotonic_ns(),
208
+ )
209
+ return self._client
210
+ except botocore.exceptions.UnknownServiceError as e:
211
+ raise AIGatewayConfigException(
212
+ "Cannot create Amazon Bedrock client; ensure boto3/botocore "
213
+ "linked from the Amazon Bedrock user guide are installed. "
214
+ "Otherwise likely missing credentials or accessing account without to "
215
+ "Amazon Bedrock Private Preview"
216
+ ) from e
217
+
218
+ def _construct_session_args(self):
219
+ session_args = {
220
+ "region_name": self.bedrock_config.aws_config.aws_region,
221
+ }
222
+
223
+ return {k: v for k, v in session_args.items() if v}
224
+
225
+ def _construct_client_args(self, session):
226
+ aws_config = self.bedrock_config.aws_config
227
+
228
+ if isinstance(aws_config, AWSRole):
229
+ role = session.client(service_name="sts").assume_role(
230
+ RoleArn=aws_config.aws_role_arn,
231
+ RoleSessionName="ai-gateway-bedrock",
232
+ DurationSeconds=aws_config.session_length_seconds,
233
+ )
234
+ return {
235
+ "aws_access_key_id": role["Credentials"]["AccessKeyId"],
236
+ "aws_secret_access_key": role["Credentials"]["SecretAccessKey"],
237
+ "aws_session_token": role["Credentials"]["SessionToken"],
238
+ }
239
+ elif isinstance(aws_config, AWSIdAndKey):
240
+ return {
241
+ "aws_access_key_id": aws_config.aws_access_key_id,
242
+ "aws_secret_access_key": aws_config.aws_secret_access_key,
243
+ "aws_session_token": aws_config.aws_session_token,
244
+ }
245
+ else:
246
+ return {}
247
+
248
+ @property
249
+ def _underlying_provider(self):
250
+ if (not self.config.model.name) or "." not in self.config.model.name:
251
+ return None
252
+ return AmazonBedrockModelProvider.of_str(self.config.model.name)
253
+
254
+ @property
255
+ def adapter_class(self) -> type[ProviderAdapter]:
256
+ provider = self._underlying_provider
257
+ if not provider:
258
+ raise AIGatewayException(
259
+ status_code=422,
260
+ detail=f"Unknown Amazon Bedrock model type {self._underlying_provider}",
261
+ )
262
+ adapter = provider.adapter_class
263
+ if not adapter:
264
+ raise AIGatewayException(
265
+ status_code=422,
266
+ detail=f"Don't know how to handle {self._underlying_provider} for Amazon Bedrock",
267
+ )
268
+ return adapter
269
+
270
+ def _request(self, body):
271
+ import botocore.exceptions
272
+
273
+ try:
274
+ response = self.get_bedrock_client().invoke_model(
275
+ body=json.dumps(body).encode(),
276
+ modelId=self.config.model.name,
277
+ # defaults
278
+ # save=False,
279
+ accept="application/json",
280
+ contentType="application/json",
281
+ )
282
+ return json.loads(response.get("body").read())
283
+
284
+ # TODO work though botocore.exceptions to make this catchable.
285
+ # except botocore.exceptions.ValidationException as e:
286
+ # raise HTTPException(status_code=422, detail=str(e)) from e
287
+
288
+ except botocore.exceptions.ReadTimeoutError as e:
289
+ raise AIGatewayException(status_code=408) from e
290
+
291
+ async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
292
+ from fastapi.encoders import jsonable_encoder
293
+
294
+ self.check_for_model_field(payload)
295
+ payload = jsonable_encoder(payload, exclude_none=True, exclude_defaults=True)
296
+ payload = self.adapter_class.completions_to_model(payload, self.config)
297
+ response = self._request(payload)
298
+ return self.adapter_class.model_to_completions(response, self.config)