Package: waii-sdk-py
Waii Python SDK DocumentationThewaii-sdk-pyis a Python library that allows you to interact with the Waii API. It provides a powerful SQL and AI capability to your Python applications.InstallationTo install thewaii-sdk-py, you can use pip:pipinstallwaii-sdk-pyImporting & Initialize the SDK>>>fromwaii_sdk_pyimportWAII>>>fromwaii_sdk_py.queryimport*>>>WAII.initialize(api_key="<your-api-key>")You can get your API key from thetweakit.waii.ai(You need to register and get access fromwaii.aifirst).Get startedFirst you can print the list of available databases:>>>print([conn.keyforconninWAII.Database.get_connections().connectors])Then, you can activate the database connection you want to use (from one of the key in the list above)>>>WAII.Database.activate_connection("snowflake://...&warehouse=COMPUTE_WH")Get Database name of the active connection>>>print([catalog.nameforcataloginWAII.Database.get_catalogs().catalogs])>>>print(WAII.Query.generate(QueryGenerationRequest(ask="How many tables are there?")).query)SELECTCOUNT(DISTINCTtable_name)FROMwaii.information_schema.tablesRun the query>>>print(WAII.Query.run(RunQueryRequest(query="SELECT COUNT(DISTINCT table_name) FROM waii.information_schema.tables")))rows=[{'COUNT(DISTINCT TABLE_NAME)':112}]more_rows=0column_definitions=[ColumnDefinition(name='COUNT(DISTINCT TABLE_NAME)',type='FIXED')]query_uuid='01afbd1e-0001-d31e-0022-ba8700a8209e'In order to know more details, you can check the following detailed API documentation.Use NotebookYou can also use the SDK in Notebooks like Jupyter, you can follow the the dochereModulesThe SDK consists of the following modules:DatabaseQuerySemanticContextHistoryEach module encapsulates a certain functionality of the Waii API.DatabaseTheDatabasemodule contains methods for handling database-related tasks.Here are some of its methods:Modify ConnectionsDatabase.modify_connections(params:ModifyDBConnectionRequest)->ModifyDBConnectionResponseThis method allows to modify the database connections.ModifyDBConnectionRequestcontains the following properties:updated: Update or add new connectionsremoved: Remove connections, by specifying the key of the connectiondefault_db_connection_key: Set the default connection (by specifying key of connection)To add connection, you need to createDBConnectionObject, which include the following fieldskey: Set it to '', it will be auto generateddb_type: Acceptable values are:snowflake,postgresql,mongodb,mongodb+srv(Mongo Atlas).description: Description of the connectionaccount_name: Account name of the connection, apply toSnowflakeusername: Username of the connection, can be None if no username needed (such as localhost database)password: Password of the connection, can be None if no password.database: Database name of the connection, you must specify it. Please note that it is case sensitive for most of the databases.warehouse: Warehouse name of the connection, apply toSnowflakerole: Role name of the connection, apply toSnowflakehost/port: Host/port of the connection, apply topostgresqlandmongodb.sample_col_values: Do you want to let Waii to sample your string/variant columns. True if you want to sample, False if you don't want to sample. Default is False. This is optional, which can help Waii to generate better queries.Examples of creatingDBConnectionObjectSnowflakeDBConnection(key='',db_type='snowflake',account_name='abcde-12345',username='username',password='password',database='database',warehouse='warehouse',role='role')PostgreSQLDBConnection(key='',db_type='postgresql',username='username',password='password',database='database',host='localhost',port=5432)MongoDBDBConnection(key='',db_type='mongodb',# or 'mongodb+srv' if you are using Mongo Atlasusername='username',password='password',database='database',host='localhost',port=27017# if you are using Mongo Atlas, you shouldn't set port)Get ConnectionsDatabase.get_connections(params:GetDBConnectionRequest=GetDBConnectionRequest())->GetDBConnectionResponseThis method fetches the list of available database connections.Response fields:connectors: List ofDBConnectionobjects (No password field)connector_status: Status of the connection, are they being indexed or not.Activate ConnectionDatabase.activate_connection(key:str)This method sets the scope of the current database connection.You can run>>>Database.get_activated_connection()To get the current activated connection.Get CatalogsDatabase.get_catalogs(params:GetCatalogRequest=GetCatalogRequest())->GetCatalogResponseThis method retrieves the list of available catalogs. It includes a hierarchical list of schemas and tables.You can run>>>WAII.Database.get_catalogs()To get catalogs, response fields:CatalogDefinition:name: Name of the catalog (database)schemas: List of schemasname: Name of the schemadescription: Description of the schemasummary: generated description of the schemacommon_questions: list of generated common questionscommon_tables: list of common tables used in the schematables: all tables belong to the schemaname: Name of the tablecolumns: List of columnsname: Name of the columntype: Type of the columndescription: Auto generated description of the columncomment: Comment of the table (fetched from underlying database)last_altered_time: Last altered time of the tablerefs: List of referential constraintsinferred_refs: List of auto-inferred referential constraintsinferred_constraints: List of auto-inferred constraints (pks, etc.)description: Auto generated table descriptionWaii automatically generate table/schema descriptions when you add them, you can fetch them by usingdescriptionfield of table and schema fromget_catalogsmethod.Description of table looks like:The CUSTOMER_ADDRESS table contains information about the addresses of customers. It includes details such as address ID, city, country, ... This table can be used to retrieve customer addresses for various purposes, such as shipping, billing, or demographic analysis.Update Table and Schema DescriptionsYou can use the following two methods to update the descriptions of tables and schemas (if you are not satisfied with the auto generated descriptions)Database.update_table_description(params:UpdateTableDescriptionRequest)->UpdateTableDescriptionResponseDatabase.update_schema_description(params:UpdateSchemaDescriptionRequest)->UpdateSchemaDescriptionResponseQueryTheQuerymodule contains methods related to SQL query handling.Important:You need to activate the database connection first before using the methods in this module. Otherwise you may trying to generate query from a wrong database.Here are some of its methods:Generate QueryQuery.generate(params:QueryGenerationRequest)->GeneratedQueryThis method generates a SQL query based on the provided parameters.Parameter fields:ask: The question you want to ask Waii to generate, such asHow many tables are there?dialect: What is the dialect of the generated query, such assnowflake,postgresql,mongodbtweak_history: We can support both asking new question, or tweak the previous question. If you want to tweak the previous question, you can set this field. (It's list ofTweakobject). Each ofTweakobject looks like:sql: The SQL query you want to tweakask: The previous question you asked.search_context: List of scopes to limit tables/schemas used to generate query (By default it searches all objects from the database). It's list ofSearchContextobject. Each ofSearchContextobject looks like:db_name: Catalog name (database name), default is*schema_name: Schema name, default is*table_name: Table name, default is*Examples:Ask a new question:>>>WAII.Query.generate(QueryGenerationRequest(ask="How many tables are there?"))Ask a complex question:>>>WAII.Query.generate(QueryGenerationRequest(ask="Give me all stores in California that have more than 1000 customers, include the store name, store address, and the number of customers."))Generate query with search context (limited to specific tables, schemas, etc.)By default, Waii search all tables in the database, if you know which tables you want to use, you can specify the search context to limit the tables used to generate query.>>>WAII.Query.generate(QueryGenerationRequest(ask="How many cars we produced in Q2",search_context=[SearchContext(db_name="*",schema_name="schema1",table_name="table1"),SearchContext(db_name="*",schema_name="schema2",table_name="*")]))The above query will only search tables fromschema1.table1andschema2.*Tweak the previous question:>>>WAII.Query.generate(QueryGenerationRequest(ask="only return 10 results",tweak_history=[Tweak(sql="select * from information_schema.tables",ask="give me all tables")]))(The above tweak will return the first 10 results of the previous query)Query.generatereturnsGeneratedQueryobject, which contains the following fields:uuid: UUID of the generated queryliked: Whether the query is liked or not (True/False). You can like a query by usingQuery.likemethod.query: The generated query texttables: Tables used in the generated querysemantic_context: Semantic context applied to the generated querydetailed_steps: Detailed steps of how the query work (in natural language)what_changed: If you do a tweak, it will tell you what changed in the querycompilation_error: If there is any compilation error, it will show here. (Waii will try to fix the compilation error automatically, but if it tried multiple times and still cannot fix it, it will show here)Generate QuestionYou can also generate questions based on your database schema, which can be useful when you want to show to your user what kind of questions can be asked to the database.>>>Query.generate_question(GenerateQuestionRequest(schema_name='PUBLIC',n_questions=10,complexity=GeneratedQuestionComplexity.hard))Parameter fields:schema_name: The schema name you want to generate questions. This is must be specified.n_questions: Number of questions you want to generate. Default is 10. You can choose a number between 1 to 30.complexity: Complexity of the questions you want to generate. Default isGeneratedQuestionComplexity.hard. You can chooseeasy,medium,hard.OutputGenerateQuestionResponse, fields:questions: List of generated questions. For each question it hasquestion(string content of the question), andcomplexity(complexity of the question)Transcode>>>WAII.Query.transcode(TranscodeQueryRequest(source_dialect="postgres",source_query="select ...;",target_dialect="snowflake",ask="..."))Parameter fields:search_context: Same asQuery.generate, you can specify the search context to limit the tables/schemas used to transcode.source_dialect: The dialect of the source query, such assnowflake,postgresql,mongodbsource_query: The source query you want to transcodetarget_dialect: The dialect of the target query, such assnowflake,postgresql,mongodbask: The intructions you want to ask Waii to add when transcode the query, such asplace tables under schema1. It is optional.Query.transcodereturnsGeneratedQueryobjectRun a query (sync)Query.run(params:RunQueryRequest)->GetQueryResultResponseRequest fields:query: The SQL query you want to runsession_id: Session ID of the query, default is None (which means Waii will generate a new session ID for you)current_schema: current schema to run the query, default is None (which means Waii will use the default schema of the connection). It isSchemaNameobject, which hasdatabase_name(optional),schema_namefields.session_parameter: a key value pairs to set the connection parameters, such as auto_commit, isolation_level, etc. Default is None (which means Waii will use the default connection parameters of the connection)Response fields:rows: List of rows returned. Each of row is json object, with column name as key, and column value as value.more_rows: Whether there are more rows or not (True/False), it can be True whenmax_returned_rowsis set and generated rows are more thanmax_returned_rowscolumn_definitions: List of column definitions. Each of column definition has name and typequery_uuid: UUID of the query (it is different from generate query UUID, this UUID is generated by underlying database and Snowflake)Example:>>>print(WAII.Query.run(RunQueryRequest(query="SELECT COUNT(DISTINCT table_name) FROM waii.information_schema.tables")))rows=[{'COUNT(DISTINCT TABLE_NAME)':112}]more_rows=0column_definitions=[ColumnDefinition(name='COUNT(DISTINCT TABLE_NAME)',type='FIXED',comment=None,sample_values=None)]query_uuid='01afbd1e-0001-d31e-0022-ba8700a8209e'Run with default schemafromwaii_sdk_py.databaseimport*>>>WAII.Query.run(RunQueryRequest(query='select current_schema(), current_database();',current_schema=SchemaName(schema_name='INFORMATION_SCHEMA')))Async submit a queryIn order to async submit query, you need 3 methods:Query.submit,Query.get_results,Query.cancelSubmitQuery.submit(params:RunQueryRequest)->RunQueryResponseRequest fields same asQuery.runResponse field:query_id: the query id you need to get result of the query, or cancel the queryGet ResultsQuery.get_results(params:GetQueryResultRequest)->GetQueryResultResponseReturn query result by providing uuid, same asQuery.run, but you need to specifyquery_idinstead ofqueryin the request.CancelQuery.cancel(params:CancelQueryRequest)->CancelQueryResponseCancel a running query by providingquery_id, it is a no-op if the query is completed or failed.LikeQuery.like(params:LikeQueryRequest)->LikeQueryResponseThis method marks a query (by specifyinguuidfor generated query, not id from run query) as "liked", which you can fetch it from history.DescribeIf you want to translate SQL to natural language, and explain step-by-step plans, you can useQuery.describemethod.Query.describe(params:DescribeQueryRequest)->DescribeQueryResponseThis method fetches a description of a SQL query.Parameters:query: The SQL query you want to describesearch_context: When describe the query, you may want to restrict which table/schema you want to use as context. Normally it is not required, and you should skip it.Output:summary: A summary of the querydetailed_steps: Detailed steps of how the query work (in natural language)tables: Tables used in the queryExample:>>>print(WAII.Query.describe(DescribeQueryRequest(query='SELECT COUNT(DISTINCT table_name) FROM information_schema.tables')))summary='How many distinct table names are there in the information_schema.tables?'detailed_steps=['Step 1: Access the information_schema.tables table.','Step 2: Count the number of distinct table names.','Step 3: Return the count of distinct table names.']tables=[TableName(table_name='TABLES',schema_name='INFORMATION_SCHEMA',database_name='WAII')]DiffThis method generates a description of the difference between two SQL queriesQuery.diff(params:DiffQueryRequest)->DiffQueryResponseParameters:Similar toQuery.describe, you need to specify two queries you want to diff, and optionally you can specify the search context.query: The SQL query you want to describeprevious_query: The previous SQL query you want to describeOutput:Same asQuery.describe, but it will also includewhat_changedfield, which tells you what changed in the query.Example:>>>print(WAII.Query.diff(DiffQueryRequest(query='SELECT COUNT(DISTINCT table_name) FROM information_schema.tables',previous_query='SELECT COUNT(DISTINCT table_name) FROM information_schema.tables WHERE table_name LIKE\'%a%\'')))summary='...'detailed_steps=[...]what_changed="The new query does not have any filter on 'table_name' column."Auto Complete (Experimental)This method allows you to automatically complete a partial query, this can be useful when you want to build a query editor with Github co-pilot-like auto complete feature.This is an experimental feature, it may not work well in some cases.Query.auto_complete(params:AutoCompleteRequest)->AutoCompleteResponseParameters:text: (required) The partial query you want to auto completecursor_offset: (required) The cursor offset of the partial query, it can be in the middle of the text (if you want to insert text in the middle. By default you should set it todialect: The dialect of the query, such assnowflake,postgresql,mongodbsearch_context: Same asQuery.describe, you can specify the search context to limit the tables/schemas used to auto complete.max_output_tokens: (required) Maximum number of tokens returned, you should set it to a lower value like 50 to make sure it can return within a few seconds.Output:textcompleted textExample:>>>print(WAII.Query.auto_complete(AutoCompleteRequest(text='select from',cursor_offset=11,max_output_tokens=50)))text='WAII.WORLD.CITY WHERE POPULATION > 1000000'Analyze Performance (Experimental)Note: this feature currently only support SnowflakeQuery.analyze_performance(params:QueryPerformanceRequest)->QueryPerformanceResponseParameters:query_id: The query id you want to analyze performance (it should be a run query id, not generated query idOutput:summary: A summary of the runtime of a query as well as recommendations on how to make the query run fasterrecommendations: List of recommendations on how to make the query run fasterquery_text: The query text of the query (based on query_id)execution_time_ms: The execution time of the query (based on query_id)compilation_time_ms: The compilation time of the query (based on query_id)This method provides a summary of the runtime of a query as well as recommendations on how to make the query run fasterExamples:>>>print(WAII.Query.analyze_performance(QueryPerformanceRequest(query_id='...')))'summary':['The metadata operation in Step 1 took 1.0 ms, which is 82% of the total execution time.',],'recommendations':["The query is already optimal as it is a simple count distinct operation on the 'table_name' column of the 'waii.information_schema.tables' table. No further optimization is required.",...],'query_text':"SELECT COUNT(DISTINCT table_name) FROM waii.information_schema.tables",'execution_time_ms':312,'compilation_time_ms':247Semantic ContextTheSemanticContextmodule contains methods related to semantic context handling.Here are its methods:Modify Semantic ContextSemanticContext.modify_semantic_context(params:ModifySemanticContextRequest)->ModifySemanticContextResponseexample:SemanticContext.modify_semantic_context(ModifySemanticContextRequest(updated=[# ... list of SemanticStatement object],deleted=[# ... list of statement id]))This method allows you to modify the semantic context, which include two fields,updatedanddeleted.Forupdated, you should include a list ofSemanticStatementobject, which include the following fields:id: (str) if you don't specify it, it will be auto generated. This is unique identifier of the semantic statement. You can use it to delete the statement later. We suggest you always include an id for the statement, so that you can find it easier. (Examples are JIRA_ID, PR_ID, CVE_ID, etc.)statement: (str), the statement you want to add. Such asrevenue is calculated by multiplying price and quantityscope: (str), the scope of the statement, such as*(which means it applies to all queries), ordb.schema.table.column(which means it applies to specific column of specific table of specific schema of specific database). This is optional, if not specified, it will be*. Whenalways_includeis False,scopewill be ignored.always_include: should we always include this statement during query generation? True/False. Default is True. This is optional. When it is False, we will use embedding/LLM to do a RAG process and match statements during query generation.lookup_summaries: Only take effect whenalways_include=False, you can specify a list of search keys for this statement, if not specified, then use statement as search key. This is optional.summarization_prompt: extract prompt from the statement, if not specified, then use statement as extract prompt. This is optional.There're several usage patterns of theSemanticStatementobject:Add description to a database object (such as table, column, etc.)Normally you want to do this to add a "description" to a table or column, so that Waii can use it to generate better query. Unless you want to apply it to the entire database, we suggest to always add a scope to the statement, such asdb.schema.table.column, so that it only applies to specific column of specific table of specific schema of specific database.example:statement=SemanticStatement(id='',statement='The CUSTOMER_ADDRESS table contains information about the addresses of customers. It includes details such as address ID, city, country, ... This table can be used to retrieve customer addresses for various purposes, such as shipping, billing, or demographic analysis.',scope='my_db.finance.customer_address',always_include=True)Add unstructured data source (such as json, or plain text document) as external knowledge baseThere're some cases you want to add unstructured data source as external knowledge base, so that Waii can use it to generate better query. You can add a statement withalways_include=False, and specifylookup_summariesto specify the "search summaries" for this statement.For example, assume you have a database of JIRA ticket, if you want to a knowledge base of all the JIRA tickets, you can add statements like this:statement=SemanticStatement(id='',statement="""{"ticket_id": "JIRA-123","summary": "User run into authentication issues when using the co-pilot feature","description": "...","assignee": "..."... (other fields)}""",always_include=False,lookup_summaries=['JIRA-123','User run into authentication issues when using the auto-pilot feature'],summarization_prompt='Output short summary (in 20 words) of the ticket ```{statement}```:')When user ask questionGive me number of tickets related to co-pilot authentication issues, Waii will use thelookup_summariesto match the statement. And bring the statement to the generated query.You can specify multiplelookup_summariesto make sure it can match the statement in different cases. If you don't specify one, we will use the statement as search key. Our recommendation is, if you think a short summary can help to match the statement, you should specify it aslookup_summaries.Summarization_promptis used to extract content from the statement, if not specified, then use the entire statement as-is during query generation. If your statement is too long, a good practice is to specify a summarization prompt, so that Waii can use it to extract content from the statement. Alternatively, you can pre-summarize the original documentation and add it asstatementfield. But that also loses the flexibility of using the original documentation.There're two placeholders forsummarization_prompt,{statement}and{ask}: You must include{statement}in the summarization prompt, which will be replaced by the statement. You can also include{ask}in the summarization prompt, which will be replaced by the ask in the query generation request.Get Semantic ContextSemanticContext.get_semantic_context(params:GetSemanticContextRequest=GetSemanticContextRequest())->GetSemanticContextResponseThis method fetches the current semantic context.Example:fromwaii_sdk_py.semantic_contextimportGetSemanticContextRequest,GetSemanticContextRequestFilters=WAII.SemanticContext.get_semantic_context(GetSemanticContextRequest(filter=GetSemanticContextRequestFilter(always_include=False),search_text='log4j CVE',limit=5))The above example searches the semantic context with the search textlog4j CVE, and limit the result to 5. If you don't specify search_text, it will apply to all statements. If you don't specify limit, it will return first 1000 statements.HistoryTheHistorymodule contains methods related to history handling.Here are some of its methods:ListHistory.list(params:GetGeneratedQueryHistoryRequest=GetGeneratedQueryHistoryRequest())->GetGeneratedQueryHistoryResponseThis method fetches the history of generated queries.For detailed information about the request and response objects and the available properties, refer to the respective Python files of each module.