opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-08-29 00:03:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: relocate wren-mdl, remove legacy WrenAI modules

Browse Source 
- wren-mdl/ → core/wren-mdl/ (MDL JSON schema kept on new main)
- wren-ai-service, wren-ui, wren-launcher, docker, deployment, wren-engine,
  misc, and root markdown files removed; preserved on legacy/v1 for
  reference.

The new main is focused on the Open Context Engine: core/, skills/,
sdks/, examples/, docs/.
This commit is contained in:
Jax Liu
2026-05-04 11:31:09 +08:00
parent 01b5b6a655
commit 4efc324501
810 changed files with 0 additions and 161579 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitmodules
-3
View File
@@ -1,3 +0,0 @@
[submodule "wren-engine"]
path = wren-engine
url = git@github.com:Canner/wren-engine.git
CODE_OF_CONDUCT.md
-53
View File
@@ -1,53 +0,0 @@
## Code of Conduct
### Our Pledge
We as members, contributors, and leaders pledge to make participation in our community a harassment-free experience for everyone, regardless of age, body size, visible or invisible disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, caste, color, religion, or sexual identity and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming, diverse, inclusive, and healthy community.
### Our Standards
Examples of behavior that contributes to a positive environment for our community include:
- Demonstrating empathy and kindness toward other people
- Being respectful of differing opinions, viewpoints, and experiences
- Giving and gracefully accepting constructive feedback
- Accepting responsibility and apologizing to those affected by our mistakes, and learning from the experience
- Focusing on what is best not just for us as individuals, but for the overall community
Examples of unacceptable behavior include:
- The use of sexualized language or imagery, and sexual attention or advances of any kind
- Trolling, insulting or derogatory comments, and personal or political attacks
- Public or private harassment
- Publishing others’ private information, such as a physical or email address, without their explicit permission
- Other conduct which could reasonably be considered inappropriate in a professional setting
### Enforcement Responsibilities
Project maintainers are responsible for clarifying and enforcing our standards of acceptable behavior and will take appropriate and fair corrective action in response to any behavior that they deem inappropriate, threatening, offensive, or harmful.
Project maintainers have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, and will communicate reasons for moderation decisions when appropriate.
### Scope
This Code of Conduct applies within all community spaces, and also applies when an individual is officially representing the community in public spaces. Examples of representing our community include using an official e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event.
### Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be reported to the project team responsible for enforcement at [contact@getwren.ai](mailto:contact@getwren.ai). All complaints will be reviewed and investigated promptly and fairly.
All project maintainers are obligated to respect the privacy and security of the reporter of any incident.
Project maintainers who do not follow or enforce the Code of Conduct in good
faith may face temporary or permanent repercussions as determined by other
members of the project's leadership.
### Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 2.1,
available at [https://www.contributor-covenant.org/version/2/1/code_of_conduct/][version]
[homepage]: http://contributor-covenant.org
[version]: https://www.contributor-covenant.org/version/2/1
CONTRIBUTING.md
-164
View File
@@ -1,164 +0,0 @@
# Contributing Guidelines
*Pull requests, bug reports, and all other forms of contribution are welcomed and highly encouraged!* :octocat:
### Contents
- [Code of Conduct](#book-code-of-conduct)
- [Overview](#mag-overview)
- [Contribution Guide of Different Services](#love_letter-contribution-guide-of-different-services)
- [Creating a New Data Source Connector](#electric_plug-creating-a-new-data-source-connector)
> **This guide serves to set clear expectations for everyone involved with the project so that we can improve it together while also creating a welcoming space for everyone to participate. Following these guidelines will help ensure a positive experience for contributors and maintainers.**
## :book: Code of Conduct
Please review our [Code of Conduct](https://github.com/Canner/WrenAI/blob/main/CODE_OF_CONDUCT.md). It is in effect at all times. We expect it to be honored by everyone who contributes to this project. Acting like an asshole will not be tolerated.
## :rocket: Get Started
1. Visit [How Wren AI works?](https://docs.getwren.ai/oss/overview/how_wrenai_works) to understand the architecture of Wren AI
1. After you understand the architecture of Wren AI, understand the scope of the services you want to contribute to.
Check each service's section under [Contribution Guide of Different Services](#love_letter-contribution-guide-of-different-services) to learn how to contribute to each service.
1. If you are dealing with UI-related tasks, such as adding a dark mode, you only need to contribute to the [Wren UI Service](#wren-ui-service).
2. If you are dealing with LLM-related tasks, such as enhancing the prompts used in the LLM pipelines, you only need to contribute to the [Wren AI Service](#wren-ai-service).
3. If you are working on data-source-related tasks, such as fixing a bug in SQL server connector, you will need to contribute to the [Wren Engine Service](#wren-engine-service).
1. If you are not sure which service to contribute to, please reach out to us in [Discord](https://discord.gg/canner) or [GitHub Issues](https://github.com/Canner/WrenAI/issues).
1. It's possible that you need to contribute to multiple services. For example, if you are adding a new data source, you will need to contribute to the [Wren UI Service](#wren-ui-service) and [Wren Engine Service](#wren-engine-service). Follow [Guide for Contributing to Multiple Services](#guide-for-contributing-to-multiple-services) to learn how to contribute to multiple services.
## :love_letter: Contribution Guide of Different Services
### Wren AI Service
Wren AI Service is responsible for LLM-related tasks like converting natural language questions into SQL queries and providing step-by-step SQL breakdowns.
To contribute to Wren AI Service, please refer to the [Wren AI Service Contributing Guide](https://github.com/Canner/WrenAI/blob/main/wren-ai-service/CONTRIBUTING.md)
### Wren UI Service
Wren UI is the client service of WrenAI. It is built with Next.js and TypeScript.
To contribute to Wren UI, you can refer to the [WrenAI/wren-ui/README.md](https://github.com/Canner/WrenAI/blob/main/wren-ui/README.md) file for instructions on how to set up the development environment and run the development server.
### Wren Engine Service
Wren Engine is the backbone of the Wren AI project. The semantic engine for LLMs, bringing business context to AI agents.
To contribute, please refer to [Wren Engine Contributing Guide](https://github.com/Canner/wren-engine/blob/main/ibis-server/docs/CONTRIBUTING.md)
## Guide for Contributing to Multiple Services
We rely on docker-compose to start all services. If you are contributing to multiple services, you could just comment out the services you'd like to start from the source code and change the `env` variables to point to the services you started by yourself.
### Example: Contributing to the [Wren UI Service](#wren-ui-service) and [Wren Engine Service](#wren-engine-service)
If you are contributing to both the [Wren UI Service](#wren-ui-service) and [Wren Engine Service](#wren-engine-service), you should comment out the `wren-engine` service in the `docker/docker-compose-dev.yml` file (note that the UI service is already excluded from `docker/docker-compose-dev.yml`). Then, adjust the environment variables in your `.env` file to point to the services you have started manually. This will ensure that your local development environment correctly interfaces with the services you are working on.
1. Prepare your `.env` file: In the `WrenAI/docker` folder, use the `.env.example` file as a template. Copy this file to create a `.env.local` file.
```sh
# assuming the current directory is wren-ui
cd ../docker
cp .env.example .env.local
```
2. Modify your `.env.local` file: Fill in the `OPENAI_API_KEY` with your OpenAI API keys before starting.
3. In the `WrenAI/docker` folder, copy `config.example.yaml` to `config.yaml` for AI service configuration. Also change `http://wren-ui:3000` to `http://host.docker.internal:3000` in `config.yaml`.
4. Start the UI and engine services from the source code.
5. Update the `env` variables in the `.env.local` file to point to the services you started manually.
6. Start the other services using docker-compose:
```sh
# current directory is WrenAI/docker
docker-compose -f docker-compose-dev.yaml --env-file .env.example up
# you can add the -d flag to run the services in the background
docker-compose -f docker-compose-dev.yaml --env-file .env.example up -d
# to stop the services, use
docker-compose -f docker-compose-dev.yaml --env-file .env.example down
```
7. Happy coding!
## :electric_plug: Creating a New Data Source Connector
To develop a new data source connector, you'll need to modify both the front-end and back-end of the Wren UI, in addition to the Wren Engine.
Below is a brief overview of a data source connector:
<img src="./misc/data_source.png" width="400">
The UI is primarily responsible for storing database connection settings, providing an interface for users to input these settings, and submitting them to the Engine, which then connects to the database.
The UI must be aware of the connection details it needs to retain, as specified by the Engine. Therefore, the implementation sequence would be as follows:
- Engine:
- Implement the new data source (you'll determine what connection information is needed and how it should be passed from the UI).
- Implement the metadata API for the UI to access.
- UI:
- Back-End:
- Safely store the connection information.
- Provide the connection information to the Engine.
- Front-End:
- Prepare an icon for the data source.
- Set up the form template for users to input the connection information.
- Update the data source list.
### Wren Engine
- To implement a new data source, please refer to [How to Add a New Data Source](https://github.com/Canner/wren-engine/blob/main/ibis-server/docs/how-to-add-data-source.md).
- After adding a new data source, you can proceed with implementing the metadata API for the UI.
Here are some previous PRs that introduced new data sources:
- [Add MSSQL data source](https://github.com/Canner/wren-engine/pull/631)
- [Add MySQL data source](https://github.com/Canner/wren-engine/pull/618)
- [Add ClickHouse data source](https://github.com/Canner/wren-engine/pull/648)
### Wren UI Guide
We'll describe what should be done in the UI for each new data source.
If you prefer to learn by example, you can refer to this Trino [issue](https://github.com/Canner/WrenAI/issues/492) and [PR](https://github.com/Canner/WrenAI/pull/535).
#### Backend
1. Define the data source in `wren-ui/src/apollo/server/dataSource.ts`
- define the `toIbisConnectionInfo` and `sensitiveProps` methods
2. Modify the ibis adaptor in `wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts`
- define an ibis connection info type for the new data source
- set up the `dataSourceUrlMap` for the new data source
3. Modify the repository in `wren-ui/src/apollo/server/repositories/projectRepository.ts`
- define the wren ui connection info type for the new data source
4. Update the graphql schema in `wren-ui/src/apollo/server/schema.ts` so that the new data source can be used in the UI
- add the new data source to the `DataSource` enum
5. Update the type definition in `wren-ui/src/apollo/server/types/dataSource.ts`
- add the new data source to the `DataSourceName` enum
#### Frontend
1. Prepare the data source's logo:
- Image size should be `40 x 40` px
- Preferably use SVG format
- Ensure the logo is centered within a `30px` container for consistent formatting
Example:
<img src="./misc/logo_template.jpg" width="120">
2. Create the data source form template:
- In `wren-ui/src/components/pages/setup/dataSources`, add a new file named `${dataSource}Properties.tsx`
- Implement the data source form template in this file
3. Set up the data source template:
- Navigate to `wren-ui/src/utils/dataSourceType.ts`
- Add new data source image, name, properties
- Update the necessary files to include the new data source template settings
4. Update the data source list:
- Add the new data source to the `DATA_SOURCES` enum in `wren-ui/src/utils/enum/dataSources.ts`
- Update relevant files in `wren-ui/src/components/pages/setup/` to include the new data source
- Ensure `wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts` handle the new data source
5. Test the new connector:
- Ensure the new data source appears in the UI
- Verify that the form works correctly
- Test the connection to the new data source
README.md
-160
View File
@@ -1,160 +0,0 @@
<p align="center" id="top">
<a href="https://getwren.ai/?utm_source=github&utm_medium=title&utm_campaign=readme">
<picture>
<source media="(prefers-color-scheme: light)" srcset="./misc/wrenai_logo.png">
<img src="./misc/wrenai_logo_white.png" width="300px">
</picture>
<h1 align="center">Wren AI - Open-Source GenBI Agent</h1>
</a>
</p>
<p align="center">
<a aria-label="Follow us on X" href="https://x.com/getwrenai">
<img alt="" src="https://img.shields.io/badge/-@getwrenai-blue?style=for-the-badge&logo=x&logoColor=white&labelColor=gray&logoWidth=20">
</a>
<a aria-label="Releases" href="https://github.com/canner/WrenAI/releases">
<img alt="" src="https://img.shields.io/github/v/release/canner/WrenAI?logo=github&label=GitHub%20Release&color=blue&style=for-the-badge">
</a>
<a aria-label="License" href="https://github.com/Canner/WrenAI/blob/main/LICENSE">
<img alt="" src="https://img.shields.io/github/license/canner/WrenAI?color=blue&style=for-the-badge">
</a>
<a aria-label="GitHub Stars" href="https://github.com/Canner/WrenAI/stargazers">
<img alt="" src="https://img.shields.io/github/stars/canner/WrenAI?style=for-the-badge&logo=github&color=blue&label=Stars">
</a>
<a href="https://docs.getwren.ai">
<img src="https://img.shields.io/badge/docs-online-brightgreen?style=for-the-badge" alt="Docs">
</a>
<a aria-label="Join the community on GitHub" href="https://discord.gg/5DvshJqG8Z">
<img alt="" src="https://img.shields.io/badge/-JOIN%20THE%20COMMUNITY-blue?style=for-the-badge&logo=discord&logoColor=white&labelColor=grey&logoWidth=20">
</a>
<a aria-label="Canner" href="https://cannerdata.com/?utm_source=github&utm_medium=badge&utm_campaign=readme">
<img src="https://img.shields.io/badge/%F0%9F%A7%A1-Made%20by%20Canner-blue?style=for-the-badge">
</a>
</p>
<p align="center">
<a href="https://trendshift.io/repositories/9263" target="_blank"><img src="https://trendshift.io/api/badge/repositories/9263" alt="Canner%2FWrenAI | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</p>
> Ask your database anything in plain English. Wren AI generates accurate SQL, charts, and BI insights — backed by a semantic layer that keeps LLM outputs grounded and trustworthy.
<p align="center">
<img width="1920" height="1080" alt="1" src="https://github.com/user-attachments/assets/bba9d37a-33e3-49ab-b7cb-32fd6dddc8d1" />
</p>
## 😍 Demos
https://github.com/user-attachments/assets/f9c1cb34-5a95-4580-8890-ec9644da4160
▶️ [Watch the full GenBI walkthrough](https://github.com/user-attachments/assets/90ad1d35-bb1e-490b-9676-b29863ff090b) — end-to-end from question to chart
## 💡 Why a Semantic Layer?
Feeding raw DDL to an LLM gets you SQL that looks right but means the wrong thing — "revenue" joins the wrong tables, "active user" uses the wrong filter. Wren AI's semantic layer (MDL) encodes your business definitions once, then every generated query is grounded in that shared understanding. The LLM doesn't guess what your metrics mean; the semantic layer tells it.
## 🤖 Features
| | What you get | Why it matters |
|--------------------|--------------|----------------|
| **Talk to Your Data** | Ask in any language → precise SQL & answers | Slash the SQL learning curve |
| **GenBI Insights** | AI-written summaries, charts & reports | Decision-ready context in one click |
| **Semantic Layer** | MDL models encode schema, metrics, joins | Keeps LLM outputs accurate & governed |
| **Embed via API** | Generate queries & charts inside your apps ([API Docs](https://wrenai.readme.io/reference/cloud-getting-started)) | Build custom agents, SaaS features, chatbots ([Streamlit Live Demo](https://huggingface.co/spaces/getWrenAI/wrenai-cloud-api-demo)) |
🤩 [Learn more about GenBI](https://getwren.ai/genbi?utm_source=github&utm_medium=content&utm_campaign=readme)
## 🔌 Data Sources
| Cloud Warehouses | Databases | Query Engines |
|-----------------|-----------|---------------|
| BigQuery | PostgreSQL | Trino |
| Snowflake | MySQL | Athena (Trino) |
| Redshift | Microsoft SQL Server | DuckDB |
| Databricks | ClickHouse | |
| | Oracle | |
Don't see yours? [Vote for it](https://github.com/Canner/WrenAI/discussions/327) — community votes drive our connector roadmap.
## 🧠 LLM Models
Wren AI works with any LLM provider you're already using:
| Cloud APIs | Platform Services | Self-hosted |
|-----------|-------------------|-------------|
| OpenAI | Azure OpenAI | Ollama |
| Anthropic | Google AI Studio (Gemini) | |
| DeepSeek | Vertex AI (Gemini + Anthropic) | |
| Groq | AWS Bedrock | |
| | Databricks | |
> [!TIP]
> For best results, use a frontier model (GPT-4o, Claude Sonnet, Gemini Pro). Wren AI works with smaller and local models too — accuracy scales with model capability. See [configuration examples](https://github.com/Canner/WrenAI/tree/main/wren-ai-service/docs/config_examples) for setup guides.
## 🚀 Getting Started
Three ways to get started — pick what fits:
| Option | Best for | Link |
|--------|----------|------|
| **Self-hosted (Docker)** | Full control, local data | [Installation guide](http://docs.getwren.ai/oss/installation?utm_source=github&utm_medium=content&utm_campaign=readme) |
| **Wren AI Cloud** | Try it without setup | [getwren.ai](https://getwren.ai/?utm_source=github&utm_medium=content&utm_campaign=readme) |
Compare [OSS vs. Cloud plans](https://docs.getwren.ai/oss/overview/cloud_vs_self_host). Full documentation at [docs.getwren.ai](https://docs.getwren.ai/oss/overview/introduction?utm_source=github&utm_medium=content&utm_campaign=readme).
<p align="center">
<img width="1920" height="1080" alt="2" src="https://github.com/user-attachments/assets/6555f539-9ef2-485d-9135-0071741fda96" />
</p>
## 🏗️ Architecture
<p align="center">
<img width="1011" height="682" alt="wrenai-architecture" src="https://github.com/user-attachments/assets/e99b999f-9912-4fa7-921a-9c86b6b83354" />
</p>
User questions flow from the Next.js UI → Apollo GraphQL → AI Service (RAG + LLM) → Wren Engine (semantic query execution) → your database. The semantic layer (MDL) sits at the center, making sure the LLM's SQL reflects your actual business definitions.
👉 [Deep dive into the design](https://getwren.ai/post/how-we-design-our-semantic-engine-for-llms-the-backbone-of-the-semantic-layer-for-llm-architecture?utm_source=github&utm_medium=content&utm_campaign=readme)
## 🧑‍💻 For Developers
WrenAI is a full-stack AI system with interesting problems at every layer — semantic modeling, RAG retrieval, LLM-driven SQL generation, and query execution across heterogeneous data sources. Here's what the stack actually looks like under the hood:
| Layer | What it does |
|-------|-------------|
| **wren-ui** | Next.js + Apollo GraphQL — semantic modeling UI and the BFF that wires everything together |
| **wren-ai-service** | Python/FastAPI pipeline — intent classification, vector retrieval from Qdrant, LLM prompting, and SQL correction loops |
| **[wren-engine](https://github.com/Canner/wren-engine)** | Rust + Apache DataFusion — the query execution core that resolves MDL semantics (metrics, joins, access controls) before SQL reaches the database |
**[wren-engine](https://github.com/Canner/wren-engine)** is a separate open-source project and the part of the stack closest to the metal. It's where MDL definitions get translated into actual query plans across 15+ data sources. If you work with Rust, DataFusion, or database connectors, it's worth a look — the codebase is approachable and there are real unsolved problems around query planning, semantic resolution, and MCP (Model Context Protocol) agent integration.
Some areas where contributions tend to have the most impact across both repos:
- **Data source connectors** — wren-engine supports 15+ sources; new connectors are always useful
- **MCP integration** — wren-engine exposes an MCP server; agent-native workflows are still early and evolving
- **SQL generation quality** — prompt engineering, correction loop heuristics, and eval harnesses in wren-ai-service
- **Semantic layer tooling** — MDL schema inference, validation, and developer ergonomics in wren-ui
## 🛠️ Contribution
1. Read [Contribution Guidelines](https://github.com/Canner/WrenAI/blob/main/CONTRIBUTING.md) for setup & PR guidelines.
2. Open an issue for bugs, feature requests, or discussion.
3. If Wren AI is useful to you, a ⭐ goes a long way — it helps more people find the project.
## ⭐️ Community
- Join 1.7k+ developers in our [Discord](https://discord.gg/5DvshJqG8Z) for real-time help and roadmap previews.
- Visit [GitHub Issues](https://github.com/Canner/WrenAI/issues) for bugs and feature requests.
- Explore our [public roadmap](https://wrenai.notion.site/) to see what's coming next.
- [Subscribe to our blog](https://www.getwren.ai/blog/?utm_source=github&utm_medium=content&utm_campaign=readme) · [Follow us on LinkedIn](https://www.linkedin.com/company/wrenai)
We follow a [Code of Conduct](./CODE_OF_CONDUCT.md) to keep the community welcoming for everyone.
## 🎉 Our Contributors
<a href="https://github.com/canner/wrenAI/graphs/contributors">
<img src="https://contrib.rocks/image?repo=Canner/WrenAI" />
</a>
<p align="right">
<a href="#top">⬆️ Back to Top</a>
</p>
SECURITY.md
-20
View File
@@ -1,20 +0,0 @@
# Security Policy
## Reporting a Vulnerability
If you believe you have found a security vulnerability in any Canner-owned repository, please report it to us through coordinated disclosure.
**Please do not report security vulnerabilities through public GitHub issues, discussions, or pull requests.**
Instead, please send an email to contact[@]cannerdata.com.
Please include as much of the information listed below as you can to help us better understand and resolve the issue:
* The type of issue (e.g., buffer overflow, SQL injection, or cross-site scripting)
* Full paths of source file(s) related to the manifestation of the issue
* The location of the affected source code (tag/branch/commit or direct URL)
* Any special configuration required to reproduce the issue
* Step-by-step instructions to reproduce the issue
* Proof-of-concept or exploit code (if possible)
* Impact of the issue, including how an attacker might exploit the issue
This information will help us triage your report more quickly.
wren-mdl/mdl.schema.json → core/wren-mdl/mdl.schema.json
View File
deployment/README.md
-4
View File
@@ -1,4 +0,0 @@
# Various deployemnt starategies of the app
- [x] [Docker](../docker/)
- [x] [Kubernetes: Kustomizations](./kustomizations/)
deployment/kustomizations/.gitignore
-2
View File
@@ -1,2 +0,0 @@
*.kustomized.yaml
charts/*
deployment/kustomizations/README.md
-101
View File
@@ -1,101 +0,0 @@
# Deployment of Wren AI to Kubernetes with Kustomization
1. Ensure you satisfy the dependencies required to deploy Wren AI.
2. Adjust the values and manifests accordingly to fit your Kubernetes environment.
3. Deploy Secrets separately.
4. Deploy the inflated kustomized app.
Note: Without authentication, once you publish this on the internet, anyone can access your app, see your data, and modify your settings!
## Dependencies used in this kustomization:
- nginx.ingress
- external-dns
- cert-manager
- kubectl kustomize
- helm (for minikube)
## Steps to deploy:
`Suggestion`: Before deploying, check out the manifests in the `deployment/kustomizations ` folder and modify them for your Kubernetes environment.
The `deployment/kustomizations` folder contains a `kustomization.yaml` file that will inflate the manifests into a `deployment/kustomizations/wrenai.kustomized.yaml` file used to deploy the app to your Kubernetes cluster.
```shell
# Clone the repository with the kustomization
git clone https://github.com/Canner/WrenAI.git
cd WrenAI
# Inflate the manifest with kustomization
kubectl kustomize deployment/kustomizations --enable-helm > deployment/kustomizations/wrenai.kustomized.yaml
# Create namespace
kubectl create namespace wren
# !!!!!!!!!!!!
# MODIFY secret-wren_example.yaml manifest file FIRST
# OPENAI_API_KEY is REQUIRED: without a valid key the wren-ai-service-deployment pod will not start
# You must update PG_URL, otherwise wren-ui will not work
#vi deployment/kustomizations/examples/secret-wren_example.yaml
kubectl apply -f deployment/kustomizations/examples/secret-wren_example.yaml -n wren
# Deploy the app:
kubectl apply -f deployment/kustomizations/wrenai.kustomized.yaml
kubectl get pods -n wren
```
### Notes on kustomization:
- `deployment/kustomizations/kustomization.yaml` is the main file responsible for versions of other apps such as Qdrant and PostgreSQL, version of your Wren AI app. It also combines resourses from the manifest such as ConfigMaps, Deployments, and Services. And example Ingress, Certificates and Secrets.
- `deployment/kustomizations/base` is the base folder that contains the core Wren AI manifests, its less likely you need to modify them, but check just in case
- `deployment/kustomizations/examples` is a place with examples of manifests must take a look and adjust to your k8s environment and your needs.
- `deployment/kustomizations/examples/secret-wren_example.yaml` is the file you would not normally include in the kustomization file as its not a best practice and especially not a good idea to include in your GitOps repo as it contains cleartext passwords. We recommend to deploy it separately. Thant's why its commented in the `kustomization.yaml` file.
- `deployment/kustomizations/examples/wrenai-ingress-example.yaml` is an example of how to deploy Ingress. You can use this as a template for your own Ingress. It contains dependancy of extarnal-dns to add your dns name to your DNS records automatically, otherwise you'll need to add it manually. Also it assumes you are using nginx.ingress, it increases timeouts, disables the owasp and modsecurity that might be enabled globaly and prevent your UI from working properly. Comment the TLS section if you do not wish to use `https` encryption. Note: without authentication, enyone can acess your app, see your data and modify your settings!
- `deployment/kustomizations/examples/certificate-wren_example.yaml` is an example of how to deploy certificates for your ingress for the Wren-UI. You can use this as a template for your own certificate. It contains dependancy of cert-manager to add your certificates automatically, otherwise you'll need to add it manually. The certificate will be used by your Ingress.
- `deployment/kustomizations/examples/certificate-qdrant_example.yaml` is an example of how to deploy certificates for your ingress for Qdrant. This is included just in case and is not required, usually you would not be publishing your Vector Database publically in internet. That's why it's commented in the `kustomization.yaml` file. You can use this as a template for your own certificate. It contains dependancy of cert-manager to add your certificates automatically, otherwise you'll need to add it manually.
- `deployment/kustomizations/patches` folder is empty, feel free to add your own patches & overlays there.
#### Wren-UI Database
Starting with wren-ui version 0.6.0 by default the postgres database is used for wren-ui in this kubernetes kustomization and will be installed in the same namespace as wren-ai.
- `postgres`: Database that will be installed in the same namespace as wren-ai. You *must* update `PG_URL` in the Secret manifest `deployment/kustomizations/examples/secret-wren_example.yaml`.
Example: `PG_URL: "postgres://postgres:postgres@wrenai-postgresql:5432/admin_ui"`
- `postgres://` This is the protocol. It tells the system that you’re connecting to a PostgreSQL database.
- `postgres:postgres` These are the username(first) and password(second) for the database respectively, separated by a colon. In this case, both the username and password are “postgres”.
- `@wren-postgresql` This is the hostname of the database server. "wren-postgresql" means the database server is running in a Kubernetes cluster and it is named "wren-postgresql" in the *same* namespace. If you are using another namespace you must provide the full hostname, example: `wren-postgresql.wrenai.svc.cluster.local`, "wrenai" is the namespace name, "svc.cluster.local" is the default domain name for Kubernetes services no need to change it.
- `:5432` This is the port number. PostgreSQL servers listen on port 5432 by default.
- `/admin_ui` This is the name of the database you’re connecting to. In this case, the database name is `admin_ui`. It can be found in the helm values file in the auth.database parameter `deployment/kustomizations/helm-values_postgresql_15.yaml`
# Minikube
Prepare your k8s environment. Then use the `Steps to deploy` section to deploy Wren AI app into your k8s.
```shell
minikube start
minikube addons enable ingress
minikube addons enable metallb
minikube kubectl -- get nodes
minikube kubectl -- get pods -A
minikube update-context
helm repo add bitnami https://charts.bitnami.com/bitnami
helm repo update
helm install external-dns bitnami/external-dns
helm install \
external-dns bitnami/external-dns \
--namespace external-dns \
--version 7.5.2 \
--create-namespace \
--set installCRDs=true
kubectl get pods -n external-dns
helm repo add jetstack https://charts.jetstack.io
helm repo update
helm install \
cert-manager jetstack/cert-manager \
--namespace cert-manager \
--version v1.13.6 \
--create-namespace \
--set installCRDs=true
kubectl get pods -n cert-manager
##########
# Use the `Steps to deploy` section to continue as you would on a production k8s cluster.
```
# GitOps Patches
In the [patches](./patches) folder you can find usefull kustomization examples files if you wish to use existing official kustomization directly from this repo as a base kustomization layer and only customize some values. It can be usefull for you GitOps workflow and can be used in conjunction with FlexCD or ArgoCD.
deployment/kustomizations/base/cm.yaml
-240
View File
@@ -1,240 +0,0 @@
apiVersion: v1
kind: ConfigMap
metadata:
name: wren-config
data:
# Wren Engine Service Port
WREN_ENGINE_PORT: "8080"
# Wren AI Service Port
WREN_AI_SERVICE_PORT: "5555"
WREN_UI_ENDPOINT: http://wren-ui-svc:3000
#Release version used by wren ui https://github.com/Canner/WrenAI/blob/main/docker/docker-compose.yaml#L85-L88
WREN_PRODUCT_VERSION: "0.12.0"
WREN_ENGINE_VERSION: "0.12.3"
WREN_AI_SERVICE_VERSION: "0.12.1"
WREN_UI_VERSION: "0.17.6"
# Document store related
QDRANT_HOST: "wren-qdrant"
# Telemetry
POSTHOG_HOST: "https://app.posthog.com"
TELEMETRY_ENABLED: "false"
# this is for telemetry to know the model, i think ai-service might be able to provide a endpoint to get the information
GENERATION_MODEL: "gpt-4o-mini-2024-07-18"
# service endpoints of AI service & engine service
WREN_ENGINE_ENDPOINT: "http://wren-engine-svc:8080"
WREN_AI_ENDPOINT: "http://wren-ai-service-svc:5555"
#WREN_AI_ENDPOINT: "http://wren-ai-service-svc.ai-system.svc.cluster.local:5555"
# "pg" for postgres as UI application database
WREN_UI_DB_TYPE: pg
#For bootstrap
WREN_ENGINE_DATA_PATH: "/app/data"
### if DB_TYPE = "postgres" you must provide PG_URL string in the *Secret* manifest file (deployment/kustomizations/examples/secret-wren_example.yaml) to connect to postgres
#DEBUG, INFO
LOGGING_LEVEL: INFO
IBIS_SERVER_ENDPOINT: http://wren-ibis-server-svc:8000
---
apiVersion: v1
kind: ConfigMap
metadata:
name: wren-ai-service-config
data:
config.yaml: |
type: llm
provider: litellm_llm
timeout: 120
models:
- alias: default
model: gpt-4.1-nano-2025-04-14
context_window_size: 1000000
kwargs:
max_tokens: 4096
n: 1
seed: 0
temperature: 0
- model: gpt-4.1-mini-2025-04-14
context_window_size: 1000000
kwargs:
max_tokens: 4096
n: 1
seed: 0
temperature: 0
- model: gpt-4.1-2025-04-14
context_window_size: 1000000
kwargs:
max_tokens: 4096
n: 1
seed: 0
temperature: 0
- model: gpt-5-nano-2025-08-07
context_window_size: 380000
kwargs:
max_completion_tokens: 4096
n: 1
seed: 0
reasoning_effort: minimal
- model: gpt-5-mini-2025-08-07
context_window_size: 380000
kwargs:
max_completion_tokens: 4096
n: 1
seed: 0
reasoning_effort: minimal
- model: gpt-5-2025-08-07
context_window_size: 380000
kwargs:
max_completion_tokens: 4096
n: 1
seed: 0
reasoning_effort: minimal
---
type: embedder
provider: litellm_embedder
models:
- model: text-embedding-3-large
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui-svc:3000
---
type: engine
provider: wren_ibis
endpoint: http://wren-ibis-server-svc:8000
---
type: document_store
provider: qdrant
location: http://wren-qdrant:6333
embedding_model_dim: 3072
timeout: 120
---
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
---
settings:
doc_endpoint: https://docs.getwren.ai
is_oss: true
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: false
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
deployment/kustomizations/base/deploy-wren-ai-service.yaml
-69
View File
@@ -1,69 +0,0 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: wren-ai-service-deployment
spec:
replicas: 1
selector:
matchLabels:
app: wren-ai-service
template:
metadata:
labels:
app: wren-ai-service
spec:
containers:
- name: wren-ai-service
image: ghcr.io/canner/wren-ai-service:latest
volumeMounts:
- name: config-volume
mountPath: /app/data
env:
- name: WREN_AI_SERVICE_PORT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_AI_SERVICE_PORT
- name: OPENAI_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: OPENAI_API_KEY
- name: QDRANT_HOST
valueFrom:
configMapKeyRef:
name: wren-config
key: QDRANT_HOST
- name: LOGGING_LEVEL
valueFrom:
configMapKeyRef:
name: wren-config
key: LOGGING_LEVEL
- name: WREN_UI_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_UI_ENDPOINT
- name: PYTHONUNBUFFERED
value: "1"
- name: LANGFUSE_PUBLIC_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LANGFUSE_PUBLIC_KEY
- name: LANGFUSE_SECRET_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LANGFUSE_SECRET_KEY
- name: CONFIG_PATH
value: /app/data/config.yaml
ports:
- containerPort: 5555
volumes:
- name: config-volume
configMap:
name: wren-ai-service-config
items:
- key: config.yaml
path: config.yaml
deployment/kustomizations/base/deploy-wren-engine.yaml
-50
View File
@@ -1,50 +0,0 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: wren-engine-deployment
spec:
replicas: 1
selector:
matchLabels:
app: wren-engine
template:
metadata:
labels:
app: wren-engine
spec:
volumes:
- name: wren-data
persistentVolumeClaim:
claimName: wren-data-pvc
initContainers:
- name: bootstrap
image: ghcr.io/canner/wren-bootstrap:0.1.4
env:
- name: DATA_PATH
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_ENGINE_DATA_PATH
- name: PG_PASSWORD
valueFrom:
secretKeyRef:
name: wrenai-postgresql
key: postgres-password
- name: PG_USERNAME
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: PG_USERNAME
volumeMounts:
- name: wren-data
mountPath: /app/data
command: ["/bin/sh", "/app/init.sh"]
containers:
- name: wren-engine
image: ghcr.io/canner/wren-engine:0.4.4
ports:
- containerPort: 8080
- containerPort: 7432
volumeMounts:
- name: wren-data
mountPath: /usr/src/app/etc
deployment/kustomizations/base/deploy-wren-ibis-server.yaml
-30
View File
@@ -1,30 +0,0 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: wren-ibis-server
spec:
replicas: 1
selector:
matchLabels:
app: wren-ibis
template:
metadata:
labels:
app: wren-ibis
spec:
containers:
- name: wren-ibis
image: ghcr.io/canner/wren-engine-ibis:0.5.0
env:
- name: WREN_ENGINE_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_ENGINE_ENDPOINT
- name: LOGGING_LEVEL
valueFrom:
configMapKeyRef:
name: wren-config
key: LOGGING_LEVEL
ports:
- containerPort: 8000
deployment/kustomizations/base/deploy-wren-ui.yaml
-119
View File
@@ -1,119 +0,0 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: wren-ui-deployment
spec:
replicas: 1
selector:
matchLabels:
app: wren-ui
template:
metadata:
labels:
app: wren-ui
spec:
containers:
- name: wren-ui
image: ghcr.io/canner/wren-ui:0.5.6
env:
- name: DB_TYPE
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_UI_DB_TYPE
- name: WREN_ENGINE_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_ENGINE_ENDPOINT
- name: WREN_AI_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_AI_ENDPOINT
- name: GENERATION_MODEL
valueFrom:
configMapKeyRef:
name: wren-config
key: GENERATION_MODEL
- name: PG_URL
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: PG_URL
# telemetry
- name: WREN_ENGINE_PORT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_ENGINE_PORT
- name: WREN_AI_SERVICE_VERSION
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_AI_SERVICE_VERSION
- name: WREN_UI_VERSION
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_UI_VERSION
- name: WREN_ENGINE_VERSION
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_ENGINE_VERSION
- name: USER_UUID
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: USER_UUID
- name: POSTHOG_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: POSTHOG_API_KEY
- name: POSTHOG_HOST
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: POSTHOG_HOST
- name: TELEMETRY_ENABLED
valueFrom:
configMapKeyRef:
name: wren-config
key: TELEMETRY_ENABLED
# client side
- name: NEXT_PUBLIC_USER_UUID
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: USER_UUID
- name: NEXT_PUBLIC_POSTHOG_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: POSTHOG_API_KEY
- name: NEXT_PUBLIC_POSTHOG_HOST
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: POSTHOG_HOST
- name: NEXT_PUBLIC_TELEMETRY_ENABLED
valueFrom:
configMapKeyRef:
name: wren-config
key: TELEMETRY_ENABLED
# configs
- name: WREN_PRODUCT_VERSION
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_PRODUCT_VERSION
- name: IBIS_SERVER_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: IBIS_SERVER_ENDPOINT
ports:
- containerPort: 3000
deployment/kustomizations/base/pvc.yaml
-11
View File
@@ -1,11 +0,0 @@
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
name: wren-data-pvc
spec:
accessModes:
- ReadWriteOnce
resources:
requests:
storage: 8Gi
# storageClassName: vsphere-retain
deployment/kustomizations/base/svc.yaml
-66
View File
@@ -1,66 +0,0 @@
# Dependancy https://external-dns.io
# You may want to add a DNS record for wren-ui.myhost.net host for your k8s Service instead of Ingress.
# Note: without authentication, enyone can acess your app, see your data and modify your settings!
# If this is the case, make sure to comment the ingress-wren_example.yaml manifest in the `kustomization.yaml` file to exclude it
# And uncomment external-dns in the Service manifest here below:
apiVersion: v1
kind: Service
metadata:
name: wren-ui-svc
#annotations:
### Dependancy external-dns
#external-dns.alpha.kubernetes.io/filter: 'include'
#external-dns.alpha.kubernetes.io/cloudflare-proxied: 'false'
#external-dns.alpha.kubernetes.io/target: wren-ui.myhost.net
spec:
selector:
app: wren-ui
ports:
- protocol: TCP
port: 3000
targetPort: 3000
name: http-ui
---
apiVersion: v1
kind: Service
metadata:
name: wren-engine-svc
spec:
selector:
app: wren-engine
ports:
- protocol: TCP
port: 8080
targetPort: 8080
name: wren-engine
- protocol: TCP
port: 7432
targetPort: 7432
name: wren-engine-sql
---
apiVersion: v1
kind: Service
metadata:
name: wren-ai-service-svc
spec:
selector:
app: wren-ai-service
ports:
- protocol: TCP
port: 5555
targetPort: 5555
name: wren-ai-service
---
apiVersion: v1
kind: Service
metadata:
name: wren-ibis-server-svc
spec:
selector:
app: wren-ibis
ports:
- protocol: TCP
port: 8000
targetPort: 8000
name: wren-ibis
deployment/kustomizations/examples/.gitignore
-1
View File
@@ -1 +0,0 @@
secret.yaml
deployment/kustomizations/examples/certificate-qdrant_example.yaml
-15
View File
@@ -1,15 +0,0 @@
# To generate valid certificates into qdrant-ai.myhost.net-tls
# Dependancy: install https://cert-manager.io
apiVersion: cert-manager.io/v1
kind: Certificate
metadata:
name: qdrant-ai.myhost.net
spec:
dnsNames:
- qdrant-ai.myhost.net
issuerRef:
group: cert-manager.io
kind: ClusterIssuer
#### Replace with the name of your issuer
name: myhost.net-prod
secretName: qdrant-ai.myhost.net-tls
deployment/kustomizations/examples/certificate-wren_example.yaml
-16
View File
@@ -1,16 +0,0 @@
# To generate valid certificates into wren-ui.myhost.net-tls
# Dependancy: install https://cert-manager.io
apiVersion: cert-manager.io/v1
kind: Certificate
metadata:
name: wren-ui.myhost.net
spec:
dnsNames:
- wren-ui.myhost.net
issuerRef:
group: cert-manager.io
kind: ClusterIssuer
### Replace with the name of your issuer, otherwise the secret will be produce randoom name and the ingress will not work.
name: myhost.net-prod
### Your ingress will be looking for this exact name:
secretName: wren-ui.myhost.net-tls
deployment/kustomizations/examples/ingress-wren_example.yaml
-47
View File
@@ -1,47 +0,0 @@
# Dependancy https://external-dns.io
# To add a DNS record for wren-ui.myhost.net host
# Note: without authentication, enyone can acess your app, see your data and modify your settings!
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
name: wren-ui-ingress
annotations:
### Dependancy external-dns
#external-dns.alpha.kubernetes.io/filter: 'include'
#external-dns.alpha.kubernetes.io/cloudflare-proxied: 'false'
##external-dns.alpha.kubernetes.io/cloudflare-proxied: 'true'
external-dns.alpha.kubernetes.io/target: ingress1.myhost.net
#external-dns.alpha.kubernetes.io/target: ingress2.myhost.net
### Dependancy nginx-ingress-controller
nginx.ingress.kubernetes.io/disable-lua: 'true'
nginx.ingress.kubernetes.io/enable-lua: 'false'
nginx.ingress.kubernetes.io/enable-vts-status: 'false'
nginx.ingress.kubernetes.io/enable-modsecurity: 'false'
#nginx.ingress.kubernetes.io/modsecurity-snippet: |
# SecRuleEngine Off
nginx.ingress.kubernetes.io/enable-owasp-modsecurity-crs: 'false'
nginx.ingress.kubernetes.io/proxy-connect-timeout: '360'
nginx.ingress.kubernetes.io/proxy-read-timeout: '360'
nginx.ingress.kubernetes.io/proxy-send-timeout: '360'
spec:
#instead you may use other ingressClassName such as AWS alb. If other than nginx ingress is used, don't forget to comment unsupported annotations above
#"nginx" or "alb"
ingressClassName: nginx
rules:
- host: wren-ui.myhost.net
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: wren-ui-svc
port:
number: 3000
### Comment TLS section if you are not going to use https
tls:
- hosts:
- wren-ui.myhost.net
secretName: wren-ui.myhost.net-tls
deployment/kustomizations/examples/secret-wren_example.yaml
-44
View File
@@ -1,44 +0,0 @@
---
apiVersion: v1
kind: Secret
metadata:
name: wrenai-secrets
type: Opaque
data:
# OPENAI_API_KEY is REQUIRED: without a valid key the wren-ai-service-deployment pod will not start
OPENAI_API_KEY: UkVRVUlSRUQ6IHNrLXByb2otYWxsLWFjY2Vzcy1wbGFjZWhvbGRlci00LXdyZW4tYWktc2VydmljZS1kZXBsb3ltZW50
# Azure openai env
AZURE_CHAT_BASE: bi9h
AZURE_CHAT_KEY: bi9h
AZURE_CHAT_VERSION: bi9h
AZURE_EMBED_BASE: bi9h
AZURE_EMBED_KEY: bi9h
AZURE_EMBED_VERSION: bi9h
# Langfuse: for LLM tracing
LANGFUSE_PUBLIC_KEY: xxxx
LANGFUSE_SECRET_KEY: xxxx
### postgres:// This is the protocol. It tells the system that you’re connecting to a PostgreSQL database.
### postgres:postgres These are the username and password for the database, separated by a colon. In this case, both the username and password are “postgres”.
### @wren-postgresql This is the hostname of the database server. "wren-postgresql" means the database server is running in a Kubernetes cluster and it is named "wren-postgresql" in the *same* namespace. If you are using another namespace you must provide the full hostname, example: `wren-postgresql.wrenai.svc.cluster.local`, "wrenai" is the namespace name, "svc.cluster.local" is the default domain name for Kubernetes services no need to change it.
### :5432 This is the port number. PostgreSQL servers listen on port 5432 by default.
### /admin_ui This is the name of the database you’re connecting to. In this case, the database name is “admin_ui”. It can be found in the helm values file in the auth.database parameter (deployment/kustomizations/helm-values_postgresql_14.yaml)
### PG_URL: "postgres://postgres:postgres@wren-postgresql:5432/admin_ui"
#Fix
PG_URL: cG9zdGdyZXM6Ly9wb3N0Z3Jlczpwb3N0Z3Jlc0B3cmVuLXBvc3RncmVzcWw6NTQzMi9hZG1pbl91aQo=
PG_USERNAME: cG9zdGdyZXM=
POSTHOG_API_KEY: cGhjX2tleS1wbGFjZWhvbGRlcg==
POSTHOG_HOST: aHR0cHM6Ly9hcHAucG9zdGhvZy5jb20=
USER_UUID: MDAwMDAwMDAtMDAwMC0wMDAwLTAwMDAtMDAwMDAwMDAwMDAw
---
apiVersion: v1
kind: Secret
metadata:
name: wrenai-postgresql
data:
postgres-password: cG9zdGdyZXM=
type: Opaque
deployment/kustomizations/helm-values-qdrant_1.11.0.yaml
-234
View File
@@ -1,234 +0,0 @@
replicaCount: 1
image:
repository: docker.io/qdrant/qdrant
pullPolicy: IfNotPresent
tag: "v1.11.0"
useUnprivilegedImage: false
imagePullSecrets: []
nameOverride: ""
fullnameOverride: ""
args: ["./config/initialize.sh"]
env: {}
# - name: QDRANT_ALLOW_RECOVERY_MODE
# value: true
# checks - Readiness and liveness checks can only be enabled for either http (REST) or grpc (multiple checks not supported)
# grpc checks are only available from k8s 1.24+ so as of per default we check http
service:
type: ClusterIP
additionalLabels: {}
annotations: {}
loadBalancerIP: ""
ports:
- name: http
port: 6333
targetPort: 6333
protocol: TCP
checksEnabled: true
- name: grpc
port: 6334
targetPort: 6334
protocol: TCP
checksEnabled: false
- name: p2p
port: 6335
targetPort: 6335
protocol: TCP
checksEnabled: false
ingress:
enabled: false
ingressClassName: "nginx"
additionalLabels: {}
annotations:
# Dependancy: https://kubernetes-sigs.github.io/external-dns
#external-dns.alpha.kubernetes.io/filter: 'include'
#external-dns.alpha.kubernetes.io/cloudflare-proxied: 'true'
external-dns.alpha.kubernetes.io/target: ingress1.myhost.net
hosts:
- host: qdrant-ai.myhost.net
paths:
- path: /
pathType: Prefix
servicePort: 6333
tls:
- hosts:
- qdrant-ai.myhost.net
secretName: qdrant-ai.myhost.net-tls
livenessProbe:
enabled: false
initialDelaySeconds: 5
periodSeconds: 5
timeoutSeconds: 1
failureThreshold: 6
successThreshold: 1
readinessProbe:
enabled: true
initialDelaySeconds: 5
periodSeconds: 5
timeoutSeconds: 1
failureThreshold: 6
successThreshold: 1
startupProbe:
enabled: false
initialDelaySeconds: 10
periodSeconds: 5
timeoutSeconds: 1
failureThreshold: 30
successThreshold: 1
additionalLabels: {}
podAnnotations: {}
podLabels: {}
resources: {}
# limits:
# cpu: 100m
# memory: 128Mi
# requests:
# cpu: 100m
# memory: 128Mi
containerSecurityContext:
runAsNonRoot: true
runAsUser: 1000
runAsGroup: 2000
allowPrivilegeEscalation: false
privileged: false
readOnlyRootFilesystem: true
podSecurityContext:
fsGroup: 3000
fsGroupChangePolicy: Always
lifecycle:
preStop:
exec:
# Sleeping before shutdown allows Qdrant to process requests that were
# in-flight before the node is removed from load-balancing.
# If using an external load balancer, you may need to increase this
# duration to be greater than the LB's health check interval.
command: ["sleep", "3"]
# If true ensures that the pre-existing files on the storage and snapshot volume are owned by the container's
# user and fsGroup
updateVolumeFsOwnership: true
nodeSelector: {}
tolerations: []
affinity: {}
topologySpreadConstraints: []
persistence:
accessModes: ["ReadWriteOnce"]
size: 20Gi
annotations: {}
# storageClassName: vsphere-retain
# If you use snapshots or the snapshot shard transfer mechanism, we recommend
# creating a separate volume of the same size as your main volume so that your
# cluster won't crash if the snapshot is too big.
snapshotPersistence:
enabled: false
accessModes: ["ReadWriteOnce"]
size: 20Gi
annotations: {}
# You can change the storageClassName to ensure snapshots are saved to cold storage.
# storageClassName: local-path
snapshotRestoration:
enabled: false
# Set pvcName if you want to restore from a separately-created PVC. Only supported for single-node clusters unless the PVC is ReadWriteMany.
# If you set snapshotPersistence.enabled and want to restore a snapshot from there, you can leave this blank to skip mounting an external volume.
pvcName: snapshots-pvc
# Must not conflict with /qdrant/snapshots or /qdrant/storage
mountPath: /qdrant/snapshot-restoration
snapshots:
# - /qdrant/snapshot-restoration/test_collection/test_collection-2022-10-24-13-56-50.snapshot:test_collection
# modification example for configuration to overwrite defaults
config:
cluster:
enabled: true
p2p:
port: 6335
consensus:
tick_period_ms: 100
sidecarContainers: []
# sidecarContainers:
# - name: my-sidecar
# image: qdrant/my-sidecar-image
# imagePullPolicy: Always
# ports:
# - name: my-port
# containerPort: 5000
# protocol: TCP
# resources:
# requests:
# memory: 10Mi
# cpu: 10m
# limits:
# memory: 100Mi
# cpu: 100m
metrics:
serviceMonitor:
enabled: false
additionalLabels: {}
scrapeInterval: 30s
scrapeTimeout: 10s
targetPort: http
targetPath: "/metrics"
## MetricRelabelConfigs to apply to samples after scraping, but before ingestion.
## ref: https://github.com/prometheus-operator/prometheus-operator/blob/main/Documentation/api.md#relabelconfig
##
metricRelabelings: []
## RelabelConfigs to apply to samples before scraping
## ref: https://github.com/prometheus-operator/prometheus-operator/blob/main/Documentation/api.md#relabelconfig
##
relabelings: []
serviceAccount:
annotations: {}
priorityClassName: ""
# We disourage changing this setting. Using the "OrderedReady" policy in a
# multi-node cluster will cause a deadlock where nodes refuse to become
# "Ready" until all nodes are running.
podManagementPolicy: Parallel
podDisruptionBudget:
enabled: false
maxUnavailable: 1
# do not enable if you are using not in 1.27
unhealthyPodEvictionPolicy: ""
# minAvailable: 1
# api key for authentication at qdrant
# false: no api key will be configured
# true: an api key will be auto-generated
# string: the given string will be set as an apikey
apiKey: false
# read-only api key for authentication at qdrant
# false: no read-only api key will be configured
# true: an read-only api key will be auto-generated
# string: the given string will be set as a read-only apikey
readOnlyApiKey: false
additionalVolumes: []
# - name: volumeName
# emptyDir: {}
additionalVolumeMounts: []
# - name: volumeName
# mountPath: "/mount/path"
deployment/kustomizations/helm-values_postgresql_15.yaml
-9
View File
@@ -1,9 +0,0 @@
auth:
secretKeys:
adminPasswordKey: postgres-password
existingSecret: wrenai-postgresql
#creates admin_ui database
global:
postgresql:
auth:
database: admin_ui
deployment/kustomizations/kustomization.yaml
-65
View File
@@ -1,65 +0,0 @@
##### Dependancies used in this kustomization:
# nginx.ingress
# external-dns
# cert-manager
# kubectl kustomize
#test output like this:
#kubectl kustomize deployment/kustomizations --enable-helm > deployment/kustomizations/wrenai.kustimized.yaml
#kubectl create namespace wren
#kubectl apply -f deployment/kustomizations/wrenai.kustimized.yaml
apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization
# Namespace for all resources here and in the Helm charts
namespace: wren
helmCharts:
### Uncomment if you are planing to use postgresql
- name: postgresql
repo: https://repo.vmware.com/bitnami-files
version: 15.5.5
releaseName: wren-postgresql
valuesFile: helm-values_postgresql_15.yaml
includeCRDs: true
# the Same Namespace
namespace: wren
- name: qdrant
repo: https://qdrant.github.io/qdrant-helm
version: 1.11.0
releaseName: wren-qdrant
valuesFile: helm-values-qdrant_1.11.0.yaml
includeCRDs: true
# The Same Namespace
namespace: wren
images:
# for the latest versions, please check here: https://github.com/Canner/WrenAI/blob/main/docker/.env.example#L23
- name: ghcr.io/canner/wren-bootstrap
newTag: 0.1.5 # WREN_BOOTSTRAP_VERSION
- name: ghcr.io/canner/wren-engine
newTag: 0.14.8 # WREN_ENGINE_VERSION
- name: ghcr.io/canner/wren-ui
newTag: 0.24.1 # WREN_UI_VERSION
- name: ghcr.io/canner/wren-ai-service
newTag: 0.19.7 # WREN_AI_SERVICE_VERSION
- name: ghcr.io/canner/wren-engine-ibis
newTag: 0.14.8 # IBIS_SERVER_VERSION
resources:
- base/cm.yaml
- base/deploy-wren-ui.yaml
- base/deploy-wren-engine.yaml
- base/deploy-wren-ibis-server.yaml
- base/deploy-wren-ai-service.yaml
- base/pvc.yaml
- base/svc.yaml
### Modify these examples first and uncomment them:
# - examples/ingress-wren_example.yaml
# - examples/certificate-wren_example.yaml
### Usually you do not need to generate a certificate for Qdrant
# - examples/certificate-qdrant_example.yaml
### Best practice is to create and deploy Secrets manually, not as part of kustomization or GitOps!
# - examples/secret-wren_example.yaml
deployment/kustomizations/patches/README.md
-6
View File
@@ -1,6 +0,0 @@
# Example of usefull Patches for Kustomization
Patches from this folder allows to utilize the official unmodified deployment/kustomization dirrectly from the repo as a base layer for your kustomization. And then add patches to update some values. This is usefull for your GitOps and can be combined with tools such as ArgoCD and FluxCD.
Patch ConfigMap, and Service if needed.
Remove Certificate and Ingress if not needed.
deployment/kustomizations/patches/cm.yaml
-35
View File
@@ -1,35 +0,0 @@
- op: replace
path: /data
value:
# Wren Engine Service Port
WREN_ENGINE_PORT: "8080"
# Wren AI Service Port
WREN_AI_SERVICE_PORT: "5555"
#Release version used by wren ui https://github.com/Canner/WrenAI/blob/main/docker/docker-compose.yaml#L85-L88
WREN_PRODUCT_VERSION: "0.12.0"
#fix:
WREN_ENGINE_VERSION: "0.12.3"
WREN_AI_SERVICE_VERSION: "0.12.1"
#fix:
WREN_UI_VERSION: "0.17.6"
# OpenAI
GENERATION_MODEL: "gpt-4o-mini"
# Telemetry
POSTHOG_HOST: "https://app.posthog.com"
TELEMETRY_ENABLED: "false"
# service endpoints of AI service & engine service
WREN_ENGINE_ENDPOINT: "http://wren-engine-svc:8080"
#fix:
WREN_AI_ENDPOINT: "http://wren-ai-service:5555"
# "pg" for postgres as application database.
#fix
DB_TYPE: pg
DATA_PATH: "/app/data"
### if DB_TYPE = "postgres" you must provide PG_URL string in the *Secret* manifest file (deployment/kustomizations/examples/secret-wren_example.yaml) to connect to postgres
deployment/kustomizations/patches/rm-certificate.yaml
-5
View File
@@ -1,5 +0,0 @@
$patch: delete
apiVersion: v1
kind: Certificate
metadata:
name: wren-ui.myhost.net
deployment/kustomizations/patches/rm-ingress.yaml
-7
View File
@@ -1,7 +0,0 @@
$patch: delete
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
name: wren-ui-ingress
annotations:
external-dns.alpha.kubernetes.io/target: ingress1.myhost.net
deployment/kustomizations/patches/service.yaml
-16
View File
@@ -1,16 +0,0 @@
- op: replace
path: /spec/ipFamilies
value:
- IPv6
- IPv4
- op: replace
path: /spec/type
value:
LoadBalancer
- op: replace
path: /spec/ipFamilyPolicy
value:
# SingleStack
PreferDualStack
docker/.env.example
-53
View File
@@ -1,53 +0,0 @@
COMPOSE_PROJECT_NAME=wrenai
PLATFORM=linux/amd64
PROJECT_DIR=.
# service port
WREN_ENGINE_PORT=8080
WREN_ENGINE_SQL_PORT=7432
WREN_AI_SERVICE_PORT=5555
WREN_UI_PORT=3000
IBIS_SERVER_PORT=8000
WREN_UI_ENDPOINT=http://wren-ui:${WREN_UI_PORT}
# ai service settings
QDRANT_HOST=qdrant
SHOULD_FORCE_DEPLOY=1
# vendor keys
OPENAI_API_KEY=
# version
# CHANGE THIS TO THE LATEST VERSION
WREN_PRODUCT_VERSION=0.29.1
WREN_ENGINE_VERSION=0.22.0
WREN_AI_SERVICE_VERSION=0.29.0
IBIS_SERVER_VERSION=0.22.0
WREN_UI_VERSION=0.32.2
WREN_BOOTSTRAP_VERSION=0.1.5
# user id (uuid v4)
USER_UUID=
# for other services
POSTHOG_API_KEY=phc_nhF32aj4xHXOZb0oqr2cn4Oy9uiWzz6CCP4KZmRq9aE
POSTHOG_HOST=https://app.posthog.com
TELEMETRY_ENABLED=true
# this is for telemetry to know the model, i think ai-service might be able to provide a endpoint to get the information.
# AI model configurations should be set in config.yaml, not in .env. See README.md for more details.
GENERATION_MODEL=gpt-4o-mini
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
# the port exposes to the host
# OPTIONAL: change the port if you have a conflict
HOST_PORT=3000
AI_SERVICE_FORWARD_PORT=5555
# Wren UI
EXPERIMENTAL_ENGINE_RUST_VERSION=false
# Wren Engine
# OPTIONAL: set if you want to use local storage for the Wren Engine
LOCAL_STORAGE=.
docker/README.md
-42
View File
@@ -1,42 +0,0 @@
## Service
- `wren-engine`: the engine service. check out example here: [wren-engine
/example](https://github.com/Canner/wren-engine/tree/main/example)
- `wren-ai-service`: the AI service.
- `qdrant`: the vector store ai service is using.
- `wren-ui`: the UI service.
- `bootstrap`: put required files to volume for engine service.
## Volume
Shared data using `data` volume.
Path structure as following:
- `/mdl`
- `*.json` (will put `sample.json` during bootstrap)
- `accounts`
- `config.properties`
## Network
- Check out [Network drivers overview](https://docs.docker.com/engine/network/drivers/) to learn more about `bridge` network driver.
## How to start with OpenAI
1. copy `.env.example` to `.env` and modify the OpenAI API key.
2. copy `config.example.yaml` to `config.yaml` for AI service configuration.
3. start all services: `docker-compose --env-file .env up -d`.
4. stop all services: `docker-compose --env-file .env down`.
### Optional
- If your port 3000 is occupied, you can modify the `HOST_PORT` in `.env`.
## How to start with custom LLM
To start with a custom LLM, the process is similar to starting with OpenAI. The main difference is that you need to modify the `config.yaml` file
that we created on the previous step. After modifying the file, you can restart the services by running `docker-compose --env-file .env up -d --force-recreate wren-ai-service`.
For detailed information on how to modify the configuration for different LLM providers and models, please refer to the [AI Service Configuration](../wren-ai-service/docs/configuration.md).
This guide provides comprehensive instructions on setting up various LLM providers, embedders, and other components of the AI service.
docker/bootstrap/Dockerfile
-4
View File
@@ -1,4 +0,0 @@
FROM busybox
WORKDIR /app
COPY init.sh ./
docker/bootstrap/init.sh
-30
View File
@@ -1,30 +0,0 @@
#!/bin/sh
# declare a variable from the environment variable: DATA_PATH
data_path=${DATA_PATH:-"./"}
# touch a empty config.properties if not exists
# put a content into config.properties if not exists
if [ ! -f ${data_path}/config.properties ]; then
echo "init config.properties"
echo "node.environment=production" >${data_path}/config.properties
fi
# after the config.properties is created, check if config properties properly set
# if not, then append default values to the config.properties
# check if wren.experimental-enable-dynamic-fields is set, otherwise append it with true
if ! grep -q "wren.experimental-enable-dynamic-fields" ${data_path}/config.properties; then
echo "wren.experimental-enable-dynamic-fields is not set, set it to true"
echo "wren.experimental-enable-dynamic-fields=true" >>${data_path}/config.properties
fi
# create a folder mdl if not exists
if [ ! -d ${data_path}/mdl ]; then
echo "create mdl folder"
mkdir ${data_path}/mdl
fi
# put an emtpy sample.json if not exists
if [ ! -f ${data_path}/mdl/sample.json ]; then
echo "init mdl/sample.json"
echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" >${data_path}/mdl/sample.json
fi
docker/config.example.yaml
-193
View File
@@ -1,193 +0,0 @@
type: llm
provider: litellm_llm
timeout: 120
models:
- alias: default
model: gpt-4.1-nano-2025-04-14
context_window_size: 1000000
kwargs:
max_tokens: 4096
n: 1
seed: 0
temperature: 0
- model: gpt-4.1-mini-2025-04-14
context_window_size: 1000000
kwargs:
max_tokens: 4096
n: 1
seed: 0
temperature: 0
- model: gpt-4.1-2025-04-14
context_window_size: 1000000
kwargs:
max_tokens: 4096
n: 1
seed: 0
temperature: 0
- model: gpt-5-nano-2025-08-07
context_window_size: 380000
kwargs:
max_completion_tokens: 4096
n: 1
seed: 0
reasoning_effort: minimal
- model: gpt-5-mini-2025-08-07
context_window_size: 380000
kwargs:
max_completion_tokens: 4096
n: 1
seed: 0
reasoning_effort: minimal
- model: gpt-5-2025-08-07
context_window_size: 380000
kwargs:
max_completion_tokens: 4096
n: 1
seed: 0
reasoning_effort: minimal
---
type: embedder
provider: litellm_embedder
models:
- model: text-embedding-3-large
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072
timeout: 120
recreate_index: true
---
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
doc_endpoint: https://docs.getwren.ai
is_oss: true
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: false
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
docker/docker-compose-dev.yaml
-95
View File
@@ -1,95 +0,0 @@
version: "3"
volumes:
data:
networks:
wren:
driver: bridge
services:
bootstrap:
image: ghcr.io/canner/wren-bootstrap:${WREN_BOOTSTRAP_VERSION}
pull_policy: always
platform: ${PLATFORM}
environment:
DATA_PATH: /app/data
volumes:
- data:/app/data
command: /bin/sh /app/init.sh
wren-engine:
image: ghcr.io/canner/wren-engine:${WREN_ENGINE_VERSION}
pull_policy: always
platform: ${PLATFORM}
expose:
- ${WREN_ENGINE_SQL_PORT}
ports:
- ${WREN_ENGINE_PORT}:${WREN_ENGINE_PORT}
volumes:
- data:/usr/src/app/etc
networks:
- wren
depends_on:
- bootstrap
wren-ai-service:
image: ghcr.io/canner/wren-ai-service:${WREN_AI_SERVICE_VERSION}
pull_policy: always
platform: ${PLATFORM}
ports:
- ${AI_SERVICE_FORWARD_PORT}:${WREN_AI_SERVICE_PORT}
environment:
WREN_UI_ENDPOINT: http://host.docker.internal:${WREN_UI_PORT}
# sometimes the console won't show print messages,
# using PYTHONUNBUFFERED: 1 can fix this
PYTHONUNBUFFERED: 1
CONFIG_PATH: /app/config.yaml
env_file:
- ${PROJECT_DIR}/.env
volumes:
- ${PROJECT_DIR}/config.yaml:/app/config.yaml:ro
- ${PROJECT_DIR}/data:/app/data:ro
networks:
- wren
depends_on:
- qdrant
ibis-server:
image: ghcr.io/canner/wren-engine-ibis:${IBIS_SERVER_VERSION}
pull_policy: always
platform: ${PLATFORM}
expose:
- 8000
ports:
- ${IBIS_SERVER_PORT}:8000
environment:
WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
LOG_LEVEL: DEBUG
networks:
- wren
qdrant:
image: qdrant/qdrant:v1.15.0
pull_policy: always
ports:
- 6333:6333
- 6334:6334
networks:
- wren
# If you want to use postgres for testing purpose, uncomment the following block
# postgres:
# image: postgres:14-alpine
# platform: ${PLATFORM}
# ports:
# - 9432:5432
# volumes:
# - data:/var/lib/postgresql/data
# environment:
# - POSTGRES_PASSWORD=secret
# - POSTGRES_USER=test
# - POSTGRES_DB=test
# - PGDATA=/var/lib/postgresql/data/pgdata
# networks:
# - wren
docker/docker-compose.yaml
-122
View File
@@ -1,122 +0,0 @@
version: "3"
volumes:
data:
networks:
wren:
driver: bridge
services:
bootstrap:
image: ghcr.io/canner/wren-bootstrap:${WREN_BOOTSTRAP_VERSION}
restart: on-failure
platform: ${PLATFORM}
environment:
DATA_PATH: /app/data
volumes:
- data:/app/data
command: /bin/sh /app/init.sh
wren-engine:
image: ghcr.io/canner/wren-engine:${WREN_ENGINE_VERSION}
restart: on-failure
platform: ${PLATFORM}
expose:
- ${WREN_ENGINE_PORT}
- ${WREN_ENGINE_SQL_PORT}
volumes:
- data:/usr/src/app/etc
- ${PROJECT_DIR}/data:/usr/src/app/data
networks:
- wren
depends_on:
- bootstrap
ibis-server:
image: ghcr.io/canner/wren-engine-ibis:${IBIS_SERVER_VERSION}
restart: on-failure
platform: ${PLATFORM}
expose:
- ${IBIS_SERVER_PORT}
environment:
WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
volumes:
- ${LOCAL_STORAGE:-.}:/usr/src/app/data
networks:
- wren
wren-ai-service:
image: ghcr.io/canner/wren-ai-service:${WREN_AI_SERVICE_VERSION}
restart: on-failure
platform: ${PLATFORM}
expose:
- ${WREN_AI_SERVICE_PORT}
ports:
- ${AI_SERVICE_FORWARD_PORT}:${WREN_AI_SERVICE_PORT}
environment:
# sometimes the console won't show print messages,
# using PYTHONUNBUFFERED: 1 can fix this
PYTHONUNBUFFERED: 1
CONFIG_PATH: /app/config.yaml
env_file:
- ${PROJECT_DIR}/.env
volumes:
- ${PROJECT_DIR}/config.yaml:/app/config.yaml:ro
- ${PROJECT_DIR}/data:/app/data:ro
networks:
- wren
depends_on:
- qdrant
qdrant:
image: qdrant/qdrant:v1.15.0
restart: on-failure
expose:
- 6333
- 6334
volumes:
- data:/qdrant/storage
networks:
- wren
wren-ui:
image: ghcr.io/canner/wren-ui:${WREN_UI_VERSION}
restart: on-failure
platform: ${PLATFORM}
environment:
DB_TYPE: sqlite
# /app is the working directory in the container
SQLITE_FILE: /app/data/db.sqlite3
WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
WREN_AI_ENDPOINT: http://wren-ai-service:${WREN_AI_SERVICE_PORT}
IBIS_SERVER_ENDPOINT: http://ibis-server:${IBIS_SERVER_PORT}
# this is for telemetry to know the model, i think ai-service might be able to provide a endpoint to get the information
GENERATION_MODEL: ${GENERATION_MODEL}
# telemetry
WREN_ENGINE_PORT: ${WREN_ENGINE_PORT}
WREN_AI_SERVICE_VERSION: ${WREN_AI_SERVICE_VERSION}
WREN_UI_VERSION: ${WREN_UI_VERSION}
WREN_ENGINE_VERSION: ${WREN_ENGINE_VERSION}
USER_UUID: ${USER_UUID}
POSTHOG_API_KEY: ${POSTHOG_API_KEY}
POSTHOG_HOST: ${POSTHOG_HOST}
TELEMETRY_ENABLED: ${TELEMETRY_ENABLED}
# client side
NEXT_PUBLIC_USER_UUID: ${USER_UUID}
NEXT_PUBLIC_POSTHOG_API_KEY: ${POSTHOG_API_KEY}
NEXT_PUBLIC_POSTHOG_HOST: ${POSTHOG_HOST}
NEXT_PUBLIC_TELEMETRY_ENABLED: ${TELEMETRY_ENABLED}
EXPERIMENTAL_ENGINE_RUST_VERSION: ${EXPERIMENTAL_ENGINE_RUST_VERSION}
# configs
WREN_PRODUCT_VERSION: ${WREN_PRODUCT_VERSION}
ports:
# HOST_PORT is the port you want to expose to the host machine
- ${HOST_PORT}:3000
volumes:
- data:/app/data
networks:
- wren
depends_on:
- wren-ai-service
- wren-engine
misc/AI-generated-understanding_recommend_questions.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 643 KiB

misc/API_doc.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 182 KiB

misc/data_source.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 41 KiB

misc/how_wrenai_works.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 97 KiB

misc/logo_template.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 11 KiB

misc/preview_ask.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 845 KiB

misc/preview_model.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 1.0 MiB

misc/workflow.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 471 KiB

misc/wren-context.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 226 KiB

misc/wren-excel.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 229 KiB

misc/wren-genbi.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 267 KiB

misc/wren-insight.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 213 KiB

misc/wren-lang.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 137 KiB

misc/wren-modeling.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 269 KiB

misc/wren-visual.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 184 KiB

misc/wrenai_logo.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 9.5 KiB

misc/wrenai_logo_white.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 7.7 KiB

misc/wrenai_view.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 58 KiB

misc/wrenai_vision.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 158 KiB

wren-ai-service/.dockerignore
-5
View File
@@ -1,5 +0,0 @@
*
!src
!entrypoint.sh
!pyproject.toml
src/eval
wren-ai-service/.pre-commit-config.yaml
-10
View File
@@ -1,10 +0,0 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.2.2
hooks:
# Run the linter.
- id: ruff
args: [ --fix ]
# Run the formatter.
- id: ruff-format
wren-ai-service/CONTRIBUTING.md
-49
View File
@@ -1,49 +0,0 @@
# Welcome to Wren AI Service contributing guide
Thank you for investing your time in contributing to our project! This document provides guidelines for contributing to the Wren AI service.
## New contributor guide
- To get an overview of the project, please read the [concepts](https://docs.getwren.ai/oss/concept/wren_ai_service).
- To set up the project for local development, please read [Environment Setup](README.md#environment-setup) and [Start the service for development](README.md#start-the-service-for-development)
- To understand the codebase more quickly, we've prepared [a codebase introduciton](docs/code_design.md) for you.
## Getting started
### Issues
#### Create a new issue
If you spot a problem, search if an issue already exists. If a related issue doesn't exist, you can open a new [issue](https://github.com/Canner/WrenAI/issues/new/choose).
#### Solve an issue
Scan through our [existing issues](https://github.com/Canner/WrenAI/issues?q=is%3Aopen+is%3Aissue+label%3Amodule%2Fai-service) to find one that interests you. As a general rule, we don't assign issues to anyone. If you find an issue to work on, you are welcome to open a PR with a fix.
### Pull Request
When you've finished with the changes, create a pull request, also known as a PR.
- Fill the description so that we can review your PR.
- Don't forget to [link PR to issue](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue) if you are solving one.
- Add the `module/ai-service` label to your PR.
- Enable the checkbox to [allow maintainer edits](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/allowing-changes-to-a-pull-request-branch-created-from-a-fork) so the branch can be updated for a merge.
Once you submit your PR, a Canner team member will review your proposal. We may ask questions or request additional information.
- We may ask for changes to be made before a PR can be merged, either using [suggested changes](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/incorporating-feedback-in-your-pull-request) or pull request comments. You can apply suggested changes directly through the UI. You can make any other changes in your fork, then commit them to your branch.
- As you update your PR and apply changes, mark each conversation as [resolved](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/commenting-on-a-pull-request#resolving-conversations).
- Be sure to add one of the prefixes to the PR title, so that our CI could automatically capture the changelog of this PR.
- `feat(wren-ai-service)`: for new features
- `chore(wren-ai-service)`: for maintenance work
- `fix(wren-ai-service)`: for bug fixes
- If you run into any merge issues, checkout this [git tutorial](https://github.com/skills/resolve-merge-conflicts) to help you resolve merge conflicts and other issues.
### Your PR is merged!
Congratulations :tada::tada: The Canner team thanks you :sparkles:.
Once your PR is merged, your contributions will be worked on the next release.
Now that you are part of the Canner community.
## How to add your preferred LLM, Embedder or Document Store
- Please read [this documentation for further details](https://docs.getwren.ai/oss/ai_service/guide/custom_llm#adding-a-custom-llm-embedder-or-document-store-to-wren-ai).
wren-ai-service/Justfile
-69
View File
@@ -1,69 +0,0 @@
GREEN := "\u{001b}[32m"
YELLOW := "\u{001b}[33m"
RESET := "\u{001b}[0m"
## todo: consider to support --override flag to override existing files
init dev='--dev':
@if [ ! -f config.yaml ]; then \
echo "{{GREEN}}config.yaml does not exist. Creating from example...{{RESET}}"; \
cp tools/config/config.example.yaml config.yaml; \
else \
echo "{{YELLOW}}config.yaml already exists. Skipping creation.{{RESET}}"; \
fi
@if [ {{dev}} = "--dev" ] || [ {{dev}} != "--non-dev" ]; then \
if [ ! -f .env.dev ]; then \
echo "{{GREEN}}.env.dev does not exist. Creating from example...{{RESET}}"; \
cp tools/config/.env.dev.example .env.dev; \
else \
echo "{{YELLOW}}.env.dev already exists. Skipping creation.{{RESET}}"; \
fi \
fi
up: prepare-files
docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env up -d
down:
docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env down
start: force_update_config
poetry run python -m src.__main__
curate_eval_data:
poetry run streamlit run eval/data_curation/app.py
prep dataset='spider1.0':
poetry run python -m eval.preparation --dataset {{dataset}}
predict dataset pipeline='ask':
poetry run python -u eval/prediction.py --file {{dataset}} --pipeline {{pipeline}}
eval prediction_result semantics='--no-semantics':
poetry run python -u eval/evaluation.py --file {{prediction_result}} {{semantics}}
test test_args='': up && down
poetry run pytest -s {{test_args}} --ignore tests/pytest/test_usecases.py
test-usecases usecases='all' lang='en':
poetry run python -m tests.pytest.test_usecases --usecases {{usecases}} --lang {{lang}}
load-test:
poetry run python -m tests.locust.locust_script
prepare-files:
# only remove files related to engine and ui
rm -rf tools/dev/etc/duckdb tools/dev/etc/mdl tools/dev/etc/config.properties tools/dev/etc/db.sqlite3 tools/dev/etc/archived
mkdir -p tools/dev/etc
cp tools/dev/config.properties.example tools/dev/etc/config.properties
mkdir -p tools/dev/etc/mdl
echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" \\
> tools/dev/etc/mdl/sample.json
force_update_config:
poetry run python -m src.force_update_config
run-sql mdl_path="" data_source="" sample_dataset="":
poetry run python tools/run_sql.py --mdl-path "{{mdl_path}}" --data-source "{{data_source}}" --sample-dataset "{{sample_dataset}}"
mdl-to-str mdl_path="":
poetry run python tools/mdl_to_str.py -p {{mdl_path}}
wren-ai-service/README.md
-117
View File
@@ -1,117 +0,0 @@
# AI Service of Wren AI
## Concepts
Please read the [documentation](https://docs.getwren.ai/oss/concept/wren_ai_service) here to understand the concepts of Wren AI Service.
## Setup for Local Development
### Prerequisites
1. **Python**: Install Python 3.12.\*
- Recommended: Use [`pyenv`](https://github.com/pyenv/pyenv?tab=readme-ov-file#installation) to manage Python versions
2. **Poetry**: Install Poetry 1.8.3
```bash
curl -sSL https://install.python-poetry.org | python3 - --version 1.8.3
```
3. **Just**: Install [Just](https://github.com/casey/just?tab=readme-ov-file#packages) command runner (version 1.36 or higher)
### Step-by-Step Setup
1. **Install Dependencies**:
```bash
poetry install
```
2. **Generate Configuration Files**:
```bash
just init
```
This creates both `.env.dev` and `config.yaml`. Use `just init --non-dev` to generate only `config.yaml`.
> For Windows, add the line `set shell:= ["bash", "-cu"]` at the start of the Justfile.
4. **Configure Environment**:
- Edit `.env.dev` to set environment variables
- Modify `config.yaml` to configure components, pipelines, and other settings
- Refer to [AI Service Configuration](./docs/configuration.md) for detailed setup instructions
5. **Set Up Development Environment** (optional):
- Install pre-commit hooks:
```bash
poetry run pre-commit install
```
- Run initial pre-commit checks:
```bash
poetry run pre-commit run --all-files
```
6. **Run Tests** (optional):
```bash
just test
```
### Starting the Service
1. **Start Required Containers**:
```bash
just up
```
2. **Launch the AI Service**:
```bash
just start
```
3. **Access the Service**:
- API Documentation: `http://WREN_AI_SERVICE_HOST:WREN_AI_SERVICE_PORT` (default: <http://localhost:5556>)
- User Interface: `http://WREN_UI_HOST:WREN_UI_PORT` (default: <http://localhost:3000>)
4. **Stop the Service**:
When finished, stop the containers:
```bash
just down
```
This setup ensures a consistent development environment and helps maintain code quality through pre-commit hooks and tests. Follow these steps to get started with local development of the Wren AI Service.
## Others
### Pipeline Evaluation
For a comprehensive understanding of how to evaluate the pipelines, please refer to the [evaluation framework](./eval/README.md). This document provides detailed guidelines on the evaluation process, including how to set up and run evaluations, interpret results, and utilize the evaluation metrics effectively. It is a valuable resource for ensuring that the evaluation is conducted accurately and that the results are meaningful.
### Estimate the Speed of the Pipeline(may be outdated)
- to run the load test
- setup `DATASET_NAME` in `.env.dev`
- adjust test config if needed
- adjust user count in `tests/locust/config_users.json`
- in wren-ai-service folder, run `just up` to start the docker containers
- in wren-ai-service folder, run `just start` to start the ai service
- run `just load-test`
- check reports in /outputs/locust folder, there are 3 files with filename **locust*report*{test_timestamp}**:
- .json: test report in json format, including info like llm provider, version
- .html: test report in html format, showing tables and charts
- .log: test log
## Contributing
Thank you for investing your time in contributing to our project! Please [read this for more information](CONTRIBUTING.md)!
wren-ai-service/docker/Dockerfile
-31
View File
@@ -1,31 +0,0 @@
# reference: https://medium.com/@albertazzir/blazing-fast-python-docker-builds-with-poetry-a78a66f5aed0
FROM python:3.12.0-bookworm as builder
RUN pip install poetry==1.8.3
ENV POETRY_NO_INTERACTION=1 \
POETRY_VIRTUALENVS_IN_PROJECT=1 \
POETRY_VIRTUALENVS_CREATE=1 \
POETRY_CACHE_DIR=/tmp/poetry_cache
WORKDIR /app
COPY pyproject.toml ./
RUN poetry install --without dev,eval,test --no-root && rm -rf $POETRY_CACHE_DIR
FROM python:3.12.0-slim-bookworm as runtime
RUN apt-get update && apt install -y netcat-traditional
ENV VIRTUAL_ENV=/app/.venv \
PATH="/app/.venv/bin:$PATH"
COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
COPY src src
COPY entrypoint.sh /app/entrypoint.sh
COPY pyproject.toml pyproject.toml
RUN chmod +x /app/entrypoint.sh
ENTRYPOINT [ "/app/entrypoint.sh" ]
wren-ai-service/docs/code_design.md
-119
View File
@@ -1,119 +0,0 @@
# Introduction to the codebase of wren-ai-service
## Table of Contents
- [Purpose](#purpose)
- [Environment Setup and Start wren-ai-service Locally](#environment-setup-and-start-wren-ai-service-locally)
- [Codebase Introduction](#codebase-introduction)
- [Entrypoint](#entrypoint)
- [Globals](#globals)
- [API endpoints](#api-endpoints)
- [Services](#services)
- [Pipelines](#pipelines)
- [Providers](#providers)
- [Others](#others)
## Purpose
This document aims to dive deep to the implementation details of wren-ai-service. We have two goals in mind while writing the document:
1. You will be more knowledgeable about how wren-ai-service works under the hood.
2. You will be more confident on what part of codebase is needed for adjustment if you would like to be Wren AI's contributor.
## Environment Setup and Start wren-ai-service Locally
If you haven't setup the environment or don't know how to run wren-ai-service locally, please refer to the [document](../README.md#setup-for-local-development) here first.
## Codebase Introduction
wren-ai-service is basically an AI service which provides REST api endpoints for access. There are 4 main concepts to wren-ai-service: `API endpoints`, `Services`, `Pipelines` and `Providers`.
1. `API endpoints`: They are entry points for users to access several kinds of RAG(retrieval-augmented-generation) systems; you can also see API endpoints as encapsulation of Services. For example, when users need to ask a question in order to get SQL, they need to call `/ask` and there is AskService under the hood for background computation.
2. `Services`: They are abstraction of business-logic concepts, such as AskService for users asking questions to get SQL results back, AskDetailsService for users to get SQL breakdown as several sub-steps in order to understand the logic behind the original SQL. Every service is composed of a series of pipelines.
3. `Pipelines`: Basically RAG systems are actually implemented here. However, not all pipelines have complete indexing, retrieval and generation components; it depends on what's the purpose of the pipeline. Also, every pipeline contains some providers such as LLM provider, which represents an LLM.
4. `Providers`: Now there are 4 kinds of providers:
- llm: representing large language models, and now we support OpenAI, Azure OpenAI, OpenAI api-compatible and Ollama models
- embedder: representing embedding models, and now we support OpenAI, Azure OpenAI, OpenAI api-compatible and Ollama models
- document store: representing vector database, and now we use Qdrant
- engine: representing data engine, which is responsible for validating generated SQL's syntax.
### Entrypoint
- The entry point of wren-ai-service is located at [`wren-ai-service/src/__main__.py`](../src/__main__.py)
- The main point of the entry point is the `lifespan` method, which is FastAPI's feature for defining startup and shutdown logic.
```python
# https://fastapi.tiangolo.com/advanced/events/#lifespan
@asynccontextmanager
async def lifespan(app: FastAPI):
# startup events
pipe_components = generate_components()
app.state.service_container = create_service_container(
pipe_components,
column_indexing_batch_size=(
int(os.getenv("COLUMN_INDEXING_BATCH_SIZE"))
if os.getenv("COLUMN_INDEXING_BATCH_SIZE")
else 50
),
table_retrieval_size=(
int(os.getenv("TABLE_RETRIEVAL_SIZE"))
if os.getenv("TABLE_RETRIEVAL_SIZE")
else 10
),
table_column_retrieval_size=(
int(os.getenv("TABLE_COLUMN_RETRIEVAL_SIZE"))
if os.getenv("TABLE_COLUMN_RETRIEVAL_SIZE")
else 1000
),
query_cache={
# the maxsize is a necessary parameter to init cache, but we don't want to expose it to the user
# so we set it to 1_000_000, which is a large number
"maxsize": 1_000_000,
"ttl": int(os.getenv("QUERY_CACHE_TTL") or 120),
},
)
app.state.service_metadata = create_service_metadata(pipe_components)
init_langfuse()
yield
# shutdown events
langfuse_context.flush()
```
- For startup logic, we initialize pipeline components, service containers(which include all services), service metadata(which is some metadata logged for traces inside [Langfuse, an open-source LLM engineering platform](https://langfuse.com/)) and Langfuse.
- For initializing pipeline components, we are in the progress of supporting multiple LLMs, namely users can choose which LLM is responsible for each pipeline.
- You still need to have `.env.dev` locally, then you can prepare `config.yaml` and run `just start`.
- For shutdown logic, we make sure all Langfuse events are transmitted successfully
### Globals
- The file is located at [`wren-ai-service/src/globals.py`](../src/globals.py)
- You can understand the details of service containers and service metadata here
- service containers(Other services are not supported in UI yet)
- SemanticsPreparationService: this is responsible for indexing [MDL](https://docs.getwren.ai/oss/engine/concept/what_is_mdl) to Qdarnt
- AskService: this is responsible for answering users' questions with SQLs, namely text-to-sql
- AskDetailsService: this is responsible for SQL breakdown to several sub-steps
- service metadata
- We will record llm's and embedding model's metadata, wren-ai-service version, etc.
### API endpoints
- All business related API endpoints are located at [`wren-ai-service/src/web/v1/routers`](../src/web/v1/routers)
- Since computation for each kind of API endpoint(ex. ask, etc.) takes several seconds, so we use FastAPI's `background_tasks`. For example, after the `ask` api is invoked, the response is immediately returned, then users need to conduct polling in order to get the latest task status; and once the status is `finished`, the result is returned correspondingly
- Each kind of API endpoint corresponds to one kind of business related task, for example, AskService, AskDetailsService
### Services
- All services are located at [`wren-ai-service/src/web/v1/services`](../src/web/v1/services)
### Pipelines
- All pipelines are located at [`wren-ai-service/src/pipelines`](../src/pipelines)
- Since all pipelines are actually RAG systems, so we classify the role of each pipeline as indexing, retrieval or generation
- The abstract class is defined at [`wren-ai-service/src/core/pipeline.py`](../src/core/pipeline.py)
### Providers
- All providers are located at [`wren-ai-service/src/providers`](../src/providers)
- The abstract classes for providers(LLM, embedding model and document store) are defined at [`wren-ai-service/src/core/provider.py`](../src/core/provider.py)
- The abstract class for engine is defined at [`wren-ai-service/src/core/engine.py`](../src/core/engine.py)
wren-ai-service/docs/config_examples/README.md
-38
View File
@@ -1,38 +0,0 @@
# MUST READ!!!
Since these config files are examples, so **please carefully read the file and comments inside**. Try to understand the purpose of each section and parameter, **don't simply copy and paste the content of these config files into your own config file. It will not work.** For more detailed information to the configurations, please [read this file](../configuration.md).
We also definitely welcome your contribution to add config files for other LLM providers.
## Qwen3 Think and No_Think Configuration
The `config.qwen3.yaml` file provides an example configuration for using Qwen3 models with their unique thinking and non-thinking capabilities. Qwen3 models support two modes:
### Thinking Mode
- Use `/think` in your prompts to enable step-by-step reasoning
- Optimized with `temperature=0.6`, `top_p=0.95`, `top_k=20`
- Best for complex problems requiring detailed reasoning
- Uses the `qwen3-thinking` alias in the pipeline configuration
### Non-Thinking Mode
- Use `/no_think` in your prompts for direct, fast responses
- Optimized with `temperature=0.7`, `top_p=0.8`, `top_k=20`
- Best for simple queries and general conversation
- Uses the `qwen3-fast` alias in the pipeline configuration
### Available Models
- `qwen/qwen3-30b-a3b`: 30B parameter MoE model (3.3B activated)
- `qwen/qwen3-32b`: 32B parameter dense model
- `qwen/qwen3-8b`: 8B parameter dense model
- `qwen/qwen3-14b`: 14B parameter dense model
### Usage Examples
```
# Enable thinking for complex reasoning
"Explain the mathematical proof for the Pythagorean theorem /think"
# Use fast mode for simple queries
"What is the capital of France? /no_think"
```
**Note**: You need to set `OPENROUTER_API_KEY` in your `~/.wrenai/.env` file to use OpenRouter as the provider for Qwen3 models.
wren-ai-service/docs/config_examples/config.anthropic.yaml
-169
View File
@@ -1,169 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
timeout: 120
models:
# put ANTHROPIC_API_KEY=<your_api_key> in ~/.wrenai/.env
- api_base: https://api.anthropic.com
model: anthropic/claude-3-7-sonnet-20250219
alias: default
timeout: 600
kwargs:
n: 1
temperature: 0
---
type: embedder
provider: litellm_embedder
models:
# anthropic embedding model is not supported yet, so we use openai embedding model as a workaround
- model: text-embedding-3-large
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: false
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.azure.yaml
-182
View File
@@ -1,182 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put AZURE_OPENAI_API_KEY=<your_api_key> in ~/.wrenai/.env
- model: azure/gpt-4 # Your Azure deployment name, put 'azure/' before deployment name
alias: default
api_base: https://endpoint.openai.azure.com # Replace with your custom Azure endpoint
api_version: 2024-02-15-preview
timeout: 120
kwargs:
temperature: 0
n: 1
seed: 0
max_tokens: 4096
---
type: embedder
provider: litellm_embedder
models:
# put AZURE_OPENAI_API_KEY=<your_api_key> in ~/.wrenai/.env
- model: azure/text-embedding-ada-002 # Your Azure deployment name, put 'azure/' before deployment name
alias: default
api_base: https://endpoint.openai.azure.com # Replace with your custom Azure endpoint
api_version: 2023-05-15
timeout: 300
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 1536 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant # Match document_store name
llm: litellm_llm.default
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_preparation
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: false
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.bedrock.yaml
-185
View File
@@ -1,185 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
timeout: 120
models:
- model: bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0 # Inference profile ID, us. -> us-east-1, You can get it from "Cross-region inference" section
api_base: https://bedrock-runtime.<region-name>.amazonaws.com
alias: default
kwargs:
temperature: 0.1
aws_region_name: <region-name>
aws_bedrock_runtime_endpoint: https://bedrock-runtime.<region-name>.amazonaws.com
# to authenticate via iam role
aws_role_name: <iam-role-arn>
## other available configs
# aws_access_key_id: <access-key-id>
# aws_secret_access_key: <secret-access-key>
# aws_session_token: <session-token>
# aws_session_name: <session-name>
# aws_profile_name: <profile-name>
# aws_web_identity_token: <web-identity-token>
# aws_sts_endpoint: <sts-endpoint>
---
type: embedder
provider: litellm_embedder
models:
- model: bedrock/amazon.titan-embed-text-v2:0
api_base: https://bedrock-runtime.<region-name>.amazonaws.com
timeout: 600
alias: default
kwargs:
temperature: 0.1
aws_region_name: <region-name>
aws_role_name: <iam-role-arn>
aws_bedrock_runtime_endpoint: https://bedrock-runtime.<region-name>.amazonaws.com
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 1024 # put your embedding model dimension here, 1024 for amazon-titan
timeout: 120
recreate_index: true
---
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: false
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.deepseek.yaml
-192
View File
@@ -1,192 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put DEEPSEEK_API_KEY=<your_api_key> in ~/.wrenai/.env
- api_base: https://api.deepseek.com/v1
model: deepseek/deepseek-reasoner
timeout: 120
kwargs:
n: 1
temperature: 0
response_format:
type: text
- api_base: https://api.deepseek.com/v1
model: deepseek/deepseek-chat
timeout: 120
kwargs:
n: 1
temperature: 0
response_format:
type: text
- api_base: https://api.deepseek.com/v1
model: deepseek/deepseek-coder
alias: default
timeout: 120
kwargs:
n: 1
temperature: 0
response_format:
type: json_object
---
type: embedder
provider: litellm_embedder
models:
# define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
# please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
- model: text-embedding-3-large # put your embedding model name here, if it is not openai embedding model, should be <provider>/<model_name>
alias: default
api_base: https://api.openai.com/v1 # change this according to your embedding model
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.deepseek/deepseek-chat
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.deepseek/deepseek-chat
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.deepseek/deepseek-reasoner
- name: followup_sql_generation_reasoning
llm: litellm_llm.deepseek/deepseek-reasoner
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
-178
View File
@@ -1,178 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put GEMINI_API_KEY=<your_api_key> in ~/.wrenai/.env
- model: gemini/gemini-2.0-flash # gemini/<gemini_model_name>
alias: default
timeout: 120
kwargs:
n: 1
temperature: 0
- model: gemini/gemini-2.0-flash # gemini/<gemini_model_name>
alias: gemini-llm-for-chart
timeout: 120
kwargs:
n: 1
temperature: 0
response_format:
type: json_object
---
type: embedder
provider: litellm_embedder
models:
# put GEMINI_API_KEY=<your_api_key> in ~/.wrenai/.env
- model: gemini/text-embedding-004 # gemini/<gemini_model_name>
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 768 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06 or <provider>.<alias> such as litellm_llm.gemini-llm-for-chart
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.gemini-llm-for-chart
- name: chart_adjustment
llm: litellm_llm.gemini-llm-for-chart
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.google_vertexai.yaml
-186
View File
@@ -1,186 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put VERTEXAI_LOCATION=<your project location> in ~/.wrenai/.env
# put VERTEXAI_PROJECT=<your project id> in ~/.wrenai/.env
# put GOOGLE_APPLICATION_CREDENTIALS=/app/data/<json_file_name> in ~/.wrenai.env
- model: vertex_ai/gemini-2.5-flash # vertex_ai/<gemini_model_name>
alias: default
timeout: 120
kwargs:
n: 1
temperature: 0
seed: 0
max_tokens: 4096
- model: vertex_ai/gemini-2.5-flash # vertex_ai/<gemini_model_name>
alias: gemini-llm-for-chart
timeout: 120
kwargs:
n: 1
temperature: 0
seed: 0
max_tokens: 4096
response_format:
type: json_object
---
type: embedder
provider: litellm_embedder
models:
# put VERTEXAI_LOCATION=<your project location> in ~/.wrenai/.env
# put VERTEXAI_PROJECT=<your project id> in ~/.wrenai/.env
# put GOOGLE_APPLICATION_CREDENTIALS=/app/data/<json_file_name> in ~/.wrenai.env
- model: vertex_ai/text-embedding-004 # vertex_ai/<gemini_model_name>
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 768 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06 or <provider>.<alias> such as litellm_llm.gemini-llm-for-chart
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.gemini-llm-for-chart
- name: chart_adjustment
llm: litellm_llm.gemini-llm-for-chart
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.grok.yaml
-174
View File
@@ -1,174 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put XAI_API_KEY=<your_api_key> in ~/.wrenai/.env
- api_base: https://api.x.ai/v1
model: xai/grok-3-latest
alias: default
timeout: 120
kwargs:
temperature: 0
max_tokens: 4096
n: 1
---
type: embedder
provider: litellm_embedder
models:
# define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
# please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
- model: text-embedding-3-large # put your embedding model name here, if it is not openai embedding model, should be <provider>/<model_name>
alias: default
api_base: https://api.openai.com/v1 # change this according to your embedding model
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.groq.yaml
-173
View File
@@ -1,173 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put GROQ_API_KEY=<your_api_key> in ~/.wrenai/.env
- api_base: https://api.groq.com/openai/v1
model: groq/llama-3.3-70b-specdec # groq/<ollama_model_name>
alias: default
timeout: 120
kwargs:
n: 1
temperature: 0
---
type: embedder
provider: litellm_embedder
models:
# define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
# please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
- model: text-embedding-3-large # put your embedding model name here, if it is not openai embedding model, should be <provider>/<model_name>
alias: default
api_base: https://api.openai.com/v1 # change this according to your embedding model
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.lm_studio.yaml
-172
View File
@@ -1,172 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put LM_STUDIO_API_KEY=<random_string> in ~/.wrenai/.env
- api_base: http://host.docker.internal:1234/v1
model: openai/phi-4 # openai/<lm_studio_model_name>
alias: default
timeout: 600
kwargs:
n: 1
temperature: 0
---
type: embedder
provider: litellm_embedder
models:
# put LM_STUDIO_API_KEY=<random_string> in ~/.wrenai/.env
- model: openai/text-embedding-nomic-embed-text-v1.5 # put your lm_studio embedder model name here, openai/<lm_studio_model_name>
alias: default
api_base: http://host.docker.internal:1234/v1
timeout: 600
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 768 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.ollama.yaml
-172
View File
@@ -1,172 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
- api_base: http://host.docker.internal:11434 # if you are using mac/windows, don't change this; if you are using linux, please search "Run Ollama in docker container" in this page: https://docs.getwren.ai/oss/ai_service/guide/custom_llm#running-wren-ai-with-your-custom-llm-embedder
model: ollama_chat/phi4:14b # ollama_chat/<ollama_model_name>
alias: default
timeout: 600
kwargs:
n: 1
temperature: 0
---
type: embedder
provider: litellm_embedder
models:
# at the moment, we are invoking ollama embedding model through ollama's openai compatible endpoint
# since there is issue invoking ollama embedding model through ollama/<ollama_model_name> using litellm: https://github.com/BerriAI/litellm/issues/7572
- model: openai/nomic-embed-text # put your ollama embedder model name here, openai/<ollama_model_name>
alias: default
api_base: http://host.docker.internal:11434/v1 # if you are using mac/windows, don't change this; if you are using linux, please search "Run Ollama in docker container" in this page: https://docs.getwren.ai/oss/ai_service/guide/custom_llm#running-wren-ai-with-your-custom-llm-embedder
timeout: 600
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 768 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.open_router.yaml
-170
View File
@@ -1,170 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
timeout: 120
models:
# put OPENROUTER_API_KEY=<your_api_key> in ~/.wrenai/.env
- api_base: https://openrouter.ai/api/v1
model: openrouter/anthropic/claude-3.7-sonnet
alias: default
timeout: 600
kwargs:
n: 1
temperature: 0
---
type: embedder
provider: litellm_embedder
models:
# put GEMINI_API_KEY=<your_api_key> in ~/.wrenai/.env
# openrouter embedding model is not supported yet, so you can use gemini embedding model as a workaround
- model: gemini/text-embedding-004
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 768 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.default
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.default
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.default
- name: followup_sql_generation_reasoning
llm: litellm_llm.default
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: false
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.qwen3.yaml
-212
View File
@@ -1,212 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, 3 steps basically:
# 1. you need to use your own llm and embedding models
# 2. fill in embedding model dimension in the document_store section
# 3. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 4. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# put OPENROUTER_API_KEY=<your_api_key> in ~/.wrenai/.env
# Qwen3 models support thinking and non-thinking modes
# Use /think and /no_think in prompts to control reasoning behavior
- api_base: https://openrouter.ai/api/v1
model: openrouter/qwen/qwen3-30b-a3b
alias: default
timeout: 600
kwargs:
n: 1
temperature: 0.6 # Recommended for thinking mode
top_p: 0.95
top_k: 20
response_format:
type: text
- api_base: https://openrouter.ai/api/v1
model: openrouter/qwen/qwen3-30b-a3b
alias: qwen3-thinking
timeout: 600
kwargs:
n: 1
temperature: 0.6 # Optimized for thinking mode
top_p: 0.95
top_k: 20
response_format:
type: text
- api_base: https://openrouter.ai/api/v1
model: openrouter/qwen/qwen3-30b-a3b
alias: qwen3-fast
timeout: 600
kwargs:
n: 1
temperature: 0.7 # Optimized for non-thinking mode
top_p: 0.8
top_k: 20
response_format:
type: text
- api_base: https://openrouter.ai/api/v1
model: openrouter/qwen/qwen3-32b
alias: qwen3-32b
timeout: 600
kwargs:
n: 1
temperature: 0.6
top_p: 0.95
top_k: 20
response_format:
type: json_object
---
type: embedder
provider: litellm_embedder
models:
# define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
# please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
- model: text-embedding-3-large # put your embedding model name here, if it is not openai embedding model, should be <provider>/<model_name>
alias: default
api_base: https://api.openai.com/v1 # change this according to your embedding model
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072 # put your embedding model dimension here
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.qwen3-fast
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.qwen3-fast
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.qwen3-thinking
- name: followup_sql_generation_reasoning
llm: litellm_llm.qwen3-thinking
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/config_examples/config.zhipu.yaml
-220
View File
@@ -1,220 +0,0 @@
# you should rename this file to config.yaml and put it in ~/.wrenai
# please pay attention to the comments starting with # and adjust the config accordingly, steps:
# 1. you need to get your Zhipu AI API key from https://open.bigmodel.cn/
# 2. set your API key in environment variable
# 3. fill in embedding model dimension in the document_store section
# 4. you need to use the correct pipe definitions based on https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
# 5. you need to fill in correct llm and embedding models in the pipe definitions
type: llm
provider: litellm_llm
models:
# GLM-4.5 with thinking disabled - Method 1: Using allowed_openai_params
- api_base: https://open.bigmodel.cn/api/paas/v4/
model: openai/glm-4.5
alias: default
timeout: 900
kwargs:
n: 1
temperature: 0.1
top_p: 0.8
extra_body:
chat_template_kwargs:
enable_thinking: false
allowed_openai_params: ["extra_body"] # Force LiteLLM to allow extra_body
# GLM-4.5 fast mode
- api_base: https://open.bigmodel.cn/api/paas/v4/
model: openai/glm-4.5
alias: glm45-fast
timeout: 900
kwargs:
n: 1
temperature: 0.1
top_p: 0.8
extra_body:
chat_template_kwargs:
enable_thinking: false
allowed_openai_params: ["extra_body"]
# GLM-4.5 with thinking enabled for complex tasks
- api_base: https://open.bigmodel.cn/api/paas/v4/
model: openai/glm-4.5
alias: glm45-thinking
timeout: 1200 # Longer timeout for thinking mode
kwargs:
n: 1
temperature: 0.3
top_p: 0.9
extra_body:
chat_template_kwargs:
enable_thinking: true
allowed_openai_params: ["extra_body"]
# GLM-4.5 for JSON responses with thinking disabled
- api_base: https://open.bigmodel.cn/api/paas/v4/
model: openai/glm-4.5
alias: glm45-json
timeout: 900
kwargs:
n: 1
temperature: 0.05
top_p: 0.7
response_format:
type: json_object
extra_body:
chat_template_kwargs:
enable_thinking: false
allowed_openai_params: ["extra_body"]
---
type: embedder
provider: litellm_embedder
models:
# define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
# GLM series doesn't have dedicated embedding models yet, so we use OpenAI embedding
# please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
- model: text-embedding-3-large
alias: default
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: engine
provider: wren_ibis
endpoint: http://ibis-server:8000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072 # text-embedding-3-large dimension
timeout: 120
recreate_index: true
---
# please change the llm and embedder names to the ones you want to use
# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.glm45-fast
# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
type: pipeline
pipes:
- name: db_schema_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: table_description_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: historical_question_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_correction
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: followup_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: sql_answer
llm: litellm_llm.glm45-fast
- name: semantics_description
llm: litellm_llm.default
- name: relationship_recommendation
llm: litellm_llm.default
engine: wren_ui
- name: question_recommendation
llm: litellm_llm.default
- name: question_recommendation_db_schema_retrieval
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: question_recommendation_sql_generation
llm: litellm_llm.default
engine: wren_ui
document_store: qdrant
- name: chart_generation
llm: litellm_llm.default
- name: chart_adjustment
llm: litellm_llm.default
- name: intent_classification
llm: litellm_llm.default
embedder: litellm_embedder.default
document_store: qdrant
- name: misleading_assistance
llm: litellm_llm.default
- name: data_assistance
llm: litellm_llm.glm45-fast
- name: sql_pairs_indexing
document_store: qdrant
embedder: litellm_embedder.default
- name: sql_pairs_retrieval
document_store: qdrant
embedder: litellm_embedder.default
llm: litellm_llm.default
- name: preprocess_sql_data
llm: litellm_llm.default
- name: sql_executor
engine: wren_ui
- name: user_guide_assistance
llm: litellm_llm.default
- name: sql_question_generation
llm: litellm_llm.default
- name: sql_generation_reasoning
llm: litellm_llm.glm45-thinking
- name: followup_sql_generation_reasoning
llm: litellm_llm.glm45-thinking
- name: sql_regeneration
llm: litellm_llm.default
engine: wren_ui
- name: instructions_indexing
embedder: litellm_embedder.default
document_store: qdrant
- name: instructions_retrieval
embedder: litellm_embedder.default
document_store: qdrant
- name: sql_functions_retrieval
engine: wren_ibis
document_store: qdrant
- name: project_meta_indexing
document_store: qdrant
- name: sql_tables_extraction
llm: litellm_llm.default
- name: sql_diagnosis
llm: litellm_llm.default
- name: sql_knowledge_retrieval
engine: wren_ibis
document_store: qdrant
---
settings:
engine_timeout: 30
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_intent_classification: true
allow_sql_generation_reasoning: true
allow_sql_functions_retrieval: true
enable_column_pruning: false
max_sql_correction_retries: 3
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
logging_level: DEBUG
development: true
historical_question_retrieval_similarity_threshold: 0.9
sql_pairs_similarity_threshold: 0.7
sql_pairs_retrieval_max_size: 10
instructions_similarity_threshold: 0.7
instructions_top_k: 10
wren-ai-service/docs/configuration.md
-143
View File
@@ -1,143 +0,0 @@
# AI Service Configuration
The AI service configuration is managed through a combination of environment variables and a configuration file, providing a flexible and secure approach to setting up the service.
1. **Environment Variables**:
- Used for configuring sensitive information such as vendor API keys
- Specify the configuration file to use
- Allow for partial settings to be configured directly, see [Settings Loading Mechanism](#settings-loading-mechanism) for more details
- Provide a way to override settings in different environments
2. **Configuration File**:
- Used for detailed configuration of components, pipelines, and other service settings
- Allows for more complex and structured configuration options
This dual approach ensures that sensitive data can be kept secure (using environment variables) while allowing for more detailed and shareable configuration through the configuration file. It also provides flexibility in deployment across different environments.
## Settings Loading Mechanism
The AI service uses a hierarchical approach to load settings, ensuring flexibility across different environments and deployment scenarios. The settings are loaded in the following order of precedence:
1. **Default Values**: These are defined as class attributes in the `Settings` class within [`config.py`](../src/config.py). They serve as the base configuration.
2. **Environment Variables**: Using [pydantic-settings](https://fastapi.tiangolo.com/advanced/settings/#pydantic-settings), the service checks for environment variables that match the setting names. If found, these override the default values. For example, `WREN_AI_SERVICE_HOST` can override the default `host` value.
3. **.env.dev File**: The service loads additional settings or overrides existing ones from a `.env.dev` file if present. This is particularly useful for development environments.
4. **config.yaml File**: This file provides the highest priority configuration. It can override all previous settings and is used to configure components, pipelines, and other detailed settings. See [Configuration File](#configuration-file) for more details.
This mechanism allows for easy configuration management across different environments, from development to production, while maintaining security for sensitive information like API keys.
## Configuration File
The configuration file (`config.yaml`) is structured into several sections, each defining different aspects of the AI service. Here's a breakdown of its main components:
1. **LLM Configuration**:
```yaml
type: llm
provider: <provider_name>
models:
- model: <model_name>
kwargs: {}
api_base: <api_endpoint>
```
This component initializes the LLM provider at runtime. You can specify multiple models with different parameters. The `kwargs` field allows for model-specific configurations. For example:
```yaml
type: llm
provider: openai_llm
models:
- model: gpt-4
kwargs:
temperature: 0
n: 1
max_tokens: 4096
response_format:
type: "json_object"
- model: gpt-4o-mini
kwargs: {}
api_base: https://api.openai.com/v1
```
For detailed parameter options, refer to the implementation of the specific LLM provider.
2. **Embedder Configuration**:
```yaml
type: embedder
provider: <provider_name>
models:
- model: <model_name>
dimension: <embedding_size>
api_base: <api_endpoint>
timeout: <timeout_in_seconds>
```
This component configures the embedder, which converts text into numerical vectors. The `provider` specifies the embedder service (e.g., OpenAI, Ollama). You can define multiple `models` with their parameters. The `dimension` parameter indicates the size of the embedding vector.
3. **Engine Configuration**:
```yaml
type: engine
provider: <provider_name>
endpoint: <engine_endpoint>
```
This component configures the engine responsible for generating SQL queries. The `provider` specifies the engine service (e.g., Wren UI).
4. **Document Store Configuration**:
```yaml
type: document_store
provider: <provider_name>
```
This component configures the document store, which is responsible for storing and retrieving embeddings. The `provider` specifies the document store service (e.g., Qdrant).
5. **Pipeline Configuration**:
```yaml
type: pipeline
pipes:
- name: <pipe_name>
llm: <provider>.<model_name>
embedder: <provider>.<model_name>
engine: <provider_name>
document_store: <provider_name>
```
This component configures each pipeline, specifying different LLM, embedder, engine, and document store combinations. For LLM and embedder, use `<provider>.<model_name>`. For engine and document store, use `<provider_name>`.
Example:
```yaml
type: pipeline
pipes:
- name: sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
```
6. **Settings**:
```yaml
settings:
host: <host_address>
port: <port_number>
column_indexing_batch_size: <batch_size>
table_retrieval_size: <retrieval_size>
table_column_retrieval_size: <column_retrieval_size>
query_cache_maxsize: <cache_size>
query_cache_ttl: <cache_ttl_in_seconds>
langfuse_host: <langfuse_endpoint>
langfuse_enable: <true/false>
logging_level: <log_level>
development: <true/false>
```
This section defines various service settings including host, port, indexing and retrieval parameters, cache settings, Langfuse configuration, logging level, and development mode.
This configuration file allows for detailed customization of the AI service components, pipelines, and overall behavior. It provides a centralized place to manage complex configurations while keeping sensitive information separate (managed through environment variables). See [Full Configuration File](../tools/config/config.full.yaml) for a complete example.
wren-ai-service/docs/imgs/shallow_trace_example.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 291 KiB

wren-ai-service/entrypoint.sh
-59
View File
@@ -1,59 +0,0 @@
#!/bin/bash
set -e
INTERVAL=1
TIMEOUT=60
# Wait for qdrant to be responsive
echo "Waiting for qdrant to start..."
current=0
while ! nc -z $QDRANT_HOST 6333; do
sleep $INTERVAL
current=$((current + INTERVAL))
if [ $current -eq $TIMEOUT ]; then
echo "Timeout: qdrant did not start within $TIMEOUT seconds"
exit 1
fi
done
echo "qdrant has started."
# Start wren-ai-service in the background
uvicorn src.__main__:app --host 0.0.0.0 --port $WREN_AI_SERVICE_PORT --loop uvloop --http httptools &
if [[ -n "$SHOULD_FORCE_DEPLOY" ]]; then
# Wait for the server to be responsive
echo "Waiting for wren-ai-service to start..."
current=0
while ! nc -z localhost $WREN_AI_SERVICE_PORT; do
sleep $INTERVAL
current=$((current + INTERVAL))
if [ $current -eq $TIMEOUT ]; then
echo "Timeout: wren-ai-service did not start within $TIMEOUT seconds"
exit 1
fi
done
echo "wren-ai-service has started."
# Wait for wren-ui to be responsive
echo "Waiting for wren-ui to start..."
current=0
while ! nc -z wren-ui $WREN_UI_PORT && ! nc -z host.docker.internal $WREN_UI_PORT; do
sleep $INTERVAL
current=$((current + INTERVAL))
if [ $current -eq $TIMEOUT ]; then
echo "Timeout: wren-ui did not start within $TIMEOUT seconds"
exit 1
fi
done
echo "wren-ui has started."
echo "Forcing deployment..."
python -m src.force_deploy
fi
# Bring wren-ai-service to the foreground
wait
wren-ai-service/eval/.gitignore
-2
View File
@@ -1,2 +0,0 @@
.env
config.yaml
wren-ai-service/eval/README.md
-160
View File
@@ -1,160 +0,0 @@
# Evaluation Framework
This document describes the evaluation framework for the Wren AI service. The evaluation framework is designed to assess the performance of the Wren AI service based on the following components:
## Requirements
1. **Install Just**: Download and install [Just](https://github.com/casey/just?tab=readme-ov-file#packages) to run the evaluation framework commands.
2. **Set up Langfuse**: Create an account on [Langfuse](https://cloud.langfuse.com) and obtain the API key and secret. Populate the `.env.dev` file with these credentials.
3. **Start Development Services**: Run `just up` to initiate the necessary development services.
4. **Configuration File**: Ensure you have a copy of `config.yaml` located in the `wren-ai-service/eval/` directory.
## Dataset Curation
The dataset curation process is used to prepare the evaluation dataset for the Wren AI service on evaluation purpose. You can follow the steps below to start the curation app:
- copy `.env.example` to `.env` and fill in the environment variables
- execute the command under the `wren-ai-service` folder: `just curate_eval_data`
## Eval Dataset Preparation(If using Spider 1.0 dataset, or Bird dataset)
```cli
just prep <dataset-name>
```
Currently, we support two datasets for evaluation:
- `spider1.0`: The Spider dataset (default if no dataset specified)
- `bird`: The Bird dataset
The command performs two main steps:
1. Downloads the specified dataset to:
```txt
wren-ai-service/tools/dev/etc/<dataset-name>
```
2. Prepares and saves evaluation datasets to:
```txt
wren-ai-service/eval/dataset
```
The output files follow these naming conventions:
- Spider dataset: `spider_<db_name>_eval_dataset.toml`
- Bird dataset: `bird_<db_name>_eval_dataset.toml`
Each evaluation dataset contains questions, SQL queries, and relevant context needed for testing the system's text-to-SQL capabilities.
## Evaluation Dataset Schema
- dataset_id(UUID)
- date
- mdl
- eval dataset
## Configure the datasource for prediction and evaluation
Before starting the prediction and evaluation process, you need to configure the datasource correctly. This ensures that the system can access the necessary data for making predictions and evaluations.
### For Spider or Bird Datasets
For the Spider or Bird datasets, a built-in datasource is used. This means that the data is stored locally and accessed through a specific path. You need to specify the `eval_data_db_path` in the `config.yaml` file. This path tells the system where to find the database files.
Here's an example of how to set this up in the `config.yaml` file:
```yaml
eval_data_db_path: "etc/bird/minidev/MINIDEV/dev_databases"
```
### Configuring BigQuery as a Datasource for Other custom MDLs
When working with custom MDLs that utilize BigQuery as their datasource, it's crucial to properly configure your system to access the necessary datasets. This involves setting specific parameters in the `config.yaml` file or the `.env.dev` file. Both methods are effective, but using the `.env.dev` file is particularly beneficial for keeping sensitive credentials secure.
#### Encoding the credentials
You can use the following command to encode the credentials:
```cli
cat <path/to/credentials.json> | base64
```
#### Configuration in `config.yaml`
To enable access to your BigQuery dataset, add the following parameters to your `config.yaml` file. This configuration will guide the system in locating and authenticating with your BigQuery resources:
```yaml
bigquery_project_id: "your_project_id"
bigquery_dataset_id: "your_dataset_id"
bigquery_credentials: "your_credentials" # this is a base64 encoded string of the credentials
```
#### Configuration in `.env.dev`
For the `.env.dev` file, you can use the following parameters:
```env
BIGQUERY_PROJECT_ID="your_project_id"
BIGQUERY_DATASET_ID="your_dataset_id"
BIGQUERY_CREDENTIALS="your_credentials" # this is a base64 encoded string of the credentials
```
## Prediction Process
The prediction process is used to produce the results of the evaluation data using the Wren AI service. It will create traces and a session on Langfuse to make the results available to the user. You can use the following command to predict the evaluation dataset under the `eval/dataset` directory:
```cli
just predict <evaluation-dataset>
```
Also, sub-pipeline predictions are supported by specifying the pipeline name:
```cli
just predict <evaluation-dataset> <pipeline-name>
```
Currently, we support the following pipelines: 'ask', 'generation', and 'retrieval'. If no pipeline name is specified, the default is the 'ask' pipeline.
## Evaluation Process
The evaluation process is used to assess the prediction results of the Wren AI service. It compares the prediction results with the ground truth and calculates the evaluation metrics. This process will also add a trace in the same session on Langfuse to make the evaluation results available to the user. You can use the following command to evaluate the prediction results under the `outputs/predictions` directory:
```cli
just eval <prediction-result>
```
Note: If you would like to enable semantics comparison between SQLs by LLM in order to improve the accuracy metric, please fill in Open AI API key in `.env` file in `wren-ai-service/eval` and add `--semantics` to the end of the command like following:
```cli
just eval <prediction-result> --semantics
```
The evaluation results will be presented on Langfuse as follows:
![shallow_trace_example](../docs/imgs/shallow_trace_example.png)
## Terms
This section describes the terms used in the evaluation framework:
- **input**: The user query used as input to the Wren AI service (e.g., "What is the total number of COVID-19 cases in the US?").
- **actual_output**: The actual SQL query generated to retrieve the answer to the user query (e.g., "SELECT SUM(cases) FROM covid19 WHERE country='US'").
- **expected_output**: The expected SQL query that should retrieve the answer to the user query (e.g., "SELECT SUM(cases) FROM covid19 WHERE country='US'").
- **retrieval_context**: The relevant context that helps the LLM generate the SQL query (e.g., "covid19.country", "covid19.cases").
- **context**: The relevant context that aligns with human expectations to generate the SQL query (e.g., "covid19.country", "covid19.cases").
## Metrics
This section describes the evaluation metrics used in the evaluation framework:
- **Accuracy**: This metrics is defined as the proportion of the correct SQL output generated by the model compared to the expected SQL output. It checks if the generated SQL query produces the correct results.
- **Answer Relevancy**: This metric helps determine how well your LLM generates relevant information based on the input it receives. It ensures the efficiency and accuracy of the model's output.
- **Faithfulness**: This metric helps determine how well your LLM generates information that is factually correct and aligned with the retrieval context, minimizing hallucinations and contradictions.
- **Contextual Relevancy**: This metric helps determine how well your retriever minimizes irrelevant information while maximizing the retrieval of relevant information. It ensures the efficiency and accuracy of the retrieval process.
- **Contextual Recall**: This metric helps determine how well the embedding model identifies and retrieves relevant information based on the given context.
- **Contextual Precision**: This metric helps determine how well the reranker places relevant nodes higher in the ranking, ensuring that users get the most pertinent results quickly.
- **QuestionToReasoningJudge**: This metric helps determine how well the LLM generates reasoning that is aligned with the question.
- **ReasoningToSqlJudge**: This metric helps determine how well the LLM generates SQL that is aligned with the reasoning.
- **SqlSemanticsJudge**: This metric helps determine how well the LLM generates SQL that is semantically equivalent to the expected SQL.
wren-ai-service/eval/__init__.py
-62
View File
@@ -1,62 +0,0 @@
from pathlib import Path
from pydantic import Field, SecretStr
from src.config import Settings
SPIDER_DESTINATION_PATH = Path("./tools/dev/etc/spider1.0")
BIRD_DESTINATION_PATH = Path("./tools/dev/etc/bird")
WREN_ENGINE_API_URL = "http://localhost:8080"
EVAL_DATASET_DESTINATION_PATH = Path("./eval/dataset")
class EvalSettings(Settings):
langfuse_project_id: str = ""
batch_size: int = 4
batch_interval: int = 1
datasource: str = "bigquery"
config_path: str = "eval/config.yaml"
openai_api_key: SecretStr = Field(alias="OPENAI_API_KEY")
allow_sql_samples: bool = True
allow_instructions: bool = True
allow_sql_functions: bool = True
eval_data_db_path: str = ""
# BigQuery
bigquery_project_id: str = Field(default="")
bigquery_dataset_id: str = Field(default="")
bigquery_credentials: SecretStr = Field(default="")
# Postgres
postgres_host: str = Field(default="postgres")
postgres_port: str = Field(default="5432")
postgres_user: str = Field(default="postgres")
postgres_password: SecretStr = Field(default="postgres")
postgres_database: str = Field(default="test")
@property
def langfuse_url(self) -> str:
if not self.langfuse_project_id:
return ""
return f"{self.langfuse_host.rstrip('/')}/project/{self.langfuse_project_id}"
def get_openai_api_key(self) -> str:
return self.openai_api_key.get_secret_value()
@property
def bigquery_info(self) -> dict:
return {
"project_id": self.bigquery_project_id,
"dataset_id": self.bigquery_dataset_id,
"credentials": self.bigquery_credentials.get_secret_value(),
}
@property
def postgres_info(self) -> dict:
return {
"host": self.postgres_host,
"port": self.postgres_port,
"user": self.postgres_user,
"password": self.postgres_password.get_secret_value(),
"database": self.postgres_database,
}
wren-ai-service/eval/add_samples_to_toml.py
-36
View File
@@ -1,36 +0,0 @@
import argparse
import tomlkit
from eval.utils import (
get_next_few_items_circular,
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--toml", type=str, help="The toml file name", required=True)
args = parser.parse_args()
if args.toml:
# read toml
with open(f"eval/dataset/{args.toml}", "r") as f:
doc = tomlkit.parse(f.read())
# get the list of question-sql pairs for generating sample values
ground_truth_list = [
{"question": element["question"], "sql": element["sql"]}
for element in doc["eval_dataset"]
]
# utilize utils.get_next_few_items_circular, put n samples in the eval dataset
new_dataset = []
for i, element in enumerate(doc["eval_dataset"]):
samples = get_next_few_items_circular(ground_truth_list, i)
element["samples"] = samples
new_dataset.append(element)
# write toml
doc["eval_dataset"] = new_dataset
with open(f"eval/dataset/added_samples_{args.toml}", "w") as f:
f.write(tomlkit.dumps(doc, sort_keys=True))
wren-ai-service/eval/data_curation/__init__.py
View File
wren-ai-service/eval/data_curation/app.py
-494
View File
@@ -1,494 +0,0 @@
import asyncio
import re
import sys
import uuid
from datetime import datetime
from pathlib import Path
import orjson
import pandas as pd
import streamlit as st
import tomlkit
from openai import AsyncClient
from streamlit_tags import st_tags
sys.path.append(f"{Path().parent.resolve()}")
from utils import (
DATA_SOURCES,
WREN_ENGINE_ENDPOINT,
WREN_IBIS_ENDPOINT,
get_contexts_from_sqls,
get_data_from_wren_engine_with_sqls,
get_question_sql_pairs,
is_sql_valid,
prettify_sql,
)
from eval import EvalSettings
from eval.utils import (
get_documents_given_contexts,
get_eval_dataset_in_toml_string,
get_openai_client,
prepare_duckdb_init_sql,
prepare_duckdb_session_sql,
)
st.set_page_config(layout="wide")
st.title("WrenAI Data Curation App")
LLM_OPTIONS = ["gpt-4o-mini", "gpt-4o"]
settings = EvalSettings()
llm_client = get_openai_client(api_key=settings.get_openai_api_key())
# 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")
)
if data_source == "bigquery":
st.session_state["connection_info"] = settings.bigquery_info
elif data_source == "duckdb":
prepare_duckdb_session_sql(WREN_ENGINE_ENDPOINT)
prepare_duckdb_init_sql(
WREN_ENGINE_ENDPOINT,
st.session_state["mdl_json"]["catalog"],
"etc/spider1.0/database",
)
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"],
WREN_ENGINE_ENDPOINT
if st.session_state["data_source"] == "duckdb"
else WREN_IBIS_ENDPOINT,
)
)
if valid:
new_context = asyncio.run(
get_contexts_from_sqls([sql], st.session_state["mdl_json"])
)[0]
document = get_documents_given_contexts(
[new_context], 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
st.session_state["llm_question_sql_pairs"][i]["document"] = document
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
st.session_state["user_question_sql_pair"]["document"] = document
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"],
"document": st.session_state["llm_question_sql_pairs"][i]["document"],
}
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"],
"document": st.session_state["user_question_sql_pair"]["document"],
}
# 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": [],
"document": [],
"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_with_sqls(
[st.session_state["user_question_sql_pair"]["sql"]],
st.session_state["data_source"],
st.session_state["mdl_json"],
st.session_state["connection_info"],
WREN_ENGINE_ENDPOINT
if st.session_state["data_source"] == "duckdb"
else WREN_IBIS_ENDPOINT,
)
)[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...")
wren-ai-service/eval/data_curation/utils.py
-274
View File
@@ -1,274 +0,0 @@
import asyncio
import base64
import logging
import os
import re
import sys
from pathlib import Path
from typing import List, Tuple
import aiohttp
import orjson
import sqlparse
import streamlit as st
from dotenv import load_dotenv
from openai import AsyncClient
# add wren-ai-service to sys.path
sys.path.append(f"{Path().parent.parent.resolve()}")
from eval.utils import (
get_contexts_from_sql,
get_data_from_wren_engine,
get_ddl_commands,
get_documents_given_contexts,
)
from src.pipelines.indexing.db_schema import DDLChunker
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", "duckdb"]
TIMEOUT_SECONDS = 60
ddl_converter = DDLChunker()
logger = logging.getLogger("wren-ai-service")
async def is_sql_valid(
sql: str,
data_source: str,
mdl_json: dict,
connection_info: dict,
api_endpoint: str,
timeout: float = TIMEOUT_SECONDS,
) -> Tuple[bool, str]:
sql = sql.rstrip(";") if sql.endswith(";") else sql
if data_source == "duckdb":
async with aiohttp.request(
"GET",
f"{api_endpoint}/v1/mdl/dry-run",
json={
"sql": remove_limit_statement(sql),
"manifest": mdl_json,
"limit": 1,
},
timeout=aiohttp.ClientTimeout(total=timeout),
) as response:
if response.status == 200:
return True, None
res = await response.json()
return False, res
else:
async with aiohttp.request(
"POST",
f"{api_endpoint}/v3/connector/{data_source}/query?dryRun=true",
json={
"sql": remove_limit_statement(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,
WREN_ENGINE_ENDPOINT
if data_source == "duckdb"
else WREN_IBIS_ENDPOINT,
)
)
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_contexts_from_sqls(
sqls: list[str],
mdl_json: dict,
api_endpoint: str = WREN_ENGINE_ENDPOINT,
timeout: float = TIMEOUT_SECONDS,
) -> list[list[str]]:
async with aiohttp.ClientSession():
tasks = []
for sql in sqls:
task = asyncio.ensure_future(
get_contexts_from_sql(
sql,
mdl_json,
api_endpoint,
timeout,
)
)
tasks.append(task)
results = await asyncio.gather(*tasks)
return 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": <question_string>,
"sql": <sql_query_string>
}},
{{
"question": <question_string>,
"sql": <sql_query_string>
}},
...
]
}}
### Custom Instructions ###
{custom_instructions}
### Input ###
Data Model: {"\n\n".join(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)
documents = get_documents_given_contexts(contexts, mdl_json)
sqls_data = await get_data_from_wren_engine_with_sqls(
sqls,
data_source,
mdl_json,
connection_info,
WREN_ENGINE_ENDPOINT
if st.session_state["data_source"] == "duckdb"
else WREN_IBIS_ENDPOINT,
)
return [
{
**quesiton_sql_pair,
"context": context,
"data": sql_data,
"document": document,
}
for quesiton_sql_pair, context, sql_data, document in zip(
question_sql_pairs, contexts, sqls_data, documents
)
]
except Exception as e:
logger.error(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_with_sqls(
sqls: List[str],
data_source: str,
mdl_json: dict,
connection_info: dict,
api_endpoint: str,
timeout: float = TIMEOUT_SECONDS,
) -> List[dict]:
assert data_source in DATA_SOURCES, f"Invalid data source: {data_source}"
async with aiohttp.ClientSession():
tasks = []
for sql in sqls:
task = asyncio.ensure_future(
get_data_from_wren_engine(
sql=sql,
mdl_json=mdl_json,
api_endpoint=api_endpoint,
data_source=data_source,
connection_info=connection_info,
timeout=timeout,
limit=50,
)
)
tasks.append(task)
return await asyncio.gather(*tasks)
def remove_limit_statement(sql: str) -> str:
pattern = r"\s*LIMIT\s+\d+(\s*;?\s*--.*|\s*;?\s*)$"
modified_sql = re.sub(pattern, "", sql, flags=re.IGNORECASE)
return modified_sql
wren-ai-service/eval/dataset/.gitignore
-2
View File
@@ -1,2 +0,0 @@
*
!.gitignore
wren-ai-service/eval/dspy_modules/__init__.py
View File
wren-ai-service/eval/dspy_modules/ask_generation.py
-33
View File
@@ -1,33 +0,0 @@
import dspy
class AskGenerationSignatureV1(dspy.Signature):
"""Given a user query that is ambiguous in nature, your task is to interpret the query in various plausible ways and \
generate three SQL statements that could potentially answer each interpreted version of the queries and within-10-words summary. \
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. \
The output should be in the following JSON format:
{
"results": [
{"sql": <SQL_QUERY_STRING_1>, "summary": <SUMMARY_STRING_1>},
{"sql": <SQL_QUERY_STRING2>, "summary": <SUMMARY_STRING_2>},
{"sql": <SQL_QUERY_STRING3>, "summary": <SUMMARY_STRING_3>}
]
}
"""
question = dspy.InputField()
context = dspy.InputField(description="List of database schema documents")
answer = dspy.OutputField()
class AskGenerationV1(dspy.Module):
def __init__(self):
super().__init__()
self.generate_answer = dspy.ChainOfThought(AskGenerationSignatureV1)
def forward(self, question, context):
prediction = self.generate_answer(question=question, context=context)
return dspy.Prediction(context=context, answer=prediction.answer)
wren-ai-service/eval/dspy_modules/prompt_optimizer.py
-183
View File
@@ -1,183 +0,0 @@
import argparse
import os
import re
import sys
from datetime import datetime
from pathlib import Path
from typing import Callable, Tuple
import dotenv
import dspy
import dspy.evaluate
import dspy.teleprompt
sys.path.append(f"{Path().parent.resolve()}")
import src.utils as utils
from eval.dspy_modules.ask_generation import AskGenerationV1
from eval.utils import parse_toml
def parse_args() -> Tuple[str]:
parser = argparse.ArgumentParser()
parser.add_argument(
"--file",
"-F",
type=str,
required=True,
help="Eval dataset file name in the eval/dataset folder",
)
parser.add_argument(
"--program_file",
"-P",
type=str,
help="Optimized program file name in the eval/optimized folder",
)
parser.add_argument(
"--optimize",
action=argparse.BooleanOptionalAction,
help="Whether to optimize the program or not",
)
parser.add_argument(
"--eval",
action=argparse.BooleanOptionalAction,
help="Whether to evaluate the program or not based on the devset of eval dataset",
)
return parser.parse_args()
optimizer_parameters = {
"evaluator": None,
"metrics": None,
"meta": None,
"predictions": None,
}
def configure_llm_provider(llm: str, api_key: str):
dspy.settings.configure(lm=dspy.LM(model=llm, api_key=api_key))
def clean_sql(sql: str) -> str:
return re.sub("[ \t\n]+", " ", sql)
def prepare_dataset(path: str, train_ratio: float = 0.5):
eval_dataset = parse_toml(f"eval/dataset/{path}")["eval_dataset"]
dspy_dataset = []
for data in eval_dataset:
dspy_dataset.append(
dspy.Example(
context=[str(doc["content"]) for doc in data["document"]],
question=str(data["question"]),
answer=clean_sql(str(data["sql"])),
).with_inputs("question", "context")
)
train_size = int(len(dspy_dataset) * train_ratio)
_train = dspy_dataset[:train_size]
_dev = dspy_dataset[train_size:]
return _train, _dev
# Validation logic: check that the predicted answer is correct.
# Also check that the retrieved context does actually contain that answer.
def validate_context_and_answer(example, pred, trace=None):
if optimizer_parameters["predictions"] is None:
# if we don't have metric from eval module, we use the default metric
answer_EM = dspy.evaluate.answer_exact_match(example, pred)
answer_PM = dspy.evaluate.answer_passage_match(example, pred)
return answer_EM and answer_PM
else:
prediction = optimizer_parameters["predictions"][0]
prediction.input = example.question
prediction.expected_output = example.answer
prediction.context = example.context
prediction["type"] = "execution"
prediction.actual_output = pred.answer
# reuse the first predict result to optimize the dspy module
optimizer_parameters["evaluator"].eval(
optimizer_parameters["meta"], [prediction]
)
sum_score = 0
for metric in optimizer_parameters["metrics"].get("metrics"):
sum_score += metric.score
return sum_score
def optimize(
module: dspy.Module,
optimizer: dspy.teleprompt.Teleprompter,
trainset: list,
metric: Callable,
):
optimizer = optimizer(metric=metric)
return optimizer.compile(module(), trainset=trainset)
def build_optimizing_module(trainset):
module = optimize(
AskGenerationV1,
dspy.teleprompt.BootstrapFewShot,
trainset=trainset,
metric=validate_context_and_answer,
)
path = f"eval/optimized/{AskGenerationV1.__name__}_optimized_{datetime.now().strftime("%Y_%m_%d_%H%M%S")}.json"
directory = os.path.dirname(path)
if directory and not os.path.exists(directory):
os.makedirs(directory)
module.save(path)
return module
if __name__ == "__main__":
args = parse_args()
path = f"eval/dataset/{args.file}"
optimized_program_file = args.program_file
should_optimize = args.optimize
should_eval = args.eval
dotenv.load_dotenv()
utils.load_env_vars()
configure_llm_provider(os.getenv("GENERATION_MODEL"), os.getenv("OPENAI_API_KEY"))
trainset, devset = prepare_dataset(path)
if optimized_program_file:
module = AskGenerationV1()
module.load(f"eval/optimized/{optimized_program_file}")
elif should_optimize:
module = build_optimizing_module(trainset)
else:
module = AskGenerationV1()
if should_eval:
evaluator = dspy.evaluate.Evaluate(
devset=devset,
metric=validate_context_and_answer,
display_progress=True,
display_table=True,
return_outputs=True,
)
results = evaluator(
module,
metric=validate_context_and_answer,
display_progress=True,
display_table=True,
return_all_scores=True,
return_outputs=True,
)
for result in results:
if isinstance(result, list):
for item in result:
if isinstance(item, tuple):
if len(item) == 3:
print(item[0])
print(f"Question: {item[0].get('question')}")
print(f"Context: {item[0].get('context')}")
print(item[1])
print(f"Answer: {item[1].get('answer')}")
print(f"Context: {item[1].get('context')}")
wren-ai-service/eval/evaluation.py
-200
View File
@@ -1,200 +0,0 @@
import argparse
import sys
from pathlib import Path
from typing import Tuple
from deepeval import evaluate
from deepeval.evaluate import TestResult
from deepeval.test_case import LLMTestCase
from langfuse import Langfuse
from langfuse.decorators import langfuse_context, observe
sys.path.append(f"{Path().parent.resolve()}")
import traceback
import eval.pipelines as pipelines
import src.providers as provider
from eval import EvalSettings
from eval.utils import parse_toml, trace_metadata
from src import utils
def formatter(prediction: dict, meta: dict) -> dict:
"""
Formats the prediction result to be used as evaluation input.
This function takes a prediction dictionary and a meta dictionary,
processes them to extract relevant information, and returns a formatted
dictionary that serves as input for evaluation. It includes details such
as input, actual and expected outputs, context, and additional metadata.
Args:
prediction (dict): A dictionary containing prediction details.
meta (dict): A dictionary containing metadata information.
Returns:
dict: A formatted dictionary containing evaluation input data.
"""
retrieval_context = [str(context) for context in prediction["retrieval_context"]]
context = [str(context) for context in prediction["context"]]
enable_spider_metrics = "spider" in meta.get("evaluation_dataset", "").lower()
enable_rewrite = any(
dataset in meta.get("evaluation_dataset", "").lower() for dataset in ["spider"]
)
return {
"input": prediction["input"],
"actual_output": (
prediction.get("actual_output", {})
.get("post_process", {})
.get("valid_generation_result", {})
.get("sql", "")
),
"expected_output": prediction["expected_output"],
"retrieval_context": retrieval_context,
"context": context,
"reasoning": prediction.get("reasoning", ""),
"additional_metadata": {
"trace_id": prediction["trace_id"],
"trace_url": prediction["trace_url"],
"catalog": meta.get("catalog", None),
"enable_spider_metrics": enable_spider_metrics,
"enable_rewrite": enable_rewrite,
},
}
def parse_args() -> Tuple[str]:
parser = argparse.ArgumentParser()
parser.add_argument(
"--file",
"-F",
type=str,
help="Eval the prediction result in the outputs/predictions directory",
)
parser.add_argument(
"--semantics",
"-S",
default=False,
action=argparse.BooleanOptionalAction,
help="Whether use the LLM(OpenAI's gpt-4o-mini) to help check semantics of sqls to improve accuracy metrics",
)
parser.add_argument(
"--training-dataset",
"-T",
default=None,
help="Use the training dataset to build a dspy optimized module",
)
return parser.parse_args()
class Evaluator:
def __init__(self, metrics: list, **kwargs):
self._score_collector = {}
self._langfuse = Langfuse()
self._metrics = metrics
self._failed_count = 0
self._post_metrics = kwargs.get("post_metrics", [])
def eval(self, meta: dict, predictions: list) -> None:
for prediction in predictions:
try:
test_case = LLMTestCase(**formatter(prediction, meta))
result = evaluate(
[test_case], self._metrics, ignore_errors=True
).test_results[0]
self._score_metrics(test_case, result)
[metric.collect(test_case, result) for metric in self._post_metrics]
except Exception:
self._failed_count += 1
traceback.print_exc()
self._average_score(meta)
def _score_metrics(self, test_case: LLMTestCase, result: TestResult) -> None:
for metric in result.metrics_data:
name = metric.name
score = metric.score or 0
self._langfuse.score(
trace_id=test_case.additional_metadata["trace_id"],
name=name,
value=score,
comment=metric.reason or metric.error,
source="EVAL",
)
if name not in self._score_collector:
self._score_collector[name] = []
self._score_collector[name].append(score)
@observe(name="Summary Trace", capture_input=False, capture_output=False)
def _average_score(self, meta: dict) -> None:
langfuse_context.update_current_trace(
session_id=meta.get("session_id"),
user_id=meta.get("user_id"),
metadata=trace_metadata(meta, type="summary"),
)
summary = {
"query_count": meta["query_count"],
"expected_batch_size": meta["expected_batch_size"],
"actual_batch_size": meta["actual_batch_size"],
"valid_eval_count": meta["actual_batch_size"] - self._failed_count,
}
langfuse_context.update_current_observation(output=summary)
for name, scores in self._score_collector.items():
langfuse_context.score_current_trace(
name=name,
value=sum(scores) / len(scores),
comment=f"Average score for {name}",
)
for metric in self._post_metrics:
langfuse_context.score_current_trace(
name=metric.__name__,
value=metric.measure(),
comment=f"Average score for {metric.__name__}",
)
if __name__ == "__main__":
args = parse_args()
settings = EvalSettings()
pipe_components = provider.generate_components(settings.components)
utils.init_langfuse(settings)
predicted_file = parse_toml(f"outputs/predictions/{args.file}")
meta = predicted_file["meta"]
predictions = predicted_file["predictions"]
dataset = parse_toml(meta["evaluation_dataset"])
metrics = pipelines.metrics_initiator(
meta["pipeline"], dataset, pipe_components, args.semantics, settings
)
evaluator = Evaluator(**metrics)
evaluator.eval(meta, predictions)
# if args.training_dataset:
# # todo: for now comment dspy related code
# optimizer_parameters["evaluator"] = evaluator
# optimizer_parameters["metrics"] = metrics
# optimizer_parameters["meta"] = meta
# optimizer_parameters["predictions"] = predictions
# configure_llm_provider(
# os.getenv("GENERATION_MODEL"), os.getenv("OPENAI_API_KEY")
# )
# trainset, devset = prepare_dataset(args.training_dataset)
# build_optimizing_module(trainset)
# else:
# evaluator.eval(meta, predictions)
langfuse_context.flush()
if meta["langfuse_url"]:
print(
f"\n\nYou can view the evaluation result in Langfuse at {meta['langfuse_url']}/sessions/{meta['session_id']}"
)
wren-ai-service/eval/mdl_to_csv.py
-125
View File
@@ -1,125 +0,0 @@
import argparse
import csv
from pathlib import Path
import orjson
def gen_eval_preparation_data_from_json_to_csv(mdl_path: str):
assert Path(mdl_path).exists(), f"File not found: {mdl_path}"
with open(mdl_path) as file:
mdl = orjson.loads(file.read())
csv_data = [
[
"table",
"table alias",
"table description",
"column",
"column alias",
"column description",
]
]
for model in mdl["models"]:
for column in model["columns"]:
csv_data.append(
[
model["name"],
model.get("properties", {}).get("displayName", ""),
model.get("properties", {}).get("description", ""),
column["name"],
column.get("properties", {}).get("displayName", ""),
column.get("properties", {}).get("description", ""),
]
)
with open(f"{Path(mdl_path).stem}.csv", "w", newline="\n") as file:
writer = csv.writer(file, quoting=csv.QUOTE_MINIMAL)
for row in csv_data:
writer.writerow(row)
def gen_new_mdl_from_csv(mdl_path: str, csv_path: str):
assert Path(mdl_path).exists(), f"File not found: {mdl_path}"
assert Path(csv_path).exists(), f"File not found: {csv_path}"
with open(mdl_path) as file:
mdl = orjson.loads(file.read())
csv_data_by_table = {}
with open(csv_path, newline="\n") as file:
csv_data = csv.reader(file)
for row in csv_data:
model_name = row[0]
if model_name not in csv_data_by_table:
csv_data_by_table[model_name] = {
"model": {
"displayName": row[1],
"description": row[2],
},
"columns": {},
}
csv_data_by_table[model_name]["columns"][row[3]] = {
"displayName": row[4],
"description": row[5],
}
new_models = []
for model in mdl["models"]:
if model["name"] in csv_data_by_table:
if "properties" not in model:
model["properties"] = {}
if csv_data_by_table[model["name"]]["model"]["displayName"]:
model["properties"]["displayName"] = csv_data_by_table[
model["name"]
]["model"]["displayName"]
if csv_data_by_table[model["name"]]["model"]["description"]:
model["properties"]["description"] = csv_data_by_table[
model["name"]
]["model"]["description"]
new_columns = []
for column in model["columns"]:
if column["name"] in csv_data_by_table[model["name"]]["columns"]:
if "properties" not in column:
column["properties"] = {}
if csv_data_by_table[model["name"]]["columns"][column["name"]][
"displayName"
]:
column["properties"]["displayName"] = csv_data_by_table[
model["name"]
]["columns"][column["name"]]["displayName"]
if csv_data_by_table[model["name"]]["columns"][column["name"]][
"description"
]:
column["properties"]["description"] = csv_data_by_table[
model["name"]
]["columns"][column["name"]]["description"]
new_columns.append(column)
model["columns"] = new_columns
new_models.append(model)
mdl["models"] = new_models
with open(f"{Path(mdl_path).stem}_new.json", "w") as file:
file.write(orjson.dumps(mdl).decode())
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--mdl-path", type=str, help="Path to the MDL JSON file", required=True
)
parser.add_argument("--csv-path", type=str, help="Path to the input CSV file")
args = parser.parse_args()
if args.mdl_path and not args.csv_path:
gen_eval_preparation_data_from_json_to_csv(args.mdl_path)
elif args.mdl_path and args.csv_path:
gen_new_mdl_from_csv(args.mdl_path, args.csv_path)
wren-ai-service/eval/metrics/__init__.py
-28
View File
@@ -1,28 +0,0 @@
from .accuracy import AccuracyMetric, AccuracyMultiCandidateMetric
from .answer_relevancy import AnswerRelevancyMetric
from .context_precision import ContextualPrecisionMetric
from .context_recall import ContextualRecallMetric
from .context_relevancy import ContextualRelevancyMetric
from .faithfulness import FaithfulnessMetric
from .llm import (
QuestionToReasoningJudge,
ReasoningToSqlJudge,
SqlSemanticsJudge,
)
from .spider.exact_match import ExactMatchAccuracy
from .spider.exec_match import ExecutionAccuracy
__all__ = [
"AccuracyMetric",
"AccuracyMultiCandidateMetric",
"AnswerRelevancyMetric",
"ContextualPrecisionMetric",
"ContextualRecallMetric",
"ContextualRelevancyMetric",
"FaithfulnessMetric",
"ExactMatchAccuracy",
"ExecutionAccuracy",
"QuestionToReasoningJudge",
"ReasoningToSqlJudge",
"SqlSemanticsJudge",
]
wren-ai-service/eval/metrics/accuracy.py
-205
View File
@@ -1,205 +0,0 @@
import asyncio
import re
import traceback
import orjson
import pandas as pd
from deepeval.evaluate import TestResult
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
from deprecated import deprecated
from eval.utils import get_data_from_wren_engine, get_openai_client
class AccuracyMetric(BaseMetric):
def __init__(self, engine_info: dict, enable_semantics_comparison: bool = False):
self.threshold = 0
self.score = 0
self.engine_info = engine_info
self.enable_semantics_comparison = enable_semantics_comparison
if self.enable_semantics_comparison:
self._openai_client = get_openai_client()
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
def _is_subset(self, expected: pd.DataFrame, actual: pd.DataFrame) -> bool:
if not set(expected.columns).issubset(set(actual.columns)):
return False
common_columns = sorted(expected.columns)
expected_sorted = expected[common_columns]
actual_sorted = actual[common_columns]
# Ensure that the data types are the same
actual_sorted = actual_sorted.astype(expected_sorted.dtypes.to_dict())
merged = pd.merge(
actual_sorted,
expected_sorted,
on=common_columns,
how="left",
indicator=True,
)
return all(merged["_merge"] == "both")
def _count_partial_matches(
self, expected: pd.DataFrame, actual: pd.DataFrame
) -> int:
intersection = set(expected.columns).intersection(set(actual.columns))
common_columns = sorted(intersection)
if not common_columns:
return 0
expected_sorted = expected[common_columns]
actual_sorted = actual[common_columns]
# Ensure that the data types are the same
actual_sorted = actual_sorted.astype(expected_sorted.dtypes.to_dict())
merged = pd.merge(
actual_sorted,
expected_sorted,
on=common_columns,
how="left",
indicator=True,
)
if all(merged["_merge"] == "both"):
return len(intersection) / len(expected.columns)
else:
return 0
def _rewrite_sql(self, sql: str) -> str:
# Pattern to match double quotes after WHERE clause, including multiple occurrences
pattern = r'(WHERE\s+.*?)(")(.+?)(")(.*)$'
replacement = r"\1'\3'\5"
# Apply the replacement repeatedly until no more changes
new_sql = re.sub(pattern, replacement, sql, flags=re.IGNORECASE | re.DOTALL)
while new_sql != sql:
sql = new_sql
new_sql = re.sub(pattern, replacement, sql, flags=re.IGNORECASE | re.DOTALL)
return sql
async def _retrieve_data(self, sql: str) -> pd.DataFrame:
response = await get_data_from_wren_engine(sql=sql, **self.engine_info)
df = pd.DataFrame(**response)
sorted_columns = sorted(df.columns)
return df[sorted_columns].sort_values(by=sorted_columns)
async def _check_sql_semantics(self, expected_sql: str, actual_sql: str):
_system_prompt = (
"### TASK ### \n"
+ "You are a great data anlyst, please carefully check the semantics of two given SQLs if they are the same. \n"
+ "The output should be a JSON format with the following schema: \n"
+ "{ \n"
+ ' "reasoning": <REASONING_STRING> \n'
+ ' "same": <BOOL> \n'
+ "}"
)
_user_prompt = (
"### QUESTION ### \n"
+ f"Expected SQL: {expected_sql} \n"
+ f"Actual SQL: {actual_sql} \n"
+ "\n"
+ "Please think step by step"
)
response = await self._openai_client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": _system_prompt},
{"role": "user", "content": _user_prompt},
],
response_format={"type": "json_object"},
)
print(
f"response of _check_sql_semantics: {response.choices[0].message.content}"
)
return 1 if orjson.loads(response.choices[0].message.content)["same"] else 0
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
try:
enable_rewrite = test_case.additional_metadata.get("enable_rewrite", False)
rewritten_expected_output = test_case.expected_output
if enable_rewrite:
rewritten_expected_output = self._rewrite_sql(test_case.expected_output)
expected_dataset = await self._retrieve_data(rewritten_expected_output)
actual_dataset = await self._retrieve_data(test_case.actual_output)
print(f"expected columns: {set(expected_dataset.columns)}")
print(f"actual columns: {set(actual_dataset.columns)}")
if expected_dataset.equals(actual_dataset) or self._is_subset(
expected_dataset, actual_dataset
):
self.success = True
self.score = 1
return self.score
self.score = self._count_partial_matches(expected_dataset, actual_dataset)
# use llm to check sql semantics
if self.score == 0 and self.enable_semantics_comparison:
# TODO: we may need to upload the sql semantics result to langfuse
print(f"before _check_sql_semantics: {self.score}")
print(f"expected sql: {rewritten_expected_output}")
print(f"actual sql: {test_case.actual_output}")
self.score = await self._check_sql_semantics(
rewritten_expected_output, test_case.actual_output
)
print(f"after _check_sql_semantics: {self.score}")
except Exception as e:
self.error = f"Error occurred while evaluating the metric: {e}"
traceback.print_exc()
# if didn't pass any of the above checks
self.success = False
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "Accuracy(column-based)"
@deprecated(
reason="We don't generate multiple candidates for Text to SQL task, so don't need this metric"
)
class AccuracyMultiCandidateMetric(BaseMetric):
def __init__(self):
self.threshold = 0
self.score = 0
self._questions = {}
def collect(self, test_case: LLMTestCase, result: TestResult):
for metric in result.metrics_data:
if metric.name != "Accuracy(column-based)":
continue
# or 0 to avoid when metric.error is exist
self._questions[test_case.input] = (
self._questions.get(test_case.input, 0) or metric.score or 0
)
def measure(self):
if not self._questions:
return 0
self.score = sum(self._questions.values()) / len(self._questions)
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "Accuracy(question-based)"
wren-ai-service/eval/metrics/answer_relevancy.py
-38
View File
@@ -1,38 +0,0 @@
import asyncio
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
from eval.utils import get_contexts_from_sql
class AnswerRelevancyMetric(BaseMetric):
def __init__(self, engine_info: dict):
self.threshold = 0
self.score = 0
self.engine_info = engine_info
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
actual_units = await get_contexts_from_sql(
sql=test_case.actual_output, **self.engine_info
)
expected_units = await get_contexts_from_sql(
sql=test_case.expected_output, **self.engine_info
)
intersection = set(actual_units) & set(expected_units)
self.score = len(intersection) / len(actual_units)
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "AnswerRelevancy(column-based)"
wren-ai-service/eval/metrics/context_precision.py
-43
View File
@@ -1,43 +0,0 @@
import asyncio
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
class ContextualPrecisionMetric(BaseMetric):
def __init__(self):
self.threshold = 0
self.score = 0
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
context = test_case.context
retrieval_context = test_case.retrieval_context
intersection = set(context) & set(retrieval_context)
intersection_count = len(intersection)
if intersection_count == 0:
self.success = False
return self.score
n = len(retrieval_context)
summation = 0
for k in range(1, n + 1):
intersection_up_to_k = len(set(context[:k]) & set(retrieval_context[:k]))
rk = len(set(context[:k]) & set(retrieval_context[k - 1 : k])) > 0
summation += (intersection_up_to_k / k) * rk
self.score = (1 / intersection_count) * summation
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "ContextualPrecision(column-based)"
wren-ai-service/eval/metrics/context_recall.py
-34
View File
@@ -1,34 +0,0 @@
import asyncio
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
from eval.utils import get_contexts_from_sql
class ContextualRecallMetric(BaseMetric):
def __init__(self, engine_info: dict):
self.threshold = 0
self.score = 0
self.engine_info = engine_info
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
expected_units = await get_contexts_from_sql(
sql=test_case.expected_output, **self.engine_info
)
intersection = set(test_case.retrieval_context) & set(expected_units)
self.score = len(intersection) / len(expected_units)
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "ContextualRecall(column-based)"
wren-ai-service/eval/metrics/context_relevancy.py
-27
View File
@@ -1,27 +0,0 @@
import asyncio
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
class ContextualRelevancyMetric(BaseMetric):
def __init__(self):
self.threshold = 0
self.score = 0
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
intersection = set(test_case.retrieval_context) & set(test_case.context)
self.score = len(intersection) / len(test_case.retrieval_context)
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "ContextualRelevancy(column-based)"
wren-ai-service/eval/metrics/faithfulness.py
-33
View File
@@ -1,33 +0,0 @@
import asyncio
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
from eval.utils import get_contexts_from_sql
class FaithfulnessMetric(BaseMetric):
def __init__(self, engine_info: dict):
self.threshold = 0
self.score = 0
self.engine_info = engine_info
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
actual_units = await get_contexts_from_sql(
sql=test_case.actual_output, **self.engine_info
)
intersection = set(actual_units) & set(test_case.retrieval_context)
self.score = len(intersection) / len(actual_units)
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "Faithfulness(column-based)"
wren-ai-service/eval/metrics/llm/__init__.py
-173
View File
@@ -1,173 +0,0 @@
import asyncio
from deepeval.metrics import BaseMetric
from deepeval.test_case import LLMTestCase
from haystack.components.builders.prompt_builder import PromptBuilder
from pydantic import BaseModel
from src.providers import LLMProvider
class EvalResult(BaseModel):
score: float
reason: str
_MODEL_KWARGS = {
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "eval_result",
"schema": EvalResult.model_json_schema(),
},
}
}
def format(response: dict) -> EvalResult:
reply = response.get("replies", [])[0]
return EvalResult.model_validate_json(reply)
class QuestionToReasoningJudge(BaseMetric):
_system_prompt = """
You are an expert evaluator. Your task is to analyze the reasoning provided for a given question and determine if it makes sense.
Provide a score in the range 0.0~1.0 and a detailed explanation for your evaluation.
"""
_test_case_prompt = """
Question:
{{ question }}
Reasoning:
{{ reasoning }}
"""
def __init__(self, llm_provider: LLMProvider, **_):
self.threshold = 0
self.score = 0
self.llm_provider = llm_provider
self.llm = llm_provider.get_generator(
system_prompt=self._system_prompt,
generation_kwargs=_MODEL_KWARGS,
)
self.prompt_builder = PromptBuilder(template=self._test_case_prompt)
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
prompt = self.prompt_builder.run(
question=test_case.input,
reasoning=test_case.reasoning,
)
response = await self.llm(prompt.get("prompt"))
result = format(response)
self.score = result.score
self.reason = result.reason
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "QuestionToReasoningJudge"
class ReasoningToSqlJudge(BaseMetric):
_system_prompt = """
You are an expert evaluator. Your task is to analyze the reasoning provided for a given SQL query and determine if it makes sense.
Provide a score in the range 0.0~1.0 and a detailed explanation for your evaluation.
"""
_test_case_prompt = """
Actual Output:
{{ actual_output }}
Reasoning:
{{ reasoning }}
"""
def __init__(self, llm_provider: LLMProvider, **_):
self.threshold = 0
self.score = 0
self.llm_provider = llm_provider
self.llm = llm_provider.get_generator(
system_prompt=self._system_prompt,
generation_kwargs=_MODEL_KWARGS,
)
self.prompt_builder = PromptBuilder(template=self._test_case_prompt)
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
prompt = self.prompt_builder.run(
actual_output=test_case.actual_output,
reasoning=test_case.reasoning,
)
response = await self.llm(prompt.get("prompt"))
result = format(response)
self.score = result.score
self.reason = result.reason
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "ReasoningToSqlJudge"
class SqlSemanticsJudge(BaseMetric):
_system_prompt = """
You are an expert evaluator. Your task is to analyze the actual SQL query and the expected SQL query and determine if they are semantically equivalent.
Provide a score in the range 0.0~1.0 and a detailed explanation for your evaluation.
"""
_test_case_prompt = """
Actual SQL:
{{ actual_sql }}
Expected SQL:
{{ expected_sql }}
"""
def __init__(self, llm_provider: LLMProvider, **_):
self.threshold = 0
self.score = 0
self.llm_provider = llm_provider
self.llm = llm_provider.get_generator(
system_prompt=self._system_prompt,
generation_kwargs=_MODEL_KWARGS,
)
self.prompt_builder = PromptBuilder(template=self._test_case_prompt)
def measure(self, test_case: LLMTestCase):
return asyncio.run(self.a_measure(test_case))
async def a_measure(self, test_case: LLMTestCase, *args, **kwargs):
prompt = self.prompt_builder.run(
actual_sql=test_case.actual_output,
expected_sql=test_case.expected_output,
)
response = await self.llm(prompt.get("prompt"))
result = format(response)
self.score = result.score
self.reason = result.reason
self.success = self.score >= self.threshold
return self.score
def is_successful(self):
return self.success
@property
def __name__(self):
return "SqlSemanticsJudge"

Some files were not shown because too many files have changed in this diff Show More