blaxel 0.1.22rc70__py3-none-any.whl → 0.2.0rc1__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 (419) hide show
  1. blaxel/__init__.py +6 -3
  2. blaxel/core/__init__.py +44 -0
  3. blaxel/{agents → core/agents}/__init__.py +30 -50
  4. blaxel/{authentication → core/authentication}/apikey.py +1 -0
  5. blaxel/{authentication → core/authentication}/clientcredentials.py +6 -2
  6. blaxel/{client → core/client}/api/agents/create_agent.py +1 -1
  7. blaxel/{client → core/client}/api/agents/update_agent.py +1 -1
  8. blaxel/{client → core/client}/api/compute/create_sandbox.py +1 -1
  9. blaxel/{client → core/client}/api/compute/create_sandbox_preview.py +1 -1
  10. blaxel/{client → core/client}/api/compute/create_sandbox_preview_token.py +1 -1
  11. blaxel/{client → core/client}/api/compute/update_sandbox.py +1 -1
  12. blaxel/{client → core/client}/api/compute/update_sandbox_preview.py +1 -1
  13. blaxel/{client → core/client}/api/functions/create_function.py +1 -1
  14. blaxel/{client → core/client}/api/functions/update_function.py +1 -1
  15. blaxel/{client → core/client}/api/integrations/create_integration_connection.py +1 -1
  16. blaxel/{client → core/client}/api/integrations/update_integration_connection.py +1 -1
  17. blaxel/{client → core/client}/api/jobs/create_job.py +1 -1
  18. blaxel/{client → core/client}/api/jobs/update_job.py +1 -1
  19. blaxel/{client → core/client}/api/knowledgebases/create_knowledgebase.py +1 -1
  20. blaxel/{client → core/client}/api/knowledgebases/update_knowledgebase.py +1 -1
  21. blaxel/{client → core/client}/api/models/create_model.py +1 -1
  22. blaxel/{client → core/client}/api/models/update_model.py +1 -1
  23. blaxel/{client → core/client}/api/policies/create_policy.py +1 -1
  24. blaxel/{client → core/client}/api/policies/update_policy.py +1 -1
  25. blaxel/{client → core/client}/api/service_accounts/create_api_key_for_service_account.py +1 -1
  26. blaxel/{client → core/client}/api/service_accounts/create_workspace_service_account.py +1 -1
  27. blaxel/{client → core/client}/api/service_accounts/update_workspace_service_account.py +1 -1
  28. blaxel/{client → core/client}/api/workspaces/check_workspace_availability.py +1 -1
  29. blaxel/{client → core/client}/api/workspaces/create_worspace.py +1 -1
  30. blaxel/{client → core/client}/api/workspaces/invite_workspace_user.py +1 -1
  31. blaxel/{client → core/client}/api/workspaces/update_workspace.py +1 -1
  32. blaxel/{client → core/client}/api/workspaces/update_workspace_user_role.py +1 -1
  33. blaxel/{client → core/client}/models/agent.py +1 -1
  34. blaxel/{client → core/client}/models/agent_spec.py +2 -2
  35. blaxel/{client → core/client}/models/core_spec.py +1 -1
  36. blaxel/{client → core/client}/models/function.py +1 -1
  37. blaxel/{client → core/client}/models/function_schema_properties.py +1 -1
  38. blaxel/{client → core/client}/models/function_spec.py +2 -2
  39. blaxel/{client → core/client}/models/integration.py +2 -2
  40. blaxel/{client → core/client}/models/integration_endpoints.py +1 -1
  41. blaxel/{client → core/client}/models/job.py +1 -1
  42. blaxel/{client → core/client}/models/job_spec.py +1 -1
  43. blaxel/{client → core/client}/models/knowledgebase.py +1 -1
  44. blaxel/{client → core/client}/models/location_response.py +1 -1
  45. blaxel/{client → core/client}/models/model.py +1 -1
  46. blaxel/{client → core/client}/models/model_spec.py +1 -1
  47. blaxel/{client → core/client}/models/policy_spec.py +2 -2
  48. blaxel/{client → core/client}/models/resource_metrics.py +3 -3
  49. blaxel/{client → core/client}/models/runtime.py +1 -1
  50. blaxel/{client → core/client}/models/sandbox.py +1 -1
  51. blaxel/{client → core/client}/models/sandbox_definition.py +1 -1
  52. blaxel/{client → core/client}/models/sandbox_spec.py +1 -1
  53. blaxel/{client → core/client}/models/store_agent.py +1 -1
  54. blaxel/{client → core/client}/models/store_configuration.py +1 -1
  55. blaxel/{client → core/client}/models/template.py +1 -1
  56. blaxel/core/common/__init__.py +6 -0
  57. blaxel/core/common/autoload.py +21 -0
  58. blaxel/{common → core/common}/internal.py +33 -62
  59. blaxel/core/common/logger.py +131 -0
  60. blaxel/{jobs → core/jobs}/__init__.py +40 -60
  61. blaxel/core/mcp/__init__.py +4 -0
  62. blaxel/{mcp → core/mcp}/client.py +13 -7
  63. blaxel/{mcp → core/mcp}/server.py +76 -17
  64. blaxel/core/models/__init__.py +52 -0
  65. blaxel/core/sandbox/__init__.py +29 -0
  66. blaxel/core/sandbox/action.py +79 -0
  67. blaxel/{sandbox → core/sandbox}/client/api/filesystem/put_filesystem_path.py +1 -1
  68. blaxel/{sandbox → core/sandbox}/client/api/network/post_network_process_pid_monitor.py +1 -1
  69. blaxel/core/sandbox/client/api/process/__init__.py +0 -0
  70. blaxel/{sandbox → core/sandbox}/client/api/process/post_process.py +1 -1
  71. blaxel/{sandbox → core/sandbox}/client/models/directory.py +2 -2
  72. blaxel/core/sandbox/filesystem.py +280 -0
  73. blaxel/core/sandbox/network.py +10 -0
  74. blaxel/{sandbox → core/sandbox}/preview.py +45 -17
  75. blaxel/core/sandbox/process.py +159 -0
  76. blaxel/{sandbox → core/sandbox}/sandbox.py +62 -5
  77. blaxel/core/sandbox/session.py +124 -0
  78. blaxel/core/sandbox/types.py +103 -0
  79. blaxel/{tools → core/tools}/__init__.py +62 -90
  80. blaxel/crewai/__init__.py +4 -0
  81. blaxel/{models/crewai.py → crewai/model.py} +4 -2
  82. blaxel/crewai/py.typed +0 -0
  83. blaxel/crewai/tools.py +26 -0
  84. blaxel/googleadk/__init__.py +4 -0
  85. blaxel/{models/googleadk.py → googleadk/model.py} +8 -2
  86. blaxel/googleadk/py.typed +0 -0
  87. blaxel/googleadk/tools.py +72 -0
  88. blaxel/langgraph/__init__.py +4 -0
  89. blaxel/{models/langchain.py → langgraph/model.py} +8 -4
  90. blaxel/langgraph/py.typed +0 -0
  91. blaxel/{tools/langchain.py → langgraph/tools.py} +7 -3
  92. blaxel/livekit/__init__.py +4 -0
  93. blaxel/{models/livekit.py → livekit/model.py} +7 -1
  94. blaxel/livekit/py.typed +0 -0
  95. blaxel/{tools/livekit.py → livekit/tools.py} +8 -1
  96. blaxel/llamaindex/__init__.py +4 -0
  97. blaxel/{models/llamaindex.py → llamaindex/model.py} +6 -3
  98. blaxel/llamaindex/py.typed +0 -0
  99. blaxel/{tools/llamaindex.py → llamaindex/tools.py} +7 -4
  100. blaxel/openai/__init__.py +4 -0
  101. blaxel/{models/openai.py → openai/model.py} +4 -2
  102. blaxel/openai/py.typed +0 -0
  103. blaxel/{tools/openai.py → openai/tools.py} +7 -3
  104. blaxel/pydantic/__init__.py +4 -0
  105. blaxel/{models/custom/pydantic → pydantic/custom}/gemini.py +0 -1
  106. blaxel/{models/pydantic.py → pydantic/model.py} +6 -4
  107. blaxel/pydantic/py.typed +0 -0
  108. blaxel/{tools/pydantic.py → pydantic/tools.py} +6 -3
  109. blaxel/telemetry/__init__.py +6 -0
  110. blaxel/telemetry/instrumentation/blaxel_core.py +124 -0
  111. blaxel/telemetry/instrumentation/blaxel_langgraph.py +74 -0
  112. blaxel/telemetry/instrumentation/blaxel_langgraph_gemini.py +360 -0
  113. blaxel/telemetry/instrumentation/blaxel_llamaindex.py +89 -0
  114. blaxel/telemetry/instrumentation/map.py +61 -0
  115. blaxel/telemetry/instrumentation/utils.py +74 -0
  116. blaxel/{common → telemetry/log}/logger.py +6 -12
  117. blaxel/{instrumentation → telemetry}/manager.py +20 -12
  118. blaxel/telemetry/py.typed +0 -0
  119. blaxel/{instrumentation → telemetry}/span.py +12 -1
  120. blaxel-0.2.0rc1.dist-info/METADATA +224 -0
  121. blaxel-0.2.0rc1.dist-info/RECORD +408 -0
  122. blaxel/common/autoload.py +0 -9
  123. blaxel/instrumentation/map.py +0 -49
  124. blaxel/mcp/__init__.py +0 -3
  125. blaxel/models/__init__.py +0 -104
  126. blaxel/sandbox/base.py +0 -67
  127. blaxel/sandbox/filesystem.py +0 -104
  128. blaxel/sandbox/process.py +0 -56
  129. blaxel/tools/crewai.py +0 -22
  130. blaxel/tools/googleadk.py +0 -66
  131. blaxel-0.1.22rc70.dist-info/METADATA +0 -169
  132. blaxel-0.1.22rc70.dist-info/RECORD +0 -379
  133. /blaxel/{authentication → core/authentication}/__init__.py +0 -0
  134. /blaxel/{authentication → core/authentication}/devicemode.py +0 -0
  135. /blaxel/{authentication → core/authentication}/oauth.py +0 -0
  136. /blaxel/{authentication → core/authentication}/types.py +0 -0
  137. /blaxel/{cache → core/cache}/__init__.py +0 -0
  138. /blaxel/{cache → core/cache}/cache.py +0 -0
  139. /blaxel/{client → core/client}/__init__.py +0 -0
  140. /blaxel/{client → core/client}/api/__init__.py +0 -0
  141. /blaxel/{client → core/client}/api/agents/__init__.py +0 -0
  142. /blaxel/{client → core/client}/api/agents/delete_agent.py +0 -0
  143. /blaxel/{client → core/client}/api/agents/get_agent.py +0 -0
  144. /blaxel/{client → core/client}/api/agents/list_agent_revisions.py +0 -0
  145. /blaxel/{client → core/client}/api/agents/list_agents.py +0 -0
  146. /blaxel/{client → core/client}/api/compute/__init__.py +0 -0
  147. /blaxel/{client → core/client}/api/compute/delete_sandbox.py +0 -0
  148. /blaxel/{client → core/client}/api/compute/delete_sandbox_preview.py +0 -0
  149. /blaxel/{client → core/client}/api/compute/delete_sandbox_preview_token.py +0 -0
  150. /blaxel/{client → core/client}/api/compute/get_sandbox.py +0 -0
  151. /blaxel/{client → core/client}/api/compute/get_sandbox_preview.py +0 -0
  152. /blaxel/{client → core/client}/api/compute/list_sandbox_preview_tokens.py +0 -0
  153. /blaxel/{client → core/client}/api/compute/list_sandbox_previews.py +0 -0
  154. /blaxel/{client → core/client}/api/compute/list_sandboxes.py +0 -0
  155. /blaxel/{client → core/client}/api/compute/start_sandbox.py +0 -0
  156. /blaxel/{client → core/client}/api/compute/stop_sandbox.py +0 -0
  157. /blaxel/{client → core/client}/api/configurations/__init__.py +0 -0
  158. /blaxel/{client → core/client}/api/configurations/get_configuration.py +0 -0
  159. /blaxel/{client → core/client}/api/default/__init__.py +0 -0
  160. /blaxel/{client → core/client}/api/default/get_template.py +0 -0
  161. /blaxel/{client → core/client}/api/default/list_mcp_hub_definitions.py +0 -0
  162. /blaxel/{client → core/client}/api/default/list_sandbox_hub_definitions.py +0 -0
  163. /blaxel/{client → core/client}/api/functions/__init__.py +0 -0
  164. /blaxel/{client → core/client}/api/functions/delete_function.py +0 -0
  165. /blaxel/{client → core/client}/api/functions/get_function.py +0 -0
  166. /blaxel/{client → core/client}/api/functions/list_function_revisions.py +0 -0
  167. /blaxel/{client → core/client}/api/functions/list_functions.py +0 -0
  168. /blaxel/{client → core/client}/api/integrations/__init__.py +0 -0
  169. /blaxel/{client → core/client}/api/integrations/delete_integration_connection.py +0 -0
  170. /blaxel/{client → core/client}/api/integrations/get_integration.py +0 -0
  171. /blaxel/{client → core/client}/api/integrations/get_integration_connection.py +0 -0
  172. /blaxel/{client → core/client}/api/integrations/get_integration_connection_model.py +0 -0
  173. /blaxel/{client → core/client}/api/integrations/get_integration_connection_model_endpoint_configurations.py +0 -0
  174. /blaxel/{client → core/client}/api/integrations/list_integration_connection_models.py +0 -0
  175. /blaxel/{client → core/client}/api/integrations/list_integration_connections.py +0 -0
  176. /blaxel/{client → core/client}/api/invitations/__init__.py +0 -0
  177. /blaxel/{client → core/client}/api/invitations/list_all_pending_invitations.py +0 -0
  178. /blaxel/{client → core/client}/api/jobs/__init__.py +0 -0
  179. /blaxel/{client → core/client}/api/jobs/delete_job.py +0 -0
  180. /blaxel/{client → core/client}/api/jobs/get_job.py +0 -0
  181. /blaxel/{client → core/client}/api/jobs/list_job_revisions.py +0 -0
  182. /blaxel/{client → core/client}/api/jobs/list_jobs.py +0 -0
  183. /blaxel/{client → core/client}/api/knowledgebases/__init__.py +0 -0
  184. /blaxel/{client → core/client}/api/knowledgebases/delete_knowledgebase.py +0 -0
  185. /blaxel/{client → core/client}/api/knowledgebases/get_knowledgebase.py +0 -0
  186. /blaxel/{client → core/client}/api/knowledgebases/list_knowledgebase_revisions.py +0 -0
  187. /blaxel/{client → core/client}/api/knowledgebases/list_knowledgebases.py +0 -0
  188. /blaxel/{client → core/client}/api/locations/__init__.py +0 -0
  189. /blaxel/{client → core/client}/api/locations/list_locations.py +0 -0
  190. /blaxel/{client → core/client}/api/models/__init__.py +0 -0
  191. /blaxel/{client → core/client}/api/models/delete_model.py +0 -0
  192. /blaxel/{client → core/client}/api/models/get_model.py +0 -0
  193. /blaxel/{client → core/client}/api/models/list_model_revisions.py +0 -0
  194. /blaxel/{client → core/client}/api/models/list_models.py +0 -0
  195. /blaxel/{client → core/client}/api/policies/__init__.py +0 -0
  196. /blaxel/{client → core/client}/api/policies/delete_policy.py +0 -0
  197. /blaxel/{client → core/client}/api/policies/get_policy.py +0 -0
  198. /blaxel/{client → core/client}/api/policies/list_policies.py +0 -0
  199. /blaxel/{client → core/client}/api/privateclusters/__init__.py +0 -0
  200. /blaxel/{client → core/client}/api/privateclusters/create_private_cluster.py +0 -0
  201. /blaxel/{client → core/client}/api/privateclusters/delete_private_cluster.py +0 -0
  202. /blaxel/{client → core/client}/api/privateclusters/get_private_cluster.py +0 -0
  203. /blaxel/{client → core/client}/api/privateclusters/get_private_cluster_health.py +0 -0
  204. /blaxel/{client → core/client}/api/privateclusters/list_private_clusters.py +0 -0
  205. /blaxel/{client → core/client}/api/privateclusters/update_private_cluster.py +0 -0
  206. /blaxel/{client → core/client}/api/privateclusters/update_private_cluster_health.py +0 -0
  207. /blaxel/{client → core/client}/api/service_accounts/__init__.py +0 -0
  208. /blaxel/{client → core/client}/api/service_accounts/delete_api_key_for_service_account.py +0 -0
  209. /blaxel/{client → core/client}/api/service_accounts/delete_workspace_service_account.py +0 -0
  210. /blaxel/{client → core/client}/api/service_accounts/get_workspace_service_accounts.py +0 -0
  211. /blaxel/{client → core/client}/api/service_accounts/list_api_keys_for_service_account.py +0 -0
  212. /blaxel/{client → core/client}/api/templates/__init__.py +0 -0
  213. /blaxel/{client → core/client}/api/templates/list_templates.py +0 -0
  214. /blaxel/{client → core/client}/api/workspaces/__init__.py +0 -0
  215. /blaxel/{client → core/client}/api/workspaces/accept_workspace_invitation.py +0 -0
  216. /blaxel/{client → core/client}/api/workspaces/decline_workspace_invitation.py +0 -0
  217. /blaxel/{client → core/client}/api/workspaces/delete_workspace.py +0 -0
  218. /blaxel/{client → core/client}/api/workspaces/get_workspace.py +0 -0
  219. /blaxel/{client → core/client}/api/workspaces/leave_workspace.py +0 -0
  220. /blaxel/{client → core/client}/api/workspaces/list_workspace_users.py +0 -0
  221. /blaxel/{client → core/client}/api/workspaces/list_workspaces.py +0 -0
  222. /blaxel/{client → core/client}/api/workspaces/remove_workspace_user.py +0 -0
  223. /blaxel/{client → core/client}/client.py +0 -0
  224. /blaxel/{client → core/client}/errors.py +0 -0
  225. /blaxel/{client → core/client}/models/__init__.py +0 -0
  226. /blaxel/{client → core/client}/models/acl.py +0 -0
  227. /blaxel/{client → core/client}/models/api_key.py +0 -0
  228. /blaxel/{client → core/client}/models/billable_time_metric.py +0 -0
  229. /blaxel/{client → core/client}/models/check_workspace_availability_body.py +0 -0
  230. /blaxel/{client → core/client}/models/configuration.py +0 -0
  231. /blaxel/{client → core/client}/models/continent.py +0 -0
  232. /blaxel/{client → core/client}/models/core_event.py +0 -0
  233. /blaxel/{client → core/client}/models/core_spec_configurations.py +0 -0
  234. /blaxel/{client → core/client}/models/country.py +0 -0
  235. /blaxel/{client → core/client}/models/create_api_key_for_service_account_body.py +0 -0
  236. /blaxel/{client → core/client}/models/create_workspace_service_account_body.py +0 -0
  237. /blaxel/{client → core/client}/models/create_workspace_service_account_response_200.py +0 -0
  238. /blaxel/{client → core/client}/models/delete_sandbox_preview_token_response_200.py +0 -0
  239. /blaxel/{client → core/client}/models/delete_workspace_service_account_response_200.py +0 -0
  240. /blaxel/{client → core/client}/models/entrypoint.py +0 -0
  241. /blaxel/{client → core/client}/models/entrypoint_env.py +0 -0
  242. /blaxel/{client → core/client}/models/flavor.py +0 -0
  243. /blaxel/{client → core/client}/models/form.py +0 -0
  244. /blaxel/{client → core/client}/models/form_config.py +0 -0
  245. /blaxel/{client → core/client}/models/form_oauth.py +0 -0
  246. /blaxel/{client → core/client}/models/form_secrets.py +0 -0
  247. /blaxel/{client → core/client}/models/function_kit.py +0 -0
  248. /blaxel/{client → core/client}/models/function_schema.py +0 -0
  249. /blaxel/{client → core/client}/models/function_schema_not.py +0 -0
  250. /blaxel/{client → core/client}/models/function_schema_or_bool.py +0 -0
  251. /blaxel/{client → core/client}/models/get_workspace_service_accounts_response_200_item.py +0 -0
  252. /blaxel/{client → core/client}/models/histogram_bucket.py +0 -0
  253. /blaxel/{client → core/client}/models/histogram_stats.py +0 -0
  254. /blaxel/{client → core/client}/models/integration_additional_infos.py +0 -0
  255. /blaxel/{client → core/client}/models/integration_connection.py +0 -0
  256. /blaxel/{client → core/client}/models/integration_connection_spec.py +0 -0
  257. /blaxel/{client → core/client}/models/integration_connection_spec_config.py +0 -0
  258. /blaxel/{client → core/client}/models/integration_connection_spec_secret.py +0 -0
  259. /blaxel/{client → core/client}/models/integration_endpoint.py +0 -0
  260. /blaxel/{client → core/client}/models/integration_endpoint_token.py +0 -0
  261. /blaxel/{client → core/client}/models/integration_headers.py +0 -0
  262. /blaxel/{client → core/client}/models/integration_model.py +0 -0
  263. /blaxel/{client → core/client}/models/integration_organization.py +0 -0
  264. /blaxel/{client → core/client}/models/integration_query_params.py +0 -0
  265. /blaxel/{client → core/client}/models/integration_repository.py +0 -0
  266. /blaxel/{client → core/client}/models/invite_workspace_user_body.py +0 -0
  267. /blaxel/{client → core/client}/models/job_execution_config.py +0 -0
  268. /blaxel/{client → core/client}/models/job_metrics.py +0 -0
  269. /blaxel/{client → core/client}/models/job_metrics_executions_chart.py +0 -0
  270. /blaxel/{client → core/client}/models/job_metrics_executions_total.py +0 -0
  271. /blaxel/{client → core/client}/models/job_metrics_tasks_chart.py +0 -0
  272. /blaxel/{client → core/client}/models/job_metrics_tasks_total.py +0 -0
  273. /blaxel/{client → core/client}/models/jobs_chart.py +0 -0
  274. /blaxel/{client → core/client}/models/jobs_chart_value.py +0 -0
  275. /blaxel/{client → core/client}/models/jobs_executions.py +0 -0
  276. /blaxel/{client → core/client}/models/jobs_network_chart.py +0 -0
  277. /blaxel/{client → core/client}/models/jobs_success_failed_chart.py +0 -0
  278. /blaxel/{client → core/client}/models/jobs_tasks.py +0 -0
  279. /blaxel/{client → core/client}/models/jobs_total.py +0 -0
  280. /blaxel/{client → core/client}/models/knowledgebase_spec.py +0 -0
  281. /blaxel/{client → core/client}/models/knowledgebase_spec_options.py +0 -0
  282. /blaxel/{client → core/client}/models/last_n_requests_metric.py +0 -0
  283. /blaxel/{client → core/client}/models/latency_metric.py +0 -0
  284. /blaxel/{client → core/client}/models/logs_response.py +0 -0
  285. /blaxel/{client → core/client}/models/logs_response_data.py +0 -0
  286. /blaxel/{client → core/client}/models/mcp_definition.py +0 -0
  287. /blaxel/{client → core/client}/models/mcp_definition_entrypoint.py +0 -0
  288. /blaxel/{client → core/client}/models/mcp_definition_form.py +0 -0
  289. /blaxel/{client → core/client}/models/memory_allocation_by_name.py +0 -0
  290. /blaxel/{client → core/client}/models/memory_allocation_metric.py +0 -0
  291. /blaxel/{client → core/client}/models/metadata.py +0 -0
  292. /blaxel/{client → core/client}/models/metadata_labels.py +0 -0
  293. /blaxel/{client → core/client}/models/metric.py +0 -0
  294. /blaxel/{client → core/client}/models/metrics.py +0 -0
  295. /blaxel/{client → core/client}/models/metrics_models.py +0 -0
  296. /blaxel/{client → core/client}/models/metrics_request_total_per_code.py +0 -0
  297. /blaxel/{client → core/client}/models/metrics_rps_per_code.py +0 -0
  298. /blaxel/{client → core/client}/models/model_private_cluster.py +0 -0
  299. /blaxel/{client → core/client}/models/o_auth.py +0 -0
  300. /blaxel/{client → core/client}/models/owner_fields.py +0 -0
  301. /blaxel/{client → core/client}/models/pending_invitation.py +0 -0
  302. /blaxel/{client → core/client}/models/pending_invitation_accept.py +0 -0
  303. /blaxel/{client → core/client}/models/pending_invitation_render.py +0 -0
  304. /blaxel/{client → core/client}/models/pending_invitation_render_invited_by.py +0 -0
  305. /blaxel/{client → core/client}/models/pending_invitation_render_workspace.py +0 -0
  306. /blaxel/{client → core/client}/models/pending_invitation_workspace_details.py +0 -0
  307. /blaxel/{client → core/client}/models/pod_template_spec.py +0 -0
  308. /blaxel/{client → core/client}/models/policy.py +0 -0
  309. /blaxel/{client → core/client}/models/policy_location.py +0 -0
  310. /blaxel/{client → core/client}/models/policy_max_tokens.py +0 -0
  311. /blaxel/{client → core/client}/models/port.py +0 -0
  312. /blaxel/{client → core/client}/models/preview.py +0 -0
  313. /blaxel/{client → core/client}/models/preview_metadata.py +0 -0
  314. /blaxel/{client → core/client}/models/preview_spec.py +0 -0
  315. /blaxel/{client → core/client}/models/preview_spec_request_headers.py +0 -0
  316. /blaxel/{client → core/client}/models/preview_spec_response_headers.py +0 -0
  317. /blaxel/{client → core/client}/models/preview_token.py +0 -0
  318. /blaxel/{client → core/client}/models/preview_token_metadata.py +0 -0
  319. /blaxel/{client → core/client}/models/preview_token_spec.py +0 -0
  320. /blaxel/{client → core/client}/models/private_cluster.py +0 -0
  321. /blaxel/{client → core/client}/models/private_location.py +0 -0
  322. /blaxel/{client → core/client}/models/repository.py +0 -0
  323. /blaxel/{client → core/client}/models/request_duration_over_time_metric.py +0 -0
  324. /blaxel/{client → core/client}/models/request_duration_over_time_metrics.py +0 -0
  325. /blaxel/{client → core/client}/models/request_total_by_origin_metric.py +0 -0
  326. /blaxel/{client → core/client}/models/request_total_by_origin_metric_request_total_by_origin.py +0 -0
  327. /blaxel/{client → core/client}/models/request_total_by_origin_metric_request_total_by_origin_and_code.py +0 -0
  328. /blaxel/{client → core/client}/models/request_total_metric.py +0 -0
  329. /blaxel/{client → core/client}/models/request_total_metric_request_total_per_code.py +0 -0
  330. /blaxel/{client → core/client}/models/request_total_metric_rps_per_code.py +0 -0
  331. /blaxel/{client → core/client}/models/request_total_response_data.py +0 -0
  332. /blaxel/{client → core/client}/models/resource.py +0 -0
  333. /blaxel/{client → core/client}/models/resource_log.py +0 -0
  334. /blaxel/{client → core/client}/models/resource_log_chart.py +0 -0
  335. /blaxel/{client → core/client}/models/resource_log_response.py +0 -0
  336. /blaxel/{client → core/client}/models/resource_metrics_request_total_per_code.py +0 -0
  337. /blaxel/{client → core/client}/models/resource_metrics_request_total_per_code_previous.py +0 -0
  338. /blaxel/{client → core/client}/models/resource_metrics_rps_per_code.py +0 -0
  339. /blaxel/{client → core/client}/models/resource_metrics_rps_per_code_previous.py +0 -0
  340. /blaxel/{client → core/client}/models/resource_trace.py +0 -0
  341. /blaxel/{client → core/client}/models/revision_configuration.py +0 -0
  342. /blaxel/{client → core/client}/models/revision_metadata.py +0 -0
  343. /blaxel/{client → core/client}/models/runtime_configuration.py +0 -0
  344. /blaxel/{client → core/client}/models/runtime_startup_probe.py +0 -0
  345. /blaxel/{client → core/client}/models/serverless_config.py +0 -0
  346. /blaxel/{client → core/client}/models/serverless_config_configuration.py +0 -0
  347. /blaxel/{client → core/client}/models/spec_configuration.py +0 -0
  348. /blaxel/{client → core/client}/models/start_sandbox.py +0 -0
  349. /blaxel/{client → core/client}/models/stop_sandbox.py +0 -0
  350. /blaxel/{client → core/client}/models/store_agent_labels.py +0 -0
  351. /blaxel/{client → core/client}/models/store_configuration_option.py +0 -0
  352. /blaxel/{client → core/client}/models/template_variable.py +0 -0
  353. /blaxel/{client → core/client}/models/time_fields.py +0 -0
  354. /blaxel/{client → core/client}/models/time_to_first_token_over_time_metrics.py +0 -0
  355. /blaxel/{client → core/client}/models/token_rate_metric.py +0 -0
  356. /blaxel/{client → core/client}/models/token_rate_metrics.py +0 -0
  357. /blaxel/{client → core/client}/models/token_total_metric.py +0 -0
  358. /blaxel/{client → core/client}/models/trace_ids_response.py +0 -0
  359. /blaxel/{client → core/client}/models/trigger.py +0 -0
  360. /blaxel/{client → core/client}/models/trigger_configuration.py +0 -0
  361. /blaxel/{client → core/client}/models/update_workspace_service_account_body.py +0 -0
  362. /blaxel/{client → core/client}/models/update_workspace_service_account_response_200.py +0 -0
  363. /blaxel/{client → core/client}/models/update_workspace_user_role_body.py +0 -0
  364. /blaxel/{client → core/client}/models/websocket_channel.py +0 -0
  365. /blaxel/{client → core/client}/models/websocket_message.py +0 -0
  366. /blaxel/{client → core/client}/models/workspace.py +0 -0
  367. /blaxel/{client → core/client}/models/workspace_labels.py +0 -0
  368. /blaxel/{client → core/client}/models/workspace_runtime.py +0 -0
  369. /blaxel/{client → core/client}/models/workspace_user.py +0 -0
  370. /blaxel/{client → core/client}/py.typed +0 -0
  371. /blaxel/{client → core/client}/types.py +0 -0
  372. /blaxel/{common → core/common}/env.py +0 -0
  373. /blaxel/{common → core/common}/settings.py +0 -0
  374. /blaxel/{sandbox/client/api/filesystem/__init__.py → core/py.typed} +0 -0
  375. /blaxel/{sandbox → core/sandbox}/client/__init__.py +0 -0
  376. /blaxel/{sandbox → core/sandbox}/client/api/__init__.py +0 -0
  377. /blaxel/{sandbox/client/api/network → core/sandbox/client/api/filesystem}/__init__.py +0 -0
  378. /blaxel/{sandbox → core/sandbox}/client/api/filesystem/delete_filesystem_path.py +0 -0
  379. /blaxel/{sandbox → core/sandbox}/client/api/filesystem/get_filesystem_path.py +0 -0
  380. /blaxel/{sandbox → core/sandbox}/client/api/filesystem/get_watch_filesystem_path.py +0 -0
  381. /blaxel/{sandbox → core/sandbox}/client/api/filesystem/get_ws_watch_filesystem_path.py +0 -0
  382. /blaxel/{sandbox/client/api/process → core/sandbox/client/api/network}/__init__.py +0 -0
  383. /blaxel/{sandbox → core/sandbox}/client/api/network/delete_network_process_pid_monitor.py +0 -0
  384. /blaxel/{sandbox → core/sandbox}/client/api/network/get_network_process_pid_ports.py +0 -0
  385. /blaxel/{sandbox → core/sandbox}/client/api/process/delete_process_identifier.py +0 -0
  386. /blaxel/{sandbox → core/sandbox}/client/api/process/delete_process_identifier_kill.py +0 -0
  387. /blaxel/{sandbox → core/sandbox}/client/api/process/get_process.py +0 -0
  388. /blaxel/{sandbox → core/sandbox}/client/api/process/get_process_identifier.py +0 -0
  389. /blaxel/{sandbox → core/sandbox}/client/api/process/get_process_identifier_logs.py +0 -0
  390. /blaxel/{sandbox → core/sandbox}/client/api/process/get_process_identifier_logs_stream.py +0 -0
  391. /blaxel/{sandbox → core/sandbox}/client/api/process/get_ws_process_identifier_logs_stream.py +0 -0
  392. /blaxel/{sandbox → core/sandbox}/client/client.py +0 -0
  393. /blaxel/{sandbox → core/sandbox}/client/errors.py +0 -0
  394. /blaxel/{sandbox → core/sandbox}/client/models/__init__.py +0 -0
  395. /blaxel/{sandbox → core/sandbox}/client/models/delete_network_process_pid_monitor_response_200.py +0 -0
  396. /blaxel/{sandbox → core/sandbox}/client/models/error_response.py +0 -0
  397. /blaxel/{sandbox → core/sandbox}/client/models/file.py +0 -0
  398. /blaxel/{sandbox → core/sandbox}/client/models/file_request.py +0 -0
  399. /blaxel/{sandbox → core/sandbox}/client/models/file_with_content.py +0 -0
  400. /blaxel/{sandbox → core/sandbox}/client/models/get_network_process_pid_ports_response_200.py +0 -0
  401. /blaxel/{sandbox → core/sandbox}/client/models/port_monitor_request.py +0 -0
  402. /blaxel/{sandbox → core/sandbox}/client/models/post_network_process_pid_monitor_response_200.py +0 -0
  403. /blaxel/{sandbox → core/sandbox}/client/models/process_logs.py +0 -0
  404. /blaxel/{sandbox → core/sandbox}/client/models/process_request.py +0 -0
  405. /blaxel/{sandbox → core/sandbox}/client/models/process_request_env.py +0 -0
  406. /blaxel/{sandbox → core/sandbox}/client/models/process_response.py +0 -0
  407. /blaxel/{sandbox → core/sandbox}/client/models/process_response_status.py +0 -0
  408. /blaxel/{sandbox → core/sandbox}/client/models/subdirectory.py +0 -0
  409. /blaxel/{sandbox → core/sandbox}/client/models/success_response.py +0 -0
  410. /blaxel/{sandbox → core/sandbox}/client/py.typed +0 -0
  411. /blaxel/{sandbox → core/sandbox}/client/types.py +0 -0
  412. /blaxel/{tools → core/tools}/common.py +0 -0
  413. /blaxel/{tools → core/tools}/types.py +0 -0
  414. /blaxel/{models/custom/langchain → langgraph/custom}/gemini.py +0 -0
  415. /blaxel/{models/custom/llamaindex → llamaindex/custom}/cohere.py +0 -0
  416. /blaxel/{instrumentation → telemetry}/exporters.py +0 -0
  417. /blaxel/{instrumentation → telemetry/log}/log.py +0 -0
  418. {blaxel-0.1.22rc70.dist-info → blaxel-0.2.0rc1.dist-info}/WHEEL +0 -0
  419. {blaxel-0.1.22rc70.dist-info → blaxel-0.2.0rc1.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,360 @@
1
+ """OpenTelemetry Google Generative AI API instrumentation"""
2
+
3
+ import logging
4
+ import os
5
+ import types
6
+ from typing import Collection
7
+
8
+ from opentelemetry import context as context_api
9
+ from opentelemetry.instrumentation.google_generativeai.config import Config
10
+ from opentelemetry.instrumentation.google_generativeai.utils import dont_throw
11
+ from opentelemetry.instrumentation.google_generativeai.version import __version__
12
+ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
13
+ from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
14
+ from opentelemetry.semconv_ai import (
15
+ SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
16
+ LLMRequestTypeValues,
17
+ SpanAttributes,
18
+ )
19
+ from opentelemetry.trace import SpanKind, get_tracer
20
+ from opentelemetry.trace.status import Status, StatusCode
21
+ from wrapt import wrap_function_wrapper
22
+
23
+ logger = logging.getLogger(__name__)
24
+
25
+ WRAPPED_METHODS = [
26
+ {
27
+ "package": "blaxel_langgraph.custom.gemini",
28
+ "object": "GeminiRestClient",
29
+ "method": "generate_content",
30
+ "span_name": "gemini.generate_content",
31
+ },
32
+ {
33
+ "package": "blaxel_langgraph.custom.gemini",
34
+ "object": "GeminiRestClient",
35
+ "method": "generate_content_async",
36
+ "span_name": "gemini.generate_content_async",
37
+ },
38
+ ]
39
+
40
+
41
+ def should_send_prompts():
42
+ return (
43
+ os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
44
+ ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
45
+
46
+
47
+ def is_streaming_response(response):
48
+ return isinstance(response, types.GeneratorType)
49
+
50
+
51
+ def is_async_streaming_response(response):
52
+ return isinstance(response, types.AsyncGeneratorType)
53
+
54
+
55
+ def _set_span_attribute(span, name, value):
56
+ if value is not None:
57
+ if value != "":
58
+ span.set_attribute(name, value)
59
+ return
60
+
61
+
62
+ def _set_input_attributes(span, args, kwargs, llm_model):
63
+ if not should_send_prompts():
64
+ return
65
+
66
+ if "contents" in kwargs:
67
+ contents = kwargs["contents"]
68
+ if isinstance(contents, list):
69
+ for i, content in enumerate(contents):
70
+ if hasattr(content, "parts"):
71
+ for part in content.parts:
72
+ if hasattr(part, "text"):
73
+ _set_span_attribute(
74
+ span,
75
+ f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
76
+ part.text,
77
+ )
78
+ _set_span_attribute(
79
+ span,
80
+ f"{SpanAttributes.LLM_PROMPTS}.{i}.role",
81
+ getattr(content, "role", "user"),
82
+ )
83
+ elif args and len(args) > 0:
84
+ prompt = ""
85
+ for arg in args:
86
+ if isinstance(arg, str):
87
+ prompt = f"{prompt}{arg}\n"
88
+ elif isinstance(arg, list):
89
+ for subarg in arg:
90
+ prompt = f"{prompt}{subarg}\n"
91
+ if prompt:
92
+ _set_span_attribute(
93
+ span,
94
+ f"{SpanAttributes.LLM_PROMPTS}.0.content",
95
+ prompt,
96
+ )
97
+ _set_span_attribute(
98
+ span,
99
+ f"{SpanAttributes.LLM_PROMPTS}.0.role",
100
+ "user",
101
+ )
102
+ elif "prompt" in kwargs:
103
+ _set_span_attribute(
104
+ span,
105
+ f"{SpanAttributes.LLM_PROMPTS}.0.content",
106
+ kwargs["prompt"],
107
+ )
108
+ _set_span_attribute(
109
+ span,
110
+ f"{SpanAttributes.LLM_PROMPTS}.0.role",
111
+ "user",
112
+ )
113
+
114
+ _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, llm_model)
115
+ _set_span_attribute(
116
+ span, SpanAttributes.LLM_REQUEST_TEMPERATURE, kwargs.get("temperature")
117
+ )
118
+ _set_span_attribute(
119
+ span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_output_tokens")
120
+ )
121
+ _set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, kwargs.get("top_p"))
122
+ _set_span_attribute(span, SpanAttributes.LLM_TOP_K, kwargs.get("top_k"))
123
+ _set_span_attribute(
124
+ span, SpanAttributes.LLM_PRESENCE_PENALTY, kwargs.get("presence_penalty")
125
+ )
126
+ _set_span_attribute(
127
+ span, SpanAttributes.LLM_FREQUENCY_PENALTY, kwargs.get("frequency_penalty")
128
+ )
129
+
130
+ return
131
+
132
+
133
+ @dont_throw
134
+ def _set_response_attributes(span, response, llm_model):
135
+ _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, llm_model)
136
+
137
+ if hasattr(response, "usage_metadata"):
138
+ _set_span_attribute(
139
+ span,
140
+ SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
141
+ response.usage_metadata.total_token_count,
142
+ )
143
+ _set_span_attribute(
144
+ span,
145
+ SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
146
+ response.usage_metadata.candidates_token_count,
147
+ )
148
+ _set_span_attribute(
149
+ span,
150
+ SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
151
+ response.usage_metadata.prompt_token_count,
152
+ )
153
+
154
+ if isinstance(response.text, list):
155
+ for index, item in enumerate(response):
156
+ prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
157
+ _set_span_attribute(span, f"{prefix}.content", item.text)
158
+ _set_span_attribute(span, f"{prefix}.role", "assistant")
159
+ elif isinstance(response.text, str):
160
+ _set_span_attribute(
161
+ span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content", response.text
162
+ )
163
+ _set_span_attribute(
164
+ span, f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant"
165
+ )
166
+ else:
167
+ if isinstance(response, list):
168
+ for index, item in enumerate(response):
169
+ prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
170
+ _set_span_attribute(span, f"{prefix}.content", item)
171
+ _set_span_attribute(span, f"{prefix}.role", "assistant")
172
+ elif isinstance(response, str):
173
+ _set_span_attribute(
174
+ span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content", response
175
+ )
176
+ _set_span_attribute(
177
+ span, f"{SpanAttributes.LLM_COMPLETIONS}.0.role", "assistant"
178
+ )
179
+
180
+ return
181
+
182
+
183
+ def _build_from_streaming_response(span, response, llm_model):
184
+ complete_response = ""
185
+ for item in response:
186
+ item_to_yield = item
187
+ complete_response += str(item.text)
188
+
189
+ yield item_to_yield
190
+
191
+ _set_response_attributes(span, complete_response, llm_model)
192
+
193
+ span.set_status(Status(StatusCode.OK))
194
+ span.end()
195
+
196
+
197
+ async def _abuild_from_streaming_response(span, response, llm_model):
198
+ complete_response = ""
199
+ async for item in response:
200
+ item_to_yield = item
201
+ complete_response += str(item.text)
202
+
203
+ yield item_to_yield
204
+
205
+ _set_response_attributes(span, complete_response, llm_model)
206
+
207
+ span.set_status(Status(StatusCode.OK))
208
+ span.end()
209
+
210
+
211
+ # @dont_throw
212
+ def _handle_request(span, args, kwargs, llm_model):
213
+ if span.is_recording():
214
+ _set_input_attributes(span, args, kwargs, llm_model)
215
+
216
+
217
+ @dont_throw
218
+ def _handle_response(span, response, llm_model):
219
+ if span.is_recording():
220
+ _set_response_attributes(span, response, llm_model)
221
+
222
+ span.set_status(Status(StatusCode.OK))
223
+
224
+
225
+ def _with_tracer_wrapper(func):
226
+ """Helper for providing tracer for wrapper functions."""
227
+
228
+ def _with_tracer(tracer, to_wrap):
229
+ def wrapper(wrapped, instance, args, kwargs):
230
+ return func(tracer, to_wrap, wrapped, instance, args, kwargs)
231
+
232
+ return wrapper
233
+
234
+ return _with_tracer
235
+
236
+
237
+ @_with_tracer_wrapper
238
+ async def _awrap(tracer, to_wrap, wrapped, instance, args, kwargs):
239
+ """Instruments and calls every function defined in TO_WRAP."""
240
+ if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
241
+ SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
242
+ ):
243
+ return await wrapped(*args, **kwargs)
244
+
245
+ llm_model = "unknown"
246
+ if hasattr(instance, "_model_id"):
247
+ llm_model = instance._model_id.replace("models/", "")
248
+ if hasattr(instance, "_model_name"):
249
+ llm_model = instance._model_name.replace(
250
+ "publishers/google/models/", ""
251
+ ).replace("models/", "")
252
+ if hasattr(instance, "model") and hasattr(instance.model, "model_name"):
253
+ llm_model = instance.model.model_name.replace("models/", "")
254
+
255
+ name = to_wrap.get("span_name")
256
+ span = tracer.start_span(
257
+ name,
258
+ kind=SpanKind.CLIENT,
259
+ attributes={
260
+ SpanAttributes.LLM_SYSTEM: "Gemini",
261
+ SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
262
+ },
263
+ )
264
+
265
+ _handle_request(span, args, kwargs, llm_model)
266
+
267
+ response = await wrapped(*args, **kwargs)
268
+
269
+ if response:
270
+ if is_streaming_response(response):
271
+ return _build_from_streaming_response(span, response, llm_model)
272
+ elif is_async_streaming_response(response):
273
+ return _abuild_from_streaming_response(span, response, llm_model)
274
+ else:
275
+ _handle_response(span, response, llm_model)
276
+
277
+ span.end()
278
+ return response
279
+
280
+
281
+ @_with_tracer_wrapper
282
+ def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
283
+ """Instruments and calls every function defined in TO_WRAP."""
284
+ if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
285
+ SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
286
+ ):
287
+ return wrapped(*args, **kwargs)
288
+
289
+ llm_model = "unknown"
290
+ if hasattr(instance, "_model_id"):
291
+ llm_model = instance._model_id.replace("models/", "")
292
+ if hasattr(instance, "_model_name"):
293
+ llm_model = instance._model_name.replace(
294
+ "publishers/google/models/", ""
295
+ ).replace("models/", "")
296
+ if hasattr(instance, "model") and hasattr(instance.model, "model_name"):
297
+ llm_model = instance.model.model_name.replace("models/", "")
298
+
299
+ name = to_wrap.get("span_name")
300
+ span = tracer.start_span(
301
+ name,
302
+ kind=SpanKind.CLIENT,
303
+ attributes={
304
+ SpanAttributes.LLM_SYSTEM: "Gemini",
305
+ SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
306
+ },
307
+ )
308
+
309
+ _handle_request(span, args, kwargs, llm_model)
310
+
311
+ response = wrapped(*args, **kwargs)
312
+
313
+ if response:
314
+ if is_streaming_response(response):
315
+ return _build_from_streaming_response(span, response, llm_model)
316
+ elif is_async_streaming_response(response):
317
+ return _abuild_from_streaming_response(span, response, llm_model)
318
+ else:
319
+ _handle_response(span, response, llm_model)
320
+
321
+ span.end()
322
+ return response
323
+
324
+
325
+ class BlaxelLanggraphGeminiInstrumentor(BaseInstrumentor):
326
+ """An instrumentor for Google Generative AI's client library."""
327
+
328
+ def __init__(self, exception_logger=None):
329
+ super().__init__()
330
+ Config.exception_logger = exception_logger
331
+
332
+ def instrumentation_dependencies(self) -> Collection[str]:
333
+ return ["blaxel_langgraph"]
334
+
335
+ def _instrument(self, **kwargs):
336
+ tracer_provider = kwargs.get("tracer_provider")
337
+ tracer = get_tracer(__name__, __version__, tracer_provider)
338
+ for wrapped_method in WRAPPED_METHODS:
339
+ wrap_package = wrapped_method.get("package")
340
+ wrap_object = wrapped_method.get("object")
341
+ wrap_method = wrapped_method.get("method")
342
+
343
+ wrap_function_wrapper(
344
+ wrap_package,
345
+ f"{wrap_object}.{wrap_method}",
346
+ (
347
+ _awrap(tracer, wrapped_method)
348
+ if wrap_method == "generate_content_async"
349
+ else _wrap(tracer, wrapped_method)
350
+ ),
351
+ )
352
+
353
+ def _uninstrument(self, **kwargs):
354
+ for wrapped_method in WRAPPED_METHODS:
355
+ wrap_package = wrapped_method.get("package")
356
+ wrap_object = wrapped_method.get("object")
357
+ unwrap(
358
+ f"{wrap_package}.{wrap_object}",
359
+ wrapped_method.get("method", ""),
360
+ )
@@ -0,0 +1,89 @@
1
+ from importlib.metadata import PackageNotFoundError
2
+ from importlib.metadata import version as package_version
3
+
4
+ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
5
+ from opentelemetry.instrumentation.llamaindex.utils import (
6
+ _with_tracer_wrapper,
7
+ process_request,
8
+ process_response,
9
+ start_as_current_span_async,
10
+ )
11
+ from opentelemetry.semconv_ai import SpanAttributes, TraceloopSpanKindValues
12
+ from opentelemetry.trace import get_tracer
13
+ from wrapt import wrap_function_wrapper
14
+
15
+ TO_INSTRUMENT = [
16
+ {
17
+ "class": "FunctionTool",
18
+ "v9_module": "llama_index.tools.function_tool",
19
+ "v10_module": "llama_index.core.tools.function_tool",
20
+ "v10_legacy_module": "llama_index.legacy.tools.function_tool",
21
+ },
22
+ {
23
+ "class": "QueryEngineTool",
24
+ "v9_module": "llama_index.tools.query_engine",
25
+ "v10_module": "llama_index.core.tools.query_engine",
26
+ "v10_legacy_module": "llama_index.legacy.tools.query_engine",
27
+ },
28
+ ]
29
+
30
+
31
+ class BlaxelLlamaIndexInstrumentor(BaseInstrumentor):
32
+ def instrumentation_dependencies(self):
33
+ return ["blaxel_llamaindex", "blaxel.core"]
34
+
35
+ def _instrument(self, **kwargs):
36
+ tracer_provider = kwargs.get("tracer_provider")
37
+ tracer = get_tracer(__name__, "0.2.0", tracer_provider)
38
+
39
+ for module in TO_INSTRUMENT:
40
+ try:
41
+ package_version("llama-index-core")
42
+ self._instrument_module(module["v10_module"], module["class"], tracer)
43
+ self._instrument_module(module["v10_legacy_module"], module["class"], tracer)
44
+
45
+ except PackageNotFoundError:
46
+ self._instrument_module(module["v9_module"], module["class"], tracer)
47
+
48
+ def _uninstrument(self, **kwargs):
49
+ pass
50
+
51
+ def _instrument_module(self, module_name, class_name, tracer):
52
+ wrap_function_wrapper(
53
+ module_name, f"{class_name}.call", query_wrapper(tracer)
54
+ )
55
+ wrap_function_wrapper(
56
+ module_name, f"{class_name}.acall", aquery_wrapper(tracer)
57
+ )
58
+
59
+
60
+ @_with_tracer_wrapper
61
+ def query_wrapper(tracer, wrapped, instance, args, kwargs):
62
+ name = instance.__class__.__name__
63
+ with tracer.start_as_current_span(f"{name}.tool") as span:
64
+ span.set_attribute(
65
+ SpanAttributes.TRACELOOP_SPAN_KIND,
66
+ TraceloopSpanKindValues.TOOL.value,
67
+ )
68
+ span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
69
+
70
+ process_request(span, args, kwargs)
71
+ res = wrapped(*args, **kwargs)
72
+ process_response(span, res)
73
+ return res
74
+
75
+
76
+ @_with_tracer_wrapper
77
+ async def aquery_wrapper(tracer, wrapped, instance, args, kwargs):
78
+ name = instance.__class__.__name__
79
+ async with start_as_current_span_async(tracer=tracer, name=f"{name}.tool") as span:
80
+ span.set_attribute(
81
+ SpanAttributes.TRACELOOP_SPAN_KIND,
82
+ TraceloopSpanKindValues.TOOL.value,
83
+ )
84
+ span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
85
+
86
+ process_request(span, args, kwargs)
87
+ res = await wrapped(*args, **kwargs)
88
+ process_response(span, res)
89
+ return res
@@ -0,0 +1,61 @@
1
+ from typing import Dict, List
2
+
3
+ from pydantic import BaseModel
4
+
5
+
6
+ class InstrumentationMapping(BaseModel):
7
+ module_path: str
8
+ class_name: str
9
+ required_packages: List[str]
10
+ ignore_if_packages: List[str]
11
+
12
+ MAPPINGS: Dict[str, InstrumentationMapping] = {
13
+ "anthropic": InstrumentationMapping(
14
+ module_path="opentelemetry.instrumentation.anthropic",
15
+ class_name="AnthropicInstrumentor",
16
+ required_packages=["anthropic"],
17
+ ignore_if_packages=[]
18
+ ),
19
+ "cohere": InstrumentationMapping(
20
+ module_path="opentelemetry.instrumentation.cohere",
21
+ class_name="CohereInstrumentor",
22
+ required_packages=["cohere"],
23
+ ignore_if_packages=[]
24
+ ),
25
+ "openai": InstrumentationMapping(
26
+ module_path="opentelemetry.instrumentation.openai",
27
+ class_name="OpenAIInstrumentor",
28
+ required_packages=["openai"],
29
+ ignore_if_packages=[]
30
+ ),
31
+ "gemini": InstrumentationMapping(
32
+ module_path="opentelemetry.instrumentation.google_generativeai",
33
+ class_name="GoogleGenerativeAIInstrumentor",
34
+ required_packages=["google-generativeai"],
35
+ ignore_if_packages=[]
36
+ ),
37
+ "blaxel.core": InstrumentationMapping(
38
+ module_path="blaxel_telemetry.instrumentation.blaxel.core",
39
+ class_name="BlaxelCoreInstrumentor",
40
+ required_packages=["blaxel.core"],
41
+ ignore_if_packages=[]
42
+ ),
43
+ "blaxel_langgraph": InstrumentationMapping(
44
+ module_path="blaxel_telemetry.instrumentation.blaxel_langgraph",
45
+ class_name="BlaxelLanggraphInstrumentor",
46
+ required_packages=["blaxel_langgraph"],
47
+ ignore_if_packages=[]
48
+ ),
49
+ "blaxel_langgraph_gemini": InstrumentationMapping(
50
+ module_path="blaxel_telemetry.instrumentation.blaxel_langgraph_gemini",
51
+ class_name="BlaxelLanggraphGeminiInstrumentor",
52
+ required_packages=["blaxel_langgraph"],
53
+ ignore_if_packages=[]
54
+ ),
55
+ "blaxel_llamaindex": InstrumentationMapping(
56
+ module_path="blaxel_telemetry.instrumentation.blaxel_llamaindex",
57
+ class_name="BlaxelLlamaIndexInstrumentor",
58
+ required_packages=["blaxel_llamaindex"],
59
+ ignore_if_packages=[]
60
+ ),
61
+ }
@@ -0,0 +1,74 @@
1
+ import dataclasses
2
+ import json
3
+ import logging
4
+ import os
5
+ import traceback
6
+ from contextlib import asynccontextmanager
7
+
8
+ from opentelemetry.semconv_ai import SpanAttributes
9
+
10
+
11
+ def dont_throw(func):
12
+ # Obtain a logger specific to the function's module
13
+ logger = logging.getLogger(func.__module__)
14
+
15
+ def wrapper(*args, **kwargs):
16
+ try:
17
+ return func(*args, **kwargs)
18
+ except Exception:
19
+ logger.debug(
20
+ "Instrumentation failed to trace in %s, error: %s",
21
+ func.__name__,
22
+ traceback.format_exc(),
23
+ )
24
+
25
+ return wrapper
26
+
27
+
28
+ def with_tracer_wrapper(func):
29
+ def _with_tracer(tracer):
30
+ def wrapper(wrapped, instance, args, kwargs):
31
+ return func(tracer, wrapped, instance, args, kwargs)
32
+
33
+ return wrapper
34
+
35
+ return _with_tracer
36
+
37
+
38
+ @asynccontextmanager
39
+ async def start_as_current_span_async(tracer, *args, **kwargs):
40
+ with tracer.start_as_current_span(*args, **kwargs) as span:
41
+ yield span
42
+
43
+ class JSONEncoder(json.JSONEncoder):
44
+ def default(self, o):
45
+ if dataclasses.is_dataclass(o):
46
+ return dataclasses.asdict(o)
47
+ elif hasattr(o, "json"):
48
+ return o.json()
49
+ elif hasattr(o, "to_json"):
50
+ return o.to_json()
51
+ return super().default(o)
52
+
53
+
54
+ def should_send_prompts():
55
+ return (
56
+ os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
57
+ ).lower() == "true"
58
+
59
+ @dont_throw
60
+ def process_request(span, args, kwargs):
61
+ if should_send_prompts():
62
+ span.set_attribute(
63
+ SpanAttributes.TRACELOOP_ENTITY_INPUT,
64
+ json.dumps({"args": args, "kwargs": kwargs}, cls=JSONEncoder),
65
+ )
66
+
67
+
68
+ @dont_throw
69
+ def process_response(span, res):
70
+ if should_send_prompts():
71
+ span.set_attribute(
72
+ SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
73
+ json.dumps(res, cls=JSONEncoder),
74
+ )
@@ -6,7 +6,7 @@ import json
6
6
  import logging
7
7
  import os
8
8
 
9
- from opentelemetry.trace import get_current_span
9
+ from opentelemetry import trace
10
10
 
11
11
 
12
12
  class JsonFormatter(logging.Formatter):
@@ -35,18 +35,12 @@ class JsonFormatter(logging.Formatter):
35
35
  self.labels_name: {}
36
36
  }
37
37
 
38
- # Get current active span - equivalent to trace.getActiveSpan() in JS
39
- current_span = get_current_span()
40
-
41
- # Check if span exists and has valid context (equivalent to 'if (currentSpan)' in JS)
42
- if current_span and current_span.get_span_context().is_valid:
38
+ # Add trace context if available
39
+ current_span = trace.get_current_span()
40
+ if current_span.is_recording():
43
41
  span_context = current_span.get_span_context()
44
- # Format trace_id and span_id as hex strings (like JS does)
45
- trace_id_hex = format(span_context.trace_id, '032x')
46
- span_id_hex = format(span_context.span_id, '016x')
47
-
48
- log_entry[self.trace_id_name] = f"{self.trace_id_prefix}{trace_id_hex}"
49
- log_entry[self.span_id_name] = f"{self.span_id_prefix}{span_id_hex}"
42
+ log_entry[self.trace_id_name] = f"{self.trace_id_prefix}{span_context.trace_id}"
43
+ log_entry[self.span_id_name] = f"{self.span_id_prefix}{span_context.span_id}"
50
44
 
51
45
  # Add task ID if available
52
46
  task_id = os.environ.get(self.task_index)
@@ -21,16 +21,16 @@ from opentelemetry.sdk.trace import TracerProvider
21
21
  from opentelemetry.sdk.trace.export import BatchSpanProcessor
22
22
  from opentelemetry.trace import NoOpTracerProvider
23
23
 
24
- from blaxel.instrumentation.exporters import (
24
+ from blaxel.core.common import Settings
25
+
26
+ from .exporters import (
25
27
  DynamicHeadersLogExporter,
26
28
  DynamicHeadersMetricExporter,
27
29
  DynamicHeadersSpanExporter,
28
30
  )
29
- from blaxel.instrumentation.span import DefaultAttributesSpanProcessor
30
-
31
- from ..common.settings import Settings
32
- from .log import AsyncLogRecordProcessor
33
- from .map import MAPPINGS
31
+ from .instrumentation.map import MAPPINGS
32
+ from .log.log import AsyncLogRecordProcessor
33
+ from .span import DefaultAttributesSpanProcessor
34
34
 
35
35
  logger = logging.getLogger(__name__)
36
36
 
@@ -116,6 +116,7 @@ class TelemetryManager:
116
116
 
117
117
  def setup_signal_handler(self):
118
118
  """Set up signal handlers for graceful shutdown."""
119
+
119
120
  def handle_signal(signum, frame):
120
121
  logger.debug(f"Received signal {signum}")
121
122
  self.shutdown()
@@ -150,11 +151,15 @@ class TelemetryManager:
150
151
  # Set up the TracerProvider
151
152
  trace_provider = TracerProvider(resource=resource)
152
153
  span_processor = BatchSpanProcessor(self.get_span_exporter())
153
- trace_provider.add_span_processor(DefaultAttributesSpanProcessor({
154
- "workload.id": self.resource_name,
155
- "workload.type": self.resource_type,
156
- "workspace": self.resource_workspace,
157
- }))
154
+ trace_provider.add_span_processor(
155
+ DefaultAttributesSpanProcessor(
156
+ {
157
+ "workload.id": self.resource_name,
158
+ "workload.type": self.resource_type,
159
+ "workspace": self.resource_workspace,
160
+ }
161
+ )
162
+ )
158
163
  trace_provider.add_span_processor(span_processor)
159
164
  trace.set_tracer_provider(trace_provider)
160
165
  self.tracer = trace_provider.get_tracer(__name__)
@@ -217,10 +222,13 @@ class TelemetryManager:
217
222
  if time.time() - start_time < timeout:
218
223
  logger.debug("Instrumentation shutdown complete")
219
224
  else:
220
- logger.warning("Shutdown timed out after 5 seconds, skipping remaining shutdown tasks")
225
+ logger.warning(
226
+ "Shutdown timed out after 5 seconds, skipping remaining shutdown tasks"
227
+ )
221
228
 
222
229
  except Exception as error:
223
230
  logger.error(f"Error during shutdown: {error}")
224
231
 
232
+
225
233
  # Create a singleton instance
226
234
  telemetry_manager = TelemetryManager()
File without changes