diff --git a/wren-ai-service/.env.dev.example b/wren-ai-service/.env.dev.example index a15ca71db..55faa58f8 100644 --- a/wren-ai-service/.env.dev.example +++ b/wren-ai-service/.env.dev.example @@ -44,8 +44,7 @@ DOCUMENT_STORE_PROVIDER=qdrant QDRANT_HOST=http://localhost:6333 -## ENGINE -ENGINE=wren_ui +ENGINE=wren_ui # wren_ui, wren_ibis, wren_engine ## when using wren_ui as the engine WREN_UI_ENDPOINT=http://localhost:3000 @@ -56,7 +55,10 @@ WREN_IBIS_SOURCE=bigquery ### this is a base64 encoded string of the MDL WREN_IBIS_MANIFEST= ### this is a base64 encode string of the connection info -WREN_IBIS_CONNECTION_INFO=e30= +WREN_IBIS_CONNECTION_INFO= + +## when using wren_engine as the engine +WREN_ENGINE_ENDPOINT=http://localhost:8080 # Evaluation DATASET_NAME=book_2 diff --git a/wren-ai-service/Makefile b/wren-ai-service/Makefile index e05e6cfbf..2aaaaae1d 100644 --- a/wren-ai-service/Makefile +++ b/wren-ai-service/Makefile @@ -30,7 +30,10 @@ eval: make dev-down demo: - cd demo; poetry run streamlit run app.py + poetry run streamlit run demo/app.py + +data_curation_app: + poetry run streamlit run eval/data_curation/app.py ## utilities related ## diff --git a/wren-ai-service/README.md b/wren-ai-service/README.md index 070304ebd..81c636917 100644 --- a/wren-ai-service/README.md +++ b/wren-ai-service/README.md @@ -28,16 +28,11 @@ The following commands can quickly start the service for development: ## Others -### Pipeline Evaluation(Deprecated, will introduce new way to evaluate the speed in the future) +### Pipeline Evaluation -- install `psql` -- fill in environment variables: `.env.dev` in the src folder and `config.properties` in the src/eval/wren-engine/etc folder -- start the docker service -- evaluation - - `make eval pipeline=ask args="--help"` - - `make eval pipeline=ask_details args="--help"` -- `make eval_visualzation` to compare between the evaluation results -- to run individual pipeline: `poetry run python -m src.pipelines.ask.[pipeline_name]` (e.g. `poetry run python -m src.pipelines.ask.retrieval_pipeline`) +- evaluation dataset curation + - copy `.env.example` file to `.env` in the `eval/data_curation` folder and fill in the environment variables + - execute the command under the `wren-ai-service` folder: `make data_curation_app` ### Speed Evaluation(Deprecated, will introduce new way to evaluate the speed in the future) diff --git a/wren-ai-service/demo/.env.example b/wren-ai-service/demo/.env.example index ab0c93773..040660e59 100644 --- a/wren-ai-service/demo/.env.example +++ b/wren-ai-service/demo/.env.example @@ -1,6 +1,8 @@ -bigquery.project-id=wrenai +bigquery.project-id= +bigquery.dataset-id= bigquery.credentials-key= -bigquery.location=asia-east1 +postgres.host= +postgres.port= +postgres.database= postgres.user= -postgres.password= -postgres.jdbc.url=jdbc:postgresql://localhost:5432/ \ No newline at end of file +postgres.password= \ No newline at end of file diff --git a/wren-ai-service/demo/app.py b/wren-ai-service/demo/app.py index 70b0b16e3..40087c170 100644 --- a/wren-ai-service/demo/app.py +++ b/wren-ai-service/demo/app.py @@ -7,10 +7,8 @@ from utils import ( DATA_SOURCES, ask, ask_details, - get_current_manifest, get_mdl_json, get_new_mdl_json, - is_current_manifest_available, prepare_duckdb, prepare_semantics, rerun_wren_engine, @@ -59,11 +57,7 @@ if __name__ == "__main__": col1, col2 = st.columns([2, 4]) with col1: - with st.expander("Current Deployed Model", expanded=True): - manifest_name, models, relationships = get_current_manifest() - st.markdown(f"Current Deployed Model: {manifest_name}") - show_er_diagram(models, relationships) - with st.expander("Deploy New Model"): + with st.expander("Deploy New Model", expanded=True): uploaded_file = st.file_uploader( f"Upload an MDL json file, and the file name must be [xxx]_[datasource]_mdl.json, now we support these datasources: {DATA_SOURCES}", type="json", @@ -154,8 +148,7 @@ if __name__ == "__main__": query = st.chat_input( "Ask a question about the database", - disabled=(not is_current_manifest_available()) - and st.session_state["semantics_preparation_status"] != "finished", + disabled=st.session_state["semantics_preparation_status"] != "finished", ) with col2: diff --git a/wren-ai-service/demo/pyproject.toml b/wren-ai-service/demo/pyproject.toml deleted file mode 100644 index 0b9b1ad46..000000000 --- a/wren-ai-service/demo/pyproject.toml +++ /dev/null @@ -1,21 +0,0 @@ -[tool.poetry] -name = "wren-ai-service-demo" -version = "0.1.0" -description = "" -authors = ["Jimmy Yeh "] -license = "AGPL-3.0" -readme = "README.md" - -[tool.poetry.dependencies] -python = "^3.12" -requests = "^2.31.0" -sqlparse = "^0.5.0" -streamlit = "^1.32.2" -watchdog = "^4.0.0" -python-dotenv = "^1.0.1" -orjson = "^3.10.3" - - -[build-system] -requires = ["poetry-core"] -build-backend = "poetry.core.masonry.api" diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py index aa702e1f8..b3ae59a48 100644 --- a/wren-ai-service/demo/utils.py +++ b/wren-ai-service/demo/utils.py @@ -1,3 +1,4 @@ +import base64 import json import os import re @@ -14,43 +15,13 @@ from dotenv import load_dotenv WREN_AI_SERVICE_BASE_URL = "http://localhost:5556" WREN_ENGINE_API_URL = "http://localhost:8080" +WREN_IBIS_API_URL = "http://localhost:8000" POLLING_INTERVAL = 0.5 DATA_SOURCES = ["duckdb", "bigquery", "postgres"] load_dotenv() -def get_current_manifest(): - response = requests.get( - f"{WREN_ENGINE_API_URL}/v1/mdl", - ) - - assert response.status_code == 200 - - manifest = response.json() - - if manifest["schema"] == "test_schema" and manifest["catalog"] == "test_catalog": - return "None", [], [] - - return ( - f"{manifest['catalog']}.{manifest['schema']}", - manifest["models"], - manifest["relationships"], - ) - - -def is_current_manifest_available(): - response = requests.get( - f"{WREN_ENGINE_API_URL}/v1/mdl", - ) - - assert response.status_code == 200 - - manifest = response.json() - - return manifest["catalog"] != "text_catalog" and manifest["schema"] != "test_schema" - - def _update_wren_engine_configs(configs: list[dict]): response = requests.patch( f"{WREN_ENGINE_API_URL}/v1/config", @@ -63,50 +34,78 @@ def _update_wren_engine_configs(configs: list[dict]): def rerun_wren_engine(mdl_json: Dict, dataset_type: str): assert dataset_type in DATA_SOURCES - if dataset_type == "bigquery": - BIGQUERY_CREDENTIALS = os.getenv("bigquery.credentials-key") - assert ( - BIGQUERY_CREDENTIALS is not None - ), "bigquery.credentials-key is not set in .env" + if dataset_type == "duckdb": + # replace the values of ENGINE to wren-ui in ../.env.dev + with open(".env.dev", "r") as f: + lines = f.readlines() + for i, line in enumerate(lines): + if line.startswith("ENGINE"): + lines[i] = "ENGINE=wren_engine\n" + break + with open(".env.dev", "w") as f: + f.writelines(lines) - _update_wren_engine_configs( - [ - {"name": "wren.datasource.type", "value": "bigquery"}, - { - "name": "bigquery.project-id", - "value": os.getenv("bigquery.project-id"), - }, - {"name": "bigquery.location", "value": os.getenv("bigquery.location")}, - {"name": "bigquery.credentials-key", "value": BIGQUERY_CREDENTIALS}, - ] - ) - elif dataset_type == "duckdb": _update_wren_engine_configs( [{"name": "wren.datasource.type", "value": "duckdb"}] ) - elif dataset_type == "postgresql": - _update_wren_engine_configs( - [ - {"name": "wren.datasource.type", "value": "postgres"}, - {"name": "postgres.user", "value": os.getenv("postgres.user")}, - {"name": "postgres.password", "value": os.getenv("postgres.password")}, - {"name": "postgres.jdbc.url", "value": os.getenv("postgres.jdbc.url")}, - ] + st.toast("Wren Engine is being re-run", icon="⏳") + + response = requests.post( + f"{WREN_ENGINE_API_URL}/v1/mdl/deploy", + json={ + "manifest": mdl_json, + "version": "latest", + }, ) - st.toast("Wren Engine is being re-run", icon="⏳") + assert response.status_code == 202, response.json() - response = requests.post( - f"{WREN_ENGINE_API_URL}/v1/mdl/deploy", - json={ - "manifest": mdl_json, - "version": "latest", - }, - ) + st.toast("Wren Engine is ready", icon="🎉") + else: + WREN_IBIS_SOURCE = dataset_type + WREN_IBIS_MANIFEST = base64.b64encode(orjson.dumps(mdl_json)).decode() + if dataset_type == "bigquery": + WREN_IBIS_CONNECTION_INFO = base64.b64encode( + orjson.dumps( + { + "project_id": os.getenv("bigquery.project-id"), + "dataset_id": os.getenv("bigquery.dataset-id"), + "credentials": os.getenv("bigquery.credentials-key"), + } + ) + ).decode() + elif dataset_type == "postgres": + WREN_IBIS_CONNECTION_INFO = base64.b64encode( + orjson.dumps( + { + "host": os.getenv("postgres.host"), + "port": int(os.getenv("postgres.port")), + "database": os.getenv("postgres.database"), + "user": os.getenv("postgres.user"), + "password": os.getenv("postgres.password"), + } + ) + ).decode() - assert response.status_code == 202 + # replace the values of WREN_IBIS_xxx to ../.env.dev + with open(".env.dev", "r") as f: + lines = f.readlines() + for i, line in enumerate(lines): + if line.startswith("ENGINE"): + lines[i] = "ENGINE=wren-ibis\n" + elif line.startswith("WREN_IBIS_SOURCE"): + lines[i] = f"WREN_IBIS_SOURCE={WREN_IBIS_SOURCE}\n" + elif line.startswith("WREN_IBIS_MANIFEST"): + lines[i] = f"WREN_IBIS_MANIFEST={WREN_IBIS_MANIFEST}\n" + elif line.startswith("WREN_IBIS_CONNECTION_INFO"): + lines[ + i + ] = f"WREN_IBIS_CONNECTION_INFO={WREN_IBIS_CONNECTION_INFO}\n" + with open(".env.dev", "w") as f: + f.writelines(lines) - st.toast("Wren Engine is ready", icon="🎉") + # wait for wren-ai-service to restart + time.sleep(5) def save_mdl_json_file(file_name: str, mdl_json: Dict): @@ -120,7 +119,7 @@ def save_mdl_json_file(file_name: str, mdl_json: Dict): def get_mdl_json(database_name: str): assert database_name in ["music", "nba", "ecommerce"] - with open(f"sample_dataset/{database_name}_duckdb_mdl.json", "r") as f: + with open(f"demo/sample_dataset/{database_name}_duckdb_mdl.json", "r") as f: mdl_json = json.load(f) return mdl_json @@ -146,20 +145,62 @@ def get_new_mdl_json(chosen_models: List[str]): @st.cache_data -def get_data_from_wren_engine(sql: str): - response = requests.get( - f"{WREN_ENGINE_API_URL}/v1/mdl/preview", - json={ - "sql": sql, - }, - ) +def get_data_from_wren_engine(sql: str, dataset_type: str): + assert dataset_type in DATA_SOURCES - assert response.status_code == 200 + if dataset_type == "duckdb": + response = requests.get( + f"{WREN_ENGINE_API_URL}/v1/mdl/preview", + json={ + "sql": sql, + }, + ) - data = response.json() - column_names = [f'{i}_{col["name"]}' for i, col in enumerate(data["columns"])] + if response.status_code != 200: + st.error(response.json()) + st.stop() - return pd.DataFrame(data["data"], columns=column_names) + data = response.json() + column_names = [f'{i}_{col["name"]}' for i, col in enumerate(data["columns"])] + + return pd.DataFrame(data["data"], columns=column_names) + else: + connection_info = { + "bigquery": { + "project_id": os.getenv("bigquery.project-id"), + "dataset_id": os.getenv("bigquery.dataset-id"), + "credentials": os.getenv("bigquery.credentials-key"), + }, + "postgres": { + "host": os.getenv("postgres.host"), + "port": int(os.getenv("postgres.port")) + if os.getenv("postgres.port") + else 5432, + "database": os.getenv("postgres.database"), + "user": os.getenv("postgres.user"), + "password": os.getenv("postgres.password"), + }, + } + + response = requests.post( + f"{WREN_IBIS_API_URL}/v2/ibis/{dataset_type}/query", + json={ + "sql": sql, + "manifestStr": base64.b64encode( + orjson.dumps(st.session_state["mdl_json"]) + ).decode(), + "connectionInfo": connection_info[dataset_type], + }, + ) + + if response.status_code != 200: + st.error(response.json()) + st.stop() + + data = response.json() + column_names = [f"{i}_{col}" for i, col in enumerate(data["columns"])] + + return pd.DataFrame(data["data"], columns=column_names) # ui related @@ -319,6 +360,7 @@ def show_asks_details_results(): st.dataframe( get_data_from_wren_engine( st.session_state["preview_sql"], + st.session_state["dataset_type"], ) ) diff --git a/wren-ai-service/docker/Dockerfile b/wren-ai-service/docker/Dockerfile index 392cdaa41..ede8861bd 100644 --- a/wren-ai-service/docker/Dockerfile +++ b/wren-ai-service/docker/Dockerfile @@ -12,7 +12,7 @@ WORKDIR /app COPY pyproject.toml ./ -RUN poetry install --without dev --no-root && rm -rf $POETRY_CACHE_DIR +RUN poetry install --without dev,eval,demo --no-root && rm -rf $POETRY_CACHE_DIR FROM python:3.12.0-slim-bookworm as runtime diff --git a/wren-ai-service/eval/.gitignore b/wren-ai-service/eval/.gitignore new file mode 100644 index 000000000..2eea525d8 --- /dev/null +++ b/wren-ai-service/eval/.gitignore @@ -0,0 +1 @@ +.env \ No newline at end of file diff --git a/wren-ai-service/eval/data_curation/.env.example b/wren-ai-service/eval/data_curation/.env.example new file mode 100644 index 000000000..e203bf033 --- /dev/null +++ b/wren-ai-service/eval/data_curation/.env.example @@ -0,0 +1,6 @@ +OPENAI_API_KEY= +WREN_IBIS_ENDPOINT=http://localhost:8000 +WREN_ENGINE_ENDPOINT=http://localhost:8080 +bigquery.project-id= +bigquery.dataset-id= +bigquery.credentials-key= \ No newline at end of file diff --git a/wren-ai-service/eval/data_curation/__init__.py b/wren-ai-service/eval/data_curation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/wren-ai-service/eval/data_curation/app.py b/wren-ai-service/eval/data_curation/app.py new file mode 100644 index 000000000..4d49b5ddd --- /dev/null +++ b/wren-ai-service/eval/data_curation/app.py @@ -0,0 +1,461 @@ +import asyncio +import os +import re +import uuid +from datetime import datetime + +import orjson +import pandas as pd +import streamlit as st +import tomlkit +from openai import AsyncClient +from streamlit_tags import st_tags +from utils import ( + DATA_SOURCES, + get_contexts_from_sqls, + get_data_from_wren_engine, + get_eval_dataset_in_toml_string, + get_openai_client, + get_question_sql_pairs, + is_sql_valid, + prettify_sql, +) + +st.set_page_config(layout="wide") +st.title("WrenAI Data Curation App") + + +LLM_OPTIONS = ["gpt-3.5-turbo", "gpt-4-turbo", "gpt-4o"] + +llm_client = get_openai_client() + + +# session states +if "llm_model" not in st.session_state: + st.session_state["llm_model"] = LLM_OPTIONS[0] +if "deployment_id" not in st.session_state: + st.session_state["deployment_id"] = str(uuid.uuid4()) +if "mdl_json" not in st.session_state: + st.session_state["mdl_json"] = None +if "custom_instructions_for_llm" not in st.session_state: + st.session_state["custom_instructions_for_llm"] = "" +if "llm_question_sql_pairs" not in st.session_state: + st.session_state["llm_question_sql_pairs"] = [] +if "user_question_sql_pair" not in st.session_state: + st.session_state["user_question_sql_pair"] = {} +if "candidate_dataset" not in st.session_state: + st.session_state["candidate_dataset"] = [] +if "data_source" not in st.session_state: + st.session_state["data_source"] = None +if "connection_info" not in st.session_state: + st.session_state["connection_info"] = None + + +# widget callbacks +def on_change_upload_eval_dataset(): + doc = tomlkit.parse(st.session_state.uploaded_eval_file.getvalue().decode("utf-8")) + assert ( + doc["mdl"] == st.session_state["mdl_json"] + ), "The model in the uploaded dataset is different from the deployed model" + st.session_state["candidate_dataset"] = doc["eval_dataset"] + + +def on_change_custom_instructions_for_llm(): + st.session_state["custom_instructions_for_llm"] = st.session_state[ + "custom_instructions_text_area" + ] + + +def on_click_generate_question_sql_pairs(llm_client: AsyncClient): + st.toast("Generating question-sql-pairs...") + st.session_state["llm_question_sql_pairs"] = asyncio.run( + get_question_sql_pairs( + llm_client, + st.session_state["llm_model"], + st.session_state["mdl_json"], + st.session_state["custom_instructions_for_llm"], + st.session_state["data_source"], + st.session_state["connection_info"], + ) + ) + + +def on_click_setup_uploaded_file(): + uploaded_file = st.session_state.get("uploaded_mdl_file") + if uploaded_file: + match = re.match( + r".+_(" + "|".join(DATA_SOURCES) + r")_mdl\.json$", + uploaded_file.name, + ) + if not match: + st.error( + f"the file name must be [xxx]_[datasource]_mdl.json, now we support these datasources: {DATA_SOURCES}" + ) + st.stop() + + data_source = match.group(1) + st.session_state["data_source"] = data_source + st.session_state["mdl_json"] = orjson.loads( + uploaded_file.getvalue().decode("utf-8") + ) + + st.session_state["connection_info"] = { + "project_id": os.getenv("bigquery.project-id"), + "dataset_id": os.getenv("bigquery.dataset-id"), + "credentials": os.getenv("bigquery.credentials-key"), + } + else: + st.session_state["data_source"] = None + st.session_state["mdl_json"] = None + st.session_state["connection_info"] = None + + +def on_change_llm_model(): + st.toast(f"Switching LLM model to {st.session_state["select_llm_model"]}") + st.session_state["llm_model"] = st.session_state["select_llm_model"] + + +def on_change_sql(i: int, key: str): + sql = st.session_state[key] + + valid, error = asyncio.run( + is_sql_valid( + sql, + st.session_state["data_source"], + st.session_state["mdl_json"], + st.session_state["connection_info"], + ) + ) + if valid: + new_context = asyncio.run( + get_contexts_from_sqls([sql], st.session_state["mdl_json"]) + ) + if i != -1: + st.session_state["llm_question_sql_pairs"][i]["sql"] = sql + st.session_state["llm_question_sql_pairs"][i]["is_valid"] = valid + st.session_state["llm_question_sql_pairs"][i]["error"] = error + if valid: + st.session_state["llm_question_sql_pairs"][i]["context"] = new_context + else: + st.session_state["user_question_sql_pair"]["sql"] = sql + st.session_state["user_question_sql_pair"]["is_valid"] = valid + st.session_state["user_question_sql_pair"]["error"] = error + if valid: + st.session_state["user_question_sql_pair"]["context"] = new_context + + +def on_click_add_candidate_dataset(i: int, categories: list): + if i != -1: + dataset_to_add = { + "categories": categories, + "question": st.session_state["llm_question_sql_pairs"][i]["question"], + "context": st.session_state["llm_question_sql_pairs"][i]["context"], + "sql": st.session_state["llm_question_sql_pairs"][i]["sql"], + } + else: + dataset_to_add = { + "categories": categories, + "question": st.session_state["user_question_sql_pair"]["question"], + "context": st.session_state["user_question_sql_pair"]["context"], + "sql": st.session_state["user_question_sql_pair"]["sql"], + } + + # reset input for user question sql pair + st.session_state["user_question_sql_pair"] = {} + st.session_state["user_question"] = "" + st.session_state["user_sql"] = "" + + should_add = True + for dataset in st.session_state["candidate_dataset"]: + if dataset == dataset_to_add: + should_add = False + break + + if should_add: + st.session_state["candidate_dataset"].append(dataset_to_add) + + +def on_change_user_question(): + if not st.session_state["user_question_sql_pair"]: + st.session_state["user_question_sql_pair"] = { + "question": st.session_state["user_question"], + "context": [], + "sql": "", + "is_valid": False, + "error": "", + } + else: + st.session_state["user_question_sql_pair"] = { + **st.session_state["user_question_sql_pair"], + "question": st.session_state["user_question"], + } + + +def on_click_remove_candidate_dataset_button(i: int): + st.session_state["candidate_dataset"].pop(i) + + +st.file_uploader( + f"Upload an MDL json file, and the file name must be [xxx]_[datasource]_mdl.json, now we support these datasources: {DATA_SOURCES}", + type="json", + key="uploaded_mdl_file", + on_change=on_click_setup_uploaded_file, +) + +st.selectbox( + label="Select which LLM model you want to use", + options=LLM_OPTIONS, + index=0, + key="select_llm_model", + on_change=on_change_llm_model, +) + +tab_create_dataset, tab_modify_dataset = st.tabs( + ["Create New Evaluation Dataset", "Modify Saved Evaluation Dataset"] +) +with tab_create_dataset: + st.markdown( + """ + ### Usage Guide + 1. Upload an MDL json file first + 2. Get question-sql-pairs given by LLM or you manually enter question and corresponding sql + 3. Do validation on each group of question, context and SQL, and move it to the candidate dataset if you think it's valid + 3. Save the candidate dataset by clicking the "Save as Evaluation Dataset" button. + """ + ) + +with tab_modify_dataset: + st.markdown( + """ + ### Usage Guide + 1. Upload an MDL json file first + 2. Upload the evaluation dataset(`.toml` file) you want to modify, and please make sure the model in the dataset is the same as the deployed model + 3. Modify the evaluation dataset the same as you create a new one + 4. Save the candidate dataset by clicking the "Save as Evaluation Dataset" button. + """ + ) + + st.warning( + "WARNING: Uploading the evaluation dataset will overwrite the current candidate dataset" + ) + st.file_uploader( + "Upload Evaluation Dataset", + type="toml", + key="uploaded_eval_file", + on_change=on_change_upload_eval_dataset, + disabled=st.session_state["mdl_json"] is None, + ) + +if st.session_state["mdl_json"] is not None: + col1, col2 = st.columns(2) + with col1: + st.markdown("### Question SQL Pairs") + tab_generated_by_llm, tab_generated_by_user = st.tabs( + ["Generated by LLM", "Generated by User"] + ) + + with tab_generated_by_llm: + st.text_area( + "Custom Instructions for generating question-sql-pairs (Optional)", + key="custom_instructions_text_area", + value=st.session_state["custom_instructions_for_llm"], + placeholder="You can specify the custom instructions on how LLM should generate question-sql-pairs here, for example: what type of questions you want to generate.", + on_change=on_change_custom_instructions_for_llm, + ) + + st.button( + "Generate 10 question-sql-pairs", + key="generate_question_sql_pairs", + on_click=on_click_generate_question_sql_pairs, + args=(llm_client,), + ) + + with st.container(border=True, height=550): + for i, question_sql_pair in enumerate( + st.session_state["llm_question_sql_pairs"] + ): + st.text_input( + f"Question {i}", + question_sql_pair["question"], + disabled=True, + key=f"question_{i}", + ) + categories = st_tags( + label=f"Categories {i}", + text="Press enter to add more", + value=[], + key=f"categories_{i}", + ) + st.multiselect( + label=f"Context {i}", + options=question_sql_pair["context"], + default=question_sql_pair["context"], + key=f"context_{i}", + help="Contexts are automatically generated based on the SQL once you save the changes of the it(ctrl+enter or command+enter)", + disabled=True, + ) + st.text_area( + f"SQL {i}", + prettify_sql(question_sql_pair["sql"]), + key=f"sql_{i}", + height=250, + on_change=on_change_sql, + args=(i, f"sql_{i}"), + ) + if st.session_state["llm_question_sql_pairs"][i]["is_valid"]: + st.success("SQL is valid") + st.dataframe( + pd.DataFrame( + question_sql_pair["data"]["data"], + columns=question_sql_pair["data"]["columns"], + ) + ) + else: + st.error( + f"SQL is invalid: {st.session_state["llm_question_sql_pairs"][i]["error"]}" + ) + + st.button( + "Move it to the candidate dataset", + key=f"move_to_dataset_{i}", + disabled=( + not st.session_state["llm_question_sql_pairs"][i][ + "is_valid" + ] + or not st.session_state[f"context_{i}"] + or not categories + ), + on_click=on_click_add_candidate_dataset, + args=( + i, + categories, + ), + ) + + st.markdown("---") + + with tab_generated_by_user: + with st.container(border=True, height=550): + st.text_input( + "Question", + disabled=False, + key="user_question", + on_change=on_change_user_question, + ) + categories = st_tags( + label="Categories", + text="Press enter to add more", + value=[], + key="user_categories", + ) + st.multiselect( + label="Context", + options=st.session_state.get("user_question_sql_pair", {}).get( + "context", [] + ), + default=st.session_state.get("user_question_sql_pair", {}).get( + "context", [] + ), + key="user_context", + help="Contexts are automatically generated based on the SQL once you save the changes of the it(ctrl+enter or command+enter)", + disabled=True, + ) + st.text_area( + "SQL", + key="user_sql", + height=250, + on_change=on_change_sql, + args=(-1, "user_sql"), + ) + + if st.session_state.get("user_question_sql_pair", {}).get( + "is_valid", False + ): + st.success("SQL is valid") + data = asyncio.run( + get_data_from_wren_engine( + [st.session_state["user_question_sql_pair"]["sql"]], + st.session_state["data_source"], + ) + )[0] + st.dataframe( + pd.DataFrame( + data["data"], + columns=data["columns"], + ) + ) + else: + st.error( + f"SQL is invalid: {st.session_state.get("user_question_sql_pair", {}).get('error', '')}" + ) + + st.button( + "Move it to the candidate dataset", + key="move_to_dataset", + disabled=( + not st.session_state.get("user_question_sql_pair", {}).get( + "is_valid", False + ) + or not st.session_state["user_context"] + or not st.session_state["user_question"] + or not categories + ), + on_click=on_click_add_candidate_dataset, + args=( + -1, + categories, + ), + ) + with col2: + st.markdown("### Candidate Dataset") + + with st.container(border=True, height=600): + for i, dataset in enumerate(st.session_state["candidate_dataset"]): + st.text_input( + f"Question {i}", + dataset["question"], + disabled=True, + key=f"candidate_dataset_question_{i}", + ) + st.multiselect( + f"Categories {i}", + options=dataset["categories"], + default=dataset["categories"], + disabled=True, + key=f"candidate_dataset_categories_{i}", + ) + st.multiselect( + f"Context {i}", + options=dataset["context"], + default=dataset["context"], + disabled=True, + key=f"candidate_dataset_context_{i}", + ) + st.markdown(f"SQL {i}") + st.code(prettify_sql(dataset["sql"]), language="sql", line_numbers=True) + st.button( + "Remove", + key=f"remove_{i}", + on_click=on_click_remove_candidate_dataset_button, + args=(i,), + ) + st.markdown("---") + + with st.popover("Save as Evaluation Dataset", use_container_width=True): + file_name = st.text_input( + "File Name", + f'eval_dataset_{datetime.today().strftime("%Y_%m_%d")}.toml', + key="eval_dataset_file_name", + ) + download_btn = st.download_button( + "Download", + get_eval_dataset_in_toml_string( + st.session_state["mdl_json"], + st.session_state["candidate_dataset"], + ), + file_name=file_name, + key="download_eval_dataset_confirmed", + disabled=not st.session_state["candidate_dataset"], + ) + if download_btn: + st.toast("Downloading the evaluation dataset...") diff --git a/wren-ai-service/eval/data_curation/utils.py b/wren-ai-service/eval/data_curation/utils.py new file mode 100644 index 000000000..0b6ab8166 --- /dev/null +++ b/wren-ai-service/eval/data_curation/utils.py @@ -0,0 +1,326 @@ +import asyncio +import base64 +import os +import sys +from datetime import datetime +from pathlib import Path +from typing import List, Tuple + +import aiohttp +import orjson +import sqlglot +import sqlparse +import streamlit as st +import tomlkit +from dotenv import load_dotenv +from openai import AsyncClient + +# in order to import the DDLConverter class from the indexing module +sys.path.append(f"{Path().parent.parent.resolve()}") +from src.pipelines.indexing.indexing import DDLConverter + +load_dotenv() + +WREN_IBIS_ENDPOINT = os.getenv("WREN_IBIS_ENDPOINT", "http://localhost:8000") +WREN_ENGINE_ENDPOINT = os.getenv("WREN_ENGINE_ENDPOINT", "http://localhost:8080") +DATA_SOURCES = ["bigquery"] +TIMEOUT_SECONDS = 60 +ddl_converter = DDLConverter() + + +def get_openai_client( + api_key: str = os.getenv("OPENAI_API_KEY"), timeout: float = TIMEOUT_SECONDS +) -> AsyncClient: + return AsyncClient( + api_key=api_key, + timeout=timeout, + ) + + +def add_quotes(sql: str) -> str: + return sqlglot.transpile(sql, read="trino", identify=True)[0] + + +async def is_sql_valid( + sql: str, + data_source: str, + mdl_json: dict, + connection_info: dict, + api_endpoint: str = WREN_IBIS_ENDPOINT, + timeout: float = TIMEOUT_SECONDS, +) -> Tuple[bool, str]: + sql = sql[:-1] if sql.endswith(";") else sql + async with aiohttp.request( + "POST", + f"{api_endpoint}/v2/connector/{data_source}/query?dryRun=true", + json={ + "sql": add_quotes(sql), + "manifestStr": base64.b64encode(orjson.dumps(mdl_json)).decode(), + "connectionInfo": connection_info, + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status == 204: + return True, None + res = await response.text() + + return False, res + + +async def get_validated_question_sql_pairs( + question_sql_pairs: list[dict], + data_source: str, + mdl_json: dict, + connection_info: dict, +) -> list[dict]: + tasks = [] + + async with aiohttp.ClientSession(): + for question_sql_pair in question_sql_pairs: + task = asyncio.ensure_future( + is_sql_valid( + question_sql_pair["sql"], + data_source, + mdl_json, + connection_info, + ) + ) + tasks.append(task) + + results = await asyncio.gather(*tasks) + return [ + { + **question_sql_pairs[i], + "context": [], + "is_valid": valid, + "error": error, + } + for i, (valid, error) in enumerate(results) + ] + + +async def get_sql_analysis( + sql: str, + mdl_json: dict, + api_endpoint: str = WREN_ENGINE_ENDPOINT, + timeout: float = TIMEOUT_SECONDS, +) -> List[dict]: + sql = sql[:-1] if sql.endswith(";") else sql + async with aiohttp.request( + "GET", + f"{api_endpoint}/v1/analysis/sql", + json={ + "sql": add_quotes(sql), + "manifest": mdl_json, + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + return await response.json() + + +async def get_contexts_from_sqls( + sqls: list[str], + mdl_json: dict, +) -> list[str]: + def _compose_contexts_of_select_type(select_items: list[dict]): + return [ + f'{expr_source['sourceDataset']}.{expr_source['expression']}' + for select_item in select_items + for expr_source in select_item["exprSources"] + ] + + def _compose_contexts_of_filter_type(filter: dict): + contexts = [] + if filter["type"] == "EXPR": + contexts += [ + f'{expr_source["sourceDataset"]}.{expr_source["expression"]}' + for expr_source in filter["exprSources"] + ] + elif filter["type"] in ("AND", "OR"): + contexts += _compose_contexts_of_filter_type(filter["left"]) + contexts += _compose_contexts_of_filter_type(filter["right"]) + + return contexts + + def _compose_contexts_of_groupby_type(groupby_keys: list[list[dict]]): + contexts = [] + for groupby_key_list in groupby_keys: + contexts += [ + f'{expr_source["sourceDataset"]}.{expr_source["expression"]}' + for groupby_key in groupby_key_list + for expr_source in groupby_key["exprSources"] + ] + return contexts + + def _compose_contexts_of_sorting_type(sortings: list[dict]): + return [ + f'{expr_source["sourceDataset"]}.{expr_source["expression"]}' + for sorting in sortings + for expr_source in sorting["exprSources"] + ] + + def _get_contexts_from_sql_analysis_results(sql_analysis_results: list[dict]): + contexts = [] + for result in sql_analysis_results: + if "selectItems" in result: + contexts += _compose_contexts_of_select_type(result["selectItems"]) + if "filter" in result: + contexts += _compose_contexts_of_filter_type(result["filter"]) + if "groupByKeys" in result: + contexts += _compose_contexts_of_groupby_type(result["groupByKeys"]) + if "sortings" in result: + contexts += _compose_contexts_of_sorting_type(result["sortings"]) + + # print( + # f'SQL ANALYSIS RESULTS: {orjson.dumps(sql_analysis_results, option=orjson.OPT_INDENT_2).decode("utf-8")}' + # ) + # print(f"CONTEXTS: {sorted(set(contexts))}") + # print("\n\n") + + return sorted(set(contexts)) + + async with aiohttp.ClientSession(): + tasks = [] + for sql in sqls: + task = asyncio.ensure_future(get_sql_analysis(sql, mdl_json)) + tasks.append(task) + + results = await asyncio.gather(*tasks) + + return [_get_contexts_from_sql_analysis_results(result) for result in results] + + +async def get_question_sql_pairs( + llm_client: AsyncClient, + llm_model: str, + mdl_json: dict, + custom_instructions: str, + data_source: str, + connection_info: dict, + num_pairs: int = 10, +) -> list[dict]: + messages = [ + { + "role": "system", + "content": "", + }, + { + "role": "user", + "content": f""" +### TASK ### +Given the database DDL, generate {num_pairs} of the questions and corresponding SQL queries. + +### Output Format ### +{{ + "results": [ + {{ + "question": , + "sql": + }}, + {{ + "question": , + "sql": + }}, + ... + ] +}} + +### Custom Instructions ### + +{custom_instructions} + +### Input ### +Data Model: {"\n\n".join(ddl_converter.get_ddl_commands(mdl_json))} + +Generate {num_pairs} of the questions and corresponding SQL queries according to the Output Format in JSON +Think step by step +""", + }, + ] + + try: + response = await llm_client.chat.completions.create( + model=llm_model, + messages=messages, + response_format={"type": "json_object"}, + max_tokens=4096, + temperature=0, + ) + + results = orjson.loads(response.choices[0].message.content)["results"] + question_sql_pairs = await get_validated_question_sql_pairs( + results, + data_source=data_source, + mdl_json=mdl_json, + connection_info=connection_info, + ) + sqls = [question_sql_pair["sql"] for question_sql_pair in question_sql_pairs] + contexts = await get_contexts_from_sqls(sqls, mdl_json) + sqls_data = await get_data_from_wren_engine(sqls, data_source) + return [ + {**quesiton_sql_pair, "context": context, "data": sql_data} + for quesiton_sql_pair, context, sql_data in zip( + question_sql_pairs, contexts, sqls_data + ) + ] + except Exception as e: + st.error(f"Error generating question-sql-pairs: {e}") + return [] + + +def prettify_sql(sql: str) -> str: + return sqlparse.format( + sql, + reindent=True, + keyword_case="upper", + ) + + +async def get_data_from_wren_engine( + sqls: List[str], + data_source: str, + api_endpoint: str = WREN_IBIS_ENDPOINT, + data_sources: list[str] = DATA_SOURCES, + timeout: float = TIMEOUT_SECONDS, +) -> List[dict]: + assert data_source in data_sources, f"Invalid data source: {data_source}" + + async def _get_data(sql: str, data_source: str): + async with aiohttp.request( + "POST", + f"{api_endpoint}/v2/connector/{data_source}/query", + json={ + "sql": add_quotes(sql), + "manifestStr": base64.b64encode( + orjson.dumps(st.session_state["mdl_json"]) + ).decode(), + "connectionInfo": st.session_state["connection_info"], + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status != 200: + return {"data": [], "columns": []} + + data = await response.json() + column_names = [f"{i}_{col}" for i, col in enumerate(data["columns"])] + + return {"data": data["data"], "columns": column_names} + + async with aiohttp.ClientSession(): + tasks = [] + for sql in sqls: + task = asyncio.ensure_future(_get_data(sql, data_source)) + tasks.append(task) + + return await asyncio.gather(*tasks) + + +@st.cache_data +def get_eval_dataset_in_toml_string(mdl: dict, dataset: list) -> str: + doc = tomlkit.document() + + doc.add("date", datetime.today().strftime("%Y_%m_%d")) + doc.add("mdl", mdl) + doc.add("eval_dataset", dataset) + + return tomlkit.dumps(doc) diff --git a/wren-ai-service/poetry.lock b/wren-ai-service/poetry.lock index a96fa40d1..bf85df856 100644 --- a/wren-ai-service/poetry.lock +++ b/wren-ai-service/poetry.lock @@ -1,5 +1,19 @@ # This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +[[package]] +name = "aiodns" +version = "3.2.0" +description = "Simple DNS resolver for asyncio" +optional = false +python-versions = "*" +files = [ + {file = "aiodns-3.2.0-py3-none-any.whl", hash = "sha256:e443c0c27b07da3174a109fd9e736d69058d808f144d3c9d56dbd1776964c5f5"}, + {file = "aiodns-3.2.0.tar.gz", hash = "sha256:62869b23409349c21b072883ec8998316b234c9a9e36675756e8e317e8768f72"}, +] + +[package.dependencies] +pycares = ">=4.0.0" + [[package]] name = "aiohttp" version = "3.9.5" @@ -86,8 +100,11 @@ files = [ ] [package.dependencies] +aiodns = {version = "*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"darwin\") and extra == \"speedups\""} aiosignal = ">=1.1.2" attrs = ">=17.3.0" +Brotli = {version = "*", optional = true, markers = "platform_python_implementation == \"CPython\" and extra == \"speedups\""} +brotlicffi = {version = "*", optional = true, markers = "platform_python_implementation != \"CPython\" and extra == \"speedups\""} frozenlist = ">=1.1.1" multidict = ">=4.5,<7.0" yarl = ">=1.0,<2.0" @@ -340,15 +357,54 @@ files = [ {file = "Brotli-1.1.0.tar.gz", hash = "sha256:81de08ac11bcb85841e440c13611c00b67d3bf82698314928d0b676362546724"}, ] +[[package]] +name = "brotlicffi" +version = "1.1.0.0" +description = "Python CFFI bindings to the Brotli library" +optional = false +python-versions = ">=3.7" +files = [ + {file = "brotlicffi-1.1.0.0-cp37-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9b7ae6bd1a3f0df532b6d67ff674099a96d22bc0948955cb338488c31bfb8851"}, + {file = "brotlicffi-1.1.0.0-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:19ffc919fa4fc6ace69286e0a23b3789b4219058313cf9b45625016bf7ff996b"}, + {file = "brotlicffi-1.1.0.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9feb210d932ffe7798ee62e6145d3a757eb6233aa9a4e7db78dd3690d7755814"}, + {file = "brotlicffi-1.1.0.0-cp37-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:84763dbdef5dd5c24b75597a77e1b30c66604725707565188ba54bab4f114820"}, + {file = "brotlicffi-1.1.0.0-cp37-abi3-win32.whl", hash = "sha256:1b12b50e07c3911e1efa3a8971543e7648100713d4e0971b13631cce22c587eb"}, + {file = "brotlicffi-1.1.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:994a4f0681bb6c6c3b0925530a1926b7a189d878e6e5e38fae8efa47c5d9c613"}, + {file = "brotlicffi-1.1.0.0-pp310-pypy310_pp73-macosx_10_9_x86_64.whl", hash = "sha256:2e4aeb0bd2540cb91b069dbdd54d458da8c4334ceaf2d25df2f4af576d6766ca"}, + {file = "brotlicffi-1.1.0.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4b7b0033b0d37bb33009fb2fef73310e432e76f688af76c156b3594389d81391"}, + {file = "brotlicffi-1.1.0.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:54a07bb2374a1eba8ebb52b6fafffa2afd3c4df85ddd38fcc0511f2bb387c2a8"}, + {file = "brotlicffi-1.1.0.0-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7901a7dc4b88f1c1475de59ae9be59799db1007b7d059817948d8e4f12e24e35"}, + {file = "brotlicffi-1.1.0.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:ce01c7316aebc7fce59da734286148b1d1b9455f89cf2c8a4dfce7d41db55c2d"}, + {file = "brotlicffi-1.1.0.0-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:246f1d1a90279bb6069de3de8d75a8856e073b8ff0b09dcca18ccc14cec85979"}, + {file = "brotlicffi-1.1.0.0-pp37-pypy37_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cc4bc5d82bc56ebd8b514fb8350cfac4627d6b0743382e46d033976a5f80fab6"}, + {file = "brotlicffi-1.1.0.0-pp37-pypy37_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:37c26ecb14386a44b118ce36e546ce307f4810bc9598a6e6cb4f7fca725ae7e6"}, + {file = "brotlicffi-1.1.0.0-pp37-pypy37_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ca72968ae4eaf6470498d5c2887073f7efe3b1e7d7ec8be11a06a79cc810e990"}, + {file = "brotlicffi-1.1.0.0-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:add0de5b9ad9e9aa293c3aa4e9deb2b61e99ad6c1634e01d01d98c03e6a354cc"}, + {file = "brotlicffi-1.1.0.0-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:9b6068e0f3769992d6b622a1cd2e7835eae3cf8d9da123d7f51ca9c1e9c333e5"}, + {file = "brotlicffi-1.1.0.0-pp38-pypy38_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8557a8559509b61e65083f8782329188a250102372576093c88930c875a69838"}, + {file = "brotlicffi-1.1.0.0-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2a7ae37e5d79c5bdfb5b4b99f2715a6035e6c5bf538c3746abc8e26694f92f33"}, + {file = "brotlicffi-1.1.0.0-pp38-pypy38_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:391151ec86bb1c683835980f4816272a87eaddc46bb91cbf44f62228b84d8cca"}, + {file = "brotlicffi-1.1.0.0-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:2f3711be9290f0453de8eed5275d93d286abe26b08ab4a35d7452caa1fef532f"}, + {file = "brotlicffi-1.1.0.0-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:1a807d760763e398bbf2c6394ae9da5815901aa93ee0a37bca5efe78d4ee3171"}, + {file = "brotlicffi-1.1.0.0-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fa8ca0623b26c94fccc3a1fdd895be1743b838f3917300506d04aa3346fd2a14"}, + {file = "brotlicffi-1.1.0.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3de0cf28a53a3238b252aca9fed1593e9d36c1d116748013339f0949bfc84112"}, + {file = "brotlicffi-1.1.0.0-pp39-pypy39_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6be5ec0e88a4925c91f3dea2bb0013b3a2accda6f77238f76a34a1ea532a1cb0"}, + {file = "brotlicffi-1.1.0.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:d9eb71bb1085d996244439154387266fd23d6ad37161f6f52f1cd41dd95a3808"}, + {file = "brotlicffi-1.1.0.0.tar.gz", hash = "sha256:b77827a689905143f87915310b93b273ab17888fd43ef350d4832c4a71083c13"}, +] + +[package.dependencies] +cffi = ">=1.0.0" + [[package]] name = "cachetools" -version = "5.3.3" +version = "5.4.0" description = "Extensible memoizing collections and decorators" optional = false python-versions = ">=3.7" files = [ - {file = "cachetools-5.3.3-py3-none-any.whl", hash = "sha256:0abad1021d3f8325b2fc1d2e9c8b9c9d57b04c3932657a72465447332c24d945"}, - {file = "cachetools-5.3.3.tar.gz", hash = "sha256:ba29e2dfa0b8b556606f097407ed1aa62080ee108ab0dc5ec9d6a723a007d105"}, + {file = "cachetools-5.4.0-py3-none-any.whl", hash = "sha256:3ae3b49a3d5e28a77a0be2b37dbcb89005058959cb2323858c2657c4a8cab474"}, + {file = "cachetools-5.4.0.tar.gz", hash = "sha256:b8adc2e7c07f105ced7bc56dbb6dfbe7c4a00acce20e2227b3f355be89bc6827"}, ] [[package]] @@ -657,37 +713,37 @@ typing-inspect = ">=0.4.0,<1" [[package]] name = "datasets" -version = "2.20.0" +version = "2.19.1" description = "HuggingFace community-driven open-source library of datasets" optional = false python-versions = ">=3.8.0" files = [ - {file = "datasets-2.20.0-py3-none-any.whl", hash = "sha256:76ac02e3bdfff824492e20678f0b6b1b6d080515957fe834b00c2ba8d6b18e5e"}, - {file = "datasets-2.20.0.tar.gz", hash = "sha256:3c4dbcd27e0f642b9d41d20ff2efa721a5e04b32b2ca4009e0fc9139e324553f"}, + {file = "datasets-2.19.1-py3-none-any.whl", hash = "sha256:f7a78d15896f45004ccac1c298f3c7121f92f91f6f2bfbd4e4f210f827e6e411"}, + {file = "datasets-2.19.1.tar.gz", hash = "sha256:0df9ef6c5e9138cdb996a07385220109ff203c204245578b69cca905eb151d3a"}, ] [package.dependencies] aiohttp = "*" dill = ">=0.3.0,<0.3.9" filelock = "*" -fsspec = {version = ">=2023.1.0,<=2024.5.0", extras = ["http"]} +fsspec = {version = ">=2023.1.0,<=2024.3.1", extras = ["http"]} huggingface-hub = ">=0.21.2" multiprocess = "*" numpy = ">=1.17" packaging = "*" pandas = "*" -pyarrow = ">=15.0.0" +pyarrow = ">=12.0.0" pyarrow-hotfix = "*" pyyaml = ">=5.1" -requests = ">=2.32.2" -tqdm = ">=4.66.3" +requests = ">=2.19.0" +tqdm = ">=4.62.1" xxhash = "*" [package.extras] apache-beam = ["apache-beam (>=2.26.0)"] audio = ["librosa", "soundfile (>=0.12.1)"] benchmarks = ["tensorflow (==2.12.0)", "torch (==2.0.1)", "transformers (==4.30.1)"] -dev = ["Pillow (>=9.4.0)", "absl-py", "elasticsearch (<8.0.0)", "faiss-cpu (>=1.6.4)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "s3fs", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "sqlalchemy", "tensorflow (>=2.6.0)", "tiktoken", "torch", "torch (>=2.0.0)", "transformers", "typing-extensions (>=4.6.1)", "zstandard"] +dev = ["Pillow (>=6.2.1)", "absl-py", "apache-beam (>=2.26.0)", "elasticsearch (<8.0.0)", "faiss-cpu (>=1.6.4)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "s3fs", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "sqlalchemy", "tensorflow (>=2.6.0)", "tiktoken", "torch", "torch (>=2.0.0)", "transformers", "typing-extensions (>=4.6.1)", "zstandard"] docs = ["s3fs", "tensorflow (>=2.6.0)", "torch", "transformers"] jax = ["jax (>=0.3.14)", "jaxlib (>=0.3.14)"] metrics-tests = ["Werkzeug (>=1.0.1)", "accelerate", "bert-score (>=0.3.6)", "jiwer", "langdetect", "mauve-text", "nltk", "requests-file (>=1.5.1)", "rouge-score", "sacrebleu", "sacremoses", "scikit-learn", "scipy", "sentencepiece", "seqeval", "six (>=1.15.0,<1.16.0)", "spacy (>=3.0.0)", "texttable (>=1.6.3)", "tldextract", "tldextract (>=3.1.0)", "toml (>=0.10.1)", "typer (<0.5.0)"] @@ -695,9 +751,9 @@ quality = ["ruff (>=0.3.0)"] s3 = ["s3fs"] tensorflow = ["tensorflow (>=2.6.0)"] tensorflow-gpu = ["tensorflow (>=2.6.0)"] -tests = ["Pillow (>=9.4.0)", "absl-py", "elasticsearch (<8.0.0)", "faiss-cpu (>=1.6.4)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "sqlalchemy", "tensorflow (>=2.6.0)", "tiktoken", "torch (>=2.0.0)", "transformers", "typing-extensions (>=4.6.1)", "zstandard"] +tests = ["Pillow (>=6.2.1)", "absl-py", "apache-beam (>=2.26.0)", "elasticsearch (<8.0.0)", "faiss-cpu (>=1.6.4)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "sqlalchemy", "tensorflow (>=2.6.0)", "tiktoken", "torch (>=2.0.0)", "transformers", "typing-extensions (>=4.6.1)", "zstandard"] torch = ["torch"] -vision = ["Pillow (>=9.4.0)"] +vision = ["Pillow (>=6.2.1)"] [[package]] name = "dill" @@ -910,13 +966,13 @@ files = [ [[package]] name = "fsspec" -version = "2024.5.0" +version = "2024.3.1" description = "File-system specification" optional = false python-versions = ">=3.8" files = [ - {file = "fsspec-2024.5.0-py3-none-any.whl", hash = "sha256:e0fdbc446d67e182f49a70b82cf7889028a63588fde6b222521f10937b2b670c"}, - {file = "fsspec-2024.5.0.tar.gz", hash = "sha256:1d021b0b0f933e3b3029ed808eb400c08ba101ca2de4b3483fbc9ca23fcee94a"}, + {file = "fsspec-2024.3.1-py3-none-any.whl", hash = "sha256:918d18d41bf73f0e2b261824baeb1b124bcf771767e3a26425cd7dec3332f512"}, + {file = "fsspec-2024.3.1.tar.gz", hash = "sha256:f39780e282d7d117ffb42bb96992f8a90795e4d0fb0f661a70ca39fe9c43ded9"}, ] [package.dependencies] @@ -927,7 +983,7 @@ abfs = ["adlfs"] adl = ["adlfs"] arrow = ["pyarrow (>=1)"] dask = ["dask", "distributed"] -dev = ["pre-commit", "ruff"] +devel = ["pytest", "pytest-cov"] dropbox = ["dropbox", "dropboxdrivefs", "requests"] full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "dask", "distributed", "dropbox", "dropboxdrivefs", "fusepy", "gcsfs", "libarchive-c", "ocifs", "panel", "paramiko", "pyarrow (>=1)", "pygit2", "requests", "s3fs", "smbprotocol", "tqdm"] fuse = ["fusepy"] @@ -944,9 +1000,6 @@ s3 = ["s3fs"] sftp = ["paramiko"] smb = ["smbprotocol"] ssh = ["paramiko"] -test = ["aiohttp (!=4.0.0a0,!=4.0.0a1)", "numpy", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "requests"] -test-downstream = ["aiobotocore (>=2.5.4,<3.0.0)", "dask-expr", "dask[dataframe,test]", "moto[server] (>4,<5)", "pytest-timeout", "xarray"] -test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas", "panel", "paramiko", "pyarrow", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "smbprotocol", "tqdm", "urllib3", "zarr", "zstandard"] tqdm = ["tqdm"] [[package]] @@ -1536,13 +1589,13 @@ socks = ["socksio (==1.*)"] [[package]] name = "huggingface-hub" -version = "0.23.4" +version = "0.23.5" description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub" optional = false python-versions = ">=3.8.0" files = [ - {file = "huggingface_hub-0.23.4-py3-none-any.whl", hash = "sha256:3a0b957aa87150addf0cc7bd71b4d954b78e749850e1e7fb29ebbd2db64ca037"}, - {file = "huggingface_hub-0.23.4.tar.gz", hash = "sha256:35d99016433900e44ae7efe1c209164a5a81dbbcd53a52f99c281dcd7ce22431"}, + {file = "huggingface_hub-0.23.5-py3-none-any.whl", hash = "sha256:d7a7d337615e11a45cc14a0ce5a605db6b038dc24af42866f731684825226e90"}, + {file = "huggingface_hub-0.23.5.tar.gz", hash = "sha256:67a9caba79b71235be3752852ca27da86bd54311d2424ca8afdb8dda056edf98"}, ] [package.dependencies] @@ -1705,18 +1758,18 @@ referencing = ">=0.31.0" [[package]] name = "langchain" -version = "0.2.7" +version = "0.2.8" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langchain-0.2.7-py3-none-any.whl", hash = "sha256:98e79e0b9a60a9c740b44d5b0135c85f649219308f30d373cf5f10d0efe18b87"}, - {file = "langchain-0.2.7.tar.gz", hash = "sha256:8742f363d2890854501e0075af04fcb470600f201dec251c9bd5841e1990e73d"}, + {file = "langchain-0.2.8-py3-none-any.whl", hash = "sha256:53e7dfe50294a14200f33bec22b4e14cb63857ccf0a5500b0d18b0fd51285d58"}, + {file = "langchain-0.2.8.tar.gz", hash = "sha256:7fecb309e3558cde4e5cf7e9ffb7c1ab3f07121c40a7ff3b0c27135f8120c296"}, ] [package.dependencies] aiohttp = ">=3.8.3,<4.0.0" -langchain-core = ">=0.2.12,<0.3.0" +langchain-core = ">=0.2.19,<0.3.0" langchain-text-splitters = ">=0.2.0,<0.3.0" langsmith = ">=0.1.17,<0.2.0" numpy = {version = ">=1.26.0,<2.0.0", markers = "python_version >= \"3.12\""} @@ -1751,13 +1804,13 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0" [[package]] name = "langchain-core" -version = "0.2.17" +version = "0.2.20" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langchain_core-0.2.17-py3-none-any.whl", hash = "sha256:f8b4f64c95b381bd3ac4322bafa9ecee3ca3970aa0aa2829ef11175bcad27249"}, - {file = "langchain_core-0.2.17.tar.gz", hash = "sha256:f9b8d0f7b5f339ac62e11c417181e87f54e9282af82d1a64c8a7427f3d5d9118"}, + {file = "langchain_core-0.2.20-py3-none-any.whl", hash = "sha256:16cc4da6f7ebf33accea7af45a70480733dc852ab291030fb6924865bd7caf76"}, + {file = "langchain_core-0.2.20.tar.gz", hash = "sha256:a66c439e085d8c75f822f7650a5551d17bada4003521173c763d875d949e4ed5"}, ] [package.dependencies] @@ -1827,13 +1880,13 @@ openai = ["openai (>=0.27.8)"] [[package]] name = "langsmith" -version = "0.1.85" +version = "0.1.88" description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform." optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langsmith-0.1.85-py3-none-any.whl", hash = "sha256:c1f94384f10cea96f7b4d33fd3db7ec180c03c7468877d50846f881d2017ff94"}, - {file = "langsmith-0.1.85.tar.gz", hash = "sha256:acff31f9e53efa48586cf8e32f65625a335c74d7c4fa306d1655ac18452296f6"}, + {file = "langsmith-0.1.88-py3-none-any.whl", hash = "sha256:460ebb7de440afd150fcea8f54ca8779821f2228cd59e149e5845c9dbe06db16"}, + {file = "langsmith-0.1.88.tar.gz", hash = "sha256:28a07dec19197f4808aa2628d5a3ccafcbe14cc137aef0e607bbd128e7907821"}, ] [package.dependencies] @@ -2631,13 +2684,13 @@ testing = ["pytest", "pytest-benchmark"] [[package]] name = "portalocker" -version = "2.10.0" +version = "2.10.1" description = "Wraps the portalocker recipe for easy usage" optional = false python-versions = ">=3.8" files = [ - {file = "portalocker-2.10.0-py3-none-any.whl", hash = "sha256:48944147b2cd42520549bc1bb8fe44e220296e56f7c3d551bc6ecce69d9b0de1"}, - {file = "portalocker-2.10.0.tar.gz", hash = "sha256:49de8bc0a2f68ca98bf9e219c81a3e6b27097c7bf505a87c5a112ce1aaeb9b81"}, + {file = "portalocker-2.10.1-py3-none-any.whl", hash = "sha256:53a5984ebc86a025552264b459b46a2086e269b21823cb572f8f28ee759e45bf"}, + {file = "portalocker-2.10.1.tar.gz", hash = "sha256:ef1bf844e878ab08aee7e40184156e1151f228f103aa5c6bd0724cc330960f8f"}, ] [package.dependencies] @@ -2821,52 +2874,42 @@ files = [ [[package]] name = "pyarrow" -version = "16.1.0" +version = "17.0.0" description = "Python library for Apache Arrow" optional = false python-versions = ">=3.8" files = [ - {file = "pyarrow-16.1.0-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:17e23b9a65a70cc733d8b738baa6ad3722298fa0c81d88f63ff94bf25eaa77b9"}, - {file = "pyarrow-16.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:4740cc41e2ba5d641071d0ab5e9ef9b5e6e8c7611351a5cb7c1d175eaf43674a"}, - {file = "pyarrow-16.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:98100e0268d04e0eec47b73f20b39c45b4006f3c4233719c3848aa27a03c1aef"}, - {file = "pyarrow-16.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f68f409e7b283c085f2da014f9ef81e885d90dcd733bd648cfba3ef265961848"}, - {file = "pyarrow-16.1.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:a8914cd176f448e09746037b0c6b3a9d7688cef451ec5735094055116857580c"}, - {file = "pyarrow-16.1.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:48be160782c0556156d91adbdd5a4a7e719f8d407cb46ae3bb4eaee09b3111bd"}, - {file = "pyarrow-16.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:9cf389d444b0f41d9fe1444b70650fea31e9d52cfcb5f818b7888b91b586efff"}, - {file = "pyarrow-16.1.0-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:d0ebea336b535b37eee9eee31761813086d33ed06de9ab6fc6aaa0bace7b250c"}, - {file = "pyarrow-16.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2e73cfc4a99e796727919c5541c65bb88b973377501e39b9842ea71401ca6c1c"}, - {file = "pyarrow-16.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bf9251264247ecfe93e5f5a0cd43b8ae834f1e61d1abca22da55b20c788417f6"}, - {file = "pyarrow-16.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ddf5aace92d520d3d2a20031d8b0ec27b4395cab9f74e07cc95edf42a5cc0147"}, - {file = "pyarrow-16.1.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:25233642583bf658f629eb230b9bb79d9af4d9f9229890b3c878699c82f7d11e"}, - {file = "pyarrow-16.1.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a33a64576fddfbec0a44112eaf844c20853647ca833e9a647bfae0582b2ff94b"}, - {file = "pyarrow-16.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:185d121b50836379fe012753cf15c4ba9638bda9645183ab36246923875f8d1b"}, - {file = "pyarrow-16.1.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:2e51ca1d6ed7f2e9d5c3c83decf27b0d17bb207a7dea986e8dc3e24f80ff7d6f"}, - {file = "pyarrow-16.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:06ebccb6f8cb7357de85f60d5da50e83507954af617d7b05f48af1621d331c9a"}, - {file = "pyarrow-16.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b04707f1979815f5e49824ce52d1dceb46e2f12909a48a6a753fe7cafbc44a0c"}, - {file = "pyarrow-16.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0d32000693deff8dc5df444b032b5985a48592c0697cb6e3071a5d59888714e2"}, - {file = "pyarrow-16.1.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:8785bb10d5d6fd5e15d718ee1d1f914fe768bf8b4d1e5e9bf253de8a26cb1628"}, - {file = "pyarrow-16.1.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:e1369af39587b794873b8a307cc6623a3b1194e69399af0efd05bb202195a5a7"}, - {file = "pyarrow-16.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:febde33305f1498f6df85e8020bca496d0e9ebf2093bab9e0f65e2b4ae2b3444"}, - {file = "pyarrow-16.1.0-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:b5f5705ab977947a43ac83b52ade3b881eb6e95fcc02d76f501d549a210ba77f"}, - {file = "pyarrow-16.1.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:0d27bf89dfc2576f6206e9cd6cf7a107c9c06dc13d53bbc25b0bd4556f19cf5f"}, - {file = "pyarrow-16.1.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0d07de3ee730647a600037bc1d7b7994067ed64d0eba797ac74b2bc77384f4c2"}, - {file = "pyarrow-16.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fbef391b63f708e103df99fbaa3acf9f671d77a183a07546ba2f2c297b361e83"}, - {file = "pyarrow-16.1.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:19741c4dbbbc986d38856ee7ddfdd6a00fc3b0fc2d928795b95410d38bb97d15"}, - {file = "pyarrow-16.1.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:f2c5fb249caa17b94e2b9278b36a05ce03d3180e6da0c4c3b3ce5b2788f30eed"}, - {file = "pyarrow-16.1.0-cp38-cp38-win_amd64.whl", hash = "sha256:e6b6d3cd35fbb93b70ade1336022cc1147b95ec6af7d36906ca7fe432eb09710"}, - {file = "pyarrow-16.1.0-cp39-cp39-macosx_10_15_x86_64.whl", hash = "sha256:18da9b76a36a954665ccca8aa6bd9f46c1145f79c0bb8f4f244f5f8e799bca55"}, - {file = "pyarrow-16.1.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:99f7549779b6e434467d2aa43ab2b7224dd9e41bdde486020bae198978c9e05e"}, - {file = "pyarrow-16.1.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f07fdffe4fd5b15f5ec15c8b64584868d063bc22b86b46c9695624ca3505b7b4"}, - {file = "pyarrow-16.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ddfe389a08ea374972bd4065d5f25d14e36b43ebc22fc75f7b951f24378bf0b5"}, - {file = "pyarrow-16.1.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:3b20bd67c94b3a2ea0a749d2a5712fc845a69cb5d52e78e6449bbd295611f3aa"}, - {file = "pyarrow-16.1.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:ba8ac20693c0bb0bf4b238751d4409e62852004a8cf031c73b0e0962b03e45e3"}, - {file = "pyarrow-16.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:31a1851751433d89a986616015841977e0a188662fcffd1a5677453f1df2de0a"}, - {file = "pyarrow-16.1.0.tar.gz", hash = "sha256:15fbb22ea96d11f0b5768504a3f961edab25eaf4197c341720c4a387f6c60315"}, + {file = "pyarrow-17.0.0-cp310-cp310-macosx_10_15_x86_64.whl", hash = "sha256:a5c8b238d47e48812ee577ee20c9a2779e6a5904f1708ae240f53ecbee7c9f07"}, + {file = "pyarrow-17.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:db023dc4c6cae1015de9e198d41250688383c3f9af8f565370ab2b4cb5f62655"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da1e060b3876faa11cee287839f9cc7cdc00649f475714b8680a05fd9071d545"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:75c06d4624c0ad6674364bb46ef38c3132768139ddec1c56582dbac54f2663e2"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:fa3c246cc58cb5a4a5cb407a18f193354ea47dd0648194e6265bd24177982fe8"}, + {file = "pyarrow-17.0.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:f7ae2de664e0b158d1607699a16a488de3d008ba99b3a7aa5de1cbc13574d047"}, + {file = "pyarrow-17.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:5984f416552eea15fd9cee03da53542bf4cddaef5afecefb9aa8d1010c335087"}, + {file = "pyarrow-17.0.0-cp311-cp311-macosx_10_15_x86_64.whl", hash = "sha256:1c8856e2ef09eb87ecf937104aacfa0708f22dfeb039c363ec99735190ffb977"}, + {file = "pyarrow-17.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2e19f569567efcbbd42084e87f948778eb371d308e137a0f97afe19bb860ccb3"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b244dc8e08a23b3e352899a006a26ae7b4d0da7bb636872fa8f5884e70acf15"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b72e87fe3e1db343995562f7fff8aee354b55ee83d13afba65400c178ab2597"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:dc5c31c37409dfbc5d014047817cb4ccd8c1ea25d19576acf1a001fe07f5b420"}, + {file = "pyarrow-17.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:e3343cb1e88bc2ea605986d4b94948716edc7a8d14afd4e2c097232f729758b4"}, + {file = "pyarrow-17.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:a27532c38f3de9eb3e90ecab63dfda948a8ca859a66e3a47f5f42d1e403c4d03"}, + {file = "pyarrow-17.0.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:9b8a823cea605221e61f34859dcc03207e52e409ccf6354634143e23af7c8d22"}, + {file = "pyarrow-17.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f1e70de6cb5790a50b01d2b686d54aaf73da01266850b05e3af2a1bc89e16053"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0071ce35788c6f9077ff9ecba4858108eebe2ea5a3f7cf2cf55ebc1dbc6ee24a"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:757074882f844411fcca735e39aae74248a1531367a7c80799b4266390ae51cc"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9ba11c4f16976e89146781a83833df7f82077cdab7dc6232c897789343f7891a"}, + {file = "pyarrow-17.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b0c6ac301093b42d34410b187bba560b17c0330f64907bfa4f7f7f2444b0cf9b"}, + {file = "pyarrow-17.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:392bc9feabc647338e6c89267635e111d71edad5fcffba204425a7c8d13610d7"}, + {file = "pyarrow-17.0.0-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:af5ff82a04b2171415f1410cff7ebb79861afc5dae50be73ce06d6e870615204"}, + {file = "pyarrow-17.0.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:edca18eaca89cd6382dfbcff3dd2d87633433043650c07375d095cd3517561d8"}, ] [package.dependencies] numpy = ">=1.16.6" +[package.extras] +test = ["cffi", "hypothesis", "pandas", "pytest", "pytz"] + [[package]] name = "pyarrow-hotfix" version = "0.6" @@ -2878,6 +2921,72 @@ files = [ {file = "pyarrow_hotfix-0.6.tar.gz", hash = "sha256:79d3e030f7ff890d408a100ac16d6f00b14d44a502d7897cd9fc3e3a534e9945"}, ] +[[package]] +name = "pycares" +version = "4.4.0" +description = "Python interface for c-ares" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pycares-4.4.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:24da119850841d16996713d9c3374ca28a21deee056d609fbbed29065d17e1f6"}, + {file = "pycares-4.4.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8f64cb58729689d4d0e78f0bfb4c25ce2f851d0274c0273ac751795c04b8798a"}, + {file = "pycares-4.4.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d33e2a1120887e89075f7f814ec144f66a6ce06a54f5722ccefc62fbeda83cff"}, + {file = "pycares-4.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c680fef1b502ee680f8f0b95a41af4ec2c234e50e16c0af5bbda31999d3584bd"}, + {file = "pycares-4.4.0-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fff16b09042ba077f7b8aa5868d1d22456f0002574d0ba43462b10a009331677"}, + {file = "pycares-4.4.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:229a1675eb33bc9afb1fc463e73ee334950ccc485bc83a43f6ae5839fb4d5fa3"}, + {file = "pycares-4.4.0-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:3aebc73e5ad70464f998f77f2da2063aa617cbd8d3e8174dd7c5b4518f967153"}, + {file = "pycares-4.4.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:6ef64649eba56448f65e26546d85c860709844d2fc22ef14d324fe0b27f761a9"}, + {file = "pycares-4.4.0-cp310-cp310-win32.whl", hash = "sha256:4afc2644423f4eef97857a9fd61be9758ce5e336b4b0bd3d591238bb4b8b03e0"}, + {file = "pycares-4.4.0-cp310-cp310-win_amd64.whl", hash = "sha256:5ed4e04af4012f875b78219d34434a6d08a67175150ac1b79eb70ab585d4ba8c"}, + {file = "pycares-4.4.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:bce8db2fc6f3174bd39b81405210b9b88d7b607d33e56a970c34a0c190da0490"}, + {file = "pycares-4.4.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:9a0303428d013ccf5c51de59c83f9127aba6200adb7fd4be57eddb432a1edd2a"}, + {file = "pycares-4.4.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afb91792f1556f97be7f7acb57dc7756d89c5a87bd8b90363a77dbf9ea653817"}, + {file = "pycares-4.4.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b61579cecf1f4d616e5ea31a6e423a16680ab0d3a24a2ffe7bb1d4ee162477ff"}, + {file = "pycares-4.4.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b7af06968cbf6851566e806bf3e72825b0e6671832a2cbe840be1d2d65350710"}, + {file = "pycares-4.4.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ceb12974367b0a68a05d52f4162b29f575d241bd53de155efe632bf2c943c7f6"}, + {file = "pycares-4.4.0-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:2eeec144bcf6a7b6f2d74d6e70cbba7886a84dd373c886f06cb137a07de4954c"}, + {file = "pycares-4.4.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e3a6f7cfdfd11eb5493d6d632e582408c8f3b429f295f8799c584c108b28db6f"}, + {file = "pycares-4.4.0-cp311-cp311-win32.whl", hash = "sha256:34736a2ffaa9c08ca9c707011a2d7b69074bbf82d645d8138bba771479b2362f"}, + {file = "pycares-4.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:eb66c30eb11e877976b7ead13632082a8621df648c408b8e15cdb91a452dd502"}, + {file = "pycares-4.4.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:fd644505a8cfd7f6584d33a9066d4e3d47700f050ef1490230c962de5dfb28c6"}, + {file = "pycares-4.4.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:52084961262232ec04bd75f5043aed7e5d8d9695e542ff691dfef0110209f2d4"}, + {file = "pycares-4.4.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0c5368206057884cde18602580083aeaad9b860e2eac14fd253543158ce1e93"}, + {file = "pycares-4.4.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:112a4979c695b1c86f6782163d7dec58d57a3b9510536dcf4826550f9053dd9a"}, + {file = "pycares-4.4.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8d186dafccdaa3409194c0f94db93c1a5d191145a275f19da6591f9499b8e7b8"}, + {file = "pycares-4.4.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:64965dc19c578a683ea73487a215a8897276224e004d50eeb21f0bc7a0b63c88"}, + {file = "pycares-4.4.0-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:ed2a38e34bec6f2586435f6ff0bc5fe11d14bebd7ed492cf739a424e81681540"}, + {file = "pycares-4.4.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:94d6962db81541eb0396d2f0dfcbb18cdb8c8b251d165efc2d974ae652c547d4"}, + {file = "pycares-4.4.0-cp312-cp312-win32.whl", hash = "sha256:1168a48a834813aa80f412be2df4abaf630528a58d15c704857448b20b1675c0"}, + {file = "pycares-4.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:db24c4e7fea4a052c6e869cbf387dd85d53b9736cfe1ef5d8d568d1ca925e977"}, + {file = "pycares-4.4.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:21a5a0468861ec7df7befa69050f952da13db5427ae41ffe4713bc96291d1d95"}, + {file = "pycares-4.4.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:22c00bf659a9fa44d7b405cf1cd69b68b9d37537899898d8cbe5dffa4016b273"}, + {file = "pycares-4.4.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:23aa3993a352491a47fcf17867f61472f32f874df4adcbb486294bd9fbe8abee"}, + {file = "pycares-4.4.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:813d661cbe2e37d87da2d16b7110a6860e93ddb11735c6919c8a3545c7b9c8d8"}, + {file = "pycares-4.4.0-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:77cf5a2fd5583c670de41a7f4a7b46e5cbabe7180d8029f728571f4d2e864084"}, + {file = "pycares-4.4.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:3eaa6681c0a3e3f3868c77aca14b7760fed35fdfda2fe587e15c701950e7bc69"}, + {file = "pycares-4.4.0-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:ad58e284a658a8a6a84af2e0b62f2f961f303cedfe551854d7bd40c3cbb61912"}, + {file = "pycares-4.4.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:bfb89ca9e3d0a9b5332deeb666b2ede9d3469107742158f4aeda5ce032d003f4"}, + {file = "pycares-4.4.0-cp38-cp38-win32.whl", hash = "sha256:f36bdc1562142e3695555d2f4ac0cb69af165eddcefa98efc1c79495b533481f"}, + {file = "pycares-4.4.0-cp38-cp38-win_amd64.whl", hash = "sha256:902461a92b6a80fd5041a2ec5235680c7cc35e43615639ec2a40e63fca2dfb51"}, + {file = "pycares-4.4.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7bddc6adba8f699728f7fc1c9ce8cef359817ad78e2ed52b9502cb5f8dc7f741"}, + {file = "pycares-4.4.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:cb49d5805cd347c404f928c5ae7c35e86ba0c58ffa701dbe905365e77ce7d641"}, + {file = "pycares-4.4.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:56cf3349fa3a2e67ed387a7974c11d233734636fe19facfcda261b411af14d80"}, + {file = "pycares-4.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8bf2eaa83a5987e48fa63302f0fe7ce3275cfda87b34d40fef9ce703fb3ac002"}, + {file = "pycares-4.4.0-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:82bba2ab77eb5addbf9758d514d9bdef3c1bfe7d1649a47bd9a0d55a23ef478b"}, + {file = "pycares-4.4.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:c6a8bde63106f162fca736e842a916853cad3c8d9d137e11c9ffa37efa818b02"}, + {file = "pycares-4.4.0-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:f5f646eec041db6ffdbcaf3e0756fb92018f7af3266138c756bb09d2b5baadec"}, + {file = "pycares-4.4.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:9dc04c54c6ea615210c1b9e803d0e2d2255f87a3d5d119b6482c8f0dfa15b26b"}, + {file = "pycares-4.4.0-cp39-cp39-win32.whl", hash = "sha256:97892cced5794d721fb4ff8765764aa4ea48fe8b2c3820677505b96b83d4ef47"}, + {file = "pycares-4.4.0-cp39-cp39-win_amd64.whl", hash = "sha256:917f08f0b5d9324e9a34211e68d27447c552b50ab967044776bbab7e42a553a2"}, + {file = "pycares-4.4.0.tar.gz", hash = "sha256:f47579d508f2f56eddd16ce72045782ad3b1b3b678098699e2b6a1b30733e1c2"}, +] + +[package.dependencies] +cffi = ">=1.5.0" + +[package.extras] +idna = ["idna (>=2.1)"] + [[package]] name = "pycparser" version = "2.22" @@ -3513,13 +3622,13 @@ files = [ [[package]] name = "requests" -version = "2.32.3" +version = "2.31.0" description = "Python HTTP for Humans." optional = false -python-versions = ">=3.8" +python-versions = ">=3.7" files = [ - {file = "requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6"}, - {file = "requests-2.32.3.tar.gz", hash = "sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760"}, + {file = "requests-2.31.0-py3-none-any.whl", hash = "sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f"}, + {file = "requests-2.31.0.tar.gz", hash = "sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1"}, ] [package.dependencies] @@ -3926,6 +4035,20 @@ watchdog = {version = ">=2.1.5", markers = "platform_system != \"Darwin\""} [package.extras] snowflake = ["snowflake-connector-python (>=2.8.0)", "snowflake-snowpark-python (>=0.9.0)"] +[[package]] +name = "streamlit-tags" +version = "1.2.8" +description = "Tags custom component for Streamlit" +optional = false +python-versions = ">=3.6" +files = [ + {file = "streamlit_tags-1.2.8-py3-none-any.whl", hash = "sha256:c71b10666f3fce67d8e0b3c089aa50dc48830d310223fb88005b08f157586f95"}, + {file = "streamlit_tags-1.2.8.tar.gz", hash = "sha256:9ea46b21f206dc73164e59e3c800a96c863c90af57afde20115f001a6d986583"}, +] + +[package.dependencies] +streamlit = ">=0.63" + [[package]] name = "tenacity" version = "8.5.0" @@ -4004,6 +4127,17 @@ files = [ {file = "toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f"}, ] +[[package]] +name = "tomlkit" +version = "0.12.5" +description = "Style preserving TOML library" +optional = false +python-versions = ">=3.7" +files = [ + {file = "tomlkit-0.12.5-py3-none-any.whl", hash = "sha256:af914f5a9c59ed9d0762c7b64d3b5d5df007448eb9cd2edc8a46b1eafead172f"}, + {file = "tomlkit-0.12.5.tar.gz", hash = "sha256:eef34fba39834d4d6b73c9ba7f3e4d1c417a4e56f89a7e96e090dd0d24b8fb3c"}, +] + [[package]] name = "toolz" version = "0.12.1" @@ -4797,4 +4931,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"] [metadata] lock-version = "2.0" python-versions = ">=3.12.*, <4.0" -content-hash = "4592f4cc15dcf7566fc314c127df2d649f20ecc05c2a672d2c982ad5df5174e1" +content-hash = "7159216c0183970d94b3b819c5460bac914e8ddbbdc96bb618791c5978899b64" diff --git a/wren-ai-service/pyproject.toml b/wren-ai-service/pyproject.toml index 58894d574..b659ed1e5 100644 --- a/wren-ai-service/pyproject.toml +++ b/wren-ai-service/pyproject.toml @@ -21,7 +21,7 @@ sqlparse = "==0.5.0" sqlglot = "==22.5.0" orjson = "==3.10.3" sf-hamilton = {version = "==1.63.0", extras = ["visualization"]} -aiohttp = "==3.9.5" +aiohttp = {extras = ["speedups"], version = "==3.9.5"} ollama-haystack = "==0.0.6" langfuse = "==2.35.0" ollama = "==0.2.1" @@ -40,6 +40,13 @@ psycopg2-binary = "==2.9.9" setuptools = "==70.0.0" locust = "==2.28.0" +[tool.poetry.group.eval.dependencies] +tomlkit = "==0.12.5" +streamlit-tags = "==1.2.8" + +[tool.poetry.group.demo.dependencies] +requests = "==2.31.0" + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" diff --git a/wren-ai-service/src/core/engine.py b/wren-ai-service/src/core/engine.py index e1c999efc..d315ff98e 100644 --- a/wren-ai-service/src/core/engine.py +++ b/wren-ai-service/src/core/engine.py @@ -15,7 +15,6 @@ class Engine(metaclass=ABCMeta): self, sql: str, session: aiohttp.ClientSession, - properties: Dict[str, Any] = {}, ) -> Tuple[bool, Optional[Dict[str, Any]]]: ... diff --git a/wren-ai-service/src/force_deploy.py b/wren-ai-service/src/force_deploy.py index ee5c4381f..2cf129d5b 100644 --- a/wren-ai-service/src/force_deploy.py +++ b/wren-ai-service/src/force_deploy.py @@ -1,12 +1,13 @@ import asyncio import os +from pathlib import Path import aiohttp import backoff +from dotenv import load_dotenv -from src.utils import load_env_vars - -load_env_vars() +if Path(".env.dev").exists(): + load_dotenv(".env.dev", override=True) @backoff.on_exception(backoff.expo, aiohttp.ClientError, max_time=60, max_tries=3) diff --git a/wren-ai-service/src/pipelines/ask/followup_generation.py b/wren-ai-service/src/pipelines/ask/followup_generation.py index 76d7c489e..e0620676b 100644 --- a/wren-ai-service/src/pipelines/ask/followup_generation.py +++ b/wren-ai-service/src/pipelines/ask/followup_generation.py @@ -17,7 +17,7 @@ from src.pipelines.ask.components.prompts import ( TEXT_TO_SQL_RULES, text_to_sql_system_prompt, ) -from src.utils import async_timer, init_providers, timer +from src.utils import async_timer, timer from src.web.v1.services.ask import AskRequest logger = logging.getLogger("wren-ai-service") @@ -234,7 +234,7 @@ class FollowUpGeneration(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/pipelines/ask/generation.py b/wren-ai-service/src/pipelines/ask/generation.py index 29e521b8b..5b9741902 100644 --- a/wren-ai-service/src/pipelines/ask/generation.py +++ b/wren-ai-service/src/pipelines/ask/generation.py @@ -19,7 +19,7 @@ from src.pipelines.ask.components.prompts import ( TEXT_TO_SQL_RULES, text_to_sql_system_prompt, ) -from src.utils import async_timer, init_providers, timer +from src.utils import async_timer, timer logger = logging.getLogger("wren-ai-service") @@ -192,7 +192,7 @@ class Generation(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/pipelines/ask/historical_question.py b/wren-ai-service/src/pipelines/ask/historical_question.py index 47059f1b9..321892858 100644 --- a/wren-ai-service/src/pipelines/ask/historical_question.py +++ b/wren-ai-service/src/pipelines/ask/historical_question.py @@ -13,7 +13,6 @@ from src.core.pipeline import BasicPipeline, async_validate from src.core.provider import DocumentStoreProvider, EmbedderProvider from src.utils import ( async_timer, - init_providers, timer, ) @@ -149,7 +148,7 @@ class HistoricalQuestion(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/pipelines/ask/retrieval.py b/wren-ai-service/src/pipelines/ask/retrieval.py index 152b6ad2e..eee0c7be3 100644 --- a/wren-ai-service/src/pipelines/ask/retrieval.py +++ b/wren-ai-service/src/pipelines/ask/retrieval.py @@ -9,7 +9,7 @@ from langfuse.decorators import observe from src.core.pipeline import BasicPipeline, async_validate from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.utils import async_timer, init_providers +from src.utils import async_timer logger = logging.getLogger("wren-ai-service") @@ -83,7 +83,7 @@ class Retrieval(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/pipelines/ask/sql_correction.py b/wren-ai-service/src/pipelines/ask/sql_correction.py index b6cd76d36..38b189a03 100644 --- a/wren-ai-service/src/pipelines/ask/sql_correction.py +++ b/wren-ai-service/src/pipelines/ask/sql_correction.py @@ -17,7 +17,7 @@ from src.pipelines.ask.components.prompts import ( TEXT_TO_SQL_RULES, text_to_sql_system_prompt, ) -from src.utils import async_timer, init_providers, timer +from src.utils import async_timer, timer logger = logging.getLogger("wren-ai-service") @@ -158,7 +158,7 @@ class SQLCorrection(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/pipelines/ask_details/generation.py b/wren-ai-service/src/pipelines/ask_details/generation.py index ad0933fdf..5caca7f5f 100644 --- a/wren-ai-service/src/pipelines/ask_details/generation.py +++ b/wren-ai-service/src/pipelines/ask_details/generation.py @@ -24,7 +24,6 @@ from src.pipelines.ask_details.components.prompts import ( ) from src.utils import ( async_timer, - init_providers, timer, ) @@ -197,7 +196,7 @@ class Generation(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context, observe - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/pipelines/indexing/indexing.py b/wren-ai-service/src/pipelines/indexing/indexing.py index bd16049e3..25f77848f 100644 --- a/wren-ai-service/src/pipelines/indexing/indexing.py +++ b/wren-ai-service/src/pipelines/indexing/indexing.py @@ -17,7 +17,7 @@ from tqdm import tqdm from src.core.pipeline import BasicPipeline, async_validate from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.utils import async_timer, init_providers, timer +from src.utils import async_timer, timer logger = logging.getLogger("wren-ai-service") @@ -124,6 +124,25 @@ class DDLConverter: logger.debug(f"original mdl_json: {mdl}") + ddl_commands = self.get_ddl_commands(mdl) + + return { + "documents": [ + Document( + id=str(i), + meta={"id": str(i)}, + content=ddl_command, + ) + for i, ddl_command in enumerate( + tqdm( + ddl_commands, + desc="indexing ddl commands into the ddl store", + ) + ) + ] + } + + def get_ddl_commands(self, mdl: Dict[str, Any]) -> List[str]: semantics = { "models": [], "relationships": mdl["relationships"], @@ -159,7 +178,7 @@ class DDLConverter: } ) - ddl_commands = ( + return ( self._convert_models_and_relationships( semantics["models"], semantics["relationships"] ) @@ -167,22 +186,6 @@ class DDLConverter: + self._convert_views(semantics["views"]) ) - return { - "documents": [ - Document( - id=str(i), - meta={"id": str(i)}, - content=ddl_command, - ) - for i, ddl_command in enumerate( - tqdm( - ddl_commands, - desc="indexing ddl commands into the ddl store", - ) - ) - ] - } - # TODO: refactor this method def _convert_models_and_relationships( self, models: List[Dict[str, Any]], relationships: List[Dict[str, Any]] @@ -458,7 +461,7 @@ class Indexing(BasicPipeline): if __name__ == "__main__": from langfuse.decorators import langfuse_context - from src.utils import init_langfuse, load_env_vars + from src.utils import init_langfuse, init_providers, load_env_vars load_env_vars() init_langfuse() diff --git a/wren-ai-service/src/prepare_mdl_json.py b/wren-ai-service/src/prepare_mdl_json.py deleted file mode 100644 index cfe34b06a..000000000 --- a/wren-ai-service/src/prepare_mdl_json.py +++ /dev/null @@ -1,37 +0,0 @@ -import argparse -import os -import shutil -from pathlib import Path - -from .utils import load_env_vars - -load_env_vars() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Prepare the MDL JSON file for the Wren engine using the Spider dataset" - ) - parser.add_argument( - "--dataset_name", - type=str, - default=os.getenv("DATASET_NAME"), - help="Database name of the Spider dataset", - ) - args = parser.parse_args() - - dataset_name = args.dataset_name - - assert Path( - f"src/eval/data/{dataset_name}_mdl.json" - ).exists(), f"File not found in src/eval/data: {dataset_name}_mdl.json" - - # remove all files in src/eval/wren-engine/etc/mdl - os.system("rm -rf tools/dev/etc/mdl/*") - os.makedirs("tools/dev/etc/mdl", exist_ok=True) - - # move the file to src/eval/wren-engine/etc/mdl - shutil.copyfile( - f"src/eval/data/{dataset_name}_mdl.json", - f"tools/dev/etc/mdl/{dataset_name}_mdl.json", - ) diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index d186abb6e..7e470f8e8 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -21,7 +21,6 @@ class WrenUI(Engine): self, sql: str, session: aiohttp.ClientSession, - properties: Dict[str, Any] = {}, ) -> Tuple[bool, Optional[Dict[str, Any]]]: async with session.post( f"{self._endpoint}/api/graphql", @@ -46,27 +45,25 @@ class WrenUI(Engine): class WrenIbis(Engine): def __init__(self, endpoint: str = os.getenv("WREN_IBIS_ENDPOINT")): self._endpoint = endpoint + self._source = os.getenv("WREN_IBIS_SOURCE") + self._manifest = os.getenv("WREN_IBIS_MANIFEST") + self._connection_info = ( + orjson.loads(base64.b64decode(os.getenv("WREN_IBIS_CONNECTION_INFO"))) + if os.getenv("WREN_IBIS_CONNECTION_INFO") + else {} + ) async def dry_run_sql( self, sql: str, session: aiohttp.ClientSession, - properties: Dict[str, Any] = { - "source": os.getenv("WREN_IBIS_SOURCE"), - "manifest": os.getenv("WREN_IBIS_MANIFEST"), - "connection_info": orjson.loads( - base64.b64decode(os.getenv("WREN_IBIS_CONNECTION_INFO")) - ) - if os.getenv("WREN_IBIS_CONNECTION_INFO") - else {}, - }, ) -> Tuple[bool, Optional[Dict[str, Any]]]: async with session.post( - f"{self._endpoint}/v2/connector/{properties.get('source', '')}/query?dryRun=true", + f"{self._endpoint}/v2/connector/{self._source}/query?dryRun=true", json={ "sql": remove_limit_statement(add_quotes(sql)), - "manifestStr": properties.get("manifest", ""), - "connectionInfo": properties.get("connection_info", {}), + "manifestStr": self._manifest, + "connectionInfo": self._connection_info, }, ) as response: if response.status == 204: @@ -74,3 +71,23 @@ class WrenIbis(Engine): res = await response.text() return False, res + + +@provider("wren_engine") +class WrenEngine(Engine): + def __init__(self, endpoint: str = os.getenv("WREN_ENGINE_ENDPOINT")): + self._endpoint = endpoint + + async def dry_run_sql( + self, + sql: str, + session: aiohttp.ClientSession, + ) -> Tuple[bool, Optional[Dict[str, Any]]]: + async with session.get( + f"{self._endpoint}/v1/mdl/dry-run", + json={"sql": remove_limit_statement(add_quotes(sql)), "limit": 1}, + ) as response: + if response.status == 200: + return True, None + res = await response.text() + return False, res diff --git a/wren-ai-service/tests/pytest/providers/test_loader.py b/wren-ai-service/tests/pytest/providers/test_loader.py index 50deb557b..5c6fe910f 100644 --- a/wren-ai-service/tests/pytest/providers/test_loader.py +++ b/wren-ai-service/tests/pytest/providers/test_loader.py @@ -3,7 +3,7 @@ from src.providers import loader def test_import_mods(): loader.import_mods("src.providers") - assert len(loader.PROVIDERS) == 9 + assert len(loader.PROVIDERS) == 10 def test_get_provider(): @@ -39,3 +39,6 @@ def test_get_provider(): provider = loader.get_provider("wren_ibis") assert provider.__name__ == "WrenIbis" + + provider = loader.get_provider("wren_engine") + assert provider.__name__ == "WrenEngine" diff --git a/wren-ai-service/tools/dev/docker-compose-dev.yaml b/wren-ai-service/tools/dev/docker-compose-dev.yaml index 3a3f0273d..1cbfea987 100644 --- a/wren-ai-service/tools/dev/docker-compose-dev.yaml +++ b/wren-ai-service/tools/dev/docker-compose-dev.yaml @@ -28,6 +28,7 @@ services: - ${IBIS_SERVER_PORT}:${IBIS_SERVER_PORT} environment: WREN_ENGINE_ENDPOINT: http://engine:${WREN_ENGINE_PORT} + LOG_LEVEL: DEBUG networks: - wren