diff --git a/docker/.env.example b/docker/.env.example index 928fba918..db2619ed3 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -21,6 +21,8 @@ IBIS_SERVER_VERSION=0.5.0 WREN_BOOTSTRAP_VERSION=0.1.4 # keys +LLM_PROVIDER=openai +DOCUMENT_STORE_PROVIDER=qdrant # CHANGE THIS TO YOUR OPENAI API KEY OPENAI_API_KEY=sk-1234567890 OPENAI_API_BASE=https://api.openai.com/v1 diff --git a/docker/docker-compose-dev.yaml b/docker/docker-compose-dev.yaml index d256d98d9..ba5668eb9 100644 --- a/docker/docker-compose-dev.yaml +++ b/docker/docker-compose-dev.yaml @@ -43,6 +43,8 @@ services: - ${AI_SERVICE_FORWARD_PORT}:${WREN_AI_SERVICE_PORT} environment: WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT} + LLM_PROVIDER: ${LLM_PROVIDER} + DOCUMENT_STORE_PROVIDER: ${DOCUMENT_STORE_PROVIDER} OPENAI_API_KEY: ${OPENAI_API_KEY} OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL} QDRANT_HOST: qdrant diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index ec4a020cb..824c41266 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -55,6 +55,8 @@ services: - ${AI_SERVICE_FORWARD_PORT}:${WREN_AI_SERVICE_PORT} environment: WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT} + LLM_PROVIDER: ${LLM_PROVIDER} + DOCUMENT_STORE_PROVIDER: ${DOCUMENT_STORE_PROVIDER} OPENAI_API_KEY: ${OPENAI_API_KEY} OPENAI_API_BASE: ${OPENAI_API_BASE} OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL} diff --git a/wren-ai-service/.env.dev.example b/wren-ai-service/.env.dev.example index c7871af4d..899b9a283 100644 --- a/wren-ai-service/.env.dev.example +++ b/wren-ai-service/.env.dev.example @@ -4,7 +4,7 @@ WREN_AI_SERVICE_PORT=5556 # app related # LLM Provider name should be mapped to Haystack's supported LLM providers: https://docs.haystack.deepset.ai/v2.0/docs/generators -LLM_PROVIDER=openai # azureopenai as well +LLM_PROVIDER=openai # azure_openai as well # Document Store Provider name should be mapped to Haystack's supported Document Store providers: see the haystack documentation's left sidebar diff --git a/wren-ai-service/.env.prod.example b/wren-ai-service/.env.prod.example index 100e59a4d..f00a43309 100644 --- a/wren-ai-service/.env.prod.example +++ b/wren-ai-service/.env.prod.example @@ -11,7 +11,7 @@ WREN_AI_SERVICE_PORT=5555 # app related # LLM Provider name should be mapped to Haystack's supported LLM providers: https://docs.haystack.deepset.ai/v2.0/docs/generators -LLM_PROVIDER=openai # azureopenai as well +LLM_PROVIDER=openai # azure_openai as well # Document Store Provider name should be mapped to Haystack's supported Document Store providers: see the haystack documentation's left sidebar diff --git a/wren-ai-service/src/pipelines/ask/components/prompts.py b/wren-ai-service/src/pipelines/ask/components/prompts.py index 5e0328604..32e5faafe 100644 --- a/wren-ai-service/src/pipelines/ask/components/prompts.py +++ b/wren-ai-service/src/pipelines/ask/components/prompts.py @@ -8,6 +8,20 @@ TEXT_TO_SQL_RULES = """ - DON'T USE "DATE_ADD" or "DATE_SUB" functions for date operations, instead use syntax like this "current_date - INTERVAL '7' DAY"! - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. + +- SHOWING ALIAS FOR COLUMNS AND TABLES IS MANDATORY. + - ALWAYS SHOW alias for columns and tables such as SELECT [column_name] AS [alias_column_name]. + - MUST USE the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for the column/table alias. + - EXAMPLE + DATABASE SCHEMA + /* {"displayName":"_orders","description":"A model representing the orders data."} */ + CREATE TABLE orders ( + -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} + ApprovedTimestamp TIMESTAMP + } + + SQL + SELECT ApprovedTimestamp AS _timestamp FROM orders AS _orders; """ diff --git a/wren-ai-service/src/pipelines/ask/generation.py b/wren-ai-service/src/pipelines/ask/generation.py index 23a9fae37..282827167 100644 --- a/wren-ai-service/src/pipelines/ask/generation.py +++ b/wren-ai-service/src/pipelines/ask/generation.py @@ -29,17 +29,6 @@ generate three SQL statements that could potentially answer each interpreted ver Provide three different interpretations and corresponding SQL queries that reflect these interpretations. Ensure that your SQL queries are diverse, covering a range of possible meanings behind the ambiguous query. -### DATABASE SCHEMA ### -{% for document in documents %} - {{ document.content }} -{% endfor %} - -### EXCLUDED STATEMETS ### -Ensure that the following excluded statements are not used in the generated queries to maintain variety and avoid repetition. -{% for doc in exclude %} - {{ doc.statement}} -{% endfor %} - ### EXAMPLES ### Consider the structure of a generic database which includes common tables like users, orders, products, and transactions. Here are the ambiguous user queries: @@ -68,6 +57,17 @@ SUMMARY 3: Above-average transactions last week. Proceed in a similar manner for the other queries. +### DATABASE SCHEMA ### +{% for document in documents %} + {{ document.content }} +{% endfor %} + +### EXCLUDED STATEMETS ### +Ensure that the following excluded statements are not used in the generated queries to maintain variety and avoid repetition. +{% for doc in exclude %} + {{ doc.statement}} +{% endfor %} + {{ alert }} ### FINAL ANSWER FORMAT ### diff --git a/wren-ai-service/src/pipelines/ask_details/components/prompts.py b/wren-ai-service/src/pipelines/ask_details/components/prompts.py index 74303dafa..3378a7fcd 100644 --- a/wren-ai-service/src/pipelines/ask_details/components/prompts.py +++ b/wren-ai-service/src/pipelines/ask_details/components/prompts.py @@ -48,4 +48,6 @@ Results: - ONLY USE the tables and columns mentioned in the original sql query. - ONLY CHOOSE columns belong to the tables mentioned in the database schema. - ONLY THE last step should not have a CTE name. +- ALWAYS SHOW alias for columns and tables such as SELECT [column_name] AS [alias_column_name]. +- MUST USE alias from the original SQL query. """ diff --git a/wren-ai-service/src/pipelines/indexing/indexing.py b/wren-ai-service/src/pipelines/indexing/indexing.py index 4272d465f..92eb363a6 100644 --- a/wren-ai-service/src/pipelines/indexing/indexing.py +++ b/wren-ai-service/src/pipelines/indexing/indexing.py @@ -197,6 +197,9 @@ class DDLConverter: for column in model["columns"]: if "relationship" not in column: if "properties" in column: + column["properties"]["alias"] = column["properties"].pop( + "displayName", "" + ) comment = f"-- {orjson.dumps(column['properties']).decode("utf-8")}\n " else: comment = "" @@ -217,6 +220,7 @@ class DDLConverter: # Add foreign key constraints based on relationships for relationship in relationships: + comment = f'-- {{"condition": {relationship["condition"]}, "joinType": {relationship["joinType"]}}}\n ' if ( table_name == relationship["models"][0] and relationship["joinType"].upper() == "MANY_TO_ONE" @@ -224,7 +228,7 @@ class DDLConverter: related_table = relationship["models"][1] fk_column = relationship["condition"].split(" = ")[0].split(".")[1] fk_constraint = f"FOREIGN KEY ({fk_column}) REFERENCES {related_table}({primary_keys_map[related_table]})" - columns_ddl.append(fk_constraint) + columns_ddl.append(f"{comment}{fk_constraint}") elif ( table_name == relationship["models"][1] and relationship["joinType"].upper() == "ONE_TO_MANY" @@ -232,7 +236,7 @@ class DDLConverter: related_table = relationship["models"][0] fk_column = relationship["condition"].split(" = ")[1].split(".")[1] fk_constraint = f"FOREIGN KEY ({fk_column}) REFERENCES {related_table}({primary_keys_map[related_table]})" - columns_ddl.append(fk_constraint) + columns_ddl.append(f"{comment}{fk_constraint}") elif ( table_name in relationship["models"] and relationship["joinType"].upper() == "ONE_TO_ONE" @@ -245,9 +249,12 @@ class DDLConverter: relationship["condition"].split(" = ")[index].split(".")[1] ) fk_constraint = f"FOREIGN KEY ({fk_column}) REFERENCES {related_table}({primary_keys_map[related_table]})" - columns_ddl.append(fk_constraint) + columns_ddl.append(f"{comment}{fk_constraint}") if "properties" in model: + model["properties"]["alias"] = model["properties"].pop( + "displayName", "" + ) comment = ( f"\n/* {orjson.dumps(model['properties']).decode("utf-8")} */\n" ) @@ -261,6 +268,8 @@ class DDLConverter: ) ddl_commands.append(create_table_ddl) + logger.debug(f"DDL Commands: {ddl_commands}") + return ddl_commands def _convert_views(self, views: List[Dict[str, Any]]) -> List[str]: diff --git a/wren-ai-service/src/providers/llm/azure_openai.py b/wren-ai-service/src/providers/llm/azure_openai.py index 75912a144..7c10c7bfa 100644 --- a/wren-ai-service/src/providers/llm/azure_openai.py +++ b/wren-ai-service/src/providers/llm/azure_openai.py @@ -1,10 +1,10 @@ import logging import os +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + import backoff import openai -from openai import AsyncAzureOpenAI, AzureOpenAI, Stream, OpenAI from haystack import Document, component -from typing import Dict, Optional, Callable, Any, List, Union, Tuple from haystack.components.embedders import ( AzureOpenAIDocumentEmbedder, AzureOpenAITextEmbedder, @@ -12,7 +12,7 @@ from haystack.components.embedders import ( from haystack.components.generators import AzureOpenAIGenerator from haystack.dataclasses import ChatMessage, StreamingChunk from haystack.utils import Secret - +from openai import AsyncAzureOpenAI, Stream from openai.types.chat import ChatCompletion, ChatCompletionChunk from tqdm import tqdm @@ -85,13 +85,13 @@ class AsyncAzureGenerator(AzureOpenAIGenerator): openai_formatted_messages = [message.to_openai_format() for message in messages] - completion: Union[Stream[ChatCompletionChunk], ChatCompletion] = ( - await self.client.chat.completions.create( - model=self.model, - messages=openai_formatted_messages, - stream=self.streaming_callback is not None, - **generation_kwargs, - ) + completion: Union[ + Stream[ChatCompletionChunk], ChatCompletion + ] = await self.client.chat.completions.create( + model=self.model, + messages=openai_formatted_messages, + stream=self.streaming_callback is not None, + **generation_kwargs, ) completions: List[ChatMessage] = [] @@ -287,7 +287,7 @@ class AsyncAzureDocumentEmbedder(AzureOpenAIDocumentEmbedder): return {"documents": documents, "meta": meta} -@provider("azureopenai") +@provider("azure_openai") class AzureOpenAILLMProvider(LLMProvider): def __init__( self, @@ -300,7 +300,6 @@ class AzureOpenAILLMProvider(LLMProvider): generation_model: str = os.getenv("AZURE_GENERATION_MODEL") or AZURE_GENERATION_MODEL, ): - logger.info(f"Using Azure OpenAI Generation Model: {generation_model}") self.chat_api_key = chat_api_key self.chat_api_base = chat_api_base @@ -316,7 +315,6 @@ class AzureOpenAILLMProvider(LLMProvider): model_kwargs: Optional[Dict[str, Any]] = AZURE_GENERATION_MODEL_KWARGS, system_prompt: Optional[str] = None, ): - return AsyncAzureGenerator( api_key=self.chat_api_key, model=self.generation_model, diff --git a/wren-ai-service/tests/pytest/providers/test_loader.py b/wren-ai-service/tests/pytest/providers/test_loader.py index 2adbf001e..b6e8ad2bb 100644 --- a/wren-ai-service/tests/pytest/providers/test_loader.py +++ b/wren-ai-service/tests/pytest/providers/test_loader.py @@ -15,5 +15,5 @@ def test_get_provider(): provider = loader.get_provider("qdrant") assert provider.__name__ == "QdrantProvider" - provider = loader.get_provider("azureopenai") + provider = loader.get_provider("azure_openai") assert provider.__name__ == "AzureOpenAILLMProvider"