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.
- shaped/__init__.py +59 -4
- shaped/autogen/__init__.py +527 -66
- shaped/autogen/api/__init__.py +4 -3
- shaped/autogen/api/engine_api.py +1467 -0
- shaped/autogen/api/{dataset_api.py → query_api.py} +159 -194
- shaped/autogen/api/table_api.py +1494 -0
- shaped/autogen/api/{model_management_api.py → view_api.py} +179 -181
- shaped/autogen/api_client.py +18 -11
- shaped/autogen/configuration.py +22 -9
- shaped/autogen/exceptions.py +25 -5
- shaped/autogen/models/__init__.py +247 -52
- shaped/autogen/models/ai_enrichment_view_config.py +123 -0
- shaped/autogen/models/{path.py → algorithm.py} +19 -19
- shaped/autogen/models/amplitude_table_config.py +106 -0
- shaped/autogen/models/ascending.py +148 -0
- shaped/autogen/models/attn_dropout_prob.py +136 -0
- shaped/autogen/models/attribute_journey.py +124 -0
- shaped/autogen/models/attribute_value.py +178 -0
- shaped/autogen/models/autoscaling_config.py +95 -0
- shaped/autogen/models/aws_pinpoint_table_config.py +108 -0
- shaped/autogen/models/batch_size.py +136 -0
- shaped/autogen/models/batch_size1.py +136 -0
- shaped/autogen/models/batch_size2.py +136 -0
- shaped/autogen/models/big_query_table_config.py +147 -0
- shaped/autogen/models/bm25.py +136 -0
- shaped/autogen/models/boosted_reorder_step.py +125 -0
- shaped/autogen/models/canary_rollout.py +99 -0
- shaped/autogen/models/candidate_attributes_retrieve_step.py +113 -0
- shaped/autogen/models/candidate_ids_retrieve_step.py +113 -0
- shaped/autogen/models/candidate_retrieval_strategy.py +41 -0
- shaped/autogen/models/clickhouse_table_config.py +146 -0
- shaped/autogen/models/column_order_retrieve_step.py +123 -0
- shaped/autogen/models/column_ordering.py +91 -0
- shaped/autogen/models/create_table_response.py +87 -0
- shaped/autogen/models/create_view_response.py +87 -0
- shaped/autogen/models/custom_table_config.py +135 -0
- shaped/autogen/models/data_compute_config.py +89 -0
- shaped/autogen/models/data_config.py +145 -0
- shaped/autogen/models/data_config_interaction_table.py +146 -0
- shaped/autogen/models/data_split_config.py +88 -0
- shaped/autogen/models/data_split_strategy.py +37 -0
- shaped/autogen/models/data_tier.py +37 -0
- shaped/autogen/models/default.py +246 -0
- shaped/autogen/models/delete_engine_response.py +87 -0
- shaped/autogen/models/delete_table_response.py +87 -0
- shaped/autogen/models/delete_view_response.py +87 -0
- shaped/autogen/models/deployment_config.py +117 -0
- shaped/autogen/models/distance_function.py +38 -0
- shaped/autogen/models/diversity_reorder_step.py +137 -0
- shaped/autogen/models/dropout_rate.py +136 -0
- shaped/autogen/models/dynamo_db_table_config.py +160 -0
- shaped/autogen/models/dynamo_db_table_config_scan_kwargs_value.py +138 -0
- shaped/autogen/models/embedder_batch_size.py +136 -0
- shaped/autogen/models/embedding_config.py +93 -0
- shaped/autogen/models/embedding_dim.py +136 -0
- shaped/autogen/models/embedding_dims.py +136 -0
- shaped/autogen/models/embedding_size.py +136 -0
- shaped/autogen/models/encoder.py +140 -0
- shaped/autogen/models/encoding_pooling_strategy.py +38 -0
- shaped/autogen/models/engine.py +109 -0
- shaped/autogen/models/engine_config_v2.py +152 -0
- shaped/autogen/models/engine_details_response.py +120 -0
- shaped/autogen/models/engine_schema.py +113 -0
- shaped/autogen/models/engine_schema_user_inner.py +146 -0
- shaped/autogen/models/entity_config.py +109 -0
- shaped/autogen/models/entity_journey.py +161 -0
- shaped/autogen/models/entity_type.py +38 -0
- shaped/autogen/models/evaluation_config.py +92 -0
- shaped/autogen/models/exploration_reorder_step.py +125 -0
- shaped/autogen/models/expression_filter_step.py +106 -0
- shaped/autogen/models/factors.py +136 -0
- shaped/autogen/models/factors1.py +136 -0
- shaped/autogen/models/feature.py +90 -0
- shaped/autogen/models/feature_type.py +60 -0
- shaped/autogen/models/file_table_config.py +112 -0
- shaped/autogen/models/filter_config.py +99 -0
- shaped/autogen/models/filter_dataset.py +164 -0
- shaped/autogen/models/filter_index_type.py +36 -0
- shaped/autogen/models/filter_retrieve_step.py +113 -0
- shaped/autogen/models/filter_step_explanation.py +165 -0
- shaped/autogen/models/filter_table.py +140 -0
- shaped/autogen/models/filter_type.py +134 -0
- shaped/autogen/models/global_filter.py +102 -0
- shaped/autogen/models/hidden_dropout_prob.py +136 -0
- shaped/autogen/models/hidden_size.py +136 -0
- shaped/autogen/models/hidden_size1.py +136 -0
- shaped/autogen/models/http_problem_response.py +115 -0
- shaped/autogen/models/http_validation_error.py +2 -2
- shaped/autogen/models/hugging_face_encoder.py +115 -0
- shaped/autogen/models/iceberg_table_config.py +154 -0
- shaped/autogen/models/index_config.py +101 -0
- shaped/autogen/models/inner_entity_id.py +144 -0
- shaped/autogen/models/inner_size.py +136 -0
- shaped/autogen/models/inner_size1.py +136 -0
- shaped/autogen/models/inner_uid.py +144 -0
- shaped/autogen/models/interaction_config.py +122 -0
- shaped/autogen/models/interaction_pooling_encoder.py +104 -0
- shaped/autogen/models/interaction_round_robin_encoder.py +104 -0
- shaped/autogen/models/item_attribute_pooling_encoder.py +111 -0
- shaped/autogen/models/journey.py +140 -0
- shaped/autogen/models/kafka_table_config.py +129 -0
- shaped/autogen/models/kinesis_table_config.py +140 -0
- shaped/autogen/models/kinesis_table_config_column_schema_value.py +136 -0
- shaped/autogen/models/label.py +90 -0
- shaped/autogen/models/label_type.py +36 -0
- shaped/autogen/models/laplace_smoothing.py +136 -0
- shaped/autogen/models/latency_scaling_policy.py +112 -0
- shaped/autogen/models/learning_rate.py +136 -0
- shaped/autogen/models/learning_rate1.py +136 -0
- shaped/autogen/models/learning_rate2.py +136 -0
- shaped/autogen/models/learning_rate3.py +136 -0
- shaped/autogen/models/lexical_search_mode.py +99 -0
- shaped/autogen/models/list_engines_response.py +95 -0
- shaped/autogen/models/list_tables_response.py +95 -0
- shaped/autogen/models/list_views_response.py +95 -0
- shaped/autogen/models/location_inner.py +138 -0
- shaped/autogen/models/loss_types.py +37 -0
- shaped/autogen/models/lr.py +136 -0
- shaped/autogen/models/lr1.py +136 -0
- shaped/autogen/models/lr2.py +136 -0
- shaped/autogen/models/max_depth.py +136 -0
- shaped/autogen/models/max_leaves.py +136 -0
- shaped/autogen/models/max_seq_length.py +136 -0
- shaped/autogen/models/max_seq_length1.py +136 -0
- shaped/autogen/models/max_seq_length2.py +136 -0
- shaped/autogen/models/mode.py +134 -0
- shaped/autogen/models/mode1.py +136 -0
- shaped/autogen/models/mode2.py +150 -0
- shaped/autogen/models/models_inner.py +308 -0
- shaped/autogen/models/mongo_db_table_config.py +147 -0
- shaped/autogen/models/mssql_table_config.py +155 -0
- shaped/autogen/models/{my_sql_dataset_config.py → my_sql_table_config.py} +45 -28
- shaped/autogen/models/n_epochs.py +136 -0
- shaped/autogen/models/n_epochs1.py +136 -0
- shaped/autogen/models/n_epochs2.py +136 -0
- shaped/autogen/models/n_estimators.py +136 -0
- shaped/autogen/models/n_heads.py +136 -0
- shaped/autogen/models/n_layers.py +136 -0
- shaped/autogen/models/neg_per_positive.py +136 -0
- shaped/autogen/models/negative_samples_count.py +136 -0
- shaped/autogen/models/ngram_tokenizer.py +103 -0
- shaped/autogen/models/no_op_config.py +117 -0
- shaped/autogen/models/num_blocks.py +136 -0
- shaped/autogen/models/num_heads.py +136 -0
- shaped/autogen/models/num_leaves.py +136 -0
- shaped/autogen/models/objective.py +40 -0
- shaped/autogen/models/objective1.py +134 -0
- shaped/autogen/models/online_store_config.py +89 -0
- shaped/autogen/models/pagination_config.py +87 -0
- shaped/autogen/models/parameter_definition.py +96 -0
- shaped/autogen/models/parameters_value.py +240 -0
- shaped/autogen/models/passthrough_score.py +104 -0
- shaped/autogen/models/personal_filter.py +104 -0
- shaped/autogen/models/pipeline_stage_explanation.py +118 -0
- shaped/autogen/models/policy.py +134 -0
- shaped/autogen/models/pool_fn.py +152 -0
- shaped/autogen/models/pooling_function.py +37 -0
- shaped/autogen/models/{postgres_dataset_config.py → postgres_table_config.py} +66 -28
- shaped/autogen/models/posthog_table_config.py +133 -0
- shaped/autogen/models/prebuilt_filter_step.py +113 -0
- shaped/autogen/models/precomputed_item_embedding.py +99 -0
- shaped/autogen/models/precomputed_user_embedding.py +99 -0
- shaped/autogen/models/query.py +136 -0
- shaped/autogen/models/query1.py +136 -0
- shaped/autogen/models/query_any_of.py +172 -0
- shaped/autogen/models/query_config.py +140 -0
- shaped/autogen/models/query_definition.py +106 -0
- shaped/autogen/models/query_encoder.py +194 -0
- shaped/autogen/models/query_explanation.py +201 -0
- shaped/autogen/models/query_request.py +121 -0
- shaped/autogen/models/query_result.py +113 -0
- shaped/autogen/models/query_table_config.py +99 -0
- shaped/autogen/models/rank_item_attribute_values_query_config.py +122 -0
- shaped/autogen/models/rank_query_config.py +167 -0
- shaped/autogen/models/rank_query_config_filter_inner.py +149 -0
- shaped/autogen/models/rank_query_config_reorder_inner.py +186 -0
- shaped/autogen/models/rank_query_config_retrieve_inner.py +265 -0
- shaped/autogen/models/recreate_rollout.py +97 -0
- shaped/autogen/models/{redshift_dataset_config.py → redshift_table_config.py} +48 -25
- shaped/autogen/models/reference_table_config.py +113 -0
- shaped/autogen/models/regularization.py +136 -0
- shaped/autogen/models/reorder_inner.py +149 -0
- shaped/autogen/models/reorder_step_explanation.py +207 -0
- shaped/autogen/models/request.py +378 -0
- shaped/autogen/models/request1.py +140 -0
- shaped/autogen/models/requests_per_second_scaling_policy.py +112 -0
- shaped/autogen/models/resource_config.py +100 -0
- shaped/autogen/models/response_get_view_details_views_view_name_get.py +134 -0
- shaped/autogen/models/result.py +132 -0
- shaped/autogen/models/result_embeddings_value.py +127 -0
- shaped/autogen/models/retrieval_scores_value.py +127 -0
- shaped/autogen/models/retrieve_inner.py +196 -0
- shaped/autogen/models/retrieve_step_explanation.py +172 -0
- shaped/autogen/models/retriever.py +196 -0
- shaped/autogen/models/retriever1.py +196 -0
- shaped/autogen/models/rollout_config.py +91 -0
- shaped/autogen/models/rudderstack_table_config.py +106 -0
- shaped/autogen/models/sampling_strategy.py +36 -0
- shaped/autogen/models/saved_query_info_response.py +103 -0
- shaped/autogen/models/saved_query_list_response.py +87 -0
- shaped/autogen/models/saved_query_request.py +115 -0
- shaped/autogen/models/schema_config.py +117 -0
- shaped/autogen/models/score.py +134 -0
- shaped/autogen/models/score_ensemble.py +127 -0
- shaped/autogen/models/score_ensemble_policy_config.py +165 -0
- shaped/autogen/models/score_ensemble_policy_config_policies_inner.py +755 -0
- shaped/autogen/models/score_step_explanation.py +224 -0
- shaped/autogen/models/search_config.py +105 -0
- shaped/autogen/models/segment_table_config.py +106 -0
- shaped/autogen/models/sequence_length.py +136 -0
- shaped/autogen/models/server_config.py +100 -0
- shaped/autogen/models/setup_engine_response.py +87 -0
- shaped/autogen/models/shaped_internal_recsys_policies_als_model_policy_als_model_policy_config.py +154 -0
- shaped/autogen/models/shaped_internal_recsys_policies_beeformer_model_policy_beeformer_model_policy_beeformer_model_policy_config.py +160 -0
- shaped/autogen/models/shaped_internal_recsys_policies_bert_model_policy_bert_model_policy_bert_model_policy_config.py +215 -0
- shaped/autogen/models/shaped_internal_recsys_policies_chronological_model_policy_chronological_model_policy_config.py +156 -0
- shaped/autogen/models/shaped_internal_recsys_policies_elsa_model_policy_elsa_model_policy_elsa_model_policy_config.py +141 -0
- shaped/autogen/models/shaped_internal_recsys_policies_gsasrec_model_policy_gsasrec_model_policy_gsas_rec_model_policy_config.py +211 -0
- shaped/autogen/models/shaped_internal_recsys_policies_item2vec_model_policy_item2_vec_model_policy_config.py +140 -0
- shaped/autogen/models/shaped_internal_recsys_policies_item_content_similarity_model_policy_item_content_similarity_model_policy_config.py +195 -0
- shaped/autogen/models/shaped_internal_recsys_policies_lightgbm_model_policy_lightgbm_model_policy_light_gbm_model_policy_config.py +245 -0
- shaped/autogen/models/shaped_internal_recsys_policies_ngram_model_policy_ngram_model_policy_config.py +119 -0
- shaped/autogen/models/shaped_internal_recsys_policies_popular_model_policy_popular_model_policy_config.py +161 -0
- shaped/autogen/models/shaped_internal_recsys_policies_random_model_policy_random_model_policy_config.py +108 -0
- shaped/autogen/models/shaped_internal_recsys_policies_recently_popular_policy_recently_popular_policy_config.py +143 -0
- shaped/autogen/models/shaped_internal_recsys_policies_rising_popular_policy_rising_popular_policy_config.py +123 -0
- shaped/autogen/models/shaped_internal_recsys_policies_sasrec_model_policy_sasrec_model_policy_sas_rec_model_policy_config.py +230 -0
- shaped/autogen/models/shaped_internal_recsys_policies_svd_model_policy_svd_model_policy_config.py +119 -0
- shaped/autogen/models/shaped_internal_recsys_policies_two_tower_model_policy_two_tower_model_policy_two_tower_model_policy_config.py +165 -0
- shaped/autogen/models/shaped_internal_recsys_policies_user_item_content_similarity_model_policy_user_item_content_similarity_model_policy_config.py +153 -0
- shaped/autogen/models/shaped_internal_recsys_policies_widedeep_model_policy_wide_deep_model_policy_config.py +137 -0
- shaped/autogen/models/shaped_internal_recsys_policies_xgboost_model_policy_xg_boost_model_policy_config.py +155 -0
- shaped/autogen/models/shopify_table_config.py +156 -0
- shaped/autogen/models/similarity_retrieve_step.py +123 -0
- shaped/autogen/models/{snowflake_dataset_config.py → snowflake_table_config.py} +47 -18
- shaped/autogen/models/sql_transform_type.py +37 -0
- shaped/autogen/models/sql_view_config.py +111 -0
- shaped/autogen/models/stemmer_tokenizer.py +105 -0
- shaped/autogen/models/step_explanation.py +137 -0
- shaped/autogen/models/steps_inner.py +179 -0
- shaped/autogen/models/strategy.py +134 -0
- shaped/autogen/models/table.py +102 -0
- shaped/autogen/models/table_deployment_type.py +38 -0
- shaped/autogen/models/table_insert_arguments.py +87 -0
- shaped/autogen/models/table_insert_response.py +87 -0
- shaped/autogen/models/text_encoding.py +148 -0
- shaped/autogen/models/text_search_retrieve_step.py +121 -0
- shaped/autogen/models/time_frequency.py +136 -0
- shaped/autogen/models/time_window.py +136 -0
- shaped/autogen/models/time_window_in_days.py +154 -0
- shaped/autogen/models/tokenizer.py +149 -0
- shaped/autogen/models/trained_model_encoder.py +99 -0
- shaped/autogen/models/training_compute_config.py +99 -0
- shaped/autogen/models/training_config.py +121 -0
- shaped/autogen/models/training_config_models_inner.py +540 -0
- shaped/autogen/models/training_strategy.py +37 -0
- shaped/autogen/models/transform_status.py +41 -0
- shaped/autogen/models/trending_mode.py +37 -0
- shaped/autogen/models/truncate_filter_step.py +106 -0
- shaped/autogen/models/tunable_bool.py +97 -0
- shaped/autogen/models/tunable_float.py +118 -0
- shaped/autogen/models/tunable_int.py +118 -0
- shaped/autogen/models/tunable_int_categorical.py +99 -0
- shaped/autogen/models/tunable_string.py +99 -0
- shaped/autogen/models/tuning_config.py +89 -0
- shaped/autogen/models/type.py +150 -0
- shaped/autogen/models/update_table_response.py +87 -0
- shaped/autogen/models/update_view_response.py +87 -0
- shaped/autogen/models/user_attribute_pooling_encoder.py +124 -0
- shaped/autogen/models/user_inner.py +134 -0
- shaped/autogen/models/val_split.py +136 -0
- shaped/autogen/models/validation_error.py +13 -3
- shaped/autogen/models/validation_error_loc_inner.py +150 -0
- shaped/autogen/models/value_type.py +7 -5
- shaped/autogen/models/vector_search_mode.py +101 -0
- shaped/autogen/models/view.py +104 -0
- shaped/autogen/models/view_details_ai.py +140 -0
- shaped/autogen/models/view_details_ai_schema_value.py +153 -0
- shaped/autogen/models/view_details_sql.py +140 -0
- shaped/autogen/models/view_status.py +41 -0
- shaped/autogen/models/weight_decay.py +136 -0
- shaped/autogen/models/whitespace_tokenizer.py +97 -0
- shaped/autogen/models/window_size.py +136 -0
- shaped/autogen/rest.py +10 -4
- shaped/cli/shaped_cli.py +163 -35
- shaped/client.py +591 -171
- shaped/config_builders.py +705 -0
- shaped/query_builder.py +781 -0
- {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/METADATA +141 -6
- shaped-2.0.4.dist-info/RECORD +296 -0
- {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/WHEEL +1 -1
- shaped-2.0.4.dist-info/entry_points.txt +2 -0
- shaped/autogen/api/model_inference_api.py +0 -2825
- shaped/autogen/models/amplitude_dataset_config.py +0 -96
- shaped/autogen/models/aws_pinpoint_dataset_config.py +0 -96
- shaped/autogen/models/big_query_dataset_config.py +0 -114
- shaped/autogen/models/complement_items_request.py +0 -99
- shaped/autogen/models/complement_items_response.py +0 -89
- shaped/autogen/models/connectors_inner.py +0 -134
- shaped/autogen/models/create_dataset_arguments.py +0 -263
- shaped/autogen/models/create_embedding_response.py +0 -87
- shaped/autogen/models/create_item_embedding_request.py +0 -89
- shaped/autogen/models/create_model_arguments.py +0 -107
- shaped/autogen/models/create_model_response.py +0 -87
- shaped/autogen/models/create_user_embedding_request.py +0 -89
- shaped/autogen/models/custom_dataset_config.py +0 -115
- shaped/autogen/models/dataset_config.py +0 -101
- shaped/autogen/models/dataset_schema_type.py +0 -47
- shaped/autogen/models/datasets_inner.py +0 -91
- shaped/autogen/models/delete_model_response.py +0 -87
- shaped/autogen/models/fetch_config.py +0 -95
- shaped/autogen/models/file_config.py +0 -105
- shaped/autogen/models/file_source_config.py +0 -89
- shaped/autogen/models/inference_config.py +0 -101
- shaped/autogen/models/insert_model_response.py +0 -87
- shaped/autogen/models/interaction.py +0 -87
- shaped/autogen/models/list_datasets_response.py +0 -95
- shaped/autogen/models/list_models_response.py +0 -95
- shaped/autogen/models/model_config.py +0 -99
- shaped/autogen/models/model_response.py +0 -95
- shaped/autogen/models/mongo_db_dataset_config.py +0 -119
- shaped/autogen/models/post_rank_request.py +0 -117
- shaped/autogen/models/rank_attribute_response.py +0 -89
- shaped/autogen/models/rank_grid_attribute_request.py +0 -91
- shaped/autogen/models/rank_grid_attribute_request1.py +0 -93
- shaped/autogen/models/rank_grid_attribute_response.py +0 -91
- shaped/autogen/models/rank_response.py +0 -91
- shaped/autogen/models/retrieve_request.py +0 -101
- shaped/autogen/models/retrieve_response.py +0 -91
- shaped/autogen/models/retriever_top_k_override.py +0 -97
- shaped/autogen/models/rudder_stack_dataset_config.py +0 -96
- shaped/autogen/models/segment_dataset_config.py +0 -96
- shaped/autogen/models/similar_item_request.py +0 -101
- shaped/autogen/models/similar_response.py +0 -89
- shaped/autogen/models/similar_users_request.py +0 -99
- shaped/autogen/models/successful_response.py +0 -87
- shaped/autogen/models/view_model_response.py +0 -99
- shaped-2.0.1.dist-info/RECORD +0 -73
- shaped-2.0.1.dist-info/entry_points.txt +0 -2
- {shaped-2.0.1.dist-info → shaped-2.0.4.dist-info}/top_level.txt +0 -0
- {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
|
-
|
|
2
|
+
import socket
|
|
3
|
+
from typing import Any, Dict, List, Optional, Union
|
|
3
4
|
|
|
4
|
-
from shaped.autogen.api.
|
|
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.
|
|
8
|
-
|
|
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.
|
|
11
|
-
from shaped.autogen.models.
|
|
12
|
-
from shaped.autogen.models.
|
|
13
|
-
from shaped.autogen.models.
|
|
14
|
-
from shaped.autogen.models.
|
|
15
|
-
from shaped.autogen.models.
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
from
|
|
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
|
-
|
|
24
|
-
|
|
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
|
-
|
|
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
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
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
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
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
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
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
|
-
|
|
113
|
-
|
|
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
|
-
|
|
301
|
+
table_insert_arguments=insert_args,
|
|
116
302
|
)
|
|
117
303
|
|
|
118
|
-
def
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
""
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
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
|
|
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
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
"""
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
.
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
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
|
|
538
|
+
def execute_saved_query(
|
|
182
539
|
self,
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
return_metadata:
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
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
|
-
|
|
202
|
-
|
|
203
|
-
|
|
204
|
-
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
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
|
)
|