diff --git a/answer_rocket/config.py b/answer_rocket/config.py index ed80be8..59acedd 100644 --- a/answer_rocket/config.py +++ b/answer_rocket/config.py @@ -9,7 +9,7 @@ from answer_rocket.graphql.client import GraphQlClient from answer_rocket.graphql.schema import UUID as GQL_UUID, MaxCopilotSkill, MaxCopilot, \ MaxMutationResponse, MaxCopilotQuestionInput, \ - MaxCreateCopilotQuestionResponse, MaxUser, MaxSkillComponent, MaxLLmPrompt, Boolean, HydratedReport + MaxCreateCopilotQuestionResponse, MaxUser, MaxLLmPrompt, Boolean, HydratedReport class Config: @@ -179,35 +179,6 @@ def get_copilot_skill(self, use_published_version: bool = True, copilot_id: str except Exception as e: return None - def get_skill_components(self) -> [MaxSkillComponent]: - """ - Retrieve all available skill components. - - Returns - ------- - list[MaxSkillComponent] | None - A list of skill components, or None if an error occurs. - """ - try: - query_args = { - } - - query_vars = { - } - - operation = self._gql_client.query(variables=query_vars) - - gql_query = operation.get_skill_components( - ) - - result = self._gql_client.submit(operation, query_args) - - skill_components = result.get_skill_components - - return skill_components - except Exception as e: - return None - def get_copilot_hydrated_reports(self, copilot_id: Optional[str] = None, override_dataset_id: Optional[str] = None, load_all_skills: bool = False) -> [HydratedReport]: """ Get hydrated reports for a copilot. diff --git a/answer_rocket/data.py b/answer_rocket/data.py index b19f840..c10043b 100644 --- a/answer_rocket/data.py +++ b/answer_rocket/data.py @@ -65,23 +65,7 @@ class ExecuteSqlQueryResult(MaxResult): """ df: DataFrame | None = None data = None # deprecated -- use df instead - -@dataclass -class ExecuteRqlQueryResult(MaxResult): - """ - Result object for RQL query execution operations. - - Attributes - ---------- - df : DataFrame | None - The result of the RQL query as a pandas DataFrame. - rql_script_response : Any | None - The RQL script response containing processing information. - """ - df = None - rql_script_response = None - -@dataclass + class DomainObjectResult(MaxResult): """ Result object for domain object retrieval operations. @@ -241,93 +225,6 @@ def execute_sql_query(self, database_id: UUID, sql_query: str, row_limit: Option return result - def execute_rql_query(self, dataset_id: UUID, rql_query: str, row_limit: Optional[int] = None, copilot_id: Optional[UUID] = None, copilot_skill_id: Optional[UUID] = None) -> ExecuteRqlQueryResult: - """ - Execute an RQL query against a dataset and return results. - - Parameters - ---------- - dataset_id : UUID - The UUID of the dataset to execute the query against. - rql_query : str - The RQL query string to execute. - row_limit : Optional[int], optional - Maximum number of rows to return in the query results. - copilot_id : Optional[UUID], optional - The UUID of the copilot. Defaults to the configured copilot_id. - copilot_skill_id : Optional[UUID], optional - The UUID of the copilot skill. Defaults to the configured copilot_skill_id. - - Returns - ------- - ExecuteRqlQueryResult - The result containing success status, error information, DataFrame, and RQL script response. - """ - try: - query_args = { - 'datasetId': dataset_id, - 'rqlQuery': rql_query, - 'rowLimit': row_limit, - 'copilotId': copilot_id or self.copilot_id, - 'copilotSkillId': copilot_skill_id or self.copilot_skill_id, - } - - query_vars = { - 'dataset_id': Arg(non_null(GQL_UUID)), - 'rql_query': Arg(non_null(String)), - 'row_limit': Arg(Int), - 'copilot_id': Arg(GQL_UUID), - 'copilot_skill_id': Arg(GQL_UUID), - } - - operation = self._gql_client.query(variables=query_vars) - - execute_rql_query = operation.execute_rql_query( - dataset_id=Variable('dataset_id'), - rql_query=Variable('rql_query'), - row_limit=Variable('row_limit'), - copilot_id=Variable('copilot_id'), - copilot_skill_id=Variable('copilot_skill_id'), - ) - - execute_rql_query.success() - execute_rql_query.code() - execute_rql_query.error() - execute_rql_query.data() - execute_rql_query.process_rql_script_response() - - result = self._gql_client.submit(operation, query_args) - - execute_rql_query_response = result.execute_rql_query - - execute_rql_query_result = ExecuteRqlQueryResult() - - execute_rql_query_result.success = execute_rql_query_response.success - execute_rql_query_result.error = execute_rql_query_response.error - execute_rql_query_result.code = execute_rql_query_response.code - - if execute_rql_query_response.success: - data = execute_rql_query_response.data - - columns = [column["name"] for column in data["columns"]] - rows = [row["data"] for row in data["rows"]] if "rows" in data else [] - - df = pd.DataFrame(rows, columns=columns) - - execute_rql_query_result.df = df - - execute_rql_query_result.rql_script_response = execute_rql_query_response.process_rql_script_response - - return execute_rql_query_result - except Exception as e: - execute_rql_query_result = ExecuteRqlQueryResult() - - execute_rql_query_result.success = False - execute_rql_query_result.error = e - execute_rql_query_result.code = RESULT_EXCEPTION_CODE - - return execute_rql_query_result - def get_database(self, database_id: UUID) -> Optional[Database]: """ Retrieve a database by its ID. diff --git a/answer_rocket/graphql/schema.py b/answer_rocket/graphql/schema.py index 977f435..601cd6f 100644 --- a/answer_rocket/graphql/schema.py +++ b/answer_rocket/graphql/schema.py @@ -1002,20 +1002,6 @@ class MaxReportResult(sgqlc.types.Type): final_message = sgqlc.types.Field(String, graphql_name='finalMessage') preview = sgqlc.types.Field(String, graphql_name='preview') - -class MaxSkillComponent(sgqlc.types.Type): - __schema__ = schema - __field_names__ = ('skill_component_id', 'node_type', 'organization', 'name', 'description', 'input_properties', 'output_properties', 'component_data') - skill_component_id = sgqlc.types.Field(sgqlc.types.non_null(UUID), graphql_name='skillComponentId') - node_type = sgqlc.types.Field(sgqlc.types.non_null(String), graphql_name='nodeType') - organization = sgqlc.types.Field(sgqlc.types.non_null(String), graphql_name='organization') - name = sgqlc.types.Field(sgqlc.types.non_null(String), graphql_name='name') - description = sgqlc.types.Field(String, graphql_name='description') - input_properties = sgqlc.types.Field(sgqlc.types.non_null(sgqlc.types.list_of(sgqlc.types.non_null('MaxSkillComponentInputProperty'))), graphql_name='inputProperties') - output_properties = sgqlc.types.Field(sgqlc.types.non_null(sgqlc.types.list_of(sgqlc.types.non_null('MaxSkillComponentOutputProperty'))), graphql_name='outputProperties') - component_data = sgqlc.types.Field(sgqlc.types.non_null(JSON), graphql_name='componentData') - - class MaxSkillComponentInputProperty(sgqlc.types.Type): __schema__ = schema __field_names__ = ('name', 'label', 'description', 'type', 'is_required', 'is_list', 'can_wire_from_output') @@ -1488,7 +1474,8 @@ class ParameterDefinition(sgqlc.types.Type): class Query(sgqlc.types.Type): __schema__ = schema - __field_names__ = ('ping', 'current_user', 'get_copilot_skill_artifact_by_path', 'get_copilots', 'get_copilot_info', 'get_copilot_skill', 'run_copilot_skill', 'get_skill_components', 'get_copilot_hydrated_reports', 'get_async_skill_run_status', 'get_max_agent_workflow', 'execute_sql_query', 'execute_rql_query', 'get_databases', 'get_database', 'get_database_tables', 'get_dataset_id', 'get_dataset', 'get_dataset2', 'get_datasets', 'get_domain_object', 'get_domain_object_by_name', 'get_grounded_value', 'get_database_kshots', 'get_database_kshot_by_id', 'get_dataset_kshots', 'get_dataset_kshot_by_id', 'run_max_sql_gen', 'run_sql_ai', 'generate_visualization', 'llmapi_config_for_sdk', 'get_max_llm_prompt', 'user_chat_threads', 'user_chat_entries', 'chat_thread', 'chat_entry', 'user', 'all_chat_entries', 'skill_memory', 'chat_completion', 'narrative_completion', 'narrative_completion_with_prompt', 'sql_completion', 'research_completion', 'chat_completion_with_prompt', 'research_completion_with_prompt', 'get_chat_artifact', 'get_chat_artifacts') + __field_names__ = ('ping', 'current_user', 'get_copilot_skill_artifact_by_path', 'get_copilots', 'get_copilot_info', 'get_copilot_skill', 'run_copilot_skill', 'get_copilot_hydrated_reports', 'get_async_skill_run_status', 'get_max_agent_workflow', 'execute_sql_query', 'execute_rql_query', 'get_databases', 'get_database', 'get_database_tables', 'get_dataset_id', 'get_dataset', 'get_dataset2', 'get_datasets', 'get_domain_object', 'get_domain_object_by_name', 'get_grounded_value', 'get_database_kshots', 'get_database_kshot_by_id', 'get_dataset_kshots', 'get_dataset_kshot_by_id', 'run_max_sql_gen', 'run_sql_ai', 'generate_visualization', 'llmapi_config_for_sdk', 'get_max_llm_prompt', 'user_chat_threads', 'user_chat_entries', 'chat_thread', 'chat_entry', 'user', 'all_chat_entries', 'skill_memory', 'chat_completion', 'narrative_completion', 'narrative_completion_with_prompt', 'sql_completion', 'research_completion', 'chat_completion_with_prompt', 'research_completion_with_prompt', 'get_chat_artifact', 'get_chat_artifacts') + ping = sgqlc.types.Field(String, graphql_name='ping') current_user = sgqlc.types.Field(MaxUser, graphql_name='currentUser') get_copilot_skill_artifact_by_path = sgqlc.types.Field(CopilotSkillArtifact, graphql_name='getCopilotSkillArtifactByPath', args=sgqlc.types.ArgDict(( @@ -1517,7 +1504,7 @@ class Query(sgqlc.types.Type): ('validate_parameters', sgqlc.types.Arg(Boolean, graphql_name='validateParameters', default=None)), )) ) - get_skill_components = sgqlc.types.Field(sgqlc.types.non_null(sgqlc.types.list_of(sgqlc.types.non_null(MaxSkillComponent))), graphql_name='getSkillComponents') + get_copilot_hydrated_reports = sgqlc.types.Field(sgqlc.types.non_null(sgqlc.types.list_of(sgqlc.types.non_null(HydratedReport))), graphql_name='getCopilotHydratedReports', args=sgqlc.types.ArgDict(( ('copilot_id', sgqlc.types.Arg(sgqlc.types.non_null(UUID), graphql_name='copilotId', default=None)), ('override_dataset_id', sgqlc.types.Arg(UUID, graphql_name='overrideDatasetId', default=None)),