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