mlflow-tracing 3.0.0__tar.gz

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 (424) hide show
  1. mlflow_tracing-3.0.0/PKG-INFO +121 -0
  2. mlflow_tracing-3.0.0/README.md +91 -0
  3. mlflow_tracing-3.0.0/mlflow/__init__.py +408 -0
  4. mlflow_tracing-3.0.0/mlflow/__main__.py +3 -0
  5. mlflow_tracing-3.0.0/mlflow/anthropic/__init__.py +40 -0
  6. mlflow_tracing-3.0.0/mlflow/anthropic/autolog.py +112 -0
  7. mlflow_tracing-3.0.0/mlflow/anthropic/chat.py +144 -0
  8. mlflow_tracing-3.0.0/mlflow/autogen/__init__.py +122 -0
  9. mlflow_tracing-3.0.0/mlflow/autogen/chat.py +142 -0
  10. mlflow_tracing-3.0.0/mlflow/azure/__init__.py +0 -0
  11. mlflow_tracing-3.0.0/mlflow/azure/client.py +319 -0
  12. mlflow_tracing-3.0.0/mlflow/bedrock/__init__.py +45 -0
  13. mlflow_tracing-3.0.0/mlflow/bedrock/_autolog.py +200 -0
  14. mlflow_tracing-3.0.0/mlflow/bedrock/chat.py +122 -0
  15. mlflow_tracing-3.0.0/mlflow/bedrock/stream.py +160 -0
  16. mlflow_tracing-3.0.0/mlflow/bedrock/utils.py +43 -0
  17. mlflow_tracing-3.0.0/mlflow/cli.py +707 -0
  18. mlflow_tracing-3.0.0/mlflow/client.py +12 -0
  19. mlflow_tracing-3.0.0/mlflow/crewai/__init__.py +79 -0
  20. mlflow_tracing-3.0.0/mlflow/crewai/autolog.py +253 -0
  21. mlflow_tracing-3.0.0/mlflow/crewai/chat.py +29 -0
  22. mlflow_tracing-3.0.0/mlflow/db.py +27 -0
  23. mlflow_tracing-3.0.0/mlflow/dspy/__init__.py +17 -0
  24. mlflow_tracing-3.0.0/mlflow/dspy/autolog.py +197 -0
  25. mlflow_tracing-3.0.0/mlflow/dspy/callback.py +398 -0
  26. mlflow_tracing-3.0.0/mlflow/dspy/constant.py +1 -0
  27. mlflow_tracing-3.0.0/mlflow/dspy/load.py +93 -0
  28. mlflow_tracing-3.0.0/mlflow/dspy/save.py +393 -0
  29. mlflow_tracing-3.0.0/mlflow/dspy/util.py +109 -0
  30. mlflow_tracing-3.0.0/mlflow/dspy/wrapper.py +226 -0
  31. mlflow_tracing-3.0.0/mlflow/entities/__init__.py +106 -0
  32. mlflow_tracing-3.0.0/mlflow/entities/_mlflow_object.py +52 -0
  33. mlflow_tracing-3.0.0/mlflow/entities/assessment.py +545 -0
  34. mlflow_tracing-3.0.0/mlflow/entities/assessment_error.py +80 -0
  35. mlflow_tracing-3.0.0/mlflow/entities/assessment_source.py +141 -0
  36. mlflow_tracing-3.0.0/mlflow/entities/dataset.py +92 -0
  37. mlflow_tracing-3.0.0/mlflow/entities/dataset_input.py +51 -0
  38. mlflow_tracing-3.0.0/mlflow/entities/dataset_summary.py +62 -0
  39. mlflow_tracing-3.0.0/mlflow/entities/document.py +48 -0
  40. mlflow_tracing-3.0.0/mlflow/entities/experiment.py +109 -0
  41. mlflow_tracing-3.0.0/mlflow/entities/experiment_tag.py +35 -0
  42. mlflow_tracing-3.0.0/mlflow/entities/file_info.py +45 -0
  43. mlflow_tracing-3.0.0/mlflow/entities/input_tag.py +35 -0
  44. mlflow_tracing-3.0.0/mlflow/entities/lifecycle_stage.py +35 -0
  45. mlflow_tracing-3.0.0/mlflow/entities/logged_model.py +228 -0
  46. mlflow_tracing-3.0.0/mlflow/entities/logged_model_input.py +26 -0
  47. mlflow_tracing-3.0.0/mlflow/entities/logged_model_output.py +32 -0
  48. mlflow_tracing-3.0.0/mlflow/entities/logged_model_parameter.py +46 -0
  49. mlflow_tracing-3.0.0/mlflow/entities/logged_model_status.py +74 -0
  50. mlflow_tracing-3.0.0/mlflow/entities/logged_model_tag.py +33 -0
  51. mlflow_tracing-3.0.0/mlflow/entities/metric.py +200 -0
  52. mlflow_tracing-3.0.0/mlflow/entities/model_registry/__init__.py +29 -0
  53. mlflow_tracing-3.0.0/mlflow/entities/model_registry/_model_registry_entity.py +13 -0
  54. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version.py +243 -0
  55. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version_deployment_job_run_state.py +44 -0
  56. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version_deployment_job_state.py +70 -0
  57. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version_search.py +25 -0
  58. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version_stages.py +25 -0
  59. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version_status.py +35 -0
  60. mlflow_tracing-3.0.0/mlflow/entities/model_registry/model_version_tag.py +35 -0
  61. mlflow_tracing-3.0.0/mlflow/entities/model_registry/prompt.py +73 -0
  62. mlflow_tracing-3.0.0/mlflow/entities/model_registry/prompt_version.py +244 -0
  63. mlflow_tracing-3.0.0/mlflow/entities/model_registry/registered_model.py +175 -0
  64. mlflow_tracing-3.0.0/mlflow/entities/model_registry/registered_model_alias.py +35 -0
  65. mlflow_tracing-3.0.0/mlflow/entities/model_registry/registered_model_deployment_job_state.py +39 -0
  66. mlflow_tracing-3.0.0/mlflow/entities/model_registry/registered_model_search.py +25 -0
  67. mlflow_tracing-3.0.0/mlflow/entities/model_registry/registered_model_tag.py +35 -0
  68. mlflow_tracing-3.0.0/mlflow/entities/multipart_upload.py +74 -0
  69. mlflow_tracing-3.0.0/mlflow/entities/param.py +49 -0
  70. mlflow_tracing-3.0.0/mlflow/entities/run.py +97 -0
  71. mlflow_tracing-3.0.0/mlflow/entities/run_data.py +84 -0
  72. mlflow_tracing-3.0.0/mlflow/entities/run_info.py +188 -0
  73. mlflow_tracing-3.0.0/mlflow/entities/run_inputs.py +59 -0
  74. mlflow_tracing-3.0.0/mlflow/entities/run_outputs.py +43 -0
  75. mlflow_tracing-3.0.0/mlflow/entities/run_status.py +41 -0
  76. mlflow_tracing-3.0.0/mlflow/entities/run_tag.py +36 -0
  77. mlflow_tracing-3.0.0/mlflow/entities/source_type.py +31 -0
  78. mlflow_tracing-3.0.0/mlflow/entities/span.py +764 -0
  79. mlflow_tracing-3.0.0/mlflow/entities/span_event.py +96 -0
  80. mlflow_tracing-3.0.0/mlflow/entities/span_status.py +102 -0
  81. mlflow_tracing-3.0.0/mlflow/entities/trace.py +317 -0
  82. mlflow_tracing-3.0.0/mlflow/entities/trace_data.py +71 -0
  83. mlflow_tracing-3.0.0/mlflow/entities/trace_info.py +215 -0
  84. mlflow_tracing-3.0.0/mlflow/entities/trace_info_v2.py +153 -0
  85. mlflow_tracing-3.0.0/mlflow/entities/trace_location.py +173 -0
  86. mlflow_tracing-3.0.0/mlflow/entities/trace_state.py +39 -0
  87. mlflow_tracing-3.0.0/mlflow/entities/trace_status.py +68 -0
  88. mlflow_tracing-3.0.0/mlflow/entities/view_type.py +51 -0
  89. mlflow_tracing-3.0.0/mlflow/environment_variables.py +849 -0
  90. mlflow_tracing-3.0.0/mlflow/exceptions.py +213 -0
  91. mlflow_tracing-3.0.0/mlflow/experiments.py +140 -0
  92. mlflow_tracing-3.0.0/mlflow/gemini/__init__.py +81 -0
  93. mlflow_tracing-3.0.0/mlflow/gemini/autolog.py +186 -0
  94. mlflow_tracing-3.0.0/mlflow/gemini/chat.py +261 -0
  95. mlflow_tracing-3.0.0/mlflow/groq/__init__.py +42 -0
  96. mlflow_tracing-3.0.0/mlflow/groq/_groq_autolog.py +56 -0
  97. mlflow_tracing-3.0.0/mlflow/langchain/__init__.py +24 -0
  98. mlflow_tracing-3.0.0/mlflow/langchain/api_request_parallel_processor.py +328 -0
  99. mlflow_tracing-3.0.0/mlflow/langchain/autolog.py +147 -0
  100. mlflow_tracing-3.0.0/mlflow/langchain/chat_agent_langgraph.py +337 -0
  101. mlflow_tracing-3.0.0/mlflow/langchain/constant.py +1 -0
  102. mlflow_tracing-3.0.0/mlflow/langchain/constants.py +1 -0
  103. mlflow_tracing-3.0.0/mlflow/langchain/databricks_dependencies.py +444 -0
  104. mlflow_tracing-3.0.0/mlflow/langchain/langchain_tracer.py +597 -0
  105. mlflow_tracing-3.0.0/mlflow/langchain/model.py +920 -0
  106. mlflow_tracing-3.0.0/mlflow/langchain/output_parsers.py +142 -0
  107. mlflow_tracing-3.0.0/mlflow/langchain/retriever_chain.py +153 -0
  108. mlflow_tracing-3.0.0/mlflow/langchain/runnables.py +527 -0
  109. mlflow_tracing-3.0.0/mlflow/langchain/utils/chat.py +402 -0
  110. mlflow_tracing-3.0.0/mlflow/langchain/utils/logging.py +671 -0
  111. mlflow_tracing-3.0.0/mlflow/langchain/utils/serialization.py +36 -0
  112. mlflow_tracing-3.0.0/mlflow/legacy_databricks_cli/__init__.py +0 -0
  113. mlflow_tracing-3.0.0/mlflow/legacy_databricks_cli/configure/__init__.py +0 -0
  114. mlflow_tracing-3.0.0/mlflow/legacy_databricks_cli/configure/provider.py +482 -0
  115. mlflow_tracing-3.0.0/mlflow/litellm/__init__.py +175 -0
  116. mlflow_tracing-3.0.0/mlflow/llama_index/__init__.py +22 -0
  117. mlflow_tracing-3.0.0/mlflow/llama_index/autolog.py +55 -0
  118. mlflow_tracing-3.0.0/mlflow/llama_index/chat.py +43 -0
  119. mlflow_tracing-3.0.0/mlflow/llama_index/constant.py +1 -0
  120. mlflow_tracing-3.0.0/mlflow/llama_index/model.py +578 -0
  121. mlflow_tracing-3.0.0/mlflow/llama_index/pyfunc_wrapper.py +325 -0
  122. mlflow_tracing-3.0.0/mlflow/llama_index/serialize_objects.py +188 -0
  123. mlflow_tracing-3.0.0/mlflow/llama_index/tracer.py +545 -0
  124. mlflow_tracing-3.0.0/mlflow/mismatch.py +34 -0
  125. mlflow_tracing-3.0.0/mlflow/mistral/__init__.py +34 -0
  126. mlflow_tracing-3.0.0/mlflow/mistral/autolog.py +61 -0
  127. mlflow_tracing-3.0.0/mlflow/mistral/chat.py +135 -0
  128. mlflow_tracing-3.0.0/mlflow/ml_package_versions.py +452 -0
  129. mlflow_tracing-3.0.0/mlflow/openai/__init__.py +57 -0
  130. mlflow_tracing-3.0.0/mlflow/openai/_agent_tracer.py +364 -0
  131. mlflow_tracing-3.0.0/mlflow/openai/api_request_parallel_processor.py +126 -0
  132. mlflow_tracing-3.0.0/mlflow/openai/autolog.py +449 -0
  133. mlflow_tracing-3.0.0/mlflow/openai/constant.py +1 -0
  134. mlflow_tracing-3.0.0/mlflow/openai/model.py +824 -0
  135. mlflow_tracing-3.0.0/mlflow/openai/utils/chat_schema.py +367 -0
  136. mlflow_tracing-3.0.0/mlflow/prompt/constants.py +20 -0
  137. mlflow_tracing-3.0.0/mlflow/prompt/promptlab_model.py +197 -0
  138. mlflow_tracing-3.0.0/mlflow/prompt/registry_utils.py +248 -0
  139. mlflow_tracing-3.0.0/mlflow/promptflow/__init__.py +495 -0
  140. mlflow_tracing-3.0.0/mlflow/protos/__init__.py +0 -0
  141. mlflow_tracing-3.0.0/mlflow/protos/assessments_pb2.py +174 -0
  142. mlflow_tracing-3.0.0/mlflow/protos/databricks_artifacts_pb2.py +489 -0
  143. mlflow_tracing-3.0.0/mlflow/protos/databricks_filesystem_service_pb2.py +196 -0
  144. mlflow_tracing-3.0.0/mlflow/protos/databricks_managed_catalog_messages_pb2.py +95 -0
  145. mlflow_tracing-3.0.0/mlflow/protos/databricks_managed_catalog_service_pb2.py +86 -0
  146. mlflow_tracing-3.0.0/mlflow/protos/databricks_pb2.py +267 -0
  147. mlflow_tracing-3.0.0/mlflow/protos/databricks_trace_server_pb2.py +374 -0
  148. mlflow_tracing-3.0.0/mlflow/protos/databricks_uc_registry_messages_pb2.py +1249 -0
  149. mlflow_tracing-3.0.0/mlflow/protos/databricks_uc_registry_service_pb2.py +170 -0
  150. mlflow_tracing-3.0.0/mlflow/protos/facet_feature_statistics_pb2.py +296 -0
  151. mlflow_tracing-3.0.0/mlflow/protos/internal_pb2.py +77 -0
  152. mlflow_tracing-3.0.0/mlflow/protos/mlflow_artifacts_pb2.py +336 -0
  153. mlflow_tracing-3.0.0/mlflow/protos/model_registry_pb2.py +1073 -0
  154. mlflow_tracing-3.0.0/mlflow/protos/scalapb/__init__.py +0 -0
  155. mlflow_tracing-3.0.0/mlflow/protos/scalapb/scalapb_pb2.py +104 -0
  156. mlflow_tracing-3.0.0/mlflow/protos/service_pb2.py +2600 -0
  157. mlflow_tracing-3.0.0/mlflow/protos/unity_catalog_oss_messages_pb2.py +457 -0
  158. mlflow_tracing-3.0.0/mlflow/protos/unity_catalog_oss_service_pb2.py +130 -0
  159. mlflow_tracing-3.0.0/mlflow/protos/unity_catalog_prompt_messages_pb2.py +447 -0
  160. mlflow_tracing-3.0.0/mlflow/protos/unity_catalog_prompt_messages_pb2_grpc.py +24 -0
  161. mlflow_tracing-3.0.0/mlflow/protos/unity_catalog_prompt_service_pb2.py +164 -0
  162. mlflow_tracing-3.0.0/mlflow/protos/unity_catalog_prompt_service_pb2_grpc.py +785 -0
  163. mlflow_tracing-3.0.0/mlflow/py.typed +0 -0
  164. mlflow_tracing-3.0.0/mlflow/pydantic_ai/__init__.py +57 -0
  165. mlflow_tracing-3.0.0/mlflow/pydantic_ai/autolog.py +173 -0
  166. mlflow_tracing-3.0.0/mlflow/runs.py +89 -0
  167. mlflow_tracing-3.0.0/mlflow/smolagents/__init__.py +66 -0
  168. mlflow_tracing-3.0.0/mlflow/smolagents/autolog.py +139 -0
  169. mlflow_tracing-3.0.0/mlflow/smolagents/chat.py +29 -0
  170. mlflow_tracing-3.0.0/mlflow/store/__init__.py +10 -0
  171. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/__init__.py +1 -0
  172. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/lineage/__init__.py +1 -0
  173. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/lineage/constants.py +2 -0
  174. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/registry/__init__.py +6 -0
  175. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/registry/prompt_info.py +75 -0
  176. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/registry/rest_store.py +1740 -0
  177. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/registry/uc_oss_rest_store.py +507 -0
  178. mlflow_tracing-3.0.0/mlflow/store/_unity_catalog/registry/utils.py +121 -0
  179. mlflow_tracing-3.0.0/mlflow/store/artifact/__init__.py +0 -0
  180. mlflow_tracing-3.0.0/mlflow/store/artifact/artifact_repo.py +466 -0
  181. mlflow_tracing-3.0.0/mlflow/store/artifact/artifact_repository_registry.py +143 -0
  182. mlflow_tracing-3.0.0/mlflow/store/artifact/azure_blob_artifact_repo.py +275 -0
  183. mlflow_tracing-3.0.0/mlflow/store/artifact/azure_data_lake_artifact_repo.py +293 -0
  184. mlflow_tracing-3.0.0/mlflow/store/artifact/cli.py +141 -0
  185. mlflow_tracing-3.0.0/mlflow/store/artifact/cloud_artifact_repo.py +331 -0
  186. mlflow_tracing-3.0.0/mlflow/store/artifact/databricks_artifact_repo.py +711 -0
  187. mlflow_tracing-3.0.0/mlflow/store/artifact/databricks_artifact_repo_resources.py +301 -0
  188. mlflow_tracing-3.0.0/mlflow/store/artifact/databricks_logged_model_artifact_repo.py +93 -0
  189. mlflow_tracing-3.0.0/mlflow/store/artifact/databricks_models_artifact_repo.py +216 -0
  190. mlflow_tracing-3.0.0/mlflow/store/artifact/databricks_sdk_artifact_repo.py +134 -0
  191. mlflow_tracing-3.0.0/mlflow/store/artifact/databricks_sdk_models_artifact_repo.py +95 -0
  192. mlflow_tracing-3.0.0/mlflow/store/artifact/dbfs_artifact_repo.py +237 -0
  193. mlflow_tracing-3.0.0/mlflow/store/artifact/ftp_artifact_repo.py +133 -0
  194. mlflow_tracing-3.0.0/mlflow/store/artifact/gcs_artifact_repo.py +289 -0
  195. mlflow_tracing-3.0.0/mlflow/store/artifact/hdfs_artifact_repo.py +208 -0
  196. mlflow_tracing-3.0.0/mlflow/store/artifact/http_artifact_repo.py +218 -0
  197. mlflow_tracing-3.0.0/mlflow/store/artifact/local_artifact_repo.py +142 -0
  198. mlflow_tracing-3.0.0/mlflow/store/artifact/mlflow_artifacts_repo.py +92 -0
  199. mlflow_tracing-3.0.0/mlflow/store/artifact/models_artifact_repo.py +259 -0
  200. mlflow_tracing-3.0.0/mlflow/store/artifact/optimized_s3_artifact_repo.py +356 -0
  201. mlflow_tracing-3.0.0/mlflow/store/artifact/presigned_url_artifact_repo.py +169 -0
  202. mlflow_tracing-3.0.0/mlflow/store/artifact/r2_artifact_repo.py +70 -0
  203. mlflow_tracing-3.0.0/mlflow/store/artifact/runs_artifact_repo.py +244 -0
  204. mlflow_tracing-3.0.0/mlflow/store/artifact/s3_artifact_repo.py +324 -0
  205. mlflow_tracing-3.0.0/mlflow/store/artifact/sftp_artifact_repo.py +142 -0
  206. mlflow_tracing-3.0.0/mlflow/store/artifact/uc_volume_artifact_repo.py +73 -0
  207. mlflow_tracing-3.0.0/mlflow/store/artifact/unity_catalog_models_artifact_repo.py +168 -0
  208. mlflow_tracing-3.0.0/mlflow/store/artifact/unity_catalog_oss_models_artifact_repo.py +168 -0
  209. mlflow_tracing-3.0.0/mlflow/store/artifact/utils/__init__.py +0 -0
  210. mlflow_tracing-3.0.0/mlflow/store/artifact/utils/models.py +135 -0
  211. mlflow_tracing-3.0.0/mlflow/store/db/__init__.py +0 -0
  212. mlflow_tracing-3.0.0/mlflow/store/db/base_sql_model.py +3 -0
  213. mlflow_tracing-3.0.0/mlflow/store/db/db_types.py +10 -0
  214. mlflow_tracing-3.0.0/mlflow/store/db/utils.py +314 -0
  215. mlflow_tracing-3.0.0/mlflow/store/db_migrations/__init__.py +0 -0
  216. mlflow_tracing-3.0.0/mlflow/store/db_migrations/env.py +84 -0
  217. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/0584bdc529eb_add_cascading_deletion_to_datasets_from_experiments.py +88 -0
  218. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py +49 -0
  219. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/0c779009ac13_add_deleted_time_field_to_runs_table.py +24 -0
  220. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/181f10493468_allow_nulls_for_metric_values.py +35 -0
  221. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/27a6a02d2cf1_add_model_version_tags_table.py +38 -0
  222. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/2b4d017a5e9b_add_model_registry_tables_to_db.py +77 -0
  223. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/2d6e25af4d3e_increase_max_param_val_length.py +33 -0
  224. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/3500859a5d39_add_model_aliases_table.py +50 -0
  225. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/39d1c3be5f05_add_is_nan_constraint_for_metrics_tables_if_necessary.py +41 -0
  226. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/400f98739977_add_logged_model_tables.py +123 -0
  227. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/4465047574b1_increase_max_dataset_schema_size.py +38 -0
  228. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/451aebb31d03_add_metric_step.py +35 -0
  229. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py +40 -0
  230. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/6953534de441_add_step_to_inputs_table.py +25 -0
  231. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/728d730b5ebd_add_registered_model_tags_table.py +38 -0
  232. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/7ac759974ad8_update_run_tags_with_larger_limit.py +36 -0
  233. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/7f2a7d5fae7d_add_datasets_inputs_input_tags_tables.py +82 -0
  234. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/84291f40a231_add_run_link_to_model_version.py +26 -0
  235. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/867495a8f9d4_add_trace_tables.py +90 -0
  236. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/89d4b8295536_create_latest_metrics_table.py +169 -0
  237. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/90e64c465722_migrate_user_column_to_tags.py +64 -0
  238. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/97727af70f4d_creation_time_last_update_time_experiments.py +25 -0
  239. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/__init__.py +0 -0
  240. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/a8c4a736bde6_allow_nulls_for_run_id.py +27 -0
  241. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/acf3f17fdcc7_add_storage_location_field_to_model_.py +29 -0
  242. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/bd07f7e963c5_create_index_on_run_uuid.py +26 -0
  243. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/c48cb773bb87_reset_default_value_for_is_nan_in_metrics_table_for_mysql.py +41 -0
  244. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/cc1f77228345_change_param_value_length_to_500.py +34 -0
  245. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py +78 -0
  246. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/df50e92ffc5e_add_experiment_tags_table.py +38 -0
  247. mlflow_tracing-3.0.0/mlflow/store/db_migrations/versions/f5a4f2784254_increase_run_tag_value_limit.py +36 -0
  248. mlflow_tracing-3.0.0/mlflow/store/entities/__init__.py +3 -0
  249. mlflow_tracing-3.0.0/mlflow/store/entities/paged_list.py +18 -0
  250. mlflow_tracing-3.0.0/mlflow/store/model_registry/__init__.py +10 -0
  251. mlflow_tracing-3.0.0/mlflow/store/model_registry/abstract_store.py +1082 -0
  252. mlflow_tracing-3.0.0/mlflow/store/model_registry/base_rest_store.py +44 -0
  253. mlflow_tracing-3.0.0/mlflow/store/model_registry/databricks_workspace_model_registry_rest_store.py +37 -0
  254. mlflow_tracing-3.0.0/mlflow/store/model_registry/dbmodels/__init__.py +0 -0
  255. mlflow_tracing-3.0.0/mlflow/store/model_registry/dbmodels/models.py +205 -0
  256. mlflow_tracing-3.0.0/mlflow/store/model_registry/file_store.py +1091 -0
  257. mlflow_tracing-3.0.0/mlflow/store/model_registry/rest_store.py +481 -0
  258. mlflow_tracing-3.0.0/mlflow/store/model_registry/sqlalchemy_store.py +1286 -0
  259. mlflow_tracing-3.0.0/mlflow/store/tracking/__init__.py +23 -0
  260. mlflow_tracing-3.0.0/mlflow/store/tracking/abstract_store.py +855 -0
  261. mlflow_tracing-3.0.0/mlflow/store/tracking/dbmodels/__init__.py +0 -0
  262. mlflow_tracing-3.0.0/mlflow/store/tracking/dbmodels/initial_models.py +243 -0
  263. mlflow_tracing-3.0.0/mlflow/store/tracking/dbmodels/models.py +1055 -0
  264. mlflow_tracing-3.0.0/mlflow/store/tracking/file_store.py +2385 -0
  265. mlflow_tracing-3.0.0/mlflow/store/tracking/rest_store.py +1152 -0
  266. mlflow_tracing-3.0.0/mlflow/store/tracking/sqlalchemy_store.py +2693 -0
  267. mlflow_tracing-3.0.0/mlflow/tracing/__init__.py +14 -0
  268. mlflow_tracing-3.0.0/mlflow/tracing/assessment.py +356 -0
  269. mlflow_tracing-3.0.0/mlflow/tracing/client.py +617 -0
  270. mlflow_tracing-3.0.0/mlflow/tracing/constant.py +91 -0
  271. mlflow_tracing-3.0.0/mlflow/tracing/destination.py +81 -0
  272. mlflow_tracing-3.0.0/mlflow/tracing/display/__init__.py +40 -0
  273. mlflow_tracing-3.0.0/mlflow/tracing/display/display_handler.py +196 -0
  274. mlflow_tracing-3.0.0/mlflow/tracing/export/async_export_queue.py +186 -0
  275. mlflow_tracing-3.0.0/mlflow/tracing/export/inference_table.py +141 -0
  276. mlflow_tracing-3.0.0/mlflow/tracing/export/mlflow_v2.py +97 -0
  277. mlflow_tracing-3.0.0/mlflow/tracing/export/mlflow_v3.py +133 -0
  278. mlflow_tracing-3.0.0/mlflow/tracing/export/utils.py +70 -0
  279. mlflow_tracing-3.0.0/mlflow/tracing/fluent.py +1417 -0
  280. mlflow_tracing-3.0.0/mlflow/tracing/processor/base_mlflow.py +204 -0
  281. mlflow_tracing-3.0.0/mlflow/tracing/processor/inference_table.py +175 -0
  282. mlflow_tracing-3.0.0/mlflow/tracing/processor/mlflow_v2.py +78 -0
  283. mlflow_tracing-3.0.0/mlflow/tracing/processor/mlflow_v3.py +48 -0
  284. mlflow_tracing-3.0.0/mlflow/tracing/processor/otel.py +62 -0
  285. mlflow_tracing-3.0.0/mlflow/tracing/provider.py +494 -0
  286. mlflow_tracing-3.0.0/mlflow/tracing/trace_manager.py +200 -0
  287. mlflow_tracing-3.0.0/mlflow/tracing/utils/__init__.py +545 -0
  288. mlflow_tracing-3.0.0/mlflow/tracing/utils/artifact_utils.py +31 -0
  289. mlflow_tracing-3.0.0/mlflow/tracing/utils/environment.py +55 -0
  290. mlflow_tracing-3.0.0/mlflow/tracing/utils/exception.py +21 -0
  291. mlflow_tracing-3.0.0/mlflow/tracing/utils/once.py +35 -0
  292. mlflow_tracing-3.0.0/mlflow/tracing/utils/otlp.py +63 -0
  293. mlflow_tracing-3.0.0/mlflow/tracing/utils/search.py +292 -0
  294. mlflow_tracing-3.0.0/mlflow/tracing/utils/timeout.py +250 -0
  295. mlflow_tracing-3.0.0/mlflow/tracing/utils/token.py +19 -0
  296. mlflow_tracing-3.0.0/mlflow/tracing/utils/truncation.py +106 -0
  297. mlflow_tracing-3.0.0/mlflow/tracing/utils/warning.py +76 -0
  298. mlflow_tracing-3.0.0/mlflow/tracking/__init__.py +39 -0
  299. mlflow_tracing-3.0.0/mlflow/tracking/_model_registry/__init__.py +1 -0
  300. mlflow_tracing-3.0.0/mlflow/tracking/_model_registry/client.py +764 -0
  301. mlflow_tracing-3.0.0/mlflow/tracking/_model_registry/fluent.py +853 -0
  302. mlflow_tracing-3.0.0/mlflow/tracking/_model_registry/registry.py +67 -0
  303. mlflow_tracing-3.0.0/mlflow/tracking/_model_registry/utils.py +249 -0
  304. mlflow_tracing-3.0.0/mlflow/tracking/_tracking_service/__init__.py +0 -0
  305. mlflow_tracing-3.0.0/mlflow/tracking/_tracking_service/client.py +883 -0
  306. mlflow_tracing-3.0.0/mlflow/tracking/_tracking_service/registry.py +56 -0
  307. mlflow_tracing-3.0.0/mlflow/tracking/_tracking_service/utils.py +262 -0
  308. mlflow_tracing-3.0.0/mlflow/tracking/artifact_utils.py +172 -0
  309. mlflow_tracing-3.0.0/mlflow/tracking/client.py +5910 -0
  310. mlflow_tracing-3.0.0/mlflow/tracking/context/__init__.py +0 -0
  311. mlflow_tracing-3.0.0/mlflow/tracking/context/abstract_context.py +35 -0
  312. mlflow_tracing-3.0.0/mlflow/tracking/context/databricks_cluster_context.py +15 -0
  313. mlflow_tracing-3.0.0/mlflow/tracking/context/databricks_command_context.py +15 -0
  314. mlflow_tracing-3.0.0/mlflow/tracking/context/databricks_job_context.py +49 -0
  315. mlflow_tracing-3.0.0/mlflow/tracking/context/databricks_notebook_context.py +41 -0
  316. mlflow_tracing-3.0.0/mlflow/tracking/context/databricks_repo_context.py +43 -0
  317. mlflow_tracing-3.0.0/mlflow/tracking/context/default_context.py +51 -0
  318. mlflow_tracing-3.0.0/mlflow/tracking/context/git_context.py +32 -0
  319. mlflow_tracing-3.0.0/mlflow/tracking/context/registry.py +98 -0
  320. mlflow_tracing-3.0.0/mlflow/tracking/context/system_environment_context.py +15 -0
  321. mlflow_tracing-3.0.0/mlflow/tracking/default_experiment/__init__.py +1 -0
  322. mlflow_tracing-3.0.0/mlflow/tracking/default_experiment/abstract_context.py +43 -0
  323. mlflow_tracing-3.0.0/mlflow/tracking/default_experiment/databricks_notebook_experiment_provider.py +44 -0
  324. mlflow_tracing-3.0.0/mlflow/tracking/default_experiment/registry.py +75 -0
  325. mlflow_tracing-3.0.0/mlflow/tracking/fluent.py +3554 -0
  326. mlflow_tracing-3.0.0/mlflow/tracking/metric_value_conversion_utils.py +93 -0
  327. mlflow_tracing-3.0.0/mlflow/tracking/multimedia.py +206 -0
  328. mlflow_tracing-3.0.0/mlflow/tracking/registry.py +86 -0
  329. mlflow_tracing-3.0.0/mlflow/tracking/request_auth/__init__.py +0 -0
  330. mlflow_tracing-3.0.0/mlflow/tracking/request_auth/abstract_request_auth_provider.py +34 -0
  331. mlflow_tracing-3.0.0/mlflow/tracking/request_auth/registry.py +60 -0
  332. mlflow_tracing-3.0.0/mlflow/tracking/request_header/__init__.py +0 -0
  333. mlflow_tracing-3.0.0/mlflow/tracking/request_header/abstract_request_header_provider.py +36 -0
  334. mlflow_tracing-3.0.0/mlflow/tracking/request_header/databricks_request_header_provider.py +38 -0
  335. mlflow_tracing-3.0.0/mlflow/tracking/request_header/default_request_header_provider.py +17 -0
  336. mlflow_tracing-3.0.0/mlflow/tracking/request_header/registry.py +79 -0
  337. mlflow_tracing-3.0.0/mlflow/types/__init__.py +21 -0
  338. mlflow_tracing-3.0.0/mlflow/types/agent.py +270 -0
  339. mlflow_tracing-3.0.0/mlflow/types/chat.py +240 -0
  340. mlflow_tracing-3.0.0/mlflow/types/llm.py +935 -0
  341. mlflow_tracing-3.0.0/mlflow/types/responses.py +139 -0
  342. mlflow_tracing-3.0.0/mlflow/types/responses_helpers.py +416 -0
  343. mlflow_tracing-3.0.0/mlflow/types/schema.py +1505 -0
  344. mlflow_tracing-3.0.0/mlflow/types/type_hints.py +647 -0
  345. mlflow_tracing-3.0.0/mlflow/types/utils.py +753 -0
  346. mlflow_tracing-3.0.0/mlflow/utils/__init__.py +283 -0
  347. mlflow_tracing-3.0.0/mlflow/utils/_capture_modules.py +256 -0
  348. mlflow_tracing-3.0.0/mlflow/utils/_capture_transformers_modules.py +75 -0
  349. mlflow_tracing-3.0.0/mlflow/utils/_spark_utils.py +201 -0
  350. mlflow_tracing-3.0.0/mlflow/utils/_unity_catalog_oss_utils.py +97 -0
  351. mlflow_tracing-3.0.0/mlflow/utils/_unity_catalog_utils.py +479 -0
  352. mlflow_tracing-3.0.0/mlflow/utils/annotations.py +206 -0
  353. mlflow_tracing-3.0.0/mlflow/utils/arguments_utils.py +16 -0
  354. mlflow_tracing-3.0.0/mlflow/utils/async_logging/__init__.py +1 -0
  355. mlflow_tracing-3.0.0/mlflow/utils/async_logging/async_artifacts_logging_queue.py +258 -0
  356. mlflow_tracing-3.0.0/mlflow/utils/async_logging/async_logging_queue.py +366 -0
  357. mlflow_tracing-3.0.0/mlflow/utils/async_logging/run_artifact.py +38 -0
  358. mlflow_tracing-3.0.0/mlflow/utils/async_logging/run_batch.py +58 -0
  359. mlflow_tracing-3.0.0/mlflow/utils/async_logging/run_operations.py +49 -0
  360. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/__init__.py +737 -0
  361. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/client.py +432 -0
  362. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/config.py +33 -0
  363. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/events.py +294 -0
  364. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/logging_and_warnings.py +328 -0
  365. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/metrics_queue.py +69 -0
  366. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/safety.py +1105 -0
  367. mlflow_tracing-3.0.0/mlflow/utils/autologging_utils/versioning.py +95 -0
  368. mlflow_tracing-3.0.0/mlflow/utils/checkpoint_utils.py +206 -0
  369. mlflow_tracing-3.0.0/mlflow/utils/class_utils.py +6 -0
  370. mlflow_tracing-3.0.0/mlflow/utils/cli_args.py +257 -0
  371. mlflow_tracing-3.0.0/mlflow/utils/conda.py +354 -0
  372. mlflow_tracing-3.0.0/mlflow/utils/credentials.py +231 -0
  373. mlflow_tracing-3.0.0/mlflow/utils/data_utils.py +17 -0
  374. mlflow_tracing-3.0.0/mlflow/utils/databricks_utils.py +1436 -0
  375. mlflow_tracing-3.0.0/mlflow/utils/docstring_utils.py +476 -0
  376. mlflow_tracing-3.0.0/mlflow/utils/doctor.py +133 -0
  377. mlflow_tracing-3.0.0/mlflow/utils/download_cloud_file_chunk.py +43 -0
  378. mlflow_tracing-3.0.0/mlflow/utils/env_manager.py +16 -0
  379. mlflow_tracing-3.0.0/mlflow/utils/env_pack.py +131 -0
  380. mlflow_tracing-3.0.0/mlflow/utils/environment.py +1011 -0
  381. mlflow_tracing-3.0.0/mlflow/utils/exception_utils.py +14 -0
  382. mlflow_tracing-3.0.0/mlflow/utils/file_utils.py +970 -0
  383. mlflow_tracing-3.0.0/mlflow/utils/git_utils.py +77 -0
  384. mlflow_tracing-3.0.0/mlflow/utils/gorilla.py +797 -0
  385. mlflow_tracing-3.0.0/mlflow/utils/import_hooks/__init__.py +363 -0
  386. mlflow_tracing-3.0.0/mlflow/utils/lazy_load.py +51 -0
  387. mlflow_tracing-3.0.0/mlflow/utils/logging_utils.py +168 -0
  388. mlflow_tracing-3.0.0/mlflow/utils/mime_type_utils.py +58 -0
  389. mlflow_tracing-3.0.0/mlflow/utils/mlflow_tags.py +103 -0
  390. mlflow_tracing-3.0.0/mlflow/utils/model_utils.py +486 -0
  391. mlflow_tracing-3.0.0/mlflow/utils/name_utils.py +346 -0
  392. mlflow_tracing-3.0.0/mlflow/utils/nfs_on_spark.py +62 -0
  393. mlflow_tracing-3.0.0/mlflow/utils/openai_utils.py +164 -0
  394. mlflow_tracing-3.0.0/mlflow/utils/os.py +12 -0
  395. mlflow_tracing-3.0.0/mlflow/utils/oss_registry_utils.py +29 -0
  396. mlflow_tracing-3.0.0/mlflow/utils/plugins.py +17 -0
  397. mlflow_tracing-3.0.0/mlflow/utils/process.py +182 -0
  398. mlflow_tracing-3.0.0/mlflow/utils/promptlab_utils.py +146 -0
  399. mlflow_tracing-3.0.0/mlflow/utils/proto_json_utils.py +743 -0
  400. mlflow_tracing-3.0.0/mlflow/utils/pydantic_utils.py +54 -0
  401. mlflow_tracing-3.0.0/mlflow/utils/request_utils.py +279 -0
  402. mlflow_tracing-3.0.0/mlflow/utils/requirements_utils.py +704 -0
  403. mlflow_tracing-3.0.0/mlflow/utils/rest_utils.py +714 -0
  404. mlflow_tracing-3.0.0/mlflow/utils/search_logged_model_utils.py +127 -0
  405. mlflow_tracing-3.0.0/mlflow/utils/search_utils.py +2074 -0
  406. mlflow_tracing-3.0.0/mlflow/utils/server_cli_utils.py +61 -0
  407. mlflow_tracing-3.0.0/mlflow/utils/spark_utils.py +15 -0
  408. mlflow_tracing-3.0.0/mlflow/utils/string_utils.py +138 -0
  409. mlflow_tracing-3.0.0/mlflow/utils/thread_utils.py +55 -0
  410. mlflow_tracing-3.0.0/mlflow/utils/time.py +54 -0
  411. mlflow_tracing-3.0.0/mlflow/utils/timeout.py +42 -0
  412. mlflow_tracing-3.0.0/mlflow/utils/uri.py +572 -0
  413. mlflow_tracing-3.0.0/mlflow/utils/validation.py +662 -0
  414. mlflow_tracing-3.0.0/mlflow/utils/virtualenv.py +458 -0
  415. mlflow_tracing-3.0.0/mlflow/utils/warnings_utils.py +25 -0
  416. mlflow_tracing-3.0.0/mlflow/utils/yaml_utils.py +179 -0
  417. mlflow_tracing-3.0.0/mlflow/version.py +24 -0
  418. mlflow_tracing-3.0.0/mlflow_tracing.egg-info/PKG-INFO +121 -0
  419. mlflow_tracing-3.0.0/mlflow_tracing.egg-info/SOURCES.txt +422 -0
  420. mlflow_tracing-3.0.0/mlflow_tracing.egg-info/dependency_links.txt +1 -0
  421. mlflow_tracing-3.0.0/mlflow_tracing.egg-info/requires.txt +7 -0
  422. mlflow_tracing-3.0.0/mlflow_tracing.egg-info/top_level.txt +1 -0
  423. mlflow_tracing-3.0.0/pyproject.toml +92 -0
  424. mlflow_tracing-3.0.0/setup.cfg +4 -0
@@ -0,0 +1,121 @@
1
+ Metadata-Version: 2.4
2
+ Name: mlflow-tracing
3
+ Version: 3.0.0
4
+ Summary: MLflow Tracing SDK is an open-source, lightweight Python package that only includes the minimum set of dependencies and functionality to instrument your code/models/agents with MLflow Tracing.
5
+ Maintainer-email: Databricks <mlflow-oss-maintainers@googlegroups.com>
6
+ Project-URL: homepage, https://mlflow.org
7
+ Project-URL: issues, https://github.com/mlflow/mlflow/issues
8
+ Project-URL: documentation, https://mlflow.org/docs/latest/index.html
9
+ Project-URL: repository, https://github.com/mlflow/mlflow
10
+ Keywords: mlflow,ai,databricks,observability,monitoring
11
+ Classifier: Development Status :: 5 - Production/Stable
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: End Users/Desktop
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Intended Audience :: Information Technology
16
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
17
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
18
+ Classifier: License :: OSI Approved :: Apache Software License
19
+ Classifier: Operating System :: OS Independent
20
+ Classifier: Programming Language :: Python :: 3.9
21
+ Requires-Python: >=3.9
22
+ Description-Content-Type: text/markdown
23
+ Requires-Dist: cachetools<7,>=5.0.0
24
+ Requires-Dist: databricks-sdk<1,>=0.20.0
25
+ Requires-Dist: opentelemetry-api<3,>=1.9.0
26
+ Requires-Dist: opentelemetry-sdk<3,>=1.9.0
27
+ Requires-Dist: packaging<26
28
+ Requires-Dist: protobuf<7,>=3.12.0
29
+ Requires-Dist: pydantic<3,>=1.10.8
30
+
31
+ # MLflow Tracing: An Open-Source SDK for Observability and Monitoring GenAI Applications🔍
32
+
33
+ [![Latest Docs](https://img.shields.io/badge/docs-latest-success.svg?style=for-the-badge)](https://mlflow.org/docs/latest/index.html)
34
+ [![Apache 2 License](https://img.shields.io/badge/license-Apache%202-brightgreen.svg?style=for-the-badge&logo=apache)](https://github.com/mlflow/mlflow/blob/master/LICENSE.txt)
35
+ [![Slack](https://img.shields.io/badge/slack-@mlflow--users-CF0E5B.svg?logo=slack&logoColor=white&labelColor=3F0E40&style=for-the-badge)](https://mlflow.org/community/#slack)
36
+ [![Twitter](https://img.shields.io/twitter/follow/MLflow?style=for-the-badge&labelColor=00ACEE&logo=twitter&logoColor=white)](https://twitter.com/MLflow)
37
+
38
+ MLflow Tracing (`mlflow-tracing`) is an open-source, lightweight Python package that only includes the minimum set of dependencies and functionality
39
+ to instrument your code/models/agents with [MLflow Tracing Feature](https://mlflow.org/docs/latest/tracing). It is designed to be a perfect fit for production environments where you want:
40
+
41
+ - **⚡️ Faster Deployment**: The package size and dependencies are significantly smaller than the full MLflow package, allowing for faster deployment times in dynamic environments such as Docker containers, serverless functions, and cloud-based applications.
42
+ - **🔧 Simplified Dependency Management**: A smaller set of dependencies means less work keeping up with dependency updates, security patches, and breaking changes from upstream libraries.
43
+ - **📦 Portability**: With the less number of dependencies, MLflow Tracing can be easily deployed across different environments and platforms, without worrying about compatibility issues.
44
+ - **🔒 Fewer Security Risks**: Each dependency potentially introduces security vulnerabilities. By reducing the number of dependencies, MLflow Tracing minimizes the attack surface and reduces the risk of security breaches.
45
+
46
+ ## ✨ Features
47
+
48
+ - [Automatic Tracing](https://mlflow.org/docs/latest/tracing/integrations/) for AI libraries (OpenAI, LangChain, DSPy, Anthropic, etc...). Follow the link for the full list of supported libraries.
49
+ - [Manual instrumentation APIs](https://mlflow.org/docs/latest/tracing/api/manual-instrumentation) such as `@trace` decorator.
50
+ - [Production Monitoring](https://mlflow.org/docs/latest/tracing/production)
51
+ - Other tracing APIs such as `mlflow.set_trace_tag`, `mlflow.search_traces`, etc.
52
+
53
+ ## 🌐 Choose Backend
54
+
55
+ The MLflow Trace package is designed to work with a remote hosted MLflow server as a backend. This allows you to log your traces to a central location, making it easier to manage and analyze your traces. There are several different options for hosting your MLflow server, including:
56
+
57
+ - [Databricks](https://docs.databricks.com/machine-learning/mlflow/managed-mlflow.html) - Databricks offers a FREE, fully managed MLflow server as a part of their platform. This is the easiest way to get started with MLflow tracing, without having to set up any infrastructure.
58
+ - [Amazon SageMaker](https://aws.amazon.com/sagemaker-ai/experiments/) - MLflow on Amazon SageMaker is a fully managed service offered as part of the SageMaker platform by AWS, including tracing and other MLflow features such as model registry.
59
+ - [Nebius](https://nebius.com/) - Nebius, a cutting-edge cloud platform for GenAI explorers, offers a fully managed MLflow server.
60
+ - [Self-hosting](https://mlflow.org/docs/latest/tracking/#tracking_setup) - MLflow is a fully open-source project, allowing you to self-host your own MLflow server and keep your data private. This is a great option if you want to have full control over your data and infrastructure.
61
+
62
+ ## 🚀 Getting Started
63
+
64
+ ### Installation
65
+
66
+ To install the MLflow Python package, run the following command:
67
+
68
+ ```bash
69
+ pip install mlflow-tracing
70
+ ```
71
+
72
+ To install from the source code, run the following command:
73
+
74
+ ```bash
75
+ pip install git+https://github.com/mlflow/mlflow.git#subdirectory=packages/tracing
76
+ ```
77
+
78
+ > **NOTE:** It is **not** recommended to co-install this package with the full MLflow package together, as it may cause version mismatches issues.
79
+
80
+ ### Connect to the MLflow Server
81
+
82
+ To connect to your MLflow server to log your traces, set the `MLFLOW_TRACKING_URI` environment variable or use the `mlflow.set_tracking_uri` function:
83
+
84
+ ```python
85
+ import mlflow
86
+
87
+ mlflow.set_tracking_uri("databricks")
88
+ # Specify the experiment to log the traces to
89
+ mlflow.set_experiment("/Path/To/Experiment")
90
+ ```
91
+
92
+ ### Start Logging Traces
93
+
94
+ ```python
95
+ import openai
96
+
97
+ client = openai.OpenAI(api_key="<your-api-key>")
98
+
99
+ # Enable auto-tracing for OpenAI
100
+ mlflow.openai.autolog()
101
+
102
+ # Call the OpenAI API as usual
103
+ response = client.chat.completions.create(
104
+ model="gpt-4.1-mini",
105
+ messages=[{"role": "user", "content": "Hello, how are you?"}],
106
+ )
107
+ ```
108
+
109
+ ## 📘 Documentation
110
+
111
+ Official documentation for MLflow Tracing can be found at [here](https://mlflow.org/docs/latest/tracing).
112
+
113
+ ## 🛑 Features _Not_ Included
114
+
115
+ The following MLflow features are not included in this package.
116
+
117
+ - MLflow tracking server and UI.
118
+ - MLflow's other tracking capabilities such as Runs, Model Registry, Projects, etc.
119
+ - Evaluate models/agents and log evaluation results.
120
+
121
+ To leverage the full feature set of MLflow, install the full package by running `pip install mlflow`.
@@ -0,0 +1,91 @@
1
+ # MLflow Tracing: An Open-Source SDK for Observability and Monitoring GenAI Applications🔍
2
+
3
+ [![Latest Docs](https://img.shields.io/badge/docs-latest-success.svg?style=for-the-badge)](https://mlflow.org/docs/latest/index.html)
4
+ [![Apache 2 License](https://img.shields.io/badge/license-Apache%202-brightgreen.svg?style=for-the-badge&logo=apache)](https://github.com/mlflow/mlflow/blob/master/LICENSE.txt)
5
+ [![Slack](https://img.shields.io/badge/slack-@mlflow--users-CF0E5B.svg?logo=slack&logoColor=white&labelColor=3F0E40&style=for-the-badge)](https://mlflow.org/community/#slack)
6
+ [![Twitter](https://img.shields.io/twitter/follow/MLflow?style=for-the-badge&labelColor=00ACEE&logo=twitter&logoColor=white)](https://twitter.com/MLflow)
7
+
8
+ MLflow Tracing (`mlflow-tracing`) is an open-source, lightweight Python package that only includes the minimum set of dependencies and functionality
9
+ to instrument your code/models/agents with [MLflow Tracing Feature](https://mlflow.org/docs/latest/tracing). It is designed to be a perfect fit for production environments where you want:
10
+
11
+ - **⚡️ Faster Deployment**: The package size and dependencies are significantly smaller than the full MLflow package, allowing for faster deployment times in dynamic environments such as Docker containers, serverless functions, and cloud-based applications.
12
+ - **🔧 Simplified Dependency Management**: A smaller set of dependencies means less work keeping up with dependency updates, security patches, and breaking changes from upstream libraries.
13
+ - **📦 Portability**: With the less number of dependencies, MLflow Tracing can be easily deployed across different environments and platforms, without worrying about compatibility issues.
14
+ - **🔒 Fewer Security Risks**: Each dependency potentially introduces security vulnerabilities. By reducing the number of dependencies, MLflow Tracing minimizes the attack surface and reduces the risk of security breaches.
15
+
16
+ ## ✨ Features
17
+
18
+ - [Automatic Tracing](https://mlflow.org/docs/latest/tracing/integrations/) for AI libraries (OpenAI, LangChain, DSPy, Anthropic, etc...). Follow the link for the full list of supported libraries.
19
+ - [Manual instrumentation APIs](https://mlflow.org/docs/latest/tracing/api/manual-instrumentation) such as `@trace` decorator.
20
+ - [Production Monitoring](https://mlflow.org/docs/latest/tracing/production)
21
+ - Other tracing APIs such as `mlflow.set_trace_tag`, `mlflow.search_traces`, etc.
22
+
23
+ ## 🌐 Choose Backend
24
+
25
+ The MLflow Trace package is designed to work with a remote hosted MLflow server as a backend. This allows you to log your traces to a central location, making it easier to manage and analyze your traces. There are several different options for hosting your MLflow server, including:
26
+
27
+ - [Databricks](https://docs.databricks.com/machine-learning/mlflow/managed-mlflow.html) - Databricks offers a FREE, fully managed MLflow server as a part of their platform. This is the easiest way to get started with MLflow tracing, without having to set up any infrastructure.
28
+ - [Amazon SageMaker](https://aws.amazon.com/sagemaker-ai/experiments/) - MLflow on Amazon SageMaker is a fully managed service offered as part of the SageMaker platform by AWS, including tracing and other MLflow features such as model registry.
29
+ - [Nebius](https://nebius.com/) - Nebius, a cutting-edge cloud platform for GenAI explorers, offers a fully managed MLflow server.
30
+ - [Self-hosting](https://mlflow.org/docs/latest/tracking/#tracking_setup) - MLflow is a fully open-source project, allowing you to self-host your own MLflow server and keep your data private. This is a great option if you want to have full control over your data and infrastructure.
31
+
32
+ ## 🚀 Getting Started
33
+
34
+ ### Installation
35
+
36
+ To install the MLflow Python package, run the following command:
37
+
38
+ ```bash
39
+ pip install mlflow-tracing
40
+ ```
41
+
42
+ To install from the source code, run the following command:
43
+
44
+ ```bash
45
+ pip install git+https://github.com/mlflow/mlflow.git#subdirectory=packages/tracing
46
+ ```
47
+
48
+ > **NOTE:** It is **not** recommended to co-install this package with the full MLflow package together, as it may cause version mismatches issues.
49
+
50
+ ### Connect to the MLflow Server
51
+
52
+ To connect to your MLflow server to log your traces, set the `MLFLOW_TRACKING_URI` environment variable or use the `mlflow.set_tracking_uri` function:
53
+
54
+ ```python
55
+ import mlflow
56
+
57
+ mlflow.set_tracking_uri("databricks")
58
+ # Specify the experiment to log the traces to
59
+ mlflow.set_experiment("/Path/To/Experiment")
60
+ ```
61
+
62
+ ### Start Logging Traces
63
+
64
+ ```python
65
+ import openai
66
+
67
+ client = openai.OpenAI(api_key="<your-api-key>")
68
+
69
+ # Enable auto-tracing for OpenAI
70
+ mlflow.openai.autolog()
71
+
72
+ # Call the OpenAI API as usual
73
+ response = client.chat.completions.create(
74
+ model="gpt-4.1-mini",
75
+ messages=[{"role": "user", "content": "Hello, how are you?"}],
76
+ )
77
+ ```
78
+
79
+ ## 📘 Documentation
80
+
81
+ Official documentation for MLflow Tracing can be found at [here](https://mlflow.org/docs/latest/tracing).
82
+
83
+ ## 🛑 Features _Not_ Included
84
+
85
+ The following MLflow features are not included in this package.
86
+
87
+ - MLflow tracking server and UI.
88
+ - MLflow's other tracking capabilities such as Runs, Model Registry, Projects, etc.
89
+ - Evaluate models/agents and log evaluation results.
90
+
91
+ To leverage the full feature set of MLflow, install the full package by running `pip install mlflow`.
@@ -0,0 +1,408 @@
1
+ """
2
+ The ``mlflow`` module provides a high-level "fluent" API for starting and managing MLflow runs.
3
+ For example:
4
+
5
+ .. code:: python
6
+
7
+ import mlflow
8
+
9
+ mlflow.start_run()
10
+ mlflow.log_param("my", "param")
11
+ mlflow.log_metric("score", 100)
12
+ mlflow.end_run()
13
+
14
+ You can also use the context manager syntax like this:
15
+
16
+ .. code:: python
17
+
18
+ with mlflow.start_run() as run:
19
+ mlflow.log_param("my", "param")
20
+ mlflow.log_metric("score", 100)
21
+
22
+ which automatically terminates the run at the end of the ``with`` block.
23
+
24
+ The fluent tracking API is not currently threadsafe. Any concurrent callers to the tracking API must
25
+ implement mutual exclusion manually.
26
+
27
+ For a lower level API, see the :py:mod:`mlflow.client` module.
28
+ """
29
+
30
+ import contextlib
31
+ from typing import TYPE_CHECKING
32
+
33
+ from mlflow.version import IS_TRACING_SDK_ONLY, VERSION
34
+
35
+ __version__ = VERSION
36
+
37
+ import mlflow.mismatch
38
+
39
+ # `check_version_mismatch` must be called here before importing any other modules
40
+ with contextlib.suppress(Exception):
41
+ mlflow.mismatch._check_version_mismatch()
42
+
43
+ if not IS_TRACING_SDK_ONLY:
44
+ from mlflow import (
45
+ artifacts, # noqa: F401
46
+ client, # noqa: F401
47
+ config, # noqa: F401
48
+ data, # noqa: F401
49
+ exceptions, # noqa: F401
50
+ genai, # noqa: F401
51
+ models, # noqa: F401
52
+ projects, # noqa: F401
53
+ tracking, # noqa: F401
54
+ )
55
+
56
+ from mlflow import tracing # noqa: F401
57
+ from mlflow.environment_variables import MLFLOW_CONFIGURE_LOGGING
58
+ from mlflow.exceptions import MlflowException
59
+ from mlflow.utils.lazy_load import LazyLoader
60
+ from mlflow.utils.logging_utils import _configure_mlflow_loggers
61
+
62
+ # Lazily load mlflow flavors to avoid excessive dependencies.
63
+ anthropic = LazyLoader("mlflow.anthropic", globals(), "mlflow.anthropic")
64
+ ag2 = LazyLoader("mlflow.ag2", globals(), "mlflow.ag2")
65
+ autogen = LazyLoader("mlflow.autogen", globals(), "mlflow.autogen")
66
+ bedrock = LazyLoader("mlflow.bedrock", globals(), "mlflow.bedrock")
67
+ catboost = LazyLoader("mlflow.catboost", globals(), "mlflow.catboost")
68
+ crewai = LazyLoader("mlflow.crewai", globals(), "mlflow.crewai")
69
+ diviner = LazyLoader("mlflow.diviner", globals(), "mlflow.diviner")
70
+ dspy = LazyLoader("mlflow.dspy", globals(), "mlflow.dspy")
71
+ gemini = LazyLoader("mlflow.gemini", globals(), "mlflow.gemini")
72
+ groq = LazyLoader("mlflow.groq", globals(), "mlflow.groq")
73
+ h2o = LazyLoader("mlflow.h2o", globals(), "mlflow.h2o")
74
+ johnsnowlabs = LazyLoader("mlflow.johnsnowlabs", globals(), "mlflow.johnsnowlabs")
75
+ keras = LazyLoader("mlflow.keras", globals(), "mlflow.keras")
76
+ langchain = LazyLoader("mlflow.langchain", globals(), "mlflow.langchain")
77
+ lightgbm = LazyLoader("mlflow.lightgbm", globals(), "mlflow.lightgbm")
78
+ litellm = LazyLoader("mlflow.litellm", globals(), "mlflow.litellm")
79
+ llama_index = LazyLoader("mlflow.llama_index", globals(), "mlflow.llama_index")
80
+ llm = LazyLoader("mlflow.llm", globals(), "mlflow.llm")
81
+ metrics = LazyLoader("mlflow.metrics", globals(), "mlflow.metrics")
82
+ mistral = LazyLoader("mlflow.mistral", globals(), "mlflow.mistral")
83
+ onnx = LazyLoader("mlflow.onnx", globals(), "mlflow.onnx")
84
+ openai = LazyLoader("mlflow.openai", globals(), "mlflow.openai")
85
+ paddle = LazyLoader("mlflow.paddle", globals(), "mlflow.paddle")
86
+ pmdarima = LazyLoader("mlflow.pmdarima", globals(), "mlflow.pmdarima")
87
+ promptflow = LazyLoader("mlflow.promptflow", globals(), "mlflow.promptflow")
88
+ prophet = LazyLoader("mlflow.prophet", globals(), "mlflow.prophet")
89
+ pydantic_ai = LazyLoader("mlflow.pydantic_ai", globals(), "mlflow.pydantic_ai")
90
+ pyfunc = LazyLoader("mlflow.pyfunc", globals(), "mlflow.pyfunc")
91
+ pyspark = LazyLoader("mlflow.pyspark", globals(), "mlflow.pyspark")
92
+ pytorch = LazyLoader("mlflow.pytorch", globals(), "mlflow.pytorch")
93
+ rfunc = LazyLoader("mlflow.rfunc", globals(), "mlflow.rfunc")
94
+ sentence_transformers = LazyLoader(
95
+ "mlflow.sentence_transformers",
96
+ globals(),
97
+ "mlflow.sentence_transformers",
98
+ )
99
+ shap = LazyLoader("mlflow.shap", globals(), "mlflow.shap")
100
+ sklearn = LazyLoader("mlflow.sklearn", globals(), "mlflow.sklearn")
101
+ smolagents = LazyLoader("mlflow.smolagents", globals(), "mlflow.smolagents")
102
+ spacy = LazyLoader("mlflow.spacy", globals(), "mlflow.spacy")
103
+ spark = LazyLoader("mlflow.spark", globals(), "mlflow.spark")
104
+ statsmodels = LazyLoader("mlflow.statsmodels", globals(), "mlflow.statsmodels")
105
+ tensorflow = LazyLoader("mlflow.tensorflow", globals(), "mlflow.tensorflow")
106
+ # TxtAI integration is defined at https://github.com/neuml/mlflow-txtai
107
+ txtai = LazyLoader("mlflow.txtai", globals(), "mlflow_txtai")
108
+ transformers = LazyLoader("mlflow.transformers", globals(), "mlflow.transformers")
109
+ xgboost = LazyLoader("mlflow.xgboost", globals(), "mlflow.xgboost")
110
+
111
+ if TYPE_CHECKING:
112
+ # Do not move this block above the lazy-loaded modules above.
113
+ # All the lazy-loaded modules above must be imported here for code completion to work in IDEs.
114
+ from mlflow import ( # noqa: F401
115
+ ag2,
116
+ anthropic,
117
+ autogen,
118
+ bedrock,
119
+ catboost,
120
+ crewai,
121
+ diviner,
122
+ dspy,
123
+ gemini,
124
+ groq,
125
+ h2o,
126
+ johnsnowlabs,
127
+ keras,
128
+ langchain,
129
+ lightgbm,
130
+ litellm,
131
+ llama_index,
132
+ llm,
133
+ metrics,
134
+ mistral,
135
+ onnx,
136
+ openai,
137
+ paddle,
138
+ pmdarima,
139
+ promptflow,
140
+ prophet,
141
+ pydantic_ai,
142
+ pyfunc,
143
+ pyspark,
144
+ pytorch,
145
+ rfunc,
146
+ sentence_transformers,
147
+ shap,
148
+ sklearn,
149
+ smolagents,
150
+ spacy,
151
+ spark,
152
+ statsmodels,
153
+ tensorflow,
154
+ transformers,
155
+ xgboost,
156
+ )
157
+
158
+ if MLFLOW_CONFIGURE_LOGGING.get() is True:
159
+ _configure_mlflow_loggers(root_module_name=__name__)
160
+
161
+ # Core modules required for mlflow-tracing
162
+ from mlflow.tracing.assessment import (
163
+ delete_assessment,
164
+ get_assessment,
165
+ log_assessment,
166
+ log_expectation,
167
+ log_feedback,
168
+ override_feedback,
169
+ update_assessment,
170
+ )
171
+ from mlflow.tracing.fluent import (
172
+ add_trace,
173
+ delete_trace_tag,
174
+ get_active_trace_id,
175
+ get_current_active_span,
176
+ get_last_active_trace_id,
177
+ get_trace,
178
+ log_trace,
179
+ search_traces,
180
+ set_trace_tag,
181
+ start_span,
182
+ start_span_no_context,
183
+ trace,
184
+ update_current_trace,
185
+ )
186
+ from mlflow.tracking import (
187
+ get_tracking_uri,
188
+ is_tracking_uri_set,
189
+ set_tracking_uri,
190
+ )
191
+ from mlflow.tracking.fluent import active_run, flush_trace_async_logging, set_experiment
192
+
193
+ # These are minimal set of APIs to be exposed via `mlflow-tracing` package.
194
+ # APIs listed here must not depend on dependencies that are not part of `mlflow-tracing` package.
195
+ __all__ = [
196
+ "MlflowException",
197
+ # Minimal tracking APIs required for tracing core functionality
198
+ "set_experiment",
199
+ "set_tracking_uri",
200
+ "get_tracking_uri",
201
+ "is_tracking_uri_set",
202
+ # NB: Tracing SDK doesn't support using Runs, however, active_run is used heavily within
203
+ # the autologging code base.
204
+ "active_run",
205
+ # Tracing APIs
206
+ "add_trace",
207
+ "delete_trace_tag",
208
+ "flush_trace_async_logging",
209
+ "get_active_trace_id",
210
+ "get_current_active_span",
211
+ "get_last_active_trace_id",
212
+ "get_trace",
213
+ "log_trace",
214
+ "search_traces",
215
+ "set_trace_tag",
216
+ "start_span",
217
+ "start_span_no_context",
218
+ "trace",
219
+ "update_current_trace",
220
+ # Assessment APIs
221
+ "get_assessment",
222
+ "delete_assessment",
223
+ "log_assessment",
224
+ "update_assessment",
225
+ "log_expectation",
226
+ "log_feedback",
227
+ "override_feedback",
228
+ ]
229
+
230
+ # Only import these modules when mlflow or mlflow-skinny is installed i.e. not importing them
231
+ # when only mlflow-tracing is installed.
232
+ if not IS_TRACING_SDK_ONLY:
233
+ from mlflow.client import MlflowClient
234
+
235
+ # For backward compatibility, we expose the following functions and classes at the top level in
236
+ # addition to `mlflow.config`.
237
+ from mlflow.config import (
238
+ disable_system_metrics_logging,
239
+ enable_system_metrics_logging,
240
+ get_registry_uri,
241
+ set_registry_uri,
242
+ set_system_metrics_node_id,
243
+ set_system_metrics_samples_before_logging,
244
+ set_system_metrics_sampling_interval,
245
+ )
246
+ from mlflow.models.evaluation.deprecated import evaluate
247
+ from mlflow.models.evaluation.validation import validate_evaluation_results
248
+ from mlflow.projects import run
249
+ from mlflow.tracking._model_registry.fluent import (
250
+ # TODO: Prompt Registry APIs are moved to the `mlflow.genai` namespace and direct
251
+ # imports from mlflow will be deprecated in the future.
252
+ delete_prompt_alias,
253
+ load_prompt,
254
+ register_model,
255
+ register_prompt,
256
+ search_model_versions,
257
+ search_prompts,
258
+ search_registered_models,
259
+ set_model_version_tag,
260
+ set_prompt_alias,
261
+ )
262
+ from mlflow.tracking.fluent import (
263
+ ActiveModel,
264
+ ActiveRun,
265
+ autolog,
266
+ clear_active_model,
267
+ create_experiment,
268
+ create_external_model,
269
+ delete_experiment,
270
+ delete_logged_model_tag,
271
+ delete_run,
272
+ delete_tag,
273
+ end_run,
274
+ finalize_logged_model,
275
+ flush_artifact_async_logging,
276
+ flush_async_logging,
277
+ get_active_model_id,
278
+ get_artifact_uri,
279
+ get_experiment,
280
+ get_experiment_by_name,
281
+ get_logged_model,
282
+ get_parent_run,
283
+ get_run,
284
+ initialize_logged_model,
285
+ last_active_run,
286
+ last_logged_model,
287
+ load_table,
288
+ log_artifact,
289
+ log_artifacts,
290
+ log_dict,
291
+ log_figure,
292
+ log_image,
293
+ log_input,
294
+ log_inputs,
295
+ log_metric,
296
+ log_metrics,
297
+ log_model_params,
298
+ log_outputs,
299
+ log_param,
300
+ log_params,
301
+ log_table,
302
+ log_text,
303
+ search_experiments,
304
+ search_logged_models,
305
+ search_runs,
306
+ set_active_model,
307
+ set_experiment_tag,
308
+ set_experiment_tags,
309
+ set_logged_model_tags,
310
+ set_tag,
311
+ set_tags,
312
+ start_run,
313
+ )
314
+ from mlflow.tracking.multimedia import Image
315
+ from mlflow.utils.async_logging.run_operations import RunOperations # noqa: F401
316
+ from mlflow.utils.credentials import login
317
+ from mlflow.utils.doctor import doctor
318
+
319
+ __all__ += [
320
+ "ActiveRun",
321
+ "ActiveModel",
322
+ "MlflowClient",
323
+ "MlflowException",
324
+ "autolog",
325
+ "clear_active_model",
326
+ "create_experiment",
327
+ "create_external_model",
328
+ "delete_experiment",
329
+ "delete_run",
330
+ "delete_tag",
331
+ "disable_system_metrics_logging",
332
+ "doctor",
333
+ "enable_system_metrics_logging",
334
+ "end_run",
335
+ "evaluate",
336
+ "finalize_logged_model",
337
+ "flush_async_logging",
338
+ "flush_artifact_async_logging",
339
+ "get_active_model_id",
340
+ "get_artifact_uri",
341
+ "get_experiment",
342
+ "get_experiment_by_name",
343
+ "get_logged_model",
344
+ "get_parent_run",
345
+ "get_registry_uri",
346
+ "get_run",
347
+ "initialize_logged_model",
348
+ "last_active_run",
349
+ "last_logged_model",
350
+ "load_table",
351
+ "log_artifact",
352
+ "log_artifacts",
353
+ "log_dict",
354
+ "log_figure",
355
+ "log_image",
356
+ "log_input",
357
+ "log_inputs",
358
+ "log_model_params",
359
+ "log_outputs",
360
+ "log_metric",
361
+ "log_metrics",
362
+ "log_param",
363
+ "log_params",
364
+ "log_table",
365
+ "log_text",
366
+ "login",
367
+ "pyfunc",
368
+ "register_model",
369
+ "run",
370
+ "search_experiments",
371
+ "search_logged_models",
372
+ "search_model_versions",
373
+ "search_registered_models",
374
+ "search_runs",
375
+ "search_prompts",
376
+ "set_active_model",
377
+ "set_experiment_tag",
378
+ "set_experiment_tags",
379
+ "set_model_version_tag",
380
+ "set_registry_uri",
381
+ "set_system_metrics_node_id",
382
+ "set_system_metrics_samples_before_logging",
383
+ "set_system_metrics_sampling_interval",
384
+ "set_tag",
385
+ "set_tags",
386
+ "start_run",
387
+ "validate_evaluation_results",
388
+ "Image",
389
+ # Prompt Registry APIs
390
+ # TODO: Prompt Registry APIs are moved to the `mlflow.genai` namespace and direct
391
+ # imports from mlflow will be deprecated in the future.
392
+ "load_prompt",
393
+ "register_prompt",
394
+ "search_prompts",
395
+ "set_prompt_alias",
396
+ "delete_prompt_alias",
397
+ "set_logged_model_tags",
398
+ "delete_logged_model_tag",
399
+ ]
400
+
401
+
402
+ # `mlflow.gateway` depends on optional dependencies such as pydantic, psutil, and has version
403
+ # restrictions for dependencies. Importing this module fails if they are not installed or
404
+ # if invalid versions of these required packages are installed.
405
+ with contextlib.suppress(Exception):
406
+ from mlflow import gateway # noqa: F401
407
+
408
+ __all__.append("gateway")
@@ -0,0 +1,3 @@
1
+ from mlflow.cli import cli
2
+
3
+ cli.main()
@@ -0,0 +1,40 @@
1
+ from mlflow.anthropic.autolog import async_patched_class_call, patched_class_call
2
+ from mlflow.utils.annotations import experimental
3
+ from mlflow.utils.autologging_utils import autologging_integration, safe_patch
4
+
5
+ FLAVOR_NAME = "anthropic"
6
+
7
+
8
+ @experimental
9
+ @autologging_integration(FLAVOR_NAME)
10
+ def autolog(
11
+ log_traces: bool = True,
12
+ disable: bool = False,
13
+ silent: bool = False,
14
+ ):
15
+ """
16
+ Enables (or disables) and configures autologging from Anthropic to MLflow.
17
+ Only synchronous calls and asynchnorous APIs are supported. Streaming is not recorded.
18
+
19
+ Args:
20
+ log_traces: If ``True``, traces are logged for Anthropic models.
21
+ If ``False``, no traces are collected during inference. Default to ``True``.
22
+ disable: If ``True``, disables the Anthropic autologging. Default to ``False``.
23
+ silent: If ``True``, suppress all event logs and warnings from MLflow during Anthropic
24
+ autologging. If ``False``, show all events and warnings.
25
+ """
26
+ from anthropic.resources import AsyncMessages, Messages
27
+
28
+ safe_patch(
29
+ FLAVOR_NAME,
30
+ Messages,
31
+ "create",
32
+ patched_class_call,
33
+ )
34
+
35
+ safe_patch(
36
+ FLAVOR_NAME,
37
+ AsyncMessages,
38
+ "create",
39
+ async_patched_class_call,
40
+ )