shaped 2.0.1__py3-none-any.whl → 2.0.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (341) hide show
  1. shaped/__init__.py +59 -4
  2. shaped/autogen/__init__.py +527 -66
  3. shaped/autogen/api/__init__.py +4 -3
  4. shaped/autogen/api/engine_api.py +1467 -0
  5. shaped/autogen/api/{dataset_api.py → query_api.py} +159 -194
  6. shaped/autogen/api/table_api.py +1494 -0
  7. shaped/autogen/api/{model_management_api.py → view_api.py} +179 -181
  8. shaped/autogen/api_client.py +18 -11
  9. shaped/autogen/configuration.py +22 -9
  10. shaped/autogen/exceptions.py +25 -5
  11. shaped/autogen/models/__init__.py +247 -52
  12. shaped/autogen/models/ai_enrichment_view_config.py +123 -0
  13. shaped/autogen/models/{path.py → algorithm.py} +19 -19
  14. shaped/autogen/models/amplitude_table_config.py +106 -0
  15. shaped/autogen/models/ascending.py +148 -0
  16. shaped/autogen/models/attn_dropout_prob.py +136 -0
  17. shaped/autogen/models/attribute_journey.py +124 -0
  18. shaped/autogen/models/attribute_value.py +178 -0
  19. shaped/autogen/models/autoscaling_config.py +95 -0
  20. shaped/autogen/models/aws_pinpoint_table_config.py +108 -0
  21. shaped/autogen/models/batch_size.py +136 -0
  22. shaped/autogen/models/batch_size1.py +136 -0
  23. shaped/autogen/models/batch_size2.py +136 -0
  24. shaped/autogen/models/big_query_table_config.py +147 -0
  25. shaped/autogen/models/bm25.py +136 -0
  26. shaped/autogen/models/boosted_reorder_step.py +125 -0
  27. shaped/autogen/models/canary_rollout.py +99 -0
  28. shaped/autogen/models/candidate_attributes_retrieve_step.py +113 -0
  29. shaped/autogen/models/candidate_ids_retrieve_step.py +113 -0
  30. shaped/autogen/models/candidate_retrieval_strategy.py +41 -0
  31. shaped/autogen/models/clickhouse_table_config.py +146 -0
  32. shaped/autogen/models/column_order_retrieve_step.py +123 -0
  33. shaped/autogen/models/column_ordering.py +91 -0
  34. shaped/autogen/models/create_table_response.py +87 -0
  35. shaped/autogen/models/create_view_response.py +87 -0
  36. shaped/autogen/models/custom_table_config.py +135 -0
  37. shaped/autogen/models/data_compute_config.py +89 -0
  38. shaped/autogen/models/data_config.py +145 -0
  39. shaped/autogen/models/data_config_interaction_table.py +146 -0
  40. shaped/autogen/models/data_split_config.py +88 -0
  41. shaped/autogen/models/data_split_strategy.py +37 -0
  42. shaped/autogen/models/data_tier.py +37 -0
  43. shaped/autogen/models/default.py +246 -0
  44. shaped/autogen/models/delete_engine_response.py +87 -0
  45. shaped/autogen/models/delete_table_response.py +87 -0
  46. shaped/autogen/models/delete_view_response.py +87 -0
  47. shaped/autogen/models/deployment_config.py +117 -0
  48. shaped/autogen/models/distance_function.py +38 -0
  49. shaped/autogen/models/diversity_reorder_step.py +137 -0
  50. shaped/autogen/models/dropout_rate.py +136 -0
  51. shaped/autogen/models/dynamo_db_table_config.py +160 -0
  52. shaped/autogen/models/dynamo_db_table_config_scan_kwargs_value.py +138 -0
  53. shaped/autogen/models/embedder_batch_size.py +136 -0
  54. shaped/autogen/models/embedding_config.py +93 -0
  55. shaped/autogen/models/embedding_dim.py +136 -0
  56. shaped/autogen/models/embedding_dims.py +136 -0
  57. shaped/autogen/models/embedding_size.py +136 -0
  58. shaped/autogen/models/encoder.py +140 -0
  59. shaped/autogen/models/encoding_pooling_strategy.py +38 -0
  60. shaped/autogen/models/engine.py +109 -0
  61. shaped/autogen/models/engine_config_v2.py +152 -0
  62. shaped/autogen/models/engine_details_response.py +120 -0
  63. shaped/autogen/models/engine_schema.py +113 -0
  64. shaped/autogen/models/engine_schema_user_inner.py +146 -0
  65. shaped/autogen/models/entity_config.py +109 -0
  66. shaped/autogen/models/entity_journey.py +161 -0
  67. shaped/autogen/models/entity_type.py +38 -0
  68. shaped/autogen/models/evaluation_config.py +92 -0
  69. shaped/autogen/models/exploration_reorder_step.py +125 -0
  70. shaped/autogen/models/expression_filter_step.py +106 -0
  71. shaped/autogen/models/factors.py +136 -0
  72. shaped/autogen/models/factors1.py +136 -0
  73. shaped/autogen/models/feature.py +90 -0
  74. shaped/autogen/models/feature_type.py +60 -0
  75. shaped/autogen/models/file_table_config.py +112 -0
  76. shaped/autogen/models/filter_config.py +99 -0
  77. shaped/autogen/models/filter_dataset.py +164 -0
  78. shaped/autogen/models/filter_index_type.py +36 -0
  79. shaped/autogen/models/filter_retrieve_step.py +113 -0
  80. shaped/autogen/models/filter_step_explanation.py +165 -0
  81. shaped/autogen/models/filter_table.py +140 -0
  82. shaped/autogen/models/filter_type.py +134 -0
  83. shaped/autogen/models/global_filter.py +102 -0
  84. shaped/autogen/models/hidden_dropout_prob.py +136 -0
  85. shaped/autogen/models/hidden_size.py +136 -0
  86. shaped/autogen/models/hidden_size1.py +136 -0
  87. shaped/autogen/models/http_problem_response.py +115 -0
  88. shaped/autogen/models/http_validation_error.py +2 -2
  89. shaped/autogen/models/hugging_face_encoder.py +115 -0
  90. shaped/autogen/models/iceberg_table_config.py +154 -0
  91. shaped/autogen/models/index_config.py +101 -0
  92. shaped/autogen/models/inner_entity_id.py +144 -0
  93. shaped/autogen/models/inner_size.py +136 -0
  94. shaped/autogen/models/inner_size1.py +136 -0
  95. shaped/autogen/models/inner_uid.py +144 -0
  96. shaped/autogen/models/interaction_config.py +122 -0
  97. shaped/autogen/models/interaction_pooling_encoder.py +104 -0
  98. shaped/autogen/models/interaction_round_robin_encoder.py +104 -0
  99. shaped/autogen/models/item_attribute_pooling_encoder.py +111 -0
  100. shaped/autogen/models/journey.py +140 -0
  101. shaped/autogen/models/kafka_table_config.py +129 -0
  102. shaped/autogen/models/kinesis_table_config.py +140 -0
  103. shaped/autogen/models/kinesis_table_config_column_schema_value.py +136 -0
  104. shaped/autogen/models/label.py +90 -0
  105. shaped/autogen/models/label_type.py +36 -0
  106. shaped/autogen/models/laplace_smoothing.py +136 -0
  107. shaped/autogen/models/latency_scaling_policy.py +112 -0
  108. shaped/autogen/models/learning_rate.py +136 -0
  109. shaped/autogen/models/learning_rate1.py +136 -0
  110. shaped/autogen/models/learning_rate2.py +136 -0
  111. shaped/autogen/models/learning_rate3.py +136 -0
  112. shaped/autogen/models/lexical_search_mode.py +99 -0
  113. shaped/autogen/models/list_engines_response.py +95 -0
  114. shaped/autogen/models/list_tables_response.py +95 -0
  115. shaped/autogen/models/list_views_response.py +95 -0
  116. shaped/autogen/models/location_inner.py +138 -0
  117. shaped/autogen/models/loss_types.py +37 -0
  118. shaped/autogen/models/lr.py +136 -0
  119. shaped/autogen/models/lr1.py +136 -0
  120. shaped/autogen/models/lr2.py +136 -0
  121. shaped/autogen/models/max_depth.py +136 -0
  122. shaped/autogen/models/max_leaves.py +136 -0
  123. shaped/autogen/models/max_seq_length.py +136 -0
  124. shaped/autogen/models/max_seq_length1.py +136 -0
  125. shaped/autogen/models/max_seq_length2.py +136 -0
  126. shaped/autogen/models/mode.py +134 -0
  127. shaped/autogen/models/mode1.py +136 -0
  128. shaped/autogen/models/mode2.py +150 -0
  129. shaped/autogen/models/models_inner.py +308 -0
  130. shaped/autogen/models/mongo_db_table_config.py +147 -0
  131. shaped/autogen/models/mssql_table_config.py +155 -0
  132. shaped/autogen/models/{my_sql_dataset_config.py → my_sql_table_config.py} +45 -28
  133. shaped/autogen/models/n_epochs.py +136 -0
  134. shaped/autogen/models/n_epochs1.py +136 -0
  135. shaped/autogen/models/n_epochs2.py +136 -0
  136. shaped/autogen/models/n_estimators.py +136 -0
  137. shaped/autogen/models/n_heads.py +136 -0
  138. shaped/autogen/models/n_layers.py +136 -0
  139. shaped/autogen/models/neg_per_positive.py +136 -0
  140. shaped/autogen/models/negative_samples_count.py +136 -0
  141. shaped/autogen/models/ngram_tokenizer.py +103 -0
  142. shaped/autogen/models/no_op_config.py +117 -0
  143. shaped/autogen/models/num_blocks.py +136 -0
  144. shaped/autogen/models/num_heads.py +136 -0
  145. shaped/autogen/models/num_leaves.py +136 -0
  146. shaped/autogen/models/objective.py +40 -0
  147. shaped/autogen/models/objective1.py +134 -0
  148. shaped/autogen/models/online_store_config.py +89 -0
  149. shaped/autogen/models/pagination_config.py +87 -0
  150. shaped/autogen/models/parameter_definition.py +96 -0
  151. shaped/autogen/models/parameters_value.py +240 -0
  152. shaped/autogen/models/passthrough_score.py +104 -0
  153. shaped/autogen/models/personal_filter.py +104 -0
  154. shaped/autogen/models/pipeline_stage_explanation.py +118 -0
  155. shaped/autogen/models/policy.py +134 -0
  156. shaped/autogen/models/pool_fn.py +152 -0
  157. shaped/autogen/models/pooling_function.py +37 -0
  158. shaped/autogen/models/{postgres_dataset_config.py → postgres_table_config.py} +66 -28
  159. shaped/autogen/models/posthog_table_config.py +133 -0
  160. shaped/autogen/models/prebuilt_filter_step.py +113 -0
  161. shaped/autogen/models/precomputed_item_embedding.py +99 -0
  162. shaped/autogen/models/precomputed_user_embedding.py +99 -0
  163. shaped/autogen/models/query.py +136 -0
  164. shaped/autogen/models/query1.py +136 -0
  165. shaped/autogen/models/query_any_of.py +172 -0
  166. shaped/autogen/models/query_config.py +140 -0
  167. shaped/autogen/models/query_definition.py +106 -0
  168. shaped/autogen/models/query_encoder.py +194 -0
  169. shaped/autogen/models/query_explanation.py +201 -0
  170. shaped/autogen/models/query_request.py +121 -0
  171. shaped/autogen/models/query_result.py +113 -0
  172. shaped/autogen/models/query_table_config.py +99 -0
  173. shaped/autogen/models/rank_item_attribute_values_query_config.py +122 -0
  174. shaped/autogen/models/rank_query_config.py +167 -0
  175. shaped/autogen/models/rank_query_config_filter_inner.py +149 -0
  176. shaped/autogen/models/rank_query_config_reorder_inner.py +186 -0
  177. shaped/autogen/models/rank_query_config_retrieve_inner.py +265 -0
  178. shaped/autogen/models/recreate_rollout.py +97 -0
  179. shaped/autogen/models/{redshift_dataset_config.py → redshift_table_config.py} +48 -25
  180. shaped/autogen/models/reference_table_config.py +113 -0
  181. shaped/autogen/models/regularization.py +136 -0
  182. shaped/autogen/models/reorder_inner.py +149 -0
  183. shaped/autogen/models/reorder_step_explanation.py +207 -0
  184. shaped/autogen/models/request.py +378 -0
  185. shaped/autogen/models/request1.py +140 -0
  186. shaped/autogen/models/requests_per_second_scaling_policy.py +112 -0
  187. shaped/autogen/models/resource_config.py +100 -0
  188. shaped/autogen/models/response_get_view_details_views_view_name_get.py +134 -0
  189. shaped/autogen/models/result.py +132 -0
  190. shaped/autogen/models/result_embeddings_value.py +127 -0
  191. shaped/autogen/models/retrieval_scores_value.py +127 -0
  192. shaped/autogen/models/retrieve_inner.py +196 -0
  193. shaped/autogen/models/retrieve_step_explanation.py +172 -0
  194. shaped/autogen/models/retriever.py +196 -0
  195. shaped/autogen/models/retriever1.py +196 -0
  196. shaped/autogen/models/rollout_config.py +91 -0
  197. shaped/autogen/models/rudderstack_table_config.py +106 -0
  198. shaped/autogen/models/sampling_strategy.py +36 -0
  199. shaped/autogen/models/saved_query_info_response.py +103 -0
  200. shaped/autogen/models/saved_query_list_response.py +87 -0
  201. shaped/autogen/models/saved_query_request.py +115 -0
  202. shaped/autogen/models/schema_config.py +117 -0
  203. shaped/autogen/models/score.py +134 -0
  204. shaped/autogen/models/score_ensemble.py +127 -0
  205. shaped/autogen/models/score_ensemble_policy_config.py +165 -0
  206. shaped/autogen/models/score_ensemble_policy_config_policies_inner.py +755 -0
  207. shaped/autogen/models/score_step_explanation.py +224 -0
  208. shaped/autogen/models/search_config.py +105 -0
  209. shaped/autogen/models/segment_table_config.py +106 -0
  210. shaped/autogen/models/sequence_length.py +136 -0
  211. shaped/autogen/models/server_config.py +100 -0
  212. shaped/autogen/models/setup_engine_response.py +87 -0
  213. shaped/autogen/models/shaped_internal_recsys_policies_als_model_policy_als_model_policy_config.py +154 -0
  214. shaped/autogen/models/shaped_internal_recsys_policies_beeformer_model_policy_beeformer_model_policy_beeformer_model_policy_config.py +160 -0
  215. shaped/autogen/models/shaped_internal_recsys_policies_bert_model_policy_bert_model_policy_bert_model_policy_config.py +215 -0
  216. shaped/autogen/models/shaped_internal_recsys_policies_chronological_model_policy_chronological_model_policy_config.py +156 -0
  217. shaped/autogen/models/shaped_internal_recsys_policies_elsa_model_policy_elsa_model_policy_elsa_model_policy_config.py +141 -0
  218. shaped/autogen/models/shaped_internal_recsys_policies_gsasrec_model_policy_gsasrec_model_policy_gsas_rec_model_policy_config.py +211 -0
  219. shaped/autogen/models/shaped_internal_recsys_policies_item2vec_model_policy_item2_vec_model_policy_config.py +140 -0
  220. shaped/autogen/models/shaped_internal_recsys_policies_item_content_similarity_model_policy_item_content_similarity_model_policy_config.py +195 -0
  221. shaped/autogen/models/shaped_internal_recsys_policies_lightgbm_model_policy_lightgbm_model_policy_light_gbm_model_policy_config.py +245 -0
  222. shaped/autogen/models/shaped_internal_recsys_policies_ngram_model_policy_ngram_model_policy_config.py +119 -0
  223. shaped/autogen/models/shaped_internal_recsys_policies_popular_model_policy_popular_model_policy_config.py +161 -0
  224. shaped/autogen/models/shaped_internal_recsys_policies_random_model_policy_random_model_policy_config.py +108 -0
  225. shaped/autogen/models/shaped_internal_recsys_policies_recently_popular_policy_recently_popular_policy_config.py +143 -0
  226. shaped/autogen/models/shaped_internal_recsys_policies_rising_popular_policy_rising_popular_policy_config.py +123 -0
  227. shaped/autogen/models/shaped_internal_recsys_policies_sasrec_model_policy_sasrec_model_policy_sas_rec_model_policy_config.py +230 -0
  228. shaped/autogen/models/shaped_internal_recsys_policies_svd_model_policy_svd_model_policy_config.py +119 -0
  229. shaped/autogen/models/shaped_internal_recsys_policies_two_tower_model_policy_two_tower_model_policy_two_tower_model_policy_config.py +165 -0
  230. shaped/autogen/models/shaped_internal_recsys_policies_user_item_content_similarity_model_policy_user_item_content_similarity_model_policy_config.py +153 -0
  231. shaped/autogen/models/shaped_internal_recsys_policies_widedeep_model_policy_wide_deep_model_policy_config.py +137 -0
  232. shaped/autogen/models/shaped_internal_recsys_policies_xgboost_model_policy_xg_boost_model_policy_config.py +155 -0
  233. shaped/autogen/models/shopify_table_config.py +156 -0
  234. shaped/autogen/models/similarity_retrieve_step.py +123 -0
  235. shaped/autogen/models/{snowflake_dataset_config.py → snowflake_table_config.py} +47 -18
  236. shaped/autogen/models/sql_transform_type.py +37 -0
  237. shaped/autogen/models/sql_view_config.py +111 -0
  238. shaped/autogen/models/stemmer_tokenizer.py +105 -0
  239. shaped/autogen/models/step_explanation.py +137 -0
  240. shaped/autogen/models/steps_inner.py +179 -0
  241. shaped/autogen/models/strategy.py +134 -0
  242. shaped/autogen/models/table.py +102 -0
  243. shaped/autogen/models/table_deployment_type.py +38 -0
  244. shaped/autogen/models/table_insert_arguments.py +87 -0
  245. shaped/autogen/models/table_insert_response.py +87 -0
  246. shaped/autogen/models/text_encoding.py +148 -0
  247. shaped/autogen/models/text_search_retrieve_step.py +121 -0
  248. shaped/autogen/models/time_frequency.py +136 -0
  249. shaped/autogen/models/time_window.py +136 -0
  250. shaped/autogen/models/time_window_in_days.py +154 -0
  251. shaped/autogen/models/tokenizer.py +149 -0
  252. shaped/autogen/models/trained_model_encoder.py +99 -0
  253. shaped/autogen/models/training_compute_config.py +99 -0
  254. shaped/autogen/models/training_config.py +121 -0
  255. shaped/autogen/models/training_config_models_inner.py +540 -0
  256. shaped/autogen/models/training_strategy.py +37 -0
  257. shaped/autogen/models/transform_status.py +41 -0
  258. shaped/autogen/models/trending_mode.py +37 -0
  259. shaped/autogen/models/truncate_filter_step.py +106 -0
  260. shaped/autogen/models/tunable_bool.py +97 -0
  261. shaped/autogen/models/tunable_float.py +118 -0
  262. shaped/autogen/models/tunable_int.py +118 -0
  263. shaped/autogen/models/tunable_int_categorical.py +99 -0
  264. shaped/autogen/models/tunable_string.py +99 -0
  265. shaped/autogen/models/tuning_config.py +89 -0
  266. shaped/autogen/models/type.py +150 -0
  267. shaped/autogen/models/update_table_response.py +87 -0
  268. shaped/autogen/models/update_view_response.py +87 -0
  269. shaped/autogen/models/user_attribute_pooling_encoder.py +124 -0
  270. shaped/autogen/models/user_inner.py +134 -0
  271. shaped/autogen/models/val_split.py +136 -0
  272. shaped/autogen/models/validation_error.py +13 -3
  273. shaped/autogen/models/validation_error_loc_inner.py +150 -0
  274. shaped/autogen/models/value_type.py +7 -5
  275. shaped/autogen/models/vector_search_mode.py +101 -0
  276. shaped/autogen/models/view.py +104 -0
  277. shaped/autogen/models/view_details_ai.py +140 -0
  278. shaped/autogen/models/view_details_ai_schema_value.py +153 -0
  279. shaped/autogen/models/view_details_sql.py +140 -0
  280. shaped/autogen/models/view_status.py +41 -0
  281. shaped/autogen/models/weight_decay.py +136 -0
  282. shaped/autogen/models/whitespace_tokenizer.py +97 -0
  283. shaped/autogen/models/window_size.py +136 -0
  284. shaped/autogen/rest.py +10 -4
  285. shaped/cli/shaped_cli.py +163 -35
  286. shaped/client.py +591 -171
  287. shaped/config_builders.py +705 -0
  288. shaped/query_builder.py +781 -0
  289. {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/METADATA +141 -6
  290. shaped-2.0.4.dist-info/RECORD +296 -0
  291. {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/WHEEL +1 -1
  292. shaped-2.0.4.dist-info/entry_points.txt +2 -0
  293. shaped/autogen/api/model_inference_api.py +0 -2825
  294. shaped/autogen/models/amplitude_dataset_config.py +0 -96
  295. shaped/autogen/models/aws_pinpoint_dataset_config.py +0 -96
  296. shaped/autogen/models/big_query_dataset_config.py +0 -114
  297. shaped/autogen/models/complement_items_request.py +0 -99
  298. shaped/autogen/models/complement_items_response.py +0 -89
  299. shaped/autogen/models/connectors_inner.py +0 -134
  300. shaped/autogen/models/create_dataset_arguments.py +0 -263
  301. shaped/autogen/models/create_embedding_response.py +0 -87
  302. shaped/autogen/models/create_item_embedding_request.py +0 -89
  303. shaped/autogen/models/create_model_arguments.py +0 -107
  304. shaped/autogen/models/create_model_response.py +0 -87
  305. shaped/autogen/models/create_user_embedding_request.py +0 -89
  306. shaped/autogen/models/custom_dataset_config.py +0 -115
  307. shaped/autogen/models/dataset_config.py +0 -101
  308. shaped/autogen/models/dataset_schema_type.py +0 -47
  309. shaped/autogen/models/datasets_inner.py +0 -91
  310. shaped/autogen/models/delete_model_response.py +0 -87
  311. shaped/autogen/models/fetch_config.py +0 -95
  312. shaped/autogen/models/file_config.py +0 -105
  313. shaped/autogen/models/file_source_config.py +0 -89
  314. shaped/autogen/models/inference_config.py +0 -101
  315. shaped/autogen/models/insert_model_response.py +0 -87
  316. shaped/autogen/models/interaction.py +0 -87
  317. shaped/autogen/models/list_datasets_response.py +0 -95
  318. shaped/autogen/models/list_models_response.py +0 -95
  319. shaped/autogen/models/model_config.py +0 -99
  320. shaped/autogen/models/model_response.py +0 -95
  321. shaped/autogen/models/mongo_db_dataset_config.py +0 -119
  322. shaped/autogen/models/post_rank_request.py +0 -117
  323. shaped/autogen/models/rank_attribute_response.py +0 -89
  324. shaped/autogen/models/rank_grid_attribute_request.py +0 -91
  325. shaped/autogen/models/rank_grid_attribute_request1.py +0 -93
  326. shaped/autogen/models/rank_grid_attribute_response.py +0 -91
  327. shaped/autogen/models/rank_response.py +0 -91
  328. shaped/autogen/models/retrieve_request.py +0 -101
  329. shaped/autogen/models/retrieve_response.py +0 -91
  330. shaped/autogen/models/retriever_top_k_override.py +0 -97
  331. shaped/autogen/models/rudder_stack_dataset_config.py +0 -96
  332. shaped/autogen/models/segment_dataset_config.py +0 -96
  333. shaped/autogen/models/similar_item_request.py +0 -101
  334. shaped/autogen/models/similar_response.py +0 -89
  335. shaped/autogen/models/similar_users_request.py +0 -99
  336. shaped/autogen/models/successful_response.py +0 -87
  337. shaped/autogen/models/view_model_response.py +0 -99
  338. shaped-2.0.1.dist-info/RECORD +0 -73
  339. shaped-2.0.1.dist-info/entry_points.txt +0 -2
  340. {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/top_level.txt +0 -0
  341. {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/zip-safe +0 -0
shaped/client.py CHANGED
@@ -1,35 +1,99 @@
1
1
  import logging
2
- from typing import Any, Dict, List, Optional
2
+ import socket
3
+ from typing import Any, Dict, List, Optional, Union
3
4
 
4
- from shaped.autogen.api.model_inference_api import ModelInferenceApi
5
+ from shaped.autogen.api.engine_api import EngineApi
6
+ from shaped.autogen.api.query_api import QueryApi
7
+ from shaped.autogen.api.table_api import TableApi
8
+ from shaped.autogen.api.view_api import ViewApi
5
9
  from shaped.autogen.api_client import ApiClient
6
10
  from shaped.autogen.configuration import Configuration
7
- from shaped.autogen.models.complement_items_request import (
8
- ComplementItemsRequest,
11
+ from shaped.autogen.models.ai_enrichment_view_config import AIEnrichmentViewConfig
12
+ from shaped.autogen.models.create_table_response import CreateTableResponse
13
+ from shaped.autogen.models.create_view_response import CreateViewResponse
14
+ from shaped.autogen.models.delete_engine_response import DeleteEngineResponse
15
+ from shaped.autogen.models.delete_table_response import DeleteTableResponse
16
+ from shaped.autogen.models.delete_view_response import DeleteViewResponse
17
+ from shaped.autogen.models.engine_config_v2 import EngineConfigV2
18
+ from shaped.autogen.models.engine_details_response import EngineDetailsResponse
19
+ from shaped.autogen.models.list_engines_response import ListEnginesResponse
20
+ from shaped.autogen.models.list_tables_response import ListTablesResponse
21
+ from shaped.autogen.models.list_views_response import ListViewsResponse
22
+ from shaped.autogen.models.parameters_value import ParametersValue
23
+ from shaped.autogen.models.query1 import Query1
24
+ from shaped.autogen.models.query_request import QueryRequest
25
+ from shaped.autogen.models.query_result import QueryResult
26
+ from shaped.autogen.models.rank_query_config import RankQueryConfig
27
+ from shaped.autogen.models.request import Request as TableRequest
28
+ from shaped.autogen.models.request1 import Request1 as ViewRequest
29
+ from shaped.autogen.models.response_get_view_details_views_view_name_get import (
30
+ ResponseGetViewDetailsViewsViewNameGet,
9
31
  )
10
- from shaped.autogen.models.inference_config import InferenceConfig
11
- from shaped.autogen.models.interaction import Interaction
12
- from shaped.autogen.models.post_rank_request import PostRankRequest
13
- from shaped.autogen.models.retrieve_request import RetrieveRequest
14
- from shaped.autogen.models.similar_item_request import SimilarItemRequest
15
- from shaped.autogen.models.similar_users_request import (
16
- SimilarUsersRequest,
17
- )
18
- from pydantic import StrictBool, StrictStr
32
+ from shaped.autogen.models.saved_query_info_response import SavedQueryInfoResponse
33
+ from shaped.autogen.models.saved_query_list_response import SavedQueryListResponse
34
+ from shaped.autogen.models.saved_query_request import SavedQueryRequest
35
+ from shaped.autogen.models.setup_engine_response import SetupEngineResponse
36
+ from shaped.autogen.models.sql_view_config import SQLViewConfig
37
+ from shaped.autogen.models.table_insert_arguments import TableInsertArguments
38
+ from shaped.autogen.models.table_insert_response import TableInsertResponse
39
+ from shaped.autogen.models.update_table_response import UpdateTableResponse
40
+ from shaped.autogen.models.update_view_response import UpdateViewResponse
41
+ from shaped.config_builders import Engine
42
+ from shaped.query_builder import RankQueryBuilder
43
+
44
+ # Type aliases for better developer experience.
45
+ ViewConfig = Union[SQLViewConfig, AIEnrichmentViewConfig]
19
46
 
20
47
 
21
48
  class Client:
22
49
  """
23
- This class provides a clean user interface to access Shaped functionality.
24
- It currently supports calls to the Model Inference API.
50
+ Client SDK for Shaped AI V2 API.
51
+
52
+ Provides access to all V2 API endpoints: Engine, Table, View, and Query APIs.
25
53
  """
26
54
 
27
- def __init__(self, api_key: str):
55
+ def __init__(self, api_key: str, base_url: str = "https://api.shaped.ai/v2"):
56
+ """
57
+ Initialize the Shaped client.
58
+
59
+ Args:
60
+ api_key: Your Shaped API key.
61
+ base_url: Base URL for the API (defaults to V2 endpoint).
62
+ """
28
63
  self._api_key = api_key
64
+ self._base_url = base_url
29
65
  api_key_dict = {"main": api_key}
30
66
  self._configuration = Configuration(api_key=api_key_dict)
67
+ if base_url != "https://api.shaped.ai/v2":
68
+ self._configuration.host = base_url
69
+
70
+ # Optimize connection pool for high-performance query endpoints.
71
+ # Increased pool size supports high-throughput scenarios (1k+ qps).
72
+ self._configuration.connection_pool_maxsize = 100
73
+
74
+ # Disable retries for query endpoints to minimize latency.
75
+ # Fail fast rather than retry for high-throughput scenarios.
76
+ self._configuration.retries = 0
77
+
78
+ # Configure socket options for TCP keep-alive.
79
+ self._configuration.socket_options = [
80
+ (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
81
+ ]
82
+
83
+ # HTTP/2 support: With urllib3 >= 2.2.3 and h2 installed,
84
+ # HTTP/2 will be automatically negotiated if the server supports it.
31
85
  self._api_client = ApiClient(self._configuration)
32
- self._model_inference_api = ModelInferenceApi(self._api_client)
86
+
87
+ # Enable response compression for high-performance query endpoints.
88
+ # Reduces bandwidth and improves performance for large query results.
89
+ self._api_client.set_default_header("Accept-Encoding", "gzip, deflate")
90
+
91
+ # Initialize V2 API instances.
92
+ self._engine_api = EngineApi(self._api_client)
93
+ self._table_api = TableApi(self._api_client)
94
+ self._view_api = ViewApi(self._api_client)
95
+ self._query_api = QueryApi(self._api_client)
96
+
33
97
  self._logger = logging.getLogger(__name__)
34
98
  if not self._logger.hasHandlers():
35
99
  console_handler = logging.StreamHandler()
@@ -39,173 +103,529 @@ class Client:
39
103
  self._logger.addHandler(console_handler)
40
104
  self._logger.setLevel(logging.INFO)
41
105
 
42
- def rank(
43
- self,
44
- model_name: StrictStr,
45
- user_id: Optional[StrictStr] = None,
46
- item_ids: Optional[List[StrictStr]] = None,
47
- interactions: Optional[List[Interaction]] = None,
48
- filter_predicate: Optional[StrictStr] = None,
49
- user_features: Optional[Dict[str, Any]] = None,
50
- item_features: Optional[Dict[str, Any]] = None,
51
- text_query: Optional[StrictStr] = None,
52
- flush_paginations: Optional[StrictBool] = None,
53
- return_metadata: Optional[StrictBool] = None,
54
- config: Optional[InferenceConfig] = None,
55
- ):
56
- """
57
- Rank returns the list of relevant item ids (from most-relevant to least) for the
58
- given request context.
59
-
60
- There are several ways to use rank depending on the combination of request
61
- arguments given.
62
-
63
- For more information visit https://docs.shaped.ai/docs/api#tag/Model-Inference
64
- """
65
- self._logger.debug("Calling rank on model %s", model_name)
66
- post_rank_request = PostRankRequest(
67
- user_id=user_id,
68
- item_ids=item_ids,
69
- interactions=interactions,
70
- filter_predicate=filter_predicate,
71
- user_features=user_features,
72
- item_features=item_features,
73
- text_query=text_query,
74
- flush_paginations=flush_paginations,
75
- return_metadata=return_metadata,
76
- config=config,
106
+ def _query_api_call(self, method: Any, *args: Any, **kwargs: Any) -> Any:
107
+ """
108
+ Helper for Query API calls that require headers.
109
+
110
+ Args:
111
+ method: The API method to call.
112
+ *args: Positional arguments for the method.
113
+ **kwargs: Keyword arguments for the method.
114
+
115
+ Returns:
116
+ Result from the API call.
117
+ """
118
+ kwargs.setdefault("_headers", {})["x-api-key"] = self._api_key
119
+ return method(*args, **kwargs)
120
+
121
+ # Engine API Methods
122
+
123
+ def create_engine(
124
+ self, engine_config: Union[EngineConfigV2, Engine, Dict[str, Any]]
125
+ ) -> SetupEngineResponse:
126
+ """
127
+ Create a new engine.
128
+
129
+ Args:
130
+ engine_config: Engine configuration (dict, EngineConfigV2, or Engine builder).
131
+
132
+ Returns:
133
+ SetupEngineResponse with engine URL.
134
+ """
135
+ self._logger.debug("Creating engine")
136
+
137
+ # Handle Engine builder
138
+ if isinstance(engine_config, Engine):
139
+ engine_config = engine_config.build()
140
+
141
+ # Convert dict to EngineConfigV2 to ensure proper model conversion.
142
+ if isinstance(engine_config, dict):
143
+ try:
144
+ engine_config = EngineConfigV2.from_dict(engine_config)
145
+ except Exception:
146
+ # If conversion fails, pass dict as-is and let API handle it.
147
+ pass
148
+ return self._engine_api.post_setup_engine_engines_post(
149
+ x_api_key=self._api_key, engine_config_v2=engine_config
77
150
  )
78
- return self._model_inference_api.post_rank_models_model_id_rank_post(
79
- model_name=model_name,
80
- x_api_key=self._api_key,
81
- post_rank_request=post_rank_request
151
+
152
+ def update_engine(
153
+ self, engine_config: Union[EngineConfigV2, Engine, Dict[str, Any]]
154
+ ) -> SetupEngineResponse:
155
+ """
156
+ Update an existing engine.
157
+
158
+ Args:
159
+ engine_config: Updated engine configuration (dict, EngineConfigV2, or Engine builder).
160
+
161
+ Returns:
162
+ SetupEngineResponse.
163
+ """
164
+ self._logger.debug("Updating engine")
165
+
166
+ # Handle Engine builder
167
+ if isinstance(engine_config, Engine):
168
+ engine_config = engine_config.build()
169
+
170
+ return self._engine_api.patch_update_engine_engines_patch(
171
+ x_api_key=self._api_key, engine_config_v2=engine_config
82
172
  )
83
173
 
84
- def retrieve(
85
- self,
86
- model_name: StrictStr,
87
- user_id_query: Optional[StrictStr] = None,
88
- text_query: Optional[StrictStr] = None,
89
- filter_predicate: Optional[StrictStr] = None,
90
- flush_paginations: Optional[StrictBool] = None,
91
- return_metadata: Optional[StrictBool] = None,
92
- config: Optional[InferenceConfig] = None,
93
- ):
94
- """
95
- Retrieve returns relevant item_ids for the given input text or user query. It
96
- can be used instead of rank if the filtering, scoring and ordering stages aren't
97
- needed for the final ranking.
98
-
99
- Typically people use this endpoint over rank if they want to reduce latency and
100
- complexity of the ranking pipeline and only need a subset of the functionality,
101
- e.g. just search but without personalization.
102
- """
103
- self._logger.debug("Calling retrieve on model %s", model_name)
104
- retrieve_request = RetrieveRequest(
105
- user_id_query=user_id_query,
106
- text_query=text_query,
107
- filter_predicate=filter_predicate,
108
- flush_paginations=flush_paginations,
109
- return_metadata=return_metadata,
110
- config=config,
174
+ def list_engines(self) -> ListEnginesResponse:
175
+ """
176
+ List all engines.
177
+
178
+ Returns:
179
+ ListEnginesResponse with list of engines.
180
+ """
181
+ self._logger.debug("Listing engines")
182
+ return self._engine_api.get_engines_engines_get(x_api_key=self._api_key)
183
+
184
+ def get_engine(self, engine_name: str) -> EngineDetailsResponse:
185
+ """
186
+ Get details for a specific engine.
187
+
188
+ Args:
189
+ engine_name: Name of the engine.
190
+
191
+ Returns:
192
+ EngineDetailsResponse with engine details.
193
+ """
194
+ self._logger.debug("Getting engine details for %s", engine_name)
195
+ return self._engine_api.get_engine_details_engines_engine_name_get(
196
+ engine_name=engine_name, x_api_key=self._api_key
197
+ )
198
+
199
+ def delete_engine(self, engine_name: str) -> DeleteEngineResponse:
200
+ """
201
+ Delete an engine.
202
+
203
+ Args:
204
+ engine_name: Name of the engine to delete.
205
+
206
+ Returns:
207
+ DeleteEngineResponse.
208
+ """
209
+ self._logger.debug("Deleting engine %s", engine_name)
210
+ return self._engine_api.delete_engine_engines_engine_name_delete(
211
+ engine_name=engine_name, x_api_key=self._api_key
212
+ )
213
+
214
+ # Table API Methods
215
+
216
+ def create_table(
217
+ self, table_config: Union[TableRequest, Dict[str, Any]]
218
+ ) -> CreateTableResponse:
219
+ """
220
+ Create a new table.
221
+
222
+ Args:
223
+ table_config: Table configuration (one of the TableConfig types) or dict.
224
+
225
+ Returns:
226
+ CreateTableResponse.
227
+ """
228
+ self._logger.debug("Creating table")
229
+ # Convert dict to Request if needed
230
+ if isinstance(table_config, dict):
231
+ request = TableRequest.from_dict(table_config)
232
+ elif not isinstance(table_config, TableRequest):
233
+ # Assume it's a specific table config (e.g. PostgresTableConfig)
234
+ # and wrap it in the Request wrapper.
235
+ try:
236
+ request = TableRequest(table_config)
237
+ except Exception:
238
+ # Fallback
239
+ request = table_config
240
+ else:
241
+ request = table_config
242
+ return self._table_api.post_create_table_tables_post(
243
+ x_api_key=self._api_key, request=request
111
244
  )
112
- return self._model_inference_api.post_retrieve_models_model_id_retrieve_post(
113
- model_name=model_name,
245
+
246
+ def update_table(
247
+ self, table_config: Union[TableRequest, Dict[str, Any]]
248
+ ) -> UpdateTableResponse:
249
+ """
250
+ Update an existing table.
251
+
252
+ Args:
253
+ table_config: Updated table configuration (one of the TableConfig types) or dict.
254
+
255
+ Returns:
256
+ UpdateTableResponse.
257
+ """
258
+ self._logger.debug("Updating table")
259
+ # Convert dict to Request if needed.
260
+ if isinstance(table_config, dict):
261
+ request = TableRequest.from_dict(table_config)
262
+ elif not isinstance(table_config, TableRequest):
263
+ try:
264
+ request = TableRequest(table_config)
265
+ except Exception:
266
+ request = table_config
267
+ else:
268
+ request = table_config
269
+ return self._table_api.patch_update_table_tables_patch(
270
+ x_api_key=self._api_key, request=request
271
+ )
272
+
273
+ def list_tables(self) -> ListTablesResponse:
274
+ """
275
+ List all tables.
276
+
277
+ Returns:
278
+ ListTablesResponse with list of tables.
279
+ """
280
+ self._logger.debug("Listing tables")
281
+ return self._table_api.get_tables_tables_get(x_api_key=self._api_key)
282
+
283
+ def insert_table_rows(
284
+ self, table_name: str, rows: List[Dict[str, Any]]
285
+ ) -> TableInsertResponse:
286
+ """
287
+ Insert rows into a table.
288
+
289
+ Args:
290
+ table_name: Name of the table.
291
+ rows: List of row dictionaries to insert.
292
+
293
+ Returns:
294
+ TableInsertResponse.
295
+ """
296
+ self._logger.debug("Inserting rows into table %s", table_name)
297
+ insert_args = TableInsertArguments(data=rows)
298
+ return self._table_api.post_table_insert_tables_table_name_insert_post(
299
+ table_name=table_name,
114
300
  x_api_key=self._api_key,
115
- retrieve_request=retrieve_request,
301
+ table_insert_arguments=insert_args,
116
302
  )
117
303
 
118
- def similar_items(
119
- self,
120
- model_name: StrictStr,
121
- item_id: StrictStr = None,
122
- user_id: Optional[StrictStr] = None,
123
- return_metadata: Optional[StrictBool] = None,
124
- flush_paginations: Optional[StrictBool] = None,
125
- filter_predicate: Optional[StrictStr] = None,
126
- config: Optional[InferenceConfig] = None,
127
- ):
128
- """
129
- Similar Items returns a list of similar items to the given input item query. If
130
- a user_id is given the response will return similar items personalized for that
131
- user.
132
- """
133
- self._logger.debug("Calling similar_items on model %s", model_name)
134
- similar_item_request = SimilarItemRequest(
135
- item_id=item_id,
136
- user_id=user_id,
137
- return_metadata=return_metadata,
138
- flush_paginations=flush_paginations,
139
- filter_predicate=filter_predicate,
140
- config=config,
141
- )
142
- return (
143
- self._model_inference_api
144
- .post_similar_items_models_model_name_similar_items_post(
145
- model_name=model_name,
146
- x_api_key=self._api_key,
147
- similar_item_request=similar_item_request,
148
- )
304
+ def delete_table(self, table_name: str) -> DeleteTableResponse:
305
+ """
306
+ Delete a table.
307
+
308
+ Args:
309
+ table_name: Name of the table to delete.
310
+
311
+ Returns:
312
+ DeleteTableResponse.
313
+ """
314
+ self._logger.debug("Deleting table %s", table_name)
315
+ return self._table_api.delete_table_route_tables_table_name_delete(
316
+ table_name=table_name, x_api_key=self._api_key
317
+ )
318
+
319
+ # View API Methods
320
+
321
+ def create_view(
322
+ self, view_config: Union[ViewConfig, Dict[str, Any]]
323
+ ) -> CreateViewResponse:
324
+ """
325
+ Create a new view.
326
+
327
+ Args:
328
+ view_config: View configuration (ViewConfig or dict).
329
+
330
+ Returns:
331
+ CreateViewResponse.
332
+ """
333
+ self._logger.debug("Creating view")
334
+ # Convert dict to Request1 if needed.
335
+ if isinstance(view_config, dict):
336
+ request = ViewRequest.from_dict(view_config)
337
+ elif not isinstance(view_config, ViewRequest):
338
+ # Try to wrap specific view config (SQLViewConfig etc)
339
+ try:
340
+ request = ViewRequest(view_config)
341
+ except Exception:
342
+ request = view_config
343
+ else:
344
+ request = view_config
345
+ return self._view_api.post_create_view_views_post(
346
+ x_api_key=self._api_key, request1=request
347
+ )
348
+
349
+ def update_view(
350
+ self, view_config: Union[ViewConfig, Dict[str, Any]]
351
+ ) -> UpdateViewResponse:
352
+ """
353
+ Update an existing view.
354
+
355
+ Args:
356
+ view_config: Updated view configuration (ViewConfig or dict).
357
+
358
+ Returns:
359
+ UpdateViewResponse.
360
+ """
361
+ self._logger.debug("Updating view")
362
+ # Convert dict to Request1 if needed.
363
+ if isinstance(view_config, dict):
364
+ request = ViewRequest.from_dict(view_config)
365
+ elif not isinstance(view_config, ViewRequest):
366
+ try:
367
+ request = ViewRequest(view_config)
368
+ except Exception:
369
+ request = view_config
370
+ else:
371
+ request = view_config
372
+ return self._view_api.patch_update_view_views_patch(
373
+ x_api_key=self._api_key, request1=request
149
374
  )
150
375
 
151
- def similar_users(
376
+ def list_views(self) -> ListViewsResponse:
377
+ """
378
+ List all views.
379
+
380
+ Returns:
381
+ ListViewsResponse with list of views.
382
+ """
383
+ self._logger.debug("Listing views")
384
+ return self._view_api.get_views_views_get(x_api_key=self._api_key)
385
+
386
+ def get_view(self, view_name: str) -> ResponseGetViewDetailsViewsViewNameGet:
387
+ """
388
+ Get details for a specific view.
389
+
390
+ Args:
391
+ view_name: Name of the view.
392
+
393
+ Returns:
394
+ ResponseGetViewDetailsViewsViewNameGet with view details.
395
+ """
396
+ self._logger.debug("Getting view details for %s", view_name)
397
+ return self._view_api.get_view_details_views_view_name_get(
398
+ view_name=view_name, x_api_key=self._api_key
399
+ )
400
+
401
+ def delete_view(self, view_name: str) -> DeleteViewResponse:
402
+ """
403
+ Delete a view.
404
+
405
+ Args:
406
+ view_name: Name of the view to delete.
407
+
408
+ Returns:
409
+ DeleteViewResponse.
410
+ """
411
+ self._logger.debug("Deleting view %s", view_name)
412
+ return self._view_api.delete_view_views_view_name_delete(
413
+ view_name=view_name, x_api_key=self._api_key
414
+ )
415
+
416
+ # Query API Methods
417
+
418
+ def execute_query(
152
419
  self,
153
- model_name: StrictStr,
154
- user_id: StrictStr = None,
155
- return_metadata: Optional[StrictBool] = None,
156
- flush_paginations: Optional[StrictBool] = None,
157
- filter_predicate: Optional[StrictStr] = None,
158
- config: Optional[InferenceConfig] = None,
159
- ):
160
- """
161
- Similar Users returns a list of similar user to the given input user query.
162
- """
163
- self._logger.debug("Calling similar_users on model %s", model_name)
164
- similar_users_request = SimilarUsersRequest(
165
- user_id=user_id,
166
- return_metadata=return_metadata,
167
- flush_paginations=flush_paginations,
168
- filter_predicate=filter_predicate,
169
- config=config,
170
- )
171
- return (
172
- self._model_inference_api
173
- .post_similar_users_models_model_name_similar_users_post(
174
- model_name=model_name,
175
- user_id=user_id,
176
- x_api_key=self._api_key,
177
- similar_users_request=similar_users_request,
420
+ engine_name: str,
421
+ query: Union[RankQueryConfig, RankQueryBuilder, Dict[str, Any], str],
422
+ parameters: Optional[Dict[str, Any]] = None,
423
+ return_metadata: bool = False,
424
+ return_explanation: bool = False,
425
+ return_journey_explanations: bool = False,
426
+ pagination_key: Optional[str] = None,
427
+ ignore_pagination: bool = False,
428
+ ) -> QueryResult:
429
+ """
430
+ Execute a query against a V2 engine.
431
+
432
+ Args:
433
+ engine_name: Name of the engine to query.
434
+ query: Query configuration. Can be:
435
+ - RankQueryConfig instance
436
+ - RankQueryBuilder instance (will be built)
437
+ - Dict representing a query config
438
+ - String SQL query
439
+ parameters: Query parameters dictionary for parameterized queries.
440
+ return_metadata: Whether to return metadata in results.
441
+ return_explanation: Whether to include detailed query execution
442
+ explanation.
443
+ return_journey_explanations: Whether to include per-entity journey
444
+ tracking in results.
445
+ pagination_key: Pagination key for continuing from a previous query.
446
+ ignore_pagination: Whether to ignore pagination and return results
447
+ from the beginning.
448
+
449
+ Returns:
450
+ QueryResult from the API.
451
+ """
452
+ self._logger.debug("Executing query")
453
+
454
+ # Convert query to proper format.
455
+ if isinstance(query, RankQueryBuilder):
456
+ query = query.build()
457
+ elif isinstance(query, dict):
458
+ # Assume it's already in the right format.
459
+ pass
460
+ elif isinstance(query, str):
461
+ # SQL query string - Query1 accepts strings directly.
462
+ pass
463
+ # If it's already a RankQueryConfig, convert to dict for QueryRequest.
464
+ elif isinstance(query, RankQueryConfig):
465
+ if hasattr(query, "model_dump"):
466
+ query = query.model_dump(by_alias=True)
467
+ else:
468
+ query = dict(query)
469
+
470
+ # Convert parameters to ParametersValue instances if needed.
471
+ converted_parameters = None
472
+ if parameters:
473
+ converted_parameters = {
474
+ k: ParametersValue(v) if not isinstance(v, ParametersValue) else v
475
+ for k, v in parameters.items()
476
+ }
477
+
478
+ # Query1 accepts both dict/QueryConfig and string.
479
+ # Convert string queries to Query1 instances.
480
+ if isinstance(query, str):
481
+ # String queries need to be wrapped in Query1
482
+ query_obj = Query1(query)
483
+ elif isinstance(query, dict):
484
+ # Convert sets to lists for JSON serialization.
485
+ def convert_sets_to_lists(obj):
486
+ """Recursively convert sets to lists for JSON serialization."""
487
+ if isinstance(obj, set):
488
+ return list(obj)
489
+ elif isinstance(obj, dict):
490
+ return {k: convert_sets_to_lists(v) for k, v in obj.items()}
491
+ elif isinstance(obj, list):
492
+ return [convert_sets_to_lists(item) for item in obj]
493
+ return obj
494
+
495
+ query_cleaned = convert_sets_to_lists(query)
496
+ # Dict queries need to be converted to Query1 via from_dict
497
+ query_obj = Query1.from_dict(query_cleaned)
498
+ elif isinstance(query, RankQueryConfig):
499
+ if hasattr(query, "model_dump"):
500
+ query_dict = query.model_dump(by_alias=True)
501
+ else:
502
+ query_dict = dict(query)
503
+
504
+ # Convert sets to lists for JSON serialization.
505
+ def convert_sets_to_lists(obj):
506
+ """Recursively convert sets to lists for JSON serialization."""
507
+ if isinstance(obj, set):
508
+ return list(obj)
509
+ elif isinstance(obj, dict):
510
+ return {k: convert_sets_to_lists(v) for k, v in obj.items()}
511
+ elif isinstance(obj, list):
512
+ return [convert_sets_to_lists(item) for item in obj]
513
+ return obj
514
+
515
+ query_dict_cleaned = convert_sets_to_lists(query_dict)
516
+ query_obj = Query1.from_dict(query_dict_cleaned)
517
+ else:
518
+ raise TypeError(
519
+ f"Invalid query type: {type(query)}. Expected RankQueryConfig, dict, or str."
178
520
  )
521
+
522
+ query_request = QueryRequest(
523
+ query=query_obj,
524
+ parameters=converted_parameters,
525
+ return_metadata=return_metadata,
526
+ return_explanation=return_explanation,
527
+ return_journey_explanations=return_journey_explanations,
528
+ pagination_key=pagination_key,
529
+ ignore_pagination=ignore_pagination,
530
+ )
531
+
532
+ return self._query_api_call(
533
+ self._query_api.execute_ad_hoc_query_query_post,
534
+ engine_name=engine_name,
535
+ query_request=query_request,
179
536
  )
180
537
 
181
- def complement_items(
538
+ def execute_saved_query(
182
539
  self,
183
- model_name: StrictStr,
184
- item_ids: List[StrictStr],
185
- user_id: Optional[StrictStr] = None,
186
- return_metadata: Optional[StrictBool] = None,
187
- filter_predicate: Optional[StrictStr] = None,
188
- config: Optional[InferenceConfig] = None,
189
- ):
190
- """
191
- Complement Items returns a list of items that complement the given input items.
192
- For example, if you have a list of items in a user's cart, you can use this
193
- endpoint to return items that complement the cart (i.e. "Complete-the-bag"
194
- use-cases).
195
- """
196
- self._logger.debug("Calling complement_items on model %s", model_name)
197
- complement_items_request = ComplementItemsRequest(
198
- item_ids=item_ids,
199
- user_id=user_id,
540
+ engine_name: str,
541
+ query_name: str,
542
+ parameters: Optional[Dict[str, Any]] = None,
543
+ return_metadata: bool = True,
544
+ return_explanation: bool = False,
545
+ return_journey_explanations: bool = False,
546
+ pagination_key: Optional[str] = None,
547
+ ignore_pagination: bool = False,
548
+ ) -> QueryResult:
549
+ """
550
+ Execute a saved query by name.
551
+
552
+ Args:
553
+ engine_name: Name of the engine containing the saved query.
554
+ query_name: Name of the saved query to execute.
555
+ parameters: Query parameters dictionary.
556
+ return_metadata: Whether to return metadata in results (defaults
557
+ to True for saved queries).
558
+ return_explanation: Whether to include detailed query execution
559
+ explanation.
560
+ return_journey_explanations: Whether to include per-entity journey
561
+ tracking in results.
562
+ pagination_key: Pagination key for continuing from a previous query.
563
+ ignore_pagination: Whether to ignore pagination and return results
564
+ from the beginning.
565
+
566
+ Returns:
567
+ QueryResult from the API.
568
+ """
569
+ self._logger.debug("Executing saved query %s", query_name)
570
+
571
+ # Convert parameters to ParametersValue instances if needed.
572
+ converted_parameters = None
573
+ if parameters:
574
+ converted_parameters = {
575
+ k: ParametersValue(v) if not isinstance(v, ParametersValue) else v
576
+ for k, v in parameters.items()
577
+ }
578
+
579
+ saved_query_request = SavedQueryRequest(
580
+ parameters=converted_parameters,
200
581
  return_metadata=return_metadata,
201
- filter_predicate=filter_predicate,
202
- config=config,
203
- )
204
- return (
205
- self._model_inference_api
206
- .post_complement_items_models_model_name_complement_items_post(
207
- model_name=model_name,
208
- x_api_key=self._api_key,
209
- complement_items_request=complement_items_request,
210
- )
582
+ return_explanation=return_explanation,
583
+ return_journey_explanations=return_journey_explanations,
584
+ pagination_key=pagination_key,
585
+ ignore_pagination=ignore_pagination,
586
+ )
587
+
588
+ return self._query_api_call(
589
+ self._query_api.execute_saved_query_queries_query_name_post,
590
+ engine_name=engine_name,
591
+ query_name=query_name,
592
+ saved_query_request=saved_query_request,
593
+ )
594
+
595
+ def get_saved_query_info(
596
+ self, engine_name: str, query_name: str
597
+ ) -> SavedQueryInfoResponse:
598
+ """
599
+ Get information about a saved query.
600
+
601
+ Args:
602
+ engine_name: Name of the engine containing the saved query.
603
+ query_name: Name of the saved query.
604
+
605
+ Returns:
606
+ SavedQueryInfoResponse with query information.
607
+ """
608
+ self._logger.debug(
609
+ "Getting saved query info for %s in engine %s", query_name, engine_name
610
+ )
611
+ return self._query_api_call(
612
+ self._query_api.get_saved_query_info_queries_query_name_get,
613
+ engine_name=engine_name,
614
+ query_name=query_name,
615
+ )
616
+
617
+ def list_saved_queries(self, engine_name: str) -> SavedQueryListResponse:
618
+ """
619
+ List all saved queries for an engine.
620
+
621
+ Args:
622
+ engine_name: Name of the engine to list saved queries for.
623
+
624
+ Returns:
625
+ SavedQueryListResponse with list of saved queries.
626
+ """
627
+ self._logger.debug("Listing saved queries for engine %s", engine_name)
628
+ return self._query_api_call(
629
+ self._query_api.list_saved_queries_queries_get,
630
+ engine_name=engine_name,
211
631
  )