├── doc ├── static │ ├── .nojekyll │ └── img │ │ ├── img.png │ │ ├── logo.png │ │ └── favicon.ico ├── .dockerignore ├── docs │ ├── #notebook-example.md# │ ├── img.png │ ├── plot_map.png │ ├── df_output.png │ ├── plot_result.png │ ├── colab-api-key.png │ ├── colab-df-plot.png │ ├── colab-secret.png │ ├── df_output_cars.png │ ├── metabase_chart.png │ ├── superset_example.png │ ├── superset-db-access.png │ ├── superset-user-access.png │ ├── waii-chart-selection.png │ ├── superset-add-connection.png │ ├── superset-db-connection.png │ ├── superset-waii-chat-mode.png │ ├── superset-waii-sql-mode.png │ ├── intro.md │ ├── install.md │ ├── async-client.md │ ├── getting-started.md │ ├── colab.md │ ├── chart-module.md │ ├── superset-waii-usage.md │ ├── history-module.md │ ├── waii-user-roles.md │ ├── multi-tenant-client-module.md │ ├── notebook.md │ ├── waii-similarity-search.md │ ├── superset-rendering.md │ └── access-rules-module.md ├── babel.config.js ├── src │ ├── pages │ │ ├── markdown-page.md │ │ ├── index.js │ │ └── index.module.css │ └── css │ │ └── custom.css ├── .gitignore ├── docker-compose.yml ├── package.json ├── docusaurus.config.js └── sidebars.js ├── pyproject.toml ├── DEVELOPMENT.md ├── setup.cfg ├── tests ├── test.py ├── test2.py ├── constants.py ├── __init__.py ├── semantic_context_test.py ├── history_test.py ├── query_test.py ├── waii_basemodel_test.py ├── chat_test.py ├── knowledge_graph_test.py ├── common_test_utils.py ├── chat_v2_test.py ├── push_db_conn.py ├── semantic_layer_dump_test.py ├── multi_tenant_client_test.py ├── async_client_test.py └── README.md ├── waii_sdk_py ├── chat │ ├── __init__.py │ └── chat.py ├── user │ ├── __init__.py │ └── user_static.py ├── chart │ ├── __init__.py │ └── chart.py ├── query │ └── __init__.py ├── utils │ ├── __init__.py │ └── utils.py ├── database │ └── __init__.py ├── history │ ├── __init__.py │ └── history.py ├── sdm │ ├── __init__.py │ └── sdm.py ├── settings │ ├── __init__.py │ └── settings.py ├── access_rules │ ├── __init__.py │ └── access_rules.py ├── semantic_context │ ├── __init__.py │ └── semantic_context.py ├── semantic_layer_dump │ ├── __init__.py │ └── semantic_layer_dump.py ├── waii_http_client │ ├── __init__.py │ └── waii_http_client.py ├── __init__.py ├── kg │ ├── __init__.py │ └── kg.py ├── common │ └── __init__.py ├── my_pydantic │ └── __init__.py └── waii_sdk_py.py ├── .github └── workflows │ └── python-publish.yml ├── add_apache_headers.py ├── README.md ├── samples └── hello_waii.py └── .gitignore /doc/static/.nojekyll: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /doc/.dockerignore: -------------------------------------------------------------------------------- 1 | */node_modules 2 | *.log 3 | -------------------------------------------------------------------------------- /doc/docs/#notebook-example.md#: -------------------------------------------------------------------------------- 1 | Full Notebook with Waii Example 2 | 3 | ```python 4 | ``` -------------------------------------------------------------------------------- /doc/docs/img.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/img.png -------------------------------------------------------------------------------- /doc/docs/plot_map.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/plot_map.png -------------------------------------------------------------------------------- /doc/docs/df_output.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/df_output.png -------------------------------------------------------------------------------- /doc/docs/plot_result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/plot_result.png -------------------------------------------------------------------------------- /doc/static/img/img.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/static/img/img.png -------------------------------------------------------------------------------- /doc/static/img/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/static/img/logo.png -------------------------------------------------------------------------------- /doc/docs/colab-api-key.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/colab-api-key.png -------------------------------------------------------------------------------- /doc/docs/colab-df-plot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/colab-df-plot.png -------------------------------------------------------------------------------- /doc/docs/colab-secret.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/colab-secret.png -------------------------------------------------------------------------------- /doc/static/img/favicon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/static/img/favicon.ico -------------------------------------------------------------------------------- /doc/docs/df_output_cars.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/df_output_cars.png -------------------------------------------------------------------------------- /doc/docs/metabase_chart.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/metabase_chart.png -------------------------------------------------------------------------------- /doc/docs/superset_example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset_example.png -------------------------------------------------------------------------------- /doc/docs/superset-db-access.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset-db-access.png -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | -------------------------------------------------------------------------------- /doc/babel.config.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | presets: [require.resolve('@docusaurus/core/lib/babel/preset')], 3 | }; 4 | -------------------------------------------------------------------------------- /doc/docs/superset-user-access.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset-user-access.png -------------------------------------------------------------------------------- /doc/docs/waii-chart-selection.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/waii-chart-selection.png -------------------------------------------------------------------------------- /doc/docs/superset-add-connection.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset-add-connection.png -------------------------------------------------------------------------------- /doc/docs/superset-db-connection.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset-db-connection.png -------------------------------------------------------------------------------- /doc/docs/superset-waii-chat-mode.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset-waii-chat-mode.png -------------------------------------------------------------------------------- /doc/docs/superset-waii-sql-mode.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waii-ai/waii-sdk-py/HEAD/doc/docs/superset-waii-sql-mode.png -------------------------------------------------------------------------------- /doc/src/pages/markdown-page.md: -------------------------------------------------------------------------------- 1 | --- 2 | title: Markdown page example 3 | --- 4 | 5 | # Markdown page example 6 | 7 | You don't need React to write simple standalone pages. 8 | -------------------------------------------------------------------------------- /doc/src/pages/index.js: -------------------------------------------------------------------------------- 1 | import React from 'react'; 2 | import { Redirect } from '@docusaurus/router'; 3 | 4 | function Home() { 5 | return ; 6 | } 7 | 8 | export default Home; 9 | -------------------------------------------------------------------------------- /doc/.gitignore: -------------------------------------------------------------------------------- 1 | # Dependencies 2 | /node_modules 3 | 4 | # Production 5 | /build 6 | 7 | # Generated files 8 | .docusaurus 9 | .cache-loader 10 | 11 | # Misc 12 | .DS_Store 13 | .env.local 14 | .env.development.local 15 | .env.test.local 16 | .env.production.local 17 | 18 | npm-debug.log* 19 | yarn-debug.log* 20 | yarn-error.log* 21 | -------------------------------------------------------------------------------- /DEVELOPMENT.md: -------------------------------------------------------------------------------- 1 | # Publish the package to PyPI 2 | 3 | ## Install required packages for publish 4 | 5 | ``` 6 | pip install build twine 7 | ``` 8 | 9 | 10 | ## 1. Update the version number in setup.cfg 11 | 12 | ``` 13 | version = 1.x.x 14 | ``` 15 | 16 | ## 2. Build the package 17 | 18 | ``` 19 | rm -rf dist/ && python3 -m build . 20 | ``` 21 | 22 | ## 3. Upload the package to PyPI 23 | 24 | ``` 25 | python -m twine upload --skip-existing dist/* 26 | ``` -------------------------------------------------------------------------------- /doc/src/pages/index.module.css: -------------------------------------------------------------------------------- 1 | /** 2 | * CSS files with the .module.css suffix will be treated as CSS modules 3 | * and scoped locally. 4 | */ 5 | 6 | .heroBanner { 7 | padding: 4rem 0; 8 | text-align: center; 9 | position: relative; 10 | overflow: hidden; 11 | } 12 | 13 | @media screen and (max-width: 966px) { 14 | .heroBanner { 15 | padding: 2rem; 16 | } 17 | } 18 | 19 | .buttons { 20 | display: flex; 21 | align-items: center; 22 | justify-content: center; 23 | } 24 | -------------------------------------------------------------------------------- /doc/docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: "3" 2 | 3 | services: 4 | docusaurus: 5 | build: . 6 | ports: 7 | - 3000:3000 8 | - 35729:35729 9 | volumes: 10 | - ./docs:/app/docs 11 | - ./website/blog:/app/website/blog 12 | - ./website/core:/app/website/core 13 | - ./website/i18n:/app/website/i18n 14 | - ./website/pages:/app/website/pages 15 | - ./website/static:/app/website/static 16 | - ./website/sidebars.json:/app/website/sidebars.json 17 | - ./website/siteConfig.js:/app/website/siteConfig.js 18 | working_dir: /app/website 19 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = waii-sdk-py 3 | version = 1.33.0 4 | description = Provides access to the Waii APIs 5 | long_description = file: README.md 6 | long_description_content_type = text/markdown 7 | author = Waii, Inc. 8 | license = Apache License 2.0 9 | url = http://www.waii.ai 10 | classifiers = 11 | Development Status :: 5 - Production/Stable 12 | Intended Audience :: Developers 13 | License :: OSI Approved :: Apache Software License 14 | Programming Language :: Python :: 3 15 | 16 | [options] 17 | packages = find: 18 | install_requires = 19 | requests==2.32.5 20 | python_requires = >=3.9 21 | -------------------------------------------------------------------------------- /tests/test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | -------------------------------------------------------------------------------- /tests/test2.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | -------------------------------------------------------------------------------- /tests/constants.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | pg_port = 5432 -------------------------------------------------------------------------------- /waii_sdk_py/chat/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .chat import * -------------------------------------------------------------------------------- /waii_sdk_py/user/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .user import * -------------------------------------------------------------------------------- /waii_sdk_py/chart/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .chart import * -------------------------------------------------------------------------------- /waii_sdk_py/query/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .query import * -------------------------------------------------------------------------------- /waii_sdk_py/utils/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .utils import * -------------------------------------------------------------------------------- /waii_sdk_py/database/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .database import * -------------------------------------------------------------------------------- /waii_sdk_py/history/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .history import * -------------------------------------------------------------------------------- /waii_sdk_py/sdm/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .sdm import * 18 | 19 | -------------------------------------------------------------------------------- /waii_sdk_py/settings/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .settings import * 18 | -------------------------------------------------------------------------------- /waii_sdk_py/access_rules/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .access_rules import * -------------------------------------------------------------------------------- /waii_sdk_py/semantic_context/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .semantic_context import * -------------------------------------------------------------------------------- /waii_sdk_py/semantic_layer_dump/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .semantic_layer_dump import * -------------------------------------------------------------------------------- /waii_sdk_py/waii_http_client/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .waii_http_client import WaiiHttpClient -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | # from .DatabaseTests import * 18 | # from .SemanticContextTests import * 19 | # from .QueryTests import * 20 | # from .HistoryTests import * 21 | -------------------------------------------------------------------------------- /waii_sdk_py/user/user_static.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from waii_sdk_py.user import UserImpl 18 | from waii_sdk_py.waii_http_client import WaiiHttpClient 19 | 20 | User = UserImpl(WaiiHttpClient.get_instance()) 21 | -------------------------------------------------------------------------------- /doc/docs/intro.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: intro 3 | title: Waii Python SDK 4 | --- 5 | 6 | Welcome to the SDK documentation for Waii - the worlds most powerful SQL API built on LLM agents. This documentation provides detailed information on how to use the SDK to interact with the system. 7 | 8 | The SDK allows enables developers to generate, explain, describe, modify, transcode, optimize and run SQL queries from text input, as well as manage the semantic layer, handle database connections, perform semantic search, and access the history of generated queries. You can also generate visualizations from the SQL you are working with. 9 | 10 | Whether you are looking to build a conversational client, analtyics tooling or database infrastructure, you have come to the right place. Waii makes it much easier to integrate advanced AI capabilities into your data-centric applications. Happy coding! 11 | -------------------------------------------------------------------------------- /waii_sdk_py/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .waii_sdk_py import WAII 18 | 19 | # Add these two imports to avoid you have to from waii_sdk_py.waii_sdk_py import AsyncWaii 20 | # (which is ugly). 21 | from .waii_sdk_py import AsyncWaii 22 | from .waii_sdk_py import Waii -------------------------------------------------------------------------------- /waii_sdk_py/kg/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .kg import ( 18 | GetKnowledgeGraphRequest, 19 | GetKnowledgeGraphResponse, 20 | KnowledgeGraph, 21 | KnowledgeGraphNode, 22 | KnowledgeGraphEdge, 23 | KnowlegeGraphNodeType, 24 | KnowledgeGraphEdgeType, 25 | KnowledgeGraphImpl, 26 | AsyncKnowledgeGraphImpl, 27 | KnowledgeGraphClient 28 | ) -------------------------------------------------------------------------------- /doc/docs/install.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: install 3 | title: Installation 4 | --- 5 | 6 | To install the `waii-sdk-py`, you can use pip: 7 | 8 | ```bash 9 | pip install waii-sdk-py 10 | ``` 11 | 12 | ## Importing & Initialize the SDK 13 | 14 | ```python 15 | >>> from waii_sdk_py import WAII 16 | >>> from waii_sdk_py.query import * 17 | >>> WAII.initialize(url='...', api_key="") 18 | ``` 19 | You can get your API key from the Waii server UI you have access to (You need to register and get access from [waii.ai](https://waii.ai) first). 20 | 21 | URL is the base URL of the WAII API, depends on the URL for the environment you are using: 22 | - If you are using prod environment (sql.waii.ai, or tweakit.waii.ai), you can use 'https://sql.waii.ai/api/' 23 | - If you are using test environment (sql.test.waii.ai, or tweakit-test.waii.ai), you can use 'https://sql.test.waii.ai/api/' 24 | - If you are using local Docker environment, you can use `http://host:port/api/` (host/port are point to Waii server) 25 | 26 | ![img.png](img.png) 27 | -------------------------------------------------------------------------------- /doc/src/css/custom.css: -------------------------------------------------------------------------------- 1 | /** 2 | * Any CSS included here will be global. The classic template 3 | * bundles Infima by default. Infima is a CSS framework designed to 4 | * work well for content-centric websites. 5 | */ 6 | 7 | /* You can override the default Infima variables here. */ 8 | :root { 9 | --ifm-color-primary: #25c2a0; 10 | --ifm-color-primary-dark: rgb(33, 175, 144); 11 | --ifm-color-primary-darker: rgb(31, 165, 136); 12 | --ifm-color-primary-darkest: rgb(26, 136, 112); 13 | --ifm-color-primary-light: rgb(70, 203, 174); 14 | --ifm-color-primary-lighter: rgb(102, 212, 189); 15 | --ifm-color-primary-lightest: rgb(146, 224, 208); 16 | --ifm-code-font-size: 95%; 17 | } 18 | 19 | .docusaurus-highlight-code-line { 20 | background-color: rgba(0, 0, 0, 0.1); 21 | display: block; 22 | margin: 0 calc(-1 * var(--ifm-pre-padding)); 23 | padding: 0 var(--ifm-pre-padding); 24 | } 25 | 26 | html[data-theme='dark'] .docusaurus-highlight-code-line { 27 | background-color: rgba(0, 0, 0, 0.3); 28 | } 29 | 30 | .footer__links { 31 | display: flex; 32 | justify-content: center; 33 | text-align: center; 34 | } 35 | 36 | .footer__links-item { 37 | flex: 1; 38 | } 39 | -------------------------------------------------------------------------------- /.github/workflows/python-publish.yml: -------------------------------------------------------------------------------- 1 | # This workflow will upload a Python Package using Twine when a release is created 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries 3 | 4 | # This workflow uses actions that are not certified by GitHub. 5 | # They are provided by a third-party and are governed by 6 | # separate terms of service, privacy policy, and support 7 | # documentation. 8 | 9 | name: Upload Python Package 10 | 11 | on: 12 | release: 13 | types: [published] 14 | 15 | permissions: 16 | contents: read 17 | 18 | jobs: 19 | deploy: 20 | 21 | runs-on: ubuntu-latest 22 | 23 | steps: 24 | - uses: actions/checkout@v3 25 | - name: Set up Python 26 | uses: actions/setup-python@v3 27 | with: 28 | python-version: '3.x' 29 | - name: Install dependencies 30 | run: | 31 | python -m pip install --upgrade pip 32 | pip install build 33 | - name: Build package 34 | run: python -m build 35 | - name: Publish package 36 | uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 37 | with: 38 | user: __token__ 39 | password: ${{ secrets.PYPI_API_TOKEN }} 40 | -------------------------------------------------------------------------------- /doc/docs/async-client.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: async-client 3 | title: Async Client 4 | --- 5 | 6 | ### Async Client 7 | 8 | This guide explains how to set up and use the `AsyncWaii` client to make asynchronous API calls. 9 | 10 | **Initialization & Imports** 11 | ```python 12 | # in order to use this, you need to make sure waii-sdk-py is >= 1.28.2 13 | # otherwise, you should do 14 | # from waii_sdk_py.waii_sdk_py import AsyncWaii 15 | from waii_sdk_py import AsyncWaii 16 | from waii_sdk_py.database import * 17 | 18 | client = AsyncWaii() 19 | await client.initialize(url='...', api_key="") 20 | ``` 21 | 22 | #### Using the Async Client 23 | 24 | The AsyncWaii client follows the same interface as the synchronous client, with some examples provided below. 25 | 26 | ###### Activating a Database Connection 27 | 28 | Use the activate_connection method to activate a database connection asynchronously: 29 | ```python 30 | # put it under the async function 31 | async def your_async_function(): 32 | conn_resp = await client.database.activate_connection("") 33 | ``` 34 | 35 | ###### Generating a Query 36 | 37 | To generate a query with the async client, pass a QueryGenerationRequest to the generate method: 38 | ```python 39 | # put it under the async function 40 | async def your_async_function(): 41 | ask = QueryGenerationRequest(ask="sample ask") 42 | query_gen_resp = await client.query.generate(ask) 43 | ``` 44 | 45 | Refer to the respective module for additional use cases, as both clients use the same interface. 46 | 47 | 48 | -------------------------------------------------------------------------------- /waii_sdk_py/utils/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | import functools 19 | import inspect 20 | from typing import TypeVar, Callable, Awaitable, Any 21 | 22 | from typing_extensions import ParamSpec 23 | 24 | T = TypeVar('T') 25 | P = ParamSpec('P') 26 | def to_async(func: Callable[P, T]) -> Callable[P, Awaitable[T]]: 27 | """Decorator to convert a sync method to async""" 28 | @functools.wraps(func) 29 | async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: 30 | return await asyncio.get_event_loop().run_in_executor( 31 | None, functools.partial(func, *args, **kwargs) 32 | ) 33 | return wrapper 34 | 35 | 36 | def wrap_methods_with_async(source_class, target_class): 37 | for name, method in inspect.getmembers(source_class, predicate=callable): 38 | if not name.startswith('_') and inspect.isroutine(method): 39 | async_method = to_async(getattr(source_class, name)) 40 | setattr(target_class, name, async_method) 41 | -------------------------------------------------------------------------------- /doc/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "docs-waii-sdk-js", 3 | "version": "0.0.0", 4 | "private": true, 5 | "scripts": { 6 | "docusaurus": "docusaurus", 7 | "start": "docusaurus start", 8 | "build": "docusaurus build", 9 | "swizzle": "docusaurus swizzle", 10 | "deploy": "docusaurus deploy", 11 | "clear": "docusaurus clear", 12 | "serve": "docusaurus serve", 13 | "write-translations": "docusaurus write-translations", 14 | "write-heading-ids": "docusaurus write-heading-ids" 15 | }, 16 | "dependencies": { 17 | "@docusaurus/core": "3.8.1", 18 | "@docusaurus/preset-classic": "3.8.1", 19 | "@mdx-js/react": "^3.0.0", 20 | "clsx": "^2.1.1", 21 | "prism-react-renderer": "^2.4.0", 22 | "react": "^18.0.0", 23 | "react-dom": "^18.0.0" 24 | }, 25 | "devDependencies": { 26 | "@docusaurus/module-type-aliases": "3.8.1", 27 | "@docusaurus/tsconfig": "3.8.1", 28 | "@docusaurus/types": "3.8.1", 29 | "postcss-selector-parser": "^7.1.0", 30 | "search-insights": "^2.17.3", 31 | "webpack-dev-server": "^5.2.2" 32 | }, 33 | "browserslist": { 34 | "production": [ 35 | ">0.5%", 36 | "not dead", 37 | "not op_mini all" 38 | ], 39 | "development": [ 40 | "last 1 chrome version", 41 | "last 1 firefox version", 42 | "last 1 safari version" 43 | ] 44 | }, 45 | "engines": { 46 | "node": ">=18.0" 47 | }, 48 | "resolutions": { 49 | "webpack-dev-server": ">=5.2.1" 50 | }, 51 | "overrides": { 52 | "shell-quote": "1.8.3", 53 | "loader-utils": "2.0.4", 54 | "immer": "9.0.6" 55 | } 56 | } 57 | -------------------------------------------------------------------------------- /add_apache_headers.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | HEADER = '''""" 4 | Copyright 2023–2025 Waii, Inc. 5 | 6 | Licensed under the Apache License, Version 2.0 (the "License"); 7 | you may not use this file except in compliance with the License. 8 | You may obtain a copy of the License at 9 | 10 | http://www.apache.org/licenses/LICENSE-2.0 11 | 12 | Unless required by applicable law or agreed to in writing, software 13 | distributed under the License is distributed on an "AS IS" BASIS, 14 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | See the License for the specific language governing permissions and 16 | limitations under the License. 17 | """ 18 | ''' 19 | 20 | EXCLUDED_DIRS = {"__pycache__", "build", "dist", ".git", "venv", ".venv", "docs"} 21 | SCRIPT_NAME = os.path.basename(__file__) 22 | 23 | def should_skip(path): 24 | return any(part in EXCLUDED_DIRS for part in path.split(os.sep)) 25 | 26 | def has_header(content): 27 | return "Licensed under the Apache License" in content 28 | 29 | def add_header(file_path): 30 | if os.path.basename(file_path) == SCRIPT_NAME: 31 | return 32 | with open(file_path, "r", encoding="utf-8") as f: 33 | content = f.read() 34 | if not has_header(content): 35 | with open(file_path, "w", encoding="utf-8") as f: 36 | f.write(HEADER + "\n" + content) 37 | print(f"Added header to: {file_path}") 38 | 39 | def walk(directory): 40 | for root, dirs, files in os.walk(directory): 41 | dirs[:] = [d for d in dirs if d not in EXCLUDED_DIRS] 42 | for file in files: 43 | if file.endswith(".py"): 44 | path = os.path.join(root, file) 45 | add_header(path) 46 | 47 | walk(".") 48 | 49 | -------------------------------------------------------------------------------- /waii_sdk_py/common/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | from typing import Optional, List, Dict, Any, Union 19 | 20 | from ..my_pydantic import WaiiBaseModel 21 | 22 | class CommonRequest(WaiiBaseModel): 23 | tags: Optional[List[str]] = None 24 | parameters: Optional[Dict[str, Any]] = None 25 | 26 | 27 | class LLMBasedRequest(WaiiBaseModel): 28 | tags: Optional[List[str]] = None 29 | parameters: Optional[Dict[str, Any]] = None 30 | model: Optional[str] = None 31 | # should we use cache? 32 | use_cache: Optional[bool] = True 33 | 34 | 35 | class CommonResponse(WaiiBaseModel): 36 | pass 37 | 38 | 39 | class CheckOperationStatusRequest(CommonRequest): 40 | op_id: str 41 | 42 | 43 | class OperationStatus(str, Enum): 44 | SUCCEEDED = "succeeded" 45 | FAILED = "failed" 46 | IN_PROGRESS = "in_progress" 47 | NOT_EXISTS = "not_exists" 48 | 49 | 50 | class CheckOperationStatusResponse(CommonResponse): 51 | op_id: str 52 | status: OperationStatus 53 | info: Union[Optional[str], Any] = None 54 | 55 | 56 | class AsyncObjectResponse(CommonResponse): 57 | uuid: str 58 | 59 | 60 | class GetObjectRequest(CommonRequest): 61 | uuid: str 62 | -------------------------------------------------------------------------------- /doc/docs/getting-started.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: getting-started 3 | title: Getting Started 4 | --- 5 | 6 | First you can print the list of available databases: 7 | 8 | Install and initialize the SDK 9 | ```python 10 | >>> from waii_sdk_py import WAII 11 | >>> from waii_sdk_py.query import * 12 | >>> WAII.initialize(url='...', api_key="") 13 | ``` 14 | (Refer to the [Installation](install.md) section for more details how to set url and api_key) 15 | 16 | ```python 17 | >>> print([conn.key for conn in WAII.Database.get_connections().connectors]) 18 | ``` 19 | 20 | Then, you can activate the database connection you want to use (from one of the key in the list above) 21 | 22 | ```python 23 | >>> WAII.Database.activate_connection("snowflake://...&warehouse=COMPUTE_WH") 24 | ``` 25 | 26 | If you need to do some operations which don't need to connect to a specific database (e.g. list users, list databases, add new databases, etc.), you can skip the activation step. 27 | 28 | ```python 29 | 30 | Get Database name of the active connection 31 | 32 | ```python 33 | >>> print([catalog.name for catalog in WAII.Database.get_catalogs().catalogs]) 34 | ``` 35 | 36 | ```python 37 | >>> print(WAII.Query.generate(QueryGenerationRequest(ask = "How many tables are there?")).query) 38 | 39 | SELECT COUNT(DISTINCT table_name) 40 | FROM waii.information_schema.tables 41 | ``` 42 | 43 | Run the query 44 | 45 | ```python 46 | >>> print(WAII.Query.run(RunQueryRequest(query = "SELECT COUNT(DISTINCT table_name) FROM waii.information_schema.tables"))) 47 | 48 | rows=[{'COUNT(DISTINCT TABLE_NAME)': 112}] more_rows=0 column_definitions=[ColumnDefinition(name='COUNT(DISTINCT TABLE_NAME)', type='FIXED')] query_uuid='01afbd1e-0001-d31e-0022-ba8700a8209e' 49 | ``` 50 | 51 | In order to know more details, you can check the following detailed API documentation. 52 | -------------------------------------------------------------------------------- /doc/docs/colab.md: -------------------------------------------------------------------------------- 1 | # Using Waii with Google Colab 2 | 3 | Go to https://colab.google/ 4 | Click "New Notebook" on the top-right corner 5 | 6 | ## Add api-key secret 7 | 8 | On the left panel, click the "key" icon 9 | 10 | ![colab-secret.png](colab-secret.png) 11 | 12 | Then enter the WAII API Key (which you get from WAII UI, click hamburger menu on the top-right corner, and click "Copy API Key"). 13 | 14 | NOTE: This is not OPENAI_API_KEY! You need get access to Waii first to get the API Key. (Contact us at https://waii.ai/) 15 | 16 | ![colab-api-key.png](colab-api-key.png) 17 | 18 | ## Install dependencies 19 | 20 | ```python 21 | !pip install waii-sdk-py 22 | ``` 23 | 24 | ## Import waii-sdk-py, and use API key 25 | 26 | ```python 27 | import pandas as pd 28 | import numpy as np 29 | import plotly.express as px 30 | 31 | from waii_sdk_py import WAII 32 | from waii_sdk_py.query import * 33 | 34 | from google.colab import userdata 35 | api_key = userdata.get('waii_api_key') 36 | 37 | WAII.initialize(url='https://tweakit.waii.ai/api/', api_key=api_key) 38 | ``` 39 | 40 | ## Get all connectors, generate query, and show chart 41 | 42 | ```python 43 | df = WAII.Query.generate(QueryGenerationRequest(ask = "Give me how many column for each table, group by table and schema name?")).run().to_pandas_df() 44 | display(df) 45 | WAII.Query.plot(df, "show it in scatter chart") 46 | ``` 47 | 48 | You should be able to see the result like 49 | ![colab-df-plot.png](colab-df-plot.png) 50 | 51 | ## Upload CSV and plot 52 | 53 | Use this code to upload CSV file 54 | 55 | ```python 56 | # upload csv and try 57 | from google.colab import files 58 | uploaded = files.upload() 59 | ``` 60 | 61 | Click `Choose Files` button, and select the CSV file to upload 62 | 63 | Once it is done, you can use 64 | ```python 65 | df = pd.read_csv('') 66 | df.head() 67 | ``` 68 | 69 | After that, you can just use WAII.Query.plot to plot the chart 70 | 71 | ```python 72 | WAII.Query.plot(df, "show me number of flights by month") 73 | ``` 74 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Waii Python SDK 2 | 3 | The `waii-sdk-py` is a Python library that allows you to interact with the Waii API. It provides a powerful SQL and AI capability to your Python applications. 4 | 5 | ## Installation 6 | 7 | To install the `waii-sdk-py`, you can use pip: 8 | 9 | ```bash 10 | pip install waii-sdk-py 11 | ``` 12 | 13 | Waii Python SDK depends on pydantic, but to avoid it conflicting with other libraries, it is not included in the dependencies. You can install it separately: 14 | 15 | You can either install v1.10.15+ 16 | ```bash 17 | pip install pydantic==1.10.15 18 | ``` 19 | 20 | Or 2.7.1+ 21 | ```bash 22 | pip install pydantic==2.7.1 23 | ``` 24 | 25 | ## Importing & Initialize the SDK 26 | 27 | ```python 28 | >>> from waii_sdk_py import WAII 29 | >>> from waii_sdk_py.query import * 30 | >>> WAII.initialize(api_key="") 31 | >>> print(WAII.version()) # Check the version of the SDK 32 | ``` 33 | You can get your API key from the `tweakit.waii.ai` (You need to register and get access from [waii.ai](waii.ai) first). 34 | 35 | ![img.png](doc/img.png) 36 | 37 | ## Get started 38 | 39 | First you can print the list of available databases: 40 | 41 | ```python 42 | >>> print([conn.key for conn in WAII.Database.get_connections().connectors]) 43 | ``` 44 | 45 | Then, you can activate the database connection you want to use (from one of the key in the list above) 46 | 47 | ```python 48 | >>> WAII.Database.activate_connection("snowflake://...&warehouse=COMPUTE_WH") 49 | ``` 50 | 51 | Get Database name of the active connection 52 | 53 | ```python 54 | >>> print([catalog.name for catalog in WAII.Database.get_catalogs().catalogs]) 55 | ``` 56 | 57 | ```python 58 | >>> print(WAII.Query.generate(QueryGenerationRequest(ask = "How many tables are there?")).query) 59 | 60 | SELECT COUNT(DISTINCT table_name) 61 | FROM waii.information_schema.tables 62 | ``` 63 | 64 | Run the query 65 | 66 | ```python 67 | >>> print(WAII.Query.run(RunQueryRequest(query = "SELECT COUNT(DISTINCT table_name) FROM waii.information_schema.tables"))) 68 | 69 | rows=[{'COUNT(DISTINCT TABLE_NAME)': 112}] more_rows=0 column_definitions=[ColumnDefinition(name='COUNT(DISTINCT TABLE_NAME)', type='FIXED')] query_uuid='01afbd1e-0001-d31e-0022-ba8700a8209e' 70 | ``` 71 | -------------------------------------------------------------------------------- /tests/semantic_context_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | import pytest 19 | from waii_sdk_py import WAII 20 | from waii_sdk_py.semantic_context import GetSemanticContextRequest 21 | 22 | class TestSemanticContext(unittest.TestCase): 23 | def setUp(self): 24 | WAII.initialize(url="http://localhost:9859/api/") 25 | result = WAII.Database.get_connections() 26 | self.result = result 27 | WAII.Database.activate_connection(result.connectors[2].key) 28 | 29 | def test_modify_semantic_context(self): 30 | # Define test parameters 31 | # params = ModifySemanticContextRequest(updated=[SemanticStatement(statement="Example statement")]) 32 | # Call the function 33 | # result = self.semantic_context.modifySemanticContext(params) 34 | # Check the result 35 | # Note: The specifics of this assertion would depend on what the function should return 36 | # self.assertIsInstance(result, ModifySemanticContextResponse) 37 | # self.assertEqual(len(result.updated), 1) 38 | # self.assertEqual(result.updated[0].statement, "Example statement") 39 | pass 40 | 41 | def test_get_semantic_context(self): 42 | # Call the function 43 | result = WAII.SemanticContext.get_semantic_context() 44 | # Check the result 45 | # Note: The specifics of this assertion would depend on what the function should return 46 | self.assertGreater(len(result.semantic_context), 0) 47 | 48 | def test_unknown_fields_error(self): 49 | with pytest.raises(ValueError): 50 | result = WAII.SemanticContext.get_semantic_context( 51 | GetSemanticContextRequest(search_txt='val') # search_txt field doesn't exist 52 | ) 53 | 54 | if __name__ == '__main__': 55 | unittest.main() 56 | -------------------------------------------------------------------------------- /doc/docs/chart-module.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: chart-module 3 | title: Chart 4 | --- 5 | 6 | **Initialization & Imports** 7 | ```python 8 | from waii_sdk_py import WAII 9 | from waii_sdk_py.chat import * 10 | from waii_sdk_py.query import * 11 | from waii_sdk_py.database import * 12 | from waii_sdk_py.semantic_context import * 13 | from waii_sdk_py.chart import * 14 | from waii_sdk_py.history import * 15 | 16 | WAII.initialize(url="https://your-waii-instance/api/", api_key="your-api-key") 17 | ``` 18 | 19 | The `Chart` module contains methods related to creating visualizations from sql data. 20 | 21 | Here are some of its methods: 22 | 23 | ### Generate Chart 24 | 25 | ```python 26 | WAII.Chart.generate_chart(df, ask, sql, chart_type, tweak_history) -> ChartGenerationResponse 27 | ``` 28 | 29 | This method generates the necessary information for a plot of `plot_type` based on the provided parameters. 30 | 31 | Parameter fields: 32 | - `df`: A pandas dataframe containing the data to plot 33 | - `ask`: The ask, either a question or instructions on what chart to draw and any other preferences 34 | - `sql`: The query used to generate the data 35 | - `chart_type`: A `chart_type` enum detailing what type of chart to draw 36 | - `tweak_history`: (array of `ChartTweak`) We can support both asking new question, or tweaking a previous visualization. If you want to tweak the previous question, you can set this field. A `ChartTweak` object looks like: 37 | - `chart_spec`: The previously drawn visualization 38 | - `ask`: The previous question you asked. 39 | **Examples:** 40 | 41 | **Draw a new visualization:** 42 | 43 | ```python 44 | from waii_sdk_py.query import RunQueryRequest 45 | query = None 46 | run_query_response = WAII.Query.run(RunQueryRequest(ask=query)) 47 | >>> chart = WAII.Chart.generate_chart(df=run_query_response.to_pandas_df(), 48 | ask="Draw a bar graph", 49 | sql=query, 50 | chart_type=ChartType.METABASE) 51 | ``` 52 | 53 | **Ask a follow-up question:** 54 | ```python 55 | >>> WAII.Chart.generate_chart(df=run_query_response.to_pandas_df(), 56 | ask="Title the graph trends per year, sort by year order descending", 57 | sql=query, 58 | chart_type=ChartType.METABASE, 59 | tweak_history=[ChartTweak(ask="Draw a bar graph", chart_spec=chart.chart_spec)]) 60 | ``` 61 | -------------------------------------------------------------------------------- /doc/docs/superset-waii-usage.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: superset-waii-usage 3 | title: Superset + Waii Usage 4 | --- 5 | 6 | # Superset + Waii Usage 7 | 8 | ## Setup Steps 9 | 10 | ### 1. Add Database Connection to Superset 11 | 12 | - The name (or Display Name) used while adding the connection to Superset will be used while adding the connection to WAII as well. 13 | 14 | ![superset-add-connection.png](superset-add-connection.png) 15 | 16 | ![superset-db-connection.png](superset-db-connection.png) 17 | 18 | You will use this name while adding the connection to Waii: [Database Module Documentation](https://doc.waii.ai/python/docs/database-module) 19 | 20 | Example: 21 | 22 | ```python 23 | WAII.Database.modify_connections(ModifyDBConnectionRequest( 24 | updated=[ 25 | DBConnection( 26 | # existing fields 27 | parameters={ 28 | "superset_dbname": "Snowflake-Waii-MovieDB" 29 | } 30 | ) 31 | ] 32 | )) 33 | ``` 34 | 35 | Replace "Snowflake-Waii-MovieDB" with the name that you used while adding the Superset connection. 36 | 37 | ### 2. Add User Permissions 38 | 39 | - Add permission for the user to access the database along with some basic permissions (Refer to [gamma role documentation](https://superset.apache.org/docs/security/#gamma)) 40 | 41 | ![superset-db-access.png](superset-db-access.png) 42 | 43 | ![superset-user-access.png](superset-user-access.png) 44 | 45 | ### 3. Create User in WAII and Add Superset Secrets 46 | 47 | Create the user in WAII and add the Superset secrets this way: 48 | 49 | ```python 50 | user = User( 51 | id="wangda@example.com", 52 | name="Wangda", 53 | tenant_id="my_tenant_id", 54 | org_id="my_org_id", 55 | roles=[WaiiRoles.WAII_TRIAL_USER] 56 | ) 57 | 58 | create_user_request = CreateUserRequest(user=user) 59 | WAII.User.create_user(create_user_request) 60 | 61 | create_secret_request = CreateSecretRequest( 62 | user=user, 63 | secret={ 64 | "superset.credential": { 65 | "superset.host": "http://", 66 | "superset.username": "wangda.t", 67 | "superset.password": "xxxx" 68 | } 69 | } 70 | ) 71 | 72 | WAII.User.create_secret(create_secret_request) 73 | ``` 74 | 75 | ### 4. Choose Superset Chart Type in WAII UI 76 | 77 | In the WAII UI, select chart type as Superset: 78 | 79 | ![waii-chart-selection.png](waii-chart-selection.png) 80 | 81 | 82 | You should now be able to generate charts in editor and chat mode: 83 | 84 | ![superset-waii-sql-mode.png](superset-waii-sql-mode.png) 85 | 86 | ![superset-waii-chat-mode.png](superset-waii-chat-mode.png) 87 | -------------------------------------------------------------------------------- /doc/docs/history-module.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: history-module 3 | title: History 4 | --- 5 | 6 | The `History` module contains methods related to history handling. 7 | 8 | **Initialization & Imports** 9 | ```python 10 | from waii_sdk_py import WAII 11 | from waii_sdk_py.chat import * 12 | from waii_sdk_py.query import * 13 | from waii_sdk_py.database import * 14 | from waii_sdk_py.semantic_context import * 15 | from waii_sdk_py.chart import * 16 | from waii_sdk_py.history import * 17 | 18 | WAII.initialize(url="https://your-waii-instance/api/", api_key="your-api-key") 19 | ``` 20 | 21 | Here are some of its methods: 22 | 23 | ### Get 24 | 25 | ```python 26 | History.get(params: GetGeneratedQueryRequest = GetGeneratedQueryRequest()) -> GetGeneratedQueryResponse 27 | ``` 28 | 29 | Request fields: 30 | - `included_types` (List[GeneratedHistoryEntryType], optional): The types of entries to include in the response. Defaults to `[GeneratedHistoryEntryType.query, GeneratedHistoryEntryType.chart, GeneratedHistoryEntryType.chat]`. 31 | - `limit` (int, optional): The maximum number of entries to return. Defaults to `1000`. 32 | - `offset` (int, optional): The number of entries to skip. Defaults to `0`. 33 | - `timestamp_sort_order` (SortOrder, optional): The order in which to sort the entries by timestamp. Defaults to `SortOrder.desc`. 34 | - `uuid_filter` (str, optional): The UUID of the entry to filter by. 35 | - `liked_query_filter` (bool, optional): The flag to filter by liked queries. Defaults to `None`, which includes both liked and unliked queries. When this is set to `True` or `False`, you should only include `GeneratedHistoryEntryType.query` in the `included_types` field. 36 | 37 | Response fields: 38 | - `entries` (List[GeneratedHistoryEntry]): The list of generated history entries, it can be a query, chart, or chat. 39 | 40 | #### Examples 41 | 42 | Get all queries (use offset and limits to paginate): 43 | 44 | ```python 45 | all_history_queries = [] 46 | 47 | offset = 0 48 | 49 | while True: 50 | queries = History.get(GetHistoryRequest(included_types=[GeneratedHistoryEntryType.query], limit=1000, offset=offset)).history 51 | if queries: 52 | all_history_queries.extend(queries) 53 | offset += 1000 54 | else: 55 | break 56 | ``` 57 | 58 | ### List (This is deprecated, use `get` instead) 59 | 60 | ```python 61 | History.list(params: GetGeneratedQueryHistoryRequest = GetGeneratedQueryHistoryRequest()) -> GetGeneratedQueryHistoryResponse 62 | ``` 63 | 64 | This method fetches the history of generated queries. 65 | 66 | For detailed information about the request and response objects and the available properties, refer to the respective Python files of each module. 67 | -------------------------------------------------------------------------------- /doc/docs/waii-user-roles.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: waii-user-roles 3 | title: Waii Roles Overview 4 | --- 5 | 6 | Waii implements a hierarchical role-based access control system. Each role builds upon the previous one, inheriting all permissions and adding new ones. Here's an overview of the different roles, from lowest to highest access level: 7 | 8 | ## 1. Waii Trial User (waii-trial-user) 9 | - The most basic role with limited read-only access. 10 | - Suitable for users trying out the system. (only available for SaaS) 11 | 12 | ## 2. Waii User (waii-user) 13 | - Builds upon the Trial User role. 14 | - Adds basic write permissions (add database connection, update semantic context for the user itself). 15 | - Appropriate for regular users of the system. 16 | 17 | ## 3. Waii API User (waii-api-user) 18 | - Extends the Waii User role. 19 | - Includes API usage capabilities. 20 | - Ideal for users who need to interact with the system programmatically. 21 | 22 | ## 4. Waii Admin User (waii-admin-user) 23 | - Builds upon the API User role. 24 | - Adds publishing capabilities for certain system elements: for example, publishing a semantic context to all the users who have access to the same database connection. 25 | - Suitable for users with administrative responsibilities within their scope. 26 | 27 | ## 5. Waii Org Admin User (waii-org-admin-user) 28 | - Extends the Admin User role. 29 | - Includes organization-level read and write permissions. (But cannot update the organization itself, or create new organizations) 30 | - Adds user management capabilities and advanced features. 31 | - Appropriate for users managing an entire organization within Waii. 32 | 33 | ## 6. Waii Super Admin User (waii-super-admin-user) 34 | - The highest level role, building upon the Org Admin User. 35 | - Adds system-wide organizational read and write permissions. 36 | - Suitable for system-wide administrators with the highest level of access. 37 | 38 | ## Examples 39 | 40 | Waii roles are defined inside the `WaiiRoles` class. Here's an example of how you can use these roles in your code: 41 | 42 | ``` 43 | class WaiiRoles: 44 | WAII_TRIAL_USER = 'waii-trial-user' 45 | WAII_USER = 'waii-user' 46 | WAII_API_USER = 'waii-api-user' 47 | WAII_ADMIN_USER = 'waii-admin-user' 48 | WAII_ORG_ADMIN_USER = 'waii-org-admin-user' 49 | WAII_SUPER_ADMIN_USER = 'waii-super-admin-user' 50 | ``` 51 | 52 | When you create a user in Waii, you can assign one of these roles to them. For example: 53 | 54 | ```python 55 | # other imports ... 56 | from waii_sdk_py.user import User, CreateUserRequest, WaiiRoles 57 | 58 | user = User(id=, name=, tenant_id=..., org_id=..., roles=[WaiiRoles.WAII_USER]) 59 | WAII.User.create_user(CreateUserRequest(user=user)) 60 | ``` -------------------------------------------------------------------------------- /doc/docs/multi-tenant-client-module.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: multi-tenant-client-module 3 | title: Multi-tenant Waii SDK Client 4 | --- 5 | The Waii SDK allows you to create multiple instances of the client, enabling you to manage different configurations and databases independently. 6 | 7 | **Initialization & Imports** 8 | 9 | ```python 10 | # in order to use this, you need to make sure waii-sdk-py is >= 1.28.2 11 | # otherwise, you should do 12 | # from waii_sdk_py import Waii 13 | from waii_sdk_py import Waii 14 | from waii_sdk_py.chat import * 15 | from waii_sdk_py.query import * 16 | from waii_sdk_py.database import * 17 | from waii_sdk_py.semantic_context import * 18 | from waii_sdk_py.chart import * 19 | ``` 20 | 21 | ### Creating Multiple Client Instances 22 | You can initialize multiple instances of the Waii SDK as shown below: 23 | 24 | ```python 25 | client1_sdk = Waii() 26 | client1_sdk.initialize(url='...', api_key="") 27 | client1_sdk = Waii() 28 | client2_sdk.initialize(url='...', api_key="") 29 | ``` 30 | 31 | ### Activating Different Database Connections 32 | With different client instances, you can activate and query separate database connections: 33 | 34 | ```python 35 | client1_sdk.database.activate_connection("") 36 | client2_sdk.database.activate_connection("") 37 | ask1 = QueryGenerationRequest(ask="") 38 | ask2 = QueryGenerationRequest(ask="") 39 | client1_sdk.query.generate(ask1) 40 | client2_sdk.query.generate(ask2) 41 | ``` 42 | You can notice that in static client for accessing query module we used to use WAII.Query but here we are using 43 | client1_sdk.query 44 | 45 | ### Module Access in Multi-Tenant Clients 46 | 47 | The table below summarizes how to access different modules in both the static and multi-tenant client setups: 48 | 49 | | Module | Static Client | Multi-Tenant Client | 50 | |-----------------|------------------------|--------------------------------| 51 | | Database | `WAII.Database` | `.database` | 52 | | History | `WAII.History` | `.history` | 53 | | Query | `WAII.Query` | `.query` | 54 | | SemanticContext | `WAII.SemanticContext` | `.semantic_context` | 55 | | User | `WAII.User` | `.user` | 56 | | Chart | `WAII.Chart` | `.chart` | 57 | | Chat | `WAII.Chat` | `.chat` | 58 | 59 | This setup provides flexibility by allowing different clients to interact with different databases and configurations, all within the same application. 60 | 61 | ### Example of using multi-tenant module 62 | 63 | See this notebook: https://github.com/waii-ai/waii-sdk-py/blob/main/doc/docs/waii-multi-user-example.ipynb 64 | 65 | 66 | -------------------------------------------------------------------------------- /samples/hello_waii.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from waii_sdk_py.query import Query, QueryGenerationRequest, RunQueryRequest 18 | from waii_sdk_py.semantic_context import SemanticContext, ModifySemanticContextRequest, SemanticStatement 19 | from waii_sdk_py.history import History, GetGeneratedQueryHistoryRequest 20 | from waii_sdk_py.database import Database, ModifyDBConnectionRequest, DBConnection 21 | 22 | # Initialize a new database object 23 | db = Database() 24 | # WAII.initialize(connection, '') 25 | 26 | # Define the connection data to be updated 27 | connection_data = DBConnection( 28 | key='', # will be generated by the system 29 | db_type='snowflake', 30 | account_name='', 31 | username='', 32 | password='', 33 | database='', 34 | warehouse='', 35 | role='' 36 | ) 37 | 38 | # Create a request object 39 | modify_request = ModifyDBConnectionRequest( 40 | updated=[connection_data] 41 | ) 42 | 43 | # Modify the connection using the Database class 44 | response = db.modify_connections(modify_request) 45 | 46 | # Activating the database connection 47 | if response.connectors: 48 | db.activate_connection(response.connectors[2].key) 49 | 50 | query = Query() 51 | # Generating a query 52 | gen_query = query.generate(QueryGenerationRequest( 53 | ask= 'show me all the tables in this database')).query 54 | 55 | print('-----Generated Query-----') 56 | print(gen_query) 57 | 58 | # Running the query 59 | if gen_query: 60 | result = query.run(RunQueryRequest(query = gen_query)).rows 61 | j = 0 62 | print('-----Result of Query-----') 63 | while j < len(result): 64 | print(result[j]['table_name']) 65 | j = j + 1 66 | 67 | 68 | # Getting the semantic context 69 | semantic_context = SemanticContext.get_semantic_context().semantic_context 70 | i = 0 71 | print('-----Semantic Context-----') 72 | #while loop prints only the semantic statement 73 | while i < len(semantic_context): 74 | print(semantic_context[i].statement) 75 | i = i + 1 76 | 77 | 78 | 79 | # Getting query history 80 | history = History() 81 | query_history = history.list().history 82 | query_history_questions = [] 83 | x = 0 84 | while x < len(query_history): 85 | query_history_questions.append(query_history[x].request) 86 | query_history_questions[x] = getattr(query_history_questions[x], 'ask') 87 | x = x + 1 88 | print('-----Query History-----') 89 | 90 | print(query_history_questions ) -------------------------------------------------------------------------------- /doc/docusaurus.config.js: -------------------------------------------------------------------------------- 1 | let lightCodeTheme, darkCodeTheme; 2 | try { 3 | const {themes} = require('prism-react-renderer'); 4 | lightCodeTheme = themes.github; 5 | darkCodeTheme = themes.dracula; 6 | } catch (e) { 7 | lightCodeTheme = { 8 | plain: { 9 | color: "#393A34", 10 | backgroundColor: "#f6f8fa" 11 | }, 12 | styles: [] 13 | }; 14 | darkCodeTheme = { 15 | plain: { 16 | color: "#F8F8F2", 17 | backgroundColor: "#282A36" 18 | }, 19 | styles: [] 20 | }; 21 | } 22 | 23 | (module.exports = { 24 | title: 'Waii Python SDK', 25 | tagline: 'World most accurate text-2-sql API', 26 | url: 'https://doc.waii.ai', 27 | baseUrl: '/python/', 28 | onBrokenLinks: 'throw', 29 | onBrokenMarkdownLinks: 'warn', 30 | favicon: 'img/favicon.ico', 31 | organizationName: 'waii', 32 | projectName: 'waii-sdk-py', 33 | 34 | presets: [ 35 | [ 36 | '@docusaurus/preset-classic', 37 | ({ 38 | docs: { 39 | sidebarPath: require.resolve('./sidebars.js'), 40 | editUrl: 'https://github.com/waii-ai/waii-sdk-js/tree/main/docs/docs-waii-sdk-py', 41 | }, 42 | theme: { 43 | customCss: require.resolve('./src/css/custom.css'), 44 | }, 45 | }), 46 | ], 47 | ], 48 | 49 | themeConfig: 50 | ({ 51 | navbar: { 52 | title: 'Waii Python SDK', 53 | logo: { 54 | alt: 'Waii Logo', 55 | src: 'img/logo.png', 56 | }, 57 | items: [ 58 | { 59 | label: 'TypeScript/JavaScript SDK', 60 | href: 'https://doc.waii.ai/js/docs/intro', 61 | }, 62 | { 63 | label: 'Python SDK', 64 | href: 'https://doc.waii.ai/python/docs/intro', 65 | }, 66 | { 67 | label: 'Java SDK', 68 | href: 'https://doc.waii.ai/java/docs/intro', 69 | }, 70 | { 71 | label: 'CLI', 72 | href: 'https://doc.waii.ai/cli/docs/intro', 73 | }, 74 | { 75 | label: 'Deployment & Architecture', 76 | href: 'https://doc.waii.ai/deployment/docs/intro', 77 | } 78 | ], 79 | }, 80 | footer: { 81 | style: 'dark', 82 | links: [ 83 | { 84 | title: 'Company', 85 | items: [ 86 | { 87 | label: 'Website', 88 | href: 'https://waii.ai/', 89 | }, 90 | { 91 | label: 'LinkedIn', 92 | href: 'https://www.linkedin.com/company/96891121/', 93 | } 94 | ], 95 | }, 96 | { 97 | title: 'Community', 98 | items: [ 99 | { 100 | label: 'Slack', 101 | href: 'https://join.slack.com/t/waiicommunity/shared_invite/zt-1xib44mr5-LBa7ub9t_vGvo66QtbUUpg', 102 | } 103 | ], 104 | }, 105 | { 106 | title: 'More', 107 | items: [ 108 | { 109 | label: 'GitHub', 110 | href: 'https://github.com/waii-ai/waii-sdk-py', 111 | }, 112 | ], 113 | }, 114 | ], 115 | copyright: `Copyright © ${new Date().getFullYear()} Waii, Inc.`, 116 | }, 117 | prism: { 118 | theme: lightCodeTheme, 119 | darkTheme: darkCodeTheme, 120 | }, 121 | }), 122 | }); 123 | -------------------------------------------------------------------------------- /doc/sidebars.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | docs: [ 3 | { 4 | type: 'doc', 5 | id: 'intro', 6 | label: 'Introduction', 7 | }, 8 | { 9 | type: 'doc', 10 | id: 'install', 11 | label: 'Installation' 12 | }, 13 | { 14 | type: 'doc', 15 | id: 'getting-started', 16 | label: 'Getting Started', 17 | }, 18 | { 19 | type: 'category', 20 | label: 'Notebooks', 21 | items: [ 22 | 'colab', 23 | 'notebook', 24 | ], 25 | }, 26 | { 27 | type: 'category', 28 | label: 'Modules', 29 | items: [ 30 | { 31 | type: 'doc', 32 | id: 'sql-query-module', 33 | label: 'SQL Query', 34 | }, 35 | { 36 | type: 'doc', 37 | id: 'semantic-context-module', 38 | label: 'Semantic Context', 39 | }, 40 | { 41 | type: 'doc', 42 | id: 'database-module', 43 | label: 'Database', 44 | }, 45 | { 46 | type: 'doc', 47 | id: 'history-module', 48 | label: 'History', 49 | }, 50 | { 51 | type: "doc", 52 | id: "chart-module", 53 | label: "Chart" 54 | }, 55 | { 56 | type: 'doc', 57 | id: 'chat-module', 58 | label: 'Chat' 59 | }, 60 | { 61 | type: 'doc', 62 | id: 'knowledge-graph-module', 63 | label: 'Knowledge Graph' 64 | }, 65 | { 66 | type: 'doc', 67 | id: 'user-module', 68 | label: 'User' 69 | }, 70 | { 71 | type: 'doc', 72 | id: 'multi-tenant-client-module', 73 | label: 'Multi-tenant Waii SDK Client' 74 | }, 75 | { 76 | type: 'doc', 77 | id: 'access-rule-module', 78 | label: 'Row-level Access Rule' 79 | }, 80 | { 81 | type: 'doc', 82 | id: 'semantic-layer-dump-module', 83 | label: 'Semantic Layer Export & Import' 84 | }, 85 | { 86 | type: 'doc', 87 | id: 'async-client', 88 | label: 'Async client' 89 | }, 90 | { 91 | type: 'doc', 92 | id: 'superset-waii-usage', 93 | label: 'Superset Usage With Waii' 94 | }, 95 | { 96 | type: 'doc', 97 | id: 'waii-user-roles', 98 | label: 'Waii User Roles' 99 | } 100 | ], 101 | }, 102 | ], 103 | }; 104 | -------------------------------------------------------------------------------- /waii_sdk_py/access_rules/access_rules.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | from typing import Optional, List 19 | 20 | from ..my_pydantic import WaiiBaseModel 21 | 22 | from waii_sdk_py.common import CommonRequest, CommonResponse 23 | from waii_sdk_py.database import TableName 24 | from waii_sdk_py.waii_http_client import WaiiHttpClient 25 | from ..user import User 26 | from waii_sdk_py.utils import wrap_methods_with_async 27 | UPDATE_TABLE_ACCESS_RULES_ENDPOINT = "update-table-access-rules" 28 | REMOVE_TABLE_ACCESS_RULES_ENDPOINT = "remove-table-access-rules" 29 | LIST_TABLE_ACCESS_RULES_ENDPOINT = "list-table-access-rules" 30 | 31 | 32 | class TableAccessRuleType(str, Enum): 33 | filter = "filter" # protect all access with a filter 34 | block = "block" # stop all access from the identified users 35 | 36 | 37 | class TableAccessRule(WaiiBaseModel): 38 | id: Optional[str] 39 | name: str 40 | table: TableName 41 | org_id: str = '*' 42 | tenant_id: str = '*' 43 | user_id: str = '*' 44 | type: TableAccessRuleType 45 | expression: Optional[str] 46 | 47 | 48 | class UpdateTableAccessRuleRequest(CommonRequest): 49 | rules: List[TableAccessRule] 50 | 51 | 52 | class RemoveTableAccessRuleRequest(CommonRequest): 53 | rules: List[str] 54 | 55 | 56 | class ListTableAccessRuleRequest(CommonRequest): 57 | table: Optional[TableName] 58 | ids: Optional[List[str]] 59 | lookup_user: Optional[User] 60 | 61 | 62 | class ListTableAccessRuleResponse(CommonResponse): 63 | rules: Optional[List[TableAccessRule]] 64 | 65 | 66 | class AccessRuleImpl: 67 | 68 | def __init__(self, http_client: WaiiHttpClient): 69 | self.http_client = http_client 70 | 71 | def update_table_access_rules( 72 | self, params: UpdateTableAccessRuleRequest 73 | ) -> CommonResponse: 74 | return self.http_client.common_fetch( 75 | UPDATE_TABLE_ACCESS_RULES_ENDPOINT, params, CommonResponse 76 | ) 77 | 78 | def remove_table_access_rules( 79 | self, params: RemoveTableAccessRuleRequest 80 | ) -> CommonResponse: 81 | return self.http_client.common_fetch( 82 | REMOVE_TABLE_ACCESS_RULES_ENDPOINT, params, CommonResponse 83 | ) 84 | 85 | def list_table_access_rules( 86 | self, params: ListTableAccessRuleRequest 87 | ) -> ListTableAccessRuleResponse: 88 | return self.http_client.common_fetch( 89 | LIST_TABLE_ACCESS_RULES_ENDPOINT, params, ListTableAccessRuleResponse 90 | ) 91 | 92 | 93 | class AsyncAccessRuleImpl: 94 | def __init__(self, http_client: WaiiHttpClient): 95 | self._access_rule_impl = AccessRuleImpl(http_client) 96 | wrap_methods_with_async(self._access_rule_impl, self) 97 | 98 | AccessRules = AccessRuleImpl(WaiiHttpClient.get_instance()) 99 | -------------------------------------------------------------------------------- /waii_sdk_py/settings/settings.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | 19 | from waii_sdk_py.waii_http_client import WaiiHttpClient 20 | from ..common import CommonRequest 21 | from ..my_pydantic import WaiiBaseModel 22 | from typing import Optional, List, Dict, Any 23 | 24 | from ..user import CommonResponse 25 | from waii_sdk_py.utils import wrap_methods_with_async 26 | UPDATE_PARAMETER_ENDPOINT = "update-parameter" 27 | LIST_PARAMETER_ENDPOINT = "list-parameters" 28 | DELETE_PARAMETER_ENDPOINT = "delete-parameter" 29 | 30 | 31 | class DeleteParameterRequest(CommonRequest): 32 | parameter: str 33 | target_org_id: Optional[str] = None 34 | target_tenant_id: Optional[str] = None 35 | target_user_id: Optional[str] = None 36 | target_connection_key: Optional[str] = None 37 | 38 | 39 | class UpdateParameterRequest(CommonRequest): 40 | parameter: str 41 | value: Any 42 | target_org_id: Optional[str] = None 43 | target_tenant_id: Optional[str] = None 44 | target_user_id: Optional[str] = None 45 | target_connection_key: Optional[str] = None 46 | 47 | 48 | class ParameterInfo(WaiiBaseModel): 49 | value: Optional[Any] 50 | possible_values: Optional[List[Any]] 51 | 52 | 53 | class ListParametersResponse(CommonResponse): 54 | parameters: Dict[str, ParameterInfo] 55 | 56 | 57 | class Parameters(str, Enum): 58 | LIKED_QUERIES_ENABLED = "PUBLIC.LIKED_QUERIES.ENABLED" 59 | LIKED_QUERIES_LEARNING_MODE = "PUBLIC.LIKED_QUERIES.LEARNING_MODE" 60 | REFLECTION_ENABLED = "PUBLIC.REFLECTION.ENABLED" 61 | GUARDRAIL_INVALID_QUESTION_CHECKER_ENABLED = "PUBLIC.GUARDRAIL.INVALID_QUESTION_CHECKER.ENABLED" 62 | QUERY_GENERATION_ANALYSIS_ENABLE_ALL = "PUBLIC.QUERY_GENERATION.ANALYSIS.ENABLE_ALL" 63 | DEEP_THINKING_ENABLED = "PUBLIC.DEEP_THINKING.ENABLED" 64 | 65 | class SettingsImpl: 66 | def __init__(self, http_client: WaiiHttpClient): 67 | self.http_client = http_client 68 | 69 | def update_parameter( 70 | self, params: UpdateParameterRequest 71 | ) -> CommonResponse: 72 | return self.http_client.common_fetch( 73 | UPDATE_PARAMETER_ENDPOINT, params, CommonResponse 74 | ) 75 | 76 | def list_parameters( 77 | self 78 | ) -> ListParametersResponse: 79 | return self.http_client.common_fetch( 80 | LIST_PARAMETER_ENDPOINT, CommonRequest(), ListParametersResponse 81 | ) 82 | 83 | def delete_parameter( 84 | self, params: DeleteParameterRequest 85 | ) -> CommonResponse: 86 | return self.http_client.common_fetch( 87 | DELETE_PARAMETER_ENDPOINT, params, CommonResponse 88 | ) 89 | 90 | 91 | class AsyncSettingsImpl: 92 | def __init__(self, http_client: WaiiHttpClient): 93 | self._settings_impl = SettingsImpl(http_client) 94 | wrap_methods_with_async(self._settings_impl, self) 95 | -------------------------------------------------------------------------------- /waii_sdk_py/semantic_layer_dump/semantic_layer_dump.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | 18 | from waii_sdk_py.common import CheckOperationStatusRequest, CheckOperationStatusResponse 19 | from waii_sdk_py.my_pydantic import WaiiBaseModel 20 | from typing import List, Optional 21 | from waii_sdk_py.database import SearchContext, SchemaDefinition 22 | from waii_sdk_py.query import LikedQuery 23 | from waii_sdk_py.semantic_context import SemanticStatement 24 | from waii_sdk_py.waii_http_client.waii_http_client import WaiiHttpClient 25 | 26 | from typing import Dict, Any, Union 27 | from enum import Enum 28 | 29 | 30 | IMPORT_SEMANTIC_DUMP = "semantic-layer/import" 31 | EXPORT_SEMANTIC_DUMP = "semantic-layer/export" 32 | IMPORT_SEMANTIC_DUMP_STATUS = "semantic-layer/import/status" 33 | EXPORT_SEMANTIC_DUMP_STATUS = "semantic-layer/export/status" 34 | 35 | class ExportSemanticLayerDumpRequest(WaiiBaseModel): 36 | db_conn_key: str 37 | search_context: List[SearchContext] = [SearchContext()] 38 | 39 | class ExportSemanticLayerDumpResponse(WaiiBaseModel): 40 | op_id: str 41 | 42 | class ImportSemanticLayerDumpRequest(WaiiBaseModel): 43 | db_conn_key: str 44 | configuration: Dict[str, Any] 45 | schema_mapping: Dict[str, str] = {} 46 | database_mapping: Dict[str, str] = {} 47 | search_context: List[SearchContext] = [SearchContext()] 48 | 49 | class ImportSemanticLayerDumpResponse(WaiiBaseModel): 50 | op_id: str 51 | 52 | 53 | 54 | class SemanticLayerDumpImpl: 55 | 56 | def __init__(self, http_client: WaiiHttpClient): 57 | self.http_client = http_client 58 | 59 | def export_dump( 60 | self, params: ExportSemanticLayerDumpRequest 61 | ) -> ExportSemanticLayerDumpResponse: 62 | return self.http_client.common_fetch( 63 | EXPORT_SEMANTIC_DUMP, 64 | params, 65 | ExportSemanticLayerDumpResponse 66 | ) 67 | 68 | def import_dump( 69 | self, params: ImportSemanticLayerDumpRequest 70 | ) -> ImportSemanticLayerDumpResponse: 71 | return self.http_client.common_fetch( 72 | IMPORT_SEMANTIC_DUMP, 73 | params, 74 | ImportSemanticLayerDumpResponse 75 | ) 76 | 77 | def import_dump_status( 78 | self, params: CheckOperationStatusRequest 79 | ) -> CheckOperationStatusResponse: 80 | return self.http_client.common_fetch( 81 | IMPORT_SEMANTIC_DUMP_STATUS, 82 | params, 83 | CheckOperationStatusResponse 84 | ) 85 | 86 | def export_dump_status( 87 | self, params: CheckOperationStatusRequest 88 | ) -> CheckOperationStatusResponse: 89 | return self.http_client.common_fetch( 90 | EXPORT_SEMANTIC_DUMP_STATUS, 91 | params, 92 | CheckOperationStatusResponse 93 | ) 94 | 95 | SemanticLayerDump = SemanticLayerDumpImpl(WaiiHttpClient.get_instance()) -------------------------------------------------------------------------------- /waii_sdk_py/my_pydantic/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | try: 18 | from pydantic.v1 import ( 19 | BaseModel, 20 | ValidationError, 21 | PrivateAttr, 22 | Field, 23 | # Add other necessary imports here 24 | ) 25 | 26 | class WaiiBaseModel(BaseModel, extra='allow'): 27 | def check_extra_fields(self): 28 | extra_fields = {k: v for k, v in self.__dict__.items() if k not in self.__class__.__fields__} 29 | if extra_fields != {}: 30 | raise ValueError(f'Cannot set unknown fields: {list(extra_fields.keys())}') 31 | for field_name, field in self.__fields__.items(): 32 | field_value = getattr(self, field_name) 33 | if field_value is None: 34 | continue 35 | if isinstance(field_value, list): 36 | for item in field_value: 37 | if isinstance(item, WaiiBaseModel): 38 | item.check_extra_fields() 39 | elif isinstance(field_value, dict): 40 | for k, v in field_value.items(): 41 | if isinstance(v, WaiiBaseModel): 42 | v.check_extra_fields() 43 | elif isinstance(field_value, WaiiBaseModel): 44 | field_value.check_extra_fields() 45 | 46 | except ImportError: 47 | try: 48 | from pydantic import ( 49 | BaseModel, 50 | ValidationError, 51 | PrivateAttr, 52 | Field, 53 | # Add other necessary imports here 54 | ) 55 | 56 | class WaiiBaseModel(BaseModel, extra='allow'): 57 | def check_extra_fields(self): 58 | if self.model_extra != {}: 59 | raise ValueError(f'Cannot set unknown fields: {list(self.model_extra.keys())}') 60 | 61 | for field_name, field in self.model_fields.items(): 62 | field_value = getattr(self, field_name) 63 | if field_value is None: 64 | continue 65 | if isinstance(field_value, list): 66 | for item in field_value: 67 | if isinstance(item, WaiiBaseModel): 68 | item.check_extra_fields() 69 | elif isinstance(field_value, dict): 70 | for k, v in field_value.items(): 71 | if isinstance(v, WaiiBaseModel): 72 | v.check_extra_fields() 73 | elif isinstance(field_value, WaiiBaseModel): 74 | field_value.check_extra_fields() 75 | 76 | 77 | except ImportError: 78 | raise ImportError("Cannot find pydantic module. Please install pydantic. You can use >= 1.10.x or >= 2.7.x") 79 | 80 | __all__ = [ 81 | "BaseModel", 82 | "WaiiBaseModel", 83 | "ValidationError", 84 | "PrivateAttr", 85 | "Field" 86 | ] 87 | -------------------------------------------------------------------------------- /waii_sdk_py/sdm/sdm.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from typing import Optional, List 18 | from waii_sdk_py.common import CommonRequest, CommonResponse 19 | from waii_sdk_py.my_pydantic import WaiiBaseModel 20 | from waii_sdk_py.waii_http_client.waii_http_client import WaiiHttpClient 21 | from waii_sdk_py.utils import wrap_methods_with_async 22 | 23 | INDEX_SDM_ENDPOINT = "index-sdm" 24 | LIST_SDMS_ENDPOINT = "list-sdms" 25 | DELETE_SDM_ENDPOINT = "delete-sdm" 26 | 27 | 28 | class IndexSdmRequest(CommonRequest): 29 | db_id: str 30 | scope: str 31 | user_id: Optional[str] = None 32 | org_id: Optional[str] = None 33 | tenant_id: Optional[str] = None 34 | sdm_name: str 35 | content: Optional[str] = None 36 | path: Optional[str] = None 37 | 38 | 39 | class IndexSdmResponse(CommonResponse): 40 | success: bool 41 | message: str 42 | 43 | 44 | class SdmObjectInfo(WaiiBaseModel): 45 | api_name: str 46 | label: str 47 | object_type: str # "sdm", "data_object", "calculated_measurement", "calculated_dimension" 48 | description: Optional[str] = None 49 | 50 | 51 | class ListSdmsRequest(CommonRequest): 52 | db_id: str 53 | scope: str 54 | user_id: Optional[str] = None 55 | org_id: Optional[str] = None 56 | tenant_id: Optional[str] = None 57 | sdm_name: Optional[str] = None 58 | 59 | 60 | class ListSdmsResponse(CommonResponse): 61 | objects: List[SdmObjectInfo] 62 | 63 | 64 | class DeleteSdmRequest(CommonRequest): 65 | db_id: str 66 | scope: str 67 | user_id: Optional[str] = None 68 | org_id: Optional[str] = None 69 | tenant_id: Optional[str] = None 70 | sdm_name: Optional[str] = None 71 | 72 | 73 | class DeleteSdmResponse(CommonResponse): 74 | success: bool 75 | message: str 76 | count: Optional[int] = None 77 | 78 | 79 | class SdmImpl: 80 | 81 | def __init__(self, http_client: WaiiHttpClient): 82 | self.http_client = http_client 83 | 84 | def index_sdm( 85 | self, params: IndexSdmRequest 86 | ) -> IndexSdmResponse: 87 | return self.http_client.common_fetch( 88 | INDEX_SDM_ENDPOINT, 89 | params, 90 | IndexSdmResponse 91 | ) 92 | 93 | def list_sdms( 94 | self, params: ListSdmsRequest 95 | ) -> ListSdmsResponse: 96 | return self.http_client.common_fetch( 97 | LIST_SDMS_ENDPOINT, 98 | params, 99 | ListSdmsResponse 100 | ) 101 | 102 | def delete_sdm( 103 | self, params: DeleteSdmRequest 104 | ) -> DeleteSdmResponse: 105 | return self.http_client.common_fetch( 106 | DELETE_SDM_ENDPOINT, 107 | params, 108 | DeleteSdmResponse 109 | ) 110 | 111 | 112 | class AsyncSdmImpl: 113 | def __init__(self, http_client: WaiiHttpClient): 114 | self._sdm_impl = SdmImpl(http_client) 115 | wrap_methods_with_async(self._sdm_impl, self) 116 | 117 | 118 | Sdm = SdmImpl(WaiiHttpClient.get_instance()) 119 | 120 | -------------------------------------------------------------------------------- /tests/history_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | import pytest 19 | from waii_sdk_py import WAII 20 | from waii_sdk_py.history import * 21 | from waii_sdk_py.query import RunQueryRequest 22 | 23 | 24 | class TestHistory(unittest.TestCase): 25 | def setUp(self): 26 | WAII.initialize(url="http://localhost:9859/api/") 27 | result = WAII.Database.get_connections() 28 | self.result = result 29 | WAII.Database.activate_connection(result.connectors[0].key) 30 | 31 | def test_list_legacy_api(self): 32 | # do a list and see if it returns something 33 | result = WAII.History.list() 34 | if not result.history: 35 | # first we need to create some queries 36 | result = WAII.Query.generate(QueryGenerationRequest(ask='how many tables are there in the database', parent_uuid=None)) 37 | WAII.Query.run(RunQueryRequest(query=result.query)) 38 | 39 | # Call the function 40 | result = WAII.History.list() 41 | # Check the result 42 | # Note: The specifics of this assertion would depend on what the function should return 43 | self.assertGreater(len(result.history), 0) 44 | 45 | params = GetGeneratedQueryHistoryRequest(limit=1, offset=1) 46 | result = WAII.History.get(params) 47 | self.assertTrue(len(result.history) >= 1) 48 | 49 | def test_get_new_api(self): 50 | result = WAII.History.get(GetHistoryRequest()) 51 | if not result.history: 52 | # Then generate some chats 53 | asks = [ 54 | "tell me how many tables are there in the database", 55 | "list number of columns for each table" 56 | ] 57 | parent_uuid = None 58 | for a in asks: 59 | chat_request = ChatRequest(ask=a, parent_uuid=parent_uuid) 60 | chat_response = WAII.Chat.chat_message(chat_request) 61 | print('Ask:', a) 62 | print(chat_response.response) 63 | if chat_response.response_data.sql: 64 | print(chat_response.response_data.sql.query) 65 | if chat_response.response_data.data: 66 | print(f"There're {len(chat_response.response_data.data.rows)} rows") 67 | parent_uuid = chat_response.chat_uuid 68 | 69 | # Then use the get method 70 | result = WAII.History.get(GetHistoryRequest()) 71 | 72 | # Check the result 73 | self.assertGreater(len(result.history), 0) 74 | 75 | total_result = len(result.history) 76 | 77 | # Test with limit 78 | result = WAII.History.get(GetHistoryRequest(limit=1)) 79 | self.assertEqual(len(result.history), 1) 80 | 81 | # Test with offset 82 | result = WAII.History.get(GetHistoryRequest(offset=1)) 83 | self.assertEqual(len(result.history), total_result - 1) 84 | 85 | # Test error with unknown fields 86 | with pytest.raises(ValueError): 87 | result = WAII.History.get( 88 | GetHistoryRequest(offset=1, lt=5, id=4) # lt and id fields don't exist 89 | ) 90 | 91 | 92 | 93 | if __name__ == '__main__': 94 | unittest.main() 95 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | .idea/ 161 | 162 | # vscode 163 | .vscode/ 164 | -------------------------------------------------------------------------------- /waii_sdk_py/kg/kg.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | from typing import Optional, List, Union 19 | from pydantic import Field 20 | 21 | from ..my_pydantic import WaiiBaseModel 22 | from ..common import LLMBasedRequest 23 | from ..database import TableDefinition, ColumnDefinition, SchemaDefinition, Constraint 24 | from ..semantic_context import SemanticStatement 25 | from waii_sdk_py.utils import wrap_methods_with_async 26 | from ..waii_http_client import WaiiHttpClient 27 | 28 | GET_KNOWLEDGE_GRAPH_ENDPOINT = "get-knowledge-graph" 29 | 30 | 31 | class GetKnowledgeGraphRequest(LLMBasedRequest): 32 | ask: str 33 | 34 | def trim(self): 35 | self.ask = self.ask.strip() 36 | 37 | 38 | class KnowlegeGraphNodeType(str, Enum): 39 | table = "table" 40 | column = "column" 41 | schema = "schema" 42 | semantic_statement = "semantic_statement" 43 | 44 | 45 | class KnowledgeGraphEdgeType(str, Enum): 46 | # join between two tables 47 | constraint = "constraint" 48 | # reference to a semantic context 49 | semantic_statement_reference = "semantic_statement_reference" 50 | table_to_column = "table_to_column" 51 | schema_to_table = "schema_to_table" 52 | 53 | 54 | class KnowledgeGraphNode(WaiiBaseModel): 55 | id: str 56 | display_name: str 57 | entity_type: str # Use str instead of enum, but can compare with KnowlegeGraphNodeType values 58 | entity: Union[TableDefinition, ColumnDefinition, SchemaDefinition, SemanticStatement] = Field(..., discriminator='entity_type') 59 | parent_entity: Optional[Union[TableDefinition, SchemaDefinition]] = Field(None, discriminator='entity_type') 60 | 61 | 62 | class KnowledgeGraphEdge(WaiiBaseModel): 63 | edge_type: str # Use str instead of enum, but can compare with KnowledgeGraphEdgeType values 64 | source_id: str 65 | target_id: str 66 | # when it is directed, it's from source to target 67 | # when it is undirected, it's just link between two nodes (source and target are interchangeable) 68 | directed: bool 69 | description: Optional[str] = None 70 | # when it is table to column or schema to table, we will keep the edge_entity to None 71 | edge_entity: Optional[Union[SemanticStatement, Constraint]] = Field(None, discriminator='entity_type') 72 | 73 | 74 | class KnowledgeGraph(WaiiBaseModel): 75 | nodes: List[KnowledgeGraphNode] 76 | edges: List[KnowledgeGraphEdge] 77 | 78 | 79 | class GetKnowledgeGraphResponse(WaiiBaseModel): 80 | graph: Optional[KnowledgeGraph] = None 81 | 82 | 83 | class KnowledgeGraphImpl: 84 | 85 | def __init__(self, http_client: WaiiHttpClient): 86 | self.http_client = http_client 87 | 88 | def get_knowledge_graph(self, params: GetKnowledgeGraphRequest) -> GetKnowledgeGraphResponse: 89 | """ 90 | Get a knowledge graph based on the provided parameters. 91 | 92 | Args: 93 | params: GetKnowledgeGraphRequest containing the query parameters 94 | 95 | Returns: 96 | GetKnowledgeGraphResponse containing the generated knowledge graph 97 | """ 98 | return self.http_client.common_fetch( 99 | GET_KNOWLEDGE_GRAPH_ENDPOINT, params, GetKnowledgeGraphResponse 100 | ) 101 | 102 | 103 | class AsyncKnowledgeGraphImpl: 104 | def __init__(self, http_client: WaiiHttpClient): 105 | self._kg_impl = KnowledgeGraphImpl(http_client) 106 | wrap_methods_with_async(self._kg_impl, self) 107 | 108 | 109 | # Use KnowledgeGraphClient to avoid conflict with KnowledgeGraph model class 110 | KnowledgeGraphClient = KnowledgeGraphImpl(WaiiHttpClient.get_instance()) -------------------------------------------------------------------------------- /waii_sdk_py/waii_http_client/waii_http_client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import requests 18 | import json 19 | from typing import TypeVar, Generic, Optional, Dict, Union, Any 20 | from collections import namedtuple 21 | from ..my_pydantic import WaiiBaseModel 22 | 23 | 24 | T = TypeVar('T') 25 | 26 | class WaiiHttpClient(Generic[T]): 27 | instance = None 28 | 29 | def __init__(self, url: str, apiKey: str, verbose=False): 30 | WaiiHttpClient.instance = self 31 | self.url = url 32 | self.apiKey = apiKey 33 | self.timeout = 150000000 34 | self.scope = '' 35 | self.orgId = '' 36 | self.userId = '' 37 | self.impersonateUserId = '' 38 | self.verbose = verbose 39 | 40 | @classmethod 41 | def get_instance(cls, url: str = None, apiKey: str = None): 42 | if cls.instance is None or (url and url != cls.instance.url) or (apiKey and apiKey != cls.instance.apiKey): 43 | cls.instance = WaiiHttpClient(url, apiKey) 44 | return cls.instance 45 | 46 | def set_scope(self, scope: str): 47 | self.scope = scope 48 | 49 | def get_scope(self): 50 | return self.scope 51 | 52 | def set_org_id(self, orgId: str): 53 | self.orgId = orgId 54 | 55 | def set_user_id(self, userId: str): 56 | self.userId = userId 57 | 58 | def set_impersonate_user_id(self, userId: str): 59 | self.impersonateUserId = userId 60 | 61 | def common_fetch( 62 | self, 63 | endpoint: str, 64 | req: Union[Union[WaiiBaseModel, dict[str, Any]]], 65 | cls: WaiiBaseModel = None, 66 | need_scope: bool = True, 67 | ret_json: bool = False 68 | ) -> Optional[T]: 69 | # check to ensure no additional fields are passed 70 | if isinstance(req, WaiiBaseModel): 71 | req.check_extra_fields() 72 | params = req.__dict__ 73 | else: 74 | params = req 75 | 76 | if need_scope: 77 | if not self.scope or self.scope.strip() == '': 78 | raise Exception("You need to activate connection first, use `WAII.Database.activate_connection(...)`") 79 | params['scope'] = self.scope 80 | params['org_id'] = self.orgId 81 | params['user_id'] = self.userId 82 | 83 | headers = {'Content-Type': 'application/json'} 84 | if self.apiKey: 85 | headers['Authorization'] = f'Bearer {self.apiKey}' 86 | if self.impersonateUserId: 87 | headers['x-waii-impersonate-user'] = self.impersonateUserId 88 | 89 | if self.verbose: 90 | # print cUrl equivalent 91 | print("calling endpoint: ", endpoint) 92 | print(f"curl -X POST '{self.url + endpoint}' -H 'Content-Type: application/json' -H 'Authorization: Bearer {self.apiKey}' -d '{json.dumps(params, default=vars)}'") 93 | 94 | response = requests.post(self.url + endpoint, headers=headers, data=json.dumps(params, default=vars), timeout=self.timeout/1000) # timeout is in seconds 95 | 96 | if response.status_code != 200: 97 | try: 98 | print(response) 99 | error = response.json() 100 | raise Exception(error.get('detail', '')) 101 | except json.JSONDecodeError: 102 | raise Exception(response.text) 103 | try: 104 | if cls: 105 | result: T = cls(**response.json()) 106 | else: 107 | if not ret_json: 108 | result: T = json.loads(response.text, object_hook=lambda d: namedtuple('X', d.keys())(*d.values())) 109 | else: 110 | result: T = response.json() 111 | return result 112 | except json.JSONDecodeError: 113 | raise Exception("Invalid response received.") 114 | -------------------------------------------------------------------------------- /tests/query_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | import pytest 19 | from waii_sdk_py import WAII 20 | from waii_sdk_py.query import * 21 | 22 | 23 | class TestQuery(unittest.TestCase): 24 | def setUp(self): 25 | WAII.initialize(url="http://localhost:9859/api/") 26 | result = WAII.Database.get_connections() 27 | self.result = result 28 | pg_connector = None 29 | for connector in result.connectors: 30 | if connector.db_type == "postgresql" and connector.database == 'waii_sdk_test': 31 | pg_connector = connector 32 | break 33 | 34 | WAII.Database.activate_connection(pg_connector.key) 35 | 36 | def test_generate(self): 37 | params = QueryGenerationRequest(ask="How many movies are there?", use_cache=False) 38 | result = WAII.Query.generate(params) 39 | self.assertIsInstance(result, GeneratedQuery) 40 | assert result.uuid is not None 41 | assert len(result.detailed_steps) > 0 42 | assert len(result.query) > 0 43 | assert 'public.movies' in result.query.lower() 44 | assert len(result.tables) > 0 45 | assert hasattr(result.confidence_score, "confidence_value") 46 | 47 | params = LikeQueryRequest(query_uuid=result.uuid, liked=True) 48 | result = WAII.Query.like(params) 49 | self.assertIsInstance(result, LikeQueryResponse) 50 | 51 | def test_tweak(self): 52 | params = QueryGenerationRequest(ask = "only return 10 results", tweak_history=[Tweak(sql="select * from information_schema.tables", ask="give me all tables")]) 53 | result = WAII.Query.generate(params) 54 | print(result) 55 | assert "limit 10" in result.query.lower() 56 | assert len(result.what_changed) > 0 57 | 58 | def test_run(self): 59 | params = RunQueryRequest(query="SELECT 42") 60 | result = WAII.Query.run(params) 61 | self.assertIsInstance(result, GetQueryResultResponse) 62 | assert len(result.column_definitions) > 0 63 | assert '42' in str(result.rows[0]) 64 | 65 | 66 | 67 | def test_submit(self): 68 | params = RunQueryRequest(query="SELECT 42") 69 | result = WAII.Query.submit(params) 70 | self.assertIsInstance(result, RunQueryResponse) 71 | assert result.query_id is not None 72 | 73 | params = GetQueryResultRequest(query_id=result.query_id) 74 | result = WAII.Query.get_results(params) 75 | self.assertIsInstance(result, GetQueryResultResponse) 76 | assert '42' in str(result.rows[0]) 77 | 78 | def test_describe(self): 79 | params = DescribeQueryRequest(query="SELECT 42") 80 | result = WAII.Query.describe(params) 81 | self.assertIsInstance(result, DescribeQueryResponse) 82 | print(result) 83 | assert len(result.detailed_steps) > 0 84 | assert len(result.summary) > 0 85 | 86 | def test_transcode(self): 87 | params = TranscodeQueryRequest(source_dialect="mysql", target_dialect="postgres", source_query="SELECT * FROM movies") 88 | result = WAII.Query.transcode(params) 89 | self.assertIsInstance(result, GeneratedQuery) 90 | assert len(result.query) > 0 91 | assert 'movies' in result.query.lower() 92 | 93 | def test_like_query(self): 94 | params = LikeQueryRequest(liked=True,ask="like test ask", query="SELECT S_STATE FROM STORE", 95 | detailed_steps=["step3", "step4"]) 96 | result = WAII.Query.like(params) 97 | self.assertIsInstance(result, LikeQueryResponse) 98 | params = QueryGenerationRequest(ask="like test ask") 99 | result = WAII.Query.generate(params) 100 | assert result.detailed_steps == ["step3", "step4"] 101 | 102 | def test_error_with_unknown_fields(self): 103 | params = QueryGenerationRequest(ask='like test ask', q_id='unknown id') 104 | 105 | with pytest.raises(ValueError): 106 | result = WAII.Query.generate(params) 107 | 108 | 109 | if __name__ == '__main__': 110 | unittest.main() 111 | -------------------------------------------------------------------------------- /waii_sdk_py/chart/chart.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | from typing import Optional, Literal, List, Union, Dict, Any 19 | 20 | from ..database import ColumnDefinition 21 | from ..my_pydantic import WaiiBaseModel 22 | 23 | from ..common import LLMBasedRequest 24 | from waii_sdk_py.utils import wrap_methods_with_async 25 | from ..waii_http_client import WaiiHttpClient 26 | 27 | GENERATE_CHART_ENDPOINT = "generate-chart" 28 | 29 | 30 | class ChartType(str, Enum): 31 | METABASE = "metabase" 32 | SUPERSET = "superset" 33 | PLOTLY = "plotly" 34 | VEGALITE = "vegalite" 35 | 36 | 37 | class SuperSetChartSpec(WaiiBaseModel): 38 | spec_type: Literal['superset'] 39 | plot_type: Optional[str] 40 | metrics: Optional[List[str]] 41 | dimensions: Optional[List[str]] 42 | chart_name: Optional[str] 43 | color_hex: Optional[str] 44 | x_axis: Optional[str] 45 | y_axis: Optional[str] 46 | grid_style: Optional[str] 47 | stacked: Optional[bool] 48 | width: Optional[int] 49 | height: Optional[int] 50 | 51 | 52 | class MetabaseChartSpec(WaiiBaseModel): 53 | spec_type: Literal['metabase'] 54 | plot_type: Optional[str] 55 | metric: Optional[str] 56 | dimension: Optional[str] 57 | name: Optional[str] 58 | color_hex: Optional[str] 59 | 60 | 61 | class PlotlyChartSpec(WaiiBaseModel): 62 | spec_type: Literal['plotly'] = 'plotly' 63 | plot: Optional[str] 64 | 65 | 66 | class VegaliteChartSpec(WaiiBaseModel): 67 | spec_type: Literal['vegalite'] = 'vegalite' 68 | chart: Optional[str] 69 | number_of_data_points: Optional[int] = None 70 | generation_message: Optional[str] = None 71 | 72 | 73 | class ChartTweak(WaiiBaseModel): 74 | ask: Optional[str] 75 | chart_spec: Optional[Union[SuperSetChartSpec, MetabaseChartSpec, PlotlyChartSpec, VegaliteChartSpec]] 76 | 77 | def __str__(self): 78 | return f"previous ask={self.ask}, previous chart_spec={self.chart_spec})\n" 79 | 80 | 81 | class ChartGenerationRequest(LLMBasedRequest): 82 | sql: Optional[str] 83 | ask: Optional[str] 84 | dataframe_rows: Optional[List[Dict[str, Any]]] 85 | dataframe_cols: Optional[List[ColumnDefinition]] 86 | chart_type: Optional[ChartType] 87 | parent_uuid: Optional[str] 88 | tweak_history: Optional[List[ChartTweak]] 89 | 90 | 91 | class ChartGenerationResponse(WaiiBaseModel): 92 | uuid: str 93 | timestamp_ms: Optional[int] 94 | chart_spec: Optional[Union[SuperSetChartSpec, MetabaseChartSpec, PlotlyChartSpec, VegaliteChartSpec]] 95 | 96 | 97 | class ChartImpl: 98 | 99 | def __init__(self, http_client: WaiiHttpClient): 100 | self.http_client = http_client 101 | 102 | def generate_chart( 103 | self, df, ask=None, sql=None, chart_type=None, parent_uuid=None, tweak_history=None, 104 | ) -> ChartGenerationResponse: 105 | 106 | #Remove duplicate columns 107 | df = df.loc[:, ~df.columns.duplicated()] 108 | 109 | cols = [] 110 | for col in df.columns: 111 | cols.append(ColumnDefinition(name=col, type=df[col][0].__class__.__name__)) 112 | 113 | params = ChartGenerationRequest(dataframe_cols=cols, 114 | dataframe_rows=df.to_dict(orient='records'), 115 | ask=ask, 116 | chart_type=chart_type, 117 | parent_uuid=parent_uuid, 118 | tweak_history=tweak_history) 119 | params.check_extra_fields() 120 | params_dict = {k: v.value if isinstance(v, Enum) else v for k, v in params.dict().items() if v is not None} 121 | 122 | return self.http_client.common_fetch( 123 | GENERATE_CHART_ENDPOINT, params_dict, ChartGenerationResponse 124 | ) 125 | 126 | 127 | class AsyncChartImpl: 128 | def __init__(self, http_client: WaiiHttpClient): 129 | self._chart_impl = ChartImpl(http_client) 130 | wrap_methods_with_async(self._chart_impl, self) 131 | 132 | 133 | Chart = ChartImpl(WaiiHttpClient.get_instance()) 134 | -------------------------------------------------------------------------------- /doc/docs/notebook.md: -------------------------------------------------------------------------------- 1 | # Using Waii with Jupyter 2 | 3 | Waii can be used in your favourite notebook environment, let's take an example with Jupyter Notebook. 4 | 5 | ## Use Jupyter Notebook 6 | 7 | First, you need to install Jupyter Notebook. You can follow the [official documentation](https://jupyter.org/install) or use the following command: 8 | 9 | ```bash 10 | pip install notebook 11 | ``` 12 | 13 | Launch Jupyter Notebook: 14 | 15 | ```bash 16 | jupyter notebook 17 | ``` 18 | 19 | If you want to use Jupyter in a virtual environment, you can follow this [guide](https://medium.com/@eleroy/jupyter-notebook-in-a-virtual-environment-virtualenv-8f3c3448247) 20 | 21 | ## Install Waii 22 | 23 | Inside your notebook, you can install Waii with the following command: 24 | 25 | ``` 26 | !pip install waii-sdk-py 27 | !pip install pandas 28 | ``` 29 | 30 | ## Import Waii 31 | 32 | ```python 33 | from waii_sdk_py import WAII 34 | from waii_sdk_py.chat import * 35 | from waii_sdk_py.query import * 36 | from waii_sdk_py.database import * 37 | from waii_sdk_py.semantic_context import * 38 | from waii_sdk_py.chart import * 39 | 40 | WAII.initialize(url="https://your-waii-instance/api/", api_key="your-api-key") 41 | 42 | # show all the connections available 43 | print([conn.key for conn in WAII.Database.get_connections().connectors]) 44 | 45 | # activate the connection you want to use, use one of the keys from the previous command 46 | # This is optional if you only want to use our preloaded playground connection 47 | WAII.Database.activate_connection('snowflake://...') 48 | ``` 49 | 50 | ## Create pandas dataframe from Snowflake query 51 | 52 | First, let's try to generate a query from a question `give me most popular language for each country`: 53 | 54 | ```python 55 | print(WAII.Query.generate(QueryGenerationRequest(ask = "give me most popular language for each country")).query) 56 | ``` 57 | 58 | You will see the following output: 59 | 60 | ```sql 61 | WITH language_percentage AS ( 62 | SELECT 63 | name AS country, 64 | language, 65 | SUM(percentage) AS total_percentage 66 | FROM waii.world.countrylanguage AS l 67 | INNER JOIN waii.world.country AS c 68 | ON countrycode = code 69 | GROUP BY 70 | name, 71 | language 72 | ) 73 | 74 | SELECT 75 | country, 76 | language 77 | FROM ( 78 | SELECT 79 | country, 80 | language, 81 | ROW_NUMBER() OVER (PARTITION BY country ORDER BY total_percentage DESC) AS rn 82 | FROM language_percentage 83 | ) 84 | WHERE 85 | rn = 1 86 | ``` 87 | 88 | You can simply add `.run()` to execute the query and get the result as a pandas dataframe: 89 | 90 | ```python 91 | df = WAII.Query.generate(QueryGenerationRequest(ask = "give me most popular language for each country")).run().to_pandas_df() 92 | display(df) 93 | ``` 94 | 95 | You should be able to see result like this: 96 | 97 | ![df_output.png](df_output.png) 98 | 99 | ## Plot the result 100 | 101 | ```python 102 | WAII.Query.plot(df, "plot number of countries for each language") 103 | ``` 104 | 105 | ![plot_result.png](plot_result.png) 106 | 107 | ## Use map to visualize the result 108 | 109 | Let's look at another example, this time we will use a map to visualize the result. 110 | 111 | ```python 112 | df = WAII.Query.generate(QueryGenerationRequest(ask = "Give me all cars, with their maker, country and continent")).run().to_pandas_df() 113 | display(df) 114 | ``` 115 | ![df_output_cars.png](df_output_cars.png) 116 | 117 | Plot it on a map: 118 | 119 | ```python 120 | WAII.Query.plot(df, "use world map to show me all the countries which have at least one car maker") 121 | ``` 122 | 123 | ![plot_map.png](plot_map.png) 124 | 125 | ## Get script for the plot 126 | 127 | If you want to use the plot in your own code, or debug when something goes wrong. You can add `return_plot_script=True` to the `plot` function: 128 | 129 | ```python 130 | code = WAII.Query.plot(df, "use world map to show me all the countries which have at least one car maker", return_plot_script=True) 131 | print(code) 132 | ``` 133 | 134 | You can see the following output: 135 | 136 | ```python 137 | import pandas as pd 138 | import plotly.express as px 139 | 140 | # Assuming 'df' is your DataFrame and 'COUNTRYNAME' is the column with country names 141 | df_unique_countries = df.drop_duplicates(subset=['COUNTRYNAME']) 142 | 143 | fig = px.choropleth(df_unique_countries, locations='COUNTRYNAME', 144 | locationmode='country names', 145 | color='COUNTRYNAME', 146 | title='Countries with at least one car maker', 147 | hover_name='COUNTRYNAME', 148 | projection='natural earth') 149 | 150 | fig.show() 151 | ``` 152 | 153 | ## Example notebook 154 | 155 | You can find a full example notebook [here](waii-example.ipynb). 156 | -------------------------------------------------------------------------------- /doc/docs/waii-similarity-search.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: waii-similarity-search 3 | title: Waii Similarity Search 4 | --- 5 | 6 | ## Waii Similarity Search 7 | 8 | A similarity search index is defined over a column. The column values and any additional meanings for the values will be ingested into a embedding storage, 9 | and these values can be used during query generation. See the database documentation for more details. 10 | 11 | 12 | `waii_similarity_search` is a virtual user defined table function (UDTF) defined within the Waii query generation engine. 13 | This function conducts an embedding search over the similarity search index of a column and chooses the best values for a query. 14 | Waii will use this function during query generation if it can't find the best column value to use or if the query requires a similarity search over certain phrases. 15 | 16 | 17 | When running the query through Waii, this UDTF will be expanded into a `Values` clause with the chosen column values. 18 | This column value selection behavior can be customized with the following similarity search index properties: 19 | 20 | - `enable_llm_rerank`: Enable extra LLM reranking step to pick the best values to use within the query. Default true 21 | 22 | The following args only apply if enable_llm_rerank is set False 23 | - `similarity_score_threshold`: Only use results with an embedding match above a certain amount 24 | - `max_matched_values`: The max match values from embedding match 25 | - `min_matched_values`: The minimum matched values from embedding match. 26 | 27 | 28 | `waii_similarity_search` accepts the following arguments: 29 | - `column`: The column to conduct similarity search over. This column must have a similarity search index 30 | - `phrase`: The phrase to do similarity search with 31 | - `limit`: Optional arg. Limits the number of matched results, this will be the min of `min_matched_values` and `limit` if both are specified 32 | 33 | `waii_similarity_search` outputs the following columns: 34 | - `val`: The value of the column that passes the similarity search 35 | - `score`: The similarity score for the row 36 | 37 | #### enable_llm_rerank 38 | This is primary defining behavior for a similarity search index. The default setting is true, this is best for categorical columns. When waii_similarity_search is used, the LLM will pick the best of the categorical values to insert into the query based on the lookup phrase used. 39 | 40 | When enable_llm_rerank is false, waii_similarity_search will skip the LLM rerank step. Instead, all of the values that pass the embedding match will be used. This embedding match operation will respect the properties of the similarity search index. 41 | 42 | ### Examples 43 | 44 | #### Similarity Search Index with LLM Rerank 45 | 46 | Given 47 | - table `brand` with columns: brand_name, brand_id 48 | - table `store` with columns: brand_id, number_of_employees 49 | - similarity search index on the brand_name column with `enable_llm_rerank` True 50 | 51 | For the user ask, "How many employees work at a Kohls" 52 | 53 | ```sql 54 | select sum(number_of_employees) 55 | from brand 56 | join table(waii_similarity_search(brand_name, 'kohls', 1)) as matched_brand(val, score) 57 | on brand.brand_name = matched_stores.val 58 | join store on brand.brand_id = store.brand_id 59 | ``` 60 | 61 | During query generation, if Waii does not have access to the best brand_name column to use, it may make the decision to use waii_similarity_search to pick the best value instead. 62 | With categorical columns, limit 1 will often be used with the similarity search function to pick the best value from the column for the lookup phrase: 'Kohls' 63 | 64 | When running this query, Waii will look through the values of brand_id and pick the best value to insert into the query. waii_similarity_search will be replaced with a `Values` clause with this value. 65 | 66 | 67 | #### Similarity Search Index without LLM Rerank 68 | 69 | Given 70 | - table `documents` with columns: document_id, summary, title 71 | - similarity search index on the document_id column with summary as an additional meaning with `enable_llm_rerank` False 72 | 73 | The similarity search index will identify the document_id values with the summaries that match the phrase. By disabling `enable_llm_rerank` we want to use this index to do similarity search over the document corpus. 74 | 75 | For the user ask, "Which documents relate to the fall of Rome" 76 | 77 | ```sql 78 | select title 79 | from documents 80 | join table(waii_similarity_search(document_id, 'fall of Rome')) as matched_documents(val, score) 81 | on document_id = matched_documents.val 82 | ``` 83 | 84 | Waii will use waii_similarity_search to find the documents which best match the subject: fall of Rome. 85 | The similarity search index is set up with summary as an additional meaning for each document_id, this will be used for the embedding match. 86 | waii_similarity_search will return all matched documents, obeying the similarity search index properties described above. -------------------------------------------------------------------------------- /tests/waii_basemodel_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | import pytest 19 | from typing import Optional, List 20 | 21 | from waii_sdk_py.my_pydantic import WaiiBaseModel 22 | 23 | class ABaseModel(WaiiBaseModel): 24 | attr_a: int 25 | attr_b: float 26 | attr_c: str 27 | 28 | class BBaseModel(WaiiBaseModel): 29 | attr_a: ABaseModel 30 | attr_b: int 31 | 32 | class CBaseModel(BBaseModel): 33 | attr_c: int 34 | attr_d: ABaseModel 35 | 36 | class DBaseModel(WaiiBaseModel): 37 | attr_a: int 38 | attr_b: List[ABaseModel] 39 | attr_c: Optional[CBaseModel] = None 40 | 41 | class EBaseModel(WaiiBaseModel): 42 | attr_a: dict[str, ABaseModel] = {} 43 | 44 | 45 | class TestWaiiBaseModel(unittest.TestCase): 46 | def test_extra_field(self): 47 | # allows to add unknown attributes at init 48 | model_a = ABaseModel(attr_a=2, attr_b=3.3, attr_c='val', attr_unknown='unknown') 49 | with pytest.raises(ValueError): 50 | # should raise exception only when check_extra_fields() is called 51 | model_a.check_extra_fields() 52 | 53 | model_b = ABaseModel(attr_a=2, attr_b=3.3, attr_c='val') 54 | # should not raise exception 55 | model_b.check_extra_fields() 56 | 57 | def test_extra_field_recursive(self): 58 | model_a = ABaseModel(attr_a=2, attr_b=3.3, attr_c='val', attr_unknown='unknown') 59 | model_b = BBaseModel(attr_a=model_a, attr_b=3) 60 | 61 | # both models should raise exceptions 62 | with pytest.raises(ValueError): 63 | model_b.check_extra_fields() 64 | with pytest.raises(ValueError): 65 | model_a.check_extra_fields() 66 | 67 | def test_extra_field_kwargs(self): 68 | # kwargs 69 | model_c = ABaseModel(**{'attr_a':2, 'attr_b':3.3, 'attr_c':'val', 'attr_d':'new_val'}) 70 | with pytest.raises(ValueError): 71 | model_c.check_extra_fields() 72 | 73 | def test_extra_field_subclass(self): 74 | model_a = ABaseModel(attr_a=2, attr_b=3.3, attr_c='val') 75 | # invalid model 76 | model_a_inv = ABaseModel(attr_a=2, attr_b=3.3, attr_c='val', unknown='unknown') 77 | model_c = CBaseModel(attr_a=model_a, attr_b=3, attr_c=4, attr_d=model_a_inv) 78 | 79 | with pytest.raises(ValueError): 80 | model_c.check_extra_fields() 81 | 82 | 83 | model_c = CBaseModel(attr_a=model_a, attr_b=3, attr_c=4, attr_d=model_a) 84 | # should not raise exception since all attributes are recursively correct 85 | model_c.check_extra_fields() 86 | 87 | def test_extra_field_list(self): 88 | model_d = DBaseModel( 89 | attr_a=2, 90 | attr_b=[ 91 | ABaseModel(attr_a=2, attr_b=3.3, attr_c='a'), 92 | # invalid 93 | ABaseModel(attr_a=2, attr_b=3.3, attr_c='val', unknown='unknown') 94 | ], 95 | ) 96 | 97 | with pytest.raises(ValueError): 98 | model_d.check_extra_fields() 99 | 100 | def test_extra_field_dict(self): 101 | model_e = EBaseModel( 102 | attr_a={'a': ABaseModel(attr_a=2, attr_b=2.2, attr_c='hello', attr_d=4)} 103 | ) 104 | 105 | with pytest.raises(ValueError): 106 | model_e.check_extra_fields() 107 | 108 | model_e = EBaseModel( 109 | attr_a={'a': ABaseModel(attr_a=2, attr_b=2.2, attr_c='hello')} # no extra attribute, should not throw error 110 | ) 111 | model_e.check_extra_fields() 112 | 113 | def test_extra_field_optional(self): 114 | model_d = DBaseModel( 115 | attr_a=2, 116 | attr_b=[ 117 | ABaseModel(attr_a=2, attr_b=3.3, attr_c='a'), 118 | ABaseModel(attr_a=2, attr_b=3.3, attr_c='val') 119 | ], 120 | attr_c=CBaseModel( 121 | attr_a=ABaseModel(attr_a=2, attr_b=3.3, attr_c='a'), 122 | attr_b=2, attr_c=3, 123 | # invalid 124 | attr_d=ABaseModel(attr_a=2, attr_b=3.3, attr_c='v', unknown='unknown') 125 | ) 126 | ) 127 | 128 | with pytest.raises(ValueError): 129 | model_d.check_extra_fields() 130 | 131 | 132 | if __name__ == '__main__': 133 | unittest.main() -------------------------------------------------------------------------------- /tests/chat_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | from unittest.mock import Mock, patch 19 | from enum import Enum 20 | from typing import Optional, List 21 | from datetime import datetime 22 | 23 | from waii_sdk_py import WAII 24 | from waii_sdk_py.chat import Chat, ChatRequest, ChatResponse, ChatResponseData, ChatModule 25 | from waii_sdk_py.query import GeneratedQuery 26 | from waii_sdk_py.database import CatalogDefinition 27 | from waii_sdk_py.semantic_context import GetSemanticContextResponse 28 | from waii_sdk_py.chart import ChartGenerationResponse, ChartType 29 | 30 | class TestChat(unittest.TestCase): 31 | def setUp(self): 32 | WAII.initialize(url="http://localhost:9859/api/") 33 | result = WAII.Database.get_connections() 34 | self.result = result 35 | snowflake_connector = None 36 | for connector in result.connectors: 37 | if connector.db_type == "snowflake" and "MOVIE_DB" in connector.key: 38 | snowflake_connector = connector 39 | break 40 | 41 | WAII.Database.activate_connection(snowflake_connector.key) 42 | 43 | self.chat = Chat 44 | 45 | def test_chat_message_query(self): 46 | 47 | chat_request = ChatRequest( 48 | ask="Show me movies data", 49 | streaming=False, 50 | parent_uuid=None, 51 | chart_type=None, 52 | modules=[ChatModule.QUERY], 53 | module_limit_in_response=1 54 | ) 55 | 56 | result = self.chat.chat_message(chat_request) 57 | 58 | self.assertIsInstance(result, ChatResponse) 59 | self.assertIsNotNone(result.response_data) 60 | self.assertIsNotNone(result.response_data.query) 61 | self.assertTrue(result.is_new) 62 | self.assertIsNotNone(result.timestamp_ms) 63 | 64 | def test_chat_message_chart(self): 65 | chat_request = ChatRequest( 66 | ask="Show me a bar chart of movie revenue", 67 | streaming=False, 68 | parent_uuid=None, 69 | chart_type=ChartType.SUPERSET, 70 | modules=[ChatModule.CHART, ChatModule.QUERY, ChatModule.DATA], 71 | module_limit_in_response=2 72 | ) 73 | 74 | result = self.chat.chat_message(chat_request) 75 | 76 | self.assertIsInstance(result, ChatResponse) 77 | self.assertIsNotNone(result.response) 78 | self.assertIsNotNone(result.response_data) 79 | self.assertIsNotNone(result.response_data.chart) 80 | self.assertIsNotNone(result.response_data.query) 81 | def test_chat_message_semantic_context(self): 82 | chat_request = ChatRequest( 83 | ask="What tables are available for movie data?", 84 | streaming=False, 85 | parent_uuid=None, 86 | chart_type=None, 87 | module_limit_in_response=1 88 | ) 89 | 90 | result = self.chat.chat_message(chat_request) 91 | 92 | self.assertIsInstance(result, ChatResponse) 93 | self.assertIsNotNone(result.response) 94 | self.assertIsNotNone(result.response_data) 95 | self.assertIsNotNone(result.response_data.semantic_context) 96 | 97 | def test_chat_conversation(self): 98 | # Initial question 99 | chat_request1 = ChatRequest( 100 | ask="What's the rating of movies?", 101 | streaming=False, 102 | parent_uuid=None, 103 | modules=[ChatModule.QUERY, ChatModule.CHART, ChatModule.DATA], 104 | module_limit_in_response=2 105 | ) 106 | 107 | result1 = self.chat.chat_message(chat_request1) 108 | 109 | self.assertIsInstance(result1, ChatResponse) 110 | self.assertIsNotNone(result1.chat_uuid) 111 | 112 | # Follow-up question 113 | chat_request2 = ChatRequest( 114 | ask="Now show me this data in a bar chart", 115 | streaming=False, 116 | parent_uuid=result1.chat_uuid, 117 | chart_type=ChartType.SUPERSET, 118 | modules=[ChatModule.CHART, ChatModule.QUERY, ChatModule.DATA], 119 | module_limit_in_response=1 120 | ) 121 | 122 | result2 = self.chat.chat_message(chat_request2) 123 | 124 | self.assertIsInstance(result2, ChatResponse) 125 | self.assertIsNotNone(result2.response_data.chart) 126 | 127 | if __name__ == '__main__': 128 | unittest.main() -------------------------------------------------------------------------------- /doc/docs/superset-rendering.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: superset-rendering 3 | title: Superset chart 4 | --- 5 | 6 | ## Waii chat response -> Superset chart iframe 7 | 8 | the method render_chart takes chat_response as input 9 | 10 | 11 | ``` 12 | { 13 | "response": ..., 14 | "response_data": { 15 | "data": ..., 16 | "query": ..., 17 | "chart": ... 18 | } 19 | } 20 | } 21 | ``` 22 | 23 | and returns permalink which can be embedded as iframe 24 | 25 | E.g, http://51.8.188.37:8088/superset/explore/p/aqYmx10WDKd/?standalone=True 26 | 27 | 28 | ``` 29 | from supersetapiclient.client import SupersetClient 30 | import os 31 | import uuid 32 | import json 33 | import re 34 | 35 | # Allow insecure transport for development 36 | os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1' 37 | 38 | class SupersetChartRenderer: 39 | """ 40 | Singleton class for rendering charts in Superset by managing datasets and permalinks. 41 | Handles dataset creation, chart specifications, and permalink generation. 42 | 43 | Environment Variables Required: 44 | SUPERSET_USERNAME: Superset login username 45 | SUPERSET_PASSWORD: Superset login password 46 | SUPERSET_HOST: Superset host URL 47 | SUPERSET_DATABASE_ID: Target database ID in Superset 48 | SUPERSET_HOST_HTTPS: Optional HTTPS host URL for permalink generation 49 | """ 50 | _instance = None 51 | 52 | def __new__(cls): 53 | """ 54 | Creates or returns the singleton instance of SupersetChartRenderer. 55 | Initializes SupersetClient with credentials from environment variables. 56 | """ 57 | if cls._instance is None: 58 | cls._instance = super(SupersetChartRenderer, cls).__new__(cls) 59 | cls._instance.client = SupersetClient(username=os.getenv("SUPERSET_USERNAME"), 60 | password=os.getenv("SUPERSET_PASSWORD"), 61 | host=os.getenv("SUPERSET_HOST")) 62 | cls._instance.superset_database_id = os.getenv("SUPERSET_DATABASE_ID") 63 | cls._instance.link_map = {} 64 | cls._instance.query_id_to_dataset_id = {} 65 | cls._instance.current_dataset_id = None 66 | return cls._instance 67 | 68 | def create_dataset(self, database_id, sql_query): 69 | """ 70 | Creates a new dataset in Superset from SQL query. 71 | 72 | Args: 73 | database_id (str): ID of target Superset database 74 | sql_query (str): SQL query to create dataset 75 | 76 | Returns: 77 | str: ID of created dataset 78 | """ 79 | dataset_name = f"waii-dataset-{str(uuid.uuid4())}" 80 | dataset = self.client.post(f"{self.client.host}/api/v1/dataset/", json={ 81 | "database": database_id, 82 | "sql": sql_query, 83 | "table_name": dataset_name, 84 | "normalize_columns": True 85 | }) 86 | 87 | return dataset.json()["id"] 88 | 89 | def create_permalink(self, dataset_id, superset_specs): 90 | """ 91 | Creates a permalink for chart visualization in Superset. 92 | 93 | Args: 94 | dataset_id (str): ID of dataset to visualize 95 | superset_specs (dict): Chart specifications for Superset 96 | 97 | Returns: 98 | str: Permalink URL for chart visualization 99 | """ 100 | superset_specs["datasource"] = f"{dataset_id}__table" 101 | permlink_payload = { 102 | "formData": superset_specs 103 | } 104 | 105 | permlink = self.client.post(f"{self.client.host}/api/v1/explore/permalink", json=permlink_payload) 106 | 107 | if os.environ['SUPERSET_HOST_HTTPS']: 108 | return f"{os.environ['SUPERSET_HOST_HTTPS']}/superset/explore/p/{permlink.json()['key']}/?standalone=True" 109 | else: 110 | return f"{os.environ['SUPERSET_HOST']}/superset/explore/p/{permlink.json()['key']}/?standalone=True" 111 | 112 | def render_chart(self, chat_response): 113 | """ 114 | Renders a chart based on chat response data. 115 | Creates dataset if query exists and generates visualization permalink. 116 | 117 | Args: 118 | chat_response: Object containing query and chart specifications 119 | 120 | Returns: 121 | str: Permalink URL for rendered chart 122 | 123 | Raises: 124 | Exception: If no query exists and no current dataset ID is set 125 | """ 126 | if not chat_response.response_data.query and not self.current_dataset_id: 127 | raise Exception("Cannot plot") 128 | 129 | if chat_response.response_data.query: 130 | dataset_id = self.create_dataset(self.superset_database_id, chat_response.response_data.query.query) 131 | self.current_dataset_id = dataset_id 132 | link = self.create_permalink(self.current_dataset_id, chat_response.response_data.chart.chart_spec.superset_specs) 133 | 134 | return link 135 | ``` 136 | -------------------------------------------------------------------------------- /tests/knowledge_graph_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | from waii_sdk_py import WAII 19 | from waii_sdk_py.kg import GetKnowledgeGraphRequest 20 | 21 | 22 | class TestKnowledgeGraph(unittest.TestCase): 23 | def setUp(self): 24 | WAII.initialize(url="http://localhost:9859/api/") 25 | result = WAII.Database.get_connections() 26 | self.result = result 27 | # Activate the first available connection 28 | if result.connectors: 29 | WAII.Database.activate_connection(result.connectors[0].key) 30 | 31 | def test_get_knowledge_graph_smoke(self): 32 | """ 33 | Smoke test to ensure knowledge graph API returns a valid response with nodes and edges 34 | """ 35 | # Create a simple request 36 | request = GetKnowledgeGraphRequest(ask="Show me table relationships") 37 | 38 | # Call the knowledge graph API 39 | response = WAII.knowledge_graph.get_knowledge_graph(request) 40 | 41 | # Basic assertions - verify we got a response 42 | self.assertIsNotNone(response, "Response should not be None") 43 | self.assertIsNotNone(response.graph, "Graph should not be None") 44 | 45 | # Verify the graph has nodes and edges 46 | self.assertIsNotNone(response.graph.nodes, "Graph should have nodes") 47 | self.assertIsNotNone(response.graph.edges, "Graph should have edges") 48 | 49 | # Verify the graph is not empty (has at least some nodes) 50 | self.assertGreater(len(response.graph.nodes), 0, "Graph should have at least one node") 51 | 52 | # Verify basic structure of nodes (if any exist) 53 | if response.graph.nodes: 54 | first_node = response.graph.nodes[0] 55 | self.assertIsNotNone(first_node.id, "Node should have an id") 56 | self.assertIsNotNone(first_node.display_name, "Node should have a display_name") 57 | self.assertIsNotNone(first_node.entity_type, "Node should have an entity_type") 58 | self.assertIsNotNone(first_node.entity, "Node should have an entity") 59 | 60 | # Print summary information about the knowledge graph 61 | print("\n" + "="*60) 62 | print("KNOWLEDGE GRAPH SUMMARY") 63 | print("="*60) 64 | print(f"Total nodes: {len(response.graph.nodes)}") 65 | print(f"Total edges: {len(response.graph.edges)}") 66 | 67 | # Count nodes by type 68 | node_types = {} 69 | for node in response.graph.nodes: 70 | node_type = node.entity_type 71 | node_types[node_type] = node_types.get(node_type, 0) + 1 72 | 73 | print("\nNode types distribution:") 74 | for node_type, count in sorted(node_types.items()): 75 | print(f" - {node_type}: {count}") 76 | 77 | # Print sample nodes (first 5 of each type) 78 | print("\nSample nodes by type:") 79 | for node_type in sorted(node_types.keys()): 80 | print(f"\n {node_type.upper()}:") 81 | nodes_of_type = [n for n in response.graph.nodes if n.entity_type == node_type][:5] 82 | for node in nodes_of_type: 83 | print(f" - {node.display_name} (id: {node.id})") 84 | 85 | # Count edges by type 86 | edge_types = {} 87 | for edge in response.graph.edges: 88 | edge_type = edge.edge_type 89 | edge_types[edge_type] = edge_types.get(edge_type, 0) + 1 90 | 91 | print("\nEdge types distribution:") 92 | for edge_type, count in sorted(edge_types.items()): 93 | print(f" - {edge_type}: {count}") 94 | 95 | # Print sample relationships (first 5) 96 | if response.graph.edges: 97 | print("\nSample relationships (first 5):") 98 | for edge in response.graph.edges[:5]: 99 | source_node = next((n for n in response.graph.nodes if n.id == edge.source_id), None) 100 | target_node = next((n for n in response.graph.nodes if n.id == edge.target_id), None) 101 | if source_node and target_node: 102 | direction = "->" if edge.directed else "<->" 103 | print(f" - {source_node.display_name} {direction} {target_node.display_name} ({edge.edge_type})") 104 | if edge.description: 105 | print(f" Description: {edge.description}") 106 | 107 | print("="*60 + "\n") 108 | 109 | 110 | if __name__ == '__main__': 111 | unittest.main() -------------------------------------------------------------------------------- /waii_sdk_py/history/history.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | from typing import List, Optional 19 | 20 | from ..chat import ChatRequest, ChatResponse 21 | from ..my_pydantic import WaiiBaseModel 22 | from ..query import GeneratedQuery, QueryGenerationRequest 23 | from ..chart import ChartGenerationRequest, ChartGenerationResponse 24 | from waii_sdk_py.utils import wrap_methods_with_async 25 | from ..waii_http_client import WaiiHttpClient 26 | 27 | LIST_ENDPOINT = "get-generated-query-history" 28 | GET_ENDPOINT = "get-history" 29 | 30 | 31 | class GeneratedHistoryEntryBase(WaiiBaseModel): 32 | history_type: str 33 | # milliseconds since epoch 34 | timestamp_ms: Optional[int] 35 | 36 | 37 | class GeneratedChartHistoryEntry(GeneratedHistoryEntryBase): 38 | request: Optional[ChartGenerationRequest] 39 | response: Optional[ChartGenerationResponse] 40 | 41 | 42 | class GeneratedChatHistoryEntry(GeneratedHistoryEntryBase): 43 | request: Optional[ChatRequest] 44 | response: Optional[ChatResponse] 45 | 46 | 47 | class GeneratedHistoryEntryType(str, Enum): 48 | query = "query" 49 | chart = "chart" 50 | chat = "chat" 51 | 52 | 53 | class GeneratedQueryHistoryEntry(GeneratedHistoryEntryBase): 54 | query: Optional[GeneratedQuery] = None 55 | request: Optional[QueryGenerationRequest] = None 56 | 57 | 58 | class GetGeneratedQueryHistoryRequest(WaiiBaseModel): 59 | limit: Optional[int] = None 60 | offset: Optional[int] = None 61 | 62 | 63 | class GetGeneratedQueryHistoryResponse(WaiiBaseModel): 64 | history: Optional[List[GeneratedQueryHistoryEntry]] = None 65 | 66 | 67 | class SortOrder(str, Enum): 68 | asc = "asc" 69 | desc = "desc" 70 | 71 | 72 | class GetHistoryResponse: 73 | def __init__(self, objs): 74 | self.history = [] 75 | if 'history' not in objs: 76 | raise Exception(f"history is required, but not found in the response, {objs}") 77 | 78 | objs = objs['history'] 79 | for h in objs: 80 | if 'history_type' not in h: 81 | raise Exception(f"history_type is required, but not found in the response, {h}") 82 | 83 | history_type = h['history_type'] 84 | 85 | if history_type == GeneratedHistoryEntryType.query: 86 | self.history.append(GeneratedQueryHistoryEntry(**h)) 87 | elif history_type == GeneratedHistoryEntryType.chart: 88 | self.history.append(GeneratedChartHistoryEntry(**h)) 89 | elif history_type == GeneratedHistoryEntryType.chat: 90 | self.history.append(GeneratedChatHistoryEntry(**h)) 91 | 92 | class GetHistoryRequest(WaiiBaseModel): 93 | # by default include query for backward compatibility 94 | included_types: Optional[List[GeneratedHistoryEntryType]] = [GeneratedHistoryEntryType.query, 95 | GeneratedHistoryEntryType.chart, 96 | GeneratedHistoryEntryType.chat] 97 | 98 | # for pagination 99 | limit: Optional[int] = 1000 100 | offset: Optional[int] = 0 101 | 102 | # latest first (default) 103 | timestamp_sort_order: Optional[SortOrder] = SortOrder.desc 104 | 105 | # filter by uuid of the entry 106 | uuid_filter: Optional[str] = None 107 | 108 | # filter by the liked query flag, by default it will include both liked and unliked queries. 109 | # when this is set to not None, you must only include query as included_types. 110 | liked_query_filter: Optional[bool] = None 111 | 112 | 113 | class HistoryImpl: 114 | def __init__(self, http_client: WaiiHttpClient): 115 | self.http_client = http_client 116 | 117 | # this is deprecated, use get() instead 118 | def list( 119 | self, 120 | params: Optional[GetGeneratedQueryHistoryRequest] = None, 121 | ) -> GetGeneratedQueryHistoryResponse: 122 | if params == None: 123 | params = GetGeneratedQueryHistoryRequest() 124 | print("This method is deprecated, use get() instead") 125 | return self.http_client.common_fetch( 126 | LIST_ENDPOINT, params, GetGeneratedQueryHistoryResponse 127 | ) 128 | 129 | def get( 130 | self, 131 | params: Optional[GetHistoryRequest] = None, 132 | ) -> GetHistoryResponse: 133 | if params == None: 134 | params = GetHistoryRequest() 135 | objs = self.http_client.common_fetch( 136 | GET_ENDPOINT, params, ret_json=True 137 | ) 138 | return GetHistoryResponse(objs) 139 | 140 | 141 | class AsyncHistoryImpl: 142 | def __init__(self, http_client: WaiiHttpClient): 143 | self._history_impl = HistoryImpl(http_client) 144 | wrap_methods_with_async(self._history_impl, self) 145 | 146 | 147 | History = HistoryImpl(WaiiHttpClient.get_instance()) 148 | -------------------------------------------------------------------------------- /tests/common_test_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import time 18 | 19 | from tests import constants 20 | from waii_sdk_py import WAII 21 | from waii_sdk_py.database import ModifyDBConnectionRequest, DBConnection 22 | from waii_sdk_py.user import ListOrganizationsRequest, ListTenantsRequest, ListUsersRequest 23 | from waii_sdk_py.waii_sdk_py import Waii 24 | 25 | 26 | def connect_db(db_conn_to_add: DBConnection, impersonate_user=None): 27 | WAII.initialize(url="http://localhost:9859/api/", verbose=True) 28 | 29 | if impersonate_user: 30 | WAII.set_impersonate_user(impersonate_user) 31 | else: 32 | WAII.clear_impersonation() 33 | 34 | # add the database connection 35 | try: 36 | WAII.Database.modify_connections( 37 | ModifyDBConnectionRequest(updated=[db_conn_to_add]) 38 | ).connectors 39 | except Exception as e: 40 | print(e) 41 | print("""In order to run the test, you need to run Waii locally, and setup the following 42 | sudo -u $USER psql postgres 43 | CREATE ROLE waii WITH LOGIN CREATEDB PASSWORD 'password'; 44 | CREATE DATABASE waii_sdk_test; 45 | GRANT ALL PRIVILEGES ON DATABASE waii_sdk_test TO waii; 46 | """) 47 | 48 | max_wait_time = 60 49 | 50 | # wait till the index is finished 51 | while True: 52 | try: 53 | WAII.Database.activate_connection(db_conn_to_add.key) 54 | catalogs = WAII.Database.get_catalogs() 55 | if catalogs.catalogs and len(catalogs.catalogs) > 0: 56 | break 57 | except Exception as e: 58 | print(e) 59 | max_wait_time -= 1 60 | if max_wait_time == 0: 61 | raise Exception("Cannot get the catalog") 62 | time.sleep(1) 63 | 64 | WAII.Database.activate_connection(db_conn_to_add.key) 65 | WAII.clear_impersonation() 66 | 67 | 68 | def load_db_conn1(): 69 | port = constants.pg_port 70 | db_connection = DBConnection( 71 | key=f"postgresql://waii@localhost:{port}/waii_sdk_test", # will be generated by the system 72 | db_type="postgresql", 73 | host="localhost", 74 | port=port, 75 | database="waii_sdk_test", 76 | username="waii", 77 | password="password" 78 | ) 79 | return db_connection 80 | 81 | 82 | def load_db_conn2(): 83 | port = constants.pg_port 84 | db_connection = DBConnection( 85 | key=f"postgresql://waii@localhost:{port}/waii_sdk_test2", # will be generated by the system 86 | db_type="postgresql", 87 | host="localhost", 88 | port=port, 89 | database="waii_sdk_test2", 90 | username="waii", 91 | password="password" 92 | ) 93 | return db_connection 94 | 95 | 96 | def check_table_existence(waii: Waii, table_name, should_exist, max_wait_time=30, check_interval=1): 97 | """ 98 | Check if a table exists or not in the database. 99 | 100 | :param table_name: Name of the table to check 101 | :param should_exist: Boolean, True if table should exist, False if it should not exist 102 | :param max_wait_time: Maximum time to wait in seconds (default 30) 103 | :param check_interval: Time to wait between checks in seconds (default 1) 104 | :return: Boolean, True if the condition is met, False if timeout occurs 105 | """ 106 | end_time = time.time() + max_wait_time 107 | 108 | while time.time() < end_time: 109 | # Refresh the database connection 110 | waii.database.refresh_db_connection() 111 | 112 | # Check for table existence 113 | result = waii.database.get_catalogs() 114 | table_exists = any( 115 | table.name.table_name.lower() == table_name.lower() 116 | for catalog in result.catalogs 117 | for schema in catalog.schemas 118 | for table in schema.tables 119 | ) 120 | 121 | # If the condition is met, return True 122 | if table_exists == should_exist: 123 | return True 124 | 125 | # Wait before next check 126 | time.sleep(check_interval) 127 | 128 | # If we've reached this point, we've timed out 129 | return False 130 | 131 | 132 | def org_exists(waii: Waii, org_id: str) -> bool: 133 | orgs = waii.user.list_orgs(ListOrganizationsRequest()).organizations 134 | for org in orgs: 135 | if org.id == org_id: 136 | return True 137 | return False 138 | 139 | 140 | def tenant_exists(waii: Waii, org_id: str, tenant_id: str) -> bool: 141 | tenants = waii.user.list_tenants(ListTenantsRequest(lookup_org_id=org_id)).tenants 142 | for tenant in tenants: 143 | if tenant.id == tenant_id: 144 | return True 145 | return False 146 | 147 | 148 | def user_exists(waii: Waii, org_id: str, user_id: str) -> bool: 149 | users = waii.user.list_users(ListUsersRequest(lookup_org_id=org_id)).users 150 | for user in users: 151 | if user.id == user_id: 152 | return True 153 | return False 154 | -------------------------------------------------------------------------------- /tests/chat_v2_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | from unittest.mock import Mock, patch 19 | from enum import Enum 20 | from typing import Optional, List 21 | from datetime import datetime 22 | 23 | from waii_sdk_py import WAII 24 | from waii_sdk_py.chat import Chat, ChatRequest, ChatResponse, ChatResponseData, ChatModule 25 | from waii_sdk_py.query import GeneratedQuery 26 | from waii_sdk_py.database import CatalogDefinition 27 | from waii_sdk_py.semantic_context import GetSemanticContextResponse 28 | from waii_sdk_py.chart import ChartGenerationResponse, ChartType 29 | 30 | class TestChatV2(unittest.TestCase): 31 | def setUp(self): 32 | WAII.initialize(url="http://localhost:9859/api/") 33 | result = WAII.Database.get_connections() 34 | self.result = result 35 | snowflake_connector = None 36 | for connector in result.connectors: 37 | if connector.db_type == "snowflake": 38 | snowflake_connector = connector 39 | break 40 | 41 | WAII.Database.activate_connection(snowflake_connector.key) 42 | 43 | self.chat = Chat 44 | 45 | def test_chat_message_query(self): 46 | 47 | chat_request = ChatRequest( 48 | ask="Show me movies data", 49 | streaming=False, 50 | parent_uuid=None, 51 | chart_type=None, 52 | modules=[ChatModule.QUERY], 53 | module_limit_in_response=1 54 | ) 55 | 56 | result = self.chat.chat_message(chat_request) 57 | 58 | self.assertIsInstance(result, ChatResponse) 59 | self.assertIsNotNone(result.response_data) 60 | self.assertIsNotNone(result.response_data.query) 61 | self.assertEqual(result.response_selected_fields, [ChatModule.QUERY]) 62 | self.assertTrue(result.is_new) 63 | self.assertIsNotNone(result.timestamp_ms) 64 | 65 | def test_chat_message_chart(self): 66 | chat_request = ChatRequest( 67 | ask="Show me a bar chart for top genres for movies", 68 | streaming=False, 69 | parent_uuid=None, 70 | chart_type=ChartType.SUPERSET, 71 | modules=[ChatModule.CHART, ChatModule.QUERY, ChatModule.DATA], 72 | module_limit_in_response=2 73 | ) 74 | 75 | result = self.chat.chat_message(chat_request) 76 | 77 | self.assertIsInstance(result, ChatResponse) 78 | self.assertIsNotNone(result.response) 79 | self.assertIsNotNone(result.response_data) 80 | self.assertIsNotNone(result.response_data.chart) 81 | self.assertIsNotNone(result.response_data.query) 82 | self.assertIn(ChatModule.CHART, result.response_selected_fields) 83 | def test_chat_message_semantic_context(self): 84 | chat_request = ChatRequest( 85 | ask="Give top genres for movies", 86 | streaming=False, 87 | parent_uuid=None, 88 | chart_type=None, 89 | modules=[ChatModule.DATA, ChatModule.QUERY], 90 | module_limit_in_response=1 91 | ) 92 | 93 | result = self.chat.chat_message(chat_request) 94 | 95 | self.assertIsInstance(result, ChatResponse) 96 | self.assertIsNotNone(result.response) 97 | self.assertIsNotNone(result.response_data) 98 | self.assertIsNotNone(result.response_data.data) 99 | self.assertIn(ChatModule.DATA, result.response_selected_fields) 100 | 101 | def test_chat_message_tables(self): 102 | chat_request = ChatRequest( 103 | ask="Show me the schema for the movies table", 104 | streaming=False, 105 | parent_uuid=None, 106 | chart_type=None, 107 | modules=[ChatModule.TABLES], 108 | module_limit_in_response=1 109 | ) 110 | 111 | result = self.chat.chat_message(chat_request) 112 | 113 | self.assertIsInstance(result, ChatResponse) 114 | self.assertIsNotNone(result.response) 115 | self.assertIsNotNone(result.response_data) 116 | self.assertIsNotNone(result.response_data.tables) 117 | 118 | def test_chat_conversation(self): 119 | # Initial question 120 | chat_request1 = ChatRequest( 121 | ask="What's the rating of movies?", 122 | streaming=False, 123 | parent_uuid=None, 124 | modules=[ChatModule.QUERY, ChatModule.CHART, ChatModule.DATA], 125 | module_limit_in_response=2 126 | ) 127 | 128 | result1 = self.chat.chat_message(chat_request1) 129 | 130 | self.assertIsInstance(result1, ChatResponse) 131 | self.assertIsNotNone(result1.chat_uuid) 132 | 133 | # Follow-up question 134 | chat_request2 = ChatRequest( 135 | ask="Now show me this data in a bar chart", 136 | streaming=False, 137 | parent_uuid=result1.chat_uuid, 138 | chart_type=ChartType.SUPERSET, 139 | modules=[ChatModule.CHART, ChatModule.DATA, ChatModule.QUERY], 140 | module_limit_in_response=1 141 | ) 142 | 143 | result2 = self.chat.chat_message(chat_request2) 144 | 145 | self.assertIsInstance(result2, ChatResponse) 146 | self.assertIsNotNone(result2.response_data.chart) 147 | self.assertIn(ChatModule.CHART, result2.response_selected_fields) 148 | 149 | if __name__ == '__main__': 150 | unittest.main() -------------------------------------------------------------------------------- /doc/docs/access-rules-module.md: -------------------------------------------------------------------------------- 1 | --- 2 | id: access-rule-module 3 | title: Table Access Rules 4 | --- 5 | 6 | 7 | The `Access Rules` module contains methods related to creating access rules to ensure secure access to data for all users 8 | 9 | **Initialization & Imports** 10 | ```python 11 | from waii_sdk_py import WAII 12 | from waii_sdk_py.chat import * 13 | from waii_sdk_py.query import * 14 | from waii_sdk_py.database import * 15 | from waii_sdk_py.semantic_context import * 16 | from waii_sdk_py.chart import * 17 | from waii_sdk_py.history import * 18 | from waii_sdk_py.access_rules import * 19 | WAII.initialize(url="https://your-waii-instance/api/", api_key="your-api-key") 20 | ``` 21 | 22 | Currently, we support access rules scoped to Tables 23 | 24 | ### TableAccessRule 25 | 26 | A TableAccessRule object has the following fields 27 | - `id`: An optional id for the rule, Waii will fill this in if left blank. This must be unique for every rule 28 | - `name`: A name for help identifying the rule 29 | - `table`: A fully qualified `TableName` object for the table that the rule is scoped to 30 | - `org_id`: The organization id that the rule applies to 31 | - `tenant_id`: The tenant id that the rule applies to, can be `*` if it applies to every tenant in the organization 32 | - `user_id`: The user id that the rule applies to, can be `*` if it applies to every user in the tenant 33 | - `type`: One of `TableAccessRuleType` 34 | - `block`: Block all access to this table for the corresponding users 35 | - `filter`: Anytime this table is referenced in a query, protect access to this table with the expression used as the `WHERE` clause to filter the query 36 | - `expression`: An optional expression used as the `WHERE` clause for a filter table access rule. Can contain variables using `{variable_name}` 37 | 38 | Variables used in access rules will be resolved using variables in a user's `User`, `Tenant`, and `Organization` 39 | We also have 5 built-in variables that can be used in any access rule (user defined variables take precedence): 40 | - `org_id`: The id of the `Organization` that the current user is a part of 41 | - `tenant_id`: The id of the `Tenant` that the current user is a part of 42 | - `user_id`: The current `User` id 43 | - `roles`: An array of the roles for which the current user owns 44 | - `permissions`: An array of permissions which the current user contains 45 | 46 | ### Examples 47 | Expression: 48 | ```python 49 | table = TableName(table_name="T", schema_name="S", database_name="D") 50 | expression = col_a > 10 51 | ``` 52 | Any query that queries table T will refer to the results of the following CTE instead of T 53 | ```python 54 | with _access_controlled_t as ( 55 | SELECT * 56 | FROM t 57 | WHERE col_a > 10 58 | ) 59 | ``` 60 | 61 | ```python 62 | table = TableName(table_name="T", schema_name="S", database_name="D") 63 | expression = ( col_b = {user_id} AND col_c IN ({permissions}) ) 64 | ``` 65 | Any query that queries table T will refer to the results of the following CTE instead of T 66 | ```python 67 | with _access_controlled_t as ( 68 | SELECT * 69 | FROM t 70 | WHERE col_b = '' AND col_c IN ('array of current user permissions') 71 | ) 72 | ``` 73 | 74 | ```python 75 | table = TableName(table_name="T", schema_name="S", database_name="D") 76 | expression = col_d IN (SELECT col_d FROM d.s.t2 WHERE t2.col_e = t1.col_e) 77 | ``` 78 | Any query that queries table T will refer to the results of the following CTE instead of T 79 | ```python 80 | with _access_controlled_t as ( 81 | SELECT * 82 | FROM t 83 | WHERE col_d IN (SELECT col_d FROM d.s.t2 WHERE t2.col_e = t1.col_e) 84 | ) 85 | ``` 86 | 87 | ### Rule Management 88 | - There can only be one rule per (org_id, tenant_id, user_id) per table 89 | - Each rule must have a unique id 90 | - Multiple rules can be scoped to a table, the rule with the tightest scope for a user will be the one that is enforced within the query 91 | 92 | ## Methods 93 | 94 | ### UpdateTableAccessRules 95 | 96 | ```python 97 | WAII.access_rules.update_table_access_rules(params: UpdateTableAccessRuleRequest) 98 | ``` 99 | 100 | This method creates or updates the access rule for a table based on the parameters 101 | 102 | Request fields: 103 | - `rules`: A list of table access rules to update 104 | 105 | Each rule will be validated and rewritten before any of the rules are saved. An error will be thrown if 106 | - the `table` is invalid 107 | - the `expression` or the derived CTE are invalid (variables will be filled in for any applicable user) 108 | - an existing rule has the same combination of (org_id, tenant_id, user_id) and table, but a different id 109 | 110 | A rule will overwrite an existing rule if it has the same (org_id, tenant_id, user_id), table, and id. 111 | Otherwise, a rule with the same id as an existing rule will be rejected 112 | 113 | ### DeleteTableAccessRules 114 | ```python 115 | WAII.access_rules.remove_table_access_rules(params: RemoveTableAccessRuleRequest) 116 | ``` 117 | Request fields: 118 | - `rules`: A list of rule ids to remove 119 | 120 | ### ListTableAccessRules 121 | ```python 122 | WAII.access_rules.list_table_access_rules(params: ListTableAccessRuleRequest) -> ListTableAccessRuleResponse 123 | ``` 124 | 125 | This method will list the active table access rules that match the fields in the request. 126 | Request Fields 127 | - `table`: Optional argument, a fully qualified table name. All rules scoped to this table will be listed in the response 128 | - `ids`: Optional argument, a list of access rule ids. All rules matching these ids will be returned in the response 129 | - `lookup_user`: Optional argument, a User object for which to get the applicable rules for. If provided, all rules returned will be the access rules that queries from this user will follow. 130 | 131 | If these arguments are left empty, all active access rules within the scope will be returned. 132 | Any combination of these arguments can be filled in, the returned rules will match all filled in arguments. 133 | 134 | Response fields 135 | - `rules`: The list of rules that match the argments 136 | 137 | -------------------------------------------------------------------------------- /waii_sdk_py/semantic_context/semantic_context.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from typing import List, Optional, Literal 18 | 19 | from ..common import LLMBasedRequest, CommonRequest, CommonResponse 20 | from ..database import SearchContext 21 | from ..my_pydantic import WaiiBaseModel 22 | from waii_sdk_py.utils import wrap_methods_with_async 23 | from ..waii_http_client import WaiiHttpClient 24 | 25 | MODIFY_ENDPOINT = 'update-semantic-context' 26 | GET_ENDPOINT = 'get-semantic-context' 27 | ENABLE_ENDPOINT = 'enable-semantic-context' 28 | DISABLE_ENDPOINT = 'disable-semantic-context' 29 | 30 | 31 | class SemanticStatementWarning(WaiiBaseModel): 32 | message: str 33 | 34 | 35 | class SemanticStatement(WaiiBaseModel): 36 | entity_type: Literal['semantic_statement'] = 'semantic_statement' 37 | id: Optional[str] 38 | statement: str 39 | labels: Optional[List[str]] 40 | scope: Optional[str] 41 | 42 | # always include this statement in the context, if data_type is specified in statement, then always_include is True 43 | # when data_type matches 44 | always_include: Optional[bool] = True 45 | 46 | # Check the application of the rule in a second step after query gen when the rule is within scope 47 | critical: Optional[bool] = False 48 | 49 | # Search keys for this statement, if not specified, then use statement as search key 50 | # you can specify multiple search keys, for example, if you have a CVE doc, you can search by CVE number, or library 51 | # name, etc. 52 | lookup_summaries: Optional[List[str]] 53 | 54 | # extract prompt from the statement, if not specified, then use statement as extract prompt 55 | summarization_prompt: Optional[str] = None 56 | 57 | # Whether this semantic statement is enabled 58 | enabled: Optional[bool] = True 59 | 60 | # Warnings associated with this semantic statement 61 | warnings: Optional[List[SemanticStatementWarning]] = None 62 | 63 | # filters for user, tenant and org 64 | user_id: Optional[str] = '*' 65 | tenant_id: Optional[str] = '*' 66 | org_id: Optional[str] = '*' 67 | 68 | semantic_constraint: Optional[str] = None 69 | 70 | 71 | class ModifySemanticContextRequest(LLMBasedRequest): 72 | updated: Optional[List[SemanticStatement]] = None 73 | deleted: Optional[List[str]] = None 74 | 75 | 76 | class GetSemanticContextRequestFilter(WaiiBaseModel): 77 | # do we want to filter "always_include" rules or not? 78 | # - None: both 79 | # - True: only return rules with always_include=True 80 | # - False: only return rules with always_include=False 81 | always_include: Optional[bool] = None 82 | 83 | # Filter by labels, scope, statement. 84 | # They are connected by "AND", and we use substring match for them 85 | # Match is case insensitive 86 | labels: Optional[List[str]] = None # labels to filter the rules 87 | scope: Optional[str] = None # scope to filter the rules 88 | statement: Optional[str] = None # statement to filter the rules 89 | 90 | class ModifySemanticContextResponse(WaiiBaseModel): 91 | updated: Optional[List[SemanticStatement]] = None 92 | deleted: Optional[List[str]] = None 93 | 94 | 95 | class GetSemanticContextRequest(LLMBasedRequest): 96 | filter: GetSemanticContextRequestFilter = GetSemanticContextRequestFilter() 97 | offset: int = 0 98 | limit: int = 1000 99 | search_text: Optional[str] = None 100 | search_context: Optional[List[SearchContext]] = None 101 | 102 | 103 | class GetSemanticContextResponse(WaiiBaseModel): 104 | semantic_context: Optional[List[SemanticStatement]] = None 105 | available_statements: Optional[int] = 0 106 | 107 | 108 | class EnableSemanticContextRequest(CommonRequest): 109 | statement_ids: List[str] 110 | 111 | 112 | class EnableSemanticContextResponse(CommonResponse): 113 | statement_ids: List[str] # successfully enabled statement ids 114 | 115 | 116 | class DisableSemanticContextRequest(CommonRequest): 117 | statement_ids: List[str] 118 | 119 | 120 | class DisableSemanticContextResponse(CommonResponse): 121 | statement_ids: List[str] # successfully disabled statement ids 122 | 123 | 124 | class SemanticContextImpl: 125 | 126 | def __init__(self, http_client: WaiiHttpClient): 127 | self.http_client = http_client 128 | 129 | def modify_semantic_context(self, params: ModifySemanticContextRequest) -> ModifySemanticContextResponse: 130 | return self.http_client.common_fetch(MODIFY_ENDPOINT, params, ModifySemanticContextResponse) 131 | 132 | def get_semantic_context(self, params: Optional[GetSemanticContextRequest] = None) -> GetSemanticContextResponse: 133 | if params == None: 134 | params = GetSemanticContextRequest() 135 | return self.http_client.common_fetch(GET_ENDPOINT, params, GetSemanticContextResponse) 136 | 137 | def enable_semantic_context(self, params: EnableSemanticContextRequest) -> EnableSemanticContextResponse: 138 | return self.http_client.common_fetch(ENABLE_ENDPOINT, params, EnableSemanticContextResponse) 139 | 140 | def disable_semantic_context(self, params: DisableSemanticContextRequest) -> DisableSemanticContextResponse: 141 | return self.http_client.common_fetch(DISABLE_ENDPOINT, params, DisableSemanticContextResponse) 142 | 143 | 144 | class AsyncSemanticContextImpl: 145 | def __init__(self, http_client: WaiiHttpClient): 146 | self._semantic_context_impl = SemanticContextImpl(http_client) 147 | wrap_methods_with_async(self._semantic_context_impl, self) 148 | 149 | 150 | SemanticContext = SemanticContextImpl(WaiiHttpClient.get_instance()) -------------------------------------------------------------------------------- /tests/push_db_conn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import requests 18 | 19 | base_url = "http://localhost:9859" 20 | 21 | # API endpoint URL 22 | modify_db_conn_endpoint = f"{base_url}/api/update-db-connect-info" 23 | 24 | update_table_endpoint = f"{base_url}/api/update-table-definitions" 25 | generate_query_endpoint = f"{base_url}/api/generate-query" 26 | 27 | 28 | def add_db_conn(): 29 | # Data payload 30 | add_conn_data = { 31 | "updated": [ 32 | { 33 | "account_name": "", 34 | "warehouse": "", 35 | "database": "push_test", 36 | "host": "test_host", 37 | "db_type": "postgresql", 38 | "push": True, 39 | } 40 | ], 41 | "validate_before_save": True, 42 | } 43 | # Adding Connection 44 | add_conn_response = requests.post(modify_db_conn_endpoint, json=add_conn_data) 45 | 46 | print("Adding db conn response") 47 | return add_conn_response 48 | 49 | 50 | def remove_conn(): 51 | # Removing connection 52 | remove_conn_data = {"removed": ["postgresql://push/test_host/push_test"]} 53 | print("remove connection") 54 | remove_conn_response = requests.post(modify_db_conn_endpoint, json=remove_conn_data) 55 | return remove_conn_response 56 | 57 | 58 | def add_table(): 59 | # Updating tables 60 | table_update_data = { 61 | "updated_tables": [ 62 | { 63 | "name": { 64 | "table_name": "Album", 65 | "schema_name": "PUBLIC", 66 | "database_name": "push_test", 67 | }, 68 | "columns": [ 69 | { 70 | "name": "AlbumId", 71 | "type": "integer", 72 | "comment": None, 73 | "sample_values": None, 74 | "description": "Unique identifier for the album.", 75 | "description_update_source": "generated", 76 | }, 77 | { 78 | "name": "Title", 79 | "type": "character varying", 80 | "comment": None, 81 | "sample_values": None, 82 | "description": "Title of the album.", 83 | "description_update_source": "generated", 84 | }, 85 | { 86 | "name": "ArtistId", 87 | "type": "integer", 88 | "comment": None, 89 | "sample_values": None, 90 | "description": "Identifier for the artist of the album.", 91 | "description_update_source": "generated", 92 | }, 93 | ], 94 | "comment": None, 95 | "last_altered_time": 0, 96 | "refs": None, 97 | "constraints": [ 98 | { 99 | "source": "database", 100 | "table": { 101 | "table_name": "Artist", 102 | "schema_name": "PUBLIC", 103 | "database_name": "push_test", 104 | }, 105 | "cols": ["ArtistId"], 106 | "constraint_type": None, 107 | "src_table": { 108 | "table_name": "Album", 109 | "schema_name": "PUBLIC", 110 | "database_name": "push_test", 111 | }, 112 | "src_cols": ["ArtistId"], 113 | "comment": None, 114 | } 115 | ], 116 | "inferred_refs": [], 117 | "inferred_constraints": [], 118 | "description": 'The "Album" table contains information about music albums, including the album ID, artist ID, and title. Users can use this table to track and organize albums by various artists, making it easier to manage and access music collections.', 119 | } 120 | ], 121 | "scope": "postgresql://push/test_host/push_test", 122 | } 123 | updating_table_response = requests.post( 124 | update_table_endpoint, json=table_update_data 125 | ) 126 | # add table definition response 127 | print("Updating table def Response") 128 | return updating_table_response 129 | 130 | 131 | def gen_query(): 132 | gen_query_data = { 133 | "ask": "give me 5 albums name", 134 | "search_context": [], 135 | "tweak_history": [], 136 | "parent_uuid": None, 137 | "model": "Automatic", 138 | "scope": "postgresql://push/test_host/push_test", 139 | "org_id": "", 140 | "user_id": "", 141 | } 142 | 143 | gen_query_response = requests.post(generate_query_endpoint, json=gen_query_data) 144 | 145 | print("Generating query response") 146 | return gen_query_response 147 | 148 | 149 | def main(): 150 | add_db_resp = add_db_conn() 151 | print(add_db_resp.status_code) 152 | 153 | remove_db_resp = remove_conn() 154 | print(remove_db_resp.status_code) 155 | add_db_resp = add_db_conn() 156 | print(add_db_resp.status_code) 157 | update_table_resp = add_table() 158 | print(update_table_resp.status_code) 159 | gen_query_resp = gen_query() 160 | print(gen_query_resp.status_code) 161 | print(gen_query_resp.json()) 162 | 163 | 164 | if __name__ == "__main__": 165 | main() 166 | -------------------------------------------------------------------------------- /tests/semantic_layer_dump_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | import pytest 19 | from waii_sdk_py import WAII 20 | from waii_sdk_py.common import OperationStatus 21 | from waii_sdk_py.database import * 22 | from waii_sdk_py.query import * 23 | from waii_sdk_py.semantic_layer_dump import * 24 | 25 | class TestSemanticLayerDump(unittest.TestCase): 26 | def connect_db1(self): 27 | result = WAII.Database.get_connections() 28 | for connector in result.connectors: 29 | if connector.db_type == "postgresql" and connector.database == 'waii_sdk_test': 30 | db_conn1 = connector 31 | break 32 | 33 | WAII.Database.activate_connection(db_conn1.key) 34 | self.db_conn1 = db_conn1 35 | 36 | def connect_db2(self): 37 | result = WAII.Database.get_connections() 38 | for connector in result.connectors: 39 | if connector.db_type == "postgresql" and connector.database == 'waii_sdk_test_copy': 40 | db_conn2 = connector 41 | break 42 | 43 | WAII.Database.activate_connection(db_conn2.key) 44 | self.db_conn2 = db_conn2 45 | 46 | def setUp(self): 47 | WAII.initialize(url="http://localhost:9859/api/") 48 | self.connect_db2() 49 | self.connect_db1() 50 | 51 | def test_export_import_semantic_layer_dump(self): 52 | self.connect_db1() 53 | 54 | catalogs: Optional[List[CatalogDefinition]] = WAII.Database.get_catalogs(GetCatalogRequest()).catalogs 55 | source_table = catalogs[0].schemas[1].tables[0].name 56 | source_table_name = source_table.table_name 57 | source_old_description = catalogs[0].schemas[1].tables[0].description 58 | source_new_description = "test description" 59 | 60 | self.connect_db2() 61 | catalogs: Optional[List[CatalogDefinition]] = WAII.Database.get_catalogs(GetCatalogRequest()).catalogs 62 | for catalog in catalogs: 63 | for schema in catalog.schemas: 64 | for table in schema.tables: 65 | if table.name.table_name == source_table_name: 66 | target_table = table.name 67 | target_table_name = target_table.table_name 68 | target_old_description = table.description 69 | target_new_description = "test description 2" 70 | 71 | # add custom table description 72 | self.connect_db1() 73 | WAII.Database.update_table_description(UpdateTableDescriptionRequest( 74 | table_name=source_table, 75 | description=source_new_description 76 | )) 77 | 78 | # change the description of the target table 79 | self.connect_db2() 80 | WAII.Database.update_table_description(UpdateTableDescriptionRequest( 81 | table_name=target_table, 82 | description=target_new_description 83 | )) 84 | 85 | # export semantic layer dump 86 | self.connect_db1() 87 | resp = WAII.SemanticLayerDump.export_dump(ExportSemanticLayerDumpRequest( 88 | db_conn_key=self.db_conn1.key 89 | )) 90 | export_op_id = resp.op_id 91 | 92 | # wait for the export to complete 93 | while True: 94 | resp = WAII.SemanticLayerDump.export_dump_status(CheckOperationStatusRequest( 95 | op_id=export_op_id 96 | )) 97 | if resp.status == OperationStatus.SUCCEEDED: 98 | break 99 | elif resp.status == OperationStatus.FAILED: 100 | raise Exception(f"Export failed: {resp.info}") 101 | time.sleep(1) 102 | 103 | # check the description of the source table from the export 104 | configuration = resp.info 105 | for table in configuration['tables']: 106 | if table['name'] == source_table_name: 107 | self.assertEqual(table['description'], source_new_description) 108 | 109 | # import the semantic layer dump 110 | self.connect_db2() 111 | resp = WAII.SemanticLayerDump.import_dump(ImportSemanticLayerDumpRequest( 112 | db_conn_key=self.db_conn2.key, 113 | configuration=resp.info 114 | )) 115 | import_op_id = resp.op_id 116 | 117 | # wait for the import to complete 118 | while True: 119 | resp = WAII.SemanticLayerDump.import_dump_status(CheckOperationStatusRequest( 120 | op_id=import_op_id 121 | )) 122 | if resp.status == OperationStatus.SUCCEEDED: 123 | break 124 | elif resp.status == OperationStatus.FAILED: 125 | raise Exception(f"Import failed: {resp.info}") 126 | time.sleep(1) 127 | 128 | # check the table description 129 | self.connect_db2() 130 | catalogs: Optional[List[CatalogDefinition]] = WAII.Database.get_catalogs(GetCatalogRequest()).catalogs 131 | for catalog in catalogs: 132 | for schema in catalog.schemas: 133 | for table in schema.tables: 134 | if table.name.table_name == target_table_name: 135 | self.assertEqual(table.description, source_new_description) 136 | 137 | # restore the table descriptions 138 | self.connect_db1() 139 | WAII.Database.update_table_description(UpdateTableDescriptionRequest( 140 | table_name=source_table, 141 | description=source_old_description 142 | )) 143 | self.connect_db2() 144 | WAII.Database.update_table_description(UpdateTableDescriptionRequest( 145 | table_name=target_table, 146 | description=target_old_description 147 | )) 148 | -------------------------------------------------------------------------------- /tests/multi_tenant_client_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import unittest 18 | 19 | from tests.common_test_utils import load_db_conn1, load_db_conn2, connect_db 20 | from waii_sdk_py.database import Database 21 | from waii_sdk_py.query import * 22 | from waii_sdk_py.waii_sdk_py import Waii, WAII 23 | """ 24 | This class is to test https://doc.waii.ai/python/docs/multi-tenant-client-module 25 | """ 26 | class MultiTenantClientTest(unittest.TestCase): 27 | def setUp(self): 28 | self.movie_conn = load_db_conn1() 29 | self.chinook_conn = load_db_conn2() 30 | 31 | connect_db(self.movie_conn) 32 | connect_db(self.chinook_conn) 33 | 34 | movie_waii = Waii() 35 | chinook_waii = Waii() 36 | movie_waii.initialize(url="http://localhost:9859/api/") 37 | chinook_waii.initialize(url="http://localhost:9859/api/") 38 | result = movie_waii.database.get_connections() 39 | WAII.initialize(url="http://localhost:9859/api/") 40 | self.movie_waii = movie_waii 41 | self.chinook_waii = chinook_waii 42 | 43 | def test_movie_generate(self): 44 | self.movie_waii.database.activate_connection(self.movie_conn.key) 45 | params = QueryGenerationRequest( 46 | ask="Give me 5 movie names sorted by movie name" 47 | ) 48 | result = self.movie_waii.query.generate(params) 49 | self.assertIsInstance(result, GeneratedQuery) 50 | assert result.uuid is not None 51 | assert len(result.detailed_steps) > 0 52 | assert len(result.query) > 0 53 | assert "public.movies" in result.query.lower() 54 | assert len(result.tables) > 0 55 | 56 | params = LikeQueryRequest(query_uuid=result.uuid, liked=True) 57 | result = self.movie_waii.query.like(params) 58 | self.assertIsInstance(result, LikeQueryResponse) 59 | 60 | def test_movie_run(self): 61 | self.movie_waii.database.activate_connection(self.movie_conn.key) 62 | params = QueryGenerationRequest( 63 | ask="Give me 5 movie names sorted by movie name" 64 | ) 65 | result = self.movie_waii.query.generate(params) 66 | params = RunQueryRequest(query=result.query) 67 | result = self.movie_waii.query.run(params) 68 | self.assertIsInstance(result, GetQueryResultResponse) 69 | assert len(result.column_definitions) > 0 70 | assert "Incept" in str(result.rows[0]) 71 | 72 | def test_chinook_generate(self): 73 | self.chinook_waii.database.activate_connection(self.chinook_conn.key) 74 | params = QueryGenerationRequest(ask="Give me 5 album names sorted by name") 75 | result = self.chinook_waii.query.generate(params) 76 | self.assertIsInstance(result, GeneratedQuery) 77 | assert result.uuid is not None 78 | assert len(result.detailed_steps) > 0 79 | assert len(result.query) > 0 80 | assert "public" in result.query.lower() 81 | assert len(result.tables) > 0 82 | 83 | params = LikeQueryRequest(query_uuid=result.uuid, liked=True) 84 | result = self.chinook_waii.query.like(params) 85 | self.assertIsInstance(result, LikeQueryResponse) 86 | 87 | def test_chinook_run(self): 88 | self.chinook_waii.database.activate_connection(self.chinook_conn.key) 89 | params = QueryGenerationRequest(ask="Give me 5 album names sorted by name") 90 | result = self.chinook_waii.query.generate(params) 91 | params = RunQueryRequest(query=result.query) 92 | result = self.chinook_waii.query.run(params) 93 | self.assertIsInstance(result, GetQueryResultResponse) 94 | assert len(result.column_definitions) > 0 95 | assert "21" in str(result.rows[0]) 96 | 97 | def test_legacy_generate(self): 98 | WAII.Database.activate_connection(self.movie_conn.key) 99 | params = QueryGenerationRequest( 100 | ask="Give me 5 movie names sorted by movie name" 101 | ) 102 | result = WAII.Query.generate(params) 103 | self.assertIsInstance(result, GeneratedQuery) 104 | assert result.uuid is not None 105 | assert len(result.detailed_steps) > 0 106 | assert len(result.query) > 0 107 | assert "public.movies" in result.query.lower() 108 | assert len(result.tables) > 0 109 | 110 | params = LikeQueryRequest(query_uuid=result.uuid, liked=True) 111 | result = WAII.Query.like(params) 112 | self.assertIsInstance(result, LikeQueryResponse) 113 | 114 | def test_legacy_run(self): 115 | WAII.Database.activate_connection(self.movie_conn.key) 116 | params = QueryGenerationRequest( 117 | ask="Give me 5 movie names sorted by movie name" 118 | ) 119 | result = WAII.Query.generate(params) 120 | params = RunQueryRequest(query=result.query) 121 | result = WAII.Query.run(params) 122 | self.assertIsInstance(result, GetQueryResultResponse) 123 | assert len(result.column_definitions) > 0 124 | assert "Incepti" in str(result.rows[0]) 125 | 126 | def test_legacy_run_without_WAII(self): 127 | Database.activate_connection(self.movie_conn.key) 128 | params = QueryGenerationRequest( 129 | ask="Give me 5 movie names sorted by movie name" 130 | ) 131 | result = Query.generate(params) 132 | params = RunQueryRequest(query=result.query) 133 | result = Query.run(params) 134 | self.assertIsInstance(result, GetQueryResultResponse) 135 | assert len(result.column_definitions) > 0 136 | assert "Incep" in str(result.rows[0]) 137 | 138 | def test_activate_connection(self): 139 | Database.activate_connection(self.movie_conn.key) 140 | result = Database.get_activated_connection() 141 | assert (self.movie_conn.key == result) 142 | -------------------------------------------------------------------------------- /tests/async_client_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | import time 19 | import unittest 20 | from unittest import IsolatedAsyncioTestCase 21 | 22 | from tests.common_test_utils import load_db_conn1 23 | from waii_sdk_py.database import TableDefinition, TableName, ColumnDefinition, UpdateTableDefinitionRequest, \ 24 | ModifyDBConnectionRequest, TableReference 25 | from waii_sdk_py.query import QueryGenerationRequest, GeneratedQuery 26 | from waii_sdk_py.waii_sdk_py import AsyncWaii 27 | 28 | 29 | async def generate_task(method, params, task_name, task_indicators, start_times): 30 | start_times.append((task_name, time.time())) 31 | task_indicators.append(f"{task_name}_started") 32 | 33 | result = await method(params) if params else await method() 34 | 35 | task_indicators.append(f"{task_name}_finished") 36 | start_times.append((task_name, time.time())) 37 | 38 | return result 39 | 40 | 41 | class TestAsyncQuery(IsolatedAsyncioTestCase): 42 | 43 | async def asyncSetUp(self): 44 | async_waii_client = AsyncWaii() 45 | await async_waii_client.initialize(url="http://localhost:9859/api/") 46 | self.db_conn = load_db_conn1() 47 | self.async_waii_client = async_waii_client 48 | result = await self.async_waii_client.database.get_connections() 49 | self.result = result 50 | pg_connector = None 51 | for connector in result.connectors: 52 | if connector.db_type == "postgresql": 53 | pg_connector = connector 54 | break 55 | 56 | await self.async_waii_client.database.activate_connection(pg_connector.key) 57 | 58 | async def test_generate_query(self): 59 | params = QueryGenerationRequest(ask="List all tables in database") 60 | 61 | start = time.time() 62 | result = await self.async_waii_client.query.generate(params) 63 | isolated_execution_time = time.time() - start 64 | 65 | task_indicators = [] 66 | start_times = [] 67 | params = QueryGenerationRequest(ask="List all tables in database?.") 68 | task1 = asyncio.create_task(generate_task( 69 | self.async_waii_client.query.generate, params, "task1", task_indicators, start_times 70 | )) 71 | params = QueryGenerationRequest(ask="List all tables in database?...") 72 | task2 = asyncio.create_task(generate_task( 73 | self.async_waii_client.query.generate, params, "task2", task_indicators, start_times 74 | )) 75 | 76 | concurrent_start = time.time() 77 | await asyncio.gather(task1, task2) 78 | concurrent_execution_time = time.time() - concurrent_start 79 | print(f"Isolated execution time: {isolated_execution_time}") 80 | print(f"Concurrent execution time: {concurrent_execution_time}") 81 | self.assertLess(concurrent_execution_time, isolated_execution_time * 1.5, 82 | "Tasks did not run concurrently, likely due to blocking HTTP client") 83 | 84 | # Validate order of task start and end times 85 | task1_start, task2_start = start_times[0][1], start_times[1][1] 86 | 87 | 88 | # Expect both tasks to overlap significantly in start times for async behavior 89 | self.assertLess(abs(task1_start - task2_start), 0.1, "Tasks should start concurrently") 90 | 91 | self.assertEqual("task1_started", task_indicators[0]) 92 | self.assertEqual("task2_started", task_indicators[1]) 93 | self.assertIn("task1_finished", task_indicators) 94 | self.assertIn("task2_finished", task_indicators) 95 | self.assertIsInstance(result, GeneratedQuery) 96 | assert result.uuid is not None 97 | assert len(result.detailed_steps) > 0 98 | assert len(result.query) > 0 99 | assert 'information_schema' in result.query.lower() 100 | assert len(result.tables) > 0 101 | assert hasattr(result.confidence_score, "confidence_value") 102 | 103 | async def test_modify_connections(self): 104 | # first try to delete te connection if it exists 105 | result = await self.async_waii_client.database.get_connections() 106 | result = result.connectors 107 | for conn in result: 108 | if ( 109 | conn.db_type == self.db_conn.db_type 110 | and conn.account_name == self.db_conn.account_name 111 | and self.db_conn.database == conn.database 112 | and self.db_conn.username == conn.username 113 | ): 114 | result = await self.async_waii_client.database.modify_connections( 115 | ModifyDBConnectionRequest(removed=[conn.key]) 116 | ) 117 | result = result.connectors 118 | break 119 | 120 | # then add the new connection 121 | new_result = await self.async_waii_client.database.modify_connections( 122 | ModifyDBConnectionRequest(updated=[self.db_conn]) 123 | ) 124 | new_result = new_result.connectors 125 | 126 | assert len(result) == len(new_result) - 1 127 | 128 | # then activate the new connection 129 | for conn in new_result: 130 | if ( 131 | conn.db_type == self.db_conn.db_type 132 | and conn.account_name == self.db_conn.account_name 133 | and self.db_conn.database == conn.database 134 | and self.db_conn.username == conn.username 135 | ): 136 | await self.async_waii_client.database.activate_connection(conn.key) 137 | break 138 | else: 139 | raise Exception("Cannot find the new connection") 140 | 141 | # list databases 142 | result = await self.async_waii_client.database.get_catalogs() 143 | assert len(result.catalogs) > 0 144 | for catalog in result.catalogs: 145 | for schema in catalog.schemas: 146 | for table in schema.tables: 147 | if len(table.refs) > 0: 148 | assert type(table.refs[0]) == TableReference 149 | 150 | 151 | 152 | 153 | 154 | 155 | -------------------------------------------------------------------------------- /waii_sdk_py/waii_sdk_py.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from contextlib import contextmanager 18 | from typing import Optional, List 19 | 20 | from .access_rules import AccessRuleImpl, AsyncAccessRuleImpl 21 | from .chart import ChartImpl, Chart, AsyncChartImpl 22 | from .chat import ChatImpl, Chat, AsyncChatImpl 23 | from .history import HistoryImpl, History, AsyncHistoryImpl 24 | from .query import QueryImpl, Query, AsyncQueryImpl 25 | from .database import DatabaseImpl, Database, AsyncDatabaseImpl 26 | from .semantic_context import SemanticContextImpl, SemanticContext, AsyncSemanticContextImpl 27 | from .settings import SettingsImpl, AsyncSettingsImpl 28 | from .user import UserImpl, AsyncUserImpl 29 | from .user.user_static import User 30 | from .waii_http_client import WaiiHttpClient 31 | import importlib.metadata 32 | from .my_pydantic import WaiiBaseModel 33 | from .semantic_layer_dump import SemanticLayerDumpImpl, SemanticLayerDump 34 | from .kg import KnowledgeGraphImpl, AsyncKnowledgeGraphImpl 35 | from .sdm import SdmImpl, Sdm, AsyncSdmImpl 36 | 37 | GET_MODELS_ENDPOINT = "get-models" 38 | 39 | 40 | class GetModelsRequest(WaiiBaseModel): 41 | pass 42 | 43 | 44 | class ModelType(WaiiBaseModel): 45 | name: str 46 | description: Optional[str] 47 | vendor: Optional[str] 48 | 49 | 50 | class GetModelsResponse(WaiiBaseModel): 51 | models: Optional[List[ModelType]] 52 | 53 | 54 | class Waii: 55 | 56 | def __init__(self, initialize_legacy_fields: bool = False): 57 | self.history = None 58 | self.query = None 59 | self.database = None 60 | self.semantic_context = None 61 | self.chat = None 62 | self.chart = None 63 | self.user = None 64 | self.access_rules = None 65 | self.settings = None 66 | self.knowledge_graph = None 67 | self.sdm = None 68 | self.initialize_legacy_fields = initialize_legacy_fields 69 | self.http_client = None 70 | 71 | 72 | def initialize(self, url: str = "https://tweakit.waii.ai/api/", api_key: str = "", verbose=False): 73 | http_client = WaiiHttpClient(url, api_key, verbose=verbose) 74 | self.http_client = http_client 75 | self.history = HistoryImpl(http_client) 76 | self.query = QueryImpl(http_client) 77 | self.database = DatabaseImpl(http_client) 78 | self.semantic_context = SemanticContextImpl(http_client) 79 | self.chat = ChatImpl(http_client) 80 | self.chart = ChartImpl(http_client) 81 | self.user = UserImpl(http_client) 82 | self.access_rules = AccessRuleImpl(http_client) 83 | self.settings = SettingsImpl(http_client) 84 | self.semantic_layer_dump = SemanticLayerDumpImpl(http_client) 85 | self.knowledge_graph = KnowledgeGraphImpl(http_client) 86 | self.sdm = SdmImpl(http_client) 87 | 88 | if self.initialize_legacy_fields: 89 | self.History = self.history 90 | self.Query = self.query 91 | self.Database = self.database 92 | self.SemanticContext = self.semantic_context 93 | self.Chart = self.chart 94 | self.Chat = self.chat 95 | self.User = self.user 96 | self.SemanticLayerDump = self.semantic_layer_dump 97 | self.KnowledgeGraph = self.knowledge_graph 98 | self.Sdm = self.sdm 99 | Query.http_client = http_client 100 | History.http_client = http_client 101 | Database.http_client = http_client 102 | SemanticContext.http_client = http_client 103 | User.http_client = http_client 104 | Chat.http_client = http_client 105 | Chart.http_client = http_client 106 | SemanticLayerDump.http_client = http_client 107 | Sdm.http_client = http_client 108 | 109 | conns = self.database.get_connections().connectors 110 | if len(conns) > 0: 111 | self.database.activate_connection(conns[0].key) 112 | 113 | @staticmethod 114 | def version(): 115 | return importlib.metadata.version('waii-sdk-py') 116 | 117 | def get_models(self, params: GetModelsRequest = GetModelsRequest()) -> GetModelsResponse: 118 | return self.http_client.common_fetch( 119 | GET_MODELS_ENDPOINT, params, GetModelsResponse 120 | ) 121 | 122 | @contextmanager 123 | def impersonate_user(self, user_id: str): 124 | try: 125 | self.http_client.set_impersonate_user_id(user_id) 126 | yield 127 | finally: 128 | self.clear_impersonation() 129 | 130 | def set_impersonate_user(self, user_id: str): 131 | self.http_client.set_impersonate_user_id(user_id) 132 | 133 | def clear_impersonation(self): 134 | self.http_client.set_impersonate_user_id('') 135 | 136 | WAII = Waii(True) 137 | 138 | class AsyncWaii: 139 | 140 | def __init__(self): 141 | self.http_client = None 142 | self.query = None 143 | self.database = None 144 | self.semantic_context = None 145 | self.chat = None 146 | self.chart = None 147 | self.user = None 148 | self.access_rules = None 149 | self.settings = None 150 | self.history = None 151 | self.knowledge_graph = None 152 | self.sdm = None 153 | 154 | 155 | 156 | 157 | async def initialize(self, url: str = "https://tweakit.waii.ai/api/", api_key: str = "", verbose=False): 158 | http_client = WaiiHttpClient(url, api_key, verbose=verbose) 159 | self.http_client = http_client 160 | self.query = AsyncQueryImpl(http_client) 161 | self.database = AsyncDatabaseImpl(http_client) 162 | self.semantic_context = AsyncSemanticContextImpl(http_client) 163 | self.chat = AsyncChatImpl(http_client) 164 | self.chart = AsyncChartImpl(http_client) 165 | self.user = AsyncUserImpl(http_client) 166 | self.access_rules = AsyncAccessRuleImpl(http_client) 167 | self.settings = AsyncSettingsImpl(http_client) 168 | self.history = AsyncHistoryImpl(http_client) 169 | self.knowledge_graph = AsyncKnowledgeGraphImpl(http_client) 170 | self.sdm = AsyncSdmImpl(http_client) 171 | result = await self.database.get_connections() 172 | 173 | conns = result.connectors 174 | if len(conns) > 0: 175 | await self.database.activate_connection(conns[0].key) 176 | 177 | @staticmethod 178 | def version(): 179 | return importlib.metadata.version('waii-sdk-py') 180 | -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- 1 | # Setup unit tests 2 | 3 | sudo -u $USER psql postgres 4 | 5 | Run the following SQL commands to create the test databases and tables: 6 | 7 | ``` 8 | CREATE ROLE waii WITH LOGIN CREATEDB PASSWORD 'password'; 9 | ``` 10 | 11 | Create DB1 12 | 13 | ``` 14 | CREATE DATABASE waii_sdk_test; 15 | \c waii_sdk_test; 16 | CREATE table public.movies (id serial primary key, title text, year integer, runtime integer, genres text[], director text, actors text[], plot text, poster text, imdb text, production text, website text, response text, created_at timestamp, updated_at timestamp); 17 | INSERT INTO public.movies (title, year, runtime, genres, director, actors, plot, poster, imdb, production, website, response, created_at, updated_at) 18 | VALUES 19 | ('The Shawshank Redemption', 1994, 142, ARRAY['Drama'], 'Frank Darabont', ARRAY['Tim Robbins', 'Morgan Freeman', 'Bob Gunton'], 'Two imprisoned men bond over a number of years, finding solace and eventual redemption through acts of common decency.', 'http://example.com/shawshank.jpg', 'tt0111161', 'Castle Rock Entertainment', 'http://www.shawshankredemption.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 20 | 21 | ('Inception', 2010, 148, ARRAY['Action', 'Adventure', 'Sci-Fi'], 'Christopher Nolan', ARRAY['Leonardo DiCaprio', 'Joseph Gordon-Levitt', 'Ellen Page'], 'A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O.', 'http://example.com/inception.jpg', 'tt1375666', 'Warner Bros. Pictures', 'http://www.inception-movie.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 22 | 23 | ('Pulp Fiction', 1994, 154, ARRAY['Crime', 'Drama'], 'Quentin Tarantino', ARRAY['John Travolta', 'Uma Thurman', 'Samuel L. Jackson'], 'The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.', 'http://example.com/pulpfiction.jpg', 'tt0110912', 'Miramax', 'http://www.pulpfiction.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 24 | 25 | ('The Matrix', 1999, 136, ARRAY['Action', 'Sci-Fi'], 'Lana Wachowski, Lilly Wachowski', ARRAY['Keanu Reeves', 'Laurence Fishburne', 'Carrie-Anne Moss'], 'A computer programmer discovers that reality as he knows it is a simulation created by machines to subjugate humanity.', 'http://example.com/matrix.jpg', 'tt0133093', 'Warner Bros. Pictures', 'http://www.whatisthematrix.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP); 26 | 27 | GRANT ALL PRIVILEGES ON DATABASE waii_sdk_test TO waii; 28 | GRANT ALL PRIVILEGES ON TABLE public.movies TO waii; 29 | 30 | CREATE DATABASE waii_sdk_test_copy; 31 | \c waii_sdk_test_copy; 32 | CREATE table public.movies (id serial primary key, title text, year integer, runtime integer, genres text[], director text, actors text[], plot text, poster text, imdb text, production text, website text, response text, created_at timestamp, updated_at timestamp); 33 | INSERT INTO public.movies (title, year, runtime, genres, director, actors, plot, poster, imdb, production, website, response, created_at, updated_at) 34 | VALUES 35 | ('The Shawshank Redemption', 1994, 142, ARRAY['Drama'], 'Frank Darabont', ARRAY['Tim Robbins', 'Morgan Freeman', 'Bob Gunton'], 'Two imprisoned men bond over a number of years, finding solace and eventual redemption through acts of common decency.', 'http://example.com/shawshank.jpg', 'tt0111161', 'Castle Rock Entertainment', 'http://www.shawshankredemption.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 36 | 37 | ('Inception', 2010, 148, ARRAY['Action', 'Adventure', 'Sci-Fi'], 'Christopher Nolan', ARRAY['Leonardo DiCaprio', 'Joseph Gordon-Levitt', 'Ellen Page'], 'A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O.', 'http://example.com/inception.jpg', 'tt1375666', 'Warner Bros. Pictures', 'http://www.inception-movie.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 38 | 39 | ('Pulp Fiction', 1994, 154, ARRAY['Crime', 'Drama'], 'Quentin Tarantino', ARRAY['John Travolta', 'Uma Thurman', 'Samuel L. Jackson'], 'The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.', 'http://example.com/pulpfiction.jpg', 'tt0110912', 'Miramax', 'http://www.pulpfiction.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 40 | 41 | ('The Matrix', 1999, 136, ARRAY['Action', 'Sci-Fi'], 'Lana Wachowski, Lilly Wachowski', ARRAY['Keanu Reeves', 'Laurence Fishburne', 'Carrie-Anne Moss'], 'A computer programmer discovers that reality as he knows it is a simulation created by machines to subjugate humanity.', 'http://example.com/matrix.jpg', 'tt0133093', 'Warner Bros. Pictures', 'http://www.whatisthematrix.com', 'True', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP); 42 | 43 | GRANT ALL PRIVILEGES ON DATABASE waii_sdk_test_copy TO waii; 44 | GRANT ALL PRIVILEGES ON TABLE public.movies TO waii; 45 | 46 | CREATE DATABASE waii_sdk_test2; 47 | \c waii_sdk_test2; 48 | CREATE table public.albums (id serial primary key, title text, year integer, artist text, genre text, tracks text[], cover text, website text, created_at timestamp, updated_at timestamp); 49 | 50 | INSERT INTO public.albums (title, year, artist, genre, tracks, cover, website, created_at, updated_at) 51 | VALUES 52 | ('Thriller', 1982, 'Michael Jackson', 'Pop', 53 | ARRAY['Wanna Be Startin'' Somethin''', 'Baby Be Mine', 'The Girl Is Mine', 'Thriller', 'Beat It', 'Billie Jean', 'Human Nature', 'P.Y.T. (Pretty Young Thing)', 'The Lady in My Life'], 54 | 'http://example.com/thriller.jpg', 'http://www.michaeljackson.com', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 55 | 56 | ('Back in Black', 1980, 'AC/DC', 'Hard Rock', 57 | ARRAY['Hells Bells', 'Shoot to Thrill', 'What Do You Do for Money Honey', 'Given the Dog a Bone', 'Let Me Put My Love Into You', 'Back in Black', 'You Shook Me All Night Long', 'Have a Drink on Me', 'Shake a Leg', 'Rock and Roll Ain''t Noise Pollution'], 58 | 'http://example.com/backinblack.jpg', 'http://www.acdc.com', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 59 | 60 | ('21', 2011, 'Adele', 'Pop', 61 | ARRAY['Rolling in the Deep', 'Rumour Has It', 'Turning Tables', 'Don''t You Remember', 'Set Fire to the Rain', 'He Won''t Go', 'Take It All', 'I''ll Be Waiting', 'One and Only', 'Lovesong', 'Someone Like You'], 62 | 'http://example.com/21.jpg', 'http://www.adele.com', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 63 | 64 | ('The Dark Side of the Moon', 1973, 'Pink Floyd', 'Progressive Rock', 65 | ARRAY['Speak to Me', 'Breathe', 'On the Run', 'Time', 'The Great Gig in the Sky', 'Money', 'Us and Them', 'Any Colour You Like', 'Brain Damage', 'Eclipse'], 66 | 'http://example.com/darksideofthemoon.jpg', 'http://www.pinkfloyd.com', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP), 67 | 68 | ('Nevermind', 1991, 'Nirvana', 'Grunge', 69 | ARRAY['Smells Like Teen Spirit', 'In Bloom', 'Come as You Are', 'Breed', 'Lithium', 'Polly', 'Territorial Pissings', 'Drain You', 'Lounge Act', 'Stay Away', 'On a Plain', 'Something in the Way'], 70 | 'http://example.com/nevermind.jpg', 'http://www.nirvana.com', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP); 71 | 72 | GRANT ALL PRIVILEGES ON DATABASE waii_sdk_test2 TO waii; 73 | GRANT ALL PRIVILEGES ON TABLE public.albums TO waii; 74 | ``` 75 | 76 | Also ensure you have DB connection to the example snowflake DB - MOVIE_DB. If you do not have this DB connection, consult with a Waii team member for info on this DB connection 77 | -------------------------------------------------------------------------------- /waii_sdk_py/chat/chat.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2023–2025 Waii, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | from typing import Optional, List, Dict, Union 19 | 20 | from waii_sdk_py.database import SearchContext 21 | 22 | from waii_sdk_py.semantic_context import SemanticStatement 23 | 24 | from ..my_pydantic import WaiiBaseModel 25 | 26 | from ..common import LLMBasedRequest, GetObjectRequest, AsyncObjectResponse, CommonRequest, CommonResponse 27 | from ..query import GetQueryResultResponse, GeneratedQuery 28 | from ..database import CatalogDefinition 29 | from ..semantic_context import GetSemanticContextResponse 30 | from ..chart import ChartGenerationResponse, ChartType 31 | from waii_sdk_py.utils import wrap_methods_with_async 32 | from ..waii_http_client import WaiiHttpClient 33 | 34 | CHAT_MESSAGE_ENDPOINT = "chat-message" 35 | SUBMIT_CHAT_MESSAGE_ENDPOINT = "submit-chat-message" 36 | GET_CHAT_RESPONSE_ENDPOINT = "get-chat-response" 37 | 38 | # Research template endpoints 39 | CREATE_RESEARCH_TEMPLATE_ENDPOINT = "create-research-template" 40 | GET_RESEARCH_TEMPLATE_ENDPOINT = "get-research-template" 41 | LIST_RESEARCH_TEMPLATES_ENDPOINT = "list-research-templates" 42 | UPDATE_RESEARCH_TEMPLATE_ENDPOINT = "update-research-template" 43 | DELETE_RESEARCH_TEMPLATE_ENDPOINT = "delete-research-template" 44 | 45 | 46 | class ChatModule(str, Enum): 47 | DATA = "data" 48 | TABLES = "tables" 49 | QUERY = "query" 50 | CHART = "chart" 51 | CONTEXT = "context" 52 | 53 | 54 | class ChatRequestMode(str, Enum): 55 | automatic = "automatic" 56 | single_turn = "single_turn" 57 | multi_turn = "multi_turn" 58 | deep_research = "deep_research" 59 | 60 | 61 | class ChatRequest(LLMBasedRequest): 62 | ask: str 63 | 64 | # should we streaming the output? 65 | # no need to support for the first implementation 66 | streaming: bool = False 67 | 68 | # link to previous conversation, pick up where conversation left off 69 | parent_uuid: Optional[str] 70 | 71 | # optional chart type, default to plotly 72 | chart_type: Optional[ChartType] 73 | 74 | modules: Optional[List[ChatModule]] 75 | 76 | # optional by default there is no limit 77 | module_limit_in_response: Optional[int] 78 | 79 | additional_context: Optional[List[SemanticStatement]] = None 80 | 81 | search_context: Optional[List[SearchContext]] = None 82 | 83 | mode: ChatRequestMode = ChatRequestMode.single_turn 84 | 85 | 86 | class ChatResponseData(WaiiBaseModel): 87 | data: Optional[GetQueryResultResponse] 88 | query: Optional[GeneratedQuery] 89 | chart: Optional[ChartGenerationResponse] 90 | semantic_context: Optional[GetSemanticContextResponse] 91 | tables: Optional[CatalogDefinition] 92 | 93 | 94 | class ChatResponseDataV2(WaiiBaseModel): 95 | data: Optional[Dict[str, GetQueryResultResponse]] = None 96 | query: Optional[Dict[str, GeneratedQuery]] = None 97 | chart: Optional[Dict[str, ChartGenerationResponse]] = None 98 | semantic_context: Optional[Dict[str, SemanticStatement]] = None 99 | tables: Optional[CatalogDefinition] = None 100 | 101 | error_info: Optional[dict] = None 102 | 103 | 104 | class ChatResponseStep(str, Enum): 105 | routing_request = "Routing Request" 106 | generating_query = "Generating Query" 107 | retrieving_context = "Retrieving Context" 108 | retrieving_tables = "Retrieving Tables" 109 | running_query = "Running Query" 110 | generating_chart = "Generating Chart" 111 | preparing_result = "Preparing Result" 112 | completed = "Completed" 113 | 114 | 115 | class ChatStatusUpdateEvent(WaiiBaseModel): 116 | title: Optional[str] 117 | summary: Optional[str] 118 | timestamp: Optional[int] 119 | step_status: Optional[str] # in-progress and completed 120 | percentage: Optional[float] 121 | 122 | 123 | class ChatResponse(WaiiBaseModel): 124 | # template response 125 | response: Optional[str] = None 126 | 127 | # use union for the two 128 | response_data: Optional[Union[ChatResponseData, ChatResponseDataV2]] = None 129 | is_new: Optional[bool] = False 130 | timestamp_ms: Optional[int] = None 131 | chat_uuid: str 132 | elapsed_time_ms: Optional[int] = None 133 | session_title: Optional[str] = None 134 | research_plan: Optional[str] = None 135 | 136 | # old way to display status, routing info, we have to keep it 137 | current_step: Optional[ChatResponseStep] = None 138 | routing_info: Optional[Dict[str, str]] = None 139 | response_selected_fields: Optional[List[ChatModule]] = None 140 | 141 | # newly added status update 142 | status_update_events: Optional[List[ChatStatusUpdateEvent]] = None 143 | 144 | 145 | class ResearchTemplate(WaiiBaseModel): 146 | template_id: Optional[str] = None 147 | title: str 148 | template: str 149 | 150 | 151 | class CreateResearchTemplateRequest(CommonRequest): 152 | research_template: ResearchTemplate 153 | 154 | 155 | class GetResearchTemplateRequest(CommonRequest): 156 | template_id: str 157 | 158 | 159 | class ListResearchTemplatesRequest(CommonRequest): 160 | limit: Optional[int] = 10 161 | search_text: Optional[str] = None 162 | 163 | 164 | class GetResearchTemplateResponse(CommonResponse): 165 | research_template: Optional[ResearchTemplate] = None 166 | 167 | 168 | class ListResearchTemplatesResponse(CommonResponse): 169 | research_templates: List[ResearchTemplate] = [] 170 | 171 | 172 | class UpdateResearchTemplateRequest(CommonRequest): 173 | research_template: ResearchTemplate 174 | 175 | 176 | class DeleteResearchTemplateRequest(CommonRequest): 177 | template_id: str 178 | 179 | class ChatImpl: 180 | 181 | def __init__(self, http_client: WaiiHttpClient): 182 | self.http_client = http_client 183 | 184 | def chat_message(self, params: ChatRequest) -> ChatResponse: 185 | return self.http_client.common_fetch(CHAT_MESSAGE_ENDPOINT, params, ChatResponse) 186 | 187 | def submit_chat_message( 188 | self, params: ChatRequest 189 | ) -> AsyncObjectResponse: 190 | return self.http_client.common_fetch( 191 | SUBMIT_CHAT_MESSAGE_ENDPOINT, params, AsyncObjectResponse 192 | ) 193 | 194 | def get_chat_response( 195 | self, params: GetObjectRequest 196 | ) -> ChatResponse: 197 | return self.http_client.common_fetch( 198 | GET_CHAT_RESPONSE_ENDPOINT, params, ChatResponse 199 | ) 200 | 201 | # Research Template Methods 202 | def create_research_template(self, params: CreateResearchTemplateRequest) -> CommonResponse: 203 | return self.http_client.common_fetch( 204 | CREATE_RESEARCH_TEMPLATE_ENDPOINT, params, CommonResponse 205 | ) 206 | 207 | def get_research_template(self, params: GetResearchTemplateRequest) -> GetResearchTemplateResponse: 208 | return self.http_client.common_fetch( 209 | GET_RESEARCH_TEMPLATE_ENDPOINT, params, GetResearchTemplateResponse 210 | ) 211 | 212 | def list_research_templates(self, params: ListResearchTemplatesRequest) -> ListResearchTemplatesResponse: 213 | return self.http_client.common_fetch( 214 | LIST_RESEARCH_TEMPLATES_ENDPOINT, params, ListResearchTemplatesResponse 215 | ) 216 | 217 | def update_research_template(self, params: UpdateResearchTemplateRequest) -> CommonResponse: 218 | return self.http_client.common_fetch( 219 | UPDATE_RESEARCH_TEMPLATE_ENDPOINT, params, CommonResponse 220 | ) 221 | 222 | def delete_research_template(self, params: DeleteResearchTemplateRequest) -> CommonResponse: 223 | return self.http_client.common_fetch( 224 | DELETE_RESEARCH_TEMPLATE_ENDPOINT, params, CommonResponse 225 | ) 226 | 227 | 228 | 229 | class AsyncChatImpl: 230 | def __init__(self, http_client: WaiiHttpClient): 231 | self._chat_impl = ChatImpl(http_client) 232 | wrap_methods_with_async(self._chat_impl, self) 233 | 234 | 235 | Chat = ChatImpl(WaiiHttpClient.get_instance()) --------------------------------------------------------------------------------